[ITensors] [ENHANCMENT] Introduce Algorithm type for selecting algorithm backends (#886)

* Use new `Algorithm` type in `contract(::MPO, ::MPS)` and `truncate!(::MPS/MPO)`.

Author: mtfishman

ITensorGPU/test/test_cumpo.jl

@@ -101,7 +101,7 @@ using ITensors, ITensorGPU, Test
     psi = randomCuMPS(sites)
     psi_out = contract(K, psi; maxdim=1)
     @test inner(phi', psi_out) ≈ inner(phi', K, psi)
-    @test_throws ArgumentError contract(K', psi, method="fakemethod")
+    @test_throws MethodError contract(K', psi, method="fakemethod")
 
     badsites = [Index(2, "Site") for n in 1:(N + 1)]
     badpsi = randomCuMPS(badsites)

NEWS.md

@@ -6,8 +6,17 @@ Note that as of Julia v1.5, in order to see deprecation warnings you will need t
 
 After we release v1 of the package, we will start following [semantic versioning](https://semver.org).
 
+ITensors v0.3.1 Release Notes
+=============================
+
+Bugs:
+
+Enhancements:
+
+- Introduce `Algorithm` type for selecting algorithm backends (#886)
+
 ITensors v0.3.0 Release Notes
-==============================
+=============================
 
 Bugs:
 

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensors"
 uuid = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>", "Miles Stoudenmire <mstoudenmire@flatironinstitute.org>"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/ITensors.jl

@@ -103,6 +103,13 @@ include("imports.jl")
 #
 include("global_variables.jl")
 
+#####################################
+# Algorithm type for selecting
+# different algorithm backends
+# (for internal or advanced usage)
+#
+include("algorithm.jl")
+
 #####################################
 # Index and IndexSet
 #

src/algorithm.jl

@@ -0,0 +1,29 @@
+"""
+    Algorithm
+
+A type representing an algorithm backend for a function.
+
+For example, ITensor provides multiple backend algorithms for contracting
+an MPO with an MPS, which internally are selected with an `Algorithm` type.
+
+This allows users to extend functions in ITensor with new algorithms, but
+use the same interface.
+"""
+struct Algorithm{Alg} end
+
+Algorithm(s) = Algorithm{Symbol(s)}()
+algorithm(::Algorithm{Alg}) where {Alg} = string(Alg)
+
+show(io::IO, alg::Algorithm) = print(io, "Algorithm type ", algorithm(alg))
+print(io::IO, ::Algorithm{Alg}) where {Alg} = print(io, Alg)
+
+"""
+    @Algorithm_str
+
+A convenience macro for writing [`Algorithm`](@ref) types, typically used when
+adding methods to a function in ITensor that supports multiple algorithm
+backends (like contracting an MPO with an MPS).
+"""
+macro Algorithm_str(s)
+  return :(Algorithm{$(Expr(:quote, Symbol(s)))})
+end

src/imports.jl

@@ -48,6 +48,7 @@ import Base:
   map!,
   ndims,
   permutedims,
+  print,
   promote_rule,
   push!,
   resize!,

src/mps/abstractmps.jl

@@ -1500,7 +1500,11 @@ Perform a truncation of all bonds of an MPS/MPO,
 using the truncation parameters (cutoff,maxdim, etc.)
 provided as keyword arguments.
 """
-function truncate!(M::AbstractMPS; kwargs...)
+function truncate!(M::AbstractMPS; alg="frobenius", kwargs...)
+  return truncate!(Algorithm(alg), M; kwargs...)
+end
+
+function truncate!(::Algorithm"frobenius", M::AbstractMPS; kwargs...)
   N = length(M)
 
   # Left-orthogonalize all tensors to make

src/mps/mpo.jl

@@ -479,28 +479,23 @@ end
 
 (A::MPO)(ψ::MPS; kwargs...) = apply(A, ψ; kwargs...)
 
-function contract(A::MPO, ψ::MPS; kwargs...)
-  method = get(kwargs, :method, "densitymatrix")
+function contract(A::MPO, ψ::MPS; alg="densitymatrix", kwargs...)
+  if haskey(kwargs, :method)
+    # Backwards compatibility, use `method`.
+    alg = get(kwargs, :method, "densitymatrix")
+  end
 
   # Keyword argument deprecations
-  if method == "DensityMatrix"
+  if alg == "DensityMatrix"
     @warn "In contract, method DensityMatrix is deprecated in favor of densitymatrix"
-    method = "densitymatrix"
+    alg = "densitymatrix"
   end
-
-  if method == "Naive"
-    @warn "In contract, method Naive is deprecated in favor of naive"
-    method = "naive"
+  if alg == "Naive"
+    @warn "In contract, `alg=\"Naive\"` is deprecated in favor of `alg=\"naive\"`"
+    alg = "naive"
   end
 
-  if method == "densitymatrix"
-    Aψ = _contract_densitymatrix(A, ψ; kwargs...)
-  elseif method == "naive"
-    Aψ = _contract_naive(A, ψ; kwargs...)
-  else
-    throw(ArgumentError("Method $method not supported"))
-  end
-  return Aψ
+  return contract(Algorithm(alg), A, ψ; kwargs...)
 end
 
 contract_mpo_mps_doc = """
@@ -555,7 +550,7 @@ contract(ψ::MPS, A::MPO; kwargs...) = contract(A, ψ; kwargs...)
 
 #@doc (@doc contract(::MPO, ::MPS)) *(::MPO, ::MPS)
 
-function _contract_densitymatrix(A::MPO, ψ::MPS; kwargs...)::MPS
+function contract(::Algorithm"densitymatrix", A::MPO, ψ::MPS; kwargs...)::MPS
   n = length(A)
   n != length(ψ) &&
     throw(DimensionMismatch("lengths of MPO ($n) and MPS ($(length(ψ))) do not match"))
@@ -636,7 +631,7 @@ function _contract_densitymatrix(A::MPO, ψ::MPS; kwargs...)::MPS
   return ψ_out
 end
 
-function _contract_naive(A::MPO, ψ::MPS; kwargs...)::MPS
+function contract(::Algorithm"naive", A::MPO, ψ::MPS; kwargs...)::MPS
   truncate = get(kwargs, :truncate, true)
 
   N = length(A)

test/algorithm.jl

@@ -0,0 +1,39 @@
+using ITensors
+using Test
+
+@testset "Algorithm" begin
+  alg = ITensors.Algorithm("X")
+
+  @test alg isa ITensors.Algorithm"X"
+  @test alg == ITensors.Algorithm"X"()
+
+  s = siteinds("S=1/2", 4)
+  A = MPO(s, "Id")
+  ψ = randomMPS(s)
+
+  @test_throws MethodError contract(alg, A, ψ)
+  @test_throws MethodError contract(A, ψ; method="X")
+  @test_throws MethodError contract(A, ψ; alg="X")
+  @test contract(ITensors.Algorithm("densitymatrix"), A, ψ) ≈ A * ψ
+  @test contract(ITensors.Algorithm("naive"), A, ψ) ≈ A * ψ
+  @test contract(A, ψ; alg="densitymatrix") ≈ A * ψ
+  @test contract(A, ψ; method="densitymatrix") ≈ A * ψ
+  @test contract(A, ψ; alg="naive") ≈ A * ψ
+  @test contract(A, ψ; method="naive") ≈ A * ψ
+
+  B = copy(A)
+  truncate!(ITensors.Algorithm("frobenius"), B)
+  @test A ≈ B
+
+  B = copy(A)
+  truncate!(B; alg="frobenius")
+  @test A ≈ B
+
+  # Custom algorithm
+  function ITensors.truncate!(::ITensors.Algorithm"my_new_algorithm", A::MPO; cutoff=1e-15)
+    return "my_new_algorithm was called with cutoff $cutoff"
+  end
+  cutoff = 1e-5
+  res = truncate!(A; alg="my_new_algorithm", cutoff=cutoff)
+  @test res == "my_new_algorithm was called with cutoff $cutoff"
+end

test/mpo.jl

@@ -203,7 +203,7 @@ end
     psi = randomMPS(sites)
     psi_out = contract(K, psi; maxdim=1)
     @test inner(phi', psi_out) ≈ inner(phi', K, psi)
-    @test_throws ArgumentError contract(K, psi; method="fakemethod")
+    @test_throws MethodError contract(K, psi; method="fakemethod")
 
     badsites = [Index(2, "Site") for n in 1:(N + 1)]
     badpsi = randomMPS(badsites)