tomojax.alignment

Most code calls tomojax.align on a tomojax.Scan. This package holds what that uses: AlignConfig (expert solver settings), alignment_plan (the configuration and levels a mode runs) and align_multires, which aligns arrays against a geometry. Schedules, losses, gauges and the single-resolution solver are in tomojax.alignment.api.

MODESCopy Anchor Link

MODES = ('pose', 'cor', 'cor-then-pose', 'full')

AlignConfigCopy Anchor Link

class AlignConfig(
    *,
    quality: QualityTier = 'fast',
    outer_iterations: int = 5,
    iterations: int = 10,
    tv_weight: float = 0.005,
    regulariser: Regulariser = 'huber_tv',
    huber_delta: float = 0.01,
    tv_prox_iterations: int = 10,
    reconstruction: Reconstruction = 'fista',
    nonnegative: bool = True,
    seed: int = 0,
    recon_rel_tol: float | None = None,
    recon_patience: int = 2,
    lr_rot: float = 0.001,
    lr_trans: float = 0.1,
    views_per_batch: int = 0,
    projector_unroll: int = 1,
    checkpoint_projector: bool = True,
    gather_dtype: str = 'auto',
    backend: ProjectorBackendInput = 'cuda',
    ray_integrator: Literal['sampled', 'exact', 'joseph', 'joseph_cubic'] = 'sampled',
    opt_method: str = 'gn',
    gn_damping: float = 0.001,
    gn_jacobian: Literal['autodiff', 'central'] = 'autodiff',
    gn_difference_step: float = 0.001,
    gn_coupling: Literal['fixed_volume', 'joint'] = 'fixed_volume',
    gn_joint_solver: Literal['stacked', 'pose_eliminated'] = 'stacked',
    gn_joint_iterations: int = 40,
    gn_joint_rtol: float = 0.0001,
    gn_volume_damping: float = 0.001,
    lbfgs_maxiter: int = 20,
    lbfgs_ftol: float = 1e-06,
    lbfgs_gtol: float = 1e-05,
    lbfgs_maxls: int = 20,
    lbfgs_memory_size: int = 10,
    w_rot: float = 0.0,
    w_trans: float = 0.0,
    schedule: str | AlignmentSchedule | None = None,
    optimise_dofs: tuple[str, ...] | None = None,
    freeze: tuple[str, ...] = tuple(),
    bounds: DofBounds | str | Mapping[str, object] = tuple(),
    gauge_policy: GaugePolicy = 'reject',
    gauge_priors: Mapping[str, object] | None = None,
    pose_model: PoseModel = 'per_view',
    pose_translation_frame: PoseTranslationFrame = 'object',
    knot_spacing: int = 8,
    degree: int = 3,
    seed_translations: bool = False,
    mask_vol: str = 'off',
    log_summary: bool = False,
    log_compact: bool = True,
    lipschitz: float | None = None,
    early_stop: bool = True,
    early_stop_rel_impr: float = 0.001,
    early_stop_patience: int = 2,
    gn_accept_only_improving: bool = True,
    gn_accept_tol: float = 0.0,
    loss: AlignmentLossConfig = L2OtsuLossSpec(),
)

Configure reconstruction, pose updates and alignment schedules.

For ordinary scans use tomojax.align and its mode/quality choices. This configuration exposes expert settings for the array-level solvers.

Parameters

