Feature/jax speed up general

autofit/aggregator/summary/aggregate_fits.py

@@ -2,8 +2,6 @@
 from enum import Enum
 from typing import Dict, List, Union
 
-from astropy.table import Table
-from astropy.io import fits
 from pathlib import Path
 
 from autofit.aggregator.search_output import SearchOutput
@@ -53,7 +51,7 @@ def __init__(self, aggregator: Union[Aggregator, List[SearchOutput]]):
     def _hdus(
         result: SearchOutput,
         hdus: List[Enum],
-    ) -> List[fits.ImageHDU]:
+    ) -> "List[fits.ImageHDU]":
         """
         Extract the HDUs from a given fits for a given search.
 
@@ -68,6 +66,8 @@ def _hdus(
         -------
         The extracted HDUs.
         """
+        from astropy.io import fits
+
         row = []
         for hdu in hdus:
             source = result.value(subplot_filename(hdu))
@@ -80,7 +80,7 @@ def _hdus(
             )
         return row
 
-    def extract_fits(self, hdus: List[Enum]) -> fits.HDUList:
+    def extract_fits(self, hdus: List[Enum]) -> "fits.HDUList":
         """
         Extract the HDUs from the fits files for every search in the aggregator.
 
@@ -95,6 +95,8 @@ def extract_fits(self, hdus: List[Enum]) -> fits.HDUList:
         -------
         The extracted HDUs.
         """
+        from astropy.io import fits
+
         output = [fits.PrimaryHDU()]
         for result in self.aggregator:
             output.extend(self._hdus(result, hdus))
@@ -117,6 +119,8 @@ def extract_csv(self, filename: str) -> List[Dict]:
         -------
         The extracted HDUs.
         """
+        from astropy.table import Table
+
         output = []
         for result in self.aggregator:
             output.append(Table.read(result.value(filename), format="fits"))
@@ -144,6 +148,8 @@ def output_to_folder(
             The name of the fits file. This is the attribute of the search output that is used to name the file.
             OR a list of names for each HDU.
         """
+        from astropy.io import fits
+
         folder.mkdir(parents=True, exist_ok=True)
 
         for i, result in enumerate(self.aggregator):

autofit/database/model/prior.py

@@ -100,7 +100,6 @@ class Prior(Object):
     def _from_object(cls, model: abstract.Prior):
         instance = cls()
         instance.cls = type(model)
-        print(model.__database_args__)
         instance._add_children(
             [(key, getattr(model, key)) for key in model.__database_args__]
         )

autofit/example/util.py

@@ -1,6 +1,5 @@
 import os
 from os import path
-import matplotlib.pyplot as plt
 import numpy as np
 from typing import Optional
 
