diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 25b770a5fb..df334852de 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -221,6 +221,27 @@ def fused_decomposition(self) -> "tuple[Any, Any] | None": """ return None + def has_long_range(self) -> bool: + """Whether the atomic energy carries a long-range module beside a fitting. + + Generic capability (concrete default ``False``): a model wrapping the + atomic model asks it whether the energy holds a long-range interaction, + e.g. before combining it with a scheme defined for a learned local + energy alone. + """ + return False + + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the graph forward reads beside the neighbor graph. + + Generic capability (concrete default: none) in the sense of + :meth:`uses_graph_lower`: a long-range module, whose interaction has no + cutoff, declares the topology it is evaluated on (see + :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`), and + every caller of the graph forward builds what is declared. + """ + return () + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -493,6 +514,8 @@ def forward_common_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict: """Graph analogue of :meth:`forward_common_atomic` on the flat node axis. @@ -529,6 +552,13 @@ def forward_common_atomic_graph( (and, from there, the descriptor's ``call_graph``). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_common_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name, built on the + node axis of ``graph``; ``None`` when none is required. + node_mask + flat mask of the nodes an enclosing model counts (real, owned and + not excluded there), combined with this model's own; ``None`` + imposes nothing further. N Returns ------- @@ -537,6 +567,8 @@ def forward_common_atomic_graph( """ graph, atype_clamped, output_mask = self._prepare_graph_inputs(graph, atype) + if node_mask is not None: + output_mask = output_mask & node_mask ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -545,6 +577,8 @@ def forward_common_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=output_mask, ) return self._finalize_atomic_ret(ret_dict, output_mask, atype) @@ -940,6 +974,7 @@ def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]: from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, + build_topologies, ) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -989,9 +1024,19 @@ def model_forward( self.get_rcut(), pair_excl=self.pair_excl, ) + atype_flat = xp.reshape(atype, (-1,)) + topologies = build_topologies( + self.required_topologies(), + xp.reshape(coord, (-1, 3)), + graph.n_node, + box, + node_mask=atype_flat >= 0, + atype=atype_flat, + pair_excl=self.pair_excl, + ) atomic_ret = self.forward_common_atomic_graph( graph, - xp.reshape(atype, (-1,)), + atype_flat, fparam=fparam, aparam=( xp.reshape( @@ -1003,6 +1048,7 @@ def model_forward( ), charge_spin=charge_spin, spin=None if spin is None else xp.reshape(spin, (-1, 3)), + topologies=topologies, ) # The graph route works on a flat node axis; restore the # per-frame layout the dense route returns. diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index b1e0e98af9..ad17d8bb1e 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -24,6 +24,10 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, + check_long_range_support, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, ) @@ -102,6 +106,12 @@ class DPAtomicModel(BaseAtomicModel): :math:`\mathcal{F}` is the fitting network, and :math:`\mathbf{y}^i` is the predicted atomic property (energy, dipole, etc.). + An energy model may carry a long-range module beside its fitting: it reads + the same descriptor and adds its own atomic energy to the fitting output, + so the atomic energy :math:`E^i` is the complete potential, including the + interactions without a cutoff. The module is evaluated on the + NeighborGraph route. + Parameters ---------- descriptor @@ -111,6 +121,8 @@ class DPAtomicModel(BaseAtomicModel): type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. + long_range + Long-range module added to the energy output of the fitting. """ @@ -121,12 +133,14 @@ def __init__( descriptor: BaseDescriptor, fitting: BaseFitting, type_map: list[str], + long_range: BaseLongRange | None = None, **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.long_range = long_range self.type_map = type_map self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd() # Structural capability: only descriptors with a native spin @@ -148,6 +162,8 @@ def __init__( self.add_spin_ebd: bool = ( self._supports_native_spin and self.descriptor.use_spin is not None ) + if self.long_range is not None: + self._check_long_range() super().init_out_stat() # === Reference of the isolated atoms === # The fitting references its output only when the preset fixes the @@ -155,6 +171,31 @@ def __init__( resolve_vacuum_ref(self.fitting_net, self.preset_out_bias) self.init_vacuum_conditions() + def _check_long_range(self) -> None: + """Validate that the long-range module can join this descriptor and fitting.""" + check_long_range_support(self.descriptor) + if not self.descriptor.uses_graph_lower(): + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower, " + "which this descriptor does not provide" + ) + if "energy" not in self.fitting_output_def().keys(): + raise ValueError( + "a long-range module adds to the energy output, which this " + "fitting does not have" + ) + if self.atom_exclude_types: + raise NotImplementedError( + "`atom_exclude_types` is not supported with a long-range module, " + "which attributes each pair energy to both atoms of the pair" + ) + if self.long_range.uses_charge_spin() and not self.add_chg_spin_ebd: + raise ValueError( + "the long-range module reads the total charge from " + "`charge_spin`, but the descriptor is not configured with the " + "charge/spin condition (`add_chg_spin_ebd`)" + ) + def init_vacuum_conditions(self) -> None: """Build the conditioning inputs of the reference atoms for the type map. @@ -214,20 +255,43 @@ def uses_graph_lower(self) -> bool: return bool(self.descriptor.uses_graph_lower()) def fused_decomposition(self) -> "tuple[DPAtomicModel, None]": - """A descriptor-fitting model is its own learned part.""" + """A descriptor-fitting model is its own learned part. + + The learned part carries the long-range module of the model, if any; + a fused pipeline that does not evaluate long-range modules declines a + learned part with one. + """ return self, None + def has_long_range(self) -> bool: + """Whether this model carries a long-range module.""" + return self.long_range is not None + + def required_topologies(self) -> tuple[str, ...]: + """The topologies of the long-range module.""" + if self.long_range is None: + return () + return self.long_range.required_topologies() + def has_message_passing_across_ranks(self) -> bool: """Delegates to this model's own descriptor.""" return bool(self.descriptor.has_message_passing_across_ranks()) def supports_edge_parallel(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.supports_edge_parallel()) + """Delegates to this model's own descriptor. + + A long-range module sums over complete systems, so a model carrying + one does not run under domain decomposition. + """ + return self.long_range is None and bool( + self.descriptor.supports_edge_parallel() + ) def dense_lower_supports_comm(self) -> bool: - """Delegates to this model's own descriptor.""" - return bool(self.descriptor.dense_lower_supports_comm()) + """Delegates to this model's own descriptor; a long-range module has no dense lower.""" + return self.long_range is None and bool( + self.descriptor.dense_lower_supports_comm() + ) def uses_compact_edge_pairs(self) -> bool: """Delegates to this model's own descriptor.""" @@ -267,6 +331,8 @@ def set_case_embd(self, case_idx: int) -> None: typically concatenated with the output of the descriptor and fed into the fitting net. """ self.fitting_net.set_case_embd(case_idx) + if self.long_range is not None: + self.long_range.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model @@ -298,6 +364,12 @@ def enable_compression( ) -> None: """Call descriptor enable_compression(). + A long-range module keeps its networks: it reads the output of the + compressed descriptor like that of the uncompressed one. A compressed + descriptor folds the charge/spin condition into its tables, so a + constraint that reads the total charge from the condition is folded + to the same state, the default condition of the descriptor. + Parameters ---------- min_nbor_dist @@ -318,6 +390,8 @@ def enable_compression( table_stride_2, check_frequency, ) + if self.long_range is not None and self.long_range.uses_charge_spin(): + self.long_range.fold_charge_state(self.descriptor.get_default_chg_spin()) def forward_atomic( self, @@ -358,6 +432,11 @@ def forward_atomic( the result dict, defined by the `FittingOutputDef`. """ + if self.long_range is not None: + raise NotImplementedError( + "a long-range module is evaluated on the NeighborGraph lower " + "only; the dense (nlist) lower cannot evaluate this model" + ) nframes, nloc, nnei = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) @@ -572,13 +651,16 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph analogue of :meth:`forward_atomic` on the flat node axis. Runs the descriptor ``call_graph`` then the fitting ``call_graph`` PER NODE and returns the raw fitting dict on the flat ``(N, *)`` axis (no reduction or masking; the wrapper handles those). ``fparam`` is gathered to nodes by - ``frame_id`` so each node sees its frame's parameter. + ``frame_id`` so each node sees its frame's parameter. The atomic energy + of the long-range module is added to the energy output of the fitting. Parameters ---------- @@ -604,6 +686,11 @@ def forward_atomic_graph( ``call_graph`` (the message-passing part). ``None`` for non-parallel inference (default). Mirrors :meth:`forward_atomic`'s ``comm_dict`` on the dense route. + topologies + the topologies of :meth:`required_topologies` by name. + node_mask + flat mask of the real and owned nodes, which carry the fields of + the long-range module. N Returns ------- @@ -635,6 +722,22 @@ def forward_atomic_graph( graph.n_local, atype.shape[0], ) + # The long-range module reads the frame conditions of the real frames, + # with the defaults the descriptor and fitting fall back on. + n_node = graph.n_node + if self.long_range is not None: + # The constraint and a compressed descriptor serve one charge + # state; compressing the atomic model folds both to it. + if self.long_range.uses_charge_spin() and self.get_dim_chg_spin() == 0: + raise ValueError( + "the descriptor folds the charge/spin condition into its " + "compressed tables, while the long-range charge constraint " + "still reads it; compress the atomic model, which folds " + "both to one state" + ) + fparam, charge_spin_frame = self._frame_conditions( + fparam, charge_spin, n_node.shape[0], graph.edge_vec + ) # === Step 2. Vacuum reference nodes === # Opted-in descriptors evaluate the reference in the same graph call. # Legacy descriptors must retain the dense reference used by folding. @@ -683,7 +786,7 @@ def forward_atomic_graph( # 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( + ret = self.fitting_net.call_graph( gg, atype, gr=rot_mat, @@ -694,6 +797,71 @@ def forward_atomic_graph( **vacuum_kwargs, **gate_kwargs, ) + # === Step 5. Long-range module === + # The module reads the descriptor of the real nodes; the reference + # nodes appended above take no part in it. + if self.long_range is not None: + if topologies is None or node_mask is None: + raise ValueError( + "the long-range module of this model reads the topologies " + f"{self.required_topologies()} and the node mask, which were " + "not given" + ) + ret["energy"] = ret["energy"] + self.long_range.call_graph( + gg, + atype, + node_mask, + n_node, + topologies, + fparam=fparam, + aparam=aparam_node, + charge_spin=charge_spin_frame, + ) + return ret + + def _frame_conditions( + self, + fparam: Array | None, + charge_spin: Array | None, + nf: int, + reference: Array, + ) -> tuple[Array | None, Array | None]: + """Frame parameters and charge/spin condition with their defaults resolved. + + Parameters + ---------- + fparam + Frame parameters with shape (nf, numb_fparam), or None. + charge_spin + Charge/spin condition with shape (nf, 2) or (1, 2), or None. + nf + The number of frames. + reference + An array giving the namespace, dtype and device of the defaults. + + Returns + ------- + tuple + The frame parameters with shape (nf, numb_fparam) and the + charge/spin condition with shape (nf, 2), each None when the model + takes no such input or it has no default. + """ + xp = array_api_compat.array_namespace(reference) + device = array_api_compat.device(reference) + if fparam is None and self.has_default_fparam(): + default_fparam = xp.asarray( + self.get_default_fparam(), dtype=reference.dtype, device=device + ) + fparam = xp.tile(xp.reshape(default_fparam, (1, -1)), (nf, 1)) + if not self.add_chg_spin_ebd: + return fparam, None + if charge_spin is None: + default_cs = self.descriptor.get_default_chg_spin() + if default_cs is None: + return fparam, None + charge_spin = xp.asarray(default_cs, dtype=reference.dtype, device=device) + charge_spin = xp.broadcast_to(xp.reshape(charge_spin, (-1, 2)), (nf, 2)) + return fparam, charge_spin def compute_or_load_stat( self, @@ -724,6 +892,10 @@ def compute_or_load_stat( self.fitting_net.compute_input_stats( wrapped_sampler, stat_file_path=stat_file_path ) + if self.long_range is not None: + self.long_range.compute_input_stats( + wrapped_sampler, stat_file_path=stat_file_path + ) self._collect_and_set_observed_type( wrapped_sampler, stat_file_path, preset_observed_type ) @@ -747,6 +919,8 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + if self.long_range is not None: + self.long_range.change_type_map(type_map=type_map) self.init_vacuum_conditions() def compute_fitting_input_stat( @@ -756,6 +930,9 @@ def compute_fitting_input_stat( ) -> None: """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + The networks of the long-range module take the inputs of the fitting + and receive their statistics as well. + Parameters ---------- sample_merged : Union[Callable[[], list[dict]], list[dict]] @@ -772,6 +949,12 @@ def compute_fitting_input_stat( protection=self.data_stat_protect, stat_file_path=stat_file_path, ) + if self.long_range is not None: + self.long_range.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -785,6 +968,12 @@ def serialize(self) -> dict: "fitting": self.fitting_net.serialize(), } ) + # A record with a long-range module is version 3, which a reader that + # cannot evaluate the module rejects; every other record keeps + # version 2. + if self.long_range is not None: + dd["@version"] = 3 + dd["long_range"] = self.long_range.serialize() return dd # for subclass overridden @@ -792,17 +981,27 @@ def serialize(self) -> dict: """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls: type | None = BaseLongRange + """The base long-range module class; ``None`` for a backend without one.""" @classmethod def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 2, 2) + check_version_compatibility(data.pop("@version", 1), 3, 2) data.pop("@class") data.pop("type") descriptor_obj = cls.base_descriptor_cls.deserialize(data.pop("descriptor")) fitting_obj = cls.base_fitting_cls.deserialize(data.pop("fitting")) data["descriptor"] = descriptor_obj data["fitting"] = fitting_obj + long_range = data.pop("long_range", None) + if long_range is not None: + if cls.base_long_range_cls is None: + raise NotImplementedError( + "this backend cannot evaluate the long-range module of the " + "model record" + ) + data["long_range"] = cls.base_long_range_cls.deserialize(long_range) obj = super().deserialize(data) return obj diff --git a/deepmd/dpmodel/atomic_model/inner_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py index 4d3c51f2d7..954c34ab96 100644 --- a/deepmd/dpmodel/atomic_model/inner_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -606,12 +606,15 @@ def forward_atomic_graph( charge_spin: Any = None, spin: Any = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Any = None, ) -> dict: """Evaluate the analytical per-atom energy on the flat node axis. - ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict`` are - accepted for pipeline-signature compatibility and ignored (the - analytical term conditions on geometry and types only). + ``fparam``/``aparam``/``charge_spin``/``spin``/``comm_dict``/ + ``topologies``/``node_mask`` are accepted for pipeline-signature + compatibility and ignored (the analytical term conditions on + geometry and types only; the caller masks its outputs). Parameters ---------- diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index e224cbdc2b..1dd43e0d45 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -34,6 +34,9 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .inner_potential import ( + InnerPotentialAtomicModel, +) from .pairtab_atomic_model import ( PairTabAtomicModel, ) @@ -112,6 +115,23 @@ def __init__( + ", ".join(f"{type(m).__name__}={dim_of(m)}" for m in models) ) + # A long-range module attributes each pair energy to both atoms of + # the pair. It is not combined with a bridging term, which defines the + # close-pair limit of a learned local energy, nor with atom exclusion, + # which removes atomic outputs. + if any(isinstance(m, InnerPotentialAtomicModel) for m in models) and any( + m.has_long_range() for m in models + ): + raise NotImplementedError( + "an inner_potential bridging term cannot be combined with a " + "sub-model carrying a long-range module" + ) + if self.atom_exclude_types and any(m.has_long_range() for m in models): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a sub-model carrying " + "a long-range module" + ) + self.models = models self.type_map = type_map self._rebuild_mapping_state() @@ -377,6 +397,8 @@ def compute_fitting_input_stats( for model in self.models: if isinstance(model, DPAtomicModel): model.fitting_net.compute_input_stats(sampled_func) + if model.long_range is not None: + model.long_range.compute_input_stats(sampled_func) def supports_native_spin(self) -> bool: """Spin-capable when ANY child consumes the spin input. @@ -389,6 +411,18 @@ def supports_native_spin(self) -> bool: """ return any(m.supports_native_spin() for m in self.models) + def has_long_range(self) -> bool: + """ANY child carrying a long-range module makes the composition carry one.""" + return any(m.has_long_range() for m in self.models) + + def required_topologies(self) -> tuple[str, ...]: + """Every topology a child reads, each named once.""" + return tuple( + dict.fromkeys( + name for model in self.models for name in model.required_topologies() + ) + ) + def forward_atomic_graph( self, graph: Any, @@ -398,14 +432,18 @@ def forward_atomic_graph( charge_spin: Array | None = None, spin: Array | None = None, comm_dict: dict | None = None, + topologies: dict[str, Any] | None = None, + node_mask: Array | None = None, ) -> dict[str, Array]: """Graph-route linear combination on the flat node axis. - Every child consumes the SAME graph, so on autograd backends the - shared ``graph.edge_vec`` leaf makes the summed energy's force and - virial exactly the sum of the children's -- one edge backward - covers the whole composition (this is what makes analytical - bridging terms compose with the learned model for free). + Every child consumes the SAME graph and topologies, so on autograd + backends the shared geometry leaves make the summed energy's force + and virial exactly the sum of the children's -- one backward covers + the whole composition (this is what makes analytical bridging terms + compose with the learned model for free). Each child evaluates its + complete energy, including its long-range module, before its weight + applies. Only constant weights are supported here (``"sum"``/``"mean"`` or a per-child float list); the distance-switched ZBL-interpolation @@ -429,6 +467,12 @@ def forward_atomic_graph( flat (N, 3) per-node spin, forwarded to every child. comm_dict MPI communication metadata, forwarded to every child. + topologies + the topologies of :meth:`required_topologies`, forwarded to every + child. + node_mask + flat mask of the nodes whose outputs count, forwarded to every + child, which combines it with its own exclusions. Returns ------- @@ -472,6 +516,8 @@ def forward_atomic_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, + node_mask=node_mask, ) contrib = ret["energy"] * ww energy = contrib if energy is None else energy + contrib diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 2dc880da4f..31fc811f71 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -90,7 +90,6 @@ OrderedPairFiLM, SpinChannels, build_angular_basis, - build_moment_indices, canonicalize_charge_spin, degree_offsets, derive_bispectrum_ranks, @@ -461,12 +460,9 @@ def __init__( # === Step 4. Lay out the flat moment payload === # Degree zero owns the leading `channels` entries of the flat layout, - # so the non-scalar block is exactly its complement and needs no - # separate degree index. The spin families, when present, are appended - # after the geometric degrees, leaving every geometric offset intact. - channel_index, harmonic_index = build_moment_indices(self.degree_channels) - self.angular_channel_index = channel_index[self.channels :] - self.angular_harmonic_index = harmonic_index[self.channels :] + # and every higher degree follows in its own harmonic-major block. The + # spin families, when present, are appended after the geometric + # degrees, leaving every geometric offset intact. self.degree_offsets = degree_offsets(self.degree_channels) # === Step 5. Initialize the output calibration state === @@ -1146,30 +1142,30 @@ def aggregate_moments( ) xp = array_api_compat.array_namespace(amplitude) - device = array_api_compat.device(amplitude) - channel_index = xp_asarray_nodetach( - xp, - self.angular_channel_index, - device=device, - ) - harmonic_index = xp_asarray_nodetach( - xp, - self.angular_harmonic_index, - device=device, - ) + n_edge = amplitude.shape[0] # Payload layout: [chi^2, chi^4, degree zero (C_0), degrees one and # above (S - C_0)]. Degree zero is the whole edge amplitude under its - # single envelope, while the non-scalar block gathers an amplitude and - # a harmonic per flat moment coordinate and carries a second envelope. + # single envelope. Degree l is the outer product of its 2l + 1 + # harmonics with the leading C_l channels of the amplitude, laid out + # harmonic-major and carrying a second envelope. The products are + # formed by broadcasting, so their backward reduces over the + # broadcast axes. + angular = [ + xp.reshape( + basis[:, degree * degree : (degree + 1) * (degree + 1), None] + * amplitude[:, None, :width], + (n_edge, (2 * degree + 1) * width), + ) + for degree, width in enumerate(self.degree_channels) + if degree > 0 + ] envelope_squared = envelope * envelope parts = [ envelope_squared[:, None], (envelope_squared * envelope_squared)[:, None], amplitude, - xp.take(amplitude, channel_index, axis=1) - * xp.take(basis, harmonic_index, axis=1) - * envelope[:, None], + xp.concat(angular, axis=1) * envelope[:, None], ] if spin_payload is not None: parts.append(spin_payload) diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py index 460c7970e6..8bd43d4cc0 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/__init__.py @@ -14,7 +14,6 @@ from .geometry import ( MAX_ANGULAR_DEGREE, build_angular_basis, - build_moment_indices, degree_offsets, derive_degree_channels, packed_l2_to_stf, @@ -41,7 +40,6 @@ "SpinChannels", "build_angular_basis", "build_bispectrum_layout", - "build_moment_indices", "canonicalize_charge_spin", "degree_offsets", "derive_bispectrum_ranks", diff --git a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py index 766a9d060c..a9378f4ea1 100644 --- a/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py +++ b/deepmd/dpmodel/descriptor/dpa4c_nn/geometry.py @@ -11,7 +11,6 @@ ) import array_api_compat -import numpy as np if TYPE_CHECKING: from collections.abc import ( @@ -236,36 +235,3 @@ def degree_offsets(degree_channels: Sequence[int]) -> tuple[int, ...]: for degree, width in enumerate(degree_channels): offsets.append(offsets[-1] + (2 * degree + 1) * int(width)) return tuple(offsets) - - -def build_moment_indices( - degree_channels: Sequence[int], -) -> tuple[np.ndarray, np.ndarray]: - """Build channel and harmonic indices for one flat edge payload. - - Every degree reads the leading channels of the shared radial map, so the - channel index of a degree is simply ``range(degree_channels[degree])``. - - Parameters - ---------- - degree_channels - Number of channels for degrees zero through ``lmax``. - - Returns - ------- - channel_index - Edge-amplitude channel indices with shape ``(S,)``. - harmonic_index - Packed harmonic indices with shape ``(S,)``, where - ``S = sum((2 * l + 1) * degree_channels[l])``. - """ - channel_index = [] - harmonic_index = [] - for degree, width in enumerate(degree_channels): - for component in range(2 * degree + 1): - channel_index.extend(range(int(width))) - harmonic_index.extend([degree * degree + component] * int(width)) - return ( - np.asarray(channel_index, dtype=np.int64), - np.asarray(harmonic_index, dtype=np.int64), - ) diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 37b0a1b731..126b51e721 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -14,6 +14,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .make_base_fitting import ( make_base_fitting, ) @@ -29,6 +32,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/dpmodel/fitting/latent_charge.py b/deepmd/dpmodel/fitting/latent_charge.py new file mode 100644 index 0000000000..4967879f1f --- /dev/null +++ b/deepmd/dpmodel/fitting/latent_charge.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .invar_fitting import ( + InvarFitting, +) + + +@fitting_check_output +class LatentChargeFitting(InvarFitting): + r"""Per-atom latent charges predicted from the invariant descriptor. + + The network maps the descriptor of an atom, together with the frame + parameters, atomic parameters and case embedding of the energy fitting it + accompanies, to ``dim_out`` latent charge channels in units of the + elementary charge :math:`e`. The charges are an intermediate field of a + long-range module rather than a model output: the module constrains and + pairs them, and the potential energy is the only quantity they are + trained through. + + The constructor arguments are those of :class:`InvarFitting` with the + output variable fixed to ``latent_charge``. + """ + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def serialize(self) -> dict: + data = super().serialize() + data["type"] = "latent_charge" + return data diff --git a/deepmd/dpmodel/long_range/__init__.py b/deepmd/dpmodel/long_range/__init__.py new file mode 100644 index 0000000000..4a4aad0d7c --- /dev/null +++ b/deepmd/dpmodel/long_range/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range modules of the atomic energy of a descriptor-fitting model. + +A long-range module is a component of the atomic energy model beside its +energy fitting: it reads the descriptor features shared with that fitting and +the topologies it declares, and contributes an energy per atom. Every consumer +of the atomic model therefore evaluates the complete potential. +""" + +from .base_long_range import ( + BaseLongRange, + check_long_range_support, + make_base_long_range, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "COULOMB_CONSTANT", + "BaseInteractionKernel", + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", + "check_long_range_support", + "make_base_long_range", +] diff --git a/deepmd/dpmodel/long_range/base_long_range.py b/deepmd/dpmodel/long_range/base_long_range.py new file mode 100644 index 0000000000..f37b821333 --- /dev/null +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, + abstractmethod, +) +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.plugin import ( + PluginVariant, + make_plugin_registry, +) + + +def check_long_range_support(descriptor: Any) -> None: + """Admit only the descriptors whose combination with long-range modules is validated. + + A long-range module reads nothing but the invariant node features every + graph descriptor returns, so this gate is a statement of validation, not + of implementation: a further descriptor is admitted by adding it here once + its combination has been tested. + + Parameters + ---------- + descriptor + The descriptor of the atomic model receiving the long-range module. + + Raises + ------ + NotImplementedError + If the descriptor is not among the validated ones. + """ + from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, + ) + + if not isinstance(descriptor, DescrptDPA4C): + raise NotImplementedError( + "long-range modules are currently validated with the DPA4C " + f"descriptor only, got {type(descriptor).__name__}" + ) + + +def make_base_long_range(t_tensor: Any) -> type: + """Make the base class of the long-range modules of one backend. + + Each backend owns its registry, so a configuration builds the long-range + module native to the backend that reads it. + + Parameters + ---------- + t_tensor + The type of the tensor, used in the type hints. + """ + + class BLR(ABC, PluginVariant, make_plugin_registry("long range")): + """Long-range contribution to the atomic energy of a descriptor-fitting model. + + A long-range module reads the invariant node features of the + descriptor and the topologies it declares, and returns an energy per + node in eV. The atomic model adds it to the output of its energy + fitting before the output statistics and masks, so every consumer of + the atomic model sees one complete potential. The module owns the + attribution of its energy to the nodes; the per-frame energy is their + sum. + """ + + @abstractmethod + def required_topologies(self) -> tuple[str, ...]: + """Names of the topologies the module reads beside the neighbor graph.""" + + @abstractmethod + def uses_charge_spin(self) -> bool: + """Whether the module reads the frame-level charge/spin condition.""" + + @abstractmethod + def call_graph( + self, + descriptor: t_tensor, + atype: t_tensor, + node_mask: t_tensor, + n_node: t_tensor, + topologies: dict[str, Any], + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, + charge_spin: t_tensor | None = None, + ) -> t_tensor: + """Evaluate the atomic energy of the module on the flat node axis. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any valid + type. + node_mask + Nodes carrying fields with shape ``(N,)``: the real and owned + nodes. + n_node + Per-frame node counts with shape ``(nf,)``. + topologies + The topologies of :meth:`required_topologies` by name. + fparam + Frame parameters with shape ``(nf, numb_fparam)``, defaults + resolved. + aparam + Atomic parameters with shape ``(N, numb_aparam)``. + charge_spin + Frame charge/spin condition with shape ``(nf, 2)``, defaults + resolved, when the model takes one. + + Returns + ------- + t_tensor + The atomic energy with shape ``(N, 1)``, in eV. + """ + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the networks of the module.""" + + def change_type_map(self, type_map: list[str]) -> None: + """Remap the type-dependent state of the module to a new type map.""" + + def set_case_embd(self, case_idx: int) -> None: + """Select the case embedding of the networks of the module.""" + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the module to one frame charge/spin condition. + + Called when the descriptor of the model is compressed, which folds + the condition into its tables, for a module whose + :meth:`uses_charge_spin` holds. + + Raises + ------ + NotImplementedError + If the module cannot be specialized. + """ + raise NotImplementedError( + f"{type(self).__name__} cannot be specialized to a charge/spin " + "condition, which the compression of the model requires" + ) + + def charge_state_constants(self) -> list[t_tensor]: + """The tensors that carry the folded charge/spin condition. + + A deployment that serves another condition writes the outputs of + :meth:`rebuild_charge_state` over them, in this order. + """ + return [] + + def rebuild_charge_state(self, charge_spin: t_tensor) -> list[t_tensor]: + """Build the tensors of :meth:`charge_state_constants` for a condition. + + Parameters + ---------- + charge_spin + The condition with shape ``(1, 2)``. + """ + return [] + + @abstractmethod + def serialize(self) -> dict: + """Serialize the module to a dict.""" + + @classmethod + def deserialize(cls, data: dict) -> "BLR": + """Deserialize a module, dispatching on its registered type.""" + if cls is BLR: + return BLR.get_class_by_type(data["type"]).deserialize(data) + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + + return BLR + + +BaseLongRange = make_base_long_range(np.ndarray) diff --git a/deepmd/dpmodel/long_range/kernel.py b/deepmd/dpmodel/long_range/kernel.py new file mode 100644 index 0000000000..3486b663ce --- /dev/null +++ b/deepmd/dpmodel/long_range/kernel.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Pair interaction kernels of the latent-charge long-range modules. + +A kernel is a function :math:`K(r)` of the distance between two charges, in +Å⁻¹, so that charges in units of :math:`e` interact with the energy +:math:`k_e q_i q_j K(r_{ij})` in eV, where :math:`k_e` is the Coulomb +constant. A kernel owns its parameters; widths are stored through their +logarithm, which keeps them positive however training moves them. +""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_asarray_nodetach, + xp_erf, +) +from deepmd.dpmodel.common import ( + PRECISION_DICT, + NativeOP, + to_numpy_array, +) +from deepmd.utils.plugin import ( + make_plugin_registry, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +COULOMB_CONSTANT = 14.3996454784255 +r"""Coulomb constant :math:`e^2 / (4 \pi \varepsilon_0)` in eV Å.""" + + +class BaseInteractionKernel(NativeOP, make_plugin_registry("interaction kernel")): + """Pair interaction kernel :math:`K(r)` in Å⁻¹ with its own parameters. + + Parameters + ---------- + trainable + Whether the parameters of the kernel are trained. + precision + The precision in which the parameters are stored; the kernel is + evaluated in the precision of its input. + """ + + PARAMETER_NAMES: tuple[str, ...] = () + """Attributes holding the trainable parameters of the kernel.""" + + def __init__(self, trainable: bool = True, precision: str = "default") -> None: + self.trainable = bool(trainable) + self.precision = precision + self.prec = PRECISION_DICT[precision.lower()] + + def call(self, r_sq: Array) -> Array: + """Evaluate the kernel on squared pair distances. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.evaluate(r_sq, self.coefficients(r_sq)) + + def coefficients(self, like: Array) -> dict[str, Array]: + """The coefficients of the kernel formula, derived from its parameters. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the coefficients. + + Returns + ------- + dict[str, Array] + Every coefficient with its own shape, differentiable with respect + to the parameters. + """ + raise NotImplementedError + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + """Evaluate the kernel formula with given coefficients. + + Parameters + ---------- + r_sq + Squared distances with shape ``(P,)``, in Ų. + coefficients + The coefficients of :meth:`coefficients`, each with its own shape + or with a leading pair axis of length ``P``. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + raise NotImplementedError + + def _parameter(self, name: str, like: Array) -> Array: + """A parameter in the namespace, dtype and device of ``like``, autograd kept.""" + xp = array_api_compat.array_namespace(like) + return xp_asarray_nodetach( + xp, + getattr(self, name), + dtype=like.dtype, + device=array_api_compat.device(like), + ) + + def serialize(self) -> dict: + return { + "@class": "InteractionKernel", + "type": self.TYPE, + "@version": 1, + "trainable": self.trainable, + "precision": self.precision, + "@variables": { + name: to_numpy_array(getattr(self, name)) + for name in self.PARAMETER_NAMES + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "BaseInteractionKernel": + if cls is BaseInteractionKernel: + return cls.get_class_by_type(data["type"]).deserialize(data) + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + data.pop("type") + variables = data.pop("@variables") + obj = cls(**data, **cls._shape_kwargs(variables)) + for name in cls.PARAMETER_NAMES: + setattr(obj, name, np.asarray(variables[name], dtype=obj.prec)) + return obj + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + """Constructor arguments that size the parameters stored in ``variables``.""" + return {} + + +# Taylor coefficients of erf(x) / x in powers of x^2, 2 / sqrt(pi) (-1)^n / +# (n! (2n + 1)), and the square of x below which LESKernel evaluates them: there +# the nine terms reach the float64 rounding of erf(x) / x. Above it the closed +# form keeps the relative error of its radial derivative within 2e-14 in +# float64 and 1e-5 in float32. +_ERF_OVER_X_SERIES = ( + 1.1283791670955126, + -0.37612638903183754, + 0.11283791670955126, + -0.026866170645131252, + 0.005223977625442188, + -0.0008548327023450852, + 0.00012055332981789664, + -1.492565035840625e-05, + 1.6462114365889246e-06, +) +_ERF_SERIES_X_SQ = 0.0625 + + +@BaseInteractionKernel.register("les") +class LESKernel(BaseInteractionKernel): + r"""Latent Ewald summation kernel :math:`K(r) = \mathrm{erf}(\alpha r) / r`. + + The kernel is the interaction of two Gaussian charge clouds of width + :math:`1/(\sqrt{2}\alpha)`: it reaches :math:`1/r` beyond a few + :math:`1/\alpha` and stays finite, :math:`2\alpha/\sqrt{\pi}`, at contact. + + Parameters + ---------- + alpha + The inverse width :math:`\alpha > 0`, in Å⁻¹. + trainable + Whether :math:`\alpha` is trained. + precision + The precision in which :math:`\log \alpha` is stored. + """ + + TYPE = "les" + PARAMETER_NAMES = ("log_alpha",) + + def __init__( + self, alpha: float = 1.0, trainable: bool = True, precision: str = "default" + ) -> None: + super().__init__(trainable, precision) + if not np.isfinite(alpha) or alpha <= 0.0: + raise ValueError(f"`alpha` must be finite and positive, got {alpha}") + self.log_alpha = np.array(np.log(alpha), dtype=self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return {"alpha": xp.exp(self._parameter("log_alpha", like))} + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + alpha = coefficients["alpha"] + x_sq = alpha**2 * r_sq + # K = alpha erf(x) / x with x = alpha r. The derivative of the closed + # form subtracts two terms that agree to second order in x, so near + # contact erf(x) / x comes from its Taylor series in x^2. Both are + # evaluated on every pair, which keeps the shapes static; each reads a + # value of its own domain where the other is selected, which keeps the + # derivatives of both finite. + near = x_sq < _ERF_SERIES_X_SQ + x_far = xp.sqrt(xp.where(near, xp.ones_like(x_sq), x_sq)) + x_sq_near = xp.where(near, x_sq, xp.zeros_like(x_sq)) + series = xp.full_like(x_sq, _ERF_OVER_X_SERIES[-1]) + for coefficient in reversed(_ERF_OVER_X_SERIES[:-1]): + series = series * x_sq_near + coefficient + return alpha * xp.where(near, series, xp_erf(x_far) / x_far) + + +@BaseInteractionKernel.register("sog") +class SOGKernel(BaseInteractionKernel): + r"""Sum-of-Gaussians kernel :math:`K(r) = \sum_k a_k \exp(-r^2 / 2\sigma_k^2)`. + + Parameters + ---------- + amplitude + The amplitudes :math:`a_k`, in Å⁻¹, of either sign: a scalar shared by + every Gaussian or one value per Gaussian. Required with ``width``; + with the default widths it defaults to + :math:`a_k = \sqrt{2/\pi} \ln b / \sigma_k`, the trapezoidal + discretization in :math:`\ln s` of + :math:`1/r = \sqrt{2/\pi} \int_0^\infty e^{-r^2/2s^2} s^{-2} \, ds`. + The default kernel then follows :math:`1/r`, less a nearly constant + offset, from a few :math:`\sigma` to a fraction of + :math:`\sigma b^{M-1}`, and stays finite at contact. The offset + exerts no force between fixed charges; through charges that respond + to the positions it does. + width + The widths :math:`\sigma_k > 0`, in Å. Defaults to the geometric + sequence :math:`\sigma_k = \sigma b^k`, :math:`k = 0, \dots, M - 1`. + n_gaussian + The number :math:`M` of Gaussians of the default widths. + base_width + The width :math:`\sigma` of the first default Gaussian, in Å. + width_ratio + The ratio :math:`b > 1` of consecutive default widths. + trainable + Whether the amplitudes and widths are trained. + precision + The precision in which the amplitudes and log-widths are stored. + """ + + TYPE = "sog" + PARAMETER_NAMES = ("amplitude", "log_width") + + def __init__( + self, + amplitude: float | list[float] | None = None, + width: list[float] | None = None, + n_gaussian: int = 12, + base_width: float = 2.180230445405648, + width_ratio: float = 1.62976708826776469, + trainable: bool = True, + precision: str = "default", + ) -> None: + super().__init__(trainable, precision) + if width is None: + if n_gaussian < 1: + raise ValueError(f"`n_gaussian` must be positive, got {n_gaussian}") + if not np.isfinite(base_width) or base_width <= 0.0: + raise ValueError(f"`base_width` must be positive, got {base_width}") + if not np.isfinite(width_ratio) or width_ratio <= 1.0: + raise ValueError( + f"`width_ratio` must be greater than one, got {width_ratio}" + ) + width = base_width * np.power( + width_ratio, np.arange(n_gaussian, dtype=np.float64) + ) + if amplitude is None: + amplitude = np.sqrt(2.0 / np.pi) * np.log(width_ratio) / width + elif amplitude is None: + raise ValueError("`amplitude` is required with explicit `width`") + width = np.asarray(width, dtype=np.float64).reshape(-1) + if width.size < 1 or not np.all(np.isfinite(width)) or np.any(width <= 0.0): + raise ValueError("`width` must hold finite positive values") + amplitude = np.asarray(amplitude, dtype=np.float64).reshape(-1) + if amplitude.size == 1: + amplitude = np.full_like(width, amplitude[0]) + if amplitude.shape != width.shape or not np.all(np.isfinite(amplitude)): + raise ValueError( + "`amplitude` must be a finite scalar or hold one value per Gaussian" + ) + self.amplitude = amplitude.astype(self.prec) + self.log_width = np.log(width).astype(self.prec) + + def coefficients(self, like: Array) -> dict[str, Array]: + xp = array_api_compat.array_namespace(like) + return { + "amplitude": self._parameter("amplitude", like), + "inv_two_width_sq": 0.5 * xp.exp(-2.0 * self._parameter("log_width", like)), + } + + def evaluate(self, r_sq: Array, coefficients: dict[str, Array]) -> Array: + xp = array_api_compat.array_namespace(r_sq) + gaussian = xp.exp(-r_sq[:, None] * coefficients["inv_two_width_sq"]) + return xp.sum(coefficients["amplitude"] * gaussian, axis=-1) + + @classmethod + def _shape_kwargs(cls, variables: dict[str, np.ndarray]) -> dict[str, Any]: + return {"n_gaussian": int(np.asarray(variables["log_width"]).size)} diff --git a/deepmd/dpmodel/long_range/latent_charge.py b/deepmd/dpmodel/long_range/latent_charge.py new file mode 100644 index 0000000000..4cc6f3ccd1 --- /dev/null +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.common import ( + GLOBAL_NP_FLOAT_PRECISION, + NativeOP, + to_numpy_array, +) +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + frame_id_from_n_node, + segment_sum, +) +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + +CHARGE_TARGET_SOURCES = ("neutral", "charge_spin", "fparam") +"""Sources of the total charge a constrained channel sums to.""" + +CHARGE_NET_NEURON = [64, 64] +"""Default hidden widths of the charge network.""" + +CHARGE_NET_SEED_INDEX = 10_000 +"""Child index of the fitting seed that seeds a charge network configured without a seed. + +The networks of the fitting and of the descriptor draw from small child +indices of their seeds, so the charge network draws a stream of its own even +when the fitting and descriptor seeds coincide. +""" + + +@BaseLongRange.register("latent_charge") +class LatentChargeLongRange(NativeOP, BaseLongRange): + r"""Energy of interacting latent charges in complete non-periodic systems. + + A charge network predicts :math:`C` latent charge channels + :math:`q_i^c` per atom, in units of :math:`e`, from the descriptor shared + with the energy fitting; padding and atoms owned by another rank carry no + charge. The energy is the direct sum over the pairs of atoms of each frame + listed by the :data:`~deepmd.dpmodel.utils.neighbor_graph.ALL_PAIRS` + topology, each unordered pair once, + + .. math:: + E = k_e \sum_{i None: + if charge_net.vacuum_ref: + raise ValueError( + "the charge network reads the descriptor of the real atoms only " + "and cannot reference isolated atoms (`vacuum_ref`)" + ) + self.charge_net = charge_net + self.kernel = kernel + self.charge_constraint = self._normalize_constraint( + charge_constraint, charge_net + ) + # The frame charge/spin condition a compressed model is specialized + # to; see :meth:`fold_charge_state`. + self.charge_state = None + + @staticmethod + def _normalize_constraint( + constraint: dict | None, charge_net: LatentChargeFitting + ) -> dict | None: + if constraint is None: + return None + source = str(constraint.get("source", "neutral")) + channel = int(constraint.get("channel", 0)) + fparam_index = int(constraint.get("fparam_index", 0)) + if source not in CHARGE_TARGET_SOURCES: + raise ValueError( + f"charge constraint source must be one of {CHARGE_TARGET_SOURCES}, " + f"got {source!r}" + ) + if not 0 <= channel < charge_net.dim_out: + raise ValueError( + f"constrained channel {channel} is out of range for " + f"{charge_net.dim_out} latent charge channel(s)" + ) + if source == "fparam" and not 0 <= fparam_index < charge_net.numb_fparam: + raise ValueError( + f"`fparam_index` {fparam_index} is out of range for " + f"{charge_net.numb_fparam} frame parameter(s)" + ) + return {"source": source, "channel": channel, "fparam_index": fparam_index} + + @classmethod + def from_config( + cls, + config: dict, + descriptor: Any, + fitting: Any, + fitting_seed: int | list[int] | None = None, + ) -> "LatentChargeLongRange": + """Build the module from its configuration. + + The charge network reads the descriptor and takes the conditioning + inputs of the energy fitting: its frame and atomic parameters, their + defaults and its case embedding. It follows the activation function, + skip connection and precision of the fitting unless its configuration + sets them, and its seed defaults to a child of the fitting seed. The + kernel parameters are stored in the precision of the charge network. + + Parameters + ---------- + config + The module configuration: ``channels``, ``charge_net``, ``kernel`` + and ``charge_constraint``. + descriptor + The descriptor of the model. + fitting + The energy fitting of the model. + fitting_seed + The seed of the energy fitting. + """ + charge_net_config = { + "neuron": CHARGE_NET_NEURON, + "activation_function": fitting.activation_function, + "resnet_dt": fitting.resnet_dt, + "precision": fitting.precision, + "seed": child_seed(fitting_seed, CHARGE_NET_SEED_INDEX), + **config.get("charge_net", {}), + } + kernel_config = { + "precision": charge_net_config["precision"], + **config["kernel"], + } + kernel_cls = BaseInteractionKernel.get_class_by_type(kernel_config.pop("type")) + charge_net = LatentChargeFitting( + var_name="latent_charge", + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + dim_out=int(config.get("channels", 1)), + numb_fparam=fitting.get_dim_fparam(), + numb_aparam=fitting.get_dim_aparam(), + dim_case_embd=fitting.dim_case_embd, + default_fparam=fitting.get_default_fparam(), + use_aparam_as_mask=fitting.use_aparam_as_mask, + mixed_types=descriptor.mixed_types(), + type_map=fitting.get_type_map(), + **charge_net_config, + ) + return cls( + charge_net, + kernel_cls(**kernel_config), + config.get("charge_constraint"), + ) + + def required_topologies(self) -> tuple[str, ...]: + return (ALL_PAIRS,) + + def uses_charge_spin(self) -> bool: + # A folded state stands in for the condition. + return ( + self.charge_constraint is not None + and self.charge_constraint["source"] == "charge_spin" + and self.charge_state is None + ) + + def fold_charge_state(self, charge_spin: list[float]) -> None: + """Specialize the constraint to one frame charge/spin condition. + + A compressed descriptor folds the condition into its tables and reads + none at run time, so a constraint that reads its total charge from the + condition takes it from this state instead. The state is a constant + of the deployed model that a deployment rewrites, together with the + tables of the descriptor, when it serves another condition. + + Parameters + ---------- + charge_spin + The condition ``[charge, spin multiplicity]``. + """ + self.charge_state = np.asarray( + charge_spin, dtype=GLOBAL_NP_FLOAT_PRECISION + ).reshape(1, 2) + + def charge_state_constants(self) -> list[Array]: + return [] if self.charge_state is None else [self.charge_state] + + def rebuild_charge_state(self, charge_spin: Array) -> list[Array]: + if self.charge_state is None: + return [] + xp = array_api_compat.array_namespace(charge_spin) + state = xp.astype(charge_spin, self.charge_state.dtype) + return [xp.reshape(state, (1, 2))] + + def call(self, *args: Any, **kwargs: Any) -> Array: + return self.call_graph(*args, **kwargs) + + def call_graph( + self, + descriptor: Array, + atype: Array, + node_mask: Array, + n_node: Array, + topologies: dict[str, Any], + fparam: Array | None = None, + aparam: Array | None = None, + charge_spin: Array | None = None, + ) -> Array: + xp = array_api_compat.array_namespace(descriptor) + n = descriptor.shape[0] + pairs = topologies[ALL_PAIRS] + frame_id = frame_id_from_n_node(n_node, n_total=n) + fparam_node = None if fparam is None else xp.take(fparam, frame_id, axis=0) + charge = self.charge_net.call_graph( + descriptor, atype, fparam=fparam_node, aparam=aparam + )["latent_charge"] + # The pair sum cancels between charges of opposite sign, so it runs in + # the precision of the geometry. + charge = xp.astype(charge, pairs.vec.dtype) + charge = xp.where(node_mask[:, None], charge, xp.zeros_like(charge)) + if self.charge_constraint is not None: + charge = self._constrain( + charge, node_mask, frame_id, n_node, fparam, charge_spin + ) + src = pairs.index[0] + dst = pairs.index[1] + pair_charge = xp.sum( + xp.take(charge, src, axis=0) * xp.take(charge, dst, axis=0), axis=-1 + ) + r_sq = xp.sum(pairs.vec * pairs.vec, axis=-1) + pair_energy = pair_charge * self._pair_kernel(r_sq, src, n) + half_energy = (0.5 * COULOMB_CONSTANT) * xp.where( + pairs.mask, pair_energy, xp.zeros_like(r_sq) + ) + energy = segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + return xp.reshape(energy, (n, 1)) + + def _pair_kernel(self, r_sq: Array, src: Array, n: int) -> Array: + """Kernel values of the pairs. + + The pairs are located on the node axis for a backend that brings the + kernel coefficients to the pairs through it. + + Parameters + ---------- + r_sq + Squared pair distances with shape ``(P,)``, in Ų. + src + First node of every pair with shape ``(P,)``. + n + The number of nodes. + + Returns + ------- + Array + The kernel values with shape ``(P,)``, in Å⁻¹. + """ + return self.kernel(r_sq) + + def _constrain( + self, + charge: Array, + node_mask: Array, + frame_id: Array, + n_node: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """Shift the constrained channel so that it sums to the frame total charge.""" + xp = array_api_compat.array_namespace(charge) + nf = n_node.shape[0] + channel = self.charge_constraint["channel"] + target = self._total_charge(nf, charge, fparam, charge_spin) + weight = xp.astype(node_mask, charge.dtype) + count = segment_sum(weight, frame_id, nf) + total = segment_sum(charge[:, channel], frame_id, nf) + # A frame without a charged atom has nothing to shift. + shift = (total - target) / xp.maximum(count, xp.ones_like(count)) + constrained = charge[:, channel] - xp.take(shift, frame_id, axis=0) * weight + return xp.concat( + [charge[:, :channel], constrained[:, None], charge[:, channel + 1 :]], + axis=-1, + ) + + def _total_charge( + self, + nf: int, + like: Array, + fparam: Array | None, + charge_spin: Array | None, + ) -> Array: + """The total charge of every frame, which the constrained channel sums to. + + Parameters + ---------- + nf + The number of frames. + like + An array giving the namespace, dtype and device of the result. + fparam, charge_spin + As in :meth:`call_graph`. + + Returns + ------- + Array + The total charges with shape ``(nf,)``, in units of :math:`e`. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + source = self.charge_constraint["source"] + if source == "charge_spin" and self.charge_state is not None: + state = xp.asarray(self.charge_state, device=device) + target = xp.broadcast_to(state[:, 0], (nf,)) + elif source == "charge_spin": + if charge_spin is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`charge_spin`, which was not given" + ) + target = charge_spin[:, 0] + elif source == "fparam": + if fparam is None: + raise ValueError( + "the charge constraint reads the total charge from " + "`fparam`, which was not given" + ) + target = fparam[:, self.charge_constraint["fparam_index"]] + else: + target = xp.zeros((nf,), dtype=like.dtype, device=device) + return xp.astype(target, like.dtype) + + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + self.charge_net.compute_input_stats( + merged, protection=protection, stat_file_path=stat_file_path + ) + + def change_type_map(self, type_map: list[str]) -> None: + self.charge_net.change_type_map(type_map=type_map) + + def set_case_embd(self, case_idx: int) -> None: + self.charge_net.set_case_embd(case_idx) + + def serialize(self) -> dict: + data = { + "@class": "LongRange", + "type": "latent_charge", + "@version": 1, + "charge_net": self.charge_net.serialize(), + "kernel": self.kernel.serialize(), + "charge_constraint": self.charge_constraint, + } + if self.charge_state is not None: + data["charge_state"] = to_numpy_array(self.charge_state).tolist() + return data + + @classmethod + def deserialize(cls, data: dict) -> "LatentChargeLongRange": + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + obj = cls( + LatentChargeFitting.deserialize(data["charge_net"]), + BaseInteractionKernel.deserialize(data["kernel"]), + data["charge_constraint"], + ) + if data.get("charge_state") is not None: + obj.fold_charge_state(np.asarray(data["charge_state"]).reshape(-1).tolist()) + return obj diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 7f3f8b9705..b5549d4adb 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -73,3 +73,5 @@ def compute_fitting_input_stats( The lazy sampled function to get data frames from different data systems. """ self.get_fitting_net().compute_input_stats(sampled_func) + if self.atomic_model.has_long_range(): + self.atomic_model.long_range.compute_input_stats(sampled_func) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b74fb6b697..f05d1d9e09 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -46,8 +46,10 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, + PairGraph, build_neighbor_graph, build_neighbor_graph_ase, + build_topologies, compact_nodes, expand_node_values, ) @@ -564,12 +566,21 @@ def _call_common_graph( # atom count the mask is all true and this is a renumbering by the # identity. ng, node_index = compact_nodes(ng, atype_flat >= 0) + atype_real = xp.take(atype_flat, node_index, axis=0) + topologies = build_topologies( + self.atomic_model.required_topologies(), + xp.take(xp.reshape(cc, (n_padded, 3)), node_index, axis=0), + ng.n_node, + bb, + atype=atype_real, + pair_excl=pair_excl, + ) # OUTPUT-AGNOSTIC standard model dict (````, ``_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.take(atype_flat, node_index, axis=0), + atype=atype_real, n_node=ng.n_node, edge_index=ng.edge_index, edge_vec=ng.edge_vec, @@ -589,6 +600,7 @@ def _call_common_graph( else None ), charge_spin=charge_spin, + topologies=topologies, ) # Public ABI is rectangular (nf, nloc, *); the lower is flat over # the real atoms. Scatter per-atom keys back onto the padded width @@ -731,6 +743,7 @@ def forward_common_atomic_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Model-level graph forward (no type cast). Analogue of the dense :meth:`forward_common_atomic`. @@ -773,6 +786,9 @@ def forward_common_atomic_graph( Per-node spin vectors, flat (N, 3), or ``None``. Forwarded unchanged to the atomic model's ``forward_common_atomic_graph`` (and, from there, the descriptor's ``call_graph``). + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -795,6 +811,7 @@ def forward_common_atomic_graph( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) return fit_output_to_model_output_graph( atomic_ret, @@ -817,6 +834,7 @@ def call_common_lower_graph( comm_dict: dict | None = None, charge_spin: Array | None = None, spin: Array | None = None, + topologies: dict[str, Any] | None = None, ) -> dict[str, Array]: """Graph-native PUBLIC lower (dpa1/se_atten concat-tebd, attention included). @@ -858,6 +876,9 @@ def call_common_lower_graph( Per-node spin vectors, flat (N, 3), or ``None``. Cast to the model precision alongside the other node inputs and forwarded unchanged to :meth:`forward_common_atomic_graph`. + topologies + The topologies the atomic model declares by + ``required_topologies()``, on the node axis of the graph. Returns ------- @@ -871,6 +892,15 @@ def call_common_lower_graph( charge_spin=charge_spin, spin=spin, ) + if topologies: + # The pair vectors are geometry like the edge vectors. + xp = array_api_compat.array_namespace(edge_vec) + topologies = { + name: PairGraph( + pairs.index, xp.astype(pairs.vec, edge_vec.dtype), pairs.mask + ) + for name, pairs in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, @@ -883,6 +913,7 @@ def call_common_lower_graph( comm_dict=comm_dict, charge_spin=cs, spin=sp, + topologies=topologies, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 66d32f8f8b..21c3be799b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.long_range import ( + BaseLongRange, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -47,6 +50,7 @@ atomic_model=DPAtomicModel, pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 38a4f4b9e4..118b09a75f 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -65,6 +65,83 @@ def get_model_components( return descriptor, fitting, fitting_type +def get_long_range( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> Any | None: + """Construct the long-range module a model configuration adds to its fitting. + + The module reads the descriptor of the model, and its networks take the + conditioning inputs of the energy fitting they accompany. + + Parameters + ---------- + data : dict + The model configuration, with the optional ``long_range`` section. + descriptor : Any + The descriptor of the model. + fitting : Any + The energy fitting of the model. + long_range_base : type or None + Backend long-range registry base class; ``None`` for a backend + without long-range modules. + backend_name : str + Backend name used in error messages. + + Returns + ------- + Any or None + The backend-native long-range module, or ``None`` when none is + configured. + + Raises + ------ + NotImplementedError + If a long-range module is configured for a backend without them. + """ + config = data.get("long_range") + if config is None: + return None + if long_range_base is None: + raise NotImplementedError( + f"long-range modules are not supported in the {backend_name} backend" + ) + config = copy.deepcopy(config) + long_range_cls = long_range_base.get_class_by_type(config.pop("type")) + return long_range_cls.from_config( + config, + descriptor, + fitting, + fitting_seed=data.get("fitting_net", {}).get("seed"), + ) + + +def _long_range_kwargs( + data: dict, + descriptor: Any, + fitting: Any, + *, + long_range_base: type | None, + backend_name: str, +) -> dict[str, Any]: + """The ``long_range`` constructor argument, absent when none is configured. + + Model classes without a long-range module keep their constructor contract. + """ + long_range = get_long_range( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ) + return {} if long_range is None else {"long_range": long_range} + + def get_standard_model( data: dict, *, @@ -72,6 +149,7 @@ def get_standard_model( fitting_base: type, model_base: type, backend_name: str, + long_range_base: type | None = None, ) -> Any: """Construct a standard model through backend registries.""" descriptor, fitting, fitting_type = get_model_components( @@ -89,6 +167,13 @@ def get_standard_model( atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=data.get("pair_exclude_types", []), preset_out_bias=data.get("preset_out_bias"), + **_long_range_kwargs( + data, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) @@ -103,6 +188,11 @@ def get_zbl_model( backend_name: str, ) -> Any: """Construct a ZBL model from backend-native atomic model classes.""" + if data.get("long_range") is not None: + raise NotImplementedError( + "a long-range module is not supported with the tabulated ZBL " + "potential (`use_srtab`)" + ) descriptor, fitting, fitting_type = get_model_components( data, descriptor_base=descriptor_base, @@ -141,6 +231,7 @@ def get_linear_atomic_model( pairtab_model: type, linear_atomic_model: type | None = None, descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, + long_range_base: type | None = None, ) -> Any: """Build the ``LinearEnergyAtomicModel`` composition from a config. @@ -181,6 +272,9 @@ def get_linear_atomic_model( back to the generic registry build. Backends use it to route family-specific model types (e.g. DPA4/SeZM) through their validated builders. + long_range_base : type, optional + Backend long-range registry base class, which builds the long-range + module of a child; ``None`` for a backend without long-range modules. Raises ------ @@ -188,8 +282,11 @@ def get_linear_atomic_model( If more than one ``inner_potential`` child is given, if an ``inner_potential`` child has no unique learned sibling, if a bridged composition does not use ``weights: "sum"``, if a child - carries a ``bridging_method`` flag, or if a child is of an + carries a ``bridging_method`` flag, if a child with a long-range + module carries its own ``pair_exclude_types``, or if a child is of an unsupported kind. + NotImplementedError + If a child with a long-range module carries ``atom_exclude_types``. """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, @@ -311,6 +408,24 @@ def get_linear_atomic_model( if descriptor_child_builder is not None: child = descriptor_child_builder(sub) if child is None: + # A long-range child is checked as the same model alone: the + # pair exclusion of the composition shapes the graph and the + # pair topology the children share, and atom exclusion is + # unsupported with a long-range module. + if sub.get("long_range") is not None: + if sub.get("pair_exclude_types"): + raise ValueError( + "a linear_ener sub-model with a long-range module " + "takes `pair_exclude_types` from the composition, " + "whose sub-models share one neighbor graph and one " + "pair topology; set it on the composition" + ) + if sub.get("atom_exclude_types"): + raise NotImplementedError( + "`atom_exclude_types` is not supported with a " + "long-range module, which attributes each pair " + "energy to both atoms of the pair" + ) descriptor, fitting, _ = get_model_components( sub, descriptor_base=descriptor_base, @@ -325,6 +440,13 @@ def get_linear_atomic_model( fitting, type_map=sub["type_map"], preset_out_bias=preset_out_bias, + **_long_range_kwargs( + sub, + descriptor, + fitting, + long_range_base=long_range_base, + backend_name=backend_name, + ), ) built[i] = child else: @@ -467,6 +589,7 @@ def __init__( pairtab_model: type | None = None, zbl_model: type | None = None, linear_atomic_model: type | None = None, + long_range_base: type | None = None, ) -> None: """Store backend-native classes used by all model construction paths.""" self.descriptor_base = descriptor_base @@ -477,6 +600,7 @@ def __init__( self.pairtab_model = pairtab_model self.zbl_model = zbl_model self.linear_atomic_model = linear_atomic_model + self.long_range_base = long_range_base def get_model_components(self, data: dict) -> tuple[Any, Any, str]: """Construct descriptor and fitting objects for this backend.""" @@ -495,6 +619,17 @@ def get_standard_model(self, data: dict) -> Any: fitting_base=self.fitting_base, model_base=self.model_base, backend_name=self.backend_name, + long_range_base=self.long_range_base, + ) + + def get_long_range(self, data: dict, descriptor: Any, fitting: Any) -> Any | None: + """Construct the long-range module of a model config for this backend.""" + return get_long_range( + data, + descriptor, + fitting, + long_range_base=self.long_range_base, + backend_name=self.backend_name, ) def get_linear_atomic_model( @@ -515,6 +650,7 @@ def get_linear_atomic_model( pairtab_model=self.pairtab_model, linear_atomic_model=self.linear_atomic_model, descriptor_child_builder=descriptor_child_builder, + long_range_base=self.long_range_base, ) def get_zbl_model(self, data: dict) -> Any: diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py index 5df83d91df..6fdd5744d6 100644 --- a/deepmd/dpmodel/model/native_spin_model.py +++ b/deepmd/dpmodel/model/native_spin_model.py @@ -79,6 +79,10 @@ class NSM(T_Model, NativeSpinModelKind): def __init__(self, *args: Any, spin: Spin, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + if self.atomic_model.has_long_range(): + raise NotImplementedError( + "native spin is not supported for a model with a long-range module" + ) self.spin = spin self.ntypes_real = self.spin.ntypes_real # Per-real-type 0/1 spin gate, derived from ``use_spin`` and hence diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8678e51398..2d86dd0698 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -118,6 +118,11 @@ def __init__( spin: Spin, ) -> None: super().__init__() + if backbone_model.atomic_model.has_long_range(): + raise NotImplementedError( + "the virtual-atom spin scheme is not supported for a model with " + "a long-range module" + ) self.backbone_model = backbone_model self.spin = spin # Spin graph-lower unsupported: carry-all routing diverges on diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 6895ce3e1d..8779648f15 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -5,8 +5,10 @@ ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity + edge padding), ``csr`` (backend-agnostic CSR construction and canonicalization), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + the -``from_dense_quartet`` legacy converter), ``segment`` (mask-aware -segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly). +``from_dense_quartet`` legacy converter), ``pair_graph`` (the complete +intra-frame pair topology of interactions without a cutoff), ``segment`` +(mask-aware segment-reduction toolkit), and ``derivatives`` (edge and pair +force/virial assembly). See the design discussion wanghan-iapcm/deepmd-kit#4. """ @@ -33,6 +35,7 @@ ) from .derivatives import ( edge_force_virial, + pair_force_virial, ) from .env import ( edge_env_mat, @@ -53,6 +56,16 @@ pad_and_guard_angles, pad_and_guard_edges, ) +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + build_topologies, + check_complete_frames, + flatten_topologies, + unflatten_topologies, + unordered_pair_index, +) from .pairs import ( center_edge_pairs, ) @@ -64,8 +77,10 @@ ) __all__ = [ + "ALL_PAIRS", "GraphLayout", "NeighborGraph", + "PairGraph", "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", @@ -77,12 +92,16 @@ "build_edge_csr", "build_neighbor_graph", "build_neighbor_graph_ase", + "build_pair_graph", + "build_topologies", "canonicalize_neighbor_graph", "center_edge_pairs", + "check_complete_frames", "compact_nodes", "edge_env_mat", "edge_force_virial", "expand_node_values", + "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", @@ -92,8 +111,11 @@ "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", + "pair_force_virial", "segment_max", "segment_mean", "segment_softmax", "segment_sum", + "unflatten_topologies", + "unordered_pair_index", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 09111fa5d5..5e3f34ec04 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -118,3 +118,71 @@ def edge_force_virial( node_frame = frame_id_from_n_node(n_node, n_total=n_out) # (N,) frame per node virial = segment_sum(atom_virial, node_frame, nf) # (nf, 3, 3) return force, atom_virial, virial + + +def pair_force_virial( + g_p: Array, + pair_vec: Array, + pair_index: Array, + pair_mask: Array, + n_node: Array, + node_capacity: int | None = None, +) -> tuple[Array, Array, Array]: + """Assemble per-node force/atom-virial and PER-FRAME virial of a pair set. + + A pair set lists every unordered pair once (see + :class:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.PairGraph`), so its + per-pair gradient ``g_p = dE/d(pair_vec)`` acts on both endpoints: + ``pair_vec = r_src - r_dst`` gives ``F_dst += g_p`` and ``F_src -= g_p``, + as for an edge. The pair virial ``w = -g_p (x) pair_vec`` is split evenly + between the two endpoints, which keeps the atomic virial independent of + the order in which a pair is listed; its per-frame sum is the sum of + ``w`` over the pairs of that frame. + + Parameters + ---------- + g_p + (P, 3) per-pair gradient ``dE/d(pair_vec)``. + pair_vec + (P, 3) per-pair displacement ``r_src - r_dst``. + pair_index + (2, P) ``[src, dst]`` node endpoints of each pair. + pair_mask + (P,) boolean real-pair mask; guard pairs (``False``) are zeroed + before any scatter. + n_node + (nf,) per-frame node counts. + node_capacity + Size of the node axis ``N``; ``None`` takes ``sum(n_node)``. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, half of each pair virial per endpoint. + virial + (nf, 3, 3) per-frame virial. + """ + xp = array_api_compat.array_namespace(g_p) + n_out = node_capacity if node_capacity is not None else int(xp.sum(n_node)) + nf = n_node.shape[0] + if isinstance(n_out, int) and n_out == 0: + device = array_api_compat.device(g_p) + return ( + xp.zeros((0, 3), dtype=g_p.dtype, device=device), + xp.zeros((0, 3, 3), dtype=g_p.dtype, device=device), + xp.zeros((nf, 3, 3), dtype=g_p.dtype, device=device), + ) + g = g_p * xp.astype(pair_mask[:, None], g_p.dtype) + # As for edges, the modulo bounds the guard endpoints under export. + src = pair_index[0] % n_out + dst = pair_index[1] % n_out + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + w_half = -0.5 * (g[:, :, None] * pair_vec[:, None, :]) # (P, 3, 3) + atom_virial = segment_sum(w_half, src, n_out) + segment_sum(w_half, dst, n_out) + # Both endpoints of a pair share a frame, so reducing the per-node virial + # over nodes equals reducing the pair virial over pairs. + node_frame = frame_id_from_n_node(n_node, n_total=n_out) + virial = segment_sum(atom_virial, node_frame, nf) + return force, atom_virial, virial diff --git a/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py new file mode 100644 index 0000000000..4e9c653128 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Complete intra-frame pair topology for interactions without a cutoff. + +A module whose interaction acts between every two atoms of a system, such as +a non-periodic long-range module, cannot read its geometry from the +short-range :class:`~deepmd.dpmodel.utils.neighbor_graph.NeighborGraph`. It +declares the :data:`ALL_PAIRS` topology instead: every unordered pair of valid +nodes sharing a frame, listed once with ``src < dst``. The pairs are generated +from the per-frame node counts alone, so their storage is the true pair count +``sum_f n_f (n_f - 1) / 2`` with neither a distance search nor padding to the +widest frame. Pairs of types excluded by the model-level pair exclusion are +left out, as they are left out of the neighbor graph. Like the edge axis of +the neighbor graph, the pair axis ends in masked guard pairs, which keep its +length a dynamic dimension of its own however few real pairs a batch holds. + +Geometry enters only through :attr:`PairGraph.vec`, which an autograd backend +makes a leaf next to the short-range ``edge_vec``; its gradient is pulled back +to the nodes by :func:`~deepmd.dpmodel.utils.neighbor_graph.derivatives.pair_force_virial`. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat + +from .graph import ( + frame_id_from_n_node, + pad_and_guard_edges, +) +from .segment import ( + segment_sum, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + +ALL_PAIRS = "all_pairs" +"""Name of the topology holding every unordered pair of valid nodes of a frame.""" + + +@dataclass +class PairGraph: + """Unordered pairs of nodes on the flat node axis of a NeighborGraph.""" + + index: Array + """(2, P) int64 [src, dst] with src < dst; both endpoints in one frame.""" + vec: Array + """(P, 3) float r_src - r_dst; the geometry leaf of the pair interaction.""" + mask: Array + """(P,) bool real (1) vs guard (0) pairs; guards join node 0 to itself.""" + + +def unordered_pair_index(n_node: Array, node_mask: Array | None = None) -> Array: + """Enumerate the unordered pairs of valid nodes within every frame. + + Within a frame of ``n`` valid nodes the pairs follow the row-major order + of the strict upper triangle, ``(0, 1), (0, 2), ..., (0, n-1), (1, 2), ...``. + Row ``i`` starts at rank ``S(i) = i (2n - 1 - i) / 2``; the row of rank + ``k`` is the floor of the smaller root of ``S(i) = k``, evaluated in + float64 and corrected by one integer step against rounding. + + The pair count is data dependent, so the enumeration runs eagerly. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + node_mask + Validity of the flat nodes with shape ``(N,)``; ``None`` takes every + node as valid. Invalid nodes appear in no pair. + + Returns + ------- + Array + The ``[src, dst]`` node indices with shape ``(2, P)``, int64. + """ + xp = array_api_compat.array_namespace(n_node) + device = array_api_compat.device(n_node) + n_node = xp.astype(n_node, xp.int64) + if node_mask is None: + count = n_node + else: + frame_id = frame_id_from_n_node(n_node, n_total=node_mask.shape[0]) + count = segment_sum(xp.astype(node_mask, xp.int64), frame_id, n_node.shape[0]) + n_pair = count * (count - 1) // 2 + total = int(xp.sum(n_pair)) + if total == 0: + return xp.zeros((2, 0), dtype=xp.int64, device=device) + pair_end = xp.cumulative_sum(n_pair) + rank = xp.arange(total, dtype=xp.int64, device=device) + pair_frame = xp.astype(xp.searchsorted(pair_end, rank, side="right"), xp.int64) + local = rank - xp.take(pair_end - n_pair, pair_frame, axis=0) + n = xp.take(count, pair_frame, axis=0) + + def row_start(row: Array) -> Array: + return row * (2 * n - 1 - row) // 2 + + b = xp.astype(2 * n - 1, xp.float64) + discriminant = b * b - 8.0 * xp.astype(local, xp.float64) + discriminant = xp.where( + discriminant > 0.0, discriminant, xp.zeros_like(discriminant) + ) + row = xp.astype(xp.floor((b - xp.sqrt(discriminant)) / 2.0), xp.int64) + row = xp.where(local < row_start(row), row - 1, row) + row = xp.where(local >= row_start(row + 1), row + 1, row) + col = row + 1 + (local - row_start(row)) + first = xp.take(xp.cumulative_sum(count) - count, pair_frame, axis=0) + src = first + row + dst = first + col + if node_mask is not None: + # The ranks above count valid nodes only; map them to node positions. + position = xp.reshape(xp.nonzero(node_mask)[0], (-1,)) + position = xp.astype(position, xp.int64) + src = xp.take(position, src, axis=0) + dst = xp.take(position, dst, axis=0) + return xp.stack([src, dst]) + + +def build_pair_graph( + coord: Array, + n_node: Array, + atype: Array, + node_mask: Array | None = None, + pair_excl: PairExcludeMask | None = None, +) -> PairGraph: + """Build the :data:`ALL_PAIRS` topology of flat node coordinates. + + Parameters + ---------- + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + A differentiable input keeps :attr:`PairGraph.vec` differentiable. + n_node + Per-frame node counts with shape ``(nf,)``. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + pair_excl + The model-level pair exclusion, or ``None``; pairs of excluded types + are left out. + + Returns + ------- + PairGraph + The unordered pairs of valid nodes and their displacements, followed + by two masked guard pairs. + """ + xp = array_api_compat.array_namespace(coord) + index = unordered_pair_index(n_node, node_mask) + if pair_excl is not None: + keep = pair_excl.build_edge_exclude_mask(index, atype) + index = xp.take(index, xp.reshape(xp.nonzero(keep)[0], (-1,)), axis=1) + vec = xp.take(coord, index[0], axis=0) - xp.take(coord, index[1], axis=0) + index, vec, mask = pad_and_guard_edges(index, vec, None) + return PairGraph(index=index, vec=vec, mask=mask) + + +def build_topologies( + names: Iterable[str], + coord: Array, + n_node: Array, + box: Array | None = None, + node_mask: Array | None = None, + *, + atype: Array, + pair_excl: PairExcludeMask | None = None, +) -> dict[str, PairGraph]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + Parameters + ---------- + names + Topology names from ``required_topologies()`` of the atomic model. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells, or ``None`` for non-periodic systems. + node_mask + Validity of the nodes with shape ``(N,)``, or ``None`` when every + node is valid. + atype + Node types with shape ``(N,)``; invalid nodes carry any type. + pair_excl + The model-level pair exclusion, or ``None``. + + Returns + ------- + dict[str, PairGraph] + Each requested topology under its name; empty when none is required. + + Raises + ------ + NotImplementedError + If a cell is given: the pair topology describes complete non-periodic + systems only. + ValueError + If a name is unknown. + """ + topologies = {} + for name in names: + if name != ALL_PAIRS: + raise ValueError(f"unknown topology {name!r}; expected {ALL_PAIRS!r}") + check_complete_frames(box) + topologies[name] = build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + return topologies + + +def check_complete_frames(box: Array | None) -> None: + """Reject a cell, since the all-pairs topology covers non-periodic frames. + + Parameters + ---------- + box + Simulation cells, or ``None`` for non-periodic systems. + + Raises + ------ + NotImplementedError + If a cell is given. + """ + if box is not None: + raise NotImplementedError( + "the all-pairs topology of the long-range modules describes " + "complete non-periodic systems only; periodic cells are not " + "supported (pass box=None for a non-periodic system)" + ) + + +def flatten_topologies(topologies: dict[str, PairGraph]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives every topology as three + inputs, ``(index, vec, mask)``, in the order of the dict, after its other + inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + ``(index, vec, mask)`` of every topology, concatenated. + """ + return tuple( + value + for pairs in topologies.values() + for value in (pairs.index, pairs.vec, pairs.mask) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, PairGraph]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + ``(index, vec, mask)`` of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph] + The topologies by name. + + Raises + ------ + ValueError + If the number of values is not three per name. + """ + if len(values) != 3 * len(names): + raise ValueError( + f"expected three inputs per topology for {tuple(names)}, got {len(values)}" + ) + return { + name: PairGraph(*values[3 * ii : 3 * ii + 3]) for ii, name in enumerate(names) + } diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index 8db31d6d5d..2ef1dc5eb9 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -40,6 +40,8 @@ class jax_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" def forward_common_atomic( self, diff --git a/deepmd/pd/model/model/__init__.py b/deepmd/pd/model/model/__init__.py index 59dce4f89b..94897bcfd1 100644 --- a/deepmd/pd/model/model/__init__.py +++ b/deepmd/pd/model/model/__init__.py @@ -64,6 +64,11 @@ def _get_standard_model_components( def get_standard_model(model_params: dict) -> BaseModel: + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d2e0a09e96..b65604c40c 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -208,6 +208,11 @@ def get_linear_model(model_params: dict) -> BaseModel: "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) + if sub.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) if any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -414,6 +419,11 @@ def get_standard_model(model_params: dict) -> BaseModel: "expands the flag), or spell the composition explicitly with " '`type: "linear_ener"` and an `inner_potential` sub-model.' ) + if model_params.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) diff --git a/deepmd/pt_expt/descriptor/dpa4c.py b/deepmd/pt_expt/descriptor/dpa4c.py index a6482fbaa6..fb4495ab11 100644 --- a/deepmd/pt_expt/descriptor/dpa4c.py +++ b/deepmd/pt_expt/descriptor/dpa4c.py @@ -158,6 +158,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.use_amp_infer = use_amp_infer() self._apply_autocast_policy() + def _load_from_state_dict( + self, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, + ) -> None: + # Back-compat: checkpoints written by DeePMD-kit 3.2 hold the gather + # tables of the angular moments, which are formed by broadcasting; + # the stale entries are dropped so that the strict load succeeds. + for name in ("angular_channel_index", "angular_harmonic_index"): + state_dict.pop(prefix + name, None) + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + def call_graph( self, graph: Any, diff --git a/deepmd/pt_expt/fitting/__init__.py b/deepmd/pt_expt/fitting/__init__.py index 63bcc6c2ce..2f42a57d48 100644 --- a/deepmd/pt_expt/fitting/__init__.py +++ b/deepmd/pt_expt/fitting/__init__.py @@ -17,6 +17,9 @@ from .invar_fitting import ( InvarFitting, ) +from .latent_charge import ( + LatentChargeFitting, +) from .polarizability_fitting import ( PolarFitting, ) @@ -33,6 +36,7 @@ "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "LatentChargeFitting", "PolarFitting", "PropertyFittingNet", "SeZMEnergyFittingNet", diff --git a/deepmd/pt_expt/fitting/latent_charge.py b/deepmd/pt_expt/fitting/latent_charge.py new file mode 100644 index 0000000000..f4f38d2f0a --- /dev/null +++ b/deepmd/pt_expt/fitting/latent_charge.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting as LatentChargeFittingDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) + + +@torch_module +class LatentChargeFitting(LatentChargeFittingDP): + """Per-atom latent charges predicted from the invariant descriptor (pt_expt).""" diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..abee40c07c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,12 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + check_complete_frames, + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -664,6 +670,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) use_graph_lower = model_uses_graph_lower(model) + topology_names = model.atomic_model.required_topologies() self._dpmodel = model self._is_spin = ( model_params.get("type") == "spin_ener" or "spin" in model_params @@ -718,6 +725,8 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } + if topology_names: + self.metadata["topologies"] = list(topology_names) if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -735,7 +744,8 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # ``communicate_extended_output``, while the graph runner emits local # public keys consumed directly by ``_eval_model_graph``. # - # Graph: (..., source_row_ptr, fparam, aparam, charge_spin) + # Graph: (..., source_row_ptr, fparam, aparam, charge_spin, + # *topology inputs) # Graph spin: (..., source_row_ptr, spin, fparam, aparam, charge_spin) # Nlist: (ext_coord, ext_atype, nlist, mapping, fparam, aparam) # Nlist spin: (ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam) @@ -805,6 +815,7 @@ def _eager_runner_graph( fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None, + *topology_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: model_ret = model.forward_common_lower_graph( atype, @@ -822,6 +833,7 @@ def _eager_runner_graph( fparam=fparam, aparam=aparam, charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), ) return _translate_energy_keys( model_ret, @@ -2387,9 +2399,10 @@ def _eval_model_graph( positional schema ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, source_row_ptr, - fparam, aparam, charge_spin)`` to the graph forward. Exported AOTI - artifacts use a dynamic edge axis, so no ``edge_capacity`` padding is - needed. The ``graph_edge_dtype`` metadata selects float32 geometry for + fparam, aparam, charge_spin)`` to the graph forward, followed by the + ``(index, vec, mask)`` of every topology the ``topologies`` metadata + names. Exported AOTI artifacts use dynamic edge and pair axes, so no + ``edge_capacity`` padding is needed. The ``graph_edge_dtype`` metadata selects float32 geometry for compressed DPA1 and float64 for generic graph descriptors. The forward returns local public keys directly, so results are reshaped without ``communicate_extended_output``. @@ -2503,6 +2516,17 @@ def _eval_model_graph( compact.source_row_ptr, compact.source_order, ) + if self.metadata.get("topologies"): + # The compact lower of a model with a long-range module reads + # the positions of complete non-periodic frames, and the zero + # communicator of a frame held whole. + check_complete_frames(box_input) + model_inputs += ( + torch.as_tensor( + coord_input.reshape(-1, 3), dtype=torch.float64, device=DEVICE + ), + torch.zeros(1, dtype=torch.int64, device="cpu"), + ) else: fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE @@ -2528,6 +2552,23 @@ def _eval_model_graph( aparam_t, charge_spin_t, ) + topology_names = self.metadata.get("topologies", ()) + if topology_names: + # Padding slots of a mixed-type batch are no atoms and join + # no pair. + model_inputs += flatten_topologies( + build_topologies( + topology_names, + torch.as_tensor( + coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE + ), + n_node_t, + box_input, + node_mask=atype_t >= 0, + atype=atype_t, + pair_excl=self._model_pair_excl(), + ) + ) if self._is_pt2: model_ret = self._pt2_runner(*model_inputs) else: diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical.py b/deepmd/pt_expt/kernels/dpa4c/canonical.py index 750bb31b25..1aaf2e1010 100644 --- a/deepmd/pt_expt/kernels/dpa4c/canonical.py +++ b/deepmd/pt_expt/kernels/dpa4c/canonical.py @@ -58,6 +58,21 @@ def canonical_model_eligible(model: Any) -> bool: return False if getattr(atomic_model, "atom_excl", None) is not None: return False + # A latent-charge long-range module joins the pipeline through the node + # positions, which its direct-sum operator reads in place of a pair list; + # that pipeline evaluates no pair potential. + if parts[0].has_long_range(): + from deepmd.pt_expt.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + long_range = parts[0].long_range + if not ( + isinstance(long_range, LatentChargeLongRange) + and long_range.direct_sum_eligible() + and parts[1] is None + ): + return False from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( mega_eligible, ) diff --git a/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py new file mode 100644 index 0000000000..69e5edb332 --- /dev/null +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact canonical deployment of compressed DPA4C with a long-range module.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +if TYPE_CHECKING: + from deepmd.pt_expt.utils.canonical_graph import ( + CanonicalGraph, + ) + + +def dpa4c_canonical_long_range_energy_force( + descriptor: Any, + fitting: Any, + long_range: Any, + graph: CanonicalGraph, + atype: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None, + ownership: torch.Tensor, + atom_bias: torch.Tensor, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate compressed DPA4C and its long-range module from a compact graph. + + The descriptor operator runs once, and its output feeds both the fused + energy fitting and the charge network of the long-range module. The + derivative of the energy with respect to the descriptor collects the + fitting term and the response of the charges, so a single descriptor + backward yields the edge gradient of both; the long-range force at fixed + charges is added to the assembled force. + + Parameters + ---------- + descriptor + Compressed pt_expt DPA4C descriptor without a spin condition. + fitting + Eligible pt_expt energy fitting network. + long_range + Long-range module whose ``direct_sum_eligible()`` holds. + graph + Canonical graph whose ``source`` has shape ``(E,)`` in ``uint32`` and + whose ``destination_row_ptr`` has shape ``(N + 1,)`` in ``int64``. + atype + Flat node atom types with shape ``(N,)``. + coord + Node positions with shape ``(N, 3)`` in Å; every frame is a complete + non-periodic system. + communicator + The address of the MPI_Comm handle of a frame decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero + or ``None`` for a single domain. Across domains the owned nodes of + every rank make up the frame. + ownership + Boolean mask selecting energy-contributing nodes with shape ``(N,)``. + atom_bias + Combined atomic energy bias with shape ``(ntypes,)`` in eV. + do_atomic_virial + Whether to return per-node virials. + + Returns + ------- + energy + Per-frame energy with shape ``(F, 1)`` in eV, fp64. + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, fp64. + force + Per-node force with shape ``(N, 3)`` in eV/Å, fp32. + virial + Per-frame virial with shape ``(F, 3, 3)`` in eV, fp32. + atom_virial + Per-node virial with shape ``(N, 3, 3)`` in eV, or an empty tensor. + """ + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + ensure_registered, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + compressed_operator_arguments, + fitting_energy_and_gradient, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + canonical_edge_force_virial, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + ensure_registered as ensure_force_registered, + ) + from deepmd.pt_expt.kernels.edge_force_virial import ( + frame_scalar_sum, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + ensure_registered as ensure_fitting_registered, + ) + + ensure_registered() + ensure_fitting_registered() + ensure_force_registered() + operator_args = ( + graph.source, + graph.destination_row_ptr, + atype, + *compressed_operator_arguments(descriptor), + int(descriptor.lmax), + *descriptor._compression_scalars, + ) + node_descriptor, state = torch.ops.deepmd.dpa4c_canonical_compress( + graph.edge_vec, *operator_args + ) + # The charge network reads the descriptor before the fused fitting + # overwrites it with its own cotangent. + ( + long_range_energy, + long_range_gradient, + long_range_force, + long_range_atom_virial, + long_range_virial, + ) = long_range.energy_force_from_coord( + node_descriptor, atype, ownership, graph.n_node, coord, communicator + ) + atom_energy, descriptor_gradient = fitting_energy_and_gradient( + fitting, node_descriptor, atype, ownership, atom_bias + ) + edge_gradient, _, _ = torch.ops.deepmd.dpa4c_canonical_compress_backward_inplace( + descriptor_gradient + long_range_gradient, + state, + graph.edge_vec, + *operator_args, + ) + force, atom_virial, virial, _ = canonical_edge_force_virial( + edge_gradient, + graph.edge_vec, + graph.destination_row_ptr, + graph.source_row_ptr, + graph.source_order, + graph.n_node, + graph.edge_vec.new_empty(0), + atype.shape[0], + do_atomic_virial, + ) + # Nodes outside the ownership carry no charge and hence no long-range + # energy. + atom_energy = atom_energy + long_range_energy + energy = frame_scalar_sum(atom_energy, graph.n_node) + virial = virial + long_range_virial.reshape(-1, 3, 3) + if do_atomic_virial: + atom_virial = atom_virial + long_range_atom_virial.reshape(-1, 3, 3) + return energy, atom_energy, force + long_range_force, virial, atom_virial diff --git a/deepmd/pt_expt/kernels/long_range.py b/deepmd/pt_expt/kernels/long_range.py new file mode 100644 index 0000000000..994b78134f --- /dev/null +++ b/deepmd/pt_expt/kernels/long_range.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the direct-sum operator of the latent-charge long range. + +The operator ``deepmd::latent_charge_direct_sum`` evaluates, for the nodes of +complete frames, the potential that the charges of every other node of the +frame create at each node and the force between the charges, without listing +the pairs: its memory is linear in the node count. The CUDA kernel is +``source/op/pt/long_range/direct_sum.cu``; the CPU kernel registered here is +the reference it is tested against. + +The operator serves inference. It applies the total-charge constraint of the +module, and returns the atomic energy, its derivative with respect to the +charges the charge network predicts, and the force and virial at fixed +charges; the caller adds the response of the charges by back-propagating the +charge gradient through the charge network. + +Given the MPI communicator of a system decomposed into domains, one per rank, +the CUDA kernel evaluates the charged nodes of its rank against those of +every rank; the CPU reference evaluates a single domain. +""" + +import math +from functools import ( + cache, +) +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + frame_id_from_n_node, + unordered_pair_index, +) +from deepmd.pt_expt.kernels.utils import ( + operator_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +__all__ = [ + "MAX_CHANNELS", + "MAX_GAUSSIANS", + "kernel_eligible", + "kernel_operator_arguments", + "latent_charge_direct_sum", + "op_available", +] + +#: Charge channels and Gaussians the CUDA kernel is instantiated for. +MAX_CHANNELS = 8 +MAX_GAUSSIANS = 32 + +_LES = 0 +_SOG = 1 + + +def op_available() -> bool: + """Whether the backend device carries the direct-sum operator. + + The CPU kernel is the reference of this module, which is registered here + so that the answer does not depend on earlier calls. + """ + # Importing the module loads the operator library. + import deepmd.pt.cxx_op # noqa: F401 + + ensure_registered() + return operator_available("latent_charge_direct_sum") + + +def kernel_eligible(kernel: Any) -> bool: + """Whether the operator evaluates this interaction kernel.""" + if isinstance(kernel, LESKernel): + return True + return isinstance(kernel, SOGKernel) and kernel.amplitude.numel() <= MAX_GAUSSIANS + + +def kernel_operator_arguments(kernel: Any) -> tuple[int, torch.Tensor]: + """Convert an interaction kernel into the code and coefficients of the operator. + + Parameters + ---------- + kernel + A pt_expt interaction kernel accepted by :func:`kernel_eligible`. + + Returns + ------- + tuple[int, torch.Tensor] + The kernel code and its float32 coefficients: ``alpha`` for the erf + kernel; the amplitudes ``a_k`` and then the rates ``c_k = log2(e) / + (2 sigma_k^2)`` of ``a_k 2^(-c_k r^2)`` for the sum of Gaussians. + + Raises + ------ + ValueError + If the operator does not evaluate the kernel. + """ + if not kernel_eligible(kernel): + raise ValueError( + f"the direct-sum operator does not evaluate {type(kernel).__name__}" + ) + if isinstance(kernel, LESKernel): + coefficient = torch.exp(kernel.log_alpha).reshape(1) + return _LES, coefficient.detach().to(torch.float32).contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width) + coefficient = torch.cat((kernel.amplitude, rate)) + return _SOG, coefficient.detach().to(torch.float32).contiguous() + + +def _kernel_and_slope( + r_sq: torch.Tensor, coefficient: torch.Tensor, kernel: int +) -> tuple[torch.Tensor, torch.Tensor]: + """K(r) and K'(r) / r at the squared distances ``r_sq``.""" + coefficient = coefficient.to(dtype=r_sq.dtype, device=r_sq.device) + if kernel == _LES: + alpha = coefficient[0] + # Coincident nodes, which carry no weight, are evaluated at a short + # distance, where the closed form stays finite. + r_sq = torch.clamp(r_sq, min=1e-12 / alpha**2) + r = torch.sqrt(r_sq) + value = torch.erf(alpha * r) / r + contact = (2.0 / math.sqrt(math.pi)) * alpha + slope = (contact * torch.exp(-(alpha**2) * r_sq) - value) / r_sq + return value, slope + amplitude, rate = coefficient.reshape(2, -1) + weight = torch.exp2(-r_sq[:, None] * rate) + return weight @ amplitude, weight @ (-2.0 * math.log(2.0) * rate * amplitude) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator over the explicit pairs of every frame, in float64.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_direct_sum evaluates a single " + "domain; a communicator of several ranks requires the CUDA kernel" + ) + n = coord.shape[0] + frame_count = n_node_per_frame.shape[0] + # Padding nodes belong to no frame and carry no charge. + frame = frame_id_from_n_node(n_node_per_frame, n_total=n) + real = torch.arange(n, dtype=torch.int64, device=coord.device) < torch.sum( + n_node_per_frame + ) + weight = (node_mask & real).to(torch.float64) + q = charge.to(torch.float64) * weight[:, None] + if channel >= 0: + columns = torch.arange(q.shape[1], dtype=torch.int64, device=coord.device) + selector = (columns == channel).to(torch.float64) + count = coord.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) + total = coord.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + shift = (total - total_charge.to(torch.float64)) / count + q = q - (shift[frame] * weight)[:, None] * selector + pairs = unordered_pair_index(n_node_per_frame) + src, dst = pairs[0], pairs[1] + d = coord[src] - coord[dst] + value, slope = _kernel_and_slope(torch.sum(d * d, dim=-1), coefficient, kernel) + potential = coord.new_zeros(q.shape) + potential = potential.index_add(0, src, q[dst] * value[:, None]) + potential = prefactor * potential.index_add(0, dst, q[src] * value[:, None]) + strength = torch.sum(q[src] * q[dst], dim=-1) * slope + pair_force = -strength[:, None] * d + force = torch.zeros_like(coord).index_add(0, src, pair_force) + force = force.index_add(0, dst, -pair_force) + half = (-0.5 * strength)[:, None, None] * d[:, :, None] * d[:, None, :] + atom_virial = coord.new_zeros(n, 3, 3).index_add(0, src, half) + atom_virial = prefactor * atom_virial.index_add(0, dst, half).reshape(n, 9) + virial = coord.new_zeros(frame_count, 9).index_add(0, frame, atom_virial) + atom_energy = 0.5 * torch.sum(q * potential, dim=-1, keepdim=True) + gradient = potential * weight[:, None] + if channel >= 0: + mean = coord.new_zeros(frame_count).index_add(0, frame, gradient[:, channel]) + gradient = gradient - ((mean / count)[frame] * weight)[:, None] * selector + return ( + atom_energy, + gradient.to(charge.dtype), + (prefactor * force).to(charge.dtype), + atom_virial.to(charge.dtype), + virial.to(charge.dtype), + ) + + +def _fake( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + n = coord.shape[0] + return ( + coord.new_empty(n, 1), + torch.empty_like(charge), + charge.new_empty(n, 3), + charge.new_empty(n, 9), + charge.new_empty(n_node_per_frame.shape[0], 9), + ) + + +@cache +def _register() -> None: + torch.library.register_fake("deepmd::latent_charge_direct_sum")(_fake) + library = torch.library.Library("deepmd", "IMPL") + library.impl("latent_charge_direct_sum", _reference, "CPU") + # The registration lives as long as the library object. + _register.library = library + + +def ensure_registered() -> None: + """Register the fake and CPU implementations when the schema is loaded.""" + if isinstance( + getattr(torch.ops.deepmd, "latent_charge_direct_sum", None), + torch._ops.OpOverloadPacket, + ): + _register() + + +def latent_charge_direct_sum( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + kernel: Any, + channel: int, + prefactor: float, + communicator: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Energy, charge gradient, force and virial of the charges of complete frames. + + Parameters + ---------- + coord + Node positions with shape ``(N, 3)`` in Å, float64. The frames occupy + contiguous spans of the node axis; nodes past ``sum(n_node_per_frame)`` + are padding and join no frame. + charge + Raw node charges of the charge network with shape ``(N, C)``, + float32, with ``C`` at most :data:`MAX_CHANNELS`. + node_mask + Nodes that carry charges with shape ``(N,)``, bool; the others carry + none. + n_node_per_frame + Node count of every frame with shape ``(F,)``. + total_charge + Total charge of every frame with shape ``(F,)``, float64, which the + constrained channel sums to. + kernel + The interaction kernel, accepted by :func:`kernel_eligible`. + channel + The constrained channel, which is shifted uniformly over the charged + nodes of every frame to sum to its total charge, or -1 for none. + prefactor + Factor of every pair term, in eV Å. + communicator + The address of the MPI_Comm handle of a system decomposed into + domains, one per rank, as an int64 CPU tensor of one element, or zero, + the default, for a single domain. Across domains the call takes one + frame, whose charged nodes are those of every rank; the outputs are + those of the charged nodes of this rank, the other nodes carry no + charge, and ``virial`` is the share of this rank. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64, each pair + energy split half to each of its nodes. + charge_gradient + Derivative of the energy with respect to ``charge`` with shape + ``(N, C)``, in eV per unit charge, through the constraint. + force + Force at fixed charges with shape ``(N, 3)``, in eV/Å. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, each + pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(F, 9)`` in eV. + """ + ensure_registered() + code, coefficient = kernel_operator_arguments(kernel) + if communicator is None: + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return torch.ops.deepmd.latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node_per_frame, + total_charge, + coefficient.to(coord.device), + communicator, + code, + int(channel), + float(prefactor), + ) diff --git a/deepmd/pt_expt/long_range/__init__.py b/deepmd/pt_expt/long_range/__init__.py new file mode 100644 index 0000000000..9c230575e7 --- /dev/null +++ b/deepmd/pt_expt/long_range/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) +from .latent_charge import ( + LatentChargeLongRange, +) + +__all__ = [ + "BaseLongRange", + "LESKernel", + "LatentChargeLongRange", + "SOGKernel", +] diff --git a/deepmd/pt_expt/long_range/base_long_range.py b/deepmd/pt_expt/long_range/base_long_range.py new file mode 100644 index 0000000000..66135372f9 --- /dev/null +++ b/deepmd/pt_expt/long_range/base_long_range.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import torch + +from deepmd.dpmodel.long_range import ( + make_base_long_range, +) + +BaseLongRange = make_base_long_range(torch.Tensor) diff --git a/deepmd/pt_expt/long_range/kernel.py b/deepmd/pt_expt/long_range/kernel.py new file mode 100644 index 0000000000..3675e6b988 --- /dev/null +++ b/deepmd/pt_expt/long_range/kernel.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP +from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP +from deepmd.pt_expt.common import ( + register_dpmodel_mapping, + to_torch_array, +) + + +def _torch_kernel(kernel_dp: type) -> type: + """Wrap a dpmodel interaction kernel as a module whose parameters are trained. + + The attributes listed in ``PARAMETER_NAMES`` are stored as + ``torch.nn.Parameter`` that require a gradient exactly when the kernel is + trainable, both at construction and when a deserialized value is assigned. + """ + + class TorchKernel(kernel_dp, torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + torch.nn.Module.__init__(self) + kernel_dp.__init__(self, *args, **kwargs) + + def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor: + return torch.nn.Module.__call__(self, *args, **kwargs) + + def forward(self, r_sq: torch.Tensor) -> torch.Tensor: + return self.call(r_sq) + + def __setattr__(self, name: str, value: Any) -> None: + if name in self.PARAMETER_NAMES and "_parameters" in self.__dict__: + if not isinstance(value, torch.nn.Parameter): + value = torch.nn.Parameter( + to_torch_array(value).detach().clone(), + requires_grad=self.trainable, + ) + if name in self._parameters: + self._parameters[name] = value + return + super().__setattr__(name, value) + + TorchKernel.__name__ = TorchKernel.__qualname__ = kernel_dp.__name__ + register_dpmodel_mapping( + kernel_dp, lambda v, cls=TorchKernel: cls.deserialize(v.serialize()) + ) + return TorchKernel + + +LESKernel = _torch_kernel(LESKernelDP) +SOGKernel = _torch_kernel(SOGKernelDP) diff --git a/deepmd/pt_expt/long_range/latent_charge.py b/deepmd/pt_expt/long_range/latent_charge.py new file mode 100644 index 0000000000..2cf9bf009f --- /dev/null +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import torch + +from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange as LatentChargeLongRangeDP, +) +from deepmd.pt_expt.common import ( + torch_module, +) +from deepmd.pt_expt.fitting.latent_charge import ( + LatentChargeFitting, +) + +from .base_long_range import ( + BaseLongRange, +) +from .kernel import ( + LESKernel, + SOGKernel, +) + + +@BaseLongRange.register("latent_charge") +@torch_module +class LatentChargeLongRange(LatentChargeLongRangeDP): + """Energy of interacting latent charges in complete non-periodic systems (pt_expt).""" + + # The dpmodel class builds the charge network and the kernel, which the + # assignment converts to these pt_expt modules; importing their classes + # registers the conversions. + charge_net: LatentChargeFitting + kernel: LESKernel | SOGKernel + + def _pair_kernel( + self, r_sq: torch.Tensor, src: torch.Tensor, n: int + ) -> torch.Tensor: + # A graph traced for a symbolic pair count plans the reduction of the + # coefficient gradients over the pairs before it sees a batch. In + # training the coefficients therefore reach the pairs through the node + # axis, which turns that reduction into a scatter over the pairs and a + # sum over the nodes, parallel in the pair count of every batch. A + # concrete pair count is reduced directly. + if not (self.training and isinstance(r_sq.shape[0], torch.SymInt)): + return self.kernel(r_sq) + coefficients = { + name: value.expand(n, *value.shape)[src] + for name, value in self.kernel.coefficients(r_sq).items() + } + return self.kernel.evaluate(r_sq, coefficients) + + def direct_sum_eligible(self) -> bool: + """Whether :meth:`energy_force_from_coord` evaluates this module. + + The direct-sum operator evaluates the kernels and channel counts it is + instantiated for. The compact lower that calls it reads no frame + input: the charges depend on the descriptor alone, and a constrained + channel sums to zero or to the total charge of a folded charge state. + """ + from deepmd.pt_expt.kernels.long_range import ( + MAX_CHANNELS, + kernel_eligible, + op_available, + ) + + return ( + op_available() + and kernel_eligible(self.kernel) + and self.charge_net.dim_out <= MAX_CHANNELS + and self.charge_net.numb_fparam == 0 + and self.charge_net.numb_aparam == 0 + and not self.uses_charge_spin() + ) + + def energy_force_from_coord( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + node_mask: torch.Tensor, + n_node: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the module for inference from the positions of the nodes. + + The direct-sum operator enumerates the pairs of every frame itself, + so the memory is linear in the node count. It applies the constraint + and returns the force at fixed charges and the derivative of the + energy with respect to the charges the network predicts; pulled back + through the charge network, that derivative is the derivative of the + energy with respect to the descriptor, through which the caller adds + the response of the charges to the positions. + + Parameters + ---------- + descriptor + Invariant node features with shape ``(N, dim_descrpt)``. + atype + Node types with shape ``(N,)``. + node_mask + Nodes carrying charges with shape ``(N,)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + coord + Node positions with shape ``(N, 3)`` in Å. + communicator + The MPI communicator of a frame decomposed into domains, one per + rank, whose charged nodes are those of every rank, in the form of + :func:`~deepmd.pt_expt.kernels.long_range.latent_charge_direct_sum`; + ``None`` for a single domain. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64. + descriptor_gradient + Derivative of the energy with respect to the descriptor through + the charges, with the shape and precision of ``descriptor``. + force + Force at fixed charges with shape ``(N, 3)`` in eV/Å, float32. + atom_virial + Per-node virial at fixed charges with shape ``(N, 9)`` in eV, + float32, each pair term split half to each of its nodes. + virial + Per-frame virial at fixed charges with shape ``(nf, 9)`` in eV, + float32; across domains, the share of the nodes of this rank. + """ + from deepmd.pt_expt.kernels.long_range import ( + latent_charge_direct_sum, + ) + + leaf = descriptor.detach().requires_grad_(True) + with torch.enable_grad(): + charge = self.charge_net.call_graph(leaf, atype)["latent_charge"] + total_charge = coord.new_zeros(n_node.shape[0], dtype=torch.float64) + channel = -1 + if self.charge_constraint is not None: + channel = self.charge_constraint["channel"] + total_charge = self._total_charge(n_node.shape[0], total_charge, None, None) + atom_energy, charge_gradient, force, atom_virial, virial = ( + latent_charge_direct_sum( + coord.to(torch.float64).contiguous(), + charge.detach().to(torch.float32).contiguous(), + node_mask.contiguous(), + n_node, + total_charge.contiguous(), + self.kernel, + channel, + COULOMB_CONSTANT, + communicator, + ) + ) + (descriptor_gradient,) = torch.autograd.grad( + charge, leaf, charge_gradient.to(charge.dtype) + ) + return ( + atom_energy, + descriptor_gradient.to(descriptor.dtype), + force, + atom_virial, + virial, + ) diff --git a/deepmd/pt_expt/model/edge_transform_output.py b/deepmd/pt_expt/model/edge_transform_output.py index a4a497d453..3e95d33a3a 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -1,10 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Autograd assembly: graph energy -> force/virial/atom_virial via grad(E, edge_vec). -torch-only. The pure-array scatter (edge_force_virial) is shared with dpmodel; -this module supplies the single backward pass that produces g_e = dE/d(edge_vec). +torch-only. The pure-array scatters (edge_force_virial, pair_force_virial) are +shared with dpmodel; this module supplies the single backward pass that +produces the gradient of the energy with respect to every geometry leaf: the +edge vectors of the neighbor graph and the pair vectors of the topologies. """ +from collections.abc import ( + Sequence, +) + import torch from deepmd.dpmodel import ( @@ -20,8 +26,10 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, + PairGraph, edge_force_virial, frame_id_from_n_node, + pair_force_virial, segment_sum, ) from deepmd.pt.utils import ( @@ -60,6 +68,7 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, + pair_graphs: Sequence[PairGraph] = (), ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -79,6 +88,12 @@ def edge_energy_deriv( reduction. A retained graph (training / double backward) keeps the fp64 leaf dtype so training numerics are unchanged. + The same backward differentiates the energy with respect to the vectors + of every pair topology in ``pair_graphs``. Their gradients are pulled back + by the shared pair_force_virial scatter and summed in: the energy reaches + the coordinates through all its geometry leaves, so its complete force and + virial are the sum of their pullbacks. + Parameters ---------- energy @@ -109,6 +124,9 @@ def edge_energy_deriv( whether to materialize the per-atom virial (else ``None`` is returned). create_graph whether the backward retains a graph (training, for second-order grad). + pair_graphs + Pair topologies whose ``vec`` are further autograd leaves of + ``energy``, on the same node axis. Returns ------- @@ -119,9 +137,9 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - (g_e,) = torch.autograd.grad( + g_e, *g_pairs = torch.autograd.grad( energy.sum() if energy.dim() else energy, - edge_vec, + (edge_vec, *(pairs.vec for pairs in pair_graphs)), create_graph=create_graph, retain_graph=True, ) @@ -197,6 +215,19 @@ def edge_energy_deriv( force, atom_virial, virial = edge_force_virial( g_e, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity ) + for g_p, pairs in zip(g_pairs, pair_graphs, strict=True): + pair_force, pair_atom_virial, pair_virial = pair_force_virial( + g_p, + pairs.vec, + pairs.index, + pairs.mask, + n_node, + node_capacity=force.shape[0], + ) + force = force + pair_force.to(force.dtype) + virial = virial + pair_virial.to(virial.dtype) + if do_atomic_virial: + atom_virial = atom_virial + pair_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -211,6 +242,7 @@ def fit_output_to_model_output_graph( n_local: torch.Tensor | None = None, force_precision: torch.dtype | None = None, spin_leaf: torch.Tensor | None = None, + topologies: dict[str, PairGraph] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -282,6 +314,10 @@ def fit_output_to_model_output_graph( ``grad(E, [edge_vec, spin])`` jointly in one call; here the two grads are separate so :func:`edge_energy_deriv`'s signature stays untouched. ``None`` (default): unchanged, no mag output. + topologies + Pair topologies read by the atomic model, by name. Their ``vec`` are + autograd leaves of ``fit_ret`` next to ``graph.edge_vec``; force and + virial include their pullbacks (see :func:`edge_energy_deriv`). Returns ------- @@ -393,6 +429,7 @@ def reduce_by_frame(data: torch.Tensor) -> torch.Tensor: do_atomic_virial=(vdef.c_differentiable and do_atomic_virial), create_graph=create_graph, force_precision=force_precision if not create_graph else None, + pair_graphs=() if topologies is None else tuple(topologies.values()), ) # force (N, 3) -> (N, 1, 3) [flat; caller unravels at I/O boundary] ff_list.append(force.reshape(N, 1, 3)) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index f7e021a684..0523084381 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -18,12 +18,16 @@ from deepmd.dpmodel.model.make_hessian_model import ( make_hessian_model, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, +) from deepmd.dpmodel.utils.neighbor_list import ( NeighborList, ) from .make_model import ( _translate_energy_keys, + check_owned_frames, fused_atom_bias, make_model, ) @@ -138,6 +142,8 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, spin: torch.Tensor | None = None, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -165,6 +171,15 @@ def forward_lower_canonical_graph( spin Per-node magnetic moments with shape ``(N, 3)`` for a native-spin model, or ``None``. Ghost rows carry their owner's moment. + coord + Node positions with shape ``(N, 3)`` in Å, float64, for a model + with a long-range module, whose frames are complete non-periodic + systems; ``None`` otherwise. + communicator + For a model with a long-range module, the address of the MPI_Comm + handle of a frame decomposed into domains, one per rank, as an + int64 CPU tensor of one element, whose owned nodes are then those + of one domain; zero or ``None`` for a complete frame. Returns ------- @@ -211,7 +226,32 @@ def forward_lower_canonical_graph( descriptor = learned.descriptor fitting = learned.fitting_net atom_bias = fused_atom_bias(self.atomic_model, learned) - if use_dpa4c: + if learned.has_long_range(): + from deepmd.pt_expt.kernels.dpa4c.canonical_long_range import ( + dpa4c_canonical_long_range_energy_force, + ) + + if coord is None: + raise ValueError( + "a model with a long-range module reads the node positions" + ) + check_owned_frames(n_node, n_local, communicator) + force_mag = None + energy, atom_energy, force, virial, atom_virial = ( + dpa4c_canonical_long_range_energy_force( + descriptor, + fitting, + learned.long_range, + graph, + atype, + coord, + communicator, + output_mask, + atom_bias, + do_atomic_virial, + ) + ) + elif use_dpa4c: energy, atom_energy, force, virial, atom_virial, force_mag = ( dpa4c_canonical_compress_energy_force( descriptor, @@ -271,6 +311,8 @@ def forward_lower_canonical_graph_exportable( destination_row_ptr: torch.Tensor, source_row_ptr: torch.Tensor, source_order: torch.Tensor, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -283,6 +325,11 @@ def forward_lower_canonical_graph_exportable( Compact graph node and edge tensors. destination_row_ptr, source_row_ptr, source_order Compact dual-CSR topology. + coord, communicator + For a model with a long-range module, sample node positions with + shape ``(N, 3)``, float64, and a sample communicator, an int64 CPU + tensor of one element, which the traced module then takes as its + ninth and tenth inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -291,9 +338,56 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced eight-input compact deployment module. + Traced compact deployment module of eight inputs, or ten with the + node positions and the communicator. """ model = self + if (coord is None) != (communicator is None): + raise ValueError( + "the compact lower of a model with a long-range module takes " + "the node positions and the communicator together" + ) + + if coord is not None: + + def fn_long_range( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + source: torch.Tensor, + edge_vec: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_row_ptr: torch.Tensor, + source_order: torch.Tensor, + coord: torch.Tensor, + communicator: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_lower_canonical_graph( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + do_atomic_virial=do_atomic_virial, + coord=coord, + communicator=communicator, + ) + + return make_fx(fn_long_range, **make_fx_kwargs)( + atype, + n_node, + n_local, + source, + edge_vec, + destination_row_ptr, + source_row_ptr, + source_order, + coord, + communicator, + ) def fn( atype: torch.Tensor, @@ -613,6 +707,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, PairGraph] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -646,6 +741,11 @@ def forward_lower_graph_exportable( and ``destination_order`` is identity. fparam, aparam, do_atomic_virial, charge_spin As in ``forward_lower``. + topologies + The sample pair topology the atomic model declares, or ``None``; + its ``index``, ``vec`` and ``mask`` become the inputs + ``pair_index``, ``pair_vec`` and ``pair_mask`` after + ``charge_spin``. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -680,6 +780,7 @@ def forward_lower_graph_exportable( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, destination_sorted=destination_sorted, + topologies=topologies, **make_fx_kwargs, ) @@ -688,6 +789,80 @@ def forward_lower_graph_exportable( do_grad_r = self.do_grad_r("energy") do_grad_c = self.do_grad_c("energy") + if topologies: + # The inputs of the deployed lower keep their names, which the + # export reads, so the traced signature spells out the one pair + # topology a long-range module declares. + if len(topologies) != 1: + raise NotImplementedError( + "the deployed graph lower takes one pair topology, got " + f"{list(topologies)}" + ) + (pairs,) = topologies.values() + + def fn_pairs( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + pair_index: torch.Tensor, + pair_vec: torch.Tensor, + pair_mask: torch.Tensor, + ) -> dict[str, torch.Tensor]: + model_ret = traced( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + pair_index, + pair_vec, + pair_mask, + ) + return _translate_energy_keys( + model_ret, + do_grad_r=do_grad_r, + do_grad_c=do_grad_c, + do_atomic_virial=do_atomic_virial, + local=True, + ) + + return make_fx(fn_pairs, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + pairs.index, + pairs.vec, + pairs.mask, + ) + def fn( atype: torch.Tensor, n_node: torch.Tensor, diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 2071875b5c..5bd6e68717 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -31,6 +31,9 @@ from deepmd.pt_expt.fitting import ( BaseFitting, ) +from deepmd.pt_expt.long_range import ( + BaseLongRange, +) from deepmd.pt_expt.model.dp_zbl_model import ( DPZBLModel, ) @@ -75,6 +78,7 @@ pairtab_model=PairTabAtomicModel, zbl_model=DPZBLModel, linear_atomic_model=LinearEnergyAtomicModel, + long_range_base=BaseLongRange, ) get_zbl_model = _model_factory.get_zbl_model @@ -168,6 +172,7 @@ def get_sezm_model(data: dict) -> BaseModel: data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) descriptor, fitting, _ = _model_factory.get_model_components(data) + long_range = _model_factory.get_long_range(data, descriptor, fitting) return DPA4EnergyModel( descriptor=descriptor, fitting=fitting, @@ -175,6 +180,7 @@ def get_sezm_model(data: dict) -> BaseModel: atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, preset_out_bias=data.get("preset_out_bias"), + **({} if long_range is None else {"long_range": long_range}), ) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index 0dfd77ccc2..9b132b2a53 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -21,8 +21,12 @@ OutputVariableDef, ) from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + build_topologies, compact_nodes, expand_node_values, + flatten_topologies, + unflatten_topologies, ) from deepmd.pt_expt.common import ( auto_wrapped_class, @@ -137,7 +141,8 @@ def _fused_energy_force_graph( """ am = model.atomic_model parts = am.fused_decomposition() - if parts is None: + if parts is None or parts[0].has_long_range(): + # The descriptor pipelines evaluate no long-range module. return None learned, pair_potential = parts desc, fit = learned.descriptor, learned.fitting_net @@ -191,6 +196,51 @@ def _fused_energy_force_graph( return ret +def check_owned_frames( + n_node: torch.Tensor, + n_local: torch.Tensor | None, + communicator: torch.Tensor | None = None, +) -> None: + """Reject halo nodes in the frames of a single domain with a long-range module. + + The module sums over every pair of atoms of a frame, so a frame that one + domain holds whole owns all its nodes; halo nodes would carry no charge + and remove their interactions from the sum. Across the domains of a + communicator, every domain owns its share of the frame instead. The + counts are read eagerly only: a traced lower receives them from callers + that check them. + + Parameters + ---------- + n_node + Per-frame node counts with shape ``(nf,)``. + n_local + Per-frame owned node counts with shape ``(nf,)``, or ``None``. + communicator + The communicator input of the compact lower, a CPU tensor holding the + address of an MPI_Comm handle or zero, or ``None`` for a single domain. + + Raises + ------ + ValueError + If a frame of a single domain carries halo nodes. + """ + from torch._subclasses.fake_tensor import ( + is_fake, + ) + + if n_local is None or n_local is n_node or is_fake(n_local): + return + if communicator is not None and int(communicator.reshape(-1)[0]) != 0: + return + if not torch.equal(n_local, n_node): + raise ValueError( + "a model with a long-range module evaluates complete frames on a " + "single domain, whose graph owns every node (n_local == n_node), " + f"but got n_local {n_local.tolist()} for n_node {n_node.tolist()}" + ) + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -356,7 +406,8 @@ class _WrapperForwardEnergyGraph: """Graph twin of :class:`_WrapperForwardEnergy` for the Hessian. Given flattened LOCAL coordinates for one frame, rebuilds the carry-all - ``NeighborGraph`` (so ``edge_vec`` tracks the coordinates through + ``NeighborGraph`` and the topologies the atomic model declares (so + ``edge_vec`` and the pair vectors track the coordinates through ``autograd.functional.hessian``'s double-backward) and returns the scalar reduced-output component. Unlike the dense wrapper this differentiates w.r.t. the ``(nloc, 3)`` LOCAL coordinates directly -- the carry-all graph @@ -399,6 +450,16 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: ng = build_neighbor_graph_for_method( self.method, cc, self.atype, self.box, self.rcut, self.pair_excl ) + # The topologies follow the same coordinates, so the Hessian holds + # every geometry path of the energy and the blocks coupling them. + topologies = build_topologies( + self.model.atomic_model.required_topologies(), + coord_flat.reshape(self.nloc, 3), + ng.n_node, + self.box, + atype=self.atype.reshape(-1), + pair_excl=self.pair_excl, + ) atomic_ret = self.model.atomic_model.forward_common_atomic_graph( ng, self.atype.reshape(-1), @@ -406,6 +467,7 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor: aparam=self.aparam, spin=self.spin, charge_spin=self.charge_spin, + topologies=topologies, ) # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all # graph (all nodes owned); reduced output = sum over the node axis. @@ -608,6 +670,7 @@ def forward_common_lower_graph( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, comm_dict: dict | None = None, + topologies: dict[str, PairGraph] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -696,6 +759,12 @@ def forward_common_lower_graph( DescrptBlockRepformers._exchange_ghosts_graph`). Only meaningful for descriptors whose ``has_message_passing_across_ranks()`` is ``True``. + topologies + The topologies the atomic model declares by + ``required_topologies()``, built on the node axis of the + graph from complete systems. Each ``vec`` becomes an autograd + leaf next to ``edge_vec``, and force and virial include its + pullback. ``None`` (default) for models declaring none. Returns ------- @@ -712,6 +781,22 @@ def forward_common_lower_graph( # make edge_vec the autograd leaf for the energy backward edge_vec = edge_vec.detach().requires_grad_(True) + if topologies: + if comm_dict is not None: + raise NotImplementedError( + "the long-range topologies describe complete systems; " + "a domain-decomposed system is not supported" + ) + check_owned_frames(n_node, n_local) + # Every topology carries a geometry leaf of its own. + topologies = { + name: PairGraph( + index=pairs.index, + vec=pairs.vec.detach().requires_grad_(True), + mask=pairs.mask, + ) + for name, pairs in topologies.items() + } if spin is not None: # second autograd leaf: force_mag = -dE/dspin (native spin). # Deliberately a SEPARATE leaf/backward from edge_vec rather @@ -751,6 +836,7 @@ def forward_common_lower_graph( charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, + topologies=topologies, ) # ``forward_common_atomic_graph`` returns flat ``(N, *)`` output. # Pass directly to the flat-N transform; no rectangular reshape needed. @@ -772,6 +858,7 @@ def forward_common_lower_graph( # dynamic-edge torch.export. See fit_output_to_model_output_graph. node_capacity=atype.shape[0], n_local=n_local, + topologies=topologies, ) def _resolve_graph_method( @@ -892,6 +979,7 @@ def call_common_ragged( method = getattr(self, "neighbor_graph_method", None) if method is None: method = resolve_neighbor_graph_method("auto", coord.device) + pair_excl = getattr(self.atomic_model, "pair_excl", None) graph = build_ragged_neighbor_graph( method, coord, @@ -899,7 +987,15 @@ def call_common_ragged( n_node, box, self.get_rcut(), - getattr(self.atomic_model, "pair_excl", None), + pair_excl, + ) + topologies = build_topologies( + self.atomic_model.required_topologies(), + coord.detach(), + graph.n_node, + box, + atype=atype, + pair_excl=pair_excl, ) predict = self.forward_common_lower_graph( atype, @@ -918,6 +1014,7 @@ def call_common_ragged( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) # The per-atom mask a rectangular batch carries is what tells the # loss each frame's real atom count. Nothing here is padded, so the @@ -1022,6 +1119,14 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) + topologies = build_topologies( + self.atomic_model.required_topologies(), + cc.detach().reshape(n_padded, 3)[node_index], + ng.n_node, + bb, + atype=atype_flat, + pair_excl=pair_excl, + ) model_predict = self.forward_common_lower_graph( atype_flat, ng.n_node, @@ -1039,6 +1144,7 @@ def _call_common_graph( aparam=ap_flat, spin=spin_flat, charge_spin=charge_spin, + topologies=topologies, ) # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom # outputs over the real atoms. Scatter them back onto the padded @@ -1249,6 +1355,7 @@ def forward_common_lower_graph_exportable( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, PairGraph] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -1289,6 +1396,13 @@ def forward_common_lower_graph_exportable( output. ``None`` (default) is the existing, unconditioned trace with no mag output -- used by every non-spin caller of this generic (descriptor-agnostic) exportable. + topologies + Sample topologies the atomic model declares by + ``required_topologies()``, or ``None``. When given, they are + traced inputs laid out by + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies` + after ``charge_spin``; a model reading topologies carries no + native spin. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -1299,8 +1413,9 @@ def forward_common_lower_graph_exportable( A traced module whose ``forward`` accepts ``(atype, n_node, n_local, edge_index, edge_vec, edge_mask, destination_order, destination_row_ptr, source_order, - source_row_ptr, fparam, aparam, charge_spin, spin)`` and - returns a dict with the same internal keys as + source_row_ptr, fparam, aparam, charge_spin, spin)`` -- or, + with topologies, ``spin`` replaced by the topology inputs -- + and returns a dict with the same internal keys as ``forward_common_lower_graph``. """ validate_graph_csr_for_export( @@ -1315,6 +1430,61 @@ def forward_common_lower_graph_exportable( ) model = self + if topologies: + topology_names = tuple(topologies) + + def fn_topologies( + atype: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor, + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + destination_order: torch.Tensor, + destination_row_ptr: torch.Tensor, + source_order: torch.Tensor, + source_row_ptr: torch.Tensor, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + *topology_vals: torch.Tensor, + ) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + destination_sorted=destination_sorted, + do_atomic_virial=do_atomic_virial, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + topologies=unflatten_topologies(topology_names, topology_vals), + ) + + return make_fx(fn_topologies, **make_fx_kwargs)( + atype, + n_node, + n_local, + edge_index, + edge_vec, + edge_mask, + destination_order, + destination_row_ptr, + source_order, + source_row_ptr, + fparam, + aparam, + charge_spin, + *flatten_topologies(topologies), + ) + # ``spin`` is a traced INPUT only when a real spin tensor is # given (native-spin models). For every non-spin caller # (``spin is None``) the traced ``fn`` must have the SAME arity as diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f1e7f27c58..f4e24bc824 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -72,6 +72,11 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) @@ -297,6 +302,10 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: atomic model and fitting. Shared weights do not make those branches interchangeable: tasks without a preset bias disable the reference, while a folded reference no longer appends isolated atoms to the graph. + + A long-range module is part of the traced computation and of its input + schema, so its identity (that of its first parameter) closes the key; + tasks share a graph only when they share the module as well. """ descriptor_id: int = 0 try: @@ -322,7 +331,17 @@ def _get_model_structure_key(model: torch.nn.Module) -> tuple[int, ...]: break except AttributeError: pass - return (int(model.has_spin()), descriptor_id, fitting_id, *vacuum_state) + long_range_id = 0 + long_range = getattr(getattr(model, "atomic_model", None), "long_range", None) + if long_range is not None: + long_range_id = next((id(p) for p in long_range.parameters()), id(long_range)) + return ( + int(model.has_spin()), + descriptor_id, + fitting_id, + *vacuum_state, + long_range_id, + ) # --------------------------------------------------------------------------- @@ -765,6 +784,11 @@ def _trace_and_compile_graph( autograd leaf. Their compiled lower returns ``force_mag`` and ``mask_mag`` through the same model-owned translation used by eager execution. + A model declaring topologies (``atomic_model.required_topologies()``) + receives each of them as three further inputs, ``(index, vec, mask)`` in + the declared order, ahead of the task buffers; its pair axis is a fourth + dynamic dimension, kept a distinct symbol like the other three. + Parameters ---------- model @@ -842,6 +866,7 @@ def _trace_and_compile_graph( # its graph builders omit the CSR metadata carried by deployment inputs. from deepmd.pt_expt.utils.serialization import ( build_synthetic_graph_inputs, + build_synthetic_topologies, check_graph_trace_torch_version, count_synthetic_graph_edges, ) @@ -899,6 +924,20 @@ def _trace_and_compile_graph( # tensors, while forward_common_lower_graph takes it after them. s_spin = s_conditioning.pop(0) if spin is not None else None s_fparam, s_aparam, s_charge_spin = s_conditioning + # The pair axis of every topology is a fourth dynamic axis, its sample + # length distinct from the other three and from the CSR offsets. + s_topologies = build_synthetic_topologies( + model, + trace_N, + _next_safe_prime( + e_max + 2, _forbidden | {trace_nf, trace_N, trace_N + 1, e_max} + ), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=_trace_device, + ) + topology_names = tuple(s_topologies) + s_topology_vals = flatten_topologies(s_topologies) + n_topology_vals = len(s_topology_vals) def fn( atype: torch.Tensor, @@ -915,8 +954,10 @@ def fn( aparam: torch.Tensor | None, charge_spin: torch.Tensor | None, spin: torch.Tensor | None, - *task_buf_vals: torch.Tensor, + *extra_vals: torch.Tensor, ) -> dict[str, torch.Tensor]: + topologies = unflatten_topologies(topology_names, extra_vals[:n_topology_vals]) + task_buf_vals = extra_vals[n_topology_vals:] # Patch task-specific buffers with the proxy tensors so make_fx records # them as FX placeholders (mirrors the dense ``_trace_and_compile``). originals: dict[str, torch.Tensor | None] = {} @@ -950,6 +991,7 @@ def fn( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) if spin is None: return _translate_energy_keys( @@ -996,6 +1038,7 @@ def fn( s_aparam, s_charge_spin, s_spin, + *s_topology_vals, *task_buf_vals_trace, ) @@ -1640,6 +1683,16 @@ def _forward_graph( aparam = aparam[node_index] if spin is not None: spin = spin[node_index] + # The topologies the atomic model declares are built here as well, + # outside the compiled region, on the node axis of the graph. + topologies = build_topologies( + _model.atomic_model.required_topologies(), + coord_3d if ragged else coord_3d.reshape(n_padded, 3)[node_index], + ng.n_node, + box_flat, + atype=atype_flat, + pair_excl=pair_excl, + ) # Lazy compile of the GRAPH lower (cached per structure key and mode). compiled_lower = self._compiled_lower_by_mode.get(self.training) @@ -1667,6 +1720,17 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) + # Each pair topology enters as (index, vec, mask), its vec a further + # leaf. + topology_vals = tuple( + tensor + for pairs in topologies.values() + for tensor in ( + pairs.index, + pairs.vec.detach().requires_grad_(True), + pairs.mask, + ) + ) if self._task_buf_order: try: @@ -1708,6 +1772,7 @@ def _forward_graph( aparam, charge_spin, spin, + *topology_vals, *task_buf_vals, ) self._report_pending_compile() @@ -3560,8 +3625,8 @@ def _learned_part_path(model: torch.nn.Module) -> str | None: def _move_learned_key(key: str, path: str, source_path: str) -> str: - """Rewrite a descriptor or fitting-net key from one learned-part path to another.""" - for part in ("descriptor.", "fitting_net."): + """Rewrite a key of the learned part from one learned-part path to another.""" + for part in ("descriptor.", "fitting_net.", "long_range."): if f".{path}{part}" in key: return key.replace(f".{path}{part}", f".{source_path}{part}", 1) return key diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 22f5a2a494..325dc0d7c5 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -14,6 +14,10 @@ import numpy as np import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + flatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -674,18 +678,69 @@ def build_synthetic_graph_inputs( ) +def build_synthetic_topologies( + model: torch.nn.Module, + n_node_total: int, + n_pair: int, + *, + dtype: torch.dtype, + device: torch.device, +) -> dict[str, PairGraph]: + """Build trace-time samples of the topologies a model declares. + + Single source of the trace-time topology inputs, shared by graph export + and compiled training. Tracing reads only their shapes, so every topology + joins fixed nodes over ``n_pair`` pairs, a length the caller keeps + distinct from the other dynamic axes. + + Parameters + ---------- + model : torch.nn.Module + The model; its atomic model declares the topologies. + n_node_total : int + Length of the node axis the pairs index. + n_pair : int + Length of the pair axis of every sample. + dtype : torch.dtype + Precision of the pair vectors. + device : torch.device + Target device. + + Returns + ------- + dict[str, PairGraph] + A sample of every declared topology; empty when none is declared. + """ + index = ( + torch.arange(2 * n_pair, dtype=torch.int64, device=device).reshape(2, n_pair) + % n_node_total + ) + return { + name: PairGraph( + index=index, + vec=torch.ones(n_pair, 3, dtype=dtype, device=device), + mask=torch.ones(n_pair, dtype=torch.bool, device=device), + ) + for name in model.atomic_model.required_topologies() + } + + def build_synthetic_canonical_graph_inputs( model: torch.nn.Module, e_max: int, *, device: torch.device, want_spin: bool = False, + want_long_range: bool = False, ) -> tuple[torch.Tensor, ...]: """Build compact canonical trace inputs for compressed CUDA descriptors. ``want_spin`` appends the per-node moment at slot 8, the last slot of the compact ABI, matching :meth:`~deepmd.pt_expt.model.native_spin_model.NativeSpinEnergyModel.forward_lower_canonical_graph_exportable`. + ``want_long_range`` appends the inputs that the compact lower of a model + with a long-range module reads: the float64 node positions at slot 8 and + the int64 communicator, a CPU tensor of one element, at slot 9. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -742,7 +797,17 @@ def build_synthetic_canonical_graph_inputs( compact.source_row_ptr, compact.source_order, ) - return compact_sample if spin is None else (*compact_sample, spin) + if spin is not None: + return (*compact_sample, spin) + if want_long_range: + # Tracing reads the shape of the positions alone; distinct values keep + # the sample nodes apart. The communicator of a single domain is zero. + coord = 0.7 * torch.arange( + 3 * atype.shape[0], dtype=torch.float64, device=device + ).reshape(-1, 3) + communicator = torch.zeros(1, dtype=torch.int64, device="cpu") + return (*compact_sample, coord, communicator) + return compact_sample def _build_canonical_graph_dynamic_shapes( @@ -750,8 +815,10 @@ def _build_canonical_graph_dynamic_shapes( ) -> tuple: """Build dynamic shapes for the compact deployment ABI. - The trailing spin slot is present only for a native-spin model, so the - sample length selects the shape tuple. + The trailing slots, the moments of a native-spin model or the positions + and the communicator of a model with a long-range module, are present + only for such models, so the sample length selects the shape tuple. The + communicator holds one address and is static. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, @@ -774,7 +841,8 @@ def _build_canonical_graph_dynamic_shapes( {0: node_dim + 1}, {0: edge_storage_dim}, ) - return shapes if len(sample_inputs) == len(shapes) else (*shapes, {0: node_dim}) + trailing = ({0: node_dim}, None)[: len(sample_inputs) - len(shapes)] + return (*shapes, *trailing) def count_synthetic_graph_edges( @@ -830,6 +898,7 @@ def count_synthetic_graph_edges( def _build_graph_dynamic_shapes( *sample_inputs: torch.Tensor | None, is_native_spin: bool = False, + topology_names: tuple[str, ...] = (), ) -> tuple: """Build dynamic-shape specifications for the graph-form forward_lower export. @@ -857,6 +926,10 @@ def _build_graph_dynamic_shapes( Whether ``sample_inputs`` follows the native-spin positional ABI (spin at slot 10) instead of the regular energy ABI (charge_spin at slot 12). + topology_names : tuple[str, ...] + Topologies whose ``(index, vec, mask)`` follow the inputs above, each + with a dynamic pair axis of its own; the guard pairs of a topology + bound it below by two. """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -873,6 +946,10 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) + topologies = [] + for name in topology_names: + npair_dim = torch.export.Dim(f"npair_{name}", min=2) + topologies += [{1: npair_dim}, {0: npair_dim}, {0: npair_dim}] if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -886,6 +963,7 @@ def _build_graph_dynamic_shapes( {0: nframes_dim} if fparam is not None else None, # fparam {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) fparam = sample_inputs[10] aparam = sample_inputs[11] @@ -898,6 +976,7 @@ def _build_graph_dynamic_shapes( # dim would make torch.export prove/reject the equality). {0: n_node_total_dim} if aparam is not None else None, # aparam {0: nframes_dim} if charge_spin is not None else None, # charge_spin + *topologies, ) @@ -1286,6 +1365,15 @@ def _probe_has_message_passing(obj: object) -> bool | None: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break meta["pair_exclude_types"] = pair_exclude_types + # Topologies the lower reads beside the neighbor graph; absent when the + # model reads none. The graph lower takes each as its ``(index, vec, + # mask)`` after its other inputs (see ``flatten_topologies``), built by the + # caller from the coordinates of each frame. The compact canonical lower + # takes the node positions instead, from which its operators enumerate the + # pairs. + topologies = model.atomic_model.required_topologies() + if topologies: + meta["topologies"] = list(topologies) return meta @@ -2066,11 +2154,16 @@ def _trace_and_export_impl( while e_sample in (_forbidden | {nframes_sample, n_sample}): e_sample += 1 if canonical: + # The compact lower of a model with a long-range module reads the + # node positions, from which its operators enumerate the pairs, + # and the communicator of the domains of a decomposed frame. + want_long_range = bool(model.atomic_model.required_topologies()) sample_inputs = build_synthetic_canonical_graph_inputs( model, e_sample, device=torch.device("cpu"), want_spin=is_native_spin, + want_long_range=want_long_range, ) traced = model.forward_lower_canonical_graph_exportable( *sample_inputs, @@ -2152,6 +2245,18 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) + # The pair axis of every topology is dynamic; its sample length + # is kept distinct from the other axes and from the CSR offsets. + topologies = build_synthetic_topologies( + model, + n_sample, + next_safe_prime( + e_sample + 2, + _forbidden | {nframes_sample, n_sample, n_sample + 1, e_sample}, + ), + dtype=edge_dtype, + device=torch.device("cpu"), + ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): # slot 10 is ``spin`` (mandatory), slots 11/12 are @@ -2177,12 +2282,16 @@ def _trace_and_export_impl( do_atomic_virial=do_atomic_virial, charge_spin=sample_inputs[12], destination_sorted=True, + topologies=topologies, tracing_mode="symbolic", _allow_non_fake_inputs=True, ) dynamic_shapes = _build_graph_dynamic_shapes( - *sample_inputs, is_native_spin=is_native_spin + *sample_inputs, + is_native_spin=is_native_spin, + topology_names=tuple(topologies), ) + sample_inputs = (*sample_inputs, *flatten_topologies(topologies)) output_keys = traced_output_keys(traced) exported = torch.export.export( traced, @@ -2513,7 +2622,81 @@ def _charge_state_descriptor(data: dict, metadata: dict) -> Any | None: return descriptor -def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: +def _charge_state_long_range(data: dict, metadata: dict) -> Any | None: + """Return the long-range module folded to the charge state, if any. + + Compressing the atomic model folds a long-range charge constraint that + reads its total charge from the charge state to the condition of the + descriptor, so the fold rebuilds the state of the module with the + artifacts of the descriptor. Only a model with a long-range module, which + the collected metadata names by its topologies, pays for rebuilding the + model. + + Parameters + ---------- + data + Serialized model dictionary, as passed to the export entry point. + metadata + Archive metadata collected from that same model. + + Returns + ------- + Any or None + The evaluated long-range module, or ``None`` when the model folds no + long-range state. + """ + if not metadata.get("topologies"): + return None + + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model = BaseModel.deserialize(data["model"]) + parts = model.atomic_model.fused_decomposition() + long_range = None if parts is None else parts[0].long_range + if long_range is None or not long_range.charge_state_constants(): + return None + model.to("cpu") + model.eval() + return long_range + + +class _ChargeStateRebuild(torch.nn.Module): + """Rebuild the charge state of a descriptor and of a long-range module. + + The outputs are those of the descriptor fold followed by the folded state + of the long-range module. + + Parameters + ---------- + descriptor + Compressed charge-conditioned descriptor, as returned by + :func:`_charge_state_descriptor`. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`. + """ + + def __init__(self, descriptor: Any, long_range: Any) -> None: + super().__init__() + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + ChargeStateFold, + ) + + self.descriptor = ChargeStateFold(descriptor) + self.long_range = long_range + + def forward(self, charge_spin: torch.Tensor) -> tuple[torch.Tensor, ...]: + return ( + *self.descriptor(charge_spin), + *self.long_range.rebuild_charge_state(charge_spin), + ) + + +def _match_charge_state_constants( + descriptor: Any, exported: Any, long_range: Any | None = None +) -> list[str]: """Name, per fold output, the constant of ``exported`` it replaces. A compressed charge-conditioned descriptor carries its frame condition in @@ -2546,6 +2729,10 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: :func:`_charge_state_descriptor`. exported Exported lower whose constants are to be named. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state is matched after the + artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2566,9 +2753,18 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: for name, tensor in exported.state_dict.items() if isinstance(tensor, torch.Tensor) } + artifacts = [ + (name, getattr(descriptor, f"compress_{name}")) + for name in CHARGE_STATE_ARTIFACTS + ] + if long_range is not None: + artifacts += [ + ("long-range charge state", state) + for state in long_range.charge_state_constants() + ] constants: list[str] = [] - for name in CHARGE_STATE_ARTIFACTS: - artifact = getattr(descriptor, f"compress_{name}").detach().cpu() + for name, tensor in artifacts: + artifact = tensor.detach().cpu() if artifact.numel() == 0: constants.append("") continue @@ -2593,6 +2789,7 @@ def _match_charge_state_constants(descriptor: Any, exported: Any) -> list[str]: def _compile_charge_state_fold( descriptor: Any, aoti_configs: dict, + long_range: Any | None = None, ) -> bytes: """Compile the rebuild of the charge-state artifacts as its own archive. @@ -2608,6 +2805,10 @@ def _compile_charge_state_fold( :func:`_charge_state_descriptor`. aoti_configs Inductor options the inference lower was compiled with. + long_range + Long-range module folded to the same state, as returned by + :func:`_charge_state_long_range`, whose state the fold rebuilds after + the artifacts of the descriptor; ``None`` when there is none. Returns ------- @@ -2634,7 +2835,12 @@ def _compile_charge_state_fold( dtype=torch.float32, device="cpu", ) - fold = torch.export.export(ChargeStateFold(descriptor), (sample,)) + module = ( + ChargeStateFold(descriptor) + if long_range is None + else _ChargeStateRebuild(descriptor, long_range) + ) + fold = torch.export.export(module, (sample,)) if _env.DEVICE.type != "cpu": from torch.export.passes import ( move_to_device_pass, @@ -2767,11 +2973,12 @@ def _deserialize_to_file_pt2( # state becomes known. charge_state_bytes: bytes | None = None if charge_state_descriptor is not None: + long_range = _charge_state_long_range(data, metadata) metadata["charge_state_constants"] = _match_charge_state_constants( - charge_state_descriptor, exported + charge_state_descriptor, exported, long_range ) charge_state_bytes = _compile_charge_state_fold( - charge_state_descriptor, aoti_configs + charge_state_descriptor, aoti_configs, long_range ) # Second artifact: with-comm. Only for descriptors whose message diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 068bef876f..e88c69fcf1 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -670,6 +670,11 @@ def __init__( type_map: list[str] | None = None, **kwargs: Any, ) -> None: + if kwargs.get("long_range") is not None: + raise NotImplementedError( + "long-range modules (`long_range`) are supported by the pt_expt " + "backend only" + ) super().__init__( descriptor=descriptor, fitting=fitting_net, type_map=type_map, **kwargs ) diff --git a/deepmd/tf2/atomic_model/dp_atomic_model.py b/deepmd/tf2/atomic_model/dp_atomic_model.py index 2d68ba81ef..966c005b10 100644 --- a/deepmd/tf2/atomic_model/dp_atomic_model.py +++ b/deepmd/tf2/atomic_model/dp_atomic_model.py @@ -41,6 +41,8 @@ class tf2_atomic_model(dpmodel_atomic_model): """The base descriptor class.""" base_fitting_cls = BaseFitting """The base fitting class.""" + base_long_range_cls = None + """No long-range module is implemented in this backend.""" @tf.autograph.experimental.do_not_convert def make_atom_mask( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..b8a729b908 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3546,6 +3546,209 @@ def fitting_variant_type_args() -> Variant: ) +# --- Long-range module configurations: --- # +long_range_args_plugin = ArgsPlugin() +interaction_kernel_args_plugin = ArgsPlugin() + + +@interaction_kernel_args_plugin.register( + "les", + doc="The latent Ewald summation kernel erf(alpha r) / r, the interaction of " + "two Gaussian charge clouds: it reaches 1/r beyond a few 1/alpha and stays " + "finite at contact.", +) +def interaction_kernel_les() -> list[Argument]: + doc_alpha = "The inverse width alpha of the kernel, in Å^-1. Must be positive." + doc_trainable = "Whether alpha is trained." + return [ + Argument("alpha", float, optional=True, default=1.0, doc=doc_alpha), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +@interaction_kernel_args_plugin.register( + "sog", + doc="The sum-of-Gaussians kernel sum_k a_k exp(-r^2 / (2 sigma_k^2)).", +) +def interaction_kernel_sog() -> list[Argument]: + doc_amplitude = ( + "The amplitudes a_k of the Gaussians, in Å^-1, of either sign: one value " + "shared by every Gaussian or one value per Gaussian. Required with " + "`width`; with the default widths it defaults to " + "sqrt(2 / pi) ln(`width_ratio`) / sigma_k, the weights of a discretized " + "1 / r, so that the default kernel follows 1 / r, less a nearly " + "constant offset, from a few `base_width` to a fraction of the widest " + "Gaussian. The offset exerts no force between fixed charges; through " + "charges that respond to the positions it does." + ) + doc_width = ( + "The widths sigma_k of the Gaussians, in Å, each positive. Defaults to " + "the geometric sequence `base_width` * `width_ratio`^k for " + "k = 0, ..., `n_gaussian` - 1." + ) + doc_n_gaussian = "The number of Gaussians of the default widths." + doc_base_width = "The width of the first default Gaussian, in Å." + doc_width_ratio = "The ratio of consecutive default widths. Must exceed one." + doc_trainable = "Whether the amplitudes and widths are trained." + return [ + Argument( + "amplitude", + [float, list[float], type(None)], + optional=True, + default=None, + doc=doc_amplitude, + ), + Argument( + "width", + [list[float], type(None)], + optional=True, + default=None, + doc=doc_width, + ), + Argument("n_gaussian", int, optional=True, default=12, doc=doc_n_gaussian), + Argument( + "base_width", + float, + optional=True, + default=2.180230445405648, + doc=doc_base_width, + ), + Argument( + "width_ratio", + float, + optional=True, + default=1.6297670882677647, + doc=doc_width_ratio, + ), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + ] + + +def latent_charge_net_args() -> list[Argument]: + doc_neuron = "The number of neurons in each hidden layer of the charge network." + doc_activation_function = f"The activation function in the charge network. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Defaults to the activation function of the energy fitting." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection. Defaults to the setting of the energy fitting.' + doc_precision = f"The precision of the charge network and kernel parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Defaults to the precision of the energy fitting." + doc_trainable = "Whether the parameters of the charge network are trained. A list sets every hidden layer and the output layer separately." + doc_seed = "Random seed for parameter initialization of the charge network. Defaults to a seed derived from the seed of the energy fitting." + return [ + Argument("neuron", list[int], optional=True, default=[64, 64], doc=doc_neuron), + Argument( + "activation_function", str, optional=True, doc=doc_activation_function + ), + Argument("resnet_dt", bool, optional=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, doc=doc_precision), + Argument( + "trainable", + [list[bool], bool], + optional=True, + default=True, + doc=doc_trainable, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] + + +def charge_constraint_args() -> list[Argument]: + doc_source = ( + "The source of the total charge of a frame: `neutral` for zero, " + "`charge_spin` for the charge of the frame charge/spin condition, which " + "requires a descriptor with `add_chg_spin_ebd`, or `fparam` for the frame " + "parameter at `fparam_index`." + ) + doc_fparam_index = "The column of the frame parameters holding the total charge, used when `source` is `fparam`." + doc_channel = "The latent charge channel that is constrained to the total charge; the other channels stay unconstrained." + return [ + Argument( + "source", + str, + optional=True, + default="neutral", + doc=doc_source, + extra_check=lambda x: x in ("neutral", "charge_spin", "fparam"), + extra_check_errmsg="must be one of 'neutral', 'charge_spin' or 'fparam'", + ), + Argument("fparam_index", int, optional=True, default=0, doc=doc_fparam_index), + Argument("channel", int, optional=True, default=0, doc=doc_channel), + ] + + +@long_range_args_plugin.register( + "latent_charge", + doc="Interaction of latent charges predicted from the descriptor, summed " + "directly over every pair of atoms of a complete non-periodic system. " + "Every channel interacts through the same kernel; half of each pair energy " + "is attributed to each atom.", +) +def long_range_latent_charge() -> list[Argument]: + doc_channels = "The number of latent charge channels predicted per atom." + doc_charge_net = ( + "The network predicting the latent charges, in units of the elementary " + "charge, from the descriptor; it takes the frame and atomic parameters " + "of the energy fitting." + ) + doc_kernel = "The pair interaction kernel of the latent charges, in Å^-1." + doc_charge_constraint = ( + "Shift one latent charge channel uniformly over the atoms of every frame " + "so that it sums to the total charge of the frame. Unconstrained if not " + "set." + ) + return [ + Argument("channels", int, optional=True, default=1, doc=doc_channels), + Argument( + "charge_net", + dict, + latent_charge_net_args(), + optional=True, + doc=doc_charge_net, + ), + Argument( + "kernel", + dict, + [], + [ + Variant( + "type", + interaction_kernel_args_plugin.get_all_argument(), + doc="The type of the kernel.", + ) + ], + doc=doc_kernel, + ), + Argument( + "charge_constraint", + dict, + charge_constraint_args(), + optional=True, + default=None, + doc=doc_charge_constraint, + ), + ] + + +def long_range_args() -> Argument: + doc_long_range = ( + "A long-range module added to the atomic energy of the energy fitting. " + "It reads the same descriptor and is evaluated on complete " + "non-periodic systems; a periodic cell is rejected. Currently validated " + "with the DPA4C descriptor." + ) + return Argument( + "long_range", + dict, + [], + [ + Variant( + "type", + long_range_args_plugin.get_all_argument(), + doc="The type of the long-range module.", + ) + ], + optional=True, + doc=supported_backends("pt_expt") + doc_long_range, + ) + + # --- Modifier configurations: --- # modifier_args_plugin = ArgsPlugin() @@ -3843,6 +4046,7 @@ def standard_model_args() -> Argument: default={}, doc=supported_backends("pt", "jax", "pd", "pt_expt", "tf2") + doc_info, ), + long_range_args(), *_bridging_method_args(), ], doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index b5bdb2ea94..2fa225cc64 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -142,6 +142,9 @@ def _is_dpa4_family(sub: dict) -> bool: _LEARNED_CHILD_KEYS = ( "descriptor", "fitting_net", + # A long-range module belongs to the learned energy; the composition then + # rejects it beside the bridging term. + "long_range", "model_branch_alias", "info", "use_compile", diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 5dd3623cba..71f2e85834 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -490,6 +490,18 @@ def _apply_pretrained_model_params( single_config["descriptor"] = deepcopy(pretrained_config["descriptor"]) if not random_fitting: single_config["fitting_net"] = deepcopy(pretrained_config["fitting_net"]) + # A long-range module reads the descriptor beside the fitting and + # is carried over with it; the trainability set in the input + # applies to its parts as to the descriptor and the fitting. + if "long_range" in pretrained_config: + long_range = deepcopy(pretrained_config["long_range"]) + for part in ("charge_net", "kernel"): + long_range.setdefault(part, {})["trainable"] = ( + single_config.get("long_range", {}) + .get(part, {}) + .get("trainable", True) + ) + single_config["long_range"] = long_range log.info( "Change the '%s' model configurations according to pretrained branch '%s'.", model_branch, diff --git a/doc/model/index.rst b/doc/model/index.rst index c6db98d169..b3069c6e7b 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -13,6 +13,7 @@ Model dpa3 dpa4 dpa4c + long-range unimol train-hybrid sel diff --git a/doc/model/long-range.md b/doc/model/long-range.md new file mode 100644 index 0000000000..95e198caba --- /dev/null +++ b/doc/model/long-range.md @@ -0,0 +1,171 @@ +# Long-range latent charges {{ pytorch_icon }} + +> [!NOTE] +> **Supported backends**: PyTorch Exportable {{ pytorch_icon }} (`dp --pt-expt`) + +A short-range model sees each atom only through the neighbors within its +cutoff. Interactions that decay slowly with distance, such as those between +charged groups of a molecule, therefore escape it. The long-range module adds +such an interaction to the energy of a standard model. A second network reads +the same descriptor as the energy fitting and predicts per-atom _latent +charges_, and every pair of atoms of a system interacts through those charges: + +```math +E = \sum_i E_i^{\text{SR}} + k_e \sum_{i [!WARNING] +> The pair sum covers complete non-periodic systems: frames without a cell, +> such as isolated molecules and clusters. A periodic cell is rejected. The +> cost of the sum grows with the square of the number of atoms of a frame. + +## Configuration + +The module is a `long_range` section of a standard energy model, beside the +descriptor and the fitting. It is currently available with the +[DPA4C](./dpa4c.md) descriptor: + +```json +"model": { + "type_map": ["C", "H", "O"], + "descriptor": { + "type": "dpa4c", + "rcut": 6.0 + }, + "fitting_net": { + "neuron": [128, 128, 128] + }, + "long_range": { + "type": "latent_charge", + "channels": 1, + "kernel": { + "type": "les", + "alpha": 1.0 + }, + "charge_constraint": { + "source": "neutral" + } + } +} +``` + +The charge network, +{ref}`charge_net `, +takes the frame and atomic parameters of the energy fitting and follows its +activation function, skip connection and precision unless they are set; its +seed defaults to one derived from the seed of the fitting. +Its width is set by +{ref}`neuron `. + +### Kernels + +{ref}`kernel ` selects the +interaction of the charges. Both kernels are trained with the model unless +their `trainable` option is false, and their widths stay positive throughout +training. + +- `les`: $K(r) = \mathrm{erf}(\alpha r)/r$, the interaction of two Gaussian + charge clouds. It approaches $1/r$ beyond a few $1/\alpha$ and stays finite, + $2\alpha/\sqrt{\pi}$, at contact. The inverse width $\alpha$ is set by + {ref}`alpha `. +- `sog`: $K(r) = \sum_k a_k \exp(-r^2/2\sigma_k^2)$, a sum of Gaussians whose + amplitudes $a_k$ may take either sign. The widths are given by + {ref}`width ` + or default to a geometric sequence, and the amplitudes by + {ref}`amplitude `, + which explicit widths require. The default amplitudes, + $a_k = \sqrt{2/\pi}\,\ln b/\sigma_k$ for the ratio $b$ of consecutive + widths, discretize $1/r$: with the default widths the kernel stays finite, + $0.46$ Å⁻¹, at contact, and its radial force between fixed charges stays + within 0.4% of that of $1/r$ from 6.5 Å to 150 Å. It differs from $1/r$ + there by a nearly constant offset $C \approx 1.3 \times 10^{-3}$ Å⁻¹, which + exerts no force between fixed charges. The charges respond to the positions, + however, and through them the offset adds + $-k_e C \sum_c (Q_c^2 - \sum_i (q_i^c)^2)/2$ to the energy of a frame, with + $Q_c$ the total charge of channel $c$, which the response of the charges + turns into a force. Below about 5 Å the kernel falls short of $1/r$ by more + than $C$, by 0.12 Å⁻¹ at 2 Å. + +Several charge channels, set by +{ref}`channels `, let +the model represent interactions that one charge per atom cannot. + +### Total charge + +{ref}`charge_constraint ` +shifts one channel uniformly over the atoms of every frame so that it sums to +the total charge of the frame. The total charge is zero (`neutral`), the +charge of the frame charge/spin condition of the descriptor (`charge_spin`), +or a column of the frame parameters (`fparam`); a frame without the input +takes its default. The constrained channel then carries the total charge, +while the other channels remain free. Without the section, no channel is +constrained. + +### Exclusions + +Pairs of the types in `pair_exclude_types` see each other neither in the +descriptor nor in the long-range interaction, while the `exclude_types` of +the descriptor act on the descriptor alone. In a `linear_ener` composition, +`pair_exclude_types` belongs to the composition, whose sub-models share its +neighbor graph and pair topology. `atom_exclude_types` is not supported with +a long-range module. + +## Training and inference + +Training uses the usual energy loss; see +[training energy models](train-energy.md). With `enable_compile`, the +long-range interaction is part of the compiled model: the pairs of each frame +are listed outside the compiled region and enter it as inputs of dynamic +length. + +`dp --pt-expt freeze` exports the model, long-range module included, to +`.pt2` or `.pte`, and `dp --pt-expt compress` compresses its descriptor into a +`.pt2` archive while the networks of the long-range module stay as they are. +A trained checkpoint (`.pt`) and the frozen model are evaluated by `dp test` +and the Python inference interface; the C++ interface and LAMMPS, with +`pair_style deepmd` and `pair_style deepmd/kk`, read the `.pt2` archive. +Since the pair sum covers the complete system, LAMMPS runs such a model with +non-periodic boundaries (`f`, `s` or `m`) in every dimension, and rejects a +periodic one. A model that the dedicated operator described below evaluates, +a compressed DPA4C model frozen for a CUDA device, also runs on several MPI +ranks: every rank gathers the positions and charges of the atoms of all ranks +over MPI and evaluates the pairs of its own atoms, so that the ranks share +the cost of the pair sum. This requires DeePMD-kit built with MPI. Any other +model, like an uncompressed one, runs on a single rank. + +Compressing a descriptor conditioned on `charge_spin` specializes the model +to one charge state at a time, and a constraint that reads its total charge +from `charge_spin` follows the same state. The compressed archive serves any +charge state: the Python interface applies the state of each call, and LAMMPS +the one given by the `charge_spin` keyword of `pair_style deepmd`. + +A compressed DPA4C model frozen for a CUDA device evaluates the pair sum with +a dedicated operator, which reads the positions of the atoms and enumerates +their pairs itself: its memory grows linearly with the number of atoms, and +its cost with their square. The operator serves the `les` kernel and the +`sog` kernel of up to 32 Gaussians, with up to 8 charge channels, a charge +network without frame or atomic parameters, and no `pair_exclude_types`. Any +other model, like an uncompressed one, lists the pairs of each frame before +calling the model, so its memory grows with the square of the number of +atoms. + +The following are not supported for a model with a long-range module: +multi-rank (domain-decomposed) inference of a model that the dedicated +operator does not evaluate, such as an uncompressed one, spin models, the +tabulated ZBL potential, and short-range bridging. diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index d98d744387..a535e0c516 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -13,7 +13,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 30 +#define DP_C_API_VERSION 31 /** * @brief Neighbor list. @@ -479,6 +479,63 @@ extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp); */ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); +/** + * @brief Evaluate a compact canonical graph and the node positions on the + * model device. + * + * A model with a long-range module reads the positions of the nodes, which + * must form a complete non-periodic system: held whole by one rank + * (``nall_nodes == nloc``), or decomposed over the ranks of ``world``, whose + * owned nodes ``[0, nloc)`` are the share of this rank. With such a model on + * several ranks, every rank calls this function on every evaluation, also + * when it holds no node. Any other model ignores the positions and the + * communicator. + * + * @param[in] dp The DP to use. + * @param[out] d_atom_energy Per-local-atom energy, shape ``(nloc)``. + * @param[out] d_force Per-node force, shape ``(nall_nodes, 3)``. + * @param[out] d_atom_virial Per-node virial, shape ``(nall_nodes, 9)``. + * @param[in] d_atype Per-node atom types, shape ``(nall_nodes)``. + * @param[in] d_coord Node positions in model-node order, shape + * ``(nall_nodes, 3)``, or NULL for a model without a long-range module. + * @param[in] world The MPI communicator (a host pointer to an ``MPI_Comm``) + * of a system decomposed over several ranks, or NULL on a single rank. + * @param[in] d_source Source node per edge storage slot. + * @param[in] d_edge_vec FP32 edge vectors, shape ``(edge_storage, 3)``. + * @param[in] d_destination_row_ptr Destination CSR offsets. + * @param[in] d_source_row_ptr Source CSR offsets. + * @param[in] d_source_order Source-grouped edge storage positions. + * @param[in] nloc Number of owned local nodes. + * @param[in] nall_nodes Total local-plus-halo node count. + * @param[in] edge_storage Number of edge storage slots. + * @note Supersedes DP_DeepPotComputeCanonicalGraphGPU, adding the node + * positions and the communicator. + * @since API version 31 + */ +extern void DP_DeepPotComputeCanonicalGraphGPU2( + DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + void* world, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + int nloc, + int nall_nodes, + int64_t edge_storage); + +/** + * @brief Query whether the model carries a long-range module, which + * evaluates complete non-periodic systems. + * @since API version 31 + */ +extern bool DP_DeepPotHasLongRange(DP_DeepPot* dp); + /** * @brief Evaluate the energy, force and virial by using a DP. (double version) * @version 2 @@ -2683,6 +2740,14 @@ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp); */ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp); +/** + * @brief Query whether any model of a DP Model Deviation carries a long-range + * module, which evaluates complete non-periodic systems. + * @param[in] dp The DP Model Deviation to use. + * @since API version 31 + */ +extern bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp); + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models under diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 681b197b0e..c5153087f3 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1406,6 +1406,41 @@ class DeepPot : public DeepBaseModel { return result; } + /** + * @brief Evaluate a compact canonical graph and the node positions on the + * model device; see DP_DeepPotComputeCanonicalGraphGPU2. + */ + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + void* world, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_DeepPotComputeCanonicalGraphGPU2( + dp, d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + } + + /** + * @brief Query whether the model carries a long-range module, which + * evaluates complete non-periodic systems. + */ + bool has_long_range() const { + const bool result = DP_DeepPotHasLongRange(dp); + DP_CHECK_OK(DP_DeepPotCheckOK, dp); + return result; + } + /** * @brief Evaluate the energy, force and virial by using this DP. * @param[out] ener The system energy. @@ -2608,6 +2643,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return dchgspin; } + /** + * @brief Query whether any of the models carries a long-range module, which + * evaluates complete non-periodic systems. + */ + bool has_long_range() const { + const bool result = DP_DeepPotModelDeviHasLongRange(dp); + DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); + return result; + } + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index c35a4a1763..e6e2904295 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1864,6 +1864,37 @@ bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp) { } } +void DP_DeepPotComputeCanonicalGraphGPU2(DP_DeepPot* dp, + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const int64_t* d_atype, + const double* d_coord, + void* world, + const uint32_t* d_source, + const float* d_edge_vec, + const int64_t* d_destination_row_ptr, + const int64_t* d_source_row_ptr, + const uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const int64_t edge_storage) { + DP_REQUIRES_OK( + dp, dp->dp.compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, + d_source, d_edge_vec, d_destination_row_ptr, d_source_row_ptr, + d_source_order, nloc, nall_nodes, edge_storage)); +} + +bool DP_DeepPotHasLongRange(DP_DeepPot* dp) { + try { + return dp->dp.has_long_range(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + // multiple frames void DP_DeepPotCompute2(DP_DeepPot* dp, const int nframes, @@ -2848,6 +2879,15 @@ int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) { return dp->dp.dim_chg_spin(); } +bool DP_DeepPotModelDeviHasLongRange(DP_DeepPotModelDevi* dp) { + try { + return dp->dp.has_long_range(); + } catch (deepmd::deepmd_exception& ex) { + dp->exception = std::string(ex.what()); + return false; + } +} + void DP_DeepPotModelDeviSetChargeSpin(DP_DeepPotModelDevi* dp, const double* charge_spin, const int numb_chg_spin) { diff --git a/source/api_c/tests/test_deeppot_long_range_hpp.cc b/source/api_c/tests/test_deeppot_long_range_hpp.cc new file mode 100644 index 0000000000..9e6804e056 --- /dev/null +++ b/source/api_c/tests/test_deeppot_long_range_hpp.cc @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test the C API with a pt_expt (.pt2) DPA4C model that carries a +// latent-charge long-range module (source/tests/infer/gen_dpa4c_long_range.py). +#include + +#include +#include +#include + +#include "../../api_cc/tests/expected_ref.h" +#include "DeepPotPTExpt.h" +#include "deepmd.hpp" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kShortRangeModelPath = + "../../tests/infer/deeppot_sea.pth"; +} // namespace + +template +class TestInferDeepPotLongRangeHPP : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector expected_f; + double expected_tot_e = 0.; + deepmd::hpp::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init(kModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_f = ref.get("default", "expected_f"); + for (const VALUETYPE value : ref.get("default", "expected_e")) { + expected_tot_e += value; + } + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotLongRangeHPP, ValueTypes); + +TYPED_TEST(TestInferDeepPotLongRangeHPP, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); +} + +TEST(TestInferDeepPotModelDeviLongRangeHPP, has_long_range) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + // A long-range model anywhere in the ensemble confines all of them. + deepmd::hpp::DeepPotModelDevi mixed; + mixed.init(std::vector{kShortRangeModelPath, kModelPath}); + EXPECT_TRUE(mixed.has_long_range()); + deepmd::hpp::DeepPotModelDevi short_range; + short_range.init( + std::vector{kShortRangeModelPath, kShortRangeModelPath}); + EXPECT_FALSE(short_range.has_long_range()); +} + +TYPED_TEST(TestInferDeepPotLongRangeHPP, nopbc) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + ASSERT_EQ(force.size(), this->expected_f.size()); + EXPECT_LT(fabs(ener - this->expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++ii) { + EXPECT_LT(fabs(force[ii] - this->expected_f[ii]), epsilon); + } +} diff --git a/source/api_cc/include/BackendPlugin.h b/source/api_cc/include/BackendPlugin.h index 4149917414..272ec8e679 100644 --- a/source/api_cc/include/BackendPlugin.h +++ b/source/api_cc/include/BackendPlugin.h @@ -15,11 +15,17 @@ class DipoleChargeModifierBase; /** * @brief C ABI symbol names exported by backend plugin libraries. + * + * The suffix of a symbol is the version of the interface it serves: that of + * the DeepPot and DeepSpin symbols changes with the virtual interface of + * DeepPotBackend and DeepSpinBackend, so that a plugin built against another + * interface is refused when it is loaded rather than called through a table + * of another layout. **/ constexpr const char* DEEPMD_DEEPPOT_PLUGIN_CREATE_SYMBOL = - "deepmd_create_deeppot_backend_v1"; + "deepmd_create_deeppot_backend_v2"; constexpr const char* DEEPMD_DEEPPOT_PLUGIN_DELETE_SYMBOL = - "deepmd_delete_deeppot_backend_v1"; + "deepmd_delete_deeppot_backend_v2"; constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_CREATE_SYMBOL = "deepmd_create_deepspin_backend_v1"; constexpr const char* DEEPMD_DEEPSPIN_PLUGIN_DELETE_SYMBOL = diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 575c7924c9..41b03b4572 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -426,6 +426,8 @@ class DeepPotBackend : public DeepBaseModelBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -437,6 +439,7 @@ class DeepPotBackend : public DeepBaseModelBackend { virtual bool supports_device_edge_inference() const; virtual bool uses_fp32_edge_vectors() const; virtual bool uses_canonical_graph_inference() const; + virtual bool has_long_range() const; }; /** @@ -894,10 +897,57 @@ class DeepPot : public DeepBaseModel { */ bool uses_canonical_graph_inference() const; + /** + * @brief Whether the model carries a long-range module. + * + * Such a model sums its long-range interaction over every pair of atoms of + * a complete non-periodic system, so it evaluates a system without a + * periodic cell, held whole by one rank or, through the compact canonical + * lower, decomposed over the ranks of an MPI communicator; the device + * entries take the positions of its atoms. + */ + bool has_long_range() const; + + void compute_canonical_graph_gpu(double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage); + + /** + * @brief Evaluate a compact canonical graph together with the positions of + * its nodes. + * + * Every pointer except ``world`` resides on the model device; ``world`` + * points to the caller's ``MPI_Comm`` in host memory. A model with a + * long-range module reads the positions and the communicator; any other + * model ignores them. + * + * @param[in] d_coord Node positions in model-node order, shape + * ``(nall_nodes, 3)``, in the length unit of the model. A model with a + * long-range module requires them, for a system without a periodic cell. + * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of a + * system decomposed over several ranks, of which the owned nodes + * ``[0, nloc)`` are the share of this rank, or NULL on a single rank, + * whose node set then holds the complete system (``nall_nodes == nloc``). + * With a long-range module and several ranks, every rank calls this + * function on every evaluation, also when it holds no node. + * + * The other parameters are those of the overload without positions. + */ void compute_canonical_graph_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -959,6 +1009,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; }; + /** + * @brief Whether any of the models carries a long-range module, which + * confines all of them to complete non-periodic systems. + **/ + bool has_long_range() const { + return std::any_of(dps.begin(), dps.end(), + [](const std::shared_ptr& model) { + return model->has_long_range(); + }); + }; + /** * @brief Fix the charge/spin condition served for the rest of the run. * Applied to every model, so that the deviation is taken between models diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 4707ee9ce6..cfa0a551ca 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -377,6 +377,8 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -413,6 +415,7 @@ class DeepPotPTExpt : public DeepPotBackend { */ bool uses_fp32_edge_vectors() const override; bool uses_canonical_graph_inference() const override; + bool has_long_range() const override; private: template @@ -434,6 +437,8 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -473,6 +478,11 @@ class DeepPotPTExpt : public DeepPotBackend { bool lower_input_is_graph_ = false; bool lower_input_is_canonical_ = false; bool graph_edge_fp32_ = false; + // The model carries a long-range module, whose lower reads the all-pairs + // topology of the frame after its other inputs: the pair lists for the graph + // lower, the node positions and the communicator of the ranks for the + // compact canonical lower. + bool long_range_ = false; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -615,6 +625,9 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] destination_row_ptr Destination CSR offsets ``(N + 1,)``. * @param[in] source_order Source-grouped edge permutation ``(E,)``. * @param[in] source_row_ptr Source CSR offsets ``(N + 1,)``. + * @param[in] coord Positions of the nodes of the single frame, shape + * ``(N, 3)`` float64, from which the all-pairs topology of a model with a + * long-range module is built; read only for such a model. */ std::vector run_model_graph( const torch::Tensor& atype, @@ -629,8 +642,18 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin); + const torch::Tensor& charge_spin, + const torch::Tensor& coord); + /** + * @brief Run a compact canonical ``.pt2`` artifact. + * + * @param[in] coord Node positions, shape ``(N, 3)`` float64, appended for a + * model with a long-range module and ignored otherwise. + * @param[in] world The MPI communicator (a pointer to an ``MPI_Comm``) of the + * ranks a decomposed system spans, or NULL for a system that one rank holds + * whole; a model with a long-range module receives it after the positions. + */ std::vector run_model_canonical_graph( const torch::Tensor& atype, const torch::Tensor& n_node, @@ -639,7 +662,18 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order); + const torch::Tensor& source_order, + const torch::Tensor& coord, + const void* world); + + /** + * @brief Evaluate the long-range module on a rank that holds no node. + * + * The compact lower runs on one halo node without edges, which carries no + * charge and contributes nothing, so that the rank joins the exchanges of + * the module with the other ranks of ``world``. + */ + void join_long_range_exchange(const void* world); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 3373d6fe71..b8491393ee 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -661,6 +661,8 @@ void DeepPotBackend::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -673,6 +675,8 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_force; (void)d_atom_virial; (void)d_atype; + (void)d_coord; + (void)world; (void)d_source; (void)d_edge_vec; (void)d_destination_row_ptr; @@ -692,6 +696,8 @@ bool DeepPotBackend::supports_device_edge_inference() const { return false; } bool DeepPotBackend::uses_canonical_graph_inference() const { return false; } +bool DeepPotBackend::has_long_range() const { return false; } + void DeepPot::compute_edges_gpu(double* d_atom_energy, double* d_force, double* d_atom_virial, @@ -759,9 +765,30 @@ void DeepPot::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { dp->compute_canonical_graph_gpu( - d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, - d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, - edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, /*d_coord=*/nullptr, + /*world=*/nullptr, d_source, d_edge_vec, d_destination_row_ptr, + d_source_row_ptr, d_source_order, nloc, nall_nodes, edge_storage); +} + +void DeepPot::compute_canonical_graph_gpu( + double* d_atom_energy, + double* d_force, + double* d_atom_virial, + const std::int64_t* d_atype, + const double* d_coord, + void* world, + const std::uint32_t* d_source, + const float* d_edge_vec, + const std::int64_t* d_destination_row_ptr, + const std::int64_t* d_source_row_ptr, + const std::uint32_t* d_source_order, + const int nloc, + const int nall_nodes, + const std::int64_t edge_storage) { + dp->compute_canonical_graph_gpu( + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, + nall_nodes, edge_storage); } bool DeepPot::uses_fp32_edge_vectors() const { @@ -776,6 +803,8 @@ bool DeepPot::uses_canonical_graph_inference() const { return dp->uses_canonical_graph_inference(); } +bool DeepPot::has_long_range() const { return dp->has_long_range(); } + int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } void DeepPot::set_charge_spin(const std::vector& charge_spin) { diff --git a/source/api_cc/src/DeepPotJAXPlugin.cc b/source/api_cc/src/DeepPotJAXPlugin.cc index ac2d5485a1..2522148833 100644 --- a/source/api_cc/src/DeepPotJAXPlugin.cc +++ b/source/api_cc/src/DeepPotJAXPlugin.cc @@ -4,7 +4,7 @@ #include "BackendPluginFactory.h" #include "DeepPotJAX.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPDPlugin.cc b/source/api_cc/src/DeepPotPDPlugin.cc index fbc920cfb5..3bbd6842de 100644 --- a/source/api_cc/src/DeepPotPDPlugin.cc +++ b/source/api_cc/src/DeepPotPDPlugin.cc @@ -4,7 +4,7 @@ #include "BackendPluginFactory.h" #include "DeepPotPD.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -13,7 +13,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 005a18e722..a76cd96fa5 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -22,6 +22,7 @@ #include "commonPTExpt.h" #include "device.h" #include "errors.h" +#include "longRangePTExpt.h" #include "neighbor_list.h" using deepmd::ptexpt::check_call_charge_spin; @@ -110,6 +111,16 @@ FrameParameterLayout resolve_frame_parameter_layout(const char* name, " for frame-specific values."); } +// The communicator input of the compact lower of a long-range model: an int64 +// CPU tensor holding the address of the ``MPI_Comm`` that ``world`` points to, +// or zero for a system that one rank holds whole. The handle itself is read +// only by the operator, with the width of ``MPI_Comm`` in its MPI library. +at::Tensor communicator_tensor(const void* world) { + return torch::full( + {1}, static_cast(reinterpret_cast(world)), + torch::TensorOptions().dtype(torch::kInt64)); +} + } // namespace void DeepPotPTExpt::translate_error(std::function f) { @@ -256,6 +267,23 @@ void DeepPotPTExpt::init(const std::string& model, lower_input_is_graph_ = false; lower_input_is_canonical_ = false; } + // A model with a long-range module reads the all-pairs topology of the + // frame after its other inputs. + long_range_ = metadata.obj_val.count("topologies") > 0; + if (long_range_) { + for (const auto& name : metadata["topologies"].as_array()) { + if (name.as_string() != "all_pairs") { + throw deepmd::deepmd_exception("unknown topology '" + name.as_string() + + "' in the .pt2 metadata; expected " + "'all_pairs'."); + } + } + if (!lower_input_is_graph_ && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module requires the graph or the " + "compact canonical lower."); + } + } if (lower_input_is_edge_) { std::cerr << "WARNING: This .pt2 uses the deprecated edge_vec lower " "schema (pt-backend SeZM/DPA4 freeze). Support will be " @@ -513,7 +541,8 @@ std::vector DeepPotPTExpt::run_model_graph( const torch::Tensor& source_row_ptr, const torch::Tensor& fparam, const torch::Tensor& aparam, - const torch::Tensor& charge_spin) { + const torch::Tensor& charge_spin, + const torch::Tensor& coord) { // Graph-input ABI: original edge payload plus destination/source CSR views. check_graph_aparam_flat(aparam, daparam, "run_model_graph"); std::vector inputs = { @@ -529,6 +558,12 @@ std::vector DeepPotPTExpt::run_model_graph( if (dchgspin > 0) { inputs.push_back(charge_spin); } + if (long_range_) { + const auto pairs = deepmd::ptexpt::buildAllPairsTopology( + coord, atype, pair_exclude_table_, ntypes, + graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64); + inputs.insert(inputs.end(), pairs.begin(), pairs.end()); + } return loader->run(inputs); } @@ -540,9 +575,32 @@ std::vector DeepPotPTExpt::run_model_canonical_graph( const torch::Tensor& edge_vec, const torch::Tensor& destination_row_ptr, const torch::Tensor& source_row_ptr, - const torch::Tensor& source_order) { - return loader->run({atype, n_node, n_local, source, edge_vec, - destination_row_ptr, source_row_ptr, source_order}); + const torch::Tensor& source_order, + const torch::Tensor& coord, + const void* world) { + std::vector inputs = {atype, n_node, + n_local, source, + edge_vec, destination_row_ptr, + source_row_ptr, source_order}; + if (long_range_) { + inputs.push_back(coord); + inputs.push_back(communicator_tensor(world)); + } + return loader->run(inputs); +} + +void DeepPotPTExpt::join_long_range_exchange(const void* world) { + const torch::Device device(torch::kCUDA, gpu_id); + const auto i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto index_type = deepmd::canonicalGraphIndexType(); + // The compact lower stores at least two edge slots; without edges both are + // padding outside every CSR row. + run_model_canonical_graph( + torch::zeros({1}, i64), torch::ones({1}, i64), torch::zeros({1}, i64), + torch::zeros({2}, i64.dtype(index_type)), + torch::zeros({2, 3}, i64.dtype(torch::kFloat32)), torch::zeros({2}, i64), + torch::zeros({2}, i64), torch::arange(2, i64).to(index_type), + torch::zeros({1, 3}, i64.dtype(torch::kFloat64)), world); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -799,6 +857,34 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, bool multi_rank = (lmp_list.nprocs > 1); bool atom_map_present = (lmp_list.mapping != nullptr); bool use_with_comm = has_comm_artifact_ && multi_rank; + // A long-range module sums over every pair of atoms of the system. A single + // rank holds all of them, so a ghost atom there is a periodic image. Across + // ranks, the compact lower gathers the owned atoms of every rank through + // the communicator; the pairs that the graph lower lists stay within one. + if (long_range_ && !multi_rank && nghost_real > 0) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a complete non-periodic " + "system, but this single-rank call carries " + + std::to_string(nghost_real) + + " ghost atoms, the periodic images of a periodic cell; use a " + "non-periodic box."); + } + if (long_range_ && multi_rank && !lower_input_is_canonical_) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + + std::to_string(lmp_list.nprocs) + + " MPI ranks only through the compact canonical lower, which a " + "compressed DPA4C model frozen for a CUDA device uses when the " + "direct-sum operator serves its long-range module; this archive uses " + "another lower, so run it on a single rank."); + } + if (long_range_ && multi_rank && lmp_list.world == nullptr) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a system decomposed over " + "several MPI ranks through their communicator, but the neighbor list " + "carries none (InputNlist::world)."); + } + const void* const world = multi_rank ? lmp_list.world : nullptr; // Whether the edge topology folds ghost neighbours onto their local owners, // which reads an owner for every extended atom out of ``mapping_``. The // edge lower folds unless the with-comm artifact carries ghost features @@ -1227,7 +1313,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // ``Dim("n_node_total", min=1)``. Such a rank contributes nothing, so // fill zero outputs and return instead of running the model. (The // tested ``nloc == 0`` empty-subdomain case has ``nall_real > 0`` -- - // ghosts within rcut -- so it still runs the model normally.) + // ghosts within rcut -- so it still runs the model normally.) A + // long-range module across ranks still joins the exchanges of the + // others. + if (long_range_ && multi_rank) { + join_long_range_exchange(world); + } ener.assign(nframes, static_cast(0)); force.assign(static_cast(nframes) * fwd_map.size() * 3, static_cast(0)); @@ -1268,19 +1359,21 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, graph_pack.atype = node_atype; graph_pack.n_node = n_node_tensor; graph_pack.n_local = n_local_tensor; + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, n_node_count).reshape({n_node_count, 3}); if (lower_input_is_canonical_) { const auto compact = compactCanonicalGraph(graph_pack); flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, - compact.source_row_ptr, compact.source_order); + compact.source_row_ptr, compact.source_order, node_coord, world); } else { flat_outputs = run_model_graph( node_atype, n_node_tensor, n_local_tensor, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor); + graph_aparam, charge_spin_tensor, node_coord); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -1456,6 +1549,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, throw deepmd::deepmd_exception( "The coordinate-based inference interface requires at least one atom."); } + // The long-range module sums over the atoms of the system alone. + if (long_range_ && !box.empty()) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates complete non-periodic " + "systems only; pass an empty box."); + } const std::size_t frame_coord_size = static_cast(natoms) * 3; if (coord.empty() || coord.size() % frame_coord_size != 0) { throw deepmd::deepmd_exception( @@ -1648,12 +1747,14 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (graph_edge_fp32_) { graph_tensors.edge_vec = graph_tensors.edge_vec.to(torch::kFloat32); } + const at::Tensor node_coord = + coord_Tensor.slice(1, 0, nloc).reshape({nloc, 3}); if (lower_input_is_canonical_) { const auto compact = compactCanonicalGraph(graph_tensors); flat_outputs = run_model_canonical_graph( compact.atype, compact.n_node, compact.n_local, compact.source, compact.edge_vec, compact.destination_row_ptr, compact.source_row_ptr, - compact.source_order); + compact.source_order, node_coord, /*world=*/nullptr); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -1665,7 +1766,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // normalizes the rectangular runtime aparam to the flat // (N, daparam) node axis of the graph ABI. extend_graph_aparam(aparam_tensor, natoms, natoms, daparam), - charge_spin_tensor); + charge_spin_tensor, node_coord); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -2221,6 +2322,14 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, const int nnode = (nall_nodes > 0) ? nall_nodes : nloc; const bool extended = nnode != nloc; const bool use_with_comm = has_comm_artifact_ && comm_nlist; + if (long_range_ && (extended || comm_nlist != nullptr)) { + throw deepmd::deepmd_exception( + "compute_edges_gpu: a model with a long-range module evaluates a " + "complete non-periodic system on a single rank (nall_nodes == nloc and " + "no comm_nlist); several ranks require the compact canonical lower, " + "which a compressed DPA4C model frozen for a CUDA device uses when " + "the direct-sum operator serves its long-range module."); + } if (lower_input_is_graph_ && extended && has_message_passing_) { throw deepmd::deepmd_exception( "compute_edges_gpu: message-passing graph-input inference requires an " @@ -2448,7 +2557,7 @@ void DeepPotPTExpt::compute_edges_gpu_impl(double* d_atom_energy, graph_pack.edge_index, graph_pack.edge_vec, graph_pack.edge_mask, graph_pack.destination_order, graph_pack.destination_row_ptr, graph_pack.source_order, graph_pack.source_row_ptr, fparam_tensor, - graph_aparam, charge_spin_tensor)); + graph_aparam, charge_spin_tensor, coord_t.reshape({nloc, 3}))); ae = out["atom_energy"].reshape({nnode}).slice(0, 0, nloc).contiguous(); force_t = out["force"].reshape({nnode, 3}).contiguous(); av = out["atom_virial"].reshape({nnode, 9}).contiguous(); @@ -2537,6 +2646,8 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2553,6 +2664,29 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( throw deepmd::deepmd_exception( "compute_canonical_graph_gpu requires a CUDA device."); } + // A long-range module reads the positions of a complete non-periodic + // system: the node set of a single rank, or the owned nodes of every rank + // of ``world``, which a rank without nodes still joins. + if (long_range_ && d_coord == nullptr && nall_nodes > 0) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module reads " + "the node positions (d_coord)."); + } + if (long_range_ && world == nullptr && nall_nodes != nloc) { + throw deepmd::deepmd_exception( + "compute_canonical_graph_gpu: a model with a long-range module on a " + "single rank evaluates a complete non-periodic system " + "(nall_nodes == nloc); a system decomposed over several ranks takes " + "their communicator."); + } + if (long_range_ && world != nullptr && nall_nodes == 0) { + translate_error([&] { + const c10::DeviceGuard device_guard(torch::Device(torch::kCUDA, gpu_id)); + join_long_range_exchange(world); + synchronize_current_accelerator_stream(); + }); + return; + } if (nloc < 0 || nall_nodes <= 0 || nloc > nall_nodes || edge_storage < 2 || static_cast(edge_storage) > std::numeric_limits::max()) { @@ -2585,12 +2719,18 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl( const_cast(d_source_order), {edge_storage}, opt_u32); auto n_node = torch::full({1}, nall_nodes, opt_i64); auto n_local = torch::full({1}, nloc, opt_i64); + const at::Tensor coord = + d_coord != nullptr + ? torch::from_blob( + const_cast(d_coord), {nall_nodes, 3}, + torch::TensorOptions().dtype(torch::kFloat64).device(device)) + : at::Tensor(); std::map output; - extract_outputs(output, - run_model_canonical_graph(atype, n_node, n_local, source, - edge_vec, destination_row_ptr, - source_row_ptr, source_order)); + extract_outputs(output, run_model_canonical_graph( + atype, n_node, n_local, source, edge_vec, + destination_row_ptr, source_row_ptr, + source_order, coord, world)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2672,6 +2812,8 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2681,9 +2823,9 @@ void DeepPotPTExpt::compute_canonical_graph_gpu( const int nall_nodes, const std::int64_t edge_storage) { compute_canonical_graph_gpu_impl( - d_atom_energy, d_force, d_atom_virial, d_atype, d_source, d_edge_vec, - d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, nall_nodes, - edge_storage); + d_atom_energy, d_force, d_atom_virial, d_atype, d_coord, world, d_source, + d_edge_vec, d_destination_row_ptr, d_source_row_ptr, d_source_order, nloc, + nall_nodes, edge_storage); } bool DeepPotPTExpt::uses_fp32_edge_vectors() const { @@ -2700,4 +2842,6 @@ bool DeepPotPTExpt::uses_canonical_graph_inference() const { return lower_input_is_canonical_; } +bool DeepPotPTExpt::has_long_range() const { return long_range_; } + #endif diff --git a/source/api_cc/src/DeepPotPTExptPlugin.cc b/source/api_cc/src/DeepPotPTExptPlugin.cc index 3499c62dcb..6843dc86fc 100644 --- a/source/api_cc/src/DeepPotPTExptPlugin.cc +++ b/source/api_cc/src/DeepPotPTExptPlugin.cc @@ -57,7 +57,7 @@ bool native_spin_backend_serves(const char* model) { } // namespace #endif -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -74,7 +74,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, #endif } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotPTPlugin.cc b/source/api_cc/src/DeepPotPTPlugin.cc index ae584b12b6..80050bb2d2 100644 --- a/source/api_cc/src/DeepPotPTPlugin.cc +++ b/source/api_cc/src/DeepPotPTPlugin.cc @@ -6,7 +6,7 @@ #include "DeepSpinPT.h" #include "DeepTensorPT.h" -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -15,7 +15,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/DeepPotTFPlugin.cc b/source/api_cc/src/DeepPotTFPlugin.cc index 5e20b6fcf7..16c3a1599c 100644 --- a/source/api_cc/src/DeepPotTFPlugin.cc +++ b/source/api_cc/src/DeepPotTFPlugin.cc @@ -46,7 +46,7 @@ void convert_pbtxt_to_pb_impl(const char* pbtxt, const char* pb) { } // namespace -extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, +extern "C" void* deepmd_create_deeppot_backend_v2(const char* model, int gpu_rank, const char* file_content, std::size_t file_content_size, @@ -55,7 +55,7 @@ extern "C" void* deepmd_create_deeppot_backend_v1(const char* model, model, gpu_rank, file_content, file_content_size, error_message); } -extern "C" void deepmd_delete_deeppot_backend_v1(void* backend) { +extern "C" void deepmd_delete_deeppot_backend_v2(void* backend) { deepmd::plugin::delete_deeppot_backend(backend); } diff --git a/source/api_cc/src/longRangePTExpt.h b/source/api_cc/src/longRangePTExpt.h new file mode 100644 index 0000000000..cbb37f49f9 --- /dev/null +++ b/source/api_cc/src/longRangePTExpt.h @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include + +#include "commonPT.h" + +namespace deepmd { +namespace ptexpt { + +/** + * @brief Build the all-pairs topology a long-range module of the graph lower + * reads. + * + * Inference-path twin of Python ``build_pair_graph`` in + * ``deepmd/dpmodel/utils/neighbor_graph/pair_graph.py`` for a single frame: + * every unordered pair of nodes once, ``[src, dst]`` with ``src < dst``, its + * vector ``r_src - r_dst``, and two masked guard pairs that keep the pair axis + * at the lower bound of the exported lower. The model-level pair exclusion is + * applied to the mask, as it is to the neighbor graph. + * + * @param coord Node positions, shape ``(N, 3)``, float64, on the model device. + * @param atype Node types, shape ``(N,)``, int64. + * @param type_mask_table Device keep table of the pair exclusion, or an + * undefined tensor when the model excludes no pair. + * @param ntypes Number of atom types. + * @param vec_dtype Precision of the pair vectors, that of the edge vectors of + * the artifact. + * @return ``{pair_index (2, P + 2) int64, pair_vec (P + 2, 3), pair_mask + * (P + 2,) bool}``. + */ +inline std::vector buildAllPairsTopology( + const torch::Tensor& coord, + const torch::Tensor& atype, + const torch::Tensor& type_mask_table, + const int ntypes, + const torch::ScalarType vec_dtype) { + const std::int64_t node_count = coord.size(0); + const auto index_options = + torch::TensorOptions().dtype(torch::kInt64).device(coord.device()); + const auto mask_options = + torch::TensorOptions().dtype(torch::kBool).device(coord.device()); + const torch::Tensor index = + torch::triu_indices(node_count, node_count, 1, index_options); + const torch::Tensor vec = (coord.index_select(0, index.select(0, 0)) - + coord.index_select(0, index.select(0, 1))) + .to(vec_dtype); + const torch::Tensor mask = + applyPairExclusion(index, torch::ones({index.size(1)}, mask_options), + atype, type_mask_table, ntypes); + return {torch::cat({index, torch::zeros({2, 2}, index_options)}, 1), + torch::cat({vec, torch::zeros({2, 3}, vec.options())}, 0), + torch::cat({mask, torch::zeros({2}, mask_options)})}; +} + +} // namespace ptexpt +} // namespace deepmd diff --git a/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc new file mode 100644 index 0000000000..fa4df07ad4 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ inference for pt_expt (.pt2) DPA4C models with a latent-charge +// long-range module, which sums over every pair of atoms of a complete +// non-periodic system. Models generated by +// source/tests/infer/gen_dpa4c_long_range.py: an uncompressed archive on the +// graph lower, which receives the pair list of the atoms, and a compressed +// archive on the compact canonical lower, which reads their positions. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +namespace { +constexpr const char* kModelPath = + "../../tests/infer/deeppot_dpa4c_long_range.pt2"; +constexpr const char* kRefPath = + "../../tests/infer/deeppot_dpa4c_long_range.expected"; +constexpr const char* kCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.pt2"; +constexpr const char* kCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_compressed.expected"; + +// The full neighbor list of a non-periodic system within ``rc``: every atom +// is local, so the list holds no ghost. +template +std::vector> nopbc_nlist(const std::vector& coord, + const double rc) { + const int natoms = static_cast(coord.size() / 3); + std::vector> nlist(natoms); + for (int ii = 0; ii < natoms; ++ii) { + for (int jj = 0; jj < natoms; ++jj) { + double r2 = 0.; + for (int dd = 0; dd < 3; ++dd) { + const double diff = coord[ii * 3 + dd] - coord[jj * 3 + dd]; + r2 += diff * diff; + } + if (jj != ii && r2 < rc * rc) { + nlist[ii].push_back(jj); + } + } + } + return nlist; +} +} // namespace + +// One reference case of an archive: per-atom energies, forces and virials, +// with the totals the interface returns. +template +struct LongRangeRef { + std::vector expected_f; + double expected_tot_e = 0.; + std::vector expected_tot_v = std::vector(9, 0.); + + void load(const deepmd_test::ExpectedRef& ref, const char* section) { + const std::vector atom_e = + ref.template get(section, "expected_e"); + const std::vector atom_v = + ref.template get(section, "expected_v"); + expected_f = ref.template get(section, "expected_f"); + for (std::size_t ii = 0; ii < atom_e.size(); ++ii) { + expected_tot_e += atom_e[ii]; + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v[dd] += atom_v[ii * 9 + dd]; + } + } + } + + void check(double ener, + const std::vector& force, + const std::vector& virial, + double epsilon) const { + ASSERT_EQ(force.size(), expected_f.size()); + ASSERT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), epsilon); + for (std::size_t ii = 0; ii < force.size(); ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), epsilon); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), epsilon); + } + } +}; + +template +class TestInferDeepPotDpa4cLongRangePtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector charge_spin_explicit = {1.0, 2.0}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + + static deepmd::DeepPot dp; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp.init(kModelPath); +#endif + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + static void TearDownTestSuite() { dp = deepmd::DeepPot(); } + + // Computes through the LAMMPS-style entry point, with the full neighbor + // list of the system and no ghost atom. + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& charge_spin) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0, {}, {}, charge_spin); + } +}; + +template +deepmd::DeepPot TestInferDeepPotDpa4cLongRangePtExpt::dp; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, has_long_range) { + EXPECT_TRUE(this->dp.has_long_range()); + EXPECT_EQ(this->dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_default) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector()); + this->ref_default.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, standalone_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(), {}, {}, + this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, lmp_nlist_explicit) { + using VALUETYPE = TypeParam; + const double epsilon = std::is_same::value ? 1e-8 : 1e-4; + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial, this->charge_spin_explicit); + this->ref_explicit.check(ener, force, virial, epsilon); +} + +// A cell would leave the pair sum incomplete: its periodic images are no +// part of it. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePtExpt, periodic_box_is_rejected) { + using VALUETYPE = TypeParam; + const std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + double ener; + std::vector force, virial; + EXPECT_THROW( + this->dp.compute(ener, force, virial, this->coord, this->atype, box), + deepmd::deepmd_exception); +} + +// The compressed archive folds the charge state into constants, which +// set_charge_spin rewrites: the tables of the descriptor and the total charge +// of the constraint together. +template +class TestInferDeepPotDpa4cLongRangeCompressedPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 0.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + LongRangeRef ref_default; + LongRangeRef ref_explicit; + deepmd::DeepPot dp; + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(kCompressedModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kCompressedModelPath + << " was not generated (gen_dpa4c_long_range.py writes it " + "on a CUDA device only)."; + } + dp.init(kCompressedModelPath); + deepmd_test::ExpectedRef ref; + ref.load(kCompressedRefPath); + ref_default.load(ref, "default"); + ref_explicit.load(ref, "explicit"); + } + + void compute_lmp(double& ener, + std::vector& force, + std::vector& virial) { + std::vector> nlist_data = nopbc_nlist(coord, dp.cutoff()); + const int nloc = static_cast(atype.size()); + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + dp.compute(ener, force, virial, coord, atype, std::vector(), 0, + inlist, 0); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangeCompressedPtExpt, + set_charge_spin_follows_the_state) { + using VALUETYPE = TypeParam; + constexpr double epsilon = 1e-4; + EXPECT_TRUE(this->dp.has_long_range()); + double ener; + std::vector force, virial; + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({1.0, 2.0}); + this->compute_lmp(ener, force, virial); + this->ref_explicit.check(ener, force, virial, epsilon); + this->dp.set_charge_spin({0.0, 1.0}); + this->compute_lmp(ener, force, virial); + this->ref_default.check(ener, force, virial, epsilon); +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 8fc4fcc4d9..2751ee7838 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -150,6 +150,12 @@ else: env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_spin.py & PID15=$! fi + # DPA4C with a latent-charge long-range module: the uncompressed archive + # takes the all-pairs topology of the graph lower; for the CUDA variant the + # script also writes the compressed archive, whose compact canonical lower + # reads the node positions. + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_long_range.py & + PID16=$! wait $PID11 wait $PID12 wait $PID13 @@ -158,6 +164,7 @@ else: if [ "${DP_VARIANT}" = "cuda" ]; then wait $PID15 fi + wait $PID16 fi if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index a03482a90a..7dd715f48c 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -966,6 +966,22 @@ void PairDeepMD::settings(int narg, char** arg) { all_force.resize(numb_models); } +void PairDeepMD::init_style() { + PairDeepBaseModel::init_style(); + // A long-range module sums over every pair of atoms of the system, which a + // periodic image would leave incomplete; the atoms of the other ranks reach + // it through the communicator. + const bool long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() + : deep_pot.has_long_range(); + if (long_range && + (domain->xperiodic || domain->yperiodic || domain->zperiodic)) { + error->all(FLERR, + "pair style deepmd: a model with a long-range module evaluates " + "a complete non-periodic system; use a non-periodic boundary " + "(f, s or m) in every dimension."); + } +} + /* ---------------------------------------------------------------------- set coeffs for one or more type pairs ------------------------------------------------------------------------- */ diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index c7c40d9b18..ac8881b553 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -46,6 +46,7 @@ class PairDeepMD : public PairDeepBaseModel { ~PairDeepMD() override; void settings(int, char**) override; void coeff(int, char**) override; + void init_style() override; void compute(int, int) override; int pack_reverse_comm(int, int, double*) override; void unpack_reverse_comm(int, int*, double*) override; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index f54873dbe1..9c24685e1a 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -34,6 +34,7 @@ PairDeepMDKokkos::PairDeepMDKokkos(LAMMPS* lmp) edge_capacity(0), edge_vec_fp32(false), canonical_graph(false), + long_range(false), device_path_ok(false), reverse_virial(false), reverse_used_host(false) { @@ -212,6 +213,10 @@ void PairDeepMDKokkos::init_style() { request->enable_full(); edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); canonical_graph = deep_pot.uses_canonical_graph_inference(); + // PairDeepMD::init_style() has confined such a model to a complete + // non-periodic system, whose atoms on other ranks it reaches through the + // communicator. + long_range = deep_pot.has_long_range(); // Force exchange transfers three values per atom. Centroid per-atom virial // uses nine values even though the Kokkos full-list request runs newton off, // so comm_reverse_off reserves the classic host buffer for the wider mode. @@ -408,6 +413,35 @@ int PairDeepMDKokkos::build_edges_device() { return nedge; } +template +const double* PairDeepMDKokkos::model_node_positions() { + atomKK->sync(execution_space, X_MASK); + auto x = atomKK->k_x.template view(); + // Without virtual atoms the model nodes are the leading local atoms, already + // in the length unit of the model. + if (!compact_graph.has_null_types && dist_unit_cvt_factor == 1.0) { + return x.data(); + } + const int nnode_model = compact_graph.nnode_model; + if ((int)d_node_position.extent(0) < 3 * nnode_model) { + d_node_position = Kokkos::View( + "deepmd/kk:node_position", 3 * nnode_model); + } + auto node_position = d_node_position; + auto model2loc = compact_graph.d_model2loc; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + Kokkos::parallel_for( + "deepmd/kk:node_position", + Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + const int i = model2loc(m); + node_position(3 * m + 0) = x(i, 0) * inv_dist; + node_position(3 * m + 1) = x(i, 1) * inv_dist; + node_position(3 * m + 2) = x(i, 2) * inv_dist; + }); + return node_position.data(); +} + template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -520,7 +554,10 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { comm_ptr = &comm_list; } - if ((canonical_graph && nnode_m > 0) || + // A long-range module across ranks exchanges the charges of every rank, so + // that every rank evaluates the model, also without nodes. + const bool long_range_across_ranks = long_range && multi_rank; + if ((canonical_graph && (nnode_m > 0 || long_range_across_ranks)) || (!canonical_graph && (nloc_m > 0 || comm_ptr != nullptr))) { // Fully device-resident inference: raw device pointers in and out. The // edge buffers are produced on the Kokkos stream and consumed by the model @@ -528,18 +565,22 @@ void PairDeepMDKokkos::compute(int eflag, int vflag) { // the two runtimes are bracketed by explicit synchronization: fence the // Kokkos work before the model reads the edges, and synchronize the device // after so the scatter sees the finished model outputs. - Kokkos::fence(); // Coordinates are model-node order: the compacted buffer when virtual atoms // are present, else the local coordinates directly (the graph lower ignores - // them; edge-input models consume them). + // them; edge-input models consume them). A long-range module reads the + // positions of its nodes in the length unit of the model. const double* coord_ptr = - compact_graph.has_null_types ? d_coord_model.data() : x.data(); + long_range + ? model_node_positions() + : (compact_graph.has_null_types ? d_coord_model.data() : x.data()); + Kokkos::fence(); try { if (canonical_graph) { deep_pot.compute_canonical_graph_gpu( d_atom_energy.data(), d_out_force.data(), d_atom_virial.data(), - compact_graph.d_model_type.data(), compact_graph.d_source.data(), - compact_graph.d_edge_vec.data(), + compact_graph.d_model_type.data(), long_range ? coord_ptr : nullptr, + long_range_across_ranks ? static_cast(&world) : nullptr, + compact_graph.d_source.data(), compact_graph.d_edge_vec.data(), compact_graph.d_destination_row_ptr.data(), compact_graph.d_source_row_ptr.data(), compact_graph.d_source_order.data(), nloc_m, nnode_m, diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index bd3689aedf..65a1ce315f 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -82,6 +82,10 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { // extended device lambdas, which CUDA forbids inside non-public members. int build_edges_device(); + // Positions of the model nodes in the length unit of the model, which a + // long-range module reads; public for the same reason. + const double* model_node_positions(); + protected: // Model node set and, for a compact canonical artifact, the graph itself. CompactCanonicalGraphKokkos compact_graph; @@ -90,6 +94,8 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { bool multi_rank; // domain-decomposed run -> extended (local+ghost) node set Kokkos::View d_coord_model; // (3 * nnode_model), NULL case + Kokkos::View + d_node_position; // (3 * nnode_model), long-range model positions // Compact edge schema: edge_index is [2 * nedge] (src rows then dst rows), // edge_vec is [3 * nedge]; offsets is the per-atom exclusive edge prefix. @@ -121,6 +127,7 @@ class PairDeepMDKokkos : public PairDeepMD, public KokkosBase { int edge_capacity; // allocated edges in d_edge_index / d_edge_vec bool edge_vec_fp32; // model graph ABI consumes edge vectors in fp32 bool canonical_graph; // compact source-only graph artifact + bool long_range; // model reads the positions of every node bool device_path_ok; // resolved once in init_style bool reverse_virial; // reverse communication operates on centroid virial bool reverse_used_host; // force reverse communication selected host staging diff --git a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py index a574cef821..a73908d6f3 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_dpa3_pt2.py @@ -51,6 +51,13 @@ "domain decomposition (nswap>0). Pass '1 1 1' for a single-rank " "reference run on the same archive (single-artifact dispatch).", ) +parser.add_argument( + "--boundary", + type=str, + default="p p p", + help="LAMMPS boundary style. Default 'p p p'; a model with a long-range " + "module evaluates complete non-periodic systems and passes 'f f f'.", +) parser.add_argument( "--pair-coeff", type=str, @@ -137,7 +144,7 @@ # (the regular artifact handles nswap==0). lammps.processors(args.processors) lammps.units("metal") -lammps.boundary("p p p") +lammps.boundary(args.boundary) lammps.atom_style("atomic") # ``atom_modify map yes`` is required when single-rank dispatch goes # through the regular artifact of a use_loc_mapping=False .pt2: the diff --git a/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py new file mode 100644 index 0000000000..5241bd8e5e --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test LAMMPS with a .pt2 DPA4C model that carries a long-range module. + +The latent-charge module sums over every pair of atoms of a complete +non-periodic system, so ``pair_style deepmd`` runs it with non-periodic +boundaries and rejects a periodic one. A compressed model also runs on +several ranks, whose atoms it gathers over MPI; an uncompressed one runs on a +single rank. The models and the reference values, without a cell, at the +default charge state and at ``charge_spin 1.0 2.0``, are generated by +source/tests/infer/gen_dpa4c_long_range.py, the compressed model for CUDA +only. +""" + +import importlib.util +import os +import shutil +from pathlib import ( + Path, +) + +import constants +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from lammps_test_utils import ( + make_atomic_lammps, + run_mpi_pair_runner, +) +from write_lmp_data import ( + write_lmp_data, +) + +pt2_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.pt2" +) +ref_file = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa4c_long_range.expected" +) +compressed_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_compressed.pt2") +# A model without a long-range module, for a model deviation beside one. +short_range_file = pt2_file.with_name("deeppot_sea.pth") +compressed_ref_file = ref_file.with_name("deeppot_dpa4c_long_range_compressed.expected") +data_file = Path(__file__).parent / "data_dpa4c_long_range_pt2.lmp" +replicated_data_file = Path(__file__).parent / "data_dpa4c_long_range_grid_pt2.lmp" +mpi_runner = Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py" + +# Guarded with try/except because the generator only runs when PyTorch is +# built; matrices that disable PyTorch still load this file at collection time. +try: + _ref = read_expected_ref(ref_file) + expected_e_default = float(np.sum(_ref["default"]["expected_e"])) + expected_f_default = _ref["default"]["expected_f"].reshape(6, 3) + expected_e_explicit = float(np.sum(_ref["explicit"]["expected_e"])) + expected_f_explicit = _ref["explicit"]["expected_f"].reshape(6, 3) +except FileNotFoundError: + expected_e_default = expected_f_default = None + expected_e_explicit = expected_f_explicit = None +try: + compressed_ref = read_expected_ref(compressed_ref_file) +except FileNotFoundError: + compressed_ref = None + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] +type_OH = np.array([1, 2, 2, 1, 2, 2]) +# LAMMPS centroid/stress/atom reports the per-atom stress, the negative of the +# per-atom virial, with the components xx, yy, zz, xy, xz, yz, yx, zx, zy; these +# columns of it give the nine components of the reference virial. +expected_v_to_lammps = [0, 6, 7, 3, 1, 8, 4, 5, 2] + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + write_lmp_data(box, coord, type_OH, data_file) + + +def teardown_module() -> None: + for path in (data_file, replicated_data_file): + if path.exists(): + os.remove(path) + + +def _lammps(boundary: str) -> PyLammps: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + return make_atomic_lammps(data_file, boundary=boundary, atom_map="yes") + + +@pytest.fixture +def lammps(): + lmp = _lammps("f f f") + yield lmp + lmp.close() + + +@pytest.fixture +def lammps_periodic(): + lmp = _lammps("p p p") + yield lmp + lmp.close() + + +def test_pair_deepmd(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_pair_deepmd_charge_spin(lammps) -> None: + lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_explicit) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_explicit[lammps.atoms[ii].id - 1] + ) + + +def test_periodic_boundary_is_rejected(lammps_periodic) -> None: + lammps_periodic.pair_style(f"deepmd {pt2_file.resolve()}") + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) + + +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_periodic_boundary_is_rejected( + lammps_periodic, long_range_first: bool, tmp_path: Path +) -> None: + """Any model of a model deviation with a long-range module confines it.""" + if not short_range_file.is_file(): + pytest.skip(f"Skip because {short_range_file} was not generated.") + models = [pt2_file, short_range_file] + if not long_range_first: + models.reverse() + lammps_periodic.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps_periodic.pair_coeff("* *") + with pytest.raises(Exception, match="long-range module"): + lammps_periodic.run(0) + + +mpi_available = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="MPI or mpi4py is not installed", +) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "section", "pair_extra"), + [ + # Both ranks own atoms. + (2, "2 1 1", "default", ""), + # Every atom lies in the first slab along z: two ranks hold ghost atoms + # alone and the last none at all, and all of them join the exchange. + (4, "1 1 4", "explicit", "charge_spin 1.0 2.0"), + ], +) +def test_pair_deepmd_mpi_compressed( + nprocs: int, processors: str, section: str, pair_extra: str +) -> None: + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + runner_args = ["--boundary", "f f f"] + if pair_extra: + runner_args += ["--pair-extra", pair_extra] + result = run_mpi_pair_runner( + mpi_runner, + data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=runner_args, + ) + expected_e = float(np.sum(compressed_ref[section]["expected_e"])) + expected_f = compressed_ref[section]["expected_f"].reshape(6, 3) + expected_v = -compressed_ref[section]["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(expected_e, rel=1e-6) + np.testing.assert_allclose(result["forces"], expected_f, rtol=0, atol=1e-5) + np.testing.assert_allclose( + result["virials"][:, expected_v_to_lammps] / constants.nktv2p, + expected_v, + rtol=0, + atol=1e-5, + ) + + +@pytest.fixture(scope="module") +def replicated_reference() -> dict: + """A grid of 168 copies of the six atoms, and its single-rank result. + + The copies lie farther apart than the short-range cutoff, so only the + long-range module couples them; the 1008 atoms fill several blocks and + chunks of the direct-sum operator. + """ + if compressed_ref is None or not compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {compressed_pt2_file} was not generated.") + shifts = np.stack( + np.meshgrid( + 18.0 * np.arange(4), 6.0 * np.arange(7), 6.0 * np.arange(6), indexing="ij" + ), + axis=-1, + ).reshape(-1, 1, 3) + grid_coord = (coord[None] + shifts).reshape(-1, 3) + grid_type = np.tile(type_OH, shifts.shape[0]) + # The box reaches beyond the atoms along z, so that the last of four slabs + # holds no atom, owned or ghost. + grid_box = np.array([0, 72, 0, 42, 0, 60, 0, 0, 0]) + write_lmp_data(grid_box, grid_coord, grid_type, replicated_data_file) + return run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=1, + runner_args=["--boundary", "f f f"], + ) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors"), + [ + # Uneven slabs that cut through copies of the six atoms. + (3, "3 1 1"), + # The last slab holds no atom at all. + (4, "1 1 4"), + ], +) +def test_pair_deepmd_mpi_compressed_matches_single_rank( + replicated_reference: dict, nprocs: int, processors: str +) -> None: + result = run_mpi_pair_runner( + mpi_runner, + replicated_data_file, + compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", "f f f"], + ) + assert result["pe"] == pytest.approx(replicated_reference["pe"], rel=1e-6) + np.testing.assert_allclose( + result["forces"], replicated_reference["forces"], rtol=0, atol=1e-5 + ) + np.testing.assert_allclose( + result["virials"] / constants.nktv2p, + replicated_reference["virials"] / constants.nktv2p, + rtol=0, + atol=1e-5, + ) + + +@mpi_available +def test_pair_deepmd_mpi_uncompressed_is_rejected() -> None: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + result = run_mpi_pair_runner( + mpi_runner, + data_file, + pt2_file, + nprocs=2, + processors="2 1 1", + runner_args=["--boundary", "f f f"], + capture=True, + ) + assert result["returncode"] != 0 + assert "compact canonical lower" in result["stdout"] + result["stderr"] diff --git a/source/op/pt/CMakeLists.txt b/source/op/pt/CMakeLists.txt index 5702e6c18c..933f0fe7c9 100644 --- a/source/op/pt/CMakeLists.txt +++ b/source/op/pt/CMakeLists.txt @@ -23,7 +23,8 @@ list( cpu/neighbor_search_cpu.cc dpa4c/ops.cc dpa4c/graph_compress_cpu.cc - ${DPA4C_CPU_ISA_SRC}) + ${DPA4C_CPU_ISA_SRC} + long_range/ops.cc) # The architecture level is pinned per unit rather than added to it: the project # may be configured with ENABLE_NATIVE_OPTIMIZATION, and a trailing -march wins, # so this is what keeps the fallback units free of instructions the dispatcher @@ -179,6 +180,8 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC} graph_fitting.cu edge_force_virial.cu + long_range/direct_sum.cu + long_range/domain_exchange.cc dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} @@ -245,6 +248,11 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) set_source_files_properties( ${DPA1_GRAPH_COMPRESS_KERNEL_SRC} ${DPA4C_GRAPH_COMPRESS_KERNEL_SRC} PROPERTIES COMPILE_OPTIONS "--use_fast_math") + # The direct sum of the long range evaluates every pair with the special + # function unit; without denormal handling each of its three calls is a single + # instruction. + set_source_files_properties(long_range/direct_sum.cu + PROPERTIES COMPILE_OPTIONS "--use_fast_math") endif() if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI}) target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}) diff --git a/source/op/pt/long_range/direct_sum.cu b/source/op/pt/long_range/direct_sum.cu new file mode 100644 index 0000000000..c12e144537 --- /dev/null +++ b/source/op/pt/long_range/direct_sum.cu @@ -0,0 +1,1272 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernel of the latent-charge long-range direct sum (schema and +// formulas in ops.cc). +// +// Every node reads every node of its frame, so the memory is linear in the +// node count while the work is quadratic. One thread owns one node and walks +// the nodes of its frame in tiles staged in shared memory; the partner of a +// thread at a given step is the same for the whole block, so each tile entry +// is a broadcast read. Each pair is evaluated by both of its nodes, which +// keeps every sum in the registers of its owner: no atomic operation is +// needed, and the result does not depend on scheduling. +// +// Positions are carried relative to the first node of their frame as a +// double-float pair: a float32 head and the float32 tail of the remainder. +// The displacement of two nodes, formed as the difference of the heads plus +// the difference of the tails, keeps the float32 precision of the +// displacement itself wherever the frame lies, while every pair evaluates in +// float32; the far field of the erf kernel forms it from the heads alone (see +// evaluate_kernel). For the erf kernel the positions are scaled by alpha, so +// that the pair loop evaluates erf(x) / x and its slope without a +// coefficient, and the reduction restores the units. +// +// A thread sums one tile in float32 and adds the tile sums with compensated +// float32 arithmetic. The partner axis is split into contiguous chunks, so +// that a launch fills the device several times over; every chunk writes its +// sums, and a second pass adds the chunks in a fixed order in float64. +// +// The charge constraint and the per-frame virial reduce over the nodes of +// every frame, in slices whose float64 sums are added in a fixed order, so +// that every output is independent of scheduling. +// +// The nodes whose sums a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one +// MPI rank. The ranks then gather the positions and charges of the charged +// nodes of every domain into one frame, in rank order (domain_exchange.h), +// and each rank evaluates its own nodes against the whole frame, so that the +// quadratic work divides among the ranks. Every rank shifts the constrained +// channel by the same sum of the gathered charges; the mean potential of the +// constraint adds the sums of the ranks in rank order, and the per-frame +// virial is the share of the owners of the rank. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "domain_exchange.h" + +namespace { + +// Nodes per block, which is also the length of a tile of partners. +constexpr int kThreads = 128; +// Partners evaluated together, which share every coefficient load. +constexpr int kBatch = 4; +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; +// Fills of the resident block slots of the device that one launch aims for. +constexpr int kWaves = 8; +// Shortest chunk of the partner axis, in tiles. +constexpr int kMinChunkTiles = 2; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Sums of one node: the potential of every channel, the force direction and +// the six independent components of the virial. +template +constexpr int sum_width() { + return kChannels + 3 + 6; +} + +// Running sum with Kahan compensation. +struct CompensatedSum { + float sum = 0.f; + float error = 0.f; + + __device__ __forceinline__ void add(float value) { + const float corrected = value - error; + const float total = sum + corrected; + error = (total - sum) - corrected; + sum = total; + } +}; + +// Index of the frame holding ``node``, or -1 for a padding node past every +// frame. ``frame_ptr`` holds ``frame_count + 1`` offsets. +__device__ __forceinline__ int frame_of(const int64_t* frame_ptr, + int frame_count, + int64_t node) { + if (node >= frame_ptr[frame_count]) { + return -1; + } + int low = 0; + int high = frame_count; + while (high - low > 1) { + const int middle = (low + high) / 2; + if (frame_ptr[middle] <= node) { + low = middle; + } else { + high = middle; + } + } + return low; +} + +// Writes the offsets of the frames on the node axis: frame_ptr[f] counts the +// nodes of the frames before f, for f in [0, frame_count]. One block scans +// the counts, each thread a contiguous run of frames. +constexpr int kOffsetThreads = 256; + +__global__ void __launch_bounds__(kOffsetThreads) + frame_offsets(const int64_t* __restrict__ n_node_per_frame, + int frame_count, + int64_t* __restrict__ frame_ptr) { + __shared__ int64_t run_total[kOffsetThreads]; + const int run = (frame_count + kOffsetThreads - 1) / kOffsetThreads; + const int begin = min(frame_count, static_cast(threadIdx.x) * run); + const int end = min(frame_count, begin + run); + int64_t total = 0; + for (int f = begin; f < end; ++f) { + total += n_node_per_frame[f]; + } + run_total[threadIdx.x] = total; + __syncthreads(); + // Inclusive scan of the run totals. + for (int stride = 1; stride < kOffsetThreads; stride *= 2) { + const int64_t addend = + threadIdx.x >= stride ? run_total[threadIdx.x - stride] : 0; + __syncthreads(); + run_total[threadIdx.x] += addend; + __syncthreads(); + } + int64_t offset = threadIdx.x == 0 ? 0 : run_total[threadIdx.x - 1]; + for (int f = begin; f < end; ++f) { + frame_ptr[f] = offset; + offset += n_node_per_frame[f]; + } + if (threadIdx.x == kOffsetThreads - 1) { + frame_ptr[frame_count] = run_total[kOffsetThreads - 1]; + } +} + +// Per-frame sums. The nodes of a frame are split into ``gridDim.y`` +// contiguous slices, one block of kFrameThreads each, and a finalizing kernel +// adds the float64 sums of the slices in slice order. +constexpr int kFrameThreads = 256; +// Nodes per slice of a frame that a slice pass aims for. +constexpr int kSliceNodes = 2048; +constexpr int kMaxSlices = 256; + +struct NodeRange { + int64_t begin; + int64_t end; +}; + +// The nodes of the slice of the frame that this block sums. +__device__ __forceinline__ NodeRange +slice_of_frame(const int64_t* __restrict__ frame_ptr) { + const int64_t frame_begin = frame_ptr[blockIdx.x]; + const int64_t frame_end = frame_ptr[blockIdx.x + 1]; + const int64_t length = (frame_end - frame_begin + gridDim.y - 1) / gridDim.y; + const int64_t begin = frame_begin + blockIdx.y * length; + return {begin, min(frame_end, begin + length)}; +} + +// Adds ``value`` over the threads of the block in a fixed order and writes the +// totals to the entry of this slice in ``slice_sum``. +template +__device__ __forceinline__ void write_slice_sum( + const double (&value)[kComponents], double* __restrict__ slice_sum) { + __shared__ double partial[kComponents][kFrameThreads]; +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] = value[k]; + } + __syncthreads(); + for (int stride = kFrameThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int k = 0; k < kComponents; ++k) { + partial[k][threadIdx.x] += partial[k][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kComponents) { + slice_sum[(static_cast(blockIdx.x) * gridDim.y + blockIdx.y) * + kComponents + + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Per-frame statistics of the charge constraint: the shift of the +// constrained channel, the number of charged nodes, and the potential of the +// charged nodes in that channel summed over them. +constexpr int kStatShift = 0; +constexpr int kStatCount = 1; +constexpr int kStatPotential = 2; +constexpr int kStatWidth = 3; + +// Sums, over one slice of one frame, the number of charged nodes and their raw +// charges in the constrained channel. +__global__ void __launch_bounds__(kFrameThreads) + charge_slices(const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + const NodeRange slice = slice_of_frame(frame_ptr); + double value[2] = {0.0, 0.0}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { + if (node_mask[node]) { + value[0] += 1.0; + value[1] += charge[node * channels + channel]; + } + } + write_slice_sum<2>(value, slice_sum); +} + +// The shift s_f = (sum_f c - Q_f) / max(n_f, 1) of the constrained channel of +// every frame, from the slice sums of charge_slices. +__global__ void charge_shift(const double* __restrict__ slice_sum, + int slice_count, + int frame_count, + const double* __restrict__ total_charge, + double* __restrict__ frame_stat) { + const int f = blockIdx.x * blockDim.x + threadIdx.x; + if (f >= frame_count) { + return; + } + double count = 0.0; + double sum = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + const double* entry = + slice_sum + (static_cast(f) * slice_count + slice) * 2; + count += entry[0]; + sum += entry[1]; + } + // A frame without a charged node has nothing to shift. + count = fmax(count, 1.0); + frame_stat[f * kStatWidth + kStatShift] = (sum - total_charge[f]) / count; + frame_stat[f * kStatWidth + kStatCount] = count; +} + +// Writes the double-float positions relative to the first node of each +// frame, scaled by alpha for the erf kernel, the frame of every node, and the +// charges of the pairs: the raw charges of the charged nodes, shifted in the +// constrained channel, and zero for every other node. ``alpha`` is a device +// scalar, or null for a kernel read in plain units; ``frame_stat`` is null +// without a constraint. +__global__ void prepare_nodes(const double* __restrict__ coord, + const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int frame_count, + int node_count, + int channels, + int channel, + const float* __restrict__ alpha, + const double* __restrict__ frame_stat, + float4* __restrict__ head, + float4* __restrict__ tail, + int* __restrict__ frame, + float* __restrict__ pair_charge) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= node_count) { + return; + } + const double scale = alpha == nullptr ? 1.0 : static_cast(*alpha); + const int f = frame_of(frame_ptr, frame_count, node); + double3 offset = make_double3(0.0, 0.0, 0.0); + if (f >= 0) { + const int64_t first = frame_ptr[f]; + offset.x = scale * (coord[3 * node + 0] - coord[3 * first + 0]); + offset.y = scale * (coord[3 * node + 1] - coord[3 * first + 1]); + offset.z = scale * (coord[3 * node + 2] - coord[3 * first + 2]); + } + const float hx = __double2float_rn(offset.x); + const float hy = __double2float_rn(offset.y); + const float hz = __double2float_rn(offset.z); + head[node] = make_float4(hx, hy, hz, 0.f); + tail[node] = make_float4(__double2float_rn(offset.x - hx), + __double2float_rn(offset.y - hy), + __double2float_rn(offset.z - hz), 0.f); + frame[node] = f; + const bool charged = f >= 0 && node_mask[node]; + const double shift = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double value = charged ? static_cast(charge[index]) : 0.0; + pair_charge[index] = + static_cast(c == channel ? value - shift : value); + } +} + +// K and K'(r) / r of a batch of pairs at the squared distances ``r2``. +// +// The erf kernel reads positions scaled by alpha, x = alpha r, and evaluates +// erf(x) / x and its slope (2 / sqrt(pi) exp(-x^2) - erf(x) / x) / x^2, with +// erf(x) = 1 - t P(t) exp(-x^2), t = 1 / (1 + p x), and the coefficients of +// Abramowitz and Stegun 7.1.26 (absolute error below 1.5e-7), so that one +// exponential serves the kernel and its slope. The two terms of the slope +// agree to second order in x, which amplifies the error of the approximation +// by 1 / x^2; below x^2 = kSeriesSquared the kernel and its slope therefore +// come from their Taylor series in x^2, +// +// erf(x) / x = 2 / sqrt(pi) sum_n (-1)^n x^(2n) / (n! (2n + 1)), +// slope = 2 / sqrt(pi) sum_{n >= 1} (-1)^n 2n x^(2n - 2) / +// (n! (2n + 1)), +// +// whose first six terms reach float32 precision there. The relative error of +// the kernel stays below 1e-6 for every x, that of its slope below 7e-6, the +// largest just above the threshold. A warp evaluates the series only for a +// batch with a pair below the threshold: a batch that pairs a node with +// itself, or a contact shorter than 0.4 / alpha, which for an alpha of 0.5 per +// angstrom or more is shorter than every bond but that of H2. Every other +// batch evaluates the closed form without a branch. +// +// Beyond x^2 = kFarSquared the float32 kernel equals 1 / x and its slope +// -1 / x^3 bit for bit: exp(-x^2) and erfc(x) fall below half a unit in the +// last place of the terms they would correct. A warp whose batch lies there +// entirely skips the evaluation of erf, which most pairs of a large system +// do, and forms the displacements of the batch from the heads alone. The +// omitted tails move a displacement by at most 2^-24 (|x_i| + |x_j|), with x +// the alpha-scaled offset of a node from the first node of its frame: the +// rounding of float32 positions over the extent of the frame, against a +// distance of at least sqrt(kFarSquared). For a cluster of 100,000 carbon +// atoms at alpha = 0.5 per angstrom, with offsets of up to 26 units, the +// potential and the force keep the error of the exact displacements. +// +// The sum of Gaussians evaluates sum_k a_k 2^(-c_k r^2) and its slope +// sum_k s_k 2^(-c_k r^2), s_k = -2 ln(2) c_k a_k, from the staged +// coefficients ``{a_k, -c_k, s_k}``. +constexpr float kFarSquared = 20.25f; +constexpr float kSeriesSquared = 0.16f; + +// A series in x^2 with the coefficients ``c``, in increasing powers, by +// Horner's rule. +template +__device__ __forceinline__ float even_series(const float (&c)[kTerms], + float x2) { + float sum = c[kTerms - 1]; +#pragma unroll + for (int n = kTerms - 2; n >= 0; --n) { + sum = fmaf(sum, x2, c[n]); + } + return sum; +} + +// Whether every pair of the batch of every thread of the warp lies in the far +// field of the erf kernel, judged from the head displacements, whose squared +// lengths it writes to ``r2``. The sum of Gaussians has no far field. +template +__device__ __forceinline__ bool far_batch(const float (&dx)[kBatch], + const float (&dy)[kBatch], + const float (&dz)[kBatch], + float (&r2)[kBatch]) { + if constexpr (kKernel == kSog) { + return false; + } else { + float nearest = std::numeric_limits::infinity(); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); + nearest = fminf(nearest, r2[b]); + } + return __all_sync(0xffffffffu, nearest >= kFarSquared); + } +} + +// erf(x) / x and its slope by the approximation of Abramowitz and Stegun at +// the squared distances ``r2``. +__device__ __forceinline__ void les_closed_form(const float (&r2)[kBatch], + float (&value)[kBatch], + float (&slope)[kBatch]) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + const float x = r2[b] * inverse; + const float t = __fdividef(1.f, fmaf(0.3275911f, x, 1.f)); + const float poly = + t * + fmaf( + t, + fmaf(t, fmaf(t, fmaf(t, 1.061405429f, -1.453152027f), 1.421413741f), + -0.284496736f), + 0.254829592f); + const float gaussian = exp2f(-1.44269504088896341f * r2[b]); + value[b] = fmaf(-poly, gaussian, 1.f) * inverse; + slope[b] = + fmaf(1.12837916709551257f, gaussian, -value[b]) * (inverse * inverse); + } +} + +template +__device__ __forceinline__ void evaluate_kernel(const float (&r2)[kBatch], + const float4* gaussians, + int gaussian_count, + float (&value)[kBatch], + float (&slope)[kBatch]) { + if constexpr (kKernel == kLes) { + float nearest = r2[0]; +#pragma unroll + for (int b = 1; b < kBatch; ++b) { + nearest = fminf(nearest, r2[b]); + } + if (!__any_sync(0xffffffffu, nearest < kSeriesSquared)) { + les_closed_form(r2, value, slope); + return; + } + constexpr float kValueSeries[6] = {1.128379167f, -0.376126389f, + 0.1128379167f, -0.02686617065f, + 0.005223977625f, -0.0008548327023f}; + constexpr float kSlopeSeries[6] = {-0.7522527781f, 0.4513516668f, + -0.1611970239f, 0.041791821f, + -0.008548327023f, 0.001446639958f}; + // The closed form reads the threshold where the series is selected. + float clamped[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + clamped[b] = fmaxf(r2[b], kSeriesSquared); + } + les_closed_form(clamped, value, slope); +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + if (r2[b] < kSeriesSquared) { + value[b] = even_series(kValueSeries, r2[b]); + slope[b] = even_series(kSlopeSeries, r2[b]); + } + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + value[b] = 0.f; + slope[b] = 0.f; + } + for (int k = 0; k < gaussian_count; ++k) { + const float4 coefficient = gaussians[k]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float weight = exp2f(coefficient.y * r2[b]); + value[b] = fmaf(coefficient.x, weight, value[b]); + slope[b] = fmaf(coefficient.z, weight, slope[b]); + } + } + } +} + +// Partners a tile holds besides those of the frame of the node. +enum Exclusion : int { + // Every partner belongs to the frame and differs from the node. + kNoExclusion = 0, + // The node itself is among the partners. + kExcludeSelf = 1, + // The partners may also lie in other frames. + kExcludeOutsideFrame = 2, +}; + +// The state of the node a thread owns. +template +struct OwnNode { + int index; + int low; // first node of its frame + int high; // past the last node of its frame + float4 head; + float4 tail; + float charge[kChannels]; +}; + +// Adds the unscaled potential, force direction and virial that the partners +// of one staged tile give the node to ``sum``. +template +__device__ __forceinline__ void accumulate_tile( + const float4* __restrict__ tile_head, + const float4* __restrict__ tile_tail, + const float* __restrict__ tile_charge, + const float4* __restrict__ gaussians, + int gaussian_count, + int start, + const OwnNode& own, + float (&sum)[sum_width()]) { +#pragma unroll 4 + for (int base = 0; base < kThreads; base += kBatch) { + float dx[kBatch]; + float dy[kBatch]; + float dz[kBatch]; + float r2[kBatch]; +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_head = tile_head[base + b]; + dx[b] = own.head.x - other_head.x; + dy[b] = own.head.y - other_head.y; + dz[b] = own.head.z - other_head.z; + } + float value[kBatch]; + float slope[kBatch]; + if (far_batch(dx, dy, dz, r2)) { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float inverse = rsqrtf(r2[b]); + value[b] = inverse; + slope[b] = -inverse * (inverse * inverse); + } + } else { +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const float4 other_tail = tile_tail[base + b]; + dx[b] += own.tail.x - other_tail.x; + dy[b] += own.tail.y - other_tail.y; + dz[b] += own.tail.z - other_tail.z; + // Both kernels are finite at r2 = 0, which a node paired with itself + // and coincident padding nodes reach with a zero weight. + r2[b] = fmaf(dx[b], dx[b], fmaf(dy[b], dy[b], dz[b] * dz[b])); + } + evaluate_kernel(r2, gaussians, gaussian_count, value, slope); + } +#pragma unroll + for (int b = 0; b < kBatch; ++b) { + const int other = start + base + b; + float weight = 1.f; + if constexpr (kExclusion == kExcludeSelf) { + weight = other != own.index ? 1.f : 0.f; + } else if constexpr (kExclusion == kExcludeOutsideFrame) { + weight = (other != own.index && + static_cast(other - own.low) < + static_cast(own.high - own.low)) + ? 1.f + : 0.f; + } + float product = 0.f; +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + const float other_charge = + kExclusion == kNoExclusion + ? tile_charge[(base + b) * kChannels + c] + : weight * tile_charge[(base + b) * kChannels + c]; + sum[c] = fmaf(other_charge, value[b], sum[c]); + product = fmaf(own.charge[c], other_charge, product); + } + const float strength = product * slope[b]; + const float sx = strength * dx[b]; + const float sy = strength * dy[b]; + const float sz = strength * dz[b]; + sum[kChannels + 0] += sx; + sum[kChannels + 1] += sy; + sum[kChannels + 2] += sz; + sum[kChannels + 3] = fmaf(sx, dx[b], sum[kChannels + 3]); + sum[kChannels + 4] = fmaf(sy, dy[b], sum[kChannels + 4]); + sum[kChannels + 5] = fmaf(sz, dz[b], sum[kChannels + 5]); + sum[kChannels + 6] = fmaf(sx, dy[b], sum[kChannels + 6]); + sum[kChannels + 7] = fmaf(sx, dz[b], sum[kChannels + 7]); + sum[kChannels + 8] = fmaf(sy, dz[b], sum[kChannels + 8]); + } + } +} + +// Sums, for every owner node of one block and every partner of one chunk, the +// unscaled potential, force direction and virial of the node, and writes them +// to ``partial[chunk][component][node - owner_begin]``. The owners are the +// nodes [owner_begin, owner_begin + owner_count). +template +__global__ void __launch_bounds__(kThreads) + direct_sum_kernel(const float4* __restrict__ head, + const float4* __restrict__ tail, + const float* __restrict__ charge, + const int* __restrict__ frame, + const int64_t* __restrict__ frame_ptr, + int owner_begin, + int owner_count, + const float* __restrict__ coefficient, + int gaussian_count, + float* __restrict__ partial) { + constexpr int kWidth = sum_width(); + __shared__ float4 tile_head[kThreads]; + __shared__ float4 tile_tail[kThreads]; + __shared__ float tile_charge[kThreads * kChannels]; + __shared__ float4 gaussians[kMaxGaussians]; + __shared__ int range[2]; + __shared__ bool single_frame; + + const int owner_end = owner_begin + owner_count; + const int first = owner_begin + blockIdx.x * kThreads; + const int node = first + threadIdx.x; + const int own_frame = node < owner_end ? frame[node] : -1; + + if constexpr (kKernel == kSog) { + for (int k = threadIdx.x; k < gaussian_count; k += kThreads) { + const float amplitude = coefficient[k]; + const float rate = coefficient[gaussian_count + k]; + gaussians[k] = make_float4(amplitude, -rate, + -1.38629436111989061f * rate * amplitude, 0.f); + } + } + // The partners of the block are the nodes of the frames its nodes belong + // to: a contiguous range, since frames are contiguous on the node axis and + // padding nodes trail them. + if (threadIdx.x == 0) { + int last = min(first + kThreads, owner_end) - 1; + while (last > first && frame[last] < 0) { + --last; + } + const int low_frame = frame[first]; + const int high_frame = frame[last]; + range[0] = low_frame < 0 ? 0 : static_cast(frame_ptr[low_frame]); + range[1] = low_frame < 0 ? 0 : static_cast(frame_ptr[high_frame + 1]); + // A block whose nodes all lie in one frame finds every partner of its + // range in that frame, which spares the frame test of every pair. + single_frame = low_frame >= 0 && low_frame == high_frame && + min(first + kThreads, owner_end) - 1 == last; + } + __syncthreads(); + + // The chunk of the partner range this block sums, in whole tiles. + const int64_t tile_count = (range[1] - range[0] + kThreads - 1) / kThreads; + const int begin = + range[0] + + static_cast(tile_count * blockIdx.y / gridDim.y) * kThreads; + const int end = min( + range[1], + range[0] + static_cast(tile_count * (blockIdx.y + 1) / gridDim.y) * + kThreads); + + OwnNode own; + own.index = node; + own.low = 0; + own.high = 0; + own.head = make_float4(0.f, 0.f, 0.f, 0.f); + own.tail = own.head; + if (own_frame >= 0) { + own.head = head[node]; + own.tail = tail[node]; + own.low = static_cast(frame_ptr[own_frame]); + own.high = static_cast(frame_ptr[own_frame + 1]); + } +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + own.charge[c] = own_frame >= 0 ? charge[node * kChannels + c] : 0.f; + } + + CompensatedSum total[kWidth]; + for (int start = begin; start < end; start += kThreads) { + const int partner = start + threadIdx.x; + const bool staged = partner < end; + // A slot past the chunk carries a zero charge and contributes nothing. + tile_head[threadIdx.x] = + staged ? head[partner] : make_float4(0.f, 0.f, 0.f, 0.f); + tile_tail[threadIdx.x] = + staged ? tail[partner] : make_float4(0.f, 0.f, 0.f, 0.f); +#pragma unroll + for (int c = 0; c < kChannels; ++c) { + tile_charge[threadIdx.x * kChannels + c] = + staged ? charge[partner * kChannels + c] : 0.f; + } + __syncthreads(); + + float sum[kWidth]; +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + sum[k] = 0.f; + } + // In a single-frame block only the tiles that overlap the block's own + // nodes hold a node itself, so the tests are confined to them and to the + // blocks that span frames; the choice is uniform over the block. + if (!single_frame) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else if (start < first + kThreads && first < start + kThreads) { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } else { + accumulate_tile( + tile_head, tile_tail, tile_charge, gaussians, gaussian_count, start, + own, sum); + } +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + total[k].add(sum[k]); + } + __syncthreads(); + } + + if (node < owner_end) { + float* out = partial + + static_cast(blockIdx.y) * kWidth * owner_count + + (node - owner_begin); +#pragma unroll + for (int k = 0; k < kWidth; ++k) { + out[static_cast(k) * owner_count] = total[k].sum; + } + } +} + +// Adds the chunk sums of one component of one owner node in chunk order, one +// thread per component and owner, and scales the result into the potential, +// the force or the per-node virial of the owner. With positions scaled by +// alpha, K(r) = alpha f(x), K'(r) / r = alpha^3 f'(x) / x and d = d' / alpha, +// so the potential carries alpha, the force alpha^2 and the virial alpha. +template +__global__ void reduce_chunks(const float* __restrict__ partial, + int owner_count, + int chunk_count, + double prefactor, + const float* __restrict__ alpha, + float* __restrict__ potential, + float* __restrict__ force, + float* __restrict__ atom_virial) { + constexpr int kWidth = sum_width(); + const int64_t index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= static_cast(kWidth) * owner_count) { + return; + } + const int component = static_cast(index / owner_count); + const int node = static_cast(index - component * int64_t(owner_count)); + double value = 0.0; + for (int chunk = 0; chunk < chunk_count; ++chunk) { + value += + partial[static_cast(chunk) * kWidth * owner_count + index]; + } + const double unit = alpha == nullptr ? 1.0 : static_cast(*alpha); + if (component < kChannels) { + potential[node * kChannels + component] = + static_cast(prefactor * unit * value); + } else if (component < kChannels + 3) { + force[node * 3 + component - kChannels] = + static_cast(-prefactor * unit * unit * value); + } else { + // Each pair term is split half to each of its two nodes. The six sums + // are xx, yy, zz, xy, xz, yz. + constexpr int kRow[6] = {0, 1, 2, 0, 0, 1}; + constexpr int kColumn[6] = {0, 1, 2, 1, 2, 2}; + const int pair = component - kChannels - 3; + const float entry = static_cast(-0.5 * prefactor * unit * value); + float* out = atom_virial + static_cast(node) * 9; + out[3 * kRow[pair] + kColumn[pair]] = entry; + out[3 * kColumn[pair] + kRow[pair]] = entry; + } +} + +// Sums, over one slice of one frame, the per-node virial and, with a +// constraint, the potential of the charged nodes in the constrained channel. +template +__global__ void __launch_bounds__(kFrameThreads) + virial_slices(const float* __restrict__ atom_virial, + const float* __restrict__ potential, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + int channels, + int channel, + double* __restrict__ slice_sum) { + constexpr int kComponents = kConstrained ? 10 : 9; + const NodeRange slice = slice_of_frame(frame_ptr); + double value[kComponents] = {}; + for (int64_t node = slice.begin + threadIdx.x; node < slice.end; + node += kFrameThreads) { +#pragma unroll + for (int k = 0; k < 9; ++k) { + value[k] += atom_virial[node * 9 + k]; + } + if constexpr (kConstrained) { + if (node_mask[node]) { + value[9] += potential[node * channels + channel]; + } + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of virial_slices into the virial of every frame and, +// with a constraint, the potential summed over its charged nodes. +template +__global__ void virial_total(const double* __restrict__ slice_sum, + int slice_count, + float* __restrict__ virial, + double* __restrict__ frame_stat) { + constexpr int kComponents = kConstrained ? 10 : 9; + if (threadIdx.x >= kComponents) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kComponents + + threadIdx.x]; + } + if (threadIdx.x < 9) { + virial[blockIdx.x * 9 + threadIdx.x] = static_cast(value); + } else { + frame_stat[blockIdx.x * kStatWidth + kStatPotential] = value; + } +} + +// Writes the atomic energy q_i . phi_i / 2 of every owner node and the +// derivative of the energy with respect to its raw charges: w_i (phi_i - +// mean_f) in the constrained channel and w_i phi_i in the others, w_i marking +// the charged nodes. ``pair_charge``, ``node_mask`` and ``frame`` start at the +// first owner. +__global__ void project_charges(const float* __restrict__ potential, + const float* __restrict__ pair_charge, + const bool* __restrict__ node_mask, + const int* __restrict__ frame, + const double* __restrict__ frame_stat, + int owner_count, + int channels, + int channel, + double* __restrict__ atom_energy, + float* __restrict__ charge_gradient) { + const int node = blockIdx.x * blockDim.x + threadIdx.x; + if (node >= owner_count) { + return; + } + const int f = frame[node]; + const bool charged = f >= 0 && node_mask[node]; + const double mean = charged && frame_stat != nullptr + ? frame_stat[f * kStatWidth + kStatPotential] / + frame_stat[f * kStatWidth + kStatCount] + : 0.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t index = static_cast(node) * channels + c; + const double phi = potential[index]; + energy += static_cast(pair_charge[index]) * phi; + const double gradient = charged ? phi : 0.0; + charge_gradient[index] = + static_cast(c == channel ? gradient - mean : gradient); + } + atom_energy[node] = 0.5 * energy; +} + +// Chunks of the partner axis for the given counts of owners and partners. The +// blocks of one launch all carry the same work, so the launch fills the +// resident block slots of the device several times over, which bounds the +// idle share of the last wave; no chunk is shorter than kMinChunkTiles tiles +// of a frame that spans every partner. +int chunk_count_for(int owner_count, + int partner_count, + int resident_blocks_per_multiprocessor) { + const int multiprocessors = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int owner_blocks = (owner_count + kThreads - 1) / kThreads; + const int partner_tiles = (partner_count + kThreads - 1) / kThreads; + const int64_t wanted_blocks = static_cast(kWaves) * + resident_blocks_per_multiprocessor * + multiprocessors; + const int wanted = + static_cast((wanted_blocks + owner_blocks - 1) / owner_blocks); + return std::max(1, std::min(wanted, partner_tiles / kMinChunkTiles)); +} + +// The staged partners, one row per node, and the sums of the owners. +struct DirectSumBuffers { + const torch::Tensor& head; + const torch::Tensor& tail; + const torch::Tensor& charge; + const torch::Tensor& frame; + const torch::Tensor& frame_ptr; + const torch::Tensor& coefficient; + torch::Tensor& potential; + torch::Tensor& force; + torch::Tensor& atom_virial; +}; + +template +int resident_blocks_per_multiprocessor() { + // The occupancy is a property of the compiled kernel and of the device; a + // process drives one device architecture. + static const int blocks = [] { + int value = 0; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &value, direct_sum_kernel, kThreads, 0)); + return std::max(1, value); + }(); + return blocks; +} + +template +void launch(const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + double prefactor, + cudaStream_t stream) { + constexpr int kWidth = sum_width(); + const int chunk_count = + chunk_count_for(owner_count, static_cast(buffers.head.size(0)), + resident_blocks_per_multiprocessor()); + auto partial = + torch::empty({static_cast(chunk_count) * owner_count * kWidth}, + buffers.charge.options()); + const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); + direct_sum_kernel<<>>( + reinterpret_cast(buffers.head.data_ptr()), + reinterpret_cast(buffers.tail.data_ptr()), + buffers.charge.data_ptr(), buffers.frame.data_ptr(), + buffers.frame_ptr.data_ptr(), owner_begin, owner_count, + buffers.coefficient.data_ptr(), + static_cast(buffers.coefficient.numel() / 2), + partial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t reduce_threads = static_cast(kWidth) * owner_count; + reduce_chunks + <<((reduce_threads + 255) / 256), 256, 0, stream>>>( + partial.data_ptr(), owner_count, chunk_count, prefactor, + kKernel == kLes ? buffers.coefficient.data_ptr() : nullptr, + buffers.potential.data_ptr(), buffers.force.data_ptr(), + buffers.atom_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void dispatch_channels(int channels, + const DirectSumBuffers& buffers, + int owner_begin, + int owner_count, + double prefactor, + cudaStream_t stream) { + switch (channels) { + case 1: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 2: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 3: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 4: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 5: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 6: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 7: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + case 8: + return launch(buffers, owner_begin, owner_count, prefactor, + stream); + default: + TORCH_CHECK(false, "latent_charge_direct_sum: unsupported channel count ", + channels); + } +} + +using DirectSumOutputs = std::tuple; + +// Evaluates the owner nodes [owner_begin, owner_begin + owner_count) against +// the nodes of their frames. Returns the outputs of the owners, row r for +// owner owner_begin + r, and the per-frame virial of the owners. Owners other +// than every node lie in a single frame. A non-null ``communicator`` spans +// the domains whose owners share that frame, over which the potential of the +// constraint is summed. +DirectSumOutputs evaluate(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { + TORCH_CHECK(coord.size(0) < std::numeric_limits::max() / kThreads, + "latent_charge_direct_sum: too many nodes"); + const int node_count = static_cast(coord.size(0)); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const int64_t channels = charge.size(1); + const auto float_options = charge.options(); + const auto double_options = coord.options(); + if (node_count == 0 || frame_count == 0) { + return {torch::zeros({owner_count, 1}, double_options), + torch::zeros({owner_count, channels}, float_options), + torch::zeros({owner_count, 3}, float_options), + torch::zeros({owner_count, 9}, float_options), + torch::zeros({frame_count, 9}, float_options)}; + } + const bool every_node = owner_begin == 0 && owner_count == node_count; + TORCH_CHECK(every_node || frame_count == 1, + "latent_charge_direct_sum: owners other than every node lie in " + "one frame"); + const bool constrained = channel >= 0; + const int constrained_channel = static_cast(channel); + // The kernels write every entry of the outputs. + auto atom_energy = torch::empty({owner_count, 1}, double_options); + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto force = torch::empty({owner_count, 3}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); + auto virial = torch::empty({frame_count, 9}, float_options); + + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + auto frame_ptr = torch::empty({frame_count + 1}, counts.options()); + auto head = torch::empty({node_count, 4}, float_options); + auto tail = torch::empty({node_count, 4}, float_options); + auto frame = torch::empty({node_count}, coord.options().dtype(torch::kInt32)); + auto pair_charge = torch::empty({node_count, channels}, float_options); + auto potential = torch::empty({owner_count, channels}, float_options); + auto frame_stat = torch::empty({constrained ? frame_count * kStatWidth : 0}, + double_options); + const auto slices_for = [frame_count](int64_t nodes) { + const int64_t average_nodes = (nodes + frame_count - 1) / frame_count; + return static_cast(std::min( + kMaxSlices, + std::max(1, (average_nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(node_count); + const dim3 slice_grid(frame_count, slice_count); + const auto stream = at::cuda::getCurrentCUDAStream(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>( + counts.data_ptr(), frame_count, frame_ptr.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained) { + auto slice_charge = torch::empty( + {static_cast(frame_count) * slice_count * 2}, double_options); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), + frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_charge.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge.data_ptr(), slice_count, frame_count, + total_charge.data_ptr(), frame_stat.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // Alpha stays on the device, so that a call never waits for it. + prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr.data_ptr(), frame_count, + node_count, static_cast(channels), constrained_channel, + kernel == kLes ? coefficient.data_ptr() : nullptr, + constrained ? frame_stat.data_ptr() : nullptr, + reinterpret_cast(head.data_ptr()), + reinterpret_cast(tail.data_ptr()), frame.data_ptr(), + pair_charge.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (owner_count > 0) { + const DirectSumBuffers buffers{head, tail, pair_charge, + frame, frame_ptr, coefficient, + potential, force, atom_virial}; + if (kernel == kLes) { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } else { + dispatch_channels(static_cast(channels), buffers, owner_begin, + owner_count, prefactor, stream.stream()); + } + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame the owners + // share, written on the device. + torch::Tensor owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_ptr = + torch::full({2}, int64_t{owner_count}, frame_ptr.options()); + owner_frame_ptr.select(0, 0).zero_(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const int owner_slice_count = slices_for(owner_count); + const dim3 owner_slice_grid(frame_count, owner_slice_count); + const int components = constrained ? 10 : 9; + auto slice_virial = torch::empty( + {static_cast(frame_count) * owner_slice_count * components}, + double_options); + if (constrained) { + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial.data_ptr(), owner_slice_count, + virial.data_ptr(), frame_stat.data_ptr()); + } else { + virial_slices<<>>( + atom_virial.data_ptr(), potential.data_ptr(), owner_mask, + owner_frame_ptr.data_ptr(), static_cast(channels), + constrained_channel, slice_virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial.data_ptr(), owner_slice_count, + virial.data_ptr(), nullptr); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained && communicator != nullptr) { + // The mean potential of the constraint runs over the charged nodes of + // every domain. + auto potential_sum = + frame_stat.view({frame_count, kStatWidth}).select(1, kStatPotential); + potential_sum.copy_(deepmd::long_range::sum_over_domains( + *communicator, potential_sum.contiguous())); + } + if (owner_count > 0) { + project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + potential.data_ptr(), + pair_charge.data_ptr() + int64_t{owner_begin} * channels, + owner_mask, frame.data_ptr() + owner_begin, + constrained ? frame_stat.data_ptr() : nullptr, owner_count, + static_cast(channels), constrained_channel, + atom_energy.data_ptr(), charge_gradient.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {atom_energy, charge_gradient, force, atom_virial, virial}; +} + +// Evaluates the frame of a system decomposed into the domains of +// ``communicator``, of which this rank holds one: its charged nodes are its +// owners, the charged nodes of every domain its partners. The nodes the rank +// does not own carry no charge, and the per-frame virial is the share of its +// owners. +DirectSumOutputs evaluate_across_domains(const torch::Tensor& coord, + const torch::Tensor& charge, + const torch::Tensor& node_mask, + const torch::Tensor& n_node_per_frame, + const torch::Tensor& total_charge, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + const torch::Tensor& communicator) { + TORCH_CHECK(n_node_per_frame.size(0) == 1, + "latent_charge_direct_sum: an evaluation across MPI ranks takes " + "one frame, got ", + n_node_per_frame.size(0)); + const int64_t node_count = coord.size(0); + // The charged nodes of the frame; the padding nodes past it carry none. + // Their count is the one value the host has to wait for before the + // exchange. + const auto in_frame = + torch::arange(node_count, n_node_per_frame.options()) < n_node_per_frame; + const auto owned = torch::nonzero(node_mask & in_frame).squeeze(1); + const auto rows = + torch::cat({coord.index_select(0, owned), + charge.index_select(0, owned).to(torch::kFloat64)}, + 1) + .contiguous(); + const auto [partners, owner_begin] = + deepmd::long_range::gather_domain_rows(communicator, rows); + const int64_t partner_count = partners.size(0); + const auto partner_coord = partners.slice(1, 0, 3).contiguous(); + const auto partner_charge = + partners.slice(1, 3).to(torch::kFloat32).contiguous(); + const auto partner_mask = torch::ones({partner_count}, node_mask.options()); + const auto partner_frame = + torch::full({1}, partner_count, n_node_per_frame.options()); + const auto [owner_energy, owner_gradient, owner_force, owner_atom_virial, + virial] = + evaluate(partner_coord, partner_charge, partner_mask, partner_frame, + total_charge, coefficient, kernel, channel, prefactor, + static_cast(owner_begin), static_cast(owned.size(0)), + &communicator); + return {torch::zeros({node_count, 1}, coord.options()) + .index_copy_(0, owned, owner_energy), + torch::zeros({node_count, charge.size(1)}, charge.options()) + .index_copy_(0, owned, owner_gradient), + torch::zeros({node_count, 3}, charge.options()) + .index_copy_(0, owned, owner_force), + torch::zeros({node_count, 9}, charge.options()) + .index_copy_(0, owned, owner_atom_virial), + virial}; +} + +DirectSumOutputs latent_charge_direct_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor) { + TORCH_CHECK(coord.is_cuda() && charge.is_cuda() && node_mask.is_cuda() && + n_node_per_frame.is_cuda() && total_charge.is_cuda() && + coefficient.is_cuda(), + "latent_charge_direct_sum: inputs must be CUDA tensors"); + TORCH_CHECK(charge.device() == coord.device() && + node_mask.device() == coord.device() && + n_node_per_frame.device() == coord.device() && + total_charge.device() == coord.device() && + coefficient.device() == coord.device(), + "latent_charge_direct_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: coord must be float64 with shape " + "(N, 3)"); + TORCH_CHECK(charge.dim() == 2 && charge.size(0) == coord.size(0) && + charge.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: charge must be float32 with shape " + "(N, C)"); + TORCH_CHECK(node_mask.dim() == 1 && node_mask.size(0) == coord.size(0) && + node_mask.scalar_type() == torch::kBool, + "latent_charge_direct_sum: node_mask must be bool with shape " + "(N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_direct_sum: n_node_per_frame must be " + "one-dimensional"); + TORCH_CHECK(total_charge.dim() == 1 && + total_charge.size(0) == n_node_per_frame.size(0) && + total_charge.scalar_type() == torch::kFloat64, + "latent_charge_direct_sum: total_charge must be float64 with " + "one entry per frame"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat32, + "latent_charge_direct_sum: coefficient must be a float32 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_direct_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_direct_sum: the constrained channel must be -1 " + "or lie in [0, ", + channels, "), got ", channel); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_direct_sum: the erf kernel takes one " + "coefficient"); + } else if (kernel == kSog) { + TORCH_CHECK(coefficient.numel() % 2 == 0 && coefficient.numel() > 0 && + coefficient.numel() <= 2 * kMaxGaussians, + "latent_charge_direct_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_direct_sum: unknown kernel ", kernel); + } + + const c10::cuda::CUDAGuard guard(coord.device()); + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + coefficient = coefficient.contiguous(); + if (deepmd::long_range::domain_of(communicator).size == 1) { + return evaluate(coord, charge, node_mask, n_node_per_frame, total_charge, + coefficient, kernel, channel, prefactor, 0, + static_cast(coord.size(0)), nullptr); + } + return evaluate_across_domains(coord, charge, node_mask, n_node_per_frame, + total_charge, coefficient, kernel, channel, + prefactor, communicator); +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.impl("latent_charge_direct_sum", torch::kCUDA, + &latent_charge_direct_sum); +} diff --git a/source/op/pt/long_range/domain_exchange.cc b/source/op/pt/long_range/domain_exchange.cc new file mode 100644 index 0000000000..24c393c5ee --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.cc @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// MPI exchanges of the long-range direct sum (contract in domain_exchange.h). +// +// With an MPI library that reads device memory (MPIX_Query_cuda_support), the +// rows travel between device buffers once the current stream has finished +// writing them; otherwise they are staged through pinned host buffers. Either +// way the host waits for the device once per exchange, before MPI reads the +// data; the results return to the device without a further wait. + +#include "domain_exchange.h" + +#ifdef USE_MPI +#include +#ifdef OMPI_MPI_H +#include +#endif +#endif +#include +#include + +#include +#include + +namespace deepmd { +namespace long_range { +namespace { + +std::int64_t address_of(const torch::Tensor& communicator) { + TORCH_CHECK(communicator.device().is_cpu() && + communicator.scalar_type() == torch::kInt64 && + communicator.numel() == 1, + "latent_charge_direct_sum: communicator must be an int64 CPU " + "tensor with one element"); + return communicator.contiguous().data_ptr()[0]; +} + +#ifdef USE_MPI +// The handle of the caller, of the width of MPI_Comm in this MPI library: an +// int in MPICH, a pointer in Open MPI. +MPI_Comm communicator_of(const torch::Tensor& communicator) { + static_assert(sizeof(std::uintptr_t) <= sizeof(std::int64_t), + "an address must fit in an int64 tensor"); + return *reinterpret_cast( + static_cast(address_of(communicator))); +} + +bool reads_device_memory() { +#ifdef NO_CUDA_AWARE + return false; +#else + return MPIX_Query_cuda_support() != 0; +#endif +} +#endif + +} // namespace + +Domain domain_of(const torch::Tensor& communicator) { + if (address_of(communicator) == 0) { + return {0, 1}; + } +#ifdef USE_MPI + int initialized = 0; + MPI_Initialized(&initialized); + TORCH_CHECK(initialized != 0, + "latent_charge_direct_sum: a communicator was given, but MPI is " + "not initialized"); + const MPI_Comm comm = communicator_of(communicator); + Domain domain{0, 1}; + MPI_Comm_rank(comm, &domain.rank); + MPI_Comm_size(comm, &domain.size); + return domain; +#else + TORCH_CHECK(false, + "latent_charge_direct_sum: an evaluation across MPI ranks " + "requires DeePMD-kit built with MPI"); +#endif +} + +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows) { +#ifdef USE_MPI + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && + rows.scalar_type() == torch::kFloat64 && rows.is_contiguous(), + "gather_domain_rows: rows must be a contiguous float64 CUDA " + "matrix"); + const MPI_Comm comm = communicator_of(communicator); + int rank = 0; + int size = 1; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + const std::int64_t width = rows.size(1); + std::int64_t local = rows.size(0); + std::vector counts(size); + MPI_Allgather(&local, 1, MPI_INT64_T, counts.data(), 1, MPI_INT64_T, comm); + std::vector receive_counts(size); + std::vector displacements(size); + std::int64_t total = 0; + std::int64_t offset = 0; + for (int r = 0; r < size; ++r) { + if (r == rank) { + offset = total; + } + TORCH_CHECK((total + counts[r]) * width <= std::numeric_limits::max(), + "gather_domain_rows: the rows of all ranks exceed the int " + "range of an MPI count"); + receive_counts[r] = static_cast(counts[r] * width); + displacements[r] = static_cast(total * width); + total += counts[r]; + } + const int send_count = static_cast(local * width); + auto gathered = torch::empty({total, width}, rows.options()); + if (reads_device_memory()) { + at::cuda::getCurrentCUDAStream().synchronize(); + MPI_Allgatherv(rows.data_ptr(), send_count, MPI_DOUBLE, + gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + } else { + const auto host = rows.options().device(torch::kCPU).pinned_memory(true); + auto host_rows = torch::empty({local, width}, host); + host_rows.copy_(rows); + auto host_gathered = torch::empty({total, width}, host); + MPI_Allgatherv(host_rows.data_ptr(), send_count, MPI_DOUBLE, + host_gathered.data_ptr(), receive_counts.data(), + displacements.data(), MPI_DOUBLE, comm); + gathered.copy_(host_gathered, /*non_blocking=*/true); + } + return {gathered, offset}; +#else + TORCH_CHECK(false, + "gather_domain_rows: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values) { +#ifdef USE_MPI + TORCH_CHECK(values.dim() == 1 && values.scalar_type() == torch::kFloat64, + "sum_over_domains: values must be a float64 vector"); + const MPI_Comm comm = communicator_of(communicator); + int size = 1; + MPI_Comm_size(comm, &size); + const torch::Tensor local = values.to(torch::kCPU).contiguous(); + const int count = static_cast(local.numel()); + std::vector every(static_cast(size) * count); + MPI_Allgather(local.data_ptr(), count, MPI_DOUBLE, every.data(), + count, MPI_DOUBLE, comm); + auto total = torch::zeros({count}, local.options().pinned_memory(true)); + double* sum = total.data_ptr(); + for (int r = 0; r < size; ++r) { + for (int k = 0; k < count; ++k) { + sum[k] += every[static_cast(r) * count + k]; + } + } + return total.to(values.device(), /*non_blocking=*/true); +#else + TORCH_CHECK(false, + "sum_over_domains: an evaluation across MPI ranks requires " + "DeePMD-kit built with MPI"); +#endif +} + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/domain_exchange.h b/source/op/pt/long_range/domain_exchange.h new file mode 100644 index 0000000000..c091589b31 --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.h @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Exchange of the long-range direct sum between the domains of a decomposed +// system. +// +// A domain is one MPI rank of the communicator that a caller hands to +// deepmd::latent_charge_direct_sum as an int64 tensor holding the address of +// its MPI_Comm handle. A zero address stands for a single domain. Both +// exchanges follow the rank order, so every rank receives the same rows and the +// same sums, independent of the reduction algorithm of the MPI library. + +#pragma once + +#include + +#include +#include + +namespace deepmd { +namespace long_range { + +struct Domain { + int rank; + int size; +}; + +// The rank of this process and the number of ranks of ``communicator``; one +// rank for a zero address. +Domain domain_of(const torch::Tensor& communicator); + +// Concatenates, in rank order, the rows of every rank. ``rows`` is a +// contiguous float64 matrix of shape (M_r, W) on the CUDA device, with the +// same W on every rank. Returns the rows of all ranks, of shape (sum_r M_r, +// W) on the device of ``rows``, and the index of the first row of this rank. +std::pair gather_domain_rows( + const torch::Tensor& communicator, const torch::Tensor& rows); + +// Sums the float64 vector ``values`` over the ranks, adding the vectors of +// the ranks in rank order. Returns the sum on the device of ``values``. +torch::Tensor sum_over_domains(const torch::Tensor& communicator, + const torch::Tensor& values); + +} // namespace long_range +} // namespace deepmd diff --git a/source/op/pt/long_range/ops.cc b/source/op/pt/long_range/ops.cc new file mode 100644 index 0000000000..dd312bc555 --- /dev/null +++ b/source/op/pt/long_range/ops.cc @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Operator schema of the latent-charge long-range direct sum. +// +// The schema is declared here, unconditionally, while the CUDA kernel +// registers itself only against a CUDA-enabled PyTorch (direct_sum.cu), so +// that the Python front end can tell "library not loaded" from "this device +// has no kernel". +// +// For the nodes of one frame, the operator takes the raw charges c_i of the +// charge network in C float32 channels, the nodes that carry them (w_i = 1, +// and 0 elsewhere) and float64 positions r_i. The pairs carry the charges +// q_i = w_i c_i; when `channel` is a channel index rather than -1, that +// channel is shifted uniformly over the charged nodes of every frame so that +// it sums to the total charge Q_f of the frame, +// +// q_i^ch = w_i (c_i^ch - s_f), s_f = (sum_{j in f} w_j c_j^ch - Q_f) / +// max(n_f, 1), +// +// with n_f charged nodes in f. With the potential of the charges +// +// phi_i = prefactor * sum_{j != i} q_j K(r_ij), +// +// the operator returns +// +// atom_energy_i = q_i . phi_i / 2 (float64) +// charge_gradient_i = dE / dc_i = w_i (phi_i - [c = ch] m_f), m_f the +// mean of phi^ch over the charged nodes of f +// force_i = -prefactor * sum_{j != i} (q_i . q_j) K'(r_ij) / r_ij +// * d_ij +// W_i = -prefactor / 2 * sum_{j != i} (q_i . q_j) +// K'(r_ij) / r_ij * d_ij (x) d_ij +// W_f = sum_{i in f} W_i +// +// in float32 unless noted, with d_ij = r_i - r_j and the sums running over +// the nodes of the frame of i. The frames occupy the leading +// sum(n_node_per_frame) <= N nodes, which the caller ensures; the nodes past +// them are padding and join none. The frame energy is the sum of the atomic +// energies. `kernel` selects K: 0 is +// erf(alpha r) / r with `coefficient` holding alpha; 1 is a sum of Gaussians +// a_k 2^(-c_k r^2), c_k = log2(e) / (2 sigma_k^2), with `coefficient` holding +// the amplitudes a_k and then the rates c_k. `total_charge` holds Q_f in +// float64 and is read only with a constrained channel. The force and the +// virial hold the charges fixed; the response of the charges is the +// caller's, through the charge gradient. +// +// `communicator` is an int64 CPU tensor of one element holding the address of +// the caller's MPI_Comm handle, or zero. A handle of several ranks decomposes +// one frame into domains, one per rank: the charged nodes of every rank are +// the charged nodes of the frame, and the sums run over all of them. The +// outputs of a rank are those of its charged nodes, and W_f is their share of +// the virial of the frame; the other nodes of the rank carry no charge. + +#include + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.def( + "latent_charge_direct_sum(Tensor coord, Tensor charge, " + "Tensor node_mask, Tensor n_node_per_frame, Tensor total_charge, " + "Tensor coefficient, Tensor communicator, int kernel, int channel, " + "float prefactor) -> " + "(Tensor atom_energy, Tensor charge_gradient, Tensor force, " + "Tensor atom_virial, Tensor virial)"); +} + +// The operator serves inference, whose caller differentiates nothing through +// it; its outputs carry no gradient. +TORCH_LIBRARY_IMPL(deepmd, Autograd, library) { + library.impl("latent_charge_direct_sum", + torch::CppFunction::makeFallthrough()); +} diff --git a/source/tests/common/dpmodel/test_long_range.py b/source/tests/common/dpmodel/test_long_range.py new file mode 100644 index 0000000000..108fd7082c --- /dev/null +++ b/source/tests/common/dpmodel/test_long_range.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Primitives of the long-range modules: pair topology, kernels and constraint.""" + +import itertools + +import numpy as np +import pytest +from scipy.special import ( + erf, +) + +from deepmd.dpmodel.fitting.latent_charge import ( + LatentChargeFitting, +) +from deepmd.dpmodel.long_range import ( + BaseInteractionKernel, + LatentChargeLongRange, + LESKernel, + SOGKernel, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + build_pair_graph, + build_topologies, + pair_force_virial, + unordered_pair_index, +) + + +def _brute_force_pairs(counts: list[int], mask: np.ndarray | None) -> np.ndarray: + pairs = [] + start = 0 + for count in counts: + nodes = [ + node for node in range(start, start + count) if mask is None or mask[node] + ] + pairs += itertools.combinations(nodes, 2) + start += count + return np.asarray(pairs, dtype=np.int64).reshape(-1, 2).T + + +@pytest.mark.parametrize( + "counts", [[0], [1], [2], [3, 0, 1, 5], [7, 200, 1, 0, 33], [1000]] +) +@pytest.mark.parametrize("masked", [False, True]) +def test_pair_index_matches_brute_force(counts: list[int], masked: bool) -> None: + """Every unordered pair of valid nodes of a frame appears once, in row-major order.""" + mask = np.random.default_rng(0).random(sum(counts)) > 0.3 if masked else None + index = unordered_pair_index(np.asarray(counts), mask) + np.testing.assert_array_equal(index, _brute_force_pairs(counts, mask)) + + +def test_pair_graph_guards_and_periodic_rejection() -> None: + """The pair axis ends in masked guards; a cell is rejected.""" + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + pairs = build_pair_graph(coord, np.array([4]), atype) + assert pairs.index.shape == (2, 8) + assert pairs.mask.tolist() == [True] * 6 + [False] * 2 + np.testing.assert_allclose( + pairs.vec[:6], coord[pairs.index[0, :6]] - coord[pairs.index[1, :6]] + ) + assert build_topologies((), coord, np.array([4]), atype=atype) == {} + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies( + ("all_pairs",), coord, np.array([4]), np.eye(3)[None], atype=atype + ) + + +def test_pair_force_virial_of_a_central_potential() -> None: + """For E = sum 1/r the assembled force is Newtonian and the virial sums.""" + coord = np.array([[0.0, 0.0, 0.0], [1.0, 0.2, 0.0], [0.0, 1.5, 0.3]]) + pairs = build_pair_graph(coord, np.array([3]), np.array([0, 1, 0])) + # d(1/r)/d(vec) on the real pairs; the zero-length guards carry no gradient. + r = np.where(pairs.mask, np.linalg.norm(pairs.vec, axis=-1), 1.0) + grad = np.where(pairs.mask[:, None], -pairs.vec / r[:, None] ** 3, 0.0) + force, atom_virial, virial = pair_force_virial( + grad, pairs.vec, pairs.index, pairs.mask, np.array([3]) + ) + expected = np.zeros((3, 3)) + for i, j in itertools.permutations(range(3), 2): + d = coord[i] - coord[j] + expected[i] += d / np.linalg.norm(d) ** 3 + np.testing.assert_allclose(force, expected, rtol=1e-12) + np.testing.assert_allclose(atom_virial.sum(0), virial[0], rtol=1e-12) + + +def test_les_kernel_closed_form() -> None: + """LES is erf(alpha r)/r, finite at contact.""" + kernel = LESKernel(alpha=0.7) + r = np.array([0.0, 0.5, 2.0, 9.0]) + values = kernel(r**2) + np.testing.assert_allclose(values[1:], erf(0.7 * r[1:]) / r[1:], rtol=1e-10) + np.testing.assert_allclose(values[0], 2 * 0.7 / np.sqrt(np.pi), rtol=1e-10) + + +def test_sog_kernel_closed_form_and_defaults() -> None: + """SOG is sum_k a_k exp(-r^2 / 2 sigma_k^2); the defaults discretize 1/r.""" + kernel = SOGKernel(amplitude=[0.7, -0.2], width=[0.8, 1.7]) + r = np.array([0.0, 1.1, 3.0]) + expected = 0.7 * np.exp(-(r**2) / (2 * 0.8**2)) - 0.2 * np.exp( + -(r**2) / (2 * 1.7**2) + ) + np.testing.assert_allclose(kernel(r**2), expected, rtol=1e-12) + default = SOGKernel(n_gaussian=3, base_width=2.0, width_ratio=1.5) + np.testing.assert_allclose(np.exp(default.log_width), [2.0, 3.0, 4.5]) + np.testing.assert_allclose( + default.amplitude, np.sqrt(2 / np.pi) * np.log(1.5) / np.array([2.0, 3.0, 4.5]) + ) + # The default kernel exerts the Coulomb force, r^2 K'(r) = -1, between a + # few times its first width and a fraction of its widest one. + default = SOGKernel() + r = np.geomspace(6.5, 150.0, 200) + step = 1e-4 * r + slope = (default((r + step) ** 2) - default((r - step) ** 2)) / (2 * step) + np.testing.assert_allclose(r**2 * slope, -1.0, rtol=0, atol=5e-3) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"type": "les", "alpha": 0.0}, + {"type": "sog", "width": [1.0, -2.0]}, + {"type": "sog", "width": [1.0, 2.0], "amplitude": [1.0, 2.0, 3.0]}, + {"type": "sog", "width": [1.0, 2.0]}, + {"type": "sog", "width_ratio": 1.0}, + {"type": "sog", "amplitude": 1.0, "width_ratio": 0.5}, + ], +) +def test_kernel_parameter_domain(kwargs: dict) -> None: + """Widths must be positive and increase by default; amplitudes match the Gaussians.""" + kernel_cls = BaseInteractionKernel.get_class_by_type(kwargs.pop("type")) + with pytest.raises(ValueError): + kernel_cls(**kwargs) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=1.3), SOGKernel(n_gaussian=4)]) +def test_kernel_round_trip(kernel: BaseInteractionKernel) -> None: + """A kernel record restores the same parameters.""" + restored = BaseInteractionKernel.deserialize(kernel.serialize()) + r_sq = np.linspace(0.0, 30.0, 7) + np.testing.assert_allclose(restored(r_sq), kernel(r_sq), rtol=0, atol=0) + + +def _long_range(constraint: dict | None, channels: int = 2) -> LatentChargeLongRange: + charge_net = LatentChargeFitting( + "latent_charge", 2, 4, channels, neuron=[4], numb_fparam=1, seed=1 + ) + return LatentChargeLongRange(charge_net, LESKernel(), constraint) + + +def test_constraint_shifts_one_channel_over_the_charged_atoms() -> None: + """The constrained channel sums to the target over the charged atoms only.""" + long_range = _long_range({"source": "fparam", "channel": 1, "fparam_index": 0}) + rng = np.random.default_rng(2) + charge = rng.normal(size=(7, 2)) + node_mask = np.array([True, True, False, True, True, True, False]) + charge[~node_mask] = 0.0 + n_node = np.array([4, 3]) + frame_id = np.array([0, 0, 0, 0, 1, 1, 1]) + fparam = np.array([[1.5], [-2.0]]) + constrained = long_range._constrain( + charge, node_mask, frame_id, n_node, fparam, None + ) + np.testing.assert_array_equal(constrained[:, 0], charge[:, 0]) + np.testing.assert_array_equal(constrained[~node_mask], 0.0) + totals = np.bincount(frame_id, weights=constrained[:, 1], minlength=2) + np.testing.assert_allclose(totals, [1.5, -2.0], rtol=1e-12) + # A frame without charged atoms is left as it is. + empty = long_range._constrain( + charge * 0.0, + np.zeros(7, dtype=bool), + frame_id, + n_node, + fparam, + None, + ) + np.testing.assert_array_equal(empty, 0.0) + + +@pytest.mark.parametrize( + "constraint", + [ + {"source": "total"}, + {"source": "neutral", "channel": 2}, + {"source": "fparam", "fparam_index": 1}, + ], +) +def test_constraint_configuration_is_validated(constraint: dict) -> None: + """Unknown sources and out-of-range channels or columns are rejected.""" + with pytest.raises(ValueError): + _long_range(constraint) + + +def test_long_range_round_trip() -> None: + """The module record holds the charge network, kernel and constraint.""" + long_range = _long_range({"source": "neutral"}) + record = long_range.serialize() + restored = LatentChargeLongRange.deserialize(record) + assert restored.serialize()["charge_constraint"] == record["charge_constraint"] + assert restored.required_topologies() == ("all_pairs",) + assert not restored.uses_charge_spin() diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 9b61ee26cf..061ada901a 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -606,6 +606,9 @@ def has_default_chg_spin(self) -> bool: def get_default_chg_spin(self): return None + def has_long_range(self) -> bool: + return False + def _compose(self, children): return LinearEnergyAtomicModel(children, type_map=["Ni", "O"], weights="sum") diff --git a/source/tests/common/test_finetune_utils.py b/source/tests/common/test_finetune_utils.py index 33547b34d0..1f2bd45116 100644 --- a/source/tests/common/test_finetune_utils.py +++ b/source/tests/common/test_finetune_utils.py @@ -156,6 +156,28 @@ def test_finetune_rule_builder_updates_single_task_config_preserving_trainable() assert rule.get_has_new_type() +def test_finetune_rule_builder_carries_long_range_preserving_trainable(): + long_range = {"type": "latent_charge", "kernel": {"type": "les", "alpha": 0.7}} + pretrained = _model_config(["O", "H"]) + pretrained["long_range"] = deepcopy(long_range) + target = _model_config(["O", "H"], trainable=False) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"trainable": False}, + "kernel": {"type": "les", "alpha": 2.0, "trainable": False}, + } + + updated, _ = finetune.FinetuneRuleBuilder( + pretrained, + target, + change_model_params=True, + ).build() + + assert updated["long_range"]["kernel"]["alpha"] == 0.7 + assert updated["long_range"]["kernel"]["trainable"] is False + assert updated["long_range"]["charge_net"]["trainable"] is False + + def test_finetune_rule_builder_random_fitting_keeps_target_fitting_net(): pretrained = _model_config(["O", "H"], descriptor_sel=[8, 16], fitting_neuron=[32]) target = _model_config(["O", "H"], descriptor_sel=[1, 1], fitting_neuron=[2]) diff --git a/source/tests/infer/gen_dpa4c_long_range.py b/source/tests/infer/gen_dpa4c_long_range.py new file mode 100644 index 0000000000..554704c7ae --- /dev/null +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the DPA4C long-range .pt2 test models and reference values. + +A latent-charge long-range module sums over every pair of atoms of a complete +non-periodic system, so the C++ interface hands the model what the pair sum +needs beyond the neighbor list: the graph lower of an uncompressed model +receives the pair list the interface builds from the positions of the atoms, +and the compact canonical lower of a compressed model reads the positions +themselves. One archive covers each lower: + +- ``deeppot_dpa4c_long_range.pt2``: uncompressed, fp64, graph lower, which + compiles on the CPU; +- ``deeppot_dpa4c_long_range_compressed.pt2``: compressed, fp32, compact + canonical lower, whose operators are CUDA only; it is generated for the + CUDA variant of the C++ suite (``DP_VARIANT=cuda``). + +Both models condition the descriptor on the charge state and constrain the +latent charges to the total charge that state names, so a state other than +the default changes the long-range energy as well as the descriptor. The +references are those of the Python interface on each archive, without a cell, +at the default state and at ``[1, 2]``; each ``.expected`` sidecar holds the +sections ``default`` and ``explicit``. The six atoms form two groups about +9 Å apart, beyond the cutoff of the descriptor, which only the long-range +module couples. +""" + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +# The gen scripts run standalone, outside pytest's package machinery, so the +# relative imports the test suite uses do not apply here. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + +LONG_RANGE = { + "type": "latent_charge", + "channels": 1, + "charge_net": {"neuron": [8, 8], "seed": 9}, + "kernel": {"type": "les", "alpha": 0.8}, + "charge_constraint": {"source": "charge_spin"}, +} + +# The uncompressed model is fp64, so that its references hold to the +# precision of the double C++ interface. +CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7}, + "long_range": LONG_RANGE, +} + +# Compression takes an fp32 descriptor and, for the compact canonical lower, +# a fitting that the fused fitting operator evaluates. +COMPRESSED_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 8, + "precision": "float32", + "seed": 7, + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + }, + "fitting_net": { + "neuron": [32, 32], + "activation_function": "silu", + "resnet_dt": False, + "precision": "float32", + "seed": 7, + }, + "long_range": LONG_RANGE, +} + +# Jitters the zero-initialized arrays, among them the charge-state head of the +# descriptor, so that the charge state reaches the descriptor as well as the +# constraint. +_JITTER_SEED = 20261018 + +_ATYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) +_COORDS = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ], + dtype=np.float64, +).reshape(1, 6, 3) +_OTHER_CHG_SPIN = np.array([[1.0, 2.0]], dtype=np.float64) + + +def _model_dict(config: dict) -> dict: + """Build the dpmodel of a configuration and jitter it in place of its zeros.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + + model = get_model(copy.deepcopy(config)) + return jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED)) + + +def _check_metadata(pt2_path: str, lower: str) -> None: + """Assert the archive declares its lower and the long-range input it needs.""" + with zipfile.ZipFile(pt2_path) as zf: + md = json.loads(zf.read("model/extra/metadata.json").decode("utf-8")) + assert md["lower_input_kind"] == lower, md["lower_input_kind"] + if lower == "graph": + assert md["topologies"] == ["all_pairs"] + assert md["dim_chg_spin"] == 2 + else: + # The compressed descriptor folds the charge state; the archive + # carries the rebuild of every constant that holds it. + assert md["dim_chg_spin"] == 0 + assert len(md["charge_state_constants"]) == 5 + + +def _write_ref(pt2_path: str, ref_path: str) -> None: + """Evaluate the reference cases on the archive and write the C++ sidecar.""" + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + sections = {} + energies = {} + for name, charge_spin in (("default", None), ("explicit", _OTHER_CHG_SPIN)): + condition = {} if charge_spin is None else {"charge_spin": charge_spin} + e, f, v, ae, av = dp.eval(_COORDS, None, _ATYPES, atomic=True, **condition) + assert np.all(np.isfinite(e)) and np.all(np.isfinite(f)) + energies[name] = float(e[0, 0]) + print(f"// {name} total energy: {e[0, 0]:.18e}") # noqa: T201 + sections[name] = { + "expected_e": ae[0, :, 0], + "expected_f": f[0], + "expected_v": av[0], + } + # The C++ cases compare the two states, which needs a real difference. + assert abs(energies["explicit"] - energies["default"]) > 1e-3, ( + "the two charge states must differ, or a call that ignored the " + "requested state would still match" + ) + write_expected_ref( + ref_path, + sections=sections, + source_script="source/tests/infer/gen_dpa4c_long_range.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + + +def _export_compressed(pt2_path: str) -> None: + """Compress the fp32 model and export it for the CUDA device.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + model = BaseModel.deserialize(_model_dict(COMPRESSED_CONFIG)) + model.atomic_model.enable_compression(0.5) + data = { + "model": model.serialize(), + "model_def_script": COMPRESSED_CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {pt2_path} ...") # noqa: T201 + deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="auto") + + +def main() -> None: + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "deeppot_dpa4c_long_range.pt2") + data = { + "model": _model_dict(CONFIG), + "model_def_script": CONFIG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {pt2_path} (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file(pt2_path, data, do_atomic_virial=True, lower_kind="graph") + _check_metadata(pt2_path, "graph") + _write_ref(pt2_path, os.path.join(base_dir, "deeppot_dpa4c_long_range.expected")) + + if os.environ.get("DP_VARIANT") == "cuda": + compressed_path = os.path.join( + base_dir, "deeppot_dpa4c_long_range_compressed.pt2" + ) + _export_compressed(compressed_path) + _check_metadata(compressed_path, "dpa4c_canonical") + _write_ref( + compressed_path, + os.path.join(base_dir, "deeppot_dpa4c_long_range_compressed.expected"), + ) + + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/pt_expt/descriptor/test_dpa4c.py b/source/tests/pt_expt/descriptor/test_dpa4c.py index 4eef728456..eaa3b539a6 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4c.py +++ b/source/tests/pt_expt/descriptor/test_dpa4c.py @@ -265,6 +265,22 @@ def test_serialization_preserves_parameters(self) -> None: original_parameters[name], ) + def test_state_dict_of_version_3_2_loads(self) -> None: + """A checkpoint holding the gather tables of DeePMD-kit 3.2 loads strictly.""" + state = self.descriptor.state_dict() + width = 32 + state["angular_channel_index"] = torch.arange(width) + state["angular_harmonic_index"] = torch.arange(width) + restored = self.build() + restored.load_state_dict(state) + torch.testing.assert_close( + self._evaluate(restored, self.coord), + self._evaluate(self.descriptor, self.coord), + ) + state["unknown"] = torch.zeros(1) + with pytest.raises(RuntimeError, match="unknown"): + restored.load_state_dict(state) + def test_deserialization_preserves_version_one_raw_radial_basis(self) -> None: data = self.descriptor.serialize() radial_basis = data["radial_basis"] diff --git a/source/tests/pt_expt/kernels/test_long_range_direct_sum.py b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py new file mode 100644 index 0000000000..b5c775e8cc --- /dev/null +++ b/source/tests/pt_expt/kernels/test_long_range_direct_sum.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Correctness of the direct-sum operator of the latent-charge long range. + +The CUDA kernel evaluates every pair in float32 from positions relative to +the first node of each frame, carried as float32 head and tail, and sums in +compensated float32 before a float64 reduction; the total-charge constraint +and the per-frame reductions sum in float64 in a fixed order. It is held to +the float64 evaluation over the explicit pairs, at a tolerance relative to +the scale of each output, and to the exact invariants of a pair sum: padding +nodes and frames do not interact, the forces of a frame cancel, and the +energy, force, charge gradient and virial follow from one another. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.pt_expt.kernels.long_range import ( + _reference, + kernel_operator_arguments, + latent_charge_direct_sum, + op_available, +) +from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + SOGKernel, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not op_available(), + reason="the direct-sum operator is CUDA only", +) + +KERNELS = { + "les": lambda: LESKernel(alpha=0.8), + "sog": lambda: SOGKernel( + amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2], precision="float32" + ), +} + + +def _system( + n_node: list[int], channels: int, padding: int, seed: int, offset: float +) -> tuple[torch.Tensor, ...]: + """Frames of random nodes far from the origin, with trailing padding nodes. + + Every seventh node carries no charge, and every frame a total charge. + """ + rng = np.random.default_rng(seed) + total = sum(n_node) + padding + coord = offset + rng.uniform(0.0, 4.0 * max(n_node) ** (1 / 3), (total, 3)) + charge = rng.normal(0.0, 0.5, (total, channels)) + device = torch.device("cuda") + node_mask = torch.ones(total, dtype=torch.bool, device=device) + node_mask[::7] = False + return ( + torch.tensor(coord, dtype=torch.float64, device=device), + torch.tensor(charge, dtype=torch.float32, device=device), + node_mask, + torch.tensor(n_node, dtype=torch.int64, device=device), + torch.tensor( + rng.integers(-2, 3, len(n_node)), dtype=torch.float64, device=device + ), + ) + + +def _evaluate( + system: tuple[torch.Tensor, ...], kernel: Any, channel: int +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + coord, charge, node_mask, n_node, total_charge = system + actual = latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node, + total_charge, + kernel, + channel, + COULOMB_CONSTANT, + ) + code, coefficient = kernel_operator_arguments(kernel) + expected = _reference( + coord, + charge.double(), + node_mask, + n_node, + total_charge, + coefficient.double(), + torch.zeros(1, dtype=torch.int64), + code, + channel, + COULOMB_CONSTANT, + ) + return actual, expected + + +NAMES = ("atom_energy", "charge_gradient", "force", "atom_virial", "virial") + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel"), + [ + ([37, 1, 0, 300], 1, 5, 0), + ([37, 1, 0, 300], 1, 5, -1), + ([2048], 3, 0, -1), + ([5, 129], 8, 3, 2), + ], +) +def test_matches_float64_pair_sum( + kernel_name: str, n_node: list[int], channels: int, padding: int, channel: int +) -> None: + """Every output agrees with the float64 pair sum.""" + system = _system(n_node, channels, padding, 7, 1.0e3) + actual, expected = _evaluate(system, KERNELS[kernel_name](), channel) + _assert_close(actual, expected) + # Padding nodes join no frame. + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("alpha", [1e-3, 0.03, 0.3]) +@pytest.mark.parametrize("channel", [0, -1]) +def test_les_matches_float64_pair_sum_at_small_alpha_r( + alpha: float, channel: int +) -> None: + """The erf kernel keeps its precision where alpha r is small. + + The two terms of the slope of erf(x) / x agree to second order in x, so + that the slope of an approximate erf loses precision in proportion to + 1 / x^2 at atomic distances and a small alpha. + """ + system = _system([96], 2, 0, 5, 0.0) + _assert_close(*_evaluate(system, LESKernel(alpha=alpha), channel)) + + +def _assert_close( + actual: tuple[torch.Tensor, ...], expected: tuple[torch.Tensor, ...] +) -> None: + """Every output agrees with the float64 pair sum, relative to its largest entry.""" + for name, value, reference in zip(NAMES, actual, expected, strict=True): + assert value.dtype == ( + torch.float64 if name == "atom_energy" else torch.float32 + ) + scale = reference.abs().max().item() + error = (value.double() - reference.double()).abs().max().item() + assert error <= 2e-5 * scale, (name, error, scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_invariants(kernel_name: str) -> None: + """Frames are independent, forces cancel, and the virial is the pair virial.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + together = latent_charge_direct_sum( + coord, + charge, + node_mask, + n_node_t, + total_charge, + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + _, _, force, atom_virial, virial = together + # Each frame alone gives the same result. + first = n_node[0] + alone = latent_charge_direct_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + KERNELS[kernel_name](), + 1, + COULOMB_CONSTANT, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + # Internal forces sum to zero per frame, and the virial of a frame is + # sum_i r_i (x) f_i. + for frame, (begin, end) in enumerate(((0, first), (first, sum(n_node)))): + frame_force = force[begin:end].double() + assert frame_force.sum(0).abs().max() <= 1e-5 * frame_force.abs().max() + position = coord[begin:end] - coord[begin] + expected = torch.einsum("ni,nj->ij", position, frame_force) + scale = expected.abs().max().item() + for value in ( + atom_virial[begin:end].double().sum(0).reshape(3, 3), + virial[frame].double().reshape(3, 3), + ): + torch.testing.assert_close(value, expected, rtol=0, atol=2e-5 * scale) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize("channel", [-1, 1]) +def test_reference_outputs_are_energy_derivatives( + kernel_name: str, channel: int +) -> None: + """The float64 force and charge gradient differentiate the energy.""" + coord, charge, node_mask, n_node, total_charge = _system([9, 6], 2, 2, 3, 0.0) + coord = coord.clone().requires_grad_(True) + charge = charge.double().requires_grad_(True) + kernel = KERNELS[kernel_name]() + code, coefficient = kernel_operator_arguments(kernel) + atom_energy, charge_gradient, force, _, _ = _reference( + coord, + charge, + node_mask, + n_node, + total_charge, + coefficient.double(), + torch.zeros(1, dtype=torch.int64), + code, + channel, + COULOMB_CONSTANT, + ) + grad_coord, grad_charge = torch.autograd.grad(atom_energy.sum(), (coord, charge)) + torch.testing.assert_close(force, -grad_coord, rtol=1e-9, atol=1e-9) + torch.testing.assert_close( + charge_gradient.detach(), grad_charge, rtol=1e-9, atol=1e-9 + ) + + +def test_bitwise_reproducible() -> None: + """Two evaluations of the same input agree bit for bit.""" + system = _system([3000], 1, 0, 5, 0.0) + first = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) + second = latent_charge_direct_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) diff --git a/source/tests/pt_expt/model/test_long_range.py b/source/tests/pt_expt/model/test_long_range.py new file mode 100644 index 0000000000..606e078245 --- /dev/null +++ b/source/tests/pt_expt/model/test_long_range.py @@ -0,0 +1,1066 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Long-range latent-charge modules on the DPA4C graph lower (pt_expt).""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.long_range import ( + COULOMB_CONSTANT, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + build_pair_graph, +) +from deepmd.infer import ( + DeepPot, +) +from deepmd.pt_expt.model.get_model import ( + get_model, +) +from deepmd.pt_expt.train.training import ( + _CompiledModel, + _get_model_structure_key, + model_change_out_bias, +) +from deepmd.pt_expt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, +) + +from ..compile_utils import ( + REQUIRES_SUPPORTED_COMPILE, +) + +KERNELS = { + "les": {"type": "les", "alpha": 0.8}, + "sog": {"type": "sog", "amplitude": [0.7, -0.2, 0.1], "width": [0.8, 1.7, 3.2]}, +} +CHANNELS = {"les": 1, "sog": 3} +NEUTRAL = {"source": "neutral"} + + +def _config(kernel: str = "les", constraint: dict | None = NEUTRAL) -> dict: + """Small float64 DPA4C model with a latent-charge long-range module.""" + return { + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4c", + "rcut": 3.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 17, + }, + "fitting_net": { + "type": "ener", + "neuron": [16, 16], + "precision": "float64", + "seed": 19, + }, + "long_range": { + "type": "latent_charge", + "channels": CHANNELS[kernel], + "charge_net": {"neuron": [8, 8], "precision": "float64", "seed": 23}, + "kernel": copy.deepcopy(KERNELS[kernel]), + "charge_constraint": copy.deepcopy(constraint), + }, + } + + +def _model(kernel: str = "les", **kwargs: object) -> torch.nn.Module: + return get_model(_config(kernel, **kwargs)).to(env.DEVICE).eval() + + +def _tensor(x: np.ndarray, dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.tensor(x, dtype=dtype, device=env.DEVICE) + + +def _batch(seed: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Three frames of 6, 3 and 5 real atoms, padded to a common width.""" + rng = np.random.default_rng(seed) + coord = rng.uniform(0.0, 6.0, (3, 6, 3)) + atype = np.array( + [[0, 1, 0, 1, 1, 0], [1, 0, 0, -1, -1, -1], [0, 0, 1, 1, 0, -1]], + dtype=np.int64, + ) + return _tensor(coord), _tensor(atype, torch.int64) + + +def _set_constant_charges(model: torch.nn.Module, charges: list[float]) -> None: + """Zero every network weight and give the atoms of each type a fixed charge. + + The energy is then the long-range pair sum alone, of charges that do not + depend on the geometry; the kernel keeps its configured parameters. + """ + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.long_range.charge_net.bias_atom_e + bias.zero_() + bias[:, 0] = torch.as_tensor(charges, dtype=bias.dtype) + + +def _reference_energy(coord: np.ndarray, charge: np.ndarray, kernel: str) -> float: + """Direct pair sum of the latent-charge energy of one frame, in eV.""" + energy = 0.0 + for i in range(len(coord)): + for j in range(i + 1, len(coord)): + r = np.linalg.norm(coord[i] - coord[j]) + if kernel == "les": + from scipy.special import ( + erf, + ) + + k = erf(0.8 * r) / r + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k = np.sum(amplitude * np.exp(-(r**2) / (2 * width**2))) + energy += COULOMB_CONSTANT * float(charge[i] @ charge[j]) * k + return energy + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_energy_matches_direct_pair_sum(kernel: str) -> None: + """With fixed charges the energy is the documented pair sum, attributed half per atom.""" + model = _model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord = np.array([[[0.0, 0.0, 0.0], [1.3, 0.0, 0.0], [0.0, 2.1, 0.4]]]) + atype = np.array([[0, 1, 1]]) + out = model(_tensor(coord), _tensor(atype, torch.int64)) + charge = np.zeros((3, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, -0.5] + expected = _reference_energy(coord[0], charge, kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-12, atol=1e-14 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, affine-strain virial and Hessian are derivatives of the complete energy.""" + model = _model(kernel) + # Atoms within the cutoff of each other make the charges depend on the + # geometry, so the derivatives hold the response of the charges. + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 2.5, (1, 4, 3)) + atype = np.array([[0, 1, 0, 1]]) + coord_t, atype_t = _tensor(coord), _tensor(atype, torch.int64) + out = model(coord_t, atype_t) + + def energy(c: np.ndarray) -> float: + return model(_tensor(c), atype_t)["energy"].item() + + h = 1e-5 + force_fd = np.zeros((4, 3)) + for i in range(4): + for d in range(3): + plus, minus = coord.copy(), coord.copy() + plus[0, i, d] += h + minus[0, i, d] -= h + force_fd[i, d] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-8 + ) + + # The virial is -dE/d(strain) at zero strain, for x -> (1 + eps) x. + virial_fd = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + strain = np.zeros((3, 3)) + strain[a, b] = h + plus = coord @ (np.eye(3) + strain) + minus = coord @ (np.eye(3) - strain) + virial_fd[a, b] = -(energy(plus) - energy(minus)) / (2 * h) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-8 + ) + + hessian_model = _model(kernel) + hessian_model.enable_hessian() + hessian = hessian_model(coord_t, atype_t)["hessian"].detach().cpu().numpy() + hessian = hessian.reshape(12, 12) + hessian_fd = np.zeros((12, 12)) + for k in range(12): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[k] += h + minus.reshape(-1)[k] -= h + f_plus = model(_tensor(plus), atype_t)["force"].detach().cpu().numpy() + f_minus = model(_tensor(minus), atype_t)["force"].detach().cpu().numpy() + hessian_fd[:, k] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-6) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_atomic_reductions_and_padding(kernel: str) -> None: + """Atomic energy and virial sum to the frame values; padding atoms are inert.""" + model = _model(kernel) + coord, atype = _batch() + out = model(coord, atype, do_atomic_virial=True) + torch.testing.assert_close(out["atom_energy"].sum(1), out["energy"]) + torch.testing.assert_close(out["atom_virial"].sum(1), out["virial"]) + assert out["force"][1, 3:].abs().max() == 0 + moved = coord.clone() + moved[1, 3:] += 50.0 + moved_out = model(moved, atype) + torch.testing.assert_close(moved_out["energy"], out["energy"], rtol=0, atol=0) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_frames_are_independent(kernel: str) -> None: + """A frame keeps its energy and force whatever other frames are batched with it.""" + model = _model(kernel) + coord, atype = _batch() + batched = model(coord, atype) + for frame, nloc in enumerate((6, 3, 5)): + alone = model(coord[frame : frame + 1, :nloc], atype[frame : frame + 1, :nloc]) + torch.testing.assert_close(alone["energy"][0], batched["energy"][frame]) + torch.testing.assert_close( + alone["force"][0], batched["force"][frame, :nloc], atol=1e-12, rtol=0 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_ragged_matches_rectangular(kernel: str) -> None: + """The flat ragged entry evaluates the same model as the padded one.""" + model = _model(kernel) + coord, atype = _batch() + padded = model(coord, atype) + keep = atype.reshape(-1) >= 0 + ragged = model.forward_ragged( + coord.reshape(-1, 3)[keep], + atype.reshape(-1)[keep], + torch.tensor([6, 3, 5], device=env.DEVICE), + ) + torch.testing.assert_close(ragged["energy"], padded["energy"]) + torch.testing.assert_close(ragged["force"], padded["force"].reshape(-1, 3)[keep]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_empty_and_single_atom_frames(kernel: str) -> None: + """Frames without pairs contribute no long-range energy and stay finite.""" + model = _model(kernel) + coord = _tensor(np.random.default_rng(5).uniform(0.0, 3.0, (3, 2, 3))) + atype = _tensor(np.array([[0, -1], [1, 0], [-1, -1]]), torch.int64) + out = model(coord, atype) + assert torch.isfinite(out["energy"]).all() + assert torch.isfinite(out["force"]).all() + single = model(coord[:1, :1], atype[:1, :1]) + torch.testing.assert_close(single["energy"][0], out["energy"][0]) + + +def test_permutation_and_rigid_motion_invariance() -> None: + """The energy is invariant under atom permutation, translation and rotation.""" + model = _model("sog") + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + out = model(coord, atype) + perm = torch.tensor([3, 0, 5, 1, 4, 2], device=env.DEVICE) + permuted = model(coord[:, perm], atype[:, perm]) + torch.testing.assert_close(permuted["energy"], out["energy"]) + torch.testing.assert_close(permuted["force"], out["force"][:, perm]) + rotation, _ = torch.linalg.qr( + _tensor(np.random.default_rng(13).normal(size=(3, 3))) + ) + moved = model(coord @ rotation.T + 3.7, atype) + torch.testing.assert_close(moved["energy"], out["energy"]) + torch.testing.assert_close(moved["force"], out["force"] @ rotation.T) + + +@pytest.mark.parametrize("source", ["charge_spin", "fparam"]) +def test_constraint_default_matches_explicit_target(source: str) -> None: + """An omitted frame condition takes its default for the constraint as well.""" + config = _config("les", constraint={"source": source, "fparam_index": 0}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [2.0, 1.0] + explicit = {"charge_spin": _tensor(np.array([[2.0, 1.0], [2.0, 1.0]]))} + other = {"charge_spin": _tensor(np.array([[0.0, 1.0], [0.0, 1.0]]))} + else: + config["fitting_net"]["numb_fparam"] = 1 + config["fitting_net"]["default_fparam"] = [2.0] + explicit = {"fparam": _tensor(np.array([[2.0], [2.0]]))} + other = {"fparam": _tensor(np.array([[0.0], [0.0]]))} + model = get_model(config).to(env.DEVICE).eval() + coord, atype = _batch() + coord, atype = coord[:2], atype[:2] + implicit_out = model(coord, atype) + explicit_out = model(coord, atype, **explicit) + torch.testing.assert_close(implicit_out["energy"], explicit_out["energy"]) + torch.testing.assert_close(implicit_out["force"], explicit_out["force"]) + other_out = model(coord, atype, **other) + assert (other_out["energy"] - implicit_out["energy"]).abs().max() > 1e-6 + + +@pytest.mark.parametrize("constraint", [None, NEUTRAL]) +def test_pair_exclusion_removes_the_interaction(constraint: dict | None) -> None: + """Pairs of excluded types do not interact; their atoms keep their charges.""" + config = _config("les", constraint=constraint) + config["pair_exclude_types"] = [[0, 1]] + model = get_model(config).to(env.DEVICE).eval() + _set_constant_charges(model, [1.0, -1.0]) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + out = model(coord, atype) + # Only the two type-0 atoms interact. The neutral constraint shifts the + # charges (1, 1, -1) of the three atoms to (2/3, 2/3, -4/3). + charge = 1.0 if constraint is None else 2.0 / 3.0 + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + expected = COULOMB_CONSTANT * charge**2 * erf / 4.0 + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-12) + assert out["force"][0, 2].abs().max() == 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_les_kernel_derivatives_down_to_contact(dtype: torch.dtype) -> None: + """The radial derivatives of LES keep the precision of the dtype near contact. + + The derivatives of erf(alpha r) / r subtract terms that agree to second + order in alpha r, which the series of the kernel avoids; the reference + differentiates in 40-digit arithmetic. + """ + import mpmath + + from deepmd.pt_expt.long_range.kernel import ( + LESKernel, + ) + + alpha = 0.8 + kernel = LESKernel(alpha=alpha) + r = torch.tensor( + np.geomspace(1e-5, 6.0, 80), dtype=dtype, device=env.DEVICE + ).requires_grad_(True) + (first,) = torch.autograd.grad(kernel(r * r).sum(), r, create_graph=True) + (second,) = torch.autograd.grad(first.sum(), r) + + def les(t: mpmath.mpf) -> mpmath.mpf: + return mpmath.erf(alpha * t) / t + + rtol = 2e-5 if dtype == torch.float32 else 1e-11 + for order, actual in ((1, first), (2, second)): + with mpmath.workdps(40): + expected = np.array( + [ + float(mpmath.diff(les, mpmath.mpf(float(t)), order)) + for t in r.detach().cpu() + ] + ) + # The second derivative changes sign near 1.2 Å, where it is the + # difference of terms much larger than itself. + np.testing.assert_allclose( + actual.detach().cpu().numpy(), + expected, + rtol=rtol, + atol=rtol * np.abs(expected).max(), + err_msg=str(order), + ) + + +def test_owned_frames_across_domains() -> None: + """Halo nodes are rejected on a single domain and kept across domains.""" + from deepmd.pt_expt.model.make_model import ( + check_owned_frames, + ) + + n_node = torch.tensor([5, 3], dtype=torch.int64) + n_local = torch.tensor([4, 3], dtype=torch.int64) + check_owned_frames(n_node, n_node.clone()) + check_owned_frames(n_node, n_local, torch.ones(1, dtype=torch.int64)) + for communicator in (None, torch.zeros(1, dtype=torch.int64)): + with pytest.raises(ValueError, match="n_local == n_node"): + check_owned_frames(n_node, n_local, communicator) + + +def test_graph_lower_rejects_halo_nodes() -> None: + """A graph whose frames carry halo nodes is rejected, not summed without them. + + The halo nodes of a graph without communication would carry no charge and + remove their interactions from the sum of a frame held whole. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + from deepmd.pt_expt.utils.graph_builder import ( + build_neighbor_graph_for_method, + ) + + model = _model("les", constraint=None) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + graph = build_neighbor_graph_for_method( + "dense", coord, atype, None, model.get_rcut(), None + ) + topologies = build_topologies( + model.atomic_model.required_topologies(), + coord.reshape(-1, 3), + graph.n_node, + atype=atype.reshape(-1), + ) + inputs = ( + atype.reshape(-1), + graph.n_node, + graph.edge_index, + graph.edge_vec, + graph.edge_mask, + ) + + def lower(n_local: torch.Tensor) -> dict[str, torch.Tensor]: + return model.forward_common_lower_graph( + inputs[0], inputs[1], n_local, *inputs[2:], topologies=topologies + ) + + lower(graph.n_node.clone()) + with pytest.raises(ValueError, match="n_local == n_node"): + lower(graph.n_node - 1) + + +def test_linear_composition_takes_the_exclusions_of_the_composition() -> None: + """A long-range child follows the pair exclusion of its composition. + + The children of a composition share its neighbor graph and pair topology, + so the exclusions of the composition give the child the semantics of the + same model alone, and an exclusion of the child itself is rejected. + """ + child = _config("les", constraint=None) + child["type"] = "standard" + composition = { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "pair_exclude_types": [[0, 1]], + } + model = get_model(copy.deepcopy(composition)).to(env.DEVICE).eval() + with torch.no_grad(): + for name, parameter in model.atomic_model.named_parameters(): + if ".kernel." not in f".{name}": + parameter.zero_() + bias = model.atomic_model.models[0].long_range.charge_net.bias_atom_e + bias[:, 0] = torch.as_tensor([1.0, -1.0], dtype=bias.dtype) + coord = _tensor(np.array([[[0.0, 0.0, 0.0], [4.0, 0.0, 0.0], [8.0, 0.0, 0.0]]])) + atype = _tensor(np.array([[0, 0, 1]]), torch.int64) + # As in test_pair_exclusion_removes_the_interaction: only the two type-0 + # atoms interact. + erf = torch.erf(torch.tensor(3.2, dtype=torch.float64)).item() + np.testing.assert_allclose( + model(coord, atype)["energy"].item(), + COULOMB_CONSTANT * erf / 4.0, + rtol=1e-12, + ) + for key, value, error, match in ( + ("pair_exclude_types", [[0, 1]], ValueError, "from the composition"), + ("atom_exclude_types", [1], NotImplementedError, "atom_exclude_types"), + ): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"][0][key] = value + with pytest.raises(error, match=match): + get_model(config) + + +@pytest.mark.parametrize("compiled", [False, True]) +def test_vacuum_reference_with_long_range(compiled: bool) -> None: + """An isolated atom keeps its reference energy; reference rows join no pair.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["type_map"] = ["O", "H"] + config["preset_out_bias"] = {"energy": {"O": -3.0, "H": 0.5}} + config["descriptor"]["add_chg_spin_ebd"] = True + config["descriptor"]["default_chg_spin"] = [0.0, 1.0] + config["fitting_net"]["vacuum_ref"] = True + model = get_model(config).to(env.DEVICE).eval() + assert model.get_fitting_net().needs_vacuum_descriptor() + rng = np.random.default_rng(11) + coord = _tensor(rng.uniform(0.0, 2.5, (2, 4, 3))) + atype = _tensor(np.array([[0, 1, 1, 0], [1, 0, -1, -1]]), torch.int64) + if compiled: + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model(coord, atype) + out = runner(coord, atype) + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + return + out = model(coord, atype) + for frame in range(2): + alone = model(coord[frame : frame + 1], atype[frame : frame + 1]) + torch.testing.assert_close(alone["energy"][0], out["energy"][frame]) + # The output statistics set the bias; an atom without neighbors and in + # its reference state then has exactly the bias of its type. + bias = model.atomic_model.out_bias.clone() + bias[0, :, 0] = torch.tensor([-3.0, 0.5]) + model.atomic_model.set_out_bias(bias) + isolated = model( + coord[:1, :1], + atype[:1, :1], + charge_spin=_tensor(np.array([[0.0, 3.0]])), + ) + torch.testing.assert_close( + isolated["energy"].reshape(()), + torch.tensor(-3.0, dtype=torch.float64, device=env.DEVICE), + rtol=0, + atol=1e-12, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_self_label_bias_adjustment_is_invariant(kernel: str) -> None: + """Labelling a sample with the model's own energies leaves its bias unchanged.""" + model = _model(kernel) + with torch.no_grad(): + model.atomic_model.long_range.charge_net.bias_atom_e[:, 0] = torch.tensor( + [0.8, -0.6] + ) + coord, atype = _batch() + coord, atype = coord[:1], atype[:1] + energy = model(coord, atype)["energy"].detach().cpu().numpy() + natoms = np.array([[6, 6, 3, 3]], dtype=np.int32) + sample = [ + { + "coord": coord.cpu().numpy(), + "atype": atype.cpu().numpy(), + "box": None, + "natoms": natoms, + "energy": energy.reshape(1, 1), + "find_energy": np.float32(1.0), + } + ] + bias_before = model.atomic_model.out_bias.detach().clone() + model_change_out_bias(model, lambda: sample, "change-by-statistic") + torch.testing.assert_close( + model.atomic_model.out_bias, bias_before, atol=1e-10, rtol=0 + ) + after = model(coord, atype)["energy"].detach().cpu().numpy() + np.testing.assert_allclose(after, energy, atol=1e-10) + + +@pytest.mark.parametrize("weights", ["sum", "mean", [0.3, 1.4]]) +def test_linear_composition_combines_complete_energies(weights: object) -> None: + """A linear composition weights the complete energy of each child. + + The second child is an ordinary model, which reads no topology. + """ + child = _config("les") + child["type"] = "standard" + other = copy.deepcopy(child) + other.pop("long_range") + other["descriptor"]["seed"] = 31 + composition = ( + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": weights, + "models": [child, other], + } + ) + .to(env.DEVICE) + .eval() + ) + standalone = [_model("les"), get_model(other).to(env.DEVICE).eval()] + for index, model in enumerate(standalone): + model.atomic_model.load_state_dict( + composition.atomic_model.models[index].state_dict() + ) + coord, atype = _batch() + out = composition(coord, atype) + if weights == "sum": + w = [1.0, 1.0] + elif weights == "mean": + w = [0.5, 0.5] + else: + w = weights + energy = sum( + ww * m(coord, atype)["energy"] for ww, m in zip(w, standalone, strict=True) + ) + force = sum( + ww * m(coord, atype)["force"] for ww, m in zip(w, standalone, strict=True) + ) + torch.testing.assert_close(out["energy"], energy) + torch.testing.assert_close(out["force"], force) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_parameter_gradients_and_freezing(kernel: str) -> None: + """Energy and force losses reach the charge network and the kernel; frozen parts stay out. + + Every parameter of the module moves the losses, except under a constraint + the common charge bias of the constrained channel, whose gradient + vanishes (see test_constraint_removes_the_common_charge_bias). + """ + coord, atype = _batch() + for constraint in (None, NEUTRAL): + model = _model(kernel, constraint=constraint).train() + out = model(coord, atype) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + long_range = model.atomic_model.long_range + bias = long_range.charge_net.nets[0].layers[-1].b + for name, parameter in long_range.named_parameters(): + assert parameter.grad is not None, name + assert torch.isfinite(parameter.grad).all(), name + gradient = parameter.grad + if constraint is not None and parameter is bias: + # Channel 0 is the constrained channel. + gradient = gradient[..., 1:] + if gradient.numel(): + assert gradient.abs().max() > 0, name + + config = _config(kernel) + config["long_range"]["kernel"]["trainable"] = False + config["long_range"]["charge_net"]["trainable"] = False + frozen = get_model(config).to(env.DEVICE).train() + names = [ + name + for name, parameter in frozen.named_parameters() + if parameter.requires_grad and "long_range" in name + ] + assert names == [] + frozen(coord, atype)["energy"].sum().backward() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_constraint_removes_the_common_charge_bias(kernel: str) -> None: + """A uniform shift of the raw charges of the constrained channel changes nothing. + + The constraint subtracts the mean shift of every frame, so the output bias + of that channel leaves energy and force unchanged, while that of another + channel does not. + """ + model = _model(kernel) + coord, atype = _batch() + reference = model(coord, atype) + bias = model.atomic_model.long_range.charge_net.nets[0].layers[-1].b + for channel in range(bias.shape[-1]): + with torch.no_grad(): + bias[..., channel] += 0.3 + out = model(coord, atype) + with torch.no_grad(): + bias[..., channel] -= 0.3 + if channel == 0: + torch.testing.assert_close(out["energy"], reference["energy"]) + torch.testing.assert_close(out["force"], reference["force"]) + else: + assert (out["energy"] - reference["energy"]).abs().max() > 1e-6 + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_serialization_round_trip(kernel: str) -> None: + """The portable record and the state dict restore the same model.""" + model = _model(kernel) + rng = np.random.default_rng(5) + with torch.no_grad(): + for parameter in model.atomic_model.long_range.parameters(): + parameter.add_(0.01 * _tensor(rng.normal(size=tuple(parameter.shape)))) + coord, atype = _batch() + expected = model(coord, atype) + record = model.serialize() + assert record["@version"] == 3 + restored = type(model).deserialize(copy.deepcopy(record)).to(env.DEVICE).eval() + torch.testing.assert_close(restored(coord, atype)["energy"], expected["energy"]) + rebuilt = _model(kernel) + rebuilt.load_state_dict(model.state_dict()) + torch.testing.assert_close(rebuilt(coord, atype)["force"], expected["force"]) + + +def test_construction_is_reproducible() -> None: + """Without a seed of its own the charge network follows the fitting seed.""" + config = _config("les") + del config["long_range"]["charge_net"]["seed"] + first, second = (get_model(copy.deepcopy(config)) for _ in range(2)) + for (name, a), (_, b) in zip( + first.named_parameters(), second.named_parameters(), strict=True + ): + torch.testing.assert_close(a, b, rtol=0, atol=0, msg=name) + + +def test_ordinary_record_keeps_version() -> None: + """A model without a long-range module keeps its record and constructor.""" + config = _config("les") + config.pop("long_range") + model = get_model(config) + record = model.serialize() + assert record["@version"] == 2 + assert "long_range" not in record + assert not model.atomic_model.has_long_range() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_checkpoint_inference_matches_direct_forward(tmp_path, kernel: str) -> None: + """DeepPot on a checkpoint builds the topology over the real atoms only.""" + model = _model(kernel) + checkpoint = tmp_path / f"long_range_{kernel}.pt" + wrapper = ModelWrapper(model, model_params=_config(kernel)) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype = _batch(seed=7) + expected = model(coord, atype, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord_np = coord.cpu().numpy() + actual = dict( + zip( + names, + dp.eval(coord_np, None, atype.cpu().numpy(), atomic=True, mixed_type=True), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + coord_np[1, 3:] += 40.0 + moved = dp.eval(coord_np, None, atype.cpu().numpy(), mixed_type=True) + np.testing.assert_allclose(moved[0], actual["energy"], rtol=0, atol=1e-12) + + +@pytest.mark.parametrize( + ("kernel", "source", "suffix"), + [("les", "neutral", ".pt2"), ("sog", "neutral", ".pte"), ("les", "fparam", ".pte")], +) +def test_frozen_model_matches_direct_forward( + tmp_path, kernel: str, source: str, suffix: str +) -> None: + """A frozen model reads the pair topology as named inputs of any length.""" + config = _config(kernel, constraint={"source": source, "fparam_index": 0}) + frame = {} + if source == "fparam": + config["fitting_net"]["numb_fparam"] = 1 + frame["fparam"] = np.array([[1.0], [-0.5], [2.0]]) + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"long_range_{kernel}{suffix}" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="graph", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + assert dp.deep_eval.metadata["topologies"] == ["all_pairs"] + if suffix == ".pte": + names = [ + node.name + for node in torch.export.load(str(frozen)).graph_module.graph.nodes + if node.op == "placeholder" + ] + for name, expected_name in zip( + names[-3:], ("pair_index", "pair_vec", "pair_mask"), strict=True + ): + assert name.startswith(expected_name), names + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + for seed, nloc in ((7, 6), (8, 4)): + coord, atype = _batch(seed=seed) + coord, atype = coord[:, :nloc], atype[:, :nloc] + expected = model( + coord, + atype, + do_atomic_virial=True, + **{key: _tensor(value) for key, value in frame.items()}, + ) + actual = dp.eval( + coord.cpu().numpy(), + None, + atype.cpu().numpy(), + atomic=True, + mixed_type=True, + **frame, + ) + for name, value in zip(names, actual, strict=True): + np.testing.assert_allclose( + value, + expected[name].detach().cpu().numpy().reshape(value.shape), + rtol=1e-10, + atol=1e-10, + err_msg=name, + ) + + +@pytest.mark.parametrize( + ("kernel", "source"), [("les", None), ("les", "neutral"), ("sog", "charge_spin")] +) +def test_coordinate_evaluation_matches_pair_graph( + kernel: str, source: str | None +) -> None: + """The evaluation from node positions matches the pair-graph evaluation. + + It returns the atomic energy, the derivative of the energy with respect to + the descriptor through the constrained charges, and the force at a fixed + descriptor, all of which the pair graph yields by automatic + differentiation. Two frames, an uncharged node and trailing padding nodes + exercise the frame sums of the constraint. + """ + from deepmd.pt_expt.kernels.long_range import ( + op_available, + ) + + if not op_available(): + pytest.skip("the direct-sum operator is not available") + config = _config(kernel, constraint=None if source is None else {"source": source}) + if source == "charge_spin": + config["descriptor"]["add_chg_spin_ebd"] = True + long_range = get_model(config).atomic_model.long_range.eval() + if source == "charge_spin": + long_range.fold_charge_state([2.0, 1.0]) + rng = np.random.default_rng(7) + n_node = torch.tensor([9, 6], device=env.DEVICE) + node_mask = torch.ones(17, dtype=torch.bool, device=env.DEVICE) + node_mask[4] = False + node_mask[15:] = False + coord = _tensor(rng.uniform(0.0, 6.0, (17, 3))) + descriptor = _tensor(rng.normal(size=(17, long_range.charge_net.dim_descrpt))) + atype = _tensor(rng.integers(0, 2, 17), torch.int64) + + atom_energy, descriptor_gradient, force, _, _ = long_range.energy_force_from_coord( + descriptor, atype, node_mask, n_node, coord + ) + + leaf_descriptor = descriptor.clone().requires_grad_(True) + leaf_coord = coord.clone().requires_grad_(True) + pairs = build_pair_graph(leaf_coord, n_node, atype, node_mask) + expected = long_range.call_graph( + leaf_descriptor, atype, node_mask, n_node, {ALL_PAIRS: pairs} + ) + expected_descriptor_gradient, expected_coord_gradient = torch.autograd.grad( + expected.sum(), (leaf_descriptor, leaf_coord) + ) + for actual, target in ( + (atom_energy, expected), + (descriptor_gradient, expected_descriptor_gradient), + (force, -expected_coord_gradient), + ): + target = target.detach() + scale = target.abs().max().item() + assert scale > 0.0 + np.testing.assert_allclose( + actual.detach().to(torch.float64).cpu().numpy(), + target.cpu().numpy(), + rtol=0.0, + atol=2e-5 * scale, + ) + + +@pytest.mark.parametrize( + ("kernel", "source"), [("les", "charge_spin"), ("sog", "neutral")] +) +def test_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str +) -> None: + """A compressed model keeps its long-range module through the freeze. + + On a CUDA device the freeze selects the compact canonical lower, whose + direct-sum operator reads the positions of the atoms; on the CPU the + graph lower reads the pair lists. A constraint that reads the total + charge from the charge state is folded with the descriptor and follows + the state the deployment serves. Both match the uncompressed model. + """ + config = _config(kernel, constraint={"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + if source == "charge_spin": + config["descriptor"].update(add_chg_spin_ebd=True, default_chg_spin=[0.0, 1.0]) + config["fitting_net"].update( + neuron=[32, 32], activation_function="silu", resnet_dt=False + ) + config["fitting_net"]["precision"] = "float32" + config["long_range"]["charge_net"] = {"neuron": [16, 16], "seed": 23} + model = get_model(config).to(env.DEVICE).eval() + generator = torch.Generator().manual_seed(5) + with torch.no_grad(): + for parameter in model.parameters(): + noise = torch.randn(parameter.shape, generator=generator) + parameter.add_(0.05 * noise.to(parameter)) + reference = copy.deepcopy(model) + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"long_range_compressed_{kernel}.pt2" + deserialize_to_file( + str(frozen), + {"model": model.serialize(), "model_def_script": config}, + do_atomic_virial=True, + lower_kind="auto", + ) + dp = DeepPot(str(frozen), auto_batch_size=False) + expected_lower = "dpa4c_canonical" if env.DEVICE.type == "cuda" else "graph" + assert dp.deep_eval.metadata["lower_input_kind"] == expected_lower + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + # The tables of the compressed descriptor approximate its networks. On the + # graph lower the charges of this model amplify that error at a charged + # state, to about 1e-4 of the scale at [1, 2]; the compact canonical lower + # stays within a few 1e-6. + tolerance = 2e-5 if expected_lower == "dpa4c_canonical" else 5e-4 + states = [[0.0, 1.0], [1.0, 2.0], [0.0, 1.0]] if source == "charge_spin" else [None] + for state in states: + condition = {} if state is None else {"charge_spin": np.array([state])} + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + do_atomic_virial=True, + **{key: _tensor(value) for key, value in condition.items()}, + ) + actual = dp.eval(coord, None, atype, atomic=True, **condition) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + assert np.abs(value - target).max() <= tolerance * scale + 1e-5, ( + name, + state, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the complete model, train and eval. + + In training the gradient of an energy and force loss, whose force part is + a second-order derivative of the energy, reaches the long-range + parameters as in eager execution. + """ + model = _model(kernel) + coord, atype = _batch() + parameters = list(model.atomic_model.long_range.parameters()) + + def loss(out: dict) -> torch.Tensor: + return out["energy"].square().sum() + out["force"].square().sum() + + for training in (True, False): + model.train(training) + compiled = _CompiledModel( + model, _get_model_structure_key(model), compile_eval=True + ) + compiled.train(training) + reference = model(coord, atype) + result = compiled(coord, atype) + assert compiled._compiled_lower_by_mode, "the compiled graph was not used" + torch.testing.assert_close(result["energy"], reference["energy"]) + torch.testing.assert_close(result["force"], reference["force"]) + torch.testing.assert_close(result["virial"], reference["virial"]) + if training: + expected = torch.autograd.grad(loss(reference), parameters) + actual = torch.autograd.grad(loss(result), parameters) + for expected_grad, actual_grad in zip(expected, actual, strict=True): + torch.testing.assert_close(actual_grad, expected_grad) + + +def test_unsupported_combinations_fail() -> None: + """Unsupported inputs, routes, combinations and descriptors are rejected.""" + model = _model("les") + coord, atype = _batch() + box = 10.0 * torch.eye(3, dtype=torch.float64, device=env.DEVICE).reshape(1, 9) + with pytest.raises(NotImplementedError, match="non-periodic"): + model(coord[:1], atype[:1], box=box) + nlist = torch.full((1, 6, 4), -1, dtype=torch.int64, device=env.DEVICE) + with pytest.raises(NotImplementedError, match="NeighborGraph"): + model.atomic_model.forward_atomic(coord[:1], atype[:1], nlist) + config = _config("les") + config["spin"] = {"use_spin": [True, False], "scheme": "native"} + with pytest.raises(NotImplementedError): + get_model(config) + config = _config("les") + config["use_srtab"] = "table.txt" + config["sw_rmin"], config["sw_rmax"] = 0.5, 1.0 + with pytest.raises(NotImplementedError, match="ZBL"): + get_model(config) + config = _config("les") + config["type_map"] = ["O", "H"] + config["bridging_method"] = "ZBL" + with pytest.raises(NotImplementedError, match="bridging"): + get_model(config) + config = _config("les") + config["atom_exclude_types"] = [1] + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model(config) + child = _config("les") + child["type"] = "standard" + with pytest.raises(NotImplementedError, match="atom_exclude_types"): + get_model( + { + "type": "linear_ener", + "type_map": child["type_map"], + "weights": "sum", + "models": [child], + "atom_exclude_types": [1], + } + ) + config = _config("les") + config["descriptor"] = { + "type": "se_e2_a", + "rcut": 3.0, + "rcut_smth": 2.0, + "sel": [8, 8], + } + with pytest.raises(NotImplementedError, match="DPA4C"): + get_model(config) + + +def test_compressing_the_descriptor_alone_is_rejected() -> None: + """A folded descriptor and a constraint reading the condition are rejected.""" + config = _config("les", constraint={"source": "charge_spin"}) + config["descriptor"].update( + n_radial=8, + precision="float32", + add_chg_spin_ebd=True, + default_chg_spin=[0.0, 1.0], + ) + config["fitting_net"]["precision"] = "float32" + model = get_model(config).to(env.DEVICE).eval() + model.atomic_model.descriptor.enable_compression(0.5) + coord, atype = _batch() + with pytest.raises(ValueError, match="folds the charge/spin condition"): + model(coord[:1], atype[:1]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_dpmodel_reference_matches(kernel: str) -> None: + """The dpmodel reference evaluates the same energy from the portable record.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _model(kernel) + coord, atype = _batch() + expected = model(coord, atype) + reference = BaseModelDP.deserialize(model.serialize()) + ret = reference.call_common( + coord.cpu().numpy(), atype.cpu().numpy(), neighbor_graph_method="dense" + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + ret["energy"].reshape(3, 6), + expected["atom_energy"].detach().cpu().numpy().reshape(3, 6), + rtol=1e-12, + atol=1e-12, + ) + + +def test_pair_graph_skips_padding() -> None: + """The pair topology joins real atoms of one frame only and ends in masked guards.""" + coord = torch.arange(15, dtype=torch.float64).reshape(5, 3) + pairs = build_pair_graph( + coord, + torch.tensor([3, 2]), + torch.tensor([0, -1, 1, 0, 1]), + torch.tensor([True, False, True, True, True]), + ) + real = pairs.index[:, pairs.mask] + assert real.tolist() == [[0, 3], [2, 4]] + assert (~pairs.mask).sum() == 2 diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 2401aa59c5..aad817ad2d 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -472,6 +472,8 @@ def _make_mock_model(self): fitting_net = MagicMock() class FakeModel(DPModelCommon): + atomic_model = MagicMock(**{"has_long_range.return_value": False}) + def get_out_bias(self): return np.array([[0.0, 0.0]])