Merge pull request #1161 from rhayes777/feature/xp_no_autofit_import

Feature/xp no autofit import

Author: Jammy2211

autofit/__init__.py

@@ -1,3 +1,4 @@
+from autoconf import jax_wrapper
 from autoconf.dictable import register_parser
 from . import conf
 

autofit/config/general.yaml

@@ -1,5 +1,3 @@
-jax:
-  use_jax: false                    # If True, PyAutoFit uses JAX internally, whereas False uses normal Numpy.
 updates:
   iterations_per_quick_update: 1e99 # Non-linear search iterations between every quick update, which just displays the maximum likelihood model fit.
   iterations_per_full_update: 1e99  # Non-linear search iterations between every full update, which outputs all visuals and result fits (e.g. model.result, search.summary), this exits the search and can be slow.  

autofit/example/analysis.py

@@ -1,8 +1,6 @@
 import numpy as np
 from typing import Dict, Optional
 
-from autoconf.jax_wrapper import numpy as xp
-
 import autofit as af
 
 from autofit.example.result import ResultExample
@@ -38,7 +36,7 @@ class Analysis(af.Analysis):
 
     LATENT_KEYS = ["gaussian.fwhm"]
 
-    def __init__(self, data: np.ndarray, noise_map: np.ndarray):
+    def __init__(self, data: np.ndarray, noise_map: np.ndarray, use_jax=False):
         """
         In this example the `Analysis` object only contains the data and noise-map. It can be easily extended,
         for more complex data-sets and model fitting problems.
@@ -51,12 +49,12 @@ def __init__(self, data: np.ndarray, noise_map: np.ndarray):
             A 1D numpy array containing the noise values of the data, used for computing the goodness of fit
             metric.
         """
-        super().__init__()
+        super().__init__(use_jax=use_jax)
 
         self.data = data
         self.noise_map = noise_map
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Determine the log likelihood of a fit of multiple profiles to the dataset.
 
@@ -98,14 +96,15 @@ def model_data_1d_from(self, instance: af.ModelInstance) -> np.ndarray:
         The model data of the profiles.
         """
 
-        xvalues = xp.arange(self.data.shape[0])
-        model_data_1d = xp.zeros(self.data.shape[0])
+        xvalues = self._xp.arange(self.data.shape[0])
+        model_data_1d = self._xp.zeros(self.data.shape[0])
 
         try:
             for profile in instance:
                 try:
                     model_data_1d += profile.model_data_from(
-                        xvalues=xvalues
+                        xvalues=xvalues,
+                        xp=self._xp
                     )
                 except AttributeError:
                     pass

autofit/example/model.py

@@ -2,8 +2,6 @@
 import numpy as np
 from typing import Tuple
 
-from autoconf.jax_wrapper import numpy as xp
-
 """
 The `Gaussian` class in this module is the model components that is fitted to data using a non-linear search. The
 inputs of its __init__ constructor are the parameters which can be fitted for.
@@ -47,7 +45,7 @@ def fwhm(self) -> float:
         the free parameters of the model which we are interested and may want to store the full samples information
         on (e.g. to create posteriors).
         """
-        return 2 * xp.sqrt(2 * xp.log(2)) * self.sigma
+        return 2 * np.sqrt(2 * np.log(2)) * self.sigma
 
     def _tree_flatten(self):
         return (self.centre, self.normalization, self.sigma), None
@@ -64,7 +62,7 @@ def __eq__(self, other):
             and self.sigma == other.sigma
         )
 
-    def model_data_from(self, xvalues: np.ndarray) -> np.ndarray:
+    def model_data_from(self, xvalues: np.ndarray, xp=np) -> np.ndarray:
         """
         Calculate the normalization of the profile on a 1D grid of Cartesian x coordinates.
 
@@ -82,7 +80,7 @@ def model_data_from(self, xvalues: np.ndarray) -> np.ndarray:
             xp.exp(-0.5 * xp.square(xp.divide(transformed_xvalues, self.sigma))),
         )
 
