jaxpint.noise#

Noise components for JaxPINT timing models.

class jaxpint.noise.EcorrNoise(ecorr_names, quantization_matrix, ecorr_epoch_slices)[source]#

Bases: NoiseComponent

Epoch-correlated noise model (ECORR).

ECORR adds a low-rank contribution to the TOA covariance matrix:

C_ecorr = U · diag(ECORR²) · Uᵀ

where U is a binary quantization matrix mapping TOAs to observing epochs (pre-computed by the bridge) and the weights are the squared ECORR values.

Parameters:

  • ecorr_names (tuple of str) – Parameter names for ECORR instances (e.g. ("ECORR1", "ECORR2")). Values must be in seconds (the bridge converts from PINT’s native microseconds).

  • quantization_matrix (array, shape (n_toas, n_epochs)) – Binary matrix mapping TOAs to epochs. Pre-computed by the bridge because epoch identification is data-dependent and not JIT-compatible.

  • ecorr_epoch_slices (tuple of (int, int)) – For each ECORR parameter, the (start_col, end_col) range in the quantization matrix’s column dimension.

ecorr_names: tuple[str, ...]#
quantization_matrix: Float[Array, 'n_toas n_epochs']#
ecorr_epoch_slices: tuple[tuple[int, int], ...]#
ecorr_weights(params)[source]#

Return ECORR² weight for each epoch column.

Parameters:

params (ParameterVector) – Must contain values for all ECORR parameters.

Returns:

weights – Squared ECORR values (seconds²), one per epoch.

Return type:

(n_epochs,)

covariance(toa_data, params)[source]#

Return the Woodbury (Ndiag, U, Phidiag) triple for ECORR noise.

ECORR is purely low-rank: Ndiag = 0. The basis is the quantization matrix mapping TOAs to observing epochs.

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for array sizing).

  • params (ParameterVector) – Current parameter values for all ECORR parameters.

Returns:

  • Ndiag ((n_toas,)) – Zero diagonal (ECORR has no white component).

  • U ((n_toas, n_epochs)) – Binary quantization matrix.

  • Phidiag ((n_epochs,)) – Squared ECORR values (seconds squared) per epoch.

Return type:

tuple[Float[Array, ‘n_toas’], Float[Array, ‘n_toas n_epochs’], Float[Array, ‘n_epochs’]]

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

Draw a random ECORR noise realization.

Draws standard-normal epoch amplitudes and projects them through the quantization matrix scaled by sqrt(ECORR squared) values.

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for array dimensions).

  • params (ParameterVector) – Current parameter values for all ECORR parameters.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – ECORR noise realization in seconds.

Return type:

(n_toas,)

class jaxpint.noise.NoiseModel(white_noise, correlated, dm_white_noise=None)[source]#

Bases: Module

Container that aggregates all noise sources into a single interface.

Provides a unified Woodbury covariance decomposition:

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

where Ndiag comes from white noise (EFAC/EQUAD-scaled TOA uncertainties) and U / Phidiag are horizontally concatenated from all correlated noise components (ECORR, red noise, etc.).

Parameters:

  • white_noise (ScaleToaError or None) – White noise model (EFAC/EQUAD). When None, raw TOA errors are used.

  • correlated (tuple of NoiseComponent) – Correlated noise components whose basis matrices and weights are concatenated to form U and Phidiag.

  • dm_white_noise (ScaleDmError | None)

white_noise: ScaleToaError | None#
correlated: tuple[NoiseComponent, ...]#
dm_white_noise: ScaleDmError | None = None#
scaled_sigma(toa_data, params)[source]#

Return noise-scaled TOA uncertainties in seconds.

Parameters:
Return type:

Float[Array, ‘n_toas’]

covariance(toa_data, params)[source]#

Return the combined Woodbury (Ndiag, U, Phidiag) triple.

Returns:

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

  • U ((n_toas, n_basis)) – Concatenated basis matrices from all correlated components. Empty (n_toas, 0) when there are no correlated sources.

  • Phidiag ((n_basis,)) – Concatenated basis weights.

