define matricize and unmatricize (#50)

Author: ogauthe

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.2.10"
+version = "0.3.0"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -19,7 +19,7 @@ BlockArrays = "1.5.0"
 EllipsisNotation = "1.8.0"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.1.1"
-TensorProducts = "0.1.0"
+TensorProducts = "0.1.5"
 TupleTools = "1.6.0"
 TypeParameterAccessors = "0.2.1, 0.3"
 julia = "1.10"

docs/Project.toml

@@ -6,4 +6,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 [compat]
 Documenter = "1.8.1"
 Literate = "2.20.1"
-TensorAlgebra = "0.2.0"
+TensorAlgebra = "0.3.0"

examples/Project.toml

@@ -2,4 +2,4 @@
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 
 [compat]
-TensorAlgebra = "0.2.0"
+TensorAlgebra = "0.3.0"

src/BaseExtensions/permutedims.jl

@@ -1,5 +1,6 @@
 # Workaround for https://github.com/JuliaLang/julia/issues/52615.
 # Fixed by https://github.com/JuliaLang/julia/pull/52623.
+# TODO remove once support for Julia 1.10 is dropped
 function _permutedims!(
   a_dest::AbstractArray{<:Any,N}, a_src::AbstractArray{<:Any,N}, perm::Tuple{Vararg{Int,N}}
 ) where {N}

src/TensorAlgebra.jl

@@ -22,8 +22,7 @@ include("MatrixAlgebra.jl")
 include("blockedtuple.jl")
 include("blockedpermutation.jl")
 include("BaseExtensions/BaseExtensions.jl")
-include("fusedims.jl")
-include("splitdims.jl")
+include("matricize.jl")
 include("contract/contract.jl")
 include("contract/output_labels.jl")
 include("contract/blockedperms.jl")

src/blockedpermutation.jl

@@ -45,6 +45,17 @@ function Base.invperm(bp::AbstractBlockPermutation)
   return blockedperm(invperm(Tuple(bp)), Val(blocklengths(bp)))
 end
 
+# interface
+
+# Bipartition a vector according to the
+# bipartitioned permutation.
+# Like `Base.permute!` block out-of-place and blocked.
+function blockpermute(v, blockedperm::AbstractBlockPermutation)
+  return tuplemortar(map(blockperm -> map(i -> v[i], blockperm), blocks(blockedperm)))
+end
+
+Base.getindex(v, perm::AbstractBlockPermutation) = blockpermute(v, perm)
+
 #
 # Constructors
 #
@@ -53,13 +64,6 @@ function blockedperm(bt::AbstractBlockTuple)
   return permmortar(blocks(bt))
 end
 
-# Bipartition a vector according to the
-# bipartitioned permutation.
-# Like `Base.permute!` block out-of-place and blocked.
-function blockpermute(v, blockedperm::AbstractBlockPermutation)
-  return map(blockperm -> map(i -> v[i], blockperm), blocks(blockedperm))
-end
-
 # blockedpermvcat((4, 3), (2, 1))
 function blockedpermvcat(
   permblocks::Tuple{Vararg{Int}}...; length::Union{Val,Nothing}=nothing

src/contract/allocate_output.jl

@@ -4,137 +4,28 @@ using Base.PermutedDimsArrays: genperm
 # i.e. `ContractAdd`?
 function output_axes(
   ::typeof(contract),
-  biperm_dest::BlockedPermutation{2},
+  biperm_dest::AbstractBlockPermutation{2},
   a1::AbstractArray,
-  biperm1::BlockedPermutation{2},
+  biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
-  biperm2::BlockedPermutation{2},
+  biperm2::AbstractBlockPermutation{2},
   α::Number=one(Bool),
 )
-  axes_codomain, axes_contracted = blockpermute(axes(a1), biperm1)
-  axes_contracted2, axes_domain = blockpermute(axes(a2), biperm2)
+  axes_codomain, axes_contracted = blocks(axes(a1)[biperm1])
+  axes_contracted2, axes_domain = blocks(axes(a2)[biperm2])
   @assert axes_contracted == axes_contracted2
   return genperm((axes_codomain..., axes_domain...), invperm(Tuple(biperm_dest)))
 end
 
-# Inner-product contraction.
-# TODO: Use `ArrayLayouts`-like `MulAdd` object,
-# i.e. `ContractAdd`?
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{0},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{1},
-  a2::AbstractArray,
-  perm2::BlockedPermutation{1},
-  α::Number=one(Bool),
-)
-  axes_contracted = blockpermute(axes(a1), perm1)
-  axes_contracted′ = blockpermute(axes(a2), perm2)
-  @assert axes_contracted == axes_contracted′
-  return ()
-end
-
-# Vec-mat.
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{1},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{1},
-  a2::AbstractArray,
-  biperm2::BlockedPermutation{2},
-  α::Number=one(Bool),
-)
-  (axes_contracted,) = blockpermute(axes(a1), perm1)
-  axes_contracted′, axes_dest = blockpermute(axes(a2), biperm2)
-  @assert axes_contracted == axes_contracted′
-  return genperm((axes_dest...,), invperm(Tuple(perm_dest)))
-end
-
-# Mat-vec.
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{1},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{2},
-  a2::AbstractArray,
-  biperm2::BlockedPermutation{1},
-  α::Number=one(Bool),
-)
-  axes_dest, axes_contracted = blockpermute(axes(a1), perm1)
-  (axes_contracted′,) = blockpermute(axes(a2), biperm2)
-  @assert axes_contracted == axes_contracted′
-  return genperm((axes_dest...,), invperm(Tuple(perm_dest)))
-end
-
-# Outer product.
-function output_axes(
-  ::typeof(contract),
-  biperm_dest::BlockedPermutation{2},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{1},
-  a2::AbstractArray,
-  perm2::BlockedPermutation{1},
-  α::Number=one(Bool),
-)
-  @assert istrivialperm(Tuple(perm1))
-  @assert istrivialperm(Tuple(perm2))
-  axes_dest = (axes(a1)..., axes(a2)...)
-  return genperm(axes_dest, invperm(Tuple(biperm_dest)))
-end
-
-# Array-scalar contraction.
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{1},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{1},
-  a2::AbstractArray,
-  perm2::BlockedPermutation{0},
-  α::Number=one(Bool),
-)
-  @assert istrivialperm(Tuple(perm1))
-  axes_dest = axes(a1)
-  return genperm(axes_dest, invperm(Tuple(perm_dest)))
-end
-
-# Scalar-array contraction.
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{1},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{0},
-  a2::AbstractArray,
-  perm2::BlockedPermutation{1},
-  α::Number=one(Bool),
-)
-  @assert istrivialperm(Tuple(perm2))
-  axes_dest = axes(a2)
-  return genperm(axes_dest, invperm(Tuple(perm_dest)))
-end
-
-# Scalar-scalar contraction.
-function output_axes(
-  ::typeof(contract),
-  perm_dest::BlockedPermutation{0},
-  a1::AbstractArray,
-  perm1::BlockedPermutation{0},
-  a2::AbstractArray,
-  perm2::BlockedPermutation{0},
-  α::Number=one(Bool),
-)
-  return ()
-end
-
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
 # i.e. `ContractAdd`?
 function allocate_output(
   ::typeof(contract),
-  biperm_dest::BlockedPermutation,
+  biperm_dest::AbstractBlockPermutation,
   a1::AbstractArray,
-  biperm1::BlockedPermutation,
+  biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
-  biperm2::BlockedPermutation,
+  biperm2::AbstractBlockPermutation,
   α::Number=one(Bool),
 )
   axes_dest = output_axes(contract, biperm_dest, a1, biperm1, a2, biperm2, α)

