Mace Refactor

.github/workflows/run_pytest.yaml

@@ -36,7 +36,7 @@ jobs:
           pip install --no-cache-dir git+https://github.com/i-pi/i-pi.git@v3.0.0-beta4
           pip install torch==2.5.1
           pip install git+https://github.com/acesuit/MACE.git@v0.3.5
-          apptainer exec oras://ghcr.io/molmod/cp2k:2024.1 ls
+          apptainer exec docker://cp2k/cp2k:2025.2_mpich_x86_64_psmp ls
           apptainer exec oras://ghcr.io/molmod/gpaw:24.1 ls
       - name: Checkout specific commit
         uses: actions/checkout@v4

install_local.sh

@@ -0,0 +1,32 @@
+# Create a psiflow installation including ModelTraining and ModelEvaluation dependencies
+
+ENV_NAME="psiflow-dev"
+micromamba env create -n $ENV_NAME python=3.14 -c conda-forge
+micromamba activate $ENV_NAME
+
+# install workqueue
+micromamba install ndcctools -c conda-forge
+
+# install psiflow
+pip install -e psiflow[dev]
+
+# install PyTorch and MACE
+pip install torch --index-url https://download.pytorch.org/whl/cu128
+pip install mace-torch
+pip install cuequivariance cuequivariance-torch cuequivariance-ops-torch-cu12
+
+# install basic PLUMED and python API
+micromamba install plumed -c conda-forge
+pip install plumed
+
+# make PLUMED visible to the py-plumed interface
+PLUMED_KERNEL="$CONDA_PREFIX/lib/libplumedKernel.so"
+echo "{\"env_vars\": {\"PLUMED_KERNEL\": \"$PLUMED_KERNEL\"}}" >> "$CONDA_PREFIX/conda-meta/state"
+
+# install simple-dftd3
+micromamba install simple-dftd3 dftd3-python -c conda-forge
+
+# install cp2k-input-tools 
+# (its dependencies are a right mess, so we fix some and pip will complain)
+pip install cp2k-input-tools
+pip install --force-reinstall Pint

psiflow/data/dataset.py

@@ -285,7 +285,7 @@ def evaluate(
         Returns:
             Dataset: A new Dataset with evaluation results.
         """
-        # TODO: WIP
+        # TODO: remove this functionality?
         from psiflow.hamiltonians import Hamiltonian
 
         if not isinstance(computable, Hamiltonian):

psiflow/execution.py

@@ -246,7 +246,7 @@ def __init__(
             )
 
         if self.executor_type == "workqueue":
-            # WQ-specific checks
+            # WQ-specific checks  TODO: what about multinode?
             ensure(
                 self.kwargs["gpus_per_task"] <= resources["gpus"],
                 self.kwargs["cores_per_task"] <= resources["cores"],
@@ -541,19 +541,10 @@ def __init__(self, **kwargs) -> None:
                 "ModelTraining is configured for CPU operation. Is this what you want?"
             )
 
-        # if self.multigpu:
-        #     # TODO: why? Think this might be a multinode thing - which I do not care about
-        #     message = (
-        #         "the max_training_time keyword does not work "
-        #         "in combination with multi-gpu training. Adjust "
-        #         "the maximum number of epochs to control the "
-        #         "duration of training"
-        #     )
-        #     assert self.max_runtime is None, message
-
-    def train_command(self, initialize: bool = False):
-        command = "psiflow-mace-train"
-        return self.wrap_in_timeout(command)
+    @property
+    def multi_gpu(self) -> bool:
+        # only for WQ
+        return (self.spec or {}).get("gpus", 0) > 1
 
     def wq_resources(self, *args, **kwargs) -> dict:
         if self.spec is None: