Note

Go to the end to download the full example code.

RAM-efficient OSEM with disk-backed TOF sinograms¶

Demonstrates a memory-efficient OSEM variant in which full TOF sinograms are never held in RAM simultaneously. Instead they are stored on disk in a subset-contiguous binary layout via numpy.memmap, and only one subset’s data is loaded at a time during each OSEM update.

Why subset-contiguous layout?

The natural sinogram axis order (e.g. (num_rad, num_views, num_planes, num_tofbins) for RVP) stores views contiguously. OSEM subsets are distributed across the view axis with stride num_subsets, so reading subset k from this layout requires many non-sequential disk seeks.

Re-organising the data to shape (num_subsets, num_rad, views_per_subset, num_planes, num_tofbins) on disk makes each subset a single contiguous block. Sequential access lets the OS read-ahead prefetch the next subset from disk while the projector runs on the current one.

Memory comparison (this toy scanner)

Before disk conversion (full sinograms in RAM):
  y:                ~15 MB
  contamination:    ~15 MB
  total:            ~30 MB

After disk conversion:
  per subset in RAM during update:
    y_k:    ~0.65 MB
    s_k:    ~0.65 MB
    total:  ~1.30 MB  (vs ~30 MB -- a 23x reduction)

On a clinical scanner (400 rad x 400 views x 837 planes x 27 TOF bins) the full sinogram is ~3.5 GB each, and the saving scales accordingly.

Helper code (in parallelproj.data):

parallelproj.data.SubsetArrayMmap – read-only wrapper around a numpy.memmap file. mmap[k] returns an owned copy of subset k that Python frees as soon as the caller deletes the reference.
parallelproj.data.to_subset_mmap() – one-time conversion from a full in-memory sinogram to the on-disk subset-contiguous format.

from __future__ import annotations

import shutil
import tempfile
from copy import copy
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

import parallelproj.operators
import parallelproj.pet_lors
import parallelproj.pet_scanners
import parallelproj.projectors
import parallelproj.tof
from parallelproj import Array, to_numpy_array
from parallelproj.functions import C1Function, C2AffineObjective, NegPoissonLogL
from parallelproj.data import to_subset_mmap


from parallelproj._examples_utils import (
    elliptic_cylinder_phantom,
    show_vol_cuts,
)

Backend and device¶

In principle we can run on any backend or device However, since we will use (numpy-based) memmaps, it is best to use the CPU as device.

import array_api_compat.numpy as xp

dev = "cpu"

Scanner and projector setup¶

num_subsets = 24
num_epochs = 5

num_rings = 5
scanner = parallelproj.pet_scanners.RegularPolygonPETScannerGeometry(
    xp,
    dev,
    radius=65.0,
    num_sides=16,
    num_lor_endpoints_per_side=12,
    lor_spacing=2.3,
    ring_positions=xp.linspace(-10, 10, num_rings, device=dev),
    symmetry_axis=2,
)

img_shape = (55, 55, 8)
voxel_size = (2.0, 2.0, 2.0)

lor_desc = parallelproj.pet_lors.RegularPolygonPETLORDescriptor(
    scanner,
    parallelproj.pet_lors.Michelogram(scanner.num_rings, max_ring_difference=2, span=1),
    radial_trim=10,
    sinogram_order=parallelproj.pet_lors.SinogramSpatialAxisOrder.RVP,
)

proj = parallelproj.projectors.RegularPolygonPETProjector(
    lor_desc, img_shape=img_shape, voxel_size=voxel_size
)

x_true = elliptic_cylinder_phantom(
    xp, dev, image_shape=img_shape, voxel_size=voxel_size
)

Attenuation and full forward model¶

x_att = 0.01 * xp.astype(x_true > 0, xp.float32)
att_sino = xp.exp(-proj(x_att))

proj.tof_parameters = parallelproj.tof.TOFParameters(
    num_tofbins=13, tofbin_width=12.0, sigma_tof=12.0
)

att_values = (
    xp.broadcast_to(xp.expand_dims(att_sino, axis=-1), proj.out_shape)
    if proj.tof
    else att_sino
)
att_op = parallelproj.operators.ElementwiseMultiplicationOperator(att_values)

