tomojax.forward.api

PROJECTION_OPERATORCopy Anchor Link

PROJECTION_OPERATOR = 'core_trilinear_ray'

CoreProjectionGeometryCopy Anchor Link

class CoreProjectionGeometry(
    grid: Grid,
    detector: Detector,
    t_all: jax.Array,
    det_grid: tuple[jax.Array, jax.Array],
    operator: ProjectionOperatorName = PROJECTION_OPERATOR,
    step_size: float | None = None,
    n_steps: int | None = None,
    gather_dtype: str = 'fp32',
    checkpoint_projector: bool = True,
    projector_unroll: int = 1,
    detector_roll_rad: object = 0.0,
    axis_rot_x_rad: object = 0.0,
    axis_rot_y_rad: object = 0.0,
    alpha_rad_max_abs: object = 0.0,
    beta_rad_max_abs: object = 0.0,
    acquisition_model: str = 'parallel',
    laminography_tilt_rad: object = 0.0,
    laminography_tilt_about: str = 'x',
)

Store resolved geometry and settings for the reference ray projector.

Parameters

• grid (Grid): Physical volume grid.
• detector (Detector): Physical detector grid in the same length unit as grid.
• t_all (jax.Array): World-from-object transforms of shape (views, 4, 4).
• det_grid (tuple of jax.Array): Flattened detector-plane u and v coordinates, each of shape (nv * nu,), with detector-u varying fastest.
• operator ({'core_trilinear_ray'}, optional): Recorded projector name.
• step_size (float or None, optional): Ray sample spacing in physical grid units; None uses the core default.
• n_steps (int or None, optional): Ray sample count; None uses the core default.
• gather_dtype (str, optional): Core gather precision; default is 'fp32'.
• checkpoint_projector (bool, optional): Recompute projector intermediates during differentiation; default is True.
• projector_unroll (int, optional): Core ray-loop unrolling factor; default is 1.
• detector_roll_rad (object, optional): Recorded global detector roll in radians; default is zero.
• axis_rot_x_rad, axis_rot_y_rad (object, optional): Recorded global axis corrections in radians; default is zero.
• alpha_rad_max_abs, beta_rad_max_abs (object, optional): Recorded maximum absolute residual pose angles in radians; default is zero.
• acquisition_model (str, optional): Recorded acquisition model; default is 'parallel'.
• laminography_tilt_rad (object, optional): Recorded nominal tilt in radians; default is zero.
• laminography_tilt_about (str, optional): Recorded nominal tilt axis; default is 'x'.

CoreProjectionGeometry.provenanceCopy Anchor Link

provenance() -> dict[str, object]

Return serializable operator metadata.

Returns

• dict: Geometry, sampling settings and setup summaries. Traced scalar values that cannot be converted to Python floats are recorded as 'dynamic'.

ProjectionOperatorNameCopy Anchor Link

ProjectionOperatorName = Literal['core_trilinear_ray']

ResidualFilterConfigCopy Anchor Link

class ResidualFilterConfig(
    kind: ResidualFilterKind = 'raw',
    weight: float = 1.0,
    sigma_px: float = 1.0,
    outer_sigma_px: float = 2.0,
)

Configure a periodic detector-plane residual filter.

Parameters

• kind ({'raw', 'lowpass_gaussian', 'bandpass_difference_of_gaussians'}, optional): Filter family; default is 'raw'. Filtering acts on the last two axes.
• weight (float, optional): Multiplicative output weight; default is 1.
• sigma_px (float, optional): Gaussian width, or inner bandpass width, in pixels; default is 1.
• outer_sigma_px (float, optional): Outer bandpass Gaussian width in pixels; default is 2.

ResidualFilterKindCopy Anchor Link

ResidualFilterKind = Literal['raw', 'lowpass_gaussian', 'bandpass_difference_of_gaussians']

ResidualFilterResultCopy Anchor Link

class ResidualFilterResult(
    residual: jax.Array,
    components: tuple[jax.Array, ...],
    configs: tuple[ResidualFilterConfig, ...],
)

Store a summed residual-filter schedule and its individual outputs.

Parameters

• residual (jax.Array): Sum of all weighted, masked filter outputs, with the input shape.
• components (tuple of jax.Array): Individual weighted, masked outputs in schedule order.
• configs (tuple of ResidualFilterConfig): Settings used for each component, retained as static JAX metadata.

ResidualLossModeCopy Anchor Link

ResidualLossMode = Literal['l2', 'pseudo_huber']

ResidualResultCopy Anchor Link

class ResidualResult(
    residual: jax.Array,
    loss: jax.Array,
    weights: jax.Array,
    valid_count: jax.Array,
)

