Use equinox module classes

.gitignore

@@ -0,0 +1,4 @@
+.ipynb_checkpoints
+__pycache__
+.coverage
+*.bak

.pre-commit-config.yaml

@@ -38,7 +38,9 @@ repos:
 
   # Make sure the Poetry project is properly maintained
   - repo: https://github.com/python-poetry/poetry
-    rev: 1.3.2
+    rev: 1.6.1
     hooks:
       - id: poetry-check
-      # - id: poetry-lock
+      - id: poetry-lock
+        args: [--no-update]
+

pyproject.toml

@@ -96,6 +96,8 @@ ignore = [
     "PLR0913", # Allow many arguments in function call
     "PLR0915", # Allow many statements
     "PLR2004", # Allow magic numbers in comparisons
+    "F401", # ignore at development stage
+    "I001", # ignore at development stage
 ]
 # exclude = []
 

python/jaxodi/api.py

@@ -0,0 +1,46 @@
+import jax.numpy as jnp
+from jaxoplanet.experimental.starry.ylm import Ylm
+
+from jaxodi.doppler_surface import DopplerSurface
+from jaxodi.utils import enforce_bounds
+from jaxodi.doppler_forward import (
+    get_kT,
+    dot_design_matrix_fixed_map_into,
+    get_flux_from_dotconv,
+)
+
+from jax.experimental import checkify
+
+
+def light_curve(surface: DopplerSurface):
+    ydeg = surface.y.ell_max
+    vsini = surface.veq * jnp.sin(surface.inc)
+
+    # check bounds at runtime
+    err, vsini = enforce_bounds(vsini, 0.0, surface.wav.vsini_max)
+    checkify.check_error(err)
+
+    kT = get_kT(
+        surface.wav.xamp, vsini, ydeg, 0, surface.wav.nk, surface.inc, surface.theta
+    )
+    spec0_int = surface.spec0.spec0_int
+    nc = surface.spec0.nc
+    _y = surface.y.todense()[:, jnp.newaxis]
+
+    flux = dot_design_matrix_fixed_map_into(
+        kT,
+        _y,
+        nc,
+        surface.wav.nw0_int,
+        surface.nt,
+        surface.wav.nk,
+        surface.wav.nw_int,
+        spec0_int,
+    )
+
+    lc_int = get_flux_from_dotconv(flux, surface.nt, surface.wav.nw_int)
+
+    # Interpolate to the `wav` grid
+    lc = jnp.dot(lc_int, surface.wav._Si2eTr)
+
+    return lc