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,