Note

Go to the end to download the full example code.

Convergence comparison: SGD vs SVRG with logcosh regularization¶

This example compares the convergence speed (per epoch) of two algorithms for minimising the regularised negative Poisson log-likelihood

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

where the edge-preserving logcosh penalty is

\[R(x) = \delta \sum_i \log\!\cosh\!\left(\frac{(Gx)_i}{\delta}\right)\]

and $G$ is the finite forward-difference operator. The $\delta$ prefactor ensures the asymptotic gradient magnitude equals 1 regardless of $\delta$, so the regularisation strength $\beta$ retains the same meaning across different choices of $\delta$. The scale $\delta$ itself controls the transition between two regimes:

Quadratic for $|(Gx)_i| \ll \delta$: $R(x) \approx \tfrac{1}{2\delta}\|Gx\|_2^2$.
Linear for $|(Gx)_i| \gg \delta$: $R(x) \approx \|Gx\|_1 - n\,\delta\log 2 \approx \|Gx\|_1$.

Setting $\delta$ well below the typical edge gradient places true edges in the linear regime (edge-preserving) while penalising smooth-region deviations quadratically. The objective is decomposed into $m$ subset functions

\[f_k(x) = \underbrace{\sum_{i \in S_k} \bar{y}_i - y_i \log \bar{y}_i}_{\text{subset data fidelity}} + \frac{\beta}{m} R(x), \qquad k = 1, \ldots, m,\]

so that $F(x) = \sum_{k=1}^m f_k(x)$ exactly. Both algorithms below exploit this splitting:

SGD – stochastic gradient descent using ordered subsets; 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 while achieving the fast per-epoch progress of SGD.

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 SGD 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,
    LogCosh,
    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 SGD and SVRG
num_subsets = 12

# if run on a CPU limit the number of epochs
num_epochs = (120 if dev == "cpu" else 240) // num_subsets

# regularisation weight beta
beta = 1.0
# delta value relative to max of ground truth image for logcosh prior
delta_rel = 0.1

# step size for SGD and SVRG updates
step_size = 1.0

# factor that scales the ground truth image (also reconstruction) and the number of counts
count_factor = 1.0

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.

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 = count_factor * elliptic_cylinder_phantom(
    xp, dev, image_shape=img_shape, voxel_size=voxel_size
)

Attenuation image and sinogram setup¶

x_att = 0.01 * xp.astype(x_true > 0, xp.float32)
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]
)

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 noisy data $y$ by adding Poisson noise.

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

Splitting of the forward model into subsets $A^k$¶

Calculate the view numbers and slices for each subset.

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
)

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]

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

    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
)

Regularisation and subset objective functions¶

The logcosh penalty $R(x) = \delta \sum_i \log\cosh\!\left((Gx)_i/\delta\right)$ is built from the FiniteForwardDifference operator $G$ and LogCosh. The full regulariser reg (weight $\beta$) is used only for the total objective evaluation. Each subset function

\[f_k(x) = \sum_{i \in S_k} \left( \bar{y}_i(x) - y_i \log \bar{y}_i(x) \right) + \frac{\beta}{m} R(x)\]

is formed by adding a LogCosh scaled by $\beta / m$ to the subset data fidelity, so that $\sum_k f_k(x) = F(x)$.

delta is set to delta_rel times the maximum of the ground truth image. With delta_rel = 0.1 edges with gradient equal to the image maximum have $|(Gx)|/\delta = 10$, placing them firmly in the linear regime ($\tanh(10) \approx 1$), while smooth-region gradients near zero remain quadratic.

G = parallelproj.operators.FiniteForwardDifference(pet_lin_op.in_shape)

delta = float(xp.max(x_true)) * delta_rel

reg = C2AffineObjective(LogCosh(delta=delta, beta=beta), G)

Setup of objective functions and sensitivity image¶

We define one subset objective $f_k$ per subset and one full objective $F$ for evaluation, as well as the sensitivity image $A^T 1$

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.

# sensitivity image (transpose applied to the all-ones vector)
adjoint_ones = pet_lin_op.adjoint(
    xp.ones(pet_lin_op.out_shape, dtype=xp.float32, device=dev)
)

# reg/m term shared by all subset objectives
reg_per_subset = C2AffineObjective(LogCosh(delta=delta, beta=beta / num_subsets), G)

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

