From bc7ac93abd25fceb2a453c4aaa5d933328e1cee3 Mon Sep 17 00:00:00 2001
From: James Osborn <osborn@alcf.anl.gov>
Date: Fri, 16 Aug 2024 13:11:03 -0500
Subject: [PATCH] add Gauge type

---
 src/gauge/gaugeUtils.nim | 57 +++++++++++++++++++++++++++++++++-------
 src/physics/qcdTypes.nim | 16 +++++++++++
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/src/gauge/gaugeUtils.nim b/src/gauge/gaugeUtils.nim
index 8afb886..a61c3dc 100644
--- a/src/gauge/gaugeUtils.nim
+++ b/src/gauge/gaugeUtils.nim
@@ -12,23 +12,39 @@ import maths, rng, physics/qcdTypes
 
 import std/[hashes, tables]
 
-#[
 type
   GroupKind* = enum
-    gkU, gkSU, gkHerm, gkAntiHerm  # traceless, real/complex
-  Gauge*[T] = object
-    u*: seq[T]
-    n*: int
+    gkU, gkAntiHerm,
+    gkSU, gkTracelessAntiHerm,
+    gkO, gkRealAntiSym,
+    gkSO, gkTracelessRealAntiSym,
+    gkSp,
+    gkHerm, gkTracelessHerm,
+    gkRealSym, gkTracelessRealSym
+  GaugeBase* {.inheritable.} = object
     group*: GroupKind
-]#
+    n*: int
+  GaugeObj*[T] = object of GaugeBase
+    u*: seq[T]
+  Gauge*[T] = ref GaugeObj[T]
 
+proc default*(g: Gauge) =
+  let nd = g.u.len
+  threads:
+    if g.group in {gkU, gkSU, gkO, gkSO}:
+      for i in 0..<nd:
+        g.u[i] := 1
+    else:
+      for i in 0..<nd:
+        g.u[i] := 0
 
-proc newGauge*(l: Layout): auto =
+proc newGaugeSeq*(l: Layout): auto =
   let nd = l.nDim
   result = newSeq[type(l.ColorMatrix())](nd)
   for i in 0..<nd:
     result[i] = l.ColorMatrix()
     result[i] := 1
+template newGauge*(l: Layout): auto = newGaugeSeq(l)
 
 proc newGauge*[T](g: seq[T]): auto =
   let nd = g.len
@@ -37,6 +53,18 @@ proc newGauge*[T](g: seq[T]): auto =
     result[i] = g[0].l.ColorMatrix()
     result[i] := g[i]
 
+proc newGauge*(l: Layout, gk: GroupKind, n: static[int] = getDefaultNc()): auto =
+  type LCM = type l.ColorMatrix(n)
+  type GaugeT = Gauge[LCM]
+  let nd = l.nDim
+  result = new GaugeT
+  result.n = n
+  result.group = gk
+  result.u.newSeq(nd)
+  for i in 0..<nd:
+    result.u[i].new(l)
+  result.default
+
 proc newGaugeS*(l: Layout): auto =
   let nd = l.nDim
   result = newSeq[type(l.ColorMatrixS())](nd)
@@ -253,6 +281,8 @@ proc plaq*[T](uu: openArray[T]): auto =
   result = pl
   toc("plaq end", flops=lo.nSites.float*float(np*(2*8*nc*nc*nc-1)))
 
+template plaq*(g: Gauge): auto = plaq(g.u)
+
 discard """
 # s[mu] = a_mu s[mu] + f_mu_nu Unu Vmu Unu^+ + b_mu_nu Unu^+ Vmu Unu
 proc staples*[T,A,F,B](staples,uu,vv:openArray[T]; aa:openArray[A];
@@ -497,7 +527,7 @@ proc wline0*(g:auto, line:openarray[int]):auto =
   ## Compute the trace of ordered product of gauge links, the Wilson Line.
   ## The line is given as a list of integers +/- 1..nd, where the sign
   ## denotes forward/backward and the number denotes the dimension.
-  tic()
+  tic("wline0")
   # echo line
   type L = Link[typeof(g[0])]
   const nc = g[0][0].ncols
@@ -1475,7 +1505,7 @@ when isMainModule:
     @[3,4,-3,-4],
   ]
 
-  echoTimers()
+  #echoTimers()
   resetTimers()
 
   let wl = g.wilsonLines [
@@ -1516,5 +1546,12 @@ when isMainModule:
     echo st[i].norm2
   ]#
 
-  echoTimers()
+  var u1 = lo.newGauge(gkU, 1)
+  echo plaq(u1)
+  var su2 = lo.newGauge(gkSU, 2)
+  echo plaq(su2)
+  var su3 = lo.newGauge(gkSU, 3)
+  echo plaq(su3)
+
+  #echoTimers()
   qexFinalize()
diff --git a/src/physics/qcdTypes.nim b/src/physics/qcdTypes.nim
index 08c6f57..007ca9f 100644
--- a/src/physics/qcdTypes.nim
+++ b/src/physics/qcdTypes.nim
@@ -50,6 +50,7 @@ type
   SColorVectorV* = Color[VectorArray[nc,SComplexV]]
   SColorMatrix* = Color[MatrixArray[nc,nc,SComplex]]
   SColorMatrixV* = Color[MatrixArray[nc,nc,SComplexV]]
+
   SLatticeReal* = Field[1,float32]
   #SLatticeRealG*[V:static[int]] = Field[V,Svec0]
   SLatticeRealV* = Field[VLEN,Svec0]
@@ -59,6 +60,7 @@ type
   SLatticeColorVectorV* = Field[VLEN,SColorVectorV]
   SLatticeColorMatrix* = Field[1,SColorMatrix]
   SLatticeColorMatrixV* = Field[VLEN,SColorMatrixV]
+
   #DComplex* = AsComplex[tuple[re,im:float64]]
   #DComplexV* = AsComplex[tuple[re,im:Dvec0]]
   DComplex* = ComplexType[float64]
@@ -69,6 +71,9 @@ type
   DColorVectorV* = Color[VectorArray[nc,DComplexV]]
   DColorMatrix* = Color[MatrixArray[nc,nc,DComplex]]
   DColorMatrixV* = Color[MatrixArray[nc,nc,DComplexV]]
+
+  ColorMatrixN*[n:static[int],T] = Color[MatrixArray[n,n,T]]
+
   DLatticeReal* = Field[1,float64]
   DLatticeRealV* = Field[VLEN,Dvec0]
   DLatticeComplex* = Field[1,DComplex]
@@ -390,6 +395,17 @@ proc blend*(r:var auto; x:ptr char; b:ptr char; blnd:int) {.inline.} =
   of -8: loop(blendm8)
   else: discard
 
+template newDComplexV*[V:static[int]](l: Layout[V]): auto =
+  ComplexType[`SimdD V`]()
+
+proc newField*[V:static[int]](l: Layout[V], T:typedesc): Field[V,T] =
+  result.new(l)
+
+proc ColorMatrix*(l: Layout, n: static[int]): auto =
+  type C = type(l.newDComplexV)
+  type CM = ColorMatrixN[n,C]
+  result = l.newField(CM)
+
 macro makeConstructors(x: untyped): untyped =
   template mp(f,r,rslt: untyped) =
     proc f*(l: Layout): r =