Merge pull request #52 from zhaiwenxi/fix/dpa-adapt-review-followups

Fix dpa-adapt review follow-ups

Author: zhaiwenxi

dpa_adapt/cli.py

@@ -96,6 +96,17 @@ def _maybe_split_list(val: str | Sequence[str] | None) -> list[str] | None:
     ]
 
 
+def _parse_batch_size(val: str) -> str | int:
+    """Parse DeePMD batch-size specs, preserving strings like ``auto:512``."""
+    text = val.strip()
+    if not text:
+        raise argparse.ArgumentTypeError("batch size must not be empty")
+    try:
+        return int(text)
+    except ValueError:
+        return text
+
+
 class _RawTextArgDefaultsHelpFormatter(
     argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
 ):
@@ -304,6 +315,9 @@ def _cmd_data_convert(args: argparse.Namespace) -> int:
     elif result["method"] == "batch_dpdata":
         _LOG.info("Output dirs  : %s", len(result["output_dirs"]))
         _LOG.info("Manifest     : %s", result["manifest"])
+    elif result["method"] == "formula":
+        _LOG.info("Output systems: %s", len(result["output_systems"]))
+        _LOG.info("Wrote deepmd/npy → %s", result["output_dir"])
     else:
         _LOG.info("Wrote deepmd/npy → %s", result["output_dir"])
     return 0
@@ -485,7 +499,7 @@ def get_parser() -> argparse.ArgumentParser:
     parser_fit.add_argument("--max-steps", type=int, default=100_000)
     parser_fit.add_argument("--learning-rate", type=float, default=1e-3)
     parser_fit.add_argument("--stop-lr", type=float, default=1e-5)
-    parser_fit.add_argument("--batch-size", default="auto:512")
+    parser_fit.add_argument("--batch-size", type=_parse_batch_size, default="auto:512")
     parser_fit.add_argument("--seed", type=int, default=42)
     parser_fit.add_argument("--output-dir", default="./dpa_output")
     parser_fit.add_argument("--save-freq", type=int, default=10_000)
@@ -523,11 +537,14 @@ def get_parser() -> argparse.ArgumentParser:
         help="(mft) Downstream head type.",
     )
     parser_fit.add_argument(
-        "--aux-batch-size", default=None, help="(mft) Batch size for aux branch."
+        "--aux-batch-size",
+        type=_parse_batch_size,
+        default=None,
+        help="(mft) Batch size for aux branch.",
     )
     parser_fit.add_argument(
         "--downstream-batch-size",
-        type=int,
+        type=_parse_batch_size,
         default=None,
         help="(mft) Batch size for downstream.",
     )

dpa_adapt/data/convert.py

@@ -191,7 +191,11 @@ def convert(
         )
         if verbose:
             _LOG.info("Formula conversion: %s systems written.", len(out))
-        return {"method": "formula", "output_systems": out}
+        return {
+            "method": "formula",
+            "output_dir": str(Path(output_dir).resolve()),
+            "output_systems": out,
+        }
 
     # --- structure glob → batch dpdata ---
     input_str = str(input_path)
@@ -681,4 +685,6 @@ def attach_labels(
         )
 
     for sys_dir, sub_vals in zip(sys_dirs, values_arr, strict=True):
+        if np.isscalar(sub_vals):
+            sub_vals = np.asarray([sub_vals])
         _attach_single(sys_dir, head, sub_vals)

dpa_adapt/data/desc_cache.py

@@ -89,8 +89,8 @@ def _system_fingerprint(system: dpdata.System) -> str:
 
 
 def _data_fingerprint(systems: list) -> str:
-    """Aggregate fingerprint for a list of systems (order-independent)."""
-    fps = sorted(_system_fingerprint(s) for s in systems)
+    """Aggregate fingerprint for a list of systems in request order."""
+    fps = [_system_fingerprint(s) for s in systems]
     h = hashlib.sha1()
     for fp in fps:
         h.update(fp.encode())