# f_k = data_fidelity_k + (beta/m) * R(x)
subset_objectives = [
    C2AffineObjective(
        NegPoissonLogL(y[sl], exact=exact_mode),
        pet_subset_linop_seq[k],
        contamination[sl],
    )
    + reg_per_subset
    for k, sl in enumerate(subset_slices)
]

full_data_fidelity = C2AffineObjective(
    NegPoissonLogL(y, exact=exact_mode), pet_lin_op, contamination
)


# also setup the full objective F for evaluation of the iterates
total_objective = full_data_fidelity + reg

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$. In this example the regulariser Hessian keeps the preconditioner denominator strictly positive everywhere, so there is no divide-by-zero risk. However, zeroing the initial image outside the FOV ensures those voxels stay at zero throughout reconstruction: because the preconditioner is proportional to $x$, a zero initialisation propagates as zero updates in every subsequent SGD / SVRG step. 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

Warm start¶

Run one SGD epoch with regularisation as a common warm-start.

x_init = xp.ones(pet_lin_op.in_shape, dtype=xp.float32, device=dev)
if fov_mask is not None:
    x_init = xp.where(fov_mask, x_init, xp.zeros_like(x_init))
for k in range(num_subsets):
    print(f"warm-start SGD subset {(k+1):03} / {num_subsets:03}", end="\r")
    _denom = adjoint_ones + 2 * reg.hessian_diag_vec_prod(x_init, x_init)
    if fov_mask is None:
        init_precond = x_init / _denom
    else:
        init_precond = xp.where(fov_mask, x_init / _denom, xp.zeros_like(x_init))
    x_init = xp.clip(
        x_init - init_precond * (num_subsets * subset_objectives[k].gradient(x_init)),
        0,
        None,
    )
print()

warm-start SGD subset 001 / 012
warm-start SGD subset 002 / 012
warm-start SGD subset 003 / 012
warm-start SGD subset 004 / 012
warm-start SGD subset 005 / 012
warm-start SGD subset 006 / 012
warm-start SGD subset 007 / 012
warm-start SGD subset 008 / 012
warm-start SGD subset 009 / 012
warm-start SGD subset 010 / 012
warm-start SGD subset 011 / 012
warm-start SGD subset 012 / 012

SGD with regularisation¶

Each SGD epoch cycles through all $m$ subsets. Because $F(x) = \sum_k f_k(x)$, the full gradient is approximated as

\[\nabla F(x) \approx m\,\nabla f_k(x)\]

and a gradient step with a diagonal preconditioner $D$ is taken:

\[x^+ = \left(x - D\, m\,\nabla f_k(x)\right)_+.\]

df_sgd = xp.zeros(num_epochs + 1, dtype=xp.float32, device=dev)
x_sgd = xp.asarray(x_init, copy=True)
df_sgd[0] = total_objective(x_sgd)
sgd_recons = xp.zeros((num_epochs + 1,) + img_shape)
sgd_recons[0, ...] = x_sgd

for i in range(num_epochs):
    if i % 2 == 0 and i <= 4:
        _denom = adjoint_ones + 2 * reg.hessian_diag_vec_prod(x_sgd, x_sgd)
        if fov_mask is None:
            sgd_precond = x_sgd / _denom
        else:
            sgd_precond = xp.where(fov_mask, x_sgd / _denom, xp.zeros_like(x_sgd))
    for k in range(num_subsets):
        print(
            f"SGD epoch {(i+1):04} / {num_epochs:04}, subset {(k+1):04} / {num_subsets:04}",
            end="\r",
        )
        approx_grad = num_subsets * subset_objectives[k].gradient(x_sgd)
        x_sgd = xp.clip(x_sgd - step_size * sgd_precond * approx_grad, 0, None)

    df_sgd[i + 1] = total_objective(x_sgd)
    sgd_recons[i + 1, ...] = x_sgd
print()

