Feature/jax messages

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/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.

autofit/mapper/prior_model/abstract.py

@@ -745,7 +745,7 @@ def physical_values_from_prior_medians(self):
         """
         return self.vector_from_unit_vector([0.5] * len(self.unique_prior_tuples))
 
-    def instance_from_vector(self, vector, ignore_assertions: bool = False):
+    def instance_from_vector(self, vector, ignore_assertions: bool = False, xp=np):
         """
         Returns a ModelInstance, which has an attribute and class instance corresponding
         to every `Model` attributed to this instance.
@@ -778,6 +778,7 @@ def instance_from_vector(self, vector, ignore_assertions: bool = False):
         return self.instance_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp
         )
 
     def has(self, cls: Union[Type, Tuple[Type, ...]]) -> bool:
@@ -1078,6 +1079,7 @@ def instance_from_prior_medians(self, ignore_assertions : bool = False):
     def log_prior_list_from_vector(
         self,
         vector: [float],
+        xp=np,
     ):
         """
         Compute the log priors of every parameter in a vector, using the Prior of every parameter.
@@ -1094,7 +1096,7 @@ def log_prior_list_from_vector(
         return list(
             map(
                 lambda prior_tuple, value: prior_tuple.prior.log_prior_from_value(
-                    value=value
+                    value=value, xp=xp
                 ),
                 self.prior_tuples_ordered_by_id,
                 vector,
@@ -1308,13 +1310,15 @@ def _instance_for_arguments(
         self,
         arguments: Dict[Prior, float],
         ignore_assertions: bool = False,
+        xp=np,
     ):
         raise NotImplementedError()
 
     def instance_for_arguments(
         self,
         arguments: Dict[Prior, float],
         ignore_assertions: bool = False,
+        xp=np,
     ):
         """
         Returns an instance of the model for a set of arguments
@@ -1338,6 +1342,7 @@ def instance_for_arguments(
         return self._instance_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp
         )
 
     def path_for_name(self, name: str) -> Tuple[str, ...]:

autofit/mapper/prior_model/array.py

@@ -57,6 +57,7 @@ def _instance_for_arguments(
         self,
         arguments: Dict[Prior, float],
         ignore_assertions: bool = False,
+        xp=np,
     ) -> np.ndarray:
         """
         Create an array where the prior at each index is replaced with the

autofit/mapper/prior_model/collection.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from collections.abc import Iterable
 
 from autofit.mapper.model import ModelInstance, assert_not_frozen
@@ -208,6 +210,7 @@ def _instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         """
         Parameters
@@ -228,6 +231,7 @@ def _instance_for_arguments(
                 value = value.instance_for_arguments(
                     arguments,
                     ignore_assertions=ignore_assertions,
+                    xp=xp
                 )
             elif isinstance(value, Prior):
                 value = arguments[value]