eztaox.models

A module of light curve models, which are the interface for modeling uni/multi-band light curves using Gaussian Processes (GPs).

Classes

MultiVarModel

An interface for modeling multivariate/mutli-band time series using GPs.

UniVarModel

A subclass of MultiVarModel for modeling univariate/single-band time series data.

Module Contents

class MultiVarModel(X: tuple[tinygp.helpers.JAXArray | numpy.typing.NDArray, tinygp.helpers.JAXArray | numpy.typing.NDArray], y: tinygp.helpers.JAXArray | numpy.typing.NDArray, yerr: tinygp.helpers.JAXArray | numpy.typing.NDArray, base_kernel: eztaox.kernels.quasisep.Quasisep, nBand: int, multiband_kernel: tinygp.kernels.quasisep.Wrapper | None = quasisep.MultibandLowRank, mean_func: collections.abc.Callable | None = None, amp_scale_func: collections.abc.Callable | None = None, lag_func: collections.abc.Callable | None = None, **kwargs)[source]

Bases: equinox.Module

An interface for modeling multivariate/mutli-band time series using GPs.

This interface only takes GP kernels that can be evaluated using the scalable method of DFM+17 <https://arxiv.org/abs/1703.09710>. This interface allows fitting for a parameterized mean function of the time series, additional variance to the measurement uncertainty, and time delays between each uni-variate/single-band time series.

Parameters:

  • X (JAXArray|NDArray) – Input data containing time and band indices as a tuple.

  • y (JAXArray|NDArray) – Observed data values.

  • yerr (JAXArray|NDArray) – Observational uncertainties.

  • base_kernel (Quasisep) – A GP kernel from the kernels.quasisep module.

  • nBand (int) – An integer number of bands in the input light curve.

  • multiband_kernel (Quasisep, optional) – A multiband kernel specifying the cross-band covariance, defaults to kernels.quasisep.MultibandLowRank.

  • mean_func (Callable, optional) – A callable mean function for the GP, defaults to None.

  • amp_scale_func (Callable, optional) – A callable amplitude scaling function, defaults to None.

  • lag_func (Callable, optional) – A callable function for time delays between bands, defaults to None.

  • **kwargs –

    Additional keyword arguments.

    • zero_mean (bool): If True, assumes zero-mean GP. Defaults to True.

    • has_jitter (bool): If True, assumes the input observational erros are underestimated. Defaults to False.

    • has_lag (bool): If True, assumes time delays between time series in each band. Defaults to False.

Raises:

TypeError – If base_kernel is not one from the kernels.quasisep module.

X: tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]
y: tinygp.helpers.JAXArray[source]
diag: tinygp.helpers.JAXArray[source]
base_kernel_def: collections.abc.Callable[source]
multiband_kernel: tinygp.kernels.quasisep.Wrapper[source]
nBand: int[source]
mean_func: collections.abc.Callable | None[source]
amp_scale_func: collections.abc.Callable | None[source]
lag_func: collections.abc.Callable | None[source]
zero_mean: bool[source]
has_jitter: bool[source]
has_lag: bool[source]
get_mean(zero_mean: bool, params: dict[str, tinygp.helpers.JAXArray], X: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]

Return the mean of the GP.

get_amp_scale(params: dict[str, tinygp.helpers.JAXArray]) → tinygp.helpers.JAXArray[source]

Return the ampltiude of GP in each individaul band.

lag_transform(has_lag: bool, params: dict[str, tinygp.helpers.JAXArray], X: tinygp.helpers.JAXArray) → tuple[tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray], tinygp.helpers.JAXArray][source]

Shift the time axis by the lag in each band.

log_prior(params: dict[str, tinygp.helpers.JAXArray]) → tinygp.helpers.JAXArray[source]

Calculate the log prior of the input parameters.

Parameters:

params (dict[str, JAXArray]) – Model parameters.

Returns:

Log prior of the input parameters.

Return type:

JAXArray

log_prob(params: dict[str, tinygp.helpers.JAXArray]) → tinygp.helpers.JAXArray[source]

Calculate the log probability of the input parameters.

Parameters:

params (dict[str, JAXArray]) – Model parameters.

Returns:

Log probability of the input parameters.

Return type:

JAXArray

sample(params: dict[str, tinygp.helpers.JAXArray]) → None[source]

A convience function for intergrating with numpyro for MCMC sampling.

Parameters:

params (dict[str, JAXArray]) – Model parameters.

pred(params: dict[str, tinygp.helpers.JAXArray], X: tinygp.helpers.JAXArray) → tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]

Make conditional GP prediction.

Parameters:

  • params (dict[str, JAXArray]) – A dictionary containing model parameters.

  • X (JAXArray) – The time and band information for creating the conditional GP prediction.

Returns:

A tuple of the mean GP prediction and its uncertainty (square root of the predicted variance).

Return type:

tuple[JAXArray, JAXArray]

_default_mean_func(params: dict[str, tinygp.helpers.JAXArray], X: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]
_default_amp_scale_func(params: dict[str, tinygp.helpers.JAXArray]) → tinygp.helpers.JAXArray[source]
_default_lag_func(params: dict[str, tinygp.helpers.JAXArray]) → tuple[tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray], tinygp.helpers.JAXArray][source]
_build_gp(params: dict[str, tinygp.helpers.JAXArray]) → tuple[tinygp.GaussianProcess, tinygp.helpers.JAXArray][source]
class UniVarModel(t: tinygp.helpers.JAXArray | numpy.typing.NDArray, y: tinygp.helpers.JAXArray | numpy.typing.NDArray, yerr: tinygp.helpers.JAXArray | numpy.typing.NDArray, kernel: eztaox.kernels.quasisep.Quasisep, mean_func: collections.abc.Callable | None = None, amp_scale_func: collections.abc.Callable | None = None, **kwargs)[source]

Bases: MultiVarModel

A subclass of MultiVarModel for modeling univariate/single-band time series data.

Parameters:

  • t (JAXArray|NDArray) – Time stamps of the input light curve.

  • y (JAXArray|NDArray) – Observed data values at the corresponding time stamps.

  • yerr (JAXArray|NDArray) – Observational uncertainties.

  • kernel (Quasisep) – A GP kernel from the eztaox.kernels.quasisep module.

  • mean_func (Callable, optional) – A callable mean function for the GP, defaults to None.

  • amp_scale_func (Callable, optional) – A callable amplitude scaling function, defaults to None.

  • **kwargs –

    Additional keyword arguments.

    • zero_mean (bool): If True, assumes zero-mean GP. Defaults to True.

    • has_jitter (bool): If True, assumes the input observational erros are underestimated. Defaults to False.

Raises:

TypeError – If kernel is not one from the kernels.quasisep module.

_default_amp_scale_func(params: dict[str, tinygp.helpers.JAXArray]) → tinygp.helpers.JAXArray[source]
lag_transform(has_lag, params, X) → tuple[tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray], tinygp.helpers.JAXArray][source]

Shift the time axis by the lag in each band.

pred(params, t) → tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]

Make conditional GP prediction.

Parameters:

  • params (dict[str, JAXArray]) – A dictionary containing model parameters.

  • t (JAXArray) – The time information for creating the conditional GP prediction.

Returns:

A tuple of the mean GP prediction and its uncertainty (square root of the predicted variance).

Return type:

tuple[JAXArray, JAXArray]