Fix rand function for odd-length MatrixProductState (#54)

* Fix rand MPS function * Remove unwanted print * Fix tests * Update src/Quantum/MatrixProductState.jl --------- Co-authored-by: Sergio Sánchez Ramírez <[email protected]>
bsc-quantic · Apr 27, 2023 · 1f62855 · 1f62855
1 parent 4e40747
commit 1f62855
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/src/Quantum/MatrixProductState.jl b/src/Quantum/MatrixProductState.jl
@@ -48,6 +48,7 @@ function MatrixProductState{Open}(arrays; χ = nothing, order = (:l, :r, :p), me
 
         invpermute!(labels, permutator)
         alias = Dict([x => y for (x, y) in zip(invpermute!(original_order, permutator), labels)])
+
         push!(tn, Tensor(data, labels; alias = alias))
     end
 
@@ -151,8 +152,8 @@ function Base.rand(rng::Random.AbstractRNG, sampler::MPSSampler{Open,T}) where {
             χl = min(χ, p^(i - 1))
             χr = min(χ, p^i)
 
-            # swap bond dims after mid
-            after_mid ? (χr, χl) : (χl, χr)
+            # swap bond dims after mid and handle midpoint for odd-length MPS
+            (isodd(n) && i == n ÷ 2 + 1) ? (χl, χl) : (after_mid ? (χr, χl) : (χl, χr))
         end
 
         # fix for first site

diff --git a/test/MatrixProductState_test.jl b/test/MatrixProductState_test.jl
@@ -55,6 +55,26 @@
                 MatrixProductState{Open}(arrays) isa TensorNetwork{MatrixProductState{Open}}
             end
 
+            @testset "rand" begin
+                function max_size(t)
+                    labels = (get(t.meta[:alias], :l, :none), get(t.meta[:alias], :r, :none))
+                    return maximum(size(t, label) for label in labels if label != :none)
+                end
+
+                # Test with χ > maximum possible χ for the given parameters
+                ψ = rand(MatrixProductState{Open}, 7, 2, 32)
+
+                @test ψ isa TensorNetwork{MatrixProductState{Open}}
+                @test length(ψ) == 7
+                @test maximum(max_size(t) for t in ψ.tensors) <= 32
+
+                # Test with χ < maximum possible χ for the given parameters
+                ψ = rand(MatrixProductState{Open}, 7, 2, 4)
+                @test ψ isa TensorNetwork{MatrixProductState{Open}}
+                @test length(ψ) == 7
+                @test maximum(max_size(t) for t in ψ.tensors) <= 4
+            end
+
             @testset "Metadata" begin
                 @testset "alias" begin
                     arrays = [rand(2, 2), rand(2, 2, 2), rand(2, 2)]