ptyrax.models.detector

Classes

BackgroundEqualWeightDetector(*args[, ...])	Detector model with learnable dark-count background.
Detector(coordinates, sampling)	Abstract detector interface mapping input fields to measured detector counts.
NoiselessEqualWeightDetector(*args, **kwargs)	Ideal noiseless detector with equal weighting across coherent modes.

class ptyrax.models.detector.BackgroundEqualWeightDetector(*args, dynamic_range=None, dark_counts=None, scale=1.0, regularization_functions=(), **kwargs)

Bases: Detector

Detector model with learnable dark-count background.

Extends the equal-weight detector by adding a trainable 2-D dark_counts array that models detector background (e.g. thermal noise, stray light). The forward model averages amplitudes across coherent modes and adds the dark counts:

\[I_{\text{det}} = \frac{1}{K}\sum_k |\psi_k| + D\]

where \(D\) is the dark-count background and \(K\) is the number of coherent modes.

Regularization functions can be attached to penalize the dark-count map during optimization (e.g. smoothness or sparsity priors).

Variables:

• coordinates – Detector coordinate system.

• sampling – Detector pixel grid specification.

• dark_counts – Trainable 2-D background array of shape (m, n).

• dynamic_range – Tuple (min, max) of the detector dynamic range.

• scale – Multiplicative scale factor (static).

• regularization_functions – Tuple of callables applied to dark_counts to compute regularization losses.

Parameters:

• dynamic_range (tuple[TypeAliasForwardRef('jaxtyping.Float'), TypeAliasForwardRef('jaxtyping.Float')])

• dark_counts (Array)

• scale (jaxtyping.Float)

• regularization_functions (tuple[Callable[[Float[Array, 'm n']], TypeAliasForwardRef('jaxtyping.Float')]])

Example

>>> detector = BackgroundEqualWeightDetector(
...     coordinates=detector_coords,
...     sampling=detector_sampling,
...     dark_counts=jnp.zeros((256, 256)),
... )
>>> predicted = detector(exit_fields)

__call__(input_fields)

Compute the detector measurement including dark-count background.

Averages the amplitude across coherent modes and adds the learned dark-count background.

Parameters:

input_fields (PyTree[ptyrax.field.CoherentField]) – PyTree of coherent fields at the detector plane, with an .amplitude attribute of shape (n_modes, m, n, 1).

Returns:

A 2-D array of shape (m, n) containing the predicted detector readout (mean amplitude plus dark counts).

Return type:

Float[Array, ‘* m n’]

coordinates: CoordinateSystem

dark_counts: Array

dynamic_range: tuple[float, float]

fourier_mask(field, *, index=0, return_details=False)

Compute the Fourier-space occupancy mask for this detector.

Determines which pixels in the source field’s Fourier representation map onto the detector, given the relative tilt between the field coordinate system and the detector coordinate system.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier occupancy is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter, used when the detector has multiple indexed orientations.

• return_details (bool) – If True, return the full tuple of intermediate arrays instead of just the source occupancy.

Returns:

If return_details is False, a 2-D array of the source Fourier occupancy (fraction of each source pixel captured by the detector). If True, a tuple of (source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray]

fourier_support_mask(field, *, index=0, return_details=False)

Compute the Fourier-space binary support mask for this detector.

Similar to fourier_mask(), but returns a binary support indicating which source Fourier pixels have any overlap with the detector, rather than a fractional occupancy.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier support is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter.

• return_details (bool) – If True, return the full tuple of intermediate arrays.

Returns:

If return_details is False, a 2-D binary array of the source Fourier support. If True, a tuple of (source_support, source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray, ndarray]

classmethod from_hdf5_state(state, *, detector_path_prefix='detector')

Alias for load_from_hdf5_state().

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays.

• detector_path_prefix (str) – Path prefix for detector parameters in the state dictionary.

Returns:

An instance of the concrete detector subclass.

Return type:

Detector

classmethod load_from_hdf5_state(state, *, detector_path_prefix='detector')

