Rename amort_init for consistency

thebuckleylab · Jun 18, 2024 · 906a398 · 906a398
1 parent 243fe37
commit 906a398
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/jpc/__init__.py b/jpc/__init__.py
@@ -3,7 +3,7 @@
 from core._init import (
     init_activities_with_ffwd as init_activities_with_ffwd,
     init_activities_from_gaussian as init_activities_from_gaussian,
-    amort_init as amort_init
+    init_activities_with_amort as init_activities_with_amort
 )
 from core._energies import (
     pc_energy_fn as pc_energy_fn,

diff --git a/jpc/_train.py b/jpc/_train.py
@@ -10,7 +10,7 @@
 )
 from jpc import (
     init_activities_with_ffwd,
-    amort_init,
+    init_activities_with_amort,
     solve_pc_activities,
     compute_pc_param_grads,
     compute_gen_param_grads,
@@ -158,15 +158,15 @@ def make_hpc_step(
     """
     gen_optim, amort_optim = optims
     gen_opt_state, amort_opt_state = opt_states
-    activities = amort_init(
+    amort_activities = init_activities_with_amort(
         amortiser=amortiser,
         generator=generator,
         output=output
     )
-    train_mse_loss = mean((input - activities[0])**2)
+    train_mse_loss = mean((input - amort_activities[0])**2)
     equilib_activities = solve_pc_activities(
         network=generator,
-        activities=activities[1:],
+        activities=amort_activities[1:],
         output=output,
         input=input,
         solver=solver,

diff --git a/jpc/core/_init.py b/jpc/core/_init.py
@@ -67,7 +67,7 @@ def init_activities_from_gaussian(
     return activities
 
 
-def amort_init(
+def init_activities_with_amort(
         amortiser: PyTree[Callable],
         generator: PyTree[Callable],
         output: ArrayLike