Parameters:
Return type:

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

scaled_dm_sigma(toa_data, params)[source]#

Return noise-scaled DM uncertainties in pc/cm³.

Parameters:
Return type:

Float[Array, ‘n_toas’]

wideband_covariance(toa_data, params)[source]#

Return wideband noise decomposition.

Returns:

  • Ndiag_toa ((n_toas,)) – TOA diagonal variance (white noise).

  • U_toa ((n_toas, n_basis)) – TOA correlated noise basis.

  • Phi_toa ((n_basis,)) – TOA correlated noise weights.

  • Ndiag_dm ((n_toas,)) – DM diagonal variance (white noise only).

Parameters:
Return type:

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

property has_correlated: bool#

True if any correlated noise components are present.

property components: tuple#

white_noise, correlated, dm_white_noise.

Type:

All non-None noise components

property component_names: tuple[str, ...]#

Unique names for all components, auto-disambiguated for duplicates.

class jaxpint.noise.PLChromNoise(fourier_basis, freqs, freq_bin_widths, tnchromamp_name, tnchromgam_name, tnchromidx_name, fref=1400.0)[source]#

Bases: NoiseComponent

Power-law chromatic noise with arbitrary chromatic index.

The raw Fourier design matrix is stored unscaled. At each call to covariance() or generate(), the basis is multiplied by (f_ref / f_obs)^α where α comes from the TNCHROMIDX parameter, making the scaling differentiable through α.

Parameters:

  • fourier_basis ((n_toas, 2 * n_freqs)) – Raw (unscaled) Fourier design matrix with alternating sin/cos columns.

  • freqs ((n_freqs,)) – Frequency array in Hz.

  • freq_bin_widths ((n_freqs,)) – Δf for each frequency bin.

  • tnchromamp_name (str) – Parameter name for the log10 amplitude.

  • tnchromgam_name (str) – Parameter name for the spectral index.

  • tnchromidx_name (str) – Parameter name for the chromatic index (α).

  • fref (float) – Reference radio frequency in MHz (default 1400.0).

fourier_basis: Float[Array, 'n_toas n_basis']#
freqs: Float[Array, 'n_freqs']#
freq_bin_widths: Float[Array, 'n_freqs']#
tnchromamp_name: str#
tnchromgam_name: str#
tnchromidx_name: str#
fref: float = 1400.0#
psd_weights(params)[source]#

Compute power-law PSD weights for the chromatic noise Fourier basis.

The power spectral density follows the convention:

P(f) = (A² / 12π²) · f_yr^(γ-3) · f^(-γ)

Each weight is P(f) · Δf, repeated twice for the sin/cos pair at that frequency.

Parameters:

params (ParameterVector) – Must contain values for TNCHROMAMP (log10 amplitude) and TNCHROMGAM (spectral index).

Returns:

weights – PSD weights for each basis column.

Return type:

(2 * n_freqs,)

covariance(toa_data, params)[source]#

Return the Woodbury (Ndiag, U, Phidiag) triple for chromatic noise.

Chromatic noise is purely low-rank: Ndiag = 0. The basis is scaled at runtime by (f_ref / f_obs)^alpha to account for the chromatic index.

Parameters:

  • toa_data (TOAData) – Observed TOA data including radio frequencies for chromatic scaling.

  • params (ParameterVector) – Current parameter values for amplitude, spectral index, and chromatic index.

Returns:

  • Ndiag ((n_toas,)) – Zero diagonal (chromatic noise has no white component).

  • U ((n_toas, 2 * n_freqs)) – Chromatically-scaled Fourier design matrix.

  • Phidiag ((2 * n_freqs,)) – Power-law PSD weights.

Return type:

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

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

Draw a random chromatic noise realization.

Draws standard-normal Fourier amplitudes and projects them through the chromatically-scaled basis matrix.

