Feature/vram info

autofit/aggregator/summary/aggregate_csv/column.py

@@ -64,20 +64,34 @@ def value(self, row: "Row"):
         result = {}
 
         if ValueType.Median in self.value_types:
-            result[""] = row.median_pdf_sample_kwargs[self.path]
+            try:
+                result[""] = row.median_pdf_sample_kwargs[self.path]
+            except KeyError:
+                result[""] = None
 
         if ValueType.MaxLogLikelihood in self.value_types:
-            result["max_lh"] = row.max_likelihood_kwargs[self.path]
+            try:
+                result["max_lh"] = row.max_likelihood_kwargs[self.path]
+            except KeyError:
+                result["max_lh"] = None
 
         if ValueType.ValuesAt1Sigma in self.value_types:
-            lower, upper = row.values_at_sigma_1_kwargs[self.path]
-            result["lower_1_sigma"] = lower
-            result["upper_1_sigma"] = upper
+            try:
+                lower, upper = row.values_at_sigma_1_kwargs[self.path]
+                result["lower_1_sigma"] = lower
+                result["upper_1_sigma"] = upper
+            except KeyError:
+                result["lower_1_sigma"] = None
+                result["upper_1_sigma"] = None
 
         if ValueType.ValuesAt3Sigma in self.value_types:
-            lower, upper = row.values_at_sigma_3_kwargs[self.path]
-            result["lower_3_sigma"] = lower
-            result["upper_3_sigma"] = upper
+            try:
+                lower, upper = row.values_at_sigma_3_kwargs[self.path]
+                result["lower_3_sigma"] = lower
+                result["upper_3_sigma"] = upper
+            except KeyError:
+                result["lower_3_sigma"] = None
+                result["upper_3_sigma"] = None
 
         return result
 

autofit/aggregator/summary/aggregate_images.py

@@ -24,32 +24,11 @@ def subplot_filename(subplot: Enum) -> str:
     )
 
 
-class SubplotFit(Enum):
-    """
-    The subplots that can be extracted from the subplot_fit image.
-
-    The values correspond to the position of the subplot in the 4x3 grid.
-    """
-
-    Data = (0, 0)
-    DataSourceScaled = (1, 0)
-    SignalToNoiseMap = (2, 0)
-    ModelData = (3, 0)
-    LensLightModelData = (0, 1)
-    LensLightSubtractedImage = (1, 1)
-    SourceModelData = (2, 1)
-    SourcePlaneZoomed = (3, 1)
-    NormalizedResidualMap = (0, 2)
-    NormalizedResidualMapOneSigma = (1, 2)
-    ChiSquaredMap = (2, 2)
-    SourcePlaneNoZoom = (3, 2)
-
-
 class SubplotFitImage:
     def __init__(
         self,
         image: Image.Image,
-        suplot_type: Type[SubplotFit] = SubplotFit,
+        suplot_type,
     ):
         """
         The subplot_fit image associated with one fit.
@@ -173,7 +152,7 @@ def output_to_folder(
         self,
         folder: Path,
         name: Union[str, List[str]],
-        subplots: List[Union[SubplotFit, List[Image.Image], Callable]],
+        subplots: List[Union[List[Image.Image], Callable]],
         subplot_width: Optional[int] = sys.maxsize,
     ):
         """