dpa_adapt/finetuner.py

@@ -841,10 +841,10 @@ class DPAFineTuner:
         Auto-detected from the checkpoint if not provided.
     downstream_task_type : str
         (MFT only) Task type of the downstream head (``"property"`` etc.).
-    aux_batch_size : str or None
+    aux_batch_size : str or int or None
         (MFT only) Batch-size spec for the auxiliary head.
-    downstream_batch_size : int or None
-        (MFT only) Batch size for the downstream head.
+    downstream_batch_size : str or int or None
+        (MFT only) Batch-size spec for the downstream head.
     """
 
     _VALID_POOLING: ClassVar[set[str]] = {"mean", "sum", "mean+std", "mean+std+max+min"}
@@ -886,8 +886,8 @@ def __init__(
         aux_prob: float = 0.5,
         type_map: list[str] | None = None,
         downstream_task_type: str = "property",
-        aux_batch_size: str | None = None,
-        downstream_batch_size: int | None = None,
+        aux_batch_size: str | int | None = None,
+        downstream_batch_size: str | int | None = None,
     ) -> None:
         if pooling not in self._VALID_POOLING:
             raise ValueError(
@@ -1041,7 +1041,7 @@ def _extract_features_cached(self, systems: list[dpdata.System]) -> np.ndarray:
         except Exception:
             # Cache read failed (e.g. corrupted file, permissions) —
             # fall through and recompute features from scratch.
-            pass
+            _LOG.debug("Descriptor cache read failed, recomputing.", exc_info=True)
 
         features = self._extract_features(systems)
         try:
@@ -1050,7 +1050,7 @@ def _extract_features_cached(self, systems: list[dpdata.System]) -> np.ndarray:
         except Exception:
             # Cache write is best-effort — silently skip on permission errors
             # or disk-full conditions; the features are already in memory.
-            pass
+            _LOG.debug("Descriptor cache write failed.", exc_info=True)
         return features
 
     def _extract_features(self, systems: list[dpdata.System]) -> np.ndarray:
@@ -1099,9 +1099,10 @@ def _resolve_type_maps(self, train_data: str | list[str]) -> list[str]:
 
         try:
             elements = read_data_type_map_union(systems)
-            validate_type_map_subset(elements, tm, label="train data")
         except ValueError:
             pass  # no atom_names — deepmd uses raw atom indices
+        else:
+            validate_type_map_subset(elements, tm, label="train data")
 
         return tm
 
@@ -1372,6 +1373,10 @@ def fit(
                     "strategy='mft' requires aux_data. "
                     "Provide auxiliary system directories for the force-field head."
                 )
+            if type_map is not None:
+                self.type_map = type_map
+                if self._mft is not None:
+                    self._mft.type_map = type_map
             return self._fit_mft(train_data, aux_data, valid_data)
 
         # ---- single-task training paradigms ----
@@ -1719,6 +1724,5 @@ def freeze(self, output_path: str = "frozen_model.pth") -> str:
         import torch
 
         torch.save(bundle, output_path)
-        _LOG = logging.getLogger("dpa_adapt")
         _LOG.info("Frozen model saved to: %s", output_path)
         return output_path

dpa_adapt/mft.py

@@ -122,8 +122,8 @@ def __init__(
         warmup_steps: int = 0,
         max_steps: int = 50000,
         batch_size: str | int = "auto:32",
-        aux_batch_size: str | None = None,
-        downstream_batch_size: int | None = None,
+        aux_batch_size: str | int | None = None,
+        downstream_batch_size: str | int | None = None,
         seed: int = 42,
         fparam_dim: int = 0,
         output_dir: str = "./mft_output",

dpa_adapt/predictor.py

@@ -53,6 +53,38 @@ def _is_mlp(est: Any) -> bool:
     return isinstance(est, MLPRegressor)
 
 
+def _rf_tree_predictions(est: Any, features: np.ndarray) -> np.ndarray:
+    """Return RF per-tree predictions with shape ``(n_trees, n_frames, dim)``."""
+    from sklearn.ensemble import (
+        RandomForestRegressor,
+    )
+    from sklearn.multioutput import (
+        MultiOutputRegressor,
+    )
+
+    if isinstance(est, MultiOutputRegressor):
+        per_output = []
+        for rf in est.estimators_:
+            if not isinstance(rf, RandomForestRegressor):
+                raise TypeError(
+                    "Expected MultiOutputRegressor(RandomForestRegressor), "
+                    f"got wrapped estimator {type(rf).__name__!r}."
+                )
+            per_output.append(
+                np.array([tree.predict(features) for tree in rf.estimators_])
+            )
+        return np.stack(per_output, axis=-1)
+
+    if isinstance(est, RandomForestRegressor):
+        tree_preds = np.array([tree.predict(features) for tree in est.estimators_])
+        return tree_preds.reshape(len(est.estimators_), -1, 1)
+
+    raise TypeError(
+        "RF uncertainty requires RandomForestRegressor or "
+        f"MultiOutputRegressor(RandomForestRegressor), got {type(est).__name__!r}."
+    )
+
+
 class DPAPredictor:
     """
     Read-only inference wrapper for a frozen DPA+sklearn bundle.