Parameters:

  • toa_data (TOAData) – Observed TOA data including radio frequencies for chromatic scaling.

  • params (ParameterVector) – Current parameter values for amplitude, spectral index, and chromatic index.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – Chromatic noise realization in seconds.

Return type:

(n_toas,)

class jaxpint.noise.PLDMNoise(fourier_basis, freqs, freq_bin_widths, tndmamp_name, tndmgam_name)[source]#

Bases: NoiseComponent

Power-law DM noise via frequency-scaled Fourier basis.

The Fourier design matrix is pre-scaled by (1400 / f_obs)² by the bridge, so that DM’s inverse-frequency-squared dependence is baked into the basis. The PSD weights depend on the amplitude (TNDMAMP) and spectral index (TNDMGAM) parameters and are computed dynamically so that they are differentiable.

Parameters:

  • fourier_basis ((n_toas, 2 * n_freqs)) – Pre-computed Fourier design matrix already scaled by (1400 / f_obs)² per TOA.

  • freqs ((n_freqs,)) – Frequency array in Hz.

  • freq_bin_widths ((n_freqs,)) – Δf for each frequency bin (used to weight the PSD).

  • tndmamp_name (str) – Parameter name for the log10 amplitude.

  • tndmgam_name (str) – Parameter name for the spectral index.

fourier_basis: Float[Array, 'n_toas n_basis']#
freqs: Float[Array, 'n_freqs']#
freq_bin_widths: Float[Array, 'n_freqs']#
tndmamp_name: str#
tndmgam_name: str#
psd_weights(params)[source]#

Compute power-law PSD weights for the DM noise Fourier basis.

The power spectral density follows the convention:

P(f) = (A² / 12π²) · f_yr^(γ-3) · f^(-γ)

Each weight is P(f) · Δf, repeated twice for the sin/cos pair at that frequency.

Parameters:

params (ParameterVector) – Must contain values for TNDMAMP (log10 amplitude) and TNDMGAM (spectral index).

Returns:

weights – PSD weights for each basis column.

Return type:

(2 * n_freqs,)

covariance(toa_data, params)[source]#

Return the Woodbury (Ndiag, U, Phidiag) triple for DM noise.

DM noise is purely low-rank: Ndiag = 0.

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for array sizing).

  • params (ParameterVector) – Current parameter values for DM amplitude and spectral index.

Returns:

  • Ndiag ((n_toas,)) – Zero diagonal (DM noise has no white component).

  • U ((n_toas, 2 * n_freqs)) – DM-scaled Fourier design matrix.

  • Phidiag ((2 * n_freqs,)) – Power-law PSD weights.

Return type:

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

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

Draw a random DM noise realization.

Draws standard-normal Fourier amplitudes and projects them through the DM-scaled basis matrix scaled by sqrt(weights).

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for basis matrix dimensions).

  • params (ParameterVector) – Current parameter values for DM amplitude and spectral index.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – DM noise realization in seconds.

Return type:

(n_toas,)

class jaxpint.noise.PLRedNoise(fourier_basis, freqs, freq_bin_widths, tnredamp_name, tnredgam_name)[source]#

Bases: NoiseComponent

Power-law red noise via alternating Fourier basis.

The Fourier design matrix F is pre-computed by the bridge from TOA times and stored as a JAX array. The PSD weights depend on the amplitude (TNREDAMP) and spectral index (TNREDGAM) parameters and are computed dynamically so that they are differentiable.

Parameters:

  • fourier_basis ((n_toas, 2 * n_freqs)) – Pre-computed Fourier design matrix with alternating sin/cos columns: [sin(2πf₁t), cos(2πf₁t), sin(2πf₂t), ...].

  • freqs ((n_freqs,)) – Frequency array in Hz.

  • freq_bin_widths ((n_freqs,)) – Δf for each frequency bin (used to weight the PSD).

  • tnredamp_name (str) – Parameter name for the log10 amplitude.

  • tnredgam_name (str) – Parameter name for the spectral index.

