From 4bec98badfce88d6440dcd0bc3d446e65be5e150 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 10 Jan 2025 15:01:15 +0100 Subject: [PATCH 1/2] Use autoopt in local operator contraction --- src/algorithms/contractions/localoperator.jl | 230 ++++++++++--------- 1 file changed, 127 insertions(+), 103 deletions(-) diff --git a/src/algorithms/contractions/localoperator.jl b/src/algorithms/contractions/localoperator.jl index 324bb028..e38bdca6 100644 --- a/src/algorithms/contractions/localoperator.jl +++ b/src/algorithms/contractions/localoperator.jl @@ -56,29 +56,29 @@ end :(env.corners[ NORTHWEST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) ]), - (:C_NW_1,), - (:C_NW_2,), + (:χ_C_NW_1,), + (:χ_C_NW_2,), ) corner_NE = tensorexpr( :(env.corners[ NORTHEAST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) ]), - (:C_NE_1,), - (:C_NE_2,), + (:χ_C_NE_1,), + (:χ_C_NE_2,), ) corner_SE = tensorexpr( :(env.corners[ SOUTHEAST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) ]), - (:C_SE_1,), - (:C_SE_2,), + (:χ_C_SE_1,), + (:χ_C_SE_2,), ) corner_SW = tensorexpr( :(env.corners[ SOUTHWEST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) ]), - (:C_SW_1,), - (:C_SW_2,), + (:χ_C_SW_1,), + (:χ_C_SW_2,), ) edges_N = map(1:gridsize[2]) do i @@ -89,11 +89,11 @@ end mod1($(cmin + i - 1), size(ket, 2)), ]), ( - (i == 1 ? :C_NW_2 : Symbol(:E_N_virtual, i - 1)), - Symbol(:E_N_top, i), - Symbol(:E_N_bot, i), + (i == 1 ? :χ_C_NW_2 : Symbol(:χ_E_N, i - 1)), + Symbol(:D_E_N_top, i), + Symbol(:D_E_N_bot, i), ), - ((i == gridsize[2] ? :C_NE_1 : Symbol(:E_N_virtual, i)),), + ((i == gridsize[2] ? :χ_C_NE_1 : Symbol(:χ_E_N, i)),), ) end @@ -105,11 +105,11 @@ end mod1($(cmax + 1), size(ket, 2)), ]), ( - (i == 1 ? :C_NE_2 : Symbol(:E_E_virtual, i - 1)), - Symbol(:E_E_top, i), - Symbol(:E_E_bot, i), + (i == 1 ? :χ_C_NE_2 : Symbol(:χ_E_E, i - 1)), + Symbol(:D_E_E_top, i), + Symbol(:D_E_E_bot, i), ), - ((i == gridsize[1] ? :C_SE_1 : Symbol(:E_E_virtual, i)),), + ((i == gridsize[1] ? :χ_C_SE_1 : Symbol(:χ_E_E, i)),), ) end @@ -121,11 +121,11 @@ end mod1($(cmin + i - 1), size(ket, 2)), ]), ( - (i == gridsize[2] ? :C_SE_2 : Symbol(:E_S_virtual, i)), - Symbol(:E_S_top, i), - Symbol(:E_S_bot, i), + (i == gridsize[2] ? :χ_C_SE_2 : Symbol(:χ_E_S, i)), + Symbol(:D_E_S_top, i), + Symbol(:D_E_S_bot, i), ), - ((i == 1 ? :C_SW_1 : Symbol(:E_S_virtual, i - 1)),), + ((i == 1 ? :χ_C_SW_1 : Symbol(:χ_E_S, i - 1)),), ) end @@ -137,44 +137,56 @@ end mod1($(cmin - 1), size(ket, 2)), ]), ( - (i == gridsize[1] ? :C_SW_2 : Symbol(:E_W_virtual, i)), - Symbol(:E_W_top, i), - Symbol(:E_W_bot, i), + (i == gridsize[1] ? :χ_C_SW_2 : Symbol(:χ_E_W, i)), + Symbol(:D_E_W_top, i), + Symbol(:D_E_W_bot, i), ), - ((i == 1 ? :C_NW_1 : Symbol(:E_W_virtual, i - 1)),), + ((i == 1 ? :χ_C_NW_1 : Symbol(:χ_E_W, i - 1)),), ) end operator = tensorexpr( - :O, ntuple(i -> Symbol(:O_out_, i), N), ntuple(i -> Symbol(:O_in_, i), N) + :O, ntuple(i -> Symbol(:d_O_out_, i), N), ntuple(i -> Symbol(:d_O_in_, i), N) ) bra = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) inds_id = findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds) physical_label = - isnothing(inds_id) ? Symbol(:physical, i, "_", j) : Symbol(:O_out_, inds_id) + isnothing(inds_id) ? Symbol(:d, i, "_", j) : Symbol(:d_O_out_, inds_id) return tensorexpr( :(bra[ mod1($(rmin + i - 1), size(bra, 1)), mod1($(cmin + j - 1), size(bra, 2)) ]), (physical_label,), ( - (i == 1 ? Symbol(:E_N_bot, j) : Symbol(:bra_vertical, i - 1, "_", j)), + ( + if i == 1 + Symbol(:D_E_N_bot, j) + else + Symbol(:D_bra_vertical, i - 1, "_", j) + end + ), ( if j == gridsize[2] - Symbol(:E_E_bot, i) + Symbol(:D_E_E_bot, i) else - Symbol(:bra_horizontal, i, "_", j) + Symbol(:D_bra_horizontal, i, "_", j) end ), ( if i == gridsize[1] - Symbol(:E_S_bot, j) + Symbol(:D_E_S_bot, j) else - Symbol(:bra_vertical, i, "_", j) + Symbol(:D_bra_vertical, i, "_", j) + end + ), + ( + if j == 1 + Symbol(:D_E_W_bot, i) + else + Symbol(:D_bra_horizontal, i, "_", j - 1) end ), - (j == 1 ? Symbol(:E_W_bot, i) : Symbol(:bra_horizontal, i, "_", j - 1)), ), ) end @@ -182,29 +194,41 @@ end ket = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) inds_id = findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds) physical_label = - isnothing(inds_id) ? Symbol(:physical, i, "_", j) : Symbol(:O_in_, inds_id) + isnothing(inds_id) ? Symbol(:d, i, "_", j) : Symbol(:d_O_in_, inds_id) return tensorexpr( :(ket[ mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) ]), (physical_label,), ( - (i == 1 ? Symbol(:E_N_top, j) : Symbol(:ket_vertical, i - 1, "_", j)), + ( + if i == 1 + Symbol(:D_E_N_top, j) + else + Symbol(:D_ket_vertical, i - 1, "_", j) + end + ), ( if j == gridsize[2] - Symbol(:E_E_top, i) + Symbol(:D_E_E_top, i) else - Symbol(:ket_horizontal, i, "_", j) + Symbol(:D_ket_horizontal, i, "_", j) end ), ( if i == gridsize[1] - Symbol(:E_S_top, j) + Symbol(:D_E_S_top, j) else - Symbol(:ket_vertical, i, "_", j) + Symbol(:D_ket_vertical, i, "_", j) + end + ), + ( + if j == 1 + Symbol(:D_E_W_top, i) + else + Symbol(:D_ket_horizontal, i, "_", j - 1) end ), - (j == 1 ? Symbol(:E_W_top, i) : Symbol(:ket_horizontal, i, "_", j - 1)), ), ) end @@ -225,20 +249,8 @@ end operator, ) - opt_ex = Expr(:tuple) - allinds = TensorOperations.getallindices(multiplication_ex) - for label in allinds - if startswith(String(label), "physical") || startswith(String(label), "O") - push!(opt_ex.args, :($label => $PEPS_PHYSICALDIM)) - elseif startswith(String(label), "ket") || startswith(String(label), "bra") - push!(opt_ex.args, :($label => $PEPS_BONDDIM)) - else - push!(opt_ex.args, :($label => $PEPS_ENVBONDDIM)) - end - end - returnex = quote - @tensor opt = $opt_ex $multiplication_ex + @autoopt @tensor opt = $multiplication_ex end return macroexpand(@__MODULE__, returnex) end @@ -276,29 +288,29 @@ end :(env.corners[ NORTHWEST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) ]), - (:C_NW_1,), - (:C_NW_2,), + (:χ_C_NW_1,), + (:χ_C_NW_2,), ) corner_NE = tensorexpr( :(env.corners[ NORTHEAST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) ]), - (:C_NE_1,), - (:C_NE_2,), + (:χ_C_NE_1,), + (:χ_C_NE_2,), ) corner_SE = tensorexpr( :(env.corners[ SOUTHEAST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) ]), - (:C_SE_1,), - (:C_SE_2,), + (:χ_C_SE_1,), + (:χ_C_SE_2,), ) corner_SW = tensorexpr( :(env.corners[ SOUTHWEST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) ]), - (:C_SW_1,), - (:C_SW_2,), + (:χ_C_SW_1,), + (:χ_C_SW_2,), ) edges_N = map(1:gridsize[2]) do i @@ -309,11 +321,11 @@ end mod1($(cmin + i - 1), size(ket, 2)), ]), ( - (i == 1 ? :C_NW_2 : Symbol(:E_N_virtual, i - 1)), - Symbol(:E_N_top, i), - Symbol(:E_N_bot, i), + (i == 1 ? :χ_C_NW_2 : Symbol(:χ_E_N, i - 1)), + Symbol(:D_E_N_top, i), + Symbol(:D_E_N_bot, i), ), - ((i == gridsize[2] ? :C_NE_1 : Symbol(:E_N_virtual, i)),), + ((i == gridsize[2] ? :χ_C_NE_1 : Symbol(:χ_E_N, i)),), ) end @@ -325,11 +337,11 @@ end mod1($(cmax + 1), size(ket, 2)), ]), ( - (i == 1 ? :C_NE_2 : Symbol(:E_E_virtual, i - 1)), - Symbol(:E_E_top, i), - Symbol(:E_E_bot, i), + (i == 1 ? :χ_C_NE_2 : Symbol(:χ_E_E, i - 1)), + Symbol(:D_E_E_top, i), + Symbol(:D_E_E_bot, i), ), - ((i == gridsize[1] ? :C_SE_1 : Symbol(:E_E_virtual, i)),), + ((i == gridsize[1] ? :χ_C_SE_1 : Symbol(:χ_E_E, i)),), ) end @@ -341,11 +353,11 @@ end mod1($(cmin + i - 1), size(ket, 2)), ]), ( - (i == gridsize[2] ? :C_SE_2 : Symbol(:E_S_virtual, i)), - Symbol(:E_S_top, i), - Symbol(:E_S_bot, i), + (i == gridsize[2] ? :χ_C_SE_2 : Symbol(:χ_E_S, i)), + Symbol(:D_E_S_top, i), + Symbol(:D_E_S_bot, i), ), - ((i == 1 ? :C_SW_1 : Symbol(:E_S_virtual, i - 1)),), + ((i == 1 ? :χ_C_SW_1 : Symbol(:χ_E_S, i - 1)),), ) end @@ -357,11 +369,11 @@ end mod1($(cmin - 1), size(ket, 2)), ]), ( - (i == gridsize[1] ? :C_SW_2 : Symbol(:E_W_virtual, i)), - Symbol(:E_W_top, i), - Symbol(:E_W_bot, i), + (i == gridsize[1] ? :χ_C_SW_2 : Symbol(:χ_E_W, i)), + Symbol(:D_E_W_top, i), + Symbol(:D_E_W_bot, i), ), - ((i == 1 ? :C_NW_1 : Symbol(:E_W_virtual, i - 1)),), + ((i == 1 ? :χ_C_NW_1 : Symbol(:χ_E_W, i - 1)),), ) end @@ -370,24 +382,36 @@ end :(bra[ mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) ]), - (Symbol(:physical, i, "_", j),), + (Symbol(:d, i, "_", j),), ( - (i == 1 ? Symbol(:E_N_bot, j) : Symbol(:bra_vertical, i - 1, "_", j)), + ( + if i == 1 + Symbol(:D_E_N_bot, j) + else + Symbol(:D_bra_vertical, i - 1, "_", j) + end + ), ( if j == gridsize[2] - Symbol(:E_E_bot, i) + Symbol(:D_E_E_bot, i) else - Symbol(:bra_horizontal, i, "_", j) + Symbol(:D_bra_horizontal, i, "_", j) end ), ( if i == gridsize[1] - Symbol(:E_S_bot, j) + Symbol(:D_E_S_bot, j) else - Symbol(:bra_vertical, i, "_", j) + Symbol(:D_bra_vertical, i, "_", j) + end + ), + ( + if j == 1 + Symbol(:D_E_W_bot, i) + else + Symbol(:D_bra_horizontal, i, "_", j - 1) end ), - (j == 1 ? Symbol(:E_W_bot, i) : Symbol(:bra_horizontal, i, "_", j - 1)), ), ) end @@ -397,24 +421,36 @@ end :(ket[ mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) ]), - (Symbol(:physical, i, "_", j),), + (Symbol(:d, i, "_", j),), ( - (i == 1 ? Symbol(:E_N_top, j) : Symbol(:ket_vertical, i - 1, "_", j)), + ( + if i == 1 + Symbol(:D_E_N_top, j) + else + Symbol(:D_ket_vertical, i - 1, "_", j) + end + ), ( if j == gridsize[2] - Symbol(:E_E_top, i) + Symbol(:D_E_E_top, i) else - Symbol(:ket_horizontal, i, "_", j) + Symbol(:D_ket_horizontal, i, "_", j) end ), ( if i == gridsize[1] - Symbol(:E_S_top, j) + Symbol(:D_E_S_top, j) else - Symbol(:ket_vertical, i, "_", j) + Symbol(:D_ket_vertical, i, "_", j) + end + ), + ( + if j == 1 + Symbol(:D_E_W_top, i) + else + Symbol(:D_ket_horizontal, i, "_", j - 1) end ), - (j == 1 ? Symbol(:E_W_top, i) : Symbol(:ket_horizontal, i, "_", j - 1)), ), ) end @@ -434,20 +470,8 @@ end map(x -> Expr(:call, :conj, x), bra)..., ) - opt_ex = Expr(:tuple) - allinds = TensorOperations.getallindices(multiplication_ex) - for label in allinds - if startswith(String(label), "physical") - push!(opt_ex.args, :($label => $PEPS_PHYSICALDIM)) - elseif startswith(String(label), "ket") || startswith(String(label), "bra") - push!(opt_ex.args, :($label => $PEPS_BONDDIM)) - else - push!(opt_ex.args, :($label => $PEPS_ENVBONDDIM)) - end - end - returnex = quote - @tensor opt = $opt_ex $multiplication_ex + @autoopt @tensor opt = $multiplication_ex end return macroexpand(@__MODULE__, returnex) end From ad757523ef303e8dddf235ae90719ca1514ec72b Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Fri, 10 Jan 2025 11:24:23 -0500 Subject: [PATCH 2/2] Refactor `contract_localoperator` --- src/algorithms/contractions/localoperator.jl | 453 +++++-------------- 1 file changed, 120 insertions(+), 333 deletions(-) diff --git a/src/algorithms/contractions/localoperator.jl b/src/algorithms/contractions/localoperator.jl index e38bdca6..3d6d3dac 100644 --- a/src/algorithms/contractions/localoperator.jl +++ b/src/algorithms/contractions/localoperator.jl @@ -3,10 +3,20 @@ import MPSKit: tensorexpr # currently need this because MPSKit restricts tensor names to symbols -MPSKit.tensorexpr(ex::Expr, inds) = Expr(:ref, ex, inds...) +_totuple(t) = t isa Tuple ? t : tuple(t) +MPSKit.tensorexpr(ex::Expr, inds::Tuple) = Expr(:ref, ex, _totuple(inds)...) function MPSKit.tensorexpr(ex::Expr, indout, indin) - return Expr(:typed_vcat, ex, Expr(:row, indout...), Expr(:row, indin...)) + return Expr( + :typed_vcat, ex, Expr(:row, _totuple(indout)...), Expr(:row, _totuple(indin)...) + ) +end + +function tensorlabel(args...) + return Symbol(ntuple(i -> iseven(i) ? :_ : args[(i + 1) >> 1], 2 * length(args) - 1)...) end +envlabel(args...) = tensorlabel(:χ, args...) +virtuallabel(args...) = tensorlabel(:D, args...) +physicallabel(args...) = tensorlabel(:d, args...) """ contract_localoperator(inds, O, peps, env) @@ -35,203 +45,151 @@ end # This implements the contraction of an operator acting on sites `inds`. # The generated function ensures that we can use @tensor to write dynamic contractions (and maximize performance). -@generated function _contract_localoperator( - inds::NTuple{N,Val}, - O::AbstractTensorMap{S,N,N}, - ket::InfinitePEPS, - bra::InfinitePEPS, - env::CTMRGEnv, -) where {S,N} - cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val - if !allunique(cartesian_inds) - throw(ArgumentError("Indices should not overlap: $cartesian_inds.")) - end - - rmin, rmax = extrema(getindex.(cartesian_inds, 1)) - cmin, cmax = extrema(getindex.(cartesian_inds, 2)) +function _contract_corner_expr(rowrange, colrange) + rmin, rmax = extrema(rowrange) + cmin, cmax = extrema(colrange) gridsize = (rmax - rmin + 1, cmax - cmin + 1) - corner_NW = tensorexpr( - :(env.corners[ - NORTHWEST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) - ]), - (:χ_C_NW_1,), - (:χ_C_NW_2,), - ) - corner_NE = tensorexpr( - :(env.corners[ - NORTHEAST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) - ]), - (:χ_C_NE_1,), - (:χ_C_NE_2,), - ) - corner_SE = tensorexpr( - :(env.corners[ - SOUTHEAST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) - ]), - (:χ_C_SE_1,), - (:χ_C_SE_2,), - ) - corner_SW = tensorexpr( - :(env.corners[ - SOUTHWEST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) - ]), - (:χ_C_SW_1,), - (:χ_C_SW_2,), - ) + C_NW = :(env.corners[NORTHWEST, mod1($(rmin - 1), end), mod1($(cmin - 1), end)]) + corner_NW = tensorexpr(C_NW, envlabel(WEST, 0), envlabel(NORTH, 0)) + + C_NE = :(env.corners[NORTHEAST, mod1($(rmin - 1), end), mod1($(cmax + 1), end)]) + corner_NE = tensorexpr(C_NE, envlabel(NORTH, gridsize[2]), envlabel(EAST, 0)) + + C_SE = :(env.corners[SOUTHEAST, mod1($(rmax + 1), end), mod1($(cmax + 1), end)]) + corner_SE = tensorexpr(C_SE, envlabel(EAST, gridsize[1]), envlabel(SOUTH, gridsize[2])) + + C_SW = :(env.corners[SOUTHWEST, mod1($(rmax + 1), end), mod1($(cmin - 1), end)]) + corner_SW = tensorexpr(C_SW, envlabel(SOUTH, 0), envlabel(WEST, gridsize[1])) + + return corner_NW, corner_NE, corner_SE, corner_SW +end + +function _contract_edge_expr(rowrange, colrange) + rmin, rmax = extrema(rowrange) + cmin, cmax = extrema(colrange) + gridsize = (rmax - rmin + 1, cmax - cmin + 1) edges_N = map(1:gridsize[2]) do i + E_N = :(env.edges[NORTH, mod1($(rmin - 1), end), mod1($(cmin + i - 1), end)]) return tensorexpr( - :(env.edges[ - NORTH, - mod1($(rmin - 1), size(ket, 1)), - mod1($(cmin + i - 1), size(ket, 2)), - ]), + E_N, ( - (i == 1 ? :χ_C_NW_2 : Symbol(:χ_E_N, i - 1)), - Symbol(:D_E_N_top, i), - Symbol(:D_E_N_bot, i), + envlabel(NORTH, i - 1), + virtuallabel(NORTH, :ket, i), + virtuallabel(NORTH, :bra, i), ), - ((i == gridsize[2] ? :χ_C_NE_1 : Symbol(:χ_E_N, i)),), + envlabel(NORTH, i), ) end edges_E = map(1:gridsize[1]) do i + E_E = :(env.edges[EAST, mod1($(rmin + i - 1), end), mod1($(cmax + 1), end)]) return tensorexpr( - :(env.edges[ - EAST, - mod1($(rmin + i - 1), size(ket, 1)), - mod1($(cmax + 1), size(ket, 2)), - ]), + E_E, ( - (i == 1 ? :χ_C_NE_2 : Symbol(:χ_E_E, i - 1)), - Symbol(:D_E_E_top, i), - Symbol(:D_E_E_bot, i), + envlabel(EAST, i - 1), + virtuallabel(EAST, :ket, i), + virtuallabel(EAST, :bra, i), ), - ((i == gridsize[1] ? :χ_C_SE_1 : Symbol(:χ_E_E, i)),), + envlabel(EAST, i), ) end edges_S = map(1:gridsize[2]) do i + E_S = :(env.edges[SOUTH, mod1($(rmax + 1), end), mod1($(cmin + i - 1), end)]) return tensorexpr( - :(env.edges[ - SOUTH, - mod1($(rmax + 1), size(ket, 1)), - mod1($(cmin + i - 1), size(ket, 2)), - ]), + E_S, ( - (i == gridsize[2] ? :χ_C_SE_2 : Symbol(:χ_E_S, i)), - Symbol(:D_E_S_top, i), - Symbol(:D_E_S_bot, i), + envlabel(SOUTH, i), + virtuallabel(SOUTH, :ket, i), + virtuallabel(SOUTH, :bra, i), ), - ((i == 1 ? :χ_C_SW_1 : Symbol(:χ_E_S, i - 1)),), + envlabel(SOUTH, i - 1), ) end edges_W = map(1:gridsize[1]) do i + E_W = :(env.edges[WEST, mod1($(rmin + i - 1), end), mod1($(cmin - 1), end)]) return tensorexpr( - :(env.edges[ - WEST, - mod1($(rmin + i - 1), size(ket, 1)), - mod1($(cmin - 1), size(ket, 2)), - ]), - ( - (i == gridsize[1] ? :χ_C_SW_2 : Symbol(:χ_E_W, i)), - Symbol(:D_E_W_top, i), - Symbol(:D_E_W_bot, i), - ), - ((i == 1 ? :χ_C_NW_1 : Symbol(:χ_E_W, i - 1)),), + E_W, + (envlabel(WEST, i), virtuallabel(WEST, :ket, i), virtuallabel(WEST, :bra, i)), + envlabel(WEST, i - 1), ) end - operator = tensorexpr( - :O, ntuple(i -> Symbol(:d_O_out_, i), N), ntuple(i -> Symbol(:d_O_in_, i), N) - ) + return edges_N, edges_E, edges_S, edges_W +end - bra = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) - inds_id = findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds) - physical_label = - isnothing(inds_id) ? Symbol(:d, i, "_", j) : Symbol(:d_O_out_, inds_id) - return tensorexpr( - :(bra[ - mod1($(rmin + i - 1), size(bra, 1)), mod1($(cmin + j - 1), size(bra, 2)) - ]), - (physical_label,), - ( - ( - if i == 1 - Symbol(:D_E_N_bot, j) - else - Symbol(:D_bra_vertical, i - 1, "_", j) - end - ), - ( - if j == gridsize[2] - Symbol(:D_E_E_bot, i) - else - Symbol(:D_bra_horizontal, i, "_", j) - end - ), - ( - if i == gridsize[1] - Symbol(:D_E_S_bot, j) - else - Symbol(:D_bra_vertical, i, "_", j) - end - ), - ( - if j == 1 - Symbol(:D_E_W_bot, i) - else - Symbol(:D_bra_horizontal, i, "_", j - 1) - end - ), - ), - ) - end +function _contract_state_expr(rowrange, colrange, cartesian_inds=nothing) + rmin, rmax = extrema(rowrange) + cmin, cmax = extrema(colrange) + gridsize = (rmax - rmin + 1, cmax - cmin + 1) - ket = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) - inds_id = findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds) - physical_label = - isnothing(inds_id) ? Symbol(:d, i, "_", j) : Symbol(:d_O_in_, inds_id) - return tensorexpr( - :(ket[ - mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) - ]), - (physical_label,), - ( + return map((:bra, :ket)) do side + return map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) + inds_id = if isnothing(cartesian_inds) + nothing + else + findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds) + end + physical_label = if isnothing(inds_id) + physicallabel(i, j) + else + physicallabel(:O, side, inds_id) + end + return tensorexpr( + :(bra[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)]), + (physical_label,), ( if i == 1 - Symbol(:D_E_N_top, j) + virtuallabel(NORTH, side, j) else - Symbol(:D_ket_vertical, i - 1, "_", j) - end - ), - ( + virtuallabel(:vertical, side, i - 1, j) + end, if j == gridsize[2] - Symbol(:D_E_E_top, i) + virtuallabel(EAST, side, i) else - Symbol(:D_ket_horizontal, i, "_", j) - end - ), - ( + virtuallabel(:horizontal, side, i, j) + end, if i == gridsize[1] - Symbol(:D_E_S_top, j) + virtuallabel(SOUTH, side, j) else - Symbol(:D_ket_vertical, i, "_", j) - end - ), - ( + virtuallabel(:vertical, side, i, j) + end, if j == 1 - Symbol(:D_E_W_top, i) + virtuallabel(WEST, side, i) else - Symbol(:D_ket_horizontal, i, "_", j - 1) - end + virtuallabel(:horizontal, side, i, j - 1) + end, ), - ), - ) + ) + end end +end + +@generated function _contract_localoperator( + inds::NTuple{N,Val}, + O::AbstractTensorMap{S,N,N}, + ket::InfinitePEPS, + bra::InfinitePEPS, + env::CTMRGEnv, +) where {S,N} + cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val + allunique(cartesian_inds) || + throw(ArgumentError("Indices should not overlap: $cartesian_inds.")) + rowrange = getindex.(cartesian_inds, 1) + colrange = getindex.(cartesian_inds, 2) + + corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange) + edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange) + operator = tensorexpr( + :O, + ntuple(i -> physicallabel(:O, :bra, i), N), + ntuple(i -> physicallabel(:O, :ket, i), N), + ) + bra, ket = _contract_state_expr(rowrange, colrange, cartesian_inds) multiplication_ex = Expr( :call, @@ -275,185 +233,14 @@ end inds::NTuple{N,Val}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv ) where {N} cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val - if !allunique(cartesian_inds) + allunique(cartesian_inds) || throw(ArgumentError("Indices should not overlap: $cartesian_inds.")) - end - - rmin, rmax = extrema(getindex.(cartesian_inds, 1)) - cmin, cmax = extrema(getindex.(cartesian_inds, 2)) + rowrange = getindex.(cartesian_inds, 1) + colrange = getindex.(cartesian_inds, 2) - gridsize = (rmax - rmin + 1, cmax - cmin + 1) - - corner_NW = tensorexpr( - :(env.corners[ - NORTHWEST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) - ]), - (:χ_C_NW_1,), - (:χ_C_NW_2,), - ) - corner_NE = tensorexpr( - :(env.corners[ - NORTHEAST, mod1($(rmin - 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) - ]), - (:χ_C_NE_1,), - (:χ_C_NE_2,), - ) - corner_SE = tensorexpr( - :(env.corners[ - SOUTHEAST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmax + 1), size(ket, 2)) - ]), - (:χ_C_SE_1,), - (:χ_C_SE_2,), - ) - corner_SW = tensorexpr( - :(env.corners[ - SOUTHWEST, mod1($(rmax + 1), size(ket, 1)), mod1($(cmin - 1), size(ket, 2)) - ]), - (:χ_C_SW_1,), - (:χ_C_SW_2,), - ) - - edges_N = map(1:gridsize[2]) do i - return tensorexpr( - :(env.edges[ - NORTH, - mod1($(rmin - 1), size(ket, 1)), - mod1($(cmin + i - 1), size(ket, 2)), - ]), - ( - (i == 1 ? :χ_C_NW_2 : Symbol(:χ_E_N, i - 1)), - Symbol(:D_E_N_top, i), - Symbol(:D_E_N_bot, i), - ), - ((i == gridsize[2] ? :χ_C_NE_1 : Symbol(:χ_E_N, i)),), - ) - end - - edges_E = map(1:gridsize[1]) do i - return tensorexpr( - :(env.edges[ - EAST, - mod1($(rmin + i - 1), size(ket, 1)), - mod1($(cmax + 1), size(ket, 2)), - ]), - ( - (i == 1 ? :χ_C_NE_2 : Symbol(:χ_E_E, i - 1)), - Symbol(:D_E_E_top, i), - Symbol(:D_E_E_bot, i), - ), - ((i == gridsize[1] ? :χ_C_SE_1 : Symbol(:χ_E_E, i)),), - ) - end - - edges_S = map(1:gridsize[2]) do i - return tensorexpr( - :(env.edges[ - SOUTH, - mod1($(rmax + 1), size(ket, 1)), - mod1($(cmin + i - 1), size(ket, 2)), - ]), - ( - (i == gridsize[2] ? :χ_C_SE_2 : Symbol(:χ_E_S, i)), - Symbol(:D_E_S_top, i), - Symbol(:D_E_S_bot, i), - ), - ((i == 1 ? :χ_C_SW_1 : Symbol(:χ_E_S, i - 1)),), - ) - end - - edges_W = map(1:gridsize[1]) do i - return tensorexpr( - :(env.edges[ - WEST, - mod1($(rmin + i - 1), size(ket, 1)), - mod1($(cmin - 1), size(ket, 2)), - ]), - ( - (i == gridsize[1] ? :χ_C_SW_2 : Symbol(:χ_E_W, i)), - Symbol(:D_E_W_top, i), - Symbol(:D_E_W_bot, i), - ), - ((i == 1 ? :χ_C_NW_1 : Symbol(:χ_E_W, i - 1)),), - ) - end - - bra = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) - return tensorexpr( - :(bra[ - mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) - ]), - (Symbol(:d, i, "_", j),), - ( - ( - if i == 1 - Symbol(:D_E_N_bot, j) - else - Symbol(:D_bra_vertical, i - 1, "_", j) - end - ), - ( - if j == gridsize[2] - Symbol(:D_E_E_bot, i) - else - Symbol(:D_bra_horizontal, i, "_", j) - end - ), - ( - if i == gridsize[1] - Symbol(:D_E_S_bot, j) - else - Symbol(:D_bra_vertical, i, "_", j) - end - ), - ( - if j == 1 - Symbol(:D_E_W_bot, i) - else - Symbol(:D_bra_horizontal, i, "_", j - 1) - end - ), - ), - ) - end - - ket = map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j) - return tensorexpr( - :(ket[ - mod1($(rmin + i - 1), size(ket, 1)), mod1($(cmin + j - 1), size(ket, 2)) - ]), - (Symbol(:d, i, "_", j),), - ( - ( - if i == 1 - Symbol(:D_E_N_top, j) - else - Symbol(:D_ket_vertical, i - 1, "_", j) - end - ), - ( - if j == gridsize[2] - Symbol(:D_E_E_top, i) - else - Symbol(:D_ket_horizontal, i, "_", j) - end - ), - ( - if i == gridsize[1] - Symbol(:D_E_S_top, j) - else - Symbol(:D_ket_vertical, i, "_", j) - end - ), - ( - if j == 1 - Symbol(:D_E_W_top, i) - else - Symbol(:D_ket_horizontal, i, "_", j - 1) - end - ), - ), - ) - end + corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange) + edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange) + bra, ket = _contract_state_expr(rowrange, colrange) multiplication_ex = Expr( :call,