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.

ParametersCopy Anchor Link

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.

RaisesCopy Anchor Link

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

ExamplesCopy Anchor Link

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.

ParametersCopy Anchor Link

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.

ParametersCopy Anchor Link

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.

ParametersCopy Anchor Link

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.

ReturnsCopy Anchor Link

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.

RaisesCopy Anchor Link

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

ExamplesCopy Anchor Link

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.

ParametersCopy Anchor Link

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.

ReturnsCopy Anchor Link

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

RaisesCopy Anchor Link

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

ExamplesCopy Anchor Link

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,)