diff --git a/.codespell-whitelist b/.codespell-whitelist new file mode 100644 index 00000000..823aed73 --- /dev/null +++ b/.codespell-whitelist @@ -0,0 +1,4 @@ +sie +childs +dout +din diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7cc05c4a..02f18f77 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,28 +1,26 @@ exclude: | (?x)^( - tests/utils.py | - .devcontainer/ + tests/utils.py ) ci: autoupdate_commit_msg: "chore: update pre-commit hooks" autofix_commit_msg: "style: pre-commit fixes" - autofix_prs: false repos: - repo: https://github.com/psf/black - rev: "23.7.0" + rev: "23.11.0" hooks: - id: black-jupyter - repo: https://github.com/asottile/blacken-docs - rev: "1.15.0" + rev: "1.16.0" hooks: - id: blacken-docs additional_dependencies: [black==23.7.0] - repo: https://github.com/pre-commit/pre-commit-hooks - rev: "v4.4.0" + rev: "v4.5.0" hooks: - id: check-added-large-files - id: check-case-conflict @@ -45,36 +43,33 @@ repos: - id: rst-inline-touching-normal - repo: https://github.com/pre-commit/mirrors-prettier - rev: "v3.0.0" + rev: "v4.0.0-alpha.4" hooks: - id: prettier types_or: [yaml, markdown, html, css, scss, javascript, json] args: [--prose-wrap=always] - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.0.277" + rev: "v0.1.7" hooks: - id: ruff args: ["--fix", "--show-fixes"] - - repo: https://github.com/pre-commit/mirrors-mypy - rev: "v1.4.1" - hooks: - - id: mypy - files: src|tests - args: [] - additional_dependencies: - - pytest + # 2023-12-11: Not use mypy for now. + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: "v1.7.1" + # hooks: + # - id: mypy + # files: src|tests + # args: [] + # additional_dependencies: + # - pytest - repo: https://github.com/codespell-project/codespell - rev: "v2.2.5" + rev: "v2.2.6" hooks: - id: codespell - - - repo: https://github.com/shellcheck-py/shellcheck-py - rev: "v0.9.0.5" - hooks: - - id: shellcheck + args: ["--write-changes", "--ignore-words", ".codespell-whitelist"] - repo: https://github.com/kynan/nbstripout rev: 0.6.1 diff --git a/docs/requirements.txt b/docs/requirements.txt index 3aee37f7..431bdaf2 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,5 +1,5 @@ +ipywidgets jupyter-book +matplotlib sphinx sphinx_rtd_theme -matplotlib -ipywidgets diff --git a/docs/source/_config.yml b/docs/source/_config.yml index 2377f227..d164b662 100644 --- a/docs/source/_config.yml +++ b/docs/source/_config.yml @@ -46,7 +46,7 @@ sphinx: - "sphinx.ext.ifconfig" - "sphinx.ext.viewcode" config: - html_theme_options: + html_theme_options: logo: image_light: ../../media/caustics_logo_white.png - image_dark: ../../media/caustics_logo.png \ No newline at end of file + image_dark: ../../media/caustics_logo.png diff --git a/docs/source/tutorials/BasicIntroduction.ipynb b/docs/source/tutorials/BasicIntroduction.ipynb index 67c288a3..715508a4 100644 --- a/docs/source/tutorials/BasicIntroduction.ipynb +++ b/docs/source/tutorials/BasicIntroduction.ipynb @@ -56,7 +56,7 @@ "source": [ "## Simulating an SIE lens\n", "\n", - "Here we will demo the very basics of lensing with a classic `SIE` lens model. We will see what it takes to make an `SIE` model, lens a backgorund `Sersic` source, and sample some examples in a simulator. Caustic simulators can generalize to very complex scenarios. In these cases there can be a lot of parameters moving through the simulator, and the order/number of parameters may change depending on what lens or source is being used. To streamline this process, caustic impliments a class called `Parametrized` which has some knowledge of the parameters moving through it, this way it can keep track of everything for you. For this to work, you must put the parameters into a `Packed` object which it can recognize, each sub function can then unpack the parameters it needs. Below we will show some examples of what this looks like." + "Here we will demo the very basics of lensing with a classic `SIE` lens model. We will see what it takes to make an `SIE` model, lens a background `Sersic` source, and sample some examples in a simulator. Caustic simulators can generalize to very complex scenarios. In these cases there can be a lot of parameters moving through the simulator, and the order/number of parameters may change depending on what lens or source is being used. To streamline this process, caustic implements a class called `Parametrized` which has some knowledge of the parameters moving through it, this way it can keep track of everything for you. For this to work, you must put the parameters into a `Packed` object which it can recognize, each sub function can then unpack the parameters it needs. Below we will show some examples of what this looks like." ] }, { @@ -155,7 +155,7 @@ "source": [ "## Where to go next?\n", "\n", - "The caustic tutorials are generally short and to the point, that way you can idenfity what you want and jump right to some useful code that demo's the particular problem you face. Below is a list of caustic tutorials and a quick description of what you will learn in each one::\n", + "The caustic tutorials are generally short and to the point, that way you can identify what you want and jump right to some useful code that demo's the particular problem you face. Below is a list of caustic tutorials and a quick description of what you will learn in each one::\n", "\n", "- `LensZoo`: here you can see all the built-in lens mass distributions in `caustic` and how they distort the same background Seric source.\n", "- `Playground`: here we demo the main visualizations of a lensing system (deflection angles, convergence, potential, time delay, magnification) in an interactive display so you can change the parameters by hand and see how the visuals change!\n", diff --git a/docs/source/tutorials/InvertLensEquation.ipynb b/docs/source/tutorials/InvertLensEquation.ipynb index db75ab0c..3e18b05f 100644 --- a/docs/source/tutorials/InvertLensEquation.ipynb +++ b/docs/source/tutorials/InvertLensEquation.ipynb @@ -102,7 +102,7 @@ "paths = CS.collections[0].get_paths()\n", "caustic_paths = []\n", "for path in paths:\n", - " # Collect the path into a descrete set of points\n", + " # Collect the path into a discrete set of points\n", " vertices = path.interpolated(5).vertices\n", " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", diff --git a/docs/source/tutorials/MultiplaneDemo.ipynb b/docs/source/tutorials/MultiplaneDemo.ipynb index e6c12d53..185c3eef 100644 --- a/docs/source/tutorials/MultiplaneDemo.ipynb +++ b/docs/source/tutorials/MultiplaneDemo.ipynb @@ -7,7 +7,7 @@ "source": [ "# Multiplane Lensing\n", "\n", - "The universe is three dimensional and filled with stuff. A light ray traveling to our telescope may encounter more than a single massive object on its way to our telescopes. This is handled by a multiplane lensing framework. Multiplane lensing involves tracing the path of a ray backwards from our telescope through each individual plane (which is treated similarly to typical single plane lensing, though extra factors account for the ray physically moving in 3D space) getting perturbed at each step until it finally lands on the source we'd like to image. For more mathmatical details see [Petkova et al. 2014](https://doi.org/10.1093/mnras/stu1860) for the formalism we use internally.\n", + "The universe is three dimensional and filled with stuff. A light ray traveling to our telescope may encounter more than a single massive object on its way to our telescopes. This is handled by a multiplane lensing framework. Multiplane lensing involves tracing the path of a ray backwards from our telescope through each individual plane (which is treated similarly to typical single plane lensing, though extra factors account for the ray physically moving in 3D space) getting perturbed at each step until it finally lands on the source we'd like to image. For more mathematical details see [Petkova et al. 2014](https://doi.org/10.1093/mnras/stu1860) for the formalism we use internally.\n", "\n", "The main concept to keep in mind is that a lot of quantities we are used to working with, such as \"reduced deflection angles\" don't really exist in multiplane lensing since these are normalized by the redshift of the source and lens, however there is no single \"lens redshift\" for multiplane! Instead we define everything with respect to results from full raytracing, once the raytracing is done we can define effective quantities (like effective reduced deflection angle) which behave similarly in intuition but are not quite the same in detail." ] @@ -169,7 +169,7 @@ "paths = CS.collections[0].get_paths()\n", "\n", "for path in paths:\n", - " # Collect the path into a descrete set of points\n", + " # Collect the path into a discrete set of points\n", " vertices = path.interpolated(5).vertices\n", " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", diff --git a/docs/source/tutorials/VisualizeCaustics.ipynb b/docs/source/tutorials/VisualizeCaustics.ipynb index 8e163d5c..bcc0758e 100644 --- a/docs/source/tutorials/VisualizeCaustics.ipynb +++ b/docs/source/tutorials/VisualizeCaustics.ipynb @@ -7,7 +7,7 @@ "source": [ "# Visualize Caustics\n", "\n", - "Here we will demonstrate how to collect caustic lines using caustic! Since caustic (the code) uses autodiff and can get exact derivatives, it is actually very acurate at computing caustics. \n", + "Here we will demonstrate how to collect caustic lines using caustic! Since caustic (the code) uses autodiff and can get exact derivatives, it is actually very accurate at computing caustics. \n", "\n", "Conceptually a caustic occurs where the magnification of a lens diverges to infinity. A convenient way to measure the magnification in the image plane is by taking the determinant ($det$) of the jacobian of the lens equation ($A$), its reciprocal is the magnification. This means that anywhere that $det(A) = 0$ is a critical line in the image plane (magnification goes to infinity). If we take this line and raytrace it back to the source plane we can see the caustics which define boundaries for lensing phenomena." ] @@ -125,7 +125,7 @@ "paths = CS.collections[0].get_paths()\n", "\n", "for path in paths:\n", - " # Collect the path into a descrete set of points\n", + " # Collect the path into a discrete set of points\n", " vertices = path.interpolated(5).vertices\n", " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", diff --git a/docs/source/tutorials/index.md b/docs/source/tutorials/index.md index 75d2e53c..53484277 100644 --- a/docs/source/tutorials/index.md +++ b/docs/source/tutorials/index.md @@ -1,5 +1,5 @@ # Tutorials -Here you will find the jupyter notebook tutorials. -It is recommended that you go through the tutorials yourself, -but for quick reference a version of each tutorial is available here. \ No newline at end of file +Here you will find the jupyter notebook tutorials. It is recommended that you go +through the tutorials yourself, but for quick reference a version of each +tutorial is available here. diff --git a/pyproject.toml b/pyproject.toml index b3216a9c..75379562 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,3 +57,7 @@ version-file = "src/caustics/_version.py" [tool.hatch.version.raw-options] local_scheme = "no-local-version" + +[tool.ruff] +# Same as Black. +line-length = 100 diff --git a/src/caustics/__init__.py b/src/caustics/__init__.py index f46fcda9..6af251f4 100644 --- a/src/caustics/__init__.py +++ b/src/caustics/__init__.py @@ -1,16 +1,23 @@ from ._version import version as VERSION # noqa -from .constants import * -from .lenses import * -from .cosmology import * -from .packed import * -from .parametrized import * -from .light import * -from .utils import * -from .sims import * -from .tests import * +from . import constants, lenses, cosmology, packed, parametrized, light, utils, sims +from .tests import test # from .demo import * __version__ = VERSION __author__ = "Ciela" + +__all__ = [ + # Modules + "constants", + "lenses", + "cosmology", + "packed", + "parametrized", + "light", + "utils", + "sims", + # Functions + "test", +] diff --git a/src/caustics/cosmology.py b/src/caustics/cosmology.py index 9f72de7b..4b2d9b03 100644 --- a/src/caustics/cosmology.py +++ b/src/caustics/cosmology.py @@ -10,6 +10,7 @@ from .utils import interp1d from .constants import G_over_c2, c_Mpc_s, km_to_Mpc from .parametrized import Parametrized, unpack +from .packed import Packed __all__ = ( "h0_default", @@ -276,7 +277,8 @@ def critical_surface_density( class FlatLambdaCDM(Cosmology): """ - Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) cosmology with no radiation. + Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) + cosmology with no radiation. """ def __init__( diff --git a/src/caustics/data/__init__.py b/src/caustics/data/__init__.py index 63c0009c..5702b575 100644 --- a/src/caustics/data/__init__.py +++ b/src/caustics/data/__init__.py @@ -1,3 +1,5 @@ -from .hdf5dataset import * -from .illustris_kappa import * -from .probes import * +from .hdf5dataset import HDF5Dataset +from .illustris_kappa import IllustrisKappaDataset +from .probes import PROBESDataset + +__all__ = ["HDF5Dataset", "IllustrisKappaDataset", "PROBESDataset"] diff --git a/src/caustics/lenses/__init__.py b/src/caustics/lenses/__init__.py index c2d6bbe1..49902a79 100644 --- a/src/caustics/lenses/__init__.py +++ b/src/caustics/lenses/__init__.py @@ -1,13 +1,31 @@ -from .base import * -from .epl import * -from .external_shear import * -from .pixelated_convergence import * -from .multiplane import * -from .nfw import * -from .point import * -from .pseudo_jaffe import * -from .sie import * -from .sis import * -from .singleplane import * -from .mass_sheet import * -from .tnfw import * +from .base import ThinLens, ThickLens +from .epl import EPL +from .external_shear import ExternalShear +from .pixelated_convergence import PixelatedConvergence +from .multiplane import Multiplane +from .nfw import NFW +from .point import Point +from .pseudo_jaffe import PseudoJaffe +from .sie import SIE +from .sis import SIS +from .singleplane import SinglePlane +from .mass_sheet import MassSheet +from .tnfw import TNFW + + +__all__ = [ + "ThinLens", + "ThickLens", + "EPL", + "ExternalShear", + "PixelatedConvergence", + "Multiplane", + "NFW", + "Point", + "PseudoJaffe", + "SIE", + "SIS", + "SinglePlane", + "MassSheet", + "TNFW", +] diff --git a/src/caustics/lenses/base.py b/src/caustics/lenses/base.py index c6a26d2c..5cd78223 100644 --- a/src/caustics/lenses/base.py +++ b/src/caustics/lenses/base.py @@ -11,6 +11,7 @@ from ..parametrized import Parametrized, unpack from .utils import get_magnification from ..utils import batch_lm +from ..packed import Packed __all__ = ("ThinLens", "ThickLens") @@ -29,7 +30,8 @@ def __init__(self, cosmology: Cosmology, name: str = None): name: string The name of the lens model. cosmology: Cosmology - An instance of a Cosmology class that describes the cosmological parameters of the model. + An instance of a Cosmology class that describes + the cosmological parametersof the model. """ super().__init__(name) self.cosmology = cosmology @@ -47,7 +49,8 @@ def jacobian_lens_equation( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the lensing equation at specified points. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the lensing equation at specified points. + This equates to a (2,2) matrix at each (x,y) point. method: autograd or fft """ @@ -57,7 +60,9 @@ def jacobian_lens_equation( elif method == "finitediff": if pixelscale is None: raise ValueError( - "Finite differences lensing jacobian requires regular grid and known pixelscale. Please include the pixelscale argument" + "Finite differences lensing jacobian requires regular grid " + "and known pixelscale. " + "Please include the pixelscale argument" ) return self._jacobian_lens_equation_finitediff( x, y, z_s, pixelscale, params, **kwargs @@ -137,7 +142,7 @@ def forward_raytrace( bxy = torch.stack((bx, by)).repeat(n_init, 1) # has shape (n_init, Dout:2) - # TODO make FOV more general so that it doesnt have to be centered on zero,zero + # TODO make FOV more general so that it doesn't have to be centered on zero,zero if fov is None: raise ValueError("fov must be given to generate initial guesses") @@ -147,7 +152,7 @@ def forward_raytrace( ) # Has shape (n_init, Din:2) # Optimize guesses in image plane - x, l, c = batch_lm( + x, l, c = batch_lm( # noqa: E741 Unused `l` variable guesses, bxy, lambda *a, **k: torch.stack( @@ -172,13 +177,16 @@ def forward_raytrace( class ThickLens(Lens): """ - Base class for modeling gravitational lenses that cannot be treated using the thin lens approximation. - It is an abstract class and should be subclassed for different types of lens models. + Base class for modeling gravitational lenses that cannot be + treated using the thin lens approximation. + It is an abstract class and should be subclassed + for different types of lens models. Attributes ---------- cosmology: Cosmology - An instance of a Cosmology class that describes the cosmological parameters of the model. + An instance of a Cosmology class that describes + the cosmological parameters of the model. """ @unpack(3) @@ -192,7 +200,8 @@ def reduced_deflection_angle( **kwargs, ) -> tuple[Tensor, Tensor]: """ - ThickLens objects do not have a reduced deflection angle since the distance D_ls is undefined + ThickLens objects do not have a reduced deflection angle + since the distance D_ls is undefined Parameters ---------- @@ -202,13 +211,21 @@ def reduced_deflection_angle( Tensor of y coordinates in the lens plane. z_s: Tensor Tensor of source redshifts. - params: (Packed, optional) + params: Packed, optional Dynamic parameter container for the lens model. Defaults to None. - Raises:: NotImplementedError + Raises + ------ + NotImplementedError """ warnings.warn( - "ThickLens objects do not have a reduced deflection angle since they have no unique lens redshift. The distance D_{ls} is undefined in the equation $\alpha_{reduced} = \frac{D_{ls}}{D_s}\alpha_{physical}$. See `effective_reduced_deflection_angle`. Now using effective_reduced_deflection_angle, please switch functions to remove this warning" + "ThickLens objects do not have a reduced deflection angle " + "since they have no unique lens redshift. " + "The distance D_{ls} is undefined in the equation " + "$\alpha_{reduced} = \frac{D_{ls}}{D_s}\alpha_{physical}$." + "See `effective_reduced_deflection_angle`. " + "Now using effective_reduced_deflection_angle, " + "please switch functions to remove this warning" ) return self.effective_reduced_deflection_angle(x, y, z_s, params) @@ -273,11 +290,14 @@ def physical_deflection_angle( Returns ------- tuple[Tensor, Tensor] - Tuple of Tensors representing the x and y components of the deflection angle, respectively. + Tuple of Tensors representing the x and y components + of the deflection angle, respectively. """ raise NotImplementedError( - "Physical deflection angles are computed with respect to a lensing plane. ThickLens objects have no unique definition of a lens plane and so cannot compute a physical_deflection_angle" + "Physical deflection angles are computed with respect to a lensing plane. " + "ThickLens objects have no unique definition of a lens plane " + "and so cannot compute a physical_deflection_angle" ) @abstractmethod @@ -308,7 +328,10 @@ def raytrace( Returns ------- - tuple[Tensor, Tensor]: Tuple of Tensors representing the x and y coordinates of the ray-traced light rays, respectively. + x: Tensor + x coordinate Tensor of the ray-traced light rays + y: Tensor + y coordinate Tensor of the ray-traced light rays """ ... @@ -341,7 +364,8 @@ def surface_density( Returns ------- Tensor - The projected mass density at the given coordinates in units of solar masses per square Megaparsec. + The projected mass density at the given coordinates + in units of solar masses per square Megaparsec. """ ... @@ -389,7 +413,8 @@ def _jacobian_effective_deflection_angle_finitediff( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the effective reduced deflection angle vector field. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the effective reduced deflection angle vector field. + This equates to a (2,2) matrix at each (x,y) point. """ # Compute deflection angles ax, ay = self.effective_reduced_deflection_angle(x, y, z_s, params) @@ -411,7 +436,8 @@ def _jacobian_effective_deflection_angle_autograd( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the effective reduced deflection angle vector field. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the effective reduced deflection angle vector field. + This equates to a (2,2) matrix at each (x,y) point. """ # Ensure the x,y coordinates track gradients x = x.detach().requires_grad_() @@ -449,7 +475,8 @@ def jacobian_effective_deflection_angle( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the effective reduced deflection angle vector field. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the effective reduced deflection angle vector field. + This equates to a (2,2) matrix at each (x,y) point. method: autograd or fft """ @@ -459,7 +486,9 @@ def jacobian_effective_deflection_angle( elif method == "finitediff": if pixelscale is None: raise ValueError( - "Finite differences lensing jacobian requires regular grid and known pixelscale. Please include the pixelscale argument" + "Finite differences lensing jacobian requires " + "regular grid and known pixelscale. " + "Please include the pixelscale argument" ) return self._jacobian_effective_deflection_angle_finitediff( x, y, z_s, pixelscale, params @@ -479,7 +508,8 @@ def _jacobian_lens_equation_finitediff( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the lensing equation at specified points. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the lensing equation at specified points. + This equates to a (2,2) matrix at each (x,y) point. """ # Build Jacobian J = self._jacobian_effective_deflection_angle_finitediff( @@ -498,7 +528,8 @@ def _jacobian_lens_equation_autograd( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the lensing equation at specified points. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the lensing equation at specified points. + This equates to a (2,2) matrix at each (x,y) point. """ # Build Jacobian J = self._jacobian_effective_deflection_angle_autograd( @@ -517,7 +548,14 @@ def effective_convergence_div( **kwargs, ) -> Tensor: """ - Using the divergence of the effective reduced delfection angle we can compute the divergence component of the effective convergence field. This field produces a single plane convergence field which reproduces as much of the deflection field as possible for a single plane. See: https://arxiv.org/pdf/2006.07383.pdf see also the `effective_convergence_curl` method. + Using the divergence of the effective reduced delfection angle + we can compute the divergence component of the effective convergence field. + This field produces a single plane convergence field + which reproduces as much of the deflection field + as possible for a single plane. + + See: https://arxiv.org/pdf/2006.07383.pdf + see also the `effective_convergence_curl` method. """ J = self.jacobian_effective_deflection_angle(x, y, z_s, params, **kwargs) return 0.5 * (J[..., 0, 0] + J[..., 1, 1]) @@ -533,7 +571,13 @@ def effective_convergence_curl( **kwargs, ) -> Tensor: """ - Use the curl of the effective reduced deflection angle vector field to compute an effective convergence which derrives specifically from the curl of the deflection field. This field is purely a result of multiplane lensing and cannot occur in single plane lensing. See: https://arxiv.org/pdf/2006.07383.pdf + Use the curl of the effective reduced deflection angle vector field + to compute an effective convergence which derives specifically + from the curl of the deflection field. + This field is purely a result of multiplane lensing + and cannot occur in single plane lensing. + + See: https://arxiv.org/pdf/2006.07383.pdf """ J = self.jacobian_effective_deflection_angle(x, y, z_s, params, **kwargs) return 0.5 * (J[..., 1, 0] - J[..., 0, 1]) @@ -752,7 +796,8 @@ def raytrace( **kwargs, ) -> tuple[Tensor, Tensor]: """ - Perform a ray-tracing operation by subtracting the deflection angles from the input coordinates. + Perform a ray-tracing operation by subtracting + the deflection angles from the input coordinates. Parameters ---------- @@ -785,7 +830,8 @@ def time_delay( **kwargs, ): """ - Compute the gravitational time delay for light passing through the lens at given coordinates. + Compute the gravitational time delay for light passing + through the lens at given coordinates. Parameters ---------- @@ -824,7 +870,8 @@ def _jacobian_deflection_angle_finitediff( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the deflection angle vector. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the deflection angle vector. + This equates to a (2,2) matrix at each (x,y) point. """ # Compute deflection angles ax, ay = self.reduced_deflection_angle(x, y, z_s, params) @@ -846,7 +893,8 @@ def _jacobian_deflection_angle_autograd( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the deflection angle vector. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the deflection angle vector. + This equates to a (2,2) matrix at each (x,y) point. """ # Ensure the x,y coordinates track gradients x = x.detach().requires_grad_() @@ -884,7 +932,8 @@ def jacobian_deflection_angle( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the deflection angle vector. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the deflection angle vector. + This equates to a (2,2) matrix at each (x,y) point. method: autograd or fft """ @@ -894,7 +943,8 @@ def jacobian_deflection_angle( elif method == "finitediff": if pixelscale is None: raise ValueError( - "Finite differences lensing jacobian requires regular grid and known pixelscale. Please include the pixelscale argument" + "Finite differences lensing jacobian requires regular grid " + "and known pixelscale. Please include the pixelscale argument" ) return self._jacobian_deflection_angle_finitediff( x, y, z_s, pixelscale, params @@ -914,7 +964,8 @@ def _jacobian_lens_equation_finitediff( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the lensing equation at specified points. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the lensing equation at specified points. + This equates to a (2,2) matrix at each (x,y) point. """ # Build Jacobian J = self._jacobian_deflection_angle_finitediff( @@ -933,7 +984,8 @@ def _jacobian_lens_equation_autograd( **kwargs, ) -> tuple[tuple[Tensor, Tensor], tuple[Tensor, Tensor]]: """ - Return the jacobian of the lensing equation at specified points. This equates to a (2,2) matrix at each (x,y) point. + Return the jacobian of the lensing equation at specified points. + This equates to a (2,2) matrix at each (x,y) point. """ # Build Jacobian J = self._jacobian_deflection_angle_autograd(x, y, z_s, params, **kwargs) diff --git a/src/caustics/lenses/epl.py b/src/caustics/lenses/epl.py index 378bf97f..a1065b2f 100644 --- a/src/caustics/lenses/epl.py +++ b/src/caustics/lenses/epl.py @@ -7,6 +7,7 @@ from ..utils import derotate, translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("EPL",) @@ -15,8 +16,10 @@ class EPL(ThinLens): """ Elliptical power law (EPL, aka singular power-law ellipsoid) profile. - This class represents a thin gravitational lens model with an elliptical power law profile. The lensing equations are solved - iteratively using an approach based on Tessore et al. 2015. + This class represents a thin gravitational lens model + with an elliptical power law profile. + The lensing equations are solved iteratively + using an approach based on Tessore et al. 2015. Attributes ---------- @@ -28,17 +31,31 @@ class EPL(ThinLens): Parameters ---------- z_l: Optional[Union[Tensor, float]] - This is the redshift of the lens. In the context of gravitational lensing, the lens is the galaxy or other mass distribution that is bending the light from a more distant source. + This is the redshift of the lens. + In the context of gravitational lensing, + the lens is the galaxy or other mass distribution + that is bending the light from a more distant source. x0 and y0: Optional[Union[Tensor, float]] - These are the coordinates of the lens center in the lens plane. The lens plane is the plane perpendicular to the line of sight in which the deflection of light by the lens is considered. + These are the coordinates of the lens center in the lens plane. + The lens plane is the plane perpendicular to the line of sight + in which the deflection of light by the lens is considered. q: Optional[Union[Tensor, float]] - This is the axis ratio of the lens, i.e., the ratio of the minor axis to the major axis of the elliptical lens. + This is the axis ratio of the lens, i.e., the ratio + of the minor axis to the major axis of the elliptical lens. phi: Optional[Union[Tensor, float]] - This is the orientation of the lens on the sky, typically given as an angle measured counter-clockwise from some reference direction. + This is the orientation of the lens on the sky, + typically given as an angle measured counter-clockwise + from some reference direction. b: Optional[Union[Tensor, float]] - This is the scale length of the lens, which sets the overall scale of the lensing effect. In some contexts, this is referred to as the Einstein radius. + This is the scale length of the lens, + which sets the overall scale of the lensing effect. + In some contexts, this is referred to as the Einstein radius. t: Optional[Union[Tensor, float]] - This is the power-law slope parameter of the lens model. In the context of the EPL model, t is equivalent to the gamma parameter minus one, where gamma is the power-law index of the radial mass distribution of the lens. + This is the power-law slope parameter of the lens model. + In the context of the EPL model, + t is equivalent to the gamma parameter minus one, + where gamma is the power-law index + of the radial mass distribution of the lens. """ @@ -66,19 +83,26 @@ def __init__( cosmology: Cosmology Cosmology object that provides cosmological distance calculations. z_l: Optional[Tensor] - Redshift of the lens. If not provided, it is considered as a free parameter. + Redshift of the lens. + If not provided, it is considered as a free parameter. x0: Optional[Tensor] - X coordinate of the lens center. If not provided, it is considered as a free parameter. + X coordinate of the lens center. + If not provided, it is considered as a free parameter. y0: Optional[Tensor] - Y coordinate of the lens center. If not provided, it is considered as a free parameter. + Y coordinate of the lens center. + If not provided, it is considered as a free parameter. q: Optional[Tensor] - Axis ratio of the lens. If not provided, it is considered as a free parameter. + Axis ratio of the lens. + If not provided, it is considered as a free parameter. phi: Optional[Tensor] - Position angle of the lens. If not provided, it is considered as a free parameter. + Position angle of the lens. + If not provided, it is considered as a free parameter. b: Optional[Tensor] - Scale length of the lens. If not provided, it is considered as a free parameter. + Scale length of the lens. + If not provided, it is considered as a free parameter. t: Optional[Tensor] - Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. + Power law slope (`gamma-1`) of the lens. + If not provided, it is considered as a free parameter. s: float Softening length for the elliptical power-law profile. n_iter: int diff --git a/src/caustics/lenses/external_shear.py b/src/caustics/lenses/external_shear.py index a2276b14..ae21b1ec 100644 --- a/src/caustics/lenses/external_shear.py +++ b/src/caustics/lenses/external_shear.py @@ -6,6 +6,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("ExternalShear",) diff --git a/src/caustics/lenses/mass_sheet.py b/src/caustics/lenses/mass_sheet.py index 1e7a79ea..41dc347d 100644 --- a/src/caustics/lenses/mass_sheet.py +++ b/src/caustics/lenses/mass_sheet.py @@ -7,6 +7,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("MassSheet",) diff --git a/src/caustics/lenses/multiplane.py b/src/caustics/lenses/multiplane.py index 59192ed2..52f7ab75 100644 --- a/src/caustics/lenses/multiplane.py +++ b/src/caustics/lenses/multiplane.py @@ -8,6 +8,7 @@ from ..cosmology import Cosmology from .base import ThickLens, ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("Multiplane",) @@ -79,14 +80,19 @@ def raytrace( \vec{x}^{i+1} = \vec{x}^i + D_{i+1,i}\left[\vec{\theta} - \sum_{j=1}^{i}\bf{\alpha}^j(\vec{x}^j)\right] - As an initialization we set the physical positions at the first lensing plane to be :math:`\vec{\theta}D_{1,0}` which is just propogation through regular space to the first plane. Note that :math:`\vec{\alpha}` is a physical deflection angle. The equation above converts straightforwardly into a recursion formula: + As an initialization we set the physical positions at the first lensing plane to be :math:`\vec{\theta}D_{1,0}` which is just propagation through regular space to the first plane. + Note that :math:`\vec{\alpha}` is a physical deflection angle. The equation above converts straightforwardly into a recursion formula: .. math:: \vec{x}^{i+1} = \vec{x}^i + D_{i+1,i}\vec{\theta}^{i} \vec{\theta}^{i+1} = \vec{\theta}^{i} - \alpha^i(\vec{x}^{i+1}) - Here we set as initialization :math:`\vec{\theta}^0 = theta` the observation angular coordinates and :math:`\vec{x}^0 = 0` the initial physical coordinates (i.e. the observation rays come from a point at the observer). The indexing of :math:`\vec{x}^i` and :math:`\vec{\theta}^i` indicates the properties at the plane :math:`i`, and 0 means the observer, 1 is the first lensing plane (infinitesimally after the plane since the deflection has been applied), and so on. Note that in the actual implementation we start at :math:`\vec{x}^1` and :math:`\vec{\theta}^0` and begin at the second step in the recursion formula. + Here we set as initialization :math:`\vec{\theta}^0 = theta` the observation angular coordinates and :math:`\vec{x}^0 = 0` the initial physical coordinates (i.e. the observation rays come from a point at the observer). + The indexing of :math:`\vec{x}^i` and :math:`\vec{\theta}^i` indicates the properties at the plane :math:`i`, + and 0 means the observer, 1 is the first lensing plane (infinitesimally after the plane since the deflection has been applied), + and so on. Note that in the actual implementation we start at :math:`\vec{x}^1` and :math:`\vec{\theta}^0` + and begin at the second step in the recursion formula. Parameters ---------- @@ -101,10 +107,10 @@ def raytrace( Returns ------- - tuple[Tensor, Tensor] + (Tensor, Tensor) The reduced deflection angle. - """ + """ # noqa: E501 return self.raytrace_z1z2(x, y, torch.zeros_like(z_s), z_s, params) @unpack(4) @@ -132,7 +138,8 @@ def raytrace_z1z2( ) X, Y = x * arcsec_to_rad * D, y * arcsec_to_rad * D - # Initial angles are observation angles (negative needed because of negative in propogation term) + # Initial angles are observation angles + # (negative needed because of negative in propagation term) theta_x, theta_y = x, y for i in lens_planes: @@ -151,7 +158,7 @@ def raytrace_z1z2( theta_x = theta_x - alpha_x theta_y = theta_y - alpha_y - # Propogate rays to next plane (basically eq 18) + # Propagate rays to next plane (basically eq 18) z_next = z_ls[i + 1] if i != lens_planes[-1] else z_end D = self.cosmology.transverse_comoving_distance_z1z2( z_ls[i], z_next, params diff --git a/src/caustics/lenses/nfw.py b/src/caustics/lenses/nfw.py index 4f70ed27..5222d73a 100644 --- a/src/caustics/lenses/nfw.py +++ b/src/caustics/lenses/nfw.py @@ -9,6 +9,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed DELTA = 200.0 @@ -37,7 +38,7 @@ class NFW(ThinLens): Softening parameter to avoid singularities at the center of the lens. Default is 0.0. use_case: str Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile - specifically cant be both batchable and differentiable. You may select which version + specifically can't be both batchable and differentiable. You may select which version you wish to use by setting this parameter to one of: batchable, differentiable. Methods diff --git a/src/caustics/lenses/pixelated_convergence.py b/src/caustics/lenses/pixelated_convergence.py index b0cb2d7c..c6db002b 100644 --- a/src/caustics/lenses/pixelated_convergence.py +++ b/src/caustics/lenses/pixelated_convergence.py @@ -10,6 +10,7 @@ from ..utils import get_meshgrid, interp2d, safe_divide, safe_log from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("PixelatedConvergence",) @@ -27,7 +28,7 @@ def __init__( shape: Optional[tuple[int, ...]] = None, convolution_mode: str = "fft", use_next_fast_len: bool = True, - padding="zero", + padding: str = "zero", name: str = None, ): """Strong lensing with user provided kappa map @@ -61,16 +62,23 @@ def __init__( The shape of the convergence map. convolution_mode: (str, optional) The convolution mode for calculating deflection angles and lensing potential. - It can be either "fft" (Fast Fourier Transform) or "conv2d" (2D convolution). Default is "fft". + It can be either "fft" (Fast Fourier Transform) or "conv2d" (2D convolution). + Default is "fft". use_next_fast_len: (bool, optional) If True, adds additional padding to speed up the FFT by calling - `scipy.fft.next_fast_len`. The speed boost can be substantial when `n_pix` is a multiple of a + `scipy.fft.next_fast_len`. + The speed boost can be substantial when `n_pix` is a multiple of a small prime number. Default is True. - padding: string - Specifies the type of padding to use. "zero" will do zero padding, "circular" will do - cyclic boundaries. "reflect" will do reflection padding. "tile" will tile the image at 2x2 which - basically identical to circular padding, but is easier. Generally you should use either "zero" - or "tile". + padding: { "zero", "circular", "reflect", "tile" } + + Specifies the type of padding to use: + "zero" will do zero padding, + "circular" will do cyclic boundaries. + "reflect" will do reflection padding. + "tile" will tile the image at 2x2 which + basically identical to circular padding, but is easier. + + Generally you should use either "zero" or "tile". """ @@ -200,7 +208,7 @@ def _unpad_conv2d(self, x: Tensor) -> Tensor: Tensor The input tensor without padding. """ - return x # torch.roll(x, (-self.padding_range * self.ax_kernel.shape[0]//4,-self.padding_range * self.ax_kernel.shape[1]//4), dims = (-2,-1))[..., :self.n_pix, :self.n_pix] #[..., 1:, 1:] + return x # noqa: E501 torch.roll(x, (-self.padding_range * self.ax_kernel.shape[0]//4,-self.padding_range * self.ax_kernel.shape[1]//4), dims = (-2,-1))[..., :self.n_pix, :self.n_pix] #[..., 1:, 1:] @property def convolution_mode(self): @@ -336,7 +344,7 @@ def _deflection_angle_conv2d( 2 * self.n_pix convergence_map_flipped = convergence_map.flip((-1, -2))[ None, None - ] # F.pad(, ((pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2), mode = self.padding_mode) + ] # noqa: E501 F.pad(, ((pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2), mode = self.padding_mode) deflection_angle_x = F.conv2d( self.ax_kernel[None, None], convergence_map_flipped, padding="same" ).squeeze() * (self.pixelscale**2 / pi) diff --git a/src/caustics/lenses/point.py b/src/caustics/lenses/point.py index 0822c1ff..b649dcf1 100644 --- a/src/caustics/lenses/point.py +++ b/src/caustics/lenses/point.py @@ -7,6 +7,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("Point",) diff --git a/src/caustics/lenses/pseudo_jaffe.py b/src/caustics/lenses/pseudo_jaffe.py index f6990841..2a254606 100644 --- a/src/caustics/lenses/pseudo_jaffe.py +++ b/src/caustics/lenses/pseudo_jaffe.py @@ -9,6 +9,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("PseudoJaffe",) diff --git a/src/caustics/lenses/sie.py b/src/caustics/lenses/sie.py index 5e9e50b5..138de35c 100644 --- a/src/caustics/lenses/sie.py +++ b/src/caustics/lenses/sie.py @@ -6,6 +6,7 @@ from ..utils import derotate, translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("SIE",) diff --git a/src/caustics/lenses/singleplane.py b/src/caustics/lenses/singleplane.py index 577d336e..114d27e7 100644 --- a/src/caustics/lenses/singleplane.py +++ b/src/caustics/lenses/singleplane.py @@ -6,6 +6,7 @@ from ..cosmology import Cosmology from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("SinglePlane",) @@ -48,7 +49,8 @@ def reduced_deflection_angle( **kwargs, ) -> tuple[Tensor, Tensor]: """ - Calculate the total deflection angle by summing the deflection angles of all individual lenses. + Calculate the total deflection angle by summing + the deflection angles of all individual lenses. Parameters ---------- @@ -85,7 +87,8 @@ def convergence( **kwargs, ) -> Tensor: """ - Calculate the total projected mass density by summing the mass densities of all individual lenses. + Calculate the total projected mass density by + summing the mass densities of all individual lenses. Parameters ---------- @@ -120,7 +123,8 @@ def potential( **kwargs, ) -> Tensor: """ - Compute the total lensing potential by summing the lensing potentials of all individual lenses. + Compute the total lensing potential by summing + the lensing potentials of all individual lenses. Parameters ----------- diff --git a/src/caustics/lenses/sis.py b/src/caustics/lenses/sis.py index bb72118f..11de22d3 100644 --- a/src/caustics/lenses/sis.py +++ b/src/caustics/lenses/sis.py @@ -6,6 +6,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed __all__ = ("SIS",) diff --git a/src/caustics/lenses/tnfw.py b/src/caustics/lenses/tnfw.py index 1132448d..d983d39e 100644 --- a/src/caustics/lenses/tnfw.py +++ b/src/caustics/lenses/tnfw.py @@ -9,6 +9,7 @@ from ..utils import translate_rotate from .base import ThinLens from ..parametrized import unpack +from ..packed import Packed DELTA = 200.0 @@ -38,7 +39,7 @@ class TNFW(ThinLens): `interpret_m_total_mass = False` on initialization of the object. However, the mass within R200 will be computed for an NFW profile, not a TNFW profile. This is in line with how - lenstronomy inteprets the mass parameter. + lenstronomy interprets the mass parameter. Parameters ----- @@ -64,12 +65,12 @@ class TNFW(ThinLens): Default is 0.0. interpret_m_total_mass: boolean Indicates how to interpret the mass variable "m". If true - the mass is intepreted as the total mass of the halo (good because it makes sense). If - false it is intepreted as what the mass would have been within R200 of a an NFW that + the mass is interpreted as the total mass of the halo (good because it makes sense). If + false it is interpreted as what the mass would have been within R200 of a an NFW that isn't truncated (good because it is easily compared with an NFW). use_case: str Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile - specifically cant be both batchable and differentiable. You may select which version + specifically can't be both batchable and differentiable. You may select which version you wish to use by setting this parameter to one of: batchable, differentiable. """ @@ -154,7 +155,8 @@ def get_concentration( **kwargs, ) -> Tensor: """ - Calculate the scale radius of the lens. This is the same formula used for the classic NFW profile. + Calculate the scale radius of the lens. + This is the same formula used for the classic NFW profile. Parameters ---------- @@ -237,7 +239,9 @@ def get_M0( **kwargs, ) -> Tensor: """ - Calculate the reference mass. This is an abstract reference mass used internally in the equations from Baltz et al. 2009. + Calculate the reference mass. + This is an abstract reference mass used internally + in the equations from Baltz et al. 2009. Parameters ---------- @@ -343,7 +347,8 @@ def convergence( **kwargs, ) -> Tensor: """ - TNFW convergence as given in Baltz et al. 2009. This is unitless since it is Sigma(x) / Sigma_crit. + TNFW convergence as given in Baltz et al. 2009. + This is unitless since it is Sigma(x) / Sigma_crit. Parameters ---------- @@ -513,7 +518,9 @@ def potential( **kwargs, ) -> Tensor: """ - Compute the lensing potential. Note that this is not a unitless potential! This is the potential as given in Baltz et al. 2009. + Compute the lensing potential. + Note that this is not a unitless potential! + This is the potential as given in Baltz et al. 2009. TODO: convert to dimensionless potential. diff --git a/src/caustics/lenses/utils.py b/src/caustics/lenses/utils.py index 3d4d52ff..aaab79c3 100644 --- a/src/caustics/lenses/utils.py +++ b/src/caustics/lenses/utils.py @@ -29,7 +29,8 @@ def get_pix_jacobian( Returns -------- - The Jacobian matrix of the image position with respect to the source position at the given point. + The Jacobian matrix of the image position with respect + to the source position at the given point. """ jac = torch.func.jacfwd(raytrace, (0, 1))(x, y, z_s) # type: ignore @@ -38,7 +39,8 @@ def get_pix_jacobian( def get_pix_magnification(raytrace, x, y, z_s) -> Tensor: """ - Computes the magnification at a single point on the lensing plane. The magnification is derived from the determinant + Computes the magnification at a single point on the lensing plane. + The magnification is derived from the determinant of the Jacobian matrix of the image position with respect to the source position. Parameters @@ -63,7 +65,8 @@ def get_pix_magnification(raytrace, x, y, z_s) -> Tensor: def get_magnification(raytrace, x, y, z_s) -> Tensor: """ - Computes the magnification over a grid on the lensing plane. This is done by calling `get_pix_magnification` + Computes the magnification over a grid on the lensing plane. + This is done by calling `get_pix_magnification` for each point on the grid. Parameters diff --git a/src/caustics/light/__init__.py b/src/caustics/light/__init__.py index 120ea22d..62020e89 100644 --- a/src/caustics/light/__init__.py +++ b/src/caustics/light/__init__.py @@ -1,4 +1,6 @@ -from .base import * -from .pixelated import * -from .probes import * -from .sersic import * +from .base import Source +from .pixelated import Pixelated +from .probes import PROBESDataset +from .sersic import Sersic + +__all__ = ["Source", "Pixelated", "PROBESDataset", "Sersic"] diff --git a/src/caustics/light/base.py b/src/caustics/light/base.py index e457a58f..ae575ece 100644 --- a/src/caustics/light/base.py +++ b/src/caustics/light/base.py @@ -4,19 +4,25 @@ from torch import Tensor from ..parametrized import Parametrized, unpack +from ..packed import Packed __all__ = ("Source",) class Source(Parametrized): """ - This is an abstract base class used to represent a source in a strong gravitational lensing system. - It provides the basic structure and required methods that any derived source class should implement. - The Source class inherits from the Parametrized class, implying that it contains parameters that can + This is an abstract base class used to represent a source + in a strong gravitational lensing system. + It provides the basic structure and required methods + that any derived source class should implement. + The Source class inherits from the Parametrized class, + implying that it contains parameters that can be optimized or manipulated. - The class introduces one abstract method, `brightness`, that must be implemented in any concrete - subclass. This method calculates the brightness of the source at given coordinates. + The class introduces one abstract method, `brightness`, + that must be implemented in any concrete subclass. + This method calculates the brightness of + the source at given coordinates. """ @abstractmethod @@ -31,25 +37,31 @@ def brightness( 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. + The x-coordinate(s) at which to calculate + the source brightness. + This could be a single value or a tensor of values. y: Tensor - The y-coordinate(s) at which to calculate the source brightness. - This could be a single value or a tensor of values. + The y-coordinate(s) at which to calculate + the source brightness. + This could be a single value or a tensor of values. - params: (Packed, optional) - Dynamic parameter container that might be required to calculate - the brightness. The exact contents will depend on the specific implementation in derived classes. + params: Packed, optional + Dynamic parameter container that might be required + to calculate the brightness. + The exact contents will depend on the specific + implementation in derived classes. Returns ------- Tensor - The brightness of the source at the given coordinate(s). The exact form of the output - will depend on the specific implementation in the derived class. + The brightness of the source at the given coordinate(s). + The exact form of the output will depend on + the specific implementation in the derived class. - Note + Notes ----- - This method must be overridden in any class that inherits from `Source`. + This method must be overridden in any class + that inherits from `Source`. """ ... diff --git a/src/caustics/light/pixelated.py b/src/caustics/light/pixelated.py index 142e53b0..667f0e92 100644 --- a/src/caustics/light/pixelated.py +++ b/src/caustics/light/pixelated.py @@ -5,25 +5,34 @@ from ..utils import interp2d from .base import Source from ..parametrized import unpack +from ..packed import Packed __all__ = ("Pixelated",) class Pixelated(Source): """ - `Pixelated` is a subclass of the abstract class `Source`. It representes the brightness profile of - source with a pixelated grid of intensity values. + ``Pixelated`` is a subclass of the abstract class ``Source``. + It represents the brightness profile of source + with a pixelated grid of intensity values. - This class provides a concrete implementation of the `brightness` method required by the `Source` - superclass. In this implementation, brightness is determined by interpolating values from the - provided source image. + This class provides a concrete implementation + of the ``brightness`` method required by the ``Source`` superclass. + In this implementation, brightness is determined + by interpolating values from the provided source image. - Attributes: - x0 (Optional[Tensor]): The x-coordinate of the source image's center. - y0 (Optional[Tensor]): The y-coordinate of the source image's center. - image (Optional[Tensor]): The source image from which brightness values will be interpolated. - pixelscale (Optional[Tensor]): The pixelscale of the source image in the lens plane in units of arcsec/pixel. - shape (Optional[tuple[int, ...]]): The shape of the source image. + Attributes + ---------- + x0 : Tensor, optional + The x-coordinate of the source image's center. + y0 : Tensor, optional + The y-coordinate of the source image's center. + image : Tensor, optional + The source image from which brightness values will be interpolated. + pixelscale : Tensor, optional + The pixelscale of the source image in the lens plane in units of arcsec/pixel. + shape : tuple of ints, optional + The shape of the source image. """ def __init__( @@ -38,13 +47,20 @@ def __init__( """ Constructs the `Pixelated` object with the given parameters. - Args: - name (str): The name of the source. - x0 (Optional[Tensor]): The x-coordinate of the source image's center. - y0 (Optional[Tensor]): The y-coordinate of the source image's center. - image (Optional[Tensor]): The source image from which brightness values will be interpolated. - pixelscale (Optional[Tensor]): The pixelscale of the source image in the lens plane in units of arcsec/pixel. - shape (Optional[tuple[int, ...]]): The shape of the source image. + Parameters + ---------- + name : str + The name of the source. + x0 : Tensor, optional + The x-coordinate of the source image's center. + y0 : Tensor, optional + The y-coordinate of the source image's center. + image : Tensor, optional + The source image from which brightness values will be interpolated. + pixelscale : Tensor, optional + The pixelscale of the source image in the lens plane in units of arcsec/pixel. + shape : tuple of ints, optional + The shape of the source image. """ if image is not None and image.ndim not in [2, 3]: raise ValueError( @@ -74,20 +90,28 @@ def brightness( **kwargs, ): """ - Implements the `brightness` method for `Pixelated`. The brightness at a given point is - determined by interpolating values from the source image. + Implements the `brightness` method for `Pixelated`. + The brightness at a given point is determined + by interpolating values from the source image. - Args: - x (Tensor): The x-coordinate(s) at which to calculate the source brightness. - This could be a single value or a tensor of values. - y (Tensor): The y-coordinate(s) at which to calculate the source brightness. - This could be a single value or a tensor of values. - params (Optional[Packed]): A dictionary containing additional parameters that might be required to - calculate the brightness. + 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. + y : Tensor + The y-coordinate(s) at which to calculate the source brightness. + This could be a single value or a tensor of values. + params : Packed, optional + A dictionary containing additional parameters that might be required to + calculate the brightness. - Returns: - Tensor: The brightness of the source at the given coordinate(s). The brightness is - determined by interpolating values from the source image. + Returns + ------- + Tensor + The brightness of the source at the given coordinate(s). + The brightness is determined by interpolating values + from the source image. """ fov_x = pixelscale * image.shape[0] fov_y = pixelscale * image.shape[1] diff --git a/src/caustics/light/sersic.py b/src/caustics/light/sersic.py index c913534d..d2f5e13b 100644 --- a/src/caustics/light/sersic.py +++ b/src/caustics/light/sersic.py @@ -5,17 +5,20 @@ from ..utils import to_elliptical, translate_rotate from .base import Source from ..parametrized import unpack +from ..packed import Packed __all__ = ("Sersic",) class Sersic(Source): """ - `Sersic` is a subclass of the abstract class `Source`. It represents a source in a strong - gravitational lensing system that follows a Sersic profile, a mathematical function that describes + `Sersic` is a subclass of the abstract class `Source`. + It represents a source in a strong gravitational lensing system + that follows a Sersic profile, a mathematical function that describes how the intensity I of a galaxy varies with distance r from its center. - The Sersic profile is often used to describe elliptical galaxies and spiral galaxies' bulges. + The Sersic profile is often used to describe + elliptical galaxies and spiral galaxies' bulges. Attributes ----------- @@ -118,24 +121,25 @@ def brightness( y: Tensor The y-coordinate(s) at which to calculate the source brightness. This could be a single value or a tensor of values. - params: (Packed, optional) + 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`. + Returns + ------- + Tensor + The brightness of the source at the given point(s). + The output tensor has the same shape as `x` and `y`. Notes ----- - The Sersic profile is defined as: I(r) = Ie * exp(-k * ((r / r_e)^(1/n) - 1)), - where Ie is the intensity at the effective radius r_e, n is the Sersic index - that describes the concentration of the source, and k is a parameter that - depends on n. In this implementation, we use elliptical coordinates ex and ey, - and the transformation from Cartesian coordinates is handled by `to_elliptical`. - The value of k can be calculated in two ways, controlled by `lenstronomy_k_mode`. - If `lenstronomy_k_mode` is True, we use the approximation from Lenstronomy, - otherwise, we use the approximation from Ciotti & Bertin (1999). + The Sersic profile is defined as: I(r) = Ie * exp(-k * ((r / r_e)^(1/n) - 1)), + where Ie is the intensity at the effective radius r_e, n is the Sersic index + that describes the concentration of the source, and k is a parameter that + depends on n. In this implementation, we use elliptical coordinates ex and ey, + and the transformation from Cartesian coordinates is handled by `to_elliptical`. + The value of k can be calculated in two ways, controlled by `lenstronomy_k_mode`. + If `lenstronomy_k_mode` is True, we use the approximation from Lenstronomy, + otherwise, we use the approximation from Ciotti & Bertin (1999). """ x, y = translate_rotate(x, y, x0, y0, phi) ex, ey = to_elliptical(x, y, q) diff --git a/src/caustics/namespace_dict.py b/src/caustics/namespace_dict.py index 7106f5df..e96026bb 100644 --- a/src/caustics/namespace_dict.py +++ b/src/caustics/namespace_dict.py @@ -124,7 +124,7 @@ class NestedNamespaceDict(_NestedNamespaceDict): nested_namespace.foo = 'Hello' nested_namespace.bar = {'baz': 'World'} nested_namespace.bar.qux = 42 - # works also in the follwoing way + # works also in the following way nested_namespace["bar.qux"] = 42 print(nested_namespace) diff --git a/src/caustics/parameter.py b/src/caustics/parameter.py index 1fcda9ff..933de529 100644 --- a/src/caustics/parameter.py +++ b/src/caustics/parameter.py @@ -8,9 +8,13 @@ class Parameter: """ - Represents a static or dynamic parameter used for strong gravitational lensing simulations in the caustics codebase. + Represents a static or dynamic parameter used + for strong gravitational lensing simulations + in the caustics codebase. - A static parameter has a fixed value, while a dynamic parameter must be passed in each time it's required. + A static parameter has a fixed value, + while a dynamic parameter must be passed + in each time it's required. Attributes ---------- @@ -56,7 +60,8 @@ def value(self, value: Union[None, Tensor, float]): value = torch.as_tensor(value) if value.shape != self.shape: raise ValueError( - f"Cannot set Parameter value with a different shape. Received {value.shape}, expected {self.shape}" + "Cannot set Parameter value with a different shape. " + f"Received {value.shape}, expected {self.shape}" ) self._value = value self._dtype = None if value is None else value.dtype diff --git a/src/caustics/parametrized.py b/src/caustics/parametrized.py index ecac2de9..e42b69e8 100644 --- a/src/caustics/parametrized.py +++ b/src/caustics/parametrized.py @@ -8,6 +8,7 @@ import re import keyword from torch import Tensor +import graphviz from .packed import Packed from .namespace_dict import NamespaceDict, NestedNamespaceDict @@ -26,11 +27,14 @@ def check_valid_name(name): class Parametrized: """ - Represents a class with Param and Parametrized attributes, typically used to construct parts of a simulator + Represents a class with Param and Parametrized attributes, + typically used to construct parts of a simulator that have parameters which need to be tracked during MCMC sampling. - This class can contain Params, Parametrized, tensor buffers or normal attributes as its attributes. - It provides functionalities to manage these attributes, ensuring that an attribute of one type isn't rebound + This class can contain Params, Parametrized, + tensor buffers or normal attributes as its attributes. + It provides functionalities to manage these attributes, + ensuring that an attribute of one type isn't rebound to be of a different type. TODO @@ -92,10 +96,13 @@ def __setattr__(self, key, value): # Create new parameter and attach it as an attribute self.add_param(key, value.value, value.shape) elif isinstance(value, Parametrized): - # Update map from attribute key to module name for __getattribute__ method + # Update map from attribute key to module name + # for __getattribute__ method self._module_key_map[value.name] = key self.add_parametrized(value, set_attr=False) - # set attr only to user defined key, not module name (self.{module.name} is still accessible, see __getattribute__ method) + # set attr only to user defined key, + # not module name (self.{module.name} is still accessible, + # see __getattribute__ method) super().__setattr__(key, value) else: super().__setattr__(key, value) @@ -122,7 +129,8 @@ def to( self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None ): """ - Moves static Params for this component and its childs to the specified device and casts them to the specified data type. + Moves static Params for this component and its childs + to the specified device and casts them to the specified data type. """ for name, p in self._params.items(): self._params[name] = p.to(device, dtype) @@ -198,13 +206,14 @@ def pack( ) -> Packed: """ Converts a list or tensor into a dict that can subsequently be unpacked - into arguments to this component and its childs. Also, add a batch dimension - to each Tensor without such a dimension. + into arguments to this component and its childs. + Also, add a batch dimension to each Tensor + without such a dimension. Parameters ---------- - x: (Union[list[Tensor], dict[str, Union[list[Tensor], Tensor, dict[str, Tensor]]], Tensor) - The input to be packed. Can be a list of tensors, a dictionary of tensors, or a single tensor. + x : list of tensor, dict of tensor, or tensor + The input to be packed. Returns ------- @@ -307,7 +316,9 @@ def unpack( if isinstance(x, dict): if self.name in x.keys() and x.get(self.name, {}): print( - f"Module {self.name} is static, the parameters {' '.join(x[self.name].keys())} passed dynamically will be ignored ignored" + f"Module {self.name} is static, " + f"the parameters {' '.join(x[self.name].keys())} " + "passed dynamically will be ignored." ) unpacked_x = [] offset = 0 @@ -327,7 +338,8 @@ def unpack( else: raise ValueError( f"Invalid data type found when unpacking parameters for {self.name}." - f"Expected argument of unpack to be a list/tuple/dict of Tensor, or simply a flattened tensor" + "Expected argument of unpack to be a list/tuple/dict of Tensor, " + "or simply a flattened tensor" f"but found {type(dynamic_x)}." ) else: # param is static @@ -498,7 +510,8 @@ def wrapped(self, *args, **kwargs): # Case 1: Params is already Packed (or no params were passed) x = args.pop(0) elif "params" in kwargs: - # Case 2: params was passed explicitly as a kwargs, i.e. user used signature "method(*leading_args, params=params)" + # Case 2: params was passed explicitly as a kwargs, + # i.e. user used signature "method(*leading_args, params=params)" x = kwargs["params"] else: # Case 3 (most common): params were passed as the trailing arguments of the method diff --git a/src/caustics/sims/__init__.py b/src/caustics/sims/__init__.py index a1fba9aa..812d627e 100644 --- a/src/caustics/sims/__init__.py +++ b/src/caustics/sims/__init__.py @@ -1,2 +1,4 @@ -from .lens_source import * -from .simulator import * +from .lens_source import Lens_Source +from .simulator import Simulator + +__all__ = ["Lens_Source", "Simulator"] diff --git a/src/caustics/sims/lens_source.py b/src/caustics/sims/lens_source.py index 80401d7a..dcb22807 100644 --- a/src/caustics/sims/lens_source.py +++ b/src/caustics/sims/lens_source.py @@ -19,11 +19,11 @@ class Lens_Source(Simulator): """Lens image of a source. - Striaghtforward simulator to sample a lensed image of a source + Straightforward simulator to sample a lensed image of a source object. Constructs a sampling grid internally based on the pixelscale and gridding parameters. It can automatically upscale and fine sample an image. This is the most straightforward - simulator to view the image if you aready have a lens and source + simulator to view the image if you already have a lens and source chosen. Example usage:: @@ -72,7 +72,7 @@ class Lens_Source(Simulator): - For arbitrary pixel integration accuracy using the quad_level parameter. This will use Gaussian quadrature to sample the image at a higher resolution, then integrate the image back to the original resolution. This is useful for high accuracy integration of the image, but is not recommended for large images as it will be slow. The quad_level and upsample_factor can be used together to achieve high accuracy integration of the image convolved with a PSF. - A `Pixelated` light source is defined by bilinear interpolation of the provided image. This means that sub-pixel integration is not required for accurate integration of the pixels. However, if you are using a PSF then you should still use upsample_factor (if your PSF is supersampled) to ensure that everything is sampled at the PSF resolution. - """ + """ # noqa: E501 def __init__( self, diff --git a/src/caustics/tests.py b/src/caustics/tests.py index e86114ad..0e26dc13 100644 --- a/src/caustics/tests.py +++ b/src/caustics/tests.py @@ -210,7 +210,10 @@ def _test_multiplane_effective_convergence(): def test(): """ - Run tests for caustics basic functionallity. Run this function to ensure that caustics is working properly. Simply call:: + Run tests for caustics basic functionality. + Run this function to ensure that caustics is working properly. + + Simply call:: >>> import caustics >>> caustics.test() diff --git a/src/caustics/utils.py b/src/caustics/utils.py index b784bd5e..be5bdcc8 100644 --- a/src/caustics/utils.py +++ b/src/caustics/utils.py @@ -31,7 +31,8 @@ def flip_axis_ratio(q, phi): def translate_rotate(x, y, x0, y0, phi: Optional[Tensor] = None): """ - Translates and rotates the points (x, y) by subtracting (x0, y0) and applying rotation angle phi. + Translates and rotates the points (x, y) by subtracting (x0, y0) + and applying rotation angle phi. Parameters ---------- @@ -239,7 +240,8 @@ def gaussian_quadrature_integrator( ): """ Performs a pixel-wise integration using Gaussian quadrature. - It takes the brightness function evaluated at the quadrature points `F` and the quadrature weights `weight` as input. + It takes the brightness function evaluated at the quadrature points `F` + and the quadrature weights `weight` as input. The result is the integrated brightness function at each pixel. @@ -285,7 +287,8 @@ def quad( Parameters ---------- F : Callable - The brightness function to be evaluated at the quadrature points. The function should take as input: F(X, Y, *args). + The brightness function to be evaluated at the quadrature points. + The function should take as input: F(X, Y, *args). pixelscale : float The scale of each pixel. X : Tensor diff --git a/tests/test_base.py b/tests/test_base.py index 12aa54af..d7ee2f85 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -36,8 +36,9 @@ def test(): assert torch.all((sp_x - bx).abs() < 1e-3) assert torch.all((sp_y - by).abs() < 1e-3) + def test_quicktest(): """ Quick test to check that the built-in `test` module is working """ - mini_test() \ No newline at end of file + mini_test() diff --git a/tests/test_batching.py b/tests/test_batching.py index 78004aea..d49e50c9 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -47,7 +47,8 @@ def test_vmapped_simulator_with_pixelated_modules(): print(x[2].shape) assert vmap(sim)(x).shape == torch.Size([2, n_pix, n_pix]) - # test tensor input: Does not work well with images since it would require unflattening the images in caustics + # test tensor input: Does not work well with images since it would require + # unflattening the images in caustics # x_tensor = torch.concat([_x.view(2, -1) for _x in x], dim=1) # print(x_tensor.shape) # assert vmap(sim)(x_tensor).shape == torch.Size([2, n_pix, n_pix]) diff --git a/tests/test_jacobian_lens_equation.py b/tests/test_jacobian_lens_equation.py index ef126551..31b57347 100644 --- a/tests/test_jacobian_lens_equation.py +++ b/tests/test_jacobian_lens_equation.py @@ -125,4 +125,7 @@ def test_multiplane_effective_convergence(): if __name__ == "__main__": - test() + test_jacobian_autograd_vs_finitediff() + test_multiplane_jacobian() + test_multiplane_jacobian_autograd_vs_finitediff() + test_multiplane_effective_convergence() diff --git a/tests/test_parameter.py b/tests/test_parameter.py index 5c6f90f1..8f1c49d8 100644 --- a/tests/test_parameter.py +++ b/tests/test_parameter.py @@ -22,7 +22,8 @@ def test_shape_error_messages(): n_pix = 20 cosmo = FlatLambdaCDM() with pytest.raises(TypeError): - # User cannot enter a list, only a tuple (because of type checking and consistency with torch) + # User cannot enter a list, only a tuple + # (because of type checking and consistency with torch) PixelatedConvergence(fov, n_pix, cosmo, shape=[8, 8]) with pytest.raises(ValueError): diff --git a/tests/test_parametrized.py b/tests/test_parametrized.py index 660321c3..6463a2b7 100644 --- a/tests/test_parametrized.py +++ b/tests/test_parametrized.py @@ -21,7 +21,7 @@ def __init__(self): sim = Sim() assert len(sim.module_params) == 2 # dynamic and static - assert len(sim.module_params.dynamic) == 0 # simulator has no dynmaic params + assert len(sim.module_params.dynamic) == 0 # simulator has no dynamic params assert len(sim.module_params.static) == 1 # and 1 static param (z_s) assert len(sim.params) == 2 # dynamic and static assert len(sim.params.dynamic) == 3 # cosmo, epl and sersic @@ -120,7 +120,8 @@ def __init__(self): def test_parametrized_name_setter_bad_names(): - # Make sure bad names are catched by our added method. Bad names are name which cannot be used as class attributes. + # Make sure bad names are caught by our added method. + # Bad names are name which cannot be used as class attributes. good_names = ["variable", "_variable", "var_iable2"] for name in good_names: Sersic(name=name) @@ -142,7 +143,7 @@ def __init__(self): self.lens2 = EPL(self.cosmo) sim = Sim() - # Current way names are updated. Could be chnaged so that all params in collision + # Current way names are updated. Could be changed so that all params in collision # Get a number assert sim.lens1.name == "EPL" assert sim.lens2.name == "EPL_1"