Introducing units and refactoring Keplerian orbit models (#61)

* adding units

* adding example of how we might use units

* adding tests for quantity_input decorator

* adding docstring to decorator

* adding doctring and tests for field

* adding tests for keplerian

* fixing edge cases test

* updating keplerian tests

* adding units to TransitOrbit

* fixing python 3.9 errors

* remove pyright from pre-commit

* less stringent test comparisons

* test tol for soln

* sketching OO interface

* default central

* improving interface for keplerian system

.pre-commit-config.yaml

@@ -19,3 +19,13 @@ repos:
     rev: "0.6.1"
     hooks:
       - id: nbstripout
+  # - repo: https://github.com/RobertCraigie/pyright-python
+  #   rev: v1.1.324
+  #   hooks:
+  #     - id: pyright
+  #       additional_dependencies:
+  #         - nox
+  #         - pytest
+  #         - jax
+  #         - equinox
+  #         - jpu>=0.0.2

pyproject.toml

@@ -15,7 +15,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
 ]
 dynamic = ["version"]
-dependencies = ["jax", "jaxlib", "jpu"]
+dependencies = ["jax", "jaxlib", "jpu>=0.0.2", "equinox"]
 
 [project.urls]
 "Homepage" = "https://jax.exoplanet.codes"

src/jaxoplanet/__init__.py

@@ -1,2 +1,2 @@
-from jaxoplanet import core as core, orbits as orbits
+from jaxoplanet import core as core, orbits as orbits, units as units
 from jaxoplanet.jaxoplanet_version import __version__ as __version__

src/jaxoplanet/light_curves.py

@@ -1,30 +1,32 @@
 from functools import partial
-from typing import NamedTuple, Optional
+from typing import Optional
 
+import equinox as eqx
 import jax.numpy as jnp
 
+from jaxoplanet import units
 from jaxoplanet.core.limb_dark import light_curve
 from jaxoplanet.proto import LightCurveOrbit
-from jaxoplanet.types import Array
+from jaxoplanet.types import Array, Quantity
+from jaxoplanet.units import unit_registry as ureg
 
 
-class LimbDarkLightCurve(NamedTuple):
+class LimbDarkLightCurve(eqx.Module):
     u: Array
 
-    @classmethod
-    def init(cls, *u: Array) -> "LimbDarkLightCurve":
+    def __init__(self, *u: Array):
         if u:
-            u = jnp.concatenate([jnp.atleast_1d(u0) for u0 in u], axis=0)
+            self.u = jnp.concatenate([jnp.atleast_1d(u0) for u0 in u], axis=0)
         else:
-            u = jnp.array([])
-        return cls(u=u)
+            self.u = jnp.array([])
 
+    @units.quantity_input(t=ureg.d, texp=ureg.s)
     def light_curve(
         self,
         orbit: LightCurveOrbit,
-        t: Array,
+        t: Quantity,
         *,
-        texp: Optional[Array] = None,
+        texp: Optional[Quantity] = None,
         oversample: Optional[int] = 7,
         texp_order: Optional[int] = 0,
         limbdark_order: Optional[int] = 10,
@@ -100,7 +102,7 @@ def light_curve(
         lc_func = partial(light_curve, self.u, order=limbdark_order)
         if orbit.shape == ():
             b /= r_star
-            lc = lc_func(b, r)
+            lc: Array = lc_func(b.magnitude, r.magnitude)
         else:
             b /= r_star[..., None]
             lc = jnp.vectorize(lc_func, signature="(k),()->(k)")(b, r)

src/jaxoplanet/orbits/__init__.py

@@ -1,6 +1,2 @@
-from jaxoplanet.orbits.keplerian import (
-    KeplerianBody as KeplerianBody,
-    KeplerianCentral as KeplerianCentral,
-    KeplerianOrbit as KeplerianOrbit,
-)
+from jaxoplanet.orbits import keplerian as keplerian
 from jaxoplanet.orbits.transit import TransitOrbit as TransitOrbit