Construct a concrete detector instance from an HDF5 state dictionary.

Extracts rotation, translation, pixel size, and optionally dark counts from the HDF5 state and builds the detector coordinate system and sampling grid.

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays, as returned by reading an HDF5 checkpoint file.

• detector_path_prefix (str) – Path prefix within the state dictionary under which detector parameters are stored.

Returns:

An instance of the concrete detector subclass populated from the state dictionary.

Raises:

• TypeError – If called directly on the abstract Detector base class.

• KeyError – If the required dark_counts key is missing from the state dictionary.

• ValueError – If the detector pixel size has fewer than one value.

Return type:

Detector

Example

>>> state = load_hdf5_to_dict("reconstruction.hdf5")
>>> detector = NoiselessEqualWeightDetector.load_from_hdf5_state(state, detector_path_prefix="detector")

regularization_functions: tuple[Callable[[Float[Array, 'm n']], float]]

sampling: SamplingGrid

scale: float

class ptyrax.models.detector.Detector(coordinates, sampling)

Bases: Module

Abstract detector interface mapping input fields to measured detector counts.

Parameters:

• coordinates (IndexSliceParameter[CoordinateSystem])

• sampling (SamplingGrid)

abstractmethod __call__(input_field, index)

Call self as a function.

Parameters:

• input_field (CoherentField)

• index (TypeAliasForwardRef('jaxtyping.Integer') | None)

Return type:

Float[Array, ‘* m n’]

coordinates: IndexSliceParameter[CoordinateSystem]

fourier_mask(field, *, index=0, return_details=False)

Compute the Fourier-space occupancy mask for this detector.

Determines which pixels in the source field’s Fourier representation map onto the detector, given the relative tilt between the field coordinate system and the detector coordinate system.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier occupancy is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter, used when the detector has multiple indexed orientations.

• return_details (bool) – If True, return the full tuple of intermediate arrays instead of just the source occupancy.

Returns:

If return_details is False, a 2-D array of the source Fourier occupancy (fraction of each source pixel captured by the detector). If True, a tuple of (source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray]

fourier_support_mask(field, *, index=0, return_details=False)

Compute the Fourier-space binary support mask for this detector.

Similar to fourier_mask(), but returns a binary support indicating which source Fourier pixels have any overlap with the detector, rather than a fractional occupancy.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier support is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter.

• return_details (bool) – If True, return the full tuple of intermediate arrays.

Returns:

If return_details is False, a 2-D binary array of the source Fourier support. If True, a tuple of (source_support, source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray, ndarray]

classmethod from_hdf5_state(state, *, detector_path_prefix='detector')

Alias for load_from_hdf5_state().

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays.

• detector_path_prefix (str) – Path prefix for detector parameters in the state dictionary.

Returns:

An instance of the concrete detector subclass.

Return type:

Detector

classmethod load_from_hdf5_state(state, *, detector_path_prefix='detector')

Construct a concrete detector instance from an HDF5 state dictionary.

Extracts rotation, translation, pixel size, and optionally dark counts from the HDF5 state and builds the detector coordinate system and sampling grid.

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays, as returned by reading an HDF5 checkpoint file.

• detector_path_prefix (str) – Path prefix within the state dictionary under which detector parameters are stored.

Returns:

An instance of the concrete detector subclass populated from the state dictionary.

Raises:

• TypeError – If called directly on the abstract Detector base class.

• KeyError – If the required dark_counts key is missing from the state dictionary.

• ValueError – If the detector pixel size has fewer than one value.

Return type:

Detector

Example

>>> state = load_hdf5_to_dict("reconstruction.hdf5")
>>> detector = NoiselessEqualWeightDetector.load_from_hdf5_state(state, detector_path_prefix="detector")

sampling: SamplingGrid

class ptyrax.models.detector.NoiselessEqualWeightDetector(*args, **kwargs)

Bases: Detector

Ideal noiseless detector with equal weighting across coherent modes.

Models a detector that sums the intensity contributions from all coherent modes with equal weight and no noise or background. The output can be returned as either amplitude (square root of summed intensity) or raw intensity, controlled by the mode parameter.

