working on docs, broken tho

docs/use-cases/README.md

@@ -0,0 +1 @@
+Here I show some use cases for `LAmatrix` as jupyter notebooks.

mkdocs.yml

@@ -0,0 +1,23 @@
+site_name: lamatrix
+nav:
+    - Home: README.md
+
+theme:
+  name: "material"
+  icon:
+      logo: material/matrix
+  palette:
+    scheme: default
+    primary: deep purple
+    accent: indigo
+repo_url: https://github.com/christinahedges/lamatrix
+plugins:
+  - search
+markdown_extensions:
+  - pymdownx.highlight:
+      anchor_linenums: true
+      line_spans: __span
+      pygments_lang_class: true
+  - pymdownx.inlinehilite
+  - pymdownx.snippets
+  - pymdownx.superfences

pyproject.toml

@@ -19,6 +19,10 @@ mypy = "^1.9.0"
 pytest = "^8.1.1"
 jupyterlab = "^4.1.4"
 matplotlib = "^3.8.3"
+mkdocs = "^1.5.3"
+mkdocs-material = "^9.5.14"
+mkdocstrings = "^0.24.1"
+pytkdocs = {version = "^0.16.1", extras = ["numpy-style"]}
 
 [build-system]
 requires = ["poetry-core"]

src/lamatrix/bound.py

@@ -1,14 +1,18 @@
 import numpy as np
 
 from .combine import StackedIndependentGenerator
+from .math import MathMixins
 
 __all__ = ["BoundedGenerator"]
 
 
-class BoundedGenerator(StackedIndependentGenerator):
-    def __init__(self, generator, bounds: list):
+class BoundedGenerator(MathMixins, StackedIndependentGenerator):
+    def __init__(self, generator, bounds: list, x_name=None):
         self.generator = generator
