From 9a90b21d05b0b9621dd94c8a48705c93afaff611 Mon Sep 17 00:00:00 2001
From: James Osborn <osborn@alcf.anl.gov>
Date: Thu, 19 Dec 2024 10:28:06 -0600
Subject: [PATCH] fix and test hisqsmear

---
 src/gauge/hisqsmear.nim | 115 ++++++++++++++++++++++++++++------------
 src/layout/layoutX.nim  |  40 ++++++++++++--
 2 files changed, 118 insertions(+), 37 deletions(-)

diff --git a/src/gauge/hisqsmear.nim b/src/gauge/hisqsmear.nim
index d9e931f..77b7749 100644
--- a/src/gauge/hisqsmear.nim
+++ b/src/gauge/hisqsmear.nim
@@ -39,7 +39,7 @@ proc smearGetForce*[T](
     var
       dsdx_dxdw = newOneOf(dsdu)
       dsdx_dxdw_dwdv = newOneOf(dsdu)
-    dsdx_dxdw.fat7lDeriv(su,dsdsu,fat7l2,sul,dsdsul,naik,info) # Second fat7
+    dsdx_dxdw.fat7lDeriv(w,dsdsu,fat7l2,w,dsdsul,naik,info) # Second fat7
     threads: # Unitary projection
       for mu in 0..<dsdx_dxdw_dwdv.len:
         for s in dsdx_dxdw_dwdv[mu]:
@@ -49,46 +49,93 @@ proc smearGetForce*[T](
   if displayPerformance: echo $(info)
   return smearedForce
 
-if isMainModule:
+when isMainModule:
+  import gauge/gaugeAction, strformat
   qexInit()
   let
     defaultLat = @[8,8,8,8]
-    hisq = newHISQ()
   var
     (lo, g, r) = setupLattice(defaultLat)
-    sg = lo.newGauge()
-    sgl = lo.newGauge()
-    f = lo.newGauge()
-    ff = lo.newGauge()
-    ffl = lo.newGauge()
-    g2 = lo.newGauge()
-    sg2 = lo.newGauge()
-    sgl2 = lo.newGauge()
     dg = lo.newGauge()
+    g2 = lo.newGauge()
     fd = lo.newGauge()
-  g.random
-  for mu in 0..<dg.len:
-    dg[mu] := 0.00001 * g[mu]
-    g2[mu] := g[mu] + dg[mu]
-    fd[mu] := 0
+    fl = lo.newGauge()
+    ll = lo.newGauge()
+    fl2 = lo.newGauge()
+    ll2 = lo.newGauge()
+    fc = lo.newGauge()
+    lc = lo.newGauge()
+    gc = GaugeActionCoeffs(plaq:1.0)
+    eps = floatParam("eps", 1e-6)
+    #warm = floatParam("warm", 1e-5)
+    hisq = newHISQ()
+
+  #g.unit
+  #g.gaussian r
+  g.random r
+  #g.warm warm, r
+  g.stagPhase
+  dg.gaussian r
+  for mu in 0..<g.len:
+    for e in g[mu]:
+      dg[mu][e] *= eps
+      g2[mu][e] := g[mu][e] + dg[mu][e]
+
+  proc check(da: float, tol: float) =
+    var ds = 0.0
+    for mu in 0..3:
+      ds += redot(dg[mu], fd[mu])
+    let r = (da-ds)/da
+    echo &"  da {da}  ds {ds}  rel {r}"
+    if abs(r) > tol*eps:
+      echo &"> ERROR rel error |{r}| > {tol*eps}"
+
+  proc checkG =
+    echo "Checking GaugeDeriv"
+    for mu in 0..<fd.len:
+      fd[mu] := 0
+    let a = gc.gaugeAction2(g)
+    let a2 = gc.gaugeAction2(g2)
+    gc.gaugeDeriv2(g, fd)
+    check(a2-a, 10)
+  checkG()
+
+  proc checkL(name: string, tol: float) =
+    echo "Checking ", name
+    for mu in 0..<fd.len:
+      fd[mu] := 0
+      fc[mu] := 0
+      lc[mu] := 0
+    resetTimers()
+    let f = hisq.smearGetForce(g, fl, ll)
+    let f2 = hisq.smearGetForce(g2, fl2, ll2)
+    gc.gaugeDeriv2(fl, fc)
+    gc.gaugeDeriv2(ll, lc)
+    f(fd, fc, lc)
+    let a = gc.gaugeAction2(fl) + gc.gaugeAction2(ll)
+    let a2 = gc.gaugeAction2(fl2) + gc.gaugeAction2(ll2)
+    check(a2-a, tol)
+
+  hisq.fat7first.oneLink = 1.0
+  hisq.fat7first.threeStaple = 0.0
+  hisq.fat7first.fiveStaple = 0.0
+  hisq.fat7first.sevenStaple = 0.0
+  hisq.fat7first.lepage = 0.0
+  hisq.fat7second.oneLink = 1.0
+  hisq.fat7second.threeStaple = 0.0
+  hisq.fat7second.fiveStaple = 0.0
+  hisq.fat7second.sevenStaple = 0.0
+  hisq.fat7second.lepage = 0.0
+  hisq.naik = 0.0
+  checkL("oneLink", 10)
+
+  hisq.fat7first.setHisqFat7(0.0, 0.0)
+  hisq.fat7second.setHisqFat7(0.0, 0.0)
+  hisq.naik = 0.0
+  checkL("all", 300)
 
-  echo g.plaq
-  echo g2.plaq
-  #hisq.smear(g,sg,sgl)
-  var force = hisq.smearGetForce(g,sg,sgl)
-  f.force(dg,dg)
-  hisq.smear(g2,sg2,sgl2)
+  hisq = newHisq()
+  checkL("all", 700)
 
-  echo f.plaq
-  echo sg.plaq
-  echo sgl.plaq
-  echo "--"
-  echo sg.plaq
-  echo sgl.plaq
-  ff.gaugeForce(sg)
-  ffl.gaugeForce(sgl)
-  echo "--"
-  echo f.plaq
-  echo sg.plaq
-  echo sgl.plaq
+  echoProf()
   qexFinalize()
diff --git a/src/layout/layoutX.nim b/src/layout/layoutX.nim
index 239adf0..1241e5e 100644
--- a/src/layout/layoutX.nim
+++ b/src/layout/layoutX.nim
@@ -9,7 +9,11 @@ export shiftX
 
 import qlayout
 
-# partition x into n blocks with geometry nx
+var noSplitInnerList = newSeq[int](0)
+proc noSplitInnerDim*(d: int) =
+  noSplitInnerList.add d
+
+# partition x into n blocks with geometry nx, only works with powers of 2
 # dist=0: prefer to split already split direction
 # dist=1: split unsplit directions first
 proc partitionGeom[T](lx,nx:var openArray[T]; x:openArray[T]; n,dist:int) =
@@ -19,7 +23,7 @@ proc partitionGeom[T](lx,nx:var openArray[T]; x:openArray[T]; n,dist:int) =
   var ww = n
   while ww > 1:
     var k = x.len-1
-    while (lx[k] and 1) > 0:
+    while (lx[k] and 1) > 0 or (k in noSplitInnerList):
       k.dec
       if k < 0:
         echo "not enough 2's in partitioned geom:"
@@ -27,7 +31,7 @@ proc partitionGeom[T](lx,nx:var openArray[T]; x:openArray[T]; n,dist:int) =
         echo " /", n
         quit(-1)
     for i in countdown(k-1,0):
-      if (lx[i] and 1) == 0:
+      if (lx[i] and 1) == 0 and (k notin noSplitInnerList):
         if dist == 0:
           if lx[i]>lx[k] or (lx[i]==lx[k] and nx[i]>nx[k]): k = i
         else:
@@ -298,6 +302,36 @@ proc paritySubset*(s: var Subset; l: Layout; par: int) =
     s.layoutSubset(l, "e")
   else:
     s.layoutSubset(l, "o")
+proc timesliceSubsetsRange*(l: Layout, i0,i1: int): seq[Subset] =
+  # find timeslice subsets within a range of outer sites [i0,i1] inclusive
+  doAssert(l.innerGeom[^1] == 1)
+  let nt = l.physGeom[^1]
+  result.newSeq(nt)
+  var x = newSeq[int32](l.nDim)
+  for i in 0..<nt:
+    result[i].lowOuter = l.nSites
+    result[i].highOuter = 0
+  for i in i0..i1:
+    let k = i * l.V
+    let t = l.coords[^1][k]
+    result[t].lowOuter = min(result[t].lowOuter, i)
+    result[t].highOuter = max(result[t].highOuter, i)
+  for i in 0..<nt:
+    var a = result[i].lowOuter
+    var b = result[i].highOuter
+    if a <= b:
+      inc b
+    else:
+      a = 0
+      b = 0
+    result[i].lowOuter = a
+    result[i].highOuter = b
+    result[i].low = a * l.V
+    result[i].high = b * l.V
+proc timesliceSubsets*(l: Layout): array[2,seq[Subset]] =
+  result[0] = l.timesliceSubsetsRange(0, l.nEvenOuter-1)
+  result[1] = l.timesliceSubsetsRange(l.nEvenOuter, l.nSitesOuter-1)
+
 template `len`*(s:Subset):untyped = s.high-s.low
 template `lenOuter`*(s:Subset):untyped = s.highOuter-s.lowOuter