• quality (QualityTier, optional): Quality policy, "fast" or "reference". Reference forces the JAX reference reconstruction settings. Default is 'fast'.
• outer_iterations (int, optional): Maximum number of outer alignment/reconstruction iterations. Default is 5.
• iterations (int, optional): Inner reconstruction iterations per outer iteration. Default is 10.
• tv_weight (float, optional): Nonnegative volume regularization weight. Default is 0.005.
• regulariser (Regulariser, optional): Volume regularization model, "tv" or "huber_tv". Default is 'huber_tv'.
• huber_delta (float, optional): Positive transition scale for Huber-TV. Default is 0.01.
• tv_prox_iterations (int, optional): Iterations per TV proximal solve. Default is 10.
• reconstruction (Reconstruction, optional): Inner reconstruction method, "fista" or "spdhg". Default is 'fista'.
• nonnegative (bool, optional): Constrain reconstructed voxel values to be nonnegative. Default is True.
• seed (int, optional): Random seed for stochastic reconstruction. Default is 0.
• recon_rel_tol (float | None, optional): Relative objective-change threshold for inner reconstruction stopping. None disables the threshold.
• recon_patience (int, optional): Successive small-change iterations required for reconstruction stopping. Default is 2.
• lr_rot (float, optional): Gradient-descent rotation step size in radians. Default is 0.001.
• lr_trans (float, optional): Gradient-descent translation step size in physical length units. Default is 0.1.
• views_per_batch (int, optional): View batch size; zero chooses it from free device memory at alignment startup. Default is 0.
• projector_unroll (int, optional): Unrolling factor for the reference projector loop. Default is 1.
• checkpoint_projector (bool, optional): Recompute intermediate projector arrays during differentiation to save memory. Default is True.
• gather_dtype (str, optional): Reference gather precision: "auto", "fp32", "fp16" or "bf16". Default is 'auto'.
• backend (ProjectorBackendInput, optional): Projector implementation: "jax" or "cuda". Default is 'cuda'.
• ray_integrator (Literal['sampled', 'exact', 'joseph', 'joseph_cubic'], optional): Ray model: "sampled", "exact", "joseph" or "joseph_cubic". Default is 'sampled'.
• opt_method (str, optional): Pose optimizer: "gd", "gn" or "lbfgs". Default is 'gn'.
• gn_damping (float, optional): Gauss-Newton pose damping; must be finite and positive for joint solves. Default is 0.001.
• gn_jacobian (Literal['autodiff', 'central'], optional): Pose derivatives, "autodiff" or "central" finite differences. Default is 'autodiff'.
• gn_difference_step (float, optional): Positive central-difference displacement as a fraction of the smallest voxel pitch. Default is 0.001.
• gn_coupling (Literal['fixed_volume', 'joint'], optional): Gauss-Newton update model: "fixed_volume" or "joint" volume and poses. Default is 'fixed_volume'.
• gn_joint_solver (Literal['stacked', 'pose_eliminated'], optional): Joint linear-system formulation, "stacked" or "pose_eliminated". Default is 'stacked'.
• gn_joint_iterations (int, optional): Positive limit on iterations of each joint linear solve. Default is 40.
• gn_joint_rtol (float, optional): Relative tolerance between zero and one for each joint linear solve. Default is 0.0001.
• gn_volume_damping (float, optional): Positive volume damping in the joint Gauss-Newton system. Default is 0.001.
• lbfgs_maxiter (int, optional): Maximum L-BFGS iterations per pose update. Default is 20.
• lbfgs_ftol (float, optional): L-BFGS objective-change stopping tolerance. Default is 1e-06.
• lbfgs_gtol (float, optional): L-BFGS gradient stopping tolerance. Default is 1e-05.
• lbfgs_maxls (int, optional): Maximum line-search steps per L-BFGS update. Default is 20.
• lbfgs_memory_size (int, optional): Number of correction pairs retained by L-BFGS. Default is 10.
• w_rot (float, optional): Rotation regularization weight for pose updates. Default is 0.0.
• w_trans (float, optional): Translation regularization weight for pose updates. Default is 0.0.
• schedule (str | AlignmentSchedule | None, optional): Named preset or explicit stage schedule. Mutually exclusive with optimise_dofs; None uses the default pose schedule.
• optimise_dofs (tuple[str, ...] | None, optional): Explicit active parameter names instead of a named schedule.
• freeze (tuple[str, ...], optional): Parameter names held at their initial values; defaults to an empty tuple.
• bounds (DofBounds | str | Mapping[str, object], optional): Bounds by parameter name as triples, a mapping or a CLI-style string. Pose rotations use radians, translations physical units, setup angles degrees and detector offsets native pixels. Defaults to no bounds.
• gauge_policy (GaugePolicy, optional): Gauge handling: "reject", "anchor_mean", "prior_required" or "diagnose_only". Default is 'reject'.
• gauge_priors (Mapping[str, object] | None, optional): Named priors required by the selected gauge policy, or None.
• pose_model (PoseModel, optional): Per-view motion representation: "per_view", "polynomial" or "spline". Default is 'per_view'.
• pose_translation_frame (PoseTranslationFrame, optional): Translation basis: "object" or lab-frame "detector". Translations are physical lengths, not pixels. Default is 'object'.
• knot_spacing (int, optional): Positive view spacing between spline knots. Default is 8.
• degree (int, optional): Polynomial degree, or spline degree 1, 2 or 3. Default is 3.
• seed_translations (bool, optional): Initialize per-view detector shifts with a global search. Default is False.
• mask_vol (str, optional): Reconstruction support used for projection: "off" or an xy cylinder ("cyl"). Default is 'off'.
• log_summary (bool, optional): Log progress after each outer iteration. Default is False.
• log_compact (bool, optional): Use a compact line when summary logging is enabled. Default is True.
• lipschitz (float | None, optional): Fixed inner-FISTA data Lipschitz constant, or None for power estimation.
• early_stop (bool, optional): Stop outer alignment iterations after sustained small improvement. Default is True.
• early_stop_rel_impr (float, optional): Relative loss-improvement threshold for outer stopping. Default is 0.001.
• early_stop_patience (int, optional): Successive small-improvement outer iterations required to stop. Default is 2.
• gn_accept_only_improving (bool, optional): Accept Gauss-Newton updates only when loss improves within tolerance. Default is True.
• gn_accept_tol (float, optional): Allowed loss increase as a fraction of the loss before a Gauss-Newton step. Default is 0.0.
• loss (AlignmentLossConfig, optional): Projection-domain loss specification or multiresolution loss schedule. Defaults to soft Otsu-weighted L2.

