Merge pull request #268 from Jammy2211/feature/sub_iterate_tracer

Feature/sub iterate tracer

autolens/imaging/fit_imaging.py

@@ -131,13 +131,15 @@ def tracer_to_inversion(self) -> TracerToInversion:
             border_relocator=self.dataset.border_relocator,
         )
 
+
         return TracerToInversion(
             dataset=dataset,
             tracer=self.tracer,
             sky=self.sky,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            run_time_dict=self.run_time_dict
         )
 
     @cached_property

autolens/lens/to_inversion.py

@@ -109,6 +109,7 @@ def lp_linear_func_list_galaxy_dict(
                 sky=self.sky,
                 settings_inversion=self.settings_inversion,
                 adapt_images=self.adapt_images,
+                run_time_dict=self.run_time_dict,
             )
 
             lp_linear_galaxy_dict_of_plane = (
@@ -170,6 +171,7 @@ def image_plane_mesh_grid_pg_list(self) -> List[List]:
                 sky=self.sky,
                 adapt_images=self.adapt_images,
                 settings_inversion=self.settings_inversion,
+                run_time_dict=self.run_time_dict,
             )
 
             image_plane_mesh_grid_list = to_inversion.image_plane_mesh_grid_list
@@ -247,6 +249,7 @@ def mapper_galaxy_dict(self) -> Dict[aa.AbstractMapper, ag.Galaxy]:
                     preloads=self.preloads,
                     adapt_images=self.adapt_images,
                     settings_inversion=self.settings_inversion,
