From ee92bebecd47e769ab9974807bc0e4b3127fa828 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 2 Mar 2026 11:47:09 +0000 Subject: [PATCH 1/7] add JAX-accelerated hessian_from via jacfwd Replaces the finite-difference hessian_from with a dual-path implementation: - xp=np (default): delegates to _hessian_via_finite_difference, no JAX import - xp=jnp: delegates to _hessian_via_jax which uses jax.jacfwd on a new deflections_yx_scalar helper, supporting both Grid2D and Grid2DIrregular Also fixes shear_yx_2d_via_hessian_from to use grid.shape[0] instead of grid.shape_slim (incompatible with Grid2DIrregular), and makes magnification_2d_via_hessian_from return a raw jax.Array when xp=jnp to avoid wrapping a traced value in ArrayIrregular. Updates CLAUDE.md to document the NumPy-default / JAX-opt-in design pattern. Co-Authored-By: Claude Sonnet 4.6 --- CLAUDE.md | 24 +++- autogalaxy/operate/deflections.py | 139 +++++++++++++------- test_autogalaxy/operate/test_deflections.py | 3 +- 3 files changed, 119 insertions(+), 47 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 7b53df091..327acf11b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -80,7 +80,25 @@ These inherit from `AnalysisDataset` → `Analysis` (in `analysis/analysis/`), w ### JAX Support -JAX is integrated via the `xp` parameter pattern throughout the codebase. Fit classes accept `xp=np` (NumPy, default) or `xp=jnp` (JAX). The `AbstractFitInversion.use_jax` property tracks which backend is active. The `AnalysisImaging.__init__` has `use_jax: bool = True`. The conftest.py forces JAX backend initialization before tests run. +The codebase is designed so that **NumPy is the default everywhere and JAX is opt-in**. JAX is never imported at module level — it is only imported locally inside functions when explicitly requested. + +The `xp` parameter pattern is the single point of control: +- `xp=np` (default throughout) — pure NumPy path, no JAX dependency at runtime +- `xp=jnp` — JAX path, imports `jax` / `jax.numpy` locally inside the function + +This means: +- **Unit tests** (`test_autogalaxy/`) always run on the NumPy path. No test should import JAX or pass `xp=jnp` unless it is explicitly testing the JAX path. +- **Integration tests** (in `autogalaxy_workspace_test/`) are where the JAX path is exercised, typically wrapped in `jax.jit` to test both correctness and compilation. +- `conftest.py` forces JAX backend initialisation before the test suite runs, but this only ensures JAX is available — it does not switch the default backend. + +`AbstractFitInversion.use_jax` tracks whether a fit was constructed with JAX. `AnalysisImaging` has `use_jax: bool = True` to opt into the JAX path for model-fitting. + +When adding a new function that should support JAX: +1. Default the parameter to `xp=np` +2. Guard any JAX imports with `if xp is not np:` and import `jax` / `jax.numpy` locally inside that branch +3. Add the NumPy implementation as the default path (finite-difference, `np.*` calls, etc.) +4. Add a JAX implementation in the guarded branch (e.g. `jax.jacfwd`, `jnp.vectorize`) +5. Verify correctness by comparing both paths in `autogalaxy_workspace_test/scripts/` ### Linear Light Profiles & Inversions @@ -101,6 +119,10 @@ Default priors, visualization settings, and general config live in `autogalaxy/c Both are mixin classes inherited by `LightProfile`, `MassProfile`, `Galaxy`, and `Galaxies`. +### Workspace Script Style + +Scripts in `autogalaxy_workspace` and `autogalaxy_workspace_test` use `"""..."""` docstring blocks as prose commentary throughout — **not** `#` comments. Every script opens with a module-level docstring (title + underline + description), and each logical section of code is preceded by a `"""..."""` block with a `__Section Name__` header explaining what follows. See any script in `autogalaxy_workspace/scripts/` for examples of this style. + ### Workspace (Examples & Notebooks) The `autogalaxy_workspace` at `/mnt/c/Users/Jammy/Code/PyAutoJAX/autogalaxy_workspace` contains runnable examples and tutorials. Key locations: diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index d9551b222..8fc5de15a 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -242,56 +242,115 @@ def magnification_2d_from(self, grid) -> aa.Array2D: return aa.Array2D(values=1 / det_jacobian, mask=grid.mask) - def hessian_from( - self, grid, buffer: float = 0.01, deflections_func=None, xp=np - ) -> Tuple: + def deflections_yx_scalar(self, y, x, pixel_scales): + """ + Returns the deflection angles at a single (y, x) arc-second coordinate as a JAX array of + shape (2,), where index 0 is the y-deflection and index 1 is the x-deflection. + + This is an internal method used by `hessian_from` to enable JAX auto-differentiation via + `jax.jacfwd`. The function must accept y and x as two separate scalar inputs (rather than + a single combined array) so that JAX treats the function as R² -> R² and computes a proper + 2x2 Jacobian matrix. + + Parameters + ---------- + y + The y arc-second coordinate (scalar). + x + The x arc-second coordinate (scalar). + pixel_scales + The pixel scales used to construct the internal (1, 1) Mask2D. + """ + import jax.numpy as jnp + + mask = aa.Mask2D.all_false(shape_native=(1, 1), pixel_scales=pixel_scales) + grid = aa.Grid2D( + values=jnp.stack((y.reshape(1), x.reshape(1)), axis=-1), mask=mask + ) + return self.deflections_yx_2d_from(grid, xp=jnp).squeeze() + + def hessian_from(self, grid, xp=np) -> Tuple: """ Returns the Hessian of the lensing object, where the Hessian is the second partial derivatives of the potential (see equation 55 https://inspirehep.net/literature/419263): `hessian_{i,j} = d^2 / dtheta_i dtheta_j` - The Hessian is computed by evaluating the 2D deflection angles around every (y,x) coordinate on the input 2D - grid map in four directions (positive y, negative y, positive x, negative x), exploiting how the deflection - angles are the derivative of the potential. + The Hessian is returned as a 4-entry tuple reflecting its structure as a 2x2 matrix: + (hessian_yy, hessian_xy, hessian_yx, hessian_xx). - By using evaluating the deflection angles around each grid coordinate, the Hessian can therefore be computed - using uniform or irregular 2D grids of (y,x). This can be slower, because x4 more deflection angle calculations - are required, however it is more flexible in and therefore used throughout **PyAutoLens** by default. + Two computational paths are available, selected via the `xp` parameter: - The Hessian is returned as a 4 entry tuple, which reflect its structure as a 2x2 matrix. + - **NumPy** (``xp=np``, default): finite-difference approximation. Deflection angles are + evaluated at four shifted positions around each grid coordinate (±y, ±x) and the + central difference is taken. JAX is not imported. + + - **JAX** (``xp=jnp``): exact derivatives via ``jax.jacfwd`` applied to + ``deflections_yx_scalar``, vectorised over the grid with ``jnp.vectorize``. + + Both paths support uniform ``Grid2D`` and irregular ``Grid2DIrregular`` grids. Parameters ---------- grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and Hessian are computed on. - buffer - The spacing in the y and x directions around each grid coordinate where deflection angles are computed and - used to estimate the derivative. + The 2D grid of (y,x) arc-second coordinates the Hessian is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). Controls which computational path is + used and the type of the returned arrays. """ - if deflections_func is None: - deflections_func = self.deflections_yx_2d_from + if xp is not np: + return self._hessian_via_jax(grid=grid, xp=xp) - grid_shift_y_up = aa.Grid2DIrregular( - values=xp.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1) + return self._hessian_via_finite_difference(grid=grid) + + def _hessian_via_jax(self, grid, xp) -> Tuple: + import jax + import jax.numpy as jnp + + pixel_scales = getattr(grid, "pixel_scales", (0.05, 0.05)) + + y = jnp.array(grid[:, 0]) + x = jnp.array(grid[:, 1]) + + def _hessian_single(y_scalar, x_scalar): + return jnp.stack( + jax.jacfwd(self.deflections_yx_scalar, argnums=(0, 1))( + y_scalar, x_scalar, pixel_scales + ) + ) + + h = jnp.vectorize(_hessian_single, signature="(),()->(i,i)")(y, x) + + # h has shape (N, 2, 2): + # h[..., 0, 0] = d(defl_y)/dy = hessian_yy + # h[..., 0, 1] = d(defl_x)/dy = hessian_xy + # h[..., 1, 0] = d(defl_y)/dx = hessian_yx + # h[..., 1, 1] = d(defl_x)/dx = hessian_xx + return ( + xp.array(h[..., 0, 0]), + xp.array(h[..., 0, 1]), + xp.array(h[..., 1, 0]), + xp.array(h[..., 1, 1]), ) + def _hessian_via_finite_difference(self, grid, buffer: float = 0.01) -> Tuple: + grid_shift_y_up = aa.Grid2DIrregular( + values=np.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1) + ) grid_shift_y_down = aa.Grid2DIrregular( - values=xp.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1) + values=np.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1) ) - grid_shift_x_left = aa.Grid2DIrregular( - values=xp.stack([grid[:, 0], grid[:, 1] - buffer], axis=1) + values=np.stack([grid[:, 0], grid[:, 1] - buffer], axis=1) ) - grid_shift_x_right = aa.Grid2DIrregular( - values=xp.stack([grid[:, 0], grid[:, 1] + buffer], axis=1) + values=np.stack([grid[:, 0], grid[:, 1] + buffer], axis=1) ) - deflections_up = deflections_func(grid=grid_shift_y_up, xp=xp) - deflections_down = deflections_func(grid=grid_shift_y_down, xp=xp) - deflections_left = deflections_func(grid=grid_shift_x_left, xp=xp) - deflections_right = deflections_func(grid=grid_shift_x_right, xp=xp) + deflections_up = self.deflections_yx_2d_from(grid=grid_shift_y_up) + deflections_down = self.deflections_yx_2d_from(grid=grid_shift_y_down) + deflections_left = self.deflections_yx_2d_from(grid=grid_shift_x_left) + deflections_right = self.deflections_yx_2d_from(grid=grid_shift_x_right) hessian_yy = 0.5 * (deflections_up[:, 0] - deflections_down[:, 0]) / buffer hessian_xy = 0.5 * (deflections_up[:, 1] - deflections_down[:, 1]) / buffer @@ -301,7 +360,7 @@ def hessian_from( return hessian_yy, hessian_xy, hessian_yx, hessian_xx def convergence_2d_via_hessian_from( - self, grid, buffer: float = 0.01 + self, grid ) -> aa.ArrayIrregular: """ Returns the convergence of the lensing object, which is computed from the 2D deflection angle map via the @@ -319,18 +378,13 @@ def convergence_2d_via_hessian_from( ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and Hessian are computed on. - buffer - The spacing in the y and x directions around each grid coordinate where deflection angles are computed and - used to estimate the derivative. """ - hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( - grid=grid, buffer=buffer - ) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) return aa.ArrayIrregular(values=0.5 * (hessian_yy + hessian_xx)) def shear_yx_2d_via_hessian_from( - self, grid, buffer: float = 0.01 + self, grid ) -> ShearYX2DIrregular: """ Returns the 2D (y,x) shear vectors of the lensing object, which are computed from the 2D deflection angle map @@ -355,19 +409,14 @@ def shear_yx_2d_via_hessian_from( ---------- grids The 2D grid of (y,x) arc-second coordinates the deflection angles and Hessian are computed on. - buffer - The spacing in the y and x directions around each grid coordinate where deflection angles are computed and - used to estimate the derivative. """ - hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( - grid=grid, buffer=buffer - ) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) gamma_1 = 0.5 * (hessian_xx - hessian_yy) gamma_2 = hessian_xy - shear_yx_2d = np.zeros(shape=(grid.shape_slim, 2)) + shear_yx_2d = np.zeros(shape=(grid.shape[0], 2)) shear_yx_2d[:, 0] = gamma_2 shear_yx_2d[:, 1] = gamma_1 @@ -375,7 +424,7 @@ def shear_yx_2d_via_hessian_from( return ShearYX2DIrregular(values=shear_yx_2d, grid=grid) def magnification_2d_via_hessian_from( - self, grid, buffer: float = 0.01, deflections_func=None, xp=np + self, grid, xp=np ) -> aa.ArrayIrregular: """ Returns the 2D magnification map of lensing object, which is computed from the 2D deflection angle map @@ -397,11 +446,13 @@ def magnification_2d_via_hessian_from( The 2D grid of (y,x) arc-second coordinates the deflection angles and magnification map are computed on. """ hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( - grid=grid, buffer=buffer, deflections_func=deflections_func, xp=xp + grid=grid, xp=xp ) det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx + if xp is not np: + return xp.array(1.0 / det_A) return aa.ArrayIrregular(values=1.0 / det_A) def contour_list_from(self, grid, contour_array): diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 45cc97509..28bd6b90c 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -111,7 +111,6 @@ def test__hessian_from(): def test__convergence_2d_via_hessian_from(): - buffer = 0.0001 grid = ag.Grid2DIrregular( values=[(1.075, -0.125), (-0.875, -0.075), (-0.925, -0.075), (0.075, 0.925)] ) @@ -120,7 +119,7 @@ def test__convergence_2d_via_hessian_from(): centre=(0.0, 0.0), ell_comps=(0.001, 0.001), einstein_radius=1.0 ) - convergence = mp.convergence_2d_via_hessian_from(grid=grid, buffer=buffer) + convergence = mp.convergence_2d_via_hessian_from(grid=grid) assert convergence.in_list[0] == pytest.approx(0.46208, 1.0e-1) assert convergence.in_list[1] == pytest.approx(0.56840, 1.0e-1) From b042da43516d82574655794746c29a9e424749f5 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 2 Mar 2026 12:09:24 +0000 Subject: [PATCH 2/7] replace jacobian path with hessian; restore jacobian_from as public API Remove precompute_jacobian decorator, jacobian_from (np.gradient), convergence_2d_via_jacobian_from, and shear_yx_2d_via_jacobian_from. All were redundant with the hessian path since A = I - H. Rewire tangential_eigen_value_from, radial_eigen_value_from, and magnification_2d_from to call hessian_from directly. Restore jacobian_from as a thin public wrapper over hessian_from that returns [[1-hxx, -hxy], [-hyx, 1-hyy]], supporting both xp=np and xp=jnp. Co-Authored-By: Claude Sonnet 4.6 --- autogalaxy/operate/deflections.py | 180 +++++--------------- test_autogalaxy/operate/test_deflections.py | 101 +++-------- 2 files changed, 65 insertions(+), 216 deletions(-) diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index 8fc5de15a..f358f72ca 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -7,7 +7,6 @@ import autoarray as aa -from autogalaxy.util.shear_field import ShearYX2D from autogalaxy.util.shear_field import ShearYX2DIrregular logger = logging.getLogger(__name__) @@ -34,17 +33,6 @@ def grid_scaled_2d_for_marching_squares_from( return aa.Grid2DIrregular(values=grid_scaled_1d) -def precompute_jacobian(func): - @wraps(func) - def wrapper(lensing_obj, grid, jacobian=None): - if jacobian is None: - jacobian = lensing_obj.jacobian_from(grid=grid) - - return func(lensing_obj, grid, jacobian) - - return wrapper - - def evaluation_grid(func): @wraps(func) def wrapper( @@ -181,8 +169,7 @@ def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: return aa.ArrayIrregular(values=fermat_potential) return aa.Array2D(values=fermat_potential, mask=grid.mask) - @precompute_jacobian - def tangential_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D: + def tangential_eigen_value_from(self, grid) -> aa.Array2D: """ Returns the tangential eigen values of lensing jacobian, which are given by the expression: @@ -193,54 +180,43 @@ def tangential_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D: grid The 2D grid of (y,x) arc-second coordinates the deflection angles and tangential eigen values are computed on. - jacobian - A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency. """ - convergence = self.convergence_2d_via_jacobian_from( - grid=grid, jacobian=jacobian - ) - - shear_yx = self.shear_yx_2d_via_jacobian_from(grid=grid, jacobian=jacobian) + convergence = self.convergence_2d_via_hessian_from(grid=grid) + shear_yx = self.shear_yx_2d_via_hessian_from(grid=grid) return aa.Array2D(values=1 - convergence - shear_yx.magnitudes, mask=grid.mask) - @precompute_jacobian - def radial_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D: + def radial_eigen_value_from(self, grid) -> aa.Array2D: """ Returns the radial eigen values of lensing jacobian, which are given by the expression: - radial_eigen_value = 1 - convergence + shear + `radial_eigen_value = 1 - convergence + shear` Parameters ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and radial eigen values are computed on. - jacobian - A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency. """ - convergence = self.convergence_2d_via_jacobian_from( - grid=grid, jacobian=jacobian - ) - - shear = self.shear_yx_2d_via_jacobian_from(grid=grid, jacobian=jacobian) + convergence = self.convergence_2d_via_hessian_from(grid=grid) + shear = self.shear_yx_2d_via_hessian_from(grid=grid) return aa.Array2D(values=1 - convergence + shear.magnitudes, mask=grid.mask) def magnification_2d_from(self, grid) -> aa.Array2D: """ Returns the 2D magnification map of lensing object, which is computed as the inverse of the determinant of the - jacobian. + lensing Jacobian, expressed via the Hessian components. Parameters ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and magnification map are computed on. """ - jacobian = self.jacobian_from(grid=grid) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) - det_jacobian = jacobian[0][0] * jacobian[1][1] - jacobian[0][1] * jacobian[1][0] + det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx - return aa.Array2D(values=1 / det_jacobian, mask=grid.mask) + return aa.Array2D(values=1 / det_A, mask=grid.mask) def deflections_yx_scalar(self, y, x, pixel_scales): """ @@ -359,6 +335,37 @@ def _hessian_via_finite_difference(self, grid, buffer: float = 0.01) -> Tuple: return hessian_yy, hessian_xy, hessian_yx, hessian_xx + def jacobian_from(self, grid, xp=np) -> List: + """ + Returns the lensing Jacobian of the lensing object as a 2x2 list of lists. + + The Jacobian is the matrix `A = I - H`, where `H` is the Hessian matrix of the + deflection angles: + + ``A = [[1 - hessian_xx, -hessian_xy], [-hessian_yx, 1 - hessian_yy]]`` + + It is computed from `hessian_from`, so it supports both uniform and irregular + grids and accepts the same `xp` parameter for JAX acceleration. + + Parameters + ---------- + grid + The 2D grid of (y,x) arc-second coordinates the Jacobian is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). Passed through to + ``hessian_from``. + """ + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) + + a11 = 1 - hessian_xx + a12 = -hessian_xy + a21 = -hessian_yx + a22 = 1 - hessian_yy + + return [[a11, a12], [a21, a22]] + def convergence_2d_via_hessian_from( self, grid ) -> aa.ArrayIrregular: @@ -833,106 +840,3 @@ def einstein_mass_angular_from( return einstein_mass_angular_list[0] - def jacobian_from(self, grid): - """ - Returns the Jacobian of the lensing object, which is computed by taking the gradient of the 2D deflection - angle map in four direction (positive y, negative y, positive x, negative x). - - By using the `np.gradient` method the Jacobian can therefore only be computed using uniform 2D grids of (y,x) - coordinates, and does not support irregular grids. For this reason, calculations by default use the Hessian, - which is slower to compute because more deflection angle calculations are necessary but more flexible in - general. - - The Jacobian is returned as a list of lists, which reflect its structure as a 2x2 matrix. - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and Jacobian are computed on. - """ - - deflections = self.deflections_yx_2d_from(grid=grid) - - # TODO : Can probably make this work on irregular grid? Is there any point? - - a11 = aa.Array2D( - values=1.0 - - np.gradient(deflections.native[:, :, 1], grid.native[0, :, 1], axis=1), - mask=grid.mask, - ) - - a12 = aa.Array2D( - values=-1.0 - * np.gradient(deflections.native[:, :, 1], grid.native[:, 0, 0], axis=0), - mask=grid.mask, - ) - - a21 = aa.Array2D( - values=-1.0 - * np.gradient(deflections.native[:, :, 0], grid.native[0, :, 1], axis=1), - mask=grid.mask, - ) - - a22 = aa.Array2D( - values=1 - - np.gradient(deflections.native[:, :, 0], grid.native[:, 0, 0], axis=0), - mask=grid.mask, - ) - - return [[a11, a12], [a21, a22]] - - @precompute_jacobian - def convergence_2d_via_jacobian_from(self, grid, jacobian=None) -> aa.Array2D: - """ - Returns the convergence of the lensing object, which is computed from the 2D deflection angle map via the - Jacobian using the expression (see equation 58 https://inspirehep.net/literature/419263): - - `convergence = 1.0 - 0.5 * (jacobian_{0,0} + jacobian_{1,1}) = 0.5 * (jacobian_xx + jacobian_yy)` - - By going via the Jacobian, the convergence must be calculated using 2D uniform grid. - - This calculation of the convergence is independent of analytic calculations defined within `MassProfile` - objects and the calculation via the Hessian. It can therefore be used as a cross-check. - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and Jacobian are computed on. - jacobian - A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency. - """ - convergence = 1 - 0.5 * (jacobian[0][0] + jacobian[1][1]) - - return aa.Array2D(values=convergence, mask=grid.mask) - - @precompute_jacobian - def shear_yx_2d_via_jacobian_from( - self, grid, jacobian=None - ) -> Union[ShearYX2D, ShearYX2DIrregular]: - """ - Returns the 2D (y,x) shear vectors of the lensing object, which are computed from the 2D deflection angle map - via the Jacobian using the expression (see equation 58 https://inspirehep.net/literature/419263): - - `shear_y = -0.5 * (jacobian_{0,1} + jacobian_{1,0} = -0.5 * (jacobian_yx + jacobian_xy)` - `shear_x = 0.5 * (jacobian_{1,1} + jacobian_{0,0} = 0.5 * (jacobian_yy + jacobian_xx)` - - By going via the Jacobian, the convergence must be calculated using 2D uniform grid. - - This calculation of the shear vectors is independent of analytic calculations defined within `MassProfile` - objects and the calculation via the Hessian. It can therefore be used as a cross-check. - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and Jacobian are computed on. - jacobian - A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency. - """ - - shear_yx_2d = np.zeros(shape=(grid.shape_slim, 2)) - shear_yx_2d[:, 0] = -0.5 * (jacobian[0][1] + jacobian[1][0]) - shear_yx_2d[:, 1] = 0.5 * (jacobian[1][1] - jacobian[0][0]) - - if isinstance(grid, aa.Grid2DIrregular): - return ShearYX2DIrregular(values=shear_yx_2d, grid=grid) - return ShearYX2D(values=shear_yx_2d, grid=grid, mask=grid.mask) diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 28bd6b90c..a10ebb6ed 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -139,51 +139,6 @@ def test__magnification_2d_via_hessian_from(): assert magnification.in_list[0] == pytest.approx(-0.56303, 1.0e-4) assert magnification.in_list[1] == pytest.approx(-2.57591, 1.0e-4) - -def test__magnification_2d_from__compare_eigen_values_and_determinant(): - grid = ag.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05) - - mp = ag.mp.Isothermal( - centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 - ) - - magnification_via_determinant = mp.magnification_2d_from(grid=grid) - tangential_eigen_value = mp.tangential_eigen_value_from(grid=grid) - - radal_eigen_value = mp.radial_eigen_value_from(grid=grid) - magnification_via_eigen_values = 1 / (tangential_eigen_value * radal_eigen_value) - - mean_error = np.mean( - magnification_via_determinant.slim - magnification_via_eigen_values.slim - ) - - assert mean_error < 1e-4 - - -def test__magnification_2d_from__compare_determinant_and_convergence_and_shear(): - grid = ag.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05) - - mp = ag.mp.Isothermal( - centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 - ) - - magnification_via_determinant = mp.magnification_2d_from(grid=grid) - - convergence = mp.convergence_2d_via_jacobian_from(grid=grid) - shear = mp.shear_yx_2d_via_jacobian_from(grid=grid) - - magnification_via_convergence_and_shear = 1 / ( - (1 - convergence) ** 2 - shear.magnitudes**2 - ) - - mean_error = np.mean( - magnification_via_determinant.slim - - magnification_via_convergence_and_shear.slim - ) - - assert mean_error < 1e-4 - - def test__tangential_critical_curve_list_from(): grid = ag.Grid2D.uniform(shape_native=(15, 15), pixel_scales=0.3) @@ -488,7 +443,15 @@ def test__einstein_mass_angular_from(): def test__jacobian_from(): - grid = ag.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05) + """ + The Jacobian is A = I - H, where H is the Hessian of the deflection angles. + + This test verifies the structure and values of `jacobian_from` by checking that: + - it returns a 2x2 list of lists; + - the convergence derived from its diagonal matches `convergence_2d_via_hessian_from`; + - the magnification derived from its determinant matches `magnification_2d_via_hessian_from`. + """ + grid = ag.Grid2DIrregular(values=[(1.0, 1.0), (2.0, 0.5)]) mp = ag.mp.Isothermal( centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 @@ -496,42 +459,24 @@ def test__jacobian_from(): jacobian = mp.jacobian_from(grid=grid) - A_12 = jacobian[0][1] - A_21 = jacobian[1][0] - - mean_error = np.mean(A_12.slim - A_21.slim) - - assert mean_error < 1e-4 - - -def test__convergence_2d_via_jacobian_from__compare_via_jacobian_and_analytic(): - grid = ag.Grid2D.uniform(shape_native=(20, 20), pixel_scales=0.05) + assert len(jacobian) == 2 + assert len(jacobian[0]) == 2 and len(jacobian[1]) == 2 - mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - - convergence_via_analytic = mp.convergence_2d_from(grid=grid) - - convergence_via_jacobian = mp.convergence_2d_via_jacobian_from(grid=grid) - - mean_error = np.mean(convergence_via_jacobian.slim - convergence_via_analytic.slim) - - assert convergence_via_jacobian.native.shape == (20, 20) - assert mean_error < 1e-1 + # convergence = 1 - 0.5 * (a11 + a22) should match convergence_2d_via_hessian_from + convergence_via_jacobian = 1 - 0.5 * (jacobian[0][0] + jacobian[1][1]) + convergence_via_hessian = mp.convergence_2d_via_hessian_from(grid=grid) - mean_error = np.mean(convergence_via_jacobian.slim - convergence_via_analytic.slim) - - assert mean_error < 1e-1 - - grid = ag.Grid2D.uniform(shape_native=(20, 20), pixel_scales=0.05) - - mp = ag.mp.Isothermal( - centre=(0.0, 0.0), ell_comps=(0.111111, 0.0), einstein_radius=2.0 + assert convergence_via_jacobian == pytest.approx( + np.array(convergence_via_hessian), rel=1e-6 ) - convergence_via_analytic = mp.convergence_2d_from(grid=grid) + # magnification = 1 / det(A) = 1 / (a11*a22 - a12*a21) + det_A = jacobian[0][0] * jacobian[1][1] - jacobian[0][1] * jacobian[1][0] + magnification_via_jacobian = 1 / det_A + magnification_via_hessian = mp.magnification_2d_via_hessian_from(grid=grid) - convergence_via_jacobian = mp.convergence_2d_via_jacobian_from(grid=grid) + assert magnification_via_jacobian == pytest.approx( + np.array(magnification_via_hessian), rel=1e-6 + ) - mean_error = np.mean(convergence_via_jacobian.slim - convergence_via_analytic.slim) - assert mean_error < 1e-1 From 8e2229c6338f364ab027356e054ac2378f640937 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 2 Mar 2026 14:46:23 +0000 Subject: [PATCH 3/7] midway through refactor and now responding to github review --- autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py index cae42f239..b3f1e89d7 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py @@ -1,8 +1,7 @@ +import numpy as np from typing import Tuple from autogalaxy.profiles.mass.dark.gnfw import gNFWSph - -import numpy as np from autogalaxy import cosmology as cosmo From eac4e9e14c22f941d677058af830586a990e72c8 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 2 Mar 2026 17:41:29 +0000 Subject: [PATCH 4/7] Add xp parameter to convergence, shear, eigenvalue, and magnification_2d hessian methods Thread xp=np through convergence_2d_via_hessian_from, shear_yx_2d_via_hessian_from, tangential_eigen_value_from, radial_eigen_value_from, and magnification_2d_from so all hessian-derived quantities support the JAX path consistently. For each method, xp is passed through to hessian_from; when xp is not numpy the result is returned as a raw jax.Array rather than an autoarray wrapper (which cannot be constructed during a jax.jit trace). Co-Authored-By: Claude Sonnet 4.6 --- autogalaxy/operate/deflections.py | 66 +++++++++++++++++++++++++------ 1 file changed, 53 insertions(+), 13 deletions(-) diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index f358f72ca..ce185c060 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -169,7 +169,7 @@ def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: return aa.ArrayIrregular(values=fermat_potential) return aa.Array2D(values=fermat_potential, mask=grid.mask) - def tangential_eigen_value_from(self, grid) -> aa.Array2D: + def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D: """ Returns the tangential eigen values of lensing jacobian, which are given by the expression: @@ -180,13 +180,20 @@ def tangential_eigen_value_from(self, grid) -> aa.Array2D: grid The 2D grid of (y,x) arc-second coordinates the deflection angles and tangential eigen values are computed on. + xp + The array module (``numpy`` or ``jax.numpy``). Passed through to ``convergence_2d_via_hessian_from`` + and ``shear_yx_2d_via_hessian_from``. When ``xp`` is not ``numpy`` the result is a raw array rather + than an ``aa.Array2D`` wrapper. """ - convergence = self.convergence_2d_via_hessian_from(grid=grid) - shear_yx = self.shear_yx_2d_via_hessian_from(grid=grid) + convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp) + shear_yx = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp) + if xp is not np: + shear_magnitudes = xp.sqrt(shear_yx[:, 0] ** 2 + shear_yx[:, 1] ** 2) + return xp.array(1 - convergence - shear_magnitudes) return aa.Array2D(values=1 - convergence - shear_yx.magnitudes, mask=grid.mask) - def radial_eigen_value_from(self, grid) -> aa.Array2D: + def radial_eigen_value_from(self, grid, xp=np) -> aa.Array2D: """ Returns the radial eigen values of lensing jacobian, which are given by the expression: @@ -196,13 +203,20 @@ def radial_eigen_value_from(self, grid) -> aa.Array2D: ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and radial eigen values are computed on. + xp + The array module (``numpy`` or ``jax.numpy``). Passed through to ``convergence_2d_via_hessian_from`` + and ``shear_yx_2d_via_hessian_from``. When ``xp`` is not ``numpy`` the result is a raw array rather + than an ``aa.Array2D`` wrapper. """ - convergence = self.convergence_2d_via_hessian_from(grid=grid) - shear = self.shear_yx_2d_via_hessian_from(grid=grid) + convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp) + shear = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp) + if xp is not np: + shear_magnitudes = xp.sqrt(shear[:, 0] ** 2 + shear[:, 1] ** 2) + return xp.array(1 - convergence + shear_magnitudes) return aa.Array2D(values=1 - convergence + shear.magnitudes, mask=grid.mask) - def magnification_2d_from(self, grid) -> aa.Array2D: + def magnification_2d_from(self, grid, xp=np) -> aa.Array2D: """ Returns the 2D magnification map of lensing object, which is computed as the inverse of the determinant of the lensing Jacobian, expressed via the Hessian components. @@ -211,11 +225,18 @@ def magnification_2d_from(self, grid) -> aa.Array2D: ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and magnification map are computed on. + xp + The array module (``numpy`` or ``jax.numpy``). Passed through to ``hessian_from``. When ``xp`` is + not ``numpy`` the result is a raw array rather than an ``aa.Array2D`` wrapper. """ - hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx + if xp is not np: + return xp.array(1 / det_A) return aa.Array2D(values=1 / det_A, mask=grid.mask) def deflections_yx_scalar(self, y, x, pixel_scales): @@ -367,7 +388,7 @@ def jacobian_from(self, grid, xp=np) -> List: return [[a11, a12], [a21, a22]] def convergence_2d_via_hessian_from( - self, grid + self, grid, xp=np ) -> aa.ArrayIrregular: """ Returns the convergence of the lensing object, which is computed from the 2D deflection angle map via the @@ -385,13 +406,23 @@ def convergence_2d_via_hessian_from( ---------- grid The 2D grid of (y,x) arc-second coordinates the deflection angles and Hessian are computed on. + xp + The array module to use for the computation (e.g. `numpy` or `jax.numpy`). Passed through to + `hessian_from`. When `xp` is not `numpy` (e.g. inside a `jax.jit` trace) the result is returned + as a raw array rather than an `aa.ArrayIrregular` wrapper. """ - hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) + + convergence = 0.5 * (hessian_yy + hessian_xx) - return aa.ArrayIrregular(values=0.5 * (hessian_yy + hessian_xx)) + if xp is not np: + return xp.array(convergence) + return aa.ArrayIrregular(values=convergence) def shear_yx_2d_via_hessian_from( - self, grid + self, grid, xp=np ) -> ShearYX2DIrregular: """ Returns the 2D (y,x) shear vectors of the lensing object, which are computed from the 2D deflection angle map @@ -416,13 +447,22 @@ def shear_yx_2d_via_hessian_from( ---------- grids The 2D grid of (y,x) arc-second coordinates the deflection angles and Hessian are computed on. + xp + The array module to use for the computation (e.g. `numpy` or `jax.numpy`). Passed through to + `hessian_from`. When `xp` is not `numpy` (e.g. inside a `jax.jit` trace) the result is returned + as a raw array of shape `(N, 2)` rather than a `ShearYX2DIrregular` wrapper. """ - hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from(grid=grid) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = self.hessian_from( + grid=grid, xp=xp + ) gamma_1 = 0.5 * (hessian_xx - hessian_yy) gamma_2 = hessian_xy + if xp is not np: + return xp.stack([gamma_2, gamma_1], axis=-1) + shear_yx_2d = np.zeros(shape=(grid.shape[0], 2)) shear_yx_2d[:, 0] = gamma_2 From c838c6a3cf32df934eae876b6f6729944a6aed73 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 3 Mar 2026 10:10:23 +0000 Subject: [PATCH 5/7] all unit tests pass --- autogalaxy/galaxy/galaxies.py | 31 +++- autogalaxy/galaxy/galaxy.py | 31 +++- autogalaxy/operate/deflections.py | 137 ++++++++---------- autogalaxy/plot/mass_plotter.py | 6 +- autogalaxy/plot/visuals/one_d.py | 10 +- autogalaxy/plot/visuals/two_d.py | 18 ++- autogalaxy/profiles/mass/abstract/abstract.py | 31 +++- autogalaxy/profiles/mass/dark/mcr_util.py | 1 + test_autogalaxy/operate/test_deflections.py | 90 ++++++++---- test_autogalaxy/plot/mat_wrap/test_visuals.py | 18 ++- 10 files changed, 246 insertions(+), 127 deletions(-) diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py index 0544899b5..0135b7357 100644 --- a/autogalaxy/galaxy/galaxies.py +++ b/autogalaxy/galaxy/galaxies.py @@ -8,10 +8,9 @@ from autogalaxy.profiles.basis import Basis from autogalaxy.profiles.light.linear import LightProfileLinear from autogalaxy.operate.image import OperateImageGalaxies -from autogalaxy.operate.deflections import OperateDeflections -class Galaxies(List, OperateImageGalaxies, OperateDeflections): +class Galaxies(List, OperateImageGalaxies): def __init__( self, galaxies: List[Galaxy], @@ -244,6 +243,34 @@ def potential_2d_from( """ return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self)) + def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: + """ + Returns the Fermat potential of the galaxies for a given grid of image-plane positions. + + This is the sum of the geometric time delay term and the lensing potential, computed + using the combined deflections and potential of all galaxies: + + .. math:: + \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2 + - \\psi(\\boldsymbol{\\theta}) + + Parameters + ---------- + grid + The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). + """ + from autogalaxy.operate.deflections import OperateDeflections + + od = OperateDeflections.from_mass_obj(self) + time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) + potential = self.potential_2d_from(grid=grid, xp=xp) + fermat_potential = time_delay - potential + if isinstance(grid, aa.Grid2DIrregular): + return aa.ArrayIrregular(values=fermat_potential) + return aa.Array2D(values=fermat_potential, mask=grid.mask) + def has(self, cls: Union[Type, Tuple[Type]]) -> bool: """ Returns a bool specifying whether any of the galaxies has a certain class type. diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py index 23e47f7f1..c0e972d96 100644 --- a/autogalaxy/galaxy/galaxy.py +++ b/autogalaxy/galaxy/galaxy.py @@ -8,7 +8,6 @@ import autofit as af from autogalaxy import exc -from autogalaxy.operate.deflections import OperateDeflections from autogalaxy.operate.image import OperateImageList from autogalaxy.profiles.geometry_profiles import GeometryProfile from autogalaxy.profiles.light.abstract import LightProfile @@ -17,7 +16,7 @@ from autogalaxy.profiles.mass.abstract.abstract import MassProfile -class Galaxy(af.ModelObject, OperateImageList, OperateDeflections): +class Galaxy(af.ModelObject, OperateImageList): """ @DynamicAttrs """ @@ -344,6 +343,34 @@ def potential_2d_from( ) return xp.zeros((grid.shape[0],)) + def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: + """ + Returns the Fermat potential of the galaxy for a given grid of image-plane positions. + + This is the sum of the geometric time delay term and the lensing potential, computed + using the galaxy's total deflections and total potential across all mass profiles: + + .. math:: + \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2 + - \\psi(\\boldsymbol{\\theta}) + + Parameters + ---------- + grid + The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). + """ + from autogalaxy.operate.deflections import OperateDeflections + + od = OperateDeflections.from_mass_obj(self) + time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) + potential = self.potential_2d_from(grid=grid, xp=xp) + fermat_potential = time_delay - potential + if isinstance(grid, aa.Grid2DIrregular): + return aa.ArrayIrregular(values=fermat_potential) + return aa.Array2D(values=fermat_potential, mask=grid.mask) + @property def half_light_radius(self): return None diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index ce185c060..870955a58 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -86,88 +86,94 @@ class OperateDeflections: The majority of methods are those which from the 2D deflection angle map compute lensing quantities like a 2D shear field, magnification map or the Einstein Radius. - The methods in `CalcLens` are passed to the mass object to provide a concise API. - Parameters ---------- deflections_yx_2d_from - The function which returns the mass object's 2D deflection angles. + A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection angles on the given + grid. Typically a bound method of a ``MassProfile``, ``Galaxy``, or ``Galaxies`` instance. """ - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - raise NotImplementedError + def __init__(self, deflections_yx_2d_from): + self.deflections_yx_2d_from = deflections_yx_2d_from - def __eq__(self, other): - return self.__dict__ == other.__dict__ and self.__class__ is other.__class__ + @classmethod + def from_mass_obj(cls, mass_obj): + """Construct from any object that has a ``deflections_yx_2d_from`` method.""" + return cls(deflections_yx_2d_from=mass_obj.deflections_yx_2d_from) - def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D: + @classmethod + def from_tracer(cls, tracer, use_multi_plane: bool = True, plane_i: int = 0, plane_j: int = -1): """ - Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions. - - This term is given by: - - .. math:: - \[\tau_{\text{geom}}(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2\] + Construct from a PyAutoLens ``Tracer`` object. - where: - - \( \boldsymbol{\theta} \) is the image-plane coordinate, - - \( \boldsymbol{\beta} = \boldsymbol{\theta} - \boldsymbol{\alpha}(\boldsymbol{\theta}) \) is the source-plane coordinate, - - \( \boldsymbol{\alpha} \) is the deflection angle at each image-plane coordinate. - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay geometric term are computed - on. - - Returns - ------- - The geometric time delay term at each grid position. - """ - deflections = self.deflections_yx_2d_from(grid=grid, xp=xp) + The ``Tracer`` type is deliberately left unannotated: ``autogalaxy`` does not + depend on ``autolens``, so no import of ``Tracer`` is performed here. Callers + (which live inside ``autolens``) are responsible for passing the correct object. - src_y = grid[:, 0] - deflections[:, 0] - src_x = grid[:, 1] - deflections[:, 1] - - delay = 0.5 * ((grid[:, 0] - src_y) ** 2 + (grid[:, 1] - src_x) ** 2) - - if isinstance(grid, aa.Grid2DIrregular): - return aa.ArrayIrregular(values=delay) - return aa.Array2D(values=delay, mask=grid.mask) + Parameters + ---------- + tracer + A PyAutoLens ``Tracer`` instance. Must expose ``deflections_yx_2d_from`` + and, when ``use_multi_plane=True``, ``deflections_between_planes_from``. + use_multi_plane + If ``True`` the stored callable is + ``tracer.deflections_between_planes_from`` with ``plane_i`` and ``plane_j`` + bound via ``functools.partial``, matching the multi-plane ray-tracing path. + If ``False`` the stored callable is ``tracer.deflections_yx_2d_from``, + i.e. the standard two-plane path. + plane_i + Index of the first plane used by ``deflections_between_planes_from``. + Ignored when ``use_multi_plane=False``. Defaults to ``0`` (image plane). + plane_j + Index of the second plane used by ``deflections_between_planes_from``. + Ignored when ``use_multi_plane=False``. Defaults to ``-1`` (source plane). + """ + if use_multi_plane: + from functools import partial + + return cls( + deflections_yx_2d_from=partial( + tracer.deflections_between_planes_from, + plane_i=plane_i, + plane_j=plane_j, + ) + ) + return cls(deflections_yx_2d_from=tracer.deflections_yx_2d_from) - def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: + def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D: """ - Returns the Fermat potential for a given grid of image-plane positions. + Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions. - This is the sum of the geometric time delay term and the gravitational (Shapiro) delay term (i.e. the lensing - potential), and is given by: + This term is given by: .. math:: - \[\phi(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2 - \psi(\boldsymbol{\theta})\] + \[\tau_{\text{geom}}(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2\] where: - \( \boldsymbol{\theta} \) is the image-plane coordinate, - \( \boldsymbol{\beta} = \boldsymbol{\theta} - \boldsymbol{\alpha}(\boldsymbol{\theta}) \) is the source-plane coordinate, - - \( \psi(\boldsymbol{\theta}) \) is the lensing potential, - - \( \phi(\boldsymbol{\theta}) \) is the Fermat potential. + - \( \boldsymbol{\alpha} \) is the deflection angle at each image-plane coordinate. Parameters ---------- grid - The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. + The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay geometric term are computed + on. Returns ------- - The Fermat potential at each grid position. + The geometric time delay term at each grid position. """ - time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid, xp=xp) - potential = self.potential_2d_from(grid=grid, xp=xp) + deflections = self.deflections_yx_2d_from(grid=grid, xp=xp) + + src_y = grid[:, 0] - deflections[:, 0] + src_x = grid[:, 1] - deflections[:, 1] - fermat_potential = time_delay_geometry_term - potential + delay = 0.5 * ((grid[:, 0] - src_y) ** 2 + (grid[:, 1] - src_x) ** 2) if isinstance(grid, aa.Grid2DIrregular): - return aa.ArrayIrregular(values=fermat_potential) - return aa.Array2D(values=fermat_potential, mask=grid.mask) + return aa.ArrayIrregular(values=delay) + return aa.Array2D(values=delay, mask=grid.mask) def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D: """ @@ -188,9 +194,6 @@ def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D: convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp) shear_yx = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp) - if xp is not np: - shear_magnitudes = xp.sqrt(shear_yx[:, 0] ** 2 + shear_yx[:, 1] ** 2) - return xp.array(1 - convergence - shear_magnitudes) return aa.Array2D(values=1 - convergence - shear_yx.magnitudes, mask=grid.mask) def radial_eigen_value_from(self, grid, xp=np) -> aa.Array2D: @@ -211,9 +214,6 @@ def radial_eigen_value_from(self, grid, xp=np) -> aa.Array2D: convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp) shear = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp) - if xp is not np: - shear_magnitudes = xp.sqrt(shear[:, 0] ** 2 + shear[:, 1] ** 2) - return xp.array(1 - convergence + shear_magnitudes) return aa.Array2D(values=1 - convergence + shear.magnitudes, mask=grid.mask) def magnification_2d_from(self, grid, xp=np) -> aa.Array2D: @@ -235,8 +235,6 @@ def magnification_2d_from(self, grid, xp=np) -> aa.Array2D: det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx - if xp is not np: - return xp.array(1 / det_A) return aa.Array2D(values=1 / det_A, mask=grid.mask) def deflections_yx_scalar(self, y, x, pixel_scales): @@ -295,10 +293,9 @@ def hessian_from(self, grid, xp=np) -> Tuple: The array module (``numpy`` or ``jax.numpy``). Controls which computational path is used and the type of the returned arrays. """ - if xp is not np: - return self._hessian_via_jax(grid=grid, xp=xp) - - return self._hessian_via_finite_difference(grid=grid) + if xp is np: + return self._hessian_via_finite_difference(grid=grid) + return self._hessian_via_jax(grid=grid, xp=xp) def _hessian_via_jax(self, grid, xp) -> Tuple: import jax @@ -417,8 +414,6 @@ def convergence_2d_via_hessian_from( convergence = 0.5 * (hessian_yy + hessian_xx) - if xp is not np: - return xp.array(convergence) return aa.ArrayIrregular(values=convergence) def shear_yx_2d_via_hessian_from( @@ -460,13 +455,7 @@ def shear_yx_2d_via_hessian_from( gamma_1 = 0.5 * (hessian_xx - hessian_yy) gamma_2 = hessian_xy - if xp is not np: - return xp.stack([gamma_2, gamma_1], axis=-1) - - shear_yx_2d = np.zeros(shape=(grid.shape[0], 2)) - - shear_yx_2d[:, 0] = gamma_2 - shear_yx_2d[:, 1] = gamma_1 + shear_yx_2d = xp.stack([gamma_2, gamma_1], axis=-1) return ShearYX2DIrregular(values=shear_yx_2d, grid=grid) @@ -498,8 +487,6 @@ def magnification_2d_via_hessian_from( det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx - if xp is not np: - return xp.array(1.0 / det_A) return aa.ArrayIrregular(values=1.0 / det_A) def contour_list_from(self, grid, contour_array): diff --git a/autogalaxy/plot/mass_plotter.py b/autogalaxy/plot/mass_plotter.py index c6bb69268..b307ca48d 100644 --- a/autogalaxy/plot/mass_plotter.py +++ b/autogalaxy/plot/mass_plotter.py @@ -118,8 +118,12 @@ def figures_2d( ) if magnification: + from autogalaxy.operate.deflections import OperateDeflections + self.mat_plot_2d.plot_array( - array=self.mass_obj.magnification_2d_from(grid=self.grid), + array=OperateDeflections.from_mass_obj(self.mass_obj).magnification_2d_from( + grid=self.grid + ), visuals_2d=self.visuals_2d_with_critical_curves, auto_labels=aplt.AutoLabels( title=f"Magnification{title_suffix}", diff --git a/autogalaxy/plot/visuals/one_d.py b/autogalaxy/plot/visuals/one_d.py index 7d69d112a..a610fe0e3 100644 --- a/autogalaxy/plot/visuals/one_d.py +++ b/autogalaxy/plot/visuals/one_d.py @@ -170,10 +170,13 @@ def add_einstein_radius( The collection of attributes that can be plotted by a `Plotter` object. """ + from autogalaxy.operate.deflections import OperateDeflections + einstein_radius = None try: - einstein_radius = mass_obj.einstein_radius_from(grid=grid) + od = OperateDeflections.from_mass_obj(mass_obj) + einstein_radius = od.einstein_radius_from(grid=grid) except (TypeError, AttributeError): pass @@ -214,11 +217,14 @@ def add_einstein_radius_errors( The mean value and errors of each attribute that are plotted in 1D by a `Plotter` object. """ + from autogalaxy.operate.deflections import OperateDeflections + einstein_radius_list = [] for mass_obj in mass_obj_list: try: - einstein_radius_list.append(mass_obj.einstein_radius_from(grid=grid)) + od = OperateDeflections.from_mass_obj(mass_obj) + einstein_radius_list.append(od.einstein_radius_from(grid=grid)) except TypeError: einstein_radius_list.append(None) diff --git a/autogalaxy/plot/visuals/two_d.py b/autogalaxy/plot/visuals/two_d.py index 759627897..3bd74e90b 100644 --- a/autogalaxy/plot/visuals/two_d.py +++ b/autogalaxy/plot/visuals/two_d.py @@ -184,18 +184,19 @@ def add_critical_curves(self, mass_obj, grid: aa.type.Grid2DLike): vis.Visuals2D A collection of attributes that can be plotted by a `Plotter` object. """ + from autogalaxy.operate.deflections import OperateDeflections - tangential_critical_curves = mass_obj.tangential_critical_curve_list_from( - grid=grid - ) + od = OperateDeflections.from_mass_obj(mass_obj) + + tangential_critical_curves = od.tangential_critical_curve_list_from(grid=grid) radial_critical_curves = None - radial_critical_curve_area_list = mass_obj.radial_critical_curve_area_list_from( + radial_critical_curve_area_list = od.radial_critical_curve_area_list_from( grid=grid ) if any([area > grid.pixel_scale for area in radial_critical_curve_area_list]): - radial_critical_curves = mass_obj.radial_critical_curve_list_from(grid=grid) + radial_critical_curves = od.radial_critical_curve_list_from(grid=grid) return self + self.__class__( tangential_critical_curves=tangential_critical_curves, @@ -226,9 +227,12 @@ def add_caustics(self, mass_obj, grid: aa.type.Grid2DLike): vis.Visuals2D A collection of attributes that can be plotted by a `Plotter` object. """ + from autogalaxy.operate.deflections import OperateDeflections + + od = OperateDeflections.from_mass_obj(mass_obj) - tangential_caustics = mass_obj.tangential_caustic_list_from(grid=grid) - radial_caustics = mass_obj.radial_caustic_list_from(grid=grid) + tangential_caustics = od.tangential_caustic_list_from(grid=grid) + radial_caustics = od.radial_caustic_list_from(grid=grid) return self + self.__class__( tangential_caustics=tangential_caustics, diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py index 5ff8f0d31..3fb130688 100644 --- a/autogalaxy/profiles/mass/abstract/abstract.py +++ b/autogalaxy/profiles/mass/abstract/abstract.py @@ -4,10 +4,9 @@ import autoarray as aa from autogalaxy.profiles.geometry_profiles import EllProfile -from autogalaxy.operate.deflections import OperateDeflections -class MassProfile(EllProfile, OperateDeflections): +class MassProfile(EllProfile): def __init__( self, centre: Tuple[float, float] = (0.0, 0.0), @@ -28,6 +27,34 @@ def __init__( def deflections_yx_2d_from(self, grid): raise NotImplementedError + def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: + """ + Returns the Fermat potential for a given grid of image-plane positions. + + This is the sum of the geometric time delay term and the gravitational (Shapiro) delay term (i.e. the lensing + potential), and is given by: + + .. math:: + \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\theta} - \\boldsymbol{\\beta}|^2 + - \\psi(\\boldsymbol{\\theta}) + + Parameters + ---------- + grid + The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). + """ + from autogalaxy.operate.deflections import OperateDeflections + + od = OperateDeflections.from_mass_obj(self) + time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) + potential = self.potential_2d_from(grid=grid, xp=xp) + fermat_potential = time_delay - potential + if isinstance(grid, aa.Grid2DIrregular): + return aa.ArrayIrregular(values=fermat_potential) + return aa.Array2D(values=fermat_potential, mask=grid.mask) + def deflections_2d_via_potential_2d_from(self, grid): potential = self.potential_2d_from(grid=grid) diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index 891b93d1b..ce6f15c95 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -133,6 +133,7 @@ def ludlow16_cosmology_jax( mass_at_200, redshift_object, redshift_source, + vmap_method="sequential" ) diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index a10ebb6ed..04ef8a2b8 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -5,11 +5,16 @@ import autogalaxy as ag -from autogalaxy.operate.deflections import grid_scaled_2d_for_marching_squares_from +from autogalaxy.operate.deflections import ( + grid_scaled_2d_for_marching_squares_from, + OperateDeflections, +) def critical_curve_via_magnification_from(mass_profile, grid): - magnification = mass_profile.magnification_2d_from(grid=grid) + magnification = OperateDeflections.from_mass_obj(mass_profile).magnification_2d_from( + grid=grid + ) inverse_magnification = 1 / magnification @@ -64,7 +69,8 @@ def test__time_delay_geometry_term_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - time_delay_geometry_term = mp.time_delay_geometry_term_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + time_delay_geometry_term = od.time_delay_geometry_term_from(grid=grid) assert time_delay_geometry_term == pytest.approx( np.array([1.92815688, 1.97625436]), 1.0e-4 @@ -93,7 +99,8 @@ def test__hessian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - hessian_yy, hessian_xy, hessian_yx, hessian_xx = mp.hessian_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = od.hessian_from(grid=grid) assert hessian_yy == pytest.approx(np.array([1.3883822, 0.694127]), 1.0e-4) assert hessian_xy == pytest.approx(np.array([-1.388124, -0.694094]), 1.0e-4) @@ -102,7 +109,7 @@ def test__hessian_from(): grid = ag.Grid2DIrregular(values=[(1.0, 0.0), (0.0, 1.0)]) - hessian_yy, hessian_xy, hessian_yx, hessian_xx = mp.hessian_from(grid=grid) + hessian_yy, hessian_xy, hessian_yx, hessian_xx = od.hessian_from(grid=grid) assert hessian_yy == pytest.approx(np.array([0.0, 1.777699]), 1.0e-4) assert hessian_xy == pytest.approx(np.array([0.0, 0.0]), 1.0e-4) @@ -119,7 +126,8 @@ def test__convergence_2d_via_hessian_from(): centre=(0.0, 0.0), ell_comps=(0.001, 0.001), einstein_radius=1.0 ) - convergence = mp.convergence_2d_via_hessian_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + convergence = od.convergence_2d_via_hessian_from(grid=grid) assert convergence.in_list[0] == pytest.approx(0.46208, 1.0e-1) assert convergence.in_list[1] == pytest.approx(0.56840, 1.0e-1) @@ -134,7 +142,8 @@ def test__magnification_2d_via_hessian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - magnification = mp.magnification_2d_via_hessian_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + magnification = od.magnification_2d_via_hessian_from(grid=grid) assert magnification.in_list[0] == pytest.approx(-0.56303, 1.0e-4) assert magnification.in_list[1] == pytest.approx(-2.57591, 1.0e-4) @@ -144,7 +153,8 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - tangential_critical_curve_list = mp.tangential_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) x_critical_tangential, y_critical_tangential = ( tangential_critical_curve_list[0][:, 1], @@ -159,7 +169,8 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - tangential_critical_curve_list = mp.tangential_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) y_centre = np.mean(tangential_critical_curve_list[0][:, 0]) x_centre = np.mean(tangential_critical_curve_list[0][:, 1]) @@ -169,7 +180,8 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.5, 1.0), einstein_radius=2.0) - tangential_critical_curve_list = mp.tangential_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) y_centre = np.mean(tangential_critical_curve_list[0][:, 0]) x_centre = np.mean(tangential_critical_curve_list[0][:, 1]) @@ -217,7 +229,8 @@ def test__radial_critical_curve_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - radial_critical_curve_list = mp.radial_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) y_centre = np.mean(radial_critical_curve_list[0][:, 0]) x_centre = np.mean(radial_critical_curve_list[0][:, 1]) @@ -227,7 +240,8 @@ def test__radial_critical_curve_list_from(): mp = ag.mp.PowerLawSph(centre=(0.5, 1.0), einstein_radius=2.0, slope=1.5) - radial_critical_curve_list = mp.radial_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) y_centre = np.mean(radial_critical_curve_list[0][:, 0]) x_centre = np.mean(radial_critical_curve_list[0][:, 1]) @@ -248,7 +262,8 @@ def test__radial_critical_curve_list_from__compare_via_magnification(): mass_profile=mp, grid=grid )[1] - radial_critical_curve_list = mp.radial_critical_curve_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) assert sum(critical_curve_radial_via_magnification) == pytest.approx( sum(radial_critical_curve_list[0]), abs=0.7 @@ -260,7 +275,8 @@ def test__tangential_caustic_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - tangential_caustic_list = mp.tangential_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + tangential_caustic_list = od.tangential_caustic_list_from(grid=grid) y_centre = np.mean(tangential_caustic_list[0][:, 0]) x_centre = np.mean(tangential_caustic_list[0][:, 1]) @@ -270,7 +286,8 @@ def test__tangential_caustic_list_from(): mp = ag.mp.IsothermalSph(centre=(0.5, 1.0), einstein_radius=2.0) - tangential_caustic_list = mp.tangential_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + tangential_caustic_list = od.tangential_caustic_list_from(grid=grid) y_centre = np.mean(tangential_caustic_list[0][:, 0]) x_centre = np.mean(tangential_caustic_list[0][:, 1]) @@ -306,7 +323,8 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - radial_caustic_list = mp.radial_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_caustic_list = od.radial_caustic_list_from(grid=grid) x_caustic_radial, y_caustic_radial = ( radial_caustic_list[0][:, 1], @@ -321,7 +339,8 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - radial_caustic_list = mp.radial_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_caustic_list = od.radial_caustic_list_from(grid=grid) y_centre = np.mean(radial_caustic_list[0][:, 0]) x_centre = np.mean(radial_caustic_list[0][:, 1]) @@ -331,7 +350,8 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.5, 1.0), einstein_radius=2.0, slope=1.5) - radial_caustic_list = mp.radial_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_caustic_list = od.radial_caustic_list_from(grid=grid) y_centre = np.mean(radial_caustic_list[0][:, 0]) x_centre = np.mean(radial_caustic_list[0][:, 1]) @@ -351,7 +371,8 @@ def test__radial_caustic_list_from___compare_via_magnification(): mass_profile=mp, grid=grid )[1] - radial_caustic_list = mp.radial_caustic_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + radial_caustic_list = od.radial_caustic_list_from(grid=grid) assert sum(radial_caustic_list[0]) == pytest.approx( sum(caustic_radial_via_magnification), 7e-1 @@ -363,7 +384,8 @@ def test__radial_critical_curve_area_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - area_within_radial_critical_curve_list = mp.radial_critical_curve_area_list_from( + od = OperateDeflections.from_mass_obj(mp) + area_within_radial_critical_curve_list = od.radial_critical_curve_area_list_from( grid=grid ) @@ -377,8 +399,9 @@ def test__tangential_critical_curve_area_list_from(): area_calc = np.pi * mp.einstein_radius**2 + od = OperateDeflections.from_mass_obj(mp) area_within_tangential_critical_curve_list = ( - mp.tangential_critical_curve_area_list_from(grid=grid) + od.tangential_critical_curve_area_list_from(grid=grid) ) assert area_within_tangential_critical_curve_list[0] == pytest.approx( @@ -391,7 +414,8 @@ def test__einstein_radius_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - einstein_radius_list = mp.einstein_radius_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_radius_list = od.einstein_radius_list_from(grid=grid) assert einstein_radius_list[0] == pytest.approx(2.0, 1e-1) @@ -399,7 +423,8 @@ def test__einstein_radius_list_from(): centre=(0.0, 0.0), einstein_radius=2.0, ell_comps=(0.0, -0.25) ) - einstein_radius_list = mp.einstein_radius_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_radius_list = od.einstein_radius_list_from(grid=grid) assert einstein_radius_list[0] == pytest.approx(1.9360, 1e-1) @@ -409,7 +434,8 @@ def test__einstein_radius_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - einstein_radius = mp.einstein_radius_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_radius = od.einstein_radius_from(grid=grid) assert einstein_radius == pytest.approx(2.0, 1e-1) @@ -417,7 +443,8 @@ def test__einstein_radius_from(): centre=(0.0, 0.0), einstein_radius=2.0, ell_comps=(0.0, -0.25) ) - einstein_radius = mp.einstein_radius_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_radius = od.einstein_radius_from(grid=grid) assert einstein_radius == pytest.approx(1.9360, 1e-1) @@ -427,7 +454,8 @@ def test__einstein_mass_angular_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - einstein_mass_angular_list = mp.einstein_mass_angular_list_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_mass_angular_list = od.einstein_mass_angular_list_from(grid=grid) assert einstein_mass_angular_list[0] == pytest.approx(np.pi * 2.0**2.0, 1e-1) @@ -437,7 +465,8 @@ def test__einstein_mass_angular_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - einstein_mass_angular = mp.einstein_mass_angular_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + einstein_mass_angular = od.einstein_mass_angular_from(grid=grid) assert einstein_mass_angular == pytest.approx(np.pi * 2.0**2.0, 1e-1) @@ -457,14 +486,15 @@ def test__jacobian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - jacobian = mp.jacobian_from(grid=grid) + od = OperateDeflections.from_mass_obj(mp) + jacobian = od.jacobian_from(grid=grid) assert len(jacobian) == 2 assert len(jacobian[0]) == 2 and len(jacobian[1]) == 2 # convergence = 1 - 0.5 * (a11 + a22) should match convergence_2d_via_hessian_from convergence_via_jacobian = 1 - 0.5 * (jacobian[0][0] + jacobian[1][1]) - convergence_via_hessian = mp.convergence_2d_via_hessian_from(grid=grid) + convergence_via_hessian = od.convergence_2d_via_hessian_from(grid=grid) assert convergence_via_jacobian == pytest.approx( np.array(convergence_via_hessian), rel=1e-6 @@ -473,7 +503,7 @@ def test__jacobian_from(): # magnification = 1 / det(A) = 1 / (a11*a22 - a12*a21) det_A = jacobian[0][0] * jacobian[1][1] - jacobian[0][1] * jacobian[1][0] magnification_via_jacobian = 1 / det_A - magnification_via_hessian = mp.magnification_2d_via_hessian_from(grid=grid) + magnification_via_hessian = od.magnification_2d_via_hessian_from(grid=grid) assert magnification_via_jacobian == pytest.approx( np.array(magnification_via_hessian), rel=1e-6 diff --git a/test_autogalaxy/plot/mat_wrap/test_visuals.py b/test_autogalaxy/plot/mat_wrap/test_visuals.py index 3ae7058ff..a3cdfb888 100644 --- a/test_autogalaxy/plot/mat_wrap/test_visuals.py +++ b/test_autogalaxy/plot/mat_wrap/test_visuals.py @@ -2,6 +2,7 @@ import pytest import autogalaxy.plot as aplt +from autogalaxy.operate.deflections import OperateDeflections directory = path.dirname(path.realpath(__file__)) @@ -36,7 +37,9 @@ def test__1d__add_einstein_radius(mp_0, grid_2d_7x7): mass_obj=mp_0, grid=grid_2d_7x7 ) - assert visuals_1d_via.einstein_radius == mp_0.einstein_radius_from(grid=grid_2d_7x7) + assert visuals_1d_via.einstein_radius == OperateDeflections.from_mass_obj( + mp_0 + ).einstein_radius_from(grid=grid_2d_7x7) def test__1d__add_einstein_radius_errors(mp_0, grid_2d_7x7): @@ -45,8 +48,9 @@ def test__1d__add_einstein_radius_errors(mp_0, grid_2d_7x7): mass_obj_list=[mp_0, mp_0], grid=grid_2d_7x7, low_limit=1.0 ) - assert visuals_1d_via.einstein_radius == mp_0.einstein_radius_from(grid=grid_2d_7x7) - assert visuals_1d_via.einstein_radius_errors[0][0] == mp_0.einstein_radius_from( + od = OperateDeflections.from_mass_obj(mp_0) + assert visuals_1d_via.einstein_radius == od.einstein_radius_from(grid=grid_2d_7x7) + assert visuals_1d_via.einstein_radius_errors[0][0] == od.einstein_radius_from( grid=grid_2d_7x7 ) @@ -57,9 +61,10 @@ def test__2d__add_critical_curve(gal_x1_mp, grid_2d_7x7): mass_obj=gal_x1_mp, grid=grid_2d_7x7, plane_index=0 ) + od = OperateDeflections.from_mass_obj(gal_x1_mp) assert ( visuals_2d_via.tangential_critical_curves[0] - == gal_x1_mp.tangential_critical_curve_list_from(grid=grid_2d_7x7)[0] + == od.tangential_critical_curve_list_from(grid=grid_2d_7x7)[0] ).all() @@ -69,11 +74,12 @@ def test__2d__add_caustic(gal_x1_mp, grid_2d_7x7): mass_obj=gal_x1_mp, grid=grid_2d_7x7, plane_index=1 ) + od = OperateDeflections.from_mass_obj(gal_x1_mp) assert ( visuals_2d_via.tangential_caustics[0] - == gal_x1_mp.tangential_caustic_list_from(grid=grid_2d_7x7)[0] + == od.tangential_caustic_list_from(grid=grid_2d_7x7)[0] ).all() assert ( visuals_2d_via.radial_caustics[0] - == gal_x1_mp.radial_caustic_list_from(grid=grid_2d_7x7)[0] + == od.radial_caustic_list_from(grid=grid_2d_7x7)[0] ).all() From eabbc610634b084970fa492cb81c4651ac9121ea Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 4 Mar 2026 17:51:34 +0000 Subject: [PATCH 6/7] implement fully tested and complete --- autogalaxy/analysis/model_util.py | 18 +- autogalaxy/operate/deflections.py | 17 +- autogalaxy/plot/mass_plotter.py | 6 +- autogalaxy/profiles/mass/abstract/mge.py | 245 ++++++++++-------- autogalaxy/profiles/mass/dark/abstract.py | 10 +- autogalaxy/profiles/mass/dark/cnfw.py | 10 +- autogalaxy/profiles/mass/dark/gnfw.py | 6 +- .../mass/dark/gnfw_virial_mass_conc.py | 26 +- autogalaxy/profiles/mass/dark/mcr_util.py | 2 +- autogalaxy/profiles/mass/dark/nfw.py | 6 +- autogalaxy/profiles/mass/stellar/chameleon.py | 7 +- autogalaxy/profiles/mass/stellar/gaussian.py | 11 +- autogalaxy/profiles/mass/stellar/sersic.py | 15 +- .../profiles/mass/stellar/sersic_core.py | 29 +-- test_autogalaxy/operate/test_deflections.py | 9 +- .../profiles/mass/abstract/test_mge.py | 222 ++++++++++------ 16 files changed, 362 insertions(+), 277 deletions(-) diff --git a/autogalaxy/analysis/model_util.py b/autogalaxy/analysis/model_util.py index 25a65eb43..567a604c8 100644 --- a/autogalaxy/analysis/model_util.py +++ b/autogalaxy/analysis/model_util.py @@ -178,22 +178,14 @@ def mge_point_model_from( # and twice the pixel scale, with a floor to avoid taking log10 of # very small or non-positive values. min_log10_sigma = -2.0 # corresponds to 0.01 arcsec - max_sigma = max(2.0 * pixel_scales, 10 ** min_log10_sigma) + max_sigma = max(2.0 * pixel_scales, 10**min_log10_sigma) max_log10_sigma = np.log10(max_sigma) - log10_sigma_list = np.linspace( - min_log10_sigma, max_log10_sigma, total_gaussians - ) - centre_0 = af.UniformPrior( - lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1 - ) - centre_1 = af.UniformPrior( - lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1 - ) + log10_sigma_list = np.linspace(min_log10_sigma, max_log10_sigma, total_gaussians) + centre_0 = af.UniformPrior(lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1) + centre_1 = af.UniformPrior(lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1) - gaussian_list = af.Collection( - af.Model(Gaussian) for _ in range(total_gaussians) - ) + gaussian_list = af.Collection(af.Model(Gaussian) for _ in range(total_gaussians)) for i, gaussian in enumerate(gaussian_list): gaussian.centre.centre_0 = centre_0 diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index 870955a58..8d53ba191 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -102,7 +102,9 @@ def from_mass_obj(cls, mass_obj): return cls(deflections_yx_2d_from=mass_obj.deflections_yx_2d_from) @classmethod - def from_tracer(cls, tracer, use_multi_plane: bool = True, plane_i: int = 0, plane_j: int = -1): + def from_tracer( + cls, tracer, use_multi_plane: bool = True, plane_i: int = 0, plane_j: int = -1 + ): """ Construct from a PyAutoLens ``Tracer`` object. @@ -384,9 +386,7 @@ def jacobian_from(self, grid, xp=np) -> List: return [[a11, a12], [a21, a22]] - def convergence_2d_via_hessian_from( - self, grid, xp=np - ) -> aa.ArrayIrregular: + def convergence_2d_via_hessian_from(self, grid, xp=np) -> aa.ArrayIrregular: """ Returns the convergence of the lensing object, which is computed from the 2D deflection angle map via the Hessian using the expression (see equation 56 https://inspirehep.net/literature/419263): @@ -416,9 +416,7 @@ def convergence_2d_via_hessian_from( return aa.ArrayIrregular(values=convergence) - def shear_yx_2d_via_hessian_from( - self, grid, xp=np - ) -> ShearYX2DIrregular: + def shear_yx_2d_via_hessian_from(self, grid, xp=np) -> ShearYX2DIrregular: """ Returns the 2D (y,x) shear vectors of the lensing object, which are computed from the 2D deflection angle map via the Hessian using the expressions (see equation 57 https://inspirehep.net/literature/419263): @@ -459,9 +457,7 @@ def shear_yx_2d_via_hessian_from( return ShearYX2DIrregular(values=shear_yx_2d, grid=grid) - def magnification_2d_via_hessian_from( - self, grid, xp=np - ) -> aa.ArrayIrregular: + def magnification_2d_via_hessian_from(self, grid, xp=np) -> aa.ArrayIrregular: """ Returns the 2D magnification map of lensing object, which is computed from the 2D deflection angle map via the Hessian using the expressions (see equation 60 https://inspirehep.net/literature/419263): @@ -866,4 +862,3 @@ def einstein_mass_angular_from( ) return einstein_mass_angular_list[0] - diff --git a/autogalaxy/plot/mass_plotter.py b/autogalaxy/plot/mass_plotter.py index b307ca48d..9536ad74a 100644 --- a/autogalaxy/plot/mass_plotter.py +++ b/autogalaxy/plot/mass_plotter.py @@ -121,9 +121,9 @@ def figures_2d( from autogalaxy.operate.deflections import OperateDeflections self.mat_plot_2d.plot_array( - array=OperateDeflections.from_mass_obj(self.mass_obj).magnification_2d_from( - grid=self.grid - ), + array=OperateDeflections.from_mass_obj( + self.mass_obj + ).magnification_2d_from(grid=self.grid), visuals_2d=self.visuals_2d_with_critical_curves, auto_labels=aplt.AutoLabels( title=f"Magnification{title_suffix}", diff --git a/autogalaxy/profiles/mass/abstract/mge.py b/autogalaxy/profiles/mass/abstract/mge.py index 5fc4682e8..5b38f3407 100644 --- a/autogalaxy/profiles/mass/abstract/mge.py +++ b/autogalaxy/profiles/mass/abstract/mge.py @@ -18,22 +18,18 @@ def __init__( ): self.mass_profile = mass_profile - @property def centre(self): return self.mass_profile.centre - @property def ell_comps(self): return self.mass_profile.ell_comps - @property def transformed_to_reference_frame_grid_from(self): return self.mass_profile.transformed_to_reference_frame_grid_from - @staticmethod def kesi(p, xp=np): """ @@ -42,14 +38,13 @@ def kesi(p, xp=np): n_list = xp.arange(0, 2 * p + 1, 1) return (2.0 * p * xp.log(10) / 3.0 + 2.0 * xp.pi * n_list * 1j) ** (0.5) - @staticmethod def eta(p, xp=np): """ see Eq.(6) of Shajib 2019 1906.00263 """ i = xp.arange(1, p, 1) - kesi_last = 1 / 2 ** p + kesi_last = 1 / 2**p k = kesi_last + xp.cumsum(xp.cumprod((p + 1 - i) / i) * kesi_last) kesi_list = xp.hstack( @@ -61,7 +56,6 @@ def eta(p, xp=np): return eta_list - @staticmethod def wofz(z, xp=np): """ @@ -109,7 +103,9 @@ def wofz(z, xp=np): real_exp = xp.clip(-xp.real(z2), None, 700.0) imag_exp = -xp.imag(z2) - w5 = xp.exp(real_exp + 1j * imag_exp) + 1j * z * t / s # clip prevents overflow error + w5 = ( + xp.exp(real_exp + 1j * imag_exp) + 1j * z * t / s + ) # clip prevents overflow error # ---------- Region 6 ---------- U6 = xp.asarray( @@ -142,7 +138,7 @@ def wofz(z, xp=np): # ---------- Region logic ---------- reg1 = (r2 >= 62.0) | ((r2 >= 30.0) & (r2 < 62.0) & (y2 >= 1e-13)) reg2 = ((r2 >= 30) & (r2 < 62) & (y2 < 1e-13)) | ( - (r2 >= 2.5) & (r2 < 30) & (y2 < 0.072) + (r2 >= 2.5) & (r2 < 30) & (y2 < 0.072) ) w = w6 @@ -151,12 +147,18 @@ def wofz(z, xp=np): return w - @aa.grid_dec.to_vector_yx @aa.grid_dec.transform def deflections_2d_via_mge_from( - self, grid: aa.type.Grid2DLike, xp=np, *, sigma_log_list, three_D: bool, - ellipticity_convention: str, func_terms: int = 28, **kwargs, + self, + grid: aa.type.Grid2DLike, + xp=np, + *, + sigma_log_list, + three_D: bool, + ellipticity_convention: str, + func_terms: int = 28, + **kwargs, ): """ Calculates the deflection angle of an arbitrary elliptical convergence / 3d density @@ -178,21 +180,25 @@ def deflections_2d_via_mge_from( sigma_log_array = xp.asarray(sigma_log_list, dtype=xp.float64) amps, sigmas = self.decompose_convergence_via_mge( - sigma_log_list=sigma_log_array, func_terms=func_terms, three_D=three_D, xp=xp) + sigma_log_list=sigma_log_array, + func_terms=func_terms, + three_D=three_D, + xp=xp, + ) q = xp.asarray(self.axis_ratio(xp), dtype=xp.float64) # Change ellipticity convention to (q**2*x**2 + y**2) (most profiles are in (x**2 + y**2/q**2)) - sigmas_factor = self.sigmas_factor_from(input_convention=ellipticity_convention, - target_convention='minor', - xp=xp) + sigmas_factor = self.sigmas_factor_from( + input_convention=ellipticity_convention, target_convention="minor", xp=xp + ) sigmas = sigmas_factor * sigma_log_array deflection_angles = ( - amps[:, None] - * sigmas[:, None] - * xp.sqrt((2.0 * xp.pi) / (1.0 - q**2.0)) - * self.zeta_from(grid=grid, sigma_log_list=sigmas, xp=xp) + amps[:, None] + * sigmas[:, None] + * xp.sqrt((2.0 * xp.pi) / (1.0 - q**2.0)) + * self.zeta_from(grid=grid, sigma_log_list=sigmas, xp=xp) ) # Add Gaussian profiles @@ -203,7 +209,6 @@ def deflections_2d_via_mge_from( xp=xp, ) - def decompose_convergence_via_mge( self, sigma_log_list, three_D: bool, func_terms: int = 28, xp=np ): @@ -225,24 +230,36 @@ def decompose_convergence_via_mge( kesis = self.kesi(func_terms, xp=xp) # kesi in Eq.(6) of 1906.08263 etas = self.eta(func_terms, xp=xp) # eta in Eqr.(6) of 1906.08263 - sigmas = xp.asarray(sigma_log_list, dtype=xp.float64) # major axis + sigmas = xp.asarray(sigma_log_list, dtype=xp.float64) # major axis log_sigmas = xp.log(sigmas) d_log_sigma = log_sigmas[1] - log_sigmas[0] if three_D == True: f_sigma = xp.sum( - etas * xp.real(self.mass_profile.density_3d_func(aa.ArrayIrregular(sigmas.reshape(-1, 1) * kesis), xp=xp)), axis=1 + etas + * xp.real( + self.mass_profile.density_3d_func( + aa.ArrayIrregular(sigmas.reshape(-1, 1) * kesis), xp=xp + ) + ), + axis=1, ) amplitude_list = f_sigma * d_log_sigma * sigmas else: f_sigma = xp.sum( - etas * xp.real(self.mass_profile.convergence_func(aa.ArrayIrregular(sigmas.reshape(-1, 1) * kesis), xp=xp)), axis=1 + etas + * xp.real( + self.mass_profile.convergence_func( + aa.ArrayIrregular(sigmas.reshape(-1, 1) * kesis), xp=xp + ) + ), + axis=1, ) amplitude_list = f_sigma * d_log_sigma / xp.sqrt(2.0 * xp.pi) - if xp==np: + if xp == np: amplitude_list[0] *= 0.5 amplitude_list[-1] *= 0.5 else: @@ -251,12 +268,10 @@ def decompose_convergence_via_mge( return amplitude_list, sigmas - def axis_ratio(self, xp=np): axis_ratio = self.mass_profile.axis_ratio(xp=xp) return xp.where(axis_ratio < 0.9999, axis_ratio, 0.9999) - def zeta_from(self, grid: aa.type.Grid2DLike, sigma_log_list, xp=np): q = xp.asarray(self.axis_ratio(xp), dtype=xp.float64) q2 = q * q @@ -268,9 +283,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike, sigma_log_list, xp=np): sigmas = xp.asarray(sigma_log_list, dtype=xp.float64)[:, None] - scale = q / ( - sigmas * xp.sqrt(xp.asarray(2.0, dtype=xp.float64) * (1.0 - q2)) - ) + scale = q / (sigmas * xp.sqrt(xp.asarray(2.0, dtype=xp.float64) * (1.0 - q2))) xs = x[None, :] * scale ys = xp.abs(y)[None, :] * scale @@ -278,15 +291,9 @@ def zeta_from(self, grid: aa.type.Grid2DLike, sigma_log_list, xp=np): z1 = xs + 1j * ys z2 = q * xs + 1j * ys / q - exp_term = xp.exp( - -(xs * xs) * (1.0 - q2) - - (ys * ys) * (1.0 / q2 - 1.0) - ) + exp_term = xp.exp(-(xs * xs) * (1.0 - q2) - (ys * ys) * (1.0 / q2 - 1.0)) - core = -1j * ( - self.wofz(z1, xp=xp) - - exp_term * self.wofz(z2, xp=xp) - ) + core = -1j * (self.wofz(z1, xp=xp) - exp_term * self.wofz(z2, xp=xp)) return xp.where(ind_pos_y[None, :], core, xp.conj(core)) @@ -294,8 +301,15 @@ def zeta_from(self, grid: aa.type.Grid2DLike, sigma_log_list, xp=np): @aa.grid_dec.to_array @aa.grid_dec.transform def convergence_2d_via_mge_from( - self, grid: aa.type.Grid2DLike, xp=np, *, - sigma_log_list, three_D: bool, ellipticity_convention: str, func_terms: int = 28, **kwargs, + self, + grid: aa.type.Grid2DLike, + xp=np, + *, + sigma_log_list, + three_D: bool, + ellipticity_convention: str, + func_terms: int = 28, + **kwargs, ): """ Calculate the projected convergence at a given set of arc-second gridded coordinates. @@ -307,20 +321,36 @@ def convergence_2d_via_mge_from( """ - eccentric_radii = self.mass_profile.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) - - - sigmas_factor = self.sigmas_factor_from(input_convention=ellipticity_convention, - target_convention='circularised', - xp=xp) + eccentric_radii = self.mass_profile.eccentric_radii_grid_from( + grid=grid, xp=xp, **kwargs + ) - return self._convergence_2d_via_mge_from(grid_radii=eccentric_radii, xp=xp, sigma_log_list=sigma_log_list, - three_D=three_D, sigmas_factor=sigmas_factor, func_terms=func_terms + sigmas_factor = self.sigmas_factor_from( + input_convention=ellipticity_convention, + target_convention="circularised", + xp=xp, ) + return self._convergence_2d_via_mge_from( + grid_radii=eccentric_radii, + xp=xp, + sigma_log_list=sigma_log_list, + three_D=three_D, + sigmas_factor=sigmas_factor, + func_terms=func_terms, + ) - def _convergence_2d_via_mge_from(self, grid_radii, xp=np, *, sigma_log_list, three_D: bool, - sigmas_factor: float = 1.0, func_terms: int = 28, **kwargs): + def _convergence_2d_via_mge_from( + self, + grid_radii, + xp=np, + *, + sigma_log_list, + three_D: bool, + sigmas_factor: float = 1.0, + func_terms: int = 28, + **kwargs, + ): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -333,7 +363,11 @@ def _convergence_2d_via_mge_from(self, grid_radii, xp=np, *, sigma_log_list, thr sigma_log_array = xp.asarray(sigma_log_list, dtype=xp.float64) amps, sigmas = self.decompose_convergence_via_mge( - sigma_log_list=sigma_log_array, func_terms=func_terms, three_D=three_D, xp=xp) + sigma_log_list=sigma_log_array, + func_terms=func_terms, + three_D=three_D, + xp=xp, + ) # Change ellipticity convention to (q*x**2 + y**2/q) (most profiles are in (x**2 + y**2/q**2)) sigmas = sigmas_factor * xp.asarray(sigmas)[:, None] @@ -351,7 +385,6 @@ def _convergence_2d_via_mge_from(self, grid_radii, xp=np, *, sigma_log_list, thr ) return convergence - def convergence_func_gaussian(self, grid_radii, sigma, intensity, xp=np): return xp.multiply( intensity, xp.exp(-0.5 * xp.square(xp.divide(grid_radii.array, sigma))) @@ -359,77 +392,78 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity, xp=np): def sigmas_factor_from(self, input_convention: str, target_convention: str, xp=np): """ - Returns the multiplicative factor required to convert a scale parameter (e.g. sigma) - defined under one ellipticity convention to another. + Returns the multiplicative factor required to convert a scale parameter (e.g. sigma) + defined under one ellipticity convention to another. - The three supported ellipticity conventions are: + The three supported ellipticity conventions are: - Major-axis convention - --------------------- - The elliptical radius is defined such that the scale parameter lies along - the major axis: + Major-axis convention + --------------------- + The elliptical radius is defined such that the scale parameter lies along + the major axis: - R^2 = x^2 + y^2 / q^2 + R^2 = x^2 + y^2 / q^2 - where q = b/a is the projected minor-to-major axis ratio (q <= 1). - In this convention, the scale parameter corresponds directly to the - semi-major axis length. + where q = b/a is the projected minor-to-major axis ratio (q <= 1). + In this convention, the scale parameter corresponds directly to the + semi-major axis length. - Circularised convention - ------------------------ - The elliptical radius is defined such that the ellipse has the same area - as a circle of radius R: + Circularised convention + ------------------------ + The elliptical radius is defined such that the ellipse has the same area + as a circle of radius R: - R^2 = q x^2 + y^2 / q + R^2 = q x^2 + y^2 / q - In this case: + In this case: - a = R / sqrt(q) - b = R sqrt(q) + a = R / sqrt(q) + b = R sqrt(q) - and the enclosed area πab = πR^2 is independent of q. - The scale parameter corresponds to the circularised radius. + and the enclosed area πab = πR^2 is independent of q. + The scale parameter corresponds to the circularised radius. - Minor-axis convention - ---------------------- - The elliptical radius is defined such that the scale parameter lies along - the minor axis: + Minor-axis convention + ---------------------- + The elliptical radius is defined such that the scale parameter lies along + the minor axis: - R^2 = q^2 x^2 + y^2 + R^2 = q^2 x^2 + y^2 - Here, the scale parameter corresponds directly to the semi-minor axis length. + Here, the scale parameter corresponds directly to the semi-minor axis length. - ------------------------------------------------------------------------------ + ------------------------------------------------------------------------------ - Parameters - ---------- - input_convention : str - The current definition of the scale parameter. - Must be one of: 'major', 'circularised', or 'minor'. + Parameters + ---------- + input_convention : str + The current definition of the scale parameter. + Must be one of: 'major', 'circularised', or 'minor'. - target_convention : str - The desired definition of the scale parameter. - Must be one of: 'major', 'circularised', or 'minor'. + target_convention : str + The desired definition of the scale parameter. + Must be one of: 'major', 'circularised', or 'minor'. - Returns - ------- - float or array-like - Multiplicative factor required to convert the scale parameter - between ellipticity conventions. - """ + Returns + ------- + float or array-like + Multiplicative factor required to convert the scale parameter + between ellipticity conventions. + """ - if target_convention == 'major': - if input_convention == 'major': + if target_convention == "major": + if input_convention == "major": return 1.0 - elif input_convention == 'circularised': + elif input_convention == "circularised": return xp.divide(1.0, xp.sqrt(self.mass_profile.axis_ratio(xp))) - elif input_convention == 'minor': + elif input_convention == "minor": return xp.divide(1.0, self.mass_profile.axis_ratio(xp)) else: raise TypeError( - "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention") + "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention" + ) - elif target_convention == 'circularised': + elif target_convention == "circularised": if input_convention == "major": return xp.sqrt(self.mass_profile.axis_ratio(xp)) elif input_convention == "circularised": @@ -438,9 +472,10 @@ def sigmas_factor_from(self, input_convention: str, target_convention: str, xp=n return xp.divide(1.0, xp.sqrt(self.mass_profile.axis_ratio(xp))) else: raise TypeError( - "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention") + "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention" + ) - elif target_convention == 'minor': + elif target_convention == "minor": if input_convention == "major": return self.mass_profile.axis_ratio(xp) elif input_convention == "circularised": @@ -449,8 +484,10 @@ def sigmas_factor_from(self, input_convention: str, target_convention: str, xp=n return 1.0 else: raise TypeError( - "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention") + "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as an input ellipticity convention" + ) else: raise TypeError( - "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as a target ellipticity convention") \ No newline at end of file + "sigmas_factor_from takes either 'major', 'circularised' or 'minor' as a target ellipticity convention" + ) diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py index adb8b359a..6fb60c2ce 100644 --- a/autogalaxy/profiles/mass/dark/abstract.py +++ b/autogalaxy/profiles/mass/dark/abstract.py @@ -67,7 +67,6 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.convergence_func(grid_radius=grid_eta, xp=xp) - def tabulate_integral(self, grid, tabulate_bins, **kwargs): """Tabulate an integral over the convergence of deflection potential of a mass profile. This is used in \ the gNFW profile classes to speed up the integration procedure. @@ -92,16 +91,15 @@ def density_3d_func(self, r, xp=np): x = r.array / self.scale_radius rho_at_scale_radius = ( - self.kappa_s / self.scale_radius + self.kappa_s / self.scale_radius ) # density parameter of 3D gNFW return ( - rho_at_scale_radius - * x ** (-self.inner_slope) - * (1.0 + x) ** (self.inner_slope - 3.0) + rho_at_scale_radius + * x ** (-self.inner_slope) + * (1.0 + x) ** (self.inner_slope - 3.0) ) - def coord_func_f(self, grid_radius: np.ndarray, xp=np) -> np.ndarray: """ Given an array `grid_radius` and a work array `f`, fill f[i] with diff --git a/autogalaxy/profiles/mass/dark/cnfw.py b/autogalaxy/profiles/mass/dark/cnfw.py index 83a6e4bdc..53cba11c7 100644 --- a/autogalaxy/profiles/mass/dark/cnfw.py +++ b/autogalaxy/profiles/mass/dark/cnfw.py @@ -169,12 +169,12 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): def density_3d_func(self, theta, xp=np): rho_at_scale_radius = ( - self.kappa_s / self.scale_radius + self.kappa_s / self.scale_radius ) # density parameter of 3D gNFW return ( - rho_at_scale_radius * self.scale_radius **3.0 - * (theta.array + self.core_radius) ** (-1.0) - * (theta.array + self.scale_radius) ** (-2.0) + rho_at_scale_radius + * self.scale_radius**3.0 + * (theta.array + self.core_radius) ** (-1.0) + * (theta.array + self.scale_radius) ** (-2.0) ) - diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py index cd3494892..438ed867f 100644 --- a/autogalaxy/profiles/mass/dark/gnfw.py +++ b/autogalaxy/profiles/mass/dark/gnfw.py @@ -11,9 +11,7 @@ class gNFW(AbstractgNFW): def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_mge_from(grid=grid, xp=xp, **kwargs) - def deflections_2d_via_mge_from( - self, grid: aa.type.Grid2DLike, xp=np, **kwargs - ): + def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): radii_min = self.scale_radius / 20000.0 radii_max = self.scale_radius * 200.0 log_sigmas = xp.linspace(xp.log(radii_min), xp.log(radii_max), 30) @@ -25,7 +23,7 @@ def deflections_2d_via_mge_from( grid=grid, xp=xp, sigma_log_list=sigmas, - ellipticity_convention='major', + ellipticity_convention="major", three_D=True, ) return deflections_via_mge diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py index 682b41719..acd84c547 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py @@ -4,15 +4,18 @@ from autogalaxy.profiles.mass.dark.gnfw import gNFWSph from autogalaxy import cosmology as cosmo + def is_jax(x): try: import jax from jax import Array from jax.core import Tracer + return isinstance(x, (Array, Tracer)) except Exception: return False + def kappa_s_and_scale_radius( cosmology, virial_mass, @@ -96,6 +99,7 @@ def kappa_s_and_scale_radius( xp = np else: from jax import numpy as jnp + xp = jnp if xp is np: @@ -112,15 +116,21 @@ def kappa_s_and_scale_radius( gamma = inner_slope concentration = (2.0 - gamma) * c_2 # gNFW concentration (your definition) - critical_density = cosmology.critical_density(redshift_object, xp=xp) # Msun / kpc^3 + critical_density = cosmology.critical_density( + redshift_object, xp=xp + ) # Msun / kpc^3 - critical_surface_density = cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_object, - redshift_1=redshift_source, - xp=xp, + critical_surface_density = ( + cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( + redshift_0=redshift_object, + redshift_1=redshift_source, + xp=xp, + ) ) # Msun / kpc^2 - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=xp) # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from( + redshift=redshift_object, xp=xp + ) # kpc / arcsec # Bryan & Norman (1998) overdensity if overdens == 0 x = cosmology.Om(redshift_object, xp=xp) - 1.0 @@ -128,7 +138,9 @@ def kappa_s_and_scale_radius( overdens = xp.where(overdens == 0, overdens_bn98, overdens) # r_vir in kpc - virial_radius = (virial_mass / (overdens * critical_density * (4.0 * xp.pi / 3.0))) ** (1.0 / 3.0) + virial_radius = ( + virial_mass / (overdens * critical_density * (4.0 * xp.pi / 3.0)) + ) ** (1.0 / 3.0) # scale radius in kpc scale_radius_kpc = virial_radius / concentration diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index ce6f15c95..37a1504ad 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -133,7 +133,7 @@ def ludlow16_cosmology_jax( mass_at_200, redshift_object, redshift_source, - vmap_method="sequential" + vmap_method="sequential", ) diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py index 8d482b4c7..cd8233339 100644 --- a/autogalaxy/profiles/mass/dark/nfw.py +++ b/autogalaxy/profiles/mass/dark/nfw.py @@ -183,7 +183,6 @@ def deflection_func(u, y, x, npow, axis_ratio, scale_radius): / ((1 - (1 - axis_ratio**2) * u) ** (npow + 0.5)) ) - @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform @@ -209,10 +208,11 @@ def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) def convergence_func(self, grid_radius: float, xp=np) -> float: grid_radius = (1.0 / self.scale_radius) * grid_radius.array + 0j return xp.real( - 2.0 * self.kappa_s * xp.array(self.coord_func_g(grid_radius=grid_radius, xp=xp)) + 2.0 + * self.kappa_s + * xp.array(self.coord_func_g(grid_radius=grid_radius, xp=xp)) ) - @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py index ddf5bb710..2ba95146e 100644 --- a/autogalaxy/profiles/mass/stellar/chameleon.py +++ b/autogalaxy/profiles/mass/stellar/chameleon.py @@ -143,12 +143,13 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): The grid of (y,x) arc-second coordinates the convergence is computed on. """ return self.convergence_func( - self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs), - xp=xp + self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp ) def convergence_func(self, grid_radius: float, xp=np) -> float: - return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius, xp=xp) + return self.mass_to_light_ratio * self.image_2d_via_radii_from( + grid_radius, xp=xp + ) @aa.grid_dec.to_array def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py index 1d5d273a0..6ca7a5765 100644 --- a/autogalaxy/profiles/mass/stellar/gaussian.py +++ b/autogalaxy/profiles/mass/stellar/gaussian.py @@ -151,12 +151,13 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ return self.convergence_func( - self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), - xp=xp + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp ) def convergence_func(self, grid_radius: float, xp=np) -> float: - return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius, xp=xp) + return self.mass_to_light_ratio * self.image_2d_via_radii_from( + grid_radius, xp=xp + ) @aa.grid_dec.to_array def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): @@ -261,7 +262,9 @@ def wofz(z, xp=np): real_exp = xp.clip(-xp.real(z2), None, 700.0) imag_exp = -xp.imag(z2) - w5 = xp.exp(real_exp + 1j * imag_exp) + 1j * z * t / s # clip prevents overflow error + w5 = ( + xp.exp(real_exp + 1j * imag_exp) + 1j * z * t / s + ) # clip prevents overflow error # ---------- Region 6 ---------- U6 = xp.asarray( diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py index 9e38a7ab1..e9bc1cbea 100644 --- a/autogalaxy/profiles/mass/stellar/sersic.py +++ b/autogalaxy/profiles/mass/stellar/sersic.py @@ -122,7 +122,6 @@ def __init__( def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_cse_from(grid=grid, xp=xp, **kwargs) - @aa.grid_dec.to_vector_yx @aa.grid_dec.transform def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): @@ -154,11 +153,9 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ return self.convergence_func( - self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), - xp=xp + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp ) - @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform @@ -182,7 +179,9 @@ def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) return self._convergence_2d_via_cse_from(grid_radii=elliptical_radii) def convergence_func(self, grid_radius: float, xp=np) -> float: - return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius, xp=xp) + return self.mass_to_light_ratio * self.image_2d_via_radii_from( + grid_radius, xp=xp + ) @aa.grid_dec.to_array def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): @@ -200,12 +199,11 @@ def image_2d_via_radii_from(self, radius: np.ndarray, xp=np): return self.intensity * xp.exp( -self.sersic_constant * ( - ((radius.array / self.effective_radius) ** (1.0 / self.sersic_index)) - - 1 + ((radius.array / self.effective_radius) ** (1.0 / self.sersic_index)) + - 1 ) ) - def decompose_convergence_via_cse( self, grid_radii: np.ndarray ) -> Tuple[List, List]: @@ -258,7 +256,6 @@ def sersic_2d(r): ) ) - return self._decompose_convergence_via_cse_from( func=sersic_2d, radii_min=radii_min, diff --git a/autogalaxy/profiles/mass/stellar/sersic_core.py b/autogalaxy/profiles/mass/stellar/sersic_core.py index e84e087df..b1dfce411 100644 --- a/autogalaxy/profiles/mass/stellar/sersic_core.py +++ b/autogalaxy/profiles/mass/stellar/sersic_core.py @@ -63,9 +63,7 @@ def __init__( def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_mge_from(grid=grid, xp=xp, **kwargs) - def deflections_2d_via_mge_from( - self, grid: aa.type.Grid2DLike, xp=np, **kwargs - ): + def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): radii_min = self.effective_radius / 50.0 radii_max = self.effective_radius * 20.0 log_sigmas = xp.linspace(xp.log(radii_min), xp.log(radii_max), 20) @@ -77,7 +75,7 @@ def deflections_2d_via_mge_from( grid=grid, xp=xp, sigma_log_list=sigmas, - ellipticity_convention='circularised', + ellipticity_convention="circularised", three_D=False, ) return deflections_via_mge @@ -125,21 +123,20 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): def convergence_func(self, grid_radius: float, xp=np) -> float: return ( - self.mass_to_light_ratio - * self.intensity_prime(xp=xp) - * (1.0 + (self.radius_break / grid_radius.array) ** self.alpha) - ** (self.gamma / self.alpha) - * xp.exp( - -self.sersic_constant - * ( - (grid_radius.array ** self.alpha + self.radius_break ** self.alpha) - / self.effective_radius ** self.alpha + self.mass_to_light_ratio + * self.intensity_prime(xp=xp) + * (1.0 + (self.radius_break / grid_radius.array) ** self.alpha) + ** (self.gamma / self.alpha) + * xp.exp( + -self.sersic_constant + * ( + (grid_radius.array**self.alpha + self.radius_break**self.alpha) + / self.effective_radius**self.alpha + ) + ** (1.0 / (self.sersic_index * self.alpha)) ) - ** (1.0 / (self.sersic_index * self.alpha)) - ) ) - def intensity_prime(self, xp=np): """Overall intensity normalisation in the rescaled Core-Sersic light profiles (electrons per second)""" return ( diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 04ef8a2b8..0421ef102 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -12,9 +12,9 @@ def critical_curve_via_magnification_from(mass_profile, grid): - magnification = OperateDeflections.from_mass_obj(mass_profile).magnification_2d_from( - grid=grid - ) + magnification = OperateDeflections.from_mass_obj( + mass_profile + ).magnification_2d_from(grid=grid) inverse_magnification = 1 / magnification @@ -148,6 +148,7 @@ def test__magnification_2d_via_hessian_from(): assert magnification.in_list[0] == pytest.approx(-0.56303, 1.0e-4) assert magnification.in_list[1] == pytest.approx(-2.57591, 1.0e-4) + def test__tangential_critical_curve_list_from(): grid = ag.Grid2D.uniform(shape_native=(15, 15), pixel_scales=0.3) @@ -508,5 +509,3 @@ def test__jacobian_from(): assert magnification_via_jacobian == pytest.approx( np.array(magnification_via_hessian), rel=1e-6 ) - - diff --git a/test_autogalaxy/profiles/mass/abstract/test_mge.py b/test_autogalaxy/profiles/mass/abstract/test_mge.py index 0977dbd2a..4babc8137 100644 --- a/test_autogalaxy/profiles/mass/abstract/test_mge.py +++ b/test_autogalaxy/profiles/mass/abstract/test_mge.py @@ -30,8 +30,8 @@ def test__gnfw_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='major', - three_D=True + ellipticity_convention="major", + three_D=True, ) assert deflections_via_integral == pytest.approx(deflections_via_mge, 1.0e-3) @@ -58,8 +58,8 @@ def test__gnfw_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='major', - three_D=True + ellipticity_convention="major", + three_D=True, ) assert deflections_via_integral == pytest.approx(deflections_via_mge, 1.0e-3) @@ -85,12 +85,11 @@ def test__sersic_deflections_yx_2d_via_mge(): ) mge_decomp = ag.mp.MGEDecomposer(mass_profile=mp) - deflections_via_mge = mge_decomp.deflections_2d_via_mge_from( grid=ag.Grid2DIrregular([[0.1625, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='circularised', + ellipticity_convention="circularised", three_D=False, ) @@ -115,7 +114,7 @@ def test__sersic_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[0.1625, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='circularised', + ellipticity_convention="circularised", three_D=False, ) @@ -124,10 +123,7 @@ def test__sersic_deflections_yx_2d_via_mge(): def test__cnfw_deflections_yx_2d_via_mge(): cnfw = ag.mp.cNFWSph( - centre=(0.3, 0.2), - kappa_s=0.05, - scale_radius=1.1, - core_radius=0.01 + centre=(0.3, 0.2), kappa_s=0.05, scale_radius=1.1, core_radius=0.01 ) radii_min = cnfw.scale_radius / 1000.0 @@ -135,8 +131,9 @@ def test__cnfw_deflections_yx_2d_via_mge(): log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), 20) sigmas = np.exp(log_sigmas) - deflections_analytic = cnfw.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), - xp=np) + deflections_analytic = cnfw.deflections_yx_2d_from( + grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np + ) mge_decomp = ag.mp.MGEDecomposer(mass_profile=cnfw) @@ -144,12 +141,13 @@ def test__cnfw_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='major', - three_D=True + ellipticity_convention="major", + three_D=True, ) assert deflections_analytic == pytest.approx(deflections_via_mge, 1.0e-3) + def test__powerlaw_deflections_yx_2d_via_mge(): mp = ag.mp.PowerLaw( centre=(-0.4, -0.2), @@ -163,8 +161,9 @@ def test__powerlaw_deflections_yx_2d_via_mge(): log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), 30) sigmas = np.exp(log_sigmas) - deflections_analytic = mp.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), - xp=np) + deflections_analytic = mp.deflections_yx_2d_from( + grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np + ) mge_decomp = ag.mp.MGEDecomposer(mass_profile=mp) @@ -172,17 +171,16 @@ def test__powerlaw_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[0.1875, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='major', - three_D=False + ellipticity_convention="major", + three_D=False, ) assert deflections_analytic == pytest.approx(deflections_via_mge, 1.0e-2) + def test__chameleon_deflections_yx_2d_via_mge(): mp = ag.mp.Chameleon( - centre=(-0.4, -0.2), - ell_comps=(0.17142, -0.285116), - intensity=2.8 + centre=(-0.4, -0.2), ell_comps=(0.17142, -0.285116), intensity=2.8 ) radii_min = mp.core_radius_0 / 10000.0 @@ -190,8 +188,9 @@ def test__chameleon_deflections_yx_2d_via_mge(): log_sigmas = np.linspace(np.log(radii_min), np.log(radii_max), 30) sigmas = np.exp(log_sigmas) - deflections_analytic = mp.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[-0.1875, 0.1625]]), - xp=np) + deflections_analytic = mp.deflections_yx_2d_from( + grid=ag.Grid2DIrregular([[-0.1875, 0.1625]]), xp=np + ) mge_decomp = ag.mp.MGEDecomposer(mass_profile=mp) @@ -199,8 +198,8 @@ def test__chameleon_deflections_yx_2d_via_mge(): grid=ag.Grid2DIrregular([[-0.1875, 0.1625]]), xp=np, sigma_log_list=sigmas, - ellipticity_convention='major', - three_D=False + ellipticity_convention="major", + three_D=False, ) assert deflections_analytic == pytest.approx(deflections_via_mge, 1.0e-3) @@ -221,9 +220,12 @@ def test__DevVaucouleurs_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[1.0, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(5.6697, 1e-3) @@ -241,9 +243,12 @@ def test__DevVaucouleurs_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(7.4455, 1e-3) @@ -261,9 +266,12 @@ def test__DevVaucouleurs_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(2.0 * 7.4455, 1e-3) @@ -281,9 +289,12 @@ def test__DevVaucouleurs_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(2.0 * 7.4455, 1e-3) @@ -301,9 +312,12 @@ def test__DevVaucouleurs_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(0.351797, 1e-3) @@ -323,9 +337,12 @@ def test__convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[1.0, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(4.9047, 1e-3) @@ -343,9 +360,12 @@ def test__convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(4.8566, 1e-3) @@ -362,9 +382,12 @@ def test__convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(2.0 * 4.8566, 1e-3) @@ -382,9 +405,12 @@ def test__convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(2.0 * 4.8566, 1e-3) @@ -402,9 +428,12 @@ def test__convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.0]]), - sigma_log_list=sigmas, ellipticity_convention='circularised', - three_D=False) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.0]]), + sigma_log_list=sigmas, + ellipticity_convention="circularised", + three_D=False, + ) assert convergence == pytest.approx(4.8566, 1e-3) @@ -423,15 +452,21 @@ def test__nfw_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=nfw) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[2.0, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='major', - three_D=True) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[2.0, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="major", + three_D=True, + ) assert convergence == pytest.approx(0.263600141, 1e-2) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.5, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='major', - three_D=True) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.5, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="major", + three_D=True, + ) assert convergence == pytest.approx(1.388511, 1e-2) @@ -439,9 +474,12 @@ def test__nfw_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=nfw) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.5, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='major', - three_D=True) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.5, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="major", + three_D=True, + ) assert convergence == pytest.approx(2.0 * 1.388511, 1e-2) @@ -454,9 +492,12 @@ def test__nfw_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=nfw) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='major', - three_D=True) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[1.0, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="major", + three_D=True, + ) assert convergence == pytest.approx(1.388511, 1e-2) @@ -474,9 +515,12 @@ def test__nfw_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=nfw) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.25, 0.0]]), - sigma_log_list=sigmas, ellipticity_convention='major', - three_D=True) + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.25, 0.0]]), + sigma_log_list=sigmas, + ellipticity_convention="major", + three_D=True, + ) assert convergence == pytest.approx(1.388511, 1e-3) @@ -497,9 +541,12 @@ def test__sersic_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.5]]), - sigma_log_list=sigmas, three_D=False, - ellipticity_convention='circularised') + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.5]]), + sigma_log_list=sigmas, + three_D=False, + ellipticity_convention="circularised", + ) assert convergence == pytest.approx(4.90657319276, 1e-3) @@ -513,9 +560,12 @@ def test__sersic_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.5]]), - sigma_log_list=sigmas, three_D=False, - ellipticity_convention='circularised') + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.5]]), + sigma_log_list=sigmas, + three_D=False, + ellipticity_convention="circularised", + ) assert convergence == pytest.approx(2.0 * 4.90657319276, 1e-3) @@ -529,9 +579,12 @@ def test__sersic_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[0.0, 1.5]]), - sigma_log_list=sigmas, three_D=False, - ellipticity_convention='circularised') + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[0.0, 1.5]]), + sigma_log_list=sigmas, + three_D=False, + ellipticity_convention="circularised", + ) assert convergence == pytest.approx(2.0 * 4.90657319276, 1e-3) @@ -546,8 +599,11 @@ def test__sersic_convergence_2d_via_mge_from(): mge_decomp = MGEDecomposer(mass_profile=mp) - convergence = mge_decomp.convergence_2d_via_mge_from(grid=ag.Grid2DIrregular([[1.0, 0.0]]), - sigma_log_list=sigmas, three_D=False, - ellipticity_convention='circularised') + convergence = mge_decomp.convergence_2d_via_mge_from( + grid=ag.Grid2DIrregular([[1.0, 0.0]]), + sigma_log_list=sigmas, + three_D=False, + ellipticity_convention="circularised", + ) assert convergence == pytest.approx(5.38066670129, 1e-3) From d4f89bdc59e27fcd43b78cde61dce0c49aa7d06e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 4 Mar 2026 18:19:06 +0000 Subject: [PATCH 7/7] rename OperateDeflections to LensCalc, move fermat_potential_from into LensCalc - Rename operate/deflections.py -> operate/lens_calc.py - Rename class OperateDeflections -> LensCalc throughout - LensCalc.__init__ now accepts optional potential_2d_from callable - from_mass_obj and from_tracer capture potential_2d_from automatically - fermat_potential_from moved into LensCalc (removed from MassProfile, Galaxy, Galaxies) - Update all imports, call sites, tests, and docs in autogalaxy Co-Authored-By: Claude Sonnet 4.6 --- autogalaxy/__init__.py | 2 +- autogalaxy/galaxy/galaxies.py | 28 ------- autogalaxy/galaxy/galaxy.py | 28 ------- .../operate/{deflections.py => lens_calc.py} | 75 ++++++++++++++++--- autogalaxy/plot/mass_plotter.py | 4 +- autogalaxy/plot/visuals/one_d.py | 8 +- autogalaxy/plot/visuals/two_d.py | 8 +- autogalaxy/profiles/mass/abstract/abstract.py | 28 ------- docs/api/source.rst | 2 +- test_autogalaxy/operate/test_deflections.py | 58 +++++++------- test_autogalaxy/plot/mat_wrap/test_visuals.py | 10 +-- 11 files changed, 109 insertions(+), 142 deletions(-) rename autogalaxy/operate/{deflections.py => lens_calc.py} (90%) diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py index a606c7242..e5137276d 100644 --- a/autogalaxy/__init__.py +++ b/autogalaxy/__init__.py @@ -69,7 +69,7 @@ from .operate.image import OperateImage from .operate.image import OperateImageList from .operate.image import OperateImageGalaxies -from .operate.deflections import OperateDeflections +from .operate.lens_calc import LensCalc from .gui.scribbler import Scribbler from .imaging.fit_imaging import FitImaging from .imaging.model.analysis import AnalysisImaging diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py index 0135b7357..20992b7fc 100644 --- a/autogalaxy/galaxy/galaxies.py +++ b/autogalaxy/galaxy/galaxies.py @@ -243,34 +243,6 @@ def potential_2d_from( """ return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self)) - def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: - """ - Returns the Fermat potential of the galaxies for a given grid of image-plane positions. - - This is the sum of the geometric time delay term and the lensing potential, computed - using the combined deflections and potential of all galaxies: - - .. math:: - \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2 - - \\psi(\\boldsymbol{\\theta}) - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. - xp - The array module (``numpy`` or ``jax.numpy``). - """ - from autogalaxy.operate.deflections import OperateDeflections - - od = OperateDeflections.from_mass_obj(self) - time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) - potential = self.potential_2d_from(grid=grid, xp=xp) - fermat_potential = time_delay - potential - if isinstance(grid, aa.Grid2DIrregular): - return aa.ArrayIrregular(values=fermat_potential) - return aa.Array2D(values=fermat_potential, mask=grid.mask) - def has(self, cls: Union[Type, Tuple[Type]]) -> bool: """ Returns a bool specifying whether any of the galaxies has a certain class type. diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py index c0e972d96..a989f94bb 100644 --- a/autogalaxy/galaxy/galaxy.py +++ b/autogalaxy/galaxy/galaxy.py @@ -343,34 +343,6 @@ def potential_2d_from( ) return xp.zeros((grid.shape[0],)) - def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: - """ - Returns the Fermat potential of the galaxy for a given grid of image-plane positions. - - This is the sum of the geometric time delay term and the lensing potential, computed - using the galaxy's total deflections and total potential across all mass profiles: - - .. math:: - \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2 - - \\psi(\\boldsymbol{\\theta}) - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. - xp - The array module (``numpy`` or ``jax.numpy``). - """ - from autogalaxy.operate.deflections import OperateDeflections - - od = OperateDeflections.from_mass_obj(self) - time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) - potential = self.potential_2d_from(grid=grid, xp=xp) - fermat_potential = time_delay - potential - if isinstance(grid, aa.Grid2DIrregular): - return aa.ArrayIrregular(values=fermat_potential) - return aa.Array2D(values=fermat_potential, mask=grid.mask) - @property def half_light_radius(self): return None diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/lens_calc.py similarity index 90% rename from autogalaxy/operate/deflections.py rename to autogalaxy/operate/lens_calc.py index 8d53ba191..55055f236 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/lens_calc.py @@ -78,28 +78,40 @@ def wrapper( return wrapper -class OperateDeflections: +class LensCalc: """ - Packages methods which manipulate the 2D deflection angle map returned from the `deflections_yx_2d_from` function - of a mass object (e.g. a `MassProfile`, `Galaxy`). + Computes lensing quantities from a deflection-angle callable and an optional potential callable. - The majority of methods are those which from the 2D deflection angle map compute lensing quantities like a 2D - shear field, magnification map or the Einstein Radius. + The deflection callable is used to compute the Hessian, Jacobian, convergence, shear, + magnification, critical curves, caustics, and Einstein radius/mass. If a potential + callable is also supplied, ``fermat_potential_from`` is available as well. Parameters ---------- deflections_yx_2d_from - A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection angles on the given - grid. Typically a bound method of a ``MassProfile``, ``Galaxy``, or ``Galaxies`` instance. + A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection + angles on the given grid. Typically a bound method of a ``MassProfile``, ``Galaxy``, + or ``Galaxies`` instance. + potential_2d_from + Optional callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D + lensing potential on the given grid. Required only for ``fermat_potential_from``. """ - def __init__(self, deflections_yx_2d_from): + def __init__(self, deflections_yx_2d_from, potential_2d_from=None): self.deflections_yx_2d_from = deflections_yx_2d_from + self.potential_2d_from = potential_2d_from @classmethod def from_mass_obj(cls, mass_obj): - """Construct from any object that has a ``deflections_yx_2d_from`` method.""" - return cls(deflections_yx_2d_from=mass_obj.deflections_yx_2d_from) + """Construct from any object that has a ``deflections_yx_2d_from`` method. + + If the object also exposes ``potential_2d_from``, it is captured so that + ``fermat_potential_from`` is available on the returned instance. + """ + return cls( + deflections_yx_2d_from=mass_obj.deflections_yx_2d_from, + potential_2d_from=getattr(mass_obj, "potential_2d_from", None), + ) @classmethod def from_tracer( @@ -130,6 +142,8 @@ def from_tracer( Index of the second plane used by ``deflections_between_planes_from``. Ignored when ``use_multi_plane=False``. Defaults to ``-1`` (source plane). """ + potential_2d_from = getattr(tracer, "potential_2d_from", None) + if use_multi_plane: from functools import partial @@ -138,9 +152,13 @@ def from_tracer( tracer.deflections_between_planes_from, plane_i=plane_i, plane_j=plane_j, - ) + ), + potential_2d_from=potential_2d_from, ) - return cls(deflections_yx_2d_from=tracer.deflections_yx_2d_from) + return cls( + deflections_yx_2d_from=tracer.deflections_yx_2d_from, + potential_2d_from=potential_2d_from, + ) def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D: """ @@ -177,6 +195,39 @@ def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D: return aa.ArrayIrregular(values=delay) return aa.Array2D(values=delay, mask=grid.mask) + def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: + """ + Returns the Fermat potential for a given grid of image-plane positions. + + This is the sum of the geometric time delay term and the gravitational (Shapiro) delay + term (i.e. the lensing potential), and is given by: + + .. math:: + \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\theta} - \\boldsymbol{\\beta}|^2 + - \\psi(\\boldsymbol{\\theta}) + + Requires that ``potential_2d_from`` was supplied at construction (e.g. via + ``LensCalc.from_mass_obj`` or ``LensCalc.from_tracer``). + + Parameters + ---------- + grid + The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. + xp + The array module (``numpy`` or ``jax.numpy``). + """ + if self.potential_2d_from is None: + raise ValueError( + "fermat_potential_from requires a potential_2d_from callable. " + "Construct LensCalc with potential_2d_from, or use from_mass_obj / from_tracer." + ) + time_delay = self.time_delay_geometry_term_from(grid=grid, xp=xp) + potential = self.potential_2d_from(grid=grid, xp=xp) + fermat_potential = time_delay - potential + if isinstance(grid, aa.Grid2DIrregular): + return aa.ArrayIrregular(values=fermat_potential) + return aa.Array2D(values=fermat_potential, mask=grid.mask) + def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D: """ Returns the tangential eigen values of lensing jacobian, which are given by the expression: diff --git a/autogalaxy/plot/mass_plotter.py b/autogalaxy/plot/mass_plotter.py index 9536ad74a..16167ad59 100644 --- a/autogalaxy/plot/mass_plotter.py +++ b/autogalaxy/plot/mass_plotter.py @@ -118,10 +118,10 @@ def figures_2d( ) if magnification: - from autogalaxy.operate.deflections import OperateDeflections + from autogalaxy.operate.lens_calc import LensCalc self.mat_plot_2d.plot_array( - array=OperateDeflections.from_mass_obj( + array=LensCalc.from_mass_obj( self.mass_obj ).magnification_2d_from(grid=self.grid), visuals_2d=self.visuals_2d_with_critical_curves, diff --git a/autogalaxy/plot/visuals/one_d.py b/autogalaxy/plot/visuals/one_d.py index a610fe0e3..f6ea4af8f 100644 --- a/autogalaxy/plot/visuals/one_d.py +++ b/autogalaxy/plot/visuals/one_d.py @@ -170,12 +170,12 @@ def add_einstein_radius( The collection of attributes that can be plotted by a `Plotter` object. """ - from autogalaxy.operate.deflections import OperateDeflections + from autogalaxy.operate.lens_calc import LensCalc einstein_radius = None try: - od = OperateDeflections.from_mass_obj(mass_obj) + od = LensCalc.from_mass_obj(mass_obj) einstein_radius = od.einstein_radius_from(grid=grid) except (TypeError, AttributeError): pass @@ -217,13 +217,13 @@ def add_einstein_radius_errors( The mean value and errors of each attribute that are plotted in 1D by a `Plotter` object. """ - from autogalaxy.operate.deflections import OperateDeflections + from autogalaxy.operate.lens_calc import LensCalc einstein_radius_list = [] for mass_obj in mass_obj_list: try: - od = OperateDeflections.from_mass_obj(mass_obj) + od = LensCalc.from_mass_obj(mass_obj) einstein_radius_list.append(od.einstein_radius_from(grid=grid)) except TypeError: einstein_radius_list.append(None) diff --git a/autogalaxy/plot/visuals/two_d.py b/autogalaxy/plot/visuals/two_d.py index 3bd74e90b..9a177463d 100644 --- a/autogalaxy/plot/visuals/two_d.py +++ b/autogalaxy/plot/visuals/two_d.py @@ -184,9 +184,9 @@ def add_critical_curves(self, mass_obj, grid: aa.type.Grid2DLike): vis.Visuals2D A collection of attributes that can be plotted by a `Plotter` object. """ - from autogalaxy.operate.deflections import OperateDeflections + from autogalaxy.operate.lens_calc import LensCalc - od = OperateDeflections.from_mass_obj(mass_obj) + od = LensCalc.from_mass_obj(mass_obj) tangential_critical_curves = od.tangential_critical_curve_list_from(grid=grid) @@ -227,9 +227,9 @@ def add_caustics(self, mass_obj, grid: aa.type.Grid2DLike): vis.Visuals2D A collection of attributes that can be plotted by a `Plotter` object. """ - from autogalaxy.operate.deflections import OperateDeflections + from autogalaxy.operate.lens_calc import LensCalc - od = OperateDeflections.from_mass_obj(mass_obj) + od = LensCalc.from_mass_obj(mass_obj) tangential_caustics = od.tangential_caustic_list_from(grid=grid) radial_caustics = od.radial_caustic_list_from(grid=grid) diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py index 3fb130688..0006f82c4 100644 --- a/autogalaxy/profiles/mass/abstract/abstract.py +++ b/autogalaxy/profiles/mass/abstract/abstract.py @@ -27,34 +27,6 @@ def __init__( def deflections_yx_2d_from(self, grid): raise NotImplementedError - def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: - """ - Returns the Fermat potential for a given grid of image-plane positions. - - This is the sum of the geometric time delay term and the gravitational (Shapiro) delay term (i.e. the lensing - potential), and is given by: - - .. math:: - \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\theta} - \\boldsymbol{\\beta}|^2 - - \\psi(\\boldsymbol{\\theta}) - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on. - xp - The array module (``numpy`` or ``jax.numpy``). - """ - from autogalaxy.operate.deflections import OperateDeflections - - od = OperateDeflections.from_mass_obj(self) - time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp) - potential = self.potential_2d_from(grid=grid, xp=xp) - fermat_potential = time_delay - potential - if isinstance(grid, aa.Grid2DIrregular): - return aa.ArrayIrregular(values=fermat_potential) - return aa.Array2D(values=fermat_potential, mask=grid.mask) - def deflections_2d_via_potential_2d_from(self, grid): potential = self.potential_2d_from(grid=grid) diff --git a/docs/api/source.rst b/docs/api/source.rst index c614f429e..01d1dad04 100644 --- a/docs/api/source.rst +++ b/docs/api/source.rst @@ -31,7 +31,7 @@ Operators :recursive: OperateImage - OperateDeflections + LensCalc Total [ag.mp] ------------- diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 0421ef102..ee4617979 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -5,14 +5,14 @@ import autogalaxy as ag -from autogalaxy.operate.deflections import ( +from autogalaxy.operate.lens_calc import ( grid_scaled_2d_for_marching_squares_from, - OperateDeflections, + LensCalc, ) def critical_curve_via_magnification_from(mass_profile, grid): - magnification = OperateDeflections.from_mass_obj( + magnification = LensCalc.from_mass_obj( mass_profile ).magnification_2d_from(grid=grid) @@ -69,7 +69,7 @@ def test__time_delay_geometry_term_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) time_delay_geometry_term = od.time_delay_geometry_term_from(grid=grid) assert time_delay_geometry_term == pytest.approx( @@ -85,7 +85,7 @@ def test__fermat_potential_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - fermat_potential = mp.fermat_potential_from(grid=grid) + fermat_potential = LensCalc.from_mass_obj(mp).fermat_potential_from(grid=grid) assert fermat_potential == pytest.approx( np.array([0.24329033, -0.82766592]), 1.0e-4 @@ -99,7 +99,7 @@ def test__hessian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) hessian_yy, hessian_xy, hessian_yx, hessian_xx = od.hessian_from(grid=grid) assert hessian_yy == pytest.approx(np.array([1.3883822, 0.694127]), 1.0e-4) @@ -126,7 +126,7 @@ def test__convergence_2d_via_hessian_from(): centre=(0.0, 0.0), ell_comps=(0.001, 0.001), einstein_radius=1.0 ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) convergence = od.convergence_2d_via_hessian_from(grid=grid) assert convergence.in_list[0] == pytest.approx(0.46208, 1.0e-1) @@ -142,7 +142,7 @@ def test__magnification_2d_via_hessian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) magnification = od.magnification_2d_via_hessian_from(grid=grid) assert magnification.in_list[0] == pytest.approx(-0.56303, 1.0e-4) @@ -154,7 +154,7 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) x_critical_tangential, y_critical_tangential = ( @@ -170,7 +170,7 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) y_centre = np.mean(tangential_critical_curve_list[0][:, 0]) @@ -181,7 +181,7 @@ def test__tangential_critical_curve_list_from(): mp = ag.mp.IsothermalSph(centre=(0.5, 1.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) tangential_critical_curve_list = od.tangential_critical_curve_list_from(grid=grid) y_centre = np.mean(tangential_critical_curve_list[0][:, 0]) @@ -230,7 +230,7 @@ def test__radial_critical_curve_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) y_centre = np.mean(radial_critical_curve_list[0][:, 0]) @@ -241,7 +241,7 @@ def test__radial_critical_curve_list_from(): mp = ag.mp.PowerLawSph(centre=(0.5, 1.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) y_centre = np.mean(radial_critical_curve_list[0][:, 0]) @@ -263,7 +263,7 @@ def test__radial_critical_curve_list_from__compare_via_magnification(): mass_profile=mp, grid=grid )[1] - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_critical_curve_list = od.radial_critical_curve_list_from(grid=grid) assert sum(critical_curve_radial_via_magnification) == pytest.approx( @@ -276,7 +276,7 @@ def test__tangential_caustic_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) tangential_caustic_list = od.tangential_caustic_list_from(grid=grid) y_centre = np.mean(tangential_caustic_list[0][:, 0]) @@ -287,7 +287,7 @@ def test__tangential_caustic_list_from(): mp = ag.mp.IsothermalSph(centre=(0.5, 1.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) tangential_caustic_list = od.tangential_caustic_list_from(grid=grid) y_centre = np.mean(tangential_caustic_list[0][:, 0]) @@ -324,7 +324,7 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_caustic_list = od.radial_caustic_list_from(grid=grid) x_caustic_radial, y_caustic_radial = ( @@ -340,7 +340,7 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_caustic_list = od.radial_caustic_list_from(grid=grid) y_centre = np.mean(radial_caustic_list[0][:, 0]) @@ -351,7 +351,7 @@ def test__radial_caustic_list_from(): mp = ag.mp.PowerLawSph(centre=(0.5, 1.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_caustic_list = od.radial_caustic_list_from(grid=grid) y_centre = np.mean(radial_caustic_list[0][:, 0]) @@ -372,7 +372,7 @@ def test__radial_caustic_list_from___compare_via_magnification(): mass_profile=mp, grid=grid )[1] - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) radial_caustic_list = od.radial_caustic_list_from(grid=grid) assert sum(radial_caustic_list[0]) == pytest.approx( @@ -385,7 +385,7 @@ def test__radial_critical_curve_area_list_from(): mp = ag.mp.PowerLawSph(centre=(0.0, 0.0), einstein_radius=2.0, slope=1.5) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) area_within_radial_critical_curve_list = od.radial_critical_curve_area_list_from( grid=grid ) @@ -400,7 +400,7 @@ def test__tangential_critical_curve_area_list_from(): area_calc = np.pi * mp.einstein_radius**2 - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) area_within_tangential_critical_curve_list = ( od.tangential_critical_curve_area_list_from(grid=grid) ) @@ -415,7 +415,7 @@ def test__einstein_radius_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_radius_list = od.einstein_radius_list_from(grid=grid) assert einstein_radius_list[0] == pytest.approx(2.0, 1e-1) @@ -424,7 +424,7 @@ def test__einstein_radius_list_from(): centre=(0.0, 0.0), einstein_radius=2.0, ell_comps=(0.0, -0.25) ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_radius_list = od.einstein_radius_list_from(grid=grid) assert einstein_radius_list[0] == pytest.approx(1.9360, 1e-1) @@ -435,7 +435,7 @@ def test__einstein_radius_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_radius = od.einstein_radius_from(grid=grid) assert einstein_radius == pytest.approx(2.0, 1e-1) @@ -444,7 +444,7 @@ def test__einstein_radius_from(): centre=(0.0, 0.0), einstein_radius=2.0, ell_comps=(0.0, -0.25) ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_radius = od.einstein_radius_from(grid=grid) assert einstein_radius == pytest.approx(1.9360, 1e-1) @@ -455,7 +455,7 @@ def test__einstein_mass_angular_list_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_mass_angular_list = od.einstein_mass_angular_list_from(grid=grid) assert einstein_mass_angular_list[0] == pytest.approx(np.pi * 2.0**2.0, 1e-1) @@ -466,7 +466,7 @@ def test__einstein_mass_angular_from(): mp = ag.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=2.0) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) einstein_mass_angular = od.einstein_mass_angular_from(grid=grid) assert einstein_mass_angular == pytest.approx(np.pi * 2.0**2.0, 1e-1) @@ -487,7 +487,7 @@ def test__jacobian_from(): centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0 ) - od = OperateDeflections.from_mass_obj(mp) + od = LensCalc.from_mass_obj(mp) jacobian = od.jacobian_from(grid=grid) assert len(jacobian) == 2 diff --git a/test_autogalaxy/plot/mat_wrap/test_visuals.py b/test_autogalaxy/plot/mat_wrap/test_visuals.py index a3cdfb888..8051898ce 100644 --- a/test_autogalaxy/plot/mat_wrap/test_visuals.py +++ b/test_autogalaxy/plot/mat_wrap/test_visuals.py @@ -2,7 +2,7 @@ import pytest import autogalaxy.plot as aplt -from autogalaxy.operate.deflections import OperateDeflections +from autogalaxy.operate.lens_calc import LensCalc directory = path.dirname(path.realpath(__file__)) @@ -37,7 +37,7 @@ def test__1d__add_einstein_radius(mp_0, grid_2d_7x7): mass_obj=mp_0, grid=grid_2d_7x7 ) - assert visuals_1d_via.einstein_radius == OperateDeflections.from_mass_obj( + assert visuals_1d_via.einstein_radius == LensCalc.from_mass_obj( mp_0 ).einstein_radius_from(grid=grid_2d_7x7) @@ -48,7 +48,7 @@ def test__1d__add_einstein_radius_errors(mp_0, grid_2d_7x7): mass_obj_list=[mp_0, mp_0], grid=grid_2d_7x7, low_limit=1.0 ) - od = OperateDeflections.from_mass_obj(mp_0) + od = LensCalc.from_mass_obj(mp_0) assert visuals_1d_via.einstein_radius == od.einstein_radius_from(grid=grid_2d_7x7) assert visuals_1d_via.einstein_radius_errors[0][0] == od.einstein_radius_from( grid=grid_2d_7x7 @@ -61,7 +61,7 @@ def test__2d__add_critical_curve(gal_x1_mp, grid_2d_7x7): mass_obj=gal_x1_mp, grid=grid_2d_7x7, plane_index=0 ) - od = OperateDeflections.from_mass_obj(gal_x1_mp) + od = LensCalc.from_mass_obj(gal_x1_mp) assert ( visuals_2d_via.tangential_critical_curves[0] == od.tangential_critical_curve_list_from(grid=grid_2d_7x7)[0] @@ -74,7 +74,7 @@ def test__2d__add_caustic(gal_x1_mp, grid_2d_7x7): mass_obj=gal_x1_mp, grid=grid_2d_7x7, plane_index=1 ) - od = OperateDeflections.from_mass_obj(gal_x1_mp) + od = LensCalc.from_mass_obj(gal_x1_mp) assert ( visuals_2d_via.tangential_caustics[0] == od.tangential_caustic_list_from(grid=grid_2d_7x7)[0]