[pull] master from deepmodeling:master

deepmd/dpmodel/infer/deep_eval.py

@@ -240,9 +240,39 @@ def eval(
         natoms, numb_test = self._get_natoms_and_nframes(
             coords, atom_types, len(atom_types.shape) > 1
         )
+        # The public evaluator accepts a superset of backend/model inputs (for
+        # example ``efield`` is TensorFlow-only). Forward only the dpmodel
+        # inputs understood here instead of leaking unrelated ``None`` values
+        # into concrete model ``call`` signatures.
+        model_kwargs = {}
+        charge_spin = kwargs.get("charge_spin")
+        if charge_spin is not None:
+            model_kwargs["charge_spin"] = charge_spin
+        if self.get_has_spin():
+            spin = kwargs.get("spin")
+            if spin is None:
+                raise ValueError("spin must be provided when evaluating a spin model")
+            spin = np.asarray(spin)
+            expected_spin_size = numb_test * natoms * 3
+            if spin.size != expected_spin_size:
+                raise ValueError(
+                    "spin must contain exactly "
+                    f"{expected_spin_size} values for {numb_test} frame(s) and "
+                    f"{natoms} atom(s), but received {spin.size}"
+                )
+            # AutoBatchSize slices only arrays with a frame axis. Normalize a
+            # flattened public-API input before batching so each model call
+            # receives the spins belonging to its coordinate frames.
+            model_kwargs["spin"] = spin.reshape(numb_test, natoms, 3)
         request_defs = self._get_request_defs(atomic)
         out = self._eval_func(self._eval_model, numb_test, natoms)(
-            coords, cells, atom_types, fparam, aparam, request_defs
+            coords,
+            cells,
+            atom_types,
+            fparam,
+            aparam,
+            request_defs,
+            **model_kwargs,
         )
         # ``AutoBatchSize.execute_all`` unwraps a single-output result out of
         # its tuple, which would make ``zip`` iterate over the array's frame
@@ -287,6 +317,12 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]:
                     OutputVariableCategory.DERV_R,
                     OutputVariableCategory.DERV_C_REDU,
                 )
+                # ``mask_mag`` is exported directly by spin graphs but does not
+                # fit the category filter. Adding all OUT variables would also
+                # request atom energy and the general atom mask at
+                # ``atomic=False``, so keep this low-cost compatibility output
+                # explicit instead of widening the category set.
+                or x.name == "mask_mag"
             ]
 
     def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable:
@@ -342,6 +378,7 @@ def _eval_model(
         fparam: Array | None,
         aparam: Array | None,
         request_defs: list[OutputVariableDef],
+        **model_kwargs: Any,
     ) -> dict[str, Array]:
         model = self.dp
 
