make compatible with new ASE 3.23 release and improve numerical robus…

…tness PLUMED test

configs/threadpool.yaml

@@ -16,7 +16,7 @@ CP2K:
 CP2K_container:
   cores_per_worker: 2
   max_evaluation_time: 0.3
-  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /usr/local/bin/entrypoint.sh mpirun -np 2 -x OMP_NUM_THREADS=1 cp2k.psmp -i cp2k.inp'
+  launch_command: 'apptainer exec -e --no-init oras://ghcr.io/molmod/cp2k:2023.2 /opt/entry.sh mpirun -np 2 -x OMP_NUM_THREADS=1 cp2k.psmp -i cp2k.inp'
   use_threadpool: true
 GPAW:
   cores_per_worker: 2

psiflow/hamiltonians/utils.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import typeguard
+from ase import Atoms
 from ase.calculators.calculator import Calculator, all_changes
 from ase.stress import full_3x3_to_voigt_6_stress
 from ase.units import fs, kJ, mol, nm
@@ -62,10 +63,13 @@ def __init__(
         self.volume = volume
         self.max_force = max_force
 
-    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+    def calculate(
+        self, atoms: Optional[Atoms] = None, properties=None, system_changes=all_changes
+    ):
         # call to base-class to set atoms attribute
         Calculator.calculate(self, atoms)
 
+        assert atoms is not None
         assert self.centers.shape[0] == len(atoms)
         forces = (-1.0) * self.force_constant * (atoms.get_positions() - self.centers)
         energy = (
@@ -133,6 +137,7 @@ def __init__(
 
     def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         Calculator.calculate(self, atoms)
+        assert atoms is not None
         if self.initialize_plumed:
             self.plumed.cmd("setNatoms", len(atoms))
             self.plumed.cmd("init")
@@ -194,6 +199,7 @@ def __init__(
     def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         # call to base-class to set atoms attribute
         Calculator.calculate(self, atoms)
+        assert atoms is not None
         assert self.hessian.shape[0] == 3 * len(atoms)
 
         pos = atoms.positions.reshape(-1)

tests/test_data.py

@@ -136,8 +136,8 @@ def test_geometry(tmp_path):
         assert state.stdout == state_.stdout
         assert state.identifier == state.identifier
 
-    state.save(tmp_path / 'geo.xyz')
-    assert state == Geometry.load(tmp_path / 'geo.xyz')
+    state.save(tmp_path / "geo.xyz")
+    assert state == Geometry.load(tmp_path / "geo.xyz")
 
 
 def test_readwrite_cycle(dataset, tmp_path):
@@ -247,7 +247,7 @@ def test_dataset_from_xyz(tmp_path, dataset):
             dataset[i].result().per_atom.positions,
             atoms_list[i].get_positions(),
         )
-        assert dataset[i].result().energy == atoms_list[i].info["energy"]
+        assert dataset[i].result().energy == atoms_list[i].calc.results["energy"]
 
 
 def test_index_element_mask():
@@ -455,9 +455,8 @@ def test_data_extract(dataset):
 H 1 1 1
 H -1 1 1
 """,
-        None,
     )
-    states = [copy.deepcopy(s) for i in range(5)]
+    states = [copy.deepcopy(s) for _ in range(5)]
 
     states[1].order = {}
     states[2].order = {"some": 4.2}

tests/test_hamiltonian.py

@@ -26,7 +26,7 @@
 from psiflow.utils import copy_app_future, copy_data_future, dump_json
 
 
-def test_get_filename_hills(tmp_path):
+def test_get_filename_hills():
     plumed_input = """
 #METAD COMMENT TO BE REMOVED
 RESTART
@@ -54,7 +54,7 @@ def test_get_filename_hills(tmp_path):
     )
 
 
-def test_einstein(context, dataset, dataset_h2):
+def test_einstein(dataset, dataset_h2):
     state = dataset[0]
     hamiltonian = EinsteinCrystal(state, force_constant=1)
     evaluated = hamiltonian.evaluate(dataset[:10])
@@ -71,7 +71,7 @@ def test_einstein(context, dataset, dataset_h2):
     # test nonperiodic evaluation
     einstein = EinsteinCrystal(dataset_h2[3], force_constant=0.1)
     data = einstein.evaluate(dataset_h2)
-    energies, forces, stress = data.get("energy", "forces", "stress")
+    energies, _, stress = data.get("energy", "forces", "stress")
     assert np.all(energies.result() >= 0)
     assert np.any(energies.result() > 0)
     assert np.all(np.isnan(stress.result()))
@@ -193,7 +193,7 @@ def test_harmonic_force(dataset):
     assert h1 != h2
 
 
-def test_plumed_evaluate(context, dataset, tmp_path):
+def test_plumed_evaluate(dataset, tmp_path):
     geometry = dataset[0].result()
     atoms = Atoms(
         numbers=geometry.per_atom.numbers,
@@ -207,11 +207,8 @@ def test_plumed_evaluate(context, dataset, tmp_path):
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: DISTANCE ATOMS=1,2 NOPBC
 RESTRAINT ARG=CV AT={center} KAPPA={kappa}
-""".format(
-        center=center, kappa=kappa / (kJ / mol)
-    )
+""".format(center=center, kappa=kappa / (kJ / mol))
     calculator = PlumedCalculator(plumed_input)
-    # energy, forces, stress = evaluate_plumed(atoms, plumed_input)
     calculator.calculate(atoms)
     energy = calculator.results["energy"]
     forces = calculator.results["forces"]
@@ -238,8 +235,8 @@ def test_plumed_evaluate(context, dataset, tmp_path):
     hills = """#! FIELDS time CV sigma_CV height biasf
 #! SET multivariate false
 #! SET kerneltype gaussian
-     1.00000     2.1     2.0     7  0
-     2.00000     2.2     2.0     7  0
+     1.00000     2.5     2.0     70  0
+     2.00000     2.6     2.0     70  0
 """
     path_hills = tmp_path / "hills"
     with open(path_hills, "w") as f:
@@ -248,17 +245,16 @@ def test_plumed_evaluate(context, dataset, tmp_path):
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: DISTANCE ATOMS=1,2 NOPBC
 METAD ARG=CV PACE=1 SIGMA=3 HEIGHT=342 FILE={}
-""".format(
-        path_hills
-    )
+""".format(path_hills)
     calculator = PlumedCalculator(plumed_input, path_hills)
-    for _i in range(10):
+
+    for _ in range(30):
         calculator.calculate(atoms)
-        energy = calculator.results["energy"]
-        # energy, _, _ = evaluate_plumed(atoms, plumed_input)
+    energy = calculator.results["energy"]
+
     sigma = 2 * np.ones(2)
-    height = np.array([7, 7]) * (kJ / mol)  # unit consistency
-    center = np.array([2.1, 2.2])
+    height = np.array([70, 70]) * (kJ / mol)  # unit consistency
+    center = np.array([2.5, 2.6])
     energy_ = np.sum(height * np.exp((distance - center) ** 2 / (-2 * sigma**2)))
     assert np.allclose(
         energy,
@@ -270,16 +266,14 @@ def test_plumed_evaluate(context, dataset, tmp_path):
         assert f.read() == hills
 
 
-def test_plumed_hamiltonian(context, dataset, tmp_path):
+def test_plumed_hamiltonian(dataset, tmp_path):
     kappa = 1
     center = 100
     plumed_input = """
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: VOLUME
 RESTRAINT ARG=CV AT={center} KAPPA={kappa}
-""".format(
-        center=center, kappa=kappa / (kJ / mol)
-    )
+""".format(center=center, kappa=kappa / (kJ / mol))
     hamiltonian = PlumedHamiltonian(plumed_input)
     evaluated = hamiltonian.evaluate(dataset).geometries().result()
     for geometry in evaluated:
@@ -308,9 +302,7 @@ def test_plumed_hamiltonian(context, dataset, tmp_path):
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: DISTANCE ATOMS=1,2 NOPBC
 METAD ARG=CV PACE=1 SIGMA=3 HEIGHT=342 FILE={}
-""".format(
-        data_future.filepath
-    )
+""".format(data_future.filepath)
     hamiltonian = PlumedHamiltonian(plumed_input, data_future)
     data = hamiltonian.evaluate(dataset)
     assert np.all(data.get("energy")[0].result() > 0)
@@ -322,7 +314,7 @@ def test_plumed_hamiltonian(context, dataset, tmp_path):
         assert data[i].result().energy > 0
 
 
-def test_json_dump(context):
+def test_json_dump():
     data = {
         "a": np.ones((3, 3, 3, 2)),
         "b": [1, 2, 3],
@@ -349,7 +341,7 @@ def test_json_dump(context):
     assert type(data_["c"]) is list
 
 
-def test_serialization(context, dataset, tmp_path, mace_model):
+def test_serialization(dataset, tmp_path, mace_model):
     hamiltonian = EinsteinCrystal(dataset[0], force_constant=1)
     evaluated = hamiltonian.evaluate(dataset[:3])
 
@@ -389,9 +381,7 @@ def test_serialization(context, dataset, tmp_path, mace_model):
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: DISTANCE ATOMS=1,2 NOPBC
 METAD ARG=CV PACE=1 SIGMA=3 HEIGHT=342 FILE={}
-""".format(
-        data_future.filepath
-    )
+""".format(data_future.filepath)
     hamiltonian = PlumedHamiltonian(plumed_input, data_future)
     evaluated = hamiltonian.evaluate(dataset[:3])
 
@@ -466,9 +456,7 @@ def test_hamiltonian_serialize(dataset):
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: VOLUME
 RESTRAINT ARG=CV AT={center} KAPPA={kappa}
-""".format(
-        center=center, kappa=kappa / (kJ / mol)
-    )
+""".format(center=center, kappa=kappa / (kJ / mol))
     plumed = PlumedHamiltonian(plumed_input)
     data = json.loads(psiflow.serialize(einstein).result())
     assert "EinsteinCrystal" in data