Refactor code and tests

ext/TenetChainRulesCoreExt.jl

@@ -28,18 +28,14 @@ ChainRulesCore.rrule(T::Type{<:Tensor}, data, inds) = T(data, inds), Tensor_pull
 @non_differentiable symdiff(s::Base.AbstractVecOrTuple{Symbol}, itrs::Base.AbstractVecOrTuple{Symbol}...)
 
 function ChainRulesCore.ProjectTo(tn::T) where {T<:AbstractTensorNetwork}
-    # TODO create function to extract extra fields
-    fields = map(fieldnames(T)) do fieldname
-        if fieldname === :tensors
-            :tensors => ProjectTo(tn.tensors)
-        else
-            fieldname => getfield(tn, fieldname)
-        end
-    end
-    ProjectTo{T}(; fields...)
+    ProjectTo{T}(; tensors = ProjectTo(tensors(tn)))
 end
 
-function (projector::ProjectTo{T})(dx::Union{T,Tangent{T}}) where {T<:AbstractTensorNetwork}
+function (projector::ProjectTo{T})(dx::T) where {T<:AbstractTensorNetwork}
+    Tangent{TensorNetwork}(tensors = projector.tensors(tensors(tn)))
+end
+
+function (projector::ProjectTo{T})(dx::Tangent{T}) where {T<:AbstractTensorNetwork}
     dx.tensors isa NoTangent && return NoTangent()
     Tangent{TensorNetwork}(tensors = projector.tensors(dx.tensors))
 end
