From d8a168f8f70b6731479b10c6e270a17e17979091 Mon Sep 17 00:00:00 2001 From: Niklas Hackelberg Date: Sun, 1 Nov 2026 18:59:03 +0100 Subject: [PATCH] Add RowBlockOp for vertically stacked blocks of rows RowBlockOp(ops, rows) stacks operators vertically, each filling its rows of the result (contiguous or interleaved); rowblocks returns the blocks for block-iterative solvers such as OS-SART. Its normal operator is the sum of the blocks' normal operators. Constructors for contiguous blocks and for splitting a matrix into interleaved, contiguous or random blocks. Version 2.5.0. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- docs/make.jl | 1 + docs/src/literate/tutorials/rowblock.jl | 34 ++++ docs/src/references.md | 2 + src/LinearOperatorCollection.jl | 1 + src/RowBlockOp.jl | 202 ++++++++++++++++++++++++ test/testAdapt.jl | 15 ++ test/testOperators.jl | 41 +++++ 8 files changed, 297 insertions(+), 1 deletion(-) create mode 100644 docs/src/literate/tutorials/rowblock.jl create mode 100644 src/RowBlockOp.jl diff --git a/Project.toml b/Project.toml index 76e9d4f..824dc4d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "LinearOperatorCollection" uuid = "a4a2c56f-fead-462a-a3ab-85921a5f2575" authors = ["Tobias Knopp and contributors"] -version = "2.4.2" +version = "2.5.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/docs/make.jl b/docs/make.jl index aa4e219..6b78630 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -37,6 +37,7 @@ makedocs( "Weighting Operator" => "generated/tutorials/weighting.md", "FFT Operator" => "generated/tutorials/fft.md", "Diagonal Operator" => "generated/tutorials/diagonal.md", + "Row-Block Operator" => "generated/tutorials/rowblock.md", "Gradient Operator" => "generated/tutorials/gradient.md", #"Sampling Operator" => "generated/tutorials/sampling.md", #"NFFT Operator" => "generated/tutorials/nfft.md", diff --git a/docs/src/literate/tutorials/rowblock.jl b/docs/src/literate/tutorials/rowblock.jl new file mode 100644 index 0000000..063d04e --- /dev/null +++ b/docs/src/literate/tutorials/rowblock.jl @@ -0,0 +1,34 @@ +# # Row-Block Operator +include("../../util.jl") #hide +# This operator stacks operators vertically: all of them act on the same input, and each one fills its own rows of +# the output. The rows of a block need not be contiguous. Block-iterative and ordered-subset methods (e.g. SART, +# OS-EM) use such a split of a system matrix into blocks of rows; as a whole the operator works like any other. +# +# A matrix split into three blocks of interleaved rows (every third row), `:contiguous` and `:random` are the other +# modes: +A = rand(12, N * N) +rop = RowBlockOp(A, 3) +typeof(rop) + +# The blocks with the indices of their rows: +[(size(op), rows) for (op, rows) in rowblocks(rop)] + +# Applied as a whole, the operator equals the matrix, and so do its adjoint and its normal operator: +x = vec(image) +rop * x ≈ A * x + +#- +adjoint(rop) * (A * x) ≈ adjoint(A) * (A * x) + +# Any operators with the same number of columns can be blocks, with their rows given explicitly. Here two weighting +# operators fill the even and the odd entries of the output: +ops = [WeightingOp(fill(1.0, N * N)), WeightingOp(fill(2.0, N * N))] +rop2 = RowBlockOp(ops, [1:2:2N*N, 2:2:2N*N]) +(rop2 * x)[1:4] ≈ [x[1], 2x[1], x[2], 2x[2]] + +# Blocks with contiguous rows work on views of the data; for other rows the data of a block is gathered into a +# buffer. The products of the blocks run in parallel with a scheduler of OhMyThreads.jl, as for the +# [`DiagOp`](@ref): +using OhMyThreads +rop_serial = RowBlockOp(A, 3; scheduler = SerialScheduler()) +typeof(rop_serial) diff --git a/docs/src/references.md b/docs/src/references.md index f525825..484f4b6 100644 --- a/docs/src/references.md +++ b/docs/src/references.md @@ -2,6 +2,8 @@ SamplingOp WeightingOp DiagOp +RowBlockOp +rowblocks GradientOp ProdOp NormalOp diff --git a/src/LinearOperatorCollection.jl b/src/LinearOperatorCollection.jl index 2333f7e..18fac52 100644 --- a/src/LinearOperatorCollection.jl +++ b/src/LinearOperatorCollection.jl @@ -109,6 +109,7 @@ include("GradientOp.jl") include("SamplingOp.jl") include("NormalOp.jl") include("DiagOp.jl") +include("RowBlockOp.jl") function promote_storage_types(A, B) A_type = storage_type(A) diff --git a/src/RowBlockOp.jl b/src/RowBlockOp.jl new file mode 100644 index 0000000..ef3d802 --- /dev/null +++ b/src/RowBlockOp.jl @@ -0,0 +1,202 @@ +export RowBlockOp, rowblocks + +mutable struct RowBlockOp{T, vecT, vecO, R, B, S} <: AbstractLinearOperator{T} + const nrow :: Int + const ncol :: Int + const symmetric :: Bool + const hermitian :: Bool + const prod! :: Function + const tprod! :: Function + const ctprod! :: Function + nprod :: Int + ntprod :: Int + nctprod :: Int + Mv :: vecT + Mtu :: vecT + const ops :: vecO + const rows :: R + const buffers :: B + const scheduler :: S +end + +LinearOperators.storage_type(::RowBlockOp{T, vecT}) where {T, vecT} = vecT + +""" + RowBlockOp(ops, rows; scheduler = DynamicScheduler()) + RowBlockOp(ops; scheduler = DynamicScheduler()) + RowBlockOp(A::AbstractMatrix, nblocks; mode = :interleaved, rng = Random.default_rng(), scheduler = DynamicScheduler()) + +create an operator out of the operators (or `Array`s) in `ops` stacked vertically, `[ops[1]; ops[2]; …]`, whose +rows are scattered over the rows of the result: the rows of `ops[k]` are the rows `rows[k]` of the operator. All +operators have the same number of columns, and `rows` is a partition of `1:nrow` into index vectors or ranges, one +per operator, of the length of that operator's number of rows. Without `rows`, the blocks are contiguous, in the +order of `ops`. + +Block-iterative solvers (e.g. ordered-subset methods) access the blocks with [`rowblocks`](@ref). Applied as a +whole, the operator multiplies every block and writes its result into its rows, and its adjoint sums the adjoints +of the blocks applied to their rows. The products of the blocks run in parallel with `scheduler` (OhMyThreads.jl); +the adjoint sums the blocks one after the other. + +Blocks with contiguous rows (`UnitRange`s) work on views of the data; other rows are gathered into and scattered +from a buffer per block (with index vectors on the device of the blocks, see `Adapt.adapt`). Reordering the data +such that every block is contiguous avoids these copies. + +The third form splits the rows of a matrix into `nblocks` blocks, as copies of the rows: `:interleaved` takes every +`nblocks`-th row (`k:nblocks:m`), `:contiguous` consecutive rows, `:random` a random partition drawn with `rng`. +""" +function RowBlockOp(ops, rows; scheduler = DynamicScheduler()) + length(ops) == length(rows) || throw(ArgumentError("$(length(ops)) operators but $(length(rows)) row sets")) + isempty(ops) && throw(ArgumentError("RowBlockOp needs at least one block")) + ncol = size(first(ops), 2) + all(op -> size(op, 2) == ncol, ops) || throw(DimensionMismatch("all blocks need the same number of columns")) + for (op, r) in zip(ops, rows) + length(r) == size(op, 1) || throw(DimensionMismatch("a block of $(size(op, 1)) rows has $(length(r)) row indices")) + end + nrow = sum(op -> size(op, 1), ops) + sort!(reduce(vcat, [collect(r) for r in rows])) == 1:nrow || + throw(ArgumentError("the rows of the blocks must be a partition of 1:$nrow")) + + S = LinearOperators.storage_type(first(ops)) + for op in ops + S = promote_type(S, LinearOperators.storage_type(op)) + end + isconcretetype(S) || throw(LinearOperatorException("Storage types cannot be promoted to a concrete type")) + + # a buffer for every block whose rows are not contiguous + buffers = [r isa AbstractUnitRange ? nothing : S(undef, length(r)) for r in rows] + + return RowBlockOp{eltype(first(ops)), S, typeof(ops), typeof(rows), typeof(buffers), typeof(scheduler)}( + nrow, ncol, false, false, + (res, x) -> rowBlockProd!(res, x, ops, rows, buffers, scheduler), + (res, y) -> rowBlockAdjProd!(res, y, transpose, ops, rows, buffers), + (res, y) -> rowBlockAdjProd!(res, y, adjoint, ops, rows, buffers), + 0, 0, 0, S(undef, 0), S(undef, 0), + ops, rows, buffers, scheduler) +end + +function RowBlockOp(ops; kwargs...) + stops = cumsum([size(op, 1) for op in ops]) + starts = stops .- [size(op, 1) for op in ops] .+ 1 + return RowBlockOp(ops, [start:stop for (start, stop) in zip(starts, stops)]; kwargs...) +end + +function RowBlockOp(A::AbstractMatrix, nblocks::Integer; mode::Symbol = :interleaved, rng = Random.default_rng(), kwargs...) + m = size(A, 1) + nblocks = clamp(nblocks, 1, m) + rows = if mode == :interleaved + [k:nblocks:m for k in 1:nblocks] + elseif mode == :contiguous + [(div((k - 1) * m, nblocks) + 1):div(k * m, nblocks) for k in 1:nblocks] + elseif mode == :random + perm = randperm(rng, m) + [sort(perm[k:nblocks:m]) for k in 1:nblocks] + else + throw(ArgumentError("mode must be :interleaved, :contiguous or :random, got :$mode")) + end + return RowBlockOp([A[r, :] for r in rows], rows; kwargs...) +end + +""" + rowblocks(op::RowBlockOp) + +the blocks of `op` as pairs `(operator, rows)`. +""" +rowblocks(op::RowBlockOp) = collect(zip(op.ops, op.rows)) + +# the rows r of a data vector y: a view for contiguous rows, else gathered into the buffer +@inline block_data(y, r::AbstractUnitRange, ::Nothing) = view(y, r) +@inline block_data(y, r, buffer) = (buffer .= view(y, r); buffer) + +function rowBlockProd!(y, x, ops, rows, buffers, scheduler) + @tasks for k in 1:length(ops) + @set scheduler = scheduler + r, buffer = rows[k], buffers[k] + if isnothing(buffer) + mul!(view(y, r), ops[k], x) + else + mul!(buffer, ops[k], x) + view(y, r) .= buffer + end + end + return y +end + +function rowBlockAdjProd!(x, y, adj, ops, rows, buffers) + for k in 1:length(ops) + yk = block_data(y, rows[k], buffers[k]) + if k == 1 + mul!(x, adj(ops[k]), yk) + else + mul!(x, adj(ops[k]), yk, true, true) + end + end + return x +end + +# Index vectors move to the device of the adapted blocks and stay integers (adapting them with `to` +# itself could convert them, e.g. with `to = CuArray{Float32}`). +function Adapt.adapt_structure(to, op::RowBlockOp) + ops = map(child -> Adapt.adapt(to, child), op.ops) + S = LinearOperators.storage_type(first(ops)) + rows = map(r -> r isa AbstractUnitRange ? r : copyto!(similar(S(undef, 0), eltype(r), length(r)), collect(r)), op.rows) + return RowBlockOp(ops, rows; scheduler = op.scheduler) +end + +### Normal Matrix Code ### + +mutable struct RowBlockNormalOp{T, vecT, V} <: AbstractLinearOperator{T} + const nrow :: Int + const ncol :: Int + const symmetric :: Bool + const hermitian :: Bool + const prod! :: Function + const tprod! :: Nothing + const ctprod! :: Nothing + nprod :: Int + ntprod :: Int + nctprod :: Int + Mv :: vecT + Mtu :: vecT + normalOps :: V + tmp :: vecT +end + +LinearOperators.storage_type(::RowBlockNormalOp{T, vecT}) where {T, vecT} = vecT + +function RowBlockNormalOp(normalOps, tmp::vecT) where {vecT} + n = size(first(normalOps), 1) + return RowBlockNormalOp{eltype(first(normalOps)), vecT, typeof(normalOps)}(n, n, false, false, + (res, x) -> rowBlockNormalProd!(res, normalOps, tmp, x), nothing, nothing, + 0, 0, 0, vecT(undef, 0), vecT(undef, 0), normalOps, tmp) +end + +function rowBlockNormalProd!(y, normalOps, tmp, x) + mul!(y, normalOps[1], x) + for k in 2:length(normalOps) + mul!(tmp, normalOps[k], x) + y .+= tmp + end + return y +end + +function Adapt.adapt_structure(to, op::RowBlockNormalOp) + return RowBlockNormalOp(map(child -> Adapt.adapt(to, child), op.normalOps), Adapt.adapt(to, op.tmp)) +end + +""" + normalOperator(op::RowBlockOp, W = nothing; kwargs...) + +the normal operator `adjoint(op) * W * op` as the sum of the normal operators of the blocks (each with the part of +the weights `W` that belongs to its rows), so that specialized normal operators of the blocks are used. +""" +function LinearOperatorCollection.normalOperator(op::RowBlockOp, W = nothing; kwargs...) + S = LinearOperators.storage_type(op) + if isnothing(W) + normalOps = [normalOperator(o; kwargs...) for o in op.ops] + else + T = promote_type(eltype(op), eltype(W)) + weights = W * fill!(S(undef, op.nrow), one(eltype(op))) + normalOps = [normalOperator(o, WeightingOp(T; weights = T.(weights[r])); kwargs...) for (o, r) in zip(op.ops, op.rows)] + end + return RowBlockNormalOp(normalOps, S(undef, op.ncol)) +end diff --git a/test/testAdapt.jl b/test/testAdapt.jl index bf6fdeb..e56a200 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -191,6 +191,21 @@ end end end + @testset "RowBlockOp and its normal operator" begin + A = rand(Float32, 3N, N) + for source in (RowBlockOp(A, 3), RowBlockOp(A, 3; mode = :contiguous), RowBlockOp(A, 3; mode = :random)) + adapted = Adapt.adapt(array_type, source) + check_adapted_storage(adapted, array_type) + # index vectors move to the device and stay integers + @test all(r isa AbstractUnitRange || (r isa array_type && eltype(r) <: Integer) for r in adapted.rows) + x = rand(Float32, N) + check_adapted_mul(source, adapted, x, array_type) + check_adapted_mul(adjoint(source), adjoint(adapted), rand(Float32, 3N), array_type) + check_adapted_mul(normalOperator(source), normalOperator(adapted), x, array_type) + check_adapted_mul(normalOperator(source), Adapt.adapt(array_type, normalOperator(source)), x, array_type) + end + end + @testset "DiagOp preserves tuple ops" begin source = DiagOp((WeightingOp(rand(Float32, N)), WeightingOp(rand(Float32, N)))) source_ops_snapshot = map(block -> copy(block.weights), source.ops) diff --git a/test/testOperators.jl b/test/testOperators.jl index bd15ca8..505c7a2 100644 --- a/test/testOperators.jl +++ b/test/testOperators.jl @@ -426,6 +426,43 @@ function testDiagOp(N=32,K=2;arrayType = Array,scheduler = DynamicScheduler()) end +function testRowBlockOp(M=24, N=16, K=3; arrayType = Array, scheduler = DynamicScheduler()) + A = rand(ComplexF64, M, N) + x = rand(ComplexF64, N) + y = rand(ComplexF64, M) + Ad, xd, yd = arrayType(A), arrayType(x), arrayType(y) + + ops = Dict( + "interleaved" => RowBlockOp(Ad, K; scheduler), + "contiguous" => RowBlockOp(Ad, K; mode = :contiguous, scheduler), + "random" => RowBlockOp(Ad, K; mode = :random, scheduler), + "given blocks" => RowBlockOp([Ad[1:10, :], Ad[11:M, :]]; scheduler)) + for (name, op) in ops + @testset "RowBlockOp $name" begin + check_storage(op, arrayType) + @test size(op) == (M, N) + @test Array(op * xd) ≈ A * x + @test Array(transpose(op) * yd) ≈ transpose(A) * y + @test Array(adjoint(op) * yd) ≈ adjoint(A) * y + @test Array(normalOperator(op) * xd) ≈ adjoint(A) * A * x + w = rand(M) + W = WeightingOp(arrayType(ComplexF64.(w))) + @test Array(normalOperator(op, W) * xd) ≈ adjoint(A) * (w .* (A * x)) + # the blocks are the rows of A + for (block, rows) in rowblocks(op) + @test Array(block) ≈ A[collect(rows), :] + end + end + end + + # interleaved rows of the matrix + @test [collect(r) for (_, r) in rowblocks(RowBlockOp(A, K))] == [collect(k:K:M) for k in 1:K] + # errors: rows that are no partition, blocks with different numbers of columns + @test_throws ArgumentError RowBlockOp([A[1:12, :], A[13:M, :]], [1:12, 12:23]) + @test_throws DimensionMismatch RowBlockOp([A[1:12, :], A[13:M, 1:N-1]]) + @test_throws ArgumentError RowBlockOp(A, K; mode = :spiral) +end + function testRadonOp(N=32;arrayType = Array) x = arrayType(rand(N, N)) xop = vec(x) @@ -481,6 +518,10 @@ end for scheduler in [DynamicScheduler(), StaticScheduler(), SerialScheduler()] @testset testDiagOp(;arrayType, scheduler = scheduler) end + @info "test RowBlockOp: $arrayType" + for scheduler in [DynamicScheduler(), SerialScheduler()] + @testset testRowBlockOp(;arrayType, scheduler = scheduler) + end @info "test RadonOp: $arrayType" @testset testRadonOp(;arrayType) # Stackoverflow for kernelabstraction end