Alignment solver reference

Use this reference when changing AlignConfig fields, in Python or as keys of a tomojax align --config TOML file. For data preparation, mode selection, and interpreting recovered parameters, start with the alignment guide. These options are experimental; the free-voxel pilot still has a failed cell.

tj.align and tomojax align in pose and cor-then-pose modes start from coupled_pose_config(): the joint coupling with gn_joint_solver="pose_eliminated", Joseph plane sampling (ray_integrator="joseph") and fp32 gathers. cor and full keep the fixed-volume coupling, also with Joseph sampling. --dry-run prints the full set for a mode. A bare AlignConfig() keeps the older alternating defaults, which the statements of defaults below describe.

The opt-in spectral_setup schedule (or a setup stage with objective_kind="projection_consistency", optimizer="spectral") initializes parallel-beam axis direction and detector-u centre from Fourier common lines, preserving bounds, frozen parameters and the supplied detector roll. It leaves the volume unchanged: compose it before a pose stage or reconstruct its returned scan. Its diagnostics say unvalidated_seed, not a successful validation fit; the default full schedule still validates and reconstructs its setup updates. For calibrated, unrolled laminography, a separate projection_consistency / spectral setup stage with ("dx", "dz") initializes smooth detector-frame motion from zero poses, also without reconstructing or claiming validation.

Gauss–Newton updates at interpolation boundariesCopy Anchor Link

The default gn_jacobian="autodiff" differentiates the ray model within its current interpolation cells. At a voxel-grid boundary this chooses a one-sided derivative. For sharp objects, an update can enter a wrong local minimum even for a small detector shift.

Set gn_jacobian="central" to form a symmetric finite-difference Jacobian for the Gauss–Newton step. The forward prediction, acceptance loss, projector AD, and matched volume adjoint stay unchanged. This option costs ten projection evaluations per view for the five Jacobian columns, plus the residual and candidate checks; it is a recovery option, not a derivative-speed claim. gn_difference_step=1e-3 is a fraction of the smallest physical voxel pitch. Rotation increments use that displacement divided by the grid's bounding sphere radius, so the stencil scales with physical size and multiresolution level. Both settings are saved in checkpoint configuration. This local stencil does not guarantee recovery from large initial motion or poor identifiability.

Small rigid-transform products and Gauss–Newton normal equations explicitly request full FP32 multiplication precision. This prevents reduced global JAX matrix-multiply settings from quantising small pose updates on CUDA.

The gn_coupling="joint" option, which tj.align uses in pose and cor-then-pose modes, is for per-view least-squares alignment. It solves a damped linearised volume-and-pose problem after each reconstruction refresh, then accepts the updated pair only after scoring its constrained nonlinear objective. This addresses slow alternation when free voxels can compensate for pose errors. AlignConfig() keeps gn_coupling="fixed_volume". In align_multires, the implicit schedule uses a joint_volume_pose stage when this option is enabled. Explicitly fixed-volume schedule stages continue to keep the volume fixed; a custom joint stage must declare objective_kind="joint_volume_pose" and optimizer="gn".

Joint GN requires opt_method="gn", pose_model="per_view" and gather_dtype="fp32". It supports sampled and exact integration, active/frozen pose parameters, pose bounds, volume masks, nonnegative voxels, pose smoothness, and Huber-TV or zero volume regularisation. It does not support nonsmooth TV or polynomial/spline pose models. gn_joint_iterations=40 bounds the inner matrix-free PCG solve; gn_joint_rtol=1e-4 controls its residual stopping criterion. gn_damping and gn_volume_damping damp pose and voxel increments respectively; they are not priors on the recovered object. The returned joint_linear_relative_residual is recomputed explicitly, so reaching the iteration limit does not imply an accurate linear solve. The optional pose column cache on the device is capped at 64 MiB. Host-staged joint solves can retain the active FP32 columns in host memory for one linear update, reading only the current view batch on the device. Planning reserves eight image vectors (twelve with spectral preconditioning) and half the remaining available host memory; without that space, columns are recomputed. Each update rebuilds the cache at its current image and poses, including its weights and residual filters.

Repeated joint calls reuse compiled objectives when array shapes and solver options match. Measurements, nominal poses, masks, detector coordinates and weights remain inputs to each call, so a new scan can reuse a program without reusing the previous scan's values. Changing shapes or program options still requires compilation; this reuse is separate from the optional JAX persistent cache for different processes. On the six-cell free-voxel pilot, reuse made warm calls faster and left cold start-up and process GPU memory unchanged.

The gn_joint_solver="pose_eliminated" method factors the pose block and runs PCG on the volume Schur system, then back-substitutes poses. It solves the same damped linear problem as the default "stacked" method, with the same joint residual threshold and nonlinear acceptance. Without pose smoothness, the factors are independent 5-by-5 blocks; with smoothness, a block-banded factorisation retains inter-view coupling with storage linear in view count. Damping and iteration budgets are unchanged. Independent dense, constrained workflow and CUDA checks pass. On the free-voxel pilot it recovers the same five of six cells as the stacked solve, faster on the tilted cells, and sampled GPU memory rises from 280 to 320 MiB. It is the default for pose and cor-then-pose alignment (AlignConfig() keeps "stacked"); it does not establish complete recovery.

Joint steps report objective_kind="joint_volume_pose", their actual projector backend, accepted line-search scale, and linear-solve diagnostics. The shared central-difference stencil is unchanged. Recovery and performance claims require the complete free-voxel pilot; passing the numerical tests alone does not establish them.

Alignment's info["lipschitz"] and checkpoint lipschitz retain the effective FISTA step bound, including smooth-TV curvature. Reusing a bound does not add a safety factor on each outer iteration. If reconstruction falls back to public FISTA, the adapter removes the existing TV contribution before that solver adds it; the bound therefore stays consistent across both paths and resume. A fallback override too small to supply a positive data bound triggers operator-norm estimation. The public FistaConfig.lipschitz itself remains a data-term bound.

Without an override, Huber-FISTA alignment initialises its data bound from max(A.T @ (A @ ones)), with a 20% margin applied once. For the nonnegative trilinear ray model this row sum bounds the normal operator's largest eigenvalue. It costs one projection/adjoint pair and accounts for actual physical sampling, including detector density and oblique rays. The old dimension-only estimate did not scale with physical length units.