harv.models.extensions package

Extension base class and built-in extensions for model components.

class harv.models.extensions.AbstractExtension

Bases: Module

Abstract base class for model extensions (jitter, trends, offsets, GP, …).

An extension can declare extra parameters (nonlinear and/or linear), modify the design matrix (e.g., append trend or offset columns), and/or modify the data covariance matrix (e.g., add jitter or a GP kernel).

Extensions are composed onto a AbstractComponentModel at construction time via the extensions field. The model calls each hook in the order the extensions are listed.

All extensions receive the raw (unit-stripped) JAX arrays. Unit handling is the model’s responsibility.

Subclasses must implement extra_params(). The design-matrix and covariance hooks have default no-op implementations that return their input unchanged.

Plot-specific behavior is handled privately by plotting helpers rather than through this base extension API.

__init__()
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

abstractmethod extra_params()

Parameters introduced by this extension.

Returns:

May include both nonlinear and linear parameters.

Return type:

tuple of ParamInfo

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

class harv.models.extensions.GP

Bases: AbstractExtension

Gaussian Process covariance extension.

Adds the kernel matrix K(t, t') to the observation covariance via the modify_covariance hook, allowing analytic marginalization of linear parameters in the presence of correlated noise.

Plotting code may choose to visualize the GP contribution, but that support is handled privately on the plotting side rather than through the extension base API.

Parameters:

  • kernel_builder (Callable[[dict[str, Any]], Any]) – Receives the full nonlinear-parameter dict (unit-stripped) and returns a callable kernel object that can be called with (X, Xp) to produce the kernel matrix. For example, from tinygp.

  • hyperparams (tuple[ParamInfo, ...]) – Nonlinear hyperparameters declared by this extension (e.g. gp_amp, gp_length_scale). These are sampled alongside other nonlinear model parameters.

  • time_unit (str) – Unit to strip times to before building the coordinate array. Default "" (dimensionless / already stripped).

Examples

>>> from harv.models.extensions.base import ParamInfo
>>> from harv.models.extensions.gp import GP; GP(
...     kernel_builder=lambda hp: hp["gp_amp"] ** 2,
...     hyperparams=(ParamInfo("gp_amp", "km/s"),),
...     time_unit="day",
... ).extra_params()[0].name
'gp_amp'
__init__(kernel_builder, hyperparams, time_unit='')
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

conditional_mean(residuals, data_times, prediction_times, data_err, hp)

Conditional-mean GP prediction at prediction_times.

Conditions on observed residuals (already mean-subtracted by the deterministic model prediction) at data_times with measurement noise data_err, then returns the GP posterior mean evaluated at prediction_times. Used by the plotting layer to overlay structured noise; the same kernel and hyperparameter conventions as modify_covariance() apply.

All time arrays are assumed to be unit-stripped to the same unit (the kernel’s coordinate unit). hp is the per-sample hyperparameter dict (e.g. from a single posterior sample).

Parameters:
Return type:

Array

extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Add the GP kernel matrix K(t, t') to the covariance.

If the incoming covariance is diagonal (1-d), it is promoted to a full (n_obs, n_obs) matrix first.

Parameters:
Return type:

Array

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

time_unit: str = ''
kernel_builder: Callable[[dict[str, Any]], Any]
hyperparams: tuple[ParamInfo, ...]
class harv.models.extensions.Jitter

Bases: AbstractExtension

Add excess variance (jitter / white noise) to the observation covariance.

The extension declares one nonlinear parameter, jitter, and modifies the covariance by adding jitter**2 in quadrature to the diagonal variances. The jitter unit must match the observation unit (e.g. "km/s" for RV, "mas" for astrometry). Unit stripping is the model’s responsibility.

Parameters:

param_unit (str) – Physical unit string for the jitter parameter metadata. Default "" (dimensionless / observation units).

Examples

>>> from harv.models.extensions import Jitter
>>> j = Jitter(param_unit="km/s")
>>> j.extra_params()[0].unit
'km/s'
__init__(param_unit='')
Parameters:

param_unit (str)

Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Add jitter**2 in quadrature to the diagonal variances.

Works on both 1-d (diagonal) and 2-d (full) covariance representations.

Parameters:
Return type:

Array

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

param_unit: str = ''
class harv.models.extensions.MonomialTrend

Bases: AbstractExtension

Append monomial trend columns to the design matrix.

Uses a standard monomial basis: dt^1, dt^2, ..., dt^order (RV) or scan-angle-projected monomials (astrometry).

Parameters:

  • order (int) – Polynomial order (number of trend terms). Must be >= 1.

  • time_unit (str) – Unit string for time – used in ParamInfo metadata.

  • obs_unit (str) – Unit string for the observations – used in ParamInfo metadata.

  • astrometry (bool) – If True, add two columns per order (RA + Dec, projected by scan angle) with exponents k + 1 to avoid degeneracy with the base proper-motion columns. Default False (single column per order).

Examples

>>> from harv.models.extensions import MonomialTrend
>>> trend = MonomialTrend(order=2)
>>> [p.name for p in trend.extra_params()]
['trend_1', 'trend_2']
__init__(order=1, time_unit='', obs_unit='', astrometry=False)
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

astrometry: bool = False
extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.

modify_design_matrix(X, data, nl_values)

Append trend columns to the design matrix.

For RV (astrometry=False):

Columns are (t - t_ref)^k for k = 1..order.

For Gaia astrometry (astrometry=True):

Two columns per order: sin(psi) * dt^(k+1) and cos(psi) * dt^(k+1) where dt = (t - t_ref) and psi is the scan angle. The exponent k + 1 avoids degeneracy with the base proper-motion (dt^1) columns.

The data object must have time and t_ref attributes (and scan_angle for astrometry mode).

Parameters:
Return type:

Array

obs_unit: str = ''
order: int = 1
time_unit: str = ''
class harv.models.extensions.MultiSurveyOffset

Bases: AbstractExtension

Per-instrument additive offsets to account for instrumental systematics.

The extension stores a pre-computed indicator matrix and appends it as extra columns to the design matrix. Each column corresponds to a non-reference instrument; the reference instrument has no explicit offset (absorbed by the systemic velocity or baseline parameter).

Parameters:

  • indicator_matrix (Array) – Binary matrix: indicator_matrix[i, j] == 1 iff observation i belongs to the j-th non-reference instrument.

  • instrument_names (tuple[str, ...]) – Ordered names of the non-reference instruments (same column order as indicator_matrix). Used as parameter names.

  • obs_unit (str) – Physical unit string for the offset parameters (must match the observation unit of the model, e.g. "km/s").

Examples

>>> import jax.numpy as jnp
>>> from harv.models.extensions import MultiSurveyOffset
>>> indicator = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
>>> ext = MultiSurveyOffset(indicator, ("espresso", "keck"), "km/s")
>>> [p.name for p in ext.extra_params()]
['espresso', 'keck']
__init__(indicator_matrix, instrument_names, obs_unit='')
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.

modify_design_matrix(X, data, nl_values)

Append indicator columns to the design matrix.

Parameters:
Return type:

Array

obs_unit: str = ''
indicator_matrix: Array
instrument_names: tuple[str, ...]
class harv.models.extensions.ParamInfo

Bases: Module

Immutable descriptor for a single model parameter.

Parameters:

  • name (str) – Parameter name. Must not contain "." (reserved for JointModel tied-parameter paths).

  • unit (str) – Unit string (e.g. "day", "km/s", "" for dimensionless).

  • linear (bool) – Whether the parameter enters the model linearly (default False).

Examples

>>> from harv.models.extensions.base import ParamInfo
>>> p = ParamInfo("period", "time")
>>> p.name
'period'
>>> p.linear
False
__init__(name, unit, linear=False)
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

linear: bool = False
name: str
unit: str

Submodules

harv.models.extensions.base module

Parameter metadata and the abstract Extension base class.

This module defines:

  • ParamInfo: a frozen description of a single model parameter (name, unit / dimension (e.g., “km/s”), and role in a model (i.e., nonlinear vs linear)).

  • AbstractExtension: the abstract base that all model extensions must subclass. Extensions can add parameters, modify the design matrix (e.g., polynomial trends, or instrumental offsets), and/or modify the data covariance (e.g., to add a Gaussian Process noise model).

class harv.models.extensions.base.AbstractExtension

Bases: Module

Abstract base class for model extensions (jitter, trends, offsets, GP, …).

An extension can declare extra parameters (nonlinear and/or linear), modify the design matrix (e.g., append trend or offset columns), and/or modify the data covariance matrix (e.g., add jitter or a GP kernel).

Extensions are composed onto a AbstractComponentModel at construction time via the extensions field. The model calls each hook in the order the extensions are listed.

All extensions receive the raw (unit-stripped) JAX arrays. Unit handling is the model’s responsibility.

Subclasses must implement extra_params(). The design-matrix and covariance hooks have default no-op implementations that return their input unchanged.

Plot-specific behavior is handled privately by plotting helpers rather than through this base extension API.

abstractmethod extra_params()

Parameters introduced by this extension.

Returns:

May include both nonlinear and linear parameters.

Return type:

tuple of ParamInfo

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.

__init__()
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

class harv.models.extensions.base.ParamInfo

Bases: Module

Immutable descriptor for a single model parameter.

Parameters:

  • name (str) – Parameter name. Must not contain "." (reserved for JointModel tied-parameter paths).

  • unit (str) – Unit string (e.g. "day", "km/s", "" for dimensionless).

  • linear (bool) – Whether the parameter enters the model linearly (default False).

Examples

>>> from harv.models.extensions.base import ParamInfo
>>> p = ParamInfo("period", "time")
>>> p.name
'period'
>>> p.linear
False
name: str
unit: str
linear: bool = False
__init__(name, unit, linear=False)
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

harv.models.extensions.gp module

Gaussian Processes (GP).

final class harv.models.extensions.gp.GP

Bases: AbstractExtension

Gaussian Process covariance extension.

Adds the kernel matrix K(t, t') to the observation covariance via the modify_covariance hook, allowing analytic marginalization of linear parameters in the presence of correlated noise.

Plotting code may choose to visualize the GP contribution, but that support is handled privately on the plotting side rather than through the extension base API.

Parameters:

  • kernel_builder (Callable[[dict[str, Any]], Any]) – Receives the full nonlinear-parameter dict (unit-stripped) and returns a callable kernel object that can be called with (X, Xp) to produce the kernel matrix. For example, from tinygp.

  • hyperparams (tuple[ParamInfo, ...]) – Nonlinear hyperparameters declared by this extension (e.g. gp_amp, gp_length_scale). These are sampled alongside other nonlinear model parameters.

  • time_unit (str) – Unit to strip times to before building the coordinate array. Default "" (dimensionless / already stripped).

Examples

>>> from harv.models.extensions.base import ParamInfo
>>> from harv.models.extensions.gp import GP; GP(
...     kernel_builder=lambda hp: hp["gp_amp"] ** 2,
...     hyperparams=(ParamInfo("gp_amp", "km/s"),),
...     time_unit="day",
... ).extra_params()[0].name
'gp_amp'
kernel_builder: Callable[[dict[str, Any]], Any]
hyperparams: tuple[ParamInfo, ...]
time_unit: str = ''
extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Add the GP kernel matrix K(t, t') to the covariance.

If the incoming covariance is diagonal (1-d), it is promoted to a full (n_obs, n_obs) matrix first.

Parameters:
Return type:

Array

conditional_mean(residuals, data_times, prediction_times, data_err, hp)

Conditional-mean GP prediction at prediction_times.

Conditions on observed residuals (already mean-subtracted by the deterministic model prediction) at data_times with measurement noise data_err, then returns the GP posterior mean evaluated at prediction_times. Used by the plotting layer to overlay structured noise; the same kernel and hyperparameter conventions as modify_covariance() apply.

All time arrays are assumed to be unit-stripped to the same unit (the kernel’s coordinate unit). hp is the per-sample hyperparameter dict (e.g. from a single posterior sample).

Parameters:
Return type:

Array

__init__(kernel_builder, hyperparams, time_unit='')
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

harv.models.extensions.jitter module

Jitter (excess variance).

Examples

>>> from harv.models.extensions import Jitter; [p.name for p in Jitter().extra_params()]
['jitter']
final class harv.models.extensions.jitter.Jitter

Bases: AbstractExtension

Add excess variance (jitter / white noise) to the observation covariance.

The extension declares one nonlinear parameter, jitter, and modifies the covariance by adding jitter**2 in quadrature to the diagonal variances. The jitter unit must match the observation unit (e.g. "km/s" for RV, "mas" for astrometry). Unit stripping is the model’s responsibility.

Parameters:

param_unit (str) – Physical unit string for the jitter parameter metadata. Default "" (dimensionless / observation units).

Examples

>>> from harv.models.extensions import Jitter
>>> j = Jitter(param_unit="km/s")
>>> j.extra_params()[0].unit
'km/s'
param_unit: str = ''
extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_covariance(cov, data, nl_values)

Add jitter**2 in quadrature to the diagonal variances.

Works on both 1-d (diagonal) and 2-d (full) covariance representations.

Parameters:
Return type:

Array

__init__(param_unit='')
Parameters:

param_unit (str)

Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

modify_design_matrix(X, data, nl_values)

Optionally append columns to the design matrix.

Parameters:

  • X (Array) – Current design matrix (base + earlier extensions).

  • data (AbstractData | None) – Observation data (unit-stripped times, etc. accessed via helpers).

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated design matrix, shape (n_obs, n_cols + n_extra). Return X unchanged if not applicable.

harv.models.extensions.multi_survey module

Multi-survey offsets.

final class harv.models.extensions.multi_survey.MultiSurveyOffset

Bases: AbstractExtension

Per-instrument additive offsets to account for instrumental systematics.

The extension stores a pre-computed indicator matrix and appends it as extra columns to the design matrix. Each column corresponds to a non-reference instrument; the reference instrument has no explicit offset (absorbed by the systemic velocity or baseline parameter).

Parameters:

  • indicator_matrix (Array) – Binary matrix: indicator_matrix[i, j] == 1 iff observation i belongs to the j-th non-reference instrument.

  • instrument_names (tuple[str, ...]) – Ordered names of the non-reference instruments (same column order as indicator_matrix). Used as parameter names.

  • obs_unit (str) – Physical unit string for the offset parameters (must match the observation unit of the model, e.g. "km/s").

Examples

>>> import jax.numpy as jnp
>>> from harv.models.extensions import MultiSurveyOffset
>>> indicator = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
>>> ext = MultiSurveyOffset(indicator, ("espresso", "keck"), "km/s")
>>> [p.name for p in ext.extra_params()]
['espresso', 'keck']
indicator_matrix: Array
instrument_names: tuple[str, ...]
obs_unit: str = ''
extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_design_matrix(X, data, nl_values)

Append indicator columns to the design matrix.

Parameters:
Return type:

Array

__init__(indicator_matrix, instrument_names, obs_unit='')
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.

harv.models.extensions.trend module

Monomial polynomial trends.

final class harv.models.extensions.trend.MonomialTrend

Bases: AbstractExtension

Append monomial trend columns to the design matrix.

Uses a standard monomial basis: dt^1, dt^2, ..., dt^order (RV) or scan-angle-projected monomials (astrometry).

Parameters:

  • order (int) – Polynomial order (number of trend terms). Must be >= 1.

  • time_unit (str) – Unit string for time – used in ParamInfo metadata.

  • obs_unit (str) – Unit string for the observations – used in ParamInfo metadata.

  • astrometry (bool) – If True, add two columns per order (RA + Dec, projected by scan angle) with exponents k + 1 to avoid degeneracy with the base proper-motion columns. Default False (single column per order).

Examples

>>> from harv.models.extensions import MonomialTrend
>>> trend = MonomialTrend(order=2)
>>> [p.name for p in trend.extra_params()]
['trend_1', 'trend_2']
order: int = 1
time_unit: str = ''
obs_unit: str = ''
astrometry: bool = False
extra_params()

Parameters introduced by this extension.

Return type:

tuple[ParamInfo, ...]

modify_design_matrix(X, data, nl_values)

Append trend columns to the design matrix.

For RV (astrometry=False):

Columns are (t - t_ref)^k for k = 1..order.

For Gaia astrometry (astrometry=True):

Two columns per order: sin(psi) * dt^(k+1) and cos(psi) * dt^(k+1) where dt = (t - t_ref) and psi is the scan angle. The exponent k + 1 avoids degeneracy with the base proper-motion (dt^1) columns.

The data object must have time and t_ref attributes (and scan_angle for astrometry mode).

Parameters:
Return type:

Array

__init__(order=1, time_unit='', obs_unit='', astrometry=False)
Parameters:
Return type:

None

static __new__(cls, *args, **kwargs)
Parameters:
Return type:

TypeVar(_ModuleT, bound= Module)

modify_covariance(cov, data, nl_values)

Optionally modify the data covariance matrix.

Parameters:

  • cov (Array) – Current covariance. Diagonal (1-d) when only measurement errors are present; full (2-d) after a GP extension adds off-diagonal structure. shape (n_obs,) or (n_obs, n_obs)

  • data (AbstractData | None) – Observation data.

  • nl_values (dict[str, Any]) – Current nonlinear parameter values (unit-stripped scalars).

Return type:

Array

Returns:

Updated covariance (same or promoted shape). Return cov unchanged if not applicable.