update

src/lamatrix/__init__.py

@@ -19,3 +19,29 @@
 from .models.gaussian import *  # noqa: E402, F401
 from .models.simple import *  # noqa: E402, F401
 from .models.spline import *  # noqa: E402, F401
+
+import json
+import numpy as np
+
+def load(filename):
+    def process(arg):
+        if isinstance(arg, dict):
+            return {key:process(item) for key, item in arg.items()}
+        if arg is None:
+            return None
+        elif isinstance(arg, str):
+            if arg == "Infinity":
+                return np.nan
+            return arg
+        elif isinstance(arg, (int, float, tuple)):
+            return arg
+        elif isinstance(arg, list):
+            return np.asarray(arg)
+
+    """Load a Generator object"""
+    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

src/lamatrix/generator.py

@@ -3,11 +3,12 @@
 import math
 from abc import ABC, abstractmethod
 from copy import deepcopy
+import json
 
 import numpy as np
 
 __all__ = [
-    "Generator",
+    "Generator"
 ]
 
 
@@ -17,7 +18,7 @@ def _validate_arg_names(self):
             if not isinstance(arg, str):
                 raise ValueError("Argument names must be strings.")
 
-    def _validate_priors(self, prior_mu, prior_sigma):
+    def _validate_priors(self, prior_mu, prior_sigma, offset_prior=None):
         if prior_mu is None:
             self.prior_mu = np.zeros(self.width)
         else:
@@ -39,6 +40,17 @@ def _validate_priors(self, prior_mu, prior_sigma):
                     self.prior_sigma = prior_sigma
             else:
                 raise ValueError("Can not parse `prior_sigma`.")
+
+
+        if 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:
@@ -49,10 +61,35 @@ def _validate_priors(self, prior_mu, prior_sigma):
     #     return new
 
     def save(self, filename: str):
-        raise NotImplementedError
-
-    def load(self, filename: str):
-        raise NotImplementedError
+        def process(arg):
+            if arg is None:
+                return None
+            elif isinstance(arg, (str, int, float, list, tuple)):
+                if arg is np.inf:
+                    return "Infinity"
+                return arg
+            elif isinstance(arg, np.ndarray):
+                arg = arg.tolist()
+                arg = [a if a != np.inf else "Infinity" for a in arg]
+                return arg
+
+
+
+        results = {attr:process(getattr(self, attr)) for attr in ["fit_mu", "fit_sigma"]}
+        kwargs = {attr:process(getattr(self, attr)) for attr in self._INIT_ATTRS}
+        type_name = type(self).__name__
+
+        data_to_store = {
+            "object_type": type_name,
+            "initializing_kwargs": kwargs,
+            "fit_results": results,
+            "equation": self.equation,
+        }
+        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)
 
     def copy(self):
         return deepcopy(self)

src/lamatrix/models/gaussian.py

@@ -10,14 +10,14 @@
     "dlnGaussian2DGenerator",
 ]
 
-
 class lnGaussian2DGenerator(MathMixins, Generator):
     def __init__(
         self,
         x_name: str = "x",
         y_name: str = "y",
         prior_mu=None,
         prior_sigma=None,
+        offset_prior=None,
         stddev_x_prior=None,
         stddev_y_prior=None,
         data_shape=None,
@@ -26,7 +26,7 @@ def __init__(
         self.y_name = y_name
         self._validate_arg_names()
         self.data_shape = data_shape
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
         self.stddev_x_prior, self.stddev_y_prior = stddev_x_prior, stddev_y_prior
@@ -206,6 +206,7 @@ def __init__(
         y_name: str = "y",
         prior_mu=None,
         prior_sigma=None,
+        offset_prior=None,
         data_shape=None,
     ):
         self.stddev_x = stddev_x
@@ -215,7 +216,7 @@ def __init__(
         self.y_name = y_name
         self._validate_arg_names()
         self.data_shape = data_shape
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
 
@@ -284,6 +285,6 @@ def table_properties(self):
 
     @property
     def _equation(self):
-        dfdx = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{(y-\\mu_y)}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
-        dfdy = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{(y-\\mu_y)}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
+        dfdx = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
+        dfdy = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_x^2}} - \\rho\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x\\sigma_y}}\\right)\\right)"
         return ["", dfdx, dfdy]

src/lamatrix/models/simple.py

@@ -25,19 +25,10 @@ def __init__(
         self._validate_arg_names()
         self.polyorder = polyorder
         self.data_shape = data_shape
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
 
-        if offset_prior is not None:
-            if not hasattr(offset_prior, "__iter__"):
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-            if not len(offset_prior) == 2:
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-
-            self.prior_mu[0] = offset_prior[0]
-            self.prior_sigma[0] = offset_prior[1]
-
     @property
     def width(self):
         return self.polyorder + 1
@@ -50,6 +41,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "prior_mu", "prior_sigma", "offset_prior", "data_shape", "polyorder"]
+
     def design_matrix(self, *args, **kwargs):
         """Build a 1D polynomial in x
 
@@ -91,13 +86,14 @@ def __init__(
         nterms: int = 1,
         prior_mu=None,
         prior_sigma=None,
+        offset_prior=None,
         data_shape=None,
     ):
         self.nterms = nterms
         self.data_shape = data_shape
         self.x_name = x_name
         self._validate_arg_names()
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
 
@@ -113,6 +109,10 @@ def nvectors(self):
     def arg_names(self):
         return {self.x_name}
 
+    @property
+    def _INIT_ATTRS(self):
+        return ["x_name", "prior_mu", "prior_sigma", "offset_prior", "data_shape", "polyorder"]
+
     def design_matrix(self, *args, **kwargs):
         """Build a 1D polynomial in x
 

src/lamatrix/models/spline.py

@@ -86,19 +86,10 @@ def __init__(
         self._validate_arg_names()
         self.splineorder = splineorder
         self.data_shape = data_shape
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
 
-        if offset_prior is not None:
-            if not hasattr(offset_prior, "__iter__"):
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-            if not len(offset_prior) == 2:
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-
-            self.prior_mu[0] = offset_prior[0]
-            self.prior_sigma[0] = offset_prior[1]
-
     @property
     def width(self):
         return len(self.knots) - self.splineorder - 1 + 1
@@ -187,7 +178,6 @@ def __init__(
         splineorder: int = 3,
         prior_mu=None,
         prior_sigma=None,
-        offset_prior=None,
         data_shape=None,
     ):
         # Check if knots are padded
@@ -203,19 +193,10 @@ def __init__(
         self._validate_arg_names()
         self.splineorder = splineorder
         self.data_shape = data_shape
-        self._validate_priors(prior_mu, prior_sigma)
+        self._validate_priors(prior_mu, prior_sigma, offset_prior=offset_prior)
         self.fit_mu = None
         self.fit_sigma = None
 
-        if offset_prior is not None:
-            if not hasattr(offset_prior, "__iter__"):
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-            if not len(offset_prior) == 2:
-                raise AttributeError("Pass offset prior as a tuple with (mu, sigma)")
-
-            self.prior_mu[0] = offset_prior[0]
-            self.prior_sigma[0] = offset_prior[1]
-
     @property
     def width(self):
         return 2