Microlensing (#259)

* init

* Minimal microlensing example

* Adding point source

* Add both images and magnification

* Fun animation

* fitting working

* mcmc working

* Add limb darkening and merge image+mag

* added limb darkening

* Add limb darkening functionality

* Improving microlensing fit example

* fixing pre-commit

* Adding PointSource to __all__

* adding tE to codespell ignore

* Fix bug and improve microlensing fit example

* style: pre-commit fixes

* Noting issue with PyLIMA

* strip nbs

* Fixed bug in u0 units

* Adjusting mcmc

* changed pointsource to starsource, fixing documentation

* style: pre-commit fixes

* Update Microlensing/simulator tutorial and toc. Remove fitting

* clean

* Update toc without conflict

* adding unit tests

* style: pre-commit fixes

* fixed precommit

* Fix weird capitalization bug pt1

* Fix weird capitalization bug pt2

* reorder tutorials

* remove excess files, clean up notebook

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.codespell-whitelist

@@ -3,3 +3,4 @@ childs
 dout
 din
 tht
+tE

docs/source/_toc.yml

@@ -16,6 +16,7 @@ chapters:
           - file: tutorials/InterfaceIntroduction_oop
           - file: tutorials/InterfaceIntroduction_func
       - file: tutorials/LensZoo
+      - file: tutorials/Microlensing
       - file: tutorials/VisualizeCaustics
       - file: tutorials/MultiplaneDemo
       - file: tutorials/InvertLensEquation

docs/source/tutorials/Microlensing.ipynb

@@ -0,0 +1,283 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "%matplotlib inline\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.animation as animation\n",
+    "import torch\n",
+    "import caustics\n",
+    "from IPython.display import HTML"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Building a Microlensing Simulator\n",
+    "\n",
+    "This documentation provides an overview of how to build your own `Simulator`. In this tutorial, we will use the ray tracing tools within the `caustics` library applied to [microlensing](https://www.microlensing-source.org/concept/). We will cover:\n",
+    "\n",
+    "1.\tCreating a custom microlensing `Simulator`\n",
+    "2.  Setting up the parameters for our custom sim\n",
+    "3.  Running batched, dynamic simulations and visualizing the results\n",
+    "\n",
+    "By the end of this guide, you will hopefully feel empowered to create your own custom `Simulator` for your use case!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setting up the `Lens`, `Source`, and `Cosmology` environment\n",
+    "\n",
+    "We will be using a Point lens model and a StarSource light source model provided by `caustics`. We also must define a cosmology environment, which, `caustics` uses to calculate cosmological evolution at large distances. Here, we define a `FlatLambdaCDM` cosmology. But since we only typically care about redshift $z\\approx0$ for microlensing, the cosmology will not actually matter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cosmology = caustics.FlatLambdaCDM()\n",
+    "point_lens = caustics.Point(cosmology=cosmology, name=\"lens\")\n",
+    "src = caustics.StarSource(name=\"source\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the pixel grid for imaging\n",
+    "n_pix = 100\n",
+    "res = 0.05\n",
+    "upsample_factor = 10\n",
+    "theta_x, theta_y = caustics.utils.meshgrid(\n",
+    "    res / upsample_factor,\n",
+    "    upsample_factor * n_pix,\n",
+    "    dtype=torch.float32,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Building our `Simulator`\n",
+    "\n",
+    "Here we define a Microlens class that extends caustics.Simulator. Generically, the child class is where you define all the methods necessary to run your simulation, which is done in the `forward` method. This allows us to run the simulator by passing a single argument, the `params` dictionary, and the `caustics` backend handles the rest.\n",
+    "\n",
+    "In our case, we use our simulator to calculate the magnification of a source at a given position, as well as the lensed image(s) (by definition, multiple images are unresolvable in microlenisng). "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Microlens(caustics.Simulator):\n",
+    "\n",
+    "    theta_x = theta_x\n",
+    "    theta_y = theta_y\n",
+    "    z_s = 0.0\n",
+    "\n",
+    "    def __init__(self, lens, src, z_s=None, name: str = \"sim\"):\n",
+    "        super().__init__(name)\n",
+    "        self.lens = lens\n",
+    "        self.src = src\n",
+    "\n",
+    "    def forward(self, params):\n",
+    "        # Compute the observed positions of the source\n",
+    "        beta_x, beta_y = self.lens.raytrace(\n",
+    "            self.theta_x, self.theta_y, self.z_s, params\n",
+    "        )\n",
+    "        # Compute the brightness of the source at the observed positions (the image)\n",
+    "        brightness = self.src.brightness(beta_x, beta_y, params)\n",
+    "        # Compute the baseline (unlensed) brightness of the source\n",
+    "        baseline_brightness = self.src.brightness(theta_x, theta_y, params)\n",
+    "        # Return the lensed image [n_pix x n_pix], and magnification\n",
+    "        return brightness, brightness.mean() / baseline_brightness.mean()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Microlensing Physical Parameters\n",
+    "th_ein = 0.3  # Einstein radius in arcsec\n",
+    "\n",
+    "# Microlensing Model Parameters\n",
+    "t0 = 0.2  # time at peak magnification\n",
+    "u0 = 0.5  # minimum impact parameter (u(t=t_0)) in units of th_ein\n",
+    "tE = 0.05  # Einstein timescale\n",
+    "rho = 2.5  # source size in units of lens Einstein radii\n",
+    "\n",
+    "gamma = 0.6  # linear limb darkening coefficient"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can visualize the structure of our simulator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim = Microlens(point_lens, src)\n",
+    "\n",
+    "# Set all variables as static variables since they are not going to change in this simulation\n",
+    "sim.lens.z_l = 0.0\n",
+    "sim.lens.y0 = -u0 * th_ein\n",
+    "sim.lens.th_ein = th_ein\n",
+    "sim.src.x0 = 0.0\n",
+    "sim.src.y0 = 0.0\n",
+    "sim.src.theta_s = th_ein * rho\n",
+    "sim.src.Ie = 5.0\n",
+    "sim.src.gamma = gamma\n",
+    "\n",
+    "sim.graph(True, True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can set up the parameters for our simulation. We can look at the order of the parameters expected by our simulator by looking at `x_order`. Here, the lens x position is the only dynamic variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim.x_order"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "B = 64  # Batch size\n",
+    "ts = torch.linspace(-6 * tE + t0, 6 * tE + t0, B).view(-1, 1)  # Create time grid\n",
+    "\n",
+    "# Calculate source position at each time in arcsec\n",
+    "x0s = (ts - t0) / (tE) * th_ein  # Shape is [B, 1]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Running the simulation and visualzing the output\n",
+    "\n",
+    "Running a batched simulation over different parameters (in our case, these correspond to the lens motion in discrete time-steps by construction) is now as easy as using `vmap` on our simulator. We can then visualize the total output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "images, magnifications = torch.vmap(sim)(x0s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# Create animation\n",
+    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 8))\n",
+    "\n",
+    "# Display the first frame of the image in the first subplot\n",
+    "img = ax1.imshow(images[0].numpy(), cmap=\"cividis\", interpolation=\"bilinear\")\n",
+    "ax1.set_title(\"Image\")\n",
+    "\n",
+    "# Set up the scatter plot for magnifications in the second subplot\n",
+    "scatter = ax2.scatter(ts[0].item(), magnifications[0].item())\n",
+    "ax2.set_xlim(ts.min().item(), ts.max().item())\n",
+    "ax2.set_ylim(magnifications.min().item() * 0.9, magnifications.max().item() * 1.1)\n",
+    "ax2.axvline(-tE / 2, color=\"r\", linestyle=\"--\")\n",
+    "ax2.axvline(tE / 2, color=\"r\", linestyle=\"--\")\n",
+    "ax2.set_xlabel(\"t\")\n",
+    "ax2.set_ylabel(\"A\")\n",
+    "\n",
+    "\n",
+    "def update(frame):\n",
+    "    \"\"\"Update function for the animation.\"\"\"\n",
+    "    # Update the image in the first subplot\n",
+    "    img.set_array(images[frame].numpy())\n",
+    "\n",
+    "    # Update the scatter plot in the second subplot\n",
+    "    ax2.clear()  # Clear the previous frame\n",
+    "    ax2.scatter(ts[: frame + 1].numpy(), magnifications[: frame + 1].numpy())\n",
+    "    ax2.set_xlim(ts.min().item(), ts.max().item())\n",
+    "    ax2.set_ylim(magnifications.min().item() * 0.9, magnifications.max().item() * 1.1)\n",
+    "    ax2.set_xlabel(\"t\")\n",
+    "    ax2.set_ylabel(\"A\")\n",
+    "    ax2.set_title(\"Light-Curve\")\n",
+    "\n",
+    "    return img, scatter\n",
+    "\n",
+    "\n",
+    "ani = animation.FuncAnimation(fig, update, frames=B, interval=60)\n",
+    "\n",
+    "plt.close()\n",
+    "\n",
+    "# Save animation as gif\n",
+    "# ani.save(\"microlensing_animation.gif\", writer='pillow', fps=16)  # Adjust 'fps' for the speed\n",
+    "\n",
+    "# Or display the animation inline\n",
+    "HTML(ani.to_jshtml())"
+   ]
+  }
+ ],
+ "metadata": {
+  "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.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/caustics/__init__.py

@@ -30,6 +30,7 @@
     Pixelated,
     PixelatedTime,
     Sersic,
+    StarSource,
 )
 from . import utils
 from .sims import LensSource, Microlens, Simulator
