define permutedims!

Author: ogauthe

src/fusiontensor/base_interface.jl

@@ -77,6 +77,10 @@ Base.imag(ft::FusionTensor) = set_data_matrix(ft, imag(data_matrix(ft)))
 
 Base.permutedims(ft::FusionTensor, args...) = fusiontensor_permutedims(ft, args...)
 
+function Base.permutedims!(ftdst::FusionTensor, ftsrc::FusionTensor, args...)
+  return fusiontensor_permutedims!(ftdst, ftsrc, args...)
+end
+
 Base.real(ft::FusionTensor{<:Real}) = ft   # same object
 Base.real(ft::FusionTensor) = set_data_matrix(ft, real(data_matrix(ft)))
 

src/fusiontensor/fusiontensor.jl

@@ -255,6 +255,8 @@ function GradedArrays.sector_type(::Type{FT}) where {FT<:FusionTensor}
   return sector_type(type_parameters(FT, 3))
 end
 
+SymmetryStyle(::Type{FT}) where {FT<:FusionTensor} = SymmetryStyle(sector_type(FT))
+
 # ==============================  FusionTensor interface  ==================================
 
 # misc access

src/fusiontensor/fusiontensoraxes.jl

@@ -3,6 +3,7 @@ using GradedArrays:
   GradedArrays,
   AbstractGradedUnitRange,
   AbstractSector,
+  SymmetryStyle,
   TrivialSector,
   dual,
   sector_type,
