feat(bridging): analytical short-range bridging for DPA4 and DPA4C

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -199,6 +199,28 @@ def uses_graph_lower(self) -> bool:
         """
         return False
 
+    def fused_decomposition(self) -> "tuple[Any, Any] | None":
+        """Returns the parts a fused energy-force pipeline evaluates, if any.
+
+        Such a pipeline computes one descriptor-fitting atomic model and,
+        inside the same edge scan, one analytical pair potential. A model that
+        is exactly their plain sum answers with ``(learned, pair_potential)``,
+        either entry ``None`` when absent; every other model answers ``None``
+        (the concrete default) and keeps the autograd lower.
+
+        Generic capability in the sense of :meth:`uses_graph_lower`: the model
+        layer asks instead of reaching into the atomic model for a descriptor,
+        so the answer stays correct for compositions.
+
+        Returns
+        -------
+        tuple[Any, Any] or None
+            The learned descriptor-fitting atomic model and the analytical
+            pair potential, either ``None`` when absent, or ``None`` for a
+            model that is not such a sum.
+        """
+        return None
+
     def supports_native_spin(self) -> bool:
         """Returns whether this atomic model consumes a per-atom spin input.
 

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -213,6 +213,10 @@ def uses_graph_lower(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.uses_graph_lower())
 
+    def fused_decomposition(self) -> "tuple[DPAtomicModel, None]":
+        """A descriptor-fitting model is its own learned part."""
+        return self, None
+
     def has_message_passing_across_ranks(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.has_message_passing_across_ranks())
@@ -369,19 +373,26 @@ def forward_atomic(
                 )
                 charge_spin = xp.tile(xp.reshape(cs_array, (1, -1)), (nframes, 1))
 
-        descriptor, rot_mat, g2, h2, sw = self.descriptor(
+        # A bridged descriptor returns a sixth element, the per-atom source
+        # gate of its bridging window; every other descriptor returns five.
+        outputs = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
             comm_dict=comm_dict,
             charge_spin=charge_spin if self.add_chg_spin_ebd else None,
         )
+        descriptor, rot_mat, g2, h2 = outputs[:4]
+        node_gate = outputs[5] if len(outputs) > 5 else None
         # The vacuum descriptor of every type is handed to a fitting that
-        # references its atoms; other fittings do not take the keyword.
+        # references its atoms; other fittings do not take the keyword. The
+        # source gate fades the learned part of the fitting output: an atom
+        # of a frozen pair keeps only the bias of its type.
         vacuum_kwargs = {}
         if self.fitting_net.needs_vacuum_descriptor():
             vacuum_kwargs["vacuum_descriptor"] = self.vacuum_descriptor()
+        gate_kwargs = {} if node_gate is None else {"node_gate": node_gate}
         ret = self.fitting_net(
             descriptor,
             atype,
@@ -391,6 +402,7 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
             **vacuum_kwargs,
+            **gate_kwargs,
         )
         return ret
 
@@ -401,7 +413,11 @@ def vacuum_descriptor(self) -> Array:
         conditioned as the neutral ground-state atom: zero charge with the
         ground-state multiplicity when the descriptor takes the charge/spin
         condition, and a spin vector of one Bohr magneton per unpaired
-        electron when it takes the native spin.
+        electron when it takes the native spin. A descriptor with a graph
+        lower is evaluated on the neighbor graph of those single-atom frames,
+        which holds no real edge; a dense descriptor on a neighbor list of its
+        own capacity filled with padding, since its equations are sized by
+        that capacity.
 
         Returns
         -------
@@ -417,17 +433,42 @@ def vacuum_descriptor(self) -> Array:
         atype = xp.reshape(
             xp.arange(ntypes, dtype=xp.int64, device=device), (ntypes, 1)
         )
+        conditions = {
+            name: xp.asarray(table, dtype=bias.dtype, device=device)
+            for name, table in self.vacuum_conditions().items()
+        }
+        if self.descriptor.uses_graph_lower():
+            from deepmd.dpmodel.utils.neighbor_graph import (
+                build_neighbor_graph,
+            )
+
+            graph = build_neighbor_graph(
+                coord, atype, None, self.descriptor.get_rcut(), with_csr=True
+            )
+            # Only descriptors whose ``call_graph`` declares a condition take
+            # the keyword; a declared condition the model does not use is None.
+            spin_kwargs = (
+                {"spin": conditions.get("spin")} if self.supports_native_spin() else {}
+            )
+            charge_spin_kwargs = (
+                {"charge_spin": conditions.get("charge_spin")}
+                if self.supports_charge_spin
+                else {}
+            )
+            return self.descriptor.call_graph(
+                graph,
+                xp.reshape(atype, (ntypes,)),
+                type_embedding=self.descriptor.graph_type_embedding_table(),
+                **spin_kwargs,
+                **charge_spin_kwargs,
+            )[0]
         nlist = xp.full(
             (ntypes, 1, self.descriptor.get_nnei()),
             -1,
             dtype=xp.int64,
             device=device,
         )
         mapping = xp.zeros((ntypes, 1), dtype=xp.int64, device=device)
-        conditions = {
-            name: xp.asarray(table, dtype=bias.dtype, device=device)
-            for name, table in self.vacuum_conditions().items()
-        }
         if "spin" in conditions:
             conditions["spin"] = conditions["spin"][:, None, :]
         descriptor = self.descriptor(
@@ -617,21 +658,31 @@ def forward_atomic_graph(
         charge_spin_kwargs = (
             {"charge_spin": charge_spin} if self.supports_charge_spin else {}
         )
-        gg, rot_mat = self.descriptor.call_graph(
+        # A bridged descriptor returns a third element, the per-node source
+        # gate of its bridging window; every other descriptor returns two.
+        outputs = self.descriptor.call_graph(
             graph,
             atype,
             type_embedding=type_embedding,
             comm_dict=comm_dict,
             **spin_kwargs,
             **charge_spin_kwargs,
         )
+        gg, rot_mat = outputs[0], outputs[1]
+        node_gate = outputs[2] if len(outputs) > 2 else None
         vacuum_kwargs = {}
         if vacuum_ref:
             vacuum_kwargs["vacuum_descriptor"] = gg[n_real:]
             gg = gg[:n_real]
             atype = atype[:n_real]
             rot_mat = None if rot_mat is None else rot_mat[:n_real]
+            node_gate = None if node_gate is None else node_gate[:n_real]
         # === Step 4. Fitting ===
+        # The source gate of a bridged descriptor fades the learned part of
+        # the fitting output: an atom of a frozen pair keeps only the bias of
+        # its type, so along the closing of a pair the learned energy follows
+        # the switch linearly while the analytical pair term takes over.
+        gate_kwargs = {} if node_gate is None else {"node_gate": node_gate}
         return self.fitting_net.call_graph(
             gg,
             atype,
@@ -641,6 +692,7 @@ def forward_atomic_graph(
             fparam=fparam_node,
             aparam=aparam_node,
             **vacuum_kwargs,
+            **gate_kwargs,
         )
 
     def compute_or_load_stat(