src/contract/blockedperms.jl

@@ -22,10 +22,10 @@ function blockedperms(::typeof(contract), dimnames_dest, dimnames1, dimnames2)
   perm_domain2 = BaseExtensions.indexin(domain, dimnames2)
 
   permblocks_dest = (perm_codomain_dest, perm_domain_dest)
-  biperm_dest = blockedpermvcat(filter(!isempty, permblocks_dest)...)
+  biperm_dest = blockedpermvcat(permblocks_dest...)
   permblocks1 = (perm_codomain1, perm_domain1)
-  biperm1 = blockedpermvcat(filter(!isempty, permblocks1)...)
+  biperm1 = blockedpermvcat(permblocks1...)
   permblocks2 = (perm_codomain2, perm_domain2)
-  biperm2 = blockedpermvcat(filter(!isempty, permblocks2)...)
+  biperm2 = blockedpermvcat(permblocks2...)
   return biperm_dest, biperm1, biperm2
 end

src/contract/contract.jl

@@ -13,11 +13,11 @@ default_contract_alg() = Matricize()
 function contract!(
   alg::Algorithm,
   a_dest::AbstractArray,
-  biperm_dest::BlockedPermutation,
+  biperm_dest::AbstractBlockPermutation,
   a1::AbstractArray,
-  biperm1::BlockedPermutation,
+  biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
-  biperm2::BlockedPermutation,
+  biperm2::AbstractBlockPermutation,
   α::Number,
   β::Number,
 )
