Refactor Ansatz to traits

Project.toml

@@ -10,6 +10,7 @@ DeltaArrays = "10b0fc19-5ccc-4427-889b-d75dd6306188"
 EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+ImmutableArrays = "667c17eb-ab9b-4487-935f-1c621bb82497"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Muscle = "21fe5c4b-a943-414d-bf3e-516f24900631"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"

docs/src/quantum/index.md

@@ -15,7 +15,7 @@ sites
 ```
 
 ```@docs
-tensors(::TensorNetwork{<:Quantum}, ::Integer)
+tensors(::TensorNetwork{Quantum}, ::Integer)
 ```
 
 ```@docs
@@ -31,14 +31,14 @@ adjoint
 ## Concatenation
 
 ```@docs
-hcat(::TensorNetwork{<:Quantum}, ::TensorNetwork{<:Quantum})
+hcat(::TensorNetwork{Quantum}, ::TensorNetwork{Quantum})
 ```
 
 ## Norm
 
 ```@docs
-LinearAlgebra.norm(::TensorNetwork{<:Quantum}, p::Real)
-LinearAlgebra.normalize!(::TensorNetwork{<:Quantum}, ::Real)
+LinearAlgebra.norm(::TensorNetwork{Quantum}, p::Real)
+LinearAlgebra.normalize!(::TensorNetwork{Quantum}, ::Real)
 ```
 
 ## Fidelity

docs/src/tensor-network.md

@@ -32,7 +32,7 @@ inds(::TensorNetwork)
 size(::TensorNetwork)
 tensors(::TensorNetwork)
 length(::TensorNetwork)
-ansatz
+domain
 ```
 
 ## Modification

ext/TenetChainRulesCoreExt.jl

@@ -4,18 +4,18 @@ using Tenet
 using ChainRulesCore
 
 function ChainRulesCore.ProjectTo(tensor::T) where {T<:Tensor}
-    ProjectTo{T}(; data = ProjectTo(tensor.data), inds = tensor.inds, meta = tensor.meta)
+    ProjectTo{T}(; data = ProjectTo(tensor.data), inds = tensor.inds)
 end
 
 function (projector::ProjectTo{T})(dx::Union{T,Tangent{T}}) where {T<:Tensor}
-    T(projector.data(dx.data), projector.inds; projector.meta...)
+    T(projector.data(dx.data), projector.inds)
 end
 
-ChainRulesCore.frule((_, Δ, _), T::Type{<:Tensor}, data, inds; meta...) = T(data, inds; meta...), T(Δ, inds; meta...)
+ChainRulesCore.frule((_, Δ, _), T::Type{<:Tensor}, data, inds) = T(data, inds), T(Δ, inds)
 
 Tensor_pullback(Δ) = (NoTangent(), Δ.data, NoTangent())
 Tensor_pullback(Δ::AbstractThunk) = Tensor_pullback(unthunk(Δ))
-ChainRulesCore.rrule(T::Type{<:Tensor}, data, inds; meta...) = T(data, inds; meta...), Tensor_pullback
+ChainRulesCore.rrule(T::Type{<:Tensor}, data, inds) = T(data, inds), Tensor_pullback
 
 # NOTE fix problem with vector generator in `contract`
 @non_differentiable Tenet.__omeinsum_sym2str(x)
@@ -35,10 +35,10 @@ function (projector::ProjectTo{T})(dx::Union{T,Tangent{T}}) where {T<:TensorNetw
     Tangent{TensorNetwork}(tensors = projector.tensors(dx.tensors))
 end
 
-function Base.:+(x::TensorNetwork{A}, Δ::Tangent{TensorNetwork}) where {A<:Ansatz}
+function Base.:+(x::TensorNetwork{D}, Δ::Tangent{TensorNetwork}) where {D}
     # TODO match tensors by indices
     tensors = map(+, x.tensors, Δ.tensors)
-    TensorNetwork{A}(tensors; x.metadata...)
+    TensorNetwork{D}(tensors; x.metadata...)
 end
 
 function ChainRulesCore.frule((_, Δ), T::Type{<:TensorNetwork}, tensors; metadata...)

ext/TenetFiniteDifferencesExt.jl

@@ -3,11 +3,11 @@ module TenetFiniteDifferencesExt
 using Tenet
 using FiniteDifferences
 
-function FiniteDifferences.to_vec(x::TensorNetwork{A}) where {A<:Ansatz}
+function FiniteDifferences.to_vec(x::TensorNetwork{D}) where {D}
     x_vec, back = to_vec(x.tensors)
     function TensorNetwork_from_vec(v)
         tensors = back(v)
-        TensorNetwork{A}(tensors; x.metadata...)
+        TensorNetwork{D}(tensors; x.metadata...)
     end
 
     return x_vec, TensorNetwork_from_vec

ext/TenetMakieExt.jl

@@ -19,7 +19,7 @@ Plot a [`TensorNetwork`](@ref) as a graph.
   - `labels` If `true`, show the labels of the tensor indices. Defaults to `false`.
   -  The rest of `kwargs` are passed to `GraphMakie.graphplot`.
 """
-function Makie.plot(tn::TensorNetwork; kwargs...)
+function Makie.plot(@nospecialize tn::TensorNetwork; kwargs...)
     f = Figure()
     ax, p = plot!(f[1, 1], tn; kwargs...)
     return Makie.FigureAxisPlot(f, ax, p)
@@ -28,7 +28,7 @@ end
 # NOTE this is a hack! we did it in order not to depend on NetworkLayout but can be unstable
 __networklayout_dim(x) = typeof(x).super.parameters |> first
 
-function Makie.plot!(f::Union{Figure,GridPosition}, tn::TensorNetwork; kwargs...)
+function Makie.plot!(f::Union{Figure,GridPosition}, @nospecialize tn::TensorNetwork; kwargs...)
     ax = if haskey(kwargs, :layout) && __networklayout_dim(kwargs[:layout]) == 3
         Axis3(f[1, 1])
     else
@@ -45,14 +45,15 @@ function Makie.plot!(f::Union{Figure,GridPosition}, tn::TensorNetwork; kwargs...
     return Makie.AxisPlot(ax, p)
 end
 
-function Makie.plot!(ax::Union{Axis,Axis3}, tn::TensorNetwork; labels = false, kwargs...)
+function Makie.plot!(ax::Union{Axis,Axis3}, @nospecialize tn::TensorNetwork; labels = false, kwargs...)
+    hypermap = Tenet.hyperflatten(tn)
     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])
 
     # TODO recognise `copytensors` by using `DeltaArray` or `Diagonal` representations
-    copytensors = findall(t -> haskey(t.meta, :dual), tensors(tn))
+    copytensors = findall(tensor -> any(flatinds -> issetequal(inds(tensor), flatinds), values(hypermap)), tensors(tn))
     ghostnodes = map(inds(tn, :open)) do ind
         # create new ghost node
         add_vertex!(graph)

ext/TenetQuacExt.jl

@@ -2,7 +2,7 @@ module TenetQuacExt
 
 using Tenet
 using Quac: Circuit, lanes, arraytype, Swap
-using Bijections
+using Bijection: Bijection
 
 function Tenet.TensorNetwork(circuit::Circuit)
     n = lanes(circuit)
@@ -28,16 +28,16 @@ function Tenet.TensorNetwork(circuit::Circuit)
             (from, to)
         end |> x -> zip(x...) |> Iterators.flatten |> collect
 
-        tensor = Tensor(array, tuple(inds...); gate = gate)
+        tensor = Tensor(array, inds)
         push!(tensors, tensor)
     end
 
-    interlayer = [
+    plug = [
         Bijection(Dict([site => first(index) for (site, index) in enumerate(wire)])),
         Bijection(Dict([site => last(index) for (site, index) in enumerate(wire)])),
     ]
 
-    return TensorNetwork{Quantum}(tensors; plug = Tenet.Operator, interlayer)
+    return TensorNetwork{Quantum}(tensors; plug = Tenet.Operator, plug)
 end
 
 end

src/Helpers.jl

@@ -72,11 +72,6 @@ Base.merge(@nospecialize(A::Type{<:NamedTuple}), @nospecialize(Bs::Type{<:NamedT
     foldl((acc, B) -> Tuple{fieldtypes(acc)...,B...}, Iterators.map(fieldtypes, Bs); init = A),
 }
 
-function superansatzes(T)
-    S = supertype(T)
-    return T === Ansatz ? (T,) : (T, superansatzes(S)...)
-end
-
 # NOTE from https://stackoverflow.com/q/54652787
 function nonunique(x)
     uniqueindexes = indexin(unique(x), x)

src/Numerics.jl

@@ -51,19 +51,17 @@ function contract(a::Tensor, b::Tensor; dims = (∩(inds(a), inds(b))))
 
     data = EinCode((_ia, _ib), _ic)(parent(a), parent(b))
 
-    # TODO merge metadata?
     return Tensor(data, ic)
 end
 
 function contract(a::Tensor; dims = nonunique(inds(a)))
     ia = inds(a)
     i = ∩(dims, ia)
 
-    ic = tuple(setdiff(ia, i isa Base.AbstractVecOrTuple ? i : (i,))...)
+    ic = setdiff(ia, i isa Base.AbstractVecOrTuple ? i : (i,))
 
     data = EinCode((String.(ia),), String.(ic))(parent(a))
 
-    # TODO merge metadata
     return Tensor(data, ic)
 end
 
@@ -79,8 +77,8 @@ contract(tensors::Tensor...; kwargs...) = reduce((x, y) -> contract(x, y; kwargs
 Alias for [`contract`](@ref).
 """
 Base.:*(a::Tensor, b::Tensor) = contract(a, b)
-Base.:*(a::T, b::Number) where {T<:Tensor} = T(parent(a) * b, inds(a); a.meta...)
-Base.:*(a::Number, b::T) where {T<:Tensor} = T(a * parent(b), inds(b); b.meta...)
+Base.:*(a::T, b::Number) where {T<:Tensor} = T(parent(a) * b, inds(a))
+Base.:*(a::Number, b::T) where {T<:Tensor} = T(a * parent(b), inds(b))
 
 LinearAlgebra.svd(t::Tensor{<:Any,2}; kwargs...) = Base.@invoke svd(t::Tensor; left_inds = (first(inds(t)),), kwargs...)