@@ -110,6 +111,10 @@ function GradedArrays.sector_type(::Type{FTA}) where {BT,FTA<:FusionTensorAxes{B
   return sector_type(type_parameters(type_parameters(BT, 3), 1))
 end
 
+function GradedArrays.SymmetryStyle(::Type{FTA}) where {FTA<:FusionTensorAxes}
+  return SymmetryStyle(sector_type(FTA))
+end
+
 function GradedArrays.checkspaces(
   ::Type{Bool}, left::FusionTensorAxes, right::FusionTensorAxes
 )

src/permutedims/permutedims.jl

@@ -3,23 +3,18 @@
 using BlockArrays: blocklengths
 using Strided: Strided, @strided
 
-using TensorAlgebra: BlockedPermutation, permmortar, blockpermute
-
-function naive_permutedims(ft, biperm::BlockedPermutation{2})
-  @assert ndims(ft) == length(biperm)
-
-  # naive permute: cast to dense, permutedims, cast to FusionTensor
-  arr = Array(ft)
-  permuted_arr = permutedims(arr, Tuple(biperm))
-  permuted = to_fusiontensor(permuted_arr, blocks(axes(ft)[biperm])...)
-  return permuted
-end
+using GradedArrays: AbelianStyle, NotAbelianStyle, SymmetryStyle, checkspaces
+using TensorAlgebra: AbstractBlockPermutation, permmortar
 
 # permutedims with 1 tuple of 2 separate tuples
 function fusiontensor_permutedims(ft, new_leg_indices::Tuple{Tuple,Tuple})
   return fusiontensor_permutedims(ft, new_leg_indices...)
 end
 
+function fusiontensor_permutedims!(ftdst, ftsrc, new_leg_indices::Tuple{Tuple,Tuple})
+  return fusiontensor_permutedims!(ftdst, ftsrc, new_leg_indices...)
+end
+
 # permutedims with 2 separate tuples
 function fusiontensor_permutedims(
   ft, new_codomain_indices::Tuple, new_domain_indices::Tuple
@@ -28,29 +23,55 @@ function fusiontensor_permutedims(
   return fusiontensor_permutedims(ft, biperm)
 end
 
-function fusiontensor_permutedims(ft, biperm::BlockedPermutation{2})
+function fusiontensor_permutedims!(
+  ftdst, ftsrc, new_codomain_indices::Tuple, new_domain_indices::Tuple
+)
+  biperm = permmortar((new_codomain_indices, new_domain_indices))
+  return fusiontensor_permutedims!(ftdst, ftsrc, biperm)
+end
+
+# permutedims with BlockedPermutation
+function fusiontensor_permutedims(ft, biperm::AbstractBlockPermutation{2})
   ndims(ft) == length(biperm) || throw(ArgumentError("Invalid permutation length"))
+  ftdst = FusionTensor{eltype(ft)}(undef, axes(ft)[biperm])
+  fusiontensor_permutedims!(ftdst, ft, biperm)
+  return ftdst
+end
+
+function fusiontensor_permutedims!(ftdst, ftsrc, biperm::AbstractBlockPermutation{2})
+  ndims(ftsrc) == length(biperm) || throw(ArgumentError("Invalid permutation length"))
+  blocklengths(axes(ftdst)) == blocklengths(biperm) ||
+    throw(ArgumentError("Destination tensor does not match bipermutation"))
+  checkspaces(axes(ftdst), axes(ftsrc)[biperm])
 
-  # early return for identity operation. Do not copy. Also handle tricky 0-dim case.
-  if ndims_codomain(ft) == first(blocklengths(biperm))  # compile time
-    if Tuple(biperm) == ntuple(identity, ndims(ft))
-      return ft
+  # early return for identity operation. Also handle tricky 0-dim case.
+  if ndims_codomain(ftdst) == ndims_codomain(ftsrc)  # compile time
+    if Tuple(biperm) == ntuple(identity, ndims(ftdst))
+      copy!(data_matrix(ftdst), data_matrix(ftsrc))
+      return nothing
     end
   end
+  return permute_data!(SymmetryStyle(ftdst), ftdst, ftsrc, Tuple(biperm))
+end
 
-  new_ft = FusionTensor{eltype(ft)}(undef, axes(ft)[biperm])
-  fusiontensor_permutedims!(new_ft, ft, Tuple(biperm))
-  return new_ft
+# ===============================   Internal   =============================================
+function permute_data!(::AbelianStyle, ftdst, ftsrc, flatperm)
+  # abelian case: all unitary blocks are 1x1 identity matrices
+  # compute_unitary is only called to get block positions
+  unitary = compute_unitary(ftdst, ftsrc, flatperm)
+  for ((old_trees, new_trees), _) in unitary
+    new_block = view(ftdst, new_trees...)
+    old_block = view(ftsrc, old_trees...)
+    @strided new_block .= permutedims(old_block, flatperm)
+  end
 end
 
-function fusiontensor_permutedims!(
-  new_ft::FusionTensor{T,N}, old_ft::FusionTensor{T,N}, flatperm::NTuple{N,Integer}
-) where {T,N}
-  foreach(m -> fill!(m, zero(T)), eachstoredblock(data_matrix(new_ft)))
-  unitary = compute_unitary(new_ft, old_ft, flatperm)
+function permute_data!(::NotAbelianStyle, ftdst, ftsrc, flatperm)
+  foreach(m -> fill!(m, zero(eltype(ftdst))), eachstoredblock(data_matrix(ftdst)))
+  unitary = compute_unitary(ftdst, ftsrc, flatperm)
   for ((old_trees, new_trees), coeff) in unitary
-    new_block = view(new_ft, new_trees...)
-    old_block = view(old_ft, old_trees...)
+    new_block = view(ftdst, new_trees...)
+    old_block = view(ftsrc, old_trees...)
     @strided new_block .+= coeff .* permutedims(old_block, flatperm)
   end
 end

test/test_permutedims.jl

@@ -1,12 +1,12 @@
 using Test: @test, @testset, @test_broken, @test_throws
+using BlockArrays: blocks
 
 using FusionTensors:
   FusionTensor,
   FusionTensorAxes,
   data_matrix,
   codomain_axis,
   domain_axis,
-  naive_permutedims,
   ndims_domain,
   ndims_codomain,
   to_fusiontensor
@@ -15,6 +15,16 @@ using TensorAlgebra: permmortar, tuplemortar
 
 include("setup.jl")
 
+function naive_permutedims(ft, biperm)
+  @assert ndims(ft) == length(biperm)
+
+  # naive permute: cast to dense, permutedims, cast to FusionTensor
+  arr = Array(ft)
+  permuted_arr = permutedims(arr, Tuple(biperm))
+  permuted = to_fusiontensor(permuted_arr, blocks(axes(ft)[biperm])...)
+  return permuted
+end
+
 @testset "Abelian permutedims" begin
   @testset "dummy" begin
     g1 = gradedrange([U1(0) => 1, U1(1) => 2, U1(2) => 3])
@@ -28,14 +38,23 @@ include("setup.jl")
 
       # test permutedims interface
       ft2 = permutedims(ft1, (1, 2), (3, 4))   # trivial with 2 tuples
-      @test ft2 === ft1  # same object
+      @test ft2 ≈ ft1
+      @test ft2 !== ft1
+      @test data_matrix(ft2) !== data_matrix(ft1)  # check copy
+      @test data_matrix(ft2) == data_matrix(ft1)  # check copy
 
       ft2 = permutedims(ft1, ((1, 2), (3, 4)))   # trivial with tuple of 2 tuples
-      @test ft2 === ft1  # same object
+      @test ft2 ≈ ft1
+      @test ft2 !== ft1
+      @test data_matrix(ft2) !== data_matrix(ft1)  # check copy
+      @test data_matrix(ft2) == data_matrix(ft1)  # check copy
 
       biperm = permmortar(((1, 2), (3, 4)))
       ft2 = permutedims(ft1, biperm)   # trivial with biperm
-      @test ft2 === ft1  # same object
+      @test ft2 ≈ ft1
+      @test ft2 !== ft1
+      @test data_matrix(ft2) !== data_matrix(ft1)  # check copy
+      @test data_matrix(ft2) == data_matrix(ft1)  # check copy
 
       ft3 = permutedims(ft1, (4,), (1, 2, 3))
       @test ft3 !== ft1