From 9ebe1ced5e303c93a768e6afc966061bcdfcce9e Mon Sep 17 00:00:00 2001
From: James Osborn <osborn@alcf.anl.gov>
Date: Sat, 29 Jun 2024 21:02:12 -0500
Subject: [PATCH] add separate learning rate for hmasses

---
 src/experimental/stagag.nim | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/experimental/stagag.nim b/src/experimental/stagag.nim
index d6aee0a..0fd20e8 100644
--- a/src/experimental/stagag.nim
+++ b/src/experimental/stagag.nim
@@ -37,6 +37,7 @@ let
   outfn = stringParam("outfn", "")
 var
   lrate = floatParam("lrate", 0.001)
+  lrateh = floatParam("lrateh", lrate)
   pt0 = floatParam("t0", 0)
   pg0 = floatParam("g0", 0)
   pf0 = floatParam("f0", 0)
@@ -70,6 +71,7 @@ echoparam(fixhmasses)
 echoparam(md)
 echoparam(upit)
 echoparam(lrate)
+echoparam(lrateh)
 echoparam(anneal)
 echoparam(alpha)
 echoparam(checkg)
@@ -1475,7 +1477,7 @@ proc getGrad(m: Met) =
     #echo "costg: ", cgstat[i].mean, " ", cst[i+1]
     echo &"costg: {m/v:8.6f} {m:10.6f} {cst[i+1]}"
 
-proc updateParams(rate: float) =
+proc updateParams(ratefac: float) =
   let eps = 1e-8
   let n = cgstat.len
   #var s = 0.0
@@ -1492,6 +1494,9 @@ proc updateParams(rate: float) =
   echo "Params: ", p
   #s = rate*sqrt(n/s)
   for i in 0..<n:
+    var rate = ratefac * lrate
+    if i>0 and i<=hmasses.len:
+      rate = ratefac * lrateh
     #let m = cgstat[i].mean
     #let d = s*g[i]
     let d = rate*g[i]
@@ -1692,8 +1697,9 @@ block:
     measure()
     if upit > 0:
       if n mod upit == 0:
-        updateParams(sqrt(float upit)*lrate)
+        updateParams(sqrt(float upit))
         lrate *= anneal
+        lrateh *= anneal
     let ttot = getElapsedTime()
     echo "End trajectory update: ", tup, "  measure: ", ttot-tup, "  total: ", ttot
   let et = getElapsedTime()