black

Author: Jammy2211

autolens/__init__.py

@@ -1,3 +1,4 @@
+from autoconf import jax_wrapper
 from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
 from autoarray import preprocess
 from autoarray.dataset.imaging.w_tilde import WTildeImaging

autolens/analysis/analysis/dataset.py

@@ -1,5 +1,6 @@
-import os
 import logging
+import numpy as np
+import os
 from typing import List, Optional
 
 from autoconf import conf
@@ -33,6 +34,7 @@ def __init__(
         preloads: aa.Preloads = None,
         raise_inversion_positions_likelihood_exception: bool = True,
         title_prefix: str = None,
+        use_jax: bool = True,
         **kwargs,
     ):
         """
@@ -75,6 +77,7 @@ def __init__(
             settings_inversion=settings_inversion,
             preloads=preloads,
             title_prefix=title_prefix,
+            use_jax=use_jax,
             **kwargs,
         )
 

autolens/analysis/positions.py

@@ -1,4 +1,3 @@
-import jax
 import numpy as np
 from typing import Optional
 from os import path
@@ -189,6 +188,8 @@ def log_likelihood_penalty_from(
 
         penalty = self.log_likelihood_penalty_factor * (max_separation - self.threshold)
 
+        import jax
+
         return jax.lax.cond(
             max_separation > self.threshold,
             lambda: penalty,

autolens/config/general.yaml

@@ -1,5 +1,3 @@
-jax:
-  use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 output:
   fit_dill: false
 test:

autolens/fixtures.py

@@ -145,15 +145,14 @@ def make_analysis_imaging_7x7():
     analysis = al.AnalysisImaging(
         dataset=make_masked_imaging_7x7(),
         settings_inversion=aa.SettingsInversion(use_w_tilde=False),
+        use_jax=False,
     )
     analysis._adapt_images = make_adapt_images_7x7()
     return analysis
 
 
 def make_analysis_interferometer_7():
-    analysis = al.AnalysisInterferometer(
-        dataset=make_interferometer_7(),
-    )
+    analysis = al.AnalysisInterferometer(dataset=make_interferometer_7(), use_jax=False)
     analysis._adapt_images = make_adapt_images_7x7()
     return analysis
 
@@ -162,4 +161,5 @@ def make_analysis_point_x2():
     return al.AnalysisPoint(
         point_dict=make_point_dict(),
         solver=al.m.MockPointSolver(model_positions=make_positions_x2()),
+        use_jax=False,
     )

autolens/imaging/model/analysis.py

@@ -46,7 +46,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
@@ -87,15 +87,14 @@ def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
 
         log_likelihood_penalty = self.log_likelihood_penalty_from(
             instance=instance,
-            xp=xp
+            xp=self._xp
         )
 
-        return self.fit_from(instance=instance, xp=xp).figure_of_merit - log_likelihood_penalty
+        return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
 
     def fit_from(
         self,
         instance: af.ModelInstance,
-        xp=np
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -133,7 +132,7 @@ def fit_from(
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
-            xp=xp
+            xp=self._xp
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):

autolens/interferometer/model/analysis.py

@@ -1,21 +1,17 @@
 import logging
 import numpy as np
-from typing import Dict, Optional, Tuple, Union
+from typing import Optional
 
 import autofit as af
 import autoarray as aa
 import autogalaxy as ag
 
-from autoarray.exc import PixelizationException
-
 from autolens.analysis.analysis.dataset import AnalysisDataset
 from autolens.analysis.positions import PositionsLH
 from autolens.interferometer.model.result import ResultInterferometer
 from autolens.interferometer.model.visualizer import VisualizerInterferometer
 from autolens.interferometer.fit_interferometer import FitInterferometer
 
-from autolens import exc
-
 logger = logging.getLogger(__name__)
 
 logger.setLevel(level="INFO")
@@ -35,6 +31,7 @@ def __init__(
         preloads: aa.Preloads = None,
         raise_inversion_positions_likelihood_exception: bool = True,
         title_prefix: str = None,
+        use_jax: bool = True,
     ):
         """
         Analysis classes are used by PyAutoFit to fit a model to a dataset via a non-linear search.
