Feature/jax speed up general

autogalaxy/analysis/analysis/analysis.py

@@ -1,26 +1,20 @@
-import json
 import logging
-from typing import Dict, List, Optional, Tuple, Union
-from os import path
-import os
-import time
+from typing import List, Optional
 
-from autoconf import conf
 import autofit as af
 import autoarray as aa
 
 from autogalaxy.galaxy.galaxy import Galaxy
 from autogalaxy.galaxy.galaxies import Galaxies
 from autogalaxy.cosmology.lensing import LensingCosmology
-from autogalaxy.cosmology.wrap import Planck15
 
 logger = logging.getLogger(__name__)
 
 logger.setLevel(level="INFO")
 
 
 class Analysis(af.Analysis):
-    def __init__(self, cosmology: LensingCosmology = Planck15):
+    def __init__(self, cosmology: LensingCosmology = None):
         """
         Fits a model to a dataset via a non-linear search.
 
@@ -34,7 +28,10 @@ def __init__(self, cosmology: LensingCosmology = Planck15):
         cosmology
             The Cosmology assumed for this analysis.
         """
-        self.cosmology = cosmology
+
+        from autogalaxy.cosmology.wrap import Planck15
+
+        self.cosmology = cosmology or Planck15()
 
     def galaxies_via_instance_from(
         self,

autogalaxy/analysis/analysis/dataset.py

@@ -1,4 +1,3 @@
-import copy
 import logging
 from typing import Optional, Union
 
@@ -10,7 +9,6 @@
 from autogalaxy.analysis.adapt_images.adapt_image_maker import AdaptImageMaker
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
 from autogalaxy.cosmology.lensing import LensingCosmology
-from autogalaxy.cosmology.wrap import Planck15
 from autogalaxy.analysis.analysis.analysis import Analysis
 from autogalaxy.analysis.result import ResultDataset
 
@@ -24,7 +22,7 @@ def __init__(
         self,
         dataset: Union[aa.Imaging, aa.Interferometer],
         adapt_image_maker: Optional[AdaptImageMaker] = None,
-        cosmology: LensingCosmology = Planck15(),
+        cosmology: LensingCosmology = None,
         settings_inversion: aa.SettingsInversion = None,
         title_prefix: str = None,
     ):

autogalaxy/analysis/plotter_interface.py

@@ -236,16 +236,11 @@ def should_plot(name):
             mapper_list = inversion.cls_list_from(cls=aa.AbstractMapper)
 
             for i, mapper in enumerate(mapper_list):
-                y = mapper.mapper_grids.source_plane_mesh_grid[:, 0].array
-                x = mapper.mapper_grids.source_plane_mesh_grid[:, 1].array
+                y = mapper.mapper_grids.source_plane_mesh_grid[:, 0]
+                x = mapper.mapper_grids.source_plane_mesh_grid[:, 1]
                 reconstruction = inversion.reconstruction_dict[mapper]
                 noise_map = inversion.reconstruction_noise_map_dict[mapper]
 
-                print(y)
-                print(x)
-                print(reconstruction)
-                print(noise_map)
-
                 with open(
                     self.image_path / f"source_plane_reconstruction_{i}.csv",
                     mode="w",

autogalaxy/convert.py

@@ -1,4 +1,3 @@
-from astropy import units
 from typing import Tuple
 
 import jax
@@ -322,6 +321,8 @@ def multipole_comps_from(k_m: float, phi_m: float, m: int) -> Tuple[float, float
     -------
     The multipole component parameters.
     """
+    from astropy import units
+
     multipole_comp_0 = k_m * jnp.sin(phi_m * float(m) * units.deg.to(units.rad))
     multipole_comp_1 = k_m * jnp.cos(phi_m * float(m) * units.deg.to(units.rad))
 

autogalaxy/cosmology/__init__.py

@@ -1,5 +1,2 @@
 from .lensing import LensingCosmology
 from .wrap import Planck15
-from .model import LambdaCDMWrap
-from .model import FlatwCDMWrap
-from .model import FlatLambdaCDMWrap

autogalaxy/cosmology/lensing.py

@@ -1,17 +1,21 @@
-from astropy import constants
-from astropy import cosmology as cosmo
 import math
 import numpy as np
 
 
-class LensingCosmology(cosmo.FLRW):
+class LensingCosmology:
     """
     Class containing specific functions for performing gravitational lensing cosmology calculations.
 
     By inheriting from the astropy `cosmo.FLRW` class this provides many additional methods for performing cosmological
     calculations.
     """
 
+    def __init__(self):
+
+        from astropy.cosmology import FLRW
+
+        self._cosmo = FLRW()
+
     def arcsec_per_kpc_from(self, redshift: float) -> float:
         """
         Angular separation in arcsec corresponding to a proper kpc at redshift `z`.
@@ -189,6 +193,8 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from(
             The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface
             density is calculated.
         """
+        from astropy import constants
+
         const = constants.c.to("kpc / s") ** 2.0 / (
             4 * math.pi * constants.G.to("kpc3 / (solMass s2)")
         )
@@ -305,6 +311,8 @@ def velocity_dispersion_from(
         redshift_1
             The redshift of the second strong lens galaxy (the source).
         """
+        from astropy import constants
+
         const = constants.c.to("kpc / s")
 
         angular_diameter_distance_to_redshift_0_kpc = (

autogalaxy/cosmology/wrap.py

@@ -1,25 +1,25 @@
-from astropy import cosmology as cosmo
+def Planck15():
+    """
+    A lazy-loading wrapper for the astropy `Planck15` cosmology class.
 
-from autogalaxy.cosmology.lensing import LensingCosmology
+    The actual class is only created (and astropy imported) when this function is called.
+    """
 
+    from astropy import cosmology as cosmo
+    from autogalaxy.cosmology.lensing import LensingCosmology
 
-class Planck15(cosmo.FlatLambdaCDM, LensingCosmology):
-    def __init__(self):
-        """
-        A wrapper for the astropy `Planck15` cosmology class.
+    class _Planck15(cosmo.FlatLambdaCDM, LensingCosmology):
+        def __init__(self):
+            Planck15_astropy = cosmo.Planck15
 
-        The only role of this class is to instantiate the `Planck15` cosmology class from astropy, but to additionally
-        inherit from `LensingCosmology`, which is a class that provides additional functionality for calculating lensing
-        specific quantities in the cosmology.
-        """
-        Planck15 = cosmo.Planck15
+            super().__init__(
+                H0=Planck15_astropy.H0,
+                Om0=Planck15_astropy.Om0,
+                Tcmb0=Planck15_astropy.Tcmb0,
+                Neff=Planck15_astropy.Neff,
+                m_nu=Planck15_astropy.m_nu,
+                Ob0=Planck15_astropy.Ob0,
+                name=Planck15_astropy.name,
+            )
 
-        super().__init__(
-            H0=Planck15.H0,
-            Om0=Planck15.Om0,
-            Tcmb0=Planck15.Tcmb0,
-            Neff=Planck15.Neff,
-            m_nu=Planck15.m_nu,
-            Ob0=Planck15.Ob0,
-            name=Planck15.name,
-        )
+    return _Planck15()

autogalaxy/ellipse/dataset_interp.py

@@ -1,5 +1,4 @@
 import numpy as np
-from scipy import interpolate
 from typing import Tuple
 
 from autoconf import cached_property
@@ -45,11 +44,13 @@ def points_interp(self) -> Tuple[np.ndarray, np.ndarray]:
         return (x, y)
 
     @cached_property
-    def mask_interp(self) -> interpolate.RegularGridInterpolator:
+    def mask_interp(self) -> "interpolate.RegularGridInterpolator":
         """
         Returns a 2D interpolation of the mask, which is used to determine whether inteprolated values use a masked
         pixel for the interpolation and thus should not be included in a fit.
         """
+        from scipy import interpolate
+
         return interpolate.RegularGridInterpolator(
             points=self.points_interp,
             values=np.float64(self.dataset.data.mask),
@@ -58,10 +59,12 @@ def mask_interp(self) -> interpolate.RegularGridInterpolator:
         )
 
     @cached_property
-    def data_interp(self) -> interpolate.RegularGridInterpolator:
+    def data_interp(self) -> "interpolate.RegularGridInterpolator":
         """
         Returns a 2D interpolation of the data, which is used to evaluate the data at any point in 2D space.
         """
+        from scipy import interpolate
+
         return interpolate.RegularGridInterpolator(
             points=self.points_interp,
             values=np.float64(self.dataset.data.native),
@@ -70,10 +73,11 @@ def data_interp(self) -> interpolate.RegularGridInterpolator:
         )
 
     @cached_property
-    def noise_map_interp(self) -> interpolate.RegularGridInterpolator:
+    def noise_map_interp(self) -> "interpolate.RegularGridInterpolator":
         """
         Returns a 2D interpolation of the noise-map, which is used to evaluate the noise-map at any point in 2D space.
         """
+        from scipy import interpolate
 
         return interpolate.RegularGridInterpolator(
             points=self.points_interp,