Merge 5829669 into 230ac30

src/MPO.jl

@@ -2,62 +2,70 @@ export MPO
 
 struct MPO{T <: AbstractArray{<:Number, 4}}
     tensors::Vector{T}
+
+    function MPO{T}(L::Int) where {T}
+        new(Vector{T}(undef, L))
+    end
+
+    MPO{T}(v::Vector{T}) where {T} = new(v)
 end
 
 # consturctors
 
-newDim(::Type{T}, dims) where {T<:AbstractArray} = T.name.wrapper{eltype(T), dims}
+MPO(::Type{T}, ::Type{S}, L::Int) where {T<:AbstractArray, S<:Number} = MPO{T{S, 4}}(L)
 
-function MPO(ψ::MPS{T}) where T <: AbstractArray{<:Number, 3}
+function MPO(ψ::MPS{T}) where {T <: AbstractArray{<:Number, 3}}
+    _verify_square(ψ)
     L = length(ψ)
-
-    S = newDim(T, 4) 
-    tensors = Vector{S}(undef, L)
+    O = MPO(T.name.wrapper, eltype(T), L)
 
     for i ∈ 1:L
         A = ψ[i]
-        dd = size(A, 2)
-        d = isqrt(dd)
+        d = isqrt(size(A, 2))
 
-        if d^2 == dd
-            @cast W[x, σ, y, η] |= A[x, (σ, η), y] (σ:d) 
-            tensors[i] = W
-        else
-            error("$dd is not a square number.")
-        end        
+        @cast W[x, σ, y, η] |= A[x, (σ, η), y] (σ:d) 
+        O[i] = W
     end 
-    MPO(tensors) 
+    O
 end 
 
 function MPS(O::MPO{T}) where T <: AbstractArray{<:Number, 4}
     L = length(O)
-
-    S = newDim(T, 3) 
-    tensors = Vector{S}(undef, L)
+    ψ = MPS(T.name.wrapper, eltype(T), L)
 
     for i ∈ 1:L
         W = O[i]
         @cast A[x, (σ, η), y] := W[x, σ, y, η]
-        tensors[i] = A     
+        ψ[i] = A     
     end 
-    MPS(tensors) 
+    ψ
 end  
 
 function Base.randn(::Type{MPO{T}}, L::Int, D::Int, d::Int) where T <: AbstractArray{<:Number, 4}
-    S = AbstractArray{eltype(T), 3} # !!! This probably can be done better !!!
+    S = newdim(T, 3)
     ψ = randn(MPS{S}, L, D, d^2) 
     MPO(ψ)
 end
 
+function _verify_square(ψ::MPS)
+    arr = [size(A, 2) for A ∈ ψ]
+    @assert isqrt.(arr) .^ 2 == arr "Incorrect MPS dimensions"
+end
+
+
 # length, comparison, element types
 
 Base.:(==)(O::MPO, U::MPO) = O.tensors == U.tensors
 Base.:(≈)(O::MPO, U::MPO)  = O.tensors ≈ U.tensors
 
-Base.eltype(::Type{MPO{T}}) where {T <: AbstractArray{S, 4}} where {S <: Number} = S
+Base.eltype(::Type{MPO{T}}) where {T} = eltype(T)
+
+Base.getindex(O::MPO, i::Int) = getindex(O.tensors, i)
+Base.setindex!(O::MPO, W::AbstractArray{<:Number, 4}, i::Int) = O.tensors[i] = W
 
-Base.getindex(O::MPO) = getindex(O.tensors, )
-Base.setindex(O::MPO, i::Int) = setindex(O.tensors, i)
+Base.iterate(O::MPO) = iterate(O.tensors)
+Base.iterate(O::MPO, state) = iterate(O.tensors, state)
+Base.lastindex(O::MPO) = lastindex(O.tensors)
 
 Base.length(O::MPO) = length(O.tensors)
 Base.size(O::MPO) = (length(O.tensors), )

src/MPS.jl

@@ -2,50 +2,50 @@ export MPS
 
 struct MPS{T <: AbstractArray{<:Number, 3}}
     tensors::Vector{T}
+
+    function MPS{T}(L::Int) where {T}
+        new(Vector{T}(undef, L))
+    end
+
+    MPS{T}(v::Vector{T}) where {T} = new(v)
 end
 
 # consturctors
+MPS(::Type{T}, ::Type{S}, L::Int) where {T<:AbstractArray, S<:Number} = MPS{T{S, 3}}(L)
 
-function Base.randn(::Type{MPS{T}}, L::Int, D::Int, d::Int) where T <: AbstractArray{<:Number, 3}
-    tensors = Vector{T}(undef, L)
-    S = eltype(T)
 
-    tensors[1] = randn(S, 1, d, D)
+function Base.randn(::Type{MPS{T}}, L::Int, D::Int, d::Int) where {T}
+    ψ = MPS{T}(L)
+    S = eltype(T)
+    ψ[1] = randn(S, 1, d, D)
     for i ∈ 2:(L-1)
