TomoJAX
Parallel-beam, laminography and cone-beam CT reconstruction and alignment in JAX.
Reconstruct parallel-beam tomography, laminography and lab cone-beam CT with JAX. TomoJAX
provides differentiable projectors, fast FBP and iterative reconstruction,
joint volume-and-pose alignment, and a CLI for taking NXtomo, APS Data
Exchange, Nikon .xtekct or TIFF data through correction, reconstruction, and
slice export.
TomoJAX is an early research library. Its strongest fit is scientists who need
control over geometry and differentiation in a Python workflow. Synchrotron
geometries use parallel rays; lab scans use a point source and flat detector
(ConeGeometry), reconstructed with FDK or the iterative solvers. Alignment
needs scan-specific validation. See the
support matrix and known limitations
before choosing it for an experiment.
64³ voxels, 90 parallel views, 40 CGLS iterations. This example uses the same Joseph model for simulation and reconstruction; it demonstrates the API, not independent reconstruction accuracy. Reproduce the figure.
tomojax align and tomojax.alignment solve the volume and each view's pose
together. Here a 96³ laminography scan with ±1° and ±2 px of random per-view
motion goes from 0.68 to 0.047 relative error, with rotations recovered to
0.0022°, in 18 s on a laptop GPU including compilation. The measurements are
analytic integrals of continuous objects. Reproduce it.
First reconstruction
On Linux, install uv. For a published release, install the CPU wheel into a fresh environment:
uv venv --python 3.12
uv pip install --python .venv/bin/python "tomojax[cpu]"
.venv/bin/tomojax simulate -o synthetic.nxs --size 32 --n-views 60
.venv/bin/tomojax recon synthetic.nxs -o recon.nxs
.venv/bin/tomojax inspect recon.nxs --preview previewsNo scan data or GPU is required. To try unreleased changes, first build and install a wheel instead of installing from the package index, then run the same three commands.
inspect describes the file, checks it and ends with Valid: yes. Open
previews/slice_x.png, slice_y.png, and slice_z.png, the volume's central
slices (projection.png is the central projection). recon.nxs contains the
projections, volume, and geometry metadata. The PNGs are display-scaled
previews; use the stored volume for quantitative analysis. This CLI example
automatically selects Fourier inversion for its axial parallel geometry; the
figure above uses the Python CGLS example.
Every command reads an input and writes -o OUTPUT, refusing to replace an
existing file without --force; tomojax <command> --help describes it, and
the CLI overview lists the commands.
For CUDA installation, wheel installation, and device checks, see installation. For your own scan, start with the quickstart.
Choose a workflow
| Task | Start here |
|---|---|
Load and correct an NXtomo, APS Data Exchange, Nikon .xtekct or TIFF scan, then reconstruct it | Quickstart |
| Supply measured geometry for parallel or laminography data, including TIFF stacks | Real scan guide |
| Import, calibrate and reconstruct a lab cone-beam CT scan, including multi-orbit scans | Lab cone-beam CT |
| Simulate data or use the Python API | Synthetic workflow, runnable examples |
| Correct per-view motion or the centre of rotation | Alignment guide (experimental; check the recovered geometry) |
| Reconstruct scans larger than the GPU, or use several GPUs | Large scans and several GPUs |
| Check supported models and limitations | Support matrix, known limitations |
Python
In Python, a Scan holds projections and their geometry, and every
operation returns a result that carries its geometry:
import tomojax as tj
scan = tj.load("scan.nxs") # or a Nikon .xtekct
recon = tj.reconstruct(scan) # choose from the geometry
recon = tj.reconstruct(scan, "cgls", iterations=50)
result = tj.align(scan, mode="cor-then-pose") # result.scan carries the corrections
tj.save("recon.nxs", tj.reconstruct(result.scan))The command-line options use the same names: the CGLS line is
tomojax recon scan.nxs -o recon.nxs --method cgls --iterations 50, and the
alignment tomojax align scan.nxs -o aligned.nxs --mode cor-then-pose.
Expert settings are a configuration of the method's own class from
tomojax.recon,
tj.reconstruct(scan, "fista", config=FistaConfig(regulariser="huber_tv")),
or the same fields in a TOML file given to --config. A long alignment can
save its progress and resume after an interruption:
tj.align(scan, checkpoint="align.ckpt"), or --checkpoint align.ckpt.
tj.Scan(projections, geometry) builds a scan from arrays, and
tj.project(geometry, volume) simulates one for any geometry. Volumes use
(x, y, z) and projections (view, v, u); detector and voxel spacings must
use the same physical length unit. The building blocks behind these are in
tomojax.recon, tomojax.alignment, tomojax.corrections,
tomojax.geometry, tomojax.io, tomojax.forward and tomojax.datasets.
Start with the complete example or
browse the documentation index.
Evidence
Measured speed, accuracy and memory against ASTRA and TIGRE, failed cases included, are in the measurements. Volumes larger than the GPU stream through it in slabs: SIRT of a whole Diamond I12 scan, 1801 views of 2160 × 2560 pixels into a 56.6 GB volume, ran on an 8 GB laptop GPU with 1.4 GB of GPU memory at 1.13 times the per-voxel time of a scan that fits (large scans).
Development and license
See CONTRIBUTING.md for environment setup, engineering conventions, tests, and package checks. Changes on this branch are recorded in the changelog.
TomoJAX is licensed under GPL-3.0-only.