import json
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Type, Union
import autoarray as aa
from autogalaxy.galaxy.galaxy import Galaxy
from autogalaxy.profiles.light.basis import Basis
from autogalaxy.profiles.light.abstract import LightProfile
from autogalaxy.profiles.light.linear import LightProfileLinear
from autogalaxy.profiles.light.snr import LightProfileSNR
from autogalaxy.operate.image import OperateImageGalaxies
from autogalaxy.operate.deflections import OperateDeflections
from autogalaxy import exc
from autogalaxy.plane import plane_util
[docs]class Plane(OperateImageGalaxies, OperateDeflections):
def __init__(
self,
galaxies,
redshift: Optional[float] = None,
run_time_dict: Optional[Dict] = None,
):
"""
A plane of galaxies where all galaxies are at the same redshift.
Parameters
----------
redshift or None
The redshift of the plane.
galaxies : [Galaxy]
The list of galaxies in this plane.
"""
if redshift is None:
if not galaxies:
raise exc.PlaneException(
"No redshift and no galaxies were input to a Plane. A redshift for the Plane therefore cannot be"
"determined"
)
elif not all(
[galaxies[0].redshift == galaxy.redshift for galaxy in galaxies]
):
redshift = np.mean([galaxy.redshift for galaxy in galaxies])
else:
redshift = galaxies[0].redshift
self.redshift = redshift
self.galaxies = galaxies
self.run_time_dict = run_time_dict
self.run_time_dict = run_time_dict
def output_to_json(self, file_path: Union[Path, str]):
with open(file_path, "w+") as f:
json.dump(self.dict(), f, indent=4)
@property
def galaxy_redshifts(self) -> List[float]:
return [galaxy.redshift for galaxy in self.galaxies]
def has(self, cls: Union[Type, Tuple[Type]]) -> bool:
if self.galaxies is not None:
return any(list(map(lambda galaxy: galaxy.has(cls=cls), self.galaxies)))
return False
[docs] def cls_list_from(self, cls: Type) -> List:
"""
Returns a list of objects in the plane which are an instance of the input `cls`.
For example:
- If the input is `cls=ag.LightProfile`, a list containing all light profiles in the plane is returned.
Returns
-------
The list of objects in the plane that inherit from input `cls`.
"""
cls_list = []
for galaxy in self.galaxies:
if galaxy.has(cls=cls):
for cls_galaxy in galaxy.cls_list_from(cls=cls):
cls_list.append(cls_galaxy)
return cls_list
def galaxies_with_cls_list_from(self, cls: Type) -> List[Galaxy]:
return list(filter(lambda galaxy: galaxy.has(cls=cls), self.galaxies))
@property
def adapt_galaxies_with_pixelization_image_list(self) -> List[aa.Array2D]:
return [
galaxy.adapt_galaxy_image
for galaxy in self.galaxies_with_cls_list_from(cls=aa.Pixelization)
]
[docs] @aa.grid_dec.grid_2d_to_structure
def image_2d_from(
self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
) -> aa.Array2D:
"""
Returns the profile-image plane image of the list of galaxies of the plane's sub-grid, by summing the
individual images of each galaxy's light profile.
If the `operated_only` input is included, the function omits light profiles which are parents of
the `LightProfileOperated` object, which signifies that the light profile represents emission that has
already had the instrument operations (e.g. PSF convolution, a Fourier transform) applied to it.
If the plane has no galaxies (or no galaxies have mass profiles) an arrays of all zeros the shape of the plane's
sub-grid is returned.
Parameters
----------
grid
The 2D (y, x) coordinates where values of the image are evaluated.
operated_only
By default, the image is the sum of light profile images (irrespective of whether they have been operatd on
or not). If this input is included as a bool, only images which are or are not already operated are summed
and returned.
"""
if self.galaxies:
return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
return np.zeros((grid.shape[0],))
def image_2d_list_from(
self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
) -> List[aa.Array2D]:
return [
galaxy.image_2d_from(grid=grid, operated_only=operated_only)
for galaxy in self.galaxies
]
[docs] def galaxy_image_2d_dict_from(
self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
) -> {Galaxy: np.ndarray}:
"""
Returns a dictionary associating every `Galaxy` object in the `Plane` with its corresponding 2D image, using
the instance of each galaxy as the dictionary keys.
This object is used for adapt features, which use the image of each galaxy in a model-fit in order to
adapt quantities like a pixelization or regularization scheme to the surface brightness of the galaxies being
fitted.
By inheriting from `OperateImageGalaxies` functions which apply operations of this dictionary are accessible,
for example convolving every image with a PSF or applying a Fourier transform to create a galaxy-visibilities
dictionary.
Parameters
----------
grid
The 2D (y,x) coordinates of the (masked) grid, in its original geometric reference frame.
Returns
-------
A dictionary associated every galaxy in the plane with its corresponding 2D image.
"""
galaxy_image_2d_dict = dict()
image_2d_list = self.image_2d_list_from(grid=grid, operated_only=operated_only)
for galaxy_index, galaxy in enumerate(self.galaxies):
galaxy_image_2d_dict[galaxy] = image_2d_list[galaxy_index]
return galaxy_image_2d_dict
def plane_image_2d_from(
self, grid: aa.type.Grid2DLike, zoom_to_brightest: bool = True
) -> aa.Array2D:
return plane_util.plane_image_from(
galaxies=self.galaxies,
grid=grid.mask.derive_grid.all_false_sub_1,
zoom_to_brightest=zoom_to_brightest,
)
@aa.grid_dec.grid_2d_to_vector_yx
@aa.grid_dec.grid_2d_to_structure
def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
if self.galaxies:
return sum(
map(lambda g: g.deflections_yx_2d_from(grid=grid), self.galaxies)
)
return np.zeros(shape=(grid.shape[0], 2))
[docs] @aa.grid_dec.grid_2d_to_structure
def convergence_2d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
"""
Returns the convergence of the list of galaxies of the plane's sub-grid, by summing the individual convergences \
of each galaxy's mass profile.
The convergence is calculated on the sub-grid and binned-up to the original grid by taking the mean
value of every set of sub-pixels, provided the *returned_binned_sub_grid* bool is `True`.
If the plane has no galaxies (or no galaxies have mass profiles) an arrays of all zeros the shape of the plane's
sub-grid is returned.
Internally data structures are treated as ndarrays, however the decorator `grid_2d_to_structure` converts
the output to an `Array2D` using the input `grid`'s attributes.
Parameters
----------
grid : Grid2D
The grid (or sub) of (y,x) arc-second coordinates at the centre of every unmasked pixel which the \
potential is calculated on.
galaxies : [Galaxy]
The galaxies whose mass profiles are used to compute the surface densities.
"""
if self.galaxies:
return sum(map(lambda g: g.convergence_2d_from(grid=grid), self.galaxies))
return np.zeros((grid.shape[0],))
[docs] @aa.grid_dec.grid_2d_to_structure
def potential_2d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
"""
Returns the potential of the list of galaxies of the plane's sub-grid, by summing the individual potentials \
of each galaxy's mass profile.
The potential is calculated on the sub-grid and binned-up to the original grid by taking the mean
value of every set of sub-pixels, provided the *returned_binned_sub_grid* bool is `True`.
If the plane has no galaxies (or no galaxies have mass profiles) an arrays of all zeros the shape of the plane's
sub-grid is returned.
Internally data structures are treated as ndarrays, however the decorator `grid_2d_to_structure` converts
the output to an `Array2D` using the input `grid`'s attributes.
Parameters
----------
grid : Grid2D
The grid (or sub) of (y,x) arc-second coordinates at the centre of every unmasked pixel which the \
potential is calculated on.
galaxies : [Galaxy]
The galaxies whose mass profiles are used to compute the surface densities.
"""
if self.galaxies:
return sum(map(lambda g: g.potential_2d_from(grid=grid), self.galaxies))
return np.zeros((grid.shape[0],))
[docs] @aa.grid_dec.grid_2d_to_structure
def traced_grid_from(self, grid: aa.type.Grid2DLike) -> 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)
@property
def perform_inversion(self) -> bool:
"""
Returns a bool specifying whether this fit object performs an inversion.
This is based on whether any of the galaxies in the `model_obj` have a `Pixelization` or `LightProfileLinear`
object, in which case an inversion is performed.
Returns
-------
A bool which is True if an inversion is performed.
"""
if self.has(cls=(aa.Pixelization, LightProfileLinear)):
return True
elif self.has(cls=Basis):
basis_list = self.cls_list_from(cls=Basis)
for basis in basis_list:
for light_profile in basis.light_profile_list:
if isinstance(light_profile, LightProfileLinear):
return True
return False
[docs] def set_snr_of_snr_light_profiles(
self,
grid: aa.type.Grid2DLike,
exposure_time: float,
background_sky_level: float = 0.0,
psf: Optional[aa.Kernel2D] = None,
):
"""
Iterate over every `LightProfileSNR` in the plane and set their `intensity` values to values which give
their input `signal_to_noise_ratio` value, which is performed as follows:
- Evaluate the image of each light profile on the input grid.
- Blur this image with a PSF, if included.
- Take the value of the brightest pixel.
- Use an input `exposure_time` and `background_sky` (e.g. from the `SimulatorImaging` object) to determine
what value of `intensity` gives the desired signal to noise ratio for the image.
Parameters
----------
grid
The (y, x) coordinates in the original reference frame of the grid.
exposure_time
The exposure time of the simulated imaging.
background_sky_level
The level of the background sky of the simulated imaging.
psf
The psf of the simulated imaging which can change the S/N of the light profile due to spreading out
the emission.
"""
for galaxy in self.galaxies:
for light_profile in galaxy.cls_list_from(cls=LightProfile):
if isinstance(light_profile, LightProfileSNR):
light_profile.set_intensity_from(
grid=grid,
exposure_time=exposure_time,
background_sky_level=background_sky_level,
psf=psf,
)