feat(lmdb): support mixed-size batches and lazy label availability

deepmd/dpmodel/loss/dos.py

@@ -159,7 +159,7 @@ def call(
             )
             diff3d = local_pred - local_label  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_dos = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos
@@ -184,7 +184,7 @@ def call(
             )
             diff3d = local_pred_cdf - local_label_cdf  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_cdf = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos

deepmd/dpmodel/loss/ener_spin.py

@@ -471,6 +471,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         return label_requirement

deepmd/dpmodel/loss/reduction.py

@@ -1,19 +1,60 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Shared masked per-frame reduction idioms for the loss modules.
-
-These helpers factor out the three per-frame reduction patterns that the
-mixed_type padding mask (PR #5738) introduced into every loss term (issue
-#5768). They are written with ``array_api_compat`` so both the dpmodel
-(numpy/jax/...) loss backend and the PyTorch loss backend can call them: the
-PyTorch backend passes torch tensors and ``array_api_compat`` dispatches to the
-torch namespace, preserving autograd and producing bit-identical results to the
-previous hand-inlined torch code.
-
-Each helper implements ONLY the masked branch. Callers keep the original
-non-masked expression in the ``else`` branch verbatim, so the "bit-identical
-for non-mixed batches" guarantee from PR #5738 is preserved (defaulting the
-mask to all-ones would change the reduction order at the ULP level and is
-deliberately NOT done here).
+"""Shared masked reduction idioms for the loss modules.
+
+These helpers factor out the reduction patterns every loss term needs once a
+per-atom mask marks which rows of a batch are real. They are written with
+``array_api_compat`` so both the dpmodel (numpy/jax/...) loss backend and the
+PyTorch loss backend can call them: the PyTorch backend passes torch tensors
+and ``array_api_compat`` dispatches to the torch namespace, preserving
+autograd.
+
+Reduction convention
+--------------------
+A batch may hold frames of unequal atom count, padded to a common width. Each
+term must therefore decide how one frame's aggregate error weighs against
+another's, and the answer differs by term because frames do not carry equally
+many labels:
+
+- **Per-atom terms** (force, atomic energy, atomic prefactor force, dos,
+  tensor) carry a number of labels proportional to the frame's atom count.
+  :func:`masked_atom_mean` pools them: it divides the summed contribution of
+  the whole batch by the batch's total label count, so every real label counts
+  once and a frame's weight is proportional to its atom count.
+- **Frame-level terms** (energy, virial) carry a fixed number of labels per
+  frame. Pooling and averaging over frames coincide there, so
+  :func:`per_frame_component_mean` reduces per frame and leaves the frame axis
+  to the caller, which applies the extensive ``1 / natoms`` weighting.
+
+Pooling is what keeps a frame's weight independent of the company it keeps.
+The alternative -- averaging each frame's own per-label mean -- gives every
+frame the same weight whatever its size, which makes a label in a small frame
+count for more than one in a large frame, by the ratio of their atom counts.
+
+Writing ``S_f`` for the summed contribution of frame ``f``, ``k`` for the
+number of frames and ``d`` for the labels each of them carries, the two
+coincide whenever that count is common to the batch:
+
+    sum_f(S_f) / (k * d)  ==  (1 / k) * sum_f(S_f / d)
+
+so a batch of uniform atom count needs no special case anywhere in this
+module, and the choice between the two is unobservable there. They part
+company only where a batch holds frames of differing real atom count, which
+arises in exactly two places:
+
+- ``mix:N`` LMDB batching, which packs frames of differing atom count by
+  construction.
+- ``mixed_type`` npy data whose ``real_atom_types.npy`` spends a different
+  number of ``-1`` rows on different frames of one system. The format permits
+  this and the documentation describes it as the way to merge frames of
+  unequal atom count, but a system written by dpdata pads every frame equally
+  and is therefore unaffected.
+
+The TensorFlow backend reaches neither case: it drops ``real_natoms_vec``
+before the feed dict and normalizes by the padded width throughout.
+
+Each helper implements ONLY the masked branch. The unmasked branch of each
+caller is a plain mean over the whole batch, which pools by construction, so
+both branches express the same convention.
 """
 
 from typing import (
@@ -28,7 +69,13 @@
 
 
 def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
-    """Idiom 1: per-atom masked mean over ``ncomp`` components, averaged over frames.
+    """Idiom 1: mean of a per-atom contribution over the batch's real labels.
+
+    The contribution of every real atom is pooled across frames before the
+    division, so the reduction is a mean over labels rather than a mean over
+    frames of per-frame means. See the module docstring for why the per-atom
+    terms weigh frames by their label count, and for the identity that makes
+    this reduce to the per-frame mean on a uniform-atom-count batch.
 
     Parameters
     ----------
@@ -45,26 +92,20 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
     Returns
     -------
     Array
-        ``mean_over_frames( sum(elem * mask) / (real_natoms * ncomp) )``.
-        An all-padding frame (zero real atoms) contributes a neutral ``0``
-        instead of ``0/0 = NaN``.
+        ``sum(elem * mask) / (ncomp * sum(mask))`` over the whole batch.
+        A batch holding no real atom contributes a neutral ``0`` instead of
+        ``0/0 = NaN``.
     """
     xp = array_api_compat.array_namespace(elem, maskf)
-    nf = elem.shape[0]
-    masked = elem * maskf[:, :, None]
-    per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1)
-    per_frame_dof = xp.sum(maskf, axis=-1) * ncomp
-    # An all-padding frame has zero real atoms, so ``per_frame_dof`` is 0 and
-    # the ratio would be 0/0 = NaN -- poisoning the frame mean and, under
-    # autograd, its gradient. Divide by a safe denominator and map those frames
-    # to a neutral per-frame value of 0. Frames with real atoms are untouched,
-    # preserving the bit-identical guarantee.
-    has_dof = per_frame_dof > 0
-    safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof))
-    per_frame = xp.where(
-        has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum)
-    )
-    return xp.mean(per_frame)
+    total = xp.sum(elem * maskf[:, :, None])
+    total_dof = xp.sum(maskf) * ncomp
+    # A batch of nothing but padding has no label to average over, and the
+    # ratio would be 0/0 = NaN -- poisoning the whole batch loss and, under
+    # autograd, its gradient. The division still runs on a safe denominator so
+    # that the discarded branch stays differentiable.
+    has_dof = total_dof > 0
+    safe_dof = xp.where(has_dof, total_dof, xp.ones_like(total_dof))
+    return xp.where(has_dof, total / safe_dof, xp.zeros_like(total))
 
 
 def masked_pair_mean(elem: Array, maskf: Array, ncomp: int) -> Array:

deepmd/dpmodel/loss/tensor.py

@@ -124,7 +124,7 @@ def call(
             diff = xp.reshape(local_pred - local_label, (-1, self.tensor_size))
             diff = diff * atomic_weight
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=tensor_size).
                 maskf = xp.astype(model_dict["mask"], diff.dtype)  # [nf, natoms]
                 diff3d = xp.reshape(
                     diff, (local_pred.shape[0], natoms, self.tensor_size)
@@ -197,6 +197,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         return label_requirement

deepmd/dpmodel/model/make_model.py

@@ -48,6 +48,8 @@
     NeighborGraph,
     build_neighbor_graph,
     build_neighbor_graph_ase,
+    compact_nodes,
+    expand_node_values,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -543,38 +545,60 @@ def _call_common_graph(
                 )
             xp = array_api_compat.array_namespace(atype)
             nf, nloc = atype.shape[:2]
+            n_padded = nf * nloc
+            atype_flat = xp.reshape(atype, (n_padded,))
+            # A batch of unequal atom counts arrives padded to a common width
+            # with phantom atoms (atype < 0). The builders leave them out of
+            # every edge, so dropping them from the node axis costs nothing and
+            # spares the network from evaluating them. On a batch of uniform
+            # atom count the mask is all true and this is a renumbering by the
+            # identity.
+            ng, node_index = compact_nodes(ng, atype_flat >= 0)
             # OUTPUT-AGNOSTIC standard model dict (``<var>``, ``<var>_redu``,
             # derivative name-holders ``None``, plus int ``mask``), like the
             # dense ``call_common``.  ``call_lower_graph`` masks virtual atoms
             # (atype<0) and sets the real int mask.
             model_predict = self.call_lower_graph(
-                atype=xp.reshape(atype, (nf * nloc,)),
+                atype=xp.take(atype_flat, node_index, axis=0),
                 n_node=ng.n_node,
                 edge_index=ng.edge_index,
                 edge_vec=ng.edge_vec,
                 edge_mask=ng.edge_mask,
                 fparam=fp,
                 # graph-lower ABI: aparam is FLAT on the node axis, (N, nda).
                 aparam=(
-                    xp.reshape(ap, (nf * nloc, ap.shape[-1]))
+                    xp.take(
+                        xp.reshape(ap, (n_padded, ap.shape[-1])), node_index, axis=0
+                    )
                     if ap is not None
                     else None
                 ),
-                spin=(xp.reshape(spin, (nf * nloc, 3)) if spin is not None else None),
+                spin=(
+                    xp.take(xp.reshape(spin, (n_padded, 3)), node_index, axis=0)
+                    if spin is not None
+                    else None
+                ),
                 charge_spin=charge_spin,
             )
-            # Public ABI is rectangular (nf, nloc, *); the lower is flat
-            # (N=nf*nloc, *).  Unravel per-atom keys here at the boundary.
-            # public call_common always passes rectangular (nf,nloc) coord/atype (N == nf*nloc), so this unravel always applies; ragged graphs reach call_lower_graph/forward_common_lower_graph directly (no unravel) and stay flat (N,*).
+            # Public ABI is rectangular (nf, nloc, *); the lower is flat over
+            # the real atoms. Scatter per-atom keys back onto the padded width
+            # here at the boundary, where a phantom slot reads zero, which is
+            # what a masked-out atom contributed there before.
+            # Only the rectangular entry reaches this scatter; the ragged
+            # one keeps the flat axis its caller handed over.
+            n_real = node_index.shape[0]
             for k in list(model_predict.keys()):
                 v = model_predict[k]
                 # per-frame reduced keys (..._redu) keep their (nf, *) shape; only node-level (N,*) keys unravel — guards the nloc==1 case where N == nf.
                 if (
                     v is not None
                     and not k.endswith("_redu")
-                    and v.shape[:1] == (nf * nloc,)
+                    and v.shape[:1] == (n_real,)
                 ):
-                    model_predict[k] = xp.reshape(v, (nf, nloc, *v.shape[1:]))
+                    model_predict[k] = xp.reshape(
+                        expand_node_values(v, node_index, n_padded),
+                        (nf, nloc, *v.shape[1:]),
+                    )
             return model_predict
 
         def call_common_lower(

deepmd/dpmodel/utils/__init__.py

@@ -10,11 +10,11 @@
     PairExcludeMask,
 )
 from .lmdb_data import (
-    DistributedSameNlocBatchSampler,
+    DistributedLmdbBatchSampler,
+    LmdbBatchSampler,
     LmdbDataReader,
     LmdbTestData,
     LmdbTestDataNlocView,
-    SameNlocBatchSampler,
     is_lmdb,
     make_neighbor_stat_data,
 )
@@ -79,11 +79,12 @@
 __all__ = [
     "AtomExcludeMask",
     "DefaultNeighborList",
-    "DistributedSameNlocBatchSampler",
+    "DistributedLmdbBatchSampler",
     "EmbeddingNet",
     "EnvMat",
     "FittingNet",
     "GraphLayout",
+    "LmdbBatchSampler",
     "LmdbDataReader",
     "LmdbTestData",
     "LmdbTestDataNlocView",
@@ -93,7 +94,6 @@
     "NeighborList",
     "NetworkCollection",
     "PairExcludeMask",
-    "SameNlocBatchSampler",
     "aggregate",
     "apply_pair_exclusion_nlist",
     "build_multiple_neighbor_list",

deepmd/dpmodel/utils/batch.py

@@ -11,7 +11,18 @@
 _DROP_KEYS = {"default_mesh", "sid", "fid"}
 
 # Keys that belong to model input (everything else is label).
-_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam", "charge_spin"}
+# ``n_node`` is an input rather than a label: it states how a ragged batch's
+# flat node axis divides into frames, which the model needs to read it at all.
+_INPUT_KEYS = {
+    "coord",
+    "atype",
+    "spin",
+    "box",
+    "fparam",
+    "aparam",
+    "charge_spin",
+    "n_node",
+}
 
 
 def normalize_batch(batch: dict[str, Any]) -> dict[str, Any]: