CTOperators: subset_blocks and the OS-SART guide

docs/Project.toml

@@ -18,8 +18,13 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
 CTOperators = {path = ".."}
+# until the releases of LinearOperatorCollection 2.5 (RowBlockOp) and RegularizedLeastSquares 0.17
+# (OSSART): their development branches
+LinearOperatorCollection = {url = "https://github.com/JuliaImageRecon/LinearOperatorCollection.jl", rev = "nh/blockrows"}
+RegularizedLeastSquares = {url = "https://github.com/JuliaImageRecon/RegularizedLeastSquares.jl", rev = "nh/ossart"}
 
 [compat]
 Documenter = "1"
-LinearOperatorCollection = "2.4"
+LinearOperatorCollection = "2.5"
+RegularizedLeastSquares = "0.17"
 Literate = "2"

docs/literate/howto/subsets.jl

@@ -2,11 +2,13 @@
 #
 # Algebraic methods such as SIRT and SART, and the ordered-subset variants of statistical methods,
 # normalize each update by the sums of the rows and columns of the system matrix, and many of them
-# cycle through subsets of the views. This guide shows the two building blocks, `rowsums`/`colsums`
-# and `subset`, with an OS-SART reconstruction of a 2D phantom.
+# cycle through subsets of the views. This guide shows the building blocks, `rowsums`/`colsums` and
+# `subset_blocks`, and an OS-SART reconstruction of a 2D phantom with the solver `OSSART` of
+# [RegularizedLeastSquares.jl](https://github.com/JuliaImageRecon/RegularizedLeastSquares.jl).
 
-using CTOperators, ImagePhantoms, CairoMakie, LinearAlgebra
+using CTOperators, ImagePhantoms, RegularizedLeastSquares, LinearOperatorCollection, CairoMakie
 using ImageGeoms: axis
+using LinearAlgebra: norm
 
 # A 2D Shepp-Logan phantom of 128² pixels and 180 parallel-beam views:
 
@@ -33,48 +35,39 @@ size(r), size(c)
 #
 # `subset(op, views)` is the projector of some of the views, sharing the image geometry.
 # Ordered-subset methods interleave the views so that consecutive subsets see the object from
-# different directions. The data of a subset are its views, the last dimension of the `(nu, nv,
-# nviews)` data:
+# different directions. `subset_blocks(op, n)` returns the projectors of `n` interleaved subsets and,
+# for each, the indices of its data in `vec(y)`. A `RowBlockOp` of
+# [LinearOperatorCollection.jl](https://github.com/JuliaImageRecon/LinearOperatorCollection.jl) stacks
+# them into one operator that equals `op`, but knows its blocks; the solver below takes it as the
+# keyword `blocks`:
 
-nsubsets = 10
-views = [s:nsubsets:nviews(geometry) for s in 1:nsubsets]
-ops = [subset(op, v) for v in views]
-ys = [y[:, :, v] for v in views];
-
-# Create the subset projectors and their normalizations once; they are reused in every iteration.
-# Rays that miss the image and voxels that no ray hits have sums of zero; a small lower bound avoids
-# dividing by zero.
-
-rs = [max.(rowsums(o), 1.0f-6) for o in ops]
-cs = [max.(colsums(o), 1.0f-6) for o in ops];
+blocks = RowBlockOp(subset_blocks(op, 10)...)
+length(rowblocks(blocks)), size(blocks) == size(op)
 
 # ## OS-SART
 #
-# One SART update per subset: back project the normalized residual of the subset and divide by the
-# sensitivity of the subset; then clip negative values. With one subset this is SIRT.
+# `OSSART` makes one update per block: it back projects the residual of the block, normalized by the
+# row sums (the length of each ray in the image), and divides by the column sums (the sensitivity of
+# each pixel to the block). Rays that miss the image and pixels that no ray hits have sums of zero
+# and get no update. `PositiveRegularization` clips negative values after each block. Without
+# `blocks`, the projector is a single block, and this is SIRT. A callback records the error of the start image
+# and after each pass:
 