@@ -110,11 +110,11 @@ end
 
 function contract(
   alg::Algorithm,
-  biperm_dest::BlockedPermutation,
+  biperm_dest::AbstractBlockPermutation,
   a1::AbstractArray,
-  biperm1::BlockedPermutation,
+  biperm1::AbstractBlockPermutation,
   a2::AbstractArray,
-  biperm2::BlockedPermutation,
+  biperm2::AbstractBlockPermutation,
   α::Number;
   kwargs...,
 )

src/contract/contract_matricize/contract.jl

@@ -1,103 +1,20 @@
 using LinearAlgebra: mul!
 
 function contract!(
-  alg::Matricize,
+  ::Matricize,
   a_dest::AbstractArray,
-  biperm_dest::BlockedPermutation,
+  biperm_dest::AbstractBlockPermutation{2},
   a1::AbstractArray,
-  biperm1::BlockedPermutation,
+  biperm1::AbstractBlockPermutation{2},
   a2::AbstractArray,
-  biperm2::BlockedPermutation,
+  biperm2::AbstractBlockPermutation{2},
   α::Number,
   β::Number,
 )
-  a_dest_mat = fusedims(a_dest, biperm_dest)
-  a1_mat = fusedims(a1, biperm1)
-  a2_mat = fusedims(a2, biperm2)
-  _mul!(a_dest_mat, a1_mat, a2_mat, α, β)
-  splitdims!(a_dest, a_dest_mat, biperm_dest)
-  return a_dest
-end
-
-# Matrix multiplication.
-function _mul!(
-  a_dest::AbstractMatrix, a1::AbstractMatrix, a2::AbstractMatrix, α::Number, β::Number
-)
-  mul!(a_dest, a1, a2, α, β)
-  return a_dest
-end
-
-# Inner product.
-function _mul!(
-  a_dest::AbstractArray{<:Any,0},
-  a1::AbstractVector,
-  a2::AbstractVector,
-  α::Number,
-  β::Number,
-)
-  a_dest[] = transpose(a1) * a2 * α + a_dest[] * β
-  return a_dest
-end
-
-# Vec-mat.
-function _mul!(
-  a_dest::AbstractVector, a1::AbstractVector, a2::AbstractMatrix, α::Number, β::Number
-)
-  mul!(transpose(a_dest), transpose(a1), a2, α, β)
-  return a_dest
-end
-
-# Mat-vec.
-function _mul!(
-  a_dest::AbstractVector, a1::AbstractMatrix, a2::AbstractVector, α::Number, β::Number
-)
-  mul!(a_dest, a1, a2, α, β)
-  return a_dest
-end
-
-# Outer product.
-function _mul!(
-  a_dest::AbstractMatrix, a1::AbstractVector, a2::AbstractVector, α::Number, β::Number
-)
-  mul!(a_dest, a1, transpose(a2), α, β)
-  return a_dest
-end
-
-# Array-scalar contraction.
-function _mul!(
-  a_dest::AbstractVector,
-  a1::AbstractVector,
-  a2::AbstractArray{<:Any,0},
-  α::Number,
-  β::Number,
-)
-  α′ = a2[] * α
-  a_dest .= a1 .* α′ .+ a_dest .* β
-  return a_dest
-end
-
-# Scalar-array contraction.
-function _mul!(
-  a_dest::AbstractVector,
-  a1::AbstractArray{<:Any,0},
-  a2::AbstractVector,
-  α::Number,
-  β::Number,
-)
-  # Preserve the ordering in case of non-commutative algebra.
-  a_dest .= a1[] .* a2 .* α .+ a_dest .* β
-  return a_dest
-end
-
-# Scalar-scalar contraction.
-function _mul!(
-  a_dest::AbstractArray{<:Any,0},
-  a1::AbstractArray{<:Any,0},
-  a2::AbstractArray{<:Any,0},
-  α::Number,
-  β::Number,
-)
-  # Preserve the ordering in case of non-commutative algebra.
-  a_dest[] = a1[] * a2[] * α + a_dest[] * β
+  a_dest_mat = matricize(a_dest, biperm_dest)
+  a1_mat = matricize(a1, biperm1)
+  a2_mat = matricize(a2, biperm2)
+  mul!(a_dest_mat, a1_mat, a2_mat, α, β)
+  unmatricize!(a_dest, a_dest_mat, biperm_dest)
   return a_dest
 end