Feature/jax speed up general

autolens/analysis/analysis/dataset.py

@@ -1,6 +1,6 @@
 import os
 import logging
-from typing import List, Optional, Union
+from typing import List, Optional
 
 from autoconf import conf
 from autoconf.dictable import output_to_json
@@ -28,7 +28,7 @@ def __init__(
         dataset,
         positions_likelihood_list: Optional[List[PositionsLH]] = None,
         adapt_image_maker: Optional[ag.AdaptImageMaker] = None,
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
         raise_inversion_positions_likelihood_exception: bool = True,
         title_prefix: str = None,

autolens/analysis/analysis/lens.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 import logging
 import numpy as np
-from typing import Dict, List, Optional, Union
+from typing import List, Optional
 
 import autofit as af
 import autoarray as aa
@@ -23,7 +23,7 @@ class AnalysisLens:
     def __init__(
         self,
         positions_likelihood_list: Optional[List[PositionsLH]] = None,
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
     ):
         """
         Analysis classes are used by PyAutoFit to fit a model to a dataset via a non-linear search.
@@ -40,7 +40,9 @@ def __init__(
         cosmology
             The Cosmology assumed for this analysis.
         """
-        self.cosmology = cosmology
+        from autogalaxy.cosmology.wrap import Planck15
+
+        self.cosmology = cosmology or Planck15()
         self.positions_likelihood_list = positions_likelihood_list
 
     def tracer_via_instance_from(

autolens/imaging/mock/mock_fit_imaging.py

@@ -7,13 +7,15 @@ class MockFitImaging(aa.m.MockFitImaging):
     def __init__(
         self,
         tracer=None,
-        dataset=aa.m.MockDataset(),
+        dataset=None,
         inversion=None,
         noise_map=None,
         grid=None,
         blurred_image=None,
     ):
 
+        dataset = dataset or aa.m.MockDataset()
+
         super().__init__(
             dataset=dataset,
             inversion=inversion,

autolens/imaging/model/analysis.py

@@ -1,21 +1,14 @@
 import logging
-import numpy as np
-from typing import Dict, Optional, Tuple
 
-from autoconf.dictable import to_dict
 
 import autofit as af
 import autogalaxy as ag
 
-from autoarray.exc import PixelizationException
-
 from autolens.analysis.analysis.dataset import AnalysisDataset
 from autolens.imaging.model.result import ResultImaging
 from autolens.imaging.model.visualizer import VisualizerImaging
 from autolens.imaging.fit_imaging import FitImaging
 
-from autolens import exc
-
 logger = logging.getLogger(__name__)
 
 logger.setLevel(level="INFO")

autolens/imaging/plot/fit_imaging_plotters.py

@@ -19,9 +19,9 @@ class FitImagingPlotter(Plotter):
     def __init__(
         self,
         fit: FitImaging,
-        mat_plot_2d: aplt.MatPlot2D = aplt.MatPlot2D(),
-        visuals_2d: aplt.Visuals2D = aplt.Visuals2D(),
-        include_2d: aplt.Include2D = aplt.Include2D(),
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
         residuals_symmetric_cmap: bool = True
     ):
         """

autolens/interferometer/model/analysis.py

@@ -30,7 +30,7 @@ def __init__(
         dataset,
         positions_likelihood_list: Optional[PositionsLH] = None,
         adapt_image_maker: Optional[ag.AdaptImageMaker] = None,
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
         raise_inversion_positions_likelihood_exception: bool = True,
         title_prefix: str = None,

autolens/interferometer/plot/fit_interferometer_plotters.py

@@ -17,12 +17,12 @@ class FitInterferometerPlotter(Plotter):
     def __init__(
         self,
         fit: FitInterferometer,
-        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(),
+        mat_plot_1d: aplt.MatPlot1D = None,
+        visuals_1d: aplt.Visuals1D = None,
+        include_1d: aplt.Include1D = None,
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
         residuals_symmetric_cmap: bool = True,
     ):
         """

autolens/lens/plot/tracer_plotters.py

@@ -17,12 +17,12 @@ 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(),
+        mat_plot_1d: aplt.MatPlot1D = None,
+        visuals_1d: aplt.Visuals1D = None,
+        include_1d: aplt.Include1D = None,
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
     ):
         """
         Plots the attributes of `Tracer` objects using the matplotlib methods `plot()` and `imshow()` and many

autolens/lens/sensitivity.py

@@ -146,9 +146,9 @@ def __init__(
         source_image: Optional[aa.Array2D] = None,
         result: Optional[SubhaloSensitivityResult] = None,
         data_subtracted: Optional[aa.Array2D] = None,
-        mat_plot_2d: aplt.MatPlot2D = aplt.MatPlot2D(),
-        visuals_2d: aplt.Visuals2D = aplt.Visuals2D(),
-        include_2d: aplt.Include2D = aplt.Include2D(),
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
     ):
         """
         Plots the simulated datasets and results of a sensitivity mapping analysis, where dark matter halos are used

autolens/lens/subhalo.py

@@ -173,9 +173,9 @@ def __init__(
         result: Optional[SubhaloGridSearchResult] = None,
         fit_imaging_with_subhalo: Optional[FitImaging] = None,
         fit_imaging_no_subhalo: Optional[FitImaging] = None,
-        mat_plot_2d: aplt.MatPlot2D = aplt.MatPlot2D(),
-        visuals_2d: aplt.Visuals2D = aplt.Visuals2D(),
-        include_2d: aplt.Include2D = aplt.Include2D(),
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
     ):
         """
         Plots the results of scanning for a dark matter subhalo in strong lens imaging.

autolens/lens/tracer.py

@@ -17,7 +17,7 @@ class Tracer(ABC, ag.OperateImageGalaxies, ag.OperateDeflections):
     def __init__(
         self,
         galaxies: Union[List[ag.Galaxy], af.ModelInstance],
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
     ):
         """
         Performs gravitational lensing ray-tracing calculations based on an input list of galaxies and a cosmology.
@@ -52,7 +52,7 @@ def __init__(
 
         self.galaxies = galaxies
 
-        self.cosmology = cosmology
+        self.cosmology = cosmology or ag.cosmo.Planck15()
 
     @property
     def galaxies_ascending_redshift(self) -> List[ag.Galaxy]:
@@ -135,7 +135,7 @@ def sliced_tracer_from(
         line_of_sight_galaxies: List[ag.Galaxy],
         source_galaxies: List[ag.Galaxy],
         planes_between_lenses: List[int],
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
     ):
         """
         Returns a tracer where the lens system is split into planes with specified redshift distances between them.
@@ -182,6 +182,7 @@ def sliced_tracer_from(
         cosmology
             The cosmology used to perform ray-tracing calculations.
         """
+        cosmology = cosmology or ag.cosmo.Planck15()
 
         lens_redshifts = tracer_util.plane_redshifts_from(galaxies=lens_galaxies)
 

autolens/lens/tracer_util.py

@@ -86,7 +86,7 @@ def planes_from(
 def traced_grid_2d_list_from(
     planes: List[List[ag.Galaxy]],
     grid: aa.type.Grid2DLike,
-    cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+    cosmology: ag.cosmo.LensingCosmology = None,
     plane_index_limit: int = Optional[None],
 ):
     """
@@ -133,6 +133,7 @@ def traced_grid_2d_list_from(
     traced_grid_list
         A list of 2D (y,x) grids each of which are the input grid ray-traced to a redshift of the input list of planes.
     """
+    cosmology = cosmology or ag.cosmo.Planck15()
 
     traced_grid_list = []
     traced_deflection_list = []
@@ -175,7 +176,7 @@ def grid_2d_at_redshift_from(
     redshift: float,
     galaxies: List[ag.Galaxy],
     grid: aa.type.Grid2DLike,
-    cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+    cosmology: ag.cosmo.LensingCosmology = None,
 ) -> aa.type.Grid2DLike:
     """
     Returns a ray-traced grid of 2D Cartesian (y,x) coordinates, which accounts for multi-plane ray-tracing, at a
@@ -217,6 +218,7 @@ def grid_2d_at_redshift_from(
     cosmology
         The cosmology used for ray-tracing from which angular diameter distances between planes are computed.
     """
+    cosmology = cosmology or ag.cosmo.Planck15()
 
     plane_redshifts = plane_redshifts_from(galaxies=galaxies)
 
@@ -254,7 +256,7 @@ def grid_2d_at_redshift_from(
 def time_delays_from(
     galaxies: List[ag.Galaxy],
     grid: aa.type.Grid2DLike,
-    cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+    cosmology: ag.cosmo.LensingCosmology = None,
 ) -> aa.type.Grid2DLike:
     """
     Returns the gravitational lensing time delay in days for a grid of 2D (y, x) coordinates.
@@ -304,6 +306,8 @@ def time_delays_from(
     -------
     The time delay at each (y, x) coordinate in the input grid, in units of days.
     """
+    cosmology = cosmology or ag.cosmo.Planck15()
+
     plane_redshifts = plane_redshifts_from(galaxies=galaxies)
 
     if len(plane_redshifts) != 2:

autolens/plot/abstract_plotters.py

@@ -2,13 +2,48 @@
 
 set_backend()
 
-from autoarray.plot.abstract_plotters import AbstractPlotter
-
 from autolens.plot.get_visuals.one_d import GetVisuals1D
 from autolens.plot.get_visuals.two_d import GetVisuals2D
 
+from autogalaxy.plot.abstract_plotters import AbstractPlotter
+
+from autogalaxy.plot.mat_plot.one_d import MatPlot1D
+from autogalaxy.plot.mat_plot.two_d import MatPlot2D
+from autogalaxy.plot.visuals.one_d import Visuals1D
+from autogalaxy.plot.visuals.two_d import Visuals2D
+from autogalaxy.plot.include.one_d import Include1D
+from autogalaxy.plot.include.two_d import Include2D
+
 
 class Plotter(AbstractPlotter):
+
+    def __init__(
+        self,
+        mat_plot_1d: MatPlot1D = None,
+        visuals_1d: Visuals1D = None,
+        include_1d: Include1D = None,
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
+    ):
+
+        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.visuals_1d = visuals_1d or Visuals1D()
+        self.include_1d = include_1d or Include1D()
+        self.mat_plot_1d = mat_plot_1d or MatPlot1D()
+
+        self.visuals_2d = visuals_2d or Visuals2D()
+        self.include_2d = include_2d or Include2D()
+        self.mat_plot_2d = mat_plot_2d or MatPlot2D()
+
     @property
     def get_1d(self):
         return GetVisuals1D(visuals=self.visuals_1d, include=self.include_1d)

autolens/point/fit/positions/image/pair.py

@@ -1,7 +1,4 @@
-import jax.numpy as jnp
 import numpy as np
-from ott.geometry import pointcloud
-from ott.solvers.linear import sinkhorn
 from scipy.optimize import linear_sum_assignment
 
 import autoarray as aa

autolens/point/model/analysis.py

@@ -33,7 +33,7 @@ def __init__(
         solver: PointSolver,
         fit_positions_cls=FitPositionsImagePairRepeat,
         image=None,
-        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+        cosmology: ag.cosmo.LensingCosmology = None,
         title_prefix: str = None,
     ):
         """

autolens/point/plot/fit_point_plotters.py

@@ -8,12 +8,12 @@ class FitPointDatasetPlotter(Plotter):
     def __init__(
         self,
         fit: FitPointDataset,
-        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(),
+        mat_plot_1d: aplt.MatPlot1D = None,
+        visuals_1d: aplt.Visuals1D = None,
+        include_1d: aplt.Include1D = None,
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
     ):
         """
         Plots the attributes of `FitPointDataset` objects using matplotlib methods and functions which customize the

autolens/point/plot/point_dataset_plotters.py

@@ -8,12 +8,12 @@ class PointDatasetPlotter(Plotter):
     def __init__(
         self,
         dataset: PointDataset,
-        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(),
+        mat_plot_1d: aplt.MatPlot1D = None,
+        visuals_1d: aplt.Visuals1D = None,
+        include_1d: aplt.Include1D = None,
+        mat_plot_2d: aplt.MatPlot2D = None,
+        visuals_2d: aplt.Visuals2D = None,
+        include_2d: aplt.Include2D = None,
     ):
         """
         Plots the attributes of `PointDataset` objects using the matplotlib methods and functions functions which