feat: Add eval-desc CLI command for descriptor evaluation with 3D output format

deepmd/__about__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+__version__ = "unknown"

deepmd/entrypoints/eval_desc.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Evaluate descriptors using trained DeePMD model."""
+
+import logging
+import os
+from pathlib import (
+    Path,
+)
+from typing import (
+    Optional,
+)
+
+import numpy as np
+
+from deepmd.common import (
+    expand_sys_str,
+)
+from deepmd.infer.deep_eval import (
+    DeepEval,
+)
+from deepmd.utils.data import (
+    DeepmdData,
+)
+
+__all__ = ["eval_desc"]
+
+log = logging.getLogger(__name__)
+
+
+def eval_desc(
+    *,
+    model: str,
+    system: str,
+    datafile: str,
+    output: str = "desc",
+    head: Optional[str] = None,
+    **kwargs,
+) -> None:
+    """Evaluate descriptors for given systems.
+
+    Parameters
+    ----------
+    model : str
+        path where model is stored
+    system : str
+        system directory
+    datafile : str
+        the path to the list of systems to process
+    output : str
+        output directory for descriptor files
+    head : Optional[str], optional
+        (Supported backend: PyTorch) Task head if in multi-task mode.
+    **kwargs
+        additional arguments
+
+    Raises
+    ------
+    RuntimeError
+        if no valid system was found
+    """
+    if datafile is not None:
+        with open(datafile) as datalist:
+            all_sys = datalist.read().splitlines()
+    else:
+        all_sys = expand_sys_str(system)
+
+    if len(all_sys) == 0:
+        raise RuntimeError("Did not find valid system")
+
+    # init model
+    dp = DeepEval(model, head=head)
+
+    # create output directory
+    output_dir = Path(output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for cc, system_path in enumerate(all_sys):
+        log.info("# -------output of dp eval_desc------- ")
+        log.info(f"# processing system : {system_path}")
+
+        # create data class
+        tmap = dp.get_type_map()
+        data = DeepmdData(
+            system_path,
+            set_prefix="set",
+            shuffle_test=False,
+            type_map=tmap,
+            sort_atoms=False,
+        )
+
+        # get test data
+        test_data = data.get_test()
+        mixed_type = data.mixed_type
+        natoms = len(test_data["type"][0])
+        nframes = test_data["box"].shape[0]
+
+        # prepare input data
+        coord = test_data["coord"].reshape([nframes, -1])
+        box = test_data["box"]
+        if not data.pbc:
+            box = None
+        if mixed_type:
+            atype = test_data["type"].reshape([nframes, -1])
+        else:
+            atype = test_data["type"][0]
+
+        # handle optional parameters
+        fparam = None
+        if dp.get_dim_fparam() > 0:
+            if "fparam" in test_data:
+                fparam = test_data["fparam"]
+
+        aparam = None
+        if dp.get_dim_aparam() > 0:
+            if "aparam" in test_data:
+                aparam = test_data["aparam"]
+
+        # evaluate descriptors
+        log.info(f"# evaluating descriptors for {nframes} frames")
+        descriptors = dp.eval_descriptor(
+            coord,
+            box,
+            atype,
+            fparam=fparam,
+            aparam=aparam,
+        )
+
+        # save descriptors
+        system_name = os.path.basename(system_path.rstrip('/'))
+        desc_file = output_dir / f"{system_name}.npy"
+        np.save(desc_file, descriptors)
+
+        log.info(f"# descriptors saved to {desc_file}")
+        log.info(f"# descriptor shape: {descriptors.shape}")
+        log.info("# ----------------------------------- ")
+
+    log.info("# eval_desc completed successfully")

deepmd/entrypoints/main.py

@@ -18,6 +18,9 @@
 from deepmd.entrypoints.doc import (
     doc_train_input,
 )
+from deepmd.entrypoints.eval_desc import (
+    eval_desc,
+)
 from deepmd.entrypoints.gui import (
     start_dpgui,
 )
@@ -65,6 +68,14 @@ def main(args: argparse.Namespace) -> None:
             strict_prefer=False,
         )
         test(**dict_args)
+    elif args.command == "eval-desc":
+        dict_args["model"] = format_model_suffix(
+            dict_args["model"],
+            feature=Backend.Feature.DEEP_EVAL,
+            preferred_backend=args.backend,
+            strict_prefer=False,
+        )
+        eval_desc(**dict_args)
     elif args.command == "doc-train-input":
         doc_train_input(**dict_args)
     elif args.command == "model-devi":

deepmd/main.py

@@ -416,6 +416,56 @@ def main_parser() -> argparse.ArgumentParser:
         help="(Supported backend: PyTorch) Task head (alias: model branch) to test if in multi-task mode.",
     )
 
+    # * eval_desc script ***************************************************************
+    parser_eval_desc = subparsers.add_parser(
+        "eval-desc",
+        parents=[parser_log],
+        help="evaluate descriptors using the model",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent(
+            """\
+        examples:
+            dp eval-desc -m graph.pb -s /path/to/system -o desc
+        """
+        ),
+    )
+    parser_eval_desc.add_argument(
+        "-m",
+        "--model",
+        default="frozen_model",
+        type=str,
+        help="Frozen model file (prefix) to import. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth.",
+    )
+    parser_eval_desc_subgroup = parser_eval_desc.add_mutually_exclusive_group()
+    parser_eval_desc_subgroup.add_argument(
+        "-s",
+        "--system",
+        default=".",
+        type=str,
+        help="The system dir. Recursively detect systems in this directory",
+    )
+    parser_eval_desc_subgroup.add_argument(
+        "-f",
+        "--datafile",
+        default=None,
+        type=str,
+        help="The path to the datafile, each line of which is a path to one data system.",
+    )
+    parser_eval_desc.add_argument(
+        "-o",
+        "--output",
+        default="desc",
+        type=str,
+        help="Output directory for descriptor files. Descriptors will be saved as desc/(system_name).npy",
+    )
+    parser_eval_desc.add_argument(
+        "--head",
+        "--model-branch",
+        default=None,
+        type=str,
+        help="(Supported backend: PyTorch) Task head (alias: model branch) to use if in multi-task mode.",
+    )
+
     # * compress model *****************************************************************
     # Compress a model, which including tabulating the embedding-net.
     # The table is composed of fifth-order polynomial coefficients and is assembled
@@ -909,6 +959,7 @@ def main(args: Optional[list[str]] = None) -> None:
 
     if args.command in (
         "test",
+        "eval-desc",
         "doc-train-input",
         "model-devi",
         "neighbor-stat",

doc/inference/python.md

@@ -19,6 +19,21 @@ e, f, v = dp.eval(coord, cell, atype)
 
 where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively.
 
+One can also evaluate the descriptors of the model:
+
+```python
+from deepmd.infer import DeepPot
+import numpy as np
+
+dp = DeepPot("graph.pb")
+coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1])
+cell = np.diag(10 * np.ones(3)).reshape([1, -1])
+atype = [1, 0, 1]
+descriptors = dp.eval_descriptor(coord, cell, atype)
+```
+
+where `descriptors` is the descriptor matrix of the system. This can also be done using the command line interface `dp eval-desc` as described in the [test documentation](../test/test.md).
+
 Furthermore, one can use the python interface to calculate model deviation.
 
 ```python

doc/test/test.md

@@ -17,3 +17,25 @@ An explanation will be provided
 ```{program-output} dp test -h
 
 ```
+
+# Evaluate descriptors
+
+The descriptors of a model can be evaluated and saved using `dp eval-desc`. A typical usage of `dp eval-desc` is
+
+```bash
+dp eval-desc -m graph.pb -s /path/to/system -o desc
+```
+
+where `-m` gives the model file, `-s` the path to the system directory (or `-f` for a datafile containing paths to systems), and `-o` the output directory where descriptor files will be saved. The descriptors for each system will be saved as `.npy` files with the format `desc/(system_name).npy`.
+
+Several other command line options can be passed to `dp eval-desc`, which can be checked with
+
+```bash
+$ dp eval-desc --help
+```
+
+An explanation will be provided
+
+```{program-output} dp eval-desc -h
+
+```