save

src/lamatrix/__init__.py

@@ -14,7 +14,7 @@
 log = logging.getLogger("tesswcs")
 log.addHandler(RichHandler(markup=True))
 
-from .combined import *  # noqa: E402, F401
+from .combine import *  # noqa: E402, F401
 from .models.astrophysical import *  # noqa: E402, F401
 from .models.gaussian import *  # noqa: E402, F401
 from .models.simple import *  # noqa: E402, F401
@@ -23,6 +23,11 @@
 import json
 import numpy as np
 
+def _load_from_dict(dict):
+    new = globals()[dict['object_type']](**dict['initializing_kwargs'])
+    _ = [setattr(new, key, value) for key, value in dict['fit_results'].items()]
+    return new
+
 def load(filename):
     def process(arg):
         if isinstance(arg, dict):
@@ -42,6 +47,9 @@ def process(arg):
     with open(filename, 'r') as json_file:
         data_loaded = json.load(json_file)
     data_loaded = {key:process(item) for key, item in data_loaded.items()}
-    new = globals()[data_loaded['object_type']](**data_loaded['initializing_kwargs'])
-    _ = [setattr(new, key, value) for key, value in data_loaded['fit_results'].items()]
-    return new
+    if 'generators' in data_loaded.keys():
+        generators = [_load_from_dict(item) for _, item in data_loaded['generators'].items()]
+        new = globals()[data_loaded['object_type']](*generators)
+        _ = [setattr(new, key, value) for key, value in data_loaded['fit_results'].items()]
+        return new
+    return _load_from_dict(data_loaded)

src/lamatrix/combined.py → src/lamatrix/combine.py

@@ -1,8 +1,8 @@
 import numpy as np
-
+import json
 from .generator import Generator
 
-__all__ = ["VStackedGenerator"]
+__all__ = ["StackedIndependentGenerator", "StackedDependentGenerator"]
 
 
 def combine_equations(*equations):
@@ -53,7 +53,7 @@ def combine_matrices(*matrices):
         return np.hstack(combined)
 
 
-class VStackedGenerator(Generator):
+class StackedIndependentGenerator(Generator):
     def __init__(self, *args, **kwargs):
         if (
             not len(np.unique([a.data_shape for a in args if a.data_shape is not None]))
@@ -62,6 +62,8 @@ def __init__(self, *args, **kwargs):
             raise ValueError("Can not have different `data_shape`.")
         self.generators = [a.copy() for a in args]
         self.data_shape = self.generators[0].data_shape
+        self.fit_mu = None
+        self.fit_sigma = None
 
     def __getitem__(self, key):
         return self.generators[key]
@@ -156,16 +158,32 @@ def _equation(self):
     def arg_names(self):
         return np.unique(np.hstack([list(g.arg_names) for g in self.generators]))
 
-
-class CombinedGenerator(Generator):
-    def __init__(self, *args, **kwargs):
-        if (
-            not len(np.unique([a.data_shape for a in args if a.data_shape is not None]))
-            <= 1
-        ):
-            raise ValueError("Can not have different `data_shape`.")
-        self.generators = [a.copy() for a in args]
-        self.data_shape = self.generators[0].data_shape
+    @property
+    def _INIT_ATTRS(self):
+        return []
+
+    def save(self, filename: str):
+        if not filename.endswith(".json"):
+            filename = filename + ".json"
+
+        # Write to a JSON file
+        with open(filename, "w") as json_file:
+            data_to_store = self._create_save_data()
+            generators_to_store = {f"generator{idx+1}":g._create_save_data() for idx, g in enumerate(self.generators)}
+            data_to_store["generators"] = generators_to_store
+            json.dump(data_to_store, json_file, indent=4)
+
+class StackedDependentGenerator(StackedIndependentGenerator):
+    # def __init__(self, *args, **kwargs):
+    #     if (
+    #         not len(np.unique([a.data_shape for a in args if a.data_shape is not None]))
+    #         <= 1
+    #     ):
+    #         raise ValueError("Can not have different `data_shape`.")
+    #     self.generators = [a.copy() for a in args]
+    #     self.data_shape = self.generators[0].data_shape
+    #     self.fit_mu = None
+    #     self.fit_sigma = None
 
     @property
     def width(self):
@@ -202,3 +220,19 @@ def prior_mu(self):
 
     def fit(self, *args, **kwargs):
         self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)
+
+    # @property
+    # def _INIT_ATTRS(self):
+    #     return []
+
+    # def save(self, filename: str):
+    #     if not filename.endswith(".json"):
+    #         filename = filename + ".json"
+
+    #     # Write to a JSON file
+    #     with open(filename, "w") as json_file:
+    #         data_to_store = self._create_save_data()
+    #         json.dump(data_to_store, json_file, indent=4)
+    #         for g in self.generators:
+    #             data_to_store = g._create_save_data()
+    #             json.dump(data_to_store, json_file, indent=4)

