fix bug in env_vars parsing and forwarding; add new examples

configs/hortense.yaml

@@ -7,6 +7,8 @@ ModelEvaluation:
   cores_per_worker: 12
   gpu: True
   max_simulation_time: 20
+  env_vars:
+    KMP_BLOCKTIME: "1"
   slurm:
     partition: "gpu_rome_a100"
     account: "2023_070"
@@ -19,6 +21,8 @@ ModelTraining:
   cores_per_worker: 12
   gpu: true
   max_training_time: 40
+  env_vars:
+    OMP_PROC_BIND: "spread"
   slurm:
     partition: "gpu_rome_a100"
     account: "2023_070"
@@ -30,6 +34,7 @@ ModelTraining:
 CP2K:
   cores_per_worker: 64
   max_evaluation_time: 30
+  memory_limit: 2GB
   launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 -bind-to core cp2k.psmp'
   slurm:
     partition: "cpu_rome"

configs/lumi.yaml

@@ -6,6 +6,7 @@ default_threads: 8
 CP2K:
   cores_per_worker: 32
   max_evaluation_time: 20
+  memory_limit: 2GB
   launch_command: 'singularity exec -e --no-init oras://ghcr.io/molmod/cp2k:2024.1 /opt/entry.sh mpirun -np 32 cp2k.psmp'
   slurm:
     partition: "standard"

examples/data/acetaldehyde.xyz

@@ -0,0 +1,9 @@
+7
+Properties=species:S:1:pos:R:3
+O	0.694151672	0.776743934	-0.455455855
+C	0.195993254	-0.270095005	-0.307053207
+C	-0.846060202	-0.538006022	0.669585079 
+H	0.515801613	-1.097661033	-0.987914453
+H	-0.589257101	-0.505600908	1.733123281 
+H	-1.553309062	0.309375207	0.558315778 
+H	-1.411674563	-1.440354174	0.5617281699

examples/data/vinyl_alcohol.xyz

@@ -0,0 +1,9 @@
+7
+Properties=species:S:1:pos:R:3
+O	1.041371715	-0.216863172	0.001603252
+C	-0.098316254	0.512294574	-0.01021628
+C	-1.225162144	-0.248210652	0.020868361
+H	-0.087363805	1.596485281	-0.07557041
+H	0.61765221	-1.094559605	-0.02702971
+H	-2.216985293	0.211688229	-0.00469380
+H	-1.115257687	-1.357478425	-0.04507284

examples/online_learning_pimd.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import psiflow
+from psiflow.reference import CP2K
+from psiflow.data import Dataset
+from psiflow.sampling import Walker
+from psiflow.models import MACE
+from psiflow.hamiltonians import MACEHamiltonian
+from psiflow.learning import Learning
+
+
+def main():
+    path_output = Path.cwd() / 'output'
+
+    with open('data/cp2k_input.txt', 'r') as f: cp2k_input = f.read()
+    cp2k = CP2K(cp2k_input)
+
+    model = MACE(
+        batch_size=4,
+        lr=0.02,
+        max_ell=3,
+        r_max=6.5,
+        energy_weight=100,
+        correlation=3,
+        max_L=1,
+        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))
+
+    state = Dataset.load('data/water_train.xyz')[0]
+    walkers = (
+        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)
+    )
+    learning = Learning(
+        cp2k,
+        path_output,
+        wandb_project='psiflow_examples',
+        wandb_group='water_learning_pimd',
+    )
+
+    model, walkers = learning.passive_learning(
+        model,
+        walkers,
+        hamiltonian=MACEHamiltonian.mace_mp0(),
+        steps=10000,
+        step=2000,
+    )
+
+    for i in range(3):
+        model, walkers = learning.active_learning(
+            model,
+            walkers,
+            steps=2000,
+        )
+
+    # PIMD phase for low-temperature walkers
+    for j, walker in enumerate(walkers[:40]):
+        walker.nbeads = 32
+    model, walkers = learning.active_learning(
+        model,
+        walkers,
+        steps=500,
+    )
+
+
+if __name__ == '__main__':
+    with psiflow.load():
+        main()

examples/proton_jump_plumed.py

