Merge pull request #21 from zirenjin/master

feat(dpa_adapt): auto-read fparam.npy for all strategies, remove condition={temperature...} for frozen_sklearn

Author: zhaiwenxi

doc/dpa_adapt/README.md

@@ -149,22 +149,19 @@ For the full option list and supported dpdata formats, see
 
 ### Context features (fparam)
 
-fparam lets you condition the model on system-level context such as temperature, pressure, or experimental conditions.
-
-**frozen_sklearn** — pass a dict of numpy arrays at fit and predict time:
+fparam lets you condition the model on system-level context such as temperature, humidity, pressure, or any per-frame scalar.  All strategies use the same interface: place `fparam.npy` of shape `(n_frames, fparam_dim)` in each `set.*/` directory alongside `coord.npy` and declare the dimension at construction.
 
 ```python
-model.fit(train_data, conditions={"temperature": T_train})
-model.predict(test_data, conditions={"temperature": T_test})
-# ConditionManager standardizes and concatenates values to the descriptor
+# works identically for frozen_sklearn, frozen_head, finetune, and mft
+model = DPAFineTuner(strategy="frozen_sklearn", fparam_dim=2)
+model.fit(train_data="data/train", target_key="property")
+# fparam.npy is read automatically — no conditions= dict needed
 ```
 
-**frozen_head / finetune / mft** — place `fparam.npy` of shape `(nframes, fparam_dim)` in each `set.*/` directory alongside `coord.npy`, then declare the dimension at construction:
-
-```python
-model = DPAFineTuner(strategy="finetune", fparam_dim=2)
-model.fit(train_data)  # reads fparam.npy automatically
-```
+| Strategy | How fparam is used |
+|---|---|
+| `frozen_sklearn` | columns are standardized via `ConditionManager` and concatenated to the descriptor |
+| `frozen_head` / `finetune` / `mft` | passed into the fitting net as `numb_fparam` |
 
 ## Inference and uncertainty
 
@@ -197,15 +194,25 @@ Uncertainty estimates can drive active learning (query most uncertain candidates
 
 ## Cross-validation
 
-Formula-grouped splitting prevents same-composition leakage between folds:
+Formula-grouped splitting prevents same-composition leakage between folds.
+`group_by` accepts `"formula"` (uses each system's directory name as the group
+key — requires directories named by formula, e.g. `H2O/`, `CH4/`) or a list
+of labels the same length as `systems`:
 
 ```python
 from dpa_adapt import cross_validate, train_test_split, load_dataset
 
 systems = load_dataset("/data/root", label_key="energy")
+
+# Case 1: directory names are formulas (e.g. data/H2O/, data/CH4/)
 train, valid, test = train_test_split(systems, group_by="formula", seed=42)
 
-result = cross_validate(model, systems, label_key="energy", cv=5, group_by="formula")
+# Case 2: directory names are not formulas (e.g. QM9's sys_0000, sys_0001, …)
+formulas = ["H2O", "H2O", "CH4", "CH4", ...]  # one label per system
+train, valid, test = train_test_split(systems, group_by=formulas, seed=42)
+
+# Cross-validate (same group_by options apply)
+result = cross_validate(model, systems, label_key="energy", cv=5, group_by=formulas)
 # → {"aggregate": {"mae_mean": ..., "rmse_std": ...}, ...}
 ```
 

dpa_adapt/config/manager.py

@@ -33,7 +33,7 @@ def _build_property_fitting_net(t) -> dict:
         }
     )
     if getattr(t, "fparam_dim", 0) > 0:
-        fn["fparam_dim"] = t.fparam_dim
+        fn["numb_fparam"] = t.fparam_dim
     return fn
 
 

dpa_adapt/finetuner.py

