reorganize examples a bit

examples/submit_lumi.sh → examples/submit/submit_lumi.sh

@@ -8,10 +8,12 @@ files=(
 	"water_cp2k_noise.py"
     "water_path_integral_md.py"
     "water_train_validate.py"
+    "water_online_learning.py"
+    "proton_jump_plumed.py"
     "alanine_replica_exchange.py"
 )
 
-curl -O https://raw.githubusercontent.com/molmod/psiflow/main/examples/lumi.yaml
+curl -O https://raw.githubusercontent.com/molmod/psiflow/main/examples/submit/lumi.yaml
 
 run_dir=$(pwd)/run_examples
 mkdir $run_dir && cp lumi.yaml $run_dir && cd $run_dir

examples/learning_path_integral_md.py → examples/water_online_learning.py

@@ -3,7 +3,7 @@
 import psiflow
 from psiflow.reference import CP2K
 from psiflow.data import Dataset
-from psiflow.sampling import Walker, randomize
+from psiflow.sampling import Walker
 from psiflow.models import MACE
 from psiflow.hamiltonians import MACEHamiltonian
 from psiflow.learning import Learning
@@ -26,42 +26,41 @@ def main():
         num_channels=24,
         patience=8,
         scheduler_patience=4,
+        max_num_epochs=200,
     )
     model.add_atomic_energy('H',  cp2k.compute_atomic_energy('H', box_size=9))
     model.add_atomic_energy('O',  cp2k.compute_atomic_energy('O', box_size=9))
 
-    data = Dataset.load('data/water_train.xyz').filter('energy')
+    state = Dataset.load('data/water_train.xyz')[0]
     walkers = (
-        Walker(data[0], temperature=300, pressure=0.1).multiply(40) +
-        Walker(data[0], temperature=450, pressure=0.1).multiply(40) +
-        Walker(data[0], temperature=600, pressure=0.1).multiply(40)
+        Walker(state, temperature=300, pressure=0.1).multiply(40) +
+        Walker(state, temperature=450, pressure=0.1).multiply(40) +
+        Walker(state, temperature=600, pressure=0.1).multiply(40)
     )
-    randomize(walkers, data)  # random initialization
     learning = Learning(
         cp2k,
         path_output,
         wandb_project='psiflow_examples',
         wandb_group='my_water_test',
-        initial_data=data,
     )
 
     model, walkers = learning.passive_learning(
-            model,
-            walkers,
-            hamiltonian=MACEHamiltonian.mace_mp0(),
-            steps=10000,
-            step=2000,
-            )
+        model,
+        walkers,
+        hamiltonian=MACEHamiltonian.mace_mp0(),
+        steps=10000,
+        step=2000,
+    )
 
     for i in range(3):
         model, walkers = learning.active_learning(
-                model,
-                walkers,
-                steps=500,
-                )
+            model,
+            walkers,
+            steps=500,
+        )
 
     # PIMD phase for low-temperature walkers
-    for j, walker in walkers[:20]:
+    for j, walker in enumerate(walkers[:40]):
         walker.nbeads = 8
     model, walkers = learning.active_learning(
         model,