Refactor matricize/unmatricize (#96)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.4.6"
 authors = ["ITensor developers <[email protected]> and contributors"]
+version = "0.5.0"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -22,11 +22,11 @@ TensorAlgebraTensorOperationsExt = "TensorOperations"
 [compat]
 ArrayLayouts = "1.10.4"
 BlockArrays = "1.7.2"
-EllipsisNotation = "1.8.0"
+EllipsisNotation = "1.8"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 TensorOperations = "5"
 TensorProducts = "0.1.5"
-TupleTools = "1.6.0"
+TupleTools = "1.6"
 TypeParameterAccessors = "0.2.1, 0.3, 0.4"
 julia = "1.10"

docs/Project.toml

@@ -3,7 +3,10 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
+[sources]
+TensorAlgebra = {path = ".."}
+
 [compat]
 Documenter = "1.8.1"
 Literate = "2.20.1"
-TensorAlgebra = "0.4"
+TensorAlgebra = "0.5"

examples/Project.toml

@@ -1,5 +1,8 @@
 [deps]
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
+[sources]
+TensorAlgebra = {path = ".."}
+
 [compat]
-TensorAlgebra = "0.4"
+TensorAlgebra = "0.5"

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,29 @@ abstract type Algorithm end
 
 Algorithm(alg::Algorithm) = alg
 
-struct Matricize <: Algorithm end
+struct Matricize{Style} <: Algorithm
+    fusion_style::Style
+end
+Matricize() = Matricize(ReshapeFusion())
 
-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 +49,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 +68,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 +95,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