From c453637e88f8d825f162de05fdee5e26147d3ee3 Mon Sep 17 00:00:00 2001
From: "Fernando P. Panadero" <fernando.p@csic.es>
Date: Mon, 5 May 2025 11:30:02 +0200
Subject: [PATCH] Flow for OBC similar as in oqcd.

---
 src/LatticeGPU.jl |  2 +-
 src/YM/YM.jl      |  2 +-
 src/YM/YMact.jl   | 48 +++++++++++++++++++++++++++++++++++++++++++++++
 src/YM/YMflow.jl  | 25 ++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/src/LatticeGPU.jl b/src/LatticeGPU.jl
index 46f5ac6..bf3995b 100644
--- a/src/LatticeGPU.jl
+++ b/src/LatticeGPU.jl
@@ -40,7 +40,7 @@ include("YM/YM.jl")
 using .YM
 export ztwist
 export YMworkspace, GaugeParm, force0_wilson!, field, field_pln, randomize!, zero!, norm2
-export force_gauge, MD!
+export force_gauge, force_gauge_flw, MD!
 export gauge_action, hamiltonian, plaquette, HMC!, OMF4!
 export Eoft_clover, Eoft_plaq, Qtop
 export FlowIntr, wfl_euler, zfl_euler, wfl_rk2, zfl_rk2, wfl_rk3, zfl_rk3
diff --git a/src/YM/YM.jl b/src/YM/YM.jl
index 16165e4..343ca5c 100644
--- a/src/YM/YM.jl
+++ b/src/YM/YM.jl
@@ -165,7 +165,7 @@ include("YMfields.jl")
 export randomize!, zero!, norm2
 
 include("YMact.jl")
-export krnl_plaq!, force_gauge, force_wilson
+export krnl_plaq!, force_gauge, force_gauge_flw, force_wilson
 
 include("YMhmc.jl")
 export gauge_action, hamiltonian, plaquette, HMC!, MD!
diff --git a/src/YM/YMact.jl b/src/YM/YMact.jl
index f9849bd..a3316e4 100644
--- a/src/YM/YMact.jl
+++ b/src/YM/YMact.jl
@@ -320,6 +320,22 @@ function krnl_force_impr_pln!(frc1, frc2, U::AbstractArray{T}, c0, c1, Ubnd, cG,
     return nothing
 end
 
+function bnd_rescale_flw!(frc1, lp::SpaceParm{N,M,BC_OPEN,D}) where {N,M,D}
+
+    @inbounds begin
+        b = Int64(CUDA.threadIdx().x)
+        r = Int64(CUDA.blockIdx().x)
+        I = point_coord((b,r), lp)
+        it = I[N]
+
+        for id in 1:N-1
+            if (((it == 1) || (it == lp.iL[4])))
+                frc1[b,id,r] = 2*frc1[b,id,r]
+            end
+        end
+    end
+    return nothing
+end
 
 ##
 ## SF
@@ -918,6 +934,38 @@ end
 force_gauge(ymws::YMworkspace, U, c0, gp, lp) = force_gauge(ymws, U, c0, gp.cG[1], gp, lp)
 force_gauge(ymws::YMworkspace, U, gp, lp) = force_gauge(ymws, U, gp.c0, gp.cG[1], gp, lp)
 
+"""
+    function force_gauge_flw(ymws::YMworkspace, U, c0, cG, gp::GaugeParm, lp::SpaceParm{N,M,BC_OPEN,D})
+
+Computes the force for the gauge flow with Open Boundaries. An aditional factor two in the boundaries
+is included, see
+
+M. Luescher, S. Schaefer: "Lattice QCD with open boundary conditions and twisted-mass reweighting", Comput.Phys.Commun. 184 (2013) 519,
+
+for more details.
+
+"""
+function force_gauge_flw(ymws::YMworkspace, U, c0, cG, gp::GaugeParm, lp::SpaceParm{N,M,BC_OPEN,D}) where {NI,N,M,D}
+
+    ztw = ztwist(gp, lp)
+    if abs(c0-1) < 1.0E-10
+        @timeit "Wilson gauge force" begin
+            force_pln!(ymws.frc1, ymws.frc2, U, gp.Ubnd, cG, ztw, lp::SpaceParm)
+        end
+    else
+        @timeit "Improved gauge force" begin
+            force_pln!(ymws.frc1, ymws.frc2, U, gp.Ubnd, cG, ztw, lp::SpaceParm, c0)
+        end
+    end
+
+    CUDA.@sync begin
+        CUDA.@cuda threads=lp.bsz blocks=lp.rsz bnd_rescale_flw!(ymws.frc1,lp::SpaceParm)
+    end
+
+    return nothing
+end
+
+
 """ 
     function force_wilson(ymws::YMworkspace, U, gp::GaugeParm, lp::SpaceParm)
 
diff --git a/src/YM/YMflow.jl b/src/YM/YMflow.jl
index 4f841c3..1c726b6 100644
--- a/src/YM/YMflow.jl
+++ b/src/YM/YMflow.jl
@@ -201,6 +201,31 @@ function flw(U, int::FlowIntr{NI,T}, ns::Int64, eps, gp::GaugeParm, lp::SpacePar
 end
 flw(U, int::FlowIntr{NI,T}, ns::Int64, gp::GaugeParm, lp::SpaceParm, ymws::YMworkspace) where {NI,T} = flw(U, int, ns, int.eps, gp, lp, ymws)
 
+function flw(U, int::FlowIntr{NI,T}, ns::Int64, eps, gp::GaugeParm, lp::SpaceParm{N,M,BC_OPEN,D}, ymws::YMworkspace) where {NI,T,N,M,D}
+    @timeit "Integrating flow equations" begin
+        for i in 1:ns
+            force_gauge_flw(ymws, U, int.c0, 1, gp, lp)
+            if int.add_zth
+                add_zth_term(ymws::YMworkspace, U, lp)
+            end
+            ymws.mom .= ymws.frc1
+            U .= expm.(U, ymws.mom, 2*eps*int.r)
+
+            for k in 1:NI
+                force_gauge_flw(ymws, U, int.c0, 1, gp, lp)
+                if int.add_zth
+                    add_zth_term(ymws::YMworkspace, U, lp)
+                end
+                ymws.mom .= int.e0[k].*ymws.mom .+ int.e1[k].*ymws.frc1
+                U .= expm.(U, ymws.mom, 2*eps)
+            end
+        end
+    end
+
+    return nothing
+end
+flw(U, int::FlowIntr{NI,T}, ns::Int64, gp::GaugeParm, lp::SpaceParm{N,M,BC_OPEN,D}, ymws::YMworkspace) where {NI,T,N,M,D} = flw(U, int, ns, int.eps, gp, lp, ymws)
+
 
 ##
 # Adaptive step size integrators