@@ -370,14 +407,15 @@ def _eval_model(
         do_atomic_virial = any(
             x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
         )
-        batch_output = model(
-            coord_input,
-            type_input,
+        # Evaluator-owned arguments take precedence over extra model inputs so
+        # callers cannot accidentally bypass normalization performed above.
+        model_kwargs.update(
             box=box_input,
             fparam=fparam_input,
             aparam=aparam_input,
             do_atomic_virial=do_atomic_virial,
         )
+        batch_output = model(coord_input, type_input, **model_kwargs)
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
 

deepmd/dpmodel/utils/lmdb_data.py

@@ -2685,20 +2685,61 @@ def get_test(self) -> dict[str, Any]:
         return self._inner.get_test(nloc=self._nloc)
 
 
+def _validate_merge_type_maps(
+    source_metadata: list[tuple[str, dict[str, Any]]],
+) -> list[str] | None:
+    """Return the shared type map required for byte-for-byte frame merging.
+
+    ``merge_lmdb`` does not decode and rewrite atom-type arrays, so every
+    source must use exactly the same index-to-species mapping. All-missing
+    legacy metadata remains supported, but mixing explicit and missing maps is
+    rejected because compatibility cannot be established.
+    """
+    source_type_maps = [(path, meta.get("type_map")) for path, meta in source_metadata]
+    explicit_type_maps = [
+        (path, list(type_map))
+        for path, type_map in source_type_maps
+        if type_map is not None
+    ]
+    if not explicit_type_maps:
+        return None
+
+    formatted_maps = ", ".join(
+        f"{path}: {list(type_map)!r}" if type_map is not None else f"{path}: missing"
+        for path, type_map in source_type_maps
+    )
+    if len(explicit_type_maps) != len(source_type_maps):
+        raise ValueError(
+            "Cannot merge LMDB datasets with mixed type_map metadata because "
+            f"raw atom-type indices cannot be validated ({formatted_maps})"
+        )
+
+    canonical_type_map = explicit_type_maps[0][1]
+    if any(type_map != canonical_type_map for _, type_map in explicit_type_maps[1:]):
+        raise ValueError(
+            "Cannot merge LMDB datasets with incompatible type_map values "
+            f"because frames are copied without remapping ({formatted_maps})"
+        )
+    return canonical_type_map
+
+
 def _copy_lmdb_source(
     src_path: str,
+    metadata: dict[str, Any],
     dst_env: lmdb.Environment,
     dst_format: str,
     frame_idx: int,
     frame_nlocs: list[int],
     frame_system_ids: list[int],
     system_id_offset: int,
-) -> tuple[int, dict, list[str] | None, int]:
-    """Copy one source under a ref-counted environment lease."""
+) -> tuple[int, dict, int]:
+    """Copy one validated source under a ref-counted environment lease.
+
+    The caller supplies metadata collected during the validation preflight so
+    this copy pass does not read and decode it a second time.
+    """
     src_env = _open_lmdb(src_path)
     try:
-        with src_env.begin() as transaction:
-            metadata = _read_metadata(transaction)
         nframes, src_format, natoms_per_type = _parse_metadata(metadata)
         fallback_natoms = sum(natoms_per_type)
         source_nlocs = metadata.get("frame_nlocs")
@@ -2741,7 +2782,6 @@ def _copy_lmdb_source(
         return (
             frame_idx,
             metadata.get("system_info", {}),
-            metadata.get("type_map"),
             system_id_offset,
         )
     finally:
@@ -2772,30 +2812,47 @@ def merge_lmdb(
     -------
     str
         Path to the created LMDB.
+
+    Raises
+    ------
+    ValueError
+        If sources use different explicit type maps, or mix explicit type-map
+        metadata with legacy metadata where the mapping is missing.
     """
     import os
     import shutil
 
+    # Validate every source before replacing or creating the destination. A
+    # type-map validation failure must not destroy an existing dataset or
+    # leave a partial output.
+    source_metadata: list[tuple[str, dict[str, Any]]] = []
+    for src_path in src_paths:
+        src_env = _open_lmdb(src_path)
+        try:
+            with src_env.begin() as txn:
+                source_metadata.append((src_path, _read_metadata(txn)))
+        finally:
+            _close_lmdb(src_path)
+    merged_type_map = _validate_merge_type_maps(source_metadata)
+
     if os.path.exists(dst_path):
         shutil.rmtree(dst_path)
-
     dst_env = lmdb.open(dst_path, map_size=map_size)
     frame_idx = 0
     fmt = "012d"
     frame_nlocs: list[int] = []
     frame_system_ids: list[int] = []
     first_system_info: dict | None = None
-    first_type_map: list[str] | None = None
     sys_id_offset = 0
     try:
-        for src_path in src_paths:
+        for src_path, metadata in source_metadata:
             (
                 frame_idx,
                 source_system_info,
-                source_type_map,
                 sys_id_offset,
             ) = _copy_lmdb_source(
                 src_path,
+                metadata,
                 dst_env,
                 fmt,
                 frame_idx,
@@ -2805,8 +2862,6 @@ def merge_lmdb(
             )
             if first_system_info is None:
                 first_system_info = source_system_info
-            if first_type_map is None:
-                first_type_map = source_type_map
 
         merged_meta = {
             "nframes": frame_idx,
@@ -2815,8 +2870,8 @@ def merge_lmdb(
             "frame_nlocs": frame_nlocs,
             "frame_system_ids": frame_system_ids,
         }
-        if first_type_map is not None:
-            merged_meta["type_map"] = first_type_map
+        if merged_type_map is not None:
+            merged_meta["type_map"] = merged_type_map
         with dst_env.begin(write=True) as transaction:
             transaction.put(
                 b"__metadata__",

deepmd/pt_expt/train/training.py

@@ -971,6 +971,13 @@ def forward(
             nloc,
             rcut,
             sel,
+            # Keep the candidate list merged, matching eager training's
+            # DefaultNeighborList contract. forward_common_lower always calls
+            # model.format_nlist, which performs the type split for non-mixed
+            # descriptors. The shared builder globally truncates to sum(sel)
+            # before either split, so distinguish_types=True would only move
+            # the same layout transform earlier without changing the final
+            # formatted neighbor list consumed by the lower model.
             distinguish_types=False,
             # model-level pair exclusion is a nlist-BUILD transform (decision
             # #18/A4); the compiled dense lower consumes a pre-excluded nlist.

source/api_cc/include/DeepSpinTF.h

@@ -259,14 +259,16 @@ class DeepSpinTF : public DeepSpinBackend {
               const int nghost,
               const std::vector<VALUETYPE>& spin,
               const int numb_types,
-              const int numb_types_spin);
+              const int numb_types_spin,
+              const int nframes);
 
   template <typename VALUETYPE>
   void extend_nlist(std::vector<VALUETYPE>& extend_dcoord,
                     std::vector<int>& extend_atype,
                     const std::vector<VALUETYPE>& dcoord_,
                     const std::vector<VALUETYPE>& dspin_,
-                    const std::vector<int>& datype_);
+                    const std::vector<int>& datype_,
+                    const int nframes);
 
   void cum_sum(std::map<int, int>&, std::map<int, int>&);