res_model = parallelproj.operators.GaussianFilterOperator(
    proj.in_shape,
    sigma=[2.0 / (2.35 * float(vs)) for vs in proj.voxel_size],
)

pet_lin_op = parallelproj.operators.CompositeLinearOperator((att_op, proj, res_model))

Simulate TOF PET data¶

noise_free_data = pet_lin_op(x_true)

contamination = xp.full(
    noise_free_data.shape,
    0.5 * float(xp.mean(noise_free_data)),
    device=dev,
    dtype=xp.float32,
)

noise_free_data += contamination

np.random.seed(1)
y = xp.asarray(
    np.random.poisson(to_numpy_array(noise_free_data)),
    device=dev,
    dtype=xp.float32,
)

del noise_free_data

Subset view / slice definitions¶

subset_views, subset_slices = proj.lor_descriptor.get_distributed_views_and_slices(
    num_subsets, len(proj.out_shape)
)

_, subset_slices_non_tof = proj.lor_descriptor.get_distributed_views_and_slices(
    num_subsets, 3
)

Convert full sinograms to disk-backed subset files and free RAM¶

to_subset_mmap gathers the non-contiguous subset slices from the full array and writes them as contiguous blocks on disk. The returned parallelproj.data.SubsetArrayMmap is a thin read-only wrapper; no sinogram data lives in RAM after the del calls.

tmpdir = Path(tempfile.mkdtemp())

y_np = to_numpy_array(y)
s_np = to_numpy_array(contamination)

print("Full sinograms in RAM before conversion:")
print(f"  y:             {y_np.nbytes / 1024**2:.1f} MB")
print(f"  contamination: {s_np.nbytes / 1024**2:.1f} MB")
print(f"  total:         {(y_np.nbytes + s_np.nbytes) / 1024**2:.1f} MB")

y_mmap = to_subset_mmap(y_np, subset_slices, tmpdir / "y.bin")
s_mmap = to_subset_mmap(s_np, subset_slices, tmpdir / "s.bin")

# the strictly positive contamination guarantees positive expected data in
# every bin, so the exact (unmodified) log-likelihood of NegPoissonLogL can
# be used -- evaluate before the full contamination array is freed
exact_mode = bool(xp.min(contamination) > 0)

del y_np, s_np, y, contamination  # full arrays no longer in RAM

print("\nOn disk, per-subset in RAM on demand (OS-managed):")
print(f"  y_k per subset:   {y_mmap.nbytes_per_subset() / 1024**2:.2f} MB")
print(f"  s_k per subset:   {s_mmap.nbytes_per_subset() / 1024**2:.2f} MB")
peak_mb = (y_mmap.nbytes_per_subset() + s_mmap.nbytes_per_subset()) / 1024**2
print(f"  peak per update:  {peak_mb:.2f} MB")

Full sinograms in RAM before conversion:
  y:             15.5 MB
  contamination: 15.5 MB
  total:         30.9 MB

On disk, per-subset in RAM on demand (OS-managed):
  y_k per subset:   0.64 MB
  s_k per subset:   0.64 MB
  peak per update:  1.29 MB

Subset forward operators and sensitivity images¶

proj.clear_cached_lor_endpoints()

pet_subset_linop_seq = []

for i in range(num_subsets):
    subset_proj = copy(proj)
    subset_proj.views = subset_views[i]

    att_values_k = (
        xp.broadcast_to(
            xp.expand_dims(att_sino[subset_slices_non_tof[i]], axis=-1),
            subset_proj.out_shape,
        )
        if subset_proj.tof
        else att_sino[subset_slices_non_tof[i]]
    )
    subset_att_op = parallelproj.operators.ElementwiseMultiplicationOperator(
        att_values_k
    )

    pet_subset_linop_seq.append(
        parallelproj.operators.CompositeLinearOperator(
            [subset_att_op, subset_proj, res_model]
        )
    )

pet_subset_linop_seq = parallelproj.operators.LinearOperatorSequence(
    pet_subset_linop_seq
)

subset_adjoint_ones = xp.zeros(
    (num_subsets,) + pet_lin_op.in_shape, dtype=xp.float32, device=dev
)
for k, op in enumerate(pet_subset_linop_seq):
    subset_adjoint_ones[k] = op.adjoint(
        xp.ones(op.out_shape, dtype=xp.float32, device=dev)
    )

FOV mask¶

The scanner’s cylindrical field of view does not cover every voxel of the image grid. Voxels outside the FOV are never intersected by any LOR, so their sensitivity \((A^T 1)_i = 0\). Dividing by zero in the EM preconditioner would produce NaN / Inf values that corrupt the reconstruction. RegularPolygonPETProjector.fov_mask() returns a boolean array that is True inside the FOV. fov_mask is set to None when every image voxel is inside the FOV (no masking needed).

cyl_mask = proj.fov_mask()
fov_mask = None if bool(xp.all(cyl_mask)) else cyl_mask
del cyl_mask

EM update¶

def em_update(
    x_cur: Array,
    data_fidelity: C1Function,
    adj_ones: Array,
    img_mask: Array | None = None,
) -> Array:
    """One EM update rewritten as a preconditioned gradient descent step.

    Computes :math:`x^+ = x - D \\nabla f(x)` where the diagonal
    preconditioner is :math:`D = \\operatorname{diag}(x / (A^T 1))`.
    Voxels outside the FOV are excluded via ``img_mask`` to avoid
    division by the zero sensitivity values in ``adj_ones``.

    Parameters
    ----------
    x_cur : Array
        Current image estimate.
    data_fidelity : C1Function
        Differentiable data-fidelity term whose gradient is evaluated at
        ``x_cur``.
    adj_ones : Array
        Sensitivity image :math:`A^T 1` (or subset variant
        :math:`(A^k)^T 1`).
    img_mask : Array or None, optional
        Boolean FOV mask (``True`` inside the FOV).  Preconditioner is
        zeroed outside the FOV so that zero-sensitivity voxels do not
        produce NaN / Inf.  Pass ``None`` when every voxel is in the FOV.

    Returns
    -------
    Array
        Updated image :math:`x^+`, same shape as ``x_cur``.
    """
    if img_mask is None:
        d = x_cur / adj_ones
    else:
        d = xp.where(img_mask, x_cur / adj_ones, xp.zeros_like(x_cur))
    return x_cur - d * data_fidelity.gradient(x_cur)

Full objective evaluated subset-by-subset¶

The negative Poisson log-likelihood is separable: \(f(x) = \sum_k f_k(x)\). We accumulate the value over subsets, loading only one subset’s data at a time.

Note

By default NegPoissonLogL evaluates a “safe epsilon” (shifted Poisson) surrogate: a tiny eps = rel_eps * mean(y) is added to the measured and the expected data. This is finite for any non-negative expectation (never nan / inf), at the price of a tiny (~``rel_eps``) bias that vanishes at the fit. Since our contamination is strictly positive, the expected data A x + s are positive in every bin and we can use exact=True (exact_mode was derived from the contamination above). Keep the default whenever the expected data can reach zero in bins with counts. Note that each per-subset instance would derive its own eps from its subset mean – pass one global eps explicitly if exact separability of the subset objectives matters.

def full_objective_from_subsets(x: Array) -> float:
    """Compute f(x) by accumulating over subsets from disk."""
    total = 0.0
    for subset_idx in range(num_subsets):
        y_sub = xp.asarray(y_mmap[subset_idx], device=dev, dtype=xp.float32)
        s_sub = xp.asarray(s_mmap[subset_idx], device=dev, dtype=xp.float32)
        total += C2AffineObjective(
            NegPoissonLogL(y_sub, exact=exact_mode),
            pet_subset_linop_seq[subset_idx],
            s_sub,
        )(x)
        del y_sub, s_sub
    return total

Warm-start: one OSEM epoch before tracking convergence¶

x_osem = xp.ones(pet_lin_op.in_shape, dtype=xp.float32, device=dev)
if fov_mask is not None:
    x_osem = xp.where(fov_mask, x_osem, xp.zeros_like(x_osem))

for k in range(num_subsets):
    y_k = xp.asarray(y_mmap[k], device=dev, dtype=xp.float32)
    s_k = xp.asarray(s_mmap[k], device=dev, dtype=xp.float32)
    df_k = C2AffineObjective(
        NegPoissonLogL(y_k, exact=exact_mode), pet_subset_linop_seq[k], s_k
    )
    x_osem = em_update(x_osem, df_k, subset_adjoint_ones[k], fov_mask)
    del df_k, y_k, s_k

/home/docs/checkouts/readthedocs.org/user_builds/parallelproj/checkouts/v2.0.0/docs/examples/03_algorithms/06_run_osem_memmap.py:320: RuntimeWarning: invalid value encountered in divide
  d = xp.where(img_mask, x_cur / adj_ones, xp.zeros_like(x_cur))

OSEM reconstruction¶

Only y_k and s_k (one subset each) reside in RAM during a subset update. They are loaded from the memory-mapped files and freed by del after each update, keeping the sinogram footprint minimal.

df_osem = xp.zeros(num_epochs, dtype=xp.float32, device=dev)

for i in range(num_epochs):
    for k in range(num_subsets):
        print(
            f"OSEM epoch {i + 1:03}/{num_epochs:03},"
            f" subset {k + 1:03}/{num_subsets:03}",
            end="\r",
        )
        # --- load subset k from disk (one sequential read) ---
        y_k = xp.asarray(y_mmap[k], device=dev, dtype=xp.float32)
        s_k = xp.asarray(s_mmap[k], device=dev, dtype=xp.float32)

        df_k = C2AffineObjective(
            NegPoissonLogL(y_k, exact=exact_mode), pet_subset_linop_seq[k], s_k
        )
        x_osem = em_update(x_osem, df_k, subset_adjoint_ones[k], fov_mask)

        # --- release subset data from RAM / GPU VRAM ---
        del df_k, y_k, s_k

    df_osem[i] = full_objective_from_subsets(x_osem)

print()

