From 1a6834d1d9fe7c25f4ea83bd29ef5589cea24214 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 6 Aug 2024 04:40:46 -0400
Subject: [PATCH] feat: make `dp neighbor-stat --type-map` optional (#4049)

In most cases, we can directly read type map from data.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- The `neighbor_stat` function now accepts an optional `type_map`,
enhancing flexibility in function calls.
- The `--type-map` argument in the main parser has been changed to
optional, defaulting to a derived type map if not provided.

- **Improvements**
- Added logging to indicate when the `type_map` is not supplied,
enhancing user awareness and debugging capabilities.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/entrypoints/neighbor_stat.py | 5 ++++-
 deepmd/main.py                      | 4 ++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/deepmd/entrypoints/neighbor_stat.py b/deepmd/entrypoints/neighbor_stat.py
index 87ec101c70..8840851b91 100644
--- a/deepmd/entrypoints/neighbor_stat.py
+++ b/deepmd/entrypoints/neighbor_stat.py
@@ -2,6 +2,7 @@
 import logging
 from typing import (
     List,
+    Optional,
 )
 
 from deepmd.backend.backend import (
@@ -21,7 +22,7 @@ def neighbor_stat(
     *,
     system: str,
     rcut: float,
-    type_map: List[str],
+    type_map: Optional[List[str]],
     mixed_type: bool = False,
     backend: str = "tensorflow",
     **kwargs,
@@ -88,6 +89,8 @@ def neighbor_stat(
         rcut=rcut,
         type_map=type_map,
     )
+    if type_map is None:
+        log.info(f"type_map: {data.get_type_map()}")
     data.get_batch()
     nei = NeighborStat(data.get_ntypes(), rcut, mixed_type=mixed_type)
     min_nbor_dist, max_nbor_size = nei.get_stat(data)
diff --git a/deepmd/main.py b/deepmd/main.py
index 964521b77e..777bfd3aa3 100644
--- a/deepmd/main.py
+++ b/deepmd/main.py
@@ -655,8 +655,8 @@ def main_parser() -> argparse.ArgumentParser:
         "--type-map",
         type=str,
         nargs="+",
-        required=True,
-        help="type map",
+        required=False,
+        help="Type map. If not provided, the type map of data will be used.",
     )
     parser_neighbor_stat.add_argument(
         "--mixed-type",