Merge branch 'feat/more_functional'

lgrcia · Jul 30, 2024 · a6a94f4 · a6a94f4
2 parents ca3375d + b726bed
commit a6a94f4
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 30 deletions.
diff --git a/docs/markdown/API.md b/docs/markdown/API.md
@@ -17,8 +17,15 @@
 
 .. currentmodule:: nuance
 
-.. autofunction:: linear_search
-.. autofunction:: periodic_search
+.. automodule:: nuance.linear_search
+   :members:
+   :show-inheritance:
+
+
+.. automodule:: nuance.periodic_search
+   :members:
+   :show-inheritance:
+
 
 .. automodule:: nuance.core
    :members:

diff --git a/docs/notebooks/motivation.ipynb b/docs/notebooks/motivation.ipynb
diff --git a/nuance/linear_search.py b/nuance/linear_search.py
@@ -2,16 +2,18 @@
 import jax.numpy as jnp
 import numpy as np
 from tqdm.auto import tqdm
+from typing import Callable
+import tinygp
 
 from nuance import DEVICES_COUNT, core
 
 
 def linear_search(
-    time,
-    flux,
-    gp=None,
-    X=None,
-    model=None,
+    time: np.ndarray,
+    flux: np.ndarray,
+    gp: tinygp.GaussianProcess | None = None,
+    X: np.ndarray | None = None,
+    model: Callable | None = None,
     positive: bool = True,
     progress: bool = True,
     backend: str | None = None,
@@ -58,7 +60,7 @@ def linear_search(
             jax.vmap(jax.vmap(solver, in_axes=(None, 0)), in_axes=(0, None))
         )
 
-    def function(epochs, durations):
+    def function(epochs: np.ndarray, durations: np.ndarray):
         """Compute the log likelihood of a transit model at different epochs and durations
 
         Parameters