OSEM epoch 001/005, subset 001/024
OSEM epoch 001/005, subset 002/024
OSEM epoch 001/005, subset 003/024
OSEM epoch 001/005, subset 004/024
OSEM epoch 001/005, subset 005/024
OSEM epoch 001/005, subset 006/024
OSEM epoch 001/005, subset 007/024
OSEM epoch 001/005, subset 008/024
OSEM epoch 001/005, subset 009/024
OSEM epoch 001/005, subset 010/024
OSEM epoch 001/005, subset 011/024
OSEM epoch 001/005, subset 012/024
OSEM epoch 001/005, subset 013/024
OSEM epoch 001/005, subset 014/024
OSEM epoch 001/005, subset 015/024
OSEM epoch 001/005, subset 016/024
OSEM epoch 001/005, subset 017/024
OSEM epoch 001/005, subset 018/024
OSEM epoch 001/005, subset 019/024
OSEM epoch 001/005, subset 020/024
OSEM epoch 001/005, subset 021/024
OSEM epoch 001/005, subset 022/024
OSEM epoch 001/005, subset 023/024
OSEM epoch 001/005, subset 024/024
OSEM epoch 002/005, subset 001/024
OSEM epoch 002/005, subset 002/024
OSEM epoch 002/005, subset 003/024
OSEM epoch 002/005, subset 004/024
OSEM epoch 002/005, subset 005/024
OSEM epoch 002/005, subset 006/024
OSEM epoch 002/005, subset 007/024
OSEM epoch 002/005, subset 008/024
OSEM epoch 002/005, subset 009/024
OSEM epoch 002/005, subset 010/024
OSEM epoch 002/005, subset 011/024
OSEM epoch 002/005, subset 012/024
OSEM epoch 002/005, subset 013/024
OSEM epoch 002/005, subset 014/024
OSEM epoch 002/005, subset 015/024
OSEM epoch 002/005, subset 016/024
OSEM epoch 002/005, subset 017/024
OSEM epoch 002/005, subset 018/024
OSEM epoch 002/005, subset 019/024
OSEM epoch 002/005, subset 020/024
OSEM epoch 002/005, subset 021/024
OSEM epoch 002/005, subset 022/024
OSEM epoch 002/005, subset 023/024
OSEM epoch 002/005, subset 024/024
OSEM epoch 003/005, subset 001/024
OSEM epoch 003/005, subset 002/024
OSEM epoch 003/005, subset 003/024
OSEM epoch 003/005, subset 004/024
OSEM epoch 003/005, subset 005/024
OSEM epoch 003/005, subset 006/024
OSEM epoch 003/005, subset 007/024
OSEM epoch 003/005, subset 008/024
OSEM epoch 003/005, subset 009/024
OSEM epoch 003/005, subset 010/024
OSEM epoch 003/005, subset 011/024
OSEM epoch 003/005, subset 012/024
OSEM epoch 003/005, subset 013/024
OSEM epoch 003/005, subset 014/024
OSEM epoch 003/005, subset 015/024
OSEM epoch 003/005, subset 016/024
OSEM epoch 003/005, subset 017/024
OSEM epoch 003/005, subset 018/024
OSEM epoch 003/005, subset 019/024
OSEM epoch 003/005, subset 020/024
OSEM epoch 003/005, subset 021/024
OSEM epoch 003/005, subset 022/024
OSEM epoch 003/005, subset 023/024
OSEM epoch 003/005, subset 024/024
OSEM epoch 004/005, subset 001/024
OSEM epoch 004/005, subset 002/024
OSEM epoch 004/005, subset 003/024
OSEM epoch 004/005, subset 004/024
OSEM epoch 004/005, subset 005/024
OSEM epoch 004/005, subset 006/024
OSEM epoch 004/005, subset 007/024
OSEM epoch 004/005, subset 008/024
OSEM epoch 004/005, subset 009/024
OSEM epoch 004/005, subset 010/024
OSEM epoch 004/005, subset 011/024
OSEM epoch 004/005, subset 012/024
OSEM epoch 004/005, subset 013/024
OSEM epoch 004/005, subset 014/024
OSEM epoch 004/005, subset 015/024
OSEM epoch 004/005, subset 016/024
OSEM epoch 004/005, subset 017/024
OSEM epoch 004/005, subset 018/024
OSEM epoch 004/005, subset 019/024
OSEM epoch 004/005, subset 020/024
OSEM epoch 004/005, subset 021/024
OSEM epoch 004/005, subset 022/024
OSEM epoch 004/005, subset 023/024
OSEM epoch 004/005, subset 024/024
OSEM epoch 005/005, subset 001/024
OSEM epoch 005/005, subset 002/024
OSEM epoch 005/005, subset 003/024
OSEM epoch 005/005, subset 004/024
OSEM epoch 005/005, subset 005/024
OSEM epoch 005/005, subset 006/024
OSEM epoch 005/005, subset 007/024
OSEM epoch 005/005, subset 008/024
OSEM epoch 005/005, subset 009/024
OSEM epoch 005/005, subset 010/024
OSEM epoch 005/005, subset 011/024
OSEM epoch 005/005, subset 012/024
OSEM epoch 005/005, subset 013/024
OSEM epoch 005/005, subset 014/024
OSEM epoch 005/005, subset 015/024
OSEM epoch 005/005, subset 016/024
OSEM epoch 005/005, subset 017/024
OSEM epoch 005/005, subset 018/024
OSEM epoch 005/005, subset 019/024
OSEM epoch 005/005, subset 020/024
OSEM epoch 005/005, subset 021/024
OSEM epoch 005/005, subset 022/024
OSEM epoch 005/005, subset 023/024
OSEM epoch 005/005, subset 024/024

Remove temporary files¶

shutil.rmtree(tmpdir)

Convergence¶

fig, ax = plt.subplots(figsize=(6, 4), layout="constrained")
ax.plot(np.arange(1, num_epochs + 1), to_numpy_array(df_osem), marker="o")
ax.set_xlabel("Epoch")
ax.set_ylabel("Negative Poisson log-likelihood")
ax.set_title(f"OSEM ({num_subsets} subsets, disk-backed sinograms)")
ax.grid(ls=":")
fig.show()

OSEM (24 subsets, disk-backed sinograms)

Reconstructed image¶

fig, axs, widgets = show_vol_cuts(
    x_osem,
    voxel_size=voxel_size,
    fig_title=f"OSEM {num_epochs} epochs (disk-backed)",
    vmin=0,
    vmax=float(xp.max(x_osem)),
)
fig.show()

Total running time of the script: (1 minutes 6.823 seconds)

Gallery generated by Sphinx-Gallery