.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/03_algorithms/06_run_osem_memmap.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_03_algorithms_06_run_osem_memmap.py: 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 :mod:`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)** .. code-block:: text 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 :mod:`parallelproj.data`): * :class:`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. * :func:`parallelproj.data.to_subset_mmap` -- one-time conversion from a full in-memory sinogram to the on-disk subset-contiguous format. .. GENERATED FROM PYTHON SOURCE LINES 51-76 .. code-block:: Python 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, ) .. GENERATED FROM PYTHON SOURCE LINES 77-83 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. .. GENERATED FROM PYTHON SOURCE LINES 83-88 .. code-block:: Python import array_api_compat.numpy as xp dev = "cpu" .. GENERATED FROM PYTHON SOURCE LINES 89-91 Scanner and projector setup --------------------------- .. GENERATED FROM PYTHON SOURCE LINES 91-125 .. code-block:: Python 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 ) .. GENERATED FROM PYTHON SOURCE LINES 126-128 Attenuation and full forward model ----------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 128-150 .. code-block:: Python 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)) .. GENERATED FROM PYTHON SOURCE LINES 151-153 Simulate TOF PET data ---------------------- .. GENERATED FROM PYTHON SOURCE LINES 153-174 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 175-177 Subset view / slice definitions -------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 177-186 .. code-block:: Python 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 ) .. GENERATED FROM PYTHON SOURCE LINES 187-194 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 :class:`parallelproj.data.SubsetArrayMmap` is a thin read-only wrapper; no sinogram data lives in RAM after the ``del`` calls. .. GENERATED FROM PYTHON SOURCE LINES 194-221 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none 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 .. GENERATED FROM PYTHON SOURCE LINES 222-224 Subset forward operators and sensitivity images ------------------------------------------------ .. GENERATED FROM PYTHON SOURCE LINES 224-263 .. code-block:: Python 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) ) .. GENERATED FROM PYTHON SOURCE LINES 264-274 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 :math:`(A^T 1)_i = 0`. Dividing by zero in the EM preconditioner would produce NaN / Inf values that corrupt the reconstruction. :meth:`.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). .. GENERATED FROM PYTHON SOURCE LINES 274-279 .. code-block:: Python cyl_mask = proj.fov_mask() fov_mask = None if bool(xp.all(cyl_mask)) else cyl_mask del cyl_mask .. GENERATED FROM PYTHON SOURCE LINES 280-282 EM update ---------- .. GENERATED FROM PYTHON SOURCE LINES 282-324 .. code-block:: Python 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) .. GENERATED FROM PYTHON SOURCE LINES 325-345 Full objective evaluated subset-by-subset ------------------------------------------ The negative Poisson log-likelihood is separable: :math:`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 :class:`.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. .. GENERATED FROM PYTHON SOURCE LINES 345-362 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 363-365 Warm-start: one OSEM epoch before tracking convergence ------------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 365-379 .. code-block:: Python 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 .. rst-class:: sphx-glr-script-out .. code-block:: none /home/docs/checkouts/readthedocs.org/user_builds/parallelproj/checkouts/stable/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)) .. GENERATED FROM PYTHON SOURCE LINES 380-386 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. .. GENERATED FROM PYTHON SOURCE LINES 386-412 .. code-block:: Python 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() .. rst-class:: sphx-glr-script-out .. code-block:: none 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 .. GENERATED FROM PYTHON SOURCE LINES 413-415 Remove temporary files ---------------------- .. GENERATED FROM PYTHON SOURCE LINES 415-418 .. code-block:: Python shutil.rmtree(tmpdir) .. GENERATED FROM PYTHON SOURCE LINES 419-421 Convergence ----------- .. GENERATED FROM PYTHON SOURCE LINES 421-430 .. code-block:: Python 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() .. image-sg:: /auto_examples/03_algorithms/images/sphx_glr_06_run_osem_memmap_001.png :alt: OSEM (24 subsets, disk-backed sinograms) :srcset: /auto_examples/03_algorithms/images/sphx_glr_06_run_osem_memmap_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 431-433 Reconstructed image -------------------- .. GENERATED FROM PYTHON SOURCE LINES 433-442 .. code-block:: Python 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() .. image-sg:: /auto_examples/03_algorithms/images/sphx_glr_06_run_osem_memmap_002.png :alt: OSEM 5 epochs (disk-backed) :srcset: /auto_examples/03_algorithms/images/sphx_glr_06_run_osem_memmap_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (1 minutes 6.369 seconds) .. _sphx_glr_download_auto_examples_03_algorithms_06_run_osem_memmap.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 06_run_osem_memmap.ipynb <06_run_osem_memmap.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 06_run_osem_memmap.py <06_run_osem_memmap.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 06_run_osem_memmap.zip <06_run_osem_memmap.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_