Merge branch 'master' into feat/pt_expt-auto-neighbor-graph-builder

Author: Shaurya2k06

.github/workflows/test_cuda.yml

@@ -125,7 +125,7 @@ jobs:
         # junit-aoti.xml = compile lane (T_A), junit-gpu.xml = GPU lane (T_B);
         # used to compare the two lanes against the pre-split serial baseline.
         if: always()
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v7
         with:
           name: cuda-pytest-junit
           path: |

deepmd/dpmodel/utils/lmdb_data.py

@@ -149,6 +149,18 @@ def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]:
     return out
 
 
+def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray:
+    """Remap real atom types while preserving negative virtual sentinels.
+
+    Positive indices retain NumPy's normal bounds checking, so malformed LMDB
+    data cannot be silently reinterpreted as a different species.
+    """
+    remapped_atype = atype.astype(np.int64, copy=True)
+    real_atom_mask = remapped_atype >= 0
+    remapped_atype[real_atom_mask] = type_remap[remapped_atype[real_atom_mask]]
+    return remapped_atype
+
+
 def is_lmdb(systems: str) -> bool:
     """Check if systems points to an LMDB dataset."""
     return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file()
@@ -466,10 +478,16 @@ def __init__(
     def _compute_natoms_vec(self, atype: np.ndarray) -> np.ndarray:
         """Compute natoms_vec from a frame's atype array.
 
+        Negative virtual types are excluded from the per-type counts, matching
+        mixed-type NPY data handling. This function also excludes positive
+        indices outside the configured type map. The leading nloc entries still
+        include every atom slot.
+
         Returns [nloc, nloc, count_type0, count_type1, ...] with length ntypes+2.
         """
         nloc = len(atype)
-        counts = np.bincount(atype, minlength=self._ntypes)[: self._ntypes]
+        real_atype = atype[(atype >= 0) & (atype < self._ntypes)]
+        counts = np.bincount(real_atype, minlength=self._ntypes)
         vec = np.empty(self._ntypes + 2, dtype=np.int64)
         vec[0] = nloc
         vec[1] = nloc
@@ -579,7 +597,7 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             frame["atype"] = frame["atype"].reshape(-1).astype(np.int64)
             # Remap atom types from LMDB's type_map to model's type_map
             if self._type_remap is not None:
-                frame["atype"] = self._type_remap[frame["atype"]].astype(np.int64)
+                frame["atype"] = _remap_atom_types(frame["atype"], self._type_remap)
         if "virial" in frame and isinstance(frame["virial"], np.ndarray):
             frame["virial"] = (
                 frame["virial"].reshape(9).astype(self._resolve_dtype("virial"))
@@ -1459,9 +1477,9 @@ def __init__(
                         and "atype" in frame
                         and isinstance(frame["atype"], np.ndarray)
                     ):
-                        frame["atype"] = self._type_remap[
-                            frame["atype"].reshape(-1)
-                        ].astype(np.int64)
+                        frame["atype"] = _remap_atom_types(
+                            frame["atype"].reshape(-1), self._type_remap
+                        )
                     self._frames.append(frame)
 
         # Shuffle if requested

deepmd/dpmodel/utils/stat.py

@@ -72,7 +72,8 @@ def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str
     Returns
     -------
     list[str]
-        Sorted list of observed element symbols.
+        Sorted list of observed element symbols. Negative virtual atom types
+        and indices outside ``type_map`` are ignored.
     """
     from deepmd.utils.econf_embd import (
         sort_element_type,
@@ -83,7 +84,7 @@ def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str
         atype = to_numpy_array(system["atype"])  # shape: [nframes, natoms]
         observed_indices.update(np.unique(atype).tolist())
     observed_types = [
-        type_map[i] for i in sorted(observed_indices) if i < len(type_map)
+        type_map[i] for i in sorted(observed_indices) if 0 <= i < len(type_map)
     ]
     return sort_element_type(observed_types)
 

doc/README

@@ -1 +1,3 @@
-To run the HTML documentation build, doxygen have to be installed.
+Building the documentation requires Python 3.11 or newer because the `docs`
+extra uses MyST-Parser 5.1 and Sphinx 8. Doxygen must also be installed to build
+the HTML documentation.

doc/backend.md

@@ -45,9 +45,8 @@ The `.pd` extension is used for model checkpoint storage, which is commonly util
 
 ### DP {{ dpmodel_icon }}
 
-:::{note}
-This backend is only for development and should not take into production.
-:::
+> [!NOTE]
+> This backend is only for development and should not take into production.
 
 - Model filename extension: `.dp`, `.yaml`, `.yml`
 

doc/conf.py

@@ -201,6 +201,7 @@
     "https://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.0/es5/tex-mml-chtml.min.js"
 )
 myst_enable_extensions = [
+    "alert",
     "dollarmath",
     "colon_fence",
     "substitution",

doc/data/system.md

@@ -50,9 +50,8 @@ In general, we always use the following convention of units:
 
 ## Mixed type
 
-:::{note}
-Only the [DPA-1](../model/train-se-atten.md) and [DPA-2](../model/dpa2.md) descriptors support this format.
-:::
+> [!NOTE]
+> Only the [DPA-1](../model/train-se-atten.md) and [DPA-2](../model/dpa2.md) descriptors support this format.
 
 In the standard data format, only those frames with the same fingerprint (i.e. the number of atoms of different elements) can be put together as a unified system.
 This may lead to sparse frame numbers in those rare systems.

doc/development/create-a-model-pt.md

@@ -1,10 +1,9 @@
 # Create a model in other backends {{ pytorch_icon }} {{ dpmodel_icon }}
 
-:::{note}
-**Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
-
-In the following context, we use the PyTorch backend as the example, while it also applies to other backends listed above.
-:::
+> [!NOTE]
+> **Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
+>
+> In the following context, we use the PyTorch backend as the example, while it also applies to other backends listed above.
 
 If you'd like to create a new model that isn't covered by the existing DeePMD-kit library, but reuse DeePMD-kit's other efficient modules such as data processing, trainer, etc, you may want to read this section.
 

doc/development/type-embedding.md

@@ -86,6 +86,5 @@ build -> _pass_filter -> _filter -> _filter_lower
 
 In `fitting net`, it takes the descriptor vector as input, whose dimension is \[natoms, $M_1\times M_2$\]. Because we need to involve information on the centric atom in this step, we need to generate a matrix named `atype_embed` (of dim [natoms, nchanl]), in which each row is the type embedding vector of the specific centric atom. The input is sorted by type of centric atom, we also know the number of a particular atom type (stored in `natoms[2+i]`), thus we get the type vector of the centric atom. In the build phase of the fitting net, it will check whether type embedding exists in `input_dict` and fetch them. After that, call `embed_atom_type` function to look up the embedding vector for the type vector of the centric atom to obtain `atype_embed`, and concat input with it ([input, atype_embed]). The modified input goes through `fitting` net\` to get predicted energy.
 
-:::{note}
-You can't apply the compression method while using atom-type embedding.
-:::
+> [!NOTE]
+> You can't apply the compression method while using atom-type embedding.

doc/env.md

@@ -1,8 +1,7 @@
 # Runtime environment variables
 
-:::{note}
-For build-time environment variables, see [Install from source code](./install/install-from-source.md).
-:::
+> [!NOTE]
+> For build-time environment variables, see [Install from source code](./install/install-from-source.md).
 
 ## All interfaces