+                    run_time_dict=self.run_time_dict,
                 )
 
                 galaxies_with_pixelization_list = galaxies.galaxies_with_cls_list_from(

autolens/lens/tracer.py

@@ -415,14 +415,19 @@ def image_2d_list_from(
         """
         Returns a list of the 2D images for each plane from a 2D grid of Cartesian (y,x) coordinates.
 
-        The image of each plane is computed by summing the images of all galaxies in that plane. If a plane has no
-        galaxies, or if the galaxies in a plane has no light profiles, a numpy array of zeros is returned.
+        The image of each plane is computed by ray-tracing the grid using te mass profiles of each galaxies and then
+        summing the images of all galaxies in that plane. If a plane has no galaxies, or if the galaxies in a plane
+        has no light profiles, a numpy array of zeros is returned.
 
         For example, if the tracer's planes contain galaxies at redshifts z=0.5, z=1.0 and z=2.0, and the galaxies
         at redshifts z=0.5 and z=1.0 have light and mass profiles, the returned list of images will be the image of the
         galaxies at z=0.5 and z=1.0, where the image at redshift z=1.0 will include the lensing effects of the galaxies
         at z=0.5. The image at redshift z=2.0 will be a numpy array of zeros.
 
+        The `plane_index` input is used to return a specific image of a plane, as opposed to a list of images
+        of all planes. This can save on computational time when only the image of a specific plane is needed,
+        and is used to perform iterative over-sampling calculations.
+
         The images output by this function do not include instrument operations, such as PSF convolution (for imaging
         data) or a Fourier transform (for interferometer data).
 
@@ -449,6 +454,12 @@ def image_2d_list_from(
             therefore is used to pass the `operated_only` input to these methods.
         """
 
+        if hasattr(grid, "over_sampling"):
+            if isinstance(grid.over_sampling, aa.OverSamplingIterate):
+                return self.image_2d_list_over_sampled_from(
+                    grid=grid, operated_only=operated_only
+                )
+
         traced_grid_list = self.traced_grid_2d_list_from(
             grid=grid, plane_index_limit=self.upper_plane_index_with_light_profile
         )
@@ -485,6 +496,144 @@ def image_2d_list_from(
 
         return image_2d_list
 
+    @over_sample
+    def image_2d_of_plane_from(
+        self,
+        grid: aa.type.Grid2DLike,
+        plane_index: int,
+        operated_only: Optional[bool] = None,
+    ) -> aa.Array2D:
+        """
+        Returns a 2D image of an input plane from a 2D grid of Cartesian (y,x) coordinates.
+
+        The image of the plane is computed by ray-tracing the grid using te mass profiles of all galaxies before the
+        input plane and then summing the images of all galaxies in that plane. If a plane has no galaxies, or if the
+        galaxies in a plane, has no light profiles, a numpy array of zeros is returned.
+
+        For example, if the tracer's planes contain galaxies at redshifts z=0.5, z=1.0 and z=2.0, and the galaxies
+        at redshifts z=0.5 and z=1.0 have light and mass profiles, the returned image for `plane_index=1` will be the
+        image of the galaxy at z=1.0, where the image at redshift z=1.0 will include the lensing effects of the
+        galaxies at z=0.5. The image at redshift z=2.0 will be ignored.
+
+        The `plane_index` input specifies which plane the image os returned for. This calculation saves computational
+        time compared to `image_2d_list_from` when only the image of a specific plane is needed. It is also used to
+        perform iterative over-sampling calculations.
+
+        The images output by this function do not include instrument operations, such as PSF convolution (for imaging
+        data) or a Fourier transform (for interferometer data).
+
+        Inherited methods in the `autogalaxy.operate.image` package can apply these operations to the images.
+        These functions may have the `operated_only` input passed to them, which is why this function includes
+        the `operated_only` input.
+
+        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 and therefore
+        that operation is not performed again.
+
+        See the `autogalaxy.profiles.light` package for details of how images are computed from a light
+        profile.
+
+        Parameters
+        ----------
+        grid
+            The 2D (y, x) coordinates where values of the image are evaluated.
+        plane_index
+            The plane index of the plane the image is computed.
+        operated_only
+            The returned list from this function contains all light profile images, and they are never operated on
+            (e.g. via the imaging PSF). However, inherited methods in the `autogalaxy.operate.image` package can
+            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.
+        """
+
+        if not self.planes[plane_index].has(cls=ag.LightProfile):
+            if isinstance(grid, aa.Grid2D):
+                return aa.Array2D(values=np.zeros(shape=grid.shape[0]), mask=grid.mask)
+            else:
+                return aa.ArrayIrregular(values=np.zeros(grid.shape[0]))
+
+        traced_grid_list = self.traced_grid_2d_list_from(
+            grid=grid, plane_index_limit=plane_index
+        )
+
+        return sum(
+            [
+                galaxy.image_2d_from(
+                    grid=traced_grid_list[plane_index],
+                    operated_only=operated_only,
+                )
+                for galaxy in self.planes[plane_index]
+            ]
+        )
+
+    def image_2d_list_over_sampled_from(
+        self,
+        grid: aa.type.Grid2DLike,
+        operated_only: Optional[bool] = None,
+    ) -> List[aa.Array2D]:
+        """
+        Returns a list of the 2D images for each plane from a 2D grid of Cartesian (y,x) coordinates where adaptive
+        and iterative over-sampling is used to compute the image.
+
+        The image of each plane is computed by iteratively ray-tracing the grid using the mass profiles of each
+        galaxies and then summing the images of all galaxies in that plane, until a threshold level of accuracy
+        defined by the over-sampling grid is met. If a plane has no galaxies, or if the galaxies in a plane
+        has no light profiles, a numpy array of zeros is returned.
+
+        For example, if the tracer's planes contain galaxies at redshifts z=0.5, z=1.0 and z=2.0, and the galaxies
+        at redshifts z=0.5 and z=1.0 have light and mass profiles, the returned list of images will be the image of the
+        galaxies at z=0.5 and z=1.0, where the image at redshift z=1.0 will include the lensing effects of the galaxies
+        at z=0.5. The image at redshift z=2.0 will be a numpy array of zeros.
+
+        The implementation of this function has to wrap a function in the iterative over sampler which performs the
+        iterative over-sampling calculation. This requires a function to be defined internally in this function
+        which meets the requirements of the over-sample.
+
+        The images output by this function do not include instrument operations, such as PSF convolution (for imaging
+        data) or a Fourier transform (for interferometer data).
+
+        Inherited methods in the `autogalaxy.operate.image` package can apply these operations to the images.
+        These functions may have the `operated_only` input passed to them, which is why this function includes
+        the `operated_only` input.
+
+        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 and therefore
+        that operation is not performed again.
+
+        See the `autogalaxy.profiles.light` package for details of how images are computed from a light
+        profile.
+
+        Parameters
+        ----------
+        grid
+            The 2D (y, x) coordinates where values of the image are evaluated, which has an iterative over-sampling
+            applied to it.
+        operated_only
+            The returned list from this function contains all light profile images, and they are never operated on
+            (e.g. via the imaging PSF). However, inherited methods in the `autogalaxy.operate.image` package can
+            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.
+        """
+
+        image_2d_list = []
+
+        for plane_index in range(len(self.planes)):
+
+            def func(obj, grid, *args, **kwargs):
+                return self.image_2d_of_plane_from(
+                    grid=grid,
+                    operated_only=operated_only,
+                    plane_index=plane_index,
+                )
+
+            image_2d = grid.over_sampler.array_via_func_from(func=func, obj=self)
+
+            image_2d_list.append(image_2d)
+
+        return image_2d_list
+
     @over_sample
     @aa.grid_dec.to_array
     @aa.profile_func

docs/api/data.rst

@@ -19,7 +19,6 @@ grids of (y,x) Cartesian coordinates (which are used for evaluating light profil
    Mask2D
    Array2D
    Grid2D
-   Grid2DIterate
    Grid2DIrregular
 
 Imaging
@@ -54,6 +53,24 @@ a fast Fourier transform to map data to the uv-plane.
    TransformerDFT
    TransformerNUFFT
 
+Over Sampling
+-------------
+
+Calculations using grids approximate a 2D line integral of the light in the galaxy which falls in each image-pixel.
+Different over sampling schemes can be used to efficiently approximate this integral and these objects can be
+applied to datasets to apply over sampling to their fit.
+
+.. currentmodule:: autolens
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: custom-class-template.rst
+   :recursive:
+
+   OverSamplingUniform
+   OverSamplingIterate
+
+
 1D Data Structures
 ------------------
 

test_autolens/imaging/test_simulate_and_fit_imaging.py

@@ -79,6 +79,84 @@ def test__perfect_fit__chi_squared_0():
         shutil.rmtree(file_path)
 
 
+def test__perfect_fit__chi_squared_0__use_grid_iterate_to_simulate_and_fit():
+    over_sampling = al.OverSamplingIterate(fractional_accuracy=0.9999, sub_steps=[2, 4, 8])
+
+    grid = al.Grid2D.uniform(shape_native=(11, 11), pixel_scales=0.2,
+                             over_sampling=over_sampling)
+
+    psf = al.Kernel2D.from_gaussian(
+        shape_native=(3, 3), pixel_scales=0.2, sigma=0.75, normalize=True
+    )
+
+    lens_galaxy = al.Galaxy(
+        redshift=0.5,
+        light=al.lp.Sersic(centre=(0.1, 0.1), intensity=0.1),
+        mass=al.mp.Isothermal(centre=(0.1, 0.1), einstein_radius=1.8),
+    )
+    source_galaxy = al.Galaxy(
+        redshift=1.0, light=al.lp.Exponential(centre=(0.1, 0.1), intensity=0.5)
+    )
+    tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+    dataset = al.SimulatorImaging(exposure_time=300.0, psf=psf, add_poisson_noise=False)
+
+    dataset = dataset.via_tracer_from(tracer=tracer, grid=grid)
+    dataset.noise_map = al.Array2D.ones(
+        shape_native=dataset.data.shape_native, pixel_scales=0.2
+    )
+
+    file_path = path.join(
+        "{}".format(path.dirname(path.realpath(__file__))),
+        "data_temp",
+        "simulate_and_fit",
+    )
+
+    try:
+        shutil.rmtree(file_path)
+    except FileNotFoundError:
+        pass
+
+    if path.exists(file_path) is False:
+        os.makedirs(file_path)
+
+    dataset.output_to_fits(
+        data_path=path.join(file_path, "data.fits"),
+        noise_map_path=path.join(file_path, "noise_map.fits"),
+        psf_path=path.join(file_path, "psf.fits"),
+    )
+
+    dataset = al.Imaging.from_fits(
+        data_path=path.join(file_path, "data.fits"),
+        noise_map_path=path.join(file_path, "noise_map.fits"),
+        psf_path=path.join(file_path, "psf.fits"),
+        pixel_scales=0.2,
+        over_sampling=over_sampling
+    )
+
+    mask = al.Mask2D.circular(
+        shape_native=dataset.data.shape_native, pixel_scales=0.2, radius=0.8
+    )
+
+    masked_dataset = dataset.apply_mask(mask=mask)
+
+    tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy])
+
+    fit = al.FitImaging(dataset=masked_dataset, tracer=tracer)
+
+    # The value is actually not zero before the blurring grid assumes a sub_size=1
+    # and does not use the iterative grid, which has a small impact on the chi-squared
+
+    assert fit.chi_squared == pytest.approx(7.451891005452524e-05, 1e-4)
+
+    file_path = path.join(
+        "{}".format(path.dirname(path.realpath(__file__))), "data_temp"
+    )
+
+    if path.exists(file_path) is True:
+        shutil.rmtree(file_path)
+
+
 def test__simulate_imaging_data_and_fit__known_likelihood():
 
     grid = al.Grid2D.uniform(shape_native=(31, 31), pixel_scales=0.2)