HYP refactored with higher order derivatives

src/field/fieldET.nim

@@ -20,8 +20,8 @@ type
     l*:Layout[V]
     elemSize*:int
   Field*[V:static[int],T] = ref FieldObj[V,T]
-  FieldArray*[V:static[int],T] = object  ## share a single alignedMem
-    shape*:seq[int]
+  FieldArray*[N:static[int],V:static[int],T] = object  ## share a single alignedMem
+    shape*:array[N,int]
     arr:seq[Field[V,T]]
   Subsetted*[F,S] = object
     field*:F
@@ -182,8 +182,8 @@ proc newFarrElem[V:static[int],T](f:var Field[V,T]; l:Layout[V]; s:alignedMem[T]
   f.s.data = cast[typeof(s.data)](cast[int](s.data) + offset*l.nSitesOuter*s.stride)
   f.elemSize = sizeOf(T)
 
-proc newFieldArray*[V:static[int],T](l:Layout[V]; t:typedesc[Field[V,T]]; n: int):FieldArray[V,T] {.noinit.} =
-  result.shape = @[n]
+proc newFieldArray*[V:static[int],T](l:Layout[V]; t:typedesc[Field[V,T]]; n: int):FieldArray[1,V,T] {.noinit.} =
+  result.shape[0] = n
   result.arr = newseq[t](n)
   var s:typeof(result.arr[0].s)
   s.newU(l.nSitesOuter*n)
@@ -194,8 +194,9 @@ template newFieldArray2*[V:static[int],T](l:Layout[V]; ty:typedesc[Field[V,T]];
     ns: array[2,int]; constraint: untyped):untyped =
   let n = ns[0]
   let m = ns[1]
-  var r {.noinit.} :FieldArray[V,T]
-  r.shape = @[n,m]
+  var r {.noinit.} :FieldArray[2,V,T]
+  r.shape[0] = n
+  r.shape[1] = m
   r.arr = newseq[ty](n*m)
   var t = 0
   for i in 0..<n:
@@ -217,7 +218,7 @@ template newFieldArray2*[V:static[int],T](l:Layout[V]; ty:typedesc[Field[V,T]];
       inc k
   r
 
-proc newOneOf*[V:static[int],T](fa:FieldArray[V,T]):FieldArray[V,T] {.noinit.} =
+proc newOneOf*[N:static[int],V:static[int],T](fa:FieldArray[N,V,T]):FieldArray[N,V,T] {.noinit.} =
   result.shape = fa.shape
   result.arr = newseq[Field[V,T]](fa.arr.len)
   var s:typeof(fa.arr[0].s)
@@ -236,6 +237,28 @@ proc newOneOf*[V:static[int],T](fa:FieldArray[V,T]):FieldArray[V,T] {.noinit.} =
     newFarrElem(result.arr[i], l, s, t)
     inc t
 
+proc copyFieldArray*[N:static[int],V:static[int],T](
+    dst: FieldArray[N,V,T];
+    src: FieldArray[N,V,T]
+  ) =
+  doAssert dst.arr.len == src.arr.len, "FieldArray size mismatch"
+  doAssert dst.shape == src.shape, "FieldArray shape mismatch"
+  for i in 0..<src.arr.len:
+    doAssert (src.arr[i] == nil) == (dst.arr[i] == nil),
+      "FieldArray nil pattern mismatch"
+    if src.arr[i] == nil:
+      continue
+    when compiles(assignM(dst.arr[i], src.arr[i])):
+      assignM(dst.arr[i], src.arr[i])
+    else:
+      dst.arr[i] := src.arr[i]
+
+template `:=`*[N:static[int],V:static[int],T](
+    dst: FieldArray[N,V,T];
+    src: FieldArray[N,V,T]
+  ) =
+  copyFieldArray(dst, src)
+
 template dataPtr*[V:static[int],T](x: Field[V,T]): auto = x.s.data
 template isWrapper*(x: SomeField): bool = false
 template isWrapper*(x: typedesc[SomeField]): bool = false

src/gauge/gaugeAction.nim

@@ -158,7 +158,7 @@ proc gaugeActionDeriv*[T](c: GaugeActionCoeffs, uu: openArray[T], f: array|seq,
   let cp = c.plaq / float(nc)
   let cr = c.rect / float(nc)
   var cs = startCornerShifts(uu)
-  var ru:FieldArray[type(u[0]).V,type(u[0]).T]  # the rect parts of 3
+  var ru:FieldArray[2,type(u[0]).V,type(u[0]).T]  # the rect parts of 3
   var sb:seq[seq[ShiftB[type(u[0][0])]]]  # backward ru
   var sf:seq[seq[ShiftB[type(u[0][0])]]]  # forward stf
   if cr!=0:

src/gauge/gaugeUtils.nim

@@ -414,6 +414,42 @@ proc contractProjectTAH*[T](g:openArray[T], f:openArray[T]) =
         let s = u[mu][e]*o[mu][e].adj
         o[mu][e].projectTAH s
 
+proc contractProjectTAHVJP*[T](g: openArray[T], f: openArray[T],
+                                fbar: openArray[T],
+                                dbar: openArray[T], ubar: openArray[T]) =
+  ## VJP of contractProjectTAH.
+  ##
+  ## Forward: F = projectTAH(U · D†)
+  ##
+  ## Given F̄ (adjoint seed), computes:
+  ##   D̄ = adjoint w.r.t. D (the input force before contraction)
+  ##   Ū = adjoint w.r.t. U (the gauge field)
+  ##
+  ## Mathematical derivation:
+  ##   dT = ⟨F̄, dF⟩ = ⟨F̄, projectTAH(U · dD†)⟩
+  ##   Since projectTAH is self-adjoint on TAH: = ⟨projectTAH(F̄), U · dD†⟩
+  ##   = Re⟨projectTAH(F̄)† · U, dD⟩
+  ##
+  ## Since projectTAH output is anti-Hermitian: X† = -X
+  ## So: D̄ = projectTAH(F̄)† · U = -projectTAH(F̄) · U
+  ##
+  ## For Ū: dT = ⟨F̄, projectTAH(dU · D†)⟩ = ⟨projectTAH(F̄), dU · D†⟩
+  ##   = Re⟨projectTAH(F̄) · D, dU⟩
+  ## So: Ū = projectTAH(F̄) · D
+  let nd = g.len
+  let u = cast[ptr cArray[T]](unsafeAddr(g[0]))
+  let d = cast[ptr cArray[T]](unsafeAddr(f[0]))
+  let fb = cast[ptr cArray[T]](unsafeAddr(fbar[0]))
+  let db = cast[ptr cArray[T]](unsafeAddr(dbar[0]))
+  let ub = cast[ptr cArray[T]](unsafeAddr(ubar[0]))
+  threads:
+    for mu in 0..<nd:
+      for e in u[mu]:
+        var pf {.noInit.}: type(fb[0][0])
+        pf.projectTAH fb[mu][e]  # projectTAH(F̄)
+        db[mu][e] := -pf * u[mu][e]  # D̄ = -projectTAH(F̄) · U
+        ub[mu][e] := pf * d[mu][e]   # Ū = projectTAH(F̄) · D
+
 type
   Link[F:ref] = object
     field:F