Add a general directsum built on a concatenation interface

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.17.4"
+version = "0.17.5"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/TensorAlgebra.jl

@@ -9,7 +9,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, contractopadd!, data, datatype, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
+            "public biperm, bipartition, cat_similar, concatenate, concatenate!, contractopadd!, data, datatype, directsum, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
         )
     )
 end
@@ -21,6 +21,8 @@ include("MatrixAlgebra.jl")
 include("bituple.jl")
 include("permutedimsadd.jl")
 include("matricize.jl")
+include("concatenate.jl")
+include("directsum.jl")
 include("diagonal.jl")
 include("to_range.jl")
 include("contract/contractalgorithm.jl")

src/concatenate.jl

@@ -0,0 +1,108 @@
+# `concatenate`/`concatenate!` over arrays, with the destination chosen from all inputs rather than
+# just the first (unlike `Base.cat`). A backend customizes allocation by overloading `cat_similar` on
+# the combined `cat_style` of the arguments, and placement by overloading `concatenate!` on the
+# destination type.
+
+import Base.Broadcast as BC
+using Base: promote_eltypeof
+
+function cat_axis(
+        a1::AbstractUnitRange, a2::AbstractUnitRange, a_rest::AbstractUnitRange...
+    )
+    return cat_axis(cat_axis(a1, a2), a_rest...)
+end
+function cat_axis(a1::AbstractUnitRange, a2::AbstractUnitRange)
+    first(a1) == first(a2) == 1 || throw(ArgumentError("Concatenated axes must start at 1"))
+    return Base.OneTo(length(a1) + length(a2))
+end
+
+cat_ndims(dims, as::AbstractArray...) = cat_ndims(Val(dims), as...)
+function cat_ndims(dims::Val, as::AbstractArray...)
+    return max(maximum(unval(dims)), maximum(ndims, as))
+end
+
+cat_axes(dims, as::AbstractArray...) = cat_axes(Val(dims), as...)
+function cat_axes(dims::Val, a::AbstractArray, as::AbstractArray...)
+    return ntuple(cat_ndims(dims, a, as...)) do dim
+        return if dim in unval(dims)
+            cat_axis(map(Base.Fix2(axes, dim), (a, as...))...)
+        else
+            axes(a, dim)
+        end
+    end
+end
+
+function cat_style(dims, as::AbstractArray...)
+    N = cat_ndims(dims, as...)
+    return typeof(BC.combine_styles(as...))(Val(N))
+end
+
+# Default destination: reuse broadcast's `similar` for the style by wrapping in a `Broadcasted`.
+function cat_similar(style, ::Type{T}, ax, args...) where {T}
+    return similar(BC.Broadcasted(style, identity, args, ax), T)
+end
+
+concatenate(dims, args...) = concatenate(Val(dims), args...)
+function concatenate(dims::Val, args...)
+    style = cat_style(dims, args...)
+    dest = cat_similar(style, promote_eltypeof(args...), cat_axes(dims, args...), args...)
+    return concatenate!(dest, dims, args...)
+end
+
+# The offset placement below is adapted from Base's `cat`:
+# https://github.com/JuliaLang/julia/blob/885b1cd875f101f227b345f681cc36879124d80d/base/abstractarray.jl#L1778-L1887
+_copy_or_fill!(A, inds, x) = fill!(view(A, inds...), x)
+_copy_or_fill!(A, inds, x::AbstractArray) = (A[inds...] = x)
+
+cat_size(A) = (1,)
+cat_size(A::AbstractArray) = size(A)
+cat_size(A, d) = 1
+cat_size(A::AbstractArray, d) = size(A, d)
+
+cat_indices(A, d) = Base.OneTo(1)
+cat_indices(A::AbstractArray, d) = axes(A, d)
+
+function __cat!(A, shape, catdims, X...)
+    return __cat_offset!(A, shape, catdims, ntuple(zero, length(shape)), X...)
+end
+function __cat_offset!(A, shape, catdims, offsets, x, X...)
+    # splitting the "work" on x from X... may reduce latency (fewer costly specializations)
+    newoffsets = __cat_offset1!(A, shape, catdims, offsets, x)
+    return __cat_offset!(A, shape, catdims, newoffsets, X...)
+end
+__cat_offset!(A, shape, catdims, offsets) = A
+function __cat_offset1!(A, shape, catdims, offsets, x)
+    inds = ntuple(length(offsets)) do i
+        return if (i <= length(catdims) && catdims[i])
+            offsets[i] .+ cat_indices(x, i)
+        else
+            1:shape[i]
+        end
+    end
+    _copy_or_fill!(A, inds, x)
+    newoffsets = ntuple(length(offsets)) do i
+        return if (i <= length(catdims) && catdims[i])
+            offsets[i] + cat_size(x, i)
+        else
+            offsets[i]
+        end
+    end
+    return newoffsets
+end
+
+dims2cat(dims) = dims2cat(Val(dims))
+function dims2cat(dims::Val)
+    d = unval(dims)
+    if any(≤(0), d)
+        throw(ArgumentError("All cat dimensions must be positive integers, but got $d"))
+    end
+    return ntuple(in(d), maximum(d))
+end
+
+# Generic placement; a backend overrides `concatenate!` on its destination type.
+function concatenate!(dest, dims, args...)
+    catdims = dims2cat(dims)
+    shape = map(length, cat_axes(dims, args...))
+    count(!iszero, catdims)::Int > 1 && zero!(dest)
+    return __cat!(dest, shape, catdims, args...)
+end