@@ -112,6 +112,33 @@ def _load_labels(
     return np.column_stack(columns)
 
 
+def _read_fparam_from_systems(
+    systems: list[dpdata.System],
+) -> dict[str, np.ndarray] | None:
+    """Auto-read fparam.npy from each system's ``set.*/`` directories.
+
+    Returns a dict mapping ``"fparam_0"``, ``"fparam_1"``, ... to 1-D
+    arrays of length ``n_frames_total``, suitable for passing as
+    ``conditions=`` to :meth:`ConditionManager.fit_transform`.
+
+    Returns ``None`` when no system has a ``set.*/fparam.npy`` file.
+    """
+    all_fparams = []
+    for system in systems:
+        source = _get_source(system)
+        if source is None:
+            continue
+        fps = sorted(Path(source).glob("set.*/fparam.npy"))
+        if not fps:
+            continue
+        arrs = [np.load(str(fp)) for fp in fps]
+        all_fparams.append(np.concatenate(arrs, axis=0))
+    if not all_fparams:
+        return None
+    combined = np.concatenate(all_fparams, axis=0)  # (n_frames, fparam_dim)
+    return {f"fparam_{i}": combined[:, i] for i in range(combined.shape[1])}
+
+
 def _read_data_type_map(system) -> list[str]:
     """Read element symbols from a dpdata System's ``atom_names``.
 
@@ -560,10 +587,12 @@ class DPAFineTuner:
     loss_function : str
         ``"mse"`` or ``"smooth_mae"`` (training paradigms).
     fparam_dim : int
-        (frozen_head / finetune / mft only) Dimensionality of per-frame
-        condition inputs (e.g. temperature, pressure). Requires
-        set.*/fparam.npy of shape (n_frames, fparam_dim) in every
-        training system. Default 0 (disabled).
+        Dimension of per-frame context features (e.g. temperature,
+        humidity).  When > 0, ``set.*/fparam.npy`` of shape
+        ``(n_frames, fparam_dim)`` is read automatically for all
+        strategies.  For ``frozen_sklearn``, fparam columns are
+        standardized and concatenated to the descriptor via
+        ``ConditionManager``.  Default 0 (disabled).
     output_dir : str
         Directory for ``input.json``, checkpoints, and logs.
     save_freq, disp_freq : int
@@ -854,7 +883,6 @@ def fit(
         target_key=None,
         labels=None,
         fmt=None,
-        conditions=None,
         aux_data=None,
     ):
         """Train the model.
@@ -879,15 +907,13 @@ def fit(
             (frozen_sklearn) Pre-computed labels.
         fmt : str, optional
             Reserved for future format support.
-        conditions : dict[str, np.ndarray], optional
-            (frozen_sklearn) Named condition arrays.
         aux_data : str | list[str], optional
             (mft only) Auxiliary training system directories.  Required when
             ``strategy='mft'``; must be absent otherwise.
         """
         if self.strategy == "frozen_sklearn":
             return self._fit_sklearn(
-                train_data, type_map, target_key, labels, fmt, conditions
+                train_data, type_map, target_key, labels, fmt
             )
 
         if self.strategy == "mft":
@@ -951,7 +977,6 @@ def _fit_sklearn(
         target_key=None,
         labels=None,
         fmt=None,
-        conditions=None,
     ):
         """Fit the frozen-sklearn pipeline (delegates to ``_FrozenSklearnPipeline``).
 
@@ -978,10 +1003,12 @@ def _fit_sklearn(
         features = self._extract_features_cached(systems)
 
         self._condition_manager = None
-        if conditions is not None:
-            self._condition_manager = ConditionManager()
-            X_cond = self._condition_manager.fit_transform(conditions)
-            features = np.concatenate([features, X_cond], axis=1)
+        if self.fparam_dim > 0:
+            conditions = _read_fparam_from_systems(systems)
+            if conditions is not None:
+                self._condition_manager = ConditionManager()
+                X_cond = self._condition_manager.fit_transform(conditions)
+                features = np.concatenate([features, X_cond], axis=1)
 
         if labels is not None:
             y = np.asarray(labels)
@@ -1019,19 +1046,19 @@ def _fit_sklearn(
         p._condition_manager = self._condition_manager
         p._fitted = True
 
-    def predict(self, data, fmt=None, conditions=None) -> DotDict:
+    def predict(self, data, fmt=None) -> DotDict:
         """
         Extract features and run the fitted sklearn predictor.
 
+        fparam is automatically read from ``set.*/fparam.npy`` when the
+        model was fit with ``fparam_dim > 0``.
+
         Parameters
         ----------
         data : str | list[str]
             Path(s) to deepmd/npy system directories.
         fmt : str, optional
             Reserved for future format support.
-        conditions : dict[str, np.ndarray], optional
-            Named condition arrays.  Required when the model was fit with
-            conditions; must be absent otherwise.
 
         Returns
         -------
@@ -1047,20 +1074,20 @@ def predict(self, data, fmt=None, conditions=None) -> DotDict:
         features = self._extract_features(systems)
 
         if self._condition_manager is not None:
+            conditions = _read_fparam_from_systems(systems)
             if conditions is None:
                 raise DPAConditionError(
-                    "This model was fit with conditions. Pass conditions= to predict()."
+                    "This model was fit with fparam but set.*/fparam.npy "
+                    "was not found in the test data."
                 )
             X_cond = self._condition_manager.transform(conditions)
             features = np.concatenate([features, X_cond], axis=1)
-        elif conditions is not None:
-            raise DPAConditionError("This model was fit without conditions.")
 
         raw = self.predictor.predict(features)
         predictions = np.asarray(raw).reshape(-1, self._task_dim)
         return DotDict({"predictions": predictions})
 
-    def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
+    def evaluate(self, data, fmt=None) -> DotDict:
         """
         Predict on ``data`` and compute evaluation metrics against stored labels.
 
@@ -1070,9 +1097,6 @@ def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
             Path(s) to deepmd/npy system directories with label files.
         fmt : str, optional
             Reserved for future format support.
-        conditions : dict[str, np.ndarray], optional
-            Named condition arrays.  Required when the model was fit with
-            conditions; must be absent otherwise.
 
         Returns
         -------
@@ -1081,7 +1105,7 @@ def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
             predictions   : np.ndarray, shape (n_frames, task_dim)
             labels        : np.ndarray, shape (n_frames, task_dim)
         """
-        result = self.predict(data, fmt=fmt, conditions=conditions)
+        result = self.predict(data, fmt=fmt)
         predictions = result.predictions
 
         systems = load_data(data, fmt=fmt)

dpa_adapt/predictor.py

@@ -120,7 +120,7 @@ def __init__(self, model_path: str, n_committee: int = 1):
             pooling=self._pooling,
         )
 
-    def fit(self, data, target_key=None, labels=None, fmt=None, conditions=None):
+    def fit(self, data, target_key=None, labels=None, fmt=None):
         """Train committee members for uncertainty estimation.
 
         Only valid when *n_committee* > 1.  Clones the frozen sklearn
@@ -140,6 +140,7 @@ def fit(self, data, target_key=None, labels=None, fmt=None, conditions=None):
 
         from dpa_adapt.finetuner import (
             _load_labels,
+            _read_fparam_from_systems,
         )
 
         if target_key is not None and labels is not None:
@@ -154,14 +155,14 @@ def fit(self, data, target_key=None, labels=None, fmt=None, conditions=None):
         features = self._extractor._extract_features(systems)
 
         if self._condition_manager is not None:
+            conditions = _read_fparam_from_systems(systems)
             if conditions is None:
                 raise DPAConditionError(
-                    "This model was fit with conditions. Pass conditions= to fit()."
+                    "This model was fit with fparam but set.*/fparam.npy "
+                    "was not found in the data."
                 )
             X_cond = self._condition_manager.transform(conditions)
             features = np.concatenate([features, X_cond], axis=1)
-        elif conditions is not None:
-            raise DPAConditionError("This model was fit without conditions.")
 
         if labels is not None:
             y = np.asarray(labels)
@@ -184,44 +185,45 @@ def fit(self, data, target_key=None, labels=None, fmt=None, conditions=None):
         preds = preds.reshape(self.n_committee, -1, self._task_dim)
         self.uncertainty_threshold_ = float(np.percentile(np.std(preds, axis=0), 95))
 
-    def _extract_and_condition(self, data, fmt, conditions):
-        """Shared feature extraction + condition concatenation."""
+    def _extract_and_condition(self, data, fmt):
+        """Shared feature extraction + fparam auto-read."""
+        from dpa_adapt.finetuner import (
+            _read_fparam_from_systems,
+        )
+
         systems = load_data(data, fmt=fmt)
-        # Load the model first so the checkpoint type_map is available, then
-        # validate before extracting features (extraction relies on the data
-        # type_map being a subset of the checkpoint's).
         if self._extractor._model is None:
             self._extractor._model = self._extractor._load_descriptor_model()
         self._extractor._validate_type_map(self._type_map, systems)
         features = self._extractor._extract_features(systems)
 
         if self._condition_manager is not None:
+            conditions = _read_fparam_from_systems(systems)
             if conditions is None:
                 raise DPAConditionError(
-                    "This model was fit with conditions. Pass conditions= to predict()."
+                    "This model was fit with fparam but set.*/fparam.npy "
+                    "was not found in the data."
                 )
             X_cond = self._condition_manager.transform(conditions)
             features = np.concatenate([features, X_cond], axis=1)
-        elif conditions is not None:
-            raise DPAConditionError("This model was fit without conditions.")
 
         return features
 
     def predict(
-        self, data, fmt=None, conditions=None, return_uncertainty=False
+        self, data, fmt=None, return_uncertainty=False
     ) -> DotDict:
         """
         Run inference on ``data``.
 
+        fparam is automatically read from ``set.*/fparam.npy`` when the
+        model was fit with fparam.
+
         Parameters
         ----------
         data : str | list[str]
             Path(s) to deepmd/npy system directories.
         fmt : str, optional
             Reserved for future format support.
-        conditions : dict[str, np.ndarray], optional
-            Named condition arrays.  Required when the model was fit with
-            conditions; must be absent otherwise.
         return_uncertainty : bool
             When True, include ``"uncertainty"`` (per-sample std) in the
             result.  Behaviour depends on estimator type and committee
@@ -233,7 +235,7 @@ def predict(
             ``predictions`` : np.ndarray, shape (n_frames, task_dim)
             ``uncertainty`` : np.ndarray, shape (n_frames, task_dim)  (if requested)
         """
-        features = self._extract_and_condition(data, fmt, conditions)
+        features = self._extract_and_condition(data, fmt)
 
         if return_uncertainty:
             return self._predict_with_uncertainty(features)
@@ -291,7 +293,7 @@ def _predict_with_uncertainty(self, features):
             f"with n_committee={self.n_committee}."
         )
 
-    def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
+    def evaluate(self, data, fmt=None) -> DotDict:
         """
         Predict on ``data`` and compute evaluation metrics against stored labels.
 
@@ -301,9 +303,6 @@ def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
             Path(s) to deepmd/npy system directories with label files.
         fmt : str, optional
             Reserved for future format support.
-        conditions : dict[str, np.ndarray], optional
-            Named condition arrays.  Required when the model was fit with
-            conditions; must be absent otherwise.
 
         Returns
         -------
@@ -319,7 +318,7 @@ def evaluate(self, data, fmt=None, conditions=None) -> DotDict:
             _load_labels,
         )
 
-        result = self.predict(data, fmt=fmt, conditions=conditions)
+        result = self.predict(data, fmt=fmt)
         predictions = result.predictions
 
         systems = load_data(data, fmt=fmt)

dpa_adapt/trainer.py

@@ -327,7 +327,7 @@ def _build_fitting_net(self) -> dict:
         # property head at [128, 240], so there is no [159, 240] checkpoint
         # head to size-match against. An explicit user value still wins.
         if self.fparam_dim > 0:
-            fn["fparam_dim"] = self.fparam_dim
+            fn["numb_fparam"] = self.fparam_dim
         if self.fitting_net_params:
             fn.update(self.fitting_net_params)
         return fn