@@ -66,6 +67,7 @@
     "Pixelated",
     "PixelatedTime",
     "Sersic",
+    "StarSource",
     "utils",
     "LensSource",
     "Microlens",

src/caustics/light/__init__.py

@@ -2,5 +2,6 @@
 from .pixelated import Pixelated
 from .sersic import Sersic
 from .pixelated_time import PixelatedTime
+from .star_source import StarSource
 
-__all__ = ["Source", "Pixelated", "PixelatedTime", "Sersic"]
+__all__ = ["Source", "Pixelated", "PixelatedTime", "Sersic", "StarSource"]

src/caustics/light/func/__init__.py

@@ -1,3 +1,4 @@
 from .sersic import brightness_sersic, k_lenstronomy, k_sersic
+from .star_source import brightness_star
 
-__all__ = ("brightness_sersic", "k_lenstronomy", "k_sersic")
+__all__ = ("brightness_sersic", "k_lenstronomy", "k_sersic", "brightness_star")

src/caustics/light/func/star_source.py

@@ -0,0 +1,16 @@
+from torch import where
+
+from ...utils import translate_rotate
+
+
+def brightness_star(x0, y0, theta_s, Ie, x, y, gamma=0.0):
+
+    x, y = translate_rotate(x, y, x0, y0)
+
+    # Calculate the radial distance from the center
+    impact_parameter = (x**2 + y**2) ** (1 / 2)
+    # linear limb darkening
+    mu = (1 - impact_parameter**2) ** (1 / 2)
+    intensity = where(impact_parameter <= theta_s, Ie * (1 - gamma * (1 - mu)), 0)
+
+    return intensity