SGD epoch 0001 / 0010, subset 0001 / 0012
SGD epoch 0001 / 0010, subset 0002 / 0012
SGD epoch 0001 / 0010, subset 0003 / 0012
SGD epoch 0001 / 0010, subset 0004 / 0012
SGD epoch 0001 / 0010, subset 0005 / 0012
SGD epoch 0001 / 0010, subset 0006 / 0012
SGD epoch 0001 / 0010, subset 0007 / 0012
SGD epoch 0001 / 0010, subset 0008 / 0012
SGD epoch 0001 / 0010, subset 0009 / 0012
SGD epoch 0001 / 0010, subset 0010 / 0012
SGD epoch 0001 / 0010, subset 0011 / 0012
SGD epoch 0001 / 0010, subset 0012 / 0012
SGD epoch 0002 / 0010, subset 0001 / 0012
SGD epoch 0002 / 0010, subset 0002 / 0012
SGD epoch 0002 / 0010, subset 0003 / 0012
SGD epoch 0002 / 0010, subset 0004 / 0012
SGD epoch 0002 / 0010, subset 0005 / 0012
SGD epoch 0002 / 0010, subset 0006 / 0012
SGD epoch 0002 / 0010, subset 0007 / 0012
SGD epoch 0002 / 0010, subset 0008 / 0012
SGD epoch 0002 / 0010, subset 0009 / 0012
SGD epoch 0002 / 0010, subset 0010 / 0012
SGD epoch 0002 / 0010, subset 0011 / 0012
SGD epoch 0002 / 0010, subset 0012 / 0012
SGD epoch 0003 / 0010, subset 0001 / 0012
SGD epoch 0003 / 0010, subset 0002 / 0012
SGD epoch 0003 / 0010, subset 0003 / 0012
SGD epoch 0003 / 0010, subset 0004 / 0012
SGD epoch 0003 / 0010, subset 0005 / 0012
SGD epoch 0003 / 0010, subset 0006 / 0012
SGD epoch 0003 / 0010, subset 0007 / 0012
SGD epoch 0003 / 0010, subset 0008 / 0012
SGD epoch 0003 / 0010, subset 0009 / 0012
SGD epoch 0003 / 0010, subset 0010 / 0012
SGD epoch 0003 / 0010, subset 0011 / 0012
SGD epoch 0003 / 0010, subset 0012 / 0012
SGD epoch 0004 / 0010, subset 0001 / 0012
SGD epoch 0004 / 0010, subset 0002 / 0012
SGD epoch 0004 / 0010, subset 0003 / 0012
SGD epoch 0004 / 0010, subset 0004 / 0012
SGD epoch 0004 / 0010, subset 0005 / 0012
SGD epoch 0004 / 0010, subset 0006 / 0012
SGD epoch 0004 / 0010, subset 0007 / 0012
SGD epoch 0004 / 0010, subset 0008 / 0012
SGD epoch 0004 / 0010, subset 0009 / 0012
SGD epoch 0004 / 0010, subset 0010 / 0012
SGD epoch 0004 / 0010, subset 0011 / 0012
SGD epoch 0004 / 0010, subset 0012 / 0012
SGD epoch 0005 / 0010, subset 0001 / 0012
SGD epoch 0005 / 0010, subset 0002 / 0012
SGD epoch 0005 / 0010, subset 0003 / 0012
SGD epoch 0005 / 0010, subset 0004 / 0012
SGD epoch 0005 / 0010, subset 0005 / 0012
SGD epoch 0005 / 0010, subset 0006 / 0012
SGD epoch 0005 / 0010, subset 0007 / 0012
SGD epoch 0005 / 0010, subset 0008 / 0012
SGD epoch 0005 / 0010, subset 0009 / 0012
SGD epoch 0005 / 0010, subset 0010 / 0012
SGD epoch 0005 / 0010, subset 0011 / 0012
SGD epoch 0005 / 0010, subset 0012 / 0012
SGD epoch 0006 / 0010, subset 0001 / 0012
SGD epoch 0006 / 0010, subset 0002 / 0012
SGD epoch 0006 / 0010, subset 0003 / 0012
SGD epoch 0006 / 0010, subset 0004 / 0012
SGD epoch 0006 / 0010, subset 0005 / 0012
SGD epoch 0006 / 0010, subset 0006 / 0012
SGD epoch 0006 / 0010, subset 0007 / 0012
SGD epoch 0006 / 0010, subset 0008 / 0012
SGD epoch 0006 / 0010, subset 0009 / 0012
SGD epoch 0006 / 0010, subset 0010 / 0012
SGD epoch 0006 / 0010, subset 0011 / 0012
SGD epoch 0006 / 0010, subset 0012 / 0012
SGD epoch 0007 / 0010, subset 0001 / 0012
SGD epoch 0007 / 0010, subset 0002 / 0012
SGD epoch 0007 / 0010, subset 0003 / 0012
SGD epoch 0007 / 0010, subset 0004 / 0012
SGD epoch 0007 / 0010, subset 0005 / 0012
SGD epoch 0007 / 0010, subset 0006 / 0012
SGD epoch 0007 / 0010, subset 0007 / 0012
SGD epoch 0007 / 0010, subset 0008 / 0012
SGD epoch 0007 / 0010, subset 0009 / 0012
SGD epoch 0007 / 0010, subset 0010 / 0012
SGD epoch 0007 / 0010, subset 0011 / 0012
SGD epoch 0007 / 0010, subset 0012 / 0012
SGD epoch 0008 / 0010, subset 0001 / 0012
SGD epoch 0008 / 0010, subset 0002 / 0012
SGD epoch 0008 / 0010, subset 0003 / 0012
SGD epoch 0008 / 0010, subset 0004 / 0012
SGD epoch 0008 / 0010, subset 0005 / 0012
SGD epoch 0008 / 0010, subset 0006 / 0012
SGD epoch 0008 / 0010, subset 0007 / 0012
SGD epoch 0008 / 0010, subset 0008 / 0012
SGD epoch 0008 / 0010, subset 0009 / 0012
SGD epoch 0008 / 0010, subset 0010 / 0012
SGD epoch 0008 / 0010, subset 0011 / 0012
SGD epoch 0008 / 0010, subset 0012 / 0012
SGD epoch 0009 / 0010, subset 0001 / 0012
SGD epoch 0009 / 0010, subset 0002 / 0012
SGD epoch 0009 / 0010, subset 0003 / 0012
SGD epoch 0009 / 0010, subset 0004 / 0012
SGD epoch 0009 / 0010, subset 0005 / 0012
SGD epoch 0009 / 0010, subset 0006 / 0012
SGD epoch 0009 / 0010, subset 0007 / 0012
SGD epoch 0009 / 0010, subset 0008 / 0012
SGD epoch 0009 / 0010, subset 0009 / 0012
SGD epoch 0009 / 0010, subset 0010 / 0012
SGD epoch 0009 / 0010, subset 0011 / 0012
SGD epoch 0009 / 0010, subset 0012 / 0012
SGD epoch 0010 / 0010, subset 0001 / 0012
SGD epoch 0010 / 0010, subset 0002 / 0012
SGD epoch 0010 / 0010, subset 0003 / 0012
SGD epoch 0010 / 0010, subset 0004 / 0012
SGD epoch 0010 / 0010, subset 0005 / 0012
SGD epoch 0010 / 0010, subset 0006 / 0012
SGD epoch 0010 / 0010, subset 0007 / 0012
SGD epoch 0010 / 0010, subset 0008 / 0012
SGD epoch 0010 / 0010, subset 0009 / 0012
SGD epoch 0010 / 0010, subset 0010 / 0012
SGD epoch 0010 / 0010, subset 0011 / 0012
SGD epoch 0010 / 0010, subset 0012 / 0012

