From 61c4c515bb767f74c2d063e99f17f850976a89d3 Mon Sep 17 00:00:00 2001 From: Alexandre Adam Date: Tue, 13 Feb 2024 22:11:37 -0500 Subject: [PATCH] Fixed imports related to issue 154 and updated requirements for dev (#155) * Fixed imports related to issue 154 and updated requirements for dev * style: pre-commit fixes * Updated numba version * Bumped numba version to 0.58, one version down from the problematic 0.59 * style: pre-commit fixes * Updated numpy version in requirements to be more flexible and accomodate dev requirements * Bumbed numpy version causing problems with python 3.11 * style: pre-commit fixes * make explicit imports * add test import * Remove empty test_imports file * Update tutorials with new import structure * fix magnification call in base --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Connor Stone --- README.md | 10 +-- docs/source/intro.md | 10 +-- docs/source/tutorials/BasicIntroduction.ipynb | 10 +-- .../source/tutorials/InvertLensEquation.ipynb | 4 +- docs/source/tutorials/LensZoo.ipynb | 40 +++++------ docs/source/tutorials/MultiplaneDemo.ipynb | 8 +-- docs/source/tutorials/Parameters.ipynb | 10 +-- docs/source/tutorials/Playground.ipynb | 8 +-- docs/source/tutorials/Simulators.ipynb | 16 ++--- docs/source/tutorials/VisualizeCaustics.ipynb | 4 +- src/caustics/__init__.py | 66 +++++++++++++++---- src/caustics/lenses/base.py | 4 +- 12 files changed, 115 insertions(+), 75 deletions(-) diff --git a/README.md b/README.md index 245b8d7b..04a25641 100644 --- a/README.md +++ b/README.md @@ -32,10 +32,10 @@ import matplotlib.pyplot as plt import caustics import torch -cosmology = caustics.cosmology.FlatLambdaCDM() -sie = caustics.lenses.SIE(cosmology=cosmology, name="lens") -src = caustics.light.Sersic(name="source") -lnslt = caustics.light.Sersic(name="lenslight") +cosmology = caustics.FlatLambdaCDM() +sie = caustics.SIE(cosmology=cosmology, name="lens") +src = caustics.Sersic(name="source") +lnslt = caustics.Sersic(name="lenslight") x = torch.tensor([ # z_s z_l x0 y0 q phi b x0 y0 q phi n Re @@ -44,7 +44,7 @@ x = torch.tensor([ 5.0, -0.2, 0.0, 0.8, 0.0, 1., 1.0, 10.0 ]) # fmt: skip -minisim = caustics.sims.Lens_Source( +minisim = caustics.Lens_Source( lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100 ) plt.imshow(minisim(x, quad_level=3), origin="lower") diff --git a/docs/source/intro.md b/docs/source/intro.md index 02b156b2..ba9c6a94 100644 --- a/docs/source/intro.md +++ b/docs/source/intro.md @@ -27,10 +27,10 @@ import matplotlib.pyplot as plt import caustics import torch -cosmology = caustics.cosmology.FlatLambdaCDM() -sie = caustics.lenses.SIE(cosmology=cosmology, name="lens") -src = caustics.light.Sersic(name="source") -lnslt = caustics.light.Sersic(name="lenslight") +cosmology = caustics.FlatLambdaCDM() +sie = caustics.SIE(cosmology=cosmology, name="lens") +src = caustics.Sersic(name="source") +lnslt = caustics.Sersic(name="lenslight") x = torch.tensor([ # z_s z_l x0 y0 q phi b x0 y0 q phi n Re @@ -39,7 +39,7 @@ x = torch.tensor([ 5.0, -0.2, 0.0, 0.8, 0.0, 1., 1.0, 10.0 ]) # fmt: skip -minisim = caustics.sims.Lens_Source( +minisim = caustics.Lens_Source( lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100 ) plt.imshow(minisim(x, quad_level=3), origin="lower") diff --git a/docs/source/tutorials/BasicIntroduction.ipynb b/docs/source/tutorials/BasicIntroduction.ipynb index 2b24631d..1f90ff07 100644 --- a/docs/source/tutorials/BasicIntroduction.ipynb +++ b/docs/source/tutorials/BasicIntroduction.ipynb @@ -44,7 +44,7 @@ "metadata": {}, "outputs": [], "source": [ - "cosmology = caustics.cosmology.FlatLambdaCDM()" + "cosmology = caustics.FlatLambdaCDM()" ] }, { @@ -62,7 +62,7 @@ "metadata": {}, "outputs": [], "source": [ - "sie = caustics.lenses.SIE(cosmology=cosmology, name=\"lens\")" + "sie = caustics.SIE(cosmology=cosmology, name=\"lens\")" ] }, { @@ -80,7 +80,7 @@ "metadata": {}, "outputs": [], "source": [ - "src = caustics.light.Sersic(name=\"source\")" + "src = caustics.Sersic(name=\"source\")" ] }, { @@ -98,7 +98,7 @@ "metadata": {}, "outputs": [], "source": [ - "lnslt = caustics.light.Sersic(name=\"lenslight\")" + "lnslt = caustics.Sersic(name=\"lenslight\")" ] }, { @@ -116,7 +116,7 @@ "metadata": {}, "outputs": [], "source": [ - "minisim = caustics.sims.Lens_Source(\n", + "minisim = caustics.Lens_Source(\n", " lens=sie, source=src, lens_light=lnslt, pixelscale=0.05, pixels_x=100\n", ")" ] diff --git a/docs/source/tutorials/InvertLensEquation.ipynb b/docs/source/tutorials/InvertLensEquation.ipynb index 4a98380f..777ec174 100644 --- a/docs/source/tutorials/InvertLensEquation.ipynb +++ b/docs/source/tutorials/InvertLensEquation.ipynb @@ -42,7 +42,7 @@ "source": [ "# initialization stuff for an SIE lens\n", "\n", - "cosmology = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", "cosmology.to(dtype=torch.float32)\n", "n_pix = 100\n", "res = 0.05\n", @@ -56,7 +56,7 @@ ")\n", "z_l = torch.tensor(0.5, dtype=torch.float32)\n", "z_s = torch.tensor(1.5, dtype=torch.float32)\n", - "lens = caustics.lenses.SIE(\n", + "lens = caustics.SIE(\n", " cosmology=cosmology,\n", " name=\"sie\",\n", " z_l=z_l,\n", diff --git a/docs/source/tutorials/LensZoo.ipynb b/docs/source/tutorials/LensZoo.ipynb index 73e1ec22..8bf0b44b 100644 --- a/docs/source/tutorials/LensZoo.ipynb +++ b/docs/source/tutorials/LensZoo.ipynb @@ -38,10 +38,10 @@ "metadata": {}, "outputs": [], "source": [ - "cosmology = caustics.cosmology.FlatLambdaCDM()\n", + "cosmology = caustics.FlatLambdaCDM()\n", "cosmology.to(dtype=torch.float32)\n", "z_s = torch.tensor(1.0)\n", - "base_sersic = caustics.light.Sersic(\n", + "base_sersic = caustics.Sersic(\n", " x0=0.1,\n", " y0=0.1,\n", " q=0.6,\n", @@ -85,13 +85,13 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.Point(\n", + "lens = caustics.Point(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", " th_ein=1.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -126,13 +126,13 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.SIS(\n", + "lens = caustics.SIS(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", " th_ein=1.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -170,7 +170,7 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.SIE(\n", + "lens = caustics.SIE(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", @@ -178,7 +178,7 @@ " phi=np.pi / 2,\n", " b=1.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -216,7 +216,7 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.EPL(\n", + "lens = caustics.EPL(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", @@ -225,7 +225,7 @@ " b=1.0,\n", " t=0.5,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -265,14 +265,14 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.NFW(\n", + "lens = caustics.NFW(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", " m=1e13,\n", " c=20.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -314,7 +314,7 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.TNFW(\n", + "lens = caustics.TNFW(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", @@ -322,7 +322,7 @@ " scale_radius=1.0,\n", " tau=3.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -364,7 +364,7 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.PseudoJaffe(\n", + "lens = caustics.PseudoJaffe(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", @@ -372,7 +372,7 @@ " core_radius=5e-1,\n", " scale_radius=15.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -410,14 +410,14 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.ExternalShear(\n", + "lens = caustics.ExternalShear(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", " gamma_1=1.0,\n", " gamma_2=-1.0,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", @@ -453,13 +453,13 @@ "metadata": {}, "outputs": [], "source": [ - "lens = caustics.lenses.MassSheet(\n", + "lens = caustics.MassSheet(\n", " cosmology=cosmology,\n", " x0=0.0,\n", " y0=0.0,\n", " surface_density=1.5,\n", ")\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens,\n", " source=base_sersic,\n", " pixelscale=res,\n", diff --git a/docs/source/tutorials/MultiplaneDemo.ipynb b/docs/source/tutorials/MultiplaneDemo.ipynb index 86a519b1..e96aa09a 100644 --- a/docs/source/tutorials/MultiplaneDemo.ipynb +++ b/docs/source/tutorials/MultiplaneDemo.ipynb @@ -40,7 +40,7 @@ "outputs": [], "source": [ "# initialization stuff for lenses\n", - "cosmology = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", "cosmology.to(dtype=torch.float32)\n", "n_pix = 100\n", "res = 0.5\n", @@ -73,7 +73,7 @@ "\n", " for _ in range(N_lenses):\n", " lenses.append(\n", - " caustics.lenses.SIE(\n", + " caustics.SIE(\n", " cosmology=cosmology,\n", " z_l=z_p,\n", " x0=np.random.uniform(-fov / 4.0, fov / 4.0),\n", @@ -85,12 +85,12 @@ " )\n", "\n", " planes.append(\n", - " caustics.lenses.SinglePlane(\n", + " caustics.SinglePlane(\n", " z_l=z_p, cosmology=cosmology, lenses=lenses, name=f\"plane_{p}\"\n", " )\n", " )\n", "\n", - "lens = caustics.lenses.Multiplane(name=\"multiplane\", cosmology=cosmology, lenses=planes)" + "lens = caustics.Multiplane(name=\"multiplane\", cosmology=cosmology, lenses=planes)" ] }, { diff --git a/docs/source/tutorials/Parameters.ipynb b/docs/source/tutorials/Parameters.ipynb index cb10c040..ddf2fcda 100644 --- a/docs/source/tutorials/Parameters.ipynb +++ b/docs/source/tutorials/Parameters.ipynb @@ -49,19 +49,19 @@ "outputs": [], "source": [ "# Flat cosmology with all dynamic parameters\n", - "cosmo = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\", h0=None, Om0=None)\n", + "cosmo = caustics.FlatLambdaCDM(name=\"cosmo\", h0=None, Om0=None)\n", "\n", "# SIE lens with q and b as static parameters\n", - "lens = caustics.lenses.SIE(cosmology=cosmo, q=0.4, b=1.0)\n", + "lens = caustics.SIE(cosmology=cosmo, q=0.4, b=1.0)\n", "\n", "# Sersic with all dynamic parameters except the sersic index, effective radius, and effective brightness\n", - "source = caustics.light.Sersic(name=\"source\", n=2.0, Re=1.0, Ie=1.0)\n", + "source = caustics.Sersic(name=\"source\", n=2.0, Re=1.0, Ie=1.0)\n", "\n", "# Sersic with all dynamic parameters except the x position, position angle, and effective radius\n", - "lens_light = caustics.light.Sersic(name=\"lenslight\", x0=0.0, phi=1.3, Re=1.0)\n", + "lens_light = caustics.Sersic(name=\"lenslight\", x0=0.0, phi=1.3, Re=1.0)\n", "\n", "# A simulator which captures all these parameters into a single DAG\n", - "sim = caustics.sims.Lens_Source(\n", + "sim = caustics.Lens_Source(\n", " lens=lens, source=source, lens_light=lens_light, pixelscale=0.05, pixels_x=100\n", ")" ] diff --git a/docs/source/tutorials/Playground.ipynb b/docs/source/tutorials/Playground.ipynb index 5516e2c3..f4e49419 100644 --- a/docs/source/tutorials/Playground.ipynb +++ b/docs/source/tutorials/Playground.ipynb @@ -49,7 +49,7 @@ ")\n", "z_l = torch.tensor(0.5, dtype=torch.float32)\n", "z_s = torch.tensor(1.5, dtype=torch.float32)\n", - "cosmology = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", "cosmology.to(dtype=torch.float32)" ] }, @@ -64,7 +64,7 @@ "\n", "\n", "def plot_lens_metrics(thx0, thy0, q, phi, b):\n", - " lens = caustics.lenses.SIE(\n", + " lens = caustics.SIE(\n", " cosmology=cosmology,\n", " z_l=z_l,\n", " x0=thx0,\n", @@ -143,7 +143,7 @@ " Re_src,\n", " Ie_src,\n", "):\n", - " lens = caustics.lenses.SIE(\n", + " lens = caustics.SIE(\n", " cosmology,\n", " z_l,\n", " x0=x0_lens,\n", @@ -152,7 +152,7 @@ " phi=phi_lens,\n", " b=b_lens,\n", " )\n", - " source = caustics.light.Sersic(\n", + " source = caustics.Sersic(\n", " x0=x0_src,\n", " y0=y0_src,\n", " q=q_src,\n", diff --git a/docs/source/tutorials/Simulators.ipynb b/docs/source/tutorials/Simulators.ipynb index eb1e6618..1174a4a1 100644 --- a/docs/source/tutorials/Simulators.ipynb +++ b/docs/source/tutorials/Simulators.ipynb @@ -49,7 +49,7 @@ ")\n", "z_l = torch.tensor(0.5, dtype=torch.float32)\n", "z_s = torch.tensor(1.5, dtype=torch.float32)\n", - "cosmology = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", "cosmology.to(dtype=torch.float32)" ] }, @@ -63,7 +63,7 @@ "# demo simulator with sersic source, SIE lens. then sample some examples. demo the model graph\n", "\n", "\n", - "class Simple_Sim(caustics.sims.Simulator):\n", + "class Simple_Sim(caustics.Simulator):\n", " def __init__(\n", " self,\n", " lens,\n", @@ -99,8 +99,8 @@ "metadata": {}, "outputs": [], "source": [ - "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", - "src = caustics.light.Sersic(name=\"src\")\n", + "sie = caustics.SIE(cosmology, name=\"sie\")\n", + "src = caustics.Sersic(name=\"src\")\n", "\n", "sim = Simple_Sim(sie, src, torch.tensor(0.8))\n", "\n", @@ -142,12 +142,12 @@ "metadata": {}, "outputs": [], "source": [ - "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "sie = caustics.SIE(cosmology, name=\"sie\")\n", "hdu = fits.open(\n", " \"https://www.legacysurvey.org/viewer/fits-cutout?ra=36.3684&dec=-25.6389&size=250&layer=ls-dr9&pixscale=0.262&bands=r\"\n", ")\n", "image_data = np.array(hdu[0].data, dtype=np.float64)\n", - "src = caustics.light.Pixelated(\n", + "src = caustics.Pixelated(\n", " name=\"ESO479_G1\", image=torch.tensor(image_data, dtype=torch.float32)\n", ")\n", "\n", @@ -206,14 +206,14 @@ "metadata": {}, "outputs": [], "source": [ - "sief = caustics.lenses.SIE(\n", + "sief = caustics.SIE(\n", " name=\"sie\",\n", " cosmology=cosmology,\n", " z_l=torch.tensor(0.5),\n", " x0=torch.tensor(0.0),\n", " y0=torch.tensor(0.0),\n", ")\n", - "srcf = caustics.light.Sersic(name=\"src\", n=torch.tensor(2.0))\n", + "srcf = caustics.Sersic(name=\"src\", n=torch.tensor(2.0))\n", "\n", "simf = Simple_Sim(sief, srcf, z_s=torch.tensor(0.8))\n", "\n", diff --git a/docs/source/tutorials/VisualizeCaustics.ipynb b/docs/source/tutorials/VisualizeCaustics.ipynb index 37fee467..51345450 100644 --- a/docs/source/tutorials/VisualizeCaustics.ipynb +++ b/docs/source/tutorials/VisualizeCaustics.ipynb @@ -42,9 +42,9 @@ "source": [ "# initialization stuff for an SIE lens\n", "\n", - "cosmology = caustics.cosmology.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", "cosmology.to(dtype=torch.float32)\n", - "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "sie = caustics.SIE(cosmology, name=\"sie\")\n", "n_pix = 100\n", "res = 0.05\n", "upsample_factor = 2\n", diff --git a/src/caustics/__init__.py b/src/caustics/__init__.py index 6af251f4..e43a8a04 100644 --- a/src/caustics/__init__.py +++ b/src/caustics/__init__.py @@ -1,23 +1,65 @@ from ._version import version as VERSION # noqa -from . import constants, lenses, cosmology, packed, parametrized, light, utils, sims +from .cosmology import ( + Cosmology, + FlatLambdaCDM, + h0_default, + critical_density_0_default, + Om0_default, +) +from .lenses import ( + ThinLens, + ThickLens, + EPL, + ExternalShear, + PixelatedConvergence, + Multiplane, + NFW, + Point, + PseudoJaffe, + SIE, + SIS, + SinglePlane, + MassSheet, + TNFW, +) +from .light import Source, Pixelated, Sersic # PROBESDataset conflicts with .data +from .data import HDF5Dataset, IllustrisKappaDataset, PROBESDataset +from . import utils +from .sims import Lens_Source, Simulator from .tests import test -# from .demo import * - __version__ = VERSION __author__ = "Ciela" __all__ = [ - # Modules - "constants", - "lenses", - "cosmology", - "packed", - "parametrized", - "light", + "Cosmology", + "FlatLambdaCDM", + "h0_default", + "critical_density_0_default", + "Om0_default", + "ThinLens", + "ThickLens", + "EPL", + "ExternalShear", + "PixelatedConvergence", + "Multiplane", + "NFW", + "Point", + "PseudoJaffe", + "SIE", + "SIS", + "SinglePlane", + "MassSheet", + "TNFW", + "Source", + "Pixelated", + "Sersic", + "HDF5Dataset", + "IllustrisKappaDataset", + "PROBESDataset", "utils", - "sims", - # Functions + "Lens_Source", + "Simulator", "test", ] diff --git a/src/caustics/lenses/base.py b/src/caustics/lenses/base.py index 3dd5d4ce..9dced0dd 100644 --- a/src/caustics/lenses/base.py +++ b/src/caustics/lenses/base.py @@ -100,9 +100,7 @@ def magnification( Tensor Gravitational magnification at the given coordinates. """ - return get_magnification( - partial(self.raytrace, params=params), x, y, z_s, **kwargs - ) + return get_magnification(partial(self.raytrace, params=params), x, y, z_s) @unpack def forward_raytrace(