Merge pull request #50 from zirenjin/master

fix(dpa-adapt): resolve pre-commit errors, desc_cache import cycle, descriptor hook accumulator

Author: zhaiwenxi

doc/conf.py

@@ -7,6 +7,10 @@
 
 # -- Path setup --------------------------------------------------------------
 
+from __future__ import (
+    annotations,
+)
+
 import datetime
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -224,3 +228,54 @@
 bibtex_bibfiles = ["../CITATIONS.bib"]
 
 remove_from_toctrees = ["autoapi/**/*", "API_CC/*", "api_c/*", "api_core/*"]
+
+
+# Auto-generated CLI reference pages (sphinx-argparse) nest a section per
+# subcommand and per argument group. Under the global ``:numbered:`` toctree
+# this explodes into unhelpful deep numbers (e.g. ``9.3.3.6.3.1.1.``). Cap the
+# section numbering at the given depth (number of dotted components); headings
+# deeper than that are left unnumbered. Only the listed pages are affected.
+from typing import (
+    TYPE_CHECKING,
+)
+
+from docutils import (
+    nodes,
+)
+
+if TYPE_CHECKING:
+    from sphinx.application import (
+        Sphinx,
+    )
+
+cli_secnumber_max_depth = {
+    "dpa_adapt/cli": 5,
+}
+
+
+def _cap_cli_secnumbers(app: Sphinx, doctree: nodes.document, docname: str) -> None:
+    """Drop section numbers below ``cli_secnumber_max_depth`` for CLI pages."""
+    max_depth = cli_secnumber_max_depth.get(docname)
+    if max_depth is None:
+        return
+    secnumbers = app.env.toc_secnumbers.get(docname)
+    if not secnumbers:
+        return
+    # The empty anchor "" holds the page chapter number (e.g. ``(9, 3)``).
+    # It must be dropped from the map, otherwise the writer falls back to it for
+    # the now-unnumbered deep sections; re-attach it to the page title instead.
+    page_number = secnumbers.get("")
+    app.env.toc_secnumbers[docname] = {
+        anchor: number
+        for anchor, number in secnumbers.items()
+        if anchor != "" and len(number) <= max_depth
+    }
+    if page_number:
+        for title in doctree.findall(nodes.title):
+            title["secnumber"] = page_number
+            break
+
+
+def setup(app: Sphinx) -> dict[str, bool]:
+    app.connect("doctree-resolved", _cap_cli_secnumbers)
+    return {"parallel_read_safe": True, "parallel_write_safe": True}

doc/dpa_adapt/input_formats.md

@@ -12,7 +12,7 @@ auto-detect the input type and route it to the correct pipeline:
 **formula table** → random doping from a POSCAR template,
 **structure files** → dpdata (auto-detect or explicit `--fmt`).
 
-## 1. SMILES Tables (CSV)
+## SMILES Tables (CSV)
 
 **Trigger:** file extension `.csv` **and** a SMILES column.
 By default, the converter reads `SMILES`/`smiles`; use `--smiles-col` for
@@ -46,7 +46,7 @@ dpaad data convert --input data.csv --output ./npy --fmt smiles \
     --split-seed 42 --conformer-seed 43
 ```
 
-## 2. Formula Tables (CSV/TXT + POSCAR Template)
+## Formula Tables (CSV/TXT + POSCAR Template)
 
 **Trigger:** `--fmt formula`. Reads a table of elemental composition formulas
 (e.g. `Ni0.65Gd0.15O2H1`) and a template POSCAR, then generates doped
@@ -91,7 +91,7 @@ dpa-adapt data convert --input compositions.txt --output ./npy --fmt formula \
     --poscar template.POSCAR --formula-col 0 --property-col 1
 ```
 
-## 3. Structure Files via dpdata
+## Structure Files via dpdata
 
 **Trigger:** inputs not routed to the SMILES or formula pipelines. This means
 `--fmt` is neither `smiles` nor `formula`; when `--fmt` is omitted, CSV inputs

dpa_adapt/cv.py

