Eval_WWSe_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gomb - Net\n",
    "### Performance test on the WSSe dataset\n",
    "### And analysis of WSSe experimental data\n",
    "\n",
    "Austin Houston"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![OpenInColab](https://colab.research.google.com/assets/colab-badge.svg)](\n",
    "    https://colab.research.google.com/github/ahoust17/Gomb-Net/blob/main/Eval_WSSe_model.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Necessary installs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# basics\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "# plotting\n",
    "import matplotlib.pylab as plt\n",
    "import matplotlib.colors as mcolors\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib import cm\n",
    "\n",
    "# colab interactive plots and drive\n",
    "drive = False\n",
    "if 'google.colab' in sys.modules:\n",
    "    from  google.colab import drive \n",
    "    from google.colab import output\n",
    "    drive.mount('/content/drive')\n",
    "    output.enable_custom_widget_manager()\n",
    "    drive = True\n",
    "else:\n",
    "    %matplotlib widget\n",
    "\n",
    "# other imports\n",
    "from scipy.ndimage import label, center_of_mass, gaussian_filter, zoom, uniform_filter\n",
    "from scipy.spatial import KDTree\n",
    "from scipy.interpolate import griddata\n",
    "from scipy.stats import norm, gaussian_kde\n",
    "from skimage.filters import threshold_otsu\n",
    "from skimage.feature import blob_log\n",
    "\n",
    "# for cropping function\n",
    "if drive:\n",
    "    print('installing DataGenSTEM')\n",
    "    !pip install ase\n",
    "    !git clone https://github.com/ahoust17/DataGenSTEM.git\n",
    "    sys.path.append('./DataGenSTEM/DataGenSTEM')\n",
    "    import data_generator as dg\n",
    "\n",
    "# for Gomb-Net\n",
    "if drive:\n",
    "    print('installing Gomb-Net')\n",
    "    !git clone https://github.com/ahoust17/Gomb-Net.git\n",
    "    sys.path.append('./Gomb-Net/GombNet')    \n",
    "from GombNet.networks import *\n",
    "from GombNet.loss_func import GombinatorialLoss\n",
    "from GombNet.utils import *\n",
    "\n",
    "import torch\n",
    "# Check if CUDA is available\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"CUDA is available. Using GPU.\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"CUDA is not available. Using CPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set mps, just for my computer\n",
    "device = torch.device('mps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, you need to add the following shared drive to your google drive:\n",
    "*** WARNING: it is a big file.  Check before you download ***\n",
    "\n",
    "\n",
    "https://drive.google.com/file/d/1DyKtrmJ8wNYQg3YEJ8_iXjz6lB_DQfwy/view?usp=sharing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the following cell after the download is complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if drive:\n",
    "    shared_folder = 'drive/My Drive/Gomb-Net files'\n",
    "else:\n",
    "    shared_folder = '/Users/austin/Desktop/gomb_beta'\n",
    "\n",
    "print('available files & directories:')\n",
    "!ls '{shared_folder}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now we can start with the actual code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's look at the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataloaders\n",
    "images_dir = str(shared_folder + '/WSSe_dataset/images')\n",
    "labels_dir = str(shared_folder + '/WSSe_dataset/labels')\n",
    "train_loader, val_loader, test_loader = get_dataloaders(images_dir, labels_dir, batch_size = 1, val_split = 0.2, test_split = 0.1, seed = 42) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_iter = 1 # 2,4,5,7,8,9, 10, 18\n",
    "test = test_loader.dataset[test_iter][0].unsqueeze(0)\n",
    "gt = test_loader.dataset[test_iter][1]\n",
    "\n",
    "fig, ax = plt.subplots(1, 7, figsize=(12, 5))\n",
    "ax[0].imshow(test[0, 0].cpu().numpy(), cmap='gray')\n",
    "ax[0].set_title('Input')\n",
    "\n",
    "titles = ['L1: S', 'L1: Se', 'L1: W', 'L2: S', 'L2: Se', 'L2: W']\n",
    "for i in range(6):\n",
    "    ax[i+1].imshow(gt[i].cpu().numpy(), cmap='gray')\n",
    "    ax[i+1].set_title(titles[i])\n",
    "for a in ax:\n",
    "    a.axis('off')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now let's look at the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize model 'skeleton'\n",
    "input_channels = 1\n",
    "num_classes = 6    # number of output classes\n",
    "num_filters = [32, 64, 128, 256, 512]\n",
    "\n",
    "model = TwoLeggedUnet(input_channels, num_classes, num_filters, dropout = 0.2)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
    "loss = GombinatorialLoss(group_size = num_classes//2, loss = 'Dice', epsilon=1e-6, class_weights = None, alpha=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the number of trainable parameters\n",
    "def count_trainable_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "num_trainable_params = count_trainable_parameters(model)\n",
    "print(f\"Number of trainable parameters: {num_trainable_params}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "visualize the training history for the pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2, sharex=True, sharey=True)\n",
    "for file in os.listdir(str(shared_folder + '/WSSe_models')):\n",
    "    if file.endswith('.npz'):\n",
    "        loss_history = np.load(str(shared_folder + '/WSSe_models/' + file))\n",
    "        train_loss = loss_history['train_loss_history']\n",
    "        val_loss = loss_history['val_loss_history']\n",
    "\n",
    "        label = file.split('_')[2]\n",
    "        print(label)\n",
    "        ax[0].plot(train_loss, label=label)\n",
    "        ax[1].plot(val_loss, label=label)\n",
    "\n",
    "ax[0].set_title('Training loss')\n",
    "ax[1].set_title('Validation loss')\n",
    "\n",
    "ax[0].legend()\n",
    "ax[1].legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "load in the pretrained weights onto our model 'skeleton'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models_dict = {}\n",
    "for file in os.listdir(str(shared_folder + '/WSSe_models')):\n",
    "    if file.endswith('.pth'):\n",
    "        if 'best' in file:\n",
    "            model_path = str(shared_folder + '/WSSe_models/' + file)\n",
    "            checkpoint = torch.load(model_path, map_location=device)\n",
    "            model.load_state_dict(checkpoint['model_state_dict'])\n",
    "            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "            model.eval()\n",
    "\n",
    "            models_dict[str('model_' + file.split('_')[2])] = model\n",
    "\n",
    "model_names = [key for key in models_dict.keys()]\n",
    "ensemble = GombSemble(models_dict)\n",
    "\n",
    "def threshold(prediction, transfrom_list, param = 1):\n",
    "    # operates on an image-wise basis\n",
    "    if 'sigmoid' in transfrom_list:\n",
    "        prediction = 1/(1 + np.exp(-prediction))\n",
    "\n",
    "    if 'threshold' in transfrom_list:\n",
    "        for i in range(prediction.shape[0]):\n",
    "            cutoff = np.std(prediction[i]) * param\n",
    "            prediction[i] = (prediction[i] > cutoff).astype(int)\n",
    "\n",
    "    if 'otsu' in transfrom_list:        \n",
    "        for i in range(prediction.shape[0]):\n",
    "            thresh = threshold_otsu(prediction[i])\n",
    "            thresh = thresh * param\n",
    "            prediction[i] = (prediction[i] > thresh).astype(int)\n",
    "            \n",
    "    return prediction\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_iter = 0 # 2,4,5,7,8,9, 10, 18\n",
    "test = test_loader.dataset[test_iter][0].unsqueeze(0)\n",
    "gt = test_loader.dataset[test_iter][1]\n",
    "\n",
    "fig, ax = plt.subplots(1, 7, figsize=(12, 5))\n",
    "ax[0].imshow(test[0, 0].cpu().numpy(), cmap='gray')\n",
    "ax[0].set_title('Input')\n",
    "\n",
    "titles = ['L1: S', 'L1: Se', 'L1: W', 'L2: S', 'L2: Se', 'L2: W']\n",
    "for i in range(6):\n",
    "    ax[i+1].imshow(gt[i].cpu().numpy(), cmap='gray')\n",
    "    ax[i+1].set_title(titles[i])\n",
    "for a in ax:\n",
    "    a.axis('off')\n",
    "fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict using the ensemble\n",
    "ensemble.predict(test, return_plot=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble.rearrange_ensemble()\n",
    "ensemble.vote(mode='mean')\n",
    "# ensemble.plot_vote()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binary = threshold(ensemble.voted_prediction, ['threshold'], param = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2,6, figsize=(12, 5))\n",
    "for i in range(6):\n",
    "    ax[0,i].set_title('GT')\n",
    "    ax[0,i].imshow(gt[i].cpu().numpy(), cmap='gray')\n",
    "\n",
    "    ax[1,i].set_title(titles[i])\n",
    "    ax[1,i].imshow(binary[i], cmap='gray')\n",
    "    \n",
    "for a in ax.ravel():\n",
    "    a.axis('off')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now, on Experimental data:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, some useful functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_atom_histograms(results, n_bins=50):\n",
    "    colors = {'W': ['blue', 'cyan'], 'Se': ['green', 'lime'], 'S': ['red', 'magenta']}\n",
    "    \n",
    "    for i, result in results.items():\n",
    "        intensities = {el: [np.concatenate(result['intensities'][el][0]) if result['intensities'][el][0] else np.array([]),\n",
    "                            np.concatenate(result['intensities'][el][1]) if result['intensities'][el][1] else np.array([])]\n",
    "                       for el in colors}\n",
    "\n",
    "        fig, axs = plt.subplots(1, 3, figsize=(20, 8))\n",
    "        # Plot the image with centroids\n",
    "        axs[0].imshow(result['image'], cmap='gray')\n",
    "        for element, colors_list in colors.items():\n",
    "            for idx, (color, layer) in enumerate(zip(colors_list, result['centroids'][element])):\n",
    "                if layer.size > 0:\n",
    "                    axs[0].scatter(layer[:, 1], layer[:, 0], color=color, label=f'{element} Layer {idx + 1} Centroids', alpha=0.6)\n",
    "        axs[0].axis('off')\n",
    "        axs[0].legend(loc='upper right', fontsize=8)\n",
    "\n",
    "        # Plot histograms and KDEs for both layers\n",
    "        for j, layer in enumerate(['Layer 1', 'Layer 2']):\n",
    "            for element, (color1, color2) in colors.items():\n",
    "                color = color1 if j == 0 else color2\n",
    "                layer_intensities = intensities[element][j]\n",
    "                if layer_intensities.size > 0:\n",
    "                    density = gaussian_kde(layer_intensities.flatten())\n",
    "                    x = np.linspace(min(layer_intensities.flatten()), max(layer_intensities.flatten()), 1000)\n",
    "                    axs[j + 1].hist(layer_intensities.flatten(), bins=n_bins, color=color, alpha=0.5, label=f'{element} {layer} Intensities', density=False)\n",
    "                    axs[j + 1].plot(x, density(x), color=color)\n",
    "            axs[j + 1].set_title(f'{layer} Intensities')\n",
    "            axs[j + 1].set_xlabel('Intensity')\n",
    "            axs[j + 1].set_ylabel('Density')\n",
    "            axs[j + 1].legend(loc='upper right', fontsize=8)\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_gaussian_mixture(metric, n_components):\n",
    "    gmm = GaussianMixture(n_components=n_components)\n",
    "    gmm.fit(metric.reshape(-1, 1))\n",
    "    means = gmm.means_.flatten()\n",
    "    stds = np.sqrt(gmm.covariances_).flatten()\n",
    "    weights = gmm.weights_\n",
    "    log_likelihood = gmm.score(metric.reshape(-1, 1)) * len(metric)\n",
    "    return means, stds, weights, log_likelihood\n",
    "\n",
    "def plot_histogram_with_gaussian(ax, metric, title, n_components, color_bounds):\n",
    "    n, bins, patches = ax.hist(metric, bins=20, density=True, edgecolor='k', alpha=1)\n",
    "    bin_centers = 0.5 * (bins[:-1] + bins[1:])\n",
    "    norm_dist = mcolors.Normalize(color_bounds[0], color_bounds[1])\n",
    "    cmap = cm.get_cmap(\"viridis\")\n",
    "    for c, p in zip(bin_centers, patches):\n",
    "        plt.setp(p, 'facecolor', cmap(norm_dist(c)))\n",
    "    ax.set_yticks([])\n",
    "\n",
    "    means, stds, weights, log_likelihood = fit_gaussian_mixture(metric, n_components)\n",
    "    xmin, xmax = ax.get_xlim()\n",
    "    x = np.linspace(xmin, xmax, 1000)\n",
    "\n",
    "    total_pdf = np.zeros_like(x)\n",
    "    for mean, std, weight in zip(means, stds, weights):\n",
    "        print('mean:', mean)\n",
    "        pdf = weight * (1 / (std * np.sqrt(2 * np.pi))) * np.exp(-0.5 * ((x - mean) / std) ** 2)\n",
    "        total_pdf += pdf\n",
    "        ax.plot(x, pdf, '--', c='k', linewidth=2)\n",
    "\n",
    "    ax.plot(x, total_pdf, 'k-', linewidth=3, alpha=0.8)\n",
    "    ax.set_title(title, fontsize=14)\n",
    "    ax.set_xlim(np.min(metric), np.max(metric))\n",
    "    return means, stds, weights, x, total_pdf, log_likelihood\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_histogram(ax, metric, color_bounds):\n",
    "    n, bins, patches = ax.hist(metric, bins=20, density=True, edgecolor='k', alpha=1)\n",
    "    bin_centers = 0.5 * (bins[:-1] + bins[1:])\n",
    "    norm_dist = mcolors.Normalize(color_bounds[0], color_bounds[1])\n",
    "    cmap = cm.get_cmap(\"viridis\")\n",
    "    for c, p in zip(bin_centers, patches):\n",
    "        plt.setp(p, 'facecolor', cmap(norm_dist(c)))\n",
    "    ax.set_yticks([])\n",
    "\n",
    "    ax.set_xlim(np.min(metric), np.max(metric))\n",
    "    return\n",
    "\n",
    "def plot_moire_distances(results, image_number=0):\n",
    "    def get_layer_data(layer):\n",
    "        return np.vstack(results[image_number][\"centroids\"][layer][0]), np.vstack(results[image_number][\"centroids\"][layer][1])\n",
    "\n",
    "    Se_0, Se_1 = get_layer_data(\"Se\")\n",
    "    W_0, W_1 = get_layer_data(\"W\")\n",
    "    S_0, S_1 = get_layer_data(\"S\")\n",
    "\n",
    "    if Se_0.shape[0] > Se_1.shape[0]:\n",
    "        Se_bottom, W_bottom, S_bottom, W_top = Se_0, W_0, S_0, W_1\n",
    "        Se_top, S_top = Se_1, S_1\n",
    "    else:\n",
    "        Se_bottom, W_bottom, S_bottom, W_top = Se_1, W_1, S_1, W_0\n",
    "        Se_top, S_top = Se_0, S_0\n",
    "    chalc_top = np.vstack((S_top, Se_top))\n",
    "    chalc_bottom = np.vstack((S_bottom, Se_bottom))\n",
    "    W_top_tree, W_bottom_tree = KDTree(W_top), KDTree(W_bottom)\n",
    "    chalc_top_tree, chalc_bottom_tree = KDTree(chalc_top), KDTree(chalc_bottom)\n",
    "\n",
    "    fig, ax = plt.subplots(2, 5, figsize=(20, 5))\n",
    "    ax[0,0].imshow(results[image_number][\"image\"].T, cmap='gray')\n",
    "    ax[0,0].set_title('Image')\n",
    "\n",
    "    grid_density = 512\n",
    "    x = np.linspace(0, 512, grid_density)\n",
    "    y = np.linspace(0, 512, grid_density)\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    points = np.vstack((X.ravel(), Y.ravel())).T\n",
    "\n",
    "    distance_arrays = []\n",
    "    \n",
    "    # W_top - W_bottom distances\n",
    "    distances = np.linalg.norm(W_top[W_top_tree.query(points)[1]] - W_bottom[W_bottom_tree.query(points)[1]], axis=1)\n",
    "    distances = distances.reshape((grid_density, grid_density)) * pixel_size\n",
    "    moire_sites = griddata((X.ravel(), Y.ravel()), distances.ravel(), (Se_bottom[:, 1], Se_bottom[:, 0]), method='cubic')\n",
    "    ax[0,1].imshow(distances, cmap='viridis', extent=(0, 512, 512, 0))\n",
    "    ax[0,1].set_title('W-W Distances')\n",
    "    color_bounds = [distances.min(), distances.max()]\n",
    "    plot_histogram(ax[1,1], moire_sites.ravel(), color_bounds)\n",
    "    distance_arrays.append(distances)\n",
    "\n",
    "    # W_top - Chalc_bottom distances\n",
    "    distances = np.linalg.norm(W_top[W_top_tree.query(points)[1]] - chalc_bottom[chalc_bottom_tree.query(points)[1]], axis=1)\n",
    "    distances = distances.reshape((grid_density, grid_density)) * pixel_size\n",
    "    moire_sites = griddata((X.ravel(), Y.ravel()), distances.ravel(), (Se_bottom[:, 1], Se_bottom[:, 0]), method='cubic')\n",
    "    ax[0,2].imshow(distances, cmap='viridis', extent=(0, 512, 512, 0))\n",
    "    ax[0,2].set_title('W-Chalc Distances')\n",
    "    color_bounds = [distances.min(), distances.max()]\n",
    "    plot_histogram(ax[1,2], moire_sites.ravel(), color_bounds)\n",
    "    distance_arrays.append(distances)\n",
    "\n",
    "    # Chalc_top - W_bottom distances\n",
    "    distances = np.linalg.norm(chalc_top[chalc_top_tree.query(points)[1]] - W_bottom[W_bottom_tree.query(points)[1]], axis=1)\n",
    "    distances = distances.reshape((grid_density, grid_density)) * pixel_size\n",
    "    moire_sites = griddata((X.ravel(), Y.ravel()), distances.ravel(), (Se_bottom[:, 1], Se_bottom[:, 0]), method='cubic')\n",
    "    ax[0,3].imshow(distances, cmap='viridis', extent=(0, 512, 512, 0))\n",
    "    ax[0,3].set_title('Chalc-W Distances')\n",
    "    color_bounds = [distances.min(), distances.max()]\n",
    "    plot_histogram(ax[1,3], moire_sites.ravel(), color_bounds)\n",
    "    distance_arrays.append(distances)\n",
    "\n",
    "    # Chalc_top - Chalc_bottom distances\n",
    "    distances = np.linalg.norm(chalc_top[chalc_top_tree.query(points)[1]] - chalc_bottom[chalc_bottom_tree.query(points)[1]], axis=1)\n",
    "    distances = distances.reshape((grid_density, grid_density)) * pixel_size\n",
    "    moire_sites = griddata((X.ravel(), Y.ravel()), distances.ravel(), (Se_bottom[:, 1], Se_bottom[:, 0]), method='cubic')\n",
    "    ax[0,4].imshow(distances, cmap='viridis', extent=(0, 512, 512, 0))\n",
    "    ax[0,4].set_title('Chalc-Chalc Distances')\n",
    "    color_bounds = [distances.min(), distances.max()]\n",
    "    plot_histogram(ax[1,4], moire_sites.ravel(), color_bounds)\n",
    "    distance_arrays.append(distances)\n",
    "\n",
    "    for a in ax.ravel()[:5]:\n",
    "        a.axis('off')\n",
    "    ax[1,0].axis('off')\n",
    "    fig.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    return distance_arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_predictions(image_number, results, n_classes):\n",
    "    # Plot the original image\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(results[image_number]['image'], cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.title('Original Image')\n",
    "\n",
    "    # Create a figure with subplots for masks and probabilities\n",
    "    fig, axs = plt.subplots(2, n_classes, figsize=(20, 10))\n",
    "\n",
    "    # Plot the masks\n",
    "    for j in range(n_classes):\n",
    "        axs[0, j].imshow(results[image_number]['mask'][j], cmap='gray')\n",
    "        axs[0, j].axis('off')\n",
    "        axs[0, j].set_title(f'Mask {j+1}')\n",
    "\n",
    "    # Plot the probabilities\n",
    "    for j in range(n_classes):\n",
    "        axs[1, j].imshow(results[image_number]['probability'][j], cmap='plasma')\n",
    "        axs[1, j].axis('off')\n",
    "        axs[1, j].set_title(f'Probability {j+1}')\n",
    "\n",
    "    plt.suptitle(f'Results for Image {image_number}', fontsize=16)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image 1 ~ 20 deg twist (or -40, however you want to think about it)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('/Users/austin/Documents/GitHub/DataGenSTEM')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git clone https://github.com/ahoust17/DataGenSTEM.git\n",
    "sys.path.append('./DataGenSTEM/DataGenSTEM')\n",
    "import DataGenSTEM.data_generator as dg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to crop the experimental images\n",
    "\n",
    "exp_data = exp_data = np.load(str(shared_folder + '/Experimental_datasets/WSSe_haadf.npz'))\n",
    "im_array = exp_data['im_array']\n",
    "pixel_size = exp_data['pixel_size']\n",
    "\n",
    "im_array = np.sum(im_array, axis = 0)\n",
    "\n",
    "# im_array = gaussian_filter(im_array, sigma=1)\n",
    "im_array = np.power(im_array,2.6)\n",
    "im_array = im_array - np.min(im_array)\n",
    "\n",
    "im_array = im_array / np.max(im_array)\n",
    "\n",
    "\n",
    "print(f\"Pixel size: {pixel_size.astype(float)} m/pix\")\n",
    "plt.figure()\n",
    "plt.imshow(im_array, cmap='gray')\n",
    "plt.axis('off')\n",
    "\n",
    "\n",
    "exp_hist, bins = np.histogram(im_array.ravel(), bins=256, range=(0.0, 1.0))\n",
    "\n",
    "selected_images = []\n",
    "for i, data in enumerate(train_loader):\n",
    "    if i >= 10:\n",
    "        break\n",
    "    images = data[0].numpy()  # Convert batch of images to numpy array\n",
    "    selected_images.append(images)\n",
    "\n",
    "selected_images = np.concatenate(selected_images, axis=0)\n",
    "selected_images = selected_images / selected_images.max()\n",
    "selected_images_raveled = selected_images.ravel()\n",
    "train_hist, _ = np.histogram(selected_images_raveled, bins=256, range=(0.0, 1.0))\n",
    "\n",
    "# Plot the histograms\n",
    "plt.figure()\n",
    "plt.hist(im_array.ravel(), bins=256, range=(0.0, 1.0), fc='k', ec='k', alpha=0.5, label='Experimental Image')\n",
    "plt.plot(bins[:-1], train_hist, 'r', alpha=0.5, label='Training Images')\n",
    "plt.legend()\n",
    "plt.xlabel('Intensity')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Histogram of Experimental Image vs Training Images')\n",
    "plt.show()\n",
    "\n",
    "n_crops = 10\n",
    "images = dg.shotgun_crop(im_array, crop_size=512, n_crops = n_crops, roi = None)\n",
    "# normalize each image in images to 0,1\n",
    "for i in range(n_crops):\n",
    "    images[i] = images[i] - np.min(images[i])\n",
    "    images[i] = images[i] / np.max(images[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select the top right of the image\n",
    "adj = 0\n",
    "manual_image = im_array[-512:, 0+200:512+200]\n",
    "manual_image = gaussian_filter(manual_image, sigma=1)\n",
    "\n",
    "manual_image -= manual_image.min()\n",
    "manual_image /= manual_image.max()\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(manual_image, cmap='gray')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from scipy.ndimage import center_of_mass\n",
    "# from scipy.spatial.distance import cdist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select input image\n",
    "im = manual_image\n",
    "test = torch.tensor(im.astype(np.float32)).unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "# show input image\n",
    "fig, ax = plt.subplots(1,1, figsize=(4,4))\n",
    "ax.imshow(test.squeeze().cpu().numpy(), cmap='gray')\n",
    "ax.axis('off')\n",
    "\n",
    "ensemble.predict(test, return_plot=False)\n",
    "ensemble.rearrange_ensemble()\n",
    "ensemble.vote(mode='max')\n",
    "ensemble.plot_vote()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binary = threshold(ensemble.voted_prediction.copy(), ['threshold'], param = 1)\n",
    "\n",
    "fig, ax = plt.subplots(1,6, figsize=(12, 5))\n",
    "for i in range(6):\n",
    "    ax[i].imshow(binary[i], cmap='gray')\n",
    "    ax[i].axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Do it batchwise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.ndimage import label\n",
    "results = {}\n",
    "for i, im in enumerate(images):\n",
    "    print(f\"Processing image {i}\")\n",
    "    # Model Prediction\n",
    "    im_tensor = torch.tensor(im.astype(np.float32)).unsqueeze(0).unsqueeze(0)\n",
    "    ensemble.predict(im_tensor, return_plot=False)\n",
    "    ensemble.rearrange_ensemble()\n",
    "    ensemble.vote(mode='max')\n",
    "    masks = threshold(ensemble.voted_prediction, ['threshold'], param = 0)\n",
    "    probability = ensemble.voted_prediction\n",
    "\n",
    "    centroids = {'W': [], 'S': [], 'Se': []}\n",
    "\n",
    "\n",
    "    for j, layer in enumerate(masks):\n",
    "        if j == 0 or j == 3:\n",
    "            element = 'S'\n",
    "        elif j == 1 or j == 4:\n",
    "            element = 'Se'\n",
    "        else:\n",
    "            element = 'W'\n",
    "\n",
    "        labeled_array, num_features = label(layer)\n",
    "        layer_centroids = np.array(center_of_mass(layer, labeled_array, range(1, num_features + 1)))\n",
    "        centroids[element].append(layer_centroids)\n",
    "\n",
    "    results[i] = {\n",
    "        \"image\": im,\n",
    "        \"mask\": masks,\n",
    "        \"probability\": probability,\n",
    "        \"centroids\": centroids,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_vegards_law(results):\n",
    "    S_atoms = []\n",
    "    Se_atoms = []\n",
    "    W_atoms = []\n",
    "\n",
    "    for result in results:\n",
    "        for element in results[result][\"centroids\"]:\n",
    "            for layer in results[result][\"centroids\"][element]:\n",
    "                if element == \"S\":\n",
    "                    S_atoms.append(layer)\n",
    "                elif element == \"Se\":\n",
    "                    Se_atoms.append(layer)\n",
    "                else:\n",
    "                    W_atoms.append(layer)\n",
    "\n",
    "    WS2_lattice = 3.15 # Angstrom, https://www.hqgraphene.com/WS2.php\n",
    "    WSe2_lattice = 3.28 # Angstrom, https://www.hqgraphene.com/WSe2.php\n",
    "    WSSe_lattice = 3.24 # Angstrom\n",
    "\n",
    "    # Vegard's law\n",
    "    def vegards_law(a1, a2, x1, x2):\n",
    "        return x1*a1 + x2*a2\n",
    "\n",
    "    x_Se = np.linspace(0,1,11)/2\n",
    "    ideal_lattice = vegards_law(WS2_lattice, WSSe_lattice, 1-x_Se, x_Se)\n",
    "\n",
    "    # plot the ideal lattice\n",
    "    plt.figure(figsize=(6,4), dpi = 300)\n",
    "    plt.plot(x_Se, ideal_lattice, 'k--', label='Vegard\\'s Law')\n",
    "\n",
    "    i = 0\n",
    "    j = 0\n",
    "    for S, Se, W in zip(S_atoms, Se_atoms, W_atoms):\n",
    "        # calculate stoichiometry:\n",
    "        n_S = len(S)\n",
    "        n_Se = len(Se)\n",
    "        \n",
    "        x_Se = n_Se / (n_S * (2 + n_Se/n_S))\n",
    "\n",
    "        # calculate the standard error of the mean\n",
    "        std_err_stoic = np.sqrt((x_Se * (1 - x_Se)) / (n_S + n_Se))\n",
    "\n",
    "        tree = KDTree(W)\n",
    "        distances, indices = tree.query(W, k=4)\n",
    "\n",
    "        # The first column of distances is zero (distance to itself), we want the next three columns\n",
    "        nearest_distances = distances[:, 1:4] * pixel_size # angstroms\n",
    "        # keep distances between 3.0 and 4.0\n",
    "        nearest_distances = nearest_distances[(nearest_distances > 3.0) & (nearest_distances < 3.25)]\n",
    "        avg_distance = np.mean(nearest_distances)\n",
    "        std_distance = np.std(nearest_distances)\n",
    "        std_error = std_distance / np.sqrt(len(nearest_distances))\n",
    "\n",
    "        if x_Se < 0.07:\n",
    "            color = '#1f77b4'\n",
    "            label = 'layer A'\n",
    "            plt.errorbar(x_Se, avg_distance, yerr=std_error, xerr=std_err_stoic,capsize=5, c=color, fmt='o', label=label if i == 0 else None)\n",
    "            i = i + 1\n",
    "        else:\n",
    "            color = '#d62728'\n",
    "            label = 'layer B'\n",
    "            plt.errorbar(x_Se, avg_distance, yerr=std_error, xerr=std_err_stoic,capsize=5, c=color, fmt='o', label=label if j == 0 else None)\n",
    "            j = j + 1\n",
    "\n",
    "    plt.xlabel('Se Fraction (x)')\n",
    "    plt.ylabel(f'Lattice Constant ($\\AA$)')\n",
    "    # plt.legend(loc='lower right', fontsize=12)\n",
    "    plt.legend(loc='upper left', fontsize=12)\n",
    "    plt.xlim(0,0.2)\n",
    "    plt.ylim(3.125, 3.185)\n",
    "    plt.xticks([0, 0.05, 0.1, 0.15, 0.2])\n",
    "    plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_vegards_law(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot all 10 crops on the same plot\n",
    "fig, ax = plt.subplots(2,5, figsize=(20, 12), sharex=True, sharey=True)\n",
    "colors = {'W': ['blue', 'cyan'], 'Se': ['green', 'lime'], 'S': ['red', 'magenta']}\n",
    "\n",
    "for i in range(10):\n",
    "    ax[i%2, i//2].imshow(images[i], cmap='gray')\n",
    "    ax[i%2, i//2].axis('off')\n",
    "\n",
    "    for element, colors in colors.items():\n",
    "        for idx, (color, layer) in enumerate(zip(colors, results[i]['centroids'][element])):\n",
    "            if layer.size > 0:\n",
    "                ax[i%2, i//2].scatter(layer[:, 1], layer[:, 0], color=color, label=f'{element} Layer {idx + 1} Centroids', alpha=0.4, s=5)\n",
    "    ax[i%2, i//2].legend(loc='upper right', fontsize=8)\n",
    "    colors = {'W': ['blue', 'cyan'], 'Se': ['green', 'lime'], 'S': ['red', 'magenta']}\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_number = 3  # Replace with the desired image number\n",
    "visualize_predictions(image_number, results, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# gettimg the moire distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_number = 0\n",
    "\n",
    "# define two layers - more Se gets calaled 'top'\n",
    "n_se_0 = len(results[image_number]['centroids']['Se'][0])\n",
    "n_se_1 = len(results[image_number]['centroids']['Se'][1])\n",
    "top_idx, bottom_idx = (1, 0) if n_se_1 >= n_se_0 else (0, 1)\n",
    "\n",
    "s_top = results[image_number]['centroids']['S'][top_idx]\n",
    "se_top = results[image_number]['centroids']['Se'][top_idx]\n",
    "w_top = results[image_number]['centroids']['W'][top_idx]\n",
    "chalc_top = np.vstack([arr for arr in [s_top, se_top] if arr.size > 0])\n",
    "\n",
    "s_bottom = results[image_number]['centroids']['S'][bottom_idx]\n",
    "se_bottom = results[image_number]['centroids']['Se'][bottom_idx]\n",
    "w_bottom = results[image_number]['centroids']['W'][bottom_idx]\n",
    "chalc_bottom = np.vstack([arr for arr in [s_bottom, se_bottom] if arr.size > 0])\n",
    "\n",
    "# clean up chalchogen sites that may overlap\n",
    "distance_cutoff = 10 # pixel\n",
    "\n",
    "tree = KDTree(w_top)\n",
    "distances, indices = tree.query(chalc_top, k=2)\n",
    "chalc_top = chalc_top[distances[:, 1] > distance_cutoff]\n",
    "\n",
    "tree = KDTree(w_bottom)\n",
    "distances, indices = tree.query(chalc_bottom, k=2)\n",
    "chalc_bottom = chalc_bottom[distances[:, 1] > distance_cutoff]\n",
    "\n",
    "# make updataed trees for distance maps\n",
    "w_top_tree, w_bottom_tree = KDTree(w_top), KDTree(w_bottom)\n",
    "chalc_top_tree, chalc_bottom_tree = KDTree(chalc_top), KDTree(chalc_bottom)\n",
    "\n",
    "\n",
    "# Create meshgrid for points\n",
    "grid_density = 512\n",
    "x = np.linspace(0, 512, grid_density)\n",
    "y = np.linspace(0, 512, grid_density)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "points = np.vstack((X.ravel(), Y.ravel())).T\n",
    "\n",
    "\n",
    "# Define a helper function to compute distances and moire sites\n",
    "def compute_distances(query_tree1, query_tree2, layer1, layer2):\n",
    "    distances = np.linalg.norm(layer1[query_tree1.query(points)[1]] - layer2[query_tree2.query(points)[1]], axis=1)\n",
    "    distances = distances.reshape((grid_density, grid_density))\n",
    "    return distances\n",
    "\n",
    "# Compute distances and moire sites for different layers\n",
    "w_w_distances = compute_distances(w_top_tree, w_bottom_tree, w_top, w_bottom)\n",
    "w_chalc_distances = compute_distances(w_top_tree, chalc_bottom_tree, w_top, chalc_bottom)\n",
    "chalc_w_distances = compute_distances(chalc_top_tree, w_bottom_tree, chalc_top, w_bottom)\n",
    "chalc_chalc_distances = compute_distances(chalc_top_tree, chalc_bottom_tree, chalc_top, chalc_bottom)\n",
    "\n",
    "# Store all distances for further use\n",
    "distance_arrays = [w_w_distances, w_chalc_distances, chalc_w_distances, chalc_chalc_distances]\n",
    "map_2H = np.sqrt(distance_arrays[1]**2 + distance_arrays[2]**2)\n",
    "\n",
    "# Plot results\n",
    "fig, ax = plt.subplots(1, 4, figsize=(20, 5))\n",
    "titles = ['W-W Distances', 'W-Chalc Distances', 'Chalc-W Distances', 'Chalc-Chalc Distances']\n",
    "for i in range(4):\n",
    "    ax[i].imshow(distance_arrays[i], cmap='viridis', extent=(0, 512, 512, 0))\n",
    "    ax[i].set_title(titles[i])\n",
    "    ax[i].axis('off')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_2H = map_2H * pixel_size\n",
    "plt.figure()\n",
    "plt.imshow(map_2H, cmap='bone')\n",
    "plt.axis('off')\n",
    "\n",
    "plt.scatter(se_top[:, 1], se_top[:, 0], s=5, c='r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moire_sites = griddata((X.ravel(), Y.ravel()), map_2H.ravel(), (se_top[:, 1], se_top[:, 0]), method='cubic')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 6))\n",
    "im1 = ax1.imshow(map_2H, cmap='inferno')\n",
    "ax1.scatter(se_top[:, 1], se_top[:, 0], s=20, c='orange', linewidths=0.5, edgecolors='k', label='Se Sites')\n",
    "fig.colorbar(im1, ax=ax1)\n",
    "ax1.axis('off')\n",
    "\n",
    "# Histogram calculations\n",
    "binedges = np.histogram_bin_edges(np.concatenate([moire_sites.ravel(), map_2H.ravel()]), bins=50)\n",
    "hist_moire_sites_top, _ = np.histogram(moire_sites.ravel(), bins=binedges, density=False)\n",
    "hist_map_2H_top, _ = np.histogram(map_2H.ravel(), bins=binedges, density=False)\n",
    "\n",
    "bin_widths = np.diff(binedges)\n",
    "hist_moire_sites_density_top = hist_moire_sites_top / (np.sum(hist_moire_sites_top) * bin_widths)\n",
    "hist_map_2H_density_top = hist_map_2H_top / (np.sum(hist_map_2H_top) * bin_widths)\n",
    "\n",
    "# Create colormap for Se_top\n",
    "norm_top = Normalize(vmin=np.min(map_2H), vmax=np.max(map_2H))\n",
    "colors_top = cm.inferno(norm_top(binedges[:-1]))\n",
    "ax2.bar(binedges[:-1], hist_map_2H_density_top, width=bin_widths, color=colors_top, edgecolor='k', alpha=1, label='map_ravel')\n",
    "ax2.step(binedges[:-1], hist_moire_sites_density_top, where='mid', color='#E69F00', linewidth=3, label='Se_site_distribution_top')\n",
    "\n",
    "ax2.legend()\n",
    "# ax2.set_xlim(0, 1)\n",
    "ax2.set_yticks([])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}\n",
    "moire_sites_total = []\n",
    "map_2H_total = []\n",
    "\n",
    "selected_images = [0,2,3,5,6,7]\n",
    "for i in selected_images:\n",
    "    im = images[i]\n",
    "    print(f\"Processing image {i}\")\n",
    "    # Model Prediction\n",
    "    im_tensor = torch.tensor(im.astype(np.float32)).unsqueeze(0).unsqueeze(0)\n",
    "    ensemble.predict(im_tensor, return_plot=False)\n",
    "    ensemble.rearrange_ensemble()\n",
    "    ensemble.vote(mode='max')\n",
    "    masks = threshold(ensemble.voted_prediction, ['threshold'], param = 0) \n",
    "    probability = ensemble.voted_prediction\n",
    "\n",
    "    centroids = {'W': [], 'S': [], 'Se': []}\n",
    "    for j, layer in enumerate(masks):\n",
    "        if j == 0 or j == 3:\n",
    "            element = 'S'\n",
    "        elif j == 1 or j == 4:\n",
    "            element = 'Se'\n",
    "        else:\n",
    "            element = 'W'\n",
    "\n",
    "        labeled_array, num_features = label(layer)\n",
    "        layer_centroids = np.array(center_of_mass(layer, labeled_array, range(1, num_features + 1)))\n",
    "        centroids[element].append(layer_centroids)\n",
    "\n",
    "    results[i] = {\n",
    "        \"image\": im,\n",
    "        \"mask\": masks,\n",
    "        \"probability\": probability,\n",
    "        \"centroids\": centroids,\n",
    "    }\n",
    "\n",
    "    # define two layers - more Se gets calaled 'top'\n",
    "    n_se_0 = len(results[i]['centroids']['Se'][0])\n",
    "    n_se_1 = len(results[i]['centroids']['Se'][1])\n",
    "    top_idx, bottom_idx = (1, 0) if n_se_1 >= n_se_0 else (0, 1)\n",
    "\n",
    "    s_top = results[i]['centroids']['S'][top_idx]\n",
    "    se_top = results[i]['centroids']['Se'][top_idx]\n",
    "    w_top = results[i]['centroids']['W'][top_idx]\n",
    "    chalc_top = np.vstack([arr for arr in [s_top, se_top] if arr.size > 0])\n",
    "\n",
    "    s_bottom = results[i]['centroids']['S'][bottom_idx]\n",
    "    se_bottom = results[i]['centroids']['Se'][bottom_idx]\n",
    "    w_bottom = results[i]['centroids']['W'][bottom_idx]\n",
    "    chalc_bottom = np.vstack([arr for arr in [s_bottom, se_bottom] if arr.size > 0])\n",
    "\n",
    "    results[i]['se_top'] = se_top\n",
    "    results[i]['se_bottom'] = se_bottom\n",
    "    results[i]['s_top'] = s_top\n",
    "    results[i]['s_bottom'] = s_bottom\n",
    "    results[i]['w_top'] = w_top\n",
    "    results[i]['w_bottom'] = w_bottom\n",
    "    \n",
    "\n",
    "    # clean up chalchogen sites that may overlap\n",
    "    distance_cutoff = 10  # pixel\n",
    "\n",
    "    tree = KDTree(w_top)\n",
    "    distances, indices = tree.query(chalc_top, k=2)\n",
    "    chalc_top = chalc_top[distances[:, 1] > distance_cutoff]\n",
    "\n",
    "    tree = KDTree(w_bottom)\n",
    "    distances, indices = tree.query(chalc_bottom, k=2)\n",
    "    chalc_bottom = chalc_bottom[distances[:, 1] > distance_cutoff]\n",
    "\n",
    "    # make updataed trees for distance maps\n",
    "    w_top_tree, w_bottom_tree = KDTree(w_top), KDTree(w_bottom)\n",
    "    chalc_top_tree, chalc_bottom_tree = KDTree(chalc_top), KDTree(chalc_bottom)\n",
    "\n",
    "    # Create meshgrid for points\n",
    "    grid_density = 512\n",
    "    x = np.linspace(0, 512, grid_density)\n",
    "    y = np.linspace(0, 512, grid_density)\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    points = np.vstack((X.ravel(), Y.ravel())).T\n",
    "\n",
    "    # Define a helper function to compute distances and moire sites\n",
    "    def compute_distances(query_tree1, query_tree2, layer1, layer2):\n",
    "        distances = np.linalg.norm(layer1[query_tree1.query(points)[1]] - layer2[query_tree2.query(points)[1]], axis=1)\n",
    "        distances = distances.reshape((grid_density, grid_density))\n",
    "        return distances\n",
    "\n",
    "    # Compute distances and moire sites for different layers\n",
    "    w_w_distances = compute_distances(w_top_tree, w_bottom_tree, w_top, w_bottom)\n",
    "    w_chalc_distances = compute_distances(w_top_tree, chalc_bottom_tree, w_top, chalc_bottom)\n",
    "    chalc_w_distances = compute_distances(chalc_top_tree, w_bottom_tree, chalc_top, w_bottom)\n",
    "    chalc_chalc_distances = compute_distances(chalc_top_tree, chalc_bottom_tree, chalc_top, chalc_bottom)\n",
    "\n",
    "    # Store all distances for further use\n",
    "    distance_arrays = [w_w_distances, w_chalc_distances, chalc_w_distances, chalc_chalc_distances]\n",
    "    map_2H = np.sqrt(distance_arrays[1]**2 + distance_arrays[2]**2)\n",
    "    map_2H = map_2H * pixel_size\n",
    "    map_2H_total.append(map_2H)\n",
    "    \n",
    "    moire_sites = griddata((X.ravel(), Y.ravel()), map_2H.ravel(), (se_top[:, 1], se_top[:, 0]), method='cubic')\n",
    "    moire_sites_total.append(moire_sites)\n",
    "\n",
    "    # store these now \n",
    "    results[i]['map_2H'] = map_2H\n",
    "    results[i]['moire_sites'] = moire_sites\n",
    "    results[i]['distance_arrays'] = distance_arrays\n",
    "\n",
    "    # Plot results  \n",
    "    plt.figure()\n",
    "    plt.imshow(results[i]['image'])\n",
    "\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 6))\n",
    "    im1 = ax1.imshow(map_2H, cmap='inferno')\n",
    "    ax1.scatter(se_top[:, 1], se_top[:, 0], s=20, c='orange', linewidths=0.5, edgecolors='k', label='Se Sites')\n",
    "    fig.colorbar(im1, ax=ax1)\n",
    "    ax1.axis('off')\n",
    "\n",
    "    # Histogram calculations\n",
    "    binedges = np.histogram_bin_edges(np.concatenate([moire_sites.ravel(), map_2H.ravel()]), bins=50)\n",
    "    hist_moire_sites_top, _ = np.histogram(moire_sites.ravel(), bins=binedges, density=False)\n",
    "    hist_map_2H_top, _ = np.histogram(map_2H.ravel(), bins=binedges, density=False)\n",
    "\n",
    "    bin_widths = np.diff(binedges)\n",
    "    hist_moire_sites_density_top = hist_moire_sites_top / (np.sum(hist_moire_sites_top) * bin_widths)\n",
    "    hist_map_2H_density_top = hist_map_2H_top / (np.sum(hist_map_2H_top) * bin_widths)\n",
    "\n",
    "    # Create colormap for Se_top\n",
    "    norm_top = Normalize(vmin=np.min(map_2H), vmax=np.max(map_2H))\n",
    "    colors_top = cm.inferno(norm_top(binedges[:-1]))\n",
    "    ax2.bar(binedges[:-1], hist_map_2H_density_top, width=bin_widths, color=colors_top, edgecolor='k', alpha=1, label='map_ravel')\n",
    "    ax2.step(binedges[:-1], hist_moire_sites_density_top, where='mid', color='#E69F00', linewidth=3, label='Se_site_distribution_top')\n",
    "    ax2.legend()\n",
    "    ax2.set_yticks([])\n",
    "    plt.tight_layout()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sites found on \"top\" vs \"bottom\":\n",
    "s_tops = []\n",
    "s_bottoms = []\n",
    "se_tops = []\n",
    "se_bottoms = []\n",
    "w_tops = []\n",
    "w_bottoms = []\n",
    "\n",
    "for i in selected_images:\n",
    "    s_tops.append(len(results[i]['s_top']))\n",
    "    s_bottoms.append(len(results[i]['s_bottom']))\n",
    "    se_tops.append(len(results[i]['se_top']))\n",
    "    se_bottoms.append(len(results[i]['se_bottom']))\n",
    "    w_tops.append(len(results[i]['w_top']))\n",
    "    w_bottoms.append(len(results[i]['w_bottom']))\n",
    "\n",
    "s_tops = np.array(s_tops)\n",
    "s_bottoms = np.array(s_bottoms)\n",
    "se_tops = np.array(se_tops)\n",
    "se_bottoms = np.array(se_bottoms)\n",
    "w_tops = np.array(w_tops)\n",
    "w_bottoms = np.array(w_bottoms)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean(se_bottoms/ (se_bottoms + s_bottoms))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moire_sites_total = np.concatenate([np.ravel(element) for element in moire_sites_total])\n",
    "map_2H_total = np.concatenate([np.ravel(element) for element in map_2H_total])\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "\n",
    "binedges = np.histogram_bin_edges(np.concatenate([moire_sites_total, map_2H_total]), bins=50)\n",
    "hist_moire_sites_top, _ = np.histogram(moire_sites_total, bins=binedges, density=False)\n",
    "hist_map_2H_top, _ = np.histogram(map_2H_total.ravel(), bins=binedges, density=False)\n",
    "\n",
    "bin_widths = np.diff(binedges)\n",
    "hist_moire_sites_density_top = hist_moire_sites_top / (np.sum(hist_moire_sites_top) * bin_widths)\n",
    "hist_map_2H_density_top = hist_map_2H_top / (np.sum(hist_map_2H_top) * bin_widths)\n",
    "\n",
    "# Create colormap for Se_top\n",
    "norm_top = Normalize(vmin=np.min(map_2H_total), vmax=np.max(map_2H_total))\n",
    "colors_top = cm.inferno(norm_top(binedges[:-1]))\n",
    "ax.bar(binedges[:-1], hist_map_2H_density_top, width=bin_widths, color=colors_top, edgecolor='k', alpha=1, label='map_ravel')\n",
    "ax.step(binedges[:-1], hist_moire_sites_density_top, where='mid', color='#E69F00', linewidth=3, label='Se_site_distribution_top')\n",
    "ax.legend()\n",
    "ax.set_yticks([])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KernelDensity\n",
    "\n",
    "# Kernel density estimation function\n",
    "def compute_kde(data, bandwidth=0.1, n_bins=100, x_min=0, x_max=10):\n",
    "    kde = KernelDensity(bandwidth=bandwidth)\n",
    "    kde.fit(data[:, np.newaxis])\n",
    "    x_d = np.linspace(x_min, x_max, n_bins)\n",
    "    log_dens = kde.score_samples(x_d[:, np.newaxis])\n",
    "    return np.exp(log_dens), x_d\n",
    "\n",
    "# Bootstrap KDE\n",
    "def bootstrap_kde(data, n_bootstrap=1000, bandwidth=0.1, n_bins=100, x_min=0, x_max=10):\n",
    "    kde_samples = np.zeros((n_bootstrap, n_bins))\n",
    "    for i in range(n_bootstrap):\n",
    "        resample = np.random.choice(data, size=len(data), replace=True)\n",
    "        kde_samples[i], x_d = compute_kde(resample, bandwidth, n_bins, x_min, x_max)\n",
    "    kde_mean = np.mean(kde_samples, axis=0)\n",
    "    kde_std = np.std(kde_samples, axis=0)\n",
    "    return kde_mean, kde_std, x_d\n",
    "\n",
    "# Normalize the x-axis\n",
    "n_bins = 100\n",
    "bandwidth = 0.1\n",
    "\n",
    "# Concatenate all histograms\n",
    "combined_atom_histogram = moire_sites_total.copy()\n",
    "combined_2H_map_histogram = map_2H_total.copy()\n",
    "\n",
    "# Apply KDE with bootstrapping to both histograms\n",
    "kde_atom_mean, kde_atom_std, x_d = bootstrap_kde(combined_atom_histogram, n_bootstrap=10, bandwidth=bandwidth, n_bins=n_bins)\n",
    "kde_2H_map_mean, kde_2H_map_std, _ = bootstrap_kde(combined_2H_map_histogram, n_bootstrap=10, bandwidth=bandwidth, n_bins=n_bins)\n",
    "\n",
    "# Compute bin edges for plotting\n",
    "binedges_total = x_d\n",
    "bin_widths_total = np.diff(x_d)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_size = exp_data['pixel_size']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binedges_total = x_d\n",
    "\n",
    "dispaly_iamge_number = 0\n",
    "lims = [0.2,4.7]\n",
    "\n",
    "map_2H = results[dispaly_iamge_number]['map_2H']\n",
    "moire_sites = results[dispaly_iamge_number]['moire_sites']\n",
    "distance_arrays = results[dispaly_iamge_number]['distance_arrays'] * pixel_size\n",
    "se_top = results[dispaly_iamge_number]['se_top']\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(18, 6), dpi=300)\n",
    "im1 = ax[0].imshow(map_2H, cmap='inferno', vmin=lims[0], vmax=lims[1])\n",
    "ax[0].scatter(se_top[:, 1], se_top[:, 0], s=120, c='green', linewidths=1, edgecolors='k', label='Se Sites')\n",
    "fig.colorbar(im1, ax=ax[0])\n",
    "ax[0].axis('off')\n",
    "\n",
    "ax[1].plot(binedges_total, kde_2H_map_mean, color='purple', label='2H Map Density')\n",
    "ax[1].fill_between(binedges_total, kde_2H_map_mean - kde_2H_map_std, kde_2H_map_mean + kde_2H_map_std, color='purple', alpha=0.3)\n",
    "\n",
    "ax[1].plot(binedges_total, kde_atom_mean, color='green', label='Se Site Density')\n",
    "ax[1].fill_between(binedges_total, kde_atom_mean - kde_atom_std, kde_atom_mean + kde_atom_std, color='green', alpha=0.3)\n",
    "\n",
    "ax[1].plot(binedges_total, kde_atom_mean - kde_2H_map_mean, color='grey', label='Difference', linestyle='--')\n",
    "ax[1].fill_between(binedges_total, kde_atom_mean - kde_2H_map_mean, np.zeros_like(kde_atom_mean - kde_2H_map_mean), color='grey', alpha=0.3)\n",
    "\n",
    "# ax.legend(loc='upper left')\n",
    "ax[1].set_xlim(lims[0], lims[1])\n",
    "ax[1].set_ylim(-0.13, 0.8)\n",
    "ax[1].axhline(y=0, color='k', linestyle='--', linewidth=1)\n",
    "# for dist in [0, 1, np.sqrt(2)]:\n",
    "    # ax[1].axvline(x=dist * 3.16, ymin=0, ymax=100, color='k', linestyle='--', linewidth=3)\n",
    "\n",
    "\n",
    "\n",
    "# plot the distance maps    \n",
    "fig, ax = plt.subplots(1,4, figsize = (20,5))\n",
    "titles = ['W-W Distances', 'W-Chalc Distances', 'Chalc-W Distances', 'Chalc-Chalc Distances']\n",
    "for i in range(4):\n",
    "    ax[i].imshow(distance_arrays[i], cmap='inferno', vmin=lims[0], vmax=lims[1])\n",
    "    ax[i].set_title(titles[i])\n",
    "    ax[i].axis('off')\n",
    "\n",
    "\n",
    "# make a new plot for the image\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 6), sharex=True, sharey=True, dpi=300)\n",
    "\n",
    "# plot the image\n",
    "ax[0].imshow(results[dispaly_iamge_number]['image'], cmap='gray')\n",
    "ax[0].scatter(se_top[:, 1], se_top[:, 0], s=60, c='limegreen', linewidths=0.5, edgecolors='k', label='Se Sites')\n",
    "ax[0].axis('off')\n",
    "ax[1].imshow(results[dispaly_iamge_number]['image'], cmap='gray')\n",
    "\n",
    "# plot each layer\n",
    "colors = {'W': ['blue', 'cyan'], 'Se': ['green', 'lime'], 'S': ['red', 'magenta']}\n",
    "for element, colors in colors.items():\n",
    "    for idx, (color, layer) in enumerate(zip(colors, results[dispaly_iamge_number]['centroids'][element])):\n",
    "        if layer.size > 0:\n",
    "            ax[1].scatter(layer[:, 1], layer[:, 0], color=color, label=f'{element} Layer {idx + 1} Centroids', alpha=0.4, s=5, linewidths=0.5, edgecolors='k')\n",
    "\n",
    "            colors = {'W': ['blue', 'cyan'], 'Se': ['green', 'lime'], 'S': ['red', 'magenta']}\n",
    "        \n",
    "ax[1].axis('off')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import entropy\n",
    "# Make sure to normalize KDEs to sum to 1 (since KL divergence expects probability distributions)\n",
    "kde_atom_mean_normalized = kde_atom_mean / kde_atom_mean.sum()\n",
    "kde_2H_map_mean_normalized = kde_2H_map_mean / kde_2H_map_mean.sum()\n",
    "\n",
    "kl_divergence = entropy(kde_atom_mean_normalized.ravel(), kde_2H_map_mean_normalized.ravel())\n",
    "print(f\"KL Divergence: {kl_divergence}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "# Assuming kde_atom_mean_normalized and kde_2H_map_mean_normalized are already defined and normalized as in the provided code\n",
    "mse = mean_squared_error(kde_2H_map_mean_normalized.ravel(), kde_atom_mean_normalized.ravel())\n",
    "print(f\"Mean Squared Error: {mse}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_image_number = 0\n",
    "\n",
    "# Define the figure and axes for 4 subplots\n",
    "fig, ax = plt.subplots(1, 4, figsize=(20, 5), sharex=True, sharey=True)\n",
    "colors = {'W': ['blue', 'cyan'], 'Se': ['red', 'magenta'], 'S': ['red', 'magenta']}\n",
    "# Image without any layers\n",
    "ax[0].imshow(results[display_image_number]['image'], cmap='gray')\n",
    "ax[0].axis('off')\n",
    "ax[0].set_title('Original Image')\n",
    "\n",
    "# Image with bottom (first) layer atoms only\n",
    "ax[1].imshow(results[display_image_number]['image'], cmap='gray')\n",
    "for element, color in colors.items():\n",
    "    layer = results[display_image_number]['centroids'][element][0]  # First layer atoms\n",
    "    if layer.size > 0:\n",
    "        ax[1].scatter(layer[:, 1], layer[:, 0], color=color[0], label=f'{element} Bottom Layer Centroids', alpha=0.6, s=60, linewidths=0.5, edgecolors='k')\n",
    "ax[1].axis('off')\n",
    "ax[1].set_title('Bottom Layer Atoms')\n",
    "\n",
    "# Image with top (second) layer atoms only\n",
    "ax[2].imshow(results[display_image_number]['image'], cmap='gray')\n",
    "for element, color in colors.items():\n",
    "    layer = results[display_image_number]['centroids'][element][1]  # Top layer atoms\n",
    "    if layer.size > 0:\n",
    "        ax[2].scatter(layer[:, 1], layer[:, 0], color=color[1], label=f'{element} Top Layer Centroids', alpha=0.6, s=60, linewidths=0.5, edgecolors='k')\n",
    "ax[2].axis('off')\n",
    "ax[2].set_title('Top Layer Atoms')\n",
    "\n",
    "# Image with all Se atoms\n",
    "ax[3].imshow(results[display_image_number]['image'], cmap='gray')\n",
    "for layer in results[display_image_number]['centroids']['Se']:  # All Se layers\n",
    "    if layer.size > 0:\n",
    "        \n",
    "        ax[3].scatter(layer[:, 1], layer[:, 0], color='lime', label='Se Sites', alpha=0.6, s=60, linewidths=0.5, edgecolors='k')\n",
    "ax[3].axis('off')\n",
    "ax[3].set_title('All Se Atoms')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reconstruction image for figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ase\n",
    "from ase.io import read, write\n",
    "from ase import Atoms, Atom\n",
    "from ase.visualize import view \n",
    "\n",
    "elements = ['W', 'Se', 'S']\n",
    "atoms = Atoms()\n",
    "\n",
    "for element in elements:\n",
    "    for z, layer in enumerate(results[display_image_number]['centroids'][element]):\n",
    "        z_height = (0.2 - z) * 200 * pixel_size  # Calculate the z-height\n",
    "        print(f\"Layer {z} for element {element}:\")\n",
    "        for atom in layer:\n",
    "            x = atom[1] * pixel_size\n",
    "            y = atom[0] * pixel_size\n",
    "            \n",
    "            if element in ['S', 'Se']:\n",
    "                # If the element is S or Se, append two atoms displaced in the z direction\n",
    "                atoms.append(Atom(element, (x, y, z_height + 1)))  # +z shift\n",
    "                atoms.append(Atom(element, (x, y, z_height - 1)))  # -z shift\n",
    "            else:\n",
    "                # For W, keep the original atom\n",
    "                atoms.append(Atom(element, (x, y, z_height)))\n",
    "\n",
    "positions = atoms.get_positions()\n",
    "min_pos = np.min(positions, axis=0)\n",
    "max_pos = np.max(positions, axis=0)\n",
    "cell_range = max_pos - min_pos\n",
    "atoms.set_cell(cell_range)\n",
    "\n",
    "# Save the Atoms object to a CIF file\n",
    "write('atoms.cif', atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view(atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "atoms = read('atoms.cif')\n",
    "xtal = atoms.copy()\n",
    "\n",
    "# Define the image size in Angstroms\n",
    "positions = xtal.get_positions()[:, :2]\n",
    "# xmin, xmax = np.min(positions[:, 0]), np.max(positions[:, 0])\n",
    "# ymin, ymax = np.min(positions[:, 1]), np.max(positions[:, 1])\n",
    "borders = 0\n",
    "xmin, xmax = 0, 512 * pixel_size\n",
    "ymin, ymax = 0, 512 * pixel_size\n",
    "axis_extent = (xmin - borders, xmax + borders, ymin - borders, ymax + borders)\n",
    "\n",
    "pixel_size = 0.106 # Angstrom/pixel, determines number of points, aka resolution of maps.  the xtal determines the fov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "atom_var = 0.22\n",
    "potential = dg.get_pseudo_potential(xtal = xtal, pixel_size = pixel_size, sigma = atom_var, axis_extent = axis_extent)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(potential, cmap='gray', extent=axis_extent)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make point spread function\n",
    "airy_disk_size = 1\n",
    "psf = dg.get_point_spread_function(airy_disk_radius = airy_disk_size, size = 32)\n",
    "psf_resize = dg.resize_image(np.array(psf), n = max(potential.shape)) # for plotting on same axes as image\n",
    "perfect_image = dg.convolve_kernel(potential, psf)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(perfect_image, cmap='gray', extent=axis_extent)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shot_noise = 0.7\n",
    "noisy_image = dg.add_poisson_noise(perfect_image, shot_noise = shot_noise)\n",
    "\n",
    "noisy_image = noisy_image - np.min(noisy_image)  # Normalize the noisy image to [0, 1]\n",
    "noisy_image = noisy_image / np.max(noisy_image)\n",
    "\n",
    "\n",
    "noisy_image = gaussian_filter(noisy_image, sigma=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].imshow(perfect_image, cmap='gray')\n",
    "ax[0].set_title('Perfect Image')\n",
    "\n",
    "ax[1].imshow(noisy_image, cmap='gray', vmin=0, vmax=1)\n",
    "ax[1].set_title('Noisy Image')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.hist(noisy_image.ravel(), bins=20, range=(0.0, 1.0), fc='k', ec='k', alpha=0.5, label='Noisy Image')\n",
    "plt.hist(results[display_image_number]['image'].ravel(), bins=20, range=(0.0, 1.0), fc='r', ec='r', alpha=0.5, label='Experimental Image')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_image = gaussian_filter(results[display_image_number]['image'],3)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 5), sharex=True, sharey=True)\n",
    "\n",
    "ax[0].imshow(results[display_image_number]['image'], cmap='gray')\n",
    "\n",
    "ax[1].imshow(noisy_image, cmap='gray', vmin=0, vmax=1)\n",
    "\n",
    "difference = results[display_image_number]['image'] - noisy_image\n",
    "diff = np.abs(difference)\n",
    "\n",
    "ax[2].imshow(diff, cmap='gray', vmin=0, vmax=1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crop_size = 100\n",
    "n_crops = 9\n",
    "crops = []\n",
    "for _ in range(n_crops):\n",
    "    x = np.random.randint(0, noisy_image.shape[1] - crop_size)\n",
    "    y = np.random.randint(0, noisy_image.shape[0] - crop_size)\n",
    "    crops.append((x, y))  # Store the coordinates as tuples\n",
    "\n",
    "# Create a subplot with 10 rows and 3 columns\n",
    "fig, ax = plt.subplots(10, 3, figsize=(15, 49), dpi=300)\n",
    "ax[0, 0].imshow(gaussian_filter(results[display_image_number]['image'],sigma=3), cmap='gray')\n",
    "ax[0, 1].imshow(noisy_image, cmap='gray', vmin=0, vmax=1)\n",
    "ax[0, 2].imshow(diff, cmap='gray', vmin=0, vmax=1)\n",
    "\n",
    "# Plot the cropped images\n",
    "for i in range(9):\n",
    "    x, y = crops[i]\n",
    "    ax[i + 1, 0].imshow(gaussian_filter(results[display_image_number]['image'][y:y + crop_size, x:x + crop_size],sigma=3), cmap='gray', vmin = 0, vmax = 1)\n",
    "    ax[i + 1, 1].imshow(noisy_image[y:y + crop_size, x:x + crop_size], cmap='gray', vmin=0, vmax=1)\n",
    "    ax[i + 1, 2].imshow(diff[y:y + crop_size, x:x + crop_size], cmap='gray', vmin=0, vmax=1)\n",
    "\n",
    "for a in ax.flat:\n",
    "    a.axis('off')\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch-nightly",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}