Raises

• ValueError: If an option is unknown, inconsistent with the selected optimizer or schedule, or outside its supported range.

Examples

from tomojax.alignment import AlignConfig
config = AlignConfig(quality="reference", iterations=5)
config.backend
# 'jax'

AlignmentPlanCopy Anchor Link

class AlignmentPlan(
    mode: AlignmentMode,
    config: AlignConfig,
    levels: tuple[int, ...],
    pose_solver: Literal['coupled', 'alternating'],
)

What an alignment run executes: its configuration and resolution levels.

Parameters

• mode (str): Alignment mode: "pose", "cor", "cor-then-pose" or "full".
• config (AlignConfig): Resolved solver configuration.
• levels (tuple of int): Coarse-to-fine resolution factors.
• pose_solver ({"coupled", "alternating"}): Joint volume/pose updates or alternating fixed-volume pose updates.

ProjectionMatchingConfigCopy Anchor Link

class ProjectionMatchingConfig(*, iterations: int = 20)

Joint translation alignment of axial parallel-beam scans.

Vertical shifts use noise-weighted axial mass profiles. Horizontal shifts use five detector rows with matched low-pass sampling and joint image/pose refinement from native correlation and first-moment starts. The full scan is retained without resampling.

Read per-view vertical standard deviations from info["vertical_uncertainty"], in detector length units. These condition on the fitted template and covariance and do not cover horizontal alignment.

Parameters

• iterations (int, optional): Maximum joint volume/pose refinement updates per start, not inner conjugate-gradient iterations. Default is 20.

align_multiresCopy Anchor Link

align_multires(
    geometry: Geometry,
    grid: Grid,
    detector: Detector,
    projections: jnp.ndarray,
    *,
    factors: Iterable[int] = (2, 1),
    config: AlignConfig | None = None,
    observer: ObserverCallback | None = None,
    resume_state: AlignMultiresResumeState | None = None,
    checkpoint_callback: AlignMultiresCheckpointCallback | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, AlignMultiresInfo]

Run multiresolution alignment.