src/directsum.jl

@@ -0,0 +1,3 @@
+# `directsum` is a plain concatenation for now, kept as its own entry point so a fusing/rotating
+# variant can later be selected by style, the way `matricize` takes a `FusionStyle`.
+directsum(dims, as...) = concatenate(dims, as...)

test/test_concatenate.jl

@@ -0,0 +1,46 @@
+using TensorAlgebra: TensorAlgebra, cat_axes, concatenate, concatenate!, directsum
+using Test: @test, @testset
+
+@testset "concatenate" begin
+    a = reshape(collect(1.0:6.0), 2, 3)
+    b = reshape(collect(7.0:12.0), 2, 3)
+
+    # Single-dim concatenation matches `vcat`/`hcat`.
+    @test concatenate(1, a, b) == vcat(a, b)
+    @test concatenate(2, a, b) == hcat(a, b)
+
+    # Multi-dim concatenation is the block-diagonal placement (off-diagonal blocks zeroed).
+    c = concatenate((1, 2), a, b)
+    ref = zeros(4, 6)
+    ref[1:2, 1:3] .= a
+    ref[3:4, 4:6] .= b
+    @test c == ref
+
+    # `concatenate!` writes into a provided destination.
+    dest = zeros(4, 6)
+    @test concatenate!(dest, (1, 2), a, b) === dest
+    @test dest == ref
+
+    # Element type is promoted across all inputs, not taken from the first.
+    @test eltype(concatenate(1, a, b .+ 0im)) == ComplexF64
+
+    # `cat_axes` computes the concatenated axes from the arguments.
+    @test cat_axes(Val((1, 2)), a, b) == (Base.OneTo(4), Base.OneTo(6))
+end
+
+@testset "directsum (forwards to concatenate)" begin
+    a = reshape(collect(1.0:6.0), 2, 3)
+    b = reshape(collect(7.0:12.0), 2, 3)
+
+    # `directsum` is exactly `concatenate`: block-concatenation, no basis rotation.
+    @test directsum((1, 2), a, b) == concatenate((1, 2), a, b)
+    @test directsum(1, a, b) == vcat(a, b)
+
+    # N-ary: each summand lands in its own diagonal hyper-block.
+    d = reshape(collect(1.0:8.0), 2, 4)
+    s = directsum((1, 2), a, b, d)
+    @test size(s) == (6, 10)
+    @test s[1:2, 1:3] == a
+    @test s[3:4, 4:6] == b
+    @test s[5:6, 7:10] == d
+end

test/test_exports.jl

@@ -38,9 +38,11 @@ using Test: @test, @testset
         append!(
             exports,
             [
-                :biperm, :bipartition, :contractopadd!, :data, :datatype, :flattenlinear,
-                :label_type, :matricizeopperm, :permutedims, :permutedims!, :scalar,
-                :similar_map, :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!,
+                :biperm, :bipartition, :cat_similar,
+                :concatenate, :concatenate!, :contractopadd!, :data, :datatype, :directsum,
+                :flattenlinear, :label_type,
+                :matricizeopperm, :permutedims, :permutedims!, :scalar, :similar_map,
+                :to_range, :tr, :tryflattenlinear, :ungrade, :zero!, :scale!,
                 :permuteddims, :PermutedDims,
             ]
         )