@@ -86,6 +83,7 @@ def __init__(
             preloads=preloads,
             raise_inversion_positions_likelihood_exception=raise_inversion_positions_likelihood_exception,
             title_prefix=title_prefix,
+            use_jax=use_jax,
         )
 
     @property
@@ -114,7 +112,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance, xp=np):
+    def log_likelihood_function(self, instance):
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the interferometer dataset.
@@ -154,15 +152,12 @@ def log_likelihood_function(self, instance, xp=np):
         """
 
         log_likelihood_penalty = self.log_likelihood_penalty_from(
-            instance=instance, xp=xp
+            instance=instance, xp=self._xp
         )
 
-        return (
-            self.fit_from(instance=instance, xp=xp).figure_of_merit
-            - log_likelihood_penalty
-        )
+        return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
 
-    def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer:
+    def fit_from(self, instance: af.ModelInstance) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
 
@@ -198,7 +193,7 @@ def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer:
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
-            xp=xp,
+            xp=self._xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):

autolens/lens/tracer.py

@@ -709,7 +709,7 @@ def deflections_between_planes_from(
             The 2D (y, x) coordinates where values of the deflections are evaluated.
         """
 
-        traced_grids_list = self.traced_grid_2d_list_from(grid=grid)
+        traced_grids_list = self.traced_grid_2d_list_from(grid=grid, xp=xp)
 
         return traced_grids_list[plane_i] - traced_grids_list[plane_j]
 

autolens/point/fit/positions/source/separations.py

@@ -93,7 +93,9 @@ def model_data(self) -> aa.Grid2DIrregular:
         at its specific redshift are used.
         """
         if len(self.tracer.planes) <= 2:
-            deflections = self.tracer.deflections_yx_2d_from(grid=self.data)
+            deflections = self.tracer.deflections_yx_2d_from(
+                grid=self.data, xp=self._xp
+            )
         else:
             deflections = self.tracer.deflections_between_planes_from(
                 grid=self.data, xp=self._xp, plane_i=0, plane_j=self.plane_index

autolens/point/model/analysis.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+import numpy as np
 
 import autofit as af
 import autogalaxy as ag
@@ -13,15 +13,6 @@
 from autolens.point.model.visualizer import VisualizerPoint
 from autolens.point.solver import PointSolver
 
-from autolens import exc
-
-try:
-    import numba
-
-    NumbaException = numba.errors.TypingError
-except ModuleNotFoundError:
-    NumbaException = ValueError
-
 
 class AnalysisPoint(AgAnalysis, AnalysisLens):
     Visualizer = VisualizerPoint
@@ -35,6 +26,7 @@ def __init__(
         image=None,
         cosmology: ag.cosmo.LensingCosmology = None,
         title_prefix: str = None,
+        use_jax: bool = True,
     ):
         """
         Fits a lens model to a point source dataset (e.g. positions, fluxes, time delays) via a non-linear search.
@@ -69,7 +61,7 @@ def __init__(
             A string that is added before the title of all figures output by visualization, for example to
             put the name of the dataset and galaxy in the title.
         """
-        super().__init__(cosmology=cosmology)
+        super().__init__(cosmology=cosmology, use_jax=use_jax)
 
         AnalysisLens.__init__(self=self, cosmology=cosmology)
 
@@ -127,9 +119,6 @@ def fit_from(
         self,
         instance,
     ) -> FitPointDataset:
-        tracer = self.tracer_via_instance_from(
-            instance=instance,
-        )
         """
         Given a model instance create a `FitPointDataset` object.
 
@@ -146,11 +135,16 @@ def fit_from(
         -------
         The fit of the lens model to the point source dataset.
         """
+        tracer = self.tracer_via_instance_from(
+            instance=instance,
+        )
+
         return FitPointDataset(
             dataset=self.dataset,
             tracer=tracer,
             solver=self.solver,
             fit_positions_cls=self.fit_positions_cls,
+            xp=self._xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):