Repository: rhayes777/PyAutoFit
Pull request #974 (2024-03-19 15:28:55)

feature/arithmetic fixes
Resolves issues with arithmetic



Patch:
diff --git a/autofit/mapper/model_object.py b/autofit/mapper/model_object.py
index 335f5c909..aea48e872 100644
--- a/autofit/mapper/model_object.py
+++ b/autofit/mapper/model_object.py
@@ -152,6 +152,7 @@ class ModelObject:
         from autofit.mapper.prior.abstract import Prior
         from autofit.mapper.prior.tuple_prior import TuplePrior
         from autofit.mapper.prior.arithmetic.compound import Compound
+        from autofit.mapper.prior.arithmetic.compound import ModifiedPrior
 
         if isinstance(d, list):
             return [
@@ -195,6 +196,12 @@ class ModelObject:
                 reference=dereference(reference, "assertion"),
                 loaded_ids=loaded_ids,
             )
+        elif type_ == "modified":
+            return ModifiedPrior.from_dict(
+                d,
+                reference=dereference(reference, "prior"),
+                loaded_ids=loaded_ids,
+            )
         elif type_ == "dict":
             return {
                 key: from_dict(
diff --git a/autofit/mapper/prior/arithmetic/compound.py b/autofit/mapper/prior/arithmetic/compound.py
index a0aa70a53..98952fc74 100644
--- a/autofit/mapper/prior/arithmetic/compound.py
+++ b/autofit/mapper/prior/arithmetic/compound.py
@@ -6,6 +6,7 @@ from typing import Optional, Dict
 
 import numpy as np
 
+from autofit.mapper.model_object import dereference
 from autofit.mapper.prior.arithmetic import ArithmeticMixin
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
 
@@ -195,6 +196,9 @@ class CompoundPrior(AbstractPriorModel, ArithmeticMixin, Compound, ABC):
         except AttributeError:
             return self._right
 
+    def __add__(self, other):
+        return ArithmeticMixin.__add__(self, other)
+
 
 class SumPrior(CompoundPrior):
     """
@@ -316,16 +320,46 @@ class PowerPrior(CompoundPrior):
         )
 
 
-class ModifiedPrior(AbstractPriorModel, ABC, ArithmeticMixin):
-    def __init__(self, prior):
+class ModifiedPrior(AbstractPriorModel, ABC, ArithmeticMixin, Compound):
+    def __init__(self, prior, name=None):
         super().__init__()
-        self._prior_name = retrieve_name(prior)
+        self._prior_name = name or retrieve_name(prior)
 
         if self._prior_name == "prior":
             self._prior_name = "prior_"
 
         self.prior = prior
 
+    def dict(self):
+        return {
+            "type": "modified",
+            "modified_type": self.__class__.__name__,
+            "name": self._prior_name,
+            "prior": self.prior.dict()
+            if isinstance(self.prior, AbstractPriorModel)
+            else self.prior,
+        }
+
+    @classmethod
+    def from_dict(
+        cls,
+        d,
+        reference: Optional[Dict[str, str]] = None,
+        loaded_ids: Optional[dict] = None,
+    ):
+        modified_type = d.pop("modified_type")
+        for subclass in cls.descendants():
+            if subclass.__name__ == modified_type:
+                return subclass(
+                    AbstractPriorModel.from_dict(
+                        d["prior"],
+                        reference=dereference(reference, "prior"),
+                        loaded_ids=loaded_ids,
+                    ),
+                    name=d["name"],
+                )
+        raise ValueError(f"Modified type {modified_type} not recognised")
+
     def __add__(self, other):
         return ArithmeticMixin.__add__(self, other)
 


Test patch:
diff --git a/test_autofit/serialise/test_regression.py b/test_autofit/serialise/test_regression.py
index 642fd710c..4f7e72d10 100644
--- a/test_autofit/serialise/test_regression.py
+++ b/test_autofit/serialise/test_regression.py
@@ -1,5 +1,7 @@
 import json
 
+import pytest
+
 from autoconf.dictable import to_dict, from_dict
 from autofit.mapper.mock.mock_model import MockClassx2Tuple
 import autofit as af
@@ -18,5 +20,28 @@ class AttributeNotInInit:
 def test_embedded():
     model = af.Collection(instance=AttributeNotInInit())
 
-    print(to_dict(model))
     from_dict(to_dict(model))
+
+
+@pytest.fixture(name="alpha")
+def make_alpha():
+    return af.GaussianPrior(mean=1, sigma=0.03)
+
+
+@pytest.fixture(name="gamma")
+def make_gamma():
+    return 0.0
+
+
+def test_serialise_compound(alpha, gamma):
+    beta = 1.0 + gamma - 2 * alpha
+    assert from_dict(to_dict(beta)).instance_from_prior_medians() == -1
+
+
+def test_add_gamma(alpha, gamma):
+    beta = 1.0 + -2 * alpha + gamma
+    assert beta.instance_from_prior_medians() == -1
+
+
+def test_negative_addition(alpha):
+    assert (-alpha + 1).instance_from_prior_medians() == 0.0