Implement an MPS method initializing the tensors to identity (copy-tensors)

src/Ansatz/MPS.jl

@@ -62,6 +62,45 @@ function MPS(arrays::Vector{<:AbstractArray}; order=defaultorder(MPS))
     return MPS(ansatz, NonCanonical())
 end
 
+"""
+    Base.identity(::Type{MPS}, n::Integer; physdim=2, maxdim=physdim^(n ÷ 2), order=MPS.defaultorder(MPS))
+
+Returns an [`MPS`](@ref) of `n` sites whose tensors are initialized to the identity attending to the
+physical dimension `physdim` and the maximum dimension `maxdim`.
+
+# Keyword Arguments
+
+  - `physdim` The physical or output dimension of each site. Default is 2
+  - `maxdim` The maximum bond dimension. Default is `physdim^(n ÷ 2)`.
+  - `order` Tensors' indices order. Default: output - left - right `(:o, :l, :r)`.
+"""
+function Base.identity(::Type{MPS}, n::Integer; physdim=2, maxdim=physdim^(n ÷ 2), order=defaultorder(MPS))
+    issetequal(order, defaultorder(MPS)) ||
+        throw(ArgumentError("order must be a permutation of $(String.(defaultorder(MPS)))"))
+
+    # Create bond dimensions until the middle of the MPS considering maxdim
+    virtualdims = physdim .^ collect(1:(n ÷ 2))
+    virtualdims = ifelse.((virtualdims .> maxdim), maxdim, virtualdims)
+    # Complete the bond dimensions of the other half of the MPS
+    virtualdims = vcat(virtualdims, reverse(n % 2 == 1 ? virtualdims : virtualdims[1:(end - 1)]))
+
+    # Create each site dimensions in default order (:o, :l, :r)
+    arraysdims = [[physdim, virtualdims[1]]]
+    append!(arraysdims, [[physdim, virtualdims[i], virtualdims[i + 1]] for i in 1:(length(virtualdims) - 1)])
+    push!(arraysdims, [physdim, virtualdims[end]])
+
+    # Create the MPS with copy-tensors according to the tensors dimensions
+    return MPS(
+        map(arraysdims) do arrdims
+            arr = zeros(ComplexF64, arrdims...)
+            deltas = [fill(i, length(arrdims)) for i in 1:physdim]
+            broadcast(delta -> arr[delta...] = 1.0, deltas)
+            arr
+        end;
+        order=order,
+    )
+end
+
 function Base.convert(::Type{MPS}, tn::Product)
     @assert socket(tn) == State()
 

test/Ansatz_test.jl

@@ -1,2 +1 @@
-@testset "Ansatz" begin
-end
+@testset "Ansatz" begin end

test/MPS_test.jl

@@ -26,6 +26,49 @@
     @test inds(ψ; at=site"3", dir=:left) == inds(ψ; at=site"2", dir=:right) !== nothing
     @test all(i -> size(ψ, inds(ψ; at=Site(i))) == 2, 1:nsites(ψ))
 
+    @testset "Base.identity" begin
+        nsites_cases = [6, 7, 6, 7]
+        physdim_cases = [3, 2, 3, 2]
+        maxdim_cases = [nothing, nothing, 9, 4] # nothing means default
+        expected_tensorsizes_cases = [
+            [(3, 3), (3, 3, 9), (3, 9, 27), (3, 27, 9), (3, 9, 3), (3, 3)],
+            [(2, 2), (2, 2, 4), (2, 4, 8), (2, 8, 8), (2, 8, 4), (2, 4, 2), (2, 2)],
+            [(3, 3), (3, 3, 9), (3, 9, 9), (3, 9, 9), (3, 9, 3), (3, 3)],
+            [(2, 2), (2, 2, 4), (2, 4, 4), (2, 4, 4), (2, 4, 4), (2, 4, 2), (2, 2)],
+        ]
+
+        for (nsites, physdim, expected_tensorsizes, maxdim) in
+            zip(nsites_cases, physdim_cases, expected_tensorsizes_cases, maxdim_cases)
+            ψ = if isnothing(maxdim)
+                identity(MPS, nsites; physdim=physdim)
+            else
+                identity(MPS, nsites; physdim=physdim, maxdim=maxdim)
+            end
+            # Test the tensor dimensions
+            obtained_tensorsizes = size.(tensors(ψ))
+            @test obtained_tensorsizes == expected_tensorsizes
+
+            # Test whether all tensors are the identity
+            alltns = tensors(ψ)
+            # - Test extreme tensors (2D) equal identity
+            diagonal_2D = [fill(i, 2) for i in 1:physdim]
+            @test all(delta -> alltns[1][delta...] == 1, diagonal_2D)
+            @test sum(alltns[1]) == physdim
+            @test all(delta -> alltns[end][delta...] == 1, diagonal_2D)
+            @test sum(alltns[end]) == physdim
+            # - Test middle tensors (3D) equal identity
+            diagonal_3D = [fill(i, 3) for i in 1:physdim]
+            @test all(tns -> all(delta -> tns[delta...] == 1, diagonal_3D), alltns[2:(end - 1)])
+            @test all(tns -> sum(tns) == physdim, alltns[2:(end - 1)])
+
+            # Test whether the contraction gives the identity
+            contracted_ψ = contract(ψ)
+            diagonal_nsitesD = [fill(i, nsites) for i in 1:physdim]
+            @test all(delta -> contracted_ψ[delta...] == 1, diagonal_nsitesD)
+            @test sum(contracted_ψ) == physdim
+        end
+    end
+
     @testset "Site" begin
         ψ = MPS([rand(2, 2), rand(2, 2, 2), rand(2, 2)])