Merge pull request #3 from zirenjin/master

dpa_tools merge

Author: zhaiwenxi

deepmd/dpa_tools/README.md

@@ -0,0 +1,221 @@
+# dpa_tools
+
+Fine-tuning, descriptor extraction, cross-validation, and data utilities for
+DPA-3 pretrained models.  Lives as a self-contained subpackage of `deepmd-kit`
+at `deepmd.dpa_tools`.
+
+## Relationship with deepmd-kit
+
+`dpa_tools` sits on top of deepmd-kit without modifying any existing module:
+
+- **Model loading**: `_backend.py` is the single choke point that imports
+  `deepmd.pt.model.model.get_model` and `deepmd.pt.train.wrapper.ModelWrapper`
+  to load DPA-3 checkpoints and extract descriptors.  No other file in
+  `dpa_tools` touches `deepmd.pt.*` directly.
+- **Training**: shells out to `dp --pt train` / `dp --pt freeze` /
+  `dp --pt test`, auto-generating `input.json` config files.
+- **Inference**: deepmd-kit's built-in `DeepProperty` handles neural-network
+  models; dpa_tools adds a lightweight frozen-descriptor + sklearn-head path.
+- **CLI**: registered as `dp dpa` subcommand group via `deepmd/main.py`.
+  Torch and all DPA dependencies are loaded lazily — only when a `dp dpa ...`
+  command actually runs.
+- **Lazy import**: `import deepmd.dpa_tools` does **not** trigger a `torch`
+  import.  `dp dpa --help` is equally lightweight.
+
+## Python API
+
+```python
+from deepmd.dpa_tools import (
+    DPAFineTuner,       # train (frozen sklearn / finetune / linear probe)
+    DPAPredictor,       # read-only inference from frozen bundles
+    MFTFineTuner,       # multi-task fine-tuning
+    DPATrainer,         # single-task dp --pt train wrapper
+    extract_descriptors, # standalone descriptor extraction
+    cross_validate,     # leak-proof cross-validation
+    train_test_split,   # formula-grouped data splitting
+    # data tools
+    convert,            # structure file → deepmd/npy
+    batch_convert,      # glob-based batch conversion
+    check_data,         # data sanity checks
+    attach_labels,      # inject external label arrays
+    load_dataset,       # label-filtered data loading
+)
+```
+
+### DPAFineTuner
+
+Four training strategies:
+
+| Strategy | Description | Best for |
+|----------|------------|----------|
+| `frozen_sklearn` | Freeze descriptor, extract once, fit sklearn head (RF/Ridge/MLP) | Small data (<1k samples), CPU inference |
+| `linear_probe` | Freeze backbone, train property fitting net only | Medium data, GPU |
+| `finetune` | Full-network fine-tuning | Larger data, GPU |
+| `scratch` | Train from random init (experimental) | Large-scale data only |
+
+```python
+model = DPAFineTuner(
+    pretrained="/path/to/DPA-3.1-3M.pt",
+    strategy="frozen_sklearn",
+    predictor="rf",
+    pooling="mean",
+)
+model.fit(train_data="/data/train", target_key="homo")
+model.predict("/data/test")
+model.freeze("model.dp-sklearn.pth")
+```
+
+### DPAPredictor
+
+```python
+pred = DPAPredictor("model.dp-sklearn.pth")
+result = pred.predict("/data/test")           # DotDict with .predictions
+metrics = pred.evaluate("/data/test")         # DotDict with .mae, .rmse, .r2
+
+# uncertainty: RF native, MLP via committee, Ridge raises
+result = pred.predict("/data/test", return_uncertainty=True)
+# → .predictions, .uncertainty
+```
+
+### MFTFineTuner
+
+Joint downstream property head + auxiliary force-field head (arXiv:2601.08486):
+
+```python
+mft = MFTFineTuner(
+    pretrained="/path/to/DPA-3.1-3M.pt",
+    downstream_task_type="property",
+    property_name="homo",
+    aux_branch="MP_traj_v024_alldata_mixu",
+)
+mft.fit(train_data="/data/qm9", aux_data="/data/spice2")
+mft.evaluate("/data/qm9_test")
+```
+
+### Descriptor extraction
+
+```python
+X = extract_descriptors(
+    "/data/systems",
+    pretrained="/path/to/DPA-3.1-3M.pt",
+    pooling="mean+std",
+)
+# → np.ndarray (n_frames, feat_dim * 2)
+```
+
+### Cross-validation
+
+Formula-grouped to prevent same-molecule leakage:
+
+```python
+from deepmd.dpa_tools import cross_validate, train_test_split
+
+systems = load_dataset("/data/root", label_key="energy")
+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")
+# → {"aggregate": {"mae_mean": ..., "rmse_std": ...}, ...}
+```
+
+### Data tools
+
+```python
+convert("POSCAR", "output_dir", fmt="vasp/poscar", type_map=["Cu", "O"])
+batch_convert("calcs/**/OUTCAR", "npy_root", fmt="vasp/outcar")
+check_data("/data/system")   # → list[Issue]
+attach_labels(system, head="bandgap", values=np.array([1.0, 2.0, 3.0]))
+```
+
+## CLI
+
+All commands live under `dp dpa` with two-level nesting:
+
+```
+dp dpa
+  extract-descriptors   extract pooled DPA descriptors to .npy
+  fit                   train a model (any strategy)
+  mft                   multi-task fine-tuning
+  cv                    cross-validate frozen_sklearn baseline
+  predict               predict with a frozen .pth bundle
+  evaluate              evaluate a frozen .pth against stored labels
+  data
+    convert             structure file → deepmd/npy
+    batch-convert       glob-based batch conversion
+    validate            sanity-check deepmd/npy directories
+    attach-labels       inject .npy labels into a system
+```
+
+`dp dpa --help` does not load torch.  The parser is pure argparse in
+`deepmd/main.py`; the handler import happens lazily in
+`deepmd/entrypoints/main.py` only when `dp dpa ...` is invoked.
+
+```bash
+dp dpa fit \
+  --train-data /data/train \
+  --pretrained /path/to/DPA-3.1-3M.pt \
+  --strategy frozen_sklearn \
+  --predictor rf \
+  --target-key homo
+
+dp dpa extract-descriptors \
+  --data /data/sys1 /data/sys2 \
+  --pretrained /path/to/DPA-3.1-3M.pt \
+  --pooling mean+std \
+  --output features.npy
+
+dp dpa mft \
+  --data /data/qm9 \
+  --aux-data /data/spice2 \
+  --pretrained /path/to/DPA-3.1-3M.pt \
+  --property-name homo
+
+dp dpa data convert --input POSCAR --output npy_dir --fmt vasp/poscar
+dp dpa data validate --data /data/sys1 /data/sys2
+```
+
+## Installation
+
+```bash
+pip install deepmd-kit[dpa-tools]
+```
+
+The `dpa-tools` extra brings in `scikit-learn`.  `torch` and `dpdata` are
+already provided by deepmd-kit's core dependencies.
+
+## Internal architecture
+
+```
+deepmd/dpa_tools/
+├── __init__.py           # public API, lazy imports (no torch at import time)
+├── _backend.py            # single choke point for deepmd.pt.* calls
+├── cli.py                 # dp dpa subcommand handlers
+├── finetuner.py           # DPAFineTuner (training + descriptor extraction)
+├── predictor.py           # DPAPredictor (read-only inference + uncertainty)
+├── mft.py                 # MFTFineTuner (multi-task fine-tuning)
+├── trainer.py             # DPATrainer (dp --pt train subprocess wrapper)
+├── cv.py                  # cross-validation + data splitting
+├── conditions.py          # scalar condition manager (T, P)
+├── config/
+│   └── manager.py         # MFT input.json generation
+├── data/
+│   ├── loader.py          # polymorphic data loading
+│   ├── dataset.py         # label-filtered loading
+│   ├── convert.py         # format conversion
+│   ├── validate.py        # data sanity checks
+│   ├── desc_cache.py      # two-tier descriptor cache
+│   ├── type_map.py        # automatic type-map resolution
+│   └── errors.py          # DPADataError
+└── utils/
+    ├── dotdict.py         # DotDict
+    └── sklearn_heads.py   # sklearn regressor factory
+```
+
+Key design points:
+- `_backend.py` is the **only** file that imports `deepmd.pt.*` — every call
+  into deepmd internals goes through it
+- `_DescriptorExtraction` encapsulates the fragile chain
+  `wrapper.model["Default"]` → `set_eval_descriptor_hook` → `forward_common`
+  → `eval_descriptor()`
+- `dp --pt train/test/freeze` always runs as a subprocess, keeping
+  dpa_tools decoupled from deepmd-kit's training entry points
+- `dpdata.System` is the universal internal data format