@@ -0,0 +1,67 @@
+from ase.units import kJ, mol
+import numpy as np
+
+import psiflow
+from psiflow.data import Dataset
+from psiflow.geometry import Geometry
+from psiflow.hamiltonians import PlumedHamiltonian, MACEHamiltonian
+from psiflow.sampling import Walker, sample, quench, Metadynamics, replica_exchange
+
+
+PLUMED_INPUT = """UNITS LENGTH=A ENERGY=kj/mol
+d_C: DISTANCE ATOMS=3,5
+d_O: DISTANCE ATOMS=1,5
+CV: COMBINE ARG=d_C,d_O COEFFICIENTS=1,-1 PERIODIC=NO
+
+"""
+
+
+def get_bias(kappa: float, center: float):
+    plumed_str = PLUMED_INPUT
+    plumed_str += '\n'
+    plumed_str += 'RESTRAINT ARG=CV KAPPA={} AT={}\n'.format(kappa, center)
+    return PlumedHamiltonian(plumed_str)
+
+
+def main():
+    aldehyd = Geometry.load('data/acetaldehyde.xyz')
+    alcohol = Geometry.load('data/vinyl_alcohol.xyz')
+
+    mace = MACEHamiltonian.mace_cc()
+    energy = mace.compute([aldehyd, alcohol], 'energy').result()
+    energy = (energy - np.min(energy)) / (kJ / mol)
+    print('E_vinyl - E_aldehyde = {:7.3f} kJ/mol'.format(energy[1] - energy[0]))
+
+    # generate initial structures using metadynamics
+    plumed_str = PLUMED_INPUT
+    plumed_str += 'METAD ARG=CV PACE=5 SIGMA=0.25 HEIGHT=5\n'
+    metadynamics = Metadynamics(plumed_str)
+
+    # create 40 identical walkers
+    walkers = Walker(
+        aldehyd,
+        hamiltonian=mace,
+        temperature=300,
+        metadynamics=metadynamics,
+    ).multiply(4)
+
+    # do MTD and create large dataset from all trajectories
+    outputs = sample(walkers, steps=2000, step=20, start=1000)
+    data_mtd = sum([o.trajectory for o in outputs], start=Dataset([]))
+
+    # initialize walkers for umbrella sampling
+    walkers = []
+    for i, center in enumerate(np.linspace(1, 3, num=16)):
+        bias = get_bias(kappa=1500, center=center)
+        hamiltonian = mace + bias
+        walker = Walker(alcohol, hamiltonian=hamiltonian, temperature=300)
+        walkers.append(walker)
+    quench(walkers, data_mtd)  # make sure initial structure is reasonable
+    replica_exchange(walkers, trial_frequency=100)  # use REX for improved sampling
+
+    outputs = sample(walkers, steps=1000, step=10)
+
+
+if __name__ == '__main__':
+    with psiflow.load() as f:
+        main()

psiflow/execution.py

@@ -170,7 +170,7 @@ def __init__(
         self,
         max_simulation_time: Optional[float] = None,
         timeout: float = (10 / 60),  # 5 seconds
-        env_vars: Optional[dict] = None,
+        env_vars: Optional[dict[str, str]] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -186,9 +186,10 @@ def __init__(
             'PYTHONUNBUFFERED': 'TRUE',
         }
         if env_vars is None:
-            env_vars = dict(default_env_vars)
+            env_vars = default_env_vars
         else:
-            env_vars.update(default_env_vars)
+            default_env_vars.update(env_vars)
+            env_vars = default_env_vars
         self.env_vars = env_vars
 
     def server_command(self):
@@ -245,7 +246,7 @@ def __init__(
         self,
         gpu=True,
         max_training_time: Optional[float] = None,
-        env_vars: Optional[dict] = None,
+        env_vars: Optional[dict[str, str]] = None,
         **kwargs,
     ) -> None:
         super().__init__(gpu=gpu, **kwargs)
@@ -266,9 +267,10 @@ def __init__(
             'PYTHONUNBUFFERED': 'TRUE',
         }
         if env_vars is None:
-            env_vars = dict(default_env_vars)
+            env_vars = default_env_vars
         else:
-            env_vars.update(default_env_vars)
+            default_env_vars.update(env_vars)
+            env_vars = default_env_vars
         self.env_vars = env_vars
 
     def train_command(self, initialize: bool = False):

psiflow/sampling/client.py

@@ -7,7 +7,6 @@ def main():
     import time
     from pathlib import Path
 
-    import torch
     from ase.io import read
     from ipi._driver.driver import run_driver
 
@@ -54,6 +53,10 @@ def main():
     assert args.address is not None
     assert args.start is not None
 
+    print("pid: {}".format(os.getpid()))
+    affinity = os.sched_getaffinity(os.getpid())
+    print("CPU affinity before function init: {}".format(affinity))
+
     template = Geometry.from_atoms(read(args.start))
     function = function_from_json(
         args.path_hamiltonian,
@@ -68,9 +71,8 @@ def main():
         verbose=True,
     )
 
-    print("pid: {}".format(os.getpid()))
-    print("CPU affinity: {}".format(os.sched_getaffinity(os.getpid())))
-    print("torch num threads: ", torch.get_num_threads())
+    affinity = os.sched_getaffinity(os.getpid())
+    print("CPU affinity after function init: {}".format(affinity))
 
     try:
         t0 = time.time()

psiflow/sampling/sampling.py

@@ -533,6 +533,8 @@ def _sample(
     command_server = definition.server_command()
     command_client = definition.client_command()
     resources = definition.wq_resources(max_nclients)
+    print('ENV VARS')
+    print(definition.env_vars)
     result = execute_ipi(
         len(walkers),
         hamiltonian_names,
@@ -543,7 +545,7 @@ def _sample(
         command_server,
         command_client,
         *plumed_list,
-        env_vars=definition.env_vars,
+        env_vars=dict(definition.env_vars),
         stdout=parsl.AUTO_LOGNAME,
         stderr=parsl.AUTO_LOGNAME,
         inputs=inputs,