Merge pull request #30 from molmod/function

Rewrite Hamiltonian API + bunch of other fixes

Dockerfile

@@ -16,6 +16,7 @@ RUN apt-get update && apt-get install -y \
     libcurl4-openssl-dev \
     libgsl-dev \
     perl \
+    fftw3 \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
@@ -57,8 +58,8 @@ ARG PARSL_VERSION
 ARG GPU_LIBRARY
 RUN /bin/bash -c -o pipefail \
     "source /opt/venv/bin/activate && \
-     pip install --no-cache-dir wandb plotly plumed 'numpy<2.0.0' && \
-     pip install --no-cache-dir git+https://github.com/i-pi/i-pi.git@66eba29 && \
+     pip install --no-cache-dir pyfftw colorcet wandb pandas plotly plumed 'numpy<2.0.0' && \
+     pip install --no-cache-dir git+https://github.com/i-pi/i-pi.git@v3.0.0-beta4 && \
      pip install --no-cache-dir torch==2.1 --index-url https://download.pytorch.org/whl/${GPU_LIBRARY} && \
      pip install --no-cache-dir git+https://github.com/acesuit/[email protected]"
 ARG DATE

build_containers.sh

@@ -39,9 +39,9 @@ while [[ $# -gt 0 ]]; do
 	esac
 done
 
-PSIFLOW_VERSION="main"
-PARSL_VERSION="benc-sander-paper"
-CCTOOLS_VERSION=7.11.0
+PSIFLOW_VERSION="function"
+PARSL_VERSION="2024.07.08"
+CCTOOLS_VERSION=7.11.1
 PLUMED_VERSION=2.9.0
 GPU_LIBRARIES=("rocm5.6" "cu118")
 

configs/threadpool.yaml

@@ -1,4 +1,5 @@
 ---
+parsl_log_level: DEBUG
 retries: 0
 ModelEvaluation:
   gpu: false

examples/alanine_replica_exchange.py

@@ -2,9 +2,8 @@
 
 import psiflow
 from psiflow.geometry import Geometry
-from psiflow.hamiltonians import get_mace_mp0
-from psiflow.sampling import Walker, sample
-from psiflow.sampling.walker import replica_exchange
+from psiflow.hamiltonians import MACEHamiltonian
+from psiflow.sampling import Walker, replica_exchange, sample
 
 
 def compute_dihedrals(positions):
@@ -70,7 +69,7 @@ def main():
 H       13.27142638      10.63298597      -1.06170510
 """
     )
-    mace = get_mace_mp0()
+    mace = MACEHamiltonian.mace_mp0()
 
     walkers = []
     for temperature in [150, 200, 250, 300, 400, 500, 600, 700, 800, 900, 1000, 1100]:

examples/h2_static_dynamic.py

@@ -2,11 +2,10 @@
 from ase.units import _c, second
 
 import psiflow
+from psiflow.free_energy import compute_frequencies, compute_harmonic
 from psiflow.geometry import Geometry
-from psiflow.hamiltonians import get_mace_mp0
-from psiflow.hamiltonians._harmonic import compute_frequencies
-from psiflow.sampling import Walker, sample
-from psiflow.tools import compute_harmonic, optimize
+from psiflow.hamiltonians import MACEHamiltonian
+from psiflow.sampling import Walker, optimize, sample
 
 
 def frequency_dynamic(start, hamiltonian):
@@ -61,7 +60,7 @@ def main():
         positions=np.array([[0, 0, 0], [0.8, 0, 0]]),
         cell=None,
     )
-    mace = get_mace_mp0()
+    mace = MACEHamiltonian.mace_mp0()
 
     dynamic = frequency_dynamic(geometry, mace)
     static = frequency_static(geometry, mace)

examples/iron_bulk_modulus.py

@@ -3,14 +3,14 @@
 
 import psiflow
 from psiflow.geometry import Geometry
-from psiflow.hamiltonians import get_mace_mp0
+from psiflow.hamiltonians import MACEHamiltonian
 from psiflow.sampling import Walker, sample
 
 
 def main():
     iron = bulk("Fe", "bcc", a=2.8)
     geometry = Geometry.from_atoms(make_supercell(iron, 3 * np.eye(3)))
-    mace = get_mace_mp0()
+    mace = MACEHamiltonian.mace_mp0()
 
     pressures = (-10 + np.arange(5) * 5) * 1e3  # in MPa
     walkers = [Walker(geometry, mace, temperature=300, pressure=p) for p in pressures]

examples/iron_harmonic_fcc_bcc.py

@@ -4,10 +4,10 @@
 
 import psiflow
 from psiflow.data import Dataset
-from psiflow.free_energy import Integration
+from psiflow.free_energy import Integration, compute_harmonic
 from psiflow.geometry import Geometry
-from psiflow.hamiltonians import Harmonic, get_mace_mp0
-from psiflow.tools import compute_harmonic, optimize
+from psiflow.hamiltonians import Harmonic, MACEHamiltonian
+from psiflow.sampling import optimize
 
 
 def main():
@@ -23,7 +23,7 @@ def main():
     theoretical = {name: None for name in geometries}
     simulated = {name: None for name in geometries}
 
-    mace = get_mace_mp0("small")
+    mace = MACEHamiltonian.mace_mp0("small")
     scaling = 0.9
     temperature = 800
     beta = 1 / (kB * temperature)

examples/submit.sh

@@ -2,7 +2,6 @@
 
 # List of filenames
 files=(
-	"formic_acid_transition.py"
 	"h2_static_dynamic.py"
 	"iron_bulk_modulus.py"
 	"iron_harmonic_fcc_bcc.py"

examples/water_path_integral_md.py

@@ -3,7 +3,7 @@
 
 import psiflow
 from psiflow.data import Dataset
-from psiflow.hamiltonians import get_mace_mp0
+from psiflow.hamiltonians import MACEHamiltonian
 from psiflow.sampling import Walker, sample
 
 
@@ -25,7 +25,7 @@ def get_OH_distances(geometries):
 
 def main():
     geometry = Dataset.load("data/h2o_32.xyz")[0]
-    mace = get_mace_mp0()
+    mace = MACEHamiltonian.mace_mp0()
 
     trajectories = []
     for i in range(6):

examples/water_train_validate.py

@@ -24,7 +24,7 @@ def main():
     target_e = data.get("per_atom_energy")
     target_f = data.get("forces")
 
-    data_predicted = hamiltonian.evaluate(data)
+    data_predicted = data.evaluate(hamiltonian)
     predict_e = data_predicted.get("per_atom_energy")
     predict_f = data_predicted.get("forces")
 

psiflow/__init__.py

@@ -1,3 +1,7 @@
+from pathlib import Path
+
+import typeguard
+
 from .config import setup_slurm_config  # noqa: F401
 from .execution import ExecutionContextLoader
 from .serialization import (  # noqa: F401
@@ -7,6 +11,25 @@
     serialize,
 )
 
+
+@typeguard.typechecked
+def resolve_and_check(path: Path) -> Path:
+    path = path.resolve()
+    if Path.cwd() in path.parents:
+        pass
+    elif path.exists() and Path.cwd().samefile(path):
+        pass
+    else:
+        raise ValueError(
+            "requested file and/or path at location: {}"
+            "\nwhich is not in the present working directory: {}"
+            "\npsiflow can only load and/or save in its present "
+            "working directory because this is the only directory"
+            " that will get bound into the container.".format(path, Path.cwd())
+        )
+    return path
+
+
 load = ExecutionContextLoader.load
 context = ExecutionContextLoader.context
 wait = ExecutionContextLoader.wait

psiflow/data/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import Computable, Dataset, aggregate_multiple, compute  # noqa: F401
+from .utils import compute_mae, compute_rmse  # noqa: F401