The forward model is:

\[I_{\text{det}} = \sum_k |\psi_k|^2\]

where \(\psi_k\) are the coherent mode fields at the detector plane.

Variables:

• coordinates – Detector coordinate system (position and orientation).

• sampling – Detector pixel grid specification.

• mode – Output mode — "amplitude" returns \(\sqrt{I_{\text{det}}}\), "intensity" returns \(I_{\text{det}}\) directly.

Example

>>> detector = NoiselessEqualWeightDetector(
...     coordinates=detector_coords,
...     sampling=detector_sampling,
...     mode="amplitude",
... )
>>> predicted = detector(exit_fields)

__call__(input_fields)

Compute the noiseless detector measurement from input fields.

Sums intensity contributions from all coherent modes and returns either the amplitude or intensity depending on the detector mode.

Parameters:

input_fields (PyTree[ptyrax.field.CoherentField]) – PyTree of coherent fields at the detector plane, with an .intensity attribute of shape (n_modes, m, n, 1).

Returns:

A 2-D array of shape (m, n) containing the predicted detector readout (amplitude or intensity).

Return type:

Float[Array, ‘* m n’]

coordinates: IndexSliceParameter[CoordinateSystem]

fourier_mask(field, *, index=0, return_details=False)

Compute the Fourier-space occupancy mask for this detector.

Determines which pixels in the source field’s Fourier representation map onto the detector, given the relative tilt between the field coordinate system and the detector coordinate system.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier occupancy is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter, used when the detector has multiple indexed orientations.

• return_details (bool) – If True, return the full tuple of intermediate arrays instead of just the source occupancy.

Returns:

If return_details is False, a 2-D array of the source Fourier occupancy (fraction of each source pixel captured by the detector). If True, a tuple of (source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray]

fourier_support_mask(field, *, index=0, return_details=False)

Compute the Fourier-space binary support mask for this detector.

Similar to fourier_mask(), but returns a binary support indicating which source Fourier pixels have any overlap with the detector, rather than a fractional occupancy.

Parameters:

• field (CoherentField) – The coherent source field whose Fourier support is evaluated against the detector geometry.

• index (jaxtyping.Integer) – Position index into the detector coordinate parameter.

• return_details (bool) – If True, return the full tuple of intermediate arrays.

Returns:

If return_details is False, a 2-D binary array of the source Fourier support. If True, a tuple of (source_support, source_occupancy, detector_mask, field_frame_target_indices).

Return type:

ndarray | tuple[ndarray, ndarray, ndarray, ndarray]

classmethod from_hdf5_state(state, *, detector_path_prefix='detector')

Alias for load_from_hdf5_state().

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays.

• detector_path_prefix (str) – Path prefix for detector parameters in the state dictionary.

Returns:

An instance of the concrete detector subclass.

Return type:

Detector

classmethod load_from_hdf5_state(state, *, detector_path_prefix='detector')

Construct a concrete detector instance from an HDF5 state dictionary.

Extracts rotation, translation, pixel size, and optionally dark counts from the HDF5 state and builds the detector coordinate system and sampling grid.

Parameters:

• state (Dict[str, ndarray]) – Dictionary mapping HDF5 dataset paths to NumPy arrays, as returned by reading an HDF5 checkpoint file.

• detector_path_prefix (str) – Path prefix within the state dictionary under which detector parameters are stored.

Returns:

An instance of the concrete detector subclass populated from the state dictionary.

Raises:

• TypeError – If called directly on the abstract Detector base class.

• KeyError – If the required dark_counts key is missing from the state dictionary.

• ValueError – If the detector pixel size has fewer than one value.

Return type:

Detector

Example

>>> state = load_hdf5_to_dict("reconstruction.hdf5")
>>> detector = NoiselessEqualWeightDetector.load_from_hdf5_state(state, detector_path_prefix="detector")

mode: Literal['amplitude', 'intensity'] = 'amplitude'

sampling: SamplingGrid