Merge pull request #269 from Jammy2211/feature/sky_simplify

Feature/sky simplify

autolens/__init__.py

@@ -5,6 +5,7 @@
 from autoarray.dataset.interferometer.dataset import (
     Interferometer,
 )
+from autoarray.dataset.dataset_model import DatasetModel
 from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
 from autoarray.operators.convolver import Convolver

autolens/aggregator/fit_imaging.py

@@ -7,6 +7,7 @@
 from autolens.analysis.preloads import Preloads
 
 from autogalaxy.aggregator.imaging import _imaging_from
+from autogalaxy.aggregator.dataset_model import _dataset_model_from
 from autogalaxy.aggregator import agg_util
 
 from autolens.aggregator.tracer import _tracer_from
@@ -58,6 +59,8 @@ def _fit_imaging_from(
 
     tracer_list = _tracer_from(fit=fit, instance=instance)
 
+    dataset_model_list = _dataset_model_from(fit=fit, instance=instance)
+
     adapt_images_list = agg_util.adapt_images_from(fit=fit)
 
     settings_inversion = settings_inversion or fit.value(name="settings_inversion")
@@ -68,8 +71,8 @@ def _fit_imaging_from(
 
     fit_dataset_list = []
 
-    for dataset, tracer, adapt_images, mesh_grids_of_planes in zip(
-        dataset_list, tracer_list, adapt_images_list, mesh_grids_of_planes_list
+    for dataset, tracer, dataset_model, adapt_images, mesh_grids_of_planes in zip(
+        dataset_list, tracer_list, dataset_model_list, adapt_images_list, mesh_grids_of_planes_list
     ):
         preloads = agg_util.preloads_from(
             preloads_cls=Preloads,
@@ -82,6 +85,7 @@ def _fit_imaging_from(
             FitImaging(
                 dataset=dataset,
                 tracer=tracer,
+                dataset_model=dataset_model,
                 adapt_images=adapt_images,
                 settings_inversion=settings_inversion,
                 preloads=preloads,

autolens/aggregator/fit_interferometer.py

@@ -4,6 +4,7 @@
 import autoarray as aa
 
 from autogalaxy.aggregator.interferometer import _interferometer_from
+from autogalaxy.aggregator.dataset_model import _dataset_model_from
 
 from autolens.interferometer.fit_interferometer import FitInterferometer
 from autolens.analysis.preloads import Preloads
@@ -60,6 +61,7 @@ def _fit_interferometer_from(
         real_space_mask=real_space_mask,
     )
     tracer_list = _tracer_from(fit=fit, instance=instance)
+    dataset_model_list = _dataset_model_from(fit=fit, instance=instance)
 
     adapt_images_list = agg_util.adapt_images_from(fit=fit)
 
@@ -71,8 +73,8 @@ def _fit_interferometer_from(
 
     fit_dataset_list = []
 
-    for dataset, tracer, adapt_images, mesh_grids_of_planes in zip(
-        dataset_list, tracer_list, adapt_images_list, mesh_grids_of_planes_list
+    for dataset, tracer, dataset_model, adapt_images, mesh_grids_of_planes in zip(
+        dataset_list, tracer_list, dataset_model_list, adapt_images_list, mesh_grids_of_planes_list
     ):
         preloads = agg_util.preloads_from(
             preloads_cls=Preloads,
@@ -85,6 +87,7 @@ def _fit_interferometer_from(
             FitInterferometer(
                 dataset=dataset,
                 tracer=tracer,
+                dataset_model=dataset_model,
                 adapt_images=adapt_images,
                 settings_inversion=settings_inversion,
                 preloads=preloads,

autolens/imaging/fit_imaging.py

@@ -21,7 +21,7 @@ def __init__(
         self,
         dataset: aa.Imaging,
         tracer: Tracer,
-        sky: Optional[ag.LightProfile] = None,
+        dataset_model : Optional[aa.DatasetModel] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: Preloads = Preloads(),
@@ -57,8 +57,8 @@ def __init__(
             The imaging dataset which is fitted by the galaxies in the tracer.
         tracer
             The tracer of galaxies whose light profile images are used to fit the imaging data.
-        sky
-            Model component used to represent the background sky emission in an image (e.g. a `Sky` light profile).
+        dataset_model
+            Attributes which allow for parts of a dataset to be treated as a model (e.g. the background sky level).
         adapt_images
             Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the
             reconstructed galaxy's morphology.
@@ -72,12 +72,11 @@ def __init__(
             decorator take to run.
         """
 
-        super().__init__(dataset=dataset, run_time_dict=run_time_dict)
+        super().__init__(dataset=dataset, dataset_model=dataset_model, run_time_dict=run_time_dict)
         AbstractFitInversion.__init__(
-            self=self, model_obj=tracer, sky=sky, settings_inversion=settings_inversion
+            self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
 
-        self.sky = sky
         self.tracer = tracer
 
         self.adapt_images = adapt_images
@@ -96,16 +95,11 @@ def blurred_image(self) -> aa.Array2D:
 
         if self.preloads.blurred_image is None:
 
-            if isinstance(self.sky, ag.lp.Sky):
-                image = self.sky.image_2d_from(grid=self.dataset.grid)
-            else:
-                image = np.zeros(self.dataset.shape_slim)
-
             return self.tracer.blurred_image_2d_from(
                 grid=self.dataset.grid,
                 convolver=self.dataset.convolver,
                 blurring_grid=self.dataset.blurring_grid,
-            ) + image
+            )
 
         return self.preloads.blurred_image
 
@@ -115,7 +109,7 @@ def profile_subtracted_image(self) -> aa.Array2D:
         Returns the dataset's image with all blurred light profile images in the fit's tracer subtracted.
         """
 
-        return self.image - self.blurred_image
+        return self.data - self.blurred_image
 
     @property
     def tracer_to_inversion(self) -> TracerToInversion:
@@ -135,7 +129,6 @@ def tracer_to_inversion(self) -> TracerToInversion:
         return TracerToInversion(
             dataset=dataset,
             tracer=self.tracer,
-            sky=self.sky,
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
@@ -302,7 +295,7 @@ def subtracted_images_of_planes_list(self) -> List[aa.Array2D]:
                 if i != galaxy_index
             ]
 
-            subtracted_image = self.image - sum(other_planes_model_images)
+            subtracted_image = self.data - sum(other_planes_model_images)
 
             subtracted_images_of_planes_list.append(subtracted_image)
 
@@ -384,7 +377,7 @@ def refit_with_new_preloads(
         return FitImaging(
             dataset=self.dataset,
             tracer=self.tracer,
-            sky=self.sky,
+            dataset_model=self.dataset_model,
             adapt_images=self.adapt_images,
             settings_inversion=settings_inversion,
             preloads=preloads,

autolens/imaging/model/analysis.py

@@ -159,7 +159,7 @@ def fit_from(
             instance=instance, run_time_dict=run_time_dict
         )
 
-        sky = self.sky_via_instance_from(instance=instance)
+        dataset_model = self.dataset_model_via_instance_from(instance=instance)
 
         adapt_images = self.adapt_images_via_instance_from(instance=instance)
 
@@ -168,7 +168,7 @@ def fit_from(
         return FitImaging(
             dataset=self.dataset,
             tracer=tracer,
-            sky=sky,
+            dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=preloads,

autolens/interferometer/fit_interferometer.py

@@ -18,6 +18,7 @@ def __init__(
         self,
         dataset: aa.Interferometer,
         tracer: Tracer,
+        dataset_model: Optional[aa.DatasetModel] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: Preloads = Preloads(),
@@ -54,6 +55,8 @@ def __init__(
             The interforometer dataset which is fitted by the galaxies in the tracer.
         tracer
             The tracer of galaxies whose light profile images are used to fit the interferometer data.
+        dataset_model
+            Attributes which allow for parts of a dataset to be treated as a model (e.g. the background sky level).
         adapt_images
             Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the
             reconstructed galaxy's morphology.
@@ -82,9 +85,9 @@ def __init__(
 
         self.run_time_dict = run_time_dict
 
-        super().__init__(dataset=dataset, run_time_dict=run_time_dict)
+        super().__init__(dataset=dataset, dataset_model=dataset_model, run_time_dict=run_time_dict)
         AbstractFitInversion.__init__(
-            self=self, model_obj=tracer, sky=None, settings_inversion=settings_inversion
+            self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
 
     @property
@@ -103,7 +106,7 @@ def profile_subtracted_visibilities(self) -> aa.Visibilities:
         Returns the interferometer dataset's visibilities with all transformed light profile images in the fit's
         tracer subtracted.
         """
-        return self.visibilities - self.profile_visibilities
+        return self.data - self.profile_visibilities
 
     @property
     def tracer_to_inversion(self) -> TracerToInversion:
@@ -219,7 +222,7 @@ def model_visibilities_of_planes_list(self) -> List[aa.Visibilities]:
         galaxy_model_visibilities_dict = self.galaxy_model_visibilities_dict
 
         model_visibilities_of_planes_list = [
-            aa.Visibilities.zeros(shape_slim=(self.dataset.visibilities.shape_slim,))
+            aa.Visibilities.zeros(shape_slim=(self.dataset.data.shape_slim,))
             for i in range(self.tracer.total_planes)
         ]
 
@@ -274,8 +277,9 @@ def refit_with_new_preloads(
             settings_inversion = self.settings_inversion
 
         return FitInterferometer(
-            dataset=self.interferometer,
+            dataset=self.dataset,
             tracer=self.tracer,
+            dataset_model=self.dataset_model,
             adapt_images=self.adapt_images,
             settings_inversion=settings_inversion,
             preloads=preloads,

autolens/interferometer/model/analysis.py

@@ -313,7 +313,7 @@ def profile_log_likelihood_function(
             instance=instance,
         )
 
-        info_dict["number_of_visibilities"] = self.dataset.visibilities.shape[0]
+        info_dict["number_of_visibilities"] = self.dataset.data.shape[0]
         info_dict["transformer_cls"] = self.dataset.transformer.__class__.__name__
 
         self.output_profiling_info(

autolens/lens/to_inversion.py

@@ -16,7 +16,6 @@ def __init__(
         self,
         dataset: Optional[Union[aa.Imaging, aa.Interferometer, aa.DatasetInterface]],
         tracer,
-        sky: Optional[Basis] = None,
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads=Preloads(),
@@ -26,7 +25,6 @@ def __init__(
 
         super().__init__(
             dataset=dataset,
-            sky=sky,
             adapt_images=adapt_images,
             settings_inversion=settings_inversion,
             preloads=preloads,
@@ -106,7 +104,6 @@ def lp_linear_func_list_galaxy_dict(
             galaxies_to_inversion = ag.GalaxiesToInversion(
                 dataset=dataset,
                 galaxies=galaxies,
-                sky=self.sky,
                 settings_inversion=self.settings_inversion,
                 adapt_images=self.adapt_images,
                 run_time_dict=self.run_time_dict,
@@ -168,7 +165,6 @@ def image_plane_mesh_grid_pg_list(self) -> List[List]:
             to_inversion = ag.GalaxiesToInversion(
                 dataset=self.dataset,
                 galaxies=galaxies,
-                sky=self.sky,
                 adapt_images=self.adapt_images,
                 settings_inversion=self.settings_inversion,
                 run_time_dict=self.run_time_dict,
@@ -245,7 +241,6 @@ def mapper_galaxy_dict(self) -> Dict[aa.AbstractMapper, ag.Galaxy]:
                 to_inversion = ag.GalaxiesToInversion(
                     dataset=self.dataset,
                     galaxies=galaxies,
-                    sky=self.sky,
                     preloads=self.preloads,
                     adapt_images=self.adapt_images,
                     settings_inversion=self.settings_inversion,

autolens/point/fit_point/positions_image.py

@@ -1,11 +1,8 @@
-from functools import partial
 from typing import Optional
 
 import autoarray as aa
 import autogalaxy as ag
 
-from autolens.point.point_dataset import PointDict
-from autolens.point.point_dataset import PointDataset
 from autolens.point.point_solver import PointSolver
 from autolens.lens.tracer import Tracer
 

test_autolens/aggregator/conftest.py

@@ -56,7 +56,12 @@ def aggregator_from(database_file, analysis, model, samples):
 
 @pytest.fixture(name="model")
 def make_model():
+
+    dataset_model = af.Model(al.DatasetModel)
+    dataset_model.background_sky_level = af.UniformPrior(lower_limit=0.5, upper_limit=1.5)
+
     return af.Collection(
+        dataset_model=dataset_model,
         galaxies=af.Collection(
             lens=af.Model(al.Galaxy, redshift=0.5, light=al.lp.Sersic),
             source=af.Model(al.Galaxy, redshift=1.0, light=al.lp.Sersic),

test_autolens/aggregator/test_aggregator_fit_imaging.py

@@ -28,6 +28,8 @@ def test__fit_imaging_randomly_drawn_via_pdf_gen_from(
             assert fit_list[0].tracer.galaxies[0].light.centre == (10.0, 10.0)
             assert fit_list[0].tracer.galaxies[1].redshift == 1.0
 
+            assert fit_list[0].dataset_model.background_sky_level == 10.0
+
     assert i == 2
 
     clean(database_file=database_file)

test_autolens/aggregator/test_aggregator_fit_interferometer.py

@@ -28,6 +28,8 @@ def test__fit_interferometer_randomly_drawn_via_pdf_gen_from(
             assert fit_list[0].tracer.galaxies[0].light.centre == (10.0, 10.0)
             assert fit_list[0].tracer.galaxies[1].redshift == 1.0
 
+            assert fit_list[0].dataset_model.background_sky_level == 10.0
+
     assert i == 2
 
     clean(database_file=database_file)