Batch operators

The batch operator is particularly useful in scenarios where the same operation needs to be applied across multiple dimensions of data simultaneously. For instance, in machine learning, it can be used to process batches of input data through a neural network, ensuring that each input in the batch undergoes the same transformations. Similarly, in image processing, the batch operator can apply filters or transformations to a set of images in parallel, significantly improving efficiency. Technically it is the optimized special case of DCAT, in which the diagonally concatenated operators are all the same. This operator is essential for tasks that require uniform processing across large datasets, enabling scalable and consistent operations.

An even more advanced version is the spreading batch operator, which allows applying different operators over one or more "spreading" batch dimensions, while the rest of the batch dimensions behaves as a "normal" batch dimension. For example, in medical imaging (e.g. CT and MRI) it is common to acquire 3D-images of the patient, and these volumetric aquisitions are repeated over a specified time. This results in a 4-dimensional data set for which, let's say, we want to apply different image processing operators over the 2D slices of the volumetric "time frames", but we want to do the same procedure along the time dimension.

AbstractOperators.BatchOp — Type

BatchOp(
	operator::AbstractOperators.AbstractOperator,
	batch_dims::NTuple{N,Int},
	batch_dim_mask::Union{NTuple{M,Symbol}, Pair{NTuple{M1,Symbol},NTuple{M2,Symbol}}};
	threaded::Bool=nthreads() > 1,
)

Creates a "simple" BatchOp from an AbstractOperator. The BatchOp can be used to apply the operator to an array over the selected batch dimensions.

Arguments

operator::AbstractOperators.AbstractOperator: The operator to be batched.
batch_dims::NTuple{N,Int}: The size of the batch dimensions. This tuple should contain the sizes of the batch dimensions in the order they appear in the operator's domain and codomain.
batch_dim_mask::Union{NTuple{M,Symbol}, Pair{NTuple{M1,Symbol},NTuple{M2,Symbol}}}: The mask specifying which dimensions are batch dimensions. The symbols can be :b for batch dimensions or :_ for dimensions on which the operator acts. If a pair is provided, the first tuple specifies the domain mask and the second tuple specifies the codomain mask. When omitted, the batch dimensions are assumed to proceed the operator's domain and codomain dimensions.
threaded::Bool: If true, the operator will execute in parallel over the batch dimensions. Default is nthreads() > 1.

Examples

julia> op = DiagOp([1.0im, 1.0im])
╲  ℂ^2 -> ℂ^2

julia> batch_op = BatchOp(op, (3, 4), (:_, :b, :b))
⟳╲  ℂ^(2, 3, 4) -> ℂ^(2, 3, 4)

julia> batch_op = BatchOp(op, (3, 4)) # we don't need to specify the mask as the default is to assume the batch dimensions are at the end
⟳╲  ℂ^(2, 3, 4) -> ℂ^(2, 3, 4)

julia> x = rand(ComplexF64, 2, 3, 4);

julia> y = similar(x);

julia> [mul!(@view(y[:, i, j]), op, @view(x[:, i, j])) for i in 1:3, j in 1:4];

julia> y == batch_op*x
true

julia> op = Variation(3, 4, 5)
Ʋ  ℝ^(3, 4, 5) -> ℝ^(60, 3)

julia> batch_op = BatchOp(op, (2, 6), (:b, :_, :_, :_, :b) => (:b, :_, :b, :_))
⟳Ʋ  ℝ^(2, 3, 4, 5, 6) -> ℝ^(2, 60, 6, 3)

source

function BatchOp(
	operators::Array{<:AbstractOperators.AbstractOperator},
	batch_size::NTuple{N,Int},
	batch_dim_mask::Union{NTuple{M,Symbol}, Pair{NTuple{M1,Symbol},NTuple{M2,Symbol}};
	threaded::Bool=nthreads() > 1,
	threading_strategy::Symbol=ThreadingStrategy.AUTO,
)

Creates a "spreading" BatchOp from an array of AbstractOperators. The BatchOp can be used to apply the operators to an array over the selected batch dimensions. The "spreading" dimensions will be used to select the operator to apply, and the batch dimensions will be used to repeat the operator application over the batch dimensions.

Arguments

operators::Array{<:AbstractOperators.AbstractOperator}: The operators to be batched.
batch_size::NTuple{N,Int}: The size of the batch dimensions.
batch_dim_mask::Union{NTuple{M,Symbol}, Pair{NTuple{M1,Symbol},NTuple{M2,Symbol}}}: The mask specifying which dimensions are batch

dimensions. The symbols can be :b for batch dimensions, :s for spreading dimensions or :_ for dimensions on which the operator acts. If the mask is a pair, the first element specifies the domain batch dimensions and the second element specifies the codomain batch dimensions. If no mask is provided, the following order is assumed: (operator dims..., spreading dims..., batch dims...)

threaded::Bool: If true, the operator will execute in parallel over the batch dimensions. Default is nthreads() > 1.
threading_strategy::Symbol: The threading strategy to use. Default is ThreadingStrategy.AUTO, which will automatically choose the

best strategy based on the size of the operators and the input arrays.

Examples

julia> ops = [i * DiagOp([1.0im, 2.0im]) for i in 1:3];

julia> batch_op = BatchOp(ops, 4) # operator dims: (2,), spreading dims: (3,), batch dims: (4,)
⟳╲  ℂ^(2, 3, 4) -> ℂ^(2, 3, 4)

julia> x = rand(ComplexF64, 2, 3, 4);

julia> y = similar(x);

julia> [mul!(@view(y[:, i, j]), ops[i], @view(x[:, i, j])) for i in 1:3, j in 1:4];

julia> y == batch_op * x
true

julia> ops = [i * Variation(3, 4, 5) for i in 1:2];

julia> batch_op = BatchOp(ops, 6, (:s, :_, :_, :_, :b) => (:s, :_, :b, :_))
⟳Ʋ  ℝ^(2, 3, 4, 5, 6) -> ℝ^(2, 60, 6, 3)

source

AbstractOperators.ThreadingStrategy — Module

Threading strategy for SpreadingBatchOp operators.

AUTO: Automatically choose the best threading strategy based on the size of the operators and the input arrays.
COPYING: Create a copy of the operators for each thread.
LOCKING: Use locks to ensure thread safety.
FIXED_OPERATOR: Use a fixed set of operators for each thread, i.e. each operator is assigned to a specific thread.

source