-function os_sart(ops, ys, rs, cs, iterations; relaxation=1.0f0)
-  x = zeros(Float32, imagesize(ops[1]))
-  errors = Float64[]
-  for _ in 1:iterations
-    for (o, yk, rk, ck) in zip(ops, ys, rs, cs)
-      x .+= relaxation .* (adjoint(o) * ((yk .- o * x) ./ rk)) ./ ck
-      x .= max.(x, 0)
-    end
-    push!(errors, norm(x - x_true) / norm(x_true))
-  end
-  x, errors
+function reconstruct(; iterations=10, kwargs...)
+  error = CompareSolutionCallback(vec(x_true), (ref, x) -> norm(x - ref) / norm(ref))
+  solver = createLinearSolver(OSSART, op; reg=PositiveRegularization(), iterations, kwargs...)
+  x = reshape(solve!(solver, vec(y); callbacks=error), n, n)
+  x, error.results
 end
-
-x_os, err_os = os_sart(ops, ys, rs, cs, 10)
-x_sirt, err_sirt = os_sart([op], [y], [max.(r, 1.0f-6)], [max.(c, 1.0f-6)], 10);
+x_os, err_os = reconstruct(; blocks)
+x_sirt, err_sirt = reconstruct();
 
 # Ten subsets reach in a few passes through the data what SIRT reaches only after many more:
 
 fig = Figure(size=(900, 300))
 ax = Axis(fig[1, 1]; xlabel="pass through all views", ylabel="relative error", title="convergence")