@@ -278,12 +310,8 @@ def _predict_with_uncertainty(self, features: np.ndarray) -> DotDict:
             for _, step in self._predictor.steps[:-1]:
                 X_t = step.transform(X_t)
             rf = self._predictor.steps[-1][1]
-            tree_preds = np.array([t.predict(X_t) for t in rf.estimators_])
-            tree_preds = tree_preds.reshape(
-                len(rf.estimators_),
-                -1,
-                self._task_dim,
-            )
+            tree_preds = _rf_tree_predictions(rf, X_t)
+            tree_preds = tree_preds.reshape(tree_preds.shape[0], -1, self._task_dim)
             return DotDict(
                 {
                     "predictions": np.mean(tree_preds, axis=0),

source/tests/dpa_adapt/test_cache.py

@@ -67,6 +67,13 @@ def test_different_data_different_fp(self, tmp_path):
         fp2 = _data_fingerprint([s2])
         assert fp1 != fp2
 
+    def test_system_order_changes_fp(self, tmp_path):
+        s1 = _make_system(tmp_path, "s1", nframes=3)
+        s2 = _make_system(tmp_path, "s2", nframes=5)
+        fp1 = _data_fingerprint([s1, s2])
+        fp2 = _data_fingerprint([s2, s1])
+        assert fp1 != fp2
+
 
 class TestCacheKey:
     def test_same_inputs_same_key(self, tmp_path):
@@ -111,6 +118,15 @@ def test_different_type_map_different_key(self, tmp_path):
         k2 = _cache_key([s], str(ckpt), None, "mean", type_map=("O", "H"))
         assert k1 != k2
 
+    def test_different_system_order_different_key(self, tmp_path):
+        s1 = _make_system(tmp_path, "s1", nframes=3)
+        s2 = _make_system(tmp_path, "s2", nframes=5)
+        ckpt = tmp_path / "dummy.pt"
+        ckpt.write_text("dummy")
+        k1 = _cache_key([s1, s2], str(ckpt), None, "mean")
+        k2 = _cache_key([s2, s1], str(ckpt), None, "mean")
+        assert k1 != k2
+
 
 class TestCacheDir:
     def test_respects_xdg(self, monkeypatch, tmp_path):

source/tests/dpa_adapt/test_cli_smoke.py

@@ -162,6 +162,106 @@ def test_maybe_split_list_accepts_string_sequences(self):
         assert _maybe_split_list("H,C, O") == ["H", "C", "O"]
         assert _maybe_split_list(None) is None
 
+    def test_batch_size_parser_preserves_deepmd_specs(self):
+        from dpa_adapt.cli import (
+            _parse_batch_size,
+        )
+
+        assert _parse_batch_size("128") == 128
+        assert _parse_batch_size("auto:512") == "auto:512"
+
+    def test_fit_accepts_downstream_auto_batch_size(self):
+        from dpa_adapt.cli import (
+            get_parser,
+        )
+
+        args = get_parser().parse_args(
+            [
+                "fit",
+                "--train-data",
+                "train",
+                "--strategy",
+                "mft",
+                "--downstream-batch-size",
+                "auto:512",
+            ]
+        )
+
+        assert args.downstream_batch_size == "auto:512"
+
+    def test_fit_batch_size_numbers_parse_to_int(self):
+        from dpa_adapt.cli import (
+            get_parser,
+        )
+
+        args = get_parser().parse_args(
+            [
+                "fit",
+                "--train-data",
+                "train",
+                "--batch-size",
+                "64",
+                "--aux-batch-size",
+                "128",
+                "--downstream-batch-size",
+                "256",
+            ]
+        )
+
+        assert args.batch_size == 64
+        assert args.aux_batch_size == 128
+        assert args.downstream_batch_size == 256
+
+
+class TestDpaDataConvertDispatch:
+    """Verify data convert handles method-specific return payloads."""
+
+    def test_formula_result_exits_cleanly(self, monkeypatch, tmp_path):
+        from argparse import (
+            Namespace,
+        )
+
+        import dpa_adapt
+        from dpa_adapt.cli import (
+            _cmd_data_convert,
+        )
+
+        out = tmp_path / "npy"
+
+        def _fake_convert(**kwargs):
+            return {
+                "method": "formula",
+                "output_dir": str(out),
+                "output_systems": [str(out / "sys_0000")],
+            }
+
+        monkeypatch.setattr(dpa_adapt, "convert", _fake_convert)
+
+        args = Namespace(
+            input=str(tmp_path / "formula.csv"),
+            output=str(out),
+            fmt="formula",
+            type_map=None,
+            property_name=None,
+            property_col="energy",
+            train_ratio=0.9,
+            smiles_col="SMILES",
+            mol_dir=None,
+            mol_template="id{row}.mol",
+            split_seed=None,
+            conformer_seed=None,
+            poscar=str(tmp_path / "POSCAR"),
+            formula_col="formula",
+            base_element=None,
+            sets=1,
+            seed=42,
+            overwrite=False,
+            validate=True,
+            strict=False,
+        )
+
+        assert _cmd_data_convert(args) == 0
+
 
 class TestInitAllExports:
     """Verify __all__ covers the key public names."""

source/tests/dpa_adapt/test_convert.py

@@ -364,6 +364,7 @@ def _fake_formula_to_npy(**kwargs):
         )
 
         assert result["method"] == "formula"
+        assert result["output_dir"] == str(out.resolve())
         assert result["output_systems"] == [fake_sys_dir]
 
     def test_formula_fmt_base_element_passed_through(self, tmp_path, monkeypatch):

source/tests/dpa_adapt/test_loader.py

@@ -316,6 +316,17 @@ def test_multi_system_all_written(self, tmp_path):
             written = np.load(parent / f"sys_{i:04d}" / "set.000" / "bandgap.npy")
             np.testing.assert_array_equal(written, values[i])
 
+    def test_multi_system_1d_values_written_as_one_frame_labels(self, tmp_path):
+        parent = tmp_path / "multi"
+        parent.mkdir()
+        for i in range(3):
+            _make_system_path(parent, name=f"sys_{i:04d}", n_frames=1)
+        values = np.array([1.0, 3.0, 5.0])
+        attach_labels(parent, head="bandgap", values=values)
+        for i in range(3):
+            written = np.load(parent / f"sys_{i:04d}" / "set.000" / "bandgap.npy")
+            np.testing.assert_array_equal(written, [values[i]])
+
     def test_multi_system_values_mismatch_raises(self, tmp_path):
         parent = tmp_path / "multi"
         parent.mkdir()