Store whitened residuals and their reduced robust loss.

Parameters

• residual (jax.Array): Masked, whitened differences with the projection shape.
• loss (jax.Array): Scalar sum of elementwise loss divided by max(valid_count, 1).
• weights (jax.Array): Elementwise robust weights, or ones for L2, with the residual shape.
• valid_count (jax.Array): Sum of mask values, or sample count when no mask was supplied.

apply_residual_filterCopy Anchor Link

apply_residual_filter(
    residual: jax.Array,
    config: ResidualFilterConfig,
    *,
    mask: jax.Array | None = None,
) -> jax.Array

Apply a weighted residual filter over the last two detector axes.

Gaussian filters use periodic boundaries. The mask is applied after filtering, not before it, so mask invalid input samples before calling when necessary.

Parameters

• residual (jax.Array): Residual array with at least two dimensions, ending in (nv, nu).
• config (ResidualFilterConfig): Filter family, width and multiplicative weight.
• mask (jax.Array or None, optional): Same-shaped multiplicative output mask; None leaves the output unmasked.

Returns

• jax.Array: Float32 filtered residual with the input shape.

Raises

• ValueError: If dimensions, filter name or mask shape are unsupported.

apply_residual_filter_scheduleCopy Anchor Link

apply_residual_filter_schedule(
    residual: jax.Array,
    configs: tuple[ResidualFilterConfig, ...],
    *,
    mask: jax.Array | None = None,
) -> ResidualFilterResult

Apply and sum a deterministic residual-filter schedule.

Parameters

• residual (jax.Array): Residual array ending in the detector axes (nv, nu).
• configs (tuple of ResidualFilterConfig): Nonempty sequence of weighted filter settings.
• mask (jax.Array or None, optional): Same-shaped multiplicative mask applied to every filtered output.

Returns

• ResidualFilterResult: Summed residual, individual components and their settings.

Raises

• ValueError: If the schedule is empty or any filter input is unsupported.

core_projection_geometry_from_inputCopy Anchor Link

core_projection_geometry_from_input(
    volume_shape: tuple[int, int, int],
    geometry_input: ProjectionArrayGeometryInput,
    *,
    gather_dtype: str = 'fp32',
    checkpoint_projector: bool = True,
    projector_unroll: int = 1,
    step_size: float | None = None,
    n_steps: int | None = None,
) -> CoreProjectionGeometry

Resolve grouped arrays into unit-spacing core projector geometry.

Parameters

• volume_shape (tuple of int): Volume dimensions (nx, ny, nz).
• geometry_input (ProjectionArrayGeometryInput): Pose and setup arrays and optional detector shape.
• gather_dtype (str, optional): Gather precision of the core projector; default is 'fp32'.
• checkpoint_projector (bool, optional): Recompute projector intermediates during differentiation; default is True.
• projector_unroll (int, optional): Core ray-loop unrolling factor; default is 1.
• step_size (float or None, optional): Ray sample spacing in unit-grid lengths; None uses the core default.
• n_steps (int or None, optional): Number of ray samples; None uses the core default. With a fixed sample count, step_size still controls the spacing.

Returns

• CoreProjectionGeometry: Unit-spacing voxel grid and detector, pose matrices and sampling settings.

Raises

• ValueError: If the volume shape or per-view geometry arrays are inconsistent.

core_projection_geometry_from_stateCopy Anchor Link

core_projection_geometry_from_state(
    volume_shape: tuple[int, int, int],
    geometry: GeometryState,
    *,
    detector_shape: tuple[int, int] | None = None,
    gather_dtype: str = 'fp32',
    checkpoint_projector: bool = True,
    projector_unroll: int = 1,
    step_size: float | None = None,
    n_steps: int | None = None,
) -> CoreProjectionGeometry

Resolve a geometry state into unit-spacing core projector geometry.

Parameters

• volume_shape (tuple of int): Volume dimensions (nx, ny, nz).
• geometry (GeometryState): Nominal acquisition, setup corrections and per-view poses.
• detector_shape (tuple of int or None, optional): Detector (rows, columns); None uses (nx, nz).
• gather_dtype (str, optional): Gather precision of the core projector; default is 'fp32'.
• checkpoint_projector (bool, optional): Recompute projector intermediates during differentiation; default is True.
• projector_unroll (int, optional): Core ray-loop unrolling factor; default is 1.
• step_size (float or None, optional): Ray sample spacing in unit-grid lengths; None uses the core default.
• n_steps (int or None, optional): Number of ray samples; None uses the core default. With a fixed sample count, step_size still controls the spacing.

