[ITensorMPS] Allow customizing apply (#1440)

Author: mtfishman

src/lib/ITensorMPS/ext/ITensorMPSChainRulesCoreExt/abstractmps.jl

@@ -1,8 +1,8 @@
 using Adapt: adapt
-using ChainRulesCore: ChainRulesCore, HasReverseMode, NoTangent, RuleConfig
+using ChainRulesCore: ChainRulesCore, HasReverseMode, NoTangent, RuleConfig, rrule_via_ad
 using ITensors:
-  ITensor, apply, dag, hassameinds, inds, inner, itensor, mapprime, replaceprime, swapprime
-using ITensors.ITensorMPS: MPO, MPS, siteinds
+  ITensors, ITensor, dag, hassameinds, inds, itensor, mapprime, replaceprime, swapprime
+using ITensors.ITensorMPS: ITensorMPS, MPO, MPS, apply, inner, siteinds
 using NDTensors: datatype
 
 function ChainRulesCore.rrule(
@@ -186,7 +186,7 @@ function ChainRulesCore.rrule(
   end
   y = typeof(x)(y_data)
   if !set_limits
-    y = ITensors.set_ortho_lims(y, ortho_lims(x))
+    y = ITensorMPS.set_ortho_lims(y, ITensorMPS.ortho_lims(x))
   end
   return y, map_pullback
 end

src/lib/ITensorMPS/ext/ITensorMPSChainRulesCoreExt/mpo.jl

@@ -1,31 +1,35 @@
 using ChainRulesCore: ChainRulesCore, NoTangent
-using ITensors: contract, hassameinds, inner, mapprime
+using ITensors: Algorithm, contract, hassameinds, inner, mapprime
 using ITensors.ITensorMPS: MPO, MPS, firstsiteinds, siteinds
 using LinearAlgebra: tr
 
-function ChainRulesCore.rrule(::typeof(contract), x1::MPO, x2::MPO; kwargs...)
-  y = contract(x1, x2; kwargs...)
+function ChainRulesCore.rrule(
+  ::typeof(contract), alg::Algorithm, x1::MPO, x2::MPO; kwargs...
+)
+  y = contract(alg, x1, x2; kwargs...)
   function contract_pullback(ȳ)
-    x̄1 = contract(ȳ, dag(x2); kwargs...)
-    x̄2 = contract(dag(x1), ȳ; kwargs...)
-    return (NoTangent(), x̄1, x̄2)
+    x̄1 = contract(alg, ȳ, dag(x2); kwargs...)
+    x̄2 = contract(alg, dag(x1), ȳ; kwargs...)
+    return (NoTangent(), NoTangent(), x̄1, x̄2)
   end
   return y, contract_pullback
 end
 
-function ChainRulesCore.rrule(::typeof(contract), x1::MPO, x2::MPS; kwargs...)
-  y = contract(x1, x2; kwargs...)
+function ChainRulesCore.rrule(
+  ::typeof(contract), alg::Algorithm, x1::MPO, x2::MPS; kwargs...
+)
+  y = contract(alg, x1, x2; kwargs...)
   function contract_pullback(ȳ)
     x̄1 = _contract(MPO, ȳ, dag(x2); kwargs...)
-    x̄2 = contract(dag(x1), ȳ; kwargs...)
-    return (NoTangent(), x̄1, x̄2)
+    x̄2 = contract(alg, dag(x1), ȳ; kwargs...)
+    return (NoTangent(), NoTangent(), x̄1, x̄2)
   end
   return y, contract_pullback
 end
 
-function ChainRulesCore.rrule(::typeof(*), x1::MPO, x2::MPO; kwargs...)
-  return ChainRulesCore.rrule(contract, x1, x2; kwargs...)
-end
+## function ChainRulesCore.rrule(::typeof(*), x1::MPO, x2::MPO; alg, kwargs...)
+##   return ChainRulesCore.rrule(contract, alg, x1, x2; kwargs...)
+## end
 
 function ChainRulesCore.rrule(::typeof(+), x1::MPO, x2::MPO; kwargs...)
   y = +(x1, x2; kwargs...)

src/lib/ITensorMPS/src/mpo.jl

@@ -596,8 +596,12 @@ Equivalent to `replaceprime(contract(A, x; kwargs...), 2 => 1)`.
 
 See also [`contract`](@ref) for details about the arguments available.
 """
-function apply(A::MPO, ψ::MPS; kwargs...)
-  Aψ = contract(A, ψ; kwargs...)
+function apply(A::MPO, ψ::MPS; alg=Algorithm"densitymatrix"(), kwargs...)
+  return apply(Algorithm(alg), A, ψ; kwargs...)
+end
+
+function apply(alg::Algorithm, A::MPO, ψ::MPS; kwargs...)
+  Aψ = contract(alg, A, ψ; kwargs...)
   return replaceprime(Aψ, 1 => 0)
 end
 

src/lib/ITensorMPS/test/Ops/runtests.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using ITensors
 using Test
 
@@ -14,3 +15,4 @@ ITensors.disable_threaded_blocksparse()
     @time include(filename)
   end
 end
+end

src/lib/ITensorMPS/test/Ops/test_ops_mpo.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using Test
 using ITensors
 using ITensors.Ops
@@ -95,3 +96,4 @@ end
   @test norm(replaceprime(H' * H, 2 => 1) - H²) ≈ 0 atol = 1e-14
   @test norm(H(H) - H²) ≈ 0 atol = 1e-14
 end
+end

src/lib/ITensorMPS/test/Ops/test_trotter.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using Test
 using ITensors
 using ITensors.Ops
@@ -36,3 +37,4 @@ end
     end
   end
 end
+end

src/lib/ITensorMPS/test/base/test_autompo.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using ITensors, Test, Random, JLD2
 using NDTensors: scalartype
 
@@ -1249,3 +1250,4 @@ end
     @test_nowarn H = MPO(os, sites)
   end
 end
+end

src/lib/ITensorMPS/test/base/test_fermions.jl

@@ -260,11 +260,11 @@ import ITensors: Out, In
       # Reference state |110⟩
       ψ110 = MPS(s, n -> n == 1 || n == 2 ? "1" : "0")
 
-      function ITensors.op(::OpName"CdagC", ::SiteType, s1::Index, s2::Index)
+      function ITensors.op(::OpName"CdagC3", ::SiteType, s1::Index, s2::Index)
         return op("Cdag", s1) * op("C", s2)
       end
 
-      os = [("CdagC", 1, 3)]
+      os = [("CdagC3", 1, 3)]
       Os = ops(os, s)
 
       # Results in -|110⟩

src/lib/ITensorMPS/test/base/test_mpo.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using Combinatorics
 using ITensors
 using NDTensors: scalartype
@@ -853,3 +854,4 @@ end
     @test maxlinkdim(H) ≤ maxlinkdim(H₁) + maxlinkdim(H₂)
   end
 end
+end

src/lib/ITensorMPS/test/base/test_mps.jl

@@ -1,3 +1,4 @@
+@eval module $(gensym())
 using Combinatorics
 using ITensors
 using ITensors: ITensorMPS
@@ -1765,11 +1766,11 @@ end
       # Reference state |110⟩
       ψ110 = MPS(s, n -> n == 1 || n == 2 ? "1" : "0")
 
-      function ITensors.op(::OpName"CdagC", ::SiteType, s1::Index, s2::Index)
+      function ITensors.op(::OpName"CdagC1", ::SiteType, s1::Index, s2::Index)
         return op("Cdag", s1) * op("C", s2)
       end
 
-      os = [("CdagC", 1, 3)]
+      os = [("CdagC1", 1, 3)]
       Os = ops(os, s)
 
       # Results in -|110⟩
@@ -1810,7 +1811,7 @@ end
       cutoff!(sweeps, 1E-12)
       energy, ψ0 = dmrg(H, ψ0, sweeps; outputlevel=0)
 
-      function ITensors.op(::OpName"CdagC", ::SiteType, s1::Index, s2::Index)
+      function ITensors.op(::OpName"CdagC2", ::SiteType, s1::Index, s2::Index)
         return op("Cdag", s1) * op("C", s2)
       end
 
@@ -1821,7 +1822,7 @@ end
       end
 
       for i in 1:(N - 1), j in (i + 1):N
-        G1 = op("CdagC", s, i, j)
+        G1 = op("CdagC2", s, i, j)
 
         @disable_warn_order begin
           G2 = op("Cdag", s, i)
@@ -2006,3 +2007,4 @@ end
     @test norm(M - Mt) ≈ 0 atol = 1e-12
   end
 end
+end