jaxpint.model#

Timing model orchestration: chains delays and sums phases.

Ports the orchestration logic from PINT’s TimingModel.delay() and TimingModel.phase() as a pure Equinox module. All fields are static (component lists don’t change during fitting), so TimingModel has no dynamic leaves – all differentiation flows through ParameterVector.values.

class jaxpint.model.TimingModel(delay_components, phase_components, dispersion_components=(), phoff_name=None)[source]#

Bases: Module

Orchestrates delay and phase components.

Parameters:

  • delay_components (tuple[DelayComponent, ]) – Delay components applied sequentially (order matters: each sees the accumulated delay from prior components).

  • phase_components (tuple[PhaseComponent, ]) – Phase components whose contributions are summed (order irrelevant).

  • phoff_name (str or None) – If set, a constant phase offset (-PHOFF) is applied after TZR subtraction. This matches PINT’s PhaseOffset semantics where the offset is applied only to observation TOAs, not the TZR.

  • dispersion_components (tuple[DispersionDelayComponent, ...])

delay_components: tuple[DelayComponent, ...]#
phase_components: tuple[PhaseComponent, ...]#
dispersion_components: tuple[DispersionDelayComponent, ...] = ()#
phoff_name: str | None = None#
compute_delay(toa_data, params)[source]#

Chain all delay components sequentially.

Each component sees the accumulated delay from prior components.

Parameters:

  • toa_data (TOAData) – Pulse time-of-arrival data.

  • params (ParameterVector) – Timing model parameters.

Returns:

Total signal delay in seconds.

Return type:

Float[Array, " n_toas"]

compute_dm(toa_data, params)[source]#

Compute total model DM at each TOA (pc/cm³).

Sums DM contributions from all DispersionDelayComponent instances in dispersion_components. Used for wideband DM residual computation.

Parameters:

  • toa_data (TOAData) – Pulse time-of-arrival data.

  • params (ParameterVector) – Timing model parameters.

Return type:

Float[Array, ‘n_toas’]

compute_phase(toa_data, params)[source]#

Compute absolute model phase for each TOA.

  1. Computes total delay via compute_delay().

  2. Sums phase contributions from all phase components.

  3. Subtracts the TZR reference phase for absolute phase (using toa_data.tzr_tdb_int/frac).

Parameters:

  • toa_data (TOAData) – Pulse time-of-arrival data.

  • params (ParameterVector) – Timing model parameters.

Returns:

Absolute pulse phase in cycles (int/frac split).

Return type:

DualFloat

property components: tuple[DelayComponent | PhaseComponent, ...]#

delay components then phase components.

Type:

All components in order

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

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

decompose_delay(toa_data, params)[source]#

Return each delay component’s individual contribution.

Runs the full sequential delay chain, recording each component’s delta (what it adds on top of the accumulated delay from prior components).

Warning

This method uses Python loops and returns a plain dict. It must not be called inside jax.jit-compiled functions. Use compute_delay() for JIT-compatible total-delay computation.

Returns:

Mapping from component name to its delay contribution in seconds, shape (n_toas,).

Return type:

dict[str, Array]

Parameters:
decompose_phase(toa_data, params)[source]#

Return each phase component’s individual contribution.

Computes total delay first, then evaluates each phase component independently (phase components are summed, not sequential).

Warning

This method uses Python loops and returns a plain dict. It must not be called inside jax.jit-compiled functions. Use compute_phase() for JIT-compatible total-phase computation.

Returns:

Mapping from component name to its phase contribution in cycles.

Return type:

dict[str, DualFloat]

Parameters:
decompose_dm(toa_data, params)[source]#

Return each dispersion component’s DM contribution.

Computes total delay first, then calls compute_dm on each DispersionDelayComponent. Component names match those used by decompose_delay() for the same components.

Warning

This method uses Python loops and returns a plain dict. It must not be called inside jax.jit-compiled functions. Use compute_dm() for JIT-compatible total-DM computation.

Returns:

Mapping from component name to its DM contribution in pc/cm³, shape (n_toas,).

Return type:

dict[str, Array]

Parameters: