jaxpint.components#

Base component types for JaxPINT timing model modules.

Component fields that store parameter names must follow the naming convention: end with _name for a single parameter name, or _names for a tuple of parameter names. This enables required_params() to discover them automatically.

class jaxpint.components.PhaseComponent[source]#

Bases: Module

Base class for components that contribute to pulse phase.

Subclasses implement __call__(self, toa_data, params, delay) -> DualFloat.

In the timing model, all PhaseComponents see the same total delay and their phase contributions are summed.

Fields that store parameter names must end with _name (single) or _names (tuple). This enables required_params().

required_params()[source]#

Parameter names this component reads from the ParameterVector.

Discovered by convention: fields ending in _name (single parameter) or _names (tuple of parameters). New component fields that hold parameter names must follow this convention.

Return type:

tuple[str, …]

class jaxpint.components.NoiseComponent[source]#

Bases: Module

Base class for stochastic noise sources.

Every noise source decomposes its covariance as:

C = diag(Ndiag) + U @ diag(Phidiag) @ Uᵀ

Subclasses must implement:

  • covariance — returns the (Ndiag, U, Phidiag) triple.

  • generate — draws a random noise realization.

The fitter combines multiple NoiseComponent instances by summing their diagonal contributions and horizontally concatenating their basis matrices and weight vectors.

Fields that store parameter names must end with _name (single) or _names (tuple). This enables required_params().

covariance(toa_data, params)[source]#

Return the Woodbury decomposition of this component’s covariance.

Returns (Ndiag, U, Phidiag) such that:

C = diag(Ndiag) + U @ diag(Phidiag) @ Uᵀ

For purely diagonal noise, return (Ndiag, None, None). For purely low-rank noise, return (zeros, U, Phidiag).

Parameters:

  • toa_data (TOAData) – Pre-extracted TOA data (TDB times, frequencies, positions, etc.).

  • params (ParameterVector) – Timing model parameters.

Returns:

  • Ndiag ((n_toas,)) – Diagonal variance contribution.

  • U ((n_toas, n_basis) or None) – Basis matrix for low-rank contribution.

  • Phidiag ((n_basis,) or None) – Basis weights for low-rank contribution.

Raises:

NotImplementedError – Must be overridden by subclasses.

Return type:

tuple[Float[Array, ‘n_toas’], Float[Array, ‘n_toas n_basis’] | None, Float[Array, ‘n_basis’] | None]

generate(toa_data, params, key)[source]#

Draw a random noise realization.

Parameters:

  • toa_data (TOAData) – Pre-extracted TOA data.

  • params (ParameterVector) – Timing model parameters (including noise parameter values).

  • key (JAX PRNG key) – Random key for reproducible sampling.

Returns:

Noise delays in seconds.

Return type:

(n_toas,)

Raises:

NotImplementedError – Must be overridden by subclasses.

required_params()[source]#

Parameter names this component reads from the ParameterVector.

Discovered by convention: fields ending in _name (single parameter) or _names (tuple of parameters). New component fields that hold parameter names must follow this convention.

Return type:

tuple[str, …]

class jaxpint.components.DelayComponent[source]#

Bases: Module

Base class for components that contribute to signal delay.

Subclasses implement __call__(self, toa_data, params, delay) -> Array.

In the timing model, DelayComponents are applied sequentially: each component sees the accumulated delay from prior components.

Fields that store parameter names must end with _name (single) or _names (tuple). This enables required_params().

required_params()[source]#

Parameter names this component reads from the ParameterVector.

Discovered by convention: fields ending in _name (single parameter) or _names (tuple of parameters). New component fields that hold parameter names must follow this convention.

Return type:

tuple[str, …]

class jaxpint.components.DispersionDelayComponent[source]#

Bases: DelayComponent

Base class for delay components that contribute to dispersion measure.

Subclasses must implement compute_dm() returning the DM contribution in pc/cm³, in addition to __call__() which returns delay in seconds. The timing model uses compute_dm to evaluate the total model DM for wideband fitting.

compute_dm(toa_data, params, delay)[source]#

Return this component’s DM contribution in pc/cm³.

Parameters:

  • toa_data (TOAData) – Pre-extracted TOA data.

  • params (ParameterVector) – Timing model parameters.

  • delay ((n_toas,)) – Accumulated signal delay in seconds from prior delay components.

Returns:

DM contribution in pc/cm³.

Return type:

(n_toas,)

Raises:

NotImplementedError – Must be overridden by subclasses.