SVRG with regularisation¶

Each SVRG epoch consists of two phases:

Anchor phase (every other epoch): compute and store all $m$ subset gradients $\tilde{g}_k = \nabla f_k(\tilde{x})$ at the current point $\tilde{x}$, then take a full gradient step using $\nabla F(\tilde{x}) = \sum_k \tilde{g}_k$.
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\]

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."""
    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_snapshot_subset_gradients: Array,
    full_snapshot_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_snapshot_subset_gradients[subset_idx])
        + full_snapshot_gradient
    )
    return xp.clip(x_cur - step_size * precond * approx_grad, 0, None)


x_svrg = xp.asarray(x_init, copy=True)
svrg_recons = xp.zeros((num_epochs + 1,) + img_shape)
svrg_recons[0, ...] = x_svrg

df_svrg = xp.zeros(num_epochs + 1, dtype=xp.float32, device=dev)
df_svrg[0] = total_objective(x_svrg)

for epoch in range(num_epochs):
    if epoch % 2 == 0:
        if epoch <= 4:
            _denom = adjoint_ones + 2 * reg.hessian_diag_vec_prod(x_svrg, x_svrg)
            if fov_mask is None:
                svrg_precond = x_svrg / _denom
            else:
                svrg_precond = xp.where(
                    fov_mask, x_svrg / _denom, xp.zeros_like(x_svrg)
                )

        stored_grads, full_grad = svrg_calc_snapshot_gradients(
            x_svrg, subset_objectives
        )
        x_svrg = xp.clip(x_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_objectives,
            stored_grads,
            full_grad,
            svrg_precond,
            step_size=step_size,
        )

    df_svrg[epoch + 1] = total_objective(x_svrg)
    svrg_recons[epoch + 1, ...] = x_svrg

SVRG epoch 0001 / 0010, subset 0001 / 0012
SVRG epoch 0001 / 0010, subset 0002 / 0012
SVRG epoch 0001 / 0010, subset 0003 / 0012
SVRG epoch 0001 / 0010, subset 0004 / 0012
SVRG epoch 0001 / 0010, subset 0005 / 0012
SVRG epoch 0001 / 0010, subset 0006 / 0012
SVRG epoch 0001 / 0010, subset 0007 / 0012
SVRG epoch 0001 / 0010, subset 0008 / 0012
SVRG epoch 0001 / 0010, subset 0009 / 0012
SVRG epoch 0001 / 0010, subset 0010 / 0012
SVRG epoch 0001 / 0010, subset 0011 / 0012
SVRG epoch 0001 / 0010, subset 0012 / 0012
SVRG epoch 0002 / 0010, subset 0001 / 0012
SVRG epoch 0002 / 0010, subset 0002 / 0012
SVRG epoch 0002 / 0010, subset 0003 / 0012
SVRG epoch 0002 / 0010, subset 0004 / 0012
SVRG epoch 0002 / 0010, subset 0005 / 0012
SVRG epoch 0002 / 0010, subset 0006 / 0012
SVRG epoch 0002 / 0010, subset 0007 / 0012
SVRG epoch 0002 / 0010, subset 0008 / 0012
SVRG epoch 0002 / 0010, subset 0009 / 0012
SVRG epoch 0002 / 0010, subset 0010 / 0012
SVRG epoch 0002 / 0010, subset 0011 / 0012
SVRG epoch 0002 / 0010, subset 0012 / 0012
SVRG epoch 0003 / 0010, subset 0001 / 0012
SVRG epoch 0003 / 0010, subset 0002 / 0012
SVRG epoch 0003 / 0010, subset 0003 / 0012
SVRG epoch 0003 / 0010, subset 0004 / 0012
SVRG epoch 0003 / 0010, subset 0005 / 0012
SVRG epoch 0003 / 0010, subset 0006 / 0012
SVRG epoch 0003 / 0010, subset 0007 / 0012
SVRG epoch 0003 / 0010, subset 0008 / 0012
SVRG epoch 0003 / 0010, subset 0009 / 0012
SVRG epoch 0003 / 0010, subset 0010 / 0012
SVRG epoch 0003 / 0010, subset 0011 / 0012
SVRG epoch 0003 / 0010, subset 0012 / 0012
SVRG epoch 0004 / 0010, subset 0001 / 0012
SVRG epoch 0004 / 0010, subset 0002 / 0012
SVRG epoch 0004 / 0010, subset 0003 / 0012
SVRG epoch 0004 / 0010, subset 0004 / 0012
SVRG epoch 0004 / 0010, subset 0005 / 0012
SVRG epoch 0004 / 0010, subset 0006 / 0012
SVRG epoch 0004 / 0010, subset 0007 / 0012
SVRG epoch 0004 / 0010, subset 0008 / 0012
SVRG epoch 0004 / 0010, subset 0009 / 0012
SVRG epoch 0004 / 0010, subset 0010 / 0012
SVRG epoch 0004 / 0010, subset 0011 / 0012
SVRG epoch 0004 / 0010, subset 0012 / 0012
SVRG epoch 0005 / 0010, subset 0001 / 0012
SVRG epoch 0005 / 0010, subset 0002 / 0012
SVRG epoch 0005 / 0010, subset 0003 / 0012
SVRG epoch 0005 / 0010, subset 0004 / 0012
SVRG epoch 0005 / 0010, subset 0005 / 0012
SVRG epoch 0005 / 0010, subset 0006 / 0012
SVRG epoch 0005 / 0010, subset 0007 / 0012
SVRG epoch 0005 / 0010, subset 0008 / 0012
SVRG epoch 0005 / 0010, subset 0009 / 0012
SVRG epoch 0005 / 0010, subset 0010 / 0012
SVRG epoch 0005 / 0010, subset 0011 / 0012
SVRG epoch 0005 / 0010, subset 0012 / 0012
SVRG epoch 0006 / 0010, subset 0001 / 0012
SVRG epoch 0006 / 0010, subset 0002 / 0012
SVRG epoch 0006 / 0010, subset 0003 / 0012
SVRG epoch 0006 / 0010, subset 0004 / 0012
SVRG epoch 0006 / 0010, subset 0005 / 0012
SVRG epoch 0006 / 0010, subset 0006 / 0012
SVRG epoch 0006 / 0010, subset 0007 / 0012
SVRG epoch 0006 / 0010, subset 0008 / 0012
SVRG epoch 0006 / 0010, subset 0009 / 0012
SVRG epoch 0006 / 0010, subset 0010 / 0012
SVRG epoch 0006 / 0010, subset 0011 / 0012
SVRG epoch 0006 / 0010, subset 0012 / 0012
SVRG epoch 0007 / 0010, subset 0001 / 0012
SVRG epoch 0007 / 0010, subset 0002 / 0012
SVRG epoch 0007 / 0010, subset 0003 / 0012
SVRG epoch 0007 / 0010, subset 0004 / 0012
SVRG epoch 0007 / 0010, subset 0005 / 0012
SVRG epoch 0007 / 0010, subset 0006 / 0012
SVRG epoch 0007 / 0010, subset 0007 / 0012
SVRG epoch 0007 / 0010, subset 0008 / 0012
SVRG epoch 0007 / 0010, subset 0009 / 0012
SVRG epoch 0007 / 0010, subset 0010 / 0012
SVRG epoch 0007 / 0010, subset 0011 / 0012
SVRG epoch 0007 / 0010, subset 0012 / 0012
SVRG epoch 0008 / 0010, subset 0001 / 0012
SVRG epoch 0008 / 0010, subset 0002 / 0012
SVRG epoch 0008 / 0010, subset 0003 / 0012
SVRG epoch 0008 / 0010, subset 0004 / 0012
SVRG epoch 0008 / 0010, subset 0005 / 0012
SVRG epoch 0008 / 0010, subset 0006 / 0012
SVRG epoch 0008 / 0010, subset 0007 / 0012
SVRG epoch 0008 / 0010, subset 0008 / 0012
SVRG epoch 0008 / 0010, subset 0009 / 0012
SVRG epoch 0008 / 0010, subset 0010 / 0012
SVRG epoch 0008 / 0010, subset 0011 / 0012
SVRG epoch 0008 / 0010, subset 0012 / 0012
SVRG epoch 0009 / 0010, subset 0001 / 0012
SVRG epoch 0009 / 0010, subset 0002 / 0012
SVRG epoch 0009 / 0010, subset 0003 / 0012
SVRG epoch 0009 / 0010, subset 0004 / 0012
SVRG epoch 0009 / 0010, subset 0005 / 0012
SVRG epoch 0009 / 0010, subset 0006 / 0012
SVRG epoch 0009 / 0010, subset 0007 / 0012
SVRG epoch 0009 / 0010, subset 0008 / 0012
SVRG epoch 0009 / 0010, subset 0009 / 0012
SVRG epoch 0009 / 0010, subset 0010 / 0012
SVRG epoch 0009 / 0010, subset 0011 / 0012
SVRG epoch 0009 / 0010, subset 0012 / 0012
SVRG epoch 0010 / 0010, subset 0001 / 0012
SVRG epoch 0010 / 0010, subset 0002 / 0012
SVRG epoch 0010 / 0010, subset 0003 / 0012
SVRG epoch 0010 / 0010, subset 0004 / 0012
SVRG epoch 0010 / 0010, subset 0005 / 0012
SVRG epoch 0010 / 0010, subset 0006 / 0012
SVRG epoch 0010 / 0010, subset 0007 / 0012
SVRG epoch 0010 / 0010, subset 0008 / 0012
SVRG epoch 0010 / 0010, subset 0009 / 0012
SVRG epoch 0010 / 0010, subset 0010 / 0012
SVRG epoch 0010 / 0010, subset 0011 / 0012
SVRG epoch 0010 / 0010, subset 0012 / 0012