fourier_basis: Float[Array, 'n_toas n_basis']#
freqs: Float[Array, 'n_freqs']#
freq_bin_widths: Float[Array, 'n_freqs']#
tnredamp_name: str#
tnredgam_name: str#
psd_weights(params)[source]#

Compute power-law PSD weights for the Fourier basis.

Returns one weight per basis column (sin and cos of each frequency get the same weight).

The power spectral density follows the convention:

P(f) = (A² / 12π²) · f_yr^(γ-3) · f^(-γ)

Each weight is P(f) · Δf, repeated twice for the sin/cos pair at that frequency.

Parameters:

params (ParameterVector) – Must contain values for TNREDAMP (log10 amplitude) and TNREDGAM (spectral index).

Returns:

weights – PSD weights for each basis column.

Return type:

(2 * n_freqs,)

covariance(toa_data, params)[source]#

Return the Woodbury (Ndiag, U, Phidiag) triple for red noise.

Red noise is purely low-rank: Ndiag = 0.

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for array sizing).

  • params (ParameterVector) – Current parameter values for amplitude and spectral index.

Returns:

  • Ndiag ((n_toas,)) – Zero diagonal (red noise has no white component).

  • U ((n_toas, 2 * n_freqs)) – Fourier design matrix.

  • Phidiag ((2 * n_freqs,)) – Power-law PSD weights.

Return type:

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

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

Draw a random red noise realization.

Draws standard-normal Fourier amplitudes and projects them through the basis matrix scaled by sqrt(weights).

Parameters:

  • toa_data (TOAData) – Observed TOA data (used for basis matrix dimensions).

  • params (ParameterVector) – Current parameter values for amplitude and spectral index.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – Red noise realization in seconds.

Return type:

(n_toas,)

class jaxpint.noise.PLSWNoise(fourier_basis, freqs, freq_bin_widths, tnswamp_name, tnswgam_name, swm, swp_name=None, raj_name='RAJ', decj_name='DECJ', pmra_name=None, pmdec_name=None, posepoch_name=None, obliquity_arcsec=None)[source]#

Bases: NoiseComponent

Power-law solar wind DM noise.

The raw Fourier design matrix is stored unscaled. At each call the basis is multiplied by geometry_pc · DMCONST / f_obs² where the geometry factor is computed from the solar wind model (SWM=0 or 1).

Parameters:

  • fourier_basis ((n_toas, 2 * n_freqs)) – Raw (unscaled) Fourier design matrix.

  • freqs ((n_freqs,)) – Frequency array in Hz.

  • freq_bin_widths ((n_freqs,)) – Δf for each frequency bin.

  • tnswamp_name (str) – Parameter name for the log10 amplitude.

  • tnswgam_name (str) – Parameter name for the spectral index.

  • swm (int) – Solar wind model (0 or 1).

  • swp_name (str or None) – Parameter name for the radial power-law index (SWM=1 only).

  • raj_name (str) – Astrometry parameter names for pulsar direction.

  • decj_name (str) – Astrometry parameter names for pulsar direction.

  • pmra_name (str or None) – Proper motion parameter names.

  • pmdec_name (str or None) – Proper motion parameter names.

  • posepoch_name (str or None) – Position epoch parameter name.

  • obliquity_arcsec (float or None) – Obliquity in arcseconds (set when using ecliptic coordinates).

fourier_basis: Float[Array, 'n_toas n_basis']#
freqs: Float[Array, 'n_freqs']#
freq_bin_widths: Float[Array, 'n_freqs']#
tnswamp_name: str#
tnswgam_name: str#
swm: int#
swp_name: str | None = None#
raj_name: str = 'RAJ'#
decj_name: str = 'DECJ'#
pmra_name: str | None = None#
pmdec_name: str | None = None#
posepoch_name: str | None = None#
obliquity_arcsec: float | None = None#
psd_weights(params)[source]#

Compute power-law PSD weights for the solar wind noise Fourier basis.

The power spectral density follows the convention:

P(f) = (A² / 12π²) · f_yr^(γ-3) · f^(-γ)

