tomojax.forward

ProjectionArrayGeometryInputCopy Anchor Link

class ProjectionArrayGeometryInput(
    theta_rad: jax.Array,
    dx_px: jax.Array,
    dz_px: jax.Array,
    alpha_rad: jax.Array | float = 0.0,
    beta_rad: jax.Array | float = 0.0,
    detector_roll_rad: jax.Array | float = 0.0,
    axis_rot_x_rad: jax.Array | float = 0.0,
    axis_rot_y_rad: jax.Array | float = 0.0,
    nominal_axis_unit: jax.Array | tuple[float, float, float] = (0.0, 0.0, 1.0),
    acquisition_model: str = 'parallel',
    laminography_tilt_rad: jax.Array | float = 0.0,
    laminography_tilt_about: str = 'x',
    detector_shape: tuple[int, int] | None = None,
)

Pose and setup arrays for reference forward projection.

joseph_l2_value_and_gradCopy Anchor Link

joseph_l2_value_and_grad(
    volume: jax.Array,
    poses: jax.Array,
    target: jax.Array,
    grid: Grid,
    detector: Detector,
    *,
    backend: Literal['jax', 'cuda'] = 'jax',
    interpolation: Literal['linear', 'cubic'] = 'linear',
) -> tuple[jax.Array, tuple[jax.Array, jax.Array]]

Return half squared error and (volume_gradient, pose_gradient).

Compute half the sum of squared projection residuals using the selected interpolation, without normalization, masking or fitted scale. Inputs follow project_joseph's contract; target is a finite real (views, nv, nu) stack. Pose gradients are with respect to the matrix entries: use a VJP through a rigid pose parameterization before taking an optimization step.

The CUDA path fuses projection and pose derivatives, retaining a residual sinogram and tile reductions, then applies the matched volume transpose. This CUDA operation provides first derivatives only; the JAX reference can be differentiated further. For other losses, compose project_joseph with ordinary JAX functions and differentiate it.

joseph_pose_normal_equationsCopy Anchor Link

joseph_pose_normal_equations(
    volume: jax.Array,
    poses: jax.Array,
    pose_directions: jax.Array,
    target: jax.Array,
    grid: Grid,
    detector: Detector,
    *,
    backend: Literal['jax', 'cuda'] = 'jax',
    interpolation: Literal['linear', 'cubic'] = 'linear',
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Return per-view least-squares loss, gradient, Gauss-Newton matrix, residual.

pose_directions has shape (views, 4, 4, P), with 1 <= P <= 16. Its last axis gives matrix-pose tangents for each view's parameters, in the caller's chosen units. For example, use vmap(jacfwd(make_pose)) for a function that maps one view's rigid parameters to its world-from-object matrix. Inputs otherwise follow project_joseph; target and directions must be finite and real, including under JIT. Directions should correspond to a rigid parameterization when optimizing physical poses.

For residual r = prediction - target and directional projection Jacobian J, return (0.5*r.T@r, J.T@r, J.T@J, r) independently per view. Shapes are (views,), (views, P), (views, P, P), and (views, nv, nu). The matrix is the Gauss-Newton approximation, not the full loss Hessian. Shared global parameters can sum these contributions; per-view parameters use separate blocks. Priors and cross-view coupling are not included.

CUDA recomputes plane samples once, reduces tile-local columns and retains the residual for a subsequent line search. It stores neither the full detector-by-parameter Jacobian nor a ray-by-plane tape. This explicit CUDA operation provides first-order information only; use the JAX reference for further differentiation. No damping, parameter step or gauge is chosen.

project_josephCopy Anchor Link

project_joseph(
    volume: jax.Array,
    poses: jax.Array,
    grid: Grid,
    detector: Detector,
    *,
    backend: Literal['jax', 'cuda'] = 'jax',
    interpolation: Literal['linear', 'cubic'] = 'linear',
) -> jax.Array

Project physical voxel-centre planes with zero-padded linear or cubic samples.

ParametersCopy Anchor Link

volume : jax.Array Finite real volume in (nx, ny, nz) order, converted to float32. poses : jax.Array Finite rigid world-from-object matrices, shape (views, 4, 4). Keep poses rigid when optimizing, for example by parameterizing rotations and translations. Shapes and static metadata are validated; array values must satisfy this contract, including under JIT. grid : Grid Physical voxel grid; its origin locates voxel (0, 0, 0)'s centre. detector : Detector Canonical physical detector grid, including detector offsets. backend : {'jax', 'cuda'} Explicit JAX reference or NVIDIA CUDA implementation; no fallback. interpolation : {'linear', 'cubic'} Bilinear or Keys cubic convolution (a=-1/2). Cubic uses a 4-by-4 stencil and can produce negative samples even for nonnegative inputs; it does not enforce positivity.

ReturnsCopy Anchor Link

jax.Array Float32 projection stack, shape (views, detector.nv, detector.nu).

RaisesCopy Anchor Link

ValueError If shapes, interpolation or backend are unsupported, including CUDA without a GPU.

NotesCopy Anchor Link

The default linear discretization matches cgls(projector_model='joseph'). It is distinct from the default trilinear ray model. CUDA supports JIT, vmap, JVP, VJP, jacfwd, jacrev and transposing a linearization, with first-order derivatives only. Use the JAX reference for higher derivatives. Derivatives follow the selected dominant-axis branch and interpolation cell, as in the reference. Axis switches are not smooth; cubic interpolation has continuous first coordinate derivatives within a fixed dominant-axis branch.

project_parallel_referenceCopy Anchor Link

project_parallel_reference(volume: jax.Array, geometry: GeometryState) -> jax.Array

Project a 3D volume from a GeometryState with the core trilinear ray projector.

project_parallel_reference_from_inputCopy Anchor Link

project_parallel_reference_from_input(
    volume: jax.Array,
    geometry_input: ProjectionArrayGeometryInput,
) -> jax.Array

Project a volume from supported pose/setup arrays using core trilinear rays.