Port changes from hamiltonian_derivatives to mpo_derivatives and reuse code from mpo_derivatives

Author: AFeuerpfeil

src/algorithms/derivatives/hamiltonian_derivatives.jl

@@ -69,12 +69,6 @@ end
 
 # Constructors
 # ------------
-const _HAM_MPS_TYPES = Union{
-    FiniteMPS{<:MPSTensor},
-    WindowMPS{<:MPSTensor},
-    InfiniteMPS{<:MPSTensor},
-}
-
 function AC_hamiltonian(
         site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:JordanMPOTensor},
         above::_HAM_MPS_TYPES, envs
@@ -83,9 +77,11 @@ function AC_hamiltonian(
     GR = rightenv(envs, site, below)
     W = operator[site]
 
+    GR_2 = GR[2:(end - 1)]
+    GL_2 = GL[2:(end - 1)]
+
     # starting
     if nonzero_length(W.C) > 0
-        GR_2 = GR[2:(end - 1)]
         @plansor starting_[-1 -2; -3 -4] ≔ W.C[-1; -3 1] * GR_2[-4 1; -2]
         starting = only(starting_)
     else
@@ -94,7 +90,6 @@ function AC_hamiltonian(
 
     # ending
     if nonzero_length(W.B) > 0
-        GL_2 = GL[2:(end - 1)]
         @plansor ending_[-1 -2; -3 -4] ≔ GL_2[-1 1; -3] * W.B[1 -2; -4]
         ending = only(ending_)
     else
@@ -139,8 +134,7 @@ function AC_hamiltonian(
     end
 
     if nonzero_length(W.A) > 0
-        @plansor GLW[-1 -2 -3; -4 -5] ≔ GL[2:(end - 1)][-1 1; -4] * W.A[1 -2; -5 -3]
-        continuing = (GLW, GR[2:(end - 1)])
+        continuing = AC_hamiltonian(GL_2, W.A, GR_2)
     else
         continuing = missing
     end
@@ -159,18 +153,21 @@ function AC2_hamiltonian(
     W1 = operator[site]
     W2 = operator[site + 1]
 
+    GR_2 = GR[2:(end - 1)]
+    GL_2 = GL[2:(end - 1)]
+
     # starting left - continuing right
     if nonzero_length(W1.C) > 0 && nonzero_length(W2.A) > 0
         @plansor CA_[-1 -2 -3; -4 -5 -6] ≔ W1.C[-1; -4 2] * W2.A[2 -2; -5 1] *
-            GR[2:(end - 1)][-6 1; -3]
+            GR_2[-6 1; -3]
         CA = only(CA_)
     else
         CA = missing
     end
 
     # continuing left - ending right
     if nonzero_length(W1.A) > 0 && nonzero_length(W2.B) > 0
-        @plansor AB_[-1 -2 -3; -4 -5 -6] ≔ GL[2:(end - 1)][-1 2; -4] * W1.A[2 -2; -5 1] *
+        @plansor AB_[-1 -2 -3; -4 -5 -6] ≔ GL_2[-1 2; -4] * W1.A[2 -2; -5 1] *
             W2.B[1 -3; -6]
         AB = only(AB_)
     else
@@ -200,10 +197,10 @@ function AC2_hamiltonian(
         if !ismissing(CA)
             I = id(storagetype(GR[1]), physicalspace(W1))
             @plansor CA[-1 -2 -3; -4 -5 -6] += (I[-1; -4] * W2.C[-2; -5 1]) *
-                GR[2:(end - 1)][-6 1; -3]
+                GR_2[-6 1; -3]
             IC = missing
         else
-            @plansor IC[-1 -2; -3 -4] ≔ W2.C[-1; -3 1] * GR[2:(end - 1)][-4 1; -2]
+            @plansor IC[-1 -2; -3 -4] ≔ W2.C[-1; -3 1] * GR_2[-4 1; -2]
         end
     else
         IC = missing
@@ -213,11 +210,11 @@ function AC2_hamiltonian(
     if nonzero_length(W1.B) > 0
         if !ismissing(AB)
             I = id(storagetype(GL[end]), physicalspace(W2))
-            @plansor AB[-1 -2 -3; -4 -5 -6] += GL[2:(end - 1)][-1 1; -4] *
+            @plansor AB[-1 -2 -3; -4 -5 -6] += GL_2[-1 1; -4] *
                 (W1.B[1 -2; -5] * I[-3; -6])
             BE = missing
         else
-            @plansor BE[-1 -2; -3 -4] ≔ GL[2:(end - 1)][-1 2; -3] * W1.B[2 -2; -4]
+            @plansor BE[-1 -2; -3 -4] ≔ GL_2[-1 2; -3] * W1.B[2 -2; -4]
         end
     else
         BE = missing
@@ -290,12 +287,9 @@ function AC2_hamiltonian(
 
     # continuing - continuing
     # TODO: MPODerivativeOperator code reuse + optimization
-    # TODO: One should only calculate these operators if necessary. One can do this by checking nonzero_length(A1) > 0 && nonzero_length(A2) > 0 and then setting AA=missing or as we discussed via bit matrix multiplication.
     ## TODO: Think about how one could and whether one should store these objects and use them for (a) advancing environments in iDMRG, (b) reuse ind backwards-sweep in IDMRG, (c) subspace expansion
     if nonzero_length(W1.A) > 0 && nonzero_length(W2.A) > 0
-        @plansor GLW[-1 -2 -3; -4 -5] ≔ GL[2:(end - 1)][-1 1; -4] * W1.A[1 -2; -5 -3]
-        @plansor GWR[-1 -2 -3; -4 -5] ≔ W2.A[-3 -5; -2 1] * GR[2:(end - 1)][-1 1; -4]
-        AA = (GLW, GWR)
+        AA = AC2_hamiltonian(GL_2, W1.A, W2.A, GR_2)
     else
         AA = missing
     end
@@ -313,10 +307,8 @@ function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor)
     ismissing(H.I) || @plansor y[-1 -2; -3] += x[-1 -2; 1] * H.I[1; -3]
     ismissing(H.C) || @plansor y[-1 -2; -3] += x[-1 2; 1] * H.C[-2 -3; 2 1]
     ismissing(H.B) || @plansor y[-1 -2; -3] += H.B[-1 -2; 1 2] * x[1 2; -3]
-
     if !ismissing(H.A)
-        GLW, GR = H.A
-        @plansor y[-1 -2; -3] += GLW[-1 -2 3; 1 2] * x[1 2; 4] * GR[4 3; -3]
+        y += H.A(x)
     end
 
     return y
@@ -333,10 +325,8 @@ function (H::JordanMPO_AC2_Hamiltonian)(x::MPOTensor)
     ismissing(H.BE) || @plansor y[-1 -2; -3 -4] += x[1 2; -3 -4] * H.BE[-1 -2; 1 2]
     ismissing(H.DE) || @plansor y[-1 -2; -3 -4] += x[-1 1; -3 -4] * H.DE[-2; 1]
     ismissing(H.EE) || @plansor y[-1 -2; -3 -4] += x[1 -2; -3 -4] * H.EE[-1; 1]
-
     if !ismissing(H.AA)
-        GLW, GWR = H.AA
-        @plansor y[-1 -2; -3 -4] += GLW[-1 -2 5; 1 2] * x[1 2; 3 4] * GWR[3 4 5; -3 -4]
+        y += H.AA(x)
     end
 
     return y

src/algorithms/derivatives/mpo_derivatives.jl

@@ -9,9 +9,15 @@ struct MPODerivativeOperator{L, O <: Tuple, R} <: DerivativeOperator
     rightenv::R
 end
 
+struct MPOContractedDerivativeOperator{L, R, N} <: DerivativeOperator
+    leftenv::L
+    rightenv::R
+end
+
 Base.length(H::MPODerivativeOperator) = length(H.operators)
+Base.length(::MPOContractedDerivativeOperator{L, R, N}) where {L, R, N} = N
 
-const MPO_C_Hamiltonian{L, R} = MPODerivativeOperator{L, Tuple{}, R}
+const MPO_C_Hamiltonian{L, R} = MPOContractedDerivativeOperator{L, Tuple{}, R}
 MPO_C_Hamiltonian(GL, GR) = MPODerivativeOperator(GL, (), GR)
 
 const MPO_AC_Hamiltonian{L, O, R} = MPODerivativeOperator{L, Tuple{O}, R}
@@ -20,8 +26,23 @@ MPO_AC_Hamiltonian(GL, O, GR) = MPODerivativeOperator(GL, (O,), GR)
 const MPO_AC2_Hamiltonian{L, O₁, O₂, R} = MPODerivativeOperator{L, Tuple{O₁, O₂}, R}
 MPO_AC2_Hamiltonian(GL, O1, O2, GR) = MPODerivativeOperator(GL, (O1, O2), GR)
 
+const MPO_Contracted_C_Hamiltonian{L, R} = MPOContractedDerivativeOperator{L, R, 0}
+MPO_Contracted_C_Hamiltonian(GL::L, GR::R) where {L, R} = MPOContractedDerivativeOperator{L,R,0}(GL, GR)
+
+const MPO_Contracted_AC_Hamiltonian{L, R} = MPOContractedDerivativeOperator{L, R, 1}
+MPO_Contracted_AC_Hamiltonian(GL::L, GR::R) where {L, R} = MPOContractedDerivativeOperator{L,R,1}(GL,  GR)
+
+const MPO_Contracted_AC2_Hamiltonian{L, R} = MPOContractedDerivativeOperator{L, R, 2}
+MPO_Contracted_AC2_Hamiltonian(GL::L, GR::R) where {L, R} = MPOContractedDerivativeOperator{L,R,2}(GL, GR)
+
 # Constructors
 # ------------
+const _HAM_MPS_TYPES = Union{
+    FiniteMPS{<:MPSTensor},
+    WindowMPS{<:MPSTensor},
+    InfiniteMPS{<:MPSTensor},
+}
+
 function C_hamiltonian(site::Int, below, operator, above, envs)
     return MPO_C_Hamiltonian(leftenv(envs, site + 1, below), rightenv(envs, site, below))
 end
@@ -35,6 +56,30 @@ function AC2_hamiltonian(site::Int, below, operator, above, envs)
         leftenv(envs, site, below), O1, O2, rightenv(envs, site + 1, below)
     )
 end
+function C_hamiltonian(site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:MPOTensor}, above::_HAM_MPS_TYPES, envs)
+    return MPO_Contracted_C_Hamiltonian(leftenv(envs, site + 1, below), rightenv(envs, site, below))
+end
+function AC_hamiltonian(site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:MPOTensor}, above::_HAM_MPS_TYPES, envs)
+    O = operator[site]
+    GL = leftenv(envs, site, below)
+    return AC_hamiltonian(GL, O, rightenv(envs, site, below))
+end
+function AC_hamiltonian(GL::MPSTensor, O::MPOTensor, GR::MPSTensor)
+    @plansor GLW[-1 -2 -3; -4 -5] ≔ GL[-1 1; -4] * O[1 -2; -5 -3]
+    return MPO_Contracted_AC_Hamiltonian(GLW, O, GR)
+end
+function AC2_hamiltonian(site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:MPOTensor}, above::_HAM_MPS_TYPES, envs)
+    O1 = operator[site]
+    O2 = operator[site + 1]
+    GL = leftenv(envs, site, below)
+    GR = rightenv(envs, site + 1, below)
+    return AC2_hamiltonian(GL, O1, O2, GR)
+end
+function AC2_hamiltonian(GL::MPSTensor, O1::MPOTensor, O2::MPOTensor, GR::MPSTensor)
+    @plansor GLW[-1 -2 -3; -4 -5] ≔ GL[-1 1; -4] * O1[1 -2; -5 -3]
+    @plansor GWR[-1 -2 -3; -4 -5] ≔ O2[-3 -5; -2 1] * GR[-1 1; -4]
+    return MPO_Contracted_AC2_Hamiltonian(GLW, GWR)
+end
 
 # Properties
 # ----------
@@ -50,6 +95,20 @@ function TensorKit.codomain(H::MPODerivativeOperator)
     V_o = prod(physicalspace, H.O; init = one(V_l))
     return V_l ⊗ V_o ⊗ V_r
 end
+function TensorKit.domain(H::MPOContractedDerivativeOperator)
+    V_l = right_virtualspace(H.leftenv)
+    V_r = left_virtualspace(H.rightenv)
+    ## TODO: How to deal with the H.O here?
+    V_o = prod(physicalspace, H.O; init = one(V_l))
+    return V_l ⊗ V_o ⊗ V_r
+end
+function TensorKit.codomain(H::MPOContractedDerivativeOperator)
+    V_l = left_virtualspace(H.leftenv)
+    V_r = right_virtualspace(H.rightenv)
+    ## TODO: How to deal with the H.O here?
+    V_o = prod(physicalspace, H.O; init = one(V_l))
+    return V_l ⊗ V_o ⊗ V_r
+end
 
 # Actions
 # -------
@@ -105,3 +164,20 @@ function (h::MPO_AC2_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPOTensor, <:MPSTen
         h.operators[2][7 -6; 4 5] * τ[5 -5; 2 3]
     return y isa AbstractBlockTensorMap ? only(y) : y
 end
+function (h::MPO_Contracted_C_Hamiltonian)(x::MPSBondTensor)
+    @plansor y[-1; -2] ≔ h.leftenv[-1 3; 1] * x[1; 2] * h.rightenv[2 3; -2]
+    return y isa AbstractBlockTensorMap ? only(y) : y
+end
+function (h::MPO_Contracted_AC_Hamiltonian)(
+        x::GenericMPSTensor{<:Any, 3}
+    )
+    @plansor y[-1 -2; -3] ≔ h.leftenv[-1 -2 3; 1 2] * x[1 2; 4] * h.rightenv[4 3; -3]
+    return y isa AbstractBlockTensorMap ? only(y) : y
+end
+function (h::MPO_Contracted_AC2_Hamiltonian)(
+        x::MPOTensor
+    )
+    @plansor y[-1 -2; -3 -4] ≔ h.leftenv[-1 -2 5; 1 2] * x[1 2; 3 4] * h.rightenv[3 4 5; -3 -4]
+    return y isa AbstractBlockTensorMap ? only(y) : y
+end
+## TODO: Which interface should we use for an inplace += version to be used in hamiltonian_derivatives.jl?