@@ -481,7 +481,7 @@ def cross_validate(
     # This reuses existing desc_mean.npy when present, extracts only missing
     # systems one-by-one.  Peak memory is one system's descriptors at a time.
     if is_cheap:
-        from dpa_adapt.data.desc_cache import (
+        from dpa_adapt.finetuner import (
             ensure_per_system_cache,
         )
 

dpa_adapt/data/desc_cache.py

@@ -8,6 +8,10 @@
 #
 # Systems are ``dpdata.System`` objects; cache keys are computed from
 # data fingerprints and resolved checkpoint metadata.
+#
+# Note: ``load_or_extract()`` and ``ensure_per_system_cache()`` live in
+# ``dpa_adapt.finetuner`` to avoid an import cycle (those functions need
+# ``DPAFineTuner``, while ``finetuner`` imports cache helpers from here).
 
 from __future__ import (
     annotations,
@@ -125,78 +129,7 @@ def _cache_key(
 
 
 # ---------------------------------------------------------------------------
-# bulk cache
-# ---------------------------------------------------------------------------
-
-
-def load_or_extract(
-    systems: list,
-    pretrained: str,
-    model_branch: str | None = None,
-    pooling: str = "mean",
-    cache: bool = True,
-    type_map: list[str] | tuple[str, ...] | None = None,
-) -> np.ndarray:
-    """Return descriptors for *systems*, using the cache when possible.
-
-    Parameters
-    ----------
-    systems : list[dpdata.System]
-        Systems to extract descriptors from.
-    pretrained : str
-        Path to the DPA checkpoint.
-    model_branch : str, optional
-        Branch name.
-    pooling : str
-        Pooling strategy.
-    cache : bool
-        If False the cache is bypassed entirely.
-    type_map : list[str] or tuple[str, ...], optional
-        Element symbols used to build the descriptor model and cache key.
-
-    Returns
-    -------
-    np.ndarray, shape ``(n_frames_total, feat_dim)``
-    """
-    if cache:
-        key = _cache_key(
-            systems,
-            pretrained,
-            model_branch,
-            pooling,
-            type_map=type_map,
-        )
-        cache_path = _cache_dir() / f"{key}.npy"
-        if cache_path.is_file():
-            _LOG.info("Descriptor cache hit: %s", cache_path.name)
-            return np.load(cache_path)
-        _LOG.info("Descriptor cache miss; extracting...")
-    else:
-        _LOG.info("Descriptor cache bypassed (cache=False).")
-
-    from dpa_adapt.finetuner import (
-        DPAFineTuner,
-    )
-
-    extractor = DPAFineTuner(
-        pretrained=pretrained,
-        model_branch=model_branch,
-        predictor="linear",
-        pooling=pooling,
-        type_map=list(type_map) if type_map else None,
-    )
-    descriptors = extractor._extract_features(systems)
-
-    if cache:
-        cache_path.parent.mkdir(parents=True, exist_ok=True)
-        np.save(cache_path, descriptors)
-        _LOG.info("Cached descriptors to %s", cache_path)
-
-    return descriptors
-
-
-# ---------------------------------------------------------------------------
-# per-system cache — used by cross_validate to avoid OOM
+# per-system cache path helpers
 # ---------------------------------------------------------------------------
 
 
@@ -216,74 +149,6 @@ def _per_system_cache_path(
     return _cache_dir() / "per_system" / f"{fp}.npy"
 
 
-def ensure_per_system_cache(
-    systems: list,
-    pretrained: str,
-    model_branch: str | None = None,
-    pooling: str = "mean",
-    type_map: list[str] | tuple[str, ...] | None = None,
-) -> None:
-    """Ensure every system has its descriptors cached to disk.
-
-    Existing cache files are reused as-is.  Missing ones are extracted one
-    system at a time for low peak memory.
-    """
-    missing: list = []
-    for system in systems:
-        if not _per_system_cache_path(
-            system,
-            pretrained,
-            model_branch,
-            pooling,
-            type_map,
-        ).is_file():
-            missing.append(system)
-
-    if not missing:
-        _LOG.info(
-            "All %d systems have per-system cache; nothing to extract.", len(systems)
-        )
-        return
-
-    import torch
-
-    from dpa_adapt.finetuner import (
-        DPAFineTuner,
-    )
-
-    _LOG.info(
-        "%d/%d systems missing per-system cache; extracting one by one...",
-        len(missing),
-        len(systems),
-    )
-
-    extractor = DPAFineTuner(
-        pretrained=pretrained,
-        model_branch=model_branch,
-        predictor="linear",
-        pooling=pooling,
-        type_map=list(type_map) if type_map else None,
-    )
-
-    for i, system in enumerate(missing):
-        cache_path = _per_system_cache_path(
-            system,
-            pretrained,
-            model_branch,
-            pooling,
-            type_map,
-        )
-        cache_path.parent.mkdir(parents=True, exist_ok=True)
-        desc = extractor._extract_features([system])
-        np.save(cache_path, desc)
-        if extractor._device is not None and extractor._device.type == "cuda":
-            torch.cuda.empty_cache()
-        if i > 0 and i % 50 == 0:
-            _LOG.info("  per-system cache: %d/%d done", i, len(missing))
-
-    _LOG.info("Per-system cache ready (%d systems).", len(systems))
-
-
 def get_per_system_descriptor(
     system: dpdata.System,
     pretrained: str,