-        self.x_name = self.generator.x_name
+        if x_name is None:
+            self.x_name = self.generator.x_name
+        else:
+            self.x_name = x_name
         if isinstance(bounds, slice):
             self._latex_bounds = self._slice_bounds_to_latex(
                 bounds.start, bounds.stop, bounds.step
@@ -22,6 +26,7 @@ def __init__(self, generator, bounds: list):
         self.bounds = bounds
         self.fit_mu = None
         self.fit_sigma = None
+        self.data_shape = None
 
     def _bounds_to_latex(self, bounds):
         bound_latex = [
@@ -74,7 +79,7 @@ def prior_sigma(self):
 
     @property
     def arg_names(self):
-        return self.generator.arg_names
+        return {*np.unique([self.x_name, *list(self.generator.arg_names)])}
 
     @property
     def width(self):
@@ -118,7 +123,7 @@ def design_matrix(self, *args, **kwargs):
             bounds_list = self.bounds
         return np.hstack(
             [
-                self.generator.design_matrix(x=x) * ((x >= b[0]) & (x < b[1]))[:, None]
+                self.generator.design_matrix(*args, **kwargs) * ((x >= b[0]) & (x < b[1]))[:, None]
                 for b in bounds_list
             ]
         )
@@ -155,9 +160,9 @@ def __getitem__(self, key):
         attrs = ["fit_mu", "fit_sigma"]
         for attr in attrs:
             setattr(
-                g, attr, getattr(self, attr).reshape((self.nbounds, len(self)))[key]
+                g, attr, getattr(self, attr).reshape((self.nbounds, self.width // self.nbounds))[key]
             )
         return g
 
     def __len__(self):
-        return self.width // self.nbounds
+        return self.nbounds

src/lamatrix/combine.py

@@ -1,6 +1,7 @@
 import json
 
 import numpy as np
+import re
 
 from . import _META_DATA
 from .generator import Generator
@@ -28,6 +29,13 @@ def combine_sigmas(*sigmas):
     if len(sigmas) == 1:
         return sigmas[0]
 
+    if (np.isfinite(sigmas[0])).any():
+        if sigmas[1][0] == np.inf:
+            sigmas[1][0] = 0
+    if (np.isfinite(sigmas[1])).any():
+        if sigmas[0][0] == np.inf:
+            sigmas[0][0] = 0
+
     # Step case: combine the first two equations and recursively call the function with the result
     combined = [(f**2 + e**2) ** 0.5 for f in sigmas[1] for e in sigmas[0]]
 
@@ -72,10 +80,27 @@ def __repr__(self):
         str1 = (
             f"{type(self).__name__}({', '.join(list(self.arg_names))})[n, {self.width}]"
         )
-        str2 = [f"\t{g.__repr__()}" for g in self.generators]
+        def add_tab_to_runs_of_tabs(repr_str):
+            pattern = r'\t+'
+            replacement = lambda match: match.group(0) + '\t'
+            result_string = re.sub(pattern, replacement, repr_str)
+            return result_string
 
+        str2 = [f"\t{add_tab_to_runs_of_tabs(g.__repr__())}" for g in self.generators]
         return "\n".join([str1, *str2])
 
+    def __add__(self, other):
+        if isinstance(other, Generator):
+            return StackedIndependentGenerator(self, other)
+        else:
+            raise ValueError("Can only combine `Generator` objects.")
+
+    def __mul__(self, other):
+        if isinstance(other, Generator):
+            return StackedDependentGenerator(self, other)
+        else:
+            raise ValueError("Can only combine `Generator` objects.")
+
     def __getitem__(self, key):
         return self.generators[key]
 
@@ -121,31 +146,42 @@ def prior_sigma(self):
                 pm[0] = 0
             else:
                 pm[0] = (
-                    np.sum(np.asarray([g.prior_sigma[0] for g in self.generators]) ** 2)
+                    np.nansum(np.asarray([g.prior_sigma[0] if g.prior_sigma[0] != np.inf else 0 for g in self.generators], dtype=float) ** 2)
                     ** 0.5
                 )
             prior_sigma.append(pm)
+        if (np.asarray([pm[0] for pm in prior_sigma]) == 0).all():
+            prior_sigma[0][0] = np.inf
         return np.hstack(prior_sigma)
 
-    @property
-    def mu(self):
-        mu = []
-        for idx, g in enumerate(self.generators):
-            pm = np.copy(g.mu)
-            if idx != 0:
-                pm[0] = 0
-            mu.append(pm)
-        return np.hstack(mu)
-
-    @property
-    def sigma(self):
-        sigma = []
-        for idx, g in enumerate(self.generators):
-            pm = np.copy(g.sigma)
-            if idx != 0:
-                pm[0] = 0
-            sigma.append(pm)
-        return np.hstack(sigma)
+    # @property
+    # def mu(self):
+    #     mu = []
+    #     for idx, g in enumerate(self.generators):
+    #         pm = np.copy(g.mu)
+    #         if idx != 0:
+    #             pm[0] = 0
+    #         mu.append(pm)
+    #     return np.hstack(mu)
+
+    # @property
+    # def sigma(self):
+    #     sigma = []
+    #     for idx, g in enumerate(self.generators):
+    #         pm = np.copy(g.sigma)
+    #         if idx != 0:
+    #             pm[0] = 0
+    #         sigma.append(pm)
+    #     return np.hstack(sigma)
+
+    def update_priors(self):
+        if self.fit_mu is None:
+            raise ValueError("Can not update priors before fitting.")
+        new = self.copy()
+        for idx in range(len(new)):
+            new[idx].prior_mu = new[idx].fit_mu.copy()
+            new[idx].prior_sigma = new[idx].fit_sigma.copy()
+        return new
 
     def fit(self, *args, **kwargs):
         self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)
@@ -158,6 +194,10 @@ def fit(self, *args, **kwargs):
             self[idx].fit_mu = mu0
             self[idx].fit_sigma = sigma0
 
+        indices = np.cumsum([0, *[g.width for g in self.generators]])
+        for idx, a, b in zip(range(len(indices) - 1), indices[:-1], indices[1:]):
+            self[idx].cov = self.cov[a:b, a:b]
+
     def __len__(self):
         return len(self.generators)
 
@@ -191,6 +231,11 @@ def save(self, filename: str):
 
 class StackedDependentGenerator(StackedIndependentGenerator):
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._prior_mu = None
+        self._prior_sigma = None
+
     @property
     def width(self):
         return np.prod([g.width for g in self.generators])
@@ -218,22 +263,28 @@ def width(self):
 
     @property
     def prior_sigma(self):
-        return combine_sigmas(*[g.prior_sigma for g in self.generators])
+        if self._prior_sigma is None:
+            return combine_sigmas(*[g.prior_sigma for g in self.generators])
+        else:
+            return self._prior_sigma
 
     @property
     def prior_mu(self):
-        return combine_mus(*[g.prior_mu for g in self.generators])
+        if self._prior_mu is None:
+            return combine_mus(*[g.prior_mu for g in self.generators])
+        else:
+            return self._prior_mu
 
     def fit(self, *args, **kwargs):
         self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)
 
-    @property
-    def mu(self):
-        return self.prior_mu if self.fit_mu is None else self.fit_mu
+    # @property
+    # def mu(self):
+    #     return self.prior_mu if self.fit_mu is None else self.fit_mu
 
-    @property
-    def sigma(self):
-        return self.prior_sigma if self.fit_sigma is None else self.fit_sigma
+    # @property
+    # def sigma(self):
+    #     return self.prior_sigma if self.fit_sigma is None else self.fit_sigma
 
     def __getitem__(self, key):
         raise AttributeError(
@@ -245,3 +296,11 @@ def gradient(self):
         raise AttributeError(
             "Can not create a gradient for a dependent stacked generator."
         )
+
+    def update_priors(self):
+        if self.fit_mu is None:
+            raise ValueError("Can not update priors before fitting.")
+        new = self.copy()
+        new._prior_mu = new.fit_mu.copy()
+        new._prior_sigma = new.fit_sigma.copy()
+        return new

src/lamatrix/generator.py

@@ -51,13 +51,13 @@ def _validate_priors(self, prior_mu, prior_sigma, offset_prior=None):
             self.prior_mu[0] = self.offset_prior[0]
             self.prior_sigma[0] = self.offset_prior[1]
 
-    # def update_priors(self):
-    #     if self.fit_mu is None:
-    #         raise ValueError("Can not update priors before fitting.")
-    #     new = self.copy()
-    #     new.prior_mu = new.fit_mu.copy()
-    #     new.prior_sigma = new.fit_sigma.copy()
-    #     return new
+    def update_priors(self):
+        if self.fit_mu is None:
+            raise ValueError("Can not update priors before fitting.")
+        new = self.copy()
+        new.prior_mu = new.fit_mu.copy()
+        new.prior_sigma = new.fit_sigma.copy()
+        return new
 
     def _create_save_data(self):
         def process(arg):
@@ -239,11 +239,15 @@ def equation(self):
         return (
             f"\\[f({func_signature}) = "
             + " + ".join(
-                [f"w_{{{coeff}}} {e}" for coeff, e in enumerate(self._equation)]
+                [f"{self._mu_letter}_{{{coeff}}} {e}" for coeff, e in enumerate(self._equation)]
             )
             + "\\]"
         )
 
+    @property
+    def _mu_letter(self):
+        return "w"
+
     @abstractmethod
     def design_matrix(self):
         """Returns a design matrix, given inputs listed in self.arg_names."""

src/lamatrix/models/gaussian.py

@@ -221,9 +221,12 @@ def table_properties(self):
 
     @property
     def _equation(self):
-        dfdx = f"\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_x^2}}"
+        dfdx = f"\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma^2}}"
         return [f"\\mathbf{{{self.x_name}}}^0", dfdx]
 
+    @property
+    def _mu_letter(self):
+        return "v"
 
 class lnGaussian2DGenerator(MathMixins, Generator):
     def __init__(
@@ -374,8 +377,8 @@ def table_properties(self):
                 for idx in range(self.width)
             ],
             ("A", self.A, None),
-            ("\\sigma_x", self.stddev_x, self.stddev_x_prior),
-            ("\\sigma_y", self.stddev_y, self.stddev_y_prior),
+            (f"\\sigma_{{{self.x_name}}}", self.stddev_x, self.stddev_x_prior),
+            (f"\\sigma_{{{self.y_name}}}", self.stddev_y, self.stddev_y_prior),
             ("\\rho", self.rho, None),
         ]
 
@@ -390,12 +393,12 @@ def _equation(self):
 
     def to_latex(self):
         eq1 = f"\\begin{{equation}}\\label{{eq:lngauss}}\\ln(G(\\mathbf{{{self.x_name}}}, \\mathbf{{{self.y_name}}})) = a + b\\mathbf{{{self.x_name}}}^2 + c\\mathbf{{{self.y_name}}}^2 + 2d\\mathbf{{{self.x_name}}}\\mathbf{{{self.y_name}}}\\end{{equation}}"
-        eq2 = "\\[ a = -\\ln(2\\pi\\sigma_x\\sigma_y\\sqrt{1-\\rho^2}) \\]"
-        eq3 = "\\[ b = \\frac{1}{2(1-\\rho^2)\\sigma_x^2}\\]"
-        eq4 = "\\[ c = \\frac{1}{2(1-\\rho^2)\\sigma_y^2}\\]"
-        eq5 = "\\[ d = \\frac{\\rho}{2(1-\\rho^2)\\sigma_x\\sigma_y}\\]"
-        eq6 = "\\[\\sigma_x = \\sqrt{-\\frac{1}{2b(1-\\rho^2)}}\\]"
-        eq7 = "\\[\\sigma_y = \\sqrt{-\\frac{1}{2c(1-\\rho^2)}}\\]"
+        eq2 = f"\\[ a = -\\ln(2\\pi\\sigma_{{{self.x_name}}}\\sigma_{{{self.y_name}}}\\sqrt{{1-\\rho^2}}) \\]"
+        eq3 = f"\\[ b = \\frac{{1}}{{2(1-\\rho^2)\\sigma_{{{self.x_name}}}^2}}\\]"
+        eq4 = f"\\[ c = \\frac{{1}}{{2(1-\\rho^2)\\sigma_{{{self.y_name}}}^2}}\\]"
+        eq5 = f"\\[ d = \\frac{{\\rho}}{{2(1-\\rho^2)\\sigma_{{{self.x_name}}}\\sigma_{{{self.y_name}}}}}\\]"
+        eq6 = f"\\[\\sigma_{{{self.x_name}}} = \\sqrt{{-\\frac{{1}}{{2b(1-\\rho^2)}}}}\\]"
+        eq7 = f"\\[\\sigma_{{{self.y_name}}} = \\sqrt{{-\\frac{{1}}{{2c(1-\\rho^2)}}}}\\]"
         eq8 = "\\[\\rho = \\sqrt{\\frac{d^2}{bc}}\\]"
         return "\n".join(
             [eq1, eq2, eq3, eq4, eq5, eq6, eq7, eq8, self._to_latex_table()]
@@ -516,6 +519,10 @@ 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{{\\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)"
+        dfdx = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_{{{self.x_name}}}^2}} - \\rho\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_{{{self.x_name}}}\\sigma_{{{self.y_name}}}}}\\right)\\right)"
+        dfdy = f"\\left(-\\frac{{1}}{{1-\\rho^2}}\\left(\\frac{{\\mathbf{{{self.y_name}}}}}{{\\sigma_{{{self.x_name}}}^2}} - \\rho\\frac{{\\mathbf{{{self.x_name}}}}}{{\\sigma_{{{self.x_name}}}\\sigma_{{{self.y_name}}}}}\\right)\\right)"
         return [f"\\mathbf{{{self.x_name}}}^0", dfdx, dfdy]
+
+    @property
+    def _mu_letter(self):
+        return "v"