Returns

• CoreProjectionGeometry: Unit-spacing voxel grid and detector, pose matrices and sampling settings.

Raises

• ValueError: If the volume shape or per-view geometry arrays are inconsistent.

masked_whitened_residualCopy Anchor Link

masked_whitened_residual(
    predicted: jax.Array,
    observed: jax.Array,
    *,
    mask: jax.Array | None = None,
    sigma: float | jax.Array = 1.0,
) -> jax.Array

Whiten differences, excluding zero-mask entries before unsafe arithmetic.

Parameters

• predicted, observed (jax.Array): Prediction and measured projections with identical shapes, usually (views, nv, nu). Values are converted to float32.
• mask (jax.Array or None, optional): Multiplicative weights of the same shape. Zero entries are excluded before subtraction and division; None keeps every sample.
• sigma (float or jax.Array, optional): Positive noise standard deviations, broadcastable to the projection shape; default is 1.

Returns

• jax.Array: Float32 (predicted-observed)/sigma multiplied by mask, or unmasked when mask is None. The shape matches the inputs.

Raises

• ValueError: If prediction, observation or mask shapes disagree.

nominal_axis_unit_from_geometryCopy Anchor Link

nominal_axis_unit_from_geometry(geometry: GeometryState) -> jax.Array

Return the nominal lab-frame rotation axis for a geometry state.

Parameters

• geometry (GeometryState): Acquisition model and nominal laminography tilt; setup corrections and residual per-view rotations are not applied here.

Returns

• jax.Array: Float32 unit vector of shape (3,). Parallel geometry returns (0, 0, 1).

pseudo_huber_lossCopy Anchor Link

pseudo_huber_loss(residual: jax.Array, *, delta: float = 1.0) -> jax.Array

Compute elementwise pseudo-Huber loss without small-residual cancellation.

Parameters

• residual (jax.Array): Residual array of any shape, converted to float32.
• delta (float, optional): Positive transition scale in residual units; default is 1.

Returns

• jax.Array: Elementwise delta**2 * (sqrt(1 + (residual/delta)**2) - 1), with the residual shape.

pseudo_huber_weightsCopy Anchor Link

pseudo_huber_weights(residual: jax.Array, *, delta: float = 1.0) -> jax.Array

Compute the iteratively reweighted least-squares weights for pseudo-Huber.

Parameters

• residual (jax.Array): Residual array of any shape, converted to float32.
• delta (float, optional): Positive transition scale in residual units; default is 1.

Returns

• jax.Array: Elementwise 1 / sqrt(1 + (residual/delta)**2), with the residual shape.

residual_lossCopy Anchor Link

residual_loss(
    predicted: jax.Array,
    observed: jax.Array,
    *,
    mask: jax.Array | None = None,
    sigma: float | jax.Array = 1.0,
    delta: float = 1.0,
    mode: ResidualLossMode = 'pseudo_huber',
) -> ResidualResult

Compute masked whitened residuals and their mean robust or L2 loss.

Parameters

• predicted, observed (jax.Array): Prediction and measured projections with identical shapes, usually (views, nv, nu). Values are converted to float32.
• mask (jax.Array or None, optional): Multiplicative weights of the same shape. Zero entries are excluded before subtraction and division; None keeps every sample.
• sigma (float or jax.Array, optional): Positive noise standard deviations, broadcastable to the projection shape; default is 1.
• delta (float, optional): Positive pseudo-Huber transition in whitened residual units; default is 1.
• mode ({'l2', 'pseudo_huber'}, optional): Elementwise loss; default is 'pseudo_huber'. L2 uses half squared error.

Returns

• ResidualResult: Residuals, scalar reduced loss, robust weights and normalization count. A nonbinary mask weights residuals and uses its sum for normalization.

Raises

• ValueError: If array shapes disagree or the loss mode is unsupported.

robust_residual_scaleCopy Anchor Link

robust_residual_scale(
    residual: jax.Array,
    *,
    mask: jax.Array | None = None,
) -> jax.Array

Estimate normal-consistent median absolute deviation under JAX transforms.

Empty, constant or nonfinite-MAD samples retain the fallback scale of one.

Parameters

• residual (jax.Array): Residual samples of any shape, converted to float32.
• mask (jax.Array or None, optional): Same-shaped sample selector; nonzero entries select samples and their magnitudes do not weight the medians. None selects all samples.

Returns

• jax.Array: Scalar 1.4826 times the selected median absolute deviation about the selected median, or one when that estimate is unusable.

Raises

• ValueError: If mask and residual shapes differ.