@@ -49,13 +45,7 @@ function Base.:+(x::T, Δ::Tangent{TensorNetwork}) where {T<:AbstractTensorNetwo
     tensors = map(+, tensors(x), Δ.tensors)
 
     # TODO create function fitted for this? or maybe standardize constructors?
-    T(map(fieldnames(T)) do fieldname
-        if fieldname === :tensors
-            tensors
-        else
-            getfield(x, fieldname)
-        end
-    end...)
+    T(tensors)
 end
 
 function ChainRulesCore.frule((_, Δ), T::Type{<:AbstractTensorNetwork}, tensors)

ext/TenetFiniteDifferencesExt.jl

@@ -5,19 +5,8 @@ using Tenet: AbstractTensorNetwork
 using FiniteDifferences
 
 function FiniteDifferences.to_vec(x::T) where {T<:AbstractTensorNetwork}
-    x_vec, back = to_vec(x.tensors)
-    function TensorNetwork_from_vec(v)
-        tensors = back(v)
-
-        # TODO create function fitted for this? or maybe standardize constructors?
-        T(map(fieldnames(T)) do fieldname
-            if fieldname === :tensors
-                tensors
-            else
-                getfield(x, fieldname)
-            end
-        end...)
-    end
+    x_vec, back = to_vec(tensors(x))
+    TensorNetwork_from_vec(v) = T(back(v))
 
     return x_vec, TensorNetwork_from_vec
 end

ext/TenetMakieExt.jl

@@ -51,17 +51,17 @@ function Makie.plot!(ax::Union{Axis,Axis3}, @nospecialize tn::AbstractTensorNetw
     tn = transform(tn, Tenet.HyperindConverter)
 
     # TODO how to mark multiedges? (i.e. parallel edges)
-    graph = SimpleGraph([Edge(tensors...) for (_, tensors) in tn.indices if length(tensors) > 1])
+    graph = SimpleGraph([Edge(tensors...) for (_, tensors) in tn.indexmap if length(tensors) > 1])
 
     # TODO recognise `copytensors` by using `DeltaArray` or `Diagonal` representations
     copytensors = findall(tensor -> any(flatinds -> issetequal(inds(tensor), flatinds), keys(hypermap)), tensors(tn))
-    ghostnodes = map(inds(tn, :open)) do ind
+    ghostnodes = map(inds(tn, :open)) do index
         # create new ghost node
         add_vertex!(graph)
         node = nv(graph)
 
         # connect ghost node
-        tensor = only(tn.indices[ind])
+        tensor = only(tn.indexmap[index])
         add_edge!(graph, node, tensor)
 
         return node

src/TensorNetwork.jl

@@ -39,7 +39,7 @@ TensorNetwork() = TensorNetwork(Tensor[])
 
 Return a shallow copy of a [`TensorNetwork`](@ref).
 """
-Base.copy(tn::T) where {T<:AbstractTensorNetwork} = TensorNetwork(copy(tn.indexmap), copy(tn.tensormap))
+Base.copy(tn::T) where {T<:AbstractTensorNetwork} = TensorNetwork(tensors(tn))
 
 Base.summary(io::IO, tn::AbstractTensorNetwork) = print(io, "$(length(tn.tensormap))-tensors $(typeof(tn))")
 Base.show(io::IO, tn::AbstractTensorNetwork) =
@@ -115,7 +115,7 @@ function Base.push!(tn::AbstractTensorNetwork, tensor::Tensor)
     end
 
     tn.tensormap[tensor] = collect(inds(tensor))
-    for index in inds(tensor)
+    for index in unique(inds(tensor))
         push!(get!(tn.indexmap, index, Tensor[]), tensor)
     end
 
@@ -174,7 +174,7 @@ Base.delete!(tn::AbstractTensorNetwork, x) = (_ = pop!(tn, x); tn)
 tryprune!(tn::AbstractTensorNetwork, i::Symbol) = (x = isempty(tn.indexmap[i]) && delete!(tn.indexmap, i); x)
 
 function Base.delete!(tn::AbstractTensorNetwork, tensor::Tensor)
-    for index in inds(tensor)
+    for index in unique(inds(tensor))
         filter!(Base.Fix1(!==, tensor), tn.indexmap[index])
         tryprune!(tn, index)
     end
@@ -212,16 +212,31 @@ function Base.replace!(tn::AbstractTensorNetwork, pair::Pair{<:Tensor,<:Tensor})
     return tn
 end
 
+function Base.replace!(tn::AbstractTensorNetwork, old_new::Pair{Symbol,Symbol}...)
+    first.(old_new) ⊆ keys(tn.indexmap) ||
+        throw(ArgumentError("set of old indices must be a subset of current indices"))
+    isdisjoint(last.(old_new), keys(tn.indexmap)) ||
+        throw(ArgumentError("set of new indices must be disjoint to current indices"))
+    for pair in old_new
+        replace!(tn, pair)
+    end
+    return tn
+end
+
 function Base.replace!(tn::AbstractTensorNetwork, old_new::Pair{Symbol,Symbol})
     old, new = old_new
     old ∈ keys(tn.indexmap) || throw(ArgumentError("index $old does not exist"))
     new ∉ keys(tn.indexmap) || throw(ArgumentError("index $new is already present"))
 
-    for tensor in tn.indexmap[old]
-        delete!(tn, tensor)
+    # NOTE `copy` because collection underneath is mutated
+    for tensor in copy(tn.indexmap[old])
+        # NOTE do not `delete!` before `push!` as indices can be lost due to `tryprune!`
         push!(tn, replace(tensor, old_new))
+        delete!(tn, tensor)
     end
 
+    delete!(tn.indexmap, old)
+
     return tn
 end
 
@@ -246,7 +261,7 @@ Return tensors whose indices match with the list of indices `i`.
 select(tn::AbstractTensorNetwork, i::Symbol) = copy(tn.indexmap[i])
 select(tn::AbstractTensorNetwork, is::AbstractVecOrTuple{Symbol}) =
     filter(tn.indexmap[first(is)]) do tensor
-        issetequal(inds(tensor), is)
+        is ⊆ inds(tensor)
     end
 
 """

src/Transformations.jl

@@ -176,9 +176,7 @@ end
 function transform!(tn::AbstractTensorNetwork, config::AntiDiagonalGauging)
     skip_inds = isempty(config.skip) ? inds(tn, set = :open) : config.skip
 
-    for idx in keys(tn.tensors)
-        tensor = tn.tensors[idx]
-
+    for tensor in keys(tn.tensormap)
         anti_diag_axes = find_anti_diag_axes(parent(tensor), atol = config.atol)
 
         for (i, j) in anti_diag_axes # loop over all anti-diagonal axes
@@ -215,56 +213,14 @@ end
 function transform!(tn::AbstractTensorNetwork, config::ColumnReduction)
     skip_inds = isempty(config.skip) ? inds(tn, set = :open) : config.skip
 
-    for tensor in tn.tensors
-        zero_columns = find_zero_columns(parent(tensor), atol = config.atol)
-        zero_columns_by_axis = [filter(x -> x[1] == d, zero_columns) for d in 1:length(size(tensor))]
-
-        # find non-zero column for each axis
-        non_zero_columns =
-            [(d, setdiff(1:size(tensor, d), [x[2] for x in zero_columns_by_axis[d]])) for d in 1:length(size(tensor))]
-
-        # remove axes that have more than one non-zero column
-        axes_to_reduce = [(d, c[1]) for (d, c) in filter(x -> length(x[2]) == 1, non_zero_columns)]
-
-        # First try to reduce the whole index if only one column is non-zeros
-        for (d, c) in axes_to_reduce # loop over all column axes
-            ix_i = inds(tensor)[d]
-
-            # do not reduce output indices
-            if ix_i ∈ skip_inds
-                continue
-            end
+    for tensor in tensors(tn)
+        for (dim, index) in enumerate(inds(tensor))
+            index ∈ skip_inds && continue
 
-            # reduce all tensors where ix_i appears
-            for (ind, t) in enumerate(tensors(tn))
-                if ix_i ∈ inds(t)
-                    # Replace the tensor with the reduced one
-                    new_tensor = selectdim(parent(t), findfirst(l -> l == ix_i, inds(t)), c)
-                    new_inds = filter(l -> l != ix_i, inds(t))
+            zeroslices = iszero.(eachslice(tensor, dims = dim))
+            any(zeroslices) || continue
 
-                    tn.tensors[ind] = Tensor(new_tensor, new_inds)
-                end
-            end
-            delete!(tn.indices, ix_i)
-        end
-
-        # Then try to reduce the dimensionality of the index in the other tensors
-        zero_columns = find_zero_columns(parent(tensor), atol = config.atol)
-        for (d, c) in zero_columns # loop over all column axes
-            ix_i = inds(tensor)[d]
-
-            # do not reduce output indices
-            if ix_i ∈ skip_inds
-                continue
-            end
-
-            # reduce all tensors where ix_i appears
-            for (ind, t) in enumerate(tensors(tn))
-                if ix_i ∈ inds(t)
-                    reduced_dims = [i == ix_i ? filter(j -> j != c, 1:size(t, i)) : (1:size(t, i)) for i in inds(t)]
-                    tn.tensors[ind] = Tensor(view(parent(t), reduced_dims...), inds(t))
-                end
-            end
+            slice!(tn, index, count(!, zeroslices) == 1 ? findfirst(!, zeroslices) : findall(!, zeroslices))
         end
     end
 
@@ -321,29 +277,6 @@ function transform!(tn::AbstractTensorNetwork, config::SplitSimplification)
     return tn
 end
 
-function find_zero_columns(x; atol = 1e-12)
-    dims = size(x)
-
-    # Create an initial set of all possible column pairs
-    zero_columns = Set((d, c) for d in 1:length(dims) for c in 1:dims[d])
-
-    # Iterate over each element in tensor
-    for index in CartesianIndices(x)
-        val = x[index]
-
-        # For each non-zero element, eliminate the corresponding column from the zero_columns set
-        if abs(val) > atol
-            for d in 1:length(dims)
-                c = index[d]
-                delete!(zero_columns, (d, c))
-            end
-        end
-    end
-
-    # Now the zero_columns set only contains column pairs where all elements are zero
-    return collect(zero_columns)
-end
-
 function find_diag_axes(x; atol = 1e-12)
     # skip 1D tensors
     ndims(parent(x)) == 1 && return []