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

\[f(x) = \sum_i \bar{y}_i - y_i \log \bar{y}_i, \qquad \bar{y}(x) = A x + s\]

subject to \(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 = \(m\) subset updates ~= one full data pass; fast empirical convergence but no convergence guarantee.

• SVRG – stochastic variance-reduced gradient with subsets; one epoch = \(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 \(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.

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

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)

Using array API: array_api_compat.torch, device: cpu

# 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"

Setup of the forward model \(\bar{y}(x) = A x + s\)

We setup a linear forward operator \(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 LinearOperator (e.g. the high-level projectors).

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,
)

setup the LOR descriptor that defines the sinogram

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
)

Attenuation image and sinogram setup

# 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))

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.

# 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))

Simulation of projection data

We setup an arbitrary ground truth \(x_{true}\) and simulate noise-free and noisy data \(y\) by adding Poisson noise.

# 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,
)

Splitting of the forward model into subsets \(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.

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
)

EM update

The EM update used in MLEM and OSEM is [1] [2] [3] [4]

\[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 \(D = \operatorname{diag}(x / (A^T 1))\):

\[x^+ = x - D \, \nabla_x f(x).\]

We implement this as a single function used by both MLEM and OSEM.

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)

Setup of objective functions and sensitivity images

We define one C2AffineObjective per subset (for OSEM and SVRG) and one for the full data (for MLEM and objective evaluation). The sensitivity image \(A^T 1\) and its per-subset counterparts \((A^k)^T 1\) are precomputed once.

# 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)

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

Setup of data-fidelity terms

We define one C2AffineObjective per subset (for OSEM and SVRG) and one for the full data (for MLEM and objective evaluation).

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 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.

# 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()

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

MLEM

One MLEM iteration uses the full data (\(A\), \(y\), \(s\)) and the full sensitivity image \(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.

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()

OSEM

One OSEM epoch cycles through all \(m\) subsets, each using the subset operator \(A^k\), subset data \(y^k\), contamination \(s^k\), and subset sensitivity \((A^k)^T 1\). Fast empirical convergence but no convergence guarantee.

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()

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

SVRG

Each SVRG epoch consists of two phases:

1. Anchor phase (every other epoch): compute and store all \(m\) subset gradients at the current point \(\tilde{x}\), then take a full gradient step.

2. Variance-reduced subset updates: for each subset \(k\), form the variance-reduced gradient

   \[g^{VR}_k = m \left( \nabla f_k(x) - \tilde{g}_k \right) + \sum_{j=1}^m \tilde{g}_j\]

   where \(\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 \(\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 \(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.

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)

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

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).

# 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()

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()

fig2, axs2, widgets2 = show_vol_cuts(
    x_svrg,
    voxel_size=voxel_size,
    fig_title="SVRG result",
    vmin=0,
    vmax=vmax,
)
fig2.show()

References

[1]

P. Dempster, N. Laird, and D. Rubin. Maximum Likelihood from Incomplete Data via the EM Algorithm. Journal of the Royal Statistical Society. Series B (Methodological), 39:1–38, 1977.

[2]

L. A. Shepp and Y. Vardi. Maximum Likelihood Reconstruction for Emission Tomography. IEEE Transactions on Medical Imaging, 1(2):113–122, 1982. doi:10.1109/TMI.1982.4307558.

[3]

Kenneth Lange and Richard Carson. EM reconstruction algorithms for emission and transmission tomography. Journal of Computer Assisted Tomography, 8(2):306–316, 1984.

[4]

H. M. Hudson and R. S. Larkin. Accelerated Image Reconstructiong using Ordered Subsets of Projection Data. IEEE transactions on medical imaging, 13(4):601–609, 1994. doi:10.1109/42.363108.