Refactor matricize/unmatricize

Author: mtfishman

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.4.6"
+version = "0.4.7"
 authors = ["ITensor developers <[email protected]> and contributors"]
 
 [deps]

src/blockedpermutation.jl

@@ -8,7 +8,7 @@ function istrivialperm(t::Tuple)
     return t == trivialperm(length(t))
 end
 
-value(::Val{N}) where {N} = N
+unval(::Val{N}) where {N} = N
 
 _flatten_tuples(t::Tuple) = t
 function _flatten_tuples(t1::Tuple, t2::Tuple, trest::Tuple...)
@@ -87,7 +87,7 @@ function blockedpermvcat(
 end
 
 function blockedpermvcat(len::Val, permblocks::Tuple{Vararg{Int}}...)
-    value(len) != sum(length.(permblocks); init = 0) &&
+    unval(len) != sum(length.(permblocks); init = 0) &&
         throw(ArgumentError("Invalid total length"))
     return permmortar(Tuple(permblocks))
 end
@@ -97,7 +97,7 @@ function _blockedperm_length(::Nothing, specified_perm::Tuple{Vararg{Int}})
 end
 
 function _blockedperm_length(vallength::Val, ::Tuple{Vararg{Int}})
-    return value(vallength)
+    return unval(vallength)
 end
 
 # blockedpermvcat((4, 3), .., 1) == blockedpermvcat((4, 3), (2,), (1,))
@@ -199,8 +199,11 @@ end
 
 blockedperm(tp::BlockedTrivialPermutation) = tp
 
+function blockedtrivialperm(blocklengths::Tuple{Vararg{Val}})
+    return BlockedTrivialPermutation{length(blocklengths), unval.(blocklengths)}()
+end
 function blockedtrivialperm(blocklengths::Tuple{Vararg{Int}})
-    return BlockedTrivialPermutation{length(blocklengths), blocklengths}()
+    return blockedtrivialperm(Val.(blocklengths))
 end
 
 function trivialperm(blockedperm::AbstractBlockTuple)

src/contract/contract.jl

@@ -4,9 +4,28 @@ abstract type Algorithm end
 
 Algorithm(alg::Algorithm) = alg
 
-struct Matricize <: Algorithm end
+struct Matricize{Style} <: Algorithm
+    fusion_style::Style
+end
 
-default_contract_alg() = Matricize()
+function default_contract_alg(a1::AbstractArray, labels1, a2::AbstractArray, labels2)
+    style1 = FusionStyle(a1)
+    style2 = FusionStyle(a2)
+    style1 == style2 || error("Styles must match.")
+    return Matricize(style1)
+end
+function default_contractadd!_alg(
+        a_dest::AbstractArray, labels_dest,
+        a1::AbstractArray, labels1,
+        a2::AbstractArray, labels2,
+        α::Number, β::Number,
+    )
+    style_dest = FusionStyle(a_dest)
+    style1 = FusionStyle(a1)
+    style2 = FusionStyle(a2)
+    style_dest == style1 == style2 || error("Styles must match.")
+    return Matricize(style_dest)
+end
 
 # Required interface if not using
 # matricized contraction.
@@ -29,7 +48,7 @@ function contract(
         labels1,
         a2::AbstractArray,
         labels2;
-        alg = default_contract_alg(),
+        alg = default_contract_alg(a1, labels1, a2, labels2),
         kwargs...,
     )
     return contract(Algorithm(alg), a1, labels1, a2, labels2; kwargs...)
@@ -48,7 +67,7 @@ function contract(
         labels1,
         a2::AbstractArray,
         labels2;
-        alg = default_contract_alg(),
+        alg = default_contract_alg(a1, labels1, a2, labels2),
         kwargs...,
     )
     return contract(Algorithm(alg), labels_dest, a1, labels1, a2, labels2; kwargs...)
@@ -75,7 +94,7 @@ function contractadd!(
         labels2,
         α::Number,
         β::Number;
-        alg = default_contract_alg(),
+        alg = default_contractadd!_alg(a_dest, labels_dest, a1, labels1, a2, labels2, α, β),
         kwargs...,
     )
     contractadd!(

src/contract/contract_matricize/contract.jl

@@ -1,7 +1,7 @@
 using LinearAlgebra: mul!
 
 function contractadd!(
-        ::Matricize,
+        alg::Matricize,
         a_dest::AbstractArray,
         biperm_dest::AbstractBlockPermutation{2},
         a1::AbstractArray,
@@ -12,11 +12,10 @@ function contractadd!(
         β::Number,
     )
     invbiperm = biperm(invperm(biperm_dest), length_codomain(biperm1))
-
     check_input(contract, a_dest, invbiperm, a1, biperm1, a2, biperm2)
-    a1_mat = matricize(a1, biperm1)
-    a2_mat = matricize(a2, biperm2)
+    a1_mat = matricize(alg.fusion_style, a1, biperm1)
+    a2_mat = matricize(alg.fusion_style, a2, biperm2)
     a_dest_mat = a1_mat * a2_mat
-    unmatricizeadd!(a_dest, a_dest_mat, invbiperm, α, β)
+    unmatricizeadd!(alg.fusion_style, a_dest, a_dest_mat, invbiperm, α, β)
     return a_dest
 end

src/matricize.jl

@@ -21,77 +21,146 @@ trivial_axis(::Tuple{}) = Base.OneTo(1)
 trivial_axis(::Tuple{Vararg{AbstractUnitRange}}) = Base.OneTo(1)
 trivial_axis(::Tuple{Vararg{AbstractBlockedUnitRange}}) = blockedrange([1])
 
+# Inner version takes a list of sub-permutations, overload this one if needed.
 function fuseaxes(
-        axes::Tuple{Vararg{AbstractUnitRange}}, blockedperm::AbstractBlockPermutation
+        axes::Tuple{Vararg{AbstractUnitRange}}, lengths::Val...
     )
-    axesblocks = blocks(axes[blockedperm])
+    axesblocks = blocks(axes[blockedtrivialperm(lengths)])
     return map(block -> isempty(block) ? trivial_axis(axes) : ⊗(block...), axesblocks)
 end
 
+# Inner version takes a list of sub-permutations, overload this one if needed.
+function fuseaxes(
+        axes::Tuple{Vararg{AbstractUnitRange}}, permblocks::Tuple{Vararg{Int}}...
+    )
+    axes′ = map(d -> axes[d], permmortar(permblocks))
+    return fuseaxes(axes′, Val.(length.(permblocks))...)
+end
+
+function fuseaxes(
+        axes::Tuple{Vararg{AbstractUnitRange}}, blockedperm::AbstractBlockPermutation
+    )
+    return fuseaxes(axes, blocks(blockedperm)...)
+end
+
+# Inner version takes a list of sub-permutations, overload this one if needed.
+function permuteblockeddims(a::AbstractArray, perm1, perm2)
+    return _permutedims(a, (perm1..., perm2...))
+end
+function permuteblockeddims!(a_dest::AbstractArray, a_src::AbstractArray, perm1, perm2)
+    return _permutedims!(a_dest, a_src, (perm1..., perm2...))
+end
+
 # TODO remove _permutedims once support for Julia 1.10 is dropped
 # define permutedims with a BlockedPermuation. Default is to flatten it.
-function permuteblockeddims(a::AbstractArray, biperm::AbstractBlockPermutation)
-    return _permutedims(a, Tuple(biperm))
+function permuteblockeddims(a::AbstractArray, biperm::AbstractBlockPermutation{2})
+    return permuteblockeddims(a, blocks(biperm)...)
 end
-
 function permuteblockeddims!(
-        a::AbstractArray, b::AbstractArray, biperm::AbstractBlockPermutation
+        a_dest::AbstractArray, a_src::AbstractArray, biperm::AbstractBlockPermutation{2}
     )
-    return _permutedims!(a, b, Tuple(biperm))
+    return permuteblockeddims!(a_dest, a_src, blocks(biperm)...)
 end
 
 # =====================================  matricize  ========================================
 # TBD settle copy/not copy convention
 # matrix factorizations assume copy
 # maybe: copy=false kwarg
 
-function matricize(a::AbstractArray, biperm_dest::AbstractBlockPermutation{2})
-    ndims(a) == length(biperm_dest) || throw(ArgumentError("Invalid bipermutation"))
-    return matricize(FusionStyle(a), a, biperm_dest)
+function matricize(a::AbstractArray, length1::Val, length2::Val)
+    return matricize(FusionStyle(a), a, length1, length2)
+end
+# This is the primary function that should be overloaded for new fusion styles.
+# This assumes the permutation was already performed.
+function matricize(style::FusionStyle, a::AbstractArray, length1::Val, length2::Val)
+    return throw(
+        MethodError(
+            matricize, Tuple{typeof(style), typeof(a), typeof(length1), typeof(length2)}
+        )
+    )
 end
 
 function matricize(
-        style::FusionStyle, a::AbstractArray, biperm_dest::AbstractBlockPermutation{2}
+        a::AbstractArray, permblock1::Tuple{Vararg{Int}}, permblock2::Tuple{Vararg{Int}}
     )
-    a_perm = permuteblockeddims(a, biperm_dest)
-    return matricize(style, a_perm, trivialperm(biperm_dest))
+    return matricize(FusionStyle(a), a, permblock1, permblock2)
 end
-
+# This is a more advanced version to overload where the permutation is actually performed.
 function matricize(
-        style::FusionStyle, a::AbstractArray, biperm_dest::BlockedTrivialPermutation{2}
+        style::FusionStyle, a::AbstractArray,
+        permblock1::NTuple{N1, Int}, permblock2::NTuple{N2, Int}
+    ) where {N1, N2}
+    ndims(a) == length(permblock1) + length(permblock2) ||
+        throw(ArgumentError("Invalid bipermutation"))
+    a_perm = permuteblockeddims(a, permblock1, permblock2)
+    return matricize(style, a_perm, Val(length(permblock1)), Val(length(permblock2)))
+end
+
+# Process inputs such as `EllipsisNotation.Ellipsis`.
+function to_permblocks(a::AbstractArray, permblocks::NTuple{2, Tuple{Vararg{Int}}})
+    isperm((permblocks[1]..., permblocks[2]...)) ||
+        throw(ArgumentError("Invalid bipermutation"))
+    return permblocks
+end
+# Like `setcomplement` is like `setdiff` but assumes t2 ⊆ t1.
+function tuplesetcomplement(t1::NTuple{N1}, t2::NTuple{N2}) where {N1, N2}
+    t2 ⊆ t1 || throw(ArgumentError("t2 must be a subset of t1"))
+    return NTuple{N1 - N2}(setdiff(t1, t2))
+end
+function to_permblocks(
+        a::AbstractArray, permblocks::Tuple{Tuple{Ellipsis}, Tuple{Vararg{Int}}}
+    )
+    permblocks1 = tuplesetcomplement(ntuple(identity, ndims(a)), permblocks[2])
+    return (permblocks1, permblocks[2])
+end
+function to_permblocks(
+        a::AbstractArray, permblocks::Tuple{Tuple{Vararg{Int}}, Tuple{Ellipsis}}
     )
-    return throw(MethodError(matricize, Tuple{typeof(style), typeof(a), typeof(biperm_dest)}))
+    permblocks2 = tuplesetcomplement(ntuple(identity, ndims(a)), permblocks[1])
+    return (permblocks[1], permblocks2)
+end
+function matricize(a::AbstractArray, permblock1, permblock2)
+    return matricize(FusionStyle(a), a, permblock1, permblock2)
+end
+function matricize(style::FusionStyle, a::AbstractArray, permblock1, permblock2)
+    return matricize(style, a, to_permblocks(a, (permblock1, permblock2))...)
 end
 
-# default is reshape
+function matricize(a::AbstractArray, biperm_dest::AbstractBlockPermutation{2})
+    return matricize(FusionStyle(a), a, biperm_dest)
+end
 function matricize(
-        ::ReshapeFusion, a::AbstractArray, biperm_dest::BlockedTrivialPermutation{2}
+        style::FusionStyle, a::AbstractArray, biperm_dest::AbstractBlockPermutation{2}
     )
-    new_axes = fuseaxes(axes(a), biperm_dest)
-    return reshape(a, new_axes...)
+    return matricize(style, a, blocks(biperm_dest)...)
 end
 
-function matricize(a::AbstractArray, permblock1::Tuple, permblock2::Tuple)
-    return matricize(a, blockedpermvcat(permblock1, permblock2; length = Val(ndims(a))))
+# default is reshape
+function matricize(::ReshapeFusion, a::AbstractArray, length1::Val, length2::Val)
+    return reshape(a, fuseaxes(axes(a), length1, length2)...)
 end
 
 # ====================================  unmatricize  =======================================
 function unmatricize(m::AbstractMatrix, axes_dest, invbiperm::AbstractBlockPermutation{2})
-    length(axes_dest) == length(invbiperm) ||
-        throw(ArgumentError("axes do not match permutation"))
     return unmatricize(FusionStyle(m), m, axes_dest, invbiperm)
 end
-
 function unmatricize(
-        ::FusionStyle, m::AbstractMatrix, axes_dest, invbiperm::AbstractBlockPermutation{2}
+        style::FusionStyle, m::AbstractMatrix, axes_dest, invbiperm::AbstractBlockPermutation{2}
     )
+    length(axes_dest) == length(invbiperm) ||
+        throw(ArgumentError("axes do not match permutation"))
     blocked_axes = axes_dest[invbiperm]
-    a12 = unmatricize(m, blocked_axes)
+    a12 = unmatricize(style, m, blocked_axes)
     biperm_dest = biperm(invperm(invbiperm), length_codomain(axes_dest))
-
     return permuteblockeddims(a12, biperm_dest)
 end
 
+function unmatricize(
+        m::AbstractMatrix,
+        blocked_axes::BlockedTuple{2, <:Any, <:Tuple{Vararg{AbstractUnitRange}}},
+    )
+    return unmatricize(FusionStyle(m), m, blocked_axes)
+end
 function unmatricize(
         ::ReshapeFusion,
         m::AbstractMatrix,
@@ -100,30 +169,49 @@ function unmatricize(
     return reshape(m, Tuple(blocked_axes)...)
 end
 
-function unmatricize(m::AbstractMatrix, blocked_axes)
-    return unmatricize(FusionStyle(m), m, blocked_axes)
-end
-
 function unmatricize(
         m::AbstractMatrix,
         codomain_axes::Tuple{Vararg{AbstractUnitRange}},
         domain_axes::Tuple{Vararg{AbstractUnitRange}},
     )
+    return unmatricize(FusionStyle(m), m, codomain_axes, domain_axes)
+end
+function unmatricize(
+        style::FusionStyle, m::AbstractMatrix,
+        codomain_axes::Tuple{Vararg{AbstractUnitRange}},
+        domain_axes::Tuple{Vararg{AbstractUnitRange}},
+    )
     blocked_axes = tuplemortar((codomain_axes, domain_axes))
-    return unmatricize(m, blocked_axes)
+    return unmatricize(style, m, blocked_axes)
 end
 
-function unmatricize!(a_dest, m::AbstractMatrix, invbiperm::AbstractBlockPermutation{2})
+function unmatricize!(
+        a_dest::AbstractArray, m::AbstractMatrix, invbiperm::AbstractBlockPermutation{2}
+    )
+    return unmatricize!(FusionStyle(m), a_dest, m, invbiperm)
+end
+function unmatricize!(
+        style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix,
+        invbiperm::AbstractBlockPermutation{2},
+    )
     ndims(a_dest) == length(invbiperm) ||
         throw(ArgumentError("destination does not match permutation"))
     blocked_axes = axes(a_dest)[invbiperm]
-    a_perm = unmatricize(m, blocked_axes)
+    a_perm = unmatricize(style, m, blocked_axes)
     biperm_dest = biperm(invperm(invbiperm), length_codomain(axes(a_dest)))
     return permuteblockeddims!(a_dest, a_perm, biperm_dest)
 end
 
-function unmatricizeadd!(a_dest, a_dest_mat, invbiperm, α, β)
-    a12 = unmatricize(a_dest_mat, axes(a_dest), invbiperm)
+function unmatricizeadd!(
+        a_dest::AbstractArray, m::AbstractMatrix, invbiperm::AbstractBlockPermutation{2},
+        α::Number, β::Number
+    )
+    return unmatricizeadd!(FusionStyle(a_dest), a_dest, m, invbiperm, α, β)
+end
+function unmatricizeadd!(
+        style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, invbiperm, α, β
+    )
+    a12 = unmatricize(style, m, axes(a_dest), invbiperm)
     a_dest .= α .* a12 .+ β .* a_dest
     return a_dest
 end

test/Project.toml

@@ -14,6 +14,9 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
+[sources]
+TensorAlgebra = {path = ".."}
+
 [compat]
 Aqua = "0.8.9"
 BlockArrays = "1.6.1"