diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py index 5d1295df2..6703d30f4 100644 --- a/autogalaxy/__init__.py +++ b/autogalaxy/__init__.py @@ -25,6 +25,8 @@ mapper_from as Mapper, ) # noqa from autoarray.inversion.pixelization.border_relocator import BorderRelocator +from autoarray.preloads import Preloads +from autoarray.preloads import mapper_indices_from from autoarray.mask.mask_1d import Mask1D # noqa from autoarray.mask.mask_2d import Mask2D # noqa from autoarray.mask.derive.zoom_2d import Zoom2D @@ -51,6 +53,7 @@ from autoarray.structures.visibilities import Visibilities # noqa from autoarray.structures.visibilities import VisibilitiesNoiseMap # noqa +from .analysis import model_util from .analysis.adapt_images.adapt_images import AdaptImages from .analysis.adapt_images.adapt_image_maker import AdaptImageMaker from . import aggregator as agg @@ -80,7 +83,6 @@ from .galaxy.galaxy import Galaxy from .galaxy.galaxies import Galaxies from .galaxy.redshift import Redshift -from .galaxy.stellar_dark_decomp import StellarDarkDecomp from .galaxy.to_inversion import AbstractToInversion from .galaxy.to_inversion import GalaxiesToInversion from .profiles.geometry_profiles import EllProfile diff --git a/autogalaxy/analysis/adapt_images/adapt_images.py b/autogalaxy/analysis/adapt_images/adapt_images.py index c3cf8a80c..3788a6870 100644 --- a/autogalaxy/analysis/adapt_images/adapt_images.py +++ b/autogalaxy/analysis/adapt_images/adapt_images.py @@ -1,5 +1,4 @@ from __future__ import annotations -import jax.numpy as jnp import numpy as np from typing import TYPE_CHECKING, Dict, Optional, Tuple @@ -132,7 +131,7 @@ def from_result(cls, result, use_model_images: bool = False) -> "AdaptImages": else: galaxy_image = result.subtracted_signal_to_noise_map_galaxy_dict[path] - minimum_galaxy_value = adapt_minimum_percent * jnp.max(galaxy_image.array) + minimum_galaxy_value = adapt_minimum_percent * xp.max(galaxy_image.array) galaxy_image[galaxy_image < minimum_galaxy_value] = minimum_galaxy_value galaxy_name_image_dict[path] = galaxy_image diff --git a/autogalaxy/analysis/model_util.py b/autogalaxy/analysis/model_util.py new file mode 100644 index 000000000..9618a22ce --- /dev/null +++ b/autogalaxy/analysis/model_util.py @@ -0,0 +1,129 @@ +import numpy as np +from typing import Optional, Tuple + +import autofit as af + + +def mge_model_from( + mask_radius: float, + total_gaussians: int = 30, + gaussian_per_basis: int = 1, + centre_prior_is_uniform: bool = True, + centre: Tuple[float, float] = (0.0, 0.0), + centre_fixed: Optional[Tuple[float, float]] = None, + use_spherical: bool = False, +) -> af.Collection: + """ + Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light + + This model is designed as a "start here" configuration for lens modeling: + + - The lens and source light are represented by a Basis object composed of many + Gaussian light profiles with fixed logarithmically spaced widths (`sigma`). + - All Gaussians within each basis share common centres and ellipticity + components, reducing degeneracy while retaining flexibility. + + - Users can combine with a lens mass model of their choiuce. + + The resulting model provides a good balance of speed, flexibility, and accuracy + for fitting most galaxy-scale strong lenses. + + This code is mostly to make the API simple for new users, hiding the technical + details of setting up an MGE. More advanced users may wish to customize the + model further. + + Parameters + ---------- + mask_radius + The outer radius (in arcseconds) of the circular mask applied to the data. + This determines the maximum Gaussian width (`sigma`) used in the lens MGE. + lens_total_gaussians + Total number of Gaussian light profiles used in the lens MGE basis. + source_total_gaussians + Total number of Gaussian light profiles used in the source MGE basis. + lens_gaussian_per_basis + Number of separate Gaussian bases to include for the lens light profile. + Each basis has `lens_total_gaussians` components. + source_gaussian_per_basis + Number of separate Gaussian bases to include for the source light profile. + Each basis has `source_total_gaussians` components. + + Returns + ------- + model : af.Collection + An `autofit.Collection` containing: + - A lens galaxy at redshift 0.5, with: + * bulge light profile: MGE basis of Gaussians + * mass profile: Isothermal ellipsoid + * external shear + - A source galaxy at redshift 1.0, with: + * bulge light profile: MGE basis of Gaussians + + Notes + ----- + - Lens light Gaussians have widths (sigma) logarithmically spaced between 0.01" + and the mask radius. + - Source light Gaussians have widths logarithmically spaced between 0.01" and 1.0". + - Gaussian centres are free parameters but tied across all components in each + basis to reduce dimensionality. + - This function is a convenience utility: it hides the technical setup of MGE + composition and provides a ready-to-use lens model for quick experimentation. + """ + + from autogalaxy.profiles.light.linear import Gaussian, GaussianSph + from autogalaxy.profiles.basis import Basis + + # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0". + log10_sigma_list = np.linspace(-4, np.log10(mask_radius), total_gaussians) + + # By defining the centre here, it creates two free parameters that are assigned below to all Gaussians. + + if centre_fixed is not None: + centre_0 = centre[0] + centre_1 = centre[1] + elif centre_prior_is_uniform: + 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 + ) + else: + centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.3) + centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.3) + + if use_spherical: + model_cls = GaussianSph + else: + model_cls = Gaussian + + bulge_gaussian_list = [] + + for j in range(gaussian_per_basis): + # A list of Gaussian model components whose parameters are customized belows. + + gaussian_list = af.Collection( + af.Model(model_cls) for _ in range(total_gaussians) + ) + + # Iterate over every Gaussian and customize its parameters. + + for i, gaussian in enumerate(gaussian_list): + gaussian.centre.centre_0 = centre_0 # All Gaussians have same y centre. + gaussian.centre.centre_1 = centre_1 # All Gaussians have same x centre. + if not use_spherical: + gaussian.ell_comps = gaussian_list[ + 0 + ].ell_comps # All Gaussians have same elliptical components. + gaussian.sigma = ( + 10 ** log10_sigma_list[i] + ) # All Gaussian sigmas are fixed to values above. + + bulge_gaussian_list += gaussian_list + + # The Basis object groups many light profiles together into a single model component. + + return af.Model( + Basis, + profile_list=bulge_gaussian_list, + ) diff --git a/autogalaxy/analysis/plotter_interface.py b/autogalaxy/analysis/plotter_interface.py index d329702db..1ee0acde9 100644 --- a/autogalaxy/analysis/plotter_interface.py +++ b/autogalaxy/analysis/plotter_interface.py @@ -168,18 +168,6 @@ def should_plot(name): mat_plot_1d=mat_plot_1d, ) - try: - if should_plot("subplot_galaxies_1d"): - galaxies_plotter.subplot_galaxies_1d() - except OverflowError: - pass - - try: - if should_plot("subplot_galaxies_1d_decomposed"): - galaxies_plotter.subplot_galaxies_1d_decomposed() - except OverflowError: - pass - if should_plot("fits_galaxy_images"): image_list = [ diff --git a/autogalaxy/config/visualize/plots.yaml b/autogalaxy/config/visualize/plots.yaml index 3d82db942..b61299e42 100644 --- a/autogalaxy/config/visualize/plots.yaml +++ b/autogalaxy/config/visualize/plots.yaml @@ -30,8 +30,6 @@ fit_imaging: {} # Settings for plots of fits to imagi galaxies: # Settings for plots of galaxies (e.g. GalaxiesPlotter). subplot_galaxies: true # Plot subplot of all quantities in each galaxies group (e.g. images, convergence)? subplot_galaxy_images: false # Plot subplot of the image of each galaxy in the model? - subplot_galaxies_1d: false # Plot subplot of all quantities in 1D of each galaxies group (e.g. images, convergence)? - subplot_galaxies_1d_decomposed: false # Plot subplot of all quantities in 1D decomposed of each galaxies group (e.g. images, convergence)? fits_galaxy_images: false # Output a .fits file containing images of every galaxy? inversion: # Settings for plots of inversions (e.g. InversionPlotter). diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py index 3a199d220..5bda71120 100644 --- a/autogalaxy/convert.py +++ b/autogalaxy/convert.py @@ -1,10 +1,8 @@ +import numpy as np from typing import Tuple -import jax -import jax.numpy as jnp - -def ell_comps_from(axis_ratio: float, angle: float) -> Tuple[float, float]: +def ell_comps_from(axis_ratio: float, angle: float, xp=np) -> Tuple[float, float]: """ Returns the elliptical components e1 and e2 of a light or mass profile from an input angle in degrees and axis ratio. @@ -23,14 +21,16 @@ def ell_comps_from(axis_ratio: float, angle: float) -> Tuple[float, float]: angle Rotation angle of light profile counter-clockwise from positive x-axis. """ - angle *= jnp.pi / 180.0 + angle *= xp.pi / 180.0 fac = (1 - axis_ratio) / (1 + axis_ratio) - ellip_y = fac * jnp.sin(2 * angle) - ellip_x = fac * jnp.cos(2 * angle) + ellip_y = fac * xp.sin(2 * angle) + ellip_x = fac * xp.cos(2 * angle) return (ellip_y, ellip_x) -def axis_ratio_and_angle_from(ell_comps: Tuple[float, float]) -> Tuple[float, float]: +def axis_ratio_and_angle_from( + ell_comps: Tuple[float, float], xp=np +) -> Tuple[float, float]: """ Returns the axis-ratio and position angle in degrees (-45 < angle < 135.0) from input elliptical components e1 and e2 of a light or mass profile. @@ -60,19 +60,24 @@ def axis_ratio_and_angle_from(ell_comps: Tuple[float, float]) -> Tuple[float, fl ell_comps The elliptical components of the light or mass profile which are converted to an angle. """ - angle = jnp.arctan2(ell_comps[0], ell_comps[1]) / 2 - angle *= 180.0 / jnp.pi + angle = xp.arctan2(ell_comps[0], ell_comps[1]) / 2 + angle *= 180.0 / xp.pi - angle = jax.lax.select(angle < -45, angle + 180, angle) + angle = xp.where(angle < -45, angle + 180, angle) - fac = jnp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2) - fac = jax.lax.min(fac, 0.999) + fac = xp.sqrt(ell_comps[1] ** 2 + ell_comps[0] ** 2) + if xp.__name__.startswith("jax"): + import jax + + fac = jax.lax.min(fac, 0.999) + else: # NumPy + fac = np.minimum(fac, 0.999) axis_ratio = (1 - fac) / (1 + fac) return axis_ratio, angle -def axis_ratio_from(ell_comps: Tuple[float, float]): +def axis_ratio_from(ell_comps: Tuple[float, float], xp=np): """ Returns the axis-ratio from input elliptical components e1 and e2 of a light or mass profile. @@ -95,11 +100,11 @@ def axis_ratio_from(ell_comps: Tuple[float, float]): ell_comps The elliptical components of the light or mass profile which are converted to an angle. """ - axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps) + axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps, xp=xp) return axis_ratio -def angle_from(ell_comps: Tuple[float, float]) -> float: +def angle_from(ell_comps: Tuple[float, float], xp=np) -> float: """ Returns the position angle in degrees (-45 < angle < 135.0) from input elliptical components e1 and e2 of a light or mass profile. @@ -128,26 +133,25 @@ def angle_from(ell_comps: Tuple[float, float]) -> float: ell_comps The elliptical components of the light or mass profile which are converted to an angle. """ - axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps) - + axis_ratio, angle = axis_ratio_and_angle_from(ell_comps=ell_comps, xp=xp) return angle -def shear_gamma_1_2_from(magnitude: float, angle: float) -> Tuple[float, float]: +def shear_gamma_1_2_from(magnitude: float, angle: float, xp=np) -> Tuple[float, float]: """ Returns the shear gamma 1 and gamma 2 values an input shear magnitude and angle in degrees. The gamma 1 and gamma 2 components of a shear are given by: - gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0) - gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0) + gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0) + gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0) Which are the values this function returns. Converting from gamma 1 and gamma 2 to magnitude and angle is given by: - magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2) - angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi + magnitude = xp.sqrt(gamma_1**2 + gamma_2**2) + angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi Parameters ---------- @@ -156,26 +160,26 @@ def shear_gamma_1_2_from(magnitude: float, angle: float) -> Tuple[float, float]: angle Rotation angle of light profile counter-clockwise from positive x-axis. """ - gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0) - gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0) + gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0) + gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0) return (gamma_1, gamma_2) def shear_magnitude_and_angle_from( - gamma_1: float, gamma_2: float + gamma_1: float, gamma_2: float, xp=np ) -> Tuple[float, float]: """ Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values. The gamma 1 and gamma 2 components of a shear are given by: - gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0) - gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0) + gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0) + gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0) Converting from gamma 1 and gamma 2 to magnitude and angle is given by: - magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2) - angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi + magnitude = xp.sqrt(gamma_1**2 + gamma_2**2) + angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi Which are the values this function returns. @@ -193,30 +197,30 @@ def shear_magnitude_and_angle_from( gamma_2 The gamma 2 component of the shear. """ - angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi - magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2) + angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi + magnitude = xp.sqrt(gamma_1**2 + gamma_2**2) - angle = jnp.where(angle < 0, angle + 180.0, angle) - angle = jnp.where( - (jnp.abs(angle - 90.0) > 45.0) & (angle > 90.0), angle - 180.0, angle + angle = xp.where(angle < 0, angle + 180.0, angle) + angle = xp.where( + (xp.abs(angle - 90.0) > 45.0) & (angle > 90.0), angle - 180.0, angle ) return magnitude, angle -def shear_magnitude_from(gamma_1: float, gamma_2: float) -> float: +def shear_magnitude_from(gamma_1: float, gamma_2: float, xp=np) -> float: """ Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values. The gamma 1 and gamma 2 components of a shear are given by: - gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0) - gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0) + gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0) + gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0) Converting from gamma 1 and gamma 2 to magnitude and angle is given by: - magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2) - angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi + magnitude = xp.sqrt(gamma_1**2 + gamma_2**2) + angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi The magnitude value is what this function returns. @@ -227,23 +231,25 @@ def shear_magnitude_from(gamma_1: float, gamma_2: float) -> float: gamma_2 The gamma 2 component of the shear. """ - magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2) + magnitude, angle = shear_magnitude_and_angle_from( + gamma_1=gamma_1, gamma_2=gamma_2, xp=xp + ) return magnitude -def shear_angle_from(gamma_1: float, gamma_2: float) -> float: +def shear_angle_from(gamma_1: float, gamma_2: float, xp=np) -> float: """ Returns the shear magnitude and angle in degrees from input shear gamma 1 and gamma 2 values. The gamma 1 and gamma 2 components of a shear are given by: - gamma_1 = magnitude * jnp.cos(2 * angle * jnp.pi / 180.0) - gamma_2 = magnitude * jnp.sin(2 * angle * jnp.pi / 180.0) + gamma_1 = magnitude * xp.cos(2 * angle * xp.pi / 180.0) + gamma_2 = magnitude * xp.sin(2 * angle * xp.pi / 180.0) Converting from gamma 1 and gamma 2 to magnitude and angle is given by: - magnitude = jnp.sqrt(gamma_1**2 + gamma_2**2) - angle = jnp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / jnp.pi + magnitude = xp.sqrt(gamma_1**2 + gamma_2**2) + angle = xp.arctan2(gamma_2, gamma_1) / 2 * 180.0 / xp.pi The angle value is what this function returns. @@ -254,12 +260,14 @@ def shear_angle_from(gamma_1: float, gamma_2: float) -> float: gamma_2 The gamma 2 component of the shear. """ - magnitude, angle = shear_magnitude_and_angle_from(gamma_1=gamma_1, gamma_2=gamma_2) + magnitude, angle = shear_magnitude_and_angle_from( + gamma_1=gamma_1, gamma_2=gamma_2, xp=xp + ) return angle def multipole_k_m_and_phi_m_from( - multipole_comps: Tuple[float, float], m: int + multipole_comps: Tuple[float, float], m: int, xp=np ) -> Tuple[float, float]: """ Returns the multipole normalization value `k_m` and angle `phi` from the multipole component parameters. @@ -290,16 +298,18 @@ def multipole_k_m_and_phi_m_from( The normalization and angle parameters of the multipole. """ phi_m = ( - jnp.arctan2(multipole_comps[0], multipole_comps[1]) * 180.0 / jnp.pi / float(m) + xp.arctan2(multipole_comps[0], multipole_comps[1]) * 180.0 / xp.pi / float(m) ) - k_m = jnp.sqrt(multipole_comps[1] ** 2 + multipole_comps[0] ** 2) + k_m = xp.sqrt(multipole_comps[1] ** 2 + multipole_comps[0] ** 2) - phi_m = jnp.where(phi_m < -90.0 / m, phi_m + 360.0 / m, phi_m) + phi_m = xp.where(phi_m < -90.0 / m, phi_m + 360.0 / m, phi_m) return k_m, phi_m -def multipole_comps_from(k_m: float, phi_m: float, m: int) -> Tuple[float, float]: +def multipole_comps_from( + k_m: float, phi_m: float, m: int, xp=np +) -> Tuple[float, float]: """ Returns the multipole component parameters from their normalization value `k_m` and angle `phi`. @@ -323,7 +333,7 @@ def multipole_comps_from(k_m: float, phi_m: float, m: int) -> Tuple[float, float """ from astropy import units - multipole_comp_0 = k_m * jnp.sin(phi_m * float(m) * units.deg.to(units.rad)) - multipole_comp_1 = k_m * jnp.cos(phi_m * float(m) * units.deg.to(units.rad)) + multipole_comp_0 = k_m * xp.sin(phi_m * float(m) * units.deg.to(units.rad)) + multipole_comp_1 = k_m * xp.cos(phi_m * float(m) * units.deg.to(units.rad)) return (multipole_comp_0, multipole_comp_1) diff --git a/autogalaxy/ellipse/ellipse/ellipse.py b/autogalaxy/ellipse/ellipse/ellipse.py index 8986d085f..2ef381ce6 100644 --- a/autogalaxy/ellipse/ellipse/ellipse.py +++ b/autogalaxy/ellipse/ellipse/ellipse.py @@ -50,7 +50,7 @@ def ellipticity(self) -> float: """ The ellipticity of the ellipse, which is the factor by which the ellipse is offset from a circle. """ - return np.sqrt(1 - self.axis_ratio**2.0) + return np.sqrt(1 - self.axis_ratio() ** 2.0) @property def minor_axis(self): @@ -168,9 +168,9 @@ def ellipse_radii_from_major_axis_from( np.sqrt( np.add( self.major_axis**2.0 - * np.sin(angles_from_x0 - self.angle_radians) ** 2.0, + * np.sin(angles_from_x0 - self.angle_radians()) ** 2.0, self.minor_axis**2.0 - * np.cos(angles_from_x0 - self.angle_radians) ** 2.0, + * np.cos(angles_from_x0 - self.angle_radians()) ** 2.0, ) ), ) diff --git a/autogalaxy/ellipse/ellipse/ellipse_multipole.py b/autogalaxy/ellipse/ellipse/ellipse_multipole.py index 89ac08517..64b26e007 100644 --- a/autogalaxy/ellipse/ellipse/ellipse_multipole.py +++ b/autogalaxy/ellipse/ellipse/ellipse_multipole.py @@ -60,8 +60,10 @@ def points_perturbed_from( angles = ellipse.angles_from_x0_from(pixel_scale=pixel_scale, n_i=n_i) radial = np.add( - self.multipole_comps[1] * np.cos(self.m * (angles - ellipse.angle_radians)), - self.multipole_comps[0] * np.sin(self.m * (angles - ellipse.angle_radians)), + self.multipole_comps[1] + * np.cos(self.m * (angles - ellipse.angle_radians())), + self.multipole_comps[0] + * np.sin(self.m * (angles - ellipse.angle_radians())), ) x = points[:, 1] + (radial * np.cos(angles)) diff --git a/autogalaxy/ellipse/model/analysis.py b/autogalaxy/ellipse/model/analysis.py index 80046e6b9..5635d8c57 100644 --- a/autogalaxy/ellipse/model/analysis.py +++ b/autogalaxy/ellipse/model/analysis.py @@ -1,8 +1,6 @@ import logging -import time -from typing import Dict, List, Optional, Tuple - -from autoconf.fitsable import hdu_list_for_output_from +import numpy as np +from typing import List, Optional import autofit as af import autoarray as aa @@ -45,7 +43,7 @@ def __init__(self, dataset: aa.Imaging, title_prefix: str = None): self.dataset = dataset self.title_prefix = title_prefix - def log_likelihood_function(self, instance: af.ModelInstance) -> float: + def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float: """ Given an instance of the model, where the model parameters are set via a non-linear search, fit the model instance to the imaging dataset. diff --git a/autogalaxy/ellipse/plot/fit_ellipse_plotters.py b/autogalaxy/ellipse/plot/fit_ellipse_plotters.py index d4dbbae7a..58141729d 100644 --- a/autogalaxy/ellipse/plot/fit_ellipse_plotters.py +++ b/autogalaxy/ellipse/plot/fit_ellipse_plotters.py @@ -143,10 +143,8 @@ def __init__( super().__init__( mat_plot_1d=mat_plot_1d, visuals_1d=visuals_1d, - include_1d=include_1d, mat_plot_2d=mat_plot_2d, visuals_2d=visuals_2d, - include_2d=include_2d, ) self.fit_pdf_list = fit_pdf_list diff --git a/autogalaxy/galaxy/galaxies.py b/autogalaxy/galaxy/galaxies.py index 8b172f51a..c21f5a863 100644 --- a/autogalaxy/galaxy/galaxies.py +++ b/autogalaxy/galaxy/galaxies.py @@ -50,7 +50,7 @@ def redshift(self): return self[0].redshift def image_2d_list_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ) -> List[aa.Array2D]: """ Returns a list of the 2D images for each galaxy from a 2D grid of Cartesian (y,x) coordinates. @@ -87,13 +87,13 @@ def image_2d_list_from( therefore is used to pass the `operated_only` input to these methods. """ return [ - galaxy.image_2d_from(grid=grid, operated_only=operated_only) + galaxy.image_2d_from(grid=grid, xp=xp, operated_only=operated_only) for galaxy in self ] @aa.grid_dec.to_array def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ) -> aa.Array2D: """ Returns the 2D image of all galaxies summed from a 2D grid of Cartesian (y,x) coordinates. @@ -114,10 +114,12 @@ def image_2d_from( apply these operations to the images, which may have the `operated_only` input passed to them. This input therefore is used to pass the `operated_only` input to these methods. """ - return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only)) + return sum( + self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only) + ) def galaxy_image_2d_dict_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ) -> {Galaxy: np.ndarray}: """ Returns a dictionary associating every `Galaxy` object with its corresponding 2D image, using the instance @@ -143,7 +145,9 @@ def galaxy_image_2d_dict_from( galaxy_image_2d_dict = dict() - image_2d_list = self.image_2d_list_from(grid=grid, operated_only=operated_only) + image_2d_list = self.image_2d_list_from( + grid=grid, xp=xp, operated_only=operated_only + ) for galaxy_index, galaxy in enumerate(self): galaxy_image_2d_dict[galaxy] = image_2d_list[galaxy_index] @@ -151,7 +155,9 @@ def galaxy_image_2d_dict_from( return galaxy_image_2d_dict @aa.grid_dec.to_vector_yx - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def deflections_yx_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D deflections angles of all galaxies from a 2D grid of Cartesian (y,x) coordinates. @@ -172,17 +178,21 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarr grid The 2D (y, x) coordinates where values of the deflections are evaluated. """ - return sum(map(lambda g: g.deflections_yx_2d_from(grid=grid), self)) + return sum(map(lambda g: g.deflections_yx_2d_from(grid=grid, xp=xp), self)) @aa.grid_dec.to_grid - def traced_grid_2d_from(self, grid: aa.type.Grid2DLike) -> aa.type.Grid2DLike: + def traced_grid_2d_from( + self, grid: aa.type.Grid2DLike, xp=np + ) -> aa.type.Grid2DLike: """ Trace this plane's grid_stacks to the next plane, using its deflection angles. """ - return grid - self.deflections_yx_2d_from(grid=grid) + return grid - self.deflections_yx_2d_from(grid=grid, xp=xp) @aa.grid_dec.to_array - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def convergence_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D convergence of all galaxies from a 2D grid of Cartesian (y,x) coordinates. @@ -203,10 +213,12 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: grid The 2D (y, x) coordinates where values of the convergence are evaluated. """ - return sum(map(lambda g: g.convergence_2d_from(grid=grid), self)) + return sum(map(lambda g: g.convergence_2d_from(grid=grid, xp=xp), self)) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def potential_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D potential of all galaxies from a 2D grid of Cartesian (y,x) coordinates. @@ -227,7 +239,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: grid The 2D (y, x) coordinates where values of the potential are evaluated. """ - return sum(map(lambda g: g.potential_2d_from(grid=grid), self)) + return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self)) def has(self, cls: Union[Type, Tuple[Type]]) -> bool: """ diff --git a/autogalaxy/galaxy/galaxy.py b/autogalaxy/galaxy/galaxy.py index 0a924c496..2558b5410 100644 --- a/autogalaxy/galaxy/galaxy.py +++ b/autogalaxy/galaxy/galaxy.py @@ -1,6 +1,5 @@ from typing import Dict, List, Optional, Type, Union -import jax.numpy as jnp import numpy as np from autoconf.dictable import instance_as_dict, to_dict @@ -161,7 +160,7 @@ def grid_radial_from(self, grid, centre, angle): ) def image_2d_list_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ) -> List[aa.Array2D]: """ Returns a list of the 2D images of the galaxy's light profiles from a 2D grid of Cartesian (y,x) coordinates. @@ -194,7 +193,7 @@ def image_2d_list_from( zeros. """ return [ - light_profile.image_2d_from(grid=grid, operated_only=operated_only) + light_profile.image_2d_from(grid=grid, xp=xp, operated_only=operated_only) for light_profile in self.cls_list_from( cls=LightProfile, cls_filtered=LightProfileLinear ) @@ -202,7 +201,10 @@ def image_2d_list_from( @aa.grid_dec.to_array def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, ) -> Union[np.ndarray, aa.Array2D]: """ Returns the 2D image of all galaxy light profiles summed from a 2D grid of Cartesian (y,x) coordinates. @@ -228,44 +230,15 @@ def image_2d_from( len(self.cls_list_from(cls=LightProfile, cls_filtered=LightProfileLinear)) > 0 ): - return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only)) - return jnp.zeros((grid.shape[0],)) - - @aa.grid_dec.to_projected - def image_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray: - """ - Returns the summed 1D image of the galaxy's light profiles using a grid of Cartesian (y,x) coordinates. - - If the galaxy has no light profiles, a grid of zeros is returned. - - See `profiles.light` module for details of how this is performed. - - The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data - structure using the input `grid`'s attributes. - - Parameters - ---------- - grid - The 1D (x,) coordinates where values of the image are evaluated. - """ - if self.has(cls=LightProfile): - image_1d_list = [] - - for light_profile in self.cls_list_from( - cls=LightProfile, cls_filtered=LightProfileLinear - ): - grid_radial = self.grid_radial_from( - grid=grid, centre=light_profile.centre, angle=light_profile.angle - ) - - image_1d_list.append(light_profile.image_1d_from(grid=grid_radial)) - - return sum(image_1d_list) - - return jnp.zeros((grid.shape[0],)) + return sum( + self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only) + ) + return xp.zeros((grid.shape[0],)) @aa.grid_dec.to_vector_yx - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def deflections_yx_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D deflection angles of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates. @@ -286,15 +259,17 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarr if self.has(cls=MassProfile): return sum( map( - lambda p: p.deflections_yx_2d_from(grid=grid), + lambda p: p.deflections_yx_2d_from(grid=grid, xp=xp), self.cls_list_from(cls=MassProfile), ) ) - return jnp.zeros((grid.shape[0], 2)) + return xp.zeros((grid.shape[0], 2)) @aa.grid_dec.to_array - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def convergence_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D convergence of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates. @@ -314,65 +289,35 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: if self.has(cls=MassProfile): return sum( map( - lambda p: p.convergence_2d_from(grid=grid), + lambda p: p.convergence_2d_from(grid=grid, xp=xp), self.cls_list_from(cls=MassProfile), ) ) - return jnp.zeros((grid.shape[0],)) + return xp.zeros((grid.shape[0],)) @aa.grid_dec.to_grid - def traced_grid_2d_from(self, grid: aa.type.Grid2DLike) -> aa.type.Grid2DLike: + def traced_grid_2d_from( + self, grid: aa.type.Grid2DLike, xp=np + ) -> aa.type.Grid2DLike: """ Trace an input grid using the galaxy's its deflection angles. """ if isinstance(grid, aa.Grid2D): return aa.Grid2D( - values=grid - self.deflections_yx_2d_from(grid=grid), + values=grid - self.deflections_yx_2d_from(grid=grid, xp=xp), mask=grid.mask, over_sample_size=grid.over_sample_size, over_sampled=grid.over_sampled - - self.deflections_yx_2d_from(grid=grid.over_sampled), + - self.deflections_yx_2d_from(grid=grid.over_sampled, xp=xp), ) - return grid - self.deflections_yx_2d_from(grid=grid) - - @aa.grid_dec.to_projected - def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> np.ndarray: - """ - Returns the summed 1D convergence of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates. - - If the galaxy has no mass profiles, a grid of zeros is returned. - - See `profiles.mass` module for details of how this is performed. - - The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data - structure using the input `grid`'s attributes. - - Parameters - ---------- - grid - The 1D (x,) coordinates where values of the convergence are evaluated. - """ - if self.has(cls=MassProfile): - convergence_1d_list = [] - - for mass_profile in self.cls_list_from(cls=MassProfile): - - grid_radial = self.grid_radial_from( - grid=grid, centre=mass_profile.centre, angle=mass_profile.angle - ) - - convergence_1d_list.append( - mass_profile.convergence_1d_from(grid=grid_radial) - ) - - return sum(convergence_1d_list) - - return jnp.zeros((grid.shape[0],)) + return grid - self.deflections_yx_2d_from(grid=grid, xp=xp) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def potential_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the summed 2D potential of the galaxy's mass profiles from a 2D grid of Cartesian (y,x) coordinates. @@ -396,40 +341,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: self.cls_list_from(cls=MassProfile), ) ) - return jnp.zeros((grid.shape[0],)) - - @aa.grid_dec.to_projected - def potential_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray: - """ - Returns the summed 1D potential of the galaxy's mass profiles using a grid of Cartesian (y,x) coordinates. - - If the galaxy has no mass profiles, a grid of zeros is returned. - - See `profiles.mass` module for details of how this is performed. - - The decorator `to_projected` converts the output arrays from ndarrays to an `Array1D` data - structure using the input `grid`'s attributes. - - Parameters - ---------- - grid - The 1D (x,) coordinates where values of the potential are evaluated. - """ - if self.has(cls=MassProfile): - potential_1d_list = [] - - for mass_profile in self.cls_list_from(cls=MassProfile): - grid_radial = self.grid_radial_from( - grid=grid, centre=mass_profile.centre, angle=mass_profile.angle - ) - - potential_1d_list.append( - mass_profile.potential_1d_from(grid=grid_radial) - ) - - return sum(potential_1d_list) - - return jnp.zeros((grid.shape[0],)) + return xp.zeros((grid.shape[0],)) @property def half_light_radius(self): diff --git a/autogalaxy/galaxy/plot/galaxies_plotters.py b/autogalaxy/galaxy/plot/galaxies_plotters.py index 56522360f..e49a8e1a7 100644 --- a/autogalaxy/galaxy/plot/galaxies_plotters.py +++ b/autogalaxy/galaxy/plot/galaxies_plotters.py @@ -319,54 +319,3 @@ def subplot_galaxy_images(self): auto_filename=f"subplot_galaxy_images" ) self.close_subplot_figure() - - def subplot_galaxies_1d(self): - """ - Output a subplot of attributes of every individual 1D attribute of the `Galaxy` object. - - For example, a 1D plot showing how the image, convergence of each component varies radially outwards. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from the light profile. This is performed by aligning a 1D grid with the major-axis of the light - profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - """ - number_subplots = len(self.galaxies) * 3 - - self.open_subplot_figure(number_subplots=number_subplots) - - for galaxy_index in range(0, len(self.galaxies)): - galaxy_plotter = self.galaxy_plotter_from(galaxy_index=galaxy_index) - - galaxy_plotter.figures_1d(image=True) - galaxy_plotter.figures_1d(convergence=True) - galaxy_plotter.figures_1d(potential=True) - - self.mat_plot_1d.output.subplot_to_figure(auto_filename="subplot_galaxies_1d") - self.close_subplot_figure() - - def subplot_galaxies_1d_decomposed(self): - """ - Output a subplot of attributes of every individual 1D attribute of the `Galaxy` object decompoed into - their different light and mass profiles. - - For example, a 1D plot showing how the image, convergence of each component varies radially outwards. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from the light profile. This is performed by aligning a 1D grid with the major-axis of the light - profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - """ - number_subplots = len(self.galaxies) * 3 - - self.open_subplot_figure(number_subplots=number_subplots) - - for galaxy_index in range(0, len(self.galaxies)): - galaxy_plotter = self.galaxy_plotter_from(galaxy_index=galaxy_index) - - galaxy_plotter.figures_1d_decomposed(image=True) - galaxy_plotter.figures_1d_decomposed(convergence=True) - galaxy_plotter.figures_1d_decomposed(potential=True) - - self.mat_plot_1d.output.subplot_to_figure( - auto_filename="subplot_galaxies_1d_decomposed" - ) - self.close_subplot_figure() diff --git a/autogalaxy/galaxy/plot/galaxy_plotters.py b/autogalaxy/galaxy/plot/galaxy_plotters.py index f03fbe83d..d1a372831 100644 --- a/autogalaxy/galaxy/plot/galaxy_plotters.py +++ b/autogalaxy/galaxy/plot/galaxy_plotters.py @@ -18,12 +18,9 @@ if TYPE_CHECKING: from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePlotter - from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePDFPlotter from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePlotter -from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePDFPlotter from autogalaxy import exc -from autogalaxy.util import error_util class GalaxyPlotter(Plotter): @@ -170,209 +167,6 @@ def mass_profile_plotter_from( + Visuals1D().add_einstein_radius(mass_obj=mass_profile, grid=self.grid), ) - @property - def decomposed_light_profile_plotter_list(self): - plotter_list = [] - - for i, light_profile in enumerate(self.galaxy.cls_list_from(cls=LightProfile)): - light_profile_plotter = self.light_profile_plotter_from( - light_profile=light_profile, one_d_only=True - ) - - plotter_list.append(light_profile_plotter) - - return [self] + plotter_list - - @property - def decomposed_mass_profile_plotter_list(self): - plotter_list = [] - - for i, mass_profile in enumerate(self.galaxy.cls_list_from(cls=MassProfile)): - mass_profile_plotter = self.mass_profile_plotter_from( - mass_profile=mass_profile, one_d_only=True - ) - - plotter_list.append(mass_profile_plotter) - - return [self] + plotter_list - - def figures_1d( - self, image: bool = False, convergence: bool = False, potential: bool = False - ): - """ - Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's - grid object. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of - each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned - with the major-axis. - - The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of - the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - convergence - Whether to make a 1D plot (via `plot`) of the convergence. - potential - Whether to make a 1D plot (via `plot`) of the potential. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if image: - image_1d = self.galaxy.image_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=image_1d, - x=image_1d.grid_radial, - visuals_1d=self.visuals_1d - + Visuals1D().add_half_light_radius(self.galaxy), - auto_labels=aplt.AutoLabels( - title="Image vs Radius", - ylabel="Image ", - xlabel="Radius", - legend=self.galaxy.__class__.__name__, - filename="image_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if convergence: - convergence_1d = self.galaxy.convergence_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=convergence_1d, - x=convergence_1d.grid_radial, - visuals_1d=self.visuals_1d - + Visuals1D().add_einstein_radius(mass_obj=self.galaxy, grid=self.grid), - auto_labels=aplt.AutoLabels( - title="Convergence vs Radius", - ylabel="Convergence ", - xlabel="Radius", - legend=self.galaxy.__class__.__name__, - filename="convergence_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if potential: - potential_1d = self.galaxy.potential_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=potential_1d, - x=potential_1d.grid_radial, - visuals_1d=self.visuals_1d, - auto_labels=aplt.AutoLabels( - title="Potential vs Radius", - ylabel="Potential ", - xlabel="Radius", - legend=self.galaxy.__class__.__name__, - filename="potential_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - def figures_1d_decomposed( - self, - image: bool = False, - convergence: bool = False, - potential: bool = False, - legend_labels: List[str] = None, - ): - """ - Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's - grid object. - - This function makes a decomposed plot shows the 1D plot of the attribute for every light or mass profile in - the galaxy, as well as their combined 1D plot. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of - each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned - with the major-axis. - - The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of - the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - legend_labels - Manually overrides the labels of the plot's legend. - """ - - if self.galaxy.has(cls=LightProfile): - multi_plotter = aplt.MultiYX1DPlotter( - plotter_list=self.decomposed_light_profile_plotter_list, - legend_labels=legend_labels, - ) - multi_plotter.plotter_list[0].mat_plot_1d.output = self.mat_plot_1d.output - - if image: - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="image_1d_decomposed" - ) - change_filename = True - - multi_plotter.figure_1d(func_name="figures_1d", figure_name="image") - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) - - if self.galaxy.has(cls=MassProfile): - multi_plotter = aplt.MultiYX1DPlotter( - plotter_list=self.decomposed_mass_profile_plotter_list, - legend_labels=legend_labels, - ) - - if convergence: - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="convergence_1d_decomposed" - ) - change_filename = True - - multi_plotter.figure_1d( - func_name="figures_1d", figure_name="convergence" - ) - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) - - if potential: - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="potential_1d_decomposed" - ) - change_filename = True - - multi_plotter.figure_1d(func_name="figures_1d", figure_name="potential") - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) - def figures_2d( self, image: bool = False, @@ -499,435 +293,3 @@ def subplot_of_mass_profiles( self.subplot_of_plotters_figure( plotter_list=mass_profile_plotters, name="deflections_x" ) - - -class GalaxyPDFPlotter(GalaxyPlotter): - def __init__( - self, - galaxy_pdf_list: List[Galaxy], - grid: aa.Grid2D, - mat_plot_1d: MatPlot1D = None, - visuals_1d: Visuals1D = None, - mat_plot_2d: MatPlot2D = None, - visuals_2d: Visuals2D = None, - sigma: Optional[float] = 3.0, - ): - """ - Plots the attributes of a list of `GalaxyProfile` objects using the matplotlib methods `plot()` and `imshow()` - and many other matplotlib functions which customize the plot's appearance. - - Figures plotted by this object average over a list galaxy profiles to computed the average value of each - attribute with errors, where the 1D regions within the errors are plotted as a shaded region to show the range - of plausible models. Therefore, the input list of galaxies is expected to represent the probability density - function of an inferred model-fit. - - The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default, - the settings passed to every matplotlib function called are those specified in - the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to - customize the figure's appearance. - - Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be - extracted from the `GalaxyProfile` and plotted via the visuals object. - - Parameters - ---------- - galaxy_profile_pdf_list - The list of galaxy profiles whose mean and error values the plotter plots. - grid - The 2D (y,x) grid of coordinates used to evaluate the galaxy profile quantities that are plotted. - mat_plot_1d - Contains objects which wrap the matplotlib function calls that make 1D plots. - visuals_1d - Contains 1D visuals that can be overlaid on 1D plots. - mat_plot_2d - Contains objects which wrap the matplotlib function calls that make 2D plots. - visuals_2d - Contains 2D visuals that can be overlaid on 2D plots. - sigma - The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of - sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF). - """ - super().__init__( - galaxy=None, - grid=grid, - mat_plot_2d=mat_plot_2d, - visuals_2d=visuals_2d, - mat_plot_1d=mat_plot_1d, - visuals_1d=visuals_1d, - ) - - self.galaxy_pdf_list = galaxy_pdf_list - self.sigma = sigma - self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2 - - @property - def light_profile_pdf_plotter_list(self) -> List[LightProfilePDFPlotter]: - """ - Returns a list of `LightProfilePDFPlotter` objects from the list of galaxies in this object. These are - typically used for plotting the individual average value plus errors of the light profiles of the - plotter's `Galaxy` (e.g. in the function `figures_1d_decomposed`). - - Returns - ------- - List[LightProfilePDFPlotter] - An object that plots the average value and errors of a list of light profiles, often used for plotting - attributes of the galaxy. - """ - return [ - self.light_profile_pdf_plotter_from(index=index) - for index in range( - len(self.galaxy_pdf_list[0].cls_list_from(cls=LightProfile)) - ) - ] - - def light_profile_pdf_plotter_from(self, index) -> LightProfilePDFPlotter: - """ - Returns the `LightProfilePDFPlotter` of a specific light profile in this plotter's list of galaxies. This is - typically used for plotting the individual average value plus errors of a light profile in plotter's galaxy - list (e.g. in the function `figures_1d_decomposed`). - - Returns - ------- - LightProfilePDFPlotter - An object that plots the average value and errors of a list of light profiles, often used for plotting - attributes of the galaxy. - """ - - from autogalaxy.profiles.plot.light_profile_plotters import ( - LightProfilePDFPlotter, - ) - - light_profile_pdf_list = [ - galaxy.cls_list_from(cls=LightProfile)[index] - for galaxy in self.galaxy_pdf_list - ] - - return LightProfilePDFPlotter( - light_profile_pdf_list=light_profile_pdf_list, - grid=self.grid, - mat_plot_2d=self.mat_plot_2d, - visuals_2d=self.visuals_2d, - mat_plot_1d=self.mat_plot_1d, - visuals_1d=self.visuals_1d, - ) - - @property - def mass_profile_pdf_plotter_list(self) -> List[MassProfilePDFPlotter]: - """ - Returns a list of `MassProfilePDFPlotter` objects from the list of galaxies in this object. These are - typically used for plotting the individual average value plus errors of the mass profiles of the - plotter's `Galaxy` (e.g. in the function `figures_1d_decomposed`). - - Returns - ------- - List[MassProfilePDFPlotter] - An object that plots the average value and errors of a list of mass profiles, often used for plotting - attributes of the galaxy. - """ - return [ - self.mass_profile_pdf_plotter_from(index=index) - for index in range( - len(self.galaxy_pdf_list[0].cls_list_from(cls=MassProfile)) - ) - ] - - def mass_profile_pdf_plotter_from(self, index) -> MassProfilePDFPlotter: - """ - Returns the `MassProfilePDFPlotter` of a specific mass profile in this plotter's list of galaxies. This is - typically used for plotting the individual average value plus errors of a mass profile in plotter's galaxy - list (e.g. in the function `figures_1d_decomposed`). - - Returns - ------- - MassProfilePDFPlotter - An object that plots the average value and errors of a list of mass profiles, often used for plotting - attributes of the galaxy. - """ - mass_profile_pdf_list = [ - galaxy.cls_list_from(cls=MassProfile)[index] - for galaxy in self.galaxy_pdf_list - ] - - return MassProfilePDFPlotter( - mass_profile_pdf_list=mass_profile_pdf_list, - grid=self.grid, - mat_plot_2d=self.mat_plot_2d, - visuals_2d=self.visuals_2d, - mat_plot_1d=self.mat_plot_1d, - visuals_1d=self.visuals_1d, - ) - - def figures_1d( - self, image: bool = False, convergence: bool = False, potential: bool = False - ): - """ - Plots the individual attributes of the plotter's list of `Galaxy` object in 1D, which are computed via the - plotter's grid object. - - This averages over a list galaxies to compute the average value of each attribute with errors, where the - 1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore, - the input list of galaxies is expected to represent the probability density function of an inferred model-fit. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of - each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned - with the major-axis. - - The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of - the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if image: - image_1d_list = [ - galaxy.image_1d_from(grid=self.grid) for galaxy in self.galaxy_pdf_list - ] - - min_index = min([image_1d.shape[0] for image_1d in image_1d_list]) - image_1d_list = [image_1d[0:min_index] for image_1d in image_1d_list] - - ( - median_image_1d, - errors_image_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=image_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_light_obj_list = Visuals1D().add_half_light_radius_errors( - light_obj_list=self.galaxy_pdf_list, low_limit=self.low_limit - ) - visuals_1d_with_shaded_region = self.visuals_1d.__class__( - shaded_region=errors_image_1d - ) - - visuals_1d = visuals_1d_via_light_obj_list + visuals_1d_with_shaded_region - - median_image_1d = aa.Array1D.no_mask( - values=median_image_1d, pixel_scales=self.grid.pixel_scale - ) - - self.mat_plot_1d.plot_yx( - y=median_image_1d, - x=median_image_1d.grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title="Image vs Radius", - ylabel="Image ", - xlabel="Radius", - legend=self.galaxy_pdf_list[0].__class__.__name__, - filename="image_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if convergence: - convergence_1d_list = [ - galaxy.convergence_1d_from(grid=self.grid) - for galaxy in self.galaxy_pdf_list - ] - - min_index = min( - [convergence_1d.shape[0] for convergence_1d in convergence_1d_list] - ) - convergence_1d_list = [ - convergence_1d[0:min_index] for convergence_1d in convergence_1d_list - ] - - ( - median_convergence_1d, - errors_convergence_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=convergence_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors( - mass_obj_list=self.galaxy_pdf_list, - grid=self.grid, - low_limit=self.low_limit, - ) - visuals_1d_with_shaded_region = self.visuals_1d.__class__( - shaded_region=errors_convergence_1d - ) - - visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region - - median_convergence_1d = aa.Array1D.no_mask( - values=median_convergence_1d, pixel_scales=self.grid.pixel_scale - ) - - self.mat_plot_1d.plot_yx( - y=median_convergence_1d, - x=median_convergence_1d.grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title="Convergence vs Radius", - ylabel="Convergence ", - xlabel="Radius", - legend=self.galaxy_pdf_list[0].__class__.__name__, - filename="convergence_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if potential: - potential_1d_list = [ - galaxy.potential_1d_from(grid=self.grid) - for galaxy in self.galaxy_pdf_list - ] - - min_index = min( - [potential_1d.shape[0] for potential_1d in potential_1d_list] - ) - potential_1d_list = [ - potential_1d[0:min_index] for potential_1d in potential_1d_list - ] - - ( - median_potential_1d, - errors_potential_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=potential_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors( - mass_obj_list=self.galaxy_pdf_list, - grid=self.grid, - low_limit=self.low_limit, - ) - visuals_1d_with_shaded_region = self.visuals_1d.__class__( - shaded_region=errors_potential_1d - ) - - visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region - - median_potential_1d = aa.Array1D.no_mask( - values=median_potential_1d, pixel_scales=self.grid.pixel_scale - ) - - self.mat_plot_1d.plot_yx( - y=median_potential_1d, - x=median_potential_1d.grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title="Potential vs Radius", - ylabel="Potential ", - xlabel="Radius", - legend=self.galaxy_pdf_list[0].__class__.__name__, - filename="potential_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - def figures_1d_decomposed( - self, - image: bool = False, - convergence: bool = False, - potential: bool = False, - legend_labels: List[str] = None, - ): - """ - Plots the individual attributes of the plotter's `Galaxy` object in 1D, which are computed via the plotter's - grid object. - - This averages over a list galaxies to compute the average value of each attribute with errors, where the - 1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore, - the input list of galaxies is expected to represent the probability density function of an inferred model-fit. - - This function makes a decomposed plot showing the 1D plot of each attribute for every light or mass profile in - the galaxy, as well as their combined 1D plot. By plotting the attribute of each profile on the same figure, - one can see how much each profile contributes to the galaxy overall. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each light profile of the galaxy. This is performed by aligning a 1D grid with the major-axis of - each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - This means that the summed 1D profile of a galaxy's quantity is the sum of each individual component aligned - with the major-axis. - - The API is such that every plottable attribute of the `Galaxy` object is an input parameter of type bool of - the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - legend_labels - Manually overrides the labels of the plot's legend. - """ - if image: - multi_plotter = aplt.MultiYX1DPlotter( - plotter_list=[self] + self.light_profile_pdf_plotter_list, - legend_labels=legend_labels, - ) - - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="image_1d_decomposed" - ) - - change_filename = True - - multi_plotter.figure_1d(func_name="figures_1d", figure_name="image") - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) - - if convergence: - multi_plotter = aplt.MultiYX1DPlotter( - plotter_list=[self] + self.mass_profile_pdf_plotter_list, - legend_labels=legend_labels, - ) - - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="convergence_1d_decomposed" - ) - - change_filename = True - - multi_plotter.figure_1d(func_name="figures_1d", figure_name="convergence") - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) - - if potential: - multi_plotter = aplt.MultiYX1DPlotter( - plotter_list=[self] + self.mass_profile_pdf_plotter_list, - legend_labels=legend_labels, - ) - - change_filename = False - - if multi_plotter.plotter_list[0].mat_plot_1d.output.filename is None: - multi_plotter.plotter_list[0].set_filename( - filename="potential_1d_decomposed" - ) - - change_filename = True - - multi_plotter.figure_1d(func_name="figures_1d", figure_name="potential") - - if change_filename: - multi_plotter.plotter_list[0].set_filename(filename=None) diff --git a/autogalaxy/galaxy/stellar_dark_decomp.py b/autogalaxy/galaxy/stellar_dark_decomp.py deleted file mode 100644 index 4e5a169b4..000000000 --- a/autogalaxy/galaxy/stellar_dark_decomp.py +++ /dev/null @@ -1,44 +0,0 @@ -from autogalaxy import exc -from autogalaxy.profiles.mass.dark.abstract import DarkProfile -from autogalaxy.profiles.mass.stellar.abstract import StellarProfile - - -class StellarDarkDecomp: - def __init__(self, galaxy): - self.galaxy = galaxy - - def stellar_mass_angular_within_circle_from(self, radius: float): - if self.galaxy.has(cls=StellarProfile): - return sum( - [ - profile.mass_angular_within_circle_from(radius=radius) - for profile in self.galaxy.cls_list_from(cls=StellarProfile) - ] - ) - else: - raise exc.GalaxyException( - "You cannot perform a stellar mass-based calculation on a galaxy which does not have a stellar " - "mass-profile " - ) - - def dark_mass_angular_within_circle_from(self, radius: float): - if self.galaxy.has(cls=DarkProfile): - return sum( - [ - profile.mass_angular_within_circle_from(radius=radius) - for profile in self.galaxy.cls_list_from(cls=DarkProfile) - ] - ) - else: - raise exc.GalaxyException( - "You cannot perform a dark mass-based calculation on a galaxy which does not have a dark mass-profile" - ) - - def stellar_fraction_at_radius_from(self, radius): - return 1.0 - self.dark_fraction_at_radius_from(radius=radius) - - def dark_fraction_at_radius_from(self, radius): - stellar_mass = self.stellar_mass_angular_within_circle_from(radius=radius) - dark_mass = self.dark_mass_angular_within_circle_from(radius=radius) - - return dark_mass / (stellar_mass + dark_mass) diff --git a/autogalaxy/galaxy/to_inversion.py b/autogalaxy/galaxy/to_inversion.py index 353e09b66..cf86a2c83 100644 --- a/autogalaxy/galaxy/to_inversion.py +++ b/autogalaxy/galaxy/to_inversion.py @@ -1,4 +1,5 @@ from __future__ import annotations +import numpy as np from typing import Dict, List, Optional, Type, Union from autoconf import cached_property @@ -23,6 +24,7 @@ def __init__( dataset: Optional[Union[aa.Imaging, aa.Interferometer, aa.DatasetInterface]], adapt_images: Optional[AdaptImages] = None, settings_inversion: aa.SettingsInversion = aa.SettingsInversion(), + xp=np, ): """ Abstract class which interfaces a dataset and input modeling object (e.g. galaxies, a tracer) with the @@ -74,6 +76,7 @@ def __init__( self.adapt_images = adapt_images self.settings_inversion = settings_inversion + self._xp = xp @property def psf(self) -> Optional[aa.Kernel2D]: @@ -188,6 +191,7 @@ def __init__( galaxies: List[Galaxy], adapt_images: Optional[AdaptImages] = None, settings_inversion: aa.SettingsInversion = aa.SettingsInversion(), + xp=np, ): """ Interfaces a dataset and input list of galaxies with the inversion module. to setup a @@ -229,6 +233,7 @@ def __init__( dataset=dataset, adapt_images=adapt_images, settings_inversion=settings_inversion, + xp=xp, ) @property @@ -304,6 +309,7 @@ def cls_light_profile_func_list_galaxy_dict_from( psf=self.dataset.psf, light_profile_list=light_profile_list, regularization=light_profile.regularization, + xp=self._xp, ) lp_linear_func_galaxy_dict[lp_linear_func] = galaxy @@ -482,11 +488,11 @@ def mapper_from( source_plane_mesh_grid=source_plane_mesh_grid, image_plane_mesh_grid=image_plane_mesh_grid, adapt_data=adapt_galaxy_image, + xp=self._xp, ) return mapper_from( - mapper_grids=mapper_grids, - regularization=regularization, + mapper_grids=mapper_grids, regularization=regularization, xp=self._xp ) @cached_property @@ -568,6 +574,7 @@ def inversion(self) -> aa.AbstractInversion: dataset=self.dataset, linear_obj_list=self.linear_obj_list, settings=self.settings_inversion, + xp=self._xp, ) inversion.linear_obj_galaxy_dict = self.linear_obj_galaxy_dict diff --git a/autogalaxy/imaging/fit_imaging.py b/autogalaxy/imaging/fit_imaging.py index 9fda04b32..d08a2d5cd 100644 --- a/autogalaxy/imaging/fit_imaging.py +++ b/autogalaxy/imaging/fit_imaging.py @@ -25,6 +25,7 @@ def __init__( dataset_model: Optional[aa.DatasetModel] = None, adapt_images: Optional[AdaptImages] = None, settings_inversion: aa.SettingsInversion = aa.SettingsInversion(), + xp=np, ): """ Fits an imaging dataset using a list of galaxies. @@ -70,6 +71,7 @@ def __init__( super().__init__( dataset=dataset, dataset_model=dataset_model, + xp=xp, ) AbstractFitInversion.__init__( self=self, @@ -94,12 +96,14 @@ def blurred_image(self) -> aa.Array2D: ): return self.galaxies.image_2d_from( grid=self.grids.lp, + xp=self._xp, ) return self.galaxies.blurred_image_2d_from( grid=self.grids.lp, psf=self.dataset.psf, blurring_grid=self.grids.blurring, + xp=self._xp, ) @property @@ -124,6 +128,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion: galaxies=self.galaxies, adapt_images=self.adapt_images, settings_inversion=self.settings_inversion, + xp=self._xp, ) @cached_property diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py index 90e3c3e47..1620e4a76 100644 --- a/autogalaxy/imaging/model/analysis.py +++ b/autogalaxy/imaging/model/analysis.py @@ -1,3 +1,4 @@ +import numpy as np from typing import Optional import autofit as af @@ -21,6 +22,7 @@ def __init__( adapt_image_maker: Optional[AdaptImageMaker] = None, cosmology: LensingCosmology = None, settings_inversion: aa.SettingsInversion = None, + preloads: aa.Preloads = None, title_prefix: str = None, ): """ @@ -58,6 +60,7 @@ def __init__( adapt_image_maker=adapt_image_maker, cosmology=cosmology, settings_inversion=settings_inversion, + preloads=preloads, title_prefix=title_prefix, ) @@ -88,7 +91,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection): return self - def log_likelihood_function(self, instance: af.ModelInstance) -> float: + def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float: """ Given an instance of the model, where the model parameters are set via a non-linear search, fit the model instance to the imaging dataset. @@ -125,12 +128,9 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float: float The log likelihood indicating how well this model instance fitted the imaging data. """ - return self.fit_from(instance=instance).figure_of_merit + return self.fit_from(instance=instance, xp=xp).figure_of_merit - def fit_from( - self, - instance: af.ModelInstance, - ) -> FitImaging: + def fit_from(self, instance: af.ModelInstance, xp=np) -> FitImaging: """ Given a model instance create a `FitImaging` object. @@ -165,6 +165,7 @@ def fit_from( dataset_model=dataset_model, adapt_images=adapt_images, settings_inversion=self.settings_inversion, + xp=xp, ) def save_attributes(self, paths: af.DirectoryPaths): diff --git a/autogalaxy/imaging/model/plotter_interface.py b/autogalaxy/imaging/model/plotter_interface.py index 5339565a5..e562be12a 100644 --- a/autogalaxy/imaging/model/plotter_interface.py +++ b/autogalaxy/imaging/model/plotter_interface.py @@ -117,8 +117,10 @@ def should_plot(name): dataset.data.native_for_fits, dataset.noise_map.native_for_fits, dataset.psf.native_for_fits, - dataset.grids.lp.over_sample_size.native_for_fits, - dataset.grids.pixelization.over_sample_size.native_for_fits, + dataset.grids.lp.over_sample_size.native_for_fits.astype("float"), + dataset.grids.pixelization.over_sample_size.native_for_fits.astype( + "float" + ), ] hdu_list = hdu_list_for_output_from( diff --git a/autogalaxy/interferometer/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index 282151050..ab4987880 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -20,6 +20,7 @@ def __init__( dataset_model: Optional[aa.DatasetModel] = None, adapt_images: Optional[AdaptImages] = None, settings_inversion: aa.SettingsInversion = aa.SettingsInversion(), + xp=np, ): """ Fits an interferometer dataset using a list of galaxies. @@ -69,9 +70,7 @@ def __init__( self.galaxies = Galaxies(galaxies=galaxies) super().__init__( - dataset=dataset, - dataset_model=dataset_model, - use_mask_in_fit=False, + dataset=dataset, dataset_model=dataset_model, use_mask_in_fit=False, xp=xp ) AbstractFitInversion.__init__( self=self, @@ -89,7 +88,7 @@ def profile_visibilities(self) -> aa.Visibilities: a Fourier transform to the sum of light profile images. """ return self.galaxies.visibilities_from( - grid=self.grids.lp, transformer=self.dataset.transformer + grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp ) @property @@ -114,6 +113,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion: galaxies=self.galaxies, adapt_images=self.adapt_images, settings_inversion=self.settings_inversion, + xp=self._xp, ) @cached_property diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py index 9abffdb33..3961ff99b 100644 --- a/autogalaxy/interferometer/model/analysis.py +++ b/autogalaxy/interferometer/model/analysis.py @@ -1,4 +1,5 @@ import logging +import numpy as np from typing import Optional from autoconf.dictable import to_dict @@ -28,6 +29,7 @@ def __init__( adapt_image_maker: Optional[AdaptImageMaker] = None, cosmology: LensingCosmology = None, settings_inversion: aa.SettingsInversion = None, + preloads: aa.Preloads = None, title_prefix: str = None, ): """ @@ -65,6 +67,7 @@ def __init__( adapt_image_maker=adapt_image_maker, cosmology=cosmology, settings_inversion=settings_inversion, + preloads=preloads, title_prefix=title_prefix, ) @@ -95,7 +98,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection): return self - def log_likelihood_function(self, instance: af.ModelInstance) -> float: + def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float: """ Given an instance of the model, where the model parameters are set via a non-linear search, fit the model instance to the interferometer dataset. @@ -131,12 +134,9 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float: float The log likelihood indicating how well this model instance fitted the interferometer data. """ - return self.fit_from(instance=instance).figure_of_merit + return self.fit_from(instance=instance, xp=xp).figure_of_merit - def fit_from( - self, - instance: af.ModelInstance, - ) -> FitInterferometer: + def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer: """ Given a model instance create a `FitInterferometer` object. @@ -167,6 +167,7 @@ def fit_from( galaxies=galaxies, adapt_images=adapt_images, settings_inversion=self.settings_inversion, + xp=xp, ) def save_attributes(self, paths: af.DirectoryPaths): diff --git a/autogalaxy/operate/deflections.py b/autogalaxy/operate/deflections.py index 6c7ef3af6..ceb7bbb41 100644 --- a/autogalaxy/operate/deflections.py +++ b/autogalaxy/operate/deflections.py @@ -1,10 +1,9 @@ -import jax -from jax import jit -import jax.numpy as jnp -from functools import wraps, partial +from functools import wraps import logging +import numpy as np from typing import List, Tuple, Union +from autoconf import conf import autoarray as aa @@ -46,18 +45,49 @@ def wrapper(lensing_obj, grid, jacobian=None): return wrapper -def one_step(r, _, theta, fun, fun_dr): - r = jnp.abs(r - fun(r, theta) / fun_dr(r, theta)) - return r, None +def evaluation_grid(func): + @wraps(func) + def wrapper( + lensing_obj, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 + ): + if hasattr(grid, "is_evaluation_grid"): + if grid.is_evaluation_grid: + return func(lensing_obj, grid, pixel_scale) + pixel_scale_ratio = grid.pixel_scale / pixel_scale -@partial(jit, static_argnums=(4,)) -def step_r(r, theta, fun, fun_dr, N=20): - one_step_partial = jax.tree_util.Partial( - one_step, theta=theta, fun=fun, fun_dr=fun_dr - ) - new_r = jax.lax.scan(one_step_partial, r, xs=jnp.arange(N))[0] - return jnp.stack([new_r * jnp.sin(theta), new_r * jnp.cos(theta)]).T + zoom = aa.Zoom2D(mask=grid.mask) + + zoom_shape_native = zoom.shape_native + shape_native = ( + int(pixel_scale_ratio * zoom_shape_native[0]), + int(pixel_scale_ratio * zoom_shape_native[1]), + ) + + max_evaluation_grid_size = conf.instance["general"]["grid"][ + "max_evaluation_grid_size" + ] + + # This is a hack to prevent the evaluation gird going beyond 1000 x 1000 pixels, which slows the code + # down a lot. Need a better moe robust way to set this up for any general lens. + + if shape_native[0] > max_evaluation_grid_size: + pixel_scale = pixel_scale_ratio / ( + shape_native[0] / float(max_evaluation_grid_size) + ) + shape_native = (max_evaluation_grid_size, max_evaluation_grid_size) + + grid = aa.Grid2D.uniform( + shape_native=shape_native, + pixel_scales=(pixel_scale, pixel_scale), + origin=zoom.offset_scaled, + ) + + grid.is_evaluation_grid = True + + return func(lensing_obj, grid, pixel_scale) + + return wrapper class OperateDeflections: @@ -79,26 +109,10 @@ class OperateDeflections: def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): raise NotImplementedError - def deflections_yx_scalar(self, y, x, pixel_scales): - - # A version of the deflection function that takes in two scalars - # and outputs a 2D vector. Needed for JAX auto differentiation. - - mask = aa.Mask2D.all_false( - shape_native=(1, 1), - pixel_scales=pixel_scales, - ) - - g = aa.Grid2D( - values=jnp.stack((y.reshape(1), x.reshape(1)), axis=-1), mask=mask - ) - - return self.deflections_yx_2d_from(g).squeeze() - def __eq__(self, other): return self.__dict__ == other.__dict__ and self.__class__ is other.__class__ - def time_delay_geometry_term_from(self, grid) -> aa.Array2D: + def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D: """ Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions. @@ -122,7 +136,7 @@ def time_delay_geometry_term_from(self, grid) -> aa.Array2D: ------- The geometric time delay term at each grid position. """ - deflections = self.deflections_yx_2d_from(grid=grid) + deflections = self.deflections_yx_2d_from(grid=grid, xp=xp) src_y = grid[:, 0] - deflections[:, 0] src_x = grid[:, 1] - deflections[:, 1] @@ -133,7 +147,7 @@ def time_delay_geometry_term_from(self, grid) -> aa.Array2D: return aa.ArrayIrregular(values=delay) return aa.Array2D(values=delay, mask=grid.mask) - def fermat_potential_from(self, grid) -> aa.Array2D: + def fermat_potential_from(self, grid, xp=np) -> aa.Array2D: """ Returns the Fermat potential for a given grid of image-plane positions. @@ -158,8 +172,8 @@ def fermat_potential_from(self, grid) -> aa.Array2D: ------- The Fermat potential at each grid position. """ - time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid) - potential = self.potential_2d_from(grid=grid) + time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid, xp=xp) + potential = self.potential_2d_from(grid=grid, xp=xp) fermat_potential = time_delay_geometry_term - potential @@ -167,27 +181,6 @@ def fermat_potential_from(self, grid) -> aa.Array2D: return aa.ArrayIrregular(values=fermat_potential) return aa.Array2D(values=fermat_potential, mask=grid.mask) - def time_delays_from(self, grid) -> aa.Array2D: - """ - Returns the 2D time delay map of lensing object, which is computed as the deflection angles in the y and x - directions multiplied by the y and x coordinates of the grid. - - Parameters - ---------- - grid - The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay are computed on. - """ - deflections_yx = self.deflections_yx_2d_from(grid=grid) - - return aa.Array2D( - values=deflections_yx[:, 0] * grid[:, 0] - + deflections_yx[:, 1] * grid[:, 1], - mask=grid.mask, - ) - - def __hash__(self): - return hash(repr(self)) - @precompute_jacobian def tangential_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D: """ @@ -277,24 +270,21 @@ def hessian_from(self, grid, buffer: float = 0.01, deflections_func=None) -> Tup if deflections_func is None: deflections_func = self.deflections_yx_2d_from - if isinstance(grid, aa.Grid2DIrregular): - grid = grid.array + grid_shift_y_up = aa.Grid2DIrregular(values=np.zeros(grid.shape)) + grid_shift_y_up[:, 0] = grid[:, 0] + buffer + grid_shift_y_up[:, 1] = grid[:, 1] - grid_shift_y_up = aa.Grid2DIrregular( - values=jnp.stack([grid[:, 0] + buffer, grid[:, 1]], axis=1) - ) + grid_shift_y_down = aa.Grid2DIrregular(values=np.zeros(grid.shape)) + grid_shift_y_down[:, 0] = grid[:, 0] - buffer + grid_shift_y_down[:, 1] = grid[:, 1] - grid_shift_y_down = aa.Grid2DIrregular( - values=jnp.stack([grid[:, 0] - buffer, grid[:, 1]], axis=1) - ) + grid_shift_x_left = aa.Grid2DIrregular(values=np.zeros(grid.shape)) + grid_shift_x_left[:, 0] = grid[:, 0] + grid_shift_x_left[:, 1] = grid[:, 1] - buffer - grid_shift_x_left = aa.Grid2DIrregular( - values=jnp.stack([grid[:, 0], grid[:, 1] - buffer], axis=1) - ) - - grid_shift_x_right = aa.Grid2DIrregular( - values=jnp.stack([grid[:, 0], grid[:, 1] + buffer], axis=1) - ) + grid_shift_x_right = aa.Grid2DIrregular(values=np.zeros(grid.shape)) + grid_shift_x_right[:, 0] = grid[:, 0] + grid_shift_x_right[:, 1] = grid[:, 1] + buffer deflections_up = deflections_func(grid=grid_shift_y_up) deflections_down = deflections_func(grid=grid_shift_y_down) @@ -375,7 +365,10 @@ def shear_yx_2d_via_hessian_from( gamma_1 = 0.5 * (hessian_xx - hessian_yy) gamma_2 = hessian_xy - shear_yx_2d = jnp.stack([gamma_2.array, gamma_1.array], axis=1) + shear_yx_2d = np.zeros(shape=(grid.shape_slim, 2)) + + shear_yx_2d[:, 0] = gamma_2 + shear_yx_2d[:, 1] = gamma_1 return ShearYX2DIrregular(values=shear_yx_2d, grid=grid) @@ -419,9 +412,9 @@ def contour_list_from(self, grid, contour_array): return grid_contour.contour_list + @evaluation_grid def tangential_critical_curve_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[aa.Grid2DIrregular]: """ Returns all tangential critical curves of the lensing system, which are computed as follows: @@ -445,9 +438,9 @@ def tangential_critical_curve_list_from( return self.contour_list_from(grid=grid, contour_array=tangential_eigen_values) + @evaluation_grid def radial_critical_curve_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[aa.Grid2DIrregular]: """ Returns all radial critical curves of the lensing system, which are computed as follows: @@ -471,9 +464,9 @@ def radial_critical_curve_list_from( return self.contour_list_from(grid=grid, contour_array=radial_eigen_values) + @evaluation_grid def tangential_caustic_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[aa.Grid2DIrregular]: """ Returns all tangential caustics of the lensing system, which are computed as follows: @@ -497,7 +490,7 @@ def tangential_caustic_list_from( """ tangential_critical_curve_list = self.tangential_critical_curve_list_from( - grid=grid + grid=grid, pixel_scale=pixel_scale ) tangential_caustic_list = [] @@ -513,9 +506,9 @@ def tangential_caustic_list_from( return tangential_caustic_list + @evaluation_grid def radial_caustic_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[aa.Grid2DIrregular]: """ Returns all radial caustics of the lensing system, which are computed as follows: @@ -538,7 +531,9 @@ def radial_caustic_list_from( caustic to be computed more accurately using a higher resolution grid. """ - radial_critical_curve_list = self.radial_critical_curve_list_from(grid=grid) + radial_critical_curve_list = self.radial_critical_curve_list_from( + grid=grid, pixel_scale=pixel_scale + ) radial_caustic_list = [] @@ -553,7 +548,10 @@ def radial_caustic_list_from( return radial_caustic_list - def radial_critical_curve_area_list_from(self, grid) -> List[float]: + @evaluation_grid + def radial_critical_curve_area_list_from( + self, grid, pixel_scale: Union[Tuple[float, float], float] + ) -> List[float]: """ Returns the surface area within each radial critical curve as a list, the calculation of which is described in the function `radial_critical_curve_list_from()`. @@ -573,13 +571,15 @@ def radial_critical_curve_area_list_from(self, grid) -> List[float]: If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the caustic to be computed more accurately using a higher resolution grid. """ - radial_critical_curve_list = self.radial_critical_curve_list_from(grid=grid) + radial_critical_curve_list = self.radial_critical_curve_list_from( + grid=grid, pixel_scale=pixel_scale + ) return self.area_within_curve_list_from(curve_list=radial_critical_curve_list) + @evaluation_grid def tangential_critical_curve_area_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[float]: """ Returns the surface area within each tangential critical curve as a list, the calculation of which is @@ -600,7 +600,7 @@ def tangential_critical_curve_area_list_from( caustic to be computed more accurately using a higher resolution grid. """ tangential_critical_curve_list = self.tangential_critical_curve_list_from( - grid=grid + grid=grid, pixel_scale=pixel_scale ) return self.area_within_curve_list_from( @@ -614,14 +614,14 @@ def area_within_curve_list_from( for curve in curve_list: x, y = curve[:, 0], curve[:, 1] - area = jnp.abs(0.5 * jnp.sum(y[:-1] * jnp.diff(x) - x[:-1] * jnp.diff(y))) + area = np.abs(0.5 * np.sum(y[:-1] * np.diff(x) - x[:-1] * np.diff(y))) area_within_each_curve_list.append(area) return area_within_each_curve_list + @evaluation_grid def einstein_radius_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ): """ Returns a list of the Einstein radii corresponding to the area within each tangential critical curve. @@ -648,14 +648,16 @@ def einstein_radius_list_from( caustic to be computed more accurately using a higher resolution grid. """ try: - area_list = self.tangential_critical_curve_area_list_from(grid=grid) - return [jnp.sqrt(area / jnp.pi) for area in area_list] + area_list = self.tangential_critical_curve_area_list_from( + grid=grid, pixel_scale=pixel_scale + ) + return [np.sqrt(area / np.pi) for area in area_list] except TypeError: raise TypeError("The grid input was unable to estimate the Einstein Radius") + @evaluation_grid def einstein_radius_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ): """ Returns the Einstein radius corresponding to the area within the tangential critical curve. @@ -697,9 +699,9 @@ def einstein_radius_from( return sum(einstein_radii_list) + @evaluation_grid def einstein_mass_angular_list_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> List[float]: """ Returns a list of the angular Einstein massses corresponding to the area within each tangential critical curve. @@ -728,12 +730,14 @@ def einstein_mass_angular_list_from( If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the caustic to be computed more accurately using a higher resolution grid. """ - einstein_radius_list = self.einstein_radius_list_from(grid=grid) - return [jnp.pi * einstein_radius**2 for einstein_radius in einstein_radius_list] + einstein_radius_list = self.einstein_radius_list_from( + grid=grid, pixel_scale=pixel_scale + ) + return [np.pi * einstein_radius**2 for einstein_radius in einstein_radius_list] + @evaluation_grid def einstein_mass_angular_from( - self, - grid, + self, grid, pixel_scale: Union[Tuple[float, float], float] = 0.05 ) -> float: """ Returns the Einstein radius corresponding to the area within the tangential critical curve. @@ -762,7 +766,9 @@ def einstein_mass_angular_from( If input, the `evaluation_grid` decorator creates the 2D grid at this resolution, therefore enabling the caustic to be computed more accurately using a higher resolution grid. """ - einstein_mass_angular_list = self.einstein_mass_angular_list_from(grid=grid) + einstein_mass_angular_list = self.einstein_mass_angular_list_from( + grid=grid, pixel_scale=pixel_scale + ) if len(einstein_mass_angular_list) > 1: logger.info( @@ -774,167 +780,6 @@ def einstein_mass_angular_from( return einstein_mass_angular_list[0] - def jacobian_stack(self, y, x, pixel_scales): - return jnp.stack( - jax.jacfwd(self.deflections_yx_scalar, argnums=(0, 1))(y, x, pixel_scales) - ) - - def jacobian_stack_vector(self, y, x, pixel_scales): - return jnp.vectorize( - jax.tree_util.Partial(self.jacobian_stack, pixel_scales=pixel_scales), - signature="(),()->(i,i)", - )(y, x) - - def convergence_mag_shear_yx(self, y, x): - J = self.jacobian_stack_vector(y, x, 0.05) - K = 0.5 * (J[..., 0, 0] + J[..., 1, 1]) - mag_shear = 0.5 * jnp.sqrt( - (J[..., 0, 1] + J[..., 1, 0]) ** 2 + (J[..., 0, 0] - J[..., 1, 1]) ** 2 - ) - return K, mag_shear - - @partial(jit, static_argnums=(0,)) - def tangential_eigen_value_yx(self, y, x): - K, mag_shear = self.convergence_mag_shear_yx(y, x) - return 1 - K - mag_shear - - @partial(jit, static_argnums=(0, 3)) - def tangential_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)): - y = r * jnp.sin(theta) + centre[0] - x = r * jnp.cos(theta) + centre[1] - return self.tangential_eigen_value_yx(y, x) - - @partial(jit, static_argnums=(0, 3)) - def grad_r_tangential_eigen_value(self, r, theta, centre=(0.0, 0.0)): - # ignore `self` with the `argnums` below - tangential_eigen_part = partial(self.tangential_eigen_value_rt, centre=centre) - return jnp.vectorize( - jax.jacfwd(tangential_eigen_part, argnums=(0,)), signature="(),()->()" - )(r, theta)[0] - - @partial(jit, static_argnums=(0,)) - def radial_eigen_value_yx(self, y, x): - K, mag_shear = self.convergence_mag_shear_yx(y, x) - return 1 - K + mag_shear - - @partial(jit, static_argnums=(0, 3)) - def radial_eigen_value_rt(self, r, theta, centre=(0.0, 0.0)): - y = r * jnp.sin(theta) + centre[0] - x = r * jnp.cos(theta) + centre[1] - return self.radial_eigen_value_yx(y, x) - - @partial(jit, static_argnums=(0, 3)) - def grad_r_radial_eigen_value(self, r, theta, centre=(0.0, 0.0)): - # ignore `self` with the `argnums` below - radial_eigen_part = partial(self.radial_eigen_value_rt, centre=centre) - return jnp.vectorize( - jax.jacfwd(radial_eigen_part, argnums=(0,)), signature="(),()->()" - )(r, theta)[0] - - def tangential_critical_curve_jax( - self, - init_r=0.1, - init_centre=(0.0, 0.0), - n_points=300, - n_steps=20, - threshold=1e-5, - ): - """ - Returns all tangential critical curves of the lensing system, which are computed as follows: - - 1) Create a set of `n_points` initial points in a circle of radius `init_r` and centred on `init_centre` - 2) Apply `n_steps` of Newton's method to these points in the "radial" direction only (i.e. keeping angle fixed). - Jax's auto differentiation is used to find the radial derivatives of the tangential eigen value function for - this step. - 3) Filter the results and only keep point that have their tangential eigen value `threshold` of 0 - - No underlying grid is needed for the method, but the quality of the results are dependent on the initial - circle of points. - - Parameters - ---------- - init_r : float - Radius of the circle of initial guess points - init_centre : tuple - centre of the circle of initial guess points as `(y, x)` - n_points : Int - Number of initial guess points to create (evenly spaced in angle around `init_centre`) - n_steps : Int - Number of iterations of Newton's method to apply - threshold : float - Only keep points whose tangential eigen value is within this value of zero (inclusive) - """ - r = jnp.ones(n_points) * init_r - theta = jnp.linspace(0, 2 * jnp.pi, n_points + 1)[:-1] - new_yx = step_r( - r, - theta, - jax.tree_util.Partial(self.tangential_eigen_value_rt, centre=init_centre), - jax.tree_util.Partial( - self.grad_r_tangential_eigen_value, centre=init_centre - ), - n_steps, - ) - new_yx = new_yx + jnp.array(init_centre) - # filter out nan values - fdx = jnp.isfinite(new_yx).all(axis=1) - new_yx = new_yx[fdx] - # filter out failed points - value = jnp.abs(self.tangential_eigen_value_yx(new_yx[:, 0], new_yx[:, 1])) - gdx = value <= threshold - return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx]) - - def radial_critical_curve_jax( - self, - init_r=0.01, - init_centre=(0.0, 0.0), - n_points=300, - n_steps=20, - threshold=1e-5, - ): - """ - Returns all radial critical curves of the lensing system, which are computed as follows: - - 1) Create a set of `n_points` initial points in a circle of radius `init_r` and centred on `init_centre` - 2) Apply `n_steps` of Newton's method to these points in the "radial" direction only (i.e. keeping angle fixed). - Jax's auto differentiation is used to find the radial derivatives of the radial eigen value function for - this step. - 3) Filter the results and only keep point that have their radial eigen value `threshold` of 0 - - No underlying grid is needed for the method, but the quality of the results are dependent on the initial - circle of points. - - Parameters - ---------- - init_r : float - Radius of the circle of initial guess points - init_centre : tuple - centre of the circle of initial guess points as `(y, x)` - n_points : Int - Number of initial guess points to create (evenly spaced in angle around `init_centre`) - n_steps : Int - Number of iterations of Newton's method to apply - threshold : float - Only keep points whose radial eigen value is within this value of zero (inclusive) - """ - r = jnp.ones(n_points) * init_r - theta = jnp.linspace(0, 2 * jnp.pi, n_points + 1)[:-1] - new_yx = step_r( - r, - theta, - jax.tree_util.Partial(self.radial_eigen_value_rt, centre=init_centre), - jax.tree_util.Partial(self.grad_r_radial_eigen_value, centre=init_centre), - n_steps, - ) - new_yx = new_yx + jnp.array(init_centre) - # filter out nan values - fdx = jnp.isfinite(new_yx).all(axis=1) - new_yx = new_yx[fdx] - # filter out failed points - value = jnp.abs(self.radial_eigen_value_yx(new_yx[:, 0], new_yx[:, 1])) - gdx = value <= threshold - return aa.structures.grids.irregular_2d.Grid2DIrregular(values=new_yx[gdx]) - def jacobian_from(self, grid): """ Returns the Jacobian of the lensing object, which is computed by taking the gradient of the 2D deflection @@ -952,24 +797,36 @@ def jacobian_from(self, grid): grid The 2D grid of (y,x) arc-second coordinates the deflection angles and Jacobian are computed on. """ - A = self.jacobian_stack_vector( - grid.array[:, 0], grid.array[:, 1], grid.pixel_scales + + 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, ) - a = jnp.eye(2).reshape(1, 2, 2) - A - return [ - [ - aa.Array2D(values=a[..., 1, 1], mask=grid.mask), - aa.Array2D(values=a[..., 1, 0], mask=grid.mask), - ], - [ - aa.Array2D(values=a[..., 0, 1], mask=grid.mask), - aa.Array2D(values=a[..., 0, 0], mask=grid.mask), - ], - ] - # transpose the result - # use `moveaxis` as grid might not be nx2 - # return jnp.moveaxis(jnp.moveaxis(a, -1, 0), -1, 0) + 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: @@ -1018,9 +875,10 @@ def shear_yx_2d_via_jacobian_from( jacobian A precomputed lensing jacobian, which is passed throughout the `CalcLens` functions for efficiency. """ - shear_y = -0.5 * (jacobian[0][1] + jacobian[1][0]).array - shear_x = 0.5 * (jacobian[1][1] - jacobian[0][0]).array - shear_yx_2d = jnp.stack([shear_y, shear_x]).T + + 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) diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py index 0f3ff45a3..56601c372 100644 --- a/autogalaxy/operate/image.py +++ b/autogalaxy/operate/image.py @@ -1,6 +1,5 @@ from __future__ import annotations -import jax -import jax.numpy as jnp +import numpy as np from typing import TYPE_CHECKING, Dict, List, Optional from autoarray import Array2D @@ -23,7 +22,7 @@ class OperateImage: """ def image_2d_from( - self, grid: aa.Grid2D, operated_only: Optional[bool] = None + self, grid: aa.Grid2D, xp=np, operated_only: Optional[bool] = None ) -> aa.Array2D: raise NotImplementedError @@ -35,11 +34,11 @@ def _blurred_image_2d_from( image_2d: aa.Array2D, blurring_image_2d: aa.Array2D, psf: aa.Kernel2D, + xp=np, ) -> aa.Array2D: values = psf.convolved_image_from( - image=image_2d, - blurring_image=blurring_image_2d, + image=image_2d, blurring_image=blurring_image_2d, xp=xp ) return Array2D(values=values, mask=image_2d.mask) @@ -48,6 +47,7 @@ def blurred_image_2d_from( grid: aa.Grid2D, blurring_grid: aa.Grid2D, psf: aa.Kernel2D = None, + xp=np, ) -> aa.Array2D: """ Evaluate the light object's 2D image from a input 2D grid of coordinates and convolve it with a PSF. @@ -72,19 +72,22 @@ def blurred_image_2d_from( LightProfileOperated, ) - image_2d_not_operated = self.image_2d_from(grid=grid, operated_only=False) + image_2d_not_operated = self.image_2d_from( + grid=grid, xp=xp, operated_only=False + ) blurring_image_2d_not_operated = self.image_2d_from( - grid=blurring_grid, operated_only=False + grid=blurring_grid, xp=xp, operated_only=False ) blurred_image_2d = self._blurred_image_2d_from( image_2d=image_2d_not_operated, blurring_image_2d=blurring_image_2d_not_operated, psf=psf, + xp=xp, ) if self.has(cls=LightProfileOperated): - image_2d_operated = self.image_2d_from(grid=grid, operated_only=True) + image_2d_operated = self.image_2d_from(grid=grid, xp=xp, operated_only=True) return blurred_image_2d + image_2d_operated return blurred_image_2d @@ -161,7 +164,7 @@ def unmasked_blurred_image_2d_from(self, grid, psf): return padded_image_2d + padded_image_2d_operated def visibilities_from( - self, grid: aa.Grid2D, transformer: aa.type.Transformer + self, grid: aa.Grid2D, transformer: aa.type.Transformer, xp=np ) -> aa.Visibilities: """ Evaluate the light object's 2D image from a input 2D grid of coordinates and transform this to an array of @@ -304,7 +307,10 @@ def unmasked_blurred_image_2d_list_from( return unmasked_blurred_image_list def visibilities_list_from( - self, grid: aa.Grid2D, transformer: aa.type.Transformer + self, + grid: aa.Grid2D, + transformer: aa.type.Transformer, + xp=np, ) -> List[aa.Array2D]: """ Evaluate the light object's list of 2D image from a input 2D grid of coordinates and transform each image to @@ -334,7 +340,7 @@ def visibilities_list_from( visibilities_list = [] for image_2d in image_2d_list: - if not jnp.any(image_2d.array): + if not xp.any(image_2d.array): visibilities = aa.Visibilities.zeros( shape_slim=(transformer.uv_wavelengths.shape[0],) ) @@ -420,7 +426,7 @@ def galaxy_blurred_image_2d_dict_from( return galaxy_blurred_image_2d_dict def galaxy_visibilities_dict_from( - self, grid, transformer + self, grid, transformer, xp=np ) -> Dict[Galaxy, aa.Visibilities]: """ Evaluate the light object's dictionary mapping galaixes to their corresponding 2D images and transform each @@ -453,7 +459,7 @@ def galaxy_visibilities_dict_from( for galaxy_key in galaxy_image_2d_dict.keys(): image_2d = galaxy_image_2d_dict[galaxy_key] - if not jnp.any(image_2d.array): + if not xp.any(image_2d.array): visibilities = aa.Visibilities.zeros( shape_slim=(transformer.uv_wavelengths.shape[0],) ) diff --git a/autogalaxy/plot/__init__.py b/autogalaxy/plot/__init__.py index 66317df0b..c33fc78c1 100644 --- a/autogalaxy/plot/__init__.py +++ b/autogalaxy/plot/__init__.py @@ -76,12 +76,9 @@ from autogalaxy.plot.visuals.two_d import Visuals2D from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePlotter -from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePDFPlotter from autogalaxy.profiles.plot.basis_plotters import BasisPlotter from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePlotter -from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePDFPlotter from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPlotter -from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPDFPlotter from autogalaxy.galaxy.plot.galaxies_plotters import GalaxiesPlotter from autogalaxy.quantity.plot.fit_quantity_plotters import FitQuantityPlotter from autogalaxy.imaging.plot.fit_imaging_plotters import FitImagingPlotter diff --git a/autogalaxy/profiles/basis.py b/autogalaxy/profiles/basis.py index 3fc08b3ee..f63ba47b3 100644 --- a/autogalaxy/profiles/basis.py +++ b/autogalaxy/profiles/basis.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Dict, List, Optional, Union @@ -76,7 +75,11 @@ def mass_profile_list(self) -> List[MassProfile]: return aa.util.misc.cls_list_from(values=self.profile_list, cls=MassProfile) def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> aa.Array2D: """ Returns the summed image of all light profiles in the basis from a 2D grid of Cartesian (y,x) coordinates. @@ -99,10 +102,12 @@ def image_2d_from( ------- The image of the light profiles in the basis summed together. """ - return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only)) + return sum( + self.image_2d_list_from(grid=grid, xp=xp, operated_only=operated_only) + ) def image_2d_list_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ) -> List[aa.Array2D]: """ Returns each image of each light profiles in the basis as a list, from a 2D grid of Cartesian (y,x) coordinates. @@ -127,14 +132,18 @@ def image_2d_list_from( """ return [ ( - light_profile.image_2d_from(grid=grid, operated_only=operated_only) + light_profile.image_2d_from( + grid=grid, xp=xp, operated_only=operated_only + ) if not isinstance(light_profile, lp_linear.LightProfileLinear) - else jnp.zeros((grid.shape[0],)) + else xp.zeros((grid.shape[0],)) ) for light_profile in self.light_profile_list ] - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D: + def convergence_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> aa.Array2D: """ Returns the summed convergence of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates. @@ -153,11 +162,16 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D: """ if len(self.mass_profile_list) > 0: return sum( - [mass.convergence_2d_from(grid=grid) for mass in self.profile_list] + [ + mass.convergence_2d_from(grid=grid, xp=xp) + for mass in self.profile_list + ] ) - return jnp.zeros((grid.shape[0],)) + return xp.zeros((grid.shape[0],)) - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D: + def potential_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> aa.Array2D: """ Returns the summed potential of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates. @@ -178,9 +192,11 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D: return sum( [mass.potential_2d_from(grid=grid) for mass in self.profile_list] ) - return jnp.zeros((grid.shape[0],)) + return xp.zeros((grid.shape[0],)) - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array2D: + def deflections_yx_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> aa.Array2D: """ Returns the summed deflections of all mass profiles in the basis from a 2D grid of Cartesian (y,x) coordinates. @@ -201,7 +217,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> aa.Array return sum( [mass.deflections_yx_2d_from(grid=grid) for mass in self.profile_list] ) - return jnp.zeros((grid.shape[0], 2)) + return xp.zeros((grid.shape[0], 2)) def lp_instance_from(self, linear_light_profile_intensity_dict: Dict): light_profile_list = [] diff --git a/autogalaxy/profiles/geometry_profiles.py b/autogalaxy/profiles/geometry_profiles.py index f0a67f5ad..21516ab64 100644 --- a/autogalaxy/profiles/geometry_profiles.py +++ b/autogalaxy/profiles/geometry_profiles.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple, Type @@ -39,36 +38,12 @@ def has(self, cls: Type) -> bool: """ return aa.util.misc.has(values=self.__dict__.values(), cls=cls) - def transformed_to_reference_frame_grid_from(self, grid, **kwargs): + def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs): raise NotImplemented() - def transformed_from_reference_frame_grid_from(self, grid, **kwargs): + def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs): raise NotImplemented() - def _radial_projected_shape_slim_from( - self, grid: aa.type.Grid1D2DLike, **kwargs - ) -> int: - """ - To make 1D plots (e.g. `image_1d_from()`) from an input 2D grid, one uses that 2D grid to radially project - the coordinates across the profile's major-axis. - - This function computes the distance from the profile centre to the edge of this 2D grid. - - If a 1D grid is input it returns the shape of this grid, as the grid itself defines the radial coordinates. - - Parameters - ---------- - grid - A 1D or 2D grid from which a 1D plot of the profile is to be created. - """ - - if isinstance(grid, aa.Grid1D): - return grid.sub_shape_slim - elif isinstance(grid, aa.Grid2DIrregular): - return grid.slim.shape[0] - - return grid.grid_2d_radial_projected_shape_slim_from(centre=self.centre) - class SphProfile(GeometryProfile): """ @@ -84,7 +59,7 @@ def __init__(self, centre: Tuple[float, float] = (0.0, 0.0)): super().__init__(centre=centre) @aa.grid_dec.to_array - def radial_grid_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def radial_grid_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs) -> np.ndarray: """ Convert a grid of (y, x) coordinates, to their radial distances from the profile centre (e.g. :math: r = sqrt(x**2 + y**2)). @@ -94,12 +69,10 @@ def radial_grid_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: grid The grid of (y, x) coordinates which are converted to radial distances. """ - return jnp.sqrt( - jnp.add(jnp.square(grid.array[:, 0]), jnp.square(grid.array[:, 1])) - ) + return xp.sqrt(xp.add(xp.square(grid.array[:, 0]), xp.square(grid.array[:, 1]))) def angle_to_profile_grid_from( - self, grid_angles: np.ndarray, **kwargs + self, grid_angles: np.ndarray, xp=np, **kwargs ) -> Tuple[np.ndarray, np.ndarray]: """ Convert a grid of angles, defined in degrees counter-clockwise from the positive x-axis, to a grid of @@ -110,11 +83,15 @@ def angle_to_profile_grid_from( grid_angles The angle theta counter-clockwise from the positive x-axis to each coordinate in radians. """ - return jnp.cos(grid_angles), jnp.sin(grid_angles) + return xp.cos(grid_angles), xp.sin(grid_angles) @aa.grid_dec.to_grid def _cartesian_grid_via_radial_from( - self, grid: aa.type.Grid2DLike, radius: jnp.ndarray, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + radius: Optional[np.ndarray] = None, + **kwargs, ) -> aa.type.Grid2DLike: """ Convert a grid of (y,x) coordinates with their specified radial distances (e.g. :math: r = x**2 + y**2) to @@ -127,18 +104,15 @@ def _cartesian_grid_via_radial_from( radius The circular radius of each coordinate from the profile center. """ - grid_angles = jnp.arctan2(grid.array[:, 0], grid.array[:, 1]) - cos_theta, sin_theta = self.angle_to_profile_grid_from(grid_angles=grid_angles) - - if isinstance(radius, jnp.ndarray): - return jnp.multiply(radius[:, None], jnp.vstack((sin_theta, cos_theta)).T) - elif isinstance(radius, np.ndarray): - return np.multiply(radius[:, None], np.vstack((sin_theta, cos_theta)).T) + grid_angles = xp.arctan2(grid.array[:, 0], grid.array[:, 1]) + cos_theta, sin_theta = self.angle_to_profile_grid_from( + grid_angles=grid_angles, xp=xp + ) - return jnp.multiply(radius.array[:, None], jnp.vstack((sin_theta, cos_theta)).T) + return xp.multiply(radius[:, None], xp.vstack((sin_theta, cos_theta)).T) @aa.grid_dec.to_grid - def transformed_to_reference_frame_grid_from(self, grid, **kwargs): + def transformed_to_reference_frame_grid_from(self, grid, xp=np, **kwargs): """ Transform a grid of (y,x) coordinates to the reference frame of the profile. @@ -149,10 +123,10 @@ def transformed_to_reference_frame_grid_from(self, grid, **kwargs): grid The (y, x) coordinates in the original reference frame of the grid. """ - return jnp.subtract(grid.array, jnp.array(self.centre)) + return xp.subtract(grid.array, xp.array(self.centre)) @aa.grid_dec.to_grid - def transformed_from_reference_frame_grid_from(self, grid, **kwargs): + def transformed_from_reference_frame_grid_from(self, grid, xp=np, **kwargs): """ Transform a grid of (y,x) coordinates from the reference frame of the profile to the original observer reference frame. @@ -164,7 +138,7 @@ def transformed_from_reference_frame_grid_from(self, grid, **kwargs): grid The (y, x) coordinates in the reference frame of the profile. """ - return jnp.add(grid.array, jnp.array(self.centre)) + return xp.add(grid.array, xp.array(self.centre)) class EllProfile(SphProfile): @@ -218,28 +192,25 @@ def __init__( self.ell_comps = ell_comps - @property - def axis_ratio(self) -> float: + def axis_ratio(self, xp=np) -> float: """ The ratio of the minor-axis to major-axis (b/a) of the ellipse defined by profile (0.0 > q > 1.0). """ - return convert.axis_ratio_from(ell_comps=self.ell_comps) + return convert.axis_ratio_from(ell_comps=self.ell_comps, xp=xp) - @property - def angle(self) -> float: + def angle(self, xp=np) -> float: """ The position angle in degrees of the major-axis of the ellipse defined by profile, defined counter clockwise from the positive x-axis (0.0 > angle > 180.0). """ - return convert.angle_from(ell_comps=self.ell_comps) + return convert.angle_from(ell_comps=self.ell_comps, xp=xp) - @property - def angle_radians(self) -> float: + def angle_radians(self, xp=np) -> float: """ The position angle in radians of the major-axis of the ellipse defined by profile, defined counter clockwise from the positive x-axis (0.0 > angle > 2pi). """ - return jnp.radians(self.angle) + return xp.radians(self.angle(xp)) @property def _cos_angle(self) -> float: @@ -249,15 +220,15 @@ def _cos_angle(self) -> float: def _sin_angle(self) -> float: return self._cos_and_sin_to_x_axis()[1] - def _cos_and_sin_to_x_axis(self, **kwargs): + def _cos_and_sin_to_x_axis(self, xp=np, **kwargs): """ Determine the sin and cosine of the angle between the profile's ellipse and the positive x-axis, counter-clockwise. """ - angle_radians = jnp.radians(self.angle) - return jnp.cos(angle_radians), jnp.sin(angle_radians) + angle_radians = xp.radians(self.angle(xp)) + return xp.cos(angle_radians), xp.sin(angle_radians) - def angle_to_profile_grid_from(self, grid_angles, **kwargs): + def angle_to_profile_grid_from(self, grid_angles, xp=np, **kwargs): """ The angle between each angle theta on the grid and the profile, in radians. @@ -266,14 +237,12 @@ def angle_to_profile_grid_from(self, grid_angles, **kwargs): grid_angles The angle theta counter-clockwise from the positive x-axis to each coordinate in radians. """ - theta_coordinate_to_profile = jnp.add(grid_angles, -self.angle_radians) - return jnp.cos(theta_coordinate_to_profile), jnp.sin( - theta_coordinate_to_profile - ) + theta_coordinate_to_profile = xp.add(grid_angles, -self.angle_radians(xp=xp)) + return xp.cos(theta_coordinate_to_profile), xp.sin(theta_coordinate_to_profile) @aa.grid_dec.to_grid def rotated_grid_from_reference_frame_from( - self, grid, angle: Optional[float] = None, **kwargs + self, grid, xp=np, angle: Optional[float] = None, **kwargs ): """ Rotate a grid of (y,x) coordinates which have been transformed to the elliptical reference frame of a profile @@ -296,16 +265,16 @@ def rotated_grid_from_reference_frame_from( """ if angle is None: - angle = self.angle + angle = self.angle(xp) return aa.util.geometry.transform_grid_2d_from_reference_frame( - grid_2d=grid, centre=(0.0, 0.0), angle=angle + grid_2d=grid, centre=(0.0, 0.0), angle=angle, xp=xp ) @aa.grid_dec.to_array def elliptical_radii_grid_from( - self, grid: aa.type.Grid2DLike, **kwargs - ) -> jnp.ndarray: + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Convert a grid of (y,x) coordinates to their elliptical radii values: :math: (x^2 + (y^2/q))^0.5 @@ -314,17 +283,17 @@ def elliptical_radii_grid_from( grid The (y, x) coordinates in the reference frame of the elliptical profile. """ - return jnp.sqrt( - jnp.add( - jnp.square(grid.array[:, 1]), - jnp.square(jnp.divide(grid.array[:, 0], self.axis_ratio)), + return xp.sqrt( + xp.add( + xp.square(grid.array[:, 1]), + xp.square(xp.divide(grid.array[:, 0], self.axis_ratio(xp))), ) ) @aa.grid_dec.to_array def eccentric_radii_grid_from( - self, grid: aa.type.Grid2DLike, **kwargs - ) -> jnp.ndarray: + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Convert a grid of (y,x) coordinates to an eccentric radius: :math: axis_ratio^0.5 (x^2 + (y^2/q))^0.5 @@ -339,14 +308,14 @@ def eccentric_radii_grid_from( The (y, x) coordinates in the reference frame of the elliptical profile. """ - grid_radii = self.elliptical_radii_grid_from(grid=grid, **kwargs) + grid_radii = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs) - return jnp.multiply(jnp.sqrt(self.axis_ratio), grid_radii.array) + return xp.multiply(xp.sqrt(self.axis_ratio(xp)), grid_radii.array) @aa.grid_dec.to_grid def transformed_to_reference_frame_grid_from( - self, grid: aa.type.Grid2DLike, **kwargs - ) -> jnp.ndarray: + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Transform a grid of (y,x) coordinates to the reference frame of the profile. @@ -358,14 +327,14 @@ def transformed_to_reference_frame_grid_from( The (y, x) coordinates in the original reference frame of the grid. """ if self.__class__.__name__.endswith("Sph"): - return super().transformed_to_reference_frame_grid_from(grid=grid) + return super().transformed_to_reference_frame_grid_from(grid=grid, xp=xp) return aa.util.geometry.transform_grid_2d_to_reference_frame( - grid_2d=grid.array, centre=self.centre, angle=self.angle + grid_2d=grid.array, centre=self.centre, angle=self.angle(xp), xp=xp ) @aa.grid_dec.to_grid def transformed_from_reference_frame_grid_from( - self, grid: aa.type.Grid2DLike, **kwargs + self, grid: aa.type.Grid2DLike, xp=np, **kwargs ) -> aa.type.Grid2DLike: """ Transform a grid of (y,x) coordinates from the reference frame of the profile to the original observer @@ -379,10 +348,10 @@ def transformed_from_reference_frame_grid_from( The (y, x) coordinates in the reference frame of the profile. """ if self.__class__.__name__.startswith("Sph"): - return super().transformed_from_reference_frame_grid_from(grid=grid) + return super().transformed_from_reference_frame_grid_from(grid=grid, xp=xp) return aa.util.geometry.transform_grid_2d_from_reference_frame( - grid_2d=grid.array, centre=self.centre, angle=self.angle + grid_2d=grid.array, centre=self.centre, angle=self.angle(xp), xp=xp ) def _eta_u(self, u, coordinates): @@ -391,7 +360,7 @@ def _eta_u(self, u, coordinates): u * ( (coordinates[1] ** 2) - + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio**2) * u)) + + (coordinates[0] ** 2 / (1 - (1 - self.axis_ratio(xp) ** 2) * u)) ) ) ) diff --git a/autogalaxy/profiles/light/abstract.py b/autogalaxy/profiles/light/abstract.py index da88cbcc4..3ae084825 100644 --- a/autogalaxy/profiles/light/abstract.py +++ b/autogalaxy/profiles/light/abstract.py @@ -42,7 +42,11 @@ def coefficient_tag(self) -> str: return "" def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> aa.Array2D: """ Returns the light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates, which may have been @@ -63,7 +67,7 @@ def image_2d_from( """ raise NotImplementedError() - def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ Returns the light profile's 2D image from a 1D grid of coordinates which are the radial distance of each coordinate from the light profile `centre`. @@ -75,33 +79,6 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: """ raise NotImplementedError() - @aa.grid_dec.project_grid - def image_1d_from( - self, grid: aa.type.Grid1D2DLike, **kwargs - ) -> aa.type.Grid1D2DLike: - """ - Returns the light profile's 1D image from a grid of Cartesian coordinates, which may have been - transformed using the light profile's geometry. - - If a 1D grid is input the image is evaluated every coordinate on the grid. If a 2D grid is input, this is - converted to a 1D grid by aligning with the major-axis of the light profile's elliptical geometry. - - Internally, this function uses a 2D grid to compute the image, which is mapped to a 1D data structure on return - via the `project_grid` decorator. This avoids code repetition by ensuring that light profiles only use - their `image_2d_from()` function to evaluate their image. - - Parameters - ---------- - grid - A 1D or 2D grid of coordinates which are used to evaluate the light profile in 1D. - - Returns - ------- - image - The 1D image of the light profile evaluated at every (x,) coordinate on the 1D transformed grid. - """ - return self.image_2d_from(grid=grid, **kwargs) - def luminosity_within_circle_from(self, radius: float) -> float: """ Integrate the light profile to compute the total luminosity within a circle of specified radius. This is diff --git a/autogalaxy/profiles/light/decorators.py b/autogalaxy/profiles/light/decorators.py index 9d4c134b6..909fc507f 100644 --- a/autogalaxy/profiles/light/decorators.py +++ b/autogalaxy/profiles/light/decorators.py @@ -27,6 +27,7 @@ def check_operated_only(func): def wrapper( obj, grid: aa.type.Grid1D2DLike, + xp=np, operated_only: Optional[bool] = None, *args, **kwargs, @@ -60,13 +61,13 @@ def wrapper( ) if operated_only is None: - return func(obj, grid, operated_only, *args, **kwargs) + return func(obj, grid, xp, operated_only, *args, **kwargs) elif operated_only: if isinstance(obj, LightProfileOperated): - return func(obj, grid, operated_only, *args, **kwargs) - return np.zeros((grid.shape[0],)) + return func(obj, grid, xp, operated_only, *args, **kwargs) + return xp.zeros((grid.shape[0],)) if not isinstance(obj, LightProfileOperated): - return func(obj, grid, operated_only, *args, **kwargs) - return np.zeros((grid.shape[0],)) + return func(obj, grid, xp, operated_only, *args, **kwargs) + return xp.zeros((grid.shape[0],)) return wrapper diff --git a/autogalaxy/profiles/light/linear/abstract.py b/autogalaxy/profiles/light/linear/abstract.py index 9663c7810..43a0946de 100644 --- a/autogalaxy/profiles/light/linear/abstract.py +++ b/autogalaxy/profiles/light/linear/abstract.py @@ -1,5 +1,4 @@ import inspect -import jax.numpy as jnp import numpy as np from typing import Dict, List, Optional @@ -146,6 +145,7 @@ def __init__( psf: Optional[aa.Kernel2D], light_profile_list: List[LightProfileLinear], regularization=Optional[aa.reg.Regularization], + xp=np, ): """ A list of linear light profiles which fits a dataset via linear algebra using the images of each linear light @@ -202,10 +202,7 @@ def __init__( """ ) - super().__init__( - grid=grid, - regularization=regularization, - ) + super().__init__(grid=grid, regularization=regularization, xp=xp) self.blurring_grid = blurring_grid self.psf = psf @@ -267,11 +264,11 @@ def mapping_matrix(self) -> np.ndarray: for pixel, light_profile in enumerate(self.light_profile_list): - image_2d = light_profile.image_2d_from(grid=self.grid).slim + image_2d = light_profile.image_2d_from(grid=self.grid, xp=self._xp).slim image_2d_list.append(image_2d.array) - return jnp.stack(image_2d_list, axis=1) + return self._xp.stack(image_2d_list, axis=1) @cached_property def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]: @@ -300,26 +297,39 @@ def operated_mapping_matrix_overrideg(self) -> Optional[np.ndarray]: n_src = len(self.light_profile_list) # allocate slim-form arrays for mapping matrices - mapping_matrix = jnp.zeros((self.grid.shape_slim, n_src)) - blurring_mapping_matrix = jnp.zeros((self.blurring_grid.shape_slim, n_src)) + mapping_matrix = self._xp.zeros((self.grid.shape_slim, n_src)) + blurring_mapping_matrix = self._xp.zeros((self.blurring_grid.shape_slim, n_src)) # build each column for pixel, light_profile in enumerate(self.light_profile_list): - # main grid mapping for this light profile - mapping_matrix = mapping_matrix.at[:, pixel].set( - light_profile.image_2d_from(grid=self.grid).array - ) + if self._xp.__name__.startswith("jax"): + # main grid mapping for this light profile + mapping_matrix = mapping_matrix.at[:, pixel].set( + light_profile.image_2d_from(grid=self.grid, xp=self._xp).array + ) - # blurring grid mapping for this light profile - blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set( - light_profile.image_2d_from(grid=self.blurring_grid).array - ) + # blurring grid mapping for this light profile + blurring_mapping_matrix = blurring_mapping_matrix.at[:, pixel].set( + light_profile.image_2d_from( + grid=self.blurring_grid, xp=self._xp + ).array + ) + + else: + + mapping_matrix[:, pixel] = light_profile.image_2d_from( + grid=self.grid, xp=self._xp + ).array + blurring_mapping_matrix[:, pixel] = light_profile.image_2d_from( + grid=self.blurring_grid, xp=self._xp + ).array return self.psf.convolved_mapping_matrix_from( mapping_matrix=mapping_matrix, mask=self.grid.mask, blurring_mapping_matrix=blurring_mapping_matrix, blurring_mask=self.blurring_grid.mask, + xp=self._xp, ) @cached_property @@ -348,14 +358,16 @@ def operated_mapping_matrix_override(self) -> Optional[np.ndarray]: blurred_image_2d_list = [] for pixel, light_profile in enumerate(self.light_profile_list): - image_2d = light_profile.image_2d_from(grid=self.grid) + image_2d = light_profile.image_2d_from(grid=self.grid, xp=self._xp) - blurring_image_2d = light_profile.image_2d_from(grid=self.blurring_grid) + blurring_image_2d = light_profile.image_2d_from( + grid=self.blurring_grid, xp=self._xp + ) blurred_image_2d = self.psf.convolved_image_from( - image=image_2d, blurring_image=blurring_image_2d + image=image_2d, blurring_image=blurring_image_2d, xp=self._xp ) blurred_image_2d_list.append(blurred_image_2d.array) - return jnp.stack(blurred_image_2d_list, axis=1) + return self._xp.stack(blurred_image_2d_list, axis=1) diff --git a/autogalaxy/profiles/light/mock/mock_light_profile.py b/autogalaxy/profiles/light/mock/mock_light_profile.py index 0b7ef9aaa..fbadb455f 100644 --- a/autogalaxy/profiles/light/mock/mock_light_profile.py +++ b/autogalaxy/profiles/light/mock/mock_light_profile.py @@ -29,7 +29,9 @@ def __init__( @aa.grid_dec.to_array @check_operated_only - def image_2d_from(self, grid, operated_only: Optional[bool] = None, **kwargs): + def image_2d_from( + self, grid, xp=np, operated_only: Optional[bool] = None, **kwargs + ): if self.image_2d is not None: return self.image_2d diff --git a/autogalaxy/profiles/light/snr/abstract.py b/autogalaxy/profiles/light/snr/abstract.py index 59e12b726..064db9b81 100644 --- a/autogalaxy/profiles/light/snr/abstract.py +++ b/autogalaxy/profiles/light/snr/abstract.py @@ -28,7 +28,11 @@ def __init__(self, signal_to_noise_ratio: float = 10.0): self.signal_to_noise_ratio = signal_to_noise_ratio def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> aa.Array2D: """ Abstract method for obtaining intensity at a grid of Cartesian (y,x) coordinates. diff --git a/autogalaxy/profiles/light/standard/chameleon.py b/autogalaxy/profiles/light/standard/chameleon.py index 2935ddf76..f780ae879 100644 --- a/autogalaxy/profiles/light/standard/chameleon.py +++ b/autogalaxy/profiles/light/standard/chameleon.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -45,16 +44,15 @@ def __init__( self.core_radius_0 = core_radius_0 self.core_radius_1 = core_radius_1 - @property - def axis_ratio(self) -> float: + def axis_ratio(self, xp=np) -> float: """ The elliptical isothermal mass profile deflection angles break down for perfectly spherical systems where `axis_ratio=1.0`, thus we remove these solutions. """ - axis_ratio = super().axis_ratio + axis_ratio = super().axis_ratio(xp=xp) return axis_ratio if axis_ratio < 0.99999 else 0.99999 - def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of each coordinate from the its `centre`. @@ -65,25 +63,25 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: The radial distances from the centre of the profile, for each coordinate on the grid. """ - axis_ratio_factor = (1.0 + self.axis_ratio) ** 2.0 + axis_ratio_factor = (1.0 + self.axis_ratio(xp)) ** 2.0 - return jnp.multiply( - self._intensity / (1 + self.axis_ratio), - jnp.add( - jnp.divide( + return xp.multiply( + self._intensity / (1 + self.axis_ratio(xp)), + xp.add( + xp.divide( 1.0, - jnp.sqrt( - jnp.add( - jnp.square(grid_radii.array), + xp.sqrt( + xp.add( + xp.square(grid_radii.array), (4.0 * self.core_radius_0**2.0) / axis_ratio_factor, ) ), ), - -jnp.divide( + -xp.divide( 1.0, - jnp.sqrt( - jnp.add( - jnp.square(grid_radii.array), + xp.sqrt( + xp.add( + xp.square(grid_radii.array), (4.0 * self.core_radius_1**2.0) / axis_ratio_factor, ) ), @@ -96,7 +94,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Chameleon light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. diff --git a/autogalaxy/profiles/light/standard/eff.py b/autogalaxy/profiles/light/standard/eff.py index 675834995..27cbd2df2 100644 --- a/autogalaxy/profiles/light/standard/eff.py +++ b/autogalaxy/profiles/light/standard/eff.py @@ -42,7 +42,7 @@ def __init__( self.effective_radius = effective_radius self.eta = eta - def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of each coordinate from the its `centre`. @@ -62,7 +62,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Eff light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. @@ -81,7 +85,7 @@ def image_2d_from( The image of the Eff evaluated at every (y,x) coordinate on the transformed grid. """ return self.image_2d_via_radii_from( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) ) @property diff --git a/autogalaxy/profiles/light/standard/gaussian.py b/autogalaxy/profiles/light/standard/gaussian.py index 6feaee7b4..985e0d1f6 100644 --- a/autogalaxy/profiles/light/standard/gaussian.py +++ b/autogalaxy/profiles/light/standard/gaussian.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -48,7 +47,7 @@ def coefficient_tag(self) -> str: f"sigma_{np.round(self.sigma, 2)}__ell_comps_{np.round(self.ell_comps, 2)}" ) - def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ Returns the 2D image of the Gaussian light profile from a grid of coordinates which are the radial distance of each coordinate from the its `centre`. @@ -60,12 +59,14 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: grid_radii The radial distances from the centre of the profile, for each coordinate on the grid. """ - return jnp.multiply( + return xp.multiply( self._intensity, - jnp.exp( + xp.exp( -0.5 - * jnp.square( - jnp.divide(grid_radii.array, self.sigma / jnp.sqrt(self.axis_ratio)) + * xp.square( + xp.divide( + grid_radii.array, self.sigma / xp.sqrt(self.axis_ratio(xp)) + ) ) ), ) @@ -75,7 +76,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Gaussian light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. @@ -95,7 +100,7 @@ def image_2d_from( """ return self.image_2d_via_radii_from( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs), xp=xp ) diff --git a/autogalaxy/profiles/light/standard/moffat.py b/autogalaxy/profiles/light/standard/moffat.py index adcc3af25..87a9412c9 100644 --- a/autogalaxy/profiles/light/standard/moffat.py +++ b/autogalaxy/profiles/light/standard/moffat.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -46,7 +45,7 @@ def __init__( self.alpha = alpha self.beta = beta - def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np) -> np.ndarray: """ Returns the 2D image of the Moffat light profile from a grid of coordinates which are the radial distance of each coordinate from the its `centre`. @@ -58,12 +57,14 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: grid_radii The radial distances from the centre of the profile, for each coordinate on the grid. """ - return jnp.multiply( + return xp.multiply( self._intensity, - jnp.power( + xp.power( 1 - + jnp.square( - jnp.divide(grid_radii.array, self.alpha / jnp.sqrt(self.axis_ratio)) + + xp.square( + xp.divide( + grid_radii.array, self.alpha / xp.sqrt(self.axis_ratio(xp)) + ) ), -self.beta, ), @@ -74,7 +75,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray) -> np.ndarray: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Moffat light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. @@ -94,7 +99,7 @@ def image_2d_from( """ return self.image_2d_via_radii_from( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) ) diff --git a/autogalaxy/profiles/light/standard/sersic.py b/autogalaxy/profiles/light/standard/sersic.py index 0c0de357a..acc5d3465 100644 --- a/autogalaxy/profiles/light/standard/sersic.py +++ b/autogalaxy/profiles/light/standard/sersic.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from numpy import seterr @@ -53,7 +52,7 @@ def elliptical_effective_radius(self) -> float: The elliptical effective radius instead describes the major-axis radius of the ellipse containing half the light, and may be more appropriate for highly flattened systems like disk galaxies. """ - return self.effective_radius / jnp.sqrt(self.axis_ratio) + return self.effective_radius / xp.sqrt(self.axis_ratio(xp)) @property def sersic_constant(self) -> float: @@ -80,7 +79,7 @@ def image_2d_via_radii_from(self, radius: np.ndarray) -> np.ndarray: grid_radii The radial distances from the centre of the profile, for each coordinate on the grid. """ - return self._intensity * jnp.exp( + return self._intensity * xp.exp( -self.sersic_constant * (((radius / self.effective_radius) ** (1.0 / self.sersic_index)) - 1) ) @@ -120,7 +119,9 @@ def __init__( sersic_index=sersic_index, ) - def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarray: + def image_2d_via_radii_from( + self, grid_radii: np.ndarray, xp=np, **kwargs + ) -> np.ndarray: """ Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of each coordinate from the its `centre`. @@ -131,14 +132,14 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra The radial distances from the centre of the profile, for each coordinate on the grid. """ seterr(all="ignore") - return jnp.multiply( + return xp.multiply( self._intensity, - jnp.exp( - jnp.multiply( + xp.exp( + xp.multiply( -self.sersic_constant, - jnp.add( - jnp.power( - jnp.divide(grid_radii.array, self.effective_radius), + xp.add( + xp.power( + xp.divide(grid_radii.array, self.effective_radius), 1.0 / self.sersic_index, ), -1, @@ -152,7 +153,11 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> aa.Array2D: """ Returns the Sersic light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. @@ -171,9 +176,9 @@ def image_2d_from( The image of the Sersic evaluated at every (y,x) coordinate on the transformed grid. """ - grid_radii = self.eccentric_radii_grid_from(grid=grid, **kwargs) + grid_radii = self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) - return self.image_2d_via_radii_from(grid_radii=grid_radii, **kwargs) + return self.image_2d_via_radii_from(grid_radii=grid_radii, xp=xp, **kwargs) class SersicSph(Sersic): diff --git a/autogalaxy/profiles/light/standard/sersic_core.py b/autogalaxy/profiles/light/standard/sersic_core.py index dbc444d04..9aaa81c9b 100644 --- a/autogalaxy/profiles/light/standard/sersic_core.py +++ b/autogalaxy/profiles/light/standard/sersic_core.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -53,8 +52,7 @@ def __init__( self.alpha = alpha self.gamma = gamma - @property - def intensity_prime(self) -> float: + def intensity_prime(self, xp=np) -> float: """ Overall intensity normalisation in the rescaled cored Sersic light profile. @@ -64,7 +62,7 @@ def intensity_prime(self) -> float: return ( self._intensity * (2.0 ** (-self.gamma / self.alpha)) - * jnp.exp( + * xp.exp( self.sersic_constant * ( ((2.0 ** (1.0 / self.alpha)) * self.radius_break) @@ -74,7 +72,9 @@ def intensity_prime(self) -> float: ) ) - def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarray: + def image_2d_via_radii_from( + self, grid_radii: np.ndarray, xp=np, **kwargs + ) -> np.ndarray: """ Returns the 2D image of the Sersic light profile from a grid of coordinates which are the radial distances of each coordinate from the its `centre`. @@ -85,27 +85,27 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray, **kwargs) -> np.ndarra The radial distances from the centre of the profile, for each coordinate on the grid. """ - return jnp.multiply( - jnp.multiply( - self.intensity_prime, - jnp.power( - jnp.add( + return xp.multiply( + xp.multiply( + self.intensity_prime(xp), + xp.power( + xp.add( 1, - jnp.power( - jnp.divide(self.radius_break, grid_radii.array), self.alpha + xp.power( + xp.divide(self.radius_break, grid_radii.array), self.alpha ), ), (self.gamma / self.alpha), ), ), - jnp.exp( - jnp.multiply( + xp.exp( + xp.multiply( -self.sersic_constant, ( - jnp.power( - jnp.divide( - jnp.add( - jnp.power(grid_radii.array, self.alpha), + xp.power( + xp.divide( + xp.add( + xp.power(grid_radii.array, self.alpha), (self.radius_break**self.alpha), ), (self.effective_radius**self.alpha), diff --git a/autogalaxy/profiles/light/standard/shapelets/cartesian.py b/autogalaxy/profiles/light/standard/shapelets/cartesian.py index 76785e471..9be1b123c 100644 --- a/autogalaxy/profiles/light/standard/shapelets/cartesian.py +++ b/autogalaxy/profiles/light/standard/shapelets/cartesian.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -64,7 +63,11 @@ def coefficient_tag(self) -> str: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Cartesian Shapelet light profile's 2D image from a 2D grid of Cartesian (y,x) coordinates. @@ -97,12 +100,12 @@ def image_2d_from( return ( shapelet_y * shapelet_x - * jnp.exp(-0.5 * (y**2 + x**2) / (self.beta**2)) + * xp.exp(-0.5 * (y**2 + x**2) / (self.beta**2)) / self.beta / ( - jnp.sqrt( + xp.sqrt( 2 ** (self.n_x + self.n_y) - * (jnp.pi) + * (xp.pi) * factorial(self.n_y) * factorial(self.n_x) ) diff --git a/autogalaxy/profiles/light/standard/shapelets/exponential.py b/autogalaxy/profiles/light/standard/shapelets/exponential.py index 9dbf5bee7..0231d7351 100644 --- a/autogalaxy/profiles/light/standard/shapelets/exponential.py +++ b/autogalaxy/profiles/light/standard/shapelets/exponential.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -65,7 +64,11 @@ def coefficient_tag(self) -> str: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Exponential Shapelet light profile's 2D image from a 2D grid of Exponential (y,x) coordinates. @@ -87,30 +90,30 @@ def image_2d_from( from jax.scipy.special import factorial radial = (grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2) / self.beta - theta = jnp.arctan(grid.array[:, 1] / grid.array[:, 0]) + theta = xp.arctan(grid.array[:, 1] / grid.array[:, 0]) prefactor = ( 1.0 - / jnp.sqrt(2 * jnp.pi) + / xp.sqrt(2 * xp.pi) / self.beta - * (self.n + 0.5) ** (-1 - jnp.abs(self.m)) + * (self.n + 0.5) ** (-1 - xp.abs(self.m)) * (-1) ** (self.n + self.m) - * jnp.sqrt( - factorial(self.n - jnp.abs(self.m)) / 2 * self.n - + 1 / factorial(self.n + jnp.abs(self.m)) + * xp.sqrt( + factorial(self.n - xp.abs(self.m)) / 2 * self.n + + 1 / factorial(self.n + xp.abs(self.m)) ) ) - laguerre = genlaguerre(n=self.n - jnp.abs(self.m), alpha=2 * jnp.abs(self.m)) + laguerre = genlaguerre(n=self.n - xp.abs(self.m), alpha=2 * xp.abs(self.m)) shapelet = laguerre(radial / (self.n + 0.5)) - return jnp.abs( + return xp.abs( prefactor - * jnp.exp(-radial / (2 * self.n + 1)) - * radial ** (jnp.abs(self.m)) + * xp.exp(-radial / (2 * self.n + 1)) + * radial ** (xp.abs(self.m)) * shapelet - * jnp.cos(self.m * theta) - + -1.0j * jnp.sin(self.m * theta) + * xp.cos(self.m * theta) + + -1.0j * xp.sin(self.m * theta) ) diff --git a/autogalaxy/profiles/light/standard/shapelets/polar.py b/autogalaxy/profiles/light/standard/shapelets/polar.py index 9733308fc..1d6c98db1 100644 --- a/autogalaxy/profiles/light/standard/shapelets/polar.py +++ b/autogalaxy/profiles/light/standard/shapelets/polar.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Optional, Tuple @@ -65,7 +64,11 @@ def coefficient_tag(self) -> str: @check_operated_only @aa.grid_dec.transform def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + operated_only: Optional[bool] = None, + **kwargs, ) -> np.ndarray: """ Returns the Polar Shapelet light profile's 2D image from a 2D grid of Polar (y,x) coordinates. @@ -86,30 +89,28 @@ def image_2d_from( from scipy.special import genlaguerre from jax.scipy.special import factorial - laguerre = genlaguerre( - n=(self.n - jnp.abs(self.m)) / 2.0, alpha=jnp.abs(self.m) - ) + laguerre = genlaguerre(n=(self.n - xp.abs(self.m)) / 2.0, alpha=xp.abs(self.m)) const = ( - ((-1) ** ((self.n - jnp.abs(self.m)) // 2)) - * jnp.sqrt( - factorial((self.n - jnp.abs(self.m)) // 2) - / factorial((self.n + jnp.abs(self.m)) // 2) + ((-1) ** ((self.n - xp.abs(self.m)) // 2)) + * xp.sqrt( + factorial((self.n - xp.abs(self.m)) // 2) + / factorial((self.n + xp.abs(self.m)) // 2) ) / self.beta - / jnp.sqrt(jnp.pi) + / xp.sqrt(xp.pi) ) rsq = (grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2) / self.beta**2 - theta = jnp.arctan2(grid.array[:, 1], grid.array[:, 0]) - radial = rsq ** (abs(self.m / 2.0)) * jnp.exp(-rsq / 2.0) * laguerre(rsq) + theta = xp.arctan2(grid.array[:, 1], grid.array[:, 0]) + radial = rsq ** (abs(self.m / 2.0)) * xp.exp(-rsq / 2.0) * laguerre(rsq) if self.m == 0: azimuthal = 1 elif self.m > 0: - azimuthal = jnp.sin((-1) * self.m * theta) + azimuthal = xp.sin((-1) * self.m * theta) else: - azimuthal = jnp.cos((-1) * self.m * theta) + azimuthal = xp.cos((-1) * self.m * theta) return const * radial * azimuthal diff --git a/autogalaxy/profiles/mass/abstract/abstract.py b/autogalaxy/profiles/mass/abstract/abstract.py index 80a2da261..44b0e542f 100644 --- a/autogalaxy/profiles/mass/abstract/abstract.py +++ b/autogalaxy/profiles/mass/abstract/abstract.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -44,28 +43,20 @@ def deflections_2d_via_potential_2d_from(self, grid): mask=grid.mask, ) - def convergence_2d_from(self, grid): + def convergence_2d_from(self, grid, xp=np): raise NotImplementedError def convergence_func(self, grid_radius: float) -> float: raise NotImplementedError - @aa.grid_dec.project_grid - def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> aa.type.Grid1D2DLike: - return self.convergence_2d_from(grid=grid) - def potential_2d_from(self, grid): raise NotImplementedError - @aa.grid_dec.project_grid - def potential_1d_from(self, grid: aa.type.Grid1D2DLike) -> aa.type.Grid1D2DLike: - return self.potential_2d_from(grid=grid) - def potential_func(self, u, y, x): raise NotImplementedError - def mass_integral(self, x): - return 2 * jnp.pi * x * self.convergence_func(grid_radius=aa.ArrayIrregular(x)) + def mass_integral(self, x, xp=np): + return 2 * xp.pi * x * self.convergence_func(grid_radius=aa.ArrayIrregular(x)) @property def ellipticity_rescale(self): diff --git a/autogalaxy/profiles/mass/abstract/cse.py b/autogalaxy/profiles/mass/abstract/cse.py index a6db34172..32af87396 100644 --- a/autogalaxy/profiles/mass/abstract/cse.py +++ b/autogalaxy/profiles/mass/abstract/cse.py @@ -1,5 +1,4 @@ from abc import ABC, abstractmethod -import jax.numpy as jnp import numpy as np from typing import Callable, List, Tuple @@ -41,12 +40,12 @@ def deflections_via_cse_from( Parameters ---------- """ - phi = jnp.sqrt(axis_ratio_squared * core_radius**2.0 + term1) + phi = np.sqrt(axis_ratio_squared * core_radius**2.0 + term1) Psi = (phi + core_radius) ** 2.0 + term2 bottom = core_radius * phi * Psi defl_x = (term3 * (phi + axis_ratio_squared * core_radius)) / bottom defl_y = (term4 * (phi + core_radius)) / bottom - return jnp.vstack((defl_y, defl_x)) + return np.vstack((defl_y, defl_x)) @abstractmethod def decompose_convergence_via_cse(self, grid_radii: np.ndarray): @@ -166,7 +165,7 @@ def _deflections_2d_via_cse_from(self, grid: np.ndarray, **kwargs) -> np.ndarray amplitude_list, core_radius_list = self.decompose_convergence_via_cse( grid_radii=self.radial_grid_from(grid=grid, **kwargs) ) - q = self.axis_ratio + q = self.axis_ratio() q2 = q**2.0 grid_y = grid.array[:, 0] grid_x = grid.array[:, 1] diff --git a/autogalaxy/profiles/mass/abstract/jax_utils.py b/autogalaxy/profiles/mass/abstract/jax_utils.py index e241488f3..de809f493 100644 --- a/autogalaxy/profiles/mass/abstract/jax_utils.py +++ b/autogalaxy/profiles/mass/abstract/jax_utils.py @@ -1,6 +1,4 @@ -import jax.numpy as jnp - -from jax import custom_jvp +# from jax import custom_jvp r1_s1 = [2.5, 2, 1.5, 1, 0.5] @@ -26,7 +24,7 @@ def reg2(z, sqrt_pi, _): for s in r2_s2: f2 = s - f2 * mz2 - return jnp.exp(mz2) + 1j * z * f1 / f2 + return xp.exp(mz2) + 1j * z * f1 / f2 r3_s1 = [5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793] @@ -45,55 +43,56 @@ def reg3(z, sqrt_pi, _): return f1 / f2 -@custom_jvp -def w_f_approx(z): - """Compute the Faddeeva function :math:`w_{\\mathrm F}(z)` using the - approximation given in Zaghloul (2017). - - :param z: complex number - :type z: ``complex`` or ``numpy.array(dtype=complex)`` - :return: :math:`w_\\mathrm{F}(z)` - :rtype: ``complex`` - """ - sqrt_pi = 1 / jnp.sqrt(jnp.pi) - i_sqrt_pi = 1j * sqrt_pi - - z_imag2 = z.imag**2 - abs_z2 = z.real**2 + z_imag2 - - # use a single partial fraction approx for all large abs(z)**2 - # to have better approx of the auto-derivatives - r1 = (abs_z2 >= 62.0) | ((abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13)) - # region bounds for 5 taken directly from Zaghloul (2017) - # https://dl.acm.org/doi/pdf/10.1145/3119904 - r2_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13) - r2_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072) - r2 = r2_1 | r2_2 - r3 = jnp.logical_not(r1) & jnp.logical_not(r2) - - # exploit symmetry to avoid overflow in some regions - r_flip = z.imag < 0 - z_adjust = jnp.where(r_flip, -z, z) - two_exp_zz = 2 * jnp.exp(-(z_adjust**2)) - - args = (z_adjust, sqrt_pi, i_sqrt_pi) - wz = jnp.empty_like(z) - wz = jnp.where(r1, reg1(*args), wz) - wz = jnp.where(r2, reg2(*args), wz) - wz = jnp.where(r3, reg3(*args), wz) - - # exploit symmetry to avoid overflow in some regions - wz = jnp.where(r_flip, two_exp_zz - wz, wz) - - return wz - - -@w_f_approx.defjvp -def w_f_approx_jvp(primals, tangents): - # define a custom jvp to avoid the issue using `jnp.where` with `jax.grad` - (z,) = primals - (z_dot,) = tangents - primal_out = w_f_approx(z) - i_sqrt_pi = 1j / jnp.sqrt(jnp.pi) - tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out) - return primal_out, tangent_out +# +# @custom_jvp +# def w_f_approx(z): +# """Compute the Faddeeva function :math:`w_{\\mathrm F}(z)` using the +# approximation given in Zaghloul (2017). +# +# :param z: complex number +# :type z: ``complex`` or ``numpy.array(dtype=complex)`` +# :return: :math:`w_\\mathrm{F}(z)` +# :rtype: ``complex`` +# """ +# sqrt_pi = 1 / xp.sqrt(xp.pi) +# i_sqrt_pi = 1j * sqrt_pi +# +# z_imag2 = z.imag**2 +# abs_z2 = z.real**2 + z_imag2 +# +# # use a single partial fraction approx for all large abs(z)**2 +# # to have better approx of the auto-derivatives +# r1 = (abs_z2 >= 62.0) | ((abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 >= 1e-13)) +# # region bounds for 5 taken directly from Zaghloul (2017) +# # https://dl.acm.org/doi/pdf/10.1145/3119904 +# r2_1 = (abs_z2 >= 30.0) & (abs_z2 < 62.0) & (z_imag2 < 1e-13) +# r2_2 = (abs_z2 >= 2.5) & (abs_z2 < 30.0) & (z_imag2 < 0.072) +# r2 = r2_1 | r2_2 +# r3 = xp.logical_not(r1) & xp.logical_not(r2) +# +# # exploit symmetry to avoid overflow in some regions +# r_flip = z.imag < 0 +# z_adjust = xp.where(r_flip, -z, z) +# two_exp_zz = 2 * xp.exp(-(z_adjust**2)) +# +# args = (z_adjust, sqrt_pi, i_sqrt_pi) +# wz = xp.empty_like(z) +# wz = xp.where(r1, reg1(*args), wz) +# wz = xp.where(r2, reg2(*args), wz) +# wz = xp.where(r3, reg3(*args), wz) +# +# # exploit symmetry to avoid overflow in some regions +# wz = xp.where(r_flip, two_exp_zz - wz, wz) +# +# return wz + + +# @w_f_approx.defjvp +# def w_f_approx_jvp(primals, tangents): +# # define a custom jvp to avoid the issue using `xp.where` with `jax.grad` +# (z,) = primals +# (z_dot,) = tangents +# primal_out = w_f_approx(z) +# i_sqrt_pi = 1j / xp.sqrt(xp.pi) +# tangent_out = z_dot * 2 * (i_sqrt_pi - z * primal_out) +# return primal_out, tangent_out diff --git a/autogalaxy/profiles/mass/abstract/mge_jax.py b/autogalaxy/profiles/mass/abstract/mge_jax.py index 81beab896..b6c2d0b83 100644 --- a/autogalaxy/profiles/mass/abstract/mge_jax.py +++ b/autogalaxy/profiles/mass/abstract/mge_jax.py @@ -1,6 +1,4 @@ -import jax.numpy as jnp - -from .jax_utils import w_f_approx +# from .jax_utils import w_f_approx class MassProfileMGE: @@ -21,12 +19,12 @@ def zeta_from(grid, amps, sigmas, axis_ratio): """ q2 = axis_ratio**2.0 - scale_factor = axis_ratio / jnp.sqrt(2.0 * (1.0 - q2)) + scale_factor = axis_ratio / xp.sqrt(2.0 * (1.0 - q2)) - xs = jnp.array((grid.array[:, 1] * scale_factor).copy()) - ys = jnp.array((grid.array[:, 0] * scale_factor).copy()) + xs = xp.array((grid.array[:, 1] * scale_factor).copy()) + ys = xp.array((grid.array[:, 0] * scale_factor).copy()) - y_sign = jnp.sign(ys) + y_sign = xp.sign(ys) ys = ys * y_sign z = xs + 1j * ys @@ -40,7 +38,7 @@ def zeta_from(grid, amps, sigmas, axis_ratio): # could try `jax.lax.scan` instead if this is too much memory w = w_f_approx(inv_sigma_ * z) wq = w_f_approx(inv_sigma_ * zq) - exp_factor = jnp.exp(inv_sigma_**2 * expv) + exp_factor = xp.exp(inv_sigma_**2 * expv) sigma_func_real = w.imag - exp_factor * wq.imag sigma_func_imag = (-w.real + exp_factor * wq.real) * y_sign @@ -53,8 +51,8 @@ def kesi(p): """ see Eq.(6) of 1906.08263 """ - n_list = jnp.arange(0, 2 * p + 1, 1) - return (2.0 * p * jnp.log(10) / 3.0 + 2.0 * jnp.pi * n_list * 1j) ** (0.5) + 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): @@ -62,15 +60,15 @@ def eta(p): see Eq.(6) of 1906.00263 """ - i = jnp.arange(1, p, 1) + i = xp.arange(1, p, 1) kesi_last = 1 / 2**p - k = kesi_last + jnp.cumsum(jnp.cumprod((p + 1 - i) / i) * kesi_last) + k = kesi_last + xp.cumsum(xp.cumprod((p + 1 - i) / i) * kesi_last) - kesi_list = jnp.hstack( - [jnp.array([0.5]), jnp.ones(p), k[::-1], jnp.array([kesi_last])] + kesi_list = xp.hstack( + [xp.array([0.5]), xp.ones(p), k[::-1], xp.array([kesi_last])] ) - coef = (-1) ** jnp.arange(0, 2 * p + 1, 1) - eta_const = 2.0 * jnp.sqrt(2.0 * jnp.pi) * 10 ** (p / 3.0) + coef = (-1) ** xp.arange(0, 2 * p + 1, 1) + eta_const = 2.0 * xp.sqrt(2.0 * xp.pi) * 10 ** (p / 3.0) eta_list = coef * kesi_list return eta_const, eta_list @@ -102,16 +100,14 @@ def _decompose_convergence_via_mge( # sigma is sampled from logspace between these radii. - log_sigmas = jnp.linspace( - jnp.log(radii_min), jnp.log(radii_max), func_gaussians - ) + log_sigmas = xp.linspace(xp.log(radii_min), xp.log(radii_max), func_gaussians) d_log_sigma = log_sigmas[1] - log_sigmas[0] - sigma_list = jnp.exp(log_sigmas) + sigma_list = xp.exp(log_sigmas) - f_sigma = eta_constant * jnp.sum( - eta_n * jnp.real(func(sigma_list.reshape(-1, 1) * kesis)), axis=1 + f_sigma = eta_constant * xp.sum( + eta_n * xp.real(func(sigma_list.reshape(-1, 1) * kesis)), axis=1 ) - amplitude_list = f_sigma * d_log_sigma / jnp.sqrt(2.0 * jnp.pi) + amplitude_list = f_sigma * d_log_sigma / xp.sqrt(2.0 * xp.pi) amplitude_list = amplitude_list.at[0].multiply(0.5) amplitude_list = amplitude_list.at[-1].multiply(0.5) @@ -137,13 +133,13 @@ def _convergence_2d_via_mge_from( inv_sigma_ = 1 / sigmas.reshape((-1,) + (1,) * grid_radii.array.ndim) amps_ = amps.reshape((-1,) + (1,) * grid_radii.array.ndim) - convergence = amps_ * jnp.exp(-0.5 * (grid_radii.array * inv_sigma_) ** 2) + convergence = amps_ * xp.exp(-0.5 * (grid_radii.array * inv_sigma_) ** 2) return convergence.sum(axis=0) def _deflections_2d_via_mge_from( self, grid, sigmas_factor=1.0, func_terms=28, func_gaussians=20 ): - axis_ratio = jnp.min(jnp.array([self.axis_ratio, 0.9999])) + axis_ratio = xp.min(xp.array([self.axis_ratio(xp), 0.9999])) amps, sigmas = self.decompose_convergence_via_mge( func_terms=func_terms, func_gaussians=func_gaussians @@ -154,8 +150,8 @@ def _deflections_2d_via_mge_from( grid=grid, amps=amps, sigmas=sigmas, axis_ratio=axis_ratio ) - angle *= jnp.sqrt((2.0 * jnp.pi) / (1.0 - axis_ratio**2.0)) + angle *= xp.sqrt((2.0 * xp.pi) / (1.0 - axis_ratio**2.0)) return self.rotated_grid_from_reference_frame_from( - jnp.vstack((-angle.imag, angle.real)).T + xp.vstack((-angle.imag, angle.real)).T ) diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy.py b/autogalaxy/profiles/mass/abstract/mge_numpy.py index 124d278cf..a4fcbd29a 100644 --- a/autogalaxy/profiles/mass/abstract/mge_numpy.py +++ b/autogalaxy/profiles/mass/abstract/mge_numpy.py @@ -265,7 +265,7 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity): def _deflections_2d_via_mge_from( self, grid, sigmas_factor=1.0, func_terms=None, func_gaussians=None ): - axis_ratio = np.array(self.axis_ratio) + axis_ratio = np.array(self.axis_ratio()) if axis_ratio > 0.9999: axis_ratio = 0.9999 diff --git a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py index 5f5b8a97b..9a9c6809c 100644 --- a/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py +++ b/autogalaxy/profiles/mass/abstract/mge_numpy_coleman.py @@ -261,7 +261,7 @@ def convergence_func_gaussian(self, grid_radii, sigma, intensity): ) def _deflections_2d_via_mge_from(self, grid, sigmas_factor=1.0): - axis_ratio = self.axis_ratio + axis_ratio = self.axis_ratio(xp) if axis_ratio > 0.9999: axis_ratio = 0.9999 diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py index 34dfc3430..59f7c2f2f 100644 --- a/autogalaxy/profiles/mass/dark/abstract.py +++ b/autogalaxy/profiles/mass/dark/abstract.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -57,7 +56,7 @@ def __init__( @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -74,7 +73,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -130,7 +129,7 @@ def gnfw_3d(r): amplitude_list *= np.sqrt(2.0 * np.pi) * sigma_list return amplitude_list, sigma_list - def coord_func_f(self, grid_radius: jnp.ndarray) -> jnp.ndarray: + 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 @@ -142,24 +141,24 @@ def coord_func_f(self, grid_radius: jnp.ndarray) -> jnp.ndarray: any Python control flow on tracer values. """ if isinstance(grid_radius, float) or isinstance(grid_radius, complex): - grid_radius = jnp.array([grid_radius]) + grid_radius = xp.array([grid_radius]) - f = jnp.ones(shape=grid_radius.shape[0], dtype="complex64") + f = xp.ones(shape=grid_radius.shape[0], dtype="complex64") # compute both branches r = grid_radius inv_r = 1.0 / r # branch for r > 1 - out_gt = (1.0 / jnp.sqrt(r**2 - 1.0)) * jnp.arccos(inv_r) + out_gt = (1.0 / xp.sqrt(r**2 - 1.0)) * xp.arccos(inv_r) # branch for r < 1 - out_lt = (1.0 / jnp.sqrt(1.0 - r**2)) * jnp.arccosh(inv_r) + out_lt = (1.0 / xp.sqrt(1.0 - r**2)) * xp.arccosh(inv_r) # combine: if r>1 pick out_gt, elif r<1 pick out_lt, else keep original f - return jnp.where(r > 1.0, out_gt, jnp.where(r < 1.0, out_lt, f)) + return xp.where(r > 1.0, out_gt, xp.where(r < 1.0, out_lt, f)) - def coord_func_g(self, grid_radius: jnp.ndarray) -> jnp.ndarray: + def coord_func_g(self, grid_radius: np.ndarray, xp=np) -> np.ndarray: """ Vectorized version of the original looped `coord_func_g_jit`. @@ -170,40 +169,42 @@ def coord_func_g(self, grid_radius: jnp.ndarray) -> jnp.ndarray: Parameters ---------- - grid_radius : jnp.ndarray + grid_radius : np.ndarray The input grid radius values. - f_r : jnp.ndarray + f_r : np.ndarray Precomputed values from `coord_func_f`. - g : jnp.ndarray + g : np.ndarray Output array (will be overwritten). Returns ------- - jnp.ndarray + np.ndarray The updated `g` array. """ # Convert single values to JAX arrays if isinstance(grid_radius, (float, complex)): - grid_radius = jnp.array([grid_radius], dtype=jnp.complex64) + grid_radius = xp.array([grid_radius], dtype=xp.complex64) # Evaluate f_r - f_r = self.coord_func_f(grid_radius=grid_radius) + f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp) - r = jnp.real(grid_radius) + r = xp.real(grid_radius) r2 = r**2 - return jnp.where( + return xp.where( r > 1.0, (1.0 - f_r) / (r2 - 1.0), - jnp.where( + xp.where( r < 1.0, (f_r - 1.0) / (1.0 - r2), 1.0 / 3.0, ), ) - def coord_func_h(self, grid_radius): - return jnp.log(grid_radius / 2.0) + self.coord_func_f(grid_radius=grid_radius) + def coord_func_h(self, grid_radius, xp=np): + return xp.log(grid_radius / 2.0) + self.coord_func_f( + grid_radius=grid_radius, xp=xp + ) def rho_at_scale_radius_solar_mass_per_kpc3( self, redshift_object, redshift_source, cosmology: LensingCosmology = None @@ -270,6 +271,7 @@ def concentration( redshift_source, redshift_of_cosmic_average_density="profile", cosmology: LensingCosmology = None, + xp=np, ): from autogalaxy.cosmology.wrap import Planck15 @@ -289,11 +291,13 @@ def concentration( ) return fsolve( - func=self.concentration_func, x0=10.0, args=(delta_concentration,) + func=self.concentration_func, + x0=10.0, + args=(delta_concentration, xp), )[0] @staticmethod - def concentration_func(concentration, delta_concentration): + def concentration_func(concentration, delta_concentration, xp=np): return ( 200.0 / 3.0 @@ -301,7 +305,7 @@ def concentration_func(concentration, delta_concentration): concentration * concentration * concentration - / (jnp.log(1 + concentration) - concentration / (1 + concentration)) + / (xp.log(1 + concentration) - concentration / (1 + concentration)) ) - delta_concentration ) @@ -312,6 +316,7 @@ def radius_at_200( redshift_source, redshift_of_cosmic_average_density="profile", cosmology: LensingCosmology = None, + xp=np, ): """ Returns `r_{200m}` for this halo in **arcseconds** @@ -325,6 +330,7 @@ def radius_at_200( redshift_source=redshift_source, redshift_of_cosmic_average_density=redshift_of_cosmic_average_density, cosmology=cosmology, + xp=xp, ) return concentration * self.scale_radius @@ -335,6 +341,7 @@ def mass_at_200_solar_masses( redshift_source, redshift_of_cosmic_average_density="profile", cosmology: LensingCosmology = None, + xp=np, ): from autogalaxy.cosmology.wrap import Planck15 @@ -364,6 +371,7 @@ def mass_at_200_solar_masses( redshift_source=redshift_source, redshift_of_cosmic_average_density=redshift_of_cosmic_average_density, cosmology=cosmology, + xp=xp, ) kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) @@ -372,11 +380,11 @@ def mass_at_200_solar_masses( return ( 200.0 - * ((4.0 / 3.0) * jnp.pi) + * ((4.0 / 3.0) * xp.pi) * cosmic_average_density * (radius_at_200_kpc**3.0) ) @property def ellipticity_rescale(self): - return 1.0 - ((1.0 - self.axis_ratio) / 2.0) + return 1.0 - ((1.0 - self.axis_ratio()) / 2.0) diff --git a/autogalaxy/profiles/mass/dark/gnfw.py b/autogalaxy/profiles/mass/dark/gnfw.py index 047154085..a1a80295b 100644 --- a/autogalaxy/profiles/mass/dark/gnfw.py +++ b/autogalaxy/profiles/mass/dark/gnfw.py @@ -7,20 +7,20 @@ class gNFW(AbstractgNFW): - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_mge_from(grid=grid, **kwargs) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_mge_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self._deflections_2d_via_mge_from( - grid=grid, sigmas_factor=self.axis_ratio + grid=grid, sigmas_factor=self.axis_ratio(xp) ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform def deflections_2d_via_integral_from( - self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs + self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -51,7 +51,7 @@ def calculate_deflection_component(npow, yx_index): deflection_grid[i] = ( 2.0 * self.kappa_s - * self.axis_ratio + * self.axis_ratio(xp) * grid[i, yx_index] * quad( self.deflection_func, @@ -61,7 +61,7 @@ def calculate_deflection_component(npow, yx_index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, + self.axis_ratio(xp), minimum_log_eta, maximum_log_eta, tabulate_bins, @@ -104,7 +104,7 @@ def calculate_deflection_component(npow, yx_index): deflection_x = calculate_deflection_component(npow=0.0, yx_index=1) return self.rotated_grid_from_reference_frame_from( - np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T) + np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), ) @staticmethod @@ -162,7 +162,9 @@ def integral_y(y, eta): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def potential_2d_from(self, grid: aa.type.Grid2DLike, tabulate_bins=1000, **kwargs): + def potential_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, tabulate_bins=1000, **kwargs + ): """ Calculate the potential at a given set of arc-second gridded coordinates. @@ -220,14 +222,14 @@ def deflection_integrand(x, kappa_radius, scale_radius, inner_slope): ) for i in range(grid.shape[0]): - potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio) * quad( + potential_grid[i] = (2.0 * self.kappa_s * self.axis_ratio(xp)) * quad( self.potential_func, a=0.0, b=1.0, args=( grid.array[i, 0], grid.array[i, 1], - self.axis_ratio, + self.axis_ratio(xp), minimum_log_eta, maximum_log_eta, tabulate_bins, @@ -295,7 +297,9 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_integral_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -316,7 +320,9 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): 4.0 * self.kappa_s * self.scale_radius, self.deflection_func_sph(eta[i]) ) - return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid) + return self._cartesian_grid_via_radial_from( + grid=grid, radius=deflection_grid, xp=xp + ) @staticmethod def deflection_integrand(y, eta, inner_slope): diff --git a/autogalaxy/profiles/mass/dark/nfw.py b/autogalaxy/profiles/mass/dark/nfw.py index ae8fc3018..dc425f69b 100644 --- a/autogalaxy/profiles/mass/dark/nfw.py +++ b/autogalaxy/profiles/mass/dark/nfw.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -48,7 +47,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_integral_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -61,7 +62,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): from scipy.integrate import quad def calculate_deflection_component(npow, index): - deflection_grid = np.array(self.axis_ratio * grid.array[:, index]) + deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index]) for i in range(grid.shape[0]): deflection_grid[i] *= ( @@ -74,7 +75,7 @@ def calculate_deflection_component(npow, index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, + self.axis_ratio(xp), self.scale_radius, ), )[0] @@ -91,7 +92,7 @@ def calculate_deflection_component(npow, index): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self._deflections_2d_via_cse_from(grid=grid, **kwargs) @staticmethod @@ -121,7 +122,7 @@ def deflection_func(u, y, x, npow, axis_ratio, scale_radius): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the projected 2D convergence from a grid of (y,x) arc second coordinates, by computing and summing the convergence of each individual cse used to decompose the mass profile. @@ -143,13 +144,15 @@ def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): def convergence_func(self, grid_radius: float) -> float: grid_radius = (1.0 / self.scale_radius) * grid_radius.array + 0j return np.real( - 2.0 * self.kappa_s * np.array(self.coord_func_g(grid_radius=grid_radius)) + 2.0 + * self.kappa_s + * np.array(self.coord_func_g(grid_radius=grid_radius, xp=xp)) ) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the potential at a given set of arc-second gridded coordinates. @@ -171,7 +174,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): args=( grid.array[i, 0], grid.array[i, 1], - self.axis_ratio, + self.axis_ratio(xp), self.kappa_s, self.scale_radius, ), @@ -254,7 +257,7 @@ def nfw_2d(r): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Analytic calculation shear from Heyrovský & Karamazov 2024 @@ -287,14 +290,14 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # Rotation for shear shear_field = self.rotated_grid_from_reference_frame_from( - grid=np.vstack((g2, g1)).T, angle=self.angle * 2 + grid=np.vstack((g2, g1)).T, angle=self.angle(xp) * 2 ) return aa.VectorYX2DIrregular(values=shear_field, grid=grid) @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Analytic calculation convergence from Heyrovský & Karamazov 2024 @@ -360,12 +363,14 @@ def __init__( scale_radius=scale_radius, ) - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_analytic_from(grid=grid, **kwargs) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_analytic_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -375,25 +380,27 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - eta = jnp.multiply( + eta = xp.multiply( 1.0 / self.scale_radius, self.radial_grid_from(grid=grid, **kwargs).array ) - deflection_grid = jnp.multiply( + deflection_grid = xp.multiply( (4.0 * self.kappa_s * self.scale_radius / eta), - self.deflection_func_sph(grid_radius=eta), + self.deflection_func_sph(grid_radius=eta, xp=xp), ) - return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid) + return self._cartesian_grid_via_radial_from( + grid=grid, radius=deflection_grid, xp=xp + ) - def deflection_func_sph(self, grid_radius): + def deflection_func_sph(self, grid_radius, xp=np): grid_radius = grid_radius + 0j - return jnp.real(self.coord_func_h(grid_radius=grid_radius)) + return xp.real(self.coord_func_h(grid_radius=grid_radius)) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the potential at a given set of arc-second gridded coordinates. @@ -407,7 +414,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid=grid, **kwargs ) + 0j return np.real( - 2.0 * self.scale_radius * self.kappa_s * self.potential_func_sph(eta) + 2.0 * self.scale_radius * self.kappa_s * self.potential_func_sph(eta, xp=xp) ) @staticmethod diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py index b33f4d0d7..c6e4fab98 100644 --- a/autogalaxy/profiles/mass/dark/nfw_truncated.py +++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py @@ -1,4 +1,4 @@ -import jax.numpy as jnp +import numpy as np from typing import Tuple import autoarray as aa @@ -28,7 +28,7 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -38,68 +38,70 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - eta = jnp.multiply( + eta = xp.multiply( 1.0 / self.scale_radius, self.radial_grid_from(grid=grid, **kwargs).array ) - deflection_grid = jnp.multiply( + deflection_grid = xp.multiply( (4.0 * self.kappa_s * self.scale_radius / eta), self.deflection_func_sph(grid_radius=eta), ) - return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_grid) + return self._cartesian_grid_via_radial_from( + grid=grid, radius=deflection_grid, xp=xp + ) - def deflection_func_sph(self, grid_radius): + def deflection_func_sph(self, grid_radius, xp=np): grid_radius = grid_radius + 0j - return jnp.real(self.coord_func_m(grid_radius=grid_radius)) + return xp.real(self.coord_func_m(grid_radius=grid_radius, xp=xp)) - def convergence_func(self, grid_radius: float) -> float: + def convergence_func(self, grid_radius: float, xp=np) -> float: grid_radius = ((1.0 / self.scale_radius) * grid_radius) + 0j - return jnp.real( - 2.0 * self.kappa_s * self.coord_func_l(grid_radius=grid_radius.array) + return xp.real( + 2.0 * self.kappa_s * self.coord_func_l(grid_radius=grid_radius.array, xp=xp) ) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) - def coord_func_k(self, grid_radius): - return jnp.log( - jnp.divide( + def coord_func_k(self, grid_radius, xp=np): + return xp.log( + xp.divide( grid_radius, - jnp.sqrt(jnp.square(grid_radius) + jnp.square(self.tau)) + self.tau, + xp.sqrt(xp.square(grid_radius) + xp.square(self.tau)) + self.tau, ) ) - def coord_func_l(self, grid_radius): - f_r = self.coord_func_f(grid_radius=grid_radius) - g_r = self.coord_func_g(grid_radius=grid_radius) - k_r = self.coord_func_k(grid_radius=grid_radius) + def coord_func_l(self, grid_radius, xp=np): + f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp) + g_r = self.coord_func_g(grid_radius=grid_radius, xp=xp) + k_r = self.coord_func_k(grid_radius=grid_radius, xp=xp) - return jnp.divide(self.tau**2.0, (self.tau**2.0 + 1.0) ** 2.0) * ( + return xp.divide(self.tau**2.0, (self.tau**2.0 + 1.0) ** 2.0) * ( ((self.tau**2.0 + 1.0) * g_r) + (2 * f_r) - - (jnp.pi / (jnp.sqrt(self.tau**2.0 + grid_radius**2.0))) + - (xp.pi / (xp.sqrt(self.tau**2.0 + grid_radius**2.0))) + ( ( (self.tau**2.0 - 1.0) - / (self.tau * (jnp.sqrt(self.tau**2.0 + grid_radius**2.0))) + / (self.tau * (xp.sqrt(self.tau**2.0 + grid_radius**2.0))) ) * k_r ) ) - def coord_func_m(self, grid_radius): - f_r = self.coord_func_f(grid_radius=grid_radius) - k_r = self.coord_func_k(grid_radius=grid_radius) + def coord_func_m(self, grid_radius, xp=np): + f_r = self.coord_func_f(grid_radius=grid_radius, xp=xp) + k_r = self.coord_func_k(grid_radius=grid_radius, xp=xp) return (self.tau**2.0 / (self.tau**2.0 + 1.0) ** 2.0) * ( ((self.tau**2.0 + 2.0 * grid_radius**2.0 - 1.0) * f_r) - + (jnp.pi * self.tau) - + ((self.tau**2.0 - 1.0) * jnp.log(self.tau)) + + (xp.pi * self.tau) + + ((self.tau**2.0 - 1.0) * xp.log(self.tau)) + ( - jnp.sqrt(grid_radius**2.0 + self.tau**2.0) - * (((self.tau**2.0 - 1.0) / self.tau) * k_r - jnp.pi) + xp.sqrt(grid_radius**2.0 + self.tau**2.0) + * (((self.tau**2.0 - 1.0) / self.tau) * k_r - xp.pi) ) ) @@ -109,6 +111,7 @@ def mass_at_truncation_radius_solar_mass( redshift_source, redshift_of_cosmic_average_density="profile", cosmology: LensingCosmology = None, + xp=np, ): from autogalaxy.cosmology.wrap import Planck15 @@ -119,14 +122,15 @@ def mass_at_truncation_radius_solar_mass( redshift_source=redshift_source, redshift_of_cosmic_average_density=redshift_of_cosmic_average_density, cosmology=cosmology, + xp=xp, ) return ( mass_at_200 * (self.tau**2.0 / (self.tau**2.0 + 1.0) ** 2.0) * ( - ((self.tau**2.0 - 1) * jnp.log(self.tau)) - + (self.tau * jnp.pi) + ((self.tau**2.0 - 1) * np.log(self.tau)) + + (self.tau * np.pi) - (self.tau**2.0 + 1) ) ) diff --git a/autogalaxy/profiles/mass/mock/mock_mass_profile.py b/autogalaxy/profiles/mass/mock/mock_mass_profile.py index c6e87d5be..26a2d36d1 100644 --- a/autogalaxy/profiles/mass/mock/mock_mass_profile.py +++ b/autogalaxy/profiles/mass/mock/mock_mass_profile.py @@ -1,3 +1,5 @@ +import numpy as np + import autogalaxy as ag @@ -19,11 +21,11 @@ def __init__( self.value = value self.value1 = value1 - def convergence_2d_from(self, grid): + def convergence_2d_from(self, grid, xp=np): return self.convergence_2d - def potential_2d_from(self, grid): + def potential_2d_from(self, grid, xp=np): return self.potential_2d - def deflections_yx_2d_from(self, grid): + def deflections_yx_2d_from(self, grid, xp=np): return self.deflections_2d diff --git a/autogalaxy/profiles/mass/point/point.py b/autogalaxy/profiles/mass/point/point.py index 581484e62..8ac85b651 100644 --- a/autogalaxy/profiles/mass/point/point.py +++ b/autogalaxy/profiles/mass/point/point.py @@ -24,7 +24,7 @@ def __init__( super().__init__(centre=centre, ell_comps=(0.0, 0.0)) self.einstein_radius = einstein_radius - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): squared_distances = np.square(grid[:, 0] - self.centre[0]) + np.square( grid[:, 1] - self.centre[1] ) @@ -35,15 +35,15 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): return convergence @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return np.zeros(shape=grid.shape[0]) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): grid_radii = self.radial_grid_from(grid=grid, **kwargs) return self._cartesian_grid_via_radial_from( - grid=grid, radius=self.einstein_radius**2 / grid_radii + grid=grid, radius=self.einstein_radius**2 / grid_radii, xp=xp ) @property diff --git a/autogalaxy/profiles/mass/point/smbh_binary.py b/autogalaxy/profiles/mass/point/smbh_binary.py index 3731dde7f..299ef7e8d 100644 --- a/autogalaxy/profiles/mass/point/smbh_binary.py +++ b/autogalaxy/profiles/mass/point/smbh_binary.py @@ -90,7 +90,7 @@ def angle_binary_radians(self) -> float: """ return self.angle_binary * np.pi / 180.0 - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -102,10 +102,10 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): The grid of (y,x) arc-second coordinates the convergence is computed on. """ return self.smbh_0.convergence_2d_from( - grid=grid - ) + self.smbh_1.convergence_2d_from(grid=grid, **kwargs) + grid=grid, xp=xp, **kwargs + ) + self.smbh_1.convergence_2d_from(grid=grid, xp=xp, **kwargs) - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected potential on a grid of (y,x) arc-second coordinates. @@ -120,7 +120,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid=grid, **kwargs ) + self.smbh_1.potential_2d_from(grid=grid) - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional deflection angles on a grid of (y,x) arc-second coordinates. diff --git a/autogalaxy/profiles/mass/sheets/external_shear.py b/autogalaxy/profiles/mass/sheets/external_shear.py index fb71e7bf0..f0c762958 100644 --- a/autogalaxy/profiles/mass/sheets/external_shear.py +++ b/autogalaxy/profiles/mass/sheets/external_shear.py @@ -1,4 +1,4 @@ -import jax.numpy as jnp +import numpy as np import autoarray as aa @@ -24,13 +24,15 @@ def __init__(self, gamma_1: float = 0.0, gamma_2: float = 0.0): self.gamma_1 = gamma_1 self.gamma_2 = gamma_2 - @property - def magnitude(self): - return convert.shear_magnitude_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2) + def magnitude(self, xp=np): + return convert.shear_magnitude_from( + gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp + ) - @property - def angle(self): - return convert.shear_angle_from(gamma_1=self.gamma_1, gamma_2=self.gamma_2) + def angle(self, xp=np): + return convert.shear_angle_from( + gamma_1=self.gamma_1, gamma_2=self.gamma_2, xp=xp + ) def convergence_func(self, grid_radius: float) -> float: return 0.0 @@ -39,24 +41,24 @@ def average_convergence_of_1_radius(self): return 0.0 @aa.grid_dec.to_array - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): shear_angle = ( - self.angle - 90 + self.angle(xp) - 90 ) ##to be onsistent with autolens deflection angle calculation - phig = jnp.deg2rad(shear_angle) - shear_amp = self.magnitude - phicoord = jnp.arctan2(grid.array[:, 0], grid.array[:, 1]) - rcoord = jnp.sqrt(grid.array[:, 0] ** 2.0 + grid.array[:, 1] ** 2.0) + phig = xp.deg2rad(shear_angle) + shear_amp = self.magnitude(xp=xp) + phicoord = xp.arctan2(grid.array[:, 0], grid.array[:, 1]) + rcoord = xp.sqrt(grid.array[:, 0] ** 2.0 + grid.array[:, 1] ** 2.0) - return -0.5 * shear_amp * rcoord**2 * jnp.cos(2 * (phicoord - phig)) + return -0.5 * shear_amp * rcoord**2 * xp.cos(2 * (phicoord - phig)) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -66,8 +68,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - deflection_y = -jnp.multiply(self.magnitude, grid.array[:, 0]) - deflection_x = jnp.multiply(self.magnitude, grid.array[:, 1]) + deflection_y = -xp.multiply(self.magnitude(xp=xp), grid.array[:, 0]) + deflection_x = xp.multiply(self.magnitude(xp=xp), grid.array[:, 1]) return self.rotated_grid_from_reference_frame_from( - jnp.vstack((deflection_y, deflection_x)).T + grid=xp.vstack((deflection_y, deflection_x)).T, + xp=xp, ) diff --git a/autogalaxy/profiles/mass/sheets/mass_sheet.py b/autogalaxy/profiles/mass/sheets/mass_sheet.py index b4c1414f4..d8bc6e278 100644 --- a/autogalaxy/profiles/mass/sheets/mass_sheet.py +++ b/autogalaxy/profiles/mass/sheets/mass_sheet.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -26,17 +25,17 @@ def convergence_func(self, grid_radius: float) -> float: return 0.0 @aa.grid_dec.to_array - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.full(shape=grid.shape[0], fill_value=self.kappa) + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.full(shape=grid.shape[0], fill_value=self.kappa) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): grid_radii = self.radial_grid_from(grid=grid, **kwargs) return self._cartesian_grid_via_radial_from( - grid=grid, radius=self.kappa * grid_radii + grid=grid, radius=self.kappa * grid_radii, xp=xp ) diff --git a/autogalaxy/profiles/mass/stellar/chameleon.py b/autogalaxy/profiles/mass/stellar/chameleon.py index 00a35a123..f2850e341 100644 --- a/autogalaxy/profiles/mass/stellar/chameleon.py +++ b/autogalaxy/profiles/mass/stellar/chameleon.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -48,12 +47,14 @@ def __init__( self.core_radius_0 = core_radius_0 self.core_radius_1 = core_radius_1 - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_analytic_from(grid=grid, **kwargs) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_analytic_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. Following Eq. (15) and (16), but the parameters are slightly different. @@ -69,64 +70,72 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): 2.0 * self.mass_to_light_ratio * self.intensity - / (1 + self.axis_ratio) - * self.axis_ratio - / jnp.sqrt(1.0 - self.axis_ratio**2.0) + / (1 + self.axis_ratio(xp)) + * self.axis_ratio(xp) + / xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0) ) - core_radius_0 = jnp.sqrt( - (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio) ** 2 + core_radius_0 = xp.sqrt( + (4.0 * self.core_radius_0**2.0) / (1.0 + self.axis_ratio(xp)) ** 2 ) - core_radius_1 = jnp.sqrt( - (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio) ** 2 + core_radius_1 = xp.sqrt( + (4.0 * self.core_radius_1**2.0) / (1.0 + self.axis_ratio(xp)) ** 2 ) psi0 = psi_from( - grid=grid, axis_ratio=self.axis_ratio, core_radius=core_radius_0 + grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=core_radius_0, xp=xp ) psi1 = psi_from( - grid=grid, axis_ratio=self.axis_ratio, core_radius=core_radius_1 + grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=core_radius_1, xp=xp ) - deflection_y0 = jnp.arctanh( - jnp.divide( - jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 0]), - jnp.add(psi0, self.axis_ratio**2.0 * core_radius_0), + deflection_y0 = xp.arctanh( + xp.divide( + xp.multiply( + xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 0] + ), + xp.add(psi0, self.axis_ratio(xp) ** 2.0 * core_radius_0), ) ) - deflection_x0 = jnp.arctan( - jnp.divide( - jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 1]), - jnp.add(psi0, core_radius_0), + deflection_x0 = xp.arctan( + xp.divide( + xp.multiply( + xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 1] + ), + xp.add(psi0, core_radius_0), ) ) - deflection_y1 = jnp.arctanh( - jnp.divide( - jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 0]), - jnp.add(psi1, self.axis_ratio**2.0 * core_radius_1), + deflection_y1 = xp.arctanh( + xp.divide( + xp.multiply( + xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 0] + ), + xp.add(psi1, self.axis_ratio(xp) ** 2.0 * core_radius_1), ) ) - deflection_x1 = jnp.arctan( - jnp.divide( - jnp.multiply(jnp.sqrt(1.0 - self.axis_ratio**2.0), grid.array[:, 1]), - jnp.add(psi1, core_radius_1), + deflection_x1 = xp.arctan( + xp.divide( + xp.multiply( + xp.sqrt(1.0 - self.axis_ratio(xp) ** 2.0), grid.array[:, 1] + ), + xp.add(psi1, core_radius_1), ) ) - deflection_y = jnp.subtract(deflection_y0, deflection_y1) - deflection_x = jnp.subtract(deflection_x0, deflection_x1) + deflection_y = xp.subtract(deflection_y0, deflection_y1) + deflection_x = xp.subtract(deflection_x0, deflection_x1) return self.rotated_grid_from_reference_frame_from( - jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T) + xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T), xp=xp ) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters ---------- @@ -141,10 +150,10 @@ def convergence_func(self, grid_radius: float) -> float: return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) - def image_2d_via_radii_from(self, grid_radii: np.ndarray): + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): """Calculate the intensity of the Chamelon light profile on a grid of radial coordinates. Parameters @@ -153,25 +162,25 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray): The radial distance from the centre of the profile. for each coordinate on the grid. """ - axis_ratio_factor = (1.0 + self.axis_ratio) ** 2.0 + axis_ratio_factor = (1.0 + self.axis_ratio(xp)) ** 2.0 - return jnp.multiply( - self.intensity / (1 + self.axis_ratio), - jnp.add( - jnp.divide( + return xp.multiply( + self.intensity / (1 + self.axis_ratio(xp)), + xp.add( + xp.divide( 1.0, - jnp.sqrt( - jnp.add( - jnp.square(grid_radii.array), + xp.sqrt( + xp.add( + xp.square(grid_radii.array), (4.0 * self.core_radius_0**2.0) / axis_ratio_factor, ) ), ), - -jnp.divide( + -xp.divide( 1.0, - jnp.sqrt( - jnp.add( - jnp.square(grid_radii.array), + xp.sqrt( + xp.add( + xp.square(grid_radii.array), (4.0 * self.core_radius_1**2.0) / axis_ratio_factor, ) ), @@ -179,9 +188,8 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray): ), ) - @property - def axis_ratio(self): - axis_ratio = super().axis_ratio + def axis_ratio(self, xp=np): + axis_ratio = super().axis_ratio(xp=xp) return axis_ratio if axis_ratio < 0.99999 else 0.99999 diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py index d9e5efc89..064fbb0fe 100644 --- a/autogalaxy/profiles/mass/stellar/gaussian.py +++ b/autogalaxy/profiles/mass/stellar/gaussian.py @@ -1,8 +1,4 @@ import numpy as np -from autoconf.jax_wrapper import use_jax - -if use_jax: - import jax from typing import Tuple @@ -43,7 +39,7 @@ def __init__( self.intensity = intensity self.sigma = sigma - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -61,7 +57,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_analytic_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -76,19 +74,22 @@ def deflections_2d_via_analytic_from(self, grid: aa.type.Grid2DLike, **kwargs): self.mass_to_light_ratio * self.intensity * self.sigma - * np.sqrt((2 * np.pi) / (1.0 - self.axis_ratio**2.0)) - * self.zeta_from(grid=grid) + * xp.sqrt((2 * np.pi) / (1.0 - self.axis_ratio(xp) ** 2.0)) + * self.zeta_from(grid=grid, xp=xp) ) return self.rotated_grid_from_reference_frame_from( - np.multiply( - 1.0, np.vstack((-1.0 * np.imag(deflections), np.real(deflections))).T - ) + xp.multiply( + 1.0, xp.vstack((-1.0 * xp.imag(deflections), xp.real(deflections))).T + ), + xp=xp, ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_integral_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -103,7 +104,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): from scipy.integrate import quad def calculate_deflection_component(npow, index): - deflection_grid = np.array(self.axis_ratio * grid.array[:, index]) + deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index]) for i in range(grid.shape[0]): deflection_grid[i] *= ( @@ -117,8 +118,8 @@ def calculate_deflection_component(npow, index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, - self.sigma / np.sqrt(self.axis_ratio), + self.axis_ratio(xp), + self.sigma / xp.sqrt(self.axis_ratio(xp)), ), )[0] ) @@ -129,7 +130,7 @@ def calculate_deflection_component(npow, index): deflection_x = calculate_deflection_component(0.0, 1) return self.rotated_grid_from_reference_frame_from( - np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T) + np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp ) @staticmethod @@ -145,7 +146,7 @@ def deflection_func(u, y, x, npow, axis_ratio, sigma): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -155,17 +156,17 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ return self.convergence_func( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) ) def convergence_func(self, grid_radius: float) -> float: return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return np.zeros(shape=grid.shape[0]) - def image_2d_via_radii_from(self, grid_radii: np.ndarray): + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): """Calculate the intensity of the Gaussian light profile on a grid of radial coordinates. Parameters @@ -180,28 +181,26 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray): np.exp( -0.5 * np.square( - np.divide(grid_radii.array, self.sigma / np.sqrt(self.axis_ratio)) + np.divide( + grid_radii.array, self.sigma / np.sqrt(self.axis_ratio(xp)) + ) ) ), ) - @property - def axis_ratio(self): - axis_ratio = super().axis_ratio - if use_jax: - return jax.lax.select(axis_ratio < 0.9999, axis_ratio, 0.9999) - else: - return axis_ratio if axis_ratio < 0.9999 else 0.9999 + def axis_ratio(self, xp=np): + axis_ratio = super().axis_ratio(xp=xp) + return xp.where(axis_ratio < 0.9999, axis_ratio, 0.9999) - def zeta_from(self, grid: aa.type.Grid2DLike): + def zeta_from(self, grid: aa.type.Grid2DLike, xp=np): from scipy.special import wofz - q2 = self.axis_ratio**2.0 + q2 = self.axis_ratio(xp) ** 2.0 ind_pos_y = grid.array[:, 0] >= 0 shape_grid = np.shape(grid) output_grid = np.zeros((shape_grid[0]), dtype=np.complex128) - scale_factor = self.axis_ratio / (self.sigma * np.sqrt(2.0 * (1.0 - q2))) + scale_factor = self.axis_ratio(xp) / (self.sigma * np.sqrt(2.0 * (1.0 - q2))) xs_0 = grid.array[:, 1][ind_pos_y] * scale_factor ys_0 = grid.array[:, 0][ind_pos_y] * scale_factor @@ -211,7 +210,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike): output_grid[ind_pos_y] = -1j * ( wofz(xs_0 + 1j * ys_0) - np.exp(-(xs_0**2.0) * (1.0 - q2) - ys_0 * ys_0 * (1.0 / q2 - 1.0)) - * wofz(self.axis_ratio * xs_0 + 1j * ys_0 / self.axis_ratio) + * wofz(self.axis_ratio(xp) * xs_0 + 1j * ys_0 / self.axis_ratio(xp)) ) output_grid[~ind_pos_y] = np.conj( @@ -219,7 +218,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike): * ( wofz(xs_1 + 1j * ys_1) - np.exp(-(xs_1**2.0) * (1.0 - q2) - ys_1 * ys_1 * (1.0 / q2 - 1.0)) - * wofz(self.axis_ratio * xs_1 + 1j * ys_1 / self.axis_ratio) + * wofz(self.axis_ratio(xp) * xs_1 + 1j * ys_1 / self.axis_ratio(xp)) ) ) diff --git a/autogalaxy/profiles/mass/stellar/sersic.py b/autogalaxy/profiles/mass/stellar/sersic.py index da2f54e73..d1bda8094 100644 --- a/autogalaxy/profiles/mass/stellar/sersic.py +++ b/autogalaxy/profiles/mass/stellar/sersic.py @@ -1,5 +1,3 @@ -import jax -import jax.numpy as jnp import numpy as np from typing import List, Tuple @@ -122,7 +120,7 @@ def __init__( self.effective_radius = effective_radius self.sersic_index = sersic_index - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_cse_from(grid=grid, **kwargs) @aa.grid_dec.to_vector_yx @@ -145,14 +143,14 @@ def deflections_2d_via_mge_from( """ return self._deflections_2d_via_mge_from( grid=grid, - sigmas_factor=np.sqrt(self.axis_ratio), + sigmas_factor=np.sqrt(self.axis_ratio()), func_terms=func_terms, func_gaussians=func_gaussians, ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the projected 2D deflection angles from a grid of (y,x) arc second coordinates, by computing and summing the convergence of each individual cse used to decompose the mass profile. @@ -171,7 +169,7 @@ def deflections_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -181,14 +179,19 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ return self.convergence_func( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) ) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform def convergence_2d_via_mge_from( - self, grid: aa.type.Grid2DLike, func_terms=28, func_gaussians=20, **kwargs + self, + grid: aa.type.Grid2DLike, + xp=np, + func_terms=28, + func_gaussians=20, + **kwargs, ): """ Calculate the projected convergence at a given set of arc-second gridded coordinates. @@ -200,7 +203,7 @@ def convergence_2d_via_mge_from( """ - eccentric_radii = self.eccentric_radii_grid_from(grid=grid, **kwargs) + eccentric_radii = self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) return self._convergence_2d_via_mge_from( grid_radii=eccentric_radii, @@ -211,7 +214,7 @@ def convergence_2d_via_mge_from( @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_via_cse_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the projected 2D convergence from a grid of (y,x) arc second coordinates, by computing and summing the convergence of each individual cse used to decompose the mass profile. @@ -234,7 +237,7 @@ def convergence_func(self, grid_radius: float) -> float: return self.mass_to_light_ratio * self.image_2d_via_radii_from(grid_radius) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return np.zeros(shape=grid.shape[0]) def image_2d_via_radii_from(self, radius: np.ndarray): @@ -313,7 +316,7 @@ def decompose_convergence_via_cse( mass_to_light_gradient=0.0, ) - scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio) + scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio()) radii_min = scaled_effective_radius / 10.0**lower_dex radii_max = scaled_effective_radius * 10.0**upper_dex @@ -354,7 +357,7 @@ def sersic_constant(self): @property def ellipticity_rescale(self): - return 1.0 - ((1.0 - self.axis_ratio) / 2.0) + return 1.0 - ((1.0 - self.axis_ratio()) / 2.0) @property def elliptical_effective_radius(self): @@ -366,13 +369,15 @@ def elliptical_effective_radius(self): The elliptical effective radius instead describes the major-axis radius of the ellipse containing \ half the light, and may be more appropriate for highly flattened systems like disk galaxies. """ - return self.effective_radius / np.sqrt(self.axis_ratio) + return self.effective_radius / np.sqrt(self.axis_ratio()) class Sersic(AbstractSersic, MassProfileMGE, MassProfileCSE): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_integral_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -387,10 +392,10 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): def calculate_deflection_component(npow, index): sersic_constant = self.sersic_constant - deflection_grid = self.axis_ratio * grid.array[:, index] + deflection_grid = self.axis_ratio() * grid.array[:, index] for i in range(grid.shape[0]): - deflection_grid = deflection_grid.at[i].multiply( + deflection_grid[i] = deflection_grid[i] * ( self.intensity * self.mass_to_light_ratio * quad( @@ -401,7 +406,7 @@ def calculate_deflection_component(npow, index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, + self.axis_ratio(), self.sersic_index, self.effective_radius, sersic_constant, @@ -415,7 +420,7 @@ def calculate_deflection_component(npow, index): deflection_x = calculate_deflection_component(0.0, 1) return self.rotated_grid_from_reference_frame_from( - np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T) + np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp ) @staticmethod diff --git a/autogalaxy/profiles/mass/stellar/sersic_core.py b/autogalaxy/profiles/mass/stellar/sersic_core.py index 496e4f7f4..415c81d18 100644 --- a/autogalaxy/profiles/mass/stellar/sersic_core.py +++ b/autogalaxy/profiles/mass/stellar/sersic_core.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -60,10 +59,10 @@ def __init__( self.alpha = alpha self.gamma = gamma - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): return self.deflections_2d_via_mge_from(grid=grid, **kwargs) - def image_2d_via_radii_from(self, grid_radii: np.ndarray): + def image_2d_via_radii_from(self, grid_radii: np.ndarray, xp=np): """ Calculate the intensity of the cored-Sersic light profile on a grid of radial coordinates. @@ -72,27 +71,27 @@ def image_2d_via_radii_from(self, grid_radii: np.ndarray): grid_radii The radial distance from the centre of the profile. for each coordinate on the grid. """ - return jnp.multiply( - jnp.multiply( - self.intensity_prime, - jnp.power( - jnp.add( + return xp.multiply( + xp.multiply( + self.intensity_prime(xp), + xp.power( + xp.add( 1, - jnp.power( - jnp.divide(self.radius_break, grid_radii.array), self.alpha + xp.power( + xp.divide(self.radius_break, grid_radii.array), self.alpha ), ), (self.gamma / self.alpha), ), ), - jnp.exp( - jnp.multiply( + xp.exp( + xp.multiply( -self.sersic_constant, ( - jnp.power( - jnp.divide( - jnp.add( - jnp.power(grid_radii.array, self.alpha), + xp.power( + xp.divide( + xp.add( + xp.power(grid_radii.array, self.alpha), (self.radius_break**self.alpha), ), (self.effective_radius**self.alpha), @@ -111,7 +110,7 @@ def decompose_convergence_via_mge(self): def core_sersic_2D(r): return ( self.mass_to_light_ratio - * self.intensity_prime + * self.intensity_prime() * (1.0 + (self.radius_break / r) ** self.alpha) ** (self.gamma / self.alpha) * np.exp( @@ -128,13 +127,12 @@ def core_sersic_2D(r): func=core_sersic_2D, radii_min=radii_min, radii_max=radii_max ) - @property - def intensity_prime(self): + def intensity_prime(self, xp=np): """Overall intensity normalisation in the rescaled Core-Sersic light profiles (electrons per second)""" return ( self.intensity * (2.0 ** (-self.gamma / self.alpha)) - * jnp.exp( + * xp.exp( self.sersic_constant * ( ((2.0 ** (1.0 / self.alpha)) * self.radius_break) diff --git a/autogalaxy/profiles/mass/stellar/sersic_gradient.py b/autogalaxy/profiles/mass/stellar/sersic_gradient.py index 92c07fd0d..3b096d9a1 100644 --- a/autogalaxy/profiles/mass/stellar/sersic_gradient.py +++ b/autogalaxy/profiles/mass/stellar/sersic_gradient.py @@ -51,7 +51,9 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_2d_via_integral_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -66,7 +68,7 @@ def deflections_2d_via_integral_from(self, grid: aa.type.Grid2DLike, **kwargs): def calculate_deflection_component(npow, index): sersic_constant = self.sersic_constant - deflection_grid = np.array(self.axis_ratio * grid.array[:, index]) + deflection_grid = np.array(self.axis_ratio() * grid.array[:, index]) for i in range(grid.shape[0]): deflection_grid[i] *= ( @@ -80,7 +82,7 @@ def calculate_deflection_component(npow, index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, + self.axis_ratio(), self.sersic_index, self.effective_radius, self.mass_to_light_gradient, @@ -125,7 +127,7 @@ def deflection_func( @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """Calculate the projected convergence at a given set of arc-second gridded coordinates. Parameters @@ -135,14 +137,14 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ return self.convergence_func( - self.eccentric_radii_grid_from(grid=grid, **kwargs) + self.eccentric_radii_grid_from(grid=grid, xp=xp, **kwargs) ) def convergence_func(self, grid_radius: float) -> float: return ( self.mass_to_light_ratio * ( - ((self.axis_ratio * grid_radius) / self.effective_radius) + ((self.axis_ratio() * grid_radius) / self.effective_radius) ** -self.mass_to_light_gradient ) * self.image_2d_via_radii_from(grid_radius) @@ -157,7 +159,7 @@ def sersic_gradient_2D(r): self.mass_to_light_ratio * self.intensity * ( - ((self.axis_ratio * r) / self.effective_radius) + ((self.axis_ratio() * r) / self.effective_radius) ** -self.mass_to_light_gradient ) * np.exp( @@ -205,7 +207,7 @@ def decompose_convergence_via_cse( mass_to_light_gradient=self.mass_to_light_gradient, ) - scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio) + scaled_effective_radius = self.effective_radius / np.sqrt(self.axis_ratio()) radii_min = scaled_effective_radius / 10.0**lower_dex radii_max = scaled_effective_radius * 10.0**upper_dex @@ -214,7 +216,7 @@ def sersic_gradient_2D(r): self.mass_to_light_ratio * self.intensity * ( - ((self.axis_ratio * r) / scaled_effective_radius) + ((self.axis_ratio() * r) / scaled_effective_radius) ** -self.mass_to_light_gradient ) * np.exp( diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py index 42a44df1f..3e9fefedb 100644 --- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py +++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_mass.py @@ -1,5 +1,4 @@ from typing import Tuple -import jax.numpy as jnp import numpy as np import autoarray as aa @@ -9,7 +8,7 @@ # The dPIEPotential and dPIEPotentialSph profiles are modified from the original `dPIEPotential` and `dPIEPotentialSph`, which were implemented to PyAutoLens by Jackson O'Donnell. -def _ci05(x, y, eps, rcore): +def _ci05(x, y, eps, rcore, xp=np): """ Returns the first derivatives of the lens potential as complex number I'* = (∂ψ/∂x + i ∂ψ/∂y) / E0 for PIEMass at any positions (x,y), see Kassiola & Kovner(1993) Eq. 4.1.2, which is the integral of Eq. 2.3.8. @@ -26,7 +25,7 @@ def _ci05(x, y, eps, rcore): complex The value of the I'* term. """ - sqe = jnp.sqrt(eps) + sqe = xp.sqrt(eps) axis_ratio = (1.0 - eps) / (1.0 + eps) cxro = (1.0 + eps) * (1.0 + eps) cyro = (1.0 - eps) * (1.0 - eps) @@ -34,24 +33,24 @@ def _ci05(x, y, eps, rcore): ##### I'* = zres = zci * ln(zis) = zci * ln(znum / zden), see Eq. 4.1.2 ##### # Define intermediate complex variables - zci = jnp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=jnp.complex128) - znum = jnp.complex128( + zci = xp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=xp.complex128) + znum = xp.complex128( axis_ratio * x - + 1j * (2.0 * sqe * jnp.sqrt(rcore * rcore + rem2) - y / axis_ratio) + + 1j * (2.0 * sqe * xp.sqrt(rcore * rcore + rem2) - y / axis_ratio) ) - zden = jnp.complex128(x + 1j * (2.0 * rcore * sqe - y)) + zden = xp.complex128(x + 1j * (2.0 * rcore * sqe - y)) # zis = znum / zden = (a+bi)/(c+di) = [(ac+bd)+(bc-ad i)] / (c^2+d^2) norm = zden.real * zden.real + zden.imag * zden.imag # |zden|^2 zis_re = (znum.real * zden.real + znum.imag * zden.imag) / norm zis_im = (znum.imag * zden.real - znum.real * zden.imag) / norm - zis = jnp.complex128(zis_re + 1j * zis_im) + zis = xp.complex128(zis_re + 1j * zis_im) # ln(zis) = ln(|zis|) + i*Arg(zis) - zis_mag = jnp.abs(zis) - zis_re = jnp.log(zis_mag) - zis_im = jnp.angle(zis) - zis = jnp.complex128(zis_re + 1j * zis_im) + zis_mag = xp.abs(zis) + zis_re = xp.log(zis_mag) + zis_im = xp.angle(zis) + zis = xp.complex128(zis_re + 1j * zis_im) # I'* = zres = zci * ln(zis) zres = zci * zis @@ -59,7 +58,7 @@ def _ci05(x, y, eps, rcore): return zres -def _ci05f(x, y, eps, rcore, rcut): +def _ci05f(x, y, eps, rcore, rcut, xp=np): """ Returns the first derivatives of the lens potential as complex number I'* = (∂ψ/∂x + i ∂ψ/∂y) / (b0 * ra / (rs - ra)) for dPIEMass at any positions (x,y), which is the integral of Eq. 2.3.8 in Kassiola & Kovner(1993). @@ -81,7 +80,7 @@ def _ci05f(x, y, eps, rcore, rcut): complex The value of the I'* term. """ - sqe = jnp.sqrt(eps) + sqe = xp.sqrt(eps) axis_ratio = (1.0 - eps) / (1.0 + eps) cxro = (1.0 + eps) * (1.0 + eps) cyro = (1.0 - eps) * (1.0 - eps) @@ -90,17 +89,16 @@ def _ci05f(x, y, eps, rcore, rcut): ##### I'* = zres_rc - zres_rcut = zci * ln(zis_rc / zis_rcut) = zci * ln((znum_rc / zden_rc) / (znum_rcut / zden_rcut)) ##### # Define intermediate complex variables - zci = jnp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=jnp.complex128) - znum_rc = jnp.complex128( + zci = xp.array(0.0 + 1j * (-0.5 * (1.0 - eps * eps) / sqe), dtype=xp.complex128) + znum_rc = xp.complex128( axis_ratio * x - + 1j * (2.0 * sqe * jnp.sqrt(rcore * rcore + rem2) - y / axis_ratio) + + 1j * (2.0 * sqe * xp.sqrt(rcore * rcore + rem2) - y / axis_ratio) ) # a + bi - zden_rc = jnp.complex128(x + 1j * (2.0 * rcore * sqe - y)) # c + di - znum_rcut = jnp.complex128( - axis_ratio * x - + 1j * (2.0 * sqe * jnp.sqrt(rcut * rcut + rem2) - y / axis_ratio) + zden_rc = xp.complex128(x + 1j * (2.0 * rcore * sqe - y)) # c + di + znum_rcut = xp.complex128( + axis_ratio * x + 1j * (2.0 * sqe * xp.sqrt(rcut * rcut + rem2) - y / axis_ratio) ) # a + ei - zden_rcut = jnp.complex128(x + 1j * (2.0 * rcut * sqe - y)) # c + fi + zden_rcut = xp.complex128(x + 1j * (2.0 * rcut * sqe - y)) # c + fi # zis_rc = znum_rc / zden_rc = (a+bi)/(c+di) # zis_rcut = znum_rcut / zden_rcut = (a+ei)/(c+fi) @@ -117,13 +115,13 @@ def _ci05f(x, y, eps, rcore, rcut): norm = cc * cc + dd * dd aaa = (aa * cc + bb * dd) / norm bbb = (bb * cc - aa * dd) / norm - zis_tot = jnp.complex128(aaa + 1j * bbb) + zis_tot = xp.complex128(aaa + 1j * bbb) # ln(zis_tot) = ln(|zis_tot|) + i*Arg(zis_tot) - zis_tot_mag = jnp.abs(zis_tot) - zr_re = jnp.log(zis_tot_mag) - zr_im = jnp.angle(zis_tot) - zr = jnp.complex128(zr_re + 1j * zr_im) + zis_tot_mag = xp.abs(zis_tot) + zr_re = xp.log(zis_tot_mag) + zr_im = xp.angle(zis_tot) + zr = xp.complex128(zr_re + 1j * zr_im) # I'* = zci * ln(zis_tot) zres = zci * zr @@ -131,7 +129,7 @@ def _ci05f(x, y, eps, rcore, rcut): return zres -def _mdci05(x, y, eps, rcore, b0): +def _mdci05(x, y, eps, rcore, b0, xp=np): """ Returns the second derivatives (Hessian matrix) of the lens potential as complex number for PIEMass at any positions (x,y): ∂²ψ/∂x² = Re(∂I*/∂x), ∂²ψ/∂y² = Im(∂I*/∂y), ∂²ψ/∂x∂y = ∂²ψ/∂y∂x = Im(∂I*/∂x) = Re(∂I*/∂y) @@ -153,13 +151,13 @@ def _mdci05(x, y, eps, rcore, b0): # I*(x,y) = b0 * ci * (-i) * (ln{ q * x + (2.0 * sqe * wrem - y * 1/q )*i} - ln{ x + (2.0 * rcore * sqe - y)*i}) # = b0 * ci * (-i) * (ln{ q * x + num1*i} - ln{ x + num2*i}) # = b0 * ci * (-i) * (ln{u(x,y)} - ln{v(x,y)}) - sqe = jnp.sqrt(eps) + sqe = xp.sqrt(eps) axis_ratio = (1.0 - eps) / (1.0 + eps) axis_ratio_inv = 1.0 / axis_ratio cxro = (1.0 + eps) * (1.0 + eps) cyro = (1.0 - eps) * (1.0 - eps) ci = 0.5 * (1.0 - eps * eps) / sqe - wrem = jnp.sqrt(rcore * rcore + x * x / cxro + y * y / cyro) # √(w(x,y)) + wrem = xp.sqrt(rcore * rcore + x * x / cxro + y * y / cyro) # √(w(x,y)) num1 = 2.0 * sqe * wrem - y * axis_ratio_inv den1 = axis_ratio * axis_ratio * x * x + num1 * num1 # |q * x + num1*i|^2 num2 = 2.0 * rcore * sqe - y @@ -259,14 +257,14 @@ def __init__( self.ra = ra self.b0 = b0 - def _ellip(self): - ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) + def _ellip(self, xp=np): + ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) MAX_ELLIP = 0.99999 - return jnp.min(jnp.array([ellip, MAX_ELLIP])) + return xp.min(xp.array([ellip, MAX_ELLIP])) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -275,7 +273,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - ellip = self._ellip() + ellip = self._ellip(xp) factor = self.b0 zis = _ci05(x=grid.array[:, 1], y=grid.array[:, 0], eps=ellip, rcore=self.ra) @@ -285,12 +283,13 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # And here we convert back to the real axes return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T), + grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T), + xp=xp, **kwargs, ) @aa.grid_dec.transform - def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): + def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs): """ Calculate the hessian matrix on a grid of (y,x) arc-second coordinates. @@ -299,7 +298,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - grid = jnp.asarray(grid) + grid = xp.asarray(grid) if grid.ndim != 2 or grid.shape[1] != 2: raise ValueError("Grid must be a 2D array with shape (n, 2)") ellip = self._ellip() @@ -310,10 +309,12 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): return hessian_yy, hessian_xy, hessian_yx, hessian_xx - def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): + def analytical_magnification_2d_from( + self, grid: "aa.type.Grid2DLike", xp=np, **kwargs + ): hessian_yy, hessian_xy, hessian_yx, hessian_xx = ( - self.analytical_hessian_2d_from(grid=grid) + self.analytical_hessian_2d_from(grid=grid, xp=np) ) det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx @@ -372,14 +373,14 @@ def __init__( self.rs = rs self.b0 = b0 - def _ellip(self): - ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) + def _ellip(self, xp=np): + ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) MAX_ELLIP = 0.99999 - return jnp.min(jnp.array([ellip, MAX_ELLIP])) + return xp.min(xp.array([ellip, MAX_ELLIP])) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -388,7 +389,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - ellip = self._ellip() + ellip = self._ellip(xp) factor = self.b0 * self.rs / (self.rs - self.ra) zis = _ci05f( x=grid.array[:, 1], @@ -404,18 +405,19 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # And here we convert back to the real axes return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T), + grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T), + xp=xp, **kwargs, ) - def _deflection_angle(self, radii): + def _deflection_angle(self, radii, xp=np): """ For a circularly symmetric dPIEPotential profile, computes the magnitude of the deflection at each radius. """ a, s = self.ra, self.rs - radii = jnp.maximum(radii, 1e-8) - f = radii / (a + jnp.sqrt(a**2 + radii**2)) - radii / ( - s + jnp.sqrt(s**2 + radii**2) + radii = xp.maximum(radii, 1e-8) + f = radii / (a + xp.sqrt(a**2 + radii**2)) - radii / ( + s + xp.sqrt(s**2 + radii**2) ) # c.f. Eliasdottir '07 eq. A23 @@ -424,7 +426,7 @@ def _deflection_angle(self, radii): alpha = self.b0 * s / (s - a) * f return alpha - def _convergence(self, radii): + def _convergence(self, radii, xp=np): radsq = radii * radii a, s = self.ra, self.rs @@ -434,12 +436,12 @@ def _convergence(self, radii): / 2 * s / (s - a) - * (1 / jnp.sqrt(a**2 + radsq) - 1 / jnp.sqrt(s**2 + radsq)) + * (1 / xp.sqrt(a**2 + radsq) - 1 / xp.sqrt(s**2 + radsq)) ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -451,14 +453,14 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the convergence is computed on. """ - ellip = self._ellip() - grid_radii = jnp.sqrt( + ellip = self._ellip(xp) + grid_radii = xp.sqrt( grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip) ) # Compute the convergence and deflection of a *circular* profile - kappa_circ = self._convergence(grid_radii) - alpha_circ = self._deflection_angle(grid_radii) + kappa_circ = self._convergence(grid_radii, xp) + alpha_circ = self._deflection_angle(grid_radii, xp) asymm_term = ( ellip * (1 - ellip) * grid.array[:, 1] ** 2 @@ -471,7 +473,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): return kappa_circ * (1 - asymm_term) + (alpha_circ / grid_radii) * asymm_term @aa.grid_dec.transform - def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): + def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs): """ Calculate the hessian matrix on a grid of (y,x) arc-second coordinates. Hessian_dPIEMass = rs * (rs - ra) * ( Hessian_PIEMass(ra) - Hessian_PIEMass(rs)) @@ -481,7 +483,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - grid = jnp.asarray(grid) + grid = xp.asarray(grid) if grid.ndim != 2 or grid.shape[1] != 2: raise ValueError("Grid must be a 2D array with shape (n, 2)") ellip = self._ellip() @@ -502,10 +504,12 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): return hessian_yy, hessian_xy, hessian_yx, hessian_xx - def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): + def analytical_magnification_2d_from( + self, grid: "aa.type.Grid2DLike", xp=np, **kwargs + ): hessian_yy, hessian_xy, hessian_yx, hessian_xx = ( - self.analytical_hessian_2d_from(grid=grid) + self.analytical_hessian_2d_from(grid=grid, xp=xp) ) det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx @@ -513,8 +517,8 @@ def analytical_magnification_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs) return aa.Array2D(values=1.0 / det_A, mask=grid.mask) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) class dPIEMassSph(dPIEMass): @@ -568,7 +572,7 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. Faster and equivalent to Eliasdottir (2007), see Eq. A19 and Eq. A20. @@ -593,7 +597,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # radii = self.radial_grid_from(grid=grid, **kwargs) # R2 = radii * radii R2 = grid.array[:, 1] * grid.array[:, 1] + grid.array[:, 0] * grid.array[:, 0] - factor = jnp.sqrt(R2 + a * a) - a - jnp.sqrt(R2 + s * s) + s + factor = xp.sqrt(R2 + a * a) - a - xp.sqrt(R2 + s * s) + s factor *= self.b0 * s / (s - a) / R2 # This is in axes aligned to the major/minor axis @@ -602,11 +606,13 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # And here we convert back to the real axes return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply(1.0, jnp.vstack((deflection_y, deflection_x)).T), **kwargs + grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), + xp=xp, + **kwargs, ) @aa.grid_dec.transform - def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): + def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", xp=np, **kwargs): """ Calculate the hessian matrix on a grid of (y,x) arc-second coordinates. Chain rule of second derivatives: @@ -620,7 +626,7 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - grid = jnp.asarray(grid) + grid = xp.asarray(grid) if grid.ndim != 2 or grid.shape[1] != 2: raise ValueError("Grid must be a 2D array with shape (n, 2)") @@ -646,9 +652,9 @@ def analytical_hessian_2d_from(self, grid: "aa.type.Grid2DLike", **kwargs): # = {( R^2 / √(R^2 + a^2)) - ( R^2 / √(R^2 + s^2)) - z} / R^2 # = p R2 = grid.array[:, 1] * grid.array[:, 1] + grid.array[:, 0] * grid.array[:, 0] - z = jnp.sqrt(R2 + a * a) - a - jnp.sqrt(R2 + s * s) + s - p = (1.0 - a / jnp.sqrt(a * a + R2)) * a / R2 - ( - 1.0 - s / jnp.sqrt(s * s + R2) + z = xp.sqrt(R2 + a * a) - a - xp.sqrt(R2 + s * s) + s + p = (1.0 - a / xp.sqrt(a * a + R2)) * a / R2 - ( + 1.0 - s / xp.sqrt(s * s + R2) ) * s / R2 X = grid.array[:, 1] * grid.array[:, 1] / R2 # x^2 / R^2 Y = grid.array[:, 0] * grid.array[:, 0] / R2 # y^2 / R^2 diff --git a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py index d8d02ce52..ccc9662fe 100644 --- a/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py +++ b/autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py @@ -1,5 +1,5 @@ from typing import Tuple -import jax.numpy as jnp +import numpy as np import autoarray as aa from autogalaxy.profiles.mass.abstract.abstract import MassProfile @@ -59,19 +59,19 @@ def __init__( self.rs = rs self.b0 = b0 - def _ellip(self): - ellip = jnp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) + def _ellip(self, xp=np): + ellip = xp.sqrt(self.ell_comps[0] ** 2 + self.ell_comps[1] ** 2) MAX_ELLIP = 0.99999 - return jnp.min(jnp.array([ellip, MAX_ELLIP])) + return xp.min(xp.array([ellip, MAX_ELLIP])) - def _deflection_angle(self, radii): + def _deflection_angle(self, radii, xp=np): """ For a circularly symmetric dPIEPotential profile, computes the magnitude of the deflection at each radius. """ a, s = self.ra, self.rs - radii = jnp.maximum(radii, 1e-8) - f = radii / (a + jnp.sqrt(a**2 + radii**2)) - radii / ( - s + jnp.sqrt(s**2 + radii**2) + radii = xp.maximum(radii, 1e-8) + f = radii / (a + xp.sqrt(a**2 + radii**2)) - radii / ( + s + xp.sqrt(s**2 + radii**2) ) # c.f. Eliasdottir '07 eq. A23 @@ -80,7 +80,7 @@ def _deflection_angle(self, radii): alpha = self.b0 * s / (s - a) * f return alpha - def _convergence(self, radii): + def _convergence(self, radii, xp=np): radsq = radii * radii a, s = self.ra, self.rs @@ -90,12 +90,12 @@ def _convergence(self, radii): / 2 * s / (s - a) - * (1 / jnp.sqrt(a**2 + radsq) - 1 / jnp.sqrt(s**2 + radsq)) + * (1 / xp.sqrt(a**2 + radsq) - 1 / xp.sqrt(s**2 + radsq)) ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -104,30 +104,28 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - ellip = self._ellip() - grid_radii = jnp.sqrt( + ellip = self._ellip(xp) + grid_radii = xp.sqrt( grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip) ) # Compute the deflection magnitude of a *non-elliptical* profile - alpha_circ = self._deflection_angle(grid_radii) + alpha_circ = self._deflection_angle(grid_radii, xp) # This is in axes aligned to the major/minor axis - deflection_y = ( - alpha_circ * jnp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii) - ) - deflection_x = ( - alpha_circ * jnp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii) - ) + deflection_y = alpha_circ * xp.sqrt(1 + ellip) * (grid.array[:, 0] / grid_radii) + deflection_x = alpha_circ * xp.sqrt(1 - ellip) * (grid.array[:, 1] / grid_radii) # And here we convert back to the real axes return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply(1.0, jnp.vstack((deflection_y, deflection_x)).T), **kwargs + grid=xp.multiply(1.0, xp.vstack((deflection_y, deflection_x)).T), + xp=xp, + **kwargs, ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -139,14 +137,14 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the convergence is computed on. """ - ellip = self._ellip() - grid_radii = jnp.sqrt( + ellip = self._ellip(xp) + grid_radii = xp.sqrt( grid.array[:, 1] ** 2 * (1 - ellip) + grid.array[:, 0] ** 2 * (1 + ellip) ) # Compute the convergence and deflection of a *circular* profile - kappa_circ = self._convergence(grid_radii) - alpha_circ = self._deflection_angle(grid_radii) + kappa_circ = self._convergence(grid_radii, xp) + alpha_circ = self._deflection_angle(grid_radii, xp) asymm_term = ( ellip * (1 - ellip) * grid.array[:, 1] ** 2 @@ -159,8 +157,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): return kappa_circ * (1 - asymm_term) + (alpha_circ / grid_radii) * asymm_term @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) class dPIEPotentialSph(dPIEPotential): @@ -218,7 +216,7 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -229,7 +227,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ radii = self.radial_grid_from(grid=grid, **kwargs) - alpha = self._deflection_angle(radii.array) + alpha = self._deflection_angle(radii.array, xp) # now we decompose the deflection into y/x components defl_y = alpha * grid.array[:, 0] / radii.array @@ -239,7 +237,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -254,8 +252,8 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): # already transformed to center on profile centre so this works radsq = grid.array[:, 0] ** 2 + grid.array[:, 1] ** 2 - return self._convergence(jnp.sqrt(radsq)) + return self._convergence(xp.sqrt(radsq), xp) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) diff --git a/autogalaxy/profiles/mass/total/isothermal.py b/autogalaxy/profiles/mass/total/isothermal.py index 2b0cd5004..2d941affc 100644 --- a/autogalaxy/profiles/mass/total/isothermal.py +++ b/autogalaxy/profiles/mass/total/isothermal.py @@ -1,4 +1,4 @@ -import jax.numpy as jnp +import numpy as np from typing import Tuple @@ -7,7 +7,7 @@ from autogalaxy.profiles.mass.total.power_law import PowerLaw -def psi_from(grid, axis_ratio, core_radius): +def psi_from(grid, axis_ratio, core_radius, xp=np): """ Returns the $\Psi$ term in expressions for the calculation of the deflection of an elliptical isothermal mass distribution. This is used in the `Isothermal` and `Chameleon` `MassProfile`'s. @@ -31,7 +31,7 @@ def psi_from(grid, axis_ratio, core_radius): The value of the Psi term. """ - return jnp.sqrt( + return xp.sqrt( (axis_ratio**2.0 * (grid.array[:, 1] ** 2.0 + core_radius**2.0)) + grid.array[:, 0] ** 2.0 + 1e-16 @@ -66,14 +66,13 @@ def __init__( slope=2.0, ) - @property - def axis_ratio(self): - axis_ratio = super().axis_ratio - return jnp.minimum(axis_ratio, 0.99999) + def axis_ratio(self, xp=np): + axis_ratio = super().axis_ratio(xp=xp) + return xp.minimum(axis_ratio, 0.99999) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -88,31 +87,36 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): factor = ( 2.0 - * self.einstein_radius_rescaled - * self.axis_ratio - / jnp.sqrt(1 - self.axis_ratio**2) + * self.einstein_radius_rescaled(xp) + * self.axis_ratio(xp) + / xp.sqrt(1 - self.axis_ratio(xp) ** 2) ) - psi = psi_from(grid=grid, axis_ratio=self.axis_ratio, core_radius=0.0) + psi = psi_from( + grid=grid, axis_ratio=self.axis_ratio(xp), core_radius=0.0, xp=xp + ) - deflection_y = jnp.arctanh( - jnp.divide( - jnp.multiply(jnp.sqrt(1 - self.axis_ratio**2), grid.array[:, 0]), psi + deflection_y = xp.arctanh( + xp.divide( + xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 0]), + psi, ) ) - deflection_x = jnp.arctan( - jnp.divide( - jnp.multiply(jnp.sqrt(1 - self.axis_ratio**2), grid.array[:, 1]), psi + deflection_x = xp.arctan( + xp.divide( + xp.multiply(xp.sqrt(1 - self.axis_ratio(xp) ** 2), grid.array[:, 1]), + psi, ) ) return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply(factor, jnp.vstack((deflection_y, deflection_x)).T), + grid=xp.multiply(factor, xp.vstack((deflection_y, deflection_x)).T), + xp=xp, **kwargs, ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the (gamma_y, gamma_x) shear vector field on a grid of (y,x) arc-second coordinates. @@ -129,23 +133,23 @@ def shear_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ - convergence = self.convergence_2d_from(grid=grid, **kwargs) + convergence = self.convergence_2d_from(grid=grid, xp=xp, **kwargs) gamma_2 = ( -2 * convergence.array - * jnp.divide( + * xp.divide( grid.array[:, 1] * grid.array[:, 0], grid.array[:, 1] ** 2 + grid.array[:, 0] ** 2, ) ) - gamma_1 = -convergence.array * jnp.divide( + gamma_1 = -convergence.array * xp.divide( grid.array[:, 1] ** 2 - grid.array[:, 0] ** 2, grid.array[:, 1] ** 2 + grid.array[:, 0] ** 2, ) shear_field = self.rotated_grid_from_reference_frame_from( - grid=jnp.vstack((gamma_2, gamma_1)).T, angle=self.angle * 2 + grid=xp.vstack((gamma_2, gamma_1)).T, xp=xp, angle=self.angle(xp) * 2 ) return aa.VectorYX2DIrregular(values=shear_field, grid=grid) @@ -170,14 +174,13 @@ def __init__( centre=centre, ell_comps=(0.0, 0.0), einstein_radius=einstein_radius ) - @property - def axis_ratio(self): + def axis_ratio(self, xp=np): return 1.0 @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the potential on a grid of (y,x) arc-second coordinates. @@ -186,12 +189,12 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - eta = self.elliptical_radii_grid_from(grid=grid, **kwargs) - return 2.0 * self.einstein_radius_rescaled * eta + eta = self.elliptical_radii_grid_from(grid=grid, xp=xp, **kwargs) + return 2.0 * self.einstein_radius_rescaled(xp) * eta @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -202,6 +205,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ return self._cartesian_grid_via_radial_from( grid=grid, - radius=jnp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled), + radius=xp.full(grid.shape[0], 2.0 * self.einstein_radius_rescaled(xp)), + xp=xp, **kwargs, ) diff --git a/autogalaxy/profiles/mass/total/jax_utils.py b/autogalaxy/profiles/mass/total/jax_utils.py index 5237f34cd..d04691ea2 100644 --- a/autogalaxy/profiles/mass/total/jax_utils.py +++ b/autogalaxy/profiles/mass/total/jax_utils.py @@ -1,11 +1,4 @@ -import jax -import jax.numpy as jnp -from jax.tree_util import Partial as partial - - -# A version of scan that will *not* re-compile partial functions when variables change -# taken from https://github.com/google/jax/issues/14743#issuecomment-1456900634 -scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll")) +import numpy as np def body_fun(carry, n, factor, ei2phi, slope): @@ -18,7 +11,7 @@ def body_fun(carry, n, factor, ei2phi, slope): return (omega_n, partial_sum), None -def omega(eiphi, slope, factor, n_terms=20): +def omega(eiphi, slope, factor, n_terms=20, xp=np): """JAX implementation of the numerical evaluation of the angular component of the complex deflection angle for the elliptical power law profile as given as given by Tessore and Metcalf 2015. Based on equation 29, and gives @@ -38,10 +31,16 @@ def omega(eiphi, slope, factor, n_terms=20): The number of terms to calculate for the series expansion, defaults to 20 (this should be sufficient most of the time) """ + + from jax.tree_util import Partial as partial + import jax + + scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll")) + # use modified scan with a partial'ed function to avoid re-compile (_, partial_sum), _ = scan( partial(body_fun, factor=factor, ei2phi=eiphi**2, slope=slope), (eiphi, eiphi), - jnp.arange(1, n_terms), + xp.arange(1, n_terms), ) return partial_sum diff --git a/autogalaxy/profiles/mass/total/power_law.py b/autogalaxy/profiles/mass/total/power_law.py index 829189acb..3742bd23f 100644 --- a/autogalaxy/profiles/mass/total/power_law.py +++ b/autogalaxy/profiles/mass/total/power_law.py @@ -1,6 +1,4 @@ -import jax.numpy as jnp -from .jax_utils import omega - +import numpy as np from typing import Tuple import autoarray as aa @@ -40,7 +38,7 @@ def __init__( ) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): alpha = self.deflections_yx_2d_from(aa.Grid2DIrregular(grid), **kwargs) alpha_x = alpha[:, 1] @@ -53,7 +51,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -68,28 +66,31 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ + from .jax_utils import omega slope = self.slope - 1.0 einstein_radius = ( - 2.0 / (self.axis_ratio**-0.5 + self.axis_ratio**0.5) + 2.0 / (self.axis_ratio(xp) ** -0.5 + self.axis_ratio(xp) ** 0.5) ) * self.einstein_radius - factor = jnp.divide(1.0 - self.axis_ratio, 1.0 + self.axis_ratio) - b = jnp.multiply(einstein_radius, jnp.sqrt(self.axis_ratio)) - angle = jnp.arctan2( - grid.array[:, 0], jnp.multiply(self.axis_ratio, grid.array[:, 1]) + factor = xp.divide(1.0 - self.axis_ratio(xp), 1.0 + self.axis_ratio(xp)) + b = xp.multiply(einstein_radius, xp.sqrt(self.axis_ratio(xp))) + angle = xp.arctan2( + grid.array[:, 0], xp.multiply(self.axis_ratio(xp), grid.array[:, 1]) ) # Note, this angle is not the position angle - z = jnp.add( - jnp.multiply(jnp.cos(angle), 1 + 0j), jnp.multiply(jnp.sin(angle), 0 + 1j) + z = xp.add( + xp.multiply(xp.cos(angle), 1 + 0j), xp.multiply(xp.sin(angle), 0 + 1j) ) - R = jnp.sqrt( - (self.axis_ratio * grid.array[:, 1]) ** 2 + grid.array[:, 0] ** 2 + 1e-16 + R = xp.sqrt( + (self.axis_ratio(xp) * grid.array[:, 1]) ** 2 + + grid.array[:, 0] ** 2 + + 1e-16 ) - zh = omega(z, slope, factor, n_terms=20) + zh = omega(z, slope, factor, n_terms=20, xp=np) complex_angle = ( - 2.0 * b / (1.0 + self.axis_ratio) * (b / R) ** (slope - 1.0) * zh + 2.0 * b / (1.0 + self.axis_ratio(xp)) * (b / R) ** (slope - 1.0) * zh ) deflection_y = complex_angle.imag @@ -101,15 +102,17 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): deflection_x *= rescale_factor return self.rotated_grid_from_reference_frame_from( - grid=jnp.vstack((deflection_y, deflection_x)).T + grid=xp.vstack((deflection_y, deflection_x)).T, xp=xp ) - def convergence_func(self, grid_radius: float) -> float: - return self.einstein_radius_rescaled * grid_radius.array ** (-(self.slope - 1)) + def convergence_func(self, grid_radius: float, xp=np) -> float: + return self.einstein_radius_rescaled(xp) * grid_radius.array ** ( + -(self.slope - 1) + ) @staticmethod - def potential_func(u, y, x, axis_ratio, slope, core_radius): - _eta_u = jnp.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u))))) + def potential_func(u, y, x, axis_ratio, slope, core_radius, xp=np): + _eta_u = xp.sqrt((u * ((x**2) + (y**2 / (1 - (1 - axis_ratio**2) * u))))) return ( (_eta_u / u) * ((3.0 - slope) * _eta_u) ** -1.0 @@ -147,15 +150,17 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - eta = self.radial_grid_from(grid=grid, **kwargs).array + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + eta = self.radial_grid_from(grid=grid, xp=xp, **kwargs).array deflection_r = ( 2.0 - * self.einstein_radius_rescaled - * jnp.divide( - jnp.power(eta, (3.0 - self.slope)), - jnp.multiply((3.0 - self.slope), eta), + * self.einstein_radius_rescaled(xp) + * xp.divide( + xp.power(eta, (3.0 - self.slope)), + xp.multiply((3.0 - self.slope), eta), ) ) - return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection_r) + return self._cartesian_grid_via_radial_from( + grid=grid, radius=deflection_r, xp=xp + ) diff --git a/autogalaxy/profiles/mass/total/power_law_broken.py b/autogalaxy/profiles/mass/total/power_law_broken.py index 4e17c5c6d..577049787 100644 --- a/autogalaxy/profiles/mass/total/power_law_broken.py +++ b/autogalaxy/profiles/mass/total/power_law_broken.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -31,7 +30,7 @@ def __init__( super().__init__(centre=centre, ell_comps=ell_comps) self.einstein_radius = einstein_radius - self.einstein_radius_elliptical = jnp.sqrt(self.axis_ratio) * einstein_radius + self.einstein_radius_elliptical = np.sqrt(self.axis_ratio()) * einstein_radius self.break_radius = break_radius self.inner_slope = inner_slope self.outer_slope = outer_slope @@ -52,13 +51,13 @@ def __init__( @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the dimensionless density kappa=Sigma/Sigma_c (eq. 1) """ # Ell radius - radius = jnp.hypot(grid.array[:, 1] * self.axis_ratio, grid.array[:, 0]) + radius = xp.hypot(grid.array[:, 1] * self.axis_ratio(xp), grid.array[:, 0]) # Inside break radius kappa_inner = self.kB * (self.break_radius / radius) ** self.inner_slope @@ -71,27 +70,28 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): ) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): - return jnp.zeros(shape=grid.shape[0]) + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): + return xp.zeros(shape=grid.shape[0]) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs): + def deflections_yx_2d_from(self, grid, xp=np, max_terms=20, **kwargs): """ Returns the complex deflection angle from eq. 18 and 19 """ + # Rotate coordinates z = grid.array[:, 1] + 1j * grid.array[:, 0] # Ell radius - R = jnp.hypot(z.real * self.axis_ratio, z.imag) + R = xp.hypot(z.real * self.axis_ratio(xp), z.imag) # Factors common to eq. 18 and 19 factors = ( 2 * self.kB * (self.break_radius**2) - / (self.axis_ratio * z * (2 - self.inner_slope)) + / (self.axis_ratio(xp) * z * (2 - self.inner_slope)) ) # Hypergeometric functions @@ -99,16 +99,26 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs): # These can also be computed with scipy.special.hyp2f1(), it's # much slower can be a useful test F1 = self.hyp2f1_series( - self.inner_slope, self.axis_ratio, R, z, max_terms=max_terms + self.inner_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp ) F2 = self.hyp2f1_series( - self.inner_slope, self.axis_ratio, self.break_radius, z, max_terms=max_terms + self.inner_slope, + self.axis_ratio(xp), + self.break_radius, + z, + max_terms=max_terms, + xp=xp, ) F3 = self.hyp2f1_series( - self.outer_slope, self.axis_ratio, R, z, max_terms=max_terms + self.outer_slope, self.axis_ratio(xp), R, z, max_terms=max_terms, xp=xp ) F4 = self.hyp2f1_series( - self.outer_slope, self.axis_ratio, self.break_radius, z, max_terms=max_terms + self.outer_slope, + self.axis_ratio(xp), + self.break_radius, + z, + max_terms=max_terms, + xp=xp, ) # theta < break radius (eq. 18) @@ -126,13 +136,14 @@ def deflections_yx_2d_from(self, grid, max_terms=20, **kwargs): ).conjugate() return self.rotated_grid_from_reference_frame_from( - grid=jnp.multiply( - 1.0, jnp.vstack((jnp.imag(deflections), jnp.real(deflections))).T - ) + grid=xp.multiply( + 1.0, xp.vstack((xp.imag(deflections), xp.real(deflections))).T + ), + xp=xp, ) @staticmethod - def hyp2f1_series(t, q, r, z, max_terms=20): + def hyp2f1_series(t, q, r, z, max_terms=20, xp=np): """ Computes eq. 26 for a radius r, slope t, axis ratio q, and coordinates z. @@ -140,13 +151,13 @@ def hyp2f1_series(t, q, r, z, max_terms=20): # u from eq. 25 q_ = (1 - q**2) / (q**2) - u = 0.5 * (1 - jnp.sqrt(1 - q_ * (r / z) ** 2)) + u = 0.5 * (1 - xp.sqrt(1 - q_ * (r / z) ** 2)) # First coefficient a_n = 1.0 # Storage for sum - F = jnp.zeros_like(z, dtype="complex64") + F = xp.zeros_like(z, dtype="complex64") for n in range(max_terms): F += a_n * (u**n) diff --git a/autogalaxy/profiles/mass/total/power_law_core.py b/autogalaxy/profiles/mass/total/power_law_core.py index af2069e17..5ba9ce9b6 100644 --- a/autogalaxy/profiles/mass/total/power_law_core.py +++ b/autogalaxy/profiles/mass/total/power_law_core.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -38,20 +37,19 @@ def __init__( self.slope = slope self.core_radius = core_radius - @property - def einstein_radius_rescaled(self): + def einstein_radius_rescaled(self, xp=np): """ Rescale the einstein radius by slope and axis_ratio, to reduce its degeneracy with other mass-profiles parameters. """ - return ((3 - self.slope) / (1 + self.axis_ratio)) * self.einstein_radius ** ( - self.slope - 1 - ) + return ( + (3 - self.slope) / (1 + self.axis_ratio(xp)) + ) * self.einstein_radius ** (self.slope - 1) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -70,7 +68,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the potential on a grid of (y,x) arc-second coordinates. @@ -91,17 +89,17 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): args=( grid.array[i, 0], grid.array[i, 1], - self.axis_ratio, + self.axis_ratio(xp), self.slope, self.core_radius, ), )[0] - return self.einstein_radius_rescaled * self.axis_ratio * potential_grid + return self.einstein_radius_rescaled(xp) * self.axis_ratio(xp) * potential_grid @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -114,9 +112,9 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): from scipy.integrate import quad def calculate_deflection_component(npow, index): - einstein_radius_rescaled = self.einstein_radius_rescaled + einstein_radius_rescaled = self.einstein_radius_rescaled(xp) - deflection_grid = np.array(self.axis_ratio * grid.array[:, index]) + deflection_grid = np.array(self.axis_ratio(xp) * grid.array[:, index]) for i in range(grid.shape[0]): deflection_grid[i] *= ( @@ -129,7 +127,7 @@ def calculate_deflection_component(npow, index): grid.array[i, 0], grid.array[i, 1], npow, - self.axis_ratio, + self.axis_ratio(xp), self.slope, self.core_radius, ), @@ -142,11 +140,11 @@ def calculate_deflection_component(npow, index): deflection_x = calculate_deflection_component(0.0, 1) return self.rotated_grid_from_reference_frame_from( - grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T) + grid=np.multiply(1.0, np.vstack((deflection_y, deflection_x)).T), xp=xp ) - def convergence_func(self, grid_radius: float) -> float: - return self.einstein_radius_rescaled * ( + def convergence_func(self, grid_radius: float, xp=np) -> float: + return self.einstein_radius_rescaled(xp) * ( self.core_radius**2 + grid_radius**2 ) ** (-(self.slope - 1) / 2.0) @@ -172,7 +170,7 @@ def deflection_func(u, y, x, npow, axis_ratio, slope, core_radius): @property def ellipticity_rescale(self): - return (1.0 + self.axis_ratio) / 2.0 + return (1.0 + self.axis_ratio()) / 2.0 @property def unit_mass(self): @@ -211,7 +209,7 @@ def __init__( @aa.grid_dec.to_vector_yx @aa.grid_dec.transform - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): + def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -222,17 +220,17 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs): """ eta = self.radial_grid_from(grid=grid, **kwargs) - deflection = jnp.multiply( - 2.0 * self.einstein_radius_rescaled, - jnp.divide( - jnp.add( - jnp.power( - jnp.add(self.core_radius**2, jnp.square(eta.array)), + deflection = xp.multiply( + 2.0 * self.einstein_radius_rescaled(xp), + xp.divide( + xp.add( + xp.power( + xp.add(self.core_radius**2, xp.square(eta.array)), (3.0 - self.slope) / 2.0, ), -self.core_radius ** (3 - self.slope), ), - jnp.multiply((3.0 - self.slope), eta.array), + xp.multiply((3.0 - self.slope), eta.array), ), ) - return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection) + return self._cartesian_grid_via_radial_from(grid=grid, radius=deflection, xp=xp) diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py index 16e5a2a88..f20f02102 100644 --- a/autogalaxy/profiles/mass/total/power_law_multipole.py +++ b/autogalaxy/profiles/mass/total/power_law_multipole.py @@ -1,4 +1,3 @@ -import jax.numpy as jnp import numpy as np from typing import Tuple @@ -10,7 +9,7 @@ def radial_and_angle_grid_from( - grid: aa.type.Grid2DLike, centre: Tuple[float, float] = (0.0, 0.0) + grid: aa.type.Grid2DLike, centre: Tuple[float, float] = (0.0, 0.0), xp=np ) -> Tuple[np.ndarray, np.ndarray]: """ Converts the input grid of Cartesian (y,x) coordinates to their correspond radial and polar grids. @@ -28,12 +27,12 @@ def radial_and_angle_grid_from( """ y, x = grid.array.T - x_shifted = jnp.subtract(x, centre[1]) - y_shifted = jnp.subtract(y, centre[0]) + x_shifted = xp.subtract(x, centre[1]) + y_shifted = xp.subtract(y, centre[0]) - radial_grid = jnp.sqrt(x_shifted**2 + y_shifted**2) + radial_grid = xp.sqrt(x_shifted**2 + y_shifted**2) - angle_grid = jnp.arctan2(y_shifted, x_shifted) + angle_grid = xp.arctan2(y_shifted, x_shifted) return radial_grid, angle_grid @@ -129,7 +128,7 @@ def __init__( self.angle_m *= units.deg.to(units.rad) def jacobian( - self, a_r: np.ndarray, a_angle: np.ndarray, polar_angle_grid: np.ndarray + self, a_r: np.ndarray, a_angle: np.ndarray, polar_angle_grid: np.ndarray, xp=np ) -> Tuple[np.ndarray, Tuple]: """ The Jacobian transformation from polar to cartesian coordinates. @@ -144,14 +143,14 @@ def jacobian( The polar angle coordinates of the input (y,x) Cartesian grid of coordinates. """ return ( - a_r * jnp.sin(polar_angle_grid) + a_angle * jnp.cos(polar_angle_grid), - a_r * jnp.cos(polar_angle_grid) - a_angle * jnp.sin(polar_angle_grid), + a_r * xp.sin(polar_angle_grid) + a_angle * xp.cos(polar_angle_grid), + a_r * xp.cos(polar_angle_grid) - a_angle * xp.sin(polar_angle_grid), ) @aa.grid_dec.to_vector_yx @aa.grid_dec.transform def deflections_yx_2d_from( - self, grid: aa.type.Grid1D2DLike, **kwargs + self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs ) -> np.ndarray: """ Calculate the deflection angles on a grid of (y,x) arc-second coordinates. @@ -164,7 +163,7 @@ def deflections_yx_2d_from( grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - radial_grid, polar_angle_grid = radial_and_angle_grid_from(grid=grid) + radial_grid, polar_angle_grid = radial_and_angle_grid_from(grid=grid, xp=xp) a_r = ( -( @@ -174,7 +173,7 @@ def deflections_yx_2d_from( ) / (self.m**2.0 - (3.0 - self.slope) ** 2.0) * self.k_m - * jnp.cos(self.m * (polar_angle_grid - self.angle_m)) + * xp.cos(self.m * (polar_angle_grid - self.angle_m)) ) a_angle = ( @@ -185,18 +184,22 @@ def deflections_yx_2d_from( ) / (self.m**2.0 - (3.0 - self.slope) ** 2.0) * self.k_m - * jnp.sin(self.m * (polar_angle_grid - self.angle_m)) + * xp.sin(self.m * (polar_angle_grid - self.angle_m)) ) - return jnp.stack( - self.jacobian(a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid), + return xp.stack( + self.jacobian( + a_r=a_r, a_angle=a_angle, polar_angle_grid=polar_angle_grid, xp=xp + ), axis=-1, ) @aa.over_sample @aa.grid_dec.to_array @aa.grid_dec.transform - def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, **kwargs) -> np.ndarray: + def convergence_2d_from( + self, grid: aa.type.Grid1D2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Returns the two dimensional projected convergence on a grid of (y,x) arc-second coordinates. @@ -205,18 +208,20 @@ def convergence_2d_from(self, grid: aa.type.Grid1D2DLike, **kwargs) -> np.ndarra grid The grid of (y,x) arc-second coordinates the convergence is computed on. """ - r, angle = radial_and_angle_grid_from(grid=grid) + r, angle = radial_and_angle_grid_from(grid=grid, xp=xp) return ( 1.0 / 2.0 * (self.einstein_radius / r) ** (self.slope - 1) * self.k_m - * jnp.cos(self.m * (angle - self.angle_m)) + * xp.cos(self.m * (angle - self.angle_m)) ) @aa.grid_dec.to_array - def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: + def potential_2d_from( + self, grid: aa.type.Grid2DLike, xp=np, **kwargs + ) -> np.ndarray: """ Calculate the potential on a grid of (y,x) arc-second coordinates. @@ -225,4 +230,4 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray: grid The grid of (y,x) arc-second coordinates the deflection angles are computed on. """ - return jnp.zeros(shape=grid.shape[0]) + return xp.zeros(shape=grid.shape[0]) diff --git a/autogalaxy/profiles/plot/light_profile_plotters.py b/autogalaxy/profiles/plot/light_profile_plotters.py index e8521bb7b..7292f5eb0 100644 --- a/autogalaxy/profiles/plot/light_profile_plotters.py +++ b/autogalaxy/profiles/plot/light_profile_plotters.py @@ -1,6 +1,3 @@ -import math -from typing import List, Optional - import autoarray as aa import autoarray.plot as aplt @@ -12,7 +9,6 @@ from autogalaxy.plot.visuals.one_d import Visuals1D from autogalaxy.plot.visuals.two_d import Visuals2D -from autogalaxy.util import error_util from autogalaxy import exc @@ -73,43 +69,11 @@ def __init__( visuals_1d=visuals_1d, ) - def figures_1d(self, image: bool = False): - """ - Plots the individual attributes of the plotter's `LightProfile` object in 1D, which are computed via the - plotter's grid object. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from the light profile. This is performed by aligning a 1D grid with the major-axis of the light - profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - The API is such that every plottable attribute of the `LightProfile` object is an input parameter of type - bool of the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if image: - image_1d = self.light_profile.image_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=image_1d, - x=image_1d.grid_radial, - visuals_1d=self.visuals_1d, - auto_labels=aplt.AutoLabels( - title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)", - yunit="", - legend=self.light_profile.__class__.__name__, - filename="image_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) + @property + def grid_2d_projected(self): + return self.grid.grid_2d_radial_projected_from( + centre=self.light_profile.centre, angle=self.light_profile.angle() + ) def figures_2d(self, image: bool = False): """ @@ -130,131 +94,3 @@ def figures_2d(self, image: bool = False): visuals_2d=self.visuals_2d, auto_labels=aplt.AutoLabels(title="Image", filename="image_2d"), ) - - -class LightProfilePDFPlotter(LightProfilePlotter): - def __init__( - self, - light_profile_pdf_list: List[LightProfile], - grid: aa.type.Grid2DLike, - mat_plot_1d: MatPlot1D = None, - visuals_1d: Visuals1D = None, - mat_plot_2d: MatPlot2D = None, - visuals_2d: Visuals2D = None, - sigma: Optional[float] = 3.0, - ): - """ - Plots the attributes of a list of `LightProfile` objects using the matplotlib methods `plot()` and `imshow()` - and many other matplotlib functions which customize the plot's appearance. - - Figures plotted by this object average over a list light profiles to computed the average value of each - attribute with errors, where the 1D regions within the errors are plotted as a shaded region to show the range - of plausible models. Therefore, the input list of galaxies is expected to represent the probability density - function of an inferred model-fit. - - The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default, - the settings passed to every matplotlib function called are those specified in - the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to - customize the figure's appearance. - - Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be - extracted from the `LightProfile` and plotted via the visuals object. - - Parameters - ---------- - light_profile_pdf_list - The list of light profiles whose mean and error values the plotter plots. - grid - The 2D (y,x) grid of coordinates used to evaluate the light profile quantities that are plotted. - mat_plot_1d - Contains objects which wrap the matplotlib function calls that make 1D plots. - visuals_1d - Contains 1D visuals that can be overlaid on 1D plots. - mat_plot_2d - Contains objects which wrap the matplotlib function calls that make 2D plots. - visuals_2d - Contains 2D visuals that can be overlaid on 2D plots. - sigma - The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of - sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF). - """ - super().__init__( - light_profile=None, - grid=grid, - mat_plot_1d=mat_plot_1d, - visuals_1d=visuals_1d, - mat_plot_2d=mat_plot_2d, - visuals_2d=visuals_2d, - ) - - self.light_profile_pdf_list = light_profile_pdf_list - self.sigma = sigma - self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2 - - def figures_1d(self, image: bool = False): - """ - Plots the individual attributes of the plotter's list of ` LightProfile` object in 1D, which are computed via - the plotter's grid object. - - This averages over a list light profiles to compute the average value of each attribute with errors, where the - 1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore, - the input list of galaxies is expected to represent the probability density function of an inferred model-fit. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each light profile. This is performed by aligning a 1D grid with the major-axis of - each light profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - The API is such that every plottable attribute of the `LightProfile` object is an input parameter of type bool - of the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - image - Whether to make a 1D plot (via `plot`) of the image. - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if image: - image_1d_list = [ - light_profile.image_1d_from(grid=self.grid) - for light_profile in self.light_profile_pdf_list - ] - - min_index = min([image_1d.shape[0] for image_1d in image_1d_list]) - image_1d_list = [image_1d[0:min_index] for image_1d in image_1d_list] - - ( - median_image_1d, - errors_image_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=image_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_light_obj_list = Visuals1D().add_half_light_radius_errors( - light_obj_list=self.light_profile_pdf_list, - low_limit=self.low_limit, - ) - - visuals_1d_with_shaded_region = Visuals1D(shaded_region=errors_image_1d) - - visuals_1d = visuals_1d_via_light_obj_list + visuals_1d_with_shaded_region - - self.mat_plot_1d.plot_yx( - y=median_image_1d, - x=image_1d_list[0].grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title=r"Image ($\mathrm{e^{-}}\,\mathrm{s^{-1}}$) vs Radius (arcsec)", - yunit="", - legend=self.light_profile_pdf_list[0].__class__.__name__, - filename="image_1d_pdf", - ), - plot_axis_type_override=plot_axis_type_override, - ) diff --git a/autogalaxy/profiles/plot/mass_profile_plotters.py b/autogalaxy/profiles/plot/mass_profile_plotters.py index c2d9ea4bd..62eaf8b88 100644 --- a/autogalaxy/profiles/plot/mass_profile_plotters.py +++ b/autogalaxy/profiles/plot/mass_profile_plotters.py @@ -71,232 +71,8 @@ def __init__( self.figures_2d = self._mass_plotter.figures_2d - def figures_1d(self, convergence: bool = False, potential: bool = False): - """ - Plots the individual attributes of the plotter's `MassProfile` object in 1D, which are computed via the - plotter's grid object. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from the mass profile. This is performed by aligning a 1D grid with the major-axis of the mass - profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - The API is such that every plottable attribute of the `MassProfile` object is an input parameter of type - bool of the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if convergence: - convergence_1d = self.mass_profile.convergence_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=convergence_1d, - x=convergence_1d.grid_radial, - visuals_1d=self.visuals_1d, - auto_labels=aplt.AutoLabels( - title="Convergence vs Radius (arcsec)", - yunit="", - legend=self.mass_profile.__class__.__name__, - filename="convergence_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if potential: - potential_1d = self.mass_profile.potential_1d_from(grid=self.grid) - - self.mat_plot_1d.plot_yx( - y=potential_1d, - x=potential_1d.grid_radial, - visuals_1d=self.visuals_1d, - auto_labels=aplt.AutoLabels( - title="Potential vs Radius (arcsec)", - yunit="", - legend=self.mass_profile.__class__.__name__, - filename="potential_1d", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - -class MassProfilePDFPlotter(MassProfilePlotter): - def __init__( - self, - mass_profile_pdf_list: List[MassProfile], - grid: aa.Grid2D, - mat_plot_1d: MatPlot1D = None, - visuals_1d: Visuals1D = None, - mat_plot_2d: MatPlot2D = None, - visuals_2d: Visuals2D = None, - sigma: Optional[float] = 3.0, - ): - """ - Plots the attributes of a list of `MassProfile` objects using the matplotlib methods `plot()` and `imshow()` - and many other matplotlib functions which customize the plot's appearance. - - Figures plotted by this object average over a list mass profiles to computed the average value of each attribute - with errors, where the 1D regions within the errors are plotted as a shaded region to show the range of - plausible models. Therefore, the input list of galaxies is expected to represent the probability density - function of an inferred model-fit. - - The `mat_plot_1d` and `mat_plot_2d` attributes wrap matplotlib function calls to make the figure. By default, - the settings passed to every matplotlib function called are those specified in - the `config/visualize/mat_wrap/*.ini` files, but a user can manually input values into `MatPlot2D` to - customize the figure's appearance. - - Overlaid on the figure are visuals, contained in the `Visuals1D` and `Visuals2D` objects. Attributes may be - extracted from the `MassProfile` and plotted via the visuals object. - - Parameters - ---------- - mass_profile_pdf_list - The list of mass profiles whose mean and error values the plotter plots. - grid - The 2D (y,x) grid of coordinates used to evaluate the mass profile quantities that are plotted. - mat_plot_1d - Contains objects which wrap the matplotlib function calls that make 1D plots. - visuals_1d - Contains 1D visuals that can be overlaid on 1D plots. - mat_plot_2d - Contains objects which wrap the matplotlib function calls that make 2D plots. - visuals_2d - Contains 2D visuals that can be overlaid on 2D plots. - sigma - The confidence interval in terms of a sigma value at which the errors are computed (e.g. a value of - sigma=3.0 uses confidence intevals at ~0.01 and 0.99 the PDF). - """ - super().__init__( - mass_profile=None, - grid=grid, - mat_plot_1d=mat_plot_1d, - visuals_1d=visuals_1d, - mat_plot_2d=mat_plot_2d, - visuals_2d=visuals_2d, + @property + def grid_2d_projected(self): + return self.grid.grid_2d_radial_projected_from( + centre=self.mass_profile.centre, angle=self.mass_profile.angle() ) - - self.mass_profile_pdf_list = mass_profile_pdf_list - self.sigma = sigma - self.low_limit = (1 - math.erf(sigma / math.sqrt(2))) / 2 - - def figures_1d(self, convergence=False, potential=False): - """ - Plots the individual attributes of the plotter's list of ` MassProfile` object in 1D, which are computed via - the plotter's grid object. - - This averages over a list mass profiles to compute the average value of each attribute with errors, where the - 1D regions within the errors are plotted as a shaded region to show the range of plausible models. Therefore, - the input list of galaxies is expected to represent the probability density function of an inferred model-fit. - - If the plotter has a 1D grid object this is used to evaluate each quantity. If it has a 2D grid, a 1D grid is - computed from each mass profile. This is performed by aligning a 1D grid with the major-axis of - each mass profile in projection, uniformly computing 1D values based on the 2D grid's size and pixel-scale. - - The API is such that every plottable attribute of the `MassProfile` object is an input parameter of type bool - of the function, which if switched to `True` means that it is plotted. - - Parameters - ---------- - convergence - Whether to make a 1D plot (via `imshow`) of the convergence. - potential - Whether to make a 1D plot (via `imshow`) of the potential. - """ - if self.mat_plot_1d.yx_plot.plot_axis_type is None: - plot_axis_type_override = "semilogy" - else: - plot_axis_type_override = None - - if convergence: - convergence_1d_list = [ - mass_profile.convergence_1d_from(grid=self.grid) - for mass_profile in self.mass_profile_pdf_list - ] - - min_index = min( - [convergence_1d.shape[0] for convergence_1d in convergence_1d_list] - ) - convergence_1d_list = [ - convergence_1d[0:min_index] for convergence_1d in convergence_1d_list - ] - - ( - median_convergence_1d, - errors_convergence_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=convergence_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors( - mass_obj_list=self.mass_profile_pdf_list, - grid=self.grid, - low_limit=self.low_limit, - ) - visuals_1d_with_shaded_region = Visuals1D( - shaded_region=errors_convergence_1d - ) - - visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region - - self.mat_plot_1d.plot_yx( - y=median_convergence_1d, - x=convergence_1d_list[0].grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title="Convergence vs Radius (arcsec)", - yunit="", - legend=self.mass_profile_pdf_list[0].__class__.__name__, - filename="convergence_1d_pdf", - ), - plot_axis_type_override=plot_axis_type_override, - ) - - if potential: - potential_1d_list = [ - mass_profile.potential_1d_from(grid=self.grid) - for mass_profile in self.mass_profile_pdf_list - ] - - min_index = min( - [potential_1d.shape[0] for potential_1d in potential_1d_list] - ) - potential_1d_list = [ - potential_1d[0:min_index] for potential_1d in potential_1d_list - ] - - ( - median_potential_1d, - errors_potential_1d, - ) = error_util.profile_1d_median_and_error_region_via_quantile( - profile_1d_list=potential_1d_list, low_limit=self.low_limit - ) - - visuals_1d_via_lensing_obj_list = Visuals1D().add_einstein_radius_errors( - mass_obj_list=self.mass_profile_pdf_list, - grid=self.grid, - low_limit=self.low_limit, - ) - visuals_1d_with_shaded_region = Visuals1D(shaded_region=errors_potential_1d) - - visuals_1d = visuals_1d_via_lensing_obj_list + visuals_1d_with_shaded_region - - self.mat_plot_1d.plot_yx( - y=median_potential_1d, - x=potential_1d_list[0].grid_radial, - visuals_1d=visuals_1d, - auto_labels=aplt.AutoLabels( - title="Potential vs Radius (arcsec)", - yunit="", - legend=self.mass_profile_pdf_list[0].__class__.__name__, - filename="potential_1d_pdf", - ), - plot_axis_type_override=plot_axis_type_override, - ) diff --git a/autogalaxy/quantity/dataset_quantity.py b/autogalaxy/quantity/dataset_quantity.py index 683f3d464..7deb42dd2 100644 --- a/autogalaxy/quantity/dataset_quantity.py +++ b/autogalaxy/quantity/dataset_quantity.py @@ -1,10 +1,12 @@ import logging import numpy as np from pathlib import Path -from typing import List, Optional, Union +from typing import Optional, Union import autoarray as aa + from autoarray.dataset.abstract.dataset import AbstractDataset +from autoarray.dataset.grids import GridsDataset logger = logging.getLogger(__name__) @@ -84,6 +86,12 @@ def __init__( over_sample_size_pixelization=over_sample_size_pixelization, ) + self.grids = GridsDataset( + mask=self.data.mask, + over_sample_size_lp=self.over_sample_size_lp, + over_sample_size_pixelization=self.over_sample_size_pixelization, + ) + @classmethod def via_signal_to_noise_map( cls, diff --git a/autogalaxy/quantity/model/analysis.py b/autogalaxy/quantity/model/analysis.py index 48621ba26..3e5710692 100644 --- a/autogalaxy/quantity/model/analysis.py +++ b/autogalaxy/quantity/model/analysis.py @@ -1,3 +1,5 @@ +import numpy as np + from autoconf.dictable import to_dict import autofit as af @@ -60,7 +62,7 @@ def __init__( self.func_str = func_str self.title_prefix = title_prefix - def log_likelihood_function(self, instance: af.ModelInstance) -> float: + def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float: """ Given an instance of the model, where the model parameters are set via a non-linear search, fit the model instance to the quantity's dataset. diff --git a/docs/api/plot.rst b/docs/api/plot.rst index a930d6872..6325e8383 100644 --- a/docs/api/plot.rst +++ b/docs/api/plot.rst @@ -32,7 +32,6 @@ Create figures and subplots showing quantities of standard **PyAutoGalaxy** obje ImagingPlotter InterferometerPlotter LightProfilePlotter - LightProfilePDFPlotter GalaxyPlotter FitImagingPlotter FitInterferometerPlotter diff --git a/test_autogalaxy/analysis/test_plotter_interface.py b/test_autogalaxy/analysis/test_plotter_interface.py index e7132c878..b883e74ae 100644 --- a/test_autogalaxy/analysis/test_plotter_interface.py +++ b/test_autogalaxy/analysis/test_plotter_interface.py @@ -29,10 +29,6 @@ def test__galaxies(masked_imaging_7x7, galaxies_7x7, plot_path, plot_patch): assert path.join(plot_path, "subplot_galaxies.png") in plot_patch.paths assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths - assert path.join(plot_path, "subplot_galaxies_1d.png") in plot_patch.paths - assert ( - path.join(plot_path, "subplot_galaxies_1d_decomposed.png") in plot_patch.paths - ) image = ag.ndarray_via_fits_from( file_path=path.join(plot_path, "galaxy_images.fits"), hdu=1 diff --git a/test_autogalaxy/config/general.yaml b/test_autogalaxy/config/general.yaml index 0b88a5865..07df488b4 100644 --- a/test_autogalaxy/config/general.yaml +++ b/test_autogalaxy/config/general.yaml @@ -22,8 +22,6 @@ hpc: adapt: adapt_minimum_percent: 0.01 adapt_noise_limit: 100000000.0 -model: - ignore_prior_limits: true numba: cache: true nopython: true diff --git a/test_autogalaxy/config/visualize.yaml b/test_autogalaxy/config/visualize.yaml index af4a18bcb..754ce9375 100644 --- a/test_autogalaxy/config/visualize.yaml +++ b/test_autogalaxy/config/visualize.yaml @@ -351,8 +351,6 @@ plots: galaxies: subplot_galaxies: true subplot_galaxy_images: true - subplot_galaxies_1d: true - subplot_galaxies_1d_decomposed: true fits_galaxy_images: true # Output a .fits file containing images of every galaxy? positions: image_with_positions: true diff --git a/test_autogalaxy/conftest.py b/test_autogalaxy/conftest.py index 05aea7a6b..a017b64e7 100644 --- a/test_autogalaxy/conftest.py +++ b/test_autogalaxy/conftest.py @@ -1,7 +1,6 @@ -import jax.numpy as jnp - - def pytest_configure(): + import jax.numpy as jnp + _ = jnp.sum(jnp.array([0.0])) # Force backend init diff --git a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py index eddee94c0..af14dc927 100644 --- a/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py +++ b/test_autogalaxy/galaxy/plot/test_galaxies_plotter.py @@ -73,11 +73,3 @@ def test__galaxies_sub_plot_output(galaxies_x2_7x7, grid_2d_7x7, plot_path, plot plotter.subplot_galaxy_images() assert path.join(plot_path, "subplot_galaxy_images.png") in plot_patch.paths - - plotter.subplot_galaxies_1d() - assert path.join(plot_path, "subplot_galaxies_1d.png") in plot_patch.paths - - plotter.subplot_galaxies_1d_decomposed() - assert ( - path.join(plot_path, "subplot_galaxies_1d_decomposed.png") in plot_patch.paths - ) diff --git a/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py b/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py index e03fd3f29..e14916496 100644 --- a/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py +++ b/test_autogalaxy/galaxy/plot/test_galaxy_plotters.py @@ -14,106 +14,6 @@ def make_galaxy_plotter_setup(): ) -def test__figures_1d__all_are_output( - gal_x1_lp_x1_mp, grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch -): - galaxy_plotter = aplt.GalaxyPlotter( - galaxy=gal_x1_lp_x1_mp, - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - galaxy_plotter.figures_1d(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d.png") in plot_patch.paths - - plot_patch.paths = [] - - galaxy_plotter = aplt.GalaxyPDFPlotter( - galaxy_pdf_list=[gal_x1_lp_x1_mp, gal_x1_lp_x1_mp, gal_x1_lp_x1_mp], - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - galaxy_plotter.figures_1d(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d.png") in plot_patch.paths - - -def test__figures_1d_decomposed__all_are_output( - gal_x1_lp_x1_mp, grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch -): - galaxy_plotter = aplt.GalaxyPlotter( - galaxy=gal_x1_lp_x1_mp, - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths - - plot_patch.paths = [] - - galaxy_plotter = aplt.GalaxyPDFPlotter( - galaxy_pdf_list=[gal_x1_lp_x1_mp, gal_x1_lp_x1_mp, gal_x1_lp_x1_mp], - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - - galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths - - -def test__figures_1d_decomposed__light_profiles_different_centres_making_offset_radial_grid( - grid_2d_7x7, mask_2d_7x7, plot_path, plot_patch -): - lp_0 = ag.lp.SersicSph(centre=(0.0, 0.0)) - lp_1 = ag.lp.SersicSph(centre=(1.0, 1.0)) - - mp_0 = ag.mp.IsothermalSph(centre=(0.0, 0.0)) - mp_1 = ag.mp.IsothermalSph(centre=(1.0, 1.0)) - - gal = ag.Galaxy(redshift=0.5, light_0=lp_0, light_1=lp_1, mass_0=mp_0, mass_1=mp_1) - - galaxy_plotter = aplt.GalaxyPlotter( - galaxy=gal, - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths - - lp_0 = ag.lp.SersicSph(centre=(0.0, 0.0)) - lp_1 = ag.lp.SersicSph(centre=(1.0, 1.0)) - - mp_0 = ag.mp.IsothermalSph(centre=(0.0, 0.0)) - mp_1 = ag.mp.IsothermalSph(centre=(1.0, 1.0)) - - gal_0 = ag.Galaxy(redshift=0.5, light_0=lp_0, mass_0=mp_0) - gal_1 = ag.Galaxy(redshift=0.5, light_1=lp_1, mass_1=mp_1) - - galaxy_plotter = aplt.GalaxyPDFPlotter( - galaxy_pdf_list=[gal_0, gal_1], - grid=grid_2d_7x7, - mat_plot_1d=aplt.MatPlot1D(output=aplt.Output(plot_path, format="png")), - ) - galaxy_plotter.figures_1d_decomposed(image=True, convergence=True, potential=True) - - assert path.join(plot_path, "image_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "convergence_1d_decomposed.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d_decomposed.png") in plot_patch.paths - - def test__figures_2d__all_are_output( gal_x1_lp_x1_mp, grid_2d_7x7, diff --git a/test_autogalaxy/galaxy/test_galaxy.py b/test_autogalaxy/galaxy/test_galaxy.py index bb1098104..6922a1a41 100644 --- a/test_autogalaxy/galaxy/test_galaxy.py +++ b/test_autogalaxy/galaxy/test_galaxy.py @@ -34,17 +34,6 @@ def test__cls_list_from(lp_0, lp_linear_0): assert cls_list == [lp_linear_0, lp_linear_0] -def test__image_1d_from(lp_0, lp_1, gal_x2_lp): - grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55]]], pixel_scales=1.0) - - lp_image = lp_0.image_1d_from(grid=grid) - lp_image += lp_1.image_1d_from(grid=grid) - - gal_image = gal_x2_lp.image_1d_from(grid=grid) - - assert lp_image == gal_image - - def test__image_2d_from(grid_2d_7x7, gal_x2_lp): lp_0_image = gal_x2_lp.light_profile_0.image_2d_from(grid=grid_2d_7x7) lp_1_image = gal_x2_lp.light_profile_1.image_2d_from(grid=grid_2d_7x7) @@ -105,32 +94,6 @@ def test__luminosity_within_circle(lp_0, lp_1, gal_x2_lp): assert gal_no_lp.luminosity_within_circle_from(radius=1.0) == None -def test__convergence_1d_from(grid_1d_7, mp_0, gal_x1_mp, mp_1, gal_x2_mp): - grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55], [2.05, -0.55]]], pixel_scales=1.0) - - mp_convergence = mp_0.convergence_1d_from(grid=grid) - mp_convergence += mp_1.convergence_1d_from(grid=grid) - - gal_convergence = gal_x2_mp.convergence_1d_from(grid=grid) - - assert (mp_convergence == gal_convergence).all() - - # Test explicitly for a profile with an offset centre and ellipticity, given the 1D to 2D projections are nasty. - - grid = ag.Grid2D.no_mask(values=[[(1.05, -0.55), (2.05, -0.55)]], pixel_scales=1.0) - - elliptical_mp = ag.mp.Isothermal( - centre=(0.5, 1.0), ell_comps=(0.2, 0.3), einstein_radius=1.0 - ) - - galaxy = ag.Galaxy(redshift=0.5, mass=elliptical_mp) - - mp_convergence = elliptical_mp.convergence_1d_from(grid=grid) - gal_convergence = galaxy.convergence_1d_from(grid=grid) - - assert (mp_convergence == gal_convergence).all() - - def test__convergence_2d_from(grid_2d_7x7, mp_0, gal_x1_mp, mp_1, gal_x2_mp): mp_0_convergence = gal_x2_mp.mass_profile_0.convergence_2d_from(grid=grid_2d_7x7) @@ -143,32 +106,6 @@ def test__convergence_2d_from(grid_2d_7x7, mp_0, gal_x1_mp, mp_1, gal_x2_mp): assert gal_convergence[0] == mp_convergence[0] -def test__potential_1d_from(grid_1d_7, mp_0, gal_x1_mp, mp_1, gal_x2_mp): - grid = ag.Grid2D.no_mask(values=[[[1.05, -0.55]]], pixel_scales=1.0) - - mp_potential = mp_0.potential_1d_from(grid=grid) - mp_potential += mp_1.potential_1d_from(grid=grid) - - gal_convergence = gal_x2_mp.potential_1d_from(grid=grid) - - assert mp_potential == gal_convergence - - # Test explicitly for a profile with an offset centre and ellipticity, given the 1D to 2D projections are nasty. - - grid = ag.Grid2D.no_mask(values=[[(1.05, -0.55), (2.05, -0.55)]], pixel_scales=1.0) - - elliptical_mp = ag.mp.Isothermal( - centre=(0.5, 1.0), ell_comps=(0.2, 0.3), einstein_radius=1.0 - ) - - galaxy = ag.Galaxy(redshift=0.5, mass=elliptical_mp) - - mp_potential = elliptical_mp.potential_1d_from(grid=grid) - gal_mp_potential = galaxy.potential_1d_from(grid=grid) - - assert (mp_potential == gal_mp_potential).all() - - def test__potential_2d_from(grid_2d_7x7, gal_x2_mp): mp_0_potential = gal_x2_mp.mass_profile_0.potential_2d_from(grid=grid_2d_7x7) diff --git a/test_autogalaxy/galaxy/test_stellar_dark_decomp.py b/test_autogalaxy/galaxy/test_stellar_dark_decomp.py deleted file mode 100644 index 616fce0ee..000000000 --- a/test_autogalaxy/galaxy/test_stellar_dark_decomp.py +++ /dev/null @@ -1,158 +0,0 @@ -import autogalaxy as ag - -from autogalaxy import exc - -import pytest -import warnings - - -def test__stellar_mass_angular_within_galaxy__is_sum_of_individual_profiles( - smp_0, smp_1 -): - galaxy = ag.Galaxy( - redshift=0.5, - stellar_0=smp_0, - non_stellar_profile=ag.mp.Isothermal(einstein_radius=1.0), - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=0.5) - - gal_mass = decomp.stellar_mass_angular_within_circle_from(radius=0.5) - - assert stellar_mass_0 == gal_mass - - galaxy = ag.Galaxy( - redshift=0.5, - stellar_0=smp_0, - stellar_1=smp_1, - non_stellar_profile=ag.mp.Isothermal(einstein_radius=1.0), - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=0.5) - - gal_mass = decomp.stellar_mass_angular_within_circle_from(radius=0.5) - - assert stellar_mass_0 + stellar_mass_1 == gal_mass - - galaxy = ag.Galaxy(redshift=0.5) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - with pytest.raises(exc.GalaxyException): - decomp.stellar_mass_angular_within_circle_from(radius=1.0) - - -@pytest.mark.filterwarnings("ignore") -def test__stellar_fraction_at_radius(dmp_0, dmp_1, smp_0, smp_1): - galaxy = ag.Galaxy(redshift=0.5, stellar_0=smp_0, dark_0=dmp_0) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=1.0) - dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=1.0) - - stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0) - - assert stellar_fraction == pytest.approx( - stellar_mass_0 / (dark_mass_0 + stellar_mass_0), 1.0e-4 - ) - - galaxy = ag.Galaxy(redshift=0.5, stellar_0=smp_0, stellar_1=smp_1, dark_0=dmp_0) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0) - stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=1.0) - - assert stellar_fraction == pytest.approx( - (stellar_mass_0 + stellar_mass_1) - / (dark_mass_0 + stellar_mass_0 + stellar_mass_1), - 1.0e-4, - ) - - galaxy = ag.Galaxy( - redshift=0.5, stellar_0=smp_0, stellar_1=smp_1, dark_0=dmp_0, dark_mass_1=dmp_1 - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_fraction = decomp.stellar_fraction_at_radius_from(radius=1.0) - dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=1.0) - - assert stellar_fraction == pytest.approx( - (stellar_mass_0 + stellar_mass_1) - / (dark_mass_0 + dark_mass_1 + stellar_mass_0 + stellar_mass_1), - 1.0e-4, - ) - - -@pytest.mark.filterwarnings("ignore") -def test__dark_mass_within_galaxy__is_sum_of_individual_profiles(dmp_0, dmp_1): - galaxy = ag.Galaxy( - redshift=0.5, - dark_0=dmp_0, - non_dark_profile=ag.mp.Isothermal(einstein_radius=1.0), - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=0.5) - - gal_mass = decomp.dark_mass_angular_within_circle_from(radius=0.5) - - assert dark_mass_0 == gal_mass - - galaxy = ag.Galaxy( - redshift=0.5, - dark_0=dmp_0, - dark_1=dmp_1, - non_dark_profile=ag.mp.Isothermal(einstein_radius=1.0), - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=0.5) - - gal_mass = decomp.dark_mass_angular_within_circle_from(radius=0.5) - - assert dark_mass_0 + dark_mass_1 == gal_mass - - galaxy = ag.Galaxy(redshift=0.5) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - with pytest.raises(exc.GalaxyException): - decomp.dark_mass_angular_within_circle_from(radius=1.0) - - -@pytest.mark.filterwarnings("ignore") -def test__dark_fraction_at_radius(dmp_0, dmp_1, smp_0, smp_1): - galaxy = ag.Galaxy(redshift=0.5, dark_0=dmp_0, stellar_0=smp_0) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - stellar_mass_0 = smp_0.mass_angular_within_circle_from(radius=1.0) - dark_mass_0 = dmp_0.mass_angular_within_circle_from(radius=1.0) - - dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0) - - assert dark_fraction == dark_mass_0 / (stellar_mass_0 + dark_mass_0) - - galaxy = ag.Galaxy(redshift=0.5, dark_0=dmp_0, dark_1=dmp_1, stellar_0=smp_0) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0) - dark_mass_1 = dmp_1.mass_angular_within_circle_from(radius=1.0) - - assert dark_fraction == pytest.approx( - (dark_mass_0 + dark_mass_1) / (stellar_mass_0 + dark_mass_0 + dark_mass_1), - 1.0e-4, - ) - - galaxy = ag.Galaxy( - redshift=0.5, dark_0=dmp_0, dark_1=dmp_1, stellar_0=smp_0, stellar_mass_1=smp_1 - ) - decomp = ag.StellarDarkDecomp(galaxy=galaxy) - - dark_fraction = decomp.dark_fraction_at_radius_from(radius=1.0) - stellar_mass_1 = smp_1.mass_angular_within_circle_from(radius=1.0) - - assert dark_fraction == pytest.approx( - (dark_mass_0 + dark_mass_1) - / (stellar_mass_0 + stellar_mass_1 + dark_mass_0 + dark_mass_1), - 1.0e-4, - ) diff --git a/test_autogalaxy/imaging/test_simulator.py b/test_autogalaxy/imaging/test_simulator.py index cc2fc99fe..64a64707c 100644 --- a/test_autogalaxy/imaging/test_simulator.py +++ b/test_autogalaxy/imaging/test_simulator.py @@ -105,7 +105,9 @@ def test__simulator__via_galaxies_from(): assert dataset.shape_native == (20, 20) assert dataset.data.native[0, 0] != imaging_via_image.data.native[0, 0] - assert dataset.data.native[10, 10] == imaging_via_image.data.native[10, 10] + assert dataset.data.native[10, 10] == pytest.approx( + imaging_via_image.data.native[10, 10], 1.0e-4 + ) assert dataset.psf == pytest.approx(imaging_via_image.psf, 1.0e-4) assert dataset.noise_map == pytest.approx(imaging_via_image.noise_map, 1.0e-4) diff --git a/test_autogalaxy/operate/test_deflections.py b/test_autogalaxy/operate/test_deflections.py index 67749599d..45cc97509 100644 --- a/test_autogalaxy/operate/test_deflections.py +++ b/test_autogalaxy/operate/test_deflections.py @@ -224,36 +224,38 @@ def test__tangential_critical_curve_list_from(): assert 0.97 < x_centre < 1.03 -def test__tangential_critical_curve_list_from__compare_via_magnification(): - grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2) - - mp = ag.mp.Isothermal( - centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294) - ) - - tangential_critical_curve_via_magnification = critical_curve_via_magnification_from( - mass_profile=mp, grid=grid - )[0] - - tangential_critical_curve_list = mp.tangential_critical_curve_list_from( - grid=grid, - ) - - assert tangential_critical_curve_list[0] == pytest.approx( - tangential_critical_curve_via_magnification, 5e-1 - ) - - tangential_critical_curve_via_magnification = critical_curve_via_magnification_from( - mass_profile=mp, grid=grid - )[0] - - tangential_critical_curve_list = mp.tangential_critical_curve_list_from( - grid=grid, - ) - - assert tangential_critical_curve_list[0] == pytest.approx( - tangential_critical_curve_via_magnification, 5e-1 - ) +# TODO : reinstate one JAX deflections in. + +# def test__tangential_critical_curve_list_from__compare_via_magnification(): +# grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2) +# +# mp = ag.mp.Isothermal( +# centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294) +# ) +# +# tangential_critical_curve_via_magnification = critical_curve_via_magnification_from( +# mass_profile=mp, grid=grid +# )[0] +# +# tangential_critical_curve_list = mp.tangential_critical_curve_list_from( +# grid=grid, +# ) +# +# assert tangential_critical_curve_list[0] == pytest.approx( +# tangential_critical_curve_via_magnification, 5e-1 +# ) +# +# tangential_critical_curve_via_magnification = critical_curve_via_magnification_from( +# mass_profile=mp, grid=grid +# )[0] +# +# tangential_critical_curve_list = mp.tangential_critical_curve_list_from( +# grid=grid, +# ) +# +# assert tangential_critical_curve_list[0] == pytest.approx( +# tangential_critical_curve_via_magnification, 5e-1 +# ) def test__radial_critical_curve_list_from(): @@ -323,24 +325,26 @@ def test__tangential_caustic_list_from(): assert 0.97 < x_centre < 1.03 -def test__tangential_caustic_list_from___compare_via_magnification(): - grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2) - - mp = ag.mp.Isothermal( - centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294) - ) - - tangential_caustic_via_magnification = caustics_via_magnification_from( - mass_profile=mp, grid=grid - )[0] - - tangential_caustic_list = mp.tangential_caustic_list_from( - grid=grid, - ) - - assert sum(tangential_caustic_list[0]) == pytest.approx( - sum(tangential_caustic_via_magnification), 5e-1 - ) +# TODO : Reinstate one JAX defleciton sin. + +# def test__tangential_caustic_list_from___compare_via_magnification(): +# grid = ag.Grid2D.uniform(shape_native=(50, 50), pixel_scales=0.2) +# +# mp = ag.mp.Isothermal( +# centre=(0.0, 0.0), einstein_radius=2, ell_comps=(0.109423, -0.019294) +# ) +# +# tangential_caustic_via_magnification = caustics_via_magnification_from( +# mass_profile=mp, grid=grid +# )[0] +# +# tangential_caustic_list = mp.tangential_caustic_list_from( +# grid=grid, +# ) +# +# assert sum(tangential_caustic_list[0]) == pytest.approx( +# sum(tangential_caustic_via_magnification), 5e-1 +# ) def test__radial_caustic_list_from(): diff --git a/test_autogalaxy/plot/mat_wrap/test_visuals.py b/test_autogalaxy/plot/mat_wrap/test_visuals.py index 1cd20d1ca..3ae7058ff 100644 --- a/test_autogalaxy/plot/mat_wrap/test_visuals.py +++ b/test_autogalaxy/plot/mat_wrap/test_visuals.py @@ -61,10 +61,6 @@ def test__2d__add_critical_curve(gal_x1_mp, grid_2d_7x7): visuals_2d_via.tangential_critical_curves[0] == gal_x1_mp.tangential_critical_curve_list_from(grid=grid_2d_7x7)[0] ).all() - assert ( - visuals_2d_via.radial_critical_curves[0] - == gal_x1_mp.radial_critical_curve_list_from(grid=grid_2d_7x7)[0] - ).all() def test__2d__add_caustic(gal_x1_mp, grid_2d_7x7): diff --git a/test_autogalaxy/profiles/light/standard/test_abstract.py b/test_autogalaxy/profiles/light/standard/test_abstract.py index 8281e187c..058c03677 100644 --- a/test_autogalaxy/profiles/light/standard/test_abstract.py +++ b/test_autogalaxy/profiles/light/standard/test_abstract.py @@ -43,40 +43,6 @@ def test__luminosity_within_centre__compare_to_gridded_calculations(): assert luminosity_grid == pytest.approx(luminosity_integral, 0.02) -def test__image_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities(): - grid_2d = ag.Grid2D.uniform( - shape_native=(5, 5), - pixel_scales=1.0, - over_sample_size=1, - ) - - lp = ag.lp.Gaussian( - centre=(0.0, 0.0), ell_comps=(0.0, 0.0), intensity=1.0, sigma=1.0 - ) - - image_1d = lp.image_1d_from(grid=grid_2d) - image_2d = lp.image_2d_from(grid=grid_2d) - - assert image_1d[0] == pytest.approx(image_2d.native.array[2, 2], 1.0e-4) - assert image_1d[1] == pytest.approx(image_2d.native.array[2, 3], 1.0e-4) - assert image_1d[2] == pytest.approx(image_2d.native.array[2, 4], 1.0e-4) - - lp = ag.lp.Gaussian( - centre=(0.2, 0.2), ell_comps=(0.3, 0.3), intensity=1.0, sigma=1.0 - ) - - image_1d = lp.image_1d_from(grid=grid_2d) - - grid_2d_projected = grid_2d.grid_2d_radial_projected_from( - centre=lp.centre, angle=lp.angle + 90.0 - ) - - image_projected = lp.image_2d_from(grid=grid_2d_projected) - - assert image_1d == pytest.approx(image_projected.array, 1.0e-4) - assert (image_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all() - - def test__decorator__oversample_uniform__numerical_values(gal_x1_lp): mask = ag.Mask2D( mask=[ diff --git a/test_autogalaxy/profiles/light/test_decorators.py b/test_autogalaxy/profiles/light/test_decorators.py index e7916359a..cf27648d1 100644 --- a/test_autogalaxy/profiles/light/test_decorators.py +++ b/test_autogalaxy/profiles/light/test_decorators.py @@ -13,7 +13,7 @@ class MockLightProfile(ag.LightProfile): @check_operated_only def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ): return np.ones(shape=(3, 3)) @@ -21,7 +21,7 @@ def image_2d_from( class MockLightProfileOperated(ag.lp_operated.LightProfileOperated): @check_operated_only def image_2d_from( - self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None + self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None ): return np.ones(shape=(3, 3)) diff --git a/test_autogalaxy/profiles/mass/abstract/test_abstract.py b/test_autogalaxy/profiles/mass/abstract/test_abstract.py index 4e8f817f4..7532e48ea 100644 --- a/test_autogalaxy/profiles/mass/abstract/test_abstract.py +++ b/test_autogalaxy/profiles/mass/abstract/test_abstract.py @@ -203,110 +203,3 @@ def test__regression__centre_of_profile_in_right_place(): assert deflections.native[2, 4, 0] < 0 assert deflections.native[1, 4, 1] > 0 assert deflections.native[1, 3, 1] < 0 - - -def test__decorators__convergence_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities(): - grid_2d = ag.Grid2D.uniform( - shape_native=(5, 5), - pixel_scales=1.0, - over_sample_size=1, - ) - - sie = ag.mp.Isothermal( - centre=(1e-6, 1e-6), ell_comps=(0.0, 0.0), einstein_radius=1.0 - ) - - convergence_1d = sie.convergence_1d_from(grid=grid_2d) - convergence_2d = sie.convergence_2d_from(grid=grid_2d) - - print(convergence_2d.native.array) - - assert convergence_1d[0] == pytest.approx(convergence_2d.native.array[2, 2], 1.0e-4) - assert convergence_1d[1] == pytest.approx(convergence_2d.native.array[2, 3], 1.0e-4) - assert convergence_1d[2] == pytest.approx(convergence_2d.native.array[2, 4], 1.0e-4) - - sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0) - - convergence_1d = sie.convergence_1d_from(grid=grid_2d) - - grid_2d_projected = grid_2d.grid_2d_radial_projected_from( - centre=sie.centre, angle=sie.angle + 90.0 - ) - - convergence_projected = sie.convergence_2d_from(grid=grid_2d_projected) - - assert convergence_1d == pytest.approx(convergence_projected.array, 1.0e-4) - assert (convergence_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all() - - -def test__decorators__convergence_1d_from__grid_2d_irregular_in__returns_1d_quantities(): - grid_2d = ag.Grid2DIrregular(values=[[1.0, 1.0], [2.0, 2.0], [4.0, 4.0]]) - - sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0) - - convergence_1d = sie.convergence_1d_from(grid=grid_2d) - convergence_2d = sie.convergence_2d_from(grid=grid_2d) - - assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4) - assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4) - assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4) - - sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0) - - convergence_1d = sie.convergence_1d_from(grid=grid_2d) - convergence_2d = sie.convergence_2d_from(grid=grid_2d) - - assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4) - assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4) - assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4) - - -def test__decorators__convergence_1d_from__grid_1d_in__returns_1d_quantities_via_projection(): - grid_1d = ag.Grid1D.no_mask(values=[1.0, 2.0, 3.0], pixel_scales=1.0) - - sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0) - - convergence_1d = sie.convergence_1d_from(grid=grid_1d) - convergence_2d = sie.convergence_2d_from(grid=grid_1d) - - assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4) - assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4) - assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4) - - sie = ag.mp.Isothermal(centre=(0.5, 0.5), ell_comps=(0.2, 0.2), einstein_radius=1.0) - - convergence_1d = sie.convergence_1d_from(grid=grid_1d) - - grid_2d_radial = grid_1d.grid_2d_radial_projected_from(angle=sie.angle + 90.0) - - convergence_2d = sie.convergence_2d_from(grid=grid_2d_radial) - - assert convergence_1d[0] == pytest.approx(convergence_2d[0].array, 1.0e-4) - assert convergence_1d[1] == pytest.approx(convergence_2d[1].array, 1.0e-4) - assert convergence_1d[2] == pytest.approx(convergence_2d[2].array, 1.0e-4) - - -def test__decorators__potential_1d_from__grid_2d_in__returns_1d_image_via_projected_quantities(): - grid_2d = ag.Grid2D.uniform(shape_native=(5, 5), pixel_scales=1.0) - - sie = ag.mp.Isothermal(centre=(0.0, 0.0), ell_comps=(0.0, 0.0), einstein_radius=1.0) - - potential_1d = sie.potential_1d_from(grid=grid_2d) - potential_2d = sie.potential_2d_from(grid=grid_2d) - - assert potential_1d[0] == pytest.approx(potential_2d.native.array[2, 2], abs=1.0e-4) - assert potential_1d[1] == pytest.approx(potential_2d.native.array[2, 3], abs=1.0e-4) - assert potential_1d[2] == pytest.approx(potential_2d.native.array[2, 4], abs=1.0e-4) - - sie = ag.mp.Isothermal(centre=(0.2, 0.2), ell_comps=(0.3, 0.3), einstein_radius=1.0) - - potential_1d = sie.potential_1d_from(grid=grid_2d) - - grid_2d_projected = grid_2d.grid_2d_radial_projected_from( - centre=sie.centre, angle=sie.angle + 90.0 - ) - - potential_projected = sie.potential_2d_from(grid=grid_2d_projected) - - assert potential_1d == pytest.approx(potential_projected.array, abs=1.0e-4) - assert (potential_1d.grid_radial == np.array([0.0, 1.0, 2.0])).all() diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py index 7d8465412..143f6f108 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py @@ -45,9 +45,9 @@ def test__mass_and_concentration_consistent_with_normal_nfw(): assert mp.centre == (1.0, 2.0) - assert mp.axis_ratio == 1.0 + assert mp.axis_ratio() == 1.0 - assert mp.angle == 0.0 + assert mp.angle() == 0.0 assert mp.inner_slope == 1.0 @@ -93,8 +93,8 @@ def test__mass_and_concentration_consistent_with_normal_nfw__scatter_0(): assert concentration_via_kappa_s == concentration_via_mass assert mp.centre == (1.0, 2.0) - assert mp.axis_ratio == 1.0 - assert mp.angle == 0.0 + assert mp.axis_ratio() == 1.0 + assert mp.angle() == 0.0 assert mp.inner_slope == 1.0 assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4) @@ -148,8 +148,8 @@ def test__same_as_above_but_elliptical(): axis_ratio, angle = ag.convert.axis_ratio_and_angle_from(ell_comps=(0.1, 0.2)) - assert mp.axis_ratio == axis_ratio - assert mp.angle == angle + assert mp.axis_ratio() == axis_ratio + assert mp.angle() == angle assert mp.inner_slope == 1.0 assert mp.scale_radius == pytest.approx(0.211578, 1.0e-4) @@ -206,8 +206,8 @@ def test__same_as_above_but_generalized_elliptical(): axis_ratio, angle = ag.convert.axis_ratio_and_angle_from(ell_comps=(0.1, 0.2)) - assert mp.axis_ratio == axis_ratio - assert mp.angle == angle + assert mp.axis_ratio() == axis_ratio + assert mp.angle() == angle assert mp.inner_slope == 2.0 assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4) diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py b/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py index a9d972325..91f621a44 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_scatter.py @@ -58,4 +58,4 @@ def test__scatter_is_nonzero(): deflections_sph = mp.deflections_yx_2d_from(grid=grid) deflections_ell = nfw_ell.deflections_yx_2d_from(grid=grid) - assert deflections_sph[0] != pytest.approx(deflections_ell[0].array, 1.0e-4) + assert deflections_sph[0] != pytest.approx(deflections_ell[0], 1.0e-4) diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py index 64adf18fb..99188189b 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py @@ -46,8 +46,8 @@ def test__duffy__mass_and_concentration_consistent_with_normal_truncated_nfw(): assert concentration_via_kappa_s == concentration_via_mass assert mp.centre == (1.0, 2.0) - assert mp.axis_ratio == 1.0 - assert mp.angle == 0.0 + assert mp.axis_ratio() == 1.0 + assert mp.angle() == 0.0 assert mp.inner_slope == 1.0 assert mp.scale_radius == pytest.approx(0.273382, 1.0e-4) @@ -93,8 +93,8 @@ def test__ludlow__mass_and_concentration_consistent_with_normal_truncated_nfw__s assert concentration_via_kappa_s == concentration_via_mass assert mp.centre == (1.0, 2.0) - assert mp.axis_ratio == 1.0 - assert mp.angle == 0.0 + assert mp.axis_ratio() == 1.0 + assert mp.angle() == 0.0 assert mp.inner_slope == 1.0 assert mp.scale_radius == pytest.approx(0.21157, 1.0e-4) diff --git a/test_autogalaxy/profiles/mass/total/test_power_law_broken.py b/test_autogalaxy/profiles/mass/total/test_power_law_broken.py index ea27bc8de..d9859a0e3 100644 --- a/test_autogalaxy/profiles/mass/total/test_power_law_broken.py +++ b/test_autogalaxy/profiles/mass/total/test_power_law_broken.py @@ -179,7 +179,7 @@ def test__deflections_yx_2d_from__compare_to_power_law(): power_law_yx_ratio = deflections[0, 0] / deflections[0, 1] - assert broken_yx_ratio == pytest.approx(power_law_yx_ratio.array, 1.0e-4) + assert broken_yx_ratio == pytest.approx(power_law_yx_ratio, 1.0e-4) mp = ag.mp.PowerLawBrokenSph( centre=(0, 0), @@ -201,4 +201,4 @@ def test__deflections_yx_2d_from__compare_to_power_law(): power_law_yx_ratio = deflections[0, 0] / deflections[0, 1] - assert broken_yx_ratio == pytest.approx(power_law_yx_ratio.array, 1.0e-4) + assert broken_yx_ratio == pytest.approx(power_law_yx_ratio, 1.0e-4) diff --git a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py index e39e6cf77..851860afa 100644 --- a/test_autogalaxy/profiles/plot/test_light_profile_plotters.py +++ b/test_autogalaxy/profiles/plot/test_light_profile_plotters.py @@ -15,46 +15,6 @@ def make_profile_plotter_setup(): ) -def test__figures_1d__all_are_output( - lp_0, - lp_1, - grid_2d_7x7, - grid_2d_irregular_7x7_list, - plot_path, - plot_patch, -): - mat_plot_1d = aplt.MatPlot1D( - half_light_radius_axvline=aplt.HalfLightRadiusAXVLine(color="r"), - einstein_radius_axvline=aplt.EinsteinRadiusAXVLine(color="r"), - output=aplt.Output(plot_path, format="png"), - ) - - light_profile_plotter = aplt.LightProfilePlotter( - light_profile=lp_0, - grid=grid_2d_7x7, - mat_plot_1d=mat_plot_1d, - ) - - light_profile_plotter.figures_1d(image=True) - - assert path.join(plot_path, "image_1d.png") in plot_patch.paths - - plot_patch.paths = [] - - lp_offset_centre = ag.lp.SersicSph(centre=(5.0, 5.0)) - - light_profile_plotter = aplt.LightProfilePDFPlotter( - light_profile_pdf_list=[lp_0, lp_1, lp_0, lp_1, lp_0, lp_offset_centre], - grid=grid_2d_7x7, - mat_plot_1d=mat_plot_1d, - sigma=2.0, - ) - - light_profile_plotter.figures_1d(image=True) - - assert path.join(plot_path, "image_1d_pdf.png") in plot_patch.paths - - def test__figures_2d__all_are_output( lp_0, grid_2d_7x7, diff --git a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py index 2b1b9bbe5..304c050c8 100644 --- a/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py +++ b/test_autogalaxy/profiles/plot/test_mass_profile_plotters.py @@ -14,47 +14,6 @@ def make_mp_plotter_setup(): ) -def test__figures_1d__all_are_output( - mp_0, - mp_1, - grid_2d_7x7, - grid_2d_irregular_7x7_list, - plot_path, - plot_patch, -): - mat_plot_1d = aplt.MatPlot1D( - half_light_radius_axvline=aplt.HalfLightRadiusAXVLine(color="r"), - einstein_radius_axvline=aplt.EinsteinRadiusAXVLine(color="r"), - output=aplt.Output(plot_path, format="png"), - ) - - mass_profile_plotter = aplt.MassProfilePlotter( - mass_profile=mp_0, - grid=grid_2d_7x7, - mat_plot_1d=mat_plot_1d, - ) - - mass_profile_plotter.figures_1d(convergence=True, potential=True) - - assert path.join(plot_path, "convergence_1d.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d.png") in plot_patch.paths - - plot_patch.paths = [] - - mp_offset_centre = ag.mp.IsothermalSph(centre=(5.0, 5.0)) - - mass_profile_plotter = aplt.MassProfilePDFPlotter( - mass_profile_pdf_list=[mp_0, mp_1, mp_0, mp_1, mp_0, mp_offset_centre], - grid=grid_2d_7x7, - mat_plot_1d=mat_plot_1d, - ) - - mass_profile_plotter.figures_1d(convergence=True, potential=True) - - assert path.join(plot_path, "convergence_1d_pdf.png") in plot_patch.paths - assert path.join(plot_path, "potential_1d_pdf.png") in plot_patch.paths - - def test__figures_2d__all_are_output( mp_0, grid_2d_7x7,