Source code for autolens.lens.plot.ray_tracing_plotters

from typing import Optional, List

import autoarray as aa
import autogalaxy as ag
import autogalaxy.plot as aplt

from autogalaxy.plot.mass_plotter import MassPlotter

from autolens.plot.abstract_plotters import Plotter
from autolens.lens.ray_tracing import Tracer


[docs]class TracerPlotter(Plotter):
[docs] def __init__( self, tracer: Tracer, grid: aa.type.Grid2DLike, mat_plot_1d: aplt.MatPlot1D = aplt.MatPlot1D(), visuals_1d: aplt.Visuals1D = aplt.Visuals1D(), include_1d: aplt.Include1D = aplt.Include1D(), mat_plot_2d: aplt.MatPlot2D = aplt.MatPlot2D(), visuals_2d: aplt.Visuals2D = aplt.Visuals2D(), include_2d: aplt.Include2D = aplt.Include2D(), ): """ Plots the attributes of `Tracer` objects using the matplotlib methods `plot()` and `imshow()` and many other matplotlib functions which customize the plot's appearance. 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, if the corresponding entry is `True` in the `Include1D` or `Include2D` object or the `config/visualize/include.ini` file. Parameters ---------- tracer The tracer the plotter plots. grid The 2D (y,x) grid of coordinates used to evaluate the tracer's light and mass 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. include_1d Specifies which attributes of the `MassProfile` are extracted and plotted as visuals for 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. include_2d Specifies which attributes of the `MassProfile` are extracted and plotted as visuals for 2D plots. """ super().__init__( mat_plot_1d=mat_plot_1d, visuals_1d=visuals_1d, include_1d=include_1d, mat_plot_2d=mat_plot_2d, include_2d=include_2d, visuals_2d=visuals_2d, ) self.tracer = tracer self.grid = grid self._mass_plotter = MassPlotter( mass_obj=self.tracer, grid=self.grid, get_visuals_2d=self.get_visuals_2d, mat_plot_2d=self.mat_plot_2d, include_2d=self.include_2d, visuals_2d=self.visuals_2d, )
def get_visuals_2d(self) -> aplt.Visuals2D: return self.get_visuals_2d_of_plane(plane_index=0) def get_visuals_2d_of_plane(self, plane_index: int) -> aplt.Visuals2D: return self.get_2d.via_tracer_from( tracer=self.tracer, grid=self.grid, plane_index=plane_index ) def plane_plotter_from(self, plane_index: int) -> aplt.PlanePlotter: """ Returns an `PlanePlotter` corresponding to a `Plane` in the `Tracer`. Returns ------- plane_index The index of the plane in the `Tracer` used to make the `PlanePlotter`. """ plane_grid = self.tracer.traced_grid_2d_list_from(grid=self.grid)[plane_index] return aplt.PlanePlotter( plane=self.tracer.planes[plane_index], grid=plane_grid, mat_plot_2d=self.mat_plot_2d, visuals_2d=self.get_visuals_2d_of_plane(plane_index=plane_index), include_2d=self.include_2d, ) def figures_2d( self, image: bool = False, source_plane: bool = False, convergence: bool = False, potential: bool = False, deflections_y: bool = False, deflections_x: bool = False, magnification: bool = False, contribution_map: bool = False, ): """ Plots the individual attributes of the plotter's `Tracer` object in 2D, which are computed via the plotter's 2D grid object. The API is such that every plottable attribute of the `Tracer` object is an input parameter of type bool of the function, which if switched to `True` means that it is plotted. Parameters ---------- image Whether or not to make a 2D plot (via `imshow`) of the image of tracer in its image-plane (e.g. after lensing). source_plane Whether or not to make a 2D plot (via `imshow`) of the image of the tracer in the source-plane (e.g. its unlensed light). convergence Whether or not to make a 2D plot (via `imshow`) of the convergence. potential Whether or not to make a 2D plot (via `imshow`) of the potential. deflections_y Whether or not to make a 2D plot (via `imshow`) of the y component of the deflection angles. deflections_x Whether or not to make a 2D plot (via `imshow`) of the x component of the deflection angles. magnification Whether or not to make a 2D plot (via `imshow`) of the magnification. contribution_map Whether or not to make a 2D plot (via `imshow`) of the contribution map. """ if image: self.mat_plot_2d.plot_array( array=self.tracer.image_2d_from(grid=self.grid), visuals_2d=self.get_visuals_2d(), auto_labels=aplt.AutoLabels(title="Image", filename="image_2d"), ) if source_plane: self.figures_2d_of_planes( plane_image=True, plane_index=len(self.tracer.planes) - 1 ) self._mass_plotter.figures_2d( convergence=convergence, potential=potential, deflections_y=deflections_y, deflections_x=deflections_x, magnification=magnification, ) if contribution_map: self.mat_plot_2d.plot_array( array=self.tracer.contribution_map, visuals_2d=self.get_visuals_2d(), auto_labels=aplt.AutoLabels( title="Contribution Map", filename="contribution_map_2d" ), ) def plane_indexes_from(self, plane_index: Optional[int]) -> List[int]: """ Returns a list of all indexes of the planes in the fit, which is iterated over in figures that plot individual figures of each plane in a tracer. Parameters ---------- plane_index A specific plane index which when input means that only a single plane index is returned. Returns ------- list A list of galaxy indexes corresponding to planes in the plane. """ if plane_index is None: return list(range(len(self.tracer.planes))) return [plane_index] def figures_2d_of_planes( self, plane_image: bool = False, plane_grid: bool = False, plane_index: Optional[int] = None, ): """ Plots source-plane images (e.g. the unlensed light) each individual `Plane` in the plotter's `Tracer` in 2D, which are computed via the plotter's 2D grid object. The API is such that every plottable attribute of the `Plane` object is an input parameter of type bool of the function, which if switched to `True` means that it is plotted. Parameters ---------- plane_image Whether or not to make a 2D plot (via `imshow`) of the image of the plane in the soure-plane (e.g. its unlensed light). plane_grid Whether or not to make a 2D plot (via `scatter`) of the lensed (y,x) coordinates of the plane in the source-plane. plane_index If input, plots for only a single plane based on its index in the tracer are created. """ plane_indexes = self.plane_indexes_from(plane_index=plane_index) for plane_index in plane_indexes: plane_plotter = self.plane_plotter_from(plane_index=plane_index) if plane_image: plane_plotter.figures_2d( plane_image=True, title_suffix=f" Of Plane {plane_index}", filename_suffix=f"_of_plane_{plane_index}", ) if plane_grid: plane_plotter.figures_2d( plane_grid=True, title_suffix=f" Of Plane {plane_index}", filename_suffix=f"_of_plane_{plane_index}", ) def subplot( self, image: bool = False, source_plane: bool = False, convergence: bool = False, potential: bool = False, deflections_y: bool = False, deflections_x: bool = False, magnification: bool = False, contribution_map: bool = False, auto_filename: str = "subplot_tracer", ): """ Plots the individual attributes of the plotter's `Tracer` object in 2D on a subplot, which are computed via the plotter's 2D grid object. The API is such that every plottable attribute of the `Tracer` object is an input parameter of type bool of the function, which if switched to `True` means that it is included on the subplot. Parameters ---------- image Whether or not to include a 2D plot (via `imshow`) of the image of tracer in its image-plane (e.g. after lensing). source_plane Whether or not to include a 2D plot (via `imshow`) of the image of the tracer in the source-plane (e.g. its unlensed light). convergence Whether or not to include a 2D plot (via `imshow`) of the convergence. potential Whether or not to include a 2D plot (via `imshow`) of the potential. deflections_y Whether or not to include a 2D plot (via `imshow`) of the y component of the deflection angles. deflections_x Whether or not to include a 2D plot (via `imshow`) of the x component of the deflection angles. magnification Whether or not to include a 2D plot (via `imshow`) of the magnification. contribution_map Whether or not to include a 2D plot (via `imshow`) of the contribution map. auto_filename The default filename of the output subplot if written to hard-disk. """ self._subplot_custom_plot( image=image, source_plane=source_plane, convergence=convergence, potential=potential, deflections_y=deflections_y, deflections_x=deflections_x, magnification=magnification, contribution_map=contribution_map, auto_labels=aplt.AutoLabels(filename=auto_filename), ) def subplot_tracer(self): """ Standard subplot of the attributes of the plotter's `Tracer` object. """ return self.subplot( image=True, source_plane=True, convergence=True, potential=True, deflections_y=True, deflections_x=True, ) def subplot_lensed_images(self): """ Subplot of the lensed image of every plane. For example, for a 2 plane `Tracer`, this creates a subplot with 2 panels, one for the image-plane image and one for the source-plane lensed image. If there are 3 planes, 3 panels are created, showing images at each plane. """ number_subplots = self.tracer.total_planes self.open_subplot_figure(number_subplots=number_subplots) for plane_index in range(0, self.tracer.total_planes): plane_plotter = self.plane_plotter_from(plane_index=plane_index) plane_plotter.figures_2d( image=True, title_suffix=f" Of Plane {plane_index}" ) self.mat_plot_2d.output.subplot_to_figure( auto_filename=f"subplot_lensed_images" ) self.close_subplot_figure() def subplot_plane_images(self): """ Subplot of the image of every plane in its own plane. For example, for a 2 plane `Tracer`, this creates a subplot with 2 panels, one for the image-plane image and one for the source-plane (e.g. unlensed) image. If there are 3 planes, 3 panels are created, showing images at each plane. """ number_subplots = 2 * self.tracer.total_planes - 1 self.open_subplot_figure(number_subplots=number_subplots) plane_plotter = self.plane_plotter_from(plane_index=0) plane_plotter.figures_2d(image=True, title_suffix=" Of Plane 0") self.mat_plot_2d.subplot_index += 1 for plane_index in range(1, self.tracer.total_planes): plane_plotter = self.plane_plotter_from(plane_index=plane_index) plane_plotter.figures_2d( image=True, title_suffix=f" Of Plane {plane_index}" ) plane_plotter.figures_2d( plane_image=True, title_suffix=f" Of Plane {plane_index}" ) self.mat_plot_2d.output.subplot_to_figure(auto_filename=f"subplot_plane_images") self.close_subplot_figure()