@@ -226,7 +205,7 @@ def output_to_folder(
     def _matrix_for_result(
         i: int,
         result: SearchOutput,
-        subplots: List[Union[SubplotFit, List[Image.Image], Callable]],
+        subplots: List[Union[List[Image.Image], Callable]],
         subplot_width: int = sys.maxsize,
     ) -> List[List[Image.Image]]:
         """
@@ -272,7 +251,8 @@ class name but using snake_case.
                 _images[subplot_type] = SubplotFitImage(
                     result.image(
                         subplot_filename(subplot_),
-                    )
+                    ),
+                    subplot_type
                 )
             return _images[subplot_type]
 

autofit/graphical/declarative/abstract.py

@@ -20,7 +20,7 @@ class AbstractDeclarativeFactor(Analysis, ABC):
     _plates: Tuple[Plate, ...] = ()
 
     def __init__(self, include_prior_factors=False, use_jax : bool = False):
-        self.include_prior_factors = include_prior_factors
+        self._include_prior_factors = include_prior_factors
 
         super().__init__(use_jax=use_jax)
 
@@ -63,7 +63,7 @@ def prior_counts(self) -> List[Tuple[Prior, int]]:
             for prior in factor.prior_model.priors:
                 counter[prior] += 1
         return [
-            (prior, count + 1 if self.include_prior_factors else count)
+            (prior, count + 1 if self._include_prior_factors else count)
             for prior, count in counter.items()
         ]
 
@@ -96,7 +96,7 @@ def graph(self) -> DeclarativeFactorGraph:
         The complete graph made by combining all factors and priors
         """
         factors = [model for model in self.model_factors]
-        if self.include_prior_factors:
+        if self._include_prior_factors:
             factors += self.prior_factors
         # noinspection PyTypeChecker
         return DeclarativeFactorGraph(factors)

autofit/graphical/declarative/collection.py

@@ -42,7 +42,7 @@ def __init__(
     def tree_flatten(self):
         return (
             (self._model_factors,),
-            (self._name, self.include_prior_factors),
+            (self._name, self._include_prior_factors),
         )
 
     @classmethod

autofit/graphical/declarative/factor/abstract.py

@@ -42,7 +42,7 @@ def __init__(
             factor, **prior_variable_dict,
             name=name or namer(self.__class__.__name__),
         )
-        self.include_prior_factors = include_prior_factors
+        self._include_prior_factors = include_prior_factors
 
     @property
     def info(self) -> str:

autofit/graphical/declarative/factor/hierarchical.py

@@ -8,6 +8,7 @@
 from autofit.messages import NormalMessage
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.tools.namer import namer
+
 from .abstract import AbstractModelFactor
 
 
@@ -19,6 +20,7 @@ def __init__(
         distribution: Type[Prior],
         optimiser=None,
         name: Optional[str] = None,
+        use_jax : bool = False,
         **kwargs,
     ):
         """
@@ -70,6 +72,7 @@ def __init__(
         self._name = name or namer(self.__class__.__name__)
         self._factors = list()
         self.optimiser = optimiser
+        self._use_jax = use_jax
 
     @property
     def name(self):
@@ -144,7 +147,7 @@ def __call__(self, **kwargs):
 
 class _HierarchicalFactor(AbstractModelFactor):
     def __init__(
-        self, distribution_model: HierarchicalFactor, drawn_prior: Prior, use_jax : bool = False
+        self, distribution_model: HierarchicalFactor, drawn_prior: Prior,
     ):
         """
         A factor that links a variable to a parameterised distribution.
@@ -159,7 +162,7 @@ def __init__(
         """
         self.distribution_model = distribution_model
         self.drawn_prior = drawn_prior
-        self.use_jax = use_jax
+        self._use_jax = distribution_model._use_jax
 
         prior_variable_dict = {prior.name: prior for prior in distribution_model.priors}
 
@@ -188,7 +191,7 @@ def variable(self):
         return self.drawn_prior
 
     def log_likelihood_function(self, instance):
-        return instance.distribution_model.message(instance.drawn_prior)
+        return instance.distribution_model.message(instance.drawn_prior, xp=self._xp)
 
     @property
     def priors(self) -> Set[Prior]:

autofit/mapper/prior/arithmetic/assertion.py

@@ -1,4 +1,5 @@
 from abc import ABC
+import numpy as np
 from typing import Optional, Dict
 
 from autofit.mapper.prior.arithmetic.compound import CompoundPrior, Compound
@@ -24,7 +25,7 @@ def __le__(self, other):
 
 
 class GreaterThanLessThanAssertion(ComparisonAssertion):
-    def _instance_for_arguments(self, arguments, ignore_assertions=False):
+    def _instance_for_arguments(self, arguments, ignore_assertions=False, xp=np):
         """
         Assert that the value in the dictionary associated with the lower
         prior is lower than the value associated with the greater prior.
@@ -55,6 +56,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         """
         Assert that the value in the dictionary associated with the lower
@@ -90,6 +92,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.assertion_1.instance_for_arguments(
             arguments,

autofit/mapper/prior/arithmetic/compound.py

@@ -212,6 +212,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -237,6 +238,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -256,6 +258,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -275,6 +278,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -294,6 +298,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -313,6 +318,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return self.left_for_arguments(
             arguments,
@@ -396,6 +402,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return -self.prior.instance_for_arguments(
             arguments,
@@ -412,6 +419,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return abs(
             self.prior.instance_for_arguments(
@@ -430,6 +438,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return np.log(
             self.prior.instance_for_arguments(
@@ -448,6 +457,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         return np.log10(
             self.prior.instance_for_arguments(

autofit/mapper/prior/gaussian.py

@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Optional
 
 from autofit.messages.normal import NormalMessage
@@ -46,6 +47,7 @@ def __init__(
         >>> prior = GaussianPrior(mean=1.0, sigma=2.0)
         >>> physical_value = prior.value_for(unit=0.5)  # Returns ~1.0 (mean)
         """
+
         super().__init__(
             message=NormalMessage(
                 mean=mean,

autofit/mapper/prior/log_gaussian.py

@@ -133,7 +133,7 @@ def value_for(self, unit: float) -> float:
     def parameter_string(self) -> str:
         return f"mean = {self.mean}, sigma = {self.sigma}"
 
-    def log_prior_from_value(self, value):
+    def log_prior_from_value(self, value, xp=np):
         if value <= 0:
             return float("-inf")
 

autofit/mapper/prior/log_uniform.py

@@ -110,7 +110,7 @@ def with_limits(cls, lower_limit: float, upper_limit: float) -> "LogUniformPrior
 
     __identifier_fields__ = ("lower_limit", "upper_limit")
 
-    def log_prior_from_value(self, value) -> float:
+    def log_prior_from_value(self, value, xp=np) -> float:
         """
         Returns the log prior of a physical value, so the log likelihood of a model evaluation can be converted to a
         posterior as log_prior + log_likelihood.

autofit/mapper/prior/uniform.py

@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Optional, Tuple
 
 from autofit.messages.normal import UniformNormalMessage
@@ -128,7 +129,7 @@ def value_for(self, unit: float) -> float:
             round(super().value_for(unit), 14)
         )
 
-    def log_prior_from_value(self, value):
+    def log_prior_from_value(self, value, xp=np):
         """
         Returns the log prior of a physical value, so the log likelihood of a model evaluation can be converted to a
         posterior as log_prior + log_likelihood.