From 7ce5d599cfde5df2ed1e32d780b0027abbc28f69 Mon Sep 17 00:00:00 2001 From: Nolan Smyth Date: Wed, 2 Oct 2024 22:51:21 +0200 Subject: [PATCH] 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 | 1 + docs/source/_toc.yml | 1 + docs/source/tutorials/Microlensing.ipynb | 283 +++++++++++++++++++++++ src/caustics/__init__.py | 2 + src/caustics/light/__init__.py | 3 +- src/caustics/light/func/__init__.py | 3 +- src/caustics/light/func/star_source.py | 16 ++ src/caustics/light/star_source.py | 165 +++++++++++++ tests/test_star_source.py | 41 ++++ 9 files changed, 513 insertions(+), 2 deletions(-) create mode 100644 docs/source/tutorials/Microlensing.ipynb create mode 100644 src/caustics/light/func/star_source.py create mode 100644 src/caustics/light/star_source.py create mode 100644 tests/test_star_source.py diff --git a/.codespell-whitelist b/.codespell-whitelist index 8e435c17..77a8d056 100644 --- a/.codespell-whitelist +++ b/.codespell-whitelist @@ -3,3 +3,4 @@ childs dout din tht +tE diff --git a/docs/source/_toc.yml b/docs/source/_toc.yml index 9ad25f2e..b80e748e 100644 --- a/docs/source/_toc.yml +++ b/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 diff --git a/docs/source/tutorials/Microlensing.ipynb b/docs/source/tutorials/Microlensing.ipynb new file mode 100644 index 00000000..61dae44e --- /dev/null +++ b/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 +} diff --git a/src/caustics/__init__.py b/src/caustics/__init__.py index ef1d231f..9133ed1d 100644 --- a/src/caustics/__init__.py +++ b/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", diff --git a/src/caustics/light/__init__.py b/src/caustics/light/__init__.py index 80d1932c..f3fc302a 100644 --- a/src/caustics/light/__init__.py +++ b/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"] diff --git a/src/caustics/light/func/__init__.py b/src/caustics/light/func/__init__.py index fc3bf6fc..7db23798 100644 --- a/src/caustics/light/func/__init__.py +++ b/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") diff --git a/src/caustics/light/func/star_source.py b/src/caustics/light/func/star_source.py new file mode 100644 index 00000000..826931c8 --- /dev/null +++ b/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 diff --git a/src/caustics/light/star_source.py b/src/caustics/light/star_source.py new file mode 100644 index 00000000..7dd2f3c1 --- /dev/null +++ b/src/caustics/light/star_source.py @@ -0,0 +1,165 @@ +# mypy: disable-error-code="operator,union-attr" +from typing import Optional, Union, Annotated + +from torch import Tensor + +from .base import Source, NameType +from ..parametrized import unpack +from ..packed import Packed +from . import func + +__all__ = ("StarSource",) + + +class StarSource(Source): + """ + `Star` is a subclass of the abstract class `Source`. + It represents a star source in a gravitational lensing system. + + The Star profile is meant to describe individual light sources. + + Attributes + ----------- + x0: Optional[Tensor] + The x-coordinate of the Star source's center. + + *Unit: arcsec* + + y0: Optional[Tensor] + The y-coordinate of the Star source's center. + + *Unit: arcsec* + + theta_s: Optional[Tensor] + The radius of the star. + + *Unit: arcsec* + + Ie: Optional[Tensor] + The intensity at the center of the star. + + *Unit: flux* + + gamma: Optional[Tensor] + The linear limb darkening coefficient. + + *Unit: unitless* + + """ + + def __init__( + self, + x0: Annotated[ + Optional[Union[Tensor, float]], + "The x-coordinate of the star source's center", + True, + ] = None, + y0: Annotated[ + Optional[Union[Tensor, float]], + "The y-coordinate of the star source's center", + True, + ] = None, + theta_s: Annotated[ + Optional[Union[Tensor, float]], + "The radius of the star source", + True, + ] = None, + Ie: Annotated[ + Optional[Union[Tensor, float]], + "The intensity at the effective radius", + True, + ] = None, + gamma: Annotated[ + Optional[Union[Tensor, float]], + "The linear limb darkening coefficient", + True, + ] = None, + name: NameType = None, + ): + """ + Constructs the `Star` object with the given parameters. + + Parameters + ---------- + name: str + The name of the source. + + x0: Optional[Tensor] + The x-coordinate of the star source's center. + + *Unit: arcsec* + + y0: Optional[Tensor] + The y-coordinate of the star source's center. + + *Unit: arcsec* + + theta_s: Optional[Tensor] + The radius of the star. + + *Unit: arcsec* + + Ie: Optional[Tensor] + The intensity at the center of the source. + + *Unit: flux* + + gamma: Optional[Tensor] + The linear limb darkening coefficient. + + *Unit: unitless* + + """ + super().__init__(name=name) + self.add_param("x0", x0) + self.add_param("y0", y0) + self.add_param("theta_s", theta_s) + self.add_param("Ie", Ie) + self.add_param("gamma", gamma) + + @unpack + def brightness( + self, + x, + y, + *args, + params: Optional["Packed"] = None, + x0: Optional[Tensor] = None, + y0: Optional[Tensor] = None, + theta_s: Optional[Tensor] = None, + Ie: Optional[Tensor] = None, + gamma: Optional[Tensor] = None, + **kwargs, + ): + """ + Implements the `brightness` method for `star`. This method calculates the + brightness of the source at the given point(s). + + Parameters + ---------- + x: Tensor + The x-coordinate(s) at which to calculate the source brightness. + This could be a single value or a tensor of values. + + *Unit: arcsec* + + y: Tensor + The y-coordinate(s) at which to calculate the source brightness. + This could be a single value or a tensor of values. + + *Unit: arcsec* + + params: Packed, optional + Dynamic parameter container. + + Returns + ------- + Tensor + The brightness of the source at the given point(s). + The output tensor has the same shape as `x` and `y`. + + *Unit: flux* + + """ + + return func.brightness_star(x0, y0, theta_s, Ie, x, y, gamma) diff --git a/tests/test_star_source.py b/tests/test_star_source.py new file mode 100644 index 00000000..254b734e --- /dev/null +++ b/tests/test_star_source.py @@ -0,0 +1,41 @@ +import caustics +import torch +from caustics.utils import meshgrid + + +def test_star_source(): + src = caustics.StarSource(name="source") + + rho = 1.5 + x = torch.tensor( + [ + 0.0, # x0 + 0.0, # y0 + 1 * rho, # theta_s + 5.0, # Ie, + 0.0, # gamma + ] + ) + theta_x = torch.tensor([0.1]) + theta_y = torch.tensor([0.2]) + # Check if brightness is constant inside source with no limb darkening + assert src.brightness(theta_x, theta_y, x) == 5.0 + # Check if brightness is zero outside source + assert src.brightness(1.01 * rho, theta_y, x) == 0.0 + # Check if function behaves as expected + assert src.brightness(theta_x, theta_y, x) == caustics.light.func.brightness_star( + 0.0, 0.0, 1.5, 5.0, theta_x, theta_y, 0.0 + ) + + # Check if brightness is finite everywhere over a grid + n_pix = 10 + res = 0.05 + upsample_factor = 2 + thx, thy = meshgrid( + res / upsample_factor, + upsample_factor * n_pix, + upsample_factor * n_pix, + dtype=torch.float32, + ) + + assert torch.all(torch.isfinite(src.brightness(thx, thy, x)))