docs: clean docstrings

nuance/kernels.py

@@ -1,8 +1,11 @@
-from tinygp import GaussianProcess, kernels
 import jax.numpy as jnp
+from tinygp import GaussianProcess, kernels
 
 
 def Rotation(sigma, period, Q0, dQ, f):
+    """
+    A kernel for a rotating star with a single mode of oscillation.
+    """
     Q1 = 1 / 2 + Q0 + dQ
     w1 = (4 * jnp.pi * Q1) / (period * jnp.sqrt(4 * Q1**2 - 1))
     s1 = sigma**2 / ((1 + f) * w1 * Q1)
@@ -15,7 +18,7 @@ def Rotation(sigma, period, Q0, dQ, f):
 
 
 def rotation(period=None, error=None, long_scale=0.5):
-    params = {
+    initial_params = {
         "log_period": jnp.log(period) if period is not None else jnp.log(1.0),
         "log_Q": jnp.log(100),
         "log_sigma": jnp.log(1e-1),
@@ -47,4 +50,4 @@ def build_gp(params, time):
 
         return GaussianProcess(kernel, time, diag=jitter2, mean=0.0)
 
-    return build_gp, params
+    return build_gp, initial_params

nuance/nuance.py

@@ -47,7 +47,8 @@ class Nuance:
     search_data : SearchData, optional
         Search data instance, by default None.
     model : callable, optional
-        Model function with signature `model(time, t0, D, P=None)`, by default None.
+        Model function with signature `model(time, t0, D, P=None)`, by default None, which set the model to
+        :code:`partial(nuance.core.transit_protopapas, c=12)`
     """
 
     time: np.ndarray
@@ -67,7 +68,7 @@ class Nuance:
     search_data: SearchData = None
     """Search data instance"""
     model: callable = None
-    """Model function with signature `model(time, t0, D, P=None)`"""
+    """Model function with signature :code:`model(time, t0, D, P=None)`"""
 
     def __post_init__(self):
         if self.model is None:
@@ -137,22 +138,23 @@ def gp_mean():
 
     def gp_optimization(self, build_gp, mask=None):
         """
-        Optimize the Gaussian Process (GP) model using the given build_gp function.
+        Optimization functions to fit a Gaussian Process given a building function.
 
         Parameters
         ----------
         build_gp : function
-            A function that returns a GP object.
+            A function that returns a tinygp.GaussianProcess object.
         mask : array-like, optional
             A boolean array to mask the data, by default None.
 
         Returns
         -------
         tuple
             A tuple containing three functions:
-            - optimize: a function that optimizes the GP model.
-            - mu: a function that returns the mean of the GP model.
-            - nll: a function that returns the negative log-likelihood of the GP model.
+
+            - :code:`optimize`: a function that optimizes the GP model.
+            - :code:`mu`: a function that returns the mean of the GP model (jax.jit-compiled).
+            - :code:`nll`: a function that returns the negative log-likelihood of the GP model (jax.jit-compiled).
         """
         if mask is None:
             mask = np.ones_like(self.time).astype(bool)

nuance/star.py

@@ -225,17 +225,21 @@ def period_grid(self, time_span, period_max=None, period_min=None, oversampling=
 
         Parameters
         ----------
+        time_span : float
+            time span in days.
         period_max : float
             maximum period in days, by default None which
             defaults to half the time span.
         period_min : float, optional
             minimum period in days, by default None which
             defaults to the roche limit period.
+        oversampling: int, optional
+            oversampling factor, by default 1.
 
         Returns
         -------
         np.ndarray
-            period grid
+            periods in days.
         """
         if period_min is None:
             period_min = self.roche_period()