src/lamatrix/generator.py

@@ -41,16 +41,15 @@ def _validate_priors(self, prior_mu, prior_sigma, offset_prior=None):
             else:
                 raise ValueError("Can not parse `prior_sigma`.")
 
-
-        if offset_prior is not None:
+        self.offset_prior = offset_prior
+        if self.offset_prior is not None:
             if not hasattr(self.offset_prior, "__iter__"):
                 raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
             if not len(self.offset_prior) == 2:
                 raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
 
             self.prior_mu[0] = self.offset_prior[0]
             self.prior_sigma[0] = self.offset_prior[1]
-        self.offset_prior = offset_prior
 
     # def update_priors(self):
     #     if self.fit_mu is None:
@@ -60,7 +59,7 @@ def _validate_priors(self, prior_mu, prior_sigma, offset_prior=None):
     #     new.prior_sigma = new.fit_sigma.copy()
     #     return new
 
-    def save(self, filename: str):
+    def _create_save_data(self):
         def process(arg):
             if arg is None:
                 return None
@@ -85,8 +84,13 @@ def process(arg):
             "fit_results": results,
             "equation": self.equation,
         }
+        return data_to_store
+
+    def save(self, filename: str):
+        data_to_store = self._create_save_data()
         if not filename.endswith(".json"):
             filename = filename + ".json"
+
         # Write to a JSON file
         with open(filename, "w") as json_file:
             json.dump(data_to_store, json_file, indent=4)
@@ -260,3 +264,9 @@ def width(self):
     def fit(self):
         """Fits the design matrix, given input vectors and data"""
         pass
+
+    @property
+    @abstractmethod
+    def _INIT_ATTRS(self):
+        """Defines the variables needed to initialize self, so that they can be stored when saved."""
+        pass

src/lamatrix/math.py

@@ -1,4 +1,4 @@
-from .combined import VStackedGenerator, CombinedGenerator
+from .combine import StackedIndependentGenerator, StackedDependentGenerator
 from .generator import Generator
 
 __all__ = ["MathMixins"]
@@ -7,12 +7,12 @@
 class MathMixins:
     def __add__(self, other):
         if isinstance(other, Generator):
-            return VStackedGenerator(self, other)
+            return StackedIndependentGenerator(self, other)
         else:
             raise ValueError("Can only combine `Generator` objects.")
 
     def __mul__(self, other):
         if isinstance(other, Generator):
-            return CombinedGenerator(self, other)
+            return StackedDependentGenerator(self, other)
         else:
             raise ValueError("Can only combine `Generator` objects.")

src/lamatrix/models/gaussian.py

@@ -64,6 +64,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name, self.y_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "y_name", "stddev_x_prior", "stddev_y_prior", "prior_mu", "prior_sigma", "offset_prior", "data_shape", "nterms"]
+
     def design_matrix(self, *args, **kwargs):
         """Build a 1D polynomial in x
 
@@ -232,6 +236,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name, self.y_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "y_name", "stddev_x", "stddev_y", "prior_mu", "rho", "prior_sigma", "offset_prior", "data_shape", "nterms"]
+
     def design_matrix(self, *args, **kwargs):
         """Build a 1D polynomial in x
 

src/lamatrix/models/simple.py

@@ -111,7 +111,7 @@ def arg_names(self):
 
     @property
     def _INIT_ATTRS(self):
-        return ["x_name", "prior_mu", "prior_sigma", "offset_prior", "data_shape", "polyorder"]
+        return ["x_name", "prior_mu", "prior_sigma", "offset_prior", "data_shape", "nterms"]
 
     def design_matrix(self, *args, **kwargs):
         """Build a 1D polynomial in x

src/lamatrix/models/spline.py

@@ -102,6 +102,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "knots", "splineorder", "prior_mu", "prior_sigma", "offset_prior", "data_shape"]
+
     def design_matrix(self, *args, **kwargs):
         """Build a 1D spline in x
 
@@ -169,13 +173,13 @@ def derivative(self):
 
 
 class dSpline1DGenerator(MathMixins, SplineMixins, Generator):
-
     def __init__(
         self,
         weights: np.ndarray,
         knots: np.ndarray,
         x_name: str = "x",
         splineorder: int = 3,
+        offset_prior=None,
         prior_mu=None,
         prior_sigma=None,
         data_shape=None,
@@ -209,6 +213,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "weights", "knots", "splineorder", "prior_mu", "prior_sigma", "offset_prior", "data_shape"]
+
     def design_matrix(self, *args, **kwargs):
         if not self.arg_names.issubset(set(kwargs.keys())):
             raise ValueError(f"Expected {self.arg_names} to be passed.")

