Jax/PyTorch Backend

.github/workflows/ci.yml

@@ -0,0 +1,46 @@
+name: stella-ci
+
+on:
+  push:
+    branches: [ main, feat/**, fix/** ]
+  pull_request:
+
+jobs:
+  test:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            backend: jax
+            python: '3.12'
+          - os: ubuntu-latest
+            backend: torch
+            python: '3.12'
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      # macOS-specific gettext step removed; both matrix jobs run on Linux now
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - name: Install dependencies (base)
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+      - name: Install backend (JAX)
+        if: matrix.backend == 'jax'
+        run: |
+          pip install -r requirements-jax.txt
+          echo "KERAS_BACKEND=jax" >> $GITHUB_ENV
+        - name: Install backend (Torch)
+        if: matrix.backend == 'torch'
+        run: |
+          pip install -r requirements-torch.txt
+          echo "KERAS_BACKEND=torch" >> $GITHUB_ENV
+      - name: Install package
+        run: |
+          python -m pip install --upgrade pip build setuptools wheel
+          pip install .
+      - name: Run tests
+        run: pytest stella/tests -q

.github/workflows/docs.yml

@@ -0,0 +1,61 @@
+name: Deploy Docs
+
+on:
+  push:
+    branches:
+      - main
+      - feat/**
+      - fix/**
+  pull_request:
+    branches:
+      - main
+      - feat/**
+      - fix/**
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: 'pages'
+  cancel-in-progress: true
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e .[dev]
+          python -m pip install mkdocs mkdocs-material mkdocstrings[python] mkdocs-jupyter mkdocs-exclude
+
+      - name: Build MkDocs site
+        run: |
+          mkdocs build --strict
+
+      - name: Upload artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: site
+
+  deploy:
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4

.github/workflows/stella-tests.yml

@@ -8,26 +8,27 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        name: [3.6, 3.7]
+        name: [3.12]
         include:
-          - name: 3.6
-            python-version: 3.6
-            pytest-command: poetry run pytest
-          - name: 3.7
-            python-version: 3.7
+          - name: 3.12
+            python-version: 3.12
             pytest-command: poetry run pytest
 
     steps:
-      - uses: actions/checkout@v1
+      - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Install package and dev deps (with JAX)
+        run: |
+          python -m pip install --upgrade setuptools wheel
+          python -m pip install -e ".[dev,jax]"
       - name: Test with pytest
         run: |
           pytest --doctest-modules --cov=com --cov-report=xml --cov-report=html

.gitignore

@@ -1,3 +1,6 @@
+mastDownload/
+stella/mastDownload/
+*.fits
 __pycache__
 */__pycache__
 */*/__pycache__
@@ -16,3 +19,8 @@ dist
 *.cppimporthash
 .rendered.*
 *.egg-info
+
+## Local examples and helpers (not part of package)
+notebooks/jax_keras.ipynb
+scripts/
+site/

.pre-commit-config.yaml

@@ -0,0 +1,9 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 24.10.0
+    hooks:
+      - id: black
+        name: black (python formatter)
+        language_version: python3
+        types: [python]
+        additional_dependencies: []

.travis.yml

@@ -1,11 +1,15 @@
-dist: xenial
+dist: focal
 
 language: python
-python: 3.7
 
 matrix:
   include:
   - os: linux
+    python: 3.10
+    env: BACKEND=jax
+  - os: osx
+    language: shell
+    env: BACKEND=torch PYTHON=3.10
 
 addons:
   apt_packages:
@@ -14,26 +18,15 @@ addons:
 cache: pip
 
 before_install:
-  - sudo apt-get install python3-pip
-  - sudo pip3 install requests
+  - if [ "$TRAVIS_OS_NAME" = "linux" ]; then sudo apt-get install -y python3-pip; fi
+  - if [ "$TRAVIS_OS_NAME" = "linux" ]; then sudo pip3 install requests; fi
 
 install:
-  - pip install tornado
-  - pip install tensorflow>2.1.0
-  - pip install scipy!=1.4.1
-  - pip install lightkurve
-  - pip install astropy
-  - pip install astroquery
-  - pip install tqdm
-  - pip install sphinx
-  - pip install nbsphinx
-  - pip install cython
-  - pip install codecov
-  - pip install coverage
-  - pip install pytest
-  - pip install pytest-cov
-  - pip install lxml
-  - pip install pybind11
+  - if [ "$TRAVIS_OS_NAME" = "osx" ]; then python3 -m pip install --upgrade pip; else pip install --upgrade pip; fi
+  - if [ "$TRAVIS_OS_NAME" = "osx" ]; then python3 -m pip install -r requirements.txt; else pip install -r requirements.txt; fi
+  - if [ "$BACKEND" = "jax" ]; then pip install -r requirements-jax.txt; export KERAS_BACKEND=jax; fi
+  - if [ "$BACKEND" = "torch" ]; then python3 -m pip install -r requirements-torch.txt; export KERAS_BACKEND=torch; fi
+  - pip install sphinx nbsphinx
   - python setup.py install
 
 script:

README.md

@@ -8,6 +8,9 @@
   <a href="https://afeinstein20.github.io/stella/"><img src="https://img.shields.io/badge/read-the_docs-3C1370.svg?style=flat"/></a>
   <a href="https://pypi.org/project/stella"><img alt="PyPI - Downloads" src="https://img.shields.io/pypi/dm/stella?color=D35968"></a>
   <a href="https://doi.org/10.21105/joss.02347">   <img src="https://joss.theoj.org/papers/10.21105/joss.02347/status.svg?color=D35968"></a>
+  <!-- Fork-specific docs status -->
+  <a href="https://github.com/benjaminpope/stella/actions/workflows/docs.yml"><img src="https://github.com/benjaminpope/stella/actions/workflows/docs.yml/badge.svg" alt="Docs Deploy"></a>
+  <a href="https://benjaminpope.github.io/stella/"><img src="https://img.shields.io/badge/docs-fork_pages-3C1370.svg?style=flat" alt="GitHub Pages"></a>
 </p>
 
 
@@ -29,6 +32,52 @@ Alternatively you can install the current development version of stella:
         cd stella
         python setup.py install
 
+Backends and models
+-------------------
+- stella uses Keras Core and works with multiple backends: JAX or PyTorch. Choose one, install, and set it before importing `keras`:
+
+```bash
+# JAX
+pip install -r requirements-jax.txt
+export KERAS_BACKEND=jax
+
+# or PyTorch
+pip install -r requirements-torch.txt
+export KERAS_BACKEND=torch
+```
+
+ Models are saved/loaded in the native Keras format (`.keras`). If you have older `.h5`
+  or legacy SavedModel models, use the included converter:
+
+```bash
+# Convert .h5 or legacy SavedModel to .keras (TensorFlow backend only for conversion)
+python scripts/convert_h5_to_keras.py /path/to/models -r --cadences 350 -o /path/to/out
+```
+
+Pipeline quickstart
+-------------------
+- See the new pipeline guide for a friendly, end-to-end example (TIC 62124646):
+  https://afeinstein20.github.io/stella/getting_started/pipeline.html
+
+Swap backends quickly
+---------------------
+- Inspect availability and devices:
+
+```python
+import stella
+stella.check_backend()
+```
+
+- Prepare a swap (restart your interpreter after):
+
+```python
+import stella
+stella.swap_backend('torch', accelerator='mps')   # Apple Silicon
+# or
+stella.swap_backend('jax', accelerator='cpu')
+```
+```
+
 <p>
 If your work uses the stella software, please cite <a href="https://ui.adsabs.harvard.edu/abs/2020JOSS....5.2347F/abstract">Feinstein, Montet, & Ansdell (2020)</a>.
 </p>

docs/api.md

@@ -0,0 +1,3 @@
+# API Reference
+
+::: stella

docs/conf.py

@@ -174,3 +174,7 @@
 if not on_rtd:  # only import and set the theme if we're building docs locally
     html_theme = 'sphinx_rtd_theme'
     html_theme_path = ['.',]
+
+# Ensure theme templates have a default 'style' context to avoid Jinja errors
+html_context = globals().get('html_context', {})
+html_context.setdefault('style', 'default')

docs/getting_started/about.md

@@ -0,0 +1,11 @@
+# About stella
+
+stella identifies flares in TESS short‑cadence data with a convolutional neural network (CNN).
+
+In its simplest form, stella takes a pre‑trained CNN (details in Feinstein et al.) and a light curve `(time, flux, flux_err)` and returns a probability light curve. Values are between 0 and 1, where 1 indicates a likely flare.
+
+You can also train your own customized CNN architecture. See the Quickstart tutorial for a walkthrough.
+
+Links:
+- Paper: Feinstein, Montet, & Ansdell (2020, JOSS)
+- Issues: https://github.com/benjaminpope/stella/issues

docs/getting_started/backends.md

@@ -0,0 +1,49 @@
+# Backends (JAX and PyTorch)
+
+stella uses Keras 3, which can run on multiple numerical backends. You can use either the JAX or the PyTorch backend with the same stella code.
+
+Quick facts
+-----------
+- Select backend via the environment variable `KERAS_BACKEND` before importing `keras`.
+- The stella code is backend‑agnostic; no changes are required besides selecting the backend and installing the backend packages.
+
+Setup: JAX
+-----------
+
+```bash
+pip install -U keras jax jaxlib
+export KERAS_BACKEND=jax
+```
+
+Setup: PyTorch
+--------------
+
+```bash
+pip install -U keras torch
+export KERAS_BACKEND=torch
+```
+
+Usage is identical
+------------------
+
+```python
+import os
+os.environ.setdefault("KERAS_BACKEND", "jax")  # or "torch"
+import keras
+m = keras.models.load_model("/path/to/model.keras", compile=False)
+y = m.predict(x)
+```
+
+Troubleshooting
+---------------
+- If you see a backend mismatch, ensure `KERAS_BACKEND` is set before any import of `keras` occurs in your Python process.
+- Some ops or layers may have different performance characteristics across backends.
+
+Inspect, swap, and benchmark
+----------------------------
+
+```python
+import stella
+stella.check_backend()
+stella.swap_backend('torch', accelerator='mps')  # prepare env for PyTorch Metal
+```