Each weight is P(f) · Δf, repeated twice for the sin/cos pair at that frequency.

Parameters:

params (ParameterVector) – Must contain values for TNSWAMP (log10 amplitude) and TNSWGAM (spectral index).

Returns:

weights – PSD weights for each basis column.

Return type:

(2 * n_freqs,)

covariance(toa_data, params)[source]#

Return the Woodbury (Ndiag, U, Phidiag) triple for solar wind noise.

Solar wind noise is purely low-rank: Ndiag = 0. The basis is scaled at runtime by the solar wind geometry factor and DMCONST / f_obs^2.

Parameters:

  • toa_data (TOAData) – Observed TOA data including Sun positions and radio frequencies.

  • params (ParameterVector) – Current parameter values for amplitude, spectral index, and astrometry parameters.

Returns:

  • Ndiag ((n_toas,)) – Zero diagonal (solar wind noise has no white component).

  • U ((n_toas, 2 * n_freqs)) – Solar-wind-geometry-scaled Fourier design matrix.

  • Phidiag ((2 * n_freqs,)) – Power-law PSD weights.

Return type:

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

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

Draw a random solar wind noise realization.

Draws standard-normal Fourier amplitudes and projects them through the SW-geometry-scaled basis matrix.

Parameters:

  • toa_data (TOAData) – Observed TOA data including Sun positions and radio frequencies.

  • params (ParameterVector) – Current parameter values for amplitude, spectral index, and astrometry parameters.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – Solar wind noise realization in seconds.

Return type:

(n_toas,)

class jaxpint.noise.ScaleToaError(efac_names, equad_names)[source]#

Bases: NoiseComponent

White noise model: EFAC/EQUAD scaling of TOA uncertainties.

Stores the names of EFAC and EQUAD parameters (static metadata). The boolean masks selecting which TOAs each parameter applies to live in TOAData.flag_masks (extracted by the bridge).

Parameters:

  • efac_names (tuple of str) – Parameter names for EFAC instances (e.g. ("EFAC1", "EFAC2")).

  • equad_names (tuple of str) – Parameter names for EQUAD instances (e.g. ("EQUAD1", "EQUAD2")). Values must be in seconds (the bridge converts from PINT’s native microseconds).

efac_names: tuple[str, ...]#
equad_names: tuple[str, ...]#
scaled_sigma(toa_data, params)[source]#

Compute noise-scaled TOA uncertainties.

Applies EQUAD in quadrature first, then multiplies by EFAC, matching PINT’s ScaleToaError.scale_toa_sigma() convention.

Parameters:

  • toa_data (TOAData) – Must contain error (seconds) and flag_masks with entries for every name in efac_names and equad_names.

  • params (ParameterVector) – Must contain values for all EFAC/EQUAD parameters.

Returns:

sigma_scaled – Scaled uncertainties in seconds.

Return type:

(n_toas,)

covariance(toa_data, params)[source]#

Compute the white noise diagonal covariance.

Returns the EFAC/EQUAD-scaled variance as the diagonal term, with no low-rank (basis) contribution.

Parameters:

  • toa_data (TOAData) – Observed TOA data including raw uncertainties and flag masks.

  • params (ParameterVector) – Current parameter values for all EFAC/EQUAD parameters.

Returns:

  • Ndiag ((n_toas,)) – Scaled variance for each TOA (seconds squared).

  • U (None) – No basis matrix (white noise is purely diagonal).

  • Phidiag (None) – No basis weights.

Return type:

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

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

Draw a random white noise realization.

Samples independent Gaussian noise scaled by the EFAC/EQUAD-adjusted TOA uncertainties.

Parameters:

  • toa_data (TOAData) – Observed TOA data including raw uncertainties and flag masks.

  • params (ParameterVector) – Current parameter values for all EFAC/EQUAD parameters.

  • key (jax.Array) – PRNG key for random sampling.

Returns:

noise – White noise realization in seconds.

Return type:

(n_toas,)

On this page
Show Source