developing

.github/workflows/black.yml

@@ -0,0 +1,14 @@
+name: black
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+      - uses: psf/black@stable
+        with:
+          options: "--check --verbose"
+          src: "./src"

.github/workflows/flake8.yml

@@ -0,0 +1,29 @@
+name: flake8
+
+on: [push, pull_request]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    name: flake8
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.9
+    - name: Cache Poetry dependencies
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/pypoetry
+        key: ${{ runner.os }}-poetry-${{ hashFiles('**/poetry.lock') }}
+        restore-keys: |
+          ${{ runner.os }}-poetry-
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install poetry
+        poetry install
+    - name: Run flake8
+      run: |
+        make flake8

.github/workflows/pytest.yml

@@ -0,0 +1,44 @@
+name: pytest
+
+on: [push, pull_request]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.9]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set CI environment variable
+      run: echo "CI=true" >> $GITHUB_ENV
+
+    - name: Configure git
+      run: |
+        git config --global user.email "christina.l.hedges@nasa.gov"
+        git config --global user.name "Christina Hedges"
+        
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Cache Poetry dependencies
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/pypoetry
+        key: ${{ runner.os }}-poetry-${{ hashFiles('**/poetry.lock') }}
+        restore-keys: |
+          ${{ runner.os }}-poetry-
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install poetry
+        poetry install
+
+    - name: Test with pytest
+      run: |
+        make pytest

README.md

@@ -0,0 +1,10 @@
+<a href="https://github.com/christinahedges/lamatrix/actions/workflows/tests.yml"><img src="https://github.com/christinahedges/lamatrix/workflows/pytest/badge.svg" alt="Test status"/></a> [![Generic badge](https://img.shields.io/badge/documentation-live-blue.svg)](https://christinahedges.github.io/lamatrix/)
+[![PyPI version](https://badge.fury.io/py/lamatrix.svg)](https://badge.fury.io/py/lamatrix)
+
+<p align="center">
+  <img src="https://github.com/christinahedges/lamatrix/blob/main/docs/images/logo.png?raw=true" width="350" alt="lamatrix logo">
+</p>
+
+# LAmatrix
+
+This package is designed to help you fit linear algebra models to data. It's designed to eventually replace `Lightkurve.design_matrix` and `Lightkurve.Corrector` objects.

pyproject.toml

@@ -17,6 +17,8 @@ black = "^24.2.0"
 isort = "^5.13.2"
 mypy = "^1.9.0"
 pytest = "^8.1.1"
+jupyterlab = "^4.1.4"
+matplotlib = "^3.8.3"
 
 [build-system]
 requires = ["poetry-core"]

src/lamatrix/__init__.py

@@ -14,4 +14,8 @@
 log = logging.getLogger("tesswcs")
 log.addHandler(RichHandler(markup=True))
 
-from .generator import *  # noqa: E402, F401
+from .combined 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
+from .models.spline import *  # noqa: E402, F401

src/lamatrix/combined.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+from .generator import Generator
+
+__all__ = ["VStackedGenerator"]
+
+
+def combine_equations(*equations):
+    # Base case: if there's only one equation, just return it
+    if len(equations) == 1:
+        return equations[0]
+
+    # Step case: combine the first two equations and recursively call the function with the result
+    combined = [f + e for f in equations[1] for e in equations[0]]
+
+    # If there are more equations left, combine further
+    if len(equations) > 2:
+        return combine_equations(combined, *equations[2:])
+    else:
+        return combined
+
+
+class VStackedGenerator(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
+
+    def __getitem__(self, key):
+        return self.generators[key]
+
+    # def __add__(self, other):
+    #     if isinstance(other, Generator):
+    #         return VStackedGenerator(*self.generators, other)
+    #     else:
+    #         raise ValueError("Can only combine `Generator` objects.")
+
+    def design_matrix(self, *args, **kwargs):
+        return np.hstack([g.design_matrix(*args, **kwargs) for g in self.generators])
+
+    @property
+    def width(self):
+        return np.sum([g.width for g in self.generators])
+
+    @property
+    def nvectors(self):
+        return len(self.arg_names)
+
+    @property
+    def prior_mu(self):
+        prior_mu = []
+        for idx, g in enumerate(self.generators):
+            pm = np.copy(g.prior_mu)
+            if idx != 0:
+                pm[0] = 0
+            else:
+                pm[0] = np.sum([g.prior_mu[0] for g in self.generators])
+            prior_mu.append(pm)
+        return np.hstack(prior_mu)
+
+    @property
+    def prior_sigma(self):
+        prior_sigma = []
+        for idx, g in enumerate(self.generators):
+            pm = np.copy(g.prior_sigma)
+            if idx != 0:
+                pm[0] = 0
+            else:
+                pm[0] = np.sum(np.asarray([g.prior_sigma[0] for g in self.generators])**2)**0.5
+            prior_sigma.append(pm)
+        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)
+
+    def fit(self, *args, **kwargs):
+        self.fit_mu, self.fit_sigma = self._fit(*args, **kwargs)
+        lengths = [g.width for g in self.generators]
+        mu, sigma = (
+            np.array_split(self.fit_mu, np.cumsum(lengths))[:-1],
+            np.array_split(self.fit_sigma, np.cumsum(lengths))[:-1],
+        )
+        for idx, mu0, sigma0 in zip(np.arange(len(mu)), mu, sigma):
+            self[idx].fit_mu = mu0
+            self[idx].fit_sigma = sigma0
+
+    def __len__(self):
+        return len(self.generators)
+
+    # @property
+    # def _equation(self):
+    #     return combine_equations(*[g._equation for g in self])
+
+    @property
+    def _equation(self):
+        return np.hstack([g._equation for g in self.generators])
+
+    @property
+    def arg_names(self):
+        return np.unique(np.hstack([list(g.arg_names) for g in self.generators]))
+
+
+class CombinedGenerator(VStackedGenerator):
+
+    @property
+    def width(self):
+        return np.prod([g.width for g in self.generators])
+
+    # def design_matrix(self, *args, **kwargs):
+    #     return np.hstack([g.design_matrix(*args, **kwargs) for g in self.generators])
+
+    @property
+    def _equation(self):
+        return combine_equations(*[g._equation for g in self])