test.json

@@ -0,0 +1,61 @@
+{
+    "object_type": "StackedIndependentGenerator",
+    "initializing_kwargs": {},
+    "fit_results": {
+        "fit_mu": null,
+        "fit_sigma": null
+    },
+    "equation": "\\[f(\\mathbf{c}, \\mathbf{r}) = w_{0}  + w_{1} \\mathbf{c}^{1} + w_{2} \\mathbf{c}^{2} + w_{3}  + w_{4} \\mathbf{r}^{1} + w_{5} \\mathbf{r}^{2} + w_{6} \\mathbf{r}^{3}\\]",
+    "generators": {
+        "generator1": {
+            "object_type": "Polynomial1DGenerator",
+            "initializing_kwargs": {
+                "x_name": "c",
+                "prior_mu": [
+                    0.0,
+                    0.0,
+                    0.0
+                ],
+                "prior_sigma": [
+                    "Infinity",
+                    "Infinity",
+                    "Infinity"
+                ],
+                "offset_prior": null,
+                "data_shape": null,
+                "polyorder": 2
+            },
+            "fit_results": {
+                "fit_mu": null,
+                "fit_sigma": null
+            },
+            "equation": "\\[f(\\mathbf{c}) = w_{0}  + w_{1} \\mathbf{c}^{1} + w_{2} \\mathbf{c}^{2}\\]"
+        },
+        "generator2": {
+            "object_type": "Polynomial1DGenerator",
+            "initializing_kwargs": {
+                "x_name": "r",
+                "prior_mu": [
+                    0.0,
+                    0.0,
+                    0.0,
+                    0.0
+                ],
+                "prior_sigma": [
+                    "Infinity",
+                    "Infinity",
+                    "Infinity",
+                    "Infinity"
+                ],
+                "offset_prior": null,
+                "data_shape": null,
+                "polyorder": 3
+            },
+            "fit_results": {
+                "fit_mu": null,
+                "fit_sigma": null
+            },
+            "equation": "\\[f(\\mathbf{r}) = w_{0}  + w_{1} \\mathbf{r}^{1} + w_{2} \\mathbf{r}^{2} + w_{3} \\mathbf{r}^{3}\\]"
+        }
+    }
+}

tests/test_generator.py

@@ -1,12 +1,12 @@
 import numpy as np
-import pytest
 
 from lamatrix import (
     Polynomial1DGenerator,
     Spline1DGenerator,
-    VStackedGenerator,
+    StackedIndependentGenerator,
     dlnGaussian2DGenerator,
     lnGaussian2DGenerator,
+    load,
 )
 
 
@@ -57,7 +57,7 @@ def Gauss2D(column, row, A, sigma_x, sigma_y, rho):
     assert np.isclose(-0.01, dg.shift_x[0], atol=dg.shift_x[1] * 2)
     assert np.isclose(0.02, dg.shift_y[0], atol=dg.shift_y[1] * 2)
 
-    c = VStackedGenerator(g, dg)
+    c = StackedIndependentGenerator(g, dg)
     c.fit(column=column, row=row, data=data, errors=errors)
     assert np.isclose(sigma_x, c[0].stddev_x[0], atol=c[0].stddev_x[1] * 2)
     assert np.isclose(sigma_y, c[0].stddev_y[0], atol=c[0].stddev_y[1] * 2)
@@ -89,7 +89,7 @@ def test_polycombine():
 
     p1 = Polynomial1DGenerator("r", data_shape=c.shape)
     p2 = Polynomial1DGenerator("c", data_shape=c.shape)
-    for p in [VStackedGenerator(p1, p2), (p1 + p2)]:
+    for p in [StackedIndependentGenerator(p1, p2), (p1 + p2)]:
         true_w = np.random.normal(0, 1, size=(p.width))
         data = p.design_matrix(c=c, r=r).dot(true_w).reshape(r.shape)
         errors = np.ones_like(r) + 10
@@ -131,3 +131,13 @@ def test_spline():
     model.fit(x=x, y=y, data=data)
 
     assert np.allclose(true_w, model.mu)
+
+def test_save():
+    p1 = Polynomial1DGenerator('c', polyorder=2)
+    p2 = Polynomial1DGenerator('r')
+    p = p1 + p2
+    p.save('test.json')
+    p = load('test.json')
+    assert p[0].x_name == 'c'
+    assert p[1].x_name == 'r'
+    assert p[0].polyorder == 2