notebooks/lens-fit.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3094d16-cdd0-42b7-b892-4fb863d75384",
   "metadata": {
    "executionInfo": {
     "elapsed": 685,
     "status": "ok",
     "timestamp": 1652455590622,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "b3094d16-cdd0-42b7-b892-4fb863d75384"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import vmap, jit\n",
    "import numpyro.distributions as dist\n",
    "from einops import rearrange, repeat\n",
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9555cb-2d63-478a-9cbe-239a22b77d19",
   "metadata": {
    "executionInfo": {
     "elapsed": 810,
     "status": "ok",
     "timestamp": 1652455593247,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "2a9555cb-2d63-478a-9cbe-239a22b77d19"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pylab as pylab\n",
    "import warnings\n",
    "import matplotlib.cbook\n",
    "\n",
    "from plot_params import params\n",
    "\n",
    "warnings.filterwarnings(\"ignore\",category=matplotlib.cbook.mplDeprecation)\n",
    "\n",
    "pylab.rcParams.update(params)\n",
    "cols_default = plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
    "\n",
    "plt.rcParams.update({\n",
    "  \"text.usetex\": True,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69acb82a-2052-4add-ad4d-00484ee29bd4",
   "metadata": {
    "executionInfo": {
     "elapsed": 13696,
     "status": "ok",
     "timestamp": 1652455630634,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "69acb82a-2052-4add-ad4d-00484ee29bd4"
   },
   "outputs": [],
   "source": [
    "from gigalens.model import PhysicalModel\n",
    "from gigalens.jax.simulator import LensSimulator\n",
    "from gigalens.simulator import SimulatorConfig\n",
    "from gigalens.jax.profiles.light import interp\n",
    "from gigalens.jax.profiles.mass import shear, sie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55484ce3-46fb-4be6-8edd-60d58f12c741",
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1652455630905,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "55484ce3-46fb-4be6-8edd-60d58f12c741"
   },
   "outputs": [],
   "source": [
    "# Color maps for main and residual plots\n",
    "cmap = 'magma'\n",
    "cmap_resid = 'RdGy'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0282ba-a574-46cc-b2e6-111655496719",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a58d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"../data/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6fb5df0-6e33-49ed-b0b1-e48b7716b0fe",
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1652455630906,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "e6fb5df0-6e33-49ed-b0b1-e48b7716b0fe"
   },
   "outputs": [],
   "source": [
    "## Image configuration\n",
    "\n",
    "image_extent = 5  # Size of image in arcmin\n",
    "n_xy = 500  # Number of pixels\n",
    "\n",
    "n_sup = 1  # Supersampling of source factor during fitting\n",
    "n_xy_s = int(n_sup * n_xy)\n",
    "\n",
    "fwhm_psf = 0.005  # ELT FWHM of 5 mas from 1903.04742"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed64376-cc39-44a7-90c9-70f2db959b5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lens_conf_idx = 0  # Whether to use lens configuration 0 or 1 (see below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8dddb9-e02b-465b-83c6-ea73dbf559ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "## MLP configuration\n",
    "\n",
    "num_dense_layers = 4\n",
    "dense_layer_width = 256\n",
    "add_positional_encoding = False\n",
    "positional_encoding_dims = 5\n",
    "add_random_fourier = True\n",
    "bandwidth_rff = 5.\n",
    "add_skip_connection = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3fa688-edb6-4f5d-91d5-1447d08094e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## RNG; set to fixed values for deterministic run\n",
    "key_model = 42  # np.random.randint(1000)\n",
    "key_train = 43  # np.random.randint(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0793c191-314a-4a59-89c7-8f6c89506ed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training configuration\n",
    "\n",
    "n_q_samples = 64  # Number of MC samples used to build ELBO expectation\n",
    "\n",
    "learning_rate_init = 1e-3\n",
    "learning_rate_variational = 1e-3\n",
    "weight_decay = 1e-4\n",
    "momentum = 0.5\n",
    "\n",
    "train_epochs = 15000\n",
    "warmup_steps = 0\n",
    "plot_every = 200"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c1371a-0a90-411d-89d5-b394de47029d",
   "metadata": {
    "id": "14c1371a-0a90-411d-89d5-b394de47029d"
   },
   "source": [
    "## Source image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f90f5e-b51e-44aa-b0ee-6d183d4e953a",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 351
    },
    "executionInfo": {
     "elapsed": 2254,
     "status": "ok",
     "timestamp": 1652455633155,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "35f90f5e-b51e-44aa-b0ee-6d183d4e953a",
    "outputId": "aec61986-3dba-4857-adb6-562b971331b0"
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from skimage.transform import resize\n",
    "from urllib import request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c1f517-3d17-4929-8b9f-2ea33b4e81a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "galaxy_url = \"https://cdn.spacetelescope.org/archives/images/large/potw2015a.jpg\"\n",
    "galaxy_name = \"NGC2906\"\n",
    "\n",
    "# galaxy_url = \"https://cdn.spacetelescope.org/archives/images/large/opo0501a.jpg\"\n",
    "# galaxy_name = \"NGC1300\"\n",
    "\n",
    "# galaxy_url = \"https://cdn.spacetelescope.org/archives/images/large/potw1801a.jpg\"\n",
    "# galaxy_name = \"UGC6093\"\n",
    "\n",
    "if not os.path.exists(\"{}/galaxy_images/{}.jpg\".format(data_dir, galaxy_name)):\n",
    "    request.urlretrieve(galaxy_url, \"{}/galaxy_images/{}.jpg\".format(data_dir, galaxy_name))\n",
    "\n",
    "im_source = np.asarray(Image.open(\"{}/galaxy_images/{}.jpg\".format(data_dir, galaxy_name)).convert('L'))\n",
    "\n",
    "im_source_temp = np.zeros((np.max(im_source.shape), np.max(im_source.shape)))\n",
    "im_source_temp[int(im_source.shape[1] / 2 - im_source.shape[0] / 2):int(im_source.shape[1] / 2 + im_source.shape[0] / 2), :] = im_source\n",
    "im_source = im_source_temp\n",
    "im_source = im_source - np.min(im_source)\n",
    "im_source = im_source / np.max(im_source)\n",
    "if galaxy_name == \"NGC1300\":\n",
    "    im_source = np.pad(im_source, (1920, 1920))\n",
    "elif galaxy_name == \"UGC6093\":\n",
    "    im_source = np.pad(im_source, (640, 640))\n",
    "else:\n",
    "    im_source = np.pad(im_source, (1024, 1024))\n",
    "im_source = resize(im_source, (n_xy_s, n_xy_s), anti_aliasing=True)\n",
    "im_source[im_source < 0.05] = 0.\n",
    "plt.imshow(im_source , cmap=cmap)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2f9c1cf-2b5e-437c-aef6-40926c5fc5a2",
   "metadata": {},
   "source": [
    "## Lensed image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e67adeb7-aac0-48bd-a688-2aec955578d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_psf = fwhm_psf / 2 ** 1.5 * np.sqrt(np.log(2))  # Convert FWHM to standard deviation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdc9d656-e7be-41a4-b308-3b5d77d29f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Make Gaussian PSF kernel\n",
    "\n",
    "delta_pix = image_extent / n_xy\n",
    "\n",
    "x_kern = np.linspace(-6, 6, 13)\n",
    "x, y = np.meshgrid(x_kern, x_kern)\n",
    "kernel = multivariate_normal(mean=np.zeros(2), cov=np.ones(2) * sigma_psf / delta_pix).pdf(np.dstack((x, y)))\n",
    "kernel /= np.sum(kernel)\n",
    "\n",
    "# plt.imshow(kernel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c0f2e83-4fda-4104-8574-a26777c96e3c",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 345
    },
    "executionInfo": {
     "elapsed": 2755,
     "status": "ok",
     "timestamp": 1652455959393,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "1c0f2e83-4fda-4104-8574-a26777c96e3c",
    "outputId": "4fc15ee2-3c95-4347-ba97-e4b590fb0b4e"
   },
   "outputs": [],
   "source": [
    "## Lensed image\n",
    "\n",
    "# Supersample to reduce artifacting\n",
    "sim_config = SimulatorConfig(delta_pix=delta_pix, num_pix=n_xy_s, supersample=8, kernel=kernel)\n",
    "coord_half = sim_config.delta_pix * sim_config.num_pix / 2\n",
    "x_lims = [-coord_half, coord_half]\n",
    "y_lims = [-coord_half, coord_half]\n",
    "phys_model = PhysicalModel([sie.SIE(), shear.Shear()], [], [interp.Interp()])\n",
    "\n",
    "lens_sim = LensSimulator(phys_model, sim_config, bs=1)\n",
    "\n",
    "if lens_conf_idx == 0:\n",
    "    truth = [[\n",
    "        {'theta_E': 1.3, 'e1': 0.2, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.0},\n",
    "        {'gamma1': -0.02, 'gamma2': 0.02}\n",
    "    ], [ \n",
    "        {'img': im_source, 'x_lims': x_lims, 'y_lims': y_lims}\n",
    "    ]]\n",
    "elif lens_conf_idx == 1:\n",
    "    truth = [[\n",
    "        {'theta_E': 1.1, 'e1': 0.1, 'e2': 0.15, 'center_x': 0.4, 'center_y': -0.1},\n",
    "        {'gamma1': -0.03, 'gamma2': 0.04}\n",
    "    ], [ \n",
    "        {'img': im_source, 'x_lims': x_lims, 'y_lims': y_lims}\n",
    "    ]]\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "    \n",
    "std_noise = 1.e-6\n",
    "image = lens_sim.simulate(truth)\n",
    "image += std_noise * np.random.randn(*image.shape)\n",
    "image = image.reshape(-1, n_sup, image.shape[-1] // n_sup, n_sup).sum((-1, -3))\n",
    "\n",
    "# Revert supersampling configuration for forward simulator\n",
    "sim_config = SimulatorConfig(delta_pix=delta_pix, num_pix=n_xy_s, supersample=1, kernel=kernel)\n",
    "lens_sim = LensSimulator(phys_model, sim_config, bs=1)\n",
    "\n",
    "plt.imshow(image.T, origin='lower' , cmap=cmap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ff94da-307e-4884-9b41-92294363405e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute SNR around lens\n",
    "np.median(image[image > 0.5 * image.max()] / std_noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666bd0ce-61fc-4f5d-8d64-20c597b2eb98",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 297,
     "status": "ok",
     "timestamp": 1652455960710,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "666bd0ce-61fc-4f5d-8d64-20c597b2eb98",
    "outputId": "e78a2bcb-cc96-445e-dd81-739fe07c6d1b",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute max SNR\n",
    "image.max() / std_noise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cad5878-5bb4-4527-b83e-14d812dc57e8",
   "metadata": {
    "id": "6cad5878-5bb4-4527-b83e-14d812dc57e8"
   },
   "source": [
    "## Lensed image rendering function wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08dd77f1-df62-4dee-9b09-7da5b31a18ed",
   "metadata": {
    "executionInfo": {
     "elapsed": 266,
     "status": "ok",
     "timestamp": 1652455967628,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "08dd77f1-df62-4dee-9b09-7da5b31a18ed"
   },
   "outputs": [],
   "source": [
    "def render_lens(inputs):\n",
    "    \n",
    "    samples_src, samples_lens_params, x_lims, y_lims = inputs\n",
    "    \n",
    "    theta_E, theta_x, theta_y, e_1, e_2, gamma_1, gamma_2 = samples_lens_params\n",
    "    \n",
    "    truth = [[\n",
    "        {'theta_E': theta_E, 'e1': e_1, 'e2': e_2, 'center_x': theta_x, 'center_y': theta_y},\n",
    "        {'gamma1': gamma_1, 'gamma2': gamma_2}\n",
    "    ], [ \n",
    "        {'img': samples_src, 'x_lims': x_lims, 'y_lims': y_lims}\n",
    "    ]]\n",
    "\n",
    "    image = lens_sim.simulate(truth)\n",
    "    lens_sim.simulate(truth)\n",
    "    \n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ee2fb61-ed30-4dbe-9bf1-5d74d8684c61",
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1652455967899,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "6ee2fb61-ed30-4dbe-9bf1-5d74d8684c61"
   },
   "outputs": [],
   "source": [
    "render_lens_vmapped = vmap(jit(render_lens))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb72e84f-70a5-4ef5-99d8-0476fc5445c0",
   "metadata": {
    "id": "bb72e84f-70a5-4ef5-99d8-0476fc5445c0"
   },
   "source": [
    "## NeRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c96cfa17-ec96-4e57-b6c2-409b47d20864",
   "metadata": {
    "executionInfo": {
     "elapsed": 314,
     "status": "ok",
     "timestamp": 1652455972364,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "c96cfa17-ec96-4e57-b6c2-409b47d20864"
   },
   "outputs": [],
   "source": [
    "from models.nerf import NeRFModel\n",
    "\n",
    "import jax\n",
    "import optax\n",
    "import flax.linen as nn\n",
    "from flax.training import train_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cb23117-3977-48b2-92ca-dc4b1baef0c7",
   "metadata": {
    "executionInfo": {
     "elapsed": 323,
     "status": "ok",
     "timestamp": 1652455975383,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "2cb23117-3977-48b2-92ca-dc4b1baef0c7"
   },
   "outputs": [],
   "source": [
    "def initialize_model(key, input_pts_shape):\n",
    "\n",
    "    model = NeRFModel(num_dense_layers=num_dense_layers,\n",
    "                      dense_layer_width=dense_layer_width,\n",
    "                      add_positional_encoding=add_positional_encoding,\n",
    "                      add_random_fourier=add_random_fourier,\n",
    "                      add_skip_connection=add_skip_connection,\n",
    "                      positional_encoding_dims=positional_encoding_dims,\n",
    "                      bandwidth_rff=bandwidth_rff)\n",
    "\n",
    "    # Initialize the model parameters\n",
    "    initial_params = jax.jit(model.init)(\n",
    "        {\"params\": key},\n",
    "        jnp.ones(input_pts_shape),\n",
    "    )\n",
    "    \n",
    "    return model, initial_params[\"params\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be6a72ed-ee75-4336-96f4-6a7596c8faf9",
   "metadata": {
    "executionInfo": {
     "elapsed": 10,
     "status": "ok",
     "timestamp": 1652455982381,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "be6a72ed-ee75-4336-96f4-6a7596c8faf9"
   },
   "outputs": [],
   "source": [
    "theta_x_lims = theta_y_lims = x_lims\n",
    "theta_x_coords, theta_y_coords = jnp.meshgrid(jnp.linspace(theta_x_lims[0], theta_x_lims[1], n_xy_s), jnp.linspace(theta_y_lims[0], theta_y_lims[1], n_xy_s))\n",
    "\n",
    "coords = jnp.array([theta_x_coords, theta_y_coords])\n",
    "image_input =  jnp.expand_dims(image, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a6b32d-82d4-49c3-b40e-12bba78941fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "\n",
    "lct = dist.transforms.SoftplusLowerCholeskyTransform()\n",
    "\n",
    "@functools.partial(jax.jit, static_argnums=3)\n",
    "def train_step(state, batch, key, learning_rate_fn):\n",
    "    \"\"\" Single train step\n",
    "    \"\"\"\n",
    "    # Unravel the inputs and targets from the batch\n",
    "    inputs, targets = batch\n",
    "    inputs = rearrange(inputs, \"c h w -> (h w) c\")\n",
    "    \n",
    "    # jitter = jax.random.normal(key, shape=(2,)) * sim_config.delta_pix / 2.  # Add some jitter to the input coordinates\n",
    "    jitter = jax.random.normal(key, shape=(2,)) * 0. # No jitter\n",
    "    \n",
    "    inputs = inputs + jitter[jnp.newaxis, :]\n",
    "    \n",
    "    key_lens, key_source = jax.random.split(key)\n",
    "    \n",
    "    # Compute the loss\n",
    "    def loss_fn(params):\n",
    "\n",
    "        model_fn = lambda x: state.apply_fn({\"params\": params}, x)\n",
    "        model_output, loc, scale_tri = model_fn(inputs)\n",
    "                        \n",
    "        model_output = rearrange(model_output, \"(h w) mus -> mus (h w)\", h=n_xy_s, w=n_xy_s)\n",
    "        \n",
    "        ## Lens param loss\n",
    "        \n",
    "        # Lens variational distribution\n",
    "        scale_tri = lct(scale_tri)\n",
    "        dist_lens = dist.MultivariateNormal(loc=loc, scale_tril=scale_tri)\n",
    "        \n",
    "        # Lens prior\n",
    "        dist_prior_lens_unif = dist.Uniform(low=jnp.array([1., -0.5, -0.5]), high=jnp.array([2., 0.5, 0.5]))\n",
    "        dist_prior_lens_norm = dist.Normal(loc=jnp.array([0., 0., 0., 0.]), scale=jnp.array([0.3, 0.3, .05, 0.05]))\n",
    "        \n",
    "        # MC samples of lens parameters\n",
    "        samples_lens = dist_lens.sample(key=key_lens, sample_shape=(n_q_samples,))\n",
    "        \n",
    "        \n",
    "        loss_q_lens = jnp.sum(dist_lens.log_prob(samples_lens)) / n_q_samples                        \n",
    "        loss_prior_lens = -jnp.sum(dist_prior_lens_unif.log_prob(samples_lens[:, :-4])) / n_q_samples \\\n",
    "                            -jnp.sum(dist_prior_lens_norm.log_prob(samples_lens[:, -4:])) / n_q_samples \\\n",
    "\n",
    "        ## Source param loss\n",
    "        \n",
    "        # Source variational distribution\n",
    "        mu, sigma = nn.softplus(model_output[0]), nn.softplus(model_output[1])\n",
    "        dist_source = dist.Normal(mu, sigma)  \n",
    "        \n",
    "        # Source prior\n",
    "        dist_prior_source = dist.HalfNormal(scale=1. * jnp.ones_like(mu))\n",
    "        \n",
    "        # Source MC samples\n",
    "        samples_source = dist_source.sample(key=key_source, sample_shape=(n_q_samples,))\n",
    "        # samples_source = repeat(im_source, \"h w -> s (h w)\", s=n_q_samples)  # If testing with true source\n",
    "        \n",
    "        loss_q_source = jnp.sum(dist_source.log_prob(samples_source)) / n_q_samples\n",
    "        loss_prior_source = -jnp.sum(dist_prior_source.log_prob(samples_source)) / n_q_samples\n",
    "                \n",
    "        samples_source = rearrange(samples_source, \"s (h w) -> s h w\", h=n_xy_s, w=n_xy_s)\n",
    "        \n",
    "        x_lims = repeat(jnp.array([-coord_half, coord_half]) + jitter[0], \"xx -> nq xx\", nq=n_q_samples)\n",
    "        y_lims = repeat(jnp.array([-coord_half, coord_half]) + jitter[1], \"yy -> nq yy\", nq=n_q_samples)\n",
    "        \n",
    "        # Render lenses\n",
    "        lenses = render_lens_vmapped((samples_source, samples_lens, x_lims, y_lims))\n",
    "        lenses = lenses.reshape(lenses.shape[0], lenses.shape[-1] // n_sup, n_sup, lenses.shape[-1] // n_sup, n_sup).sum((-1, -3))\n",
    "        \n",
    "        # Compute log-probability\n",
    "        loss_likelihood = -jnp.sum(dist.Normal(targets, std_noise).log_prob(lenses)) / n_q_samples\n",
    "            \n",
    "        # Return total loss\n",
    "        return  loss_q_source + loss_prior_source + loss_q_lens + loss_prior_lens + loss_likelihood\n",
    "    \n",
    "    # Transform the loss function to get the loss value and the gradients\n",
    "    train_loss, gradients = jax.value_and_grad(loss_fn)(state.params)\n",
    "    \n",
    "    # Updated state\n",
    "    new_state = state.apply_gradients(grads=gradients)\n",
    "    train_loss = jnp.mean(train_loss)\n",
    "        \n",
    "    return train_loss, new_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48f06301-5657-450d-bba8-810da1e2987f",
   "metadata": {
    "id": "48f06301-5657-450d-bba8-810da1e2987f"
   },
   "outputs": [],
   "source": [
    "lr_schedule = optax.warmup_cosine_decay_schedule(init_value=0, peak_value=learning_rate_init, warmup_steps=warmup_steps, decay_steps=train_epochs - warmup_steps)\n",
    "lr_schedule_variational = optax.warmup_cosine_decay_schedule(init_value=0, peak_value=learning_rate_variational, warmup_steps=warmup_steps, decay_steps=train_epochs - warmup_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e66f93b-b222-4a5d-b851-b76d1829831c",
   "metadata": {
    "id": "8e66f93b-b222-4a5d-b851-b76d1829831c"
   },
   "outputs": [],
   "source": [
    "def map_nested_fn(fn):\n",
    "    '''Recursively apply `fn` to the key-value pairs of a nested dict'''\n",
    "    def map_fn(nested_dict):\n",
    "        return {k: (map_fn(v) if isinstance(v, dict) else fn(k, v)) for k, v in nested_dict.items()}\n",
    "    return map_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5cf91f-e13b-408f-a98f-d71271839850",
   "metadata": {
    "executionInfo": {
     "elapsed": 637,
     "status": "ok",
     "timestamp": 1652456236936,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "eb5cf91f-e13b-408f-a98f-d71271839850"
   },
   "outputs": [],
   "source": [
    "# Random Number Generator\n",
    "key, rng = jax.random.split(jax.random.PRNGKey(key_model))\n",
    "\n",
    "# Initialize the Model\n",
    "model, params = initialize_model(key, (n_xy_s * n_xy_s, 2))\n",
    "\n",
    "# Define the Optimizer\n",
    "optimizer = optax.adamw(learning_rate=lr_schedule, weight_decay=weight_decay)\n",
    "optimizer_variational = optax.adabelief(learning_rate=lr_schedule_variational)\n",
    "\n",
    "## Training state\n",
    "label_fn = map_nested_fn(lambda k, _: \"mlp\" if k in [\"w\", \"b\"] else \"var_params\")\n",
    "tx = optax.multi_transform({'mlp': optimizer, 'var_params': optimizer_variational}, label_fn)\n",
    "params = params.unfreeze()\n",
    "state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)\n",
    "\n",
    "# Train and validation loop\n",
    "def train_and_eval(state, train_step_fn, learning_rate_fn):\n",
    "        \n",
    "    train_loss_history, reconstructed_src_mu_history, reconstructed_src_std_history, reconstructed_lens_history = [], [], [], []\n",
    "    loc_history, scale_tri_history = [], []\n",
    "    key, rng = jax.random.split(jax.random.PRNGKey(key_train))\n",
    "    \n",
    "    for epoch in tqdm(range(train_epochs)):\n",
    "                \n",
    "        train_batch = coords, image_input\n",
    "\n",
    "        # Perform the Training Step\n",
    "        key, rng = jax.random.split(key)\n",
    "        train_loss, state = train_step_fn(state, train_batch, key, learning_rate_fn)\n",
    "                      \n",
    "        train_loss_history.append(np.asarray(np.mean(train_loss)))\n",
    "\n",
    "        # Plot the result every plot interval\n",
    "        if epoch % plot_every == 0:    \n",
    "            \n",
    "            model_fn = lambda x: state.apply_fn({\"params\": state.params}, x)\n",
    "            model_output, loc, scale_tri = model_fn(rearrange(coords, \"c h w -> (h w) c\"))\n",
    "            model_output = rearrange(model_output, \"(h w) c -> c h w\", h=n_xy_s, w=n_xy_s)\n",
    "                        \n",
    "            lenses = render_lens((nn.softplus(model_output[0]), loc, x_lims, y_lims))\n",
    "            lenses = lenses.reshape(lenses.shape[-1] // n_sup, n_sup, lenses.shape[-1] // n_sup, n_sup).sum((-1, -3))\n",
    "\n",
    "            reconstructed_lens_history.append(np.asarray(lenses.T))\n",
    "            reconstructed_src_mu_history.append(np.asarray(nn.softplus(model_output[0])))\n",
    "            reconstructed_src_std_history.append(np.asarray(nn.softplus(model_output[1])))\n",
    "            \n",
    "            loc_history.append(np.asarray(loc))\n",
    "            scale_tri_history.append(np.asarray(scale_tri))\n",
    "    \n",
    "    hist = {\n",
    "        \"train_loss\": train_loss_history,\n",
    "        \"reconstructed_lens\": reconstructed_lens_history,\n",
    "        \"reconstructed_src_mu\": reconstructed_src_mu_history,\n",
    "        \"reconstructed_src_std\": reconstructed_src_std_history,\n",
    "        \"loc_history\": loc_history,\n",
    "        \"scale_tri_history\": scale_tri_history\n",
    "    }\n",
    "    \n",
    "    return state, hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ab0df4d-61d5-443e-a90b-6cd35a57db7f",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4ab0df4d-61d5-443e-a90b-6cd35a57db7f",
    "outputId": "2f1e1dae-32f7-46a6-f6d0-d07610d091fe"
   },
   "outputs": [],
   "source": [
    "state, hist = train_and_eval(state, train_step, lr_schedule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22336120",
   "metadata": {},
   "outputs": [],
   "source": [
    "snr = ((hist[\"reconstructed_src_mu\"][-1])) / ((hist[\"reconstructed_src_std\"])[-1]) < 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "123da083-3dc6-4105-94c5-98cc7393fc6e",
   "metadata": {},
   "source": [
    "## Paper Figure 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d9b39d3-6fae-450c-9b43-9c82c46a6078",
   "metadata": {},
   "outputs": [],
   "source": [
    "truths = [truth[0][0]['theta_E'], \n",
    "                     truth[0][0]['center_x'], \n",
    "                     truth[0][0]['center_y'], \n",
    "                     truth[0][0]['e1'],\n",
    "                     truth[0][0]['e2'],\n",
    "                     truth[0][1]['gamma1'],\n",
    "                     truth[0][1]['gamma2']]\n",
    "\n",
    "labels = [r\"$\\theta_\\mathrm{E}$\",\n",
    "          r\"$\\theta_x$\",\n",
    "          r\"$\\theta_y$\",\n",
    "          r\"$\\epsilon_1$\",\n",
    "          r\"$\\epsilon_2$\",\n",
    "          r\"$\\gamma_1$\",\n",
    "          r\"$\\gamma_2$\",]          \n",
    "\n",
    "suptitles = [r\"\\bf{Einstein radius\\,[$''$]}\",\n",
    "            r\"\\bf{Lens-source offsets\\,[$''$]}\",\n",
    "            r\"\",\n",
    "            r\"\\bf{Lens eccentricities}\",\n",
    "            r\"\",\n",
    "            r\"\\bf{External shear}\",\n",
    "            r\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0c60bb4-cce2-4bce-935b-d814917184b6",
   "metadata": {
    "executionInfo": {
     "elapsed": 479,
     "status": "ok",
     "timestamp": 1652454506204,
     "user": {
      "displayName": "Siddharth Mishra-Sharma",
      "userId": "12134882132898528741"
     },
     "user_tz": 240
    },
    "id": "e0c60bb4-cce2-4bce-935b-d814917184b6"
   },
   "outputs": [],
   "source": [
    "from matplotlib.patches import Ellipse\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "def plot_fisher_single(i ,cov, mu, ax, lims, truths, labels) :\n",
    "    nb=128\n",
    "    sigma=np.sqrt(cov[i,i])\n",
    "    x_arr=mu[i]-8*sigma+16*sigma*np.arange(nb)/(nb-1.)\n",
    "    p_arr=np.exp(-(x_arr-mu[i])**2/(2*sigma**2))\n",
    "    ax.plot(x_arr,p_arr, color=cols_default[0])    \n",
    "    ax.axvline(truths[i], color='darkgrey', lw=1.5)\n",
    "    ax.set_xlim(lims[0][i], lims[1][i])\n",
    "    ax.set_ylim(-0.02, 1.2)\n",
    "    ax.set_title(labels[i], fontsize=18.5, y=1.02)\n",
    "    \n",
    "    if i == 0:\n",
    "        custom_lines = [Line2D([0], [0], color=cols_default[0], lw=4),\n",
    "                Line2D([0], [0], color='darkgrey', lw=4)]\n",
    "        \n",
    "        ax.legend(custom_lines, [\"Posterior\", \"Truth\"], bbox_to_anchor=[6, 0.5], frameon=True, fancybox=True, fontsize=19)\n",
    "\n",
    "\n",
    "def plot_fisher_two(i1, i2, cov, mu, ax, lims, truths, labels) :\n",
    "    covar=np.zeros([2,2])\n",
    "    covar[0,0]=cov[i1,i1]\n",
    "    covar[0,1]=cov[i1,i2]\n",
    "    covar[1,0]=cov[i2,i1]\n",
    "    covar[1,1]=cov[i2,i2]\n",
    "    sig0=np.sqrt(covar[0,0])\n",
    "    sig1=np.sqrt(covar[1,1])\n",
    "\n",
    "    w,v=np.linalg.eigh(covar)\n",
    "    \n",
    "    angle=180*np.arctan2(v[1,0],v[0,0])/np.pi\n",
    "    a_1s=np.sqrt(2.3*w[0])\n",
    "    b_1s=np.sqrt(2.3*w[1])\n",
    "    a_2s=np.sqrt(6.17*w[0])\n",
    "    b_2s=np.sqrt(6.17*w[1])\n",
    "\n",
    "    centre=np.array([mu[i1],mu[i2]])\n",
    "        \n",
    "    e_1s=Ellipse(xy=centre,width=2*a_1s,height=2*b_1s,angle=angle, color=cols_default[0], lw=2.)\n",
    "    e_1s.set_fill(False)\n",
    "    \n",
    "    e_2s=Ellipse(xy=centre,width=2*a_2s,height=2*b_2s,angle=angle, ls='--', color=cols_default[0], lw=2.)\n",
    "    e_2s.set_fill(False)\n",
    "\n",
    "    ax.add_artist(e_1s)\n",
    "    ax.add_artist(e_2s)\n",
    "    \n",
    "    ax.axvline(truths[i1], color='darkgrey', lw=1.5)\n",
    "    ax.axhline(truths[i2], color='darkgrey', lw=1.5)\n",
    "\n",
    "    ax.set_xlim(lims[0][i1], lims[1][i1])\n",
    "    ax.set_ylim(lims[0][i2], lims[1][i2])\n",
    "    \n",
    "    ax.set_xlabel(labels[i1])\n",
    "    ax.set_ylabel(labels[i2])\n",
    "\n",
    "def plot_fisher_all(mu, cov, lims, truths, labels, suptitles): \n",
    "    n_params = len(mu)\n",
    "\n",
    "    fig=plt.figure(figsize=(12, 12))\n",
    "    plt.subplots_adjust(hspace=0, wspace=0)\n",
    "    for i in np.arange(n_params):\n",
    "        i_col=i\n",
    "        for j in np.arange(n_params-i)+i :\n",
    "            i_row=j\n",
    "            iplot=i_col+n_params*i_row+1\n",
    "\n",
    "            ax=fig.add_subplot(n_params,n_params,iplot)\n",
    "            if i==j :\n",
    "                plot_fisher_single(i, cov, mu, ax, lims, truths, labels)\n",
    "                ax.text(0.5, 1.5, suptitles[i], \n",
    "                         horizontalalignment='left' if not i==0 else \"center\",\n",
    "                         verticalalignment='center',\n",
    "                         transform = ax.transAxes,\n",
    "                         fontsize=18)\n",
    "\n",
    "            else :\n",
    "                plot_fisher_two(i, j, cov, mu, ax, lims, truths, labels)\n",
    "\n",
    "            if i_row!=n_params-1 :\n",
    "                ax.get_xaxis().set_visible(False)\n",
    "\n",
    "            if i_col!=0 :\n",
    "                ax.get_yaxis().set_visible(False)\n",
    "\n",
    "            if i_col==0 and i_row==0 :\n",
    "                ax.get_yaxis().set_visible(False)\n",
    "                            \n",
    "            ax.locator_params(nbins=2)\n",
    "            [l.set_rotation(45) for l in ax.get_xticklabels()]\n",
    "            [l.set_rotation(45) for l in ax.get_yticklabels()]\n",
    "\n",
    "    fig.align_labels()\n",
    "    \n",
    "    plt.savefig(\"figures/joint_lens.pdf\", bbox_inches='tight')\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "i = -1\n",
    "\n",
    "chol = lct(hist['scale_tri_history'][i])\n",
    "cov = np.dot(chol, chol.T)\n",
    "mu = hist['loc_history'][i]\n",
    "\n",
    "lims = [mu - 6 * np.sqrt(np.diag(cov)), mu + 6 * np.sqrt(np.diag(cov))]\n",
    "\n",
    "plot_fisher_all(mu, cov, lims, truths, labels, suptitles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "becbf59b-42bb-4065-b2b7-6b07b3da7388",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Save corner plot data\n",
    "# np.savez(\"../data/corner_data.npz\",\n",
    "#          mu=mu, cov=cov, lims=lims, truths=truths, labels=labels, suptitles=suptitles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2405f36-9a30-466d-844a-c3eacfb10f98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# corner_data = np.load(\"../data/corner_data.npz\")\n",
    "# mu, cov, lims, truths, labels, suptitles = corner_data[\"mu\"], corner_data[\"cov\"], corner_data[\"lims\"], corner_data[\"truths\"], corner_data[\"labels\"], corner_data[\"suptitles\"]\n",
    "# plot_fisher_all(mu, cov, lims, truths, labels, suptitles)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dfdfaa8-5375-4b26-a194-99e2e676b761",
   "metadata": {},
   "source": [
    "## Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "172e14db-4b2f-4952-9e19-fe93cb87879a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from celluloid import Camera\n",
    "\n",
    "epoch = -1\n",
    "\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig, ax = plt.subplots(1, 6, figsize=(19, 4.5))\n",
    "\n",
    "camera = Camera(fig)\n",
    "\n",
    "for i in range(len(hist[\"reconstructed_lens\"])):\n",
    "\n",
    "    source_coord_limit = 1.\n",
    "\n",
    "    im_list = []\n",
    "\n",
    "    im_list.append(ax[0].imshow(im_source * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    ax[0].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[0].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    ax[0].set_title(r\"\\bf{True source}\")\n",
    "\n",
    "    im_list.append(ax[1].imshow((hist[\"reconstructed_src_mu\"][i]) * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    ax[1].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[1].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    ax[1].set_title(r\"\\bf{Predicted source}\")\n",
    "\n",
    "    im_list.append(ax[2].imshow((hist[\"reconstructed_src_mu\"][i] - im_source) * ~snr * lens_sim.conversion_factor / std_noise, vmin=-4, vmax=4, cmap=cmap_resid, extent=[*x_lims, *y_lims]))\n",
    "    ax[2].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[2].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    ax[2].set_title(r\"\\bf{Source residuals}\")\n",
    "\n",
    "    im_list.append(ax[3].imshow(image_input[0].T / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "    im_list.append(ax[4].imshow((hist[\"reconstructed_lens\"][i]) / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "    im_list.append(ax[5].imshow((hist[\"reconstructed_lens\"][i] - image_input[0].T) / std_noise, cmap=cmap_resid, vmin=-4, vmax=4, extent=[*x_lims, *y_lims]))\n",
    "    \n",
    "    # ax[6].plot(np.arange(plot_every * (i)), np.array(hist[\"train_loss\"][:plot_every * (i)]) / (n_xy * n_xy), color=cols_default[0])\n",
    "    \n",
    "    ax[3].set_title(r\"\\bf{True lensed image}\")\n",
    "    ax[4].set_title(r\"\\bf{Predicted lensed image}\")\n",
    "    ax[5].set_title(r\"\\bf{Lensed image residuals}\")\n",
    "    # ax[6].set_title(r\"\\bf{Loss}\")\n",
    "\n",
    "    for i_lens in [3, 4]:\n",
    "        ax[i_lens].plot([-2.2, -1.2], [-2.2, -2.2], c=\"white\", lw=1.5, ls=\"-\")\n",
    "        ax[i_lens].text(-1.7, -1.85, \"$1''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "    for i_src in [0, 1]:\n",
    "        ax[i_src].plot([-0.85, -0.35], [-0.85, -0.85], c=\"white\", lw=1.5, ls=\"-\")\n",
    "        ax[i_src].text(-0.6, -0.75, \"$0.5''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "    # labels = [\"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\", \"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\"]\n",
    "    # for j in range(6):\n",
    "    #     divider = make_axes_locatable(ax[j])\n",
    "    #     cax = divider.append_axes('bottom', size='5%', pad=0.2)\n",
    "    #     cbar = fig.colorbar(im_list[j], cax=cax, orientation=\"horizontal\", label=labels[j])\n",
    "\n",
    "\n",
    "    # labels = [r'\\emph{(a)}',r'\\emph{(b)}',r'\\emph{(c)}',r'\\emph{(d)}',r'\\emph{(e)}',r'\\emph{(f)}']\n",
    "    # colors = ['w','w','k','w','w','k']\n",
    "    # for i in range(6):\n",
    "    #     ax[i].text(0.1, 0.9, labels[i], \n",
    "    #          horizontalalignment='center',\n",
    "    #          verticalalignment='center',\n",
    "    #          transform = ax[i].transAxes,\n",
    "    #          color=colors[i], fontsize=17)\n",
    "\n",
    "    [axx.xaxis.set_visible(False) for axx in ax.flatten()];\n",
    "    [axx.yaxis.set_visible(False) for axx in ax.flatten()];\n",
    "    plt.tight_layout();\n",
    "    plt.show()\n",
    "    camera.snap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cea6caa8-7c2c-4945-9703-3dba74f657ed",
   "metadata": {
    "id": "cea6caa8-7c2c-4945-9703-3dba74f657ed",
    "outputId": "44ccea08-a5d5-40e4-d2cd-fa7a2d99425a"
   },
   "outputs": [],
   "source": [
    "animation = camera.animate()\n",
    "animation.save('figures/render.gif', fps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c718e7a7-7ca7-4ac7-869e-ca64867eaf98",
   "metadata": {},
   "source": [
    "## Paper Figure 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72000944-a01d-4a35-b098-f46c8637c22d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig, ax = plt.subplots(1, 6, figsize=(19, 4.5))\n",
    "\n",
    "i = -1\n",
    "\n",
    "source_coord_limit = 1.\n",
    "\n",
    "im_list = []\n",
    "\n",
    "im_list.append(ax[0].imshow(im_source / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "ax[0].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[0].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[0].set_title(r\"\\bf{True source}\")\n",
    "\n",
    "im_list.append(ax[1].imshow((hist[\"reconstructed_src_mu\"][i]) / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "ax[1].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[1].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[1].set_title(r\"\\bf{Predicted source}\")\n",
    "\n",
    "im_list.append(ax[2].imshow((hist[\"reconstructed_src_mu\"][i] - im_source) * lens_sim.conversion_factor / std_noise, vmin=-4, vmax=4, cmap=cmap_resid, extent=[*x_lims, *y_lims]))\n",
    "ax[2].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[2].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[2].set_title(r\"\\bf{Source residuals}\")\n",
    "\n",
    "im_list.append(ax[3].imshow(image_input[0].T / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "im_list.append(ax[4].imshow((hist[\"reconstructed_lens\"][i]) / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "im_list.append(ax[5].imshow((hist[\"reconstructed_lens\"][i] - image_input[0].T) / std_noise, cmap=cmap_resid, vmin=-4, vmax=4, extent=[*x_lims, *y_lims]))\n",
    "\n",
    "ax[3].set_title(r\"\\bf{True lensed image}\")\n",
    "ax[4].set_title(r\"\\bf{Predicted lensed image}\")\n",
    "ax[5].set_title(r\"\\bf{Lensed image residuals}\")\n",
    "\n",
    "for i_lens in [3, 4]:\n",
    "    ax[i_lens].plot([-2.2, -1.2], [-2.2, -2.2], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[i_lens].text(-1.7, -1.85, \"$1''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "for i_src in [0, 1]:\n",
    "    ax[i_src].plot([-0.85, -0.35], [-0.85, -0.85], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[i_src].text(-0.6, -0.75, \"$0.5''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "labels = [\"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\", \"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\"]\n",
    "for i in range(6):\n",
    "    divider = make_axes_locatable(ax[i])\n",
    "    cax = divider.append_axes('bottom', size='5%', pad=0.2)\n",
    "    cbar = fig.colorbar(im_list[i], cax=cax, orientation=\"horizontal\", label=labels[i])\n",
    "\n",
    "\n",
    "labels = [r'\\emph{(a)}',r'\\emph{(b)}',r'\\emph{(c)}',r'\\emph{(d)}',r'\\emph{(e)}',r'\\emph{(f)}']\n",
    "colors = ['w','w','k','w','w','k']\n",
    "for i in range(6):\n",
    "    ax[i].text(0.1, 0.9, labels[i], \n",
    "         horizontalalignment='center',\n",
    "         verticalalignment='center',\n",
    "         transform = ax[i].transAxes,\n",
    "         color=colors[i], fontsize=17)\n",
    "\n",
    "[axx.xaxis.set_visible(False) for axx in ax.flatten()];\n",
    "[axx.yaxis.set_visible(False) for axx in ax.flatten()];\n",
    "\n",
    "plt.tight_layout()\n",
    "# plt.savefig(\"figures/reconstruction.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a1fffc-a8a2-470b-acfa-b6b442dc09c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig, ax = plt.subplots(1, 6, figsize=(19, 4.5))\n",
    "\n",
    "i = -1\n",
    "\n",
    "source_coord_limit = 1.\n",
    "\n",
    "im_list = []\n",
    "\n",
    "im_list.append(ax[0].imshow(im_source * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "ax[0].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[0].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[0].set_title(r\"\\bf{True source}\")\n",
    "\n",
    "im_list.append(ax[1].imshow((hist[\"reconstructed_src_mu\"][i]) * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "ax[1].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[1].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[1].set_title(r\"\\bf{Predicted source}\")\n",
    "\n",
    "im_list.append(ax[2].imshow((hist[\"reconstructed_src_mu\"][i] - im_source) * ~snr * lens_sim.conversion_factor / std_noise, vmin=-4, vmax=4, cmap=cmap_resid, extent=[*x_lims, *y_lims]))\n",
    "ax[2].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[2].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[2].set_title(r\"\\bf{Source residuals}\")\n",
    "\n",
    "im_list.append(ax[3].imshow(image_input[0].T / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "im_list.append(ax[4].imshow((hist[\"reconstructed_lens\"][i]) / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "im_list.append(ax[5].imshow((hist[\"reconstructed_lens\"][i] - image_input[0].T) / std_noise, cmap=cmap_resid, vmin=-4, vmax=4, extent=[*x_lims, *y_lims]))\n",
    "\n",
    "ax[3].set_title(r\"\\bf{True lensed image}\")\n",
    "ax[4].set_title(r\"\\bf{Predicted lensed image}\")\n",
    "ax[5].set_title(r\"\\bf{Lensed image residuals}\")\n",
    "\n",
    "for i_lens in [3, 4]:\n",
    "    ax[i_lens].plot([-2.2, -1.2], [-2.2, -2.2], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[i_lens].text(-1.7, -1.85, \"$1''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "for i_src in [0, 1]:\n",
    "    ax[i_src].plot([-0.85, -0.35], [-0.85, -0.85], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[i_src].text(-0.6, -0.75, \"$0.5''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "labels = [\"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\", \"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\"]\n",
    "for i in range(6):\n",
    "    divider = make_axes_locatable(ax[i])\n",
    "    cax = divider.append_axes('bottom', size='5%', pad=0.2)\n",
    "    cbar = fig.colorbar(im_list[i], cax=cax, orientation=\"horizontal\", label=labels[i])\n",
    "\n",
    "\n",
    "labels = [r'\\emph{(a)}',r'\\emph{(b)}',r'\\emph{(c)}',r'\\emph{(d)}',r'\\emph{(e)}',r'\\emph{(f)}']\n",
    "colors = ['w','w','k','w','w','k']\n",
    "for i in range(6):\n",
    "    ax[i].text(0.1, 0.9, labels[i], \n",
    "         horizontalalignment='center',\n",
    "         verticalalignment='center',\n",
    "         transform = ax[i].transAxes,\n",
    "         color=colors[i], fontsize=17)\n",
    "\n",
    "[axx.xaxis.set_visible(False) for axx in ax.flatten()];\n",
    "[axx.yaxis.set_visible(False) for axx in ax.flatten()];\n",
    "\n",
    "plt.tight_layout()\n",
    "# plt.savefig(\"figures/reconstruction.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a51975-2d26-4e70-9eae-b5302360e390",
   "metadata": {},
   "source": [
    "## Save plotting data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2216b40-64f2-4e22-961e-617fd167ddb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53edb319-2c9d-470a-a4d4-7b206e5f8c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save figure data\n",
    "np.savez(\"../data/fig1_data.npz\",\n",
    "         reconstructed_lens=hist[\"reconstructed_lens\"][i], \n",
    "         reconstructed_src_mu=hist[\"reconstructed_src_mu\"][i],\n",
    "         snr=snr,\n",
    "         im_source=im_source,\n",
    "         image_input=image_input[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0295670c-90d9-40a6-8a76-4462d1973d17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save figure data\n",
    "\n",
    "i = -1\n",
    "np.savez(\"../data/fig1_data_no_pos.npz\",\n",
    "         reconstructed_lens=hist[\"reconstructed_lens\"][i], \n",
    "         reconstructed_src_mu=hist[\"reconstructed_src_mu\"][i],\n",
    "         snr=snr,\n",
    "         im_source=im_source,\n",
    "         image_input=image_input[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db6ad360-6f6f-4e14-b2c3-0a1ed4592481",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save figure data\n",
    "\n",
    "i = -1\n",
    "np.savez(\"../data/fig1_data_rff.npz\",\n",
    "         reconstructed_lens=hist[\"reconstructed_lens\"][i], \n",
    "         reconstructed_src_mu=hist[\"reconstructed_src_mu\"][i],\n",
    "         snr=snr,\n",
    "         im_source=im_source,\n",
    "         image_input=image_input[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dba1619b-62e5-4ce1-82d2-f8407e365f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save figure data\n",
    "# Made with lens config 1\n",
    "\n",
    "i = -1\n",
    "np.savez(\"../data/fig1_data_ngc1300.npz\",\n",
    "         reconstructed_lens=hist[\"reconstructed_lens\"][i], \n",
    "         reconstructed_src_mu=hist[\"reconstructed_src_mu\"][i],\n",
    "         snr=snr,\n",
    "         im_source=im_source,\n",
    "         image_input=image_input[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c369915-f95c-41eb-bbba-47b30cfb1fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save figure data\n",
    "# Made with lens config 0\n",
    "\n",
    "i = -1\n",
    "np.savez(\"../data/fig1_data_ugc6093.npz\",\n",
    "         reconstructed_lens=hist[\"reconstructed_lens\"][i], \n",
    "         reconstructed_src_mu=hist[\"reconstructed_src_mu\"][i],\n",
    "         snr=snr,\n",
    "         im_source=im_source,\n",
    "         image_input=image_input[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1709a76-5fc1-46d0-a019-4f235d8f380d",
   "metadata": {},
   "source": [
    "## Exploring variations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e08fa3a-a061-4abc-ae2c-6eb115bbec82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig, ax = plt.subplots(2, 4, figsize=(9.8 * 1.3, 5.8 * 1.3))\n",
    "\n",
    "source_coord_limit = 1.\n",
    "snr = np.load(\"../data/fig1_data_rff.npz\")['snr']\n",
    "\n",
    "im_list_src = []\n",
    "im_list_lens = []\n",
    "labels = ['_no_pos', '', '_rff']\n",
    "titles = [\"No mapping\", \"Pos. enc. $L_\\mathrm{max} = 5$\", \"RFF $B = 5$\"]\n",
    "\n",
    "im_list_src.append(ax[0, 0].imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i_src - 1]))['im_source'] * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "ax[0,0].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "ax[0,0].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "ax[0,0].set_title(r\"\\bf{Ground truth}\", y=1.02)\n",
    "\n",
    "im_list_lens.append(ax[1,0].imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i_src - 1]))['image_input'].T / std_noise, cmap=cmap, extent=[*x_lims, *y_lims]))\n",
    "# ax[1,0].set_title(r\"\\bf{Ground truth}\")\n",
    "\n",
    "for i in range(1, 4):\n",
    "    im_list_src.append(ax[0,i].imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i - 1]))['reconstructed_src_mu'] * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    ax[0,i].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[0,i].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    ax[0,i].set_title(r\"\\bf{{{}}}\".format(titles[i - 1]), y=1.02)\n",
    "    \n",
    "for i in range(1, 4):\n",
    "    im_list_lens.append(ax[1,i].imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i - 1]))['reconstructed_lens'] / std_noise, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    # ax[1,i].set_title(r\"\\bf{{{}}}\".format(titles[i - 1]))\n",
    "\n",
    "for i_lens in range(4):\n",
    "    ax[1, i_lens].plot([-2.2, -1.2], [-2.2, -2.2], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[1, i_lens].text(-1.7, -1.85, \"$1''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "for i_src in range(4):\n",
    "    ax[0, i_src].plot([-0.85, -0.35], [-0.85, -0.85], c=\"white\", lw=1.5, ls=\"-\")\n",
    "    ax[0, i_src].text(-0.6, -0.75, \"$0.5''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "            \n",
    "    # Inset\n",
    "    axins = ax[0, i_src].inset_axes([0.6, 0.6, 0.37, 0.37])\n",
    "    \n",
    "    if i_src == 0:\n",
    "        axins.imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i_src - 1]))['im_source'] * ~snr / std_noise * lens_sim.conversion_factor, extent=[*x_lims, *y_lims], cmap=cmap, vmin=0)\n",
    "    else:\n",
    "        axins.imshow(np.load(\"../data/fig1_data{}.npz\".format(labels[i_src - 1]))['reconstructed_src_mu'] * ~snr / std_noise * lens_sim.conversion_factor, extent=[*x_lims, *y_lims], cmap=cmap, vmin=0)\n",
    "\n",
    "    x1, x2, y1, y2 = -0.4, -0.1, -0.4, -0.1\n",
    "    axins.set_xlim(x1, x2)\n",
    "    axins.set_ylim(y1, y2)\n",
    "    axins.set_xticklabels([])\n",
    "    axins.set_yticklabels([])\n",
    "\n",
    "    axins.xaxis.set_visible(False)\n",
    "    axins.yaxis.set_visible(False)\n",
    "\n",
    "    for side in ['bottom', 'top', 'right', 'left']:\n",
    "        axins.spines[side].set_color('k')\n",
    "\n",
    "    ax[0, i_src].indicate_inset_zoom(axins, edgecolor=\"k\", alpha=0.75)\n",
    "    [x.set_linewidth(1.) for x in axins.spines.values()]\n",
    "\n",
    "\n",
    "labels = [\"Image / noise\", \"Mean / noise\", \"Mean / noise\", \"Mean / noise\"]\n",
    "for i in range(4):\n",
    "    for j in range(1, 2):\n",
    "        divider = make_axes_locatable(ax[j, i])\n",
    "        cax = divider.append_axes('bottom', size='5%', pad=0.2)\n",
    "        if j == 0:\n",
    "            cbar = fig.colorbar(im_list_src[i], cax=cax, orientation=\"horizontal\", label=labels[i])\n",
    "        else:\n",
    "            cbar = fig.colorbar(im_list_lens[i], cax=cax, orientation=\"horizontal\", label=labels[i])\n",
    "\n",
    "[axx.xaxis.set_visible(False) for axx in ax.flatten()];\n",
    "[axx.yaxis.set_visible(False) for axx in ax.flatten()];\n",
    "\n",
    "for side in ['bottom', 'top', 'right', 'left']:\n",
    "    ax[0,0].spines[side].set_color(cols_default[0])\n",
    "    ax[1,0].spines[side].set_color(cols_default[0])\n",
    "\n",
    "ax[0,0].yaxis.set_visible(True)\n",
    "ax[1,0].yaxis.set_visible(True)\n",
    "ax[0,0].yaxis.set_ticks([])\n",
    "ax[1,0].yaxis.set_ticks([])\n",
    "\n",
    "ax[0,0].set_ylabel(r'\\bf{Source}' + '\\n' + r'\\bf{reconstruction}')\n",
    "ax[1,0].set_ylabel(r'\\bf{Lens}' + '\\n' + r'\\bf{reconstruction}')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"figures/variations.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c252f543-d324-4815-aaac-3c8358e60d78",
   "metadata": {},
   "source": [
    "## Other sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c018d67-4e13-41f3-9c47-98d2fae45a49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "fig, ax = plt.subplots(2, 6, figsize=(19, 7.5))\n",
    "\n",
    "i = -1\n",
    "\n",
    "source_coord_limit = 1.\n",
    "\n",
    "im_list = []\n",
    "\n",
    "for ii in range(2):\n",
    "    \n",
    "    if ii == 0:\n",
    "        file = np.load(\"../data/fig1_data_ugc6093.npz\")\n",
    "        vmax_lens = file[\"reconstructed_lens\"].max() / std_noise\n",
    "    elif ii == 1:\n",
    "        file = np.load(\"../data/fig1_data_ngc1300.npz\")\n",
    "\n",
    "    snr = file['snr']\n",
    "    \n",
    "    im_list.append(ax[ii, 0].imshow(file['im_source'] * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    ax[ii, 0].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[ii, 0].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    \n",
    "\n",
    "    im_list.append(ax[ii, 1].imshow((file[\"reconstructed_src_mu\"]) * ~snr / std_noise * lens_sim.conversion_factor, cmap=cmap, vmin=0, extent=[*x_lims, *y_lims]))\n",
    "    ax[ii, 1].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[ii, 1].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    \n",
    "\n",
    "    im_list.append(ax[ii, 2].imshow((file[\"reconstructed_src_mu\"] - file['im_source']) * ~snr * lens_sim.conversion_factor / std_noise, vmin=-4, vmax=4, cmap=cmap_resid, extent=[*x_lims, *y_lims]))\n",
    "    ax[ii, 2].set_xlim(-source_coord_limit, source_coord_limit)\n",
    "    ax[ii, 2].set_ylim(-source_coord_limit, source_coord_limit)\n",
    "    \n",
    "\n",
    "    im_list.append(ax[ii, 3].imshow(file['image_input'].T / std_noise, cmap=cmap, extent=[*x_lims, *y_lims], vmax=vmax_lens))\n",
    "    im_list.append(ax[ii, 4].imshow((file[\"reconstructed_lens\"]) / std_noise, cmap=cmap, extent=[*x_lims, *y_lims], vmax=vmax_lens))\n",
    "    im_list.append(ax[ii, 5].imshow((file[\"reconstructed_lens\"] - file['image_input'].T) / std_noise, cmap=cmap_resid, vmin=-4, vmax=4, extent=[*x_lims, *y_lims]))\n",
    "    \n",
    "    if ii == 0:\n",
    "        ax[ii, 0].set_title(r\"\\bf{True source}\")\n",
    "        ax[ii, 1].set_title(r\"\\bf{Predicted source}\")\n",
    "        ax[ii, 2].set_title(r\"\\bf{Source residuals}\")\n",
    "        ax[ii, 3].set_title(r\"\\bf{True lensed image}\")\n",
    "        ax[ii, 4].set_title(r\"\\bf{Predicted lensed image}\")\n",
    "        ax[ii, 5].set_title(r\"\\bf{Lensed image residuals}\")\n",
    "\n",
    "    for i_lens in [3, 4]:\n",
    "        ax[ii, i_lens].plot([-2.2, -1.2], [-2.2, -2.2], c=\"white\", lw=1.5, ls=\"-\")\n",
    "        ax[ii, i_lens].text(-1.7, -1.85, \"$1''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "    for i_src in [0, 1]:\n",
    "        ax[ii, i_src].plot([-0.85, -0.35], [-0.85, -0.85], c=\"white\", lw=1.5, ls=\"-\")\n",
    "        ax[ii, i_src].text(-0.6, -0.75, \"$0.5''$\", va=\"center\", ha=\"center\", color=\"white\")\n",
    "\n",
    "    labels = [\"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\", \"Image / noise\", \"Mean / noise\", \"(Mean - true) / noise\"]\n",
    "    \n",
    "    if ii == 1:\n",
    "        for i in range(6, 12):\n",
    "            divider = make_axes_locatable(ax[ii, i - 6])\n",
    "            cax = divider.append_axes('bottom', size='5%', pad=0.2)\n",
    "            cbar = fig.colorbar(im_list[i], cax=cax, orientation=\"horizontal\", label=labels[i - 6])\n",
    "\n",
    "\n",
    "# labels = [r'\\emph{(a)}',r'\\emph{(b)}',r'\\emph{(c)}',r'\\emph{(d)}',r'\\emph{(e)}',r'\\emph{(f)}']\n",
    "# colors = ['w','w','k','w','w','k']\n",
    "# for i in range(6):\n",
    "#     ax[i].text(0.1, 0.9, labels[i], \n",
    "#          horizontalalignment='center',\n",
    "#          verticalalignment='center',\n",
    "#          transform = ax[i].transAxes,\n",
    "#          color=colors[i], fontsize=17)\n",
    "\n",
    "[axx.xaxis.set_visible(False) for axx in ax.flatten()];\n",
    "[axx.yaxis.set_visible(False) for axx in ax.flatten()];\n",
    "\n",
    "ax[0,0].yaxis.set_visible(True)\n",
    "ax[1,0].yaxis.set_visible(True)\n",
    "ax[0,0].yaxis.set_ticks([])\n",
    "ax[1,0].yaxis.set_ticks([])\n",
    "\n",
    "ax[0,0].set_ylabel(r'\\bf{NGC1300}')\n",
    "ax[1,0].set_ylabel(r'\\bf{UGC6093}')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"figures/reconstruction_sources.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f3656e-2020-4d9b-87da-f85926df9f6f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "nerf_probabilistic_gigalens.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "380fdaf81a572e76ac4453c0d0c435d6f7877243e2c66ba4adbe8d143938e87f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}