-        tensors[i] = randn(S, D, d, D)
-    end
-    tensors[end] = randn(S, D, d, 1)
-    MPS(tensors) 
-end
-
-function MPS(ψ::Vector{T}) where T <: AbstractVector{<:Number}
-    L = length(ψ)
-    tensors = Vector{T}(undef, L)
-
-    for i ∈ 1:L
-        tensors[i] = reshape(copy(ψ[i]), 1, :, 1)
+        ψ[i] = randn(S, D, d, D)
     end
-    MPS(tensors)
+    ψ[end] = randn(S, D, d, 1)
+    ψ
 end
 
 # length, comparison, element types
 
 Base.:(==)(ϕ::MPS, ψ::MPS) = ψ.tensors == ϕ.tensors
 Base.:(≈)(ϕ::MPS, ψ::MPS)  = ψ.tensors ≈ ϕ.tensors
 
-Base.eltype(::Type{MPS{T}}) where {T <: AbstractArray{S, 3}} where {S <: Number} = S
+Base.eltype(::Type{MPS{T}}) where {T} = eltype(T)
 
 Base.getindex(ψ::MPS, i::Int) = getindex(ψ.tensors, i)
-Base.setindex(ψ::MPS, i::Int) = setindex(ψ.tensors, i)
+Base.setindex!(ψ::MPS, A::AbstractArray{<:Number, 3}, i::Int) = ψ.tensors[i] = A
+Base.iterate(ψ::MPS) = iterate(ψ.tensors)
+Base.iterate(ψ::MPS, state) = iterate(ψ.tensors, state)
+Base.lastindex(ψ::MPS) = lastindex(ψ.tensors)
 
 Base.length(mps::MPS) = length(mps.tensors)
 Base.size(ψ::MPS) = (length(ψ.tensors), )
 
-Base.copy(ψ::MPS) = MPS(copy(ψ.tensors))
+Base.copy(ψ::MPS{T}) where {T} = MPS{T}(copy(ψ.tensors))
 
 # printing
 
-function Base.show(ψ::MPS{T}) where T <: AbstractArray{<:Number, 3}
+function Base.show(::IO, ψ::MPS)
     L = length(ψ)
     σ_list = [size(ψ[i], 2) for i ∈ 1:L] 
     χ_list = [size(ψ[i][:, 1, :]) for i ∈ 1:L]

src/compressions.jl

@@ -1,32 +1,21 @@
 export truncate!, canonise!
 
-function canonise!(ψ::MPS{T}) where T <: AbstractArray{<:Number, 3}
-    canonise!(ψ, "left")
-    canonise!(ψ, "right")
-end 
-
-function canonise!(ψ::MPS{T}, type::String) where T <: AbstractArray{<:Number, 3}
-    if type == "right"
-        _left_sweep_SVD!(ψ, typemax(Int))
-    elseif type == "left"
-        _right_sweep_SVD!(ψ, typemax(Int))
-    else
-        error("Choose: left or right")    
-    end    
+function canonise!(ψ::MPS)
+    canonise!(ψ, :left)
+    canonise!(ψ, :right)
 end
 
-function truncate!(ψ::MPS{T}, type::String, Dcut::Int) where T <: AbstractArray{<:Number, 3}
-    if type == "right"
-        _left_sweep_SVD!(ψ, Dcut)
-    elseif type == "left"
-        _right_sweep_SVD!(ψ, Dcut)
-    else
-        error("Choose: left or right")    
-    end    
-end
+canonise!(ψ::MPS, s::Symbol) = canonise!(ψ, Val(s))
+
+canonise!(ψ::MPS, ::Val{:right}) = _left_sweep_SVD!(ψ)
+canonise!(ψ::MPS, ::Val{:left}) = _right_sweep_SVD!(ψ)
+
+truncate!(ψ::MPS, s::Symbol, Dcut::Int) = truncate!(ψ, Val(s), Dcut)
+truncate!(ψ::MPS, ::Val{:right}, Dcut::Int) = _left_sweep_SVD!(ψ, Dcut)
+truncate!(ψ::MPS, ::Val{:left}, Dcut::Int) = _right_sweep_SVD!(ψ, Dcut)
 
-function _right_sweep_SVD!(ψ::MPS{T}, Dcut::Int) where T <: AbstractArray{<:Number, 3}
-    Σ = V = ones(eltype(T), 1, 1)
+function _right_sweep_SVD!(ψ::MPS, Dcut::Int=typemax(Int))
+    Σ = V = ones(eltype(ψ), 1, 1)
 
     for i ∈ 1:length(ψ)
 
