.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/03_algorithms/00_run_mlem_osem_svrg.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_00_run_mlem_osem_svrg.py: Convergence comparison: MLEM vs OSEM vs SVRG ============================================= This example compares the convergence speed (per epoch) of three algorithms for minimising the negative Poisson log-likelihood .. math:: f(x) = \sum_i \bar{y}_i - y_i \log \bar{y}_i, \qquad \bar{y}(x) = A x + s subject to :math:`x \geq 0`: * **MLEM** -- expectation-maximisation on the full data; one iteration = one full data pass; guaranteed to converge but slow per epoch. * **OSEM** -- ordered-subsets EM; one epoch = :math:`m` subset updates ~= one full data pass; fast empirical convergence but *no* convergence guarantee. * **SVRG** -- stochastic variance-reduced gradient with subsets; one epoch = :math:`m` variance-reduced subset updates; provably convergent like MLEM while achieving the fast per-epoch progress of OSEM. .. note:: For SVRG, one epoch requires **two** full data passes: one to compute the snapshot gradients at the anchor point, and one for the :math:`m` variance-reduced subset updates. The epoch axis in the convergence plot therefore understates the true computational cost of SVRG relative to OSEM by a factor of roughly two. .. GENERATED FROM PYTHON SOURCE LINES 31-49 .. code-block:: Python from __future__ import annotations from collections.abc import Sequence import matplotlib.pyplot as plt from copy import copy import numpy as np import parallelproj.operators import parallelproj.tof import parallelproj.pet_scanners import parallelproj.pet_lors import parallelproj.projectors from parallelproj import to_numpy_array, Array from parallelproj.functions import NegPoissonLogL, C2AffineObjective, C1Function from parallelproj._examples_utils import show_vol_cuts from parallelproj._examples_utils import elliptic_cylinder_phantom .. GENERATED FROM PYTHON SOURCE LINES 50-56 .. code-block:: Python from parallelproj._examples_utils import suggest_array_backend_and_device # To use a specific backend and/or device, replace the None arguments, e.g.: # xp, dev = suggest_array_backend_and_device(backend="numpy", dev="cpu") or by setting xp and dev manually xp, dev = suggest_array_backend_and_device(None, None) .. rst-class:: sphx-glr-script-out .. code-block:: none Using array API: array_api_compat.torch, device: cpu .. GENERATED FROM PYTHON SOURCE LINES 57-68 .. code-block:: Python # number of subsets for OSEM and SVRG num_subsets = 24 # if run on a CPU limit the number of "OSEM / SVRG epochs" num_epochs_mlem = 120 if dev == "cpu" else 1200 num_epochs = num_epochs_mlem // num_subsets # run MLEM only on GPU - it is too slow on CPU for the number of iterations used here run_mlem = dev != "cpu" .. GENERATED FROM PYTHON SOURCE LINES 69-78 Setup of the forward model :math:`\bar{y}(x) = A x + s` -------------------------------------------------------- We setup a linear forward operator :math:`A` consisting of an image-based resolution model, a non-TOF PET projector and an attenuation model .. note:: The OSEM implementation below works with all linear operators that subclass :class:`.LinearOperator` (e.g. the high-level projectors). .. GENERATED FROM PYTHON SOURCE LINES 78-91 .. code-block:: Python 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, ) .. GENERATED FROM PYTHON SOURCE LINES 92-93 setup the LOR descriptor that defines the sinogram .. GENERATED FROM PYTHON SOURCE LINES 93-113 .. code-block:: Python 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 ) # setup a simple test image containing a few "hot rods" x_true = elliptic_cylinder_phantom( xp, dev, image_shape=img_shape, voxel_size=voxel_size ) .. GENERATED FROM PYTHON SOURCE LINES 114-116 Attenuation image and sinogram setup ------------------------------------ .. GENERATED FROM PYTHON SOURCE LINES 116-122 .. code-block:: Python # setup an attenuation image x_att = 0.01 * xp.astype(x_true > 0, xp.float32) # calculate the attenuation sinogram att_sino = xp.exp(-proj(x_att)) .. GENERATED FROM PYTHON SOURCE LINES 123-129 Complete PET forward model setup -------------------------------- We combine an image-based resolution model, a non-TOF or TOF PET projector and an attenuation model into a single linear operator. .. GENERATED FROM PYTHON SOURCE LINES 129-152 .. code-block:: Python # enable TOF - comment if you want to run non-TOF proj.tof_parameters = parallelproj.tof.TOFParameters( num_tofbins=13, tofbin_width=12.0, sigma_tof=12.0 ) # For TOF, att_sino has no TOF-bins dimension while the projector output does. # broadcast_to adds a trailing singleton via expand_dims and broadcasts it over # the TOF-bins axis without copying data (zero-stride view). 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] ) # compose all 3 operators into a single linear operator pet_lin_op = parallelproj.operators.CompositeLinearOperator((att_op, proj, res_model)) .. GENERATED FROM PYTHON SOURCE LINES 153-158 Simulation of projection data ----------------------------- We setup an arbitrary ground truth :math:`x_{true}` and simulate noise-free and noisy data :math:`y` by adding Poisson noise. .. GENERATED FROM PYTHON SOURCE LINES 158-180 .. code-block:: Python # simulated noise-free data noise_free_data = pet_lin_op(x_true) # generate a constant contamination sinogram contamination = xp.full( noise_free_data.shape, 0.5 * float(xp.mean(noise_free_data)), device=dev, dtype=xp.float32, ) noise_free_data += contamination # add Poisson noise np.random.seed(1) y = xp.asarray( np.random.poisson(to_numpy_array(noise_free_data)), device=dev, dtype=xp.float32, ) .. GENERATED FROM PYTHON SOURCE LINES 181-188 Splitting of the forward model into subsets :math:`A^k` ------------------------------------------------------- Calculate the view numbers and slices for each subset. We will use the subset views to setup a sequence of projectors projecting only a subset of views. The slices can be used to extract the corresponding subsets from full data or corrections sinograms. .. GENERATED FROM PYTHON SOURCE LINES 188-238 .. 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 ) # clear the cached LOR endpoints since we will create many copies of the projector proj.clear_cached_lor_endpoints() pet_subset_linop_seq = [] # we setup a sequence of subset forward operators each constisting of # (1) image-based resolution model # (2) subset projector # (3) multiplication with the corresponding subset of the attenuation sinogram for i in range(num_subsets): # make a copy of the full projector and reset the views to project subset_proj = copy(proj) subset_proj.views = subset_views[i] # same TOF/non-TOF broadcasting as for the full operator above 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 ) # add the resolution model and multiplication with a subset of the attenuation sinogram 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 ) .. GENERATED FROM PYTHON SOURCE LINES 239-255 EM update --------- The EM update used in MLEM and OSEM is :footcite:p:`Dempster1977` :footcite:p:`Shepp1982` :footcite:p:`Lange1984` :footcite:p:`Hudson1994` .. math:: x^+ = \frac{x}{A^T 1} A^T \frac{y}{A x + s} which can be rewritten as a preconditioned gradient descent step with diagonal preconditioner :math:`D = \operatorname{diag}(x / (A^T 1))`: .. math:: x^+ = x - D \, \nabla_x f(x). We implement this as a single function used by both MLEM and OSEM. .. GENERATED FROM PYTHON SOURCE LINES 255-297 .. 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 298-305 Setup of objective functions and sensitivity images --------------------------------------------------- We define one :class:`.C2AffineObjective` per subset (for OSEM and SVRG) and one for the full data (for MLEM and objective evaluation). The sensitivity image :math:`A^T 1` and its per-subset counterparts :math:`(A^k)^T 1` are precomputed once. .. GENERATED FROM PYTHON SOURCE LINES 305-318 .. code-block:: Python # calculate A_k^T 1 for all subsets k, stored as (num_subsets, *in_shape) 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) ) # full sensitivity image A^T 1 = sum of all subset sensitivities adjoint_ones = xp.sum(subset_adjoint_ones, axis=0) .. GENERATED FROM PYTHON SOURCE LINES 319-329 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 329-334 .. 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 335-353 Setup of data-fidelity terms ---------------------------- We define one :class:`.C2AffineObjective` per subset (for OSEM and SVRG) and one for the full data (for MLEM and objective evaluation). .. 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`` is positive in every bin and we can use ``exact=True`` to evaluate the unmodified log-likelihood instead (bins with ``y = 0`` are still handled exactly). Keep the default whenever the expected data can reach zero in bins with counts, e.g. with zero contamination and a mismatched forward model. .. GENERATED FROM PYTHON SOURCE LINES 353-384 .. code-block:: Python # the strictly positive contamination guarantees A x + s > 0 in every bin, # so the exact (unmodified) log-likelihood can be used exact_mode = bool(xp.min(contamination) > 0) subset_data_fidelities = [ C2AffineObjective( NegPoissonLogL(y[sl], exact=exact_mode), pet_subset_linop_seq[k], contamination[sl], ) for k, sl in enumerate(subset_slices) ] full_data_fidelity = C2AffineObjective( NegPoissonLogL(y, exact=exact_mode), pet_lin_op, contamination ) # run 1 OSEM epoch as a common warm-start for MLEM, OSEM and SVRG x_init = xp.ones(pet_lin_op.in_shape, dtype=xp.float32, device=dev) # zero out voxels not seen by every subset so they don't carry a spurious non-zero value if fov_mask is not None: x_init = xp.where(fov_mask, x_init, xp.zeros_like(x_init)) for k in range(len(subset_slices)): print(f"warm-start OSEM subset {(k+1):03} / {num_subsets:03}", end="\r") x_init = em_update( x_init, subset_data_fidelities[k], subset_adjoint_ones[k], fov_mask ) print() .. rst-class:: sphx-glr-script-out .. code-block:: none warm-start OSEM subset 001 / 024 warm-start OSEM subset 002 / 024 warm-start OSEM subset 003 / 024 warm-start OSEM subset 004 / 024 warm-start OSEM subset 005 / 024 warm-start OSEM subset 006 / 024 warm-start OSEM subset 007 / 024 warm-start OSEM subset 008 / 024 warm-start OSEM subset 009 / 024 warm-start OSEM subset 010 / 024 warm-start OSEM subset 011 / 024 warm-start OSEM subset 012 / 024 warm-start OSEM subset 013 / 024 warm-start OSEM subset 014 / 024 warm-start OSEM subset 015 / 024 warm-start OSEM subset 016 / 024 warm-start OSEM subset 017 / 024 warm-start OSEM subset 018 / 024 warm-start OSEM subset 019 / 024 warm-start OSEM subset 020 / 024 warm-start OSEM subset 021 / 024 warm-start OSEM subset 022 / 024 warm-start OSEM subset 023 / 024 warm-start OSEM subset 024 / 024 .. GENERATED FROM PYTHON SOURCE LINES 385-396 MLEM ---- One MLEM iteration uses the full data (:math:`A`, :math:`y`, :math:`s`) and the full sensitivity image :math:`A^T 1`. It is guaranteed to converge to the maximum-likelihood solution but requires one full data pass per iteration. .. note:: MLEM is only run when ``run_mlem`` is ``True`` (i.e. on CUDA devices). On CPU it is skipped because the large number of iterations is too slow. .. GENERATED FROM PYTHON SOURCE LINES 396-405 .. code-block:: Python if run_mlem: df_mlem = xp.zeros(num_epochs_mlem, dtype=xp.float32, device=dev) x_mlem = xp.asarray(x_init, copy=True) for i in range(num_epochs_mlem): print(f"MLEM epoch {(i + 1):04} / {num_epochs_mlem:04}", end="\r") x_mlem = em_update(x_mlem, full_data_fidelity, adjoint_ones, fov_mask) df_mlem[i] = full_data_fidelity(x_mlem) print() .. GENERATED FROM PYTHON SOURCE LINES 406-413 OSEM ---- One OSEM epoch cycles through all :math:`m` subsets, each using the subset operator :math:`A^k`, subset data :math:`y^k`, contamination :math:`s^k`, and subset sensitivity :math:`(A^k)^T 1`. Fast empirical convergence but no convergence guarantee. .. GENERATED FROM PYTHON SOURCE LINES 413-429 .. code-block:: Python df_osem = xp.zeros(num_epochs, dtype=xp.float32, device=dev) x_osem = xp.asarray(x_init, copy=True) for i in range(num_epochs): for k in range(len(subset_slices)): print( f"OSEM epoch {(i+1):04} / {num_epochs:04}, subset {(k+1):04} / {num_subsets:04}", end="\r", ) x_osem = em_update( x_osem, subset_data_fidelities[k], subset_adjoint_ones[k], fov_mask ) df_osem[i] = full_data_fidelity(x_osem) print() .. rst-class:: sphx-glr-script-out .. code-block:: none OSEM epoch 0001 / 0005, subset 0001 / 0024 OSEM epoch 0001 / 0005, subset 0002 / 0024 OSEM epoch 0001 / 0005, subset 0003 / 0024 OSEM epoch 0001 / 0005, subset 0004 / 0024 OSEM epoch 0001 / 0005, subset 0005 / 0024 OSEM epoch 0001 / 0005, subset 0006 / 0024 OSEM epoch 0001 / 0005, subset 0007 / 0024 OSEM epoch 0001 / 0005, subset 0008 / 0024 OSEM epoch 0001 / 0005, subset 0009 / 0024 OSEM epoch 0001 / 0005, subset 0010 / 0024 OSEM epoch 0001 / 0005, subset 0011 / 0024 OSEM epoch 0001 / 0005, subset 0012 / 0024 OSEM epoch 0001 / 0005, subset 0013 / 0024 OSEM epoch 0001 / 0005, subset 0014 / 0024 OSEM epoch 0001 / 0005, subset 0015 / 0024 OSEM epoch 0001 / 0005, subset 0016 / 0024 OSEM epoch 0001 / 0005, subset 0017 / 0024 OSEM epoch 0001 / 0005, subset 0018 / 0024 OSEM epoch 0001 / 0005, subset 0019 / 0024 OSEM epoch 0001 / 0005, subset 0020 / 0024 OSEM epoch 0001 / 0005, subset 0021 / 0024 OSEM epoch 0001 / 0005, subset 0022 / 0024 OSEM epoch 0001 / 0005, subset 0023 / 0024 OSEM epoch 0001 / 0005, subset 0024 / 0024 OSEM epoch 0002 / 0005, subset 0001 / 0024 OSEM epoch 0002 / 0005, subset 0002 / 0024 OSEM epoch 0002 / 0005, subset 0003 / 0024 OSEM epoch 0002 / 0005, subset 0004 / 0024 OSEM epoch 0002 / 0005, subset 0005 / 0024 OSEM epoch 0002 / 0005, subset 0006 / 0024 OSEM epoch 0002 / 0005, subset 0007 / 0024 OSEM epoch 0002 / 0005, subset 0008 / 0024 OSEM epoch 0002 / 0005, subset 0009 / 0024 OSEM epoch 0002 / 0005, subset 0010 / 0024 OSEM epoch 0002 / 0005, subset 0011 / 0024 OSEM epoch 0002 / 0005, subset 0012 / 0024 OSEM epoch 0002 / 0005, subset 0013 / 0024 OSEM epoch 0002 / 0005, subset 0014 / 0024 OSEM epoch 0002 / 0005, subset 0015 / 0024 OSEM epoch 0002 / 0005, subset 0016 / 0024 OSEM epoch 0002 / 0005, subset 0017 / 0024 OSEM epoch 0002 / 0005, subset 0018 / 0024 OSEM epoch 0002 / 0005, subset 0019 / 0024 OSEM epoch 0002 / 0005, subset 0020 / 0024 OSEM epoch 0002 / 0005, subset 0021 / 0024 OSEM epoch 0002 / 0005, subset 0022 / 0024 OSEM epoch 0002 / 0005, subset 0023 / 0024 OSEM epoch 0002 / 0005, subset 0024 / 0024 OSEM epoch 0003 / 0005, subset 0001 / 0024 OSEM epoch 0003 / 0005, subset 0002 / 0024 OSEM epoch 0003 / 0005, subset 0003 / 0024 OSEM epoch 0003 / 0005, subset 0004 / 0024 OSEM epoch 0003 / 0005, subset 0005 / 0024 OSEM epoch 0003 / 0005, subset 0006 / 0024 OSEM epoch 0003 / 0005, subset 0007 / 0024 OSEM epoch 0003 / 0005, subset 0008 / 0024 OSEM epoch 0003 / 0005, subset 0009 / 0024 OSEM epoch 0003 / 0005, subset 0010 / 0024 OSEM epoch 0003 / 0005, subset 0011 / 0024 OSEM epoch 0003 / 0005, subset 0012 / 0024 OSEM epoch 0003 / 0005, subset 0013 / 0024 OSEM epoch 0003 / 0005, subset 0014 / 0024 OSEM epoch 0003 / 0005, subset 0015 / 0024 OSEM epoch 0003 / 0005, subset 0016 / 0024 OSEM epoch 0003 / 0005, subset 0017 / 0024 OSEM epoch 0003 / 0005, subset 0018 / 0024 OSEM epoch 0003 / 0005, subset 0019 / 0024 OSEM epoch 0003 / 0005, subset 0020 / 0024 OSEM epoch 0003 / 0005, subset 0021 / 0024 OSEM epoch 0003 / 0005, subset 0022 / 0024 OSEM epoch 0003 / 0005, subset 0023 / 0024 OSEM epoch 0003 / 0005, subset 0024 / 0024 OSEM epoch 0004 / 0005, subset 0001 / 0024 OSEM epoch 0004 / 0005, subset 0002 / 0024 OSEM epoch 0004 / 0005, subset 0003 / 0024 OSEM epoch 0004 / 0005, subset 0004 / 0024 OSEM epoch 0004 / 0005, subset 0005 / 0024 OSEM epoch 0004 / 0005, subset 0006 / 0024 OSEM epoch 0004 / 0005, subset 0007 / 0024 OSEM epoch 0004 / 0005, subset 0008 / 0024 OSEM epoch 0004 / 0005, subset 0009 / 0024 OSEM epoch 0004 / 0005, subset 0010 / 0024 OSEM epoch 0004 / 0005, subset 0011 / 0024 OSEM epoch 0004 / 0005, subset 0012 / 0024 OSEM epoch 0004 / 0005, subset 0013 / 0024 OSEM epoch 0004 / 0005, subset 0014 / 0024 OSEM epoch 0004 / 0005, subset 0015 / 0024 OSEM epoch 0004 / 0005, subset 0016 / 0024 OSEM epoch 0004 / 0005, subset 0017 / 0024 OSEM epoch 0004 / 0005, subset 0018 / 0024 OSEM epoch 0004 / 0005, subset 0019 / 0024 OSEM epoch 0004 / 0005, subset 0020 / 0024 OSEM epoch 0004 / 0005, subset 0021 / 0024 OSEM epoch 0004 / 0005, subset 0022 / 0024 OSEM epoch 0004 / 0005, subset 0023 / 0024 OSEM epoch 0004 / 0005, subset 0024 / 0024 OSEM epoch 0005 / 0005, subset 0001 / 0024 OSEM epoch 0005 / 0005, subset 0002 / 0024 OSEM epoch 0005 / 0005, subset 0003 / 0024 OSEM epoch 0005 / 0005, subset 0004 / 0024 OSEM epoch 0005 / 0005, subset 0005 / 0024 OSEM epoch 0005 / 0005, subset 0006 / 0024 OSEM epoch 0005 / 0005, subset 0007 / 0024 OSEM epoch 0005 / 0005, subset 0008 / 0024 OSEM epoch 0005 / 0005, subset 0009 / 0024 OSEM epoch 0005 / 0005, subset 0010 / 0024 OSEM epoch 0005 / 0005, subset 0011 / 0024 OSEM epoch 0005 / 0005, subset 0012 / 0024 OSEM epoch 0005 / 0005, subset 0013 / 0024 OSEM epoch 0005 / 0005, subset 0014 / 0024 OSEM epoch 0005 / 0005, subset 0015 / 0024 OSEM epoch 0005 / 0005, subset 0016 / 0024 OSEM epoch 0005 / 0005, subset 0017 / 0024 OSEM epoch 0005 / 0005, subset 0018 / 0024 OSEM epoch 0005 / 0005, subset 0019 / 0024 OSEM epoch 0005 / 0005, subset 0020 / 0024 OSEM epoch 0005 / 0005, subset 0021 / 0024 OSEM epoch 0005 / 0005, subset 0022 / 0024 OSEM epoch 0005 / 0005, subset 0023 / 0024 OSEM epoch 0005 / 0005, subset 0024 / 0024 .. GENERATED FROM PYTHON SOURCE LINES 430-456 SVRG ---- Each SVRG epoch consists of two phases: 1. **Anchor phase** (every other epoch): compute and store all :math:`m` subset gradients at the current point :math:`\tilde{x}`, then take a full gradient step. 2. **Variance-reduced subset updates**: for each subset :math:`k`, form the variance-reduced gradient .. math:: g^{VR}_k = m \left( \nabla f_k(x) - \tilde{g}_k \right) + \sum_{j=1}^m \tilde{g}_j where :math:`\tilde{g}_k = \nabla f_k(\tilde{x})` are the stored anchor gradients. .. note:: The anchor phase requires one full pass through all subsets to compute :math:`\tilde{g}_1, \ldots, \tilde{g}_m`. Therefore, each SVRG epoch that includes an anchor phase costs **two full data passes** (one for the snapshot, one for the :math:`m` subset updates), compared to one full data pass for OSEM or MLEM. The epoch axis in the convergence plot understates SVRG's computational cost relative to OSEM by a factor of roughly two. .. GENERATED FROM PYTHON SOURCE LINES 456-534 .. code-block:: Python def svrg_calc_snapshot_gradients( x_cur: Array, subset_obj_functions: Sequence[C1Function], ) -> tuple[Array, Array]: """Store all subset gradients at the current anchor point and return their sum. Returns ------- stored_grads : Array, shape (m, *x_cur.shape) Stacked subset gradients evaluated at the anchor point. full_grad : Array, shape x_cur.shape Sum of all subset gradients. """ m = len(subset_obj_functions) stored_grads = xp.zeros((m,) + x_cur.shape, dtype=x_cur.dtype, device=dev) for k, df in enumerate(subset_obj_functions): stored_grads[k] = df.gradient(x_cur) full_grad = xp.sum(stored_grads, axis=0) return stored_grads, full_grad def svrg_update( x_cur: Array, subset_idx: int, subset_obj_functions: Sequence[C1Function], stored_subset_gradients: Array, full_gradient: Array, precond: Array, step_size: float = 1.0, ) -> Array: """Single SVRG subset update with variance-reduced gradient.""" m = len(subset_obj_functions) grad_k = subset_obj_functions[subset_idx].gradient(x_cur) approx_grad = m * (grad_k - stored_subset_gradients[subset_idx]) + full_gradient return xp.clip(x_cur - step_size * precond * approx_grad, 0, None) # start SVRG from the same warm-start as OSEM and MLEM x_svrg = xp.asarray(x_init, copy=True) svrg_step_size = 1.0 df_svrg = xp.zeros(num_epochs, dtype=xp.float32, device=dev) for epoch in range(num_epochs): if epoch % 2 == 0: if epoch <= 4: if fov_mask is None: svrg_precond = x_svrg / adjoint_ones else: svrg_precond = xp.where( fov_mask, x_svrg / adjoint_ones, xp.zeros_like(x_svrg) ) stored_grads, full_grad = svrg_calc_snapshot_gradients( x_svrg, subset_data_fidelities ) x_svrg = xp.clip(x_svrg - svrg_step_size * svrg_precond * full_grad, 0, None) for k in range(num_subsets): print( f"SVRG epoch {(epoch+1):04} / {num_epochs:04}, subset {(k+1):04} / {num_subsets:04}", end="\r", ) x_svrg = svrg_update( x_svrg, k, subset_data_fidelities, stored_grads, full_grad, svrg_precond, step_size=svrg_step_size, ) df_svrg[epoch] = full_data_fidelity(x_svrg) .. rst-class:: sphx-glr-script-out .. code-block:: none SVRG epoch 0001 / 0005, subset 0001 / 0024 SVRG epoch 0001 / 0005, subset 0002 / 0024 SVRG epoch 0001 / 0005, subset 0003 / 0024 SVRG epoch 0001 / 0005, subset 0004 / 0024 SVRG epoch 0001 / 0005, subset 0005 / 0024 SVRG epoch 0001 / 0005, subset 0006 / 0024 SVRG epoch 0001 / 0005, subset 0007 / 0024 SVRG epoch 0001 / 0005, subset 0008 / 0024 SVRG epoch 0001 / 0005, subset 0009 / 0024 SVRG epoch 0001 / 0005, subset 0010 / 0024 SVRG epoch 0001 / 0005, subset 0011 / 0024 SVRG epoch 0001 / 0005, subset 0012 / 0024 SVRG epoch 0001 / 0005, subset 0013 / 0024 SVRG epoch 0001 / 0005, subset 0014 / 0024 SVRG epoch 0001 / 0005, subset 0015 / 0024 SVRG epoch 0001 / 0005, subset 0016 / 0024 SVRG epoch 0001 / 0005, subset 0017 / 0024 SVRG epoch 0001 / 0005, subset 0018 / 0024 SVRG epoch 0001 / 0005, subset 0019 / 0024 SVRG epoch 0001 / 0005, subset 0020 / 0024 SVRG epoch 0001 / 0005, subset 0021 / 0024 SVRG epoch 0001 / 0005, subset 0022 / 0024 SVRG epoch 0001 / 0005, subset 0023 / 0024 SVRG epoch 0001 / 0005, subset 0024 / 0024 SVRG epoch 0002 / 0005, subset 0001 / 0024 SVRG epoch 0002 / 0005, subset 0002 / 0024 SVRG epoch 0002 / 0005, subset 0003 / 0024 SVRG epoch 0002 / 0005, subset 0004 / 0024 SVRG epoch 0002 / 0005, subset 0005 / 0024 SVRG epoch 0002 / 0005, subset 0006 / 0024 SVRG epoch 0002 / 0005, subset 0007 / 0024 SVRG epoch 0002 / 0005, subset 0008 / 0024 SVRG epoch 0002 / 0005, subset 0009 / 0024 SVRG epoch 0002 / 0005, subset 0010 / 0024 SVRG epoch 0002 / 0005, subset 0011 / 0024 SVRG epoch 0002 / 0005, subset 0012 / 0024 SVRG epoch 0002 / 0005, subset 0013 / 0024 SVRG epoch 0002 / 0005, subset 0014 / 0024 SVRG epoch 0002 / 0005, subset 0015 / 0024 SVRG epoch 0002 / 0005, subset 0016 / 0024 SVRG epoch 0002 / 0005, subset 0017 / 0024 SVRG epoch 0002 / 0005, subset 0018 / 0024 SVRG epoch 0002 / 0005, subset 0019 / 0024 SVRG epoch 0002 / 0005, subset 0020 / 0024 SVRG epoch 0002 / 0005, subset 0021 / 0024 SVRG epoch 0002 / 0005, subset 0022 / 0024 SVRG epoch 0002 / 0005, subset 0023 / 0024 SVRG epoch 0002 / 0005, subset 0024 / 0024 SVRG epoch 0003 / 0005, subset 0001 / 0024 SVRG epoch 0003 / 0005, subset 0002 / 0024 SVRG epoch 0003 / 0005, subset 0003 / 0024 SVRG epoch 0003 / 0005, subset 0004 / 0024 SVRG epoch 0003 / 0005, subset 0005 / 0024 SVRG epoch 0003 / 0005, subset 0006 / 0024 SVRG epoch 0003 / 0005, subset 0007 / 0024 SVRG epoch 0003 / 0005, subset 0008 / 0024 SVRG epoch 0003 / 0005, subset 0009 / 0024 SVRG epoch 0003 / 0005, subset 0010 / 0024 SVRG epoch 0003 / 0005, subset 0011 / 0024 SVRG epoch 0003 / 0005, subset 0012 / 0024 SVRG epoch 0003 / 0005, subset 0013 / 0024 SVRG epoch 0003 / 0005, subset 0014 / 0024 SVRG epoch 0003 / 0005, subset 0015 / 0024 SVRG epoch 0003 / 0005, subset 0016 / 0024 SVRG epoch 0003 / 0005, subset 0017 / 0024 SVRG epoch 0003 / 0005, subset 0018 / 0024 SVRG epoch 0003 / 0005, subset 0019 / 0024 SVRG epoch 0003 / 0005, subset 0020 / 0024 SVRG epoch 0003 / 0005, subset 0021 / 0024 SVRG epoch 0003 / 0005, subset 0022 / 0024 SVRG epoch 0003 / 0005, subset 0023 / 0024 SVRG epoch 0003 / 0005, subset 0024 / 0024 SVRG epoch 0004 / 0005, subset 0001 / 0024 SVRG epoch 0004 / 0005, subset 0002 / 0024 SVRG epoch 0004 / 0005, subset 0003 / 0024 SVRG epoch 0004 / 0005, subset 0004 / 0024 SVRG epoch 0004 / 0005, subset 0005 / 0024 SVRG epoch 0004 / 0005, subset 0006 / 0024 SVRG epoch 0004 / 0005, subset 0007 / 0024 SVRG epoch 0004 / 0005, subset 0008 / 0024 SVRG epoch 0004 / 0005, subset 0009 / 0024 SVRG epoch 0004 / 0005, subset 0010 / 0024 SVRG epoch 0004 / 0005, subset 0011 / 0024 SVRG epoch 0004 / 0005, subset 0012 / 0024 SVRG epoch 0004 / 0005, subset 0013 / 0024 SVRG epoch 0004 / 0005, subset 0014 / 0024 SVRG epoch 0004 / 0005, subset 0015 / 0024 SVRG epoch 0004 / 0005, subset 0016 / 0024 SVRG epoch 0004 / 0005, subset 0017 / 0024 SVRG epoch 0004 / 0005, subset 0018 / 0024 SVRG epoch 0004 / 0005, subset 0019 / 0024 SVRG epoch 0004 / 0005, subset 0020 / 0024 SVRG epoch 0004 / 0005, subset 0021 / 0024 SVRG epoch 0004 / 0005, subset 0022 / 0024 SVRG epoch 0004 / 0005, subset 0023 / 0024 SVRG epoch 0004 / 0005, subset 0024 / 0024 SVRG epoch 0005 / 0005, subset 0001 / 0024 SVRG epoch 0005 / 0005, subset 0002 / 0024 SVRG epoch 0005 / 0005, subset 0003 / 0024 SVRG epoch 0005 / 0005, subset 0004 / 0024 SVRG epoch 0005 / 0005, subset 0005 / 0024 SVRG epoch 0005 / 0005, subset 0006 / 0024 SVRG epoch 0005 / 0005, subset 0007 / 0024 SVRG epoch 0005 / 0005, subset 0008 / 0024 SVRG epoch 0005 / 0005, subset 0009 / 0024 SVRG epoch 0005 / 0005, subset 0010 / 0024 SVRG epoch 0005 / 0005, subset 0011 / 0024 SVRG epoch 0005 / 0005, subset 0012 / 0024 SVRG epoch 0005 / 0005, subset 0013 / 0024 SVRG epoch 0005 / 0005, subset 0014 / 0024 SVRG epoch 0005 / 0005, subset 0015 / 0024 SVRG epoch 0005 / 0005, subset 0016 / 0024 SVRG epoch 0005 / 0005, subset 0017 / 0024 SVRG epoch 0005 / 0005, subset 0018 / 0024 SVRG epoch 0005 / 0005, subset 0019 / 0024 SVRG epoch 0005 / 0005, subset 0020 / 0024 SVRG epoch 0005 / 0005, subset 0021 / 0024 SVRG epoch 0005 / 0005, subset 0022 / 0024 SVRG epoch 0005 / 0005, subset 0023 / 0024 SVRG epoch 0005 / 0005, subset 0024 / 0024 .. GENERATED FROM PYTHON SOURCE LINES 535-543 Convergence comparison ---------------------- We plot the negative Poisson log-likelihood vs epoch for OSEM and SVRG, and overlay two horizontal reference lines showing where MLEM stands after ``num_epochs`` and ``num_epochs_mlem`` iterations respectively. One epoch of OSEM or SVRG corresponds to one cycle through all subsets (roughly one full data pass for OSEM, roughly two for SVRG). .. GENERATED FROM PYTHON SOURCE LINES 543-631 .. code-block:: Python # data-pass counts # OSEM: 1 full data pass per epoch # SVRG: 2 passes on anchor epochs (snapshot + subset updates), # 1 pass on non-anchor epochs (subset updates only) # MLEM: 1 pass per iteration epochs = np.arange(1, num_epochs + 1) osem_passes = epochs.copy() svrg_passes_per_epoch = np.where(np.arange(num_epochs) % 2 == 0, 2, 1) svrg_cumulative_passes = np.cumsum(svrg_passes_per_epoch) max_passes = int(svrg_cumulative_passes[-1]) if run_mlem: df_mlem_trimmed = to_numpy_array(df_mlem[:max_passes]) if run_mlem: df_min = min(float(xp.min(df_mlem)), float(xp.min(df_osem)), float(xp.min(df_svrg))) else: df_min = min(float(xp.min(df_osem)), float(xp.min(df_svrg))) df_max = float(df_osem[0]) # use first OSEM epoch as upper limit osem_label = f"OSEM ({num_subsets} subsets)" svrg_label = f"SVRG ({num_subsets} subsets, step={svrg_step_size:.1f})" mlem_label = "MLEM" fig, axs = plt.subplots(1, 2, figsize=(12, 4), layout="constrained") # --- left: vs epoch --- axs[0].plot(epochs, to_numpy_array(df_osem), label=osem_label, marker="o") axs[0].plot(epochs, to_numpy_array(df_svrg), label=svrg_label, marker="o") if run_mlem: axs[0].axhline( float(df_mlem[50 - 1]), label=f"{mlem_label} (50 iter.)", ls="--", color="gray", ) axs[0].axhline( float(df_mlem[100 - 1]), label=f"{mlem_label} (100 iter.)", ls="--", color="gray", ) axs[0].axhline( float(df_mlem[-1]), label=f"{mlem_label} ({num_epochs_mlem} iter.)", ls="--", color="black", ) axs[0].set_ylim(df_min, df_max) axs[0].set_xlabel("Epoch") axs[0].set_ylabel("Negative Poisson log-likelihood") axs[0].legend() axs[0].grid(ls=":") # --- right: vs full data passes --- axs[1].plot(osem_passes, to_numpy_array(df_osem), label=osem_label, marker="o") axs[1].plot( svrg_cumulative_passes, to_numpy_array(df_svrg), label=svrg_label, marker="o" ) if run_mlem: axs[1].axhline( float(df_mlem[50 - 1]), label=f"{mlem_label} (50 iter.)", ls=":", color="gray", ) axs[1].axhline( float(df_mlem[100 - 1]), label=f"{mlem_label} (100 iter.)", ls="--", color="gray", ) axs[1].axhline( float(df_mlem[-1]), label=f"{mlem_label} ({num_epochs_mlem} iter.)", ls="--", color="black", ) axs[1].set_ylim(df_min, df_max) axs[1].set_xlabel("Full data passes") axs[1].set_ylabel("Negative Poisson log-likelihood") axs[1].legend() axs[1].grid(ls=":") fig.show() .. image-sg:: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_001.png :alt: 00 run mlem osem svrg :srcset: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 632-643 .. code-block:: Python vmax = float(xp.max(x_svrg)) fig, axs, widgets = show_vol_cuts( x_osem, voxel_size=voxel_size, fig_title="OSEM result", vmin=0, vmax=vmax, ) fig.show() .. image-sg:: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_002.png :alt: OSEM result :srcset: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 644-653 .. code-block:: Python fig2, axs2, widgets2 = show_vol_cuts( x_svrg, voxel_size=voxel_size, fig_title="SVRG result", vmin=0, vmax=vmax, ) fig2.show() .. image-sg:: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_003.png :alt: SVRG result :srcset: /auto_examples/03_algorithms/images/sphx_glr_00_run_mlem_osem_svrg_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 654-657 .. rubric:: References .. footbibliography:: .. rst-class:: sphx-glr-timing **Total running time of the script:** (2 minutes 21.925 seconds) .. _sphx_glr_download_auto_examples_03_algorithms_00_run_mlem_osem_svrg.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 00_run_mlem_osem_svrg.ipynb <00_run_mlem_osem_svrg.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 00_run_mlem_osem_svrg.py <00_run_mlem_osem_svrg.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 00_run_mlem_osem_svrg.zip <00_run_mlem_osem_svrg.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_