@@ -37,6 +36,8 @@ def plot_profile_1d(
     output_format
         Determines where the plot is displayed on your screen ("show") or output to the hard-disk as a png ("png").
     """
+    import matplotlib.pyplot as plt
+
     plt.errorbar(
         x=xvalues, y=profile_1d, yerr=errors, color=color, ecolor="k", elinewidth=1, capsize=2
     )

autofit/example/visualize.py

@@ -1,5 +1,4 @@
 import os
-import matplotlib.pyplot as plt
 import numpy as np
 from typing import List
 
@@ -39,6 +38,8 @@ def visualize_before_fit(
             The model which is fitted to the data, which may be used to customize the visualization.
         """
 
+        import matplotlib.pyplot as plt
+
         xvalues = np.arange(analysis.data.shape[0])
 
         plt.errorbar(
@@ -94,6 +95,8 @@ def visualize(
             which may change which images are output.
         """
 
+        import matplotlib.pyplot as plt
+
         xvalues = np.arange(analysis.data.shape[0])
         model_data_1d_list = []
 

autofit/graphical/expectation_propagation/visualise.py

@@ -3,8 +3,6 @@
 from pathlib import Path
 import warnings
 
-import matplotlib.pyplot as plt
-
 from autofit.graphical.expectation_propagation.history import EPHistory
 
 logger = logging.getLogger(__name__)
@@ -34,6 +32,8 @@ def __call__(self):
         """
         Save a plot of Evidence and KL Divergence for the ep_history
         """
+        import matplotlib.pyplot as plt
+
         fig, (evidence_plot, kl_plot) = plt.subplots(2)
         fig.suptitle("Evidence and KL Divergence")
         evidence_plot.plot(self.ep_history.evidences(), label="evidence")

autofit/graphical/factor_graphs/transform.py

@@ -2,7 +2,6 @@
 from typing import Dict, Tuple, Optional, List
 
 import numpy as np
-from scipy.linalg import cho_factor
 
 from autoconf import cached_property
 from autofit.graphical.factor_graphs.abstract import AbstractNode, Value, FactorValue
@@ -63,12 +62,16 @@ def from_scales(cls, scales):
 
     @classmethod
     def from_covariances(cls, covs):
+        from scipy.linalg import cho_factor
+
         return cls(
             {v: InvCholeskyTransform(cho_factor(cov)) for v, cov in covs.items()}
         )
 
     @classmethod
     def from_inv_covariances(cls, inv_covs):
+        from scipy.linalg import cho_factor
+
         return cls(
             {
                 v: CholeskyOperator(cho_factor(inv_cov))

autofit/graphical/laplace/line_search.py

@@ -11,7 +11,6 @@
 from typing import Optional, Dict, Tuple
 
 import numpy as np
-from scipy.optimize import _linesearch as linesearch
 
 from autoconf import cached_property
 from autofit.graphical.factor_graphs.abstract import (
@@ -280,6 +279,8 @@ def line_search_wolfe1(
     gval : array
         Gradient of `f` at the final point
     """
+    from scipy.optimize import _linesearch as linesearch
+
     derphi0 = state.derphi(0)
     old_fval = state.value
     stepsize, _, _ = linesearch.scalar_search_wolfe1(
@@ -339,6 +340,8 @@ def line_search_wolfe2(
     gval : array
         Gradient of `f` at the final point
     """
+    from scipy.optimize import _linesearch as linesearch
+
     derphi0 = state.derphi(0)
     old_fval = state.value
     stepsize, _, _, _ = linesearch.scalar_search_wolfe2(
@@ -364,6 +367,8 @@ def line_search_wolfe2(
 def line_search(
     state: OptimisationState, old_state: Optional[FactorValue] = None, **kwargs
 ) -> Tuple[Optional[float], OptimisationState]:
+    from scipy.optimize import _linesearch as linesearch
+
     stepsize, next_state = line_search_wolfe1(state, old_state, **kwargs)
 
     if stepsize is None:

autofit/graphical/utils.py

@@ -1,15 +1,16 @@
+from __future__ import annotations
 import logging
 import warnings
 from collections import abc
 from enum import Enum
 from functools import reduce
 from operator import mul
-from typing import Iterable, Tuple, TypeVar, Dict, NamedTuple, Optional, Union
+from typing import Iterable, Tuple, TypeVar, Dict, NamedTuple, Optional, Union, TYPE_CHECKING
 
 import numpy as np
-import six
-from scipy.linalg import block_diag
-from scipy.optimize import OptimizeResult
+
+if TYPE_CHECKING:
+    from scipy.optimize import OptimizeResult
 
 from autofit.mapper.variable import Variable, VariableData
 from autofit.non_linear.result import Result
@@ -50,6 +51,9 @@ def is_variable(v, *args):
 
 
 def is_iterable(arg):
+
+    import six
+
     return isinstance(arg, abc.Iterable) and not isinstance(
         arg, six.string_types
     )
@@ -369,6 +373,9 @@ def unflatten(self, arr: np.ndarray, ndim=None) -> Dict[str, np.ndarray]:
         })
 
     def flatten2d(self, values: Dict[Variable, np.ndarray]) -> np.ndarray:
+
+        from scipy.linalg import block_diag
+
         assert all(np.shape(values[k]) == shape * 2 for k, shape in self.items())
 
         return block_diag(
@@ -386,7 +393,9 @@ def size(self):
         return self.splits[-1]
 
 
+
 class OptResult(NamedTuple):
+
     mode: Dict[Variable, np.ndarray]
     hess_inv: Dict[Variable, np.ndarray]
     log_norm: float

autofit/interpolator/covariance.py

@@ -1,5 +1,4 @@
 import numpy as np
-import scipy
 from typing import List, Dict
 
 from autofit.non_linear.samples.pdf import SamplesPDF
@@ -102,8 +101,10 @@ def inverse_covariance_matrix(self) -> np.ndarray:
         """
         Calculate the inverse covariance matrix of the samples
         """
+        from scipy.linalg import inv
+
         matrices = [
-            scipy.linalg.inv(samples.covariance_matrix) for samples in self.samples_list
+            inv(samples.covariance_matrix) for samples in self.samples_list
         ]
         return self._subsume(matrices)
 

autofit/interpolator/linear.py

@@ -1,4 +1,3 @@
-from scipy.stats import linregress
 from .abstract import AbstractInterpolator
 from autofit.interpolator.linear_relationship import LinearRelationship
 
@@ -10,6 +9,9 @@ class LinearInterpolator(AbstractInterpolator):
 
     @staticmethod
     def _relationship(x, y):
+
+        from scipy.stats import linregress
+
         slope, intercept, r, p, std_err = linregress(x, y)
 
         return LinearRelationship(slope, intercept)

autofit/interpolator/spline.py

@@ -1,4 +1,3 @@
-from scipy.interpolate import CubicSpline
 from .abstract import AbstractInterpolator
 
 
@@ -9,4 +8,7 @@ class SplineInterpolator(AbstractInterpolator):
 
     @staticmethod
     def _relationship(x, y):
+
+        from scipy.interpolate import CubicSpline
+
         return CubicSpline(x, y)

autofit/jax_wrapper.py

@@ -28,7 +28,6 @@ def jit(function, *args, **kwargs):
     def grad(function, *args, **kwargs):
         return jax.grad(function, *args, **kwargs)
 
-    from jax._src.scipy.special import erfinv
 
 else:
 
@@ -41,7 +40,6 @@ def grad(function, *args, **kwargs):
     """)
 
     import numpy  # noqa
-    from scipy.special.cython_special import erfinv  # noqa
 
     def jit(function, *_, **__):
         return function

autofit/mapper/model_object.py

@@ -150,7 +150,6 @@ def from_dict(
         from autofit.mapper.prior_model.collection import Collection
         from autofit.mapper.prior_model.prior_model import Model
         from autofit.mapper.prior.abstract import Prior
-        from autofit.mapper.prior.gaussian import GaussianPrior
         from autofit.mapper.prior.tuple_prior import TuplePrior
         from autofit.mapper.prior.arithmetic.compound import Compound
         from autofit.mapper.prior.arithmetic.compound import ModifiedPrior