-lines!(ax, err_sirt; label="SIRT")
-lines!(ax, err_os; label="OS-SART, 10 subsets")
+lines!(ax, 0:length(err_sirt)-1, err_sirt; label="SIRT")
+lines!(ax, 0:length(err_os)-1, err_os; label="OS-SART, 10 subsets")
 axislegend(ax)
 for (j, (name, img)) in enumerate((("SIRT", x_sirt), ("OS-SART", x_os)))
   heatmap(fig[1, j+1], img; colormap=:grays, colorrange=(0, 0.4),
@@ -84,5 +77,7 @@ save("os_sart.png", fig) #hide
 nothing #hide
 # ![OS-SART and SIRT](os_sart.png)
 
-# The same pattern serves other ordered-subset methods, for example OS-EM for emission data, which
-# divides by `colsums` of the subset as well.
+# On a GPU, the blocks of interleaved views gather their data with index vectors on the device;
+# `subset_blocks(op, n; mode=:contiguous)` takes consecutive views instead, whose data needs no
+# copies. The same blocks serve other ordered-subset methods, for example OS-EM for emission data,
+# which divides by the column sums of the block as well.

src/CTOperators.jl

@@ -15,7 +15,7 @@ export RayList, ConeBeam, ArcConeBeam, ParallelBeam
 export circular_cone_beam, helical_cone_beam, parallel_beam
 export datasize, nrays, ray, raytype, nviews, detectorsize, subset, projection_geometry
 export AbstractProjector, AdjointCTOperator, SiddonProjector, JosephProjector, LaunchConfig
-export geometry, image_geometry, imagesize, rowsums, colsums
+export geometry, image_geometry, imagesize, rowsums, colsums, subset_blocks
 
 """
     AbstractCTOperator{T}

src/operators/operators.jl

@@ -146,6 +146,44 @@ indices), for ordered-subset methods. It shares the image geometry and options w
 subset(op::AbstractProjector{T,<:AbstractDetectorGeometry}, views) where {T} =
   similar_projector(op, subset(geometry(op), views))
 
+"""
+    subset_blocks(op::AbstractProjector, nsubsets; mode=:interleaved) -> (ops, rows)
+
+The projectors of `nsubsets` subsets of the views of `op` (detector geometries) and, for each, the
+indices of its data in `vec(op * x)`: the blocks of rows of ordered-subset methods. With
+LinearOperatorCollection.jl, `RowBlockOp(subset_blocks(op, 10)...)` stacks them into one operator,
+e.g. for the keyword `blocks` of the OS-SART solver of RegularizedLeastSquares.jl:
+`createLinearSolver(OSSART, op; blocks=RowBlockOp(subset_blocks(op, 10)...))`.
+
+`mode=:interleaved` takes every `nsubsets`-th view, so that consecutive subsets see the object from
+different directions; `:contiguous` takes consecutive views, whose data indices are ranges (no
+copies of the data, cheapest on GPUs). Index vectors live on the device of the projector.
+"""
+function subset_blocks(op::AbstractProjector{T,<:AbstractDetectorGeometry}, nsubsets::Integer;
+                       mode::Symbol=:interleaved) where {T}
+  nv = nviews(geometry(op))
+  n = clamp(nsubsets, 1, nv)
+  views = if mode == :interleaved
+    [k:n:nv for k in 1:n]
+  elseif mode == :contiguous
+    [(div((k - 1) * nv, n) + 1):div(k * nv, n) for k in 1:n]
+  else
+    throw(ArgumentError("mode must be :interleaved or :contiguous, got :$mode"))
+  end
+  cells = prod(detectorsize(geometry(op)))
+  [subset(op, v) for v in views], [data_rows(op, v, cells) for v in views]
+end
+
+# the indices of the data of the views v in the flat data of op: a range for consecutive views,
+# else an index vector on the device of op
+function data_rows(op, v, cells)
+  if step(v) == 1
+    return ((first(v) - 1) * cells + 1):(last(v) * cells)
+  end
+  idx = [(k - 1) * cells + j for k in v for j in 1:cells]
+  copyto!(KA.allocate(KA.get_backend(op), Int, length(idx)), idx)
+end
+
 Adapt.adapt_structure(to, op::AbstractProjector) = similar_projector(op, adapt(to, geometry(op)))
 
 """

test/operators/operators.jl

@@ -8,3 +8,26 @@
   @test isapprox(mul!(copy(y), op, x, 2.0, -0.5), 2 * (op * x) - 0.5 * y)
   @test isapprox(mul!(vec(copy(x)), adjoint(op), vec(y), 1.0, 1.0), vec(adjoint(op) * y + x))
 end
+
+# The blocks of subset_blocks reproduce the projector: every block's projection is the full
+# projection in its rows, the rows partition the data, and the back projections of all blocks add
+# up to the full back projection; on every device (index vectors on the device).
+@testitem "subset_blocks" tags = [:operators, :gpu] setup = [Backends] begin
+  image = ImageGeom(; dims=(6, 6, 4))
+  g = circular_cone_beam(; angles=range(0, 2pi, 8)[1:7], dso=25, dsd=50, nu=8, nv=5, du=1.0)
+  for arrayType in arrayTypes, mode in (:interleaved, :contiguous), n in (1, 3, 7, 20)
+    op = JosephProjector(image, g; eltype=Float64, storage=arrayType)
+    x, y = rand(6, 6, 4), rand(datasize(op)...)
+    full, back = vec(Array(op * adapt(arrayType, x))), zeros(6, 6, 4)
+    ops, rows = subset_blocks(op, n; mode)
+    @test length(ops) == length(rows) == min(n, 7)
+    @test sort(reduce(vcat, [collect(Array(r)) for r in rows])) == 1:length(full)
+    for (o, r) in zip(ops, rows)
+      @test vec(Array(o * adapt(arrayType, x))) ≈ full[collect(Array(r))]
+      back .+= Array(adjoint(o) * adapt(arrayType, reshape(vec(y)[collect(Array(r))], datasize(o))))
+    end
+    @test back ≈ Array(adjoint(op) * adapt(arrayType, y))
+    mode == :contiguous && @test all(r -> r isa AbstractUnitRange, rows)
+  end
+  @test_throws ArgumentError subset_blocks(JosephProjector(image, g), 2; mode=:random)
+end