-    def f(self, x: float):
+    def f(self, x: float, xp=np):
         return (
             self.normalization
             / (self.sigma * xp.sqrt(2 * math.pi))
@@ -137,7 +135,7 @@ def __init__(
         self.normalization = normalization
         self.rate = rate
 
-    def model_data_from(self, xvalues: np.ndarray) -> np.ndarray:
+    def model_data_from(self, xvalues: np.ndarray, xp=np) -> np.ndarray:
         """
         Calculate the 1D Gaussian profile on a 1D grid of Cartesian x coordinates.
 

autofit/graphical/declarative/abstract.py

@@ -19,9 +19,11 @@ class AbstractDeclarativeFactor(Analysis, ABC):
     optimiser: AbstractFactorOptimiser
     _plates: Tuple[Plate, ...] = ()
 
-    def __init__(self, include_prior_factors=False):
+    def __init__(self, include_prior_factors=False, use_jax : bool = False):
         self.include_prior_factors = include_prior_factors
 
+        super().__init__(use_jax=use_jax)
+
     @property
     @abstractmethod
     def name(self):

autofit/graphical/declarative/collection.py

@@ -11,17 +11,15 @@
 from autofit.mapper.model import ModelInstance
 from autofit.mapper.prior_model.prior_model import Model
 
-from autoconf.jax_wrapper import register_pytree_node_class
 from ...non_linear.combined_result import CombinedResult
 
-
-@register_pytree_node_class
 class FactorGraphModel(AbstractDeclarativeFactor):
     def __init__(
         self,
         *model_factors: Union[AbstractDeclarativeFactor, HierarchicalFactor],
         name=None,
         include_prior_factors=True,
+        use_jax : bool = False
     ):
         """
         A collection of factors that describe models, which can be
@@ -36,6 +34,7 @@ def __init__(
         """
         super().__init__(
             include_prior_factors=include_prior_factors,
+            use_jax=use_jax,
         )
         self._model_factors = list(model_factors)
         self._name = name or namer(self.__class__.__name__)

autofit/graphical/declarative/factor/analysis.py

@@ -10,8 +10,6 @@
 from autofit.non_linear.paths.abstract import AbstractPaths
 from .abstract import AbstractModelFactor
 
-from autoconf.jax_wrapper import register_pytree_node_class
-
 
 class FactorCallable:
     def __init__(
@@ -45,8 +43,6 @@ def __call__(self, **kwargs: np.ndarray) -> float:
         instance = self.prior_model.instance_for_arguments(arguments)
         return self.analysis.log_likelihood_function(instance)
 
-
-@register_pytree_node_class
 class AnalysisFactor(AbstractModelFactor):
     @property
     def prior_model(self):

autofit/graphical/declarative/factor/hierarchical.py

@@ -144,7 +144,7 @@ def __call__(self, **kwargs):
 
 class _HierarchicalFactor(AbstractModelFactor):
     def __init__(
-        self, distribution_model: HierarchicalFactor, drawn_prior: Prior,
+        self, distribution_model: HierarchicalFactor, drawn_prior: Prior, use_jax : bool = False
     ):
         """
         A factor that links a variable to a parameterised distribution.
@@ -159,6 +159,7 @@ def __init__(
         """
         self.distribution_model = distribution_model
         self.drawn_prior = drawn_prior
+        self.use_jax = use_jax
 
         prior_variable_dict = {prior.name: prior for prior in distribution_model.priors}
 

autofit/graphical/factor_graphs/factor.py

@@ -1,6 +1,5 @@
 from copy import deepcopy
 from inspect import getfullargspec
-import jax
 from typing import Tuple, Dict, Any, Callable, Union, List, Optional, TYPE_CHECKING
 
 import numpy as np
@@ -285,6 +284,8 @@ def _set_jacobians(
         numerical_jacobian=True,
         jacfwd=True,
     ):
+        import jax
+
         self._vjp = vjp
         self._jacfwd = jacfwd
         if vjp or factor_vjp:
@@ -327,6 +328,7 @@ def __call__(self, values: VariableData) -> FactorValue:
         return self._cache[key]
 
     def _jax_factor_vjp(self, *args) -> Tuple[Any, Callable]:
+        import jax
         return jax.vjp(self._factor, *args)
 
     _factor_vjp = _jax_factor_vjp

autofit/graphical/laplace/newton.py

@@ -240,7 +240,7 @@ def take_quasi_newton_step(
 ) -> Tuple[Optional[float], OptimisationState]:
     """ """
     state.search_direction = search_direction(state, **(search_direction_kws or {}))
-    if state.search_direction.vecnorm(np.Inf) == 0:
+    if state.search_direction.vecnorm(np.inf) == 0:
         # if gradient is zero then at maximum already
         return 0.0, state