deepmd/dpa_tools/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DPA tools — fine-tuning, descriptor extraction, cross-validation, and data
+utilities for DPA-3 pretrained models.
+"""
+
+__version__ = "0.1.0"
+
+from .conditions import ConditionManager, DPAConditionError
+from .cv import cross_validate, train_test_split
+from .data import attach_labels, batch_convert, check_data, convert, load_dataset
+from .finetuner import DPAFineTuner, extract_descriptors
+from .mft import MFTFineTuner
+from .predictor import DPAPredictor
+from .trainer import DPATrainer
+
+__all__ = [
+    "ConditionManager",
+    "DPAConditionError",
+    "DPAFineTuner",
+    "DPAPredictor",
+    "DPATrainer",
+    "MFTFineTuner",
+    "attach_labels",
+    "batch_convert",
+    "check_data",
+    "convert",
+    "cross_validate",
+    "extract_descriptors",
+    "load_dataset",
+    "train_test_split",
+]

deepmd/dpa_tools/_backend.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Single chokepoint for all ``deepmd`` internal API and ``torch`` calls.
+
+Every import from ``deepmd.pt.*``, ``deepmd.utils.model_branch_dict``, or
+``torch`` that is needed by the rest of ``deepmd.dpa_tools`` must go through
+this module.  No other file in ``dpa_tools`` may import those packages directly.
+
+All functions that load ``torch`` or ``deepmd.pt`` keep the import inside the
+function body so that importing this module is cheap.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+# ``get_model_dict`` is backend-agnostic and lightweight — safe at module level.
+from deepmd.utils.model_branch_dict import get_model_dict as _get_model_dict
+
+
+# ---------------------------------------------------------------------------
+# torch I/O
+# ---------------------------------------------------------------------------
+
+
+def load_torch_file(path: str, map_location: str = "cpu") -> dict[str, Any]:
+    """Load a PyTorch checkpoint or frozen bundle.
+
+    Always uses ``weights_only=False`` because deepmd checkpoints carry
+    ``_extra_state`` (non-tensor metadata) and dpa_tools frozen bundles
+    carry ``sklearn`` pipeline objects.
+    """
+    import torch
+
+    return torch.load(path, map_location=map_location, weights_only=False)
+
+
+# ---------------------------------------------------------------------------
+# model construction
+# ---------------------------------------------------------------------------
+
+
+def build_model_from_config(input_param: dict[str, Any]):
+    """Build a (non-JIT) DPA model from an input-parameter dict.
+
+    Returns a ``ModelWrapper`` whose inner model is accessible as
+    ``wrapper.model["Default"]``.
+    """
+    from deepmd.pt.model.model import get_model
+    from deepmd.pt.train.wrapper import ModelWrapper
+
+    model = get_model(input_param)
+    return ModelWrapper(model)
+
+
+# ---------------------------------------------------------------------------
+# multi-task branch helpers
+# ---------------------------------------------------------------------------
+
+
+def resolve_model_branch(model_dict: dict[str, Any]) -> tuple[dict[str, str], str]:
+    """Resolve multi-task model-branch aliases.
+
+    Returns ``(alias_dict, model_dict)`` — the same tuple shape as the
+    upstream ``get_model_dict``.
+    """
+    return _get_model_dict(model_dict)
+
+
+# ---------------------------------------------------------------------------
+# device
+# ---------------------------------------------------------------------------
+
+
+def get_torch_device() -> Any:
+    """Return ``torch.device("cuda")`` if a GPU is available, else CPU."""
+    import torch
+
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# ---------------------------------------------------------------------------
+# descriptor extraction (the fragile chain)
+# ---------------------------------------------------------------------------
+
+
+class _DescriptorExtraction:
+    """Thin wrapper around a loaded model that runs a *single* forward pass
+    with ``eval_descriptor_hook`` enabled and returns per-atom descriptors.
+
+    This is the lowest-level building block.  Callers (like
+    ``DPAFineTuner._extract_features``) are responsible for pooling,
+    batching, and tensor creation.
+    """
+
+    def __init__(self, wrapper) -> None:
+        inner = wrapper.model["Default"]
+        self._inner_model = inner
+        self._atomic_model = inner.atomic_model
+
+    def _enable_hook(self) -> None:
+        self._atomic_model.set_eval_descriptor_hook(True)
+
+    def _disable_hook(self) -> None:
+        self._atomic_model.set_eval_descriptor_hook(False)
+
+    def _clear_accumulator(self) -> None:
+        self._atomic_model.eval_descriptor_list.clear()
+
+    def _run_forward(self, coord, atype, box):
+        """Run ``forward_common`` and return per-atom descriptors (detached).
+
+        Parameters
+        ----------
+        coord : torch.Tensor
+            (n_frames, n_atoms*3), float64, requires_grad.
+        atype : torch.Tensor
+            (n_frames, n_atoms), int64.
+        box : torch.Tensor
+            (n_frames, 9), float64.
+
+        Returns
+        -------
+        torch.Tensor
+            (n_frames, n_atoms, feat_dim), detached.
+        """
+        self._clear_accumulator()
+        self._inner_model.forward_common(coord, atype, box)
+        return self._atomic_model.eval_descriptor().detach()