Moving the object rigidly, and every pose with the inverse motion, predicts the same data, so the poses are only determined up to that motion. The result is the estimate with the least per-view motion: no common rotation and no rigid shift left in the poses (info["gauge"] records what was removed). In a parallel beam with detector-frame poses, a schedule that reports the detector centre also takes the poses' constant u shift into it.

Parameters

• geometry (Geometry): Nominal scan geometry with view angles and any existing per-view poses.
• grid (Grid): Reconstruction grid with physical voxel spacing in (x, y, z) order.
• detector (Detector): Detector shape and placement in the same length unit as the grid.
• projections (array-like): Measured line integrals with shape (views, nv, nu).
• factors (iterable of int, optional): Coarse-to-fine resolution factors. Default is (2, 1).
• config (AlignConfig or None, optional): Expert solver settings. None uses AlignConfig defaults.
• observer (callable or None, optional): Receives (volume, pose_params, outer_stat) after outer iterations. Returns "continue", "advance_level", "stop_run" or None to continue.
• resume_state (AlignMultiresResumeState or None, optional): Previously saved progress; None starts a new solve.
• checkpoint_callback (callable or None, optional): Receives a AlignMultiresResumeState to persist after outer iterations.

Returns

• volume (jax.Array): Reconstructed volume in (x, y, z) order, in the returned motion gauge.
• pose_params (jax.Array): Per-view corrections with columns (alpha, beta, phi, dx, dz, dy). Angles are radians, translations physical lengths in the configured frame.
• info (AlignMultiresInfo): Loss history, work, stopping, geometry and gauge diagnostics.

Raises

• ValueError: If shapes, solver settings or requested degrees of freedom are incompatible.

Examples

import jax.numpy as jnp
import numpy as np
import tomojax as tj
from tomojax.alignment import AlignConfig, align_multires
from tomojax.recon import project
grid = tj.Grid(4, 4, 2, 1.0, 1.0, 1.0)
detector = tj.Detector(4, 2, 1.0, 1.0)
geometry = tj.ParallelGeometry(grid, detector, np.linspace(0, 180, 8, endpoint=False))
data = project(geometry, jnp.ones((4, 4, 2)))
config = AlignConfig(quality="reference", outer_iterations=1, iterations=1)
volume, poses, info = align_multires(
    geometry,
    grid,
    detector,
    data,
    factors=(1,),
    config=config,
)
volume.shape
# (4, 4, 2)

alignment_planCopy Anchor Link

alignment_plan(
    mode: str,
    grid: Grid,
    *,
    quality: str = 'fast',
    geometry: Geometry | None = None,
    levels: Iterable[int] | None = None,
    freeze: Iterable[str] = (),
    config: AlignConfig | None = None,
) -> AlignmentPlan

Return the configuration and levels that run mode at quality.

config supplies expert settings and is used as given; without it the quality profile's defaults apply. The mode sets the schedule unless config names one (or optimise_dofs). freeze keeps the named parameters (for example "dy") at their initial values. When geometry is a cone beam, fast coupled plans limit each inexact linear update to ten CG steps; nonlinear acceptance and outer convergence still decide progress. Without geometry context, budgets are conservative.

Parameters

• mode (str): Alignment mode: "pose", "cor", "cor-then-pose" or "full".
• grid (Grid): Full-resolution volume grid used to choose automatic levels.
• quality ({"fast", "reference"}, optional): Solver quality policy. Default is "fast".
• geometry (Geometry or None, optional): Scan context for geometry-specific iteration budgets. None uses conservative budgets without beam-specific information.
• levels (iterable of int or None, optional): Explicit coarse-to-fine factors. None chooses them from the mode and grid.
• freeze (iterable of str, optional): Parameter names to hold fixed, in addition to config.freeze.
• config (AlignConfig or None, optional): Expert settings. None chooses the quality policy's defaults.

Returns

• AlignmentPlan: Configuration, resolution levels and pose-solver choice for the run.

Raises

• ValueError: If mode, quality or requested configuration settings are invalid.

Examples

from tomojax import Grid
from tomojax.alignment import alignment_plan
plan = alignment_plan("pose", Grid(16, 16, 16, 1, 1, 1), quality="reference")
plan.levels
# (1,)