@@ -43,12 +32,12 @@ function _right_sweep_SVD!(ψ::MPS{T}, Dcut::Int) where T <: AbstractArray{<:Num
         # create new    
         d = size(ψ[i], 2)
         @cast B[x, σ, y] |= U[(x, σ), y] (σ:d)
-        ψ.tensors[i] = B
+        ψ[i] = B
     end
 end
 
-function _left_sweep_SVD!(ψ::MPS{T}, Dcut::Int) where T <: AbstractArray{<:Number, 3}
-    Σ = U = ones(eltype(T), 1, 1)
+function _left_sweep_SVD!(ψ::MPS, Dcut::Int=typemax(Int))
+    Σ = U = ones(eltype(ψ), 1, 1)
 
     for i ∈ length(ψ):-1:1
 
@@ -65,6 +54,6 @@ function _left_sweep_SVD!(ψ::MPS{T}, Dcut::Int) where T <: AbstractArray{<:Numb
         # create new 
         d = size(ψ[i], 2)
         @cast A[x, σ, y] |= V'[x, (σ, y)] (σ:d)
-        ψ.tensors[i] = A
+        ψ[i] = A
     end
 end

src/contractions.jl

@@ -1,9 +1,8 @@
 
 export norm 
 
-function LinearAlgebra.dot(ϕ::MPS{T}, ψ::MPS{T}) where T <: AbstractArray{<:Number, 3} 
-    S = eltype(ψ)
-    C = ones(S, 1, 1)
+function LinearAlgebra.dot(ϕ::MPS, ψ::MPS)
+    C = ones(eltype(ψ), 1, 1)
 
     for i ∈ 1:length(ψ)
         M = ψ[i]
@@ -13,11 +12,11 @@ function LinearAlgebra.dot(ϕ::MPS{T}, ψ::MPS{T}) where T <: AbstractArray{<:Nu
     return C[1]
 end
 
-function LinearAlgebra.norm(ψ::MPS{T}) where T <: AbstractArray{<:Number, 3}
-    return sqrt(abs(dot(ψ,ψ)))
+function LinearAlgebra.norm(ψ::MPS)
+    return sqrt(abs(dot(ψ, ψ)))
 end
 
-function LinearAlgebra.dot(ϕ::MPS{T}, O::Vector{S}, ψ::MPS{T}) where {T <: AbstractArray{<:Number, 3}, S <: AbstractMatrix}
+function LinearAlgebra.dot(ϕ::MPS, O::Vector{S}, ψ::MPS) where {S <: AbstractMatrix}
     R = eltype(ψ)
     C = ones(R, 1, 1)
 
@@ -30,19 +29,20 @@ function LinearAlgebra.dot(ϕ::MPS{T}, O::Vector{S}, ψ::MPS{T}) where {T <: Abs
     return C[1]
 end
 
-function dot(O::MPO{T}, ψ::MPS{S}) where {T <: AbstractArray{<:Number, 4}, S <: AbstractArray{<:Number, 3}}
-    tensors = copy(ψ.tensors)
+function dot(O::MPO, ψ::MPS{T}) where {T}
+    L = length(ψ)
+    ϕ = MPS{T}(L)
 
-    for i in 1:length(O)
+    for i in 1:L
         W = O[i]
         M = ψ[i]
 
         @reduce N[(x, a), (y, b), σ] := sum(η) W[x, σ, y, η] * M[a, η, b]      
-        tensors[i] = N
+        ϕ[i] = N
     end
-    MPS(tensors)
+    ϕ
 end
 
-function Base.:(*)(O::MPO{T}, ψ::MPS{S}) where {T <: AbstractArray{<:Number, 4}, S <: AbstractArray{<:Number, 3}}
+function Base.:(*)(O::MPO, ψ::MPS)
     return dot(O, ψ)
 end

src/utils.jl

@@ -1,3 +1,3 @@
 export newdim
 
-# S = newdims(T, 4)
+newdim(::Type{T}, dims) where {T<:AbstractArray} = T.name.wrapper{eltype(T), dims}

test/compressions.jl

@@ -15,7 +15,7 @@ T = Array{ComplexF64, 3}
 χ = randn(MPS{T}, sites, D, d)
 
 @testset "Canonisation (left)" begin
-    canonise!(ψ, "left")  
+    canonise!(ψ, :left)  
     show(ψ)  
 
     is_left_normalized = true
@@ -32,7 +32,7 @@ T = Array{ComplexF64, 3}
 end
 
 @testset "Canonisation (right)" begin
-    canonise!(ϕ, "right")  
+    canonise!(ϕ, :right)  
     show(ϕ)
 
     is_right_normalized = true
@@ -55,13 +55,13 @@ end
 end
 
 @testset "Truncation (SVD, right)" begin
-    truncate!(ψ, "right", Dcut)  
+    truncate!(ψ, :right, Dcut)  
     show(ψ)
     @test dot(ψ, ψ) ≈ 1     
 end
 
 @testset "Truncation (SVD, left)" begin
-    truncate!(ψ, "left", Dcut)  
+    truncate!(ψ, :left, Dcut)  
     show(ψ)
     @test dot(ψ, ψ) ≈ 1     
 end

test/runtests.jl

@@ -12,6 +12,6 @@ for my_test in my_tests
     include(my_test)
 end
 
-if CUDA.functional()
+if CUDA.functional() && false #skip cuda for now
     include("cuda.jl")
 end