diff --git a/nuance/kernels.py b/nuance/kernels.py index 3bbf3f2..4fe2fe6 100644 --- a/nuance/kernels.py +++ b/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 diff --git a/nuance/nuance.py b/nuance/nuance.py index 2fc3fb0..ed6ec46 100644 --- a/nuance/nuance.py +++ b/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,12 +138,12 @@ 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. @@ -150,9 +151,10 @@ def gp_optimization(self, build_gp, mask=None): ------- 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) diff --git a/nuance/star.py b/nuance/star.py index c1c6289..9cc1392 100644 --- a/nuance/star.py +++ b/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()