diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 720833b5cd..90ba87a14c 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -222,6 +222,67 @@ 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 topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of :meth:`required_topologies`. + + Generic capability (concrete default: none) in the sense of + :meth:`required_topologies`: a module whose topology depends on its + own state, such as the wave-number cutoff of a periodic long-range + module, which follows its kernel, gives it here, by topology name. + """ + return {} + + def build_topologies( + self, + coord: Any, + n_node: Any, + box: Any, + *, + atype: Any, + node_mask: Any = None, + pair_excl: Any = None, + ) -> dict[str, Any]: + """Build the topologies of :meth:`required_topologies` for one batch. + + See :func:`~deepmd.dpmodel.utils.neighbor_graph.build_topologies`; + the build parameters are those of :meth:`topology_params`. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + build_topologies, + ) + + return build_topologies( + self.required_topologies(), + coord, + n_node, + box, + node_mask, + atype=atype, + pair_excl=pair_excl, + topology_params=self.topology_params(), + ) + def supports_native_spin(self) -> bool: """Returns whether this atomic model consumes a per-atom spin input. @@ -494,6 +555,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. @@ -530,6 +593,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 ------- @@ -538,6 +608,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, @@ -546,6 +618,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) @@ -1014,9 +1088,18 @@ def model_forward( self.get_rcut(), pair_excl=self.pair_excl, ) + atype_flat = xp.reshape(atype, (-1,)) + topologies = self.build_topologies( + xp.reshape(coord, (-1, 3)), + graph.n_node, + box, + atype=atype_flat, + node_mask=atype_flat >= 0, + pair_excl=self.pair_excl, + ) atomic_ret = self.forward_common_atomic_graph( graph, - xp.reshape(atype, (-1,)), + atype_flat, fparam=fparam, aparam=( xp.reshape( @@ -1028,6 +1111,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..564b6fd56a 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,43 @@ 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.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + ) + + 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.pair_exclude_types and any( + name != ALL_PAIRS for name in self.long_range.required_topologies() + ): + raise NotImplementedError( + "`pair_exclude_types` is not supported with a long-range module " + "that serves periodic systems, whose reciprocal-space sum is " + "global and excludes no 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 +267,49 @@ 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 topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the long-range module.""" + if self.long_range is None: + return {} + return self.long_range.topology_params() + 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 +349,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 +382,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 +408,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 +450,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 +669,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 +704,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 +740,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 +804,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 +815,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 +910,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 +937,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 +948,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 +967,37 @@ 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 compute_long_range_input_stat( + self, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics of the long-range module alone. + + A fine-tuning run that adds a long-range module to a pretrained model + calls it for the new module, whose statistics, unlike those of the + inherited descriptor and fitting, the pretrained model cannot provide. + + Parameters + ---------- + sampled_func + The lazy sampled function to get data frames from different data + systems. + stat_file_path + The path to the stat file. + """ + self.long_range.compute_input_stats( + self._make_wrapped_sampler(sampled_func), + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def serialize(self) -> dict: dd = super().serialize() @@ -785,6 +1011,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 +1024,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..1356569b8e 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, +) from deepmd.dpmodel.utils.nlist import ( build_multiple_neighbor_list, get_multiple_nlist_key, @@ -34,6 +37,9 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .inner_potential import ( + InnerPotentialAtomicModel, +) from .pairtab_atomic_model import ( PairTabAtomicModel, ) @@ -112,6 +118,40 @@ 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" + ) + # The composition builds every topology once for all sub-models: the + # pairs of a periodic batch are guards only, which would silence a + # module serving non-periodic systems alone, and a reciprocal-space + # topology is built with the parameters of a single module. + declared = [m.required_topologies() for m in models if m.has_long_range()] + periodic = any(name != ALL_PAIRS for names in declared for name in names) + if len(declared) > 1 and periodic: + raise NotImplementedError( + "a sub-model whose long-range module serves periodic systems " + "takes no other sub-model carrying a long-range module beside it" + ) + if self.pair_exclude_types and periodic: + raise NotImplementedError( + "`pair_exclude_types` is not supported with a sub-model whose " + "long-range module serves periodic systems, whose " + "reciprocal-space sum is global and excludes no pair" + ) + self.models = models self.type_map = type_map self._rebuild_mapping_state() @@ -377,6 +417,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 +431,30 @@ 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 topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies of the sub-models. + + At most one sub-model declares a reciprocal-space topology, the only + kind with build parameters, so that the parameters of every topology + come from one sub-model. + """ + params: dict[str, dict[str, Any]] = {} + for model in self.models: + params.update(model.topology_params()) + return params + def forward_atomic_graph( self, graph: Any, @@ -398,14 +464,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 +499,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 +548,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..d00617fc1c --- /dev/null +++ b/deepmd/dpmodel/long_range/base_long_range.py @@ -0,0 +1,195 @@ +# 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.""" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the topologies, by name; none by default.""" + return {} + + @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..911dfba5dd --- /dev/null +++ b/deepmd/dpmodel/long_range/kernel.py @@ -0,0 +1,419 @@ +# 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 kspace(self, k_sq: Array) -> Array: + r"""Fourier transform of the kernel on squared wave numbers. + + :math:`\hat K(k) = \int K(r) e^{-i \mathbf{k} \cdot \mathbf{r}} d^3r`, + so that the energy of charges in a periodic cell of volume ``V`` reads + :math:`(k_e / 2V) \sum_{\mathbf{k} \neq 0} \hat K(k) |\hat\rho(\mathbf{k})|^2`. + + Parameters + ---------- + k_sq + Squared wave numbers with shape ``(nk,)`` or ``(rows, nk)``, in + Å⁻²; every entry must be positive. + + Returns + ------- + Array + The transformed kernel with the shape of ``k_sq``, in ų. + """ + raise NotImplementedError + + def contact_value(self, like: Array) -> Array: + r"""The kernel at contact, :math:`K(0)`, in Å⁻¹. + + The k-space lattice sum of a periodic system contains the self terms, + which the evaluation subtracts as :math:`\frac{1}{2} K(0) \sum_i q_i^2`. + + Parameters + ---------- + like + An array giving the namespace, dtype and device of the result. + """ + raise NotImplementedError + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + """Wave number at which the kernel weight falls below ``tolerance``. + + The Gaussian tail of the Fourier-transformed kernel makes the k-space + truncation exponentially convergent; the k grid of a periodic + evaluation ends at this wave number. + + Parameters + ---------- + tolerance + The relative weight of the slowest-decaying mode at the cutoff. + """ + 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), + ) + + @staticmethod + def _row_values(value: Array, k_sq: Array) -> Array: + """A scalar broadcast to shape ``(..., 1)`` over the leading axes of ``k_sq``. + + The gradient of a parameter of the transform sums over every wave + number. Broadcast over the rows, it sums along every row first, a + reduction a compiled graph spreads over the rows, and then over the + rows, rather than over every wave number in one block. + """ + xp = array_api_compat.array_namespace(k_sq) + return xp.broadcast_to(value, (*k_sq.shape[:-1], 1)) + + 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) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + # erf(alpha r)/r is the interaction of two Gaussian charge clouds of + # width 1/(sqrt(2) alpha); its transform is the damped Coulomb + # 4 pi exp(-k^2 / 4 alpha^2) / k^2. + alpha = self._row_values(xp.exp(self._parameter("log_alpha", k_sq)), k_sq) + return 4.0 * np.pi * xp.exp(-k_sq / (4.0 * alpha * alpha)) / k_sq + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + alpha = xp.exp(self._parameter("log_alpha", like)) + return 2.0 * alpha / np.sqrt(np.pi) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # exp(-k^2 / 4 alpha^2) < tolerance + alpha = float(np.exp(to_numpy_array(self.log_alpha).astype(np.float64))) + return 2.0 * alpha * np.sqrt(np.log(1.0 / tolerance)) + + +@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) + + def kspace(self, k_sq: Array) -> Array: + xp = array_api_compat.array_namespace(k_sq) + amplitude = self._parameter("amplitude", k_sq) + width = xp.exp(self._parameter("log_width", k_sq)) + # FT of a_k exp(-r^2 / 2 sigma_k^2) is a_k (2 pi)^{3/2} sigma_k^3 + # exp(-sigma_k^2 k^2 / 2); for the default amplitudes + # a_k = sqrt(2/pi) ln b / sigma_k the coefficient is 4 pi ln b sigma_k^2, + # whose ln-sigma trapezoid sum approaches 4 pi / k^2. The mode loop is + # unrolled: a (nk, M) intermediate breaks inductor's double backward of + # the k-space virial (NaN), while per-mode (nk,) terms compile fine. + terms = [] + for m in range(width.shape[0]): + weight = self._row_values(amplitude[m] * width[m] ** 3, k_sq) + rate = self._row_values(0.5 * width[m] ** 2, k_sq) + terms.append((2.0 * np.pi) ** 1.5 * weight * xp.exp(-rate * k_sq)) + return sum(terms[1:], terms[0]) + + def contact_value(self, like: Array) -> Array: + xp = array_api_compat.array_namespace(like) + amplitude = self._parameter("amplitude", like) + return xp.sum(amplitude) + + def kspace_cutoff(self, tolerance: float = 1e-8) -> float: + # the slowest-decaying mode is the NARROWEST Gaussian: + # exp(-sigma_min^2 k^2 / 2) < tolerance + sigma_min = float( + np.exp(np.min(to_numpy_array(self.log_width).astype(np.float64))) + ) + return np.sqrt(2.0 * np.log(1.0 / tolerance)) / sigma_min + + @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..4a2d16b2f4 --- /dev/null +++ b/deepmd/dpmodel/long_range/latent_charge.py @@ -0,0 +1,729 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import math +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, + K_GRID, + PME_MESH, + frame_id_from_n_node, + k_grid_bounds_of_cells, + 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, +) +from .reciprocal import ( + kspace_energy, + mesh_energy, +) + +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.""" + +PME_ORDER = 8 +"""Default order of the B-splines that spread the charges to the mesh.""" + +PME_SPACING_FRACTION = 0.75 +"""Default mesh spacing, as a fraction of the Nyquist spacing of the wave-number cutoff. + +At order 8 the mesh evaluation of a periodic frame then stays within a few +parts in a million of the sum over the wave numbers, in the reciprocal energy +and in the force. +""" + +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, periodic or not. + + 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. For a non-periodic batch 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; for a periodic batch (with ``periodic=True``) it is + the lattice sum of the kernel over the images of the cell, evaluated in + reciprocal space (see :mod:`~deepmd.dpmodel.long_range.reciprocal`), + + .. math:: + E_{\rm non-periodic} = 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 + ) + self.periodic = bool(periodic) + self.kspace_tolerance = float(kspace_tolerance) + if ( + not np.isfinite(self.kspace_tolerance) + or not 0.0 < self.kspace_tolerance < 1.0 + ): + raise ValueError( + f"`kspace_tolerance` must be in (0, 1), got {self.kspace_tolerance}" + ) + # The frame charge/spin condition a compressed model is specialized + # to; see :meth:`fold_charge_state`. + self.charge_state = None + if k_grid_mode is None: + k_grid_mode = "pme" if self.periodic else "ragged" + self.k_grid_mode = str(k_grid_mode) + if self.k_grid_mode not in ("ragged", "fixed", "pme"): + raise ValueError( + f"`k_grid_mode` must be 'ragged', 'fixed' or 'pme', got {self.k_grid_mode!r}" + ) + if self.k_grid_mode in ("fixed", "pme") and not self.periodic: + raise ValueError( + f"`k_grid_mode='{self.k_grid_mode}'` requires `periodic=True`" + ) + if k_grid_bounds is not None and self.k_grid_mode != "fixed": + raise ValueError( + "`k_grid_bounds` is only meaningful with `k_grid_mode='fixed'`" + ) + self.decouple_k_grid = bool(decouple_k_grid) + if self.decouple_k_grid and not self.periodic: + raise ValueError("`decouple_k_grid` requires `periodic=True`") + # A decoupled grid reads the kernel once, at construction: the k-point + # selection stays at the initial kernel however training moves it, + # while the kernel values in the graph still use the live parameters. + self.frozen_k_max = ( + float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + if self.decouple_k_grid + else None + ) + # The per-axis bounds of the shared integer grid of the fixed mode, + # zero until the data statistics size it (see + # :meth:`compute_input_stats`): an array of fixed shape, so that a + # backend keeping arrays in its state dict sees the same keys and + # shapes before and after the sizing, and a checkpoint restores the + # sized grid. + if k_grid_bounds is not None: + bounds = np.asarray(k_grid_bounds, dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0) or not np.any(bounds): + raise ValueError( + "`k_grid_bounds` must be three non-negative integers, not " + f"all zero, got {list(k_grid_bounds)}" + ) + self.k_grid_bounds = bounds + elif self.k_grid_mode == "fixed": + self.k_grid_bounds = np.zeros(3, dtype=np.int64) + else: + self.k_grid_bounds = None + # The mesh of a periodic frame, on which training in the "pme" mode + # and every frozen model evaluate the reciprocal sum. + if not self.periodic and (pme_spacing is not None or pme_order is not None): + raise ValueError("`pme_spacing` and `pme_order` require `periodic=True`") + self.pme_order = PME_ORDER if pme_order is None else int(pme_order) + if self.pme_order < 2: + raise ValueError(f"`pme_order` must be at least 2, got {self.pme_order}") + if pme_spacing is not None: + pme_spacing = float(pme_spacing) + if not np.isfinite(pme_spacing) or pme_spacing <= 0.0: + raise ValueError( + f"`pme_spacing` must be finite and positive, got {pme_spacing}" + ) + self.pme_spacing = pme_spacing + + @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"), + periodic=config.get("periodic", False), + kspace_tolerance=config.get("kspace_tolerance", 1e-8), + k_grid_mode=config.get("k_grid_mode"), + k_grid_bounds=config.get("k_grid_bounds"), + decouple_k_grid=config.get("decouple_k_grid", False), + pme_spacing=config.get("pme_spacing"), + pme_order=config.get("pme_order"), + ) + + def _k_max(self) -> float: + """The wave-number cutoff selecting the k points of a periodic batch. + + A coupled grid reads the current kernel parameters, so a trained + kernel re-derives its grid every batch; a decoupled grid keeps the + cutoff of the initial parameters. + """ + if self.decouple_k_grid: + return self.frozen_k_max + return float(self.kernel.kspace_cutoff(self.kspace_tolerance)) + + def mesh_spacing(self) -> float: + r"""The largest spacing of the mesh of a periodic frame, in Å. + + The configured ``pme_spacing``, or :data:`PME_SPACING_FRACTION` of the + Nyquist spacing :math:`\pi / k_{max}` of the wave-number cutoff. + """ + if self.pme_spacing is not None: + return self.pme_spacing + return PME_SPACING_FRACTION * math.pi / self._k_max() + + def use_mesh_evaluation(self) -> None: + """Evaluate the reciprocal sum of every periodic batch on the mesh. + + A frozen model takes the mesh of every frame as an input of dynamic + size, sized by the cell of the frame, so that it serves any cell at + the accuracy of :meth:`mesh_spacing`; the explicit wave numbers of + the other modes have a count that no input layout fixes. The spacing + is fixed at the current kernel parameters, a constant of the frozen + model. + """ + if self.periodic: + self.pme_spacing = self.mesh_spacing() + self.k_grid_mode = "pme" + + def topology_params(self) -> dict[str, dict[str, Any]]: + """The build parameters of the reciprocal-space topology of a periodic batch. + + The wave-number cutoff reads the kernel parameters at build time, so + that a trained kernel moves the k grid and the mesh with it, unless + ``decouple_k_grid`` keeps the initial cutoff. The k grid holds the + wave numbers of every frame within the cutoff, or the points of the + fixed integer grid with those beyond it masked; the mesh resolves + every cell of the batch at :meth:`mesh_spacing`. + + Raises + ------ + ValueError + If the fixed grid has not been sized. + """ + if not self.periodic: + return {} + if self.k_grid_mode == "pme": + return {PME_MESH: {"spacing": self.mesh_spacing()}} + params = {"k_max": self._k_max()} + if self.k_grid_mode == "fixed": + if not self._k_grid_sized(): + raise ValueError( + "the fixed k grid is not sized yet: run the data " + "statistics (`compute_input_stats`) on periodic " + "frames, or configure `k_grid_bounds`" + ) + params["bounds"] = to_numpy_array(self.k_grid_bounds) + return {K_GRID: params} + + def required_topologies(self) -> tuple[str, ...]: + if self.periodic: + if self.k_grid_mode == "pme": + return (ALL_PAIRS, PME_MESH) + return (ALL_PAIRS, K_GRID) + 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 + ) + energy = self._pair_branch(charge, pairs, n) + if self.periodic: + # The lattice sum of a periodic frame and its self term cancel to + # a small fraction of either, so the reciprocal-space energy is + # formed in float64 whatever the precision of the geometry. + charge_64 = xp.astype(charge, xp.float64) + if self.k_grid_mode == "pme": + reciprocal = mesh_energy( + charge_64, + topologies[PME_MESH], + frame_id, + self.kernel, + self.pme_order, + ) + else: + reciprocal = kspace_energy(charge_64, topologies[K_GRID], self.kernel) + energy = energy + xp.astype(reciprocal, energy.dtype) + return xp.reshape(energy, (n, 1)) + + def _pair_branch(self, charge: Array, pairs: Any, n: int) -> Array: + """Direct pair sum over the non-periodic frames, per atom.""" + xp = array_api_compat.array_namespace(charge) + 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) + ) + return segment_sum(half_energy, src, n) + segment_sum(half_energy, dst, n) + + 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: + if self.k_grid_mode == "fixed" and not self._k_grid_sized(): + # Materialize the samples once for the grid sizing and the + # networks; a lazy sampler is shared by the callers. + sampled = merged() if callable(merged) else merged + self._size_k_grid(sampled) + merged = sampled + self.charge_net.compute_input_stats( + merged, protection=protection, stat_file_path=stat_file_path + ) + + def _size_k_grid(self, sampled: list[dict]) -> None: + r"""Size the fixed integer k grid by the largest sampled cell. + + Every sampled frame needs :math:`|n_d| \le k_{max} |a_d| / 2\pi` per + axis at the wave-number cutoff of the kernel; the grid takes the + elementwise maximum of the bounds over the frames. The kernel + parameters are read at statistics time, so the grid is sized for the + kernel of that moment; frames that outgrow it later, by larger cells + or a kernel whose cutoff grows, are truncated to the grid. + """ + k_max = self._k_max() + bounds = None + for frame in sampled: + box = frame.get("box") + if box is None: + continue + box = np.asarray(to_numpy_array(box), dtype=np.float64) + if box.size == 0: + continue + frame_max = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + bounds = frame_max if bounds is None else np.maximum(bounds, frame_max) + if bounds is None: + raise ValueError( + "the fixed k grid takes its size from the data statistics, " + "but the sampled frames hold no simulation cell; provide " + "periodic training data or configure `k_grid_bounds`" + ) + self.k_grid_bounds = bounds.astype(np.int64) + + def _k_grid_sized(self) -> bool: + """Whether the fixed integer grid has its bounds.""" + return bool(np.any(to_numpy_array(self.k_grid_bounds))) + + 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.periodic: + data["periodic"] = True + data["kspace_tolerance"] = self.kspace_tolerance + data["k_grid_mode"] = self.k_grid_mode + if self.k_grid_mode == "fixed" and self._k_grid_sized(): + data["k_grid_bounds"] = to_numpy_array(self.k_grid_bounds).tolist() + if self.decouple_k_grid: + data["decouple_k_grid"] = True + data["frozen_k_max"] = self.frozen_k_max + if self.periodic: + data["pme_order"] = self.pme_order + if self.pme_spacing is not None: + data["pme_spacing"] = self.pme_spacing + 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"], + periodic=data.get("periodic", False), + kspace_tolerance=data.get("kspace_tolerance", 1e-8), + # A periodic record without a mode is one of the explicit wave + # numbers. + k_grid_mode=data.get("k_grid_mode", "ragged"), + k_grid_bounds=data.get("k_grid_bounds"), + decouple_k_grid=data.get("decouple_k_grid", False), + pme_spacing=data.get("pme_spacing"), + pme_order=data.get("pme_order"), + ) + if data.get("frozen_k_max") is not None: + # The constructor derives the frozen cutoff from the (possibly + # trained) kernel; the record holds the value the grid was + # actually built with. + obj.frozen_k_max = float(data["frozen_k_max"]) + 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/long_range/reciprocal.py b/deepmd/dpmodel/long_range/reciprocal.py new file mode 100644 index 0000000000..27b74e00e4 --- /dev/null +++ b/deepmd/dpmodel/long_range/reciprocal.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +r"""Reciprocal-space energy of latent charges in periodic cells. + +For the charges :math:`q_i^c` of a periodic cell of volume :math:`V` and a +pair kernel :math:`K` with Fourier transform :math:`\hat K`, the energy of the +lattice of images, + +.. math:: + E = k_e \Big[ \frac{1}{2V} \sum_{\mathbf{k} \neq \mathbf{0}} \sum_c + \hat K(\mathbf{k}) \, |\hat\rho_c(\mathbf{k})|^2 + - \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \Big], + \qquad + \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i \mathbf{k} \cdot \mathbf{r}_i}, + +is the sum over the wave numbers of the cell minus the interaction of every +charge with itself. The :math:`\mathbf{k} = \mathbf{0}` mode is dropped, the +tin-foil convention, which for a channel with a net charge amounts to a +uniform neutralizing background. The energy is attributed to the atoms as +:math:`E_i = \frac{1}{2} \sum_c q_i^c \phi_i^c`, with the potential +:math:`\phi_i^c = \partial E / \partial q_i^c` that the other charges and the +images create at atom :math:`i`, as the pair sum of a non-periodic system +splits every pair term evenly between its atoms. + +Two evaluations serve the topologies of +:mod:`~deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph`: +:func:`kspace_energy` sums over explicit wave numbers, and +:func:`mesh_energy` evaluates the same sum by the smooth particle-mesh Ewald +method on a mesh. +""" + +import math +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, + xp_add_at, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, +) + +from .kernel import ( + COULOMB_CONSTANT, + BaseInteractionKernel, +) + + +def _self_potential( + kernel: BaseInteractionKernel, charge: Array, periodic_node: Array +) -> Array: + """The potential of every charge at itself, k_e K(0) q, on periodic nodes.""" + xp = array_api_compat.array_namespace(charge) + return ( + COULOMB_CONSTANT + * kernel.contact_value(charge) + * charge + * xp.astype(periodic_node, charge.dtype)[:, None] + ) + + +def _strained_rows(xp: Any, rows: Array, strain: Array) -> Array: + r"""Wave numbers of the strained cells, to first order in the strain. + + The strain :math:`\varepsilon` of a frame maps its cell :math:`B` to + :math:`B S^\mathsf{T}`, :math:`S = I + \varepsilon`, and a wave number, + a row :math:`\mathbf{k}`, to :math:`\mathbf{k} S^{-1}`. The strain of a + topology is zero, and the energy is differentiated once with respect to + it, for the virial; the value and the first derivative at zero strain + are those of :math:`S^{-1} = I - \varepsilon`, which maps the rows here + and keeps linear algebra out of the energy and its derivatives. + + Parameters + ---------- + rows + Wave numbers or reciprocal lattice vectors with shape ``(nf, ..., 3)``. + strain + The strain of every frame with shape ``(nf, 3, 3)``. + + Returns + ------- + Array + The strained wave numbers, of the shape of ``rows``. + """ + eps = xp.reshape(strain, (strain.shape[0],) + (1,) * (rows.ndim - 2) + (3, 3)) + return rows - xp.sum(rows[..., :, None] * eps, axis=-2) + + +def _strained_volume(volume: Array, strain: Array) -> Array: + r"""Volumes of the strained cells, :math:`V \det S = V (1 + \mathrm{tr}\,\varepsilon)` to first order.""" + return volume * (1.0 + strain[:, 0, 0] + strain[:, 1, 1] + strain[:, 2, 2]) + + +def kspace_energy( + charge: Array, + k_grid: KGrid, + kernel: BaseInteractionKernel, +) -> Array: + r"""Energy of the periodic frames over explicit wave numbers, per atom. + + The structure factors of the wave numbers of every frame come from the + phases of its nodes, and the potential at a node from the same phases: + :math:`\phi_i = \frac{k_e}{V} \sum_{\mathbf{k}} \hat K(\mathbf{k}) + [\cos(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Re}\,\hat\rho + + \sin(\mathbf{k} \cdot \mathbf{r}_i) \mathrm{Im}\,\hat\rho] - k_e K(0) q_i`. + Every frame pairs the wave numbers of its row of :attr:`KGrid.vec` with + the nodes of its row of :attr:`KGrid.slot`, in products batched over the + frames and the blocks of their wave numbers. The geometry is strained by + :attr:`KGrid.strain` (see :class:`KGrid`). + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + k_grid + The wave numbers and the nodes of the periodic frames. + kernel + The pair kernel. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + without wave numbers. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + n_node, n_channel = charge.shape + nf, n_slot = k_grid.slot.shape + # === Step 1. The nodes of every frame, (nf, ns, 3) and (nf, ns, C) === + real = k_grid.slot >= 0 + index = xp.reshape(xp.where(real, k_grid.slot, xp.zeros_like(k_grid.slot)), (-1,)) + coord = xp.take(xp.astype(k_grid.coord, dtype), index, axis=0) + coord = xp.reshape(coord, (nf, n_slot, 3)) + charge_slot = xp.reshape(xp.take(charge, index, axis=0), (nf, n_slot, n_channel)) + charge_slot = charge_slot * xp.astype(real, dtype)[:, :, None] + # === Step 2. Strained wave numbers and volumes: k' = S^-T k, V' = V det S === + # Every product over the wave numbers of a frame broadcasts the per-frame + # factor over the blocks, (nf, nb, kb, .), so that its sum runs over the + # points of a block. + strain = xp.astype(k_grid.strain, dtype) + k_vec = xp.astype(k_grid.vec, dtype) + strained_k = _strained_rows(xp, k_vec, strain) + volume = _strained_volume(xp.astype(k_grid.volume, dtype), strain) + # === Step 3. Structure factors, (nf, nb, kb, C) === + # The phase k' . r' = k . r does not change under the strain. + phase = xp.matmul(k_vec, xp.matrix_transpose(coord)[:, None, :, :]) + cos_phase = xp.cos(phase) + sin_phase = xp.sin(phase) + rho_real = xp.matmul(cos_phase, charge_slot[:, None, :, :]) + rho_imag = xp.matmul(sin_phase, charge_slot[:, None, :, :]) + # === Step 4. Potential at the nodes of every frame, (nf, ns, C) === + weight = COULOMB_CONSTANT * kernel.kspace(xp.sum(strained_k * strained_k, axis=-1)) + weight = weight / volume[:, None, None] + weight = xp.where(k_grid.mask, weight, xp.zeros_like(weight))[..., None] + block_potential = xp.matmul(xp.matrix_transpose(cos_phase), weight * rho_real) + block_potential = block_potential + xp.matmul( + xp.matrix_transpose(sin_phase), weight * rho_imag + ) + potential = xp.sum(block_potential, axis=1) + # The self term belongs to the lattice sum of a frame with wave numbers; + # a non-periodic batch evaluated by a periodic module carries none. + has_points = xp.astype(xp.any(k_grid.mask, axis=(1, 2)), dtype)[:, None, None] + potential = ( + potential + - COULOMB_CONSTANT * kernel.contact_value(charge) * charge_slot * has_points + ) + slot_energy = 0.5 * xp.sum(charge_slot * potential, axis=-1) + # === Step 5. Energy on the node axis === + # A padding slot carries no charge and adds a zero to the first node. + energy = xp.zeros((n_node,), dtype=dtype, device=device) + return xp_add_at(energy, index, xp.reshape(slot_energy, (-1,))) + + +def bspline_integer_weights(order: int) -> np.ndarray: + r"""The cardinal B-spline of order ``order`` at :math:`1, 2, \dots, p - 1`. + + These are the weights of the assignment factor + :math:`D(\theta) = \sum_{l=0}^{p-2} M_p(l + 1) e^{i \theta l}` of the + smooth particle-mesh Ewald method. + """ + return np.asarray( + [ + sum( + (-1.0) ** j * math.comb(order, j) * max(l + 1.0 - j, 0.0) ** (order - 1) + for j in range(order + 1) + ) + / math.factorial(order - 1) + for l in range(order - 1) + ] + ) + + +def _bspline(xp: Any, x: Array, order: int) -> Array: + r"""The cardinal B-spline :math:`M_p(x)` of order :math:`p`, elementwise. + + :math:`M_p(x) = \frac{1}{(p-1)!} \sum_{j=0}^{p} (-1)^j \binom{p}{j} + \max(x - j, 0)^{p-1}`, a piecewise polynomial supported on + :math:`[0, p]` with :math:`p - 2` continuous derivatives. + """ + value = None + for j in range(order + 1): + term = ((-1.0) ** j * math.comb(order, j)) * xp.maximum( + x - j, xp.zeros_like(x) + ) ** (order - 1) + value = term if value is None else value + term + return value / math.factorial(order - 1) + + +def _inverse3(xp: Any, m: Array) -> Array: + """Inverse of a batch of 3x3 matrices, its columns the cross products of the rows over the determinant.""" + a, b, c = m[..., 0, :], m[..., 1, :], m[..., 2, :] + bc = xp.linalg.cross(b, c) + ca = xp.linalg.cross(c, a) + ab = xp.linalg.cross(a, b) + det = xp.sum(a * bc, axis=-1) + return xp.stack([bc, ca, ab], axis=-1) / det[..., None, None] + + +def _stencil( + xp: Any, frac: Array, dims: tuple[Any, Any, Any], order: int +) -> tuple[Array, Array]: + r"""The mesh points every node spreads to and their weights. + + Along axis :math:`d` the scaled coordinate :math:`u = n_d s_d` of the + fractional coordinate :math:`s_d \in [0, 1)` reaches the ``order`` points + :math:`\lfloor u \rfloor - p + 1, \dots, \lfloor u \rfloor` with the + weights :math:`M_p(u - m)`; the stencil of a node is the outer product of + its three axes. + + Returns + ------- + tuple[Array, Array] + The flat mesh indices and the weights, both with shape + ``(N, order**3)``. + """ + device = array_api_compat.device(frac) + n_node = frac.shape[0] + shift = xp.arange(order, dtype=frac.dtype, device=device)[None, :] + offset = xp.arange(order, dtype=xp.int64, device=device)[None, :] + index = None + weight = None + for d, n_d in enumerate(dims): + u = frac[:, d] * n_d + base = xp.floor(u) + # The distance of point base - p + 1 + j from u is t + p - 1 - j. + w = _bspline(xp, (u - base)[:, None] + (order - 1) - shift, order) + m = (xp.astype(base, xp.int64)[:, None] - (order - 1) + offset) % n_d + if index is None: + index, weight = m, w + else: + index = xp.reshape(index[:, :, None] * n_d + m[:, None, :], (n_node, -1)) + weight = xp.reshape(weight[:, :, None] * w[:, None, :], (n_node, -1)) + return index, weight + + +def _assignment_factor( + m_axes: list[Array], dims: tuple[Any, Any, Any], order: int +) -> list[Array]: + r"""The inverse squared assignment factor :math:`|D_d(m)|^{-2}` of every axis. + + The B-spline spreading multiplies every mode of the charge mesh by + :math:`\prod_d D_d(m_d)`, :math:`D_d(m) = \sum_l M_p(l + 1) + e^{2\pi i m l / n_d}`, which the influence function divides out. + """ + xp = array_api_compat.array_namespace(m_axes[0]) + weights = bspline_integer_weights(order) + factors = [] + for m, n_d in zip(m_axes, dims, strict=True): + theta = (2.0 * math.pi) * m / n_d + real = xp.zeros_like(theta) + imag = xp.zeros_like(theta) + for l, w in enumerate(weights): + real = real + float(w) * xp.cos(theta * l) + imag = imag + float(w) * xp.sin(theta * l) + factors.append(1.0 / (real * real + imag * imag)) + return factors + + +def mesh_energy( + charge: Array, + mesh: MeshGrid, + frame_id: Array, + kernel: BaseInteractionKernel, + order: int, +) -> Array: + r"""Energy of the periodic frames by the smooth particle-mesh Ewald method, per atom. + + The charges spread to the mesh of every frame with B-splines of order + ``order`` along the fractional coordinates; a real FFT gives the half + spectrum :math:`Q(\mathbf{m})` of the charge mesh. The influence function + + .. math:: + G(\mathbf{m}) = \frac{k_e}{V} \hat K(\mathbf{k}_\mathbf{m}) + \prod_d |D_d(m_d)|^{-2}, + \qquad \mathbf{k}_\mathbf{m} = 2\pi B^{-1} \mathbf{m}, + + corrects for the spreading, so that the potential on the mesh is the + inverse transform of :math:`G Q` and the potential at a node its + interpolation by the same B-splines. :math:`|\mathbf{k}|^2` is the + quadratic form :math:`(2\pi)^2 \mathbf{m}^\mathsf{T} W \mathbf{m}` of the + reciprocal metric :math:`W = A^\mathsf{T} A`, :math:`A = S^{-\mathsf{T}} + B^{-1}` of the strained cell, and the strain scales the volume by + :math:`\det S`; the fractional coordinates do not depend on it. The zero + mode carries no weight, nor do the modes at the Nyquist frequency of an + axis, which have no unique wave number. + + Parameters + ---------- + charge + Node charges with shape ``(N, C)``, zero on the nodes that carry none. + mesh + The cells and the mesh of the periodic frames. + frame_id + The frame of every node with shape ``(N,)``. + kernel + The pair kernel. + order + The order of the B-splines. + + Returns + ------- + Array + The energy of every node with shape ``(N,)``, in eV; zero in a frame + masked out by :attr:`MeshGrid.mask`. + """ + xp = array_api_compat.array_namespace(charge) + device = array_api_compat.device(charge) + dtype = charge.dtype + nf = mesh.box.shape[0] + n_node, n_channel = charge.shape + m_full = [xp.astype(m, dtype) for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)] + n1, n2, n3 = (m.shape[0] for m in m_full) + n_mesh = n1 * n2 * n3 + box = xp.astype(mesh.box, dtype) + volume = xp.astype(mesh.volume, dtype) + # === Step 1. Stencils of the nodes on the mesh === + box_inv = _inverse3(xp, box) + # frac_d = sum_j r_j (B^-1)_jd, elementwise rather than a batched product + # of 3x3 matrices. + frac = xp.sum( + xp.astype(mesh.coord, dtype)[:, :, None] * xp.take(box_inv, frame_id, axis=0), + axis=1, + ) + frac = frac - xp.floor(frac) + index, weight = _stencil(xp, frac, (n1, n2, n3), order) + # Flat targets on the frame- and channel-major meshes, (N, C, order**3). + channel = xp.arange(n_channel, dtype=xp.int64, device=device) + first = (frame_id[:, None] * n_channel + channel[None, :]) * n_mesh + target = first[:, :, None] + index[:, None, :] + # === Step 2. Charge mesh and its half spectrum === + q_mesh = xp.zeros((nf * n_channel * n_mesh,), dtype=dtype, device=device) + q_mesh = xp_add_at( + q_mesh, + xp.reshape(target, (-1,)), + xp.reshape(charge[:, :, None] * weight[:, None, :], (-1,)), + ) + q_mesh = xp.reshape(q_mesh, (nf * n_channel, n1, n2, n3)) + spectrum = xp.fft.rfftn(q_mesh, axes=(-3, -2, -1)) + # === Step 3. Influence function on the half spectrum, (nf, n1, n2, h3) === + m3 = m_full[2][: n3 // 2 + 1] + m_half = [m_full[0], m_full[1], m3] + strain = xp.astype(mesh.strain, dtype) + # The rows of A^T = B^-T S^-1 are the reciprocal vectors of the strained + # cell, strained like wave numbers; W = A^T A. + a_t = _strained_rows(xp, xp.matrix_transpose(box_inv), strain) + metric = xp.sum(a_t[:, :, None, :] * a_t[:, None, :, :], axis=-1) + volume = _strained_volume(volume, strain) + g1 = m_half[0][:, None, None] + g2 = m_half[1][None, :, None] + g3 = m_half[2][None, None, :] + # Six monomials of the mode indices times six entries of the metric of + # every frame, elementwise. + monomials = (g1 * g1, g2 * g2, g3 * g3, g1 * g2, g1 * g3, g2 * g3) + entries = ( + metric[:, 0, 0], + metric[:, 1, 1], + metric[:, 2, 2], + 2.0 * metric[:, 0, 1], + 2.0 * metric[:, 0, 2], + 2.0 * metric[:, 1, 2], + ) + k_sq = None + for entry, monomial in zip(entries, monomials, strict=True): + term = entry[:, None, None, None] * monomial[None, ...] + k_sq = term if k_sq is None else k_sq + term + k_sq = (2.0 * math.pi) ** 2 * k_sq + # The zero mode and the Nyquist planes carry no weight. + keep = ( + (2.0 * xp.abs(g1) != n1) & (2.0 * xp.abs(g2) != n2) & (2.0 * xp.abs(g3) != n3) + ) + keep = keep & ((g1 != 0.0) | (g2 != 0.0) | (g3 != 0.0)) + keep = keep[None, ...] & mesh.mask[:, None, None, None] + safe_k_sq = xp.where(keep, k_sq, xp.ones_like(k_sq)) + shape = safe_k_sq.shape + # One row per frame and first mode index, along which the gradients of + # the kernel parameters sum first. + rows = xp.reshape(safe_k_sq, (nf * n1, n2 * (n3 // 2 + 1))) + transform = xp.reshape(kernel.kspace(rows), shape) + f1, f2, f3 = _assignment_factor(m_half, (n1, n2, n3), order) + factor = f1[:, None, None] * f2[None, :, None] * f3[None, None, :] + influence = (COULOMB_CONSTANT * transform * factor[None, ...]) / volume[ + :, None, None, None + ] + influence = xp.where(keep, influence, xp.zeros_like(influence)) + # === Step 4. Potential on the mesh and at the nodes === + spectrum = xp.reshape(spectrum, (nf, n_channel, n1, n2, n3 // 2 + 1)) + potential_mesh = xp.fft.irfftn( + spectrum * influence[:, None, ...], + s=(n1, n2, n3), + axes=(-3, -2, -1), + norm="forward", + ) + stencil_potential = xp.take( + xp.reshape(potential_mesh, (-1,)), xp.reshape(target, (-1,)), axis=0 + ) + potential = xp.sum( + xp.reshape(stencil_potential, (n_node, n_channel, -1)) * weight[:, None, :], + axis=-1, + ) + periodic_node = xp.take(mesh.mask, frame_id, axis=0) + potential = potential - _self_potential(kernel, charge, periodic_node) + return 0.5 * xp.sum(charge * potential, axis=-1) 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 607f7dfb83..4788e27b29 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -48,6 +48,7 @@ NeighborGraph, build_neighbor_graph, build_neighbor_graph_ase, + cast_topology_geometry, compact_nodes, expand_node_values, ) @@ -564,12 +565,20 @@ 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 = self.atomic_model.build_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 +598,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 +741,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 +784,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 +809,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 +832,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 +874,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 +890,12 @@ def call_common_lower_graph( charge_spin=charge_spin, spin=spin, ) + if topologies: + # The topologies carry geometry like the edge vectors. + topologies = { + name: cast_topology_geometry(topology, edge_vec.dtype) + for name, topology in topologies.items() + } model_predict = self.forward_common_atomic_graph( atype, n_node, @@ -883,6 +908,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 ae7ed8bae3..c1bb8c3fe1 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -122,6 +122,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..f13d80cb46 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -5,8 +5,12 @@ ``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), +``reciprocal_graph`` (the reciprocal-space topologies of periodic frames), +``topology`` (building and laying out the topologies a model declares), +``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 +37,7 @@ ) from .derivatives import ( edge_force_virial, + pair_force_virial, ) from .env import ( edge_env_mat, @@ -53,19 +58,56 @@ pad_and_guard_angles, pad_and_guard_edges, ) +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + unordered_pair_index, +) from .pairs import ( center_edge_pairs, ) +from .reciprocal_graph import ( + K_BLOCK, + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + mesh_dims_of_cells, +) from .segment import ( segment_max, segment_mean, segment_softmax, segment_sum, ) +from .topology import ( + TOPOLOGY_TYPES, + build_topologies, + cast_topology_geometry, + flatten_topologies, + topology_flat_fields, + unflatten_topologies, +) __all__ = [ + "ALL_PAIRS", + "K_BLOCK", + "K_GRID", + "PME_MESH", + "TOPOLOGY_TYPES", "GraphLayout", + "KGrid", + "MeshGrid", "NeighborGraph", + "PairGraph", "angle_padding_fraction", "angle_to_edge_sum", "angle_to_node_sum", @@ -75,25 +117,41 @@ "attach_edge_csr", "build_angle_index", "build_edge_csr", + "build_fixed_k_grid", + "build_k_grid", + "build_mesh_grid", "build_neighbor_graph", "build_neighbor_graph_ase", + "build_pair_graph", + "build_topologies", "canonicalize_neighbor_graph", + "cast_topology_geometry", "center_edge_pairs", + "check_complete_frames", + "check_periodic_cells", "compact_nodes", "edge_env_mat", "edge_force_virial", + "enumerate_k_grid_bounds", "expand_node_values", + "flatten_topologies", "frame_id_from_n_node", "from_dense_quartet", "graph_angle_cos", "graph_from_dense_quartet", + "k_grid_bounds_of_cells", + "mesh_dims_of_cells", "neighbor_graph_from_ijs", "node_ownership_mask", "node_validity_mask", "pad_and_guard_angles", "pad_and_guard_edges", + "pair_force_virial", "segment_max", "segment_mean", "segment_softmax", "segment_sum", + "topology_flat_fields", + "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..e474abd4ed --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/pair_graph.py @@ -0,0 +1,224 @@ +# 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, + ClassVar, +) + +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 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.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec",) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("vec",) + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "index": {1: "entry"}, + "vec": {0: "entry"}, + "mask": {0: "entry"}, + } + """Dynamic axes of every field: the pair axis.""" + + +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) + + +# Guard pairs of the pair topology of a batch without non-periodic frames. +_GUARD_PAIRS = 4 + + +def empty_pair_graph(like: Array) -> PairGraph: + """The pair topology of a batch without non-periodic frames: guard pairs only. + + The guard pairs number :data:`_GUARD_PAIRS`, unlike the two rows of the + index and the three components of the vectors: a compiled graph that met + a pair axis of the length of one of those would identify the two, and + split at the reductions over the pairs. + """ + xp = array_api_compat.array_namespace(like) + device = array_api_compat.device(like) + index = xp.zeros((2, 0), dtype=xp.int64, device=device) + vec = xp.zeros((0, 3), dtype=like.dtype, device=device) + index, vec, mask = pad_and_guard_edges(index, vec, None, min_edges=_GUARD_PAIRS) + return PairGraph(index=index, vec=vec, mask=mask) + + +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; a module evaluates periodic " + "cells with `periodic: true` (pass box=None for a non-periodic " + "system)" + ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py new file mode 100644 index 0000000000..c337cf40f4 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reciprocal-space topologies of the periodic frames of a batch. + +A long-range module that serves periodic systems evaluates the lattice sum of +its pair kernel in reciprocal space, which needs the cell of every frame and +the absolute node positions rather than pair displacements. It declares one of +two topologies beside +:data:`~deepmd.dpmodel.utils.neighbor_graph.pair_graph.ALL_PAIRS`: + +- :data:`K_GRID` lists the wave numbers of every frame explicitly, ragged over + the frames, for a direct sum over them; +- :data:`PME_MESH` carries the cells and the mode indices of a real-space mesh + shared by the frames, for a smooth particle-mesh evaluation. + +Both carry the node positions and a zero strain as geometry: an autograd +backend makes them leaves, whose gradients are the force and the negative +virial of the reciprocal-space energy. A batch without cells comes back with +masked topologies, so a module serving periodic systems also serves +non-periodic batches through its pair topology. +""" + +from __future__ import ( + annotations, +) + +import math +import warnings +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + +K_GRID = "k_grid" +"""Name of the topology holding the reciprocal lattice points of periodic frames.""" + +PME_MESH = "pme_mesh" +"""Name of the topology carrying the cells and the mesh of periodic frames.""" + + +@dataclass +class KGrid: + r"""Reciprocal lattice points of the periodic frames of a batch, frame by frame. + + Row ``f`` of :attr:`vec` holds the wave numbers :math:`\mathbf{k} = 2\pi + B_f^{-1}\mathbf{n}` of frame ``f``, for integer :math:`\mathbf{n}`, built + eagerly outside any compiled region, in blocks of :data:`K_BLOCK` points + padded with masked guard points to the blocks of the longest row; row + ``f`` of :attr:`slot` holds the nodes of frame ``f``, padded with ``-1``. + A frame thus pairs its own wave numbers with its own nodes alone, at a + cost of the frame count times the padded wave-number count and the + padded node count, and every sum over the wave numbers of a frame runs + over the points of a block first. The block, point and slot axes are dynamic dimensions of + their own, of at least two entries. + + The energy is evaluated on the strained geometry of :attr:`strain`: + positions by ``S = I + strain``, wave numbers by :math:`S^{-T}` and volumes + by :math:`\det S`, so that the gradient with respect to the strain is the + negative virial; the phases :math:`\mathbf{k} \cdot \mathbf{r}` do not + change under the strain. The strain is zero, and the evaluation takes it + to first order, which fixes the energy and that gradient. The graph lower + makes :attr:`coord` and :attr:`strain` autograd leaves at its inputs. + """ + + vec: Array + """(nf, nb, kb, 3) float the wave numbers of every frame in blocks, in Å⁻¹.""" + mask: Array + """(nf, nb, kb) bool real (1) vs guard (0) points.""" + slot: Array + """(nf, ns) int64 the nodes of every frame, ``-1`` past its last node.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ("vec", "volume", "coord", "strain") + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "vec": {0: "frame", 1: "entry", 2: "k_point"}, + "mask": {0: "frame", 1: "entry", 2: "k_point"}, + "slot": {0: "frame", 1: "slot"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + } + """Dynamic axes of every field: the frame, block, point, slot and node axes.""" + + +@dataclass +class MeshGrid: + r"""Cells and mesh of the periodic frames of a batch. + + The frames share a real-space mesh of :math:`n_1 \times n_2 \times n_3` + points spanning each cell, whose integer mode indices along every axis + are the fields :attr:`mode_1`, :attr:`mode_2` and :attr:`mode_3`: their + lengths are the mesh dimensions, so a compiled or exported graph takes the + mesh as dynamic dimensions. Every dimension is even. A batch without a + cell comes back with identity cells, a minimal mesh and an all-false + :attr:`mask`. + + The geometry contract is that of :class:`KGrid`: the energy is evaluated + with the positions strained by ``S = I + strain`` and the cells by + :math:`B S^{\mathsf{T}}`, the lattice vectors being the rows of + :math:`B`, so that the fractional coordinates are invariant and the + gradient with respect to the strain is the negative virial. The volumes + of the unstrained cells are computed eagerly, and the evaluation scales + them by :math:`\det S`, to first order in the strain. + """ + + box: Array + """(nf, 3, 3) float simulation cells; rows are the lattice vectors, in Å.""" + volume: Array + """(nf,) float the cell volumes, in ų.""" + coord: Array + """(N, 3) float node positions, in Å.""" + strain: Array + """(nf, 3, 3) float zero strain of every frame.""" + mask: Array + """(nf,) bool periodic (1) vs non-periodic (0) frames.""" + mode_1: Array + """(n1,) float mode indices of the first mesh axis, in FFT order.""" + mode_2: Array + """(n2,) float mode indices of the second mesh axis, in FFT order.""" + mode_3: Array + """(n3,) float mode indices of the third mesh axis, in FFT order.""" + + GEOMETRY: ClassVar[tuple[str, ...]] = ( + "box", + "volume", + "coord", + "strain", + "mode_1", + "mode_2", + "mode_3", + ) + """Fields in the precision of the geometry.""" + LEAVES: ClassVar[tuple[str, ...]] = ("coord", "strain") + """Fields an autograd backend makes leaves of the energy.""" + AXES: ClassVar[dict[str, dict[int, str]]] = { + "box": {0: "frame"}, + "volume": {0: "frame"}, + "coord": {0: "node"}, + "strain": {0: "frame"}, + "mask": {0: "frame"}, + "mode_1": {0: "mesh_1"}, + "mode_2": {0: "mesh_2"}, + "mode_3": {0: "mesh_3"}, + } + """Dynamic axes of every field: the frame axis, the node axis, the mesh axes.""" + + +def check_periodic_cells(box: Array) -> None: + """Reject a cell that spans no lattice. + + A cell whose volume is a vanishing fraction of the product of the lengths + of its lattice vectors, or that is not finite, spans no lattice; an + all-zero cell is no marker of a non-periodic frame, which takes no cell. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``; rows are + the lattice vectors. + + Raises + ------ + ValueError + If a cell spans no lattice. + """ + cells = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + volume = np.abs(np.linalg.det(cells)) + spanned = volume > 1e-6 * np.prod(np.linalg.norm(cells, axis=-1), axis=-1) + if not np.all(spanned): + f = int(np.argmin(spanned)) + raise ValueError( + f"the cell of periodic frame {f} has no volume; a periodic frame " + "needs three independent lattice vectors, and a non-periodic " + "system takes no cell" + ) + + +def _cells(box: Array) -> Array: + """The cells with shape ``(nf, 3, 3)``, from a ``(nf, 9)`` or ``(nf, 3, 3)`` array.""" + xp = array_api_compat.array_namespace(box) + return xp.reshape(box, (box.shape[0], 3, 3)) + + +def _volume(box: Array) -> Array: + """The cell volumes |a_1 . (a_2 x a_3)| with shape ``(nf,)``.""" + xp = array_api_compat.array_namespace(box) + triple = box[:, 0, :] * xp.linalg.cross(box[:, 1, :], box[:, 2, :]) + return xp.abs(xp.sum(triple, axis=-1)) + + +def _zero_strain(box: Array) -> Array: + xp = array_api_compat.array_namespace(box) + return xp.zeros( + (box.shape[0], 3, 3), dtype=box.dtype, device=array_api_compat.device(box) + ) + + +def k_grid_bounds_of_cells(box: Array, k_max: float) -> np.ndarray: + r"""Per-axis integer wave-number bounds the cells need at ``k_max``. + + A wave number :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` with + :math:`|\mathbf{k}| \le k_{max}` has :math:`|n_d| \le k_{max} |a_d| / 2\pi` + along every lattice vector :math:`\mathbf{a}_d` (Cauchy-Schwarz), so the + bound of every frame is the ceiling of the right-hand side. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + + Returns + ------- + np.ndarray + The per-frame, per-axis integer bounds with shape ``(nf, 3)``. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + length = np.linalg.norm(cell, axis=-1) + return np.ceil(k_max * length / (2.0 * math.pi)).astype(np.int64) + + +def enumerate_k_grid_bounds(bounds: Array | list[int]) -> np.ndarray: + r"""Enumerate the integer wave-number indices within per-axis bounds. + + The grid holds every :math:`\mathbf{n} \neq \mathbf{0}` with + :math:`|n_d| \le n_d^{max}`, in lexicographic order with the first axis + outermost. The zero vector is left out: the transform of a Coulomb-like + kernel is singular there, and the tin-foil convention drops that mode. + + Parameters + ---------- + bounds + The per-axis bounds :math:`n_d^{max}` with shape ``(3,)``. + + Returns + ------- + np.ndarray + The integer indices with shape ``(nk, 3)``, int64. + + Raises + ------ + ValueError + If the bounds are not three non-negative integers. + """ + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(-1) + if bounds.shape != (3,) or np.any(bounds < 0): + raise ValueError( + f"k-grid bounds must be three non-negative integers, got {bounds}" + ) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=np.int64) for n in bounds], indexing="ij" + ) + n_grid = np.stack(grids, axis=-1).reshape(-1, 3) + return n_grid[np.any(n_grid != 0, axis=-1)] + + +K_BLOCK = 128 +"""Wave numbers per block of a row of a :class:`KGrid`.""" +# A guard point carries a nonzero wave number, whose kernel transform stays +# finite, and no weight. +_GUARD_K = (1e-6, 0.0, 0.0) +# The fewest entries of the block and slot axes, which keeps each a dynamic +# dimension of its own however few points or nodes a batch holds. +_MIN_ENTRIES = 2 + + +def _k_blocks(rows: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Lay the wave numbers of every frame out in blocks, padded with guard points. + + Returns the wave numbers with shape ``(nf, nb, K_BLOCK, 3)`` and their + mask with shape ``(nf, nb, K_BLOCK)``. + """ + longest = max([row.shape[0] for row in rows], default=0) + blocks = max(-(-longest // K_BLOCK), _MIN_ENTRIES) + vec = np.tile(np.asarray(_GUARD_K), (len(rows), blocks * K_BLOCK, 1)) + mask = np.zeros((len(rows), blocks * K_BLOCK), dtype=bool) + for f, row in enumerate(rows): + vec[f, : row.shape[0]] = row + mask[f, : row.shape[0]] = True + return ( + vec.reshape(len(rows), blocks, K_BLOCK, 3), + mask.reshape(len(rows), blocks, K_BLOCK), + ) + + +def _frame_slots(n_node: Array) -> np.ndarray: + """The nodes of every frame, one row each, padded with ``-1`` to the longest frame.""" + counts = np.asarray(to_numpy_array(n_node), dtype=np.int64).reshape(-1) + width = max(int(counts.max(initial=0)), _MIN_ENTRIES) + first = np.cumsum(counts) - counts + column = np.arange(width, dtype=np.int64)[None, :] + return np.where(column < counts[:, None], first[:, None] + column, -1) + + +def _k_grid( + box: Array, vec: np.ndarray, mask: np.ndarray, slot: np.ndarray, coord: Array +) -> KGrid: + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + return KGrid( + vec=xp.astype(xp.asarray(vec, device=device), box.dtype), + mask=xp.asarray(mask, device=device), + slot=xp.asarray(slot, device=device), + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + ) + + +def build_k_grid(box: Array, k_max: float, coord: Array, n_node: Array) -> KGrid: + r"""Enumerate the reciprocal lattice points of the periodic frames. + + With the lattice vectors as the rows of the cell :math:`B`, the wave + numbers of a frame are :math:`\mathbf{k} = 2\pi B^{-1}\mathbf{n}` over + integer :math:`\mathbf{n}` with :math:`0 < |\mathbf{k}| \le k_{max}`, so + that :math:`\mathbf{a}_d \cdot \mathbf{k} = 2\pi n_d`. Their count is data + dependent, so the enumeration runs eagerly. The points of every frame + follow the lexicographic order of their integer indices, and the rows of + the frames end in masked guard points up to the blocks of the longest + row. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The reciprocal lattice points of every frame, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + bounds = k_grid_bounds_of_cells(cells, k_max) + rows = [] + for f, cell in enumerate(cells): + n_grid = enumerate_k_grid_bounds(bounds[f]).astype(np.float64) + k = n_grid @ (2.0 * math.pi * np.linalg.inv(cell)).T + rows.append(k[np.sum(k * k, axis=-1) <= k_max * k_max]) + vec, mask = _k_blocks(rows) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def build_fixed_k_grid( + box: Array, bounds: Array, k_max: float, coord: Array, n_node: Array +) -> KGrid: + r"""Evaluate a fixed integer k grid on the cells of the periodic frames. + + Every frame shares the integer indices within ``bounds`` (see + :func:`enumerate_k_grid_bounds`), so the blocks of a row have the static + count of the grid. The wave numbers of a frame are :math:`\mathbf{k} = + 2\pi B^{-1}\mathbf{n}`, and the points beyond the cutoff + :math:`|\mathbf{k}| \le k_{max}` and the guard points that complete the + last block are masked out. + + A frame whose cell needs indices beyond the grid, a larger cell than the + grid was sized for or a kernel whose cutoff has grown since, is evaluated + on the grid alone: its sum is truncated to the points of the grid, which + is issued as a warning. Its forces and virial remain the derivatives of + that energy. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + bounds + The per-axis bounds :math:`n_d^{max}` of the shared integer grid with + shape ``(3,)``. + k_max + The wave-number cutoff, in Å⁻¹. + coord + Node positions with shape ``(N, 3)``. + n_node + Per-frame node counts with shape ``(nf,)``; the frames occupy + contiguous spans of the node axis. + + Returns + ------- + KGrid + The wave numbers of every frame over the shared grid, one row each. + """ + box = _cells(box) + cells = np.asarray(to_numpy_array(box), dtype=np.float64) + n_grid = enumerate_k_grid_bounds(bounds) + bounds = np.asarray(to_numpy_array(bounds), dtype=np.int64).reshape(3) + required = k_grid_bounds_of_cells(cells, k_max) + beyond = np.any(required > bounds[None, :], axis=1) + if np.any(beyond): + f = int(np.argmax(beyond)) + warnings.warn( + f"frame {f} needs integer k-grid bounds {required[f].tolist()}, " + f"beyond the fixed grid {bounds.tolist()}: its k-space sum is " + "truncated to the grid; enlarge `k_grid_bounds` for the full sum", + stacklevel=2, + ) + rec = 2.0 * math.pi * np.linalg.inv(cells) + k = np.einsum("jm,fdm->fjd", n_grid.astype(np.float64), rec) + vec, real = _k_blocks(list(k)) + mask = real & (np.sum(vec * vec, axis=-1) <= k_max * k_max) + return _k_grid(box, vec, mask, _frame_slots(n_node), coord) + + +def empty_k_grid(nf: int, coord: Array) -> KGrid: + """The k grid of a batch without periodic frames: guard points and no nodes.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + eye = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + vec, mask = _k_blocks([np.zeros((0, 3), dtype=np.float64)] * nf) + return _k_grid( + eye, vec, mask, np.full((nf, _MIN_ENTRIES), -1, dtype=np.int64), coord + ) + + +def _even_fft_size(n: int) -> int: + """The smallest even integer of at least ``n`` with no prime factor above 7.""" + m = max(2, n + n % 2) + while True: + rest = m + for p in (2, 3, 5, 7): + while rest % p == 0: + rest //= p + if rest == 1: + return m + m += 2 + + +def mesh_dims_of_cells(box: Array, spacing: float) -> tuple[int, int, int]: + r"""The mesh that resolves every cell at ``spacing``. + + Along each lattice vector :math:`\mathbf{a}_d` the mesh holds the + smallest even FFT-friendly number of points (no prime factor above 7) of + at least :math:`\lceil \max_f |\mathbf{a}_d| / h \rceil`, the maximum + running over the frames, so that no frame is evaluated with a spacing + coarser than :math:`h`. The C++ interface sizes the mesh of a frozen + model by the same rule. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing :math:`h`, in Å. + + Returns + ------- + tuple[int, int, int] + The mesh dimensions. + """ + cell = np.asarray(to_numpy_array(box), dtype=np.float64).reshape(-1, 3, 3) + # The length is formed term by term, as the C++ interface forms it, so + # that both round the ceiling identically. + length = np.sqrt( + cell[..., 0] * cell[..., 0] + + cell[..., 1] * cell[..., 1] + + cell[..., 2] * cell[..., 2] + ) + longest = np.max(length, axis=0) + n1, n2, n3 = (_even_fft_size(math.ceil(float(a) / spacing)) for a in longest) + return n1, n2, n3 + + +def _mesh_modes(xp: object, n: int, dtype: object, device: object) -> Array: + """The integer mode indices of an axis of ``n`` points, in FFT order.""" + m = xp.arange(n, dtype=dtype, device=device) + return xp.where(m > n // 2, m - n, m) + + +def mesh_grid_of_dims( + box: Array, coord: Array, dims: tuple[int, int, int], mask: Array +) -> MeshGrid: + """Pack cells, node positions and a mesh of the given dimensions. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)``; rows are the lattice + vectors. + coord + Node positions with shape ``(N, 3)``. + dims + The even mesh dimensions. + mask + Periodic (1) vs non-periodic (0) frames with shape ``(nf,)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames. + """ + xp = array_api_compat.array_namespace(box) + device = array_api_compat.device(box) + mode_1, mode_2, mode_3 = (_mesh_modes(xp, n, box.dtype, device) for n in dims) + return MeshGrid( + box=box, + volume=xp.astype(_volume(box), box.dtype), + coord=coord, + strain=_zero_strain(box), + mask=mask, + mode_1=mode_1, + mode_2=mode_2, + mode_3=mode_3, + ) + + +def build_mesh_grid(box: Array, spacing: float, coord: Array) -> MeshGrid: + """Pack the cells of the periodic frames and their shared mesh. + + Parameters + ---------- + box + Simulation cells with shape ``(nf, 3, 3)`` or ``(nf, 9)``; rows are + the lattice vectors. + spacing + The largest mesh spacing, in Å (see :func:`mesh_dims_of_cells`). + coord + Node positions with shape ``(N, 3)``. + + Returns + ------- + MeshGrid + The cells, the mesh and the geometry of the frames, every frame + periodic. + """ + xp = array_api_compat.array_namespace(box) + box = _cells(box) + mask = xp.ones((box.shape[0],), dtype=xp.bool, device=array_api_compat.device(box)) + return mesh_grid_of_dims(box, coord, mesh_dims_of_cells(box, spacing), mask) + + +def empty_mesh_grid(nf: int, coord: Array) -> MeshGrid: + """The mesh grid of a batch without periodic frames: every frame masked.""" + xp = array_api_compat.array_namespace(coord) + device = array_api_compat.device(coord) + box = xp.tile(xp.eye(3, dtype=coord.dtype, device=device)[None, ...], (nf, 1, 1)) + mask = xp.zeros((nf,), dtype=xp.bool, device=device) + return mesh_grid_of_dims(box, coord, (2, 2, 2), mask) diff --git a/deepmd/dpmodel/utils/neighbor_graph/topology.py b/deepmd/dpmodel/utils/neighbor_graph/topology.py new file mode 100644 index 0000000000..9a504fc7f2 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/topology.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Topologies an atomic model declares beside its neighbor graph. + +An interaction that the short-range neighbor graph cannot describe reads a +topology of its own, declared by name (see +:meth:`~deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel.required_topologies`): + +- :data:`~.pair_graph.ALL_PAIRS`, the pairs of the non-periodic frames; +- :data:`~.reciprocal_graph.K_GRID` and :data:`~.reciprocal_graph.PME_MESH`, + the reciprocal space of the periodic frames. + +Every topology is a dataclass of arrays that declares which of its fields +carry geometry (``GEOMETRY``), which an autograd backend makes leaves of the +energy (``LEAVES``), and the dynamic axes of every field (``AXES``), so that +casting, re-leafing and the inputs of a compiled or exported graph lower are +derived from the declaration alone. +""" + +from __future__ import ( + annotations, +) + +import dataclasses +from typing import ( + TYPE_CHECKING, + Any, +) + +import array_api_compat + +from .pair_graph import ( + ALL_PAIRS, + PairGraph, + build_pair_graph, + check_complete_frames, + empty_pair_graph, +) +from .reciprocal_graph import ( + K_GRID, + PME_MESH, + KGrid, + MeshGrid, + build_fixed_k_grid, + build_k_grid, + build_mesh_grid, + check_periodic_cells, + empty_k_grid, + empty_mesh_grid, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterable, + Sequence, + ) + + from deepmd.dpmodel.array_api import ( + Array, + ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + Topology = PairGraph | KGrid | MeshGrid + +TOPOLOGY_TYPES: dict[str, type] = { + ALL_PAIRS: PairGraph, + K_GRID: KGrid, + PME_MESH: MeshGrid, +} +"""The dataclass of every topology name.""" + + +def _build_reciprocal( + name: str, + box: Array, + coord: Array, + n_node: Array, + params: dict[str, Any], +) -> KGrid | MeshGrid: + """Build a reciprocal-space topology of periodic frames from its parameters.""" + required = {K_GRID: "k_max", PME_MESH: "spacing"}[name] + if required not in params: + raise ValueError( + f"the {name!r} topology requires its `{required}` build parameter" + ) + check_periodic_cells(box) + if name == PME_MESH: + return build_mesh_grid(box, params["spacing"], coord) + if "bounds" in params: + return build_fixed_k_grid(box, params["bounds"], params["k_max"], coord, n_node) + return build_k_grid(box, params["k_max"], coord, n_node) + + +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, + topology_params: dict[str, dict[str, Any]] | None = None, +) -> dict[str, Topology]: + """Build the topologies an atomic model declares, beside its neighbor graph. + + A batch is periodic when it carries cells. A declaration of the pairs + alone serves non-periodic batches only. Beside a reciprocal-space + topology, the pairs of a periodic batch come back guard-only and the + reciprocal-space topology of a non-periodic batch masked, so that one + model serves both. + + Parameters + ---------- + names + The declared topology names. + coord + Node coordinates with shape ``(N, 3)``, frame-major over ``n_node``. + n_node + Per-frame node counts with shape ``(nf,)``. + box + Simulation cells with shape ``(nf, 9)`` or ``(nf, 3, 3)``, 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``. + topology_params + Build parameters by name: ``{"k_max": ...}`` for the k grid, with + ``"bounds"`` for a fixed integer grid (see + :func:`~.reciprocal_graph.build_fixed_k_grid`), and + ``{"spacing": ...}`` for the mesh. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + Each declared topology under its name; empty when none is declared. + + Raises + ------ + NotImplementedError + If cells are given for a declaration of the pairs alone, or a pair + exclusion for a periodic batch, whose reciprocal-space sum is global. + ValueError + If a name is unknown, or a build parameter is missing. + """ + names = tuple(names) + unknown = [name for name in names if name not in TOPOLOGY_TYPES] + if unknown: + raise ValueError( + f"unknown topology {unknown[0]!r}; expected one of {tuple(TOPOLOGY_TYPES)}" + ) + periodic = any(name in (K_GRID, PME_MESH) for name in names) + if box is None: + nf = n_node.shape[0] + empty = {K_GRID: empty_k_grid, PME_MESH: empty_mesh_grid} + return { + name: ( + build_pair_graph(coord, n_node, atype, node_mask, pair_excl) + if name == ALL_PAIRS + else empty[name](nf, coord) + ) + for name in names + } + if not periodic: + if names: + check_complete_frames(box) + return {} + if pair_excl is not None: + raise NotImplementedError( + "the reciprocal-space sum of a periodic system is global; " + "`pair_exclude_types` cannot exclude pairs from it" + ) + params = topology_params or {} + return { + name: ( + empty_pair_graph(coord) + if name == ALL_PAIRS + else _build_reciprocal(name, box, coord, n_node, params.get(name, {})) + ) + for name in names + } + + +def topology_flat_fields(topology: Topology) -> tuple[Array, ...]: + """The fields of a topology in their declaration order.""" + return tuple( + getattr(topology, field.name) for field in dataclasses.fields(topology) + ) + + +def flatten_topologies(topologies: dict[str, Topology]) -> tuple[Array, ...]: + """Lay topologies out as positional inputs of a graph lower. + + A compiled or exported graph lower receives the fields of every topology, + in the order of the dict and of the declaration of each dataclass, after + its other inputs; :func:`unflatten_topologies` restores them. + + Parameters + ---------- + topologies + The topologies by name. + + Returns + ------- + tuple + The fields of every topology, concatenated. + """ + return tuple( + value + for topology in topologies.values() + for value in topology_flat_fields(topology) + ) + + +def unflatten_topologies( + names: Sequence[str], values: Sequence[Array] +) -> dict[str, Topology]: + """Restore topologies laid out by :func:`flatten_topologies`. + + Parameters + ---------- + names + The topology names, in the order they were laid out. + values + The fields of every topology, concatenated. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + The topologies by name. + + Raises + ------ + ValueError + If a name is unknown, or the number of values does not match the + layout of the names. + """ + topologies = {} + offset = 0 + for name in names: + if name not in TOPOLOGY_TYPES: + raise ValueError(f"unknown topology {name!r}") + cls = TOPOLOGY_TYPES[name] + arity = len(dataclasses.fields(cls)) + if offset + arity > len(values): + break + topologies[name] = cls(*values[offset : offset + arity]) + offset += arity + if len(topologies) != len(names) or offset != len(values): + raise ValueError( + f"{len(values)} flattened values given, which do not lay out the " + f"topologies {tuple(names)}" + ) + return topologies + + +def cast_topology_geometry(topology: Topology, dtype: Any) -> Topology: + """A copy of a topology with its geometry fields cast to ``dtype``.""" + xp = array_api_compat.array_namespace(getattr(topology, topology.GEOMETRY[0])) + return dataclasses.replace( + topology, + **{ + name: xp.astype(getattr(topology, name), dtype) + for name in topology.GEOMETRY + }, + ) 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..c4f0771b01 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,14 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + check_complete_frames, + check_periodic_cells, + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -664,6 +672,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 +727,9 @@ 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) + self.metadata["topology_params"] = model.atomic_model.topology_params() if use_graph_lower: from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, @@ -735,7 +747,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 +818,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 +836,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 +2402,11 @@ 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 + fields of every topology the ``topologies`` metadata names, built with + its ``topology_params``. Exported AOTI artifacts use dynamic edge, + pair and mesh 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 +2520,31 @@ def _eval_model_graph( compact.source_row_ptr, compact.source_order, ) + topology_names = self.metadata.get("topologies", ()) + if topology_names: + # The compact lower of a model with a long-range module reads + # the positions and the cells of the frames, and the zero + # communicator of frames held whole; a module that serves no + # periodic system takes no cell, and the zero cell that marks + # a non-periodic frame is never one a caller passes. + if PME_MESH not in topology_names: + check_complete_frames(box_input) + elif box_input is not None: + check_periodic_cells(box_input) + box_t = ( + torch.zeros(nframes, 9, dtype=torch.float64, device="cpu") + if box_input is None + else torch.as_tensor( + np.asarray(box_input, dtype=np.float64).reshape(nframes, 9) + ) + ) + model_inputs += ( + torch.as_tensor( + coord_input.reshape(-1, 3), dtype=torch.float64, device=DEVICE + ), + torch.zeros(1, dtype=torch.int64, device="cpu"), + box_t, + ) else: fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE @@ -2528,6 +2570,37 @@ 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. The cell is torch like the coordinates: the + # topologies it namespaces feed the torch graph lower + # directly (a numpy cell would produce numpy leaves). + coord_flat_t = torch.as_tensor( + coord_input.reshape(-1, 3), dtype=edge_dtype, device=DEVICE + ) + box_t = ( + torch.as_tensor( + np.asarray(box_input).reshape(n_node_t.shape[0], 3, 3), + dtype=edge_dtype, + device=DEVICE, + ) + if box_input is not None + else None + ) + # The parameters are those of the archive, or those of the + # live model of a checkpoint. + topologies = build_topologies( + topology_names, + coord_flat_t, + n_node_t, + box_t, + node_mask=atype_t >= 0, + atype=atype_t, + pair_excl=self._model_pair_excl(), + topology_params=self.metadata.get("topology_params"), + ) + model_inputs += flatten_topologies(topologies) 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..48eef91e74 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 and the cells, which its operator reads in place of a pair + # list or a mesh; 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.operator_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..3bb786ff58 --- /dev/null +++ b/deepmd/pt_expt/kernels/dpa4c/canonical_long_range.py @@ -0,0 +1,160 @@ +# 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, + box: 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 ``operator_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 Å. + 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. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the CPU, + a zero row marking a complete non-periodic frame, or ``None`` when no + frame is periodic. + 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, box + ) + 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; the halo nodes of a periodic frame are images, which the + # reciprocal-space sum holds itself. + 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..89db40ad93 --- /dev/null +++ b/deepmd/pt_expt/kernels/long_range.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Bindings for the operator of the latent-charge long range. + +The operator ``deepmd::latent_charge_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. A non-periodic +frame is summed directly over its pairs without listing them, so that the +memory is linear in the node count; a periodic frame is summed over the images +of its cell by the smooth particle-mesh Ewald method, on a mesh sized by the +cell. The CUDA kernels are in ``source/op/pt/long_range``; the CPU kernel +registered here is the reference they are 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", + "MAX_ORDER", + "kernel_eligible", + "kernel_operator_arguments", + "latent_charge_sum", + "op_available", +] + +#: Charge channels and Gaussians the CUDA kernels are instantiated for. +MAX_CHANNELS = 8 +MAX_GAUSSIANS = 32 +#: The largest order of the B-splines of the mesh. +MAX_ORDER = 12 + +_LES = 0 +_SOG = 1 + + +def op_available() -> bool: + """Whether the backend device carries the 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_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 float64 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 long-range operator does not evaluate {type(kernel).__name__}" + ) + if isinstance(kernel, LESKernel): + coefficient = torch.exp(kernel.log_alpha.to(torch.float64)).reshape(1) + return _LES, coefficient.detach().contiguous() + rate = (0.5 / math.log(2.0)) * torch.exp(-2.0 * kernel.log_width.to(torch.float64)) + coefficient = torch.cat((kernel.amplitude.to(torch.float64), rate)) + return _SOG, coefficient.detach().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 _charges_of_frames( + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + channel: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The charges of the interactions in float64, shifted in the constrained channel. + + Returns the charges with shape ``(N, C)``, the weights of the charged + nodes and the frame of every node with shape ``(N,)``, and the number of + charged nodes of every frame, at least one, with shape ``(F,)``. + """ + n = charge.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=charge.device) < torch.sum( + n_node_per_frame + ) + weight = (node_mask & real).to(torch.float64) + q = charge.to(torch.float64) * weight[:, None] + count = weight.new_zeros(frame_count).index_add(0, frame, weight) + count = torch.clamp(count, min=1.0) + if channel >= 0: + total = weight.new_zeros(frame_count).index_add(0, frame, q[:, channel]) + shift = (total - total_charge.to(torch.float64)) / count + q = q.clone() + q[:, channel] = q[:, channel] - shift[frame] * weight + return q, weight, frame, count + + +def _direct_reference( + coord: torch.Tensor, + q: torch.Tensor, + n_node_per_frame: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial over the explicit pairs of every frame.""" + n = coord.shape[0] + 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 = prefactor * 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) + return potential, force, atom_virial + + +def _interaction_kernel(coefficient: torch.Tensor, kernel: int) -> Any: + """The dpmodel interaction kernel of the coefficients of the operator.""" + from deepmd.dpmodel.long_range.kernel import LESKernel as LESKernelDP + from deepmd.dpmodel.long_range.kernel import SOGKernel as SOGKernelDP + + values = coefficient.detach().to("cpu", torch.float64).numpy() + if kernel == _LES: + return LESKernelDP(alpha=float(values[0]), precision="float64") + amplitude, rate = values.reshape(2, -1) + width = (0.5 / (math.log(2.0) * rate)) ** 0.5 + return SOGKernelDP( + amplitude=amplitude.tolist(), width=width.tolist(), precision="float64" + ) + + +def _mesh_reference( + coord: torch.Tensor, + q: torch.Tensor, + weight: torch.Tensor, + n_node_per_frame: torch.Tensor, + box: torch.Tensor, + coefficient: torch.Tensor, + kernel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Potential, force and atomic virial on the mesh of every periodic frame. + + Every frame is evaluated on its own mesh by the evaluation of the training + path, and differentiated with respect to its charges, its positions and + its strain. The atomic virial takes ``F (x) r`` at every node and shares + the rest of the virial of the frame evenly over its charged nodes. + """ + from deepmd.dpmodel.long_range.kernel import ( + COULOMB_CONSTANT, + ) + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + interaction = _interaction_kernel(coefficient, kernel) + potential = torch.zeros_like(q) + force = torch.zeros_like(coord) + atom_virial = coord.new_zeros(coord.shape[0], 3, 3) + begin = 0 + for f, length in enumerate(n_node_per_frame.tolist()): + end = begin + length + frame_q = q[begin:end].detach().requires_grad_(True) + frame_coord = coord[begin:end].detach().requires_grad_(True) + with torch.enable_grad(): + mesh = build_mesh_grid(box[f].reshape(1, 3, 3), spacing, frame_coord) + mesh.strain.requires_grad_(True) + frame_id = torch.zeros(length, dtype=torch.int64, device=coord.device) + energy = mesh_energy(frame_q, mesh, frame_id, interaction, order).sum() + energy = energy * (prefactor / COULOMB_CONSTANT) + d_q, d_coord, d_strain = torch.autograd.grad( + energy, (frame_q, frame_coord, mesh.strain) + ) + frame_force = -d_coord + moment = frame_force[:, :, None] * coord[begin:end, None, :] + frame_weight = weight[begin:end] + count = torch.clamp(torch.sum(frame_weight), min=1.0) + rest = (-d_strain[0] - torch.sum(moment, dim=0)) / count + potential[begin:end] = d_q + force[begin:end] = frame_force + atom_virial[begin:end] = moment + rest * frame_weight[:, None, None] + begin = end + return potential, force, atom_virial.reshape(-1, 9) + + +def _reference( + coord: torch.Tensor, + charge: torch.Tensor, + node_mask: torch.Tensor, + n_node_per_frame: torch.Tensor, + total_charge: torch.Tensor, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate the operator in float64, over the explicit pairs or the meshes.""" + if int(communicator.reshape(-1)[0]) != 0: + raise NotImplementedError( + "the CPU reference of latent_charge_sum evaluates a single domain; " + "a communicator of several ranks requires the CUDA kernel" + ) + frame_count = n_node_per_frame.shape[0] + q, weight, frame, count = _charges_of_frames( + charge, node_mask, n_node_per_frame, total_charge, channel + ) + has_cell = ( + torch.zeros(frame_count, dtype=torch.bool, device="cpu") + if box is None + else torch.any(box != 0.0, dim=-1) + ) + if bool(torch.any(has_cell)) != bool(torch.all(has_cell)): + raise ValueError( + "latent_charge_sum: every frame of a call is periodic, or none is" + ) + if bool(torch.any(has_cell)): + potential, force, atom_virial = _mesh_reference( + coord, + q, + weight, + n_node_per_frame, + box.to(coord.device), + coefficient, + kernel, + prefactor, + order, + spacing, + ) + else: + potential, force, atom_virial = _direct_reference( + coord, q, n_node_per_frame, coefficient, kernel, prefactor + ) + 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.clone() + gradient[:, channel] = gradient[:, channel] - (mean / count)[frame] * weight + return ( + atom_energy, + gradient.to(charge.dtype), + 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, + box: torch.Tensor | None, + coefficient: torch.Tensor, + communicator: torch.Tensor, + kernel: int, + channel: int, + prefactor: float, + order: int, + spacing: 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_sum")(_fake) + library = torch.library.Library("deepmd", "IMPL") + library.impl("latent_charge_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_sum", None), + torch._ops.OpOverloadPacket, + ): + _register() + + +def latent_charge_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, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> 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. + box + The cells of periodic frames with shape ``(F, 9)``, float64 on the + CPU, the rows of every row-major 3x3 cell being its lattice vectors, + or ``None`` for non-periodic frames; a zero row marks a non-periodic + frame. A periodic frame is evaluated on a mesh of spacing at most + ``spacing``, with B-splines of order ``order``. + order + The order of the B-splines of the mesh, at most :data:`MAX_ORDER`. + spacing + The largest mesh spacing, in Å. + + Returns + ------- + atom_energy + Per-node energy with shape ``(N, 1)`` in eV, float64, half of the + interaction of the charges of the node with all others. + 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: half of + every pair term to each of its nodes in a non-periodic frame; in a + periodic frame ``F (x) r`` with the rest of the virial of the cell + shared evenly over the charged 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_sum( + coord, + charge, + node_mask, + n_node_per_frame, + total_charge, + box, + coefficient.to(coord.device), + communicator, + code, + int(channel), + float(prefactor), + int(order), + float(spacing), + ) 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..b79dffef4c --- /dev/null +++ b/deepmd/pt_expt/long_range/latent_charge.py @@ -0,0 +1,177 @@ +# 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, periodic or not (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 operator_eligible(self) -> bool: + """Whether :meth:`energy_force_from_coord` evaluates this module. + + The operator evaluates the kernels, channel counts and B-spline orders + 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, + MAX_ORDER, + kernel_eligible, + op_available, + ) + + return ( + op_available() + and kernel_eligible(self.kernel) + and self.charge_net.dim_out <= MAX_CHANNELS + and self.pme_order <= MAX_ORDER + 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, + box: 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 operator sums a non-periodic frame directly over its pairs, which + it enumerates itself, so that the memory is linear in the node count, + and a periodic frame on the mesh of its cell. 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_sum`; + ``None`` for a single domain. + box + The cells of the frames with shape ``(nf, 9)``, float64 on the + CPU, a zero row marking a non-periodic frame, or ``None`` when no + frame is periodic. A module that does not serve periodic systems + evaluates every frame as non-periodic. + + 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_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_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, + # A module that does not serve periodic systems sums every frame + # directly. + box if self.periodic else None, + self.pme_order, + self.mesh_spacing() if self.periodic else 1.0, + ) + (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..f149053489 100644 --- a/deepmd/pt_expt/model/edge_transform_output.py +++ b/deepmd/pt_expt/model/edge_transform_output.py @@ -1,10 +1,17 @@ # 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 geometry leaves of the topologies. """ +import dataclasses +from collections.abc import ( + Sequence, +) + import torch from deepmd.dpmodel import ( @@ -19,9 +26,14 @@ get_deriv_name_mag, ) from deepmd.dpmodel.utils.neighbor_graph import ( + KGrid, + MeshGrid, NeighborGraph, + PairGraph, edge_force_virial, frame_id_from_n_node, + node_validity_mask, + pair_force_virial, segment_sum, ) from deepmd.pt.utils import ( @@ -44,6 +56,91 @@ ) +def releaf_topologies( + topologies: dict[str, PairGraph | KGrid | MeshGrid], +) -> dict[str, PairGraph | KGrid | MeshGrid]: + """Make the geometry fields of every topology autograd leaves. + + The leaves are made from the inputs of a traced graph lower, so that the + backward roots at its inputs: a leaf created inside a ``make_fx`` trace + has a NaN gradient under inductor. + """ + return { + name: dataclasses.replace( + topology, + **{ + leaf: getattr(topology, leaf).detach().requires_grad_(True) + for leaf in topology.LEAVES + }, + ) + for name, topology in topologies.items() + } + + +def _cell_force_virial( + g_coord: torch.Tensor, + g_strain: torch.Tensor, + coord: torch.Tensor, + n_node: torch.Tensor, + n_local: torch.Tensor | None, + do_atomic_virial: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """Force and virial of a reciprocal-space energy from its geometry leaves. + + The gradient with respect to the node positions is the negative force, + and that with respect to the zero strain of every frame the negative + virial. The atomic virial takes ``F_i (x) r_i`` at every node, the force + index first as for the edge virial, and shares the rest of the virial of + the frame, the part of the cell, evenly over its owned nodes, so that the + nodes of a frame sum to its virial; the halo nodes, images that carry no + charge, take no share. + + Parameters + ---------- + g_coord + (N, 3) gradient of the energy with respect to the positions. + g_strain + (nf, 3, 3) gradient of the energy with respect to the strain. + coord + (N, 3) node positions. + n_node + (nf,) per-frame node counts. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned. + do_atomic_virial + Whether to return the per-node virial. + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, or ``None``. + virial + (nf, 3, 3) per-frame virial. + """ + force = -g_coord + virial = -g_strain + if not do_atomic_virial: + return force, None, virial + n_total = force.shape[0] + frame_id = frame_id_from_n_node(n_node, n_total) + atom_virial = force[:, :, None] * coord.to(force.dtype)[:, None, :] + remainder = virial - segment_sum(atom_virial, frame_id, n_node.shape[0]) + if n_local is None: + owned = node_validity_mask(n_node, n_total) + count = n_node + else: + owned = node_ownership_mask(n_node, n_local, n_total) + count = n_local + share = owned.to(force.dtype) / torch.clamp( + count.to(force.dtype)[frame_id], min=1.0 + ) + atom_virial = atom_virial + remainder[frame_id] * share[:, None, None] + return force, atom_virial, virial + + def edge_energy_deriv( energy: torch.Tensor, edge_vec: torch.Tensor, @@ -60,6 +157,8 @@ def edge_energy_deriv( do_atomic_virial: bool = False, create_graph: bool = False, force_precision: torch.dtype | None = None, + topologies: Sequence[PairGraph | KGrid | MeshGrid] = (), + n_local: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """Return (force, atom_virial_or_None, virial) from a graph energy. @@ -79,6 +178,15 @@ 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 geometry + leaves of every topology in ``topologies``, and their pullbacks are summed + in: the energy reaches the coordinates through all its geometry leaves, so + its complete force and virial are the sum of their pullbacks. The pair + vectors of a pair topology are pulled back by the shared + pair_force_virial scatter; a reciprocal-space topology carries the node + positions, whose gradient is its force, and a zero strain, whose gradient + is its negative virial (see :func:`_cell_force_virial`). + Parameters ---------- energy @@ -109,6 +217,12 @@ 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). + topologies + Topologies on the same node axis whose ``LEAVES`` are further + autograd leaves of ``energy``. + n_local + (nf,) per-frame owned node counts, or ``None`` when every node is + owned; the owned nodes share the virial of a cell. Returns ------- @@ -119,9 +233,12 @@ def edge_energy_deriv( virial (nf, 3, 3) per-frame virial (always computed). """ - (g_e,) = torch.autograd.grad( + leaves = [ + getattr(topology, name) for topology in topologies for name in topology.LEAVES + ] + g_e, *g_leaves = torch.autograd.grad( energy.sum() if energy.dim() else energy, - edge_vec, + (edge_vec, *leaves), create_graph=create_graph, retain_graph=True, ) @@ -197,6 +314,29 @@ 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 ) + offset = 0 + for topology in topologies: + grads = g_leaves[offset : offset + len(topology.LEAVES)] + offset += len(topology.LEAVES) + if isinstance(topology, PairGraph): + (g_p,) = grads + leaf_force, leaf_atom_virial, leaf_virial = pair_force_virial( + g_p, + topology.vec, + topology.index, + topology.mask, + n_node, + node_capacity=force.shape[0], + ) + else: + g_coord, g_strain = grads + leaf_force, leaf_atom_virial, leaf_virial = _cell_force_virial( + g_coord, g_strain, topology.coord, n_node, n_local, do_atomic_virial + ) + force = force + leaf_force.to(force.dtype) + virial = virial + leaf_virial.to(virial.dtype) + if do_atomic_virial: + atom_virial = atom_virial + leaf_atom_virial.to(atom_virial.dtype) return force, (atom_virial if do_atomic_virial else None), virial @@ -211,6 +351,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 | KGrid | MeshGrid] | None = None, ) -> dict[str, torch.Tensor]: """Graph analogue of the dense pt_expt ``fit_output_to_model_output``. @@ -282,6 +423,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 + Topologies read by the atomic model, by name. Their geometry leaves + 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 +538,8 @@ 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, + topologies=() if topologies is None else tuple(topologies.values()), + n_local=n_local, ) # 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..5ee8794186 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 ( + flatten_topologies, +) 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,9 @@ def forward_lower_canonical_graph( *, do_atomic_virial: bool, spin: torch.Tensor | None = None, + coord: torch.Tensor | None = None, + communicator: torch.Tensor | None = None, + box: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Evaluate an eligible compressed canonical deployment graph. @@ -165,6 +172,21 @@ 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; ``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. + box + For a model with a long-range module, the cells of the frames + with shape ``(nf, 9)``, float64 on the CPU, the rows of every + row-major 3x3 cell its lattice vectors; a zero row marks a + complete non-periodic frame, and ``None`` non-periodic frames. + The halo nodes of a periodic frame are images of its owned + nodes. Returns ------- @@ -211,7 +233,33 @@ 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, box) + 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, + box, + 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 +319,9 @@ 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, + box: torch.Tensor | None = None, *, do_atomic_virial: bool, **make_fx_kwargs: Any, @@ -283,6 +334,12 @@ 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, box + For a model with a long-range module, sample node positions with + shape ``(N, 3)``, float64, a sample communicator, an int64 CPU + tensor of one element, and sample cells with shape ``(nf, 9)``, + float64 on the CPU, which the traced module then takes as its + ninth to eleventh inputs; ``None`` otherwise. do_atomic_virial Whether the traced output includes per-node virial. **make_fx_kwargs @@ -291,9 +348,59 @@ def forward_lower_canonical_graph_exportable( Returns ------- torch.nn.Module - Traced eight-input compact deployment module. + Traced compact deployment module of eight inputs, or eleven with + the node positions, the communicator and the cells. """ model = self + if not (coord is None) == (communicator is None) == (box is None): + raise ValueError( + "the compact lower of a model with a long-range module takes " + "the node positions, the communicator and the cells 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, + box: 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, + box=box, + ) + + 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, + box, + ) def fn( atype: torch.Tensor, @@ -613,6 +720,7 @@ def forward_lower_graph_exportable( do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, destination_sorted: bool = False, + topologies: dict[str, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """Trace ``forward_common_lower_graph`` into an exportable module with @@ -646,6 +754,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 topologies the atomic model declares, or ``None``; + their fields follow ``charge_spin`` as further inputs, in the + layout of + :func:`~deepmd.dpmodel.utils.neighbor_graph.flatten_topologies`. **make_fx_kwargs Extra keyword arguments forwarded to ``make_fx`` (e.g. ``tracing_mode="symbolic"``). @@ -680,6 +793,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 +802,67 @@ def forward_lower_graph_exportable( do_grad_r = self.do_grad_r("energy") do_grad_c = self.do_grad_c("energy") + if topologies: + # The topology inputs follow ``charge_spin``, their fields in the + # layout of flatten_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]: + 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, + *topology_vals, + ) + 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_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), + ) + 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..25e88487a7 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -23,6 +23,8 @@ from deepmd.dpmodel.utils.neighbor_graph import ( compact_nodes, expand_node_values, + flatten_topologies, + unflatten_topologies, ) from deepmd.pt_expt.common import ( auto_wrapped_class, @@ -43,6 +45,7 @@ from .edge_transform_output import ( fit_output_to_model_output_graph, + releaf_topologies, ) from .transform_output import ( fit_output_to_model_output, @@ -137,7 +140,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 +195,76 @@ 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, + periodic_marks: torch.Tensor | None = None, +) -> None: + """Reject halo nodes in the non-periodic frames of a single domain with a long-range module. + + The module sums over every pair of atoms of a non-periodic 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 halo nodes of a periodic frame are images of its atoms, + which carry no charge: the reciprocal-space sum holds the images itself. + 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. + periodic_marks + A tensor with a nonzero entry when the frames are periodic, such as + their cells or the frame mask of a reciprocal-space topology, or + ``None`` for non-periodic frames. It is read only after the fake + check, so that a traced lower computes nothing on it. + + Raises + ------ + ValueError + If a non-periodic 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 periodic_marks is not None and bool(torch.any(periodic_marks != 0)): + return + if not torch.equal(n_local, n_node): + raise ValueError( + "a model with a long-range module evaluates complete non-periodic " + "frames on a single domain, whose graph owns every node (n_local " + f"== n_node), but got n_local {n_local.tolist()} for n_node " + f"{n_node.tolist()}" + ) + + +def _reciprocal_masks(topologies: dict[str, Any]) -> torch.Tensor | None: + """The mask of the reciprocal-space topology, true on periodic entries, or ``None``.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PairGraph, + ) + + masks = [ + topology.mask + for topology in topologies.values() + if not isinstance(topology, PairGraph) + ] + return masks[0] if masks else None + + def _pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor: """Append a single ``-1`` column to ``nlist`` for export-time tracing. @@ -356,7 +430,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 +474,15 @@ 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 = self.model.atomic_model.build_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 +490,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 +693,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, Any] | None = None, ) -> dict[str, torch.Tensor]: """Graph-native lower with autograd force/virial (dpa1/se_atten concat-tebd, attention included). @@ -696,6 +782,13 @@ 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. The geometry fields every topology names in its + ``LEAVES`` become autograd leaves next to ``edge_vec``, and + force and virial include their pullbacks. ``None`` (default) + for models declaring none. Returns ------- @@ -712,6 +805,17 @@ 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, periodic_marks=_reciprocal_masks(topologies) + ) + # Every topology carries geometry leaves of its own. + topologies = releaf_topologies(topologies) 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 +855,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 +877,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 +998,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 +1006,14 @@ def call_common_ragged( n_node, box, self.get_rcut(), - getattr(self.atomic_model, "pair_excl", None), + pair_excl, + ) + topologies = self.atomic_model.build_topologies( + coord.detach(), + graph.n_node, + box, + atype=atype, + pair_excl=pair_excl, ) predict = self.forward_common_lower_graph( atype, @@ -918,6 +1032,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 +1137,13 @@ def _call_common_graph( spin_flat = ( spin.reshape(n_padded, 3)[node_index] if spin is not None else None ) + topologies = self.atomic_model.build_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 +1161,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 +1372,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, Any] | None = None, **make_fx_kwargs: Any, ) -> torch.nn.Module: """make_fx trace of ``forward_common_lower_graph`` with ``edge_vec`` @@ -1289,6 +1413,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 +1430,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 +1447,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 22fbc3e97c..bb327bb3be 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -72,6 +72,10 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, + unflatten_topologies, +) from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) @@ -112,6 +116,9 @@ from deepmd.pt_expt.model import ( get_model, ) +from deepmd.pt_expt.model.edge_transform_output import ( + releaf_topologies, +) from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) @@ -297,6 +304,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 +333,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 +786,12 @@ 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 the fields of each, laid out by :func:`flatten_topologies`, as + further inputs ahead of the task buffers; their dynamic axes, the pairs, + the wave numbers or the axes of the mesh, are symbols distinct from the + other three. + Parameters ---------- model @@ -842,6 +869,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 +927,22 @@ 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 dynamic axes of the topologies take sample lengths 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, + nframes=trace_nf, + forbidden=_forbidden | {e_max}, + ) + 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 +959,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 +996,7 @@ def fn( aparam=aparam, charge_spin=charge_spin, spin=spin, + topologies=topologies, ) if spin is None: return _translate_energy_keys( @@ -996,6 +1043,7 @@ def fn( s_aparam, s_charge_spin, s_spin, + *s_topology_vals, *task_buf_vals_trace, ) @@ -1640,6 +1688,15 @@ 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 = _model.atomic_model.build_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 +1724,9 @@ def _forward_graph( edge_vec = ng.edge_vec.detach().requires_grad_(True) if spin is not None: spin = spin.detach().requires_grad_(True) + # The topologies enter as their fields, the geometry of every one a + # further leaf. + topology_vals = flatten_topologies(releaf_topologies(topologies)) if self._task_buf_order: try: @@ -1708,6 +1768,7 @@ def _forward_graph( aparam, charge_spin, spin, + *topology_vals, *task_buf_vals, ) self._report_pending_compile() @@ -2257,6 +2318,9 @@ def initialize_statistics( for model_key in self.model_keys } new_state = {} + # The atomic models, by module path, whose long-range module + # takes its initialization rather than pretrained weights. + fresh_long_range: dict[str, str] = {} for key in target_state: if key == "_extra_state": new_state[key] = target_state[key] @@ -2283,17 +2347,39 @@ def initialize_statistics( finetune_rule.get_random_fitting() and ".descriptor." not in key ) - if use_random: + if use_random or pretrained_key not in pretrained_state: new_state[key] = target_state[key] - elif pretrained_key in pretrained_state: - new_state[key] = pretrained_state[pretrained_key] + if ".long_range." in key: + owner = key[: key.index(".long_range.")] + fresh_long_range[owner] = model_key else: - new_state[key] = target_state[key] + new_state[key] = pretrained_state[pretrained_key] break if not matched: new_state[key] = target_state[key] self._unwrapped.load_state_dict(new_state) + # A long-range module added to the pretrained model takes the + # input statistics of the fine-tuning data, as in training from + # scratch, before the bias adjustment evaluates the model. + for owner, model_key in fresh_long_range.items(): + atomic_model = self._unwrapped.get_submodule(owner) + + def long_range_stat( + _atomic_model: torch.nn.Module = atomic_model, + _model_key: str = model_key, + ) -> None: + _atomic_model.compute_long_range_input_stat( + self._sample_funcs[_model_key] + ) + + self._run_stat_on_chief( + long_range_stat, + operation=f"long-range statistics for task {model_key!r}", + ) + if self.is_distributed: + self._broadcast_model_stat(atomic_model.long_range) + # Per-branch bias adjustment (rank 0 only, then broadcast) for model_key in self.model_keys: finetune_rule = finetune_links[model_key] @@ -3560,8 +3646,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 1bec3c9b31..6d511315a0 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import contextlib import ctypes +import dataclasses import json import logging import os @@ -14,6 +15,9 @@ import numpy as np import torch +from deepmd.dpmodel.utils.neighbor_graph import ( + flatten_topologies, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -688,18 +692,120 @@ def build_synthetic_graph_inputs( ) +def build_synthetic_topologies( + model: torch.nn.Module, + n_node_total: int, + n_entry: int, + *, + dtype: torch.dtype, + device: torch.device, + nframes: int, + forbidden: set[int] | frozenset[int] = frozenset(), +) -> dict[str, Any]: + """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. The pair + topology lists ``n_entry`` pairs, a length the caller keeps distinct from + the other dynamic axes; a mesh takes the three smallest even dimensions, + and a k grid the three smallest lengths for its slots, its blocks and the + points of a block, of at least four that no other axis and no entry of + ``forbidden`` has. Every length thus differs from the two rows of a pair + index and the three components of a vector, and from every other dynamic + axis, which a compiled graph would otherwise identify with it. + + Parameters + ---------- + model : torch.nn.Module + The model; its atomic model declares the topologies. + n_node_total : int + Length of the node axis. + n_entry : int + Length of the pair axis. + dtype : torch.dtype + Precision of the geometry. + device : torch.device + Target device. + nframes : int + Length of the frame axis. + forbidden : set[int] + Sizes of the static dimensions of the model, which the mesh and k-grid + axes avoid. + + Returns + ------- + dict[str, PairGraph | KGrid | MeshGrid] + A sample of every declared topology; empty when none is declared. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + K_GRID, + PME_MESH, + KGrid, + PairGraph, + ) + from deepmd.dpmodel.utils.neighbor_graph.reciprocal_graph import ( + mesh_grid_of_dims, + ) + + taken = set(forbidden) | {nframes, n_node_total, n_node_total + 1, n_entry} + mesh_dims = [] + for size in range(4, 4 + 2 * len(taken) + 6, 2): + if size not in taken and len(mesh_dims) < 3: + mesh_dims.append(size) + free = sorted(set(range(4, len(taken) + 10)) - taken - set(mesh_dims)) + n_slot, n_block, n_point = free[:3] + coord = torch.ones(n_node_total, 3, dtype=dtype, device=device) + # Contiguous like the inputs at run time: a compiled graph reads the + # strides of its samples. + box = torch.eye(3, dtype=dtype, device=device).repeat(nframes, 1, 1) + topologies = {} + for name in model.atomic_model.required_topologies(): + if name == K_GRID: + slot = torch.arange(nframes * n_slot, dtype=torch.int64, device=device) + blocks = (nframes, n_block, n_point) + topologies[name] = KGrid( + # Nonzero wave numbers: the kernel transform carries 1/k^2. + vec=torch.ones(*blocks, 3, dtype=dtype, device=device), + mask=torch.ones(*blocks, dtype=torch.bool, device=device), + slot=(slot % n_node_total).reshape(nframes, n_slot), + volume=torch.ones(nframes, dtype=dtype, device=device), + coord=coord, + strain=torch.zeros(nframes, 3, 3, dtype=dtype, device=device), + ) + elif name == PME_MESH: + topologies[name] = mesh_grid_of_dims( + box, + coord, + tuple(mesh_dims), + torch.ones(nframes, dtype=torch.bool, device=device), + ) + else: + index = torch.arange(2 * n_entry, dtype=torch.int64, device=device) + topologies[name] = PairGraph( + index=index.reshape(2, n_entry) % n_node_total, + vec=torch.ones(n_entry, 3, dtype=dtype, device=device), + mask=torch.ones(n_entry, dtype=torch.bool, device=device), + ) + return 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, the + int64 communicator, a CPU tensor of one element, at slot 9, and the + float64 cells of the frames, a CPU tensor of shape ``(nf, 9)``, at slot 10. """ from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -756,7 +862,18 @@ 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") + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64, device="cpu") + return (*compact_sample, coord, communicator, box) + return compact_sample def _build_canonical_graph_dynamic_shapes( @@ -764,8 +881,11 @@ 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, + the communicator and the cells 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; the cells follow + the frame axis. """ from deepmd.pt_expt.utils.canonical_graph import ( UINT32_MAX, @@ -788,7 +908,10 @@ 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, {0: nframes_dim})[ + : len(sample_inputs) - len(shapes) + ] + return (*shapes, *trailing) def count_synthetic_graph_edges( @@ -841,9 +964,53 @@ def count_synthetic_graph_edges( return int(edge_mask.sum()) +def _topology_dynamic_shapes( + topology_names: tuple[str, ...], shared: dict[str, Any] +) -> list[dict[int, Any]]: + """Dynamic-shape specifications of the fields of the topologies of a graph lower. + + Every field follows the ``AXES`` its dataclass declares. The frame and + node axes are those of the other inputs in ``shared``; the entry axis of + every topology is a dimension of its own, of at least the two guard + entries, and so are the slot axis and the point axis of a k grid; a mesh + axis is twice a dimension, since the mesh dimensions are even. + + Parameters + ---------- + topology_names : tuple[str, ...] + The topologies, in the order of their inputs. + shared : dict[str, Any] + The dimensions of the ``"frame"`` and ``"node"`` axes. + + Returns + ------- + list[dict[int, Any]] + One specification per field, in the layout of flatten_topologies. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + TOPOLOGY_TYPES, + ) + + specs = [] + for name in topology_names: + cls = TOPOLOGY_TYPES[name] + dims = dict(shared) + dims["entry"] = torch.export.Dim(f"{name}_entries", min=2) + dims["slot"] = torch.export.Dim(f"{name}_slots", min=2) + dims["k_point"] = torch.export.Dim(f"{name}_block_points", min=2) + for axis in ("mesh_1", "mesh_2", "mesh_3"): + dims[axis] = 2 * torch.export.Dim(f"{name}_{axis}_half", min=1) + specs += [ + {dim: dims[axis] for dim, axis in cls.AXES[field.name].items()} + for field in dataclasses.fields(cls) + ] + return specs + + 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. @@ -871,6 +1038,9 @@ 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 fields follow the inputs above (see + :func:`_topology_dynamic_shapes`). """ nframes_dim = torch.export.Dim("nframes", min=1) n_node_total_dim = torch.export.Dim("n_node_total", min=1) @@ -887,6 +1057,9 @@ def _build_graph_dynamic_shapes( {0: nedge_dim}, # source_order: (E,) {0: n_node_total_dim + 1}, # source_row_ptr: (N + 1,) ) + topologies = _topology_dynamic_shapes( + topology_names, {"frame": nframes_dim, "node": n_node_total_dim} + ) if is_native_spin: spin = sample_inputs[10] fparam = sample_inputs[11] @@ -900,6 +1073,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] @@ -912,6 +1086,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, ) @@ -1300,6 +1475,16 @@ 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, and the parameters + # they are built with; absent when the model reads none. The graph lower + # takes the fields of each after its other inputs (see + # ``flatten_topologies``), built by the caller from the coordinates and + # the cell of each frame. The compact canonical lower takes the node + # positions instead, from which its operators evaluate the long range. + topologies = model.atomic_model.required_topologies() + if topologies: + meta["topologies"] = list(topologies) + meta["topology_params"] = model.atomic_model.topology_params() return meta @@ -1911,6 +2096,15 @@ def _trace_and_export_impl( # pt_expt BaseModel registry returns this backend's torch class. model = BaseModel.deserialize(data["model"]) model.eval() + # A frozen model evaluates the reciprocal sum of a periodic frame on the + # mesh of its cell, whatever the mode it was trained in. + from deepmd.dpmodel.long_range.latent_charge import ( + LatentChargeLongRange, + ) + + for module in model.modules(): + if isinstance(module, LatentChargeLongRange): + module.use_mesh_evaluation() # Fold the vacuum reference on the CPU tracing device. vacuum_descriptor # evaluates an isolated atom of every type through the descriptor; on CUDA # that path can allocate several GiB even when export already holds most of @@ -2080,11 +2274,17 @@ 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 and the cells, from which its operator evaluates + # the long range, 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, @@ -2166,6 +2366,20 @@ def _trace_and_export_impl( device=torch.device("cpu"), want_spin=is_native_spin, ) + # The pair axis 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"), + nframes=nframes_sample, + forbidden=_forbidden | {e_sample}, + ) if is_native_spin: # Native-spin ABI (NativeSpinEnergyModel.forward_lower_graph_exportable): # slot 10 is ``spin`` (mandatory), slots 11/12 are @@ -2191,12 +2405,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, @@ -2527,7 +2745,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 @@ -2560,6 +2852,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 ------- @@ -2580,9 +2876,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 @@ -2607,6 +2912,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. @@ -2622,6 +2928,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 ------- @@ -2648,7 +2958,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, @@ -2781,11 +3096,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 8a7131e4f3..0e3ac1873a 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3546,6 +3546,306 @@ 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. " + "For a non-periodic system the kernel is summed directly over every pair " + "of atoms; for a periodic system (``periodic: true``) it is summed over " + "the images of the cell in reciprocal space. Every channel interacts " + "through the same kernel; every atom carries half of the interaction of " + "its charges with all others.", +) +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." + ) + doc_periodic = ( + "Evaluate the module on periodic systems as well: the kernel is summed " + "over the images of the cell in reciprocal space, with the " + "zero-wave-number mode dropped (the tin-foil convention, which places " + "a uniform neutralizing background under a channel with a net " + "charge), less the interaction of every charge with itself. " + "Non-periodic systems keep the direct sum." + ) + doc_kspace_tolerance = ( + "The relative weight of the slowest-decaying kernel mode at the " + "wave-number cutoff k_max of a periodic system." + ) + doc_k_grid_mode = ( + "How training evaluates the reciprocal sum of a periodic batch: " + "'pme', the default of a periodic module, by the smooth particle-mesh " + "Ewald method on a mesh sized by the cells of the batch at " + "`pme_spacing`, whose cost grows nearly linearly with the system; " + "'ragged' over the wave numbers " + "of every frame within k_max, enumerated per batch, so that the k axis " + "is a dynamic dimension; or 'fixed' over one integer grid shared by " + "the frames, sized by the largest cell of the data statistics or by " + "`k_grid_bounds`, with the points a frame does not reach masked, so " + "that the k axis has a static length (a frame whose cell needs points " + "beyond the grid is truncated to it, with a warning). A frozen model " + "evaluates every periodic frame on the mesh of its cell, whatever the " + "mode. 'fixed' and 'pme' require `periodic: true`." + ) + doc_k_grid_bounds = ( + "The per-axis integer bounds [n1, n2, n3] of the fixed k grid, in " + "place of the data statistics; only meaningful with " + "`k_grid_mode: fixed`." + ) + doc_decouple_k_grid = ( + "Decouple the wave-number cutoff of a periodic system from the kernel " + "parameters: it is derived once from the initial parameters and kept " + "for the life of the model, so that the k grid and the mesh stay put " + "while the kernel trains (the kernel values still follow the live " + "parameters). The frozen cutoff is serialized with the model. " + "Requires `periodic: true`." + ) + doc_pme_spacing = ( + "The largest spacing of the mesh of a periodic frame, in Angstrom; " + "the default is three quarters of the Nyquist spacing of the " + "wave-number cutoff, 0.75 pi / k_max. Requires `periodic: true`." + ) + doc_pme_order = ( + "The order of the B-splines that spread the charges to the mesh, 8 " + "by default. Requires `periodic: true`." + ) + 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, + ), + Argument("periodic", bool, optional=True, default=False, doc=doc_periodic), + Argument( + "kspace_tolerance", + float, + optional=True, + default=1e-8, + doc=doc_kspace_tolerance, + ), + Argument( + "k_grid_mode", + [str, None], + optional=True, + default=None, + doc=doc_k_grid_mode, + extra_check=lambda x: x in (None, "ragged", "fixed", "pme"), + extra_check_errmsg="must be 'ragged', 'fixed' or 'pme'", + ), + Argument( + "k_grid_bounds", + list[int], + optional=True, + default=None, + doc=doc_k_grid_bounds, + ), + Argument( + "decouple_k_grid", + bool, + optional=True, + default=False, + doc=doc_decouple_k_grid, + ), + Argument( + "pme_spacing", + [float, None], + optional=True, + default=None, + doc=doc_pme_spacing, + ), + Argument( + "pme_order", + [int, None], + optional=True, + default=None, + doc=doc_pme_order, + ), + ] + + +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 systems: " + "non-periodic ones directly, periodic ones in reciprocal space when " + "the module serves them. 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() @@ -3845,6 +4145,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..5da7d4353e --- /dev/null +++ b/doc/model/long-range.md @@ -0,0 +1,247 @@ +# 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 sum covers complete systems. Without `periodic: true`, a periodic cell +> is rejected. The cost of the non-periodic 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. + +### Periodic systems + +With +{ref}`periodic ` enabled, +a frame given a cell is evaluated over the images of the cell. The sum +converges in reciprocal space, where the Fourier transform of the kernel, + +- `les`: $\hat K(\mathbf{k}) = (4\pi/k^2)\, e^{-k^2/4\alpha^2}$, or +- `sog`: $\hat K(\mathbf{k}) = \sum_k a_k (2\pi)^{3/2} \sigma_k^3\, e^{-\sigma_k^2 k^2/2}$, + which with the default amplitudes approaches the Coulomb transform + $4\pi/k^2$, + +decays like a Gaussian. Over the wave numbers $\mathbf{k} = 2\pi B^{-1} +\mathbf{n}$, $\mathbf{n} \in \mathbb{Z}^3 \setminus \{\mathbf{0}\}$, of a cell +$B$ whose rows are the lattice vectors, + +```math +E = k_e \left[\frac{1}{2V}\sum_{\mathbf{k}\neq\mathbf{0}} \sum_{c=1}^{C} +\hat K(\mathbf{k})\, |\hat\rho_c(\mathbf{k})|^2 +- \frac{1}{2} K(0) \sum_{i,c} (q_i^c)^2 \right], +\qquad \hat\rho_c(\mathbf{k}) = \sum_i q_i^c e^{i\mathbf{k}\cdot\mathbf{r}_i}, +``` + +with the cell volume $V$ and the kernel value at contact $K(0)$, whose term +removes the interaction of every charge with itself. The +$\mathbf{k}=\mathbf{0}$ mode is dropped, the tin-foil convention, which places +a uniform neutralizing background under a channel with a net charge. The +wave-number cutoff $k_{max}$ is where the weight of the slowest-decaying mode +falls below +{ref}`kspace_tolerance `, +so the truncation converges exponentially; it follows the kernel as it trains, +unless +{ref}`decouple_k_grid ` +keeps the cutoff of the initial kernel. + +{ref}`k_grid_mode ` +selects how training evaluates the sum: + +- `pme`, the default: by the smooth particle-mesh Ewald method. The charges + are spread with B-splines of order + {ref}`pme_order ` to a + mesh that resolves every cell of the batch at + {ref}`pme_spacing `, + whose transform gives the sum; the cost grows nearly linearly with the + number of atoms. With the default order 8 and spacing $0.75\pi/k_{max}$, the + energy and the forces stay within a few parts in a million of the sum over + the wave numbers. +- `ragged`: over the wave numbers of every frame within $k_{max}$, listed for + every batch; their number grows with the volume of the cell. +- `fixed`: over one integer grid shared by the frames of a batch, sized by the + largest cell of the data statistics or by + {ref}`k_grid_bounds `; + a frame whose cell needs a wave number beyond the grid is truncated to it, + with a warning. + +A frozen model evaluates every periodic frame on the mesh of its cell, +whatever its training mode, so a model trained in another mode can differ +from its frozen form by the accuracy of the mesh. Forces and virials hold the +complete periodic term, and a model with `periodic: true` still evaluates +non-periodic frames by the direct sum. In a `linear_ener` composition, a +sub-model whose module serves periodic systems takes no other sub-model with +a long-range module beside it. + +### 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. The reciprocal-space sum of a periodic +system is global and excludes no pair, so `pair_exclude_types` is not +supported with a module that serves periodic systems, nor +`atom_exclude_types` with any 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, the wave +numbers or the mesh of each batch are built outside the compiled region and +enter it as inputs of dynamic size. + +`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 sum covers the complete system, LAMMPS runs a model whose module +serves no periodic system with non-periodic boundaries (`f`, `s` or `m`) in +every dimension, and a model with `periodic: true` with boundaries periodic in +every dimension or in none. Several models given to `pair_style deepmd` for +the model deviation evaluate a non-periodic system together only when every +one of them carries a long-range module. 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 its own atoms against all of them, +so that the ranks share the cost of the direct sum, while each evaluates the +mesh of a periodic system whole. This requires DeePMD-kit built with MPI. Any +other model, like an uncompressed one, runs on a single rank; +`pair_style deepmd/kk` evaluates an uncompressed model on non-periodic systems +only. + +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 long-range +term with a dedicated operator, which reads the positions of the atoms and +the cell. It enumerates the pairs of a non-periodic system itself, so that its +memory grows linearly with the number of atoms and its cost with their +square, and evaluates a periodic system on the mesh of its cell. The operator +serves the `les` kernel and the `sog` kernel of up to 32 Gaussians, with up to +8 charge channels, B-splines of order up to 12, a charge network without +frame or atomic parameters, and no `pair_exclude_types`. Any other model, +like an uncompressed one, receives the pairs or the mesh of each frame from +the interface, so the memory of a non-periodic evaluation 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..4a866538ae 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,71 @@ extern bool DP_DeepPotUsesFP32EdgeVectors(DP_DeepPot* dp); */ extern bool DP_DeepPotUsesCanonicalGraphInference(DP_DeepPot* dp); +/** + * @brief Evaluate a compact canonical graph, the node positions and the cell + * on the model device. + * + * A model with a long-range module reads the positions of the nodes, which + * must form a complete system: held whole by one rank, or decomposed over the + * ranks of ``world``, whose owned nodes ``[0, nloc)`` are the share of this + * rank. A non-periodic system, which every such model serves, comes without a + * cell, and one rank then holds every node (``nall_nodes == nloc``); a + * periodic system, which a model whose module serves periodic systems + * evaluates over the images of its cell, comes with it, and the halo nodes of + * a single rank are images of its owned nodes. 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, the communicator and + * the cell. + * + * @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] box The cell of a periodic system in host memory, nine values in + * the length unit of the model, row-major with the lattice vectors as rows, + * or NULL for a non-periodic system. + * @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, the communicator and the cell. + * @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 double* box, + 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 systems; see DP_DeepPotComputeCanonicalGraphGPU2. + * @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 +2748,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 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..5ff9a7833e 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1406,6 +1406,42 @@ class DeepPot : public DeepBaseModel { return result; } + /** + * @brief Evaluate a compact canonical graph, the node positions and the + * cell 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 double* box, + 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, box, + 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 systems; see DP_DeepPotComputeCanonicalGraphGPU2. + */ + 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 +2644,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi { return dchgspin; } + /** + * @brief Query whether any of the models carries a long-range module, which + * evaluates complete 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..6d0736e08a 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -1864,6 +1864,38 @@ 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 double* box, + 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, box, 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 +2880,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..c02fbe4365 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -426,6 +426,9 @@ class DeepPotBackend : public DeepBaseModelBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -437,6 +440,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 +898,65 @@ 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, or over the images of the cell of a + * periodic one when its module serves periodic systems; it evaluates a + * system 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 and the cell, which is NULL for a + * non-periodic system. + */ + 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. + * @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`` + * without a cell). With a long-range module and several ranks, every rank + * calls this function on every evaluation, also when it holds no node. + * @param[in] box The cell of a periodic system in host memory, nine values + * in the length unit of the model, row-major with the lattice vectors as + * rows, or NULL for a non-periodic system. A model with a long-range module + * that serves periodic systems reads it; the halo nodes of a single rank + * are then images of the owned nodes. + * + * 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 double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -959,6 +1018,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 systems (see DeepPot::has_long_range). + **/ + 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 @@ -1134,5 +1204,14 @@ class DeepPotModelDevi : public DeepBaseModelDevi { protected: std::vector> dps; + + private: + /** + * @brief Check that every model accepts the cell handed to the ensemble. + * Without one, the neighbor-list entries serve a model with a long-range + * module beside models that take the cell of the region. + **/ + template + void check_ensemble_box(const std::vector& dbox) const; }; } // namespace deepmd diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 4707ee9ce6..bc9b685021 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -377,6 +377,9 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -413,6 +416,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 +438,9 @@ class DeepPotPTExpt : public DeepPotBackend { double* d_force, double* d_atom_virial, const std::int64_t* d_atype, + const double* d_coord, + void* world, + const double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -473,6 +480,15 @@ 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 topologies + // of the frame after its other inputs: the pair lists, and the mesh of the + // cell of a module that serves periodic systems, for the graph lower; the + // node positions, the communicator of the ranks and the cell for the compact + // canonical lower. + bool long_range_ = false; + // The module serves periodic systems, on a mesh of at most this spacing. + bool long_range_periodic_ = false; + double mesh_spacing_ = 0.0; NeighborListData nlist_data; at::Tensor mapping_tensor; // cached mapping tensor (LAMMPS path) std::vector mapping_; // cached mapping vector (LAMMPS path) @@ -615,6 +631,11 @@ 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 topologies of a model with a + * long-range module are built; read only for such a model. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; read only for such a model. */ std::vector run_model_graph( const torch::Tensor& atype, @@ -629,8 +650,22 @@ 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, + const std::vector& box); + /** + * @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. + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one; a model with a long-range module receives + * it after the communicator. + */ std::vector run_model_canonical_graph( const torch::Tensor& atype, const torch::Tensor& n_node, @@ -639,7 +674,30 @@ 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, + const std::vector& box); + + /** + * @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``, in the cell ``box`` of the + * system, empty for a non-periodic one. + */ + void join_long_range_exchange(const void* world, + const std::vector& box); + + /** + * @brief Check the cell given to a model with a long-range module. + * + * @param[in] box The row-major cell of a periodic frame, nine values, or + * empty for a non-periodic one, which every long-range module serves; a + * cell requires a module that serves periodic systems. + */ + void check_long_range_box(const std::vector& box) const; /** * @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..d3ec820acf 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -661,6 +661,9 @@ 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 double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -673,6 +676,9 @@ void DeepPotBackend::compute_canonical_graph_gpu( (void)d_force; (void)d_atom_virial; (void)d_atype; + (void)d_coord; + (void)world; + (void)box; (void)d_source; (void)d_edge_vec; (void)d_destination_row_ptr; @@ -692,6 +698,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,11 +767,34 @@ 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_atom_energy, d_force, d_atom_virial, d_atype, /*d_coord=*/nullptr, + /*world=*/nullptr, /*box=*/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 double* box, + 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, box, + 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 { return dp->uses_fp32_edge_vectors(); } @@ -776,6 +807,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) { @@ -921,6 +954,25 @@ template void DeepPotModelDevi::compute( const std::vector& aparam, const std::vector& charge_spin); +// LAMMPS hands a model with a long-range module no cell for a non-periodic +// system, whereas the neighbor-list entries of other models, as of the +// TensorFlow backend, read the cell of the region. +template +void DeepPotModelDevi::check_ensemble_box( + const std::vector& dbox) const { + if (!dbox.empty() || !has_long_range()) { + return; + } + for (const auto& model : dps) { + if (!model->has_long_range()) { + throw deepmd::deepmd_exception( + "an ensemble that holds a model with a long-range module evaluates " + "a non-periodic system only when every model carries one; use " + "models with long-range modules throughout, or a periodic system."); + } + } +} + template void DeepPotModelDevi::compute(std::vector& all_energy, std::vector>& all_force, @@ -937,6 +989,7 @@ void DeepPotModelDevi::compute(std::vector& all_energy, if (numb_models == 0) { return; } + check_ensemble_box(dbox); all_energy.resize(numb_models); all_force.resize(numb_models); all_virial.resize(numb_models); @@ -999,6 +1052,7 @@ void DeepPotModelDevi::compute( all_virial.resize(numb_models); all_atom_energy.resize(numb_models); all_atom_virial.resize(numb_models); + check_ensemble_box(dbox); for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_, 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..751d793830 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -22,6 +23,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 +112,28 @@ 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)); +} + +// The cell input of the compact lower of a long-range model: a float64 CPU +// tensor of shape (1, 9), the rows of the row-major cell its lattice vectors, +// zero for a non-periodic frame. +at::Tensor cell_tensor(const std::vector& box) { + auto cell = + torch::zeros({1, 9}, torch::TensorOptions().dtype(torch::kFloat64)); + if (!box.empty()) { + std::copy(box.begin(), box.begin() + 9, cell.data_ptr()); + } + return cell; +} + } // namespace void DeepPotPTExpt::translate_error(std::function f) { @@ -256,6 +280,31 @@ 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, and, when the module serves periodic + // systems, the mesh of the cell, sized by the spacing of the metadata. + long_range_ = metadata.obj_val.count("topologies") > 0; + long_range_periodic_ = false; + if (long_range_) { + for (const auto& name : metadata["topologies"].as_array()) { + if (name.as_string() == "pme_mesh") { + long_range_periodic_ = true; + } else if (name.as_string() != "all_pairs") { + throw deepmd::deepmd_exception("unknown topology '" + name.as_string() + + "' in the .pt2 metadata; expected " + "'all_pairs' or 'pme_mesh'."); + } + } + if (long_range_periodic_) { + mesh_spacing_ = + metadata["topology_params"]["pme_mesh"]["spacing"].as_double(); + } + 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 +562,9 @@ 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, + const std::vector& box) { // 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 +580,21 @@ std::vector DeepPotPTExpt::run_model_graph( if (dchgspin > 0) { inputs.push_back(charge_spin); } + if (long_range_) { + // A periodic frame holds no pair; the mesh describes its interactions. + const auto vec_dtype = graph_edge_fp32_ ? torch::kFloat32 : torch::kFloat64; + const auto pairs = + box.empty() + ? deepmd::ptexpt::buildAllPairsTopology( + coord, atype, pair_exclude_table_, ntypes, vec_dtype) + : deepmd::ptexpt::buildGuardPairTopology(coord.device(), vec_dtype); + inputs.insert(inputs.end(), pairs.begin(), pairs.end()); + if (long_range_periodic_) { + const auto mesh = deepmd::ptexpt::buildMeshTopology( + coord, box, mesh_spacing_, vec_dtype); + inputs.insert(inputs.end(), mesh.begin(), mesh.end()); + } + } return loader->run(inputs); } @@ -540,9 +606,70 @@ 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, + const std::vector& box) { + 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)); + inputs.push_back(cell_tensor(box)); + } + return loader->run(inputs); +} + +void DeepPotPTExpt::check_long_range_box(const std::vector& box) const { + if (!box.empty() && box.size() != 9) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates one frame with a cell of " + "nine values, or without one, but got " + + std::to_string(box.size()) + " values."); + } + if (!box.empty() && !long_range_periodic_) { + throw deepmd::deepmd_exception( + "the long-range module of this model serves complete non-periodic " + "systems only; pass no cell (in LAMMPS, a non-periodic boundary f, s " + "or m in every dimension), or train the module with periodic: true."); + } + if (!box.empty()) { + // A cell whose volume is a vanishing fraction of the product of the + // lengths of its lattice vectors, or not finite, spans no lattice. + const double volume = + std::abs(box[0] * (box[4] * box[8] - box[5] * box[7]) - + box[1] * (box[3] * box[8] - box[5] * box[6]) + + box[2] * (box[3] * box[7] - box[4] * box[6])); + double lengths = 1.0; + for (int d = 0; d < 3; ++d) { + lengths *= + std::sqrt(box[3 * d] * box[3 * d] + box[3 * d + 1] * box[3 * d + 1] + + box[3 * d + 2] * box[3 * d + 2]); + } + if (!(volume > 1e-6 * lengths)) { + throw deepmd::deepmd_exception( + "the cell of a periodic system needs three independent lattice " + "vectors, but its volume is " + + std::to_string(volume) + "; pass no cell for a non-periodic system."); + } + } +} + +void DeepPotPTExpt::join_long_range_exchange(const void* world, + const std::vector& box) { + 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, box); } std::vector DeepPotPTExpt::run_model_with_comm( @@ -799,6 +926,40 @@ 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 a non-periodic + // system, which the caller marks with an empty box, and over the images of + // the cell of a periodic one. A single rank holds every atom, 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. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } + if (long_range_ && long_range_box.empty() && !multi_rank && nghost_real > 0) { + throw deepmd::deepmd_exception( + "a model with a long-range module evaluates a complete non-periodic " + "system without a box, but this single-rank call carries " + + std::to_string(nghost_real) + + " ghost atoms, the periodic images of a periodic cell; pass the cell " + "of a periodic system to a model whose module serves one."); + } + 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 " + "long-range 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 +1388,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, long_range_box); + } ener.assign(nframes, static_cast(0)); force.assign(static_cast(nframes) * fwd_map.size() * 3, static_cast(0)); @@ -1268,19 +1434,22 @@ 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, + long_range_box); } 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, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -1473,6 +1642,12 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, check_call_charge_spin(charge_spin, /*nframes=*/1, settable_chgspin, dchgspin > 0, default_chg_spin_, chg_spin_table_ranges_); + // A long-range module sums over the atoms of a non-periodic system, which + // the empty box marks, and over the images of the cell of a periodic one. + const std::vector long_range_box(box.begin(), box.end()); + if (long_range_) { + check_long_range_box(long_range_box); + } // The .pt2 model only contains forward_common_lower, which requires // nlist as input. We must build the nlist in C++ and fold back the // extended-region outputs to local atoms. @@ -1648,12 +1823,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, long_range_box); } else { flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.n_local, @@ -1665,7 +1842,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, long_range_box); } } else { // Model-level pair exclusion is a BUILD-time transform (decision @@ -2221,6 +2398,15 @@ 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); periodic systems and several ranks require the " + "compact canonical lower, which a compressed DPA4C model frozen for a " + "CUDA device uses when the long-range operator serves its module, or " + "the evaluation from a neighbor list (pair_style deepmd)."); + } 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 +2634,8 @@ 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}), + /*box=*/{})); 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 +2724,9 @@ 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 double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2553,6 +2743,36 @@ 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 system and the + // cell of a periodic one: the node set of a single rank, whose halo nodes + // are images of a periodic cell, or the owned nodes of every rank of + // ``world``, which a rank without nodes still joins. + const std::vector long_range_box = + box != nullptr ? std::vector(box, box + 9) + : std::vector(); + if (long_range_) { + check_long_range_box(long_range_box); + } + 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 && box == 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), or a periodic one with its cell; 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, long_range_box); + 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 +2805,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, long_range_box)); auto atom_energy = output["atom_energy"] .reshape({nall_nodes}) .slice(0, 0, nloc) @@ -2672,6 +2898,9 @@ 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 double* box, const std::uint32_t* d_source, const float* d_edge_vec, const std::int64_t* d_destination_row_ptr, @@ -2681,9 +2910,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, box, + 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 +2929,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..352f85a643 --- /dev/null +++ b/source/api_cc/src/longRangePTExpt.h @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include + +#include +#include +#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)})}; +} + +/** + * @brief Build the all-pairs topology of a periodic frame, which holds no + * pair: the reciprocal-space topology describes the interactions. + * + * Counterpart of Python ``empty_pair_graph`` for the exported lower: two + * masked guard pairs, the lower bound of its pair axis. + */ +inline std::vector buildGuardPairTopology( + const torch::Device& device, const torch::ScalarType vec_dtype) { + const auto options = torch::TensorOptions().device(device); + return {torch::zeros({2, 2}, options.dtype(torch::kInt64)), + torch::zeros({2, 3}, options.dtype(vec_dtype)), + torch::zeros({2}, options.dtype(torch::kBool))}; +} + +/** + * @brief The smallest even integer of at least ``n`` without a prime factor + * above 7. + */ +inline std::int64_t evenFftSize(std::int64_t n) { + std::int64_t m = std::max(2, n + n % 2); + while (true) { + std::int64_t rest = m; + for (const std::int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +/** + * @brief Build the mesh topology a periodic long-range module of the graph + * lower reads. + * + * Twin of Python ``build_mesh_grid`` in + * ``deepmd/dpmodel/utils/neighbor_graph/reciprocal_graph.py`` for a single + * frame: along every lattice vector ``a_d`` the mesh holds the smallest even + * number of points without a prime factor above 7 of at least + * ``ceil(|a_d| / spacing)``, with the length formed term by term as Python + * forms it, so that both size the mesh identically. Without a cell, the frame + * is non-periodic: the cell is the identity, the mesh minimal and the frame + * masked. + * + * @param coord Node positions, shape ``(N, 3)``, on the model device. + * @param box The row-major cell, the rows its lattice vectors, nine values in + * the length unit of the model, or empty for a non-periodic frame. + * @param spacing The largest mesh spacing, in the length unit of the model. + * @param dtype Precision of the geometry, that of the edge vectors. + * @return ``{box (1, 3, 3), volume (1,), coord (N, 3), strain (1, 3, 3), + * mask (1,) bool, mode_1 (n1,), mode_2 (n2,), mode_3 (n3,)}``. + */ +inline std::vector buildMeshTopology( + const torch::Tensor& coord, + const std::vector& box, + const double spacing, + const torch::ScalarType dtype) { + const auto options = + torch::TensorOptions().dtype(dtype).device(coord.device()); + const bool periodic = !box.empty(); + std::array cell = {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + std::array dims = {2, 2, 2}; + if (periodic) { + for (int k = 0; k < 9; ++k) { + cell[k] = box[k]; + } + for (int d = 0; d < 3; ++d) { + const double* v = cell.data() + 3 * d; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + dims[d] = + evenFftSize(static_cast(std::ceil(length / spacing))); + } + } + const double volume = + std::abs(cell[0] * (cell[4] * cell[8] - cell[5] * cell[7]) - + cell[1] * (cell[3] * cell[8] - cell[5] * cell[6]) + + cell[2] * (cell[3] * cell[7] - cell[4] * cell[6])); + // The copy detaches the tensor from the local array, which a conversion + // to float64 on the host would alias. + const auto cell_tensor = + torch::from_blob(cell.data(), {1, 3, 3}, torch::kFloat64) + .to(options, /*non_blocking=*/false, /*copy=*/true); + std::vector fields = { + cell_tensor, torch::full({1}, volume, options), coord.to(dtype), + torch::zeros({1, 3, 3}, options), + torch::full({1}, periodic, options.dtype(torch::kBool))}; + for (const std::int64_t n : dims) { + // The integer frequencies of an axis of n points, in FFT order. + const auto m = torch::arange(n, options); + fields.push_back(torch::where(m > n / 2, m - n, m)); + } + return fields; +} + +} // 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..7ca2b916f2 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa4c_long_range_ptexpt.cc @@ -0,0 +1,370 @@ +// 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, and over the images of the cell of a periodic one when +// the module serves periodic systems. Models generated by +// source/tests/infer/gen_dpa4c_long_range.py: uncompressed archives on the +// graph lower, which receives the pair list or the mesh of the cell, and +// compressed archives on the compact canonical lower, which reads the +// positions of the atoms and the cell. +#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"; +constexpr const char* kPeriodicModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.pt2"; +constexpr const char* kPeriodicRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic.expected"; +constexpr const char* kPeriodicCompressedModelPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_compressed.pt2"; +constexpr const char* kPeriodicCompressedRefPath = + "../../tests/infer/deeppot_dpa4c_long_range_periodic_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 module that serves no periodic system rejects a cell. +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); +} + +// A module that serves periodic systems evaluates a periodic system over the +// images of its cell, through the coordinates and the cell or through the +// neighbor list of a LAMMPS-style call, whose ghost atoms are images, and +// keeps the direct sum of a system without a cell. ``Archive`` selects the +// uncompressed archive on the graph lower or, on a CUDA device, the +// compressed one on the compact canonical lower. +template +class TestInferDeepPotDpa4cLongRangePeriodicPtExpt : 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 box = {13.0, 0.0, 0.0, 0.6, 12.5, + 0.0, -0.4, 0.8, 13.5}; + LongRangeRef ref_periodic; + LongRangeRef ref_cluster; + deepmd::DeepPot dp; + + void load(const char* model_path, const char* ref_path) { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::ifstream model_file(model_path); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << model_path + << " was not generated (gen_dpa4c_long_range.py writes the " + "compressed archives on a CUDA device only)."; + } + dp.init(model_path); + deepmd_test::ExpectedRef ref; + ref.load(ref_path); + ref_periodic.load(ref, "periodic"); + ref_cluster.load(ref, "cluster"); + } + + // Computes the periodic system through the LAMMPS-style entry point, with + // the ghost atoms of the cell, and folds their forces back onto their + // owners. + void compute_lmp_periodic(double& ener, + std::vector& force, + std::vector& virial) { + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector> nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, static_cast(dp.cutoff())); + const int nloc = static_cast(atype.size()); + const int nall = static_cast(atype_cpy.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); + inlist.mapping = mapping.data(); + std::vector force_all; + dp.compute(ener, force_all, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + _fold_back(force, force_all, mapping, nloc, nall, 3); + } + + void check_all(const double epsilon) { + EXPECT_TRUE(dp.has_long_range()); + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + ref_periodic.check(ener, force, virial, epsilon); + compute_lmp_periodic(ener, force, virial); + ref_periodic.check(ener, force, virial, epsilon); + dp.compute(ener, force, virial, coord, atype, std::vector()); + ref_cluster.check(ener, force, virial, epsilon); + } +}; + +TYPED_TEST_SUITE(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, graph_lower) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(std::is_same::value ? 1e-8 : 1e-4); +} + +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, compact_lower) { + this->load(kPeriodicCompressedModelPath, kPeriodicCompressedRefPath); + if (this->IsSkipped()) { + return; + } + this->check_all(1e-4); +} + +// A cell without volume marks no system: it is rejected rather than taken as +// a non-periodic one. +TYPED_TEST(TestInferDeepPotDpa4cLongRangePeriodicPtExpt, + degenerate_cell_is_rejected) { + using VALUETYPE = TypeParam; + this->load(kPeriodicModelPath, kPeriodicRefPath); + if (this->IsSkipped()) { + return; + } + double ener; + std::vector force, virial; + EXPECT_THROW(this->dp.compute(ener, force, virial, this->coord, this->atype, + std::vector(9, 0.)), + deepmd::deepmd_exception); +} 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..2ed86f2da2 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -178,6 +178,11 @@ double PairDeepMD::eval_energy_with_fparam( dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -297,6 +302,11 @@ void PairDeepMD::compute(int eflag, int vflag) { dbox[7] = domain->h[3] / dist_unit_cvt_factor; // zy dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx + // A long-range module reads the cell of a periodic system only; the empty + // box marks a non-periodic one. + if (long_range && domain->nonperiodic) { + dbox.clear(); + } // get coord for (int ii = 0; ii < nall; ++ii) { @@ -966,6 +976,24 @@ 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 a non-periodic + // system, or over the images of the cell of a system periodic in every + // dimension; the atoms of the other ranks reach it through the + // communicator. The model checks which of the two its module serves. + long_range = numb_models > 1 ? deep_pot_model_devi.has_long_range() + : deep_pot.has_long_range(); + const int periodic_dimensions = + domain->xperiodic + domain->yperiodic + domain->zperiodic; + if (long_range && periodic_dimensions != 0 && periodic_dimensions != 3) { + error->all(FLERR, + "pair style deepmd: a model with a long-range module evaluates " + "a system periodic in every dimension or in none; use p in " + "every dimension, or 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..69466f740a 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; @@ -58,6 +59,9 @@ class PairDeepMD : public PairDeepBaseModel { // geometry fields are unused) for the device-resident message-passing path, // where ghost features are exchanged across ranks inside the forward pass. deepmd_compat::InputNlist make_comm_nlist(); + // The model carries a long-range module, which reads the cell of a + // periodic system and no cell for a non-periodic one. + bool long_range = false; private: CommBrickDeepMD* commdata_; diff --git a/source/lmp/pair_deepmd_kokkos.cpp b/source/lmp/pair_deepmd_kokkos.cpp index f54873dbe1..cd9f102fff 100644 --- a/source/lmp/pair_deepmd_kokkos.cpp +++ b/source/lmp/pair_deepmd_kokkos.cpp @@ -2,6 +2,7 @@ #ifdef LMP_KOKKOS #include "pair_deepmd_kokkos.h" +#include #include #include #include @@ -212,6 +213,14 @@ void PairDeepMDKokkos::init_style() { request->enable_full(); edge_vec_fp32 = deep_pot.uses_fp32_edge_vectors(); canonical_graph = deep_pot.uses_canonical_graph_inference(); + // The device edge path passes a long-range module no cell, so a periodic + // system reaches such a module through the compact canonical lower alone. + if (long_range && !canonical_graph && !domain->nonperiodic) { + error->all(FLERR, + "pair style deepmd/kk evaluates a periodic system with a " + "long-range module only through a compressed DPA4C model frozen " + "for a CUDA device; use pair_style deepmd for this model."); + } // 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 +417,36 @@ 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(); + 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 bool has_null_types = compact_graph.has_null_types; + const double inv_dist = 1.0 / dist_unit_cvt_factor; + // Relative to the origin of the box, as the plain pair style passes them. + const double x0 = domain->boxlo[0]; + const double y0 = domain->boxlo[1]; + const double z0 = domain->boxlo[2]; + Kokkos::parallel_for( + "deepmd/kk:node_position", + Kokkos::RangePolicy(0, nnode_model), + KOKKOS_LAMBDA(const int m) { + // Without virtual atoms the model nodes are the leading atoms. + const int i = has_null_types ? model2loc(m) : m; + node_position(3 * m + 0) = (x(i, 0) - x0) * inv_dist; + node_position(3 * m + 1) = (x(i, 1) - y0) * inv_dist; + node_position(3 * m + 2) = (x(i, 2) - z0) * inv_dist; + }); + return node_position.data(); +} + template void PairDeepMDKokkos::compute(int eflag, int vflag) { if (!device_path_ok) { @@ -520,7 +559,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,17 +570,33 @@ 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()); + // A long-range module reads the cell of a periodic system, row-major with + // the lattice vectors as rows, in the length unit of the model, and no + // cell for a non-periodic one. + const double* h = domain->h; + std::array cell = {h[0], 0.0, 0.0, h[5], h[1], + 0.0, h[4], h[3], h[2]}; + for (double& entry : cell) { + entry /= dist_unit_cvt_factor; + } + const double* cell_ptr = + long_range && !domain->nonperiodic ? cell.data() : nullptr; + 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_model_type.data(), long_range ? coord_ptr : nullptr, + long_range_across_ranks ? static_cast(&world) : nullptr, + cell_ptr, 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(), diff --git a/source/lmp/pair_deepmd_kokkos.h b/source/lmp/pair_deepmd_kokkos.h index bd3689aedf..4db0f380e6 100644 --- a/source/lmp/pair_deepmd_kokkos.h +++ b/source/lmp/pair_deepmd_kokkos.h @@ -82,6 +82,11 @@ 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 relative to the origin of the box, 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 +95,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. 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..4829a91a7d --- /dev/null +++ b/source/lmp/tests/test_lammps_dpa4c_long_range_pt2.py @@ -0,0 +1,439 @@ +# 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, and over the images of the cell of a periodic one when +it serves periodic systems, so ``pair_style deepmd`` runs a model whose module +serves no periodic system with non-periodic boundaries only, and any +long-range model with boundaries periodic in every dimension or in none. 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 are generated by source/tests/infer/gen_dpa4c_long_range.py, the +compressed models for CUDA only: without a cell, at the default charge state +and at ``charge_spin 1.0 2.0``, for a module that serves no periodic system, +and in its cell and without one for a module that does. +""" + +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") +periodic_pt2_file = pt2_file.with_name("deeppot_dpa4c_long_range_periodic.pt2") +periodic_compressed_pt2_file = pt2_file.with_name( + "deeppot_dpa4c_long_range_periodic_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" +periodic_data_file = Path(__file__).parent / "data_dpa4c_long_range_periodic_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 + + +def _read_ref(path: Path) -> dict | None: + try: + return read_expected_ref(path) + except FileNotFoundError: + return None + + +compressed_ref = _read_ref(compressed_ref_file) +periodic_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic.expected") +) +periodic_compressed_ref = _read_ref( + ref_file.with_name("deeppot_dpa4c_long_range_periodic_compressed.expected") +) + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +# The triclinic cell of the periodic references: xlo xhi ylo yhi zlo zhi xy xz +# yz of the lattice vectors (13, 0, 0), (0.6, 12.5, 0) and (-0.4, 0.8, 13.5). +periodic_box = np.array([0, 13, 0, 12.5, 0, 13.5, 0.6, -0.4, 0.8]) +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) + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) + + +def teardown_module() -> None: + for path in (data_file, periodic_data_file, replicated_data_file): + if path.exists(): + os.remove(path) + + +def _lammps(boundary: str, data: Path = data_file) -> PyLammps: + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + return make_atomic_lammps(data, 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) + + +@pytest.mark.parametrize("long_range_first", [True, False]) +def test_model_deviation_without_cell_needs_long_range_throughout( + lammps, long_range_first: bool, tmp_path: Path +) -> None: + """A non-periodic model deviation passes no cell, which every model must serve.""" + 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.pair_style( + "deepmd " + + " ".join(str(model.resolve()) for model in models) + + f" out_freq 1 out_file {tmp_path / 'model_devi.out'}" + ) + lammps.pair_coeff("* *") + with pytest.raises(Exception, match="every model carries one"): + lammps.run(0) + + +def test_model_deviation_of_long_range_models(lammps, tmp_path: Path) -> None: + """Models that all carry a long-range module serve a non-periodic deviation.""" + md_file = tmp_path / "model_devi.out" + lammps.pair_style( + f"deepmd {pt2_file.resolve()} {pt2_file.resolve()} " + f"out_freq 1 out_file {md_file}" + ) + 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] + ) + md = np.loadtxt(md_file) + assert md[1:] == pytest.approx(np.zeros(md.shape[0] - 1), abs=1e-10) + + +def _check_reference(lmp: PyLammps, ref: dict, section: str) -> None: + """Compare the energy and the forces of a run with a reference section.""" + lmp.run(0) + assert lmp.eval("pe") == pytest.approx(float(np.sum(ref[section]["expected_e"]))) + expected_f = ref[section]["expected_f"].reshape(6, 3) + for ii in range(6): + assert lmp.atoms[ii].force == pytest.approx( + expected_f[lmp.atoms[ii].id - 1], abs=1e-8 + ) + + +@pytest.mark.parametrize( + ("boundary", "section"), [("p p p", "periodic"), ("f f f", "cluster")] +) +def test_pair_deepmd_periodic_module(boundary: str, section: str) -> None: + """A module that serves periodic systems sums over the images of the cell.""" + if periodic_ref is None or not periodic_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_pt2_file} was not generated.") + lmp = _lammps(boundary, periodic_data_file) + lmp.pair_style(f"deepmd {periodic_pt2_file.resolve()}") + lmp.pair_coeff("* *") + _check_reference(lmp, periodic_ref, section) + lmp.run(1) + lmp.close() + + +@pytest.mark.parametrize("model", [pt2_file, periodic_pt2_file]) +def test_partially_periodic_boundary_is_rejected(model: Path) -> None: + if not model.is_file(): + pytest.skip(f"Skip because {model} was not generated.") + lmp = _lammps("p p f", periodic_data_file) + lmp.pair_style(f"deepmd {model.resolve()}") + lmp.pair_coeff("* *") + with pytest.raises(Exception, match="every dimension"): + lmp.run(0) + lmp.close() + + +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, + ) + + +@mpi_available +@pytest.mark.parametrize( + ("nprocs", "processors", "boundary", "section"), + [ + (2, "2 1 1", "p p p", "periodic"), + # Every atom lies in the lower half along z; the ranks above hold + # ghost atoms alone. + (4, "1 1 4", "p p p", "periodic"), + (2, "1 2 1", "f f f", "cluster"), + ], +) +def test_pair_deepmd_mpi_periodic_compressed( + nprocs: int, processors: str, boundary: str, section: str +) -> None: + if periodic_compressed_ref is None or not periodic_compressed_pt2_file.is_file(): + pytest.skip(f"Skip because {periodic_compressed_pt2_file} was not generated.") + if not periodic_data_file.is_file(): + write_lmp_data(periodic_box, coord, type_OH, periodic_data_file) + result = run_mpi_pair_runner( + mpi_runner, + periodic_data_file, + periodic_compressed_pt2_file, + nprocs=nprocs, + processors=processors, + runner_args=["--boundary", boundary], + ) + ref = periodic_compressed_ref[section] + expected_v = -ref["expected_v"].reshape(6, 9) + assert result["pe"] == pytest.approx(float(np.sum(ref["expected_e"])), rel=1e-6) + np.testing.assert_allclose( + result["forces"], ref["expected_f"].reshape(6, 3), 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 long-range 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..9bbda9d61d 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,10 @@ 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 + long_range/latent_charge_sum.cu + long_range/mesh_sum.cu dpa1_graph_energy_force.cu dpa4/so2_conv.cu ${DPA4_SO2_CONV_KERNEL_SRC} @@ -213,7 +218,8 @@ if(USE_CUDA_TOOLKIT AND DEEPMD_TORCH_HAS_CUDA) set_property(TARGET deepmd_op_pt PROPERTY CUDA_ARCHITECTURES ${DEEPMD_PT_CUDA_ARCHITECTURES}) endif() - target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas CUDA::cublasLt) + target_link_libraries(deepmd_op_pt PRIVATE CUDA::cublas CUDA::cublasLt + CUDA::cufft) # libtorch headers require at least C++17, and PyTorch 2.14+ requires C++20. # Preserve an explicitly higher CUDA standard requested by the user. set(DEEPMD_PT_CUDA_STANDARD 17) @@ -245,6 +251,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/charge_sum.h b/source/op/pt/long_range/charge_sum.h new file mode 100644 index 0000000000..04245d3287 --- /dev/null +++ b/source/op/pt/long_range/charge_sum.h @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// The two evaluations behind deepmd::latent_charge_sum (schema and formulas in +// ops.cc): the direct sum over the pairs of non-periodic frames +// (direct_sum.cu) and the smooth particle-mesh Ewald sum of periodic frames +// (mesh_sum.cu). Both take float64 positions and float32 raw charges of the +// nodes of the frames on one CUDA device, apply the constraint, and evaluate +// the owner nodes [owner_begin, owner_begin + owner_count) against every node +// of their frames. Owners other than every node lie in a single frame, whose +// nodes are then the charged nodes of every domain of ``communicator``, over +// which the per-frame sums of the owners are added. + +#pragma once + +#include + +#include +#include + +#include "domain_exchange.h" + +namespace deepmd { +namespace long_range { + +// The atomic energy (owners, 1) float64, the charge gradient (owners, C), the +// force (owners, 3), the atomic virial (owners, 9) and the virial of the +// owners of every frame (F, 9), float32. +using ChargeSumOutputs = std::tuple; + +// The direct sum. ``coefficient`` holds the float64 coefficients of the +// kernel on the device. +ChargeSumOutputs direct_sum(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); + +// The mesh sum. ``box`` holds the cell of every frame, float64 on the host; +// ``coefficient`` the float64 coefficients of the kernel on the device. +ChargeSumOutputs mesh_sum(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& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator); + +} // namespace long_range +} // namespace 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..2e0f128d85 --- /dev/null +++ b/source/op/pt/long_range/direct_sum.cu @@ -0,0 +1,979 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the latent-charge long range of non-periodic frames, by the +// direct sum over their pairs (schema and formulas in ops.cc, the entry in +// latent_charge_sum.cu). +// +// 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, which the entry gathers with those of every rank into one frame; +// 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 "charge_sum.h" +#include "frame_sums.cuh" + +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 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; + } +}; + +// 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 double* __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 : *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 double* __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 = static_cast(coefficient[k]); + const float rate = static_cast(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 double* __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 : *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 float4* head; + const float4* tail; + const float* charge; + const int* frame; + const int64_t* frame_ptr; + int node_count; + const double* coefficient; + int gaussian_count; + float* potential; + float* force; + float* 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, buffers.node_count, + resident_blocks_per_multiprocessor()); + auto partial = + torch::empty({static_cast(chunk_count) * owner_count * kWidth}, + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA, c10::cuda::current_device())); + const dim3 grid((owner_count + kThreads - 1) / kThreads, chunk_count); + direct_sum_kernel<<>>( + buffers.head, buffers.tail, buffers.charge, buffers.frame, + buffers.frame_ptr, owner_begin, owner_count, buffers.coefficient, + buffers.gaussian_count, 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 : nullptr, buffers.potential, + buffers.force, buffers.atom_virial); + 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_sum: unsupported channel count ", + channels); + } +} + +} // namespace + +namespace deepmd { +namespace long_range { + +// 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. +ChargeSumOutputs direct_sum(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_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_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(); + 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 int owner_slice_count = slices_for(owner_count); + const int components = constrained ? 10 : 9; + Workspace workspace; + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t head_at = workspace.reserve(node_count); + const int64_t tail_at = workspace.reserve(node_count); + const int64_t frame_at = workspace.reserve(node_count); + const int64_t pair_charge_at = + workspace.reserve(int64_t{node_count} * channels); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t slice_virial_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * components); + workspace.allocate(coord.device()); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* head = workspace.at(head_at); + auto* tail = workspace.at(tail_at); + auto* frame = workspace.at(frame_at); + auto* pair_charge = workspace.at(pair_charge_at); + auto* potential = workspace.at(potential_at); + auto* frame_stat = + constrained ? workspace.at(frame_stat_at) : nullptr; + auto* slice_charge = workspace.at(slice_charge_at); + auto* slice_virial = workspace.at(slice_virial_at); + + const auto stream = at::cuda::getCurrentCUDAStream(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (constrained) { + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, + static_cast(channels), constrained_channel, slice_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // Alpha stays on the device, so that a call never waits for it. + const double* alpha = + kernel == kLes ? coefficient.data_ptr() : nullptr; + prepare_nodes<<<(node_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>(coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frame_count, + node_count, static_cast(channels), + constrained_channel, alpha, frame_stat, head, tail, + frame, pair_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (owner_count > 0) { + const DirectSumBuffers buffers{head, + tail, + pair_charge, + frame, + frame_ptr, + node_count, + coefficient.data_ptr(), + static_cast(coefficient.numel() / 2), + potential, + force.data_ptr(), + atom_virial.data_ptr()}; + 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_tensor; + const int64_t* owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const dim3 owner_slice_grid(frame_count, owner_slice_count); + if (constrained) { + virial_slices<<>>( + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial, owner_slice_count, virial.data_ptr(), frame_stat); + } else { + virial_slices<<>>( + atom_virial.data_ptr(), potential, owner_mask, owner_frame_ptr, + static_cast(channels), constrained_channel, slice_virial); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + virial_total<<>>( + slice_virial, 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 stat = + torch::from_blob(frame_stat, {frame_count, kStatWidth}, double_options); + auto potential_sum = stat.select(1, kStatPotential); + potential_sum.copy_( + sum_over_domains(*communicator, potential_sum.contiguous())); + } + if (owner_count > 0) { + project_charges<<<(owner_count + kThreads - 1) / kThreads, kThreads, 0, + stream>>>( + potential, pair_charge + int64_t{owner_begin} * channels, owner_mask, + frame + owner_begin, frame_stat, 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}; +} + +} // namespace long_range +} // namespace deepmd 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..2ae3d06e05 --- /dev/null +++ b/source/op/pt/long_range/domain_exchange.cc @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// MPI exchanges of the latent-charge long range (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_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_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_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..f423974f26 --- /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 latent-charge long range between the domains of a +// decomposed system. +// +// A domain is one MPI rank of the communicator that a caller hands to +// deepmd::latent_charge_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/frame_sums.cuh b/source/op/pt/long_range/frame_sums.cuh new file mode 100644 index 0000000000..d3f8d40af6 --- /dev/null +++ b/source/op/pt/long_range/frame_sums.cuh @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Per-frame bookkeeping shared by the CUDA units of the latent-charge long +// range: the workspace of a call, the offsets of the frames on the node axis, +// the sums over the nodes of every frame in a fixed order, and the statistics +// of the charge constraint. +// +// The kernels are defined in an unnamed namespace, so that every unit that +// includes this header compiles its own copy. + +#pragma once + +#include +#include + +#include +#include + +namespace { + +// The intermediate buffers of one call, carved out of a single device +// allocation: ``reserve`` lays out a buffer of ``count`` elements and returns +// its byte offset, and ``at`` its address once ``allocate`` has made the +// allocation. Every buffer starts on a 256-byte boundary. +class Workspace { + public: + template + int64_t reserve(int64_t count) { + bytes_ = (bytes_ + kAlignment - 1) / kAlignment * kAlignment; + const int64_t offset = bytes_; + bytes_ += count * static_cast(sizeof(T)); + return offset; + } + + void allocate(const torch::Device& device) { + storage_ = torch::empty( + {std::max(bytes_, 1)}, + torch::TensorOptions().dtype(torch::kUInt8).device(device)); + } + + template + T* at(int64_t offset) const { + return reinterpret_cast(storage_.data_ptr() + offset); + } + + private: + static constexpr int64_t kAlignment = 256; + int64_t bytes_ = 0; + torch::Tensor storage_; +}; + +// 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; +} + +} // namespace diff --git a/source/op/pt/long_range/latent_charge_sum.cu b/source/op/pt/long_range/latent_charge_sum.cu new file mode 100644 index 0000000000..4763c72fcc --- /dev/null +++ b/source/op/pt/long_range/latent_charge_sum.cu @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA entry of deepmd::latent_charge_sum (schema and formulas in ops.cc). +// +// A call evaluates its frames by the direct sum when they carry no cell and +// by the mesh sum when they do (charge_sum.h). Across the domains of a +// decomposed system, the ranks 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 charged nodes against that frame; the outputs +// return to the node axis of the rank, where the nodes it does not own carry +// no charge. + +#include +#include +#include + +#include +#include +#include + +#include "charge_sum.h" +#include "domain_exchange.h" + +namespace { + +using deepmd::long_range::ChargeSumOutputs; + +constexpr int kMaxChannels = 8; +constexpr int kMaxGaussians = 32; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// The evaluation of the frames, by the direct or the mesh sum, of the owner +// nodes [owner_begin, owner_begin + owner_count). ``box`` is defined for +// periodic frames only. +struct Evaluation { + const torch::Tensor& box; + const torch::Tensor& coefficient; + int64_t kernel; + int64_t channel; + double prefactor; + int64_t order; + double spacing; + bool periodic; + + ChargeSumOutputs operator()(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, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) const { + if (periodic) { + return deepmd::long_range::mesh_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, box, + coefficient, kernel, channel, prefactor, order, spacing, owner_begin, + owner_count, communicator); + } + return deepmd::long_range::direct_sum( + coord, charge, node_mask, n_node_per_frame, total_charge, coefficient, + kernel, channel, prefactor, owner_begin, owner_count, communicator); + } +}; + +// 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. +ChargeSumOutputs evaluate_across_domains(const Evaluation& evaluation, + 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& communicator) { + TORCH_CHECK(n_node_per_frame.size(0) == 1, + "latent_charge_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] = + evaluation(partner_coord, partner_charge, partner_mask, partner_frame, + total_charge, 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}; +} + +// Whether the frames carry cells: true when ``box`` is given and none of its +// rows is zero, false when it is absent or all of its rows are; a call that +// mixes the two is rejected. +bool periodic_frames(const c10::optional& box) { + if (!box.has_value()) { + return false; + } + const auto cells = box->contiguous(); + const double* value = cells.data_ptr(); + const int64_t frame_count = cells.size(0); + int64_t periodic = 0; + for (int64_t f = 0; f < frame_count; ++f) { + periodic += std::any_of(value + 9 * f, value + 9 * (f + 1), + [](double entry) { return entry != 0.0; }); + } + TORCH_CHECK(periodic == 0 || periodic == frame_count, + "latent_charge_sum: every frame of a call is periodic, or none " + "is"); + return periodic > 0; +} + +ChargeSumOutputs latent_charge_sum(torch::Tensor coord, + torch::Tensor charge, + torch::Tensor node_mask, + torch::Tensor n_node_per_frame, + torch::Tensor total_charge, + c10::optional box, + torch::Tensor coefficient, + torch::Tensor communicator, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing) { + 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_sum: inputs must be CUDA tensors, except box and " + "communicator"); + 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_sum: inputs must share one device"); + TORCH_CHECK(coord.dim() == 2 && coord.size(1) == 3 && + coord.scalar_type() == torch::kFloat64, + "latent_charge_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_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_sum: node_mask must be bool with shape (N,)"); + TORCH_CHECK(n_node_per_frame.dim() == 1, + "latent_charge_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_sum: total_charge must be float64 with one entry " + "per frame"); + TORCH_CHECK( + !box.has_value() || + (box->device().is_cpu() && box->scalar_type() == torch::kFloat64 && + box->dim() == 2 && box->size(0) == n_node_per_frame.size(0) && + box->size(1) == 9), + "latent_charge_sum: box must be a float64 CPU tensor with shape " + "(F, 9)"); + TORCH_CHECK( + coefficient.dim() == 1 && coefficient.scalar_type() == torch::kFloat64, + "latent_charge_sum: coefficient must be a float64 vector"); + const int64_t channels = charge.size(1); + TORCH_CHECK(channels >= 1 && channels <= kMaxChannels, + "latent_charge_sum: the channel count must lie in [1, ", + kMaxChannels, "], got ", channels); + TORCH_CHECK(channel >= -1 && channel < channels, + "latent_charge_sum: the constrained channel must be -1 or lie " + "in [0, ", + channels, "), got ", channel); + if (kernel == kLes) { + TORCH_CHECK(coefficient.numel() == 1, + "latent_charge_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_sum: the Gaussian kernel takes two " + "coefficients per Gaussian, for at most ", + kMaxGaussians, " Gaussians"); + } else { + TORCH_CHECK(false, "latent_charge_sum: unknown kernel ", kernel); + } + TORCH_CHECK(coord.size(0) < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + + const c10::cuda::CUDAGuard guard(coord.device()); + coord = coord.contiguous(); + charge = charge.contiguous(); + node_mask = node_mask.contiguous(); + total_charge = total_charge.contiguous(); + const bool periodic = periodic_frames(box); + const torch::Tensor cells = periodic ? box->contiguous() : torch::Tensor(); + coefficient = coefficient.contiguous(); + const Evaluation evaluation{cells, coefficient, kernel, channel, + prefactor, order, spacing, periodic}; + if (deepmd::long_range::domain_of(communicator).size == 1) { + return evaluation(coord, charge, node_mask, n_node_per_frame, total_charge, + 0, static_cast(coord.size(0)), nullptr); + } + return evaluate_across_domains(evaluation, coord, charge, node_mask, + n_node_per_frame, total_charge, communicator); +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(deepmd, library) { + library.impl("latent_charge_sum", torch::kCUDA, &latent_charge_sum); +} diff --git a/source/op/pt/long_range/mesh_sum.cu b/source/op/pt/long_range/mesh_sum.cu new file mode 100644 index 0000000000..130cbdcb11 --- /dev/null +++ b/source/op/pt/long_range/mesh_sum.cu @@ -0,0 +1,1375 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// CUDA kernels of the latent-charge long range of periodic frames, by the +// smooth particle-mesh Ewald method (schema and formulas in ops.cc, the +// evaluation in deepmd/dpmodel/long_range/reciprocal.py). +// +// Every frame owns a mesh of n1 x n2 x n3 points spanning its cell. The +// charges spread to the mesh with cardinal B-splines of order p along the +// fractional coordinates s = r B^-1, the lattice vectors being the rows of the +// cell B: the scaled coordinate u_d = n_d s_d reaches the p points +// floor(u_d) - p + 1, ..., floor(u_d) with the weights M_p(u_d - m). A real +// FFT gives the half spectrum Q(m) of the charge mesh, and the influence +// function +// +// G(m) = prefactor * K^(k_m) / V * prod_d |D_d(m_d)|^-2, k_m = 2 pi B^-1 m, +// +// with the Fourier transform K^ of the kernel and the assignment factor +// D_d(m) = sum_l M_p(l + 1) exp(2 pi i m l / n_d), turns it into the +// potential mesh, the unnormalized inverse transform of G Q. The potential at +// a node is its interpolation by the B-splines of the node, less the potential +// prefactor * K(0) q of the charge at itself, and the force at fixed charges +// the negative gradient of that interpolation times the charges. The zero mode +// and the modes at the Nyquist frequency of an axis, which have no unique wave +// number, carry no weight. +// +// The virial of a frame is that of the strain of its cell, under which the +// fractional coordinates stay put while k_m and the volume change: +// +// W = sum_m e_m (1 + 2 dln K^ / d(k^2) k_m (x) k_m), e_m = G |Q|^2 / 2, +// +// the sum over the full spectrum, every interior plane of the half spectrum +// counting twice. Its per-node share is F_i (x) r_i, the force index first, +// plus the rest of W shared evenly over the charged nodes. +// +// The charges spread in 64-bit fixed point, so that the mesh, and every +// output, does not depend on the order of the atomic additions. The sums over +// the modes and over the nodes of a frame add float64 sums in a fixed order. +// The meshes, the transforms and the influence function are float32, against +// a float64 throughput of a small fraction of the float32 one on many +// devices; the B-spline weights are evaluated in float64, whose rounding to +// float32 leaves their sum unbiased. The forces then carry the rounding of +// float32, and the energy of a frame that of the reciprocal sum including the +// interaction of every charge with itself, which the self term cancels: a few +// 1e-7 of the self energy. +// +// The nodes whose outputs a call evaluates, its owners, are either every node +// or, across the domains of a decomposed system, the charged nodes of one MPI +// rank, whose frame then holds the charged nodes of every rank. Every rank +// spreads every charge and transforms the whole mesh, and interpolates only at +// its owners; the per-frame sums of the owners are added over the domains. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "charge_sum.h" +#include "frame_sums.cuh" + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxGaussians = 32; +constexpr int kMaxOrder = 12; + +enum KernelKind : int { kLes = 0, kSog = 1 }; + +// Resolution of the fixed-point charge mesh: 2^-40 e, far below the float32 +// precision of the raw charges, for a sum of at most 2^23 e at a point. +constexpr float kFixedScaleFloat = 1099511627776.f; // 2^40 + +constexpr double kPi = 3.14159265358979323846; + +// The cell of one frame and its mesh. +struct FrameMesh { + // B^-1, row-major: s_d = sum_j r_j inverse[3 j + d]. + double inverse[9]; + double volume; + // -n_d (B^-1)_jd, row-major in (j, d): the force at fixed charges from the + // derivatives of the energy with respect to the scaled coordinates u_d. + float gradient_to_force[9]; + int n[3]; + // The first point of the meshes of the frame, channel-major, in the mesh + // buffer of the call. + int64_t offset; +}; + +// The smallest even number of at least ``n`` without a prime factor above 7. +int64_t even_fft_size(int64_t n) { + int64_t m = std::max(2, n + n % 2); + while (true) { + int64_t rest = m; + for (const int64_t p : {2, 3, 5, 7}) { + while (rest % p == 0) { + rest /= p; + } + } + if (rest == 1) { + return m; + } + m += 2; + } +} + +FrameMesh frame_mesh(const double* cell, double spacing) { + FrameMesh mesh{}; + const double a = cell[0], b = cell[1], c = cell[2]; + const double d = cell[3], e = cell[4], f = cell[5]; + const double g = cell[6], h = cell[7], i = cell[8]; + const double det = + a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g); + // A volume that is a vanishing fraction of the product of the lengths of + // the lattice vectors, or not finite, spans no lattice. + const double lengths = std::sqrt(a * a + b * b + c * c) * + std::sqrt(d * d + e * e + f * f) * + std::sqrt(g * g + h * h + i * i); + TORCH_CHECK(std::abs(det) > 1e-6 * lengths, + "latent_charge_sum: a periodic cell is singular"); + const double adjugate[9] = {e * i - f * h, c * h - b * i, b * f - c * e, + f * g - d * i, a * i - c * g, c * d - a * f, + d * h - e * g, b * g - a * h, a * e - b * d}; + for (int k = 0; k < 9; ++k) { + mesh.inverse[k] = adjugate[k] / det; + } + mesh.volume = std::abs(det); + // The squared lengths of the lattice vectors are formed term by term, as + // the Python front end forms them, so that both round the mesh size + // identically. + for (int axis = 0; axis < 3; ++axis) { + const double* v = cell + 3 * axis; + const double length = std::sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + const double points = std::ceil(length / spacing); + TORCH_CHECK(points < 1 << 20, + "latent_charge_sum: the mesh of a cell exceeds 2^20 points " + "along a lattice vector"); + mesh.n[axis] = + static_cast(even_fft_size(static_cast(points))); + } + TORCH_CHECK(int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2] < + std::numeric_limits::max(), + "latent_charge_sum: the mesh of a cell exceeds 2^31 points"); + for (int j = 0; j < 3; ++j) { + for (int d = 0; d < 3; ++d) { + mesh.gradient_to_force[3 * j + d] = + static_cast(-mesh.n[d] * mesh.inverse[3 * j + d]); + } + } + return mesh; +} + +__host__ __device__ int64_t mesh_points(const FrameMesh& mesh) { + return int64_t{mesh.n[0]} * mesh.n[1] * mesh.n[2]; +} + +// Nodes per block of the kernels that spread and interpolate, one warp each. +constexpr int kStencilWarps = 4; +constexpr int kWarp = 32; + +// The stencil of one node along one axis of its mesh: for each of its p +// points on the axis, its offset on the flat mesh of a frame along that axis, +// the index of the point times the stride of the axis, and its B-spline +// weight and slope, evaluated in float64 and stored in float32, the precision +// of the outputs: rounding a recursion of float32 steps instead would bias +// the sum of the weights. +template +struct AxisStencil { + int offset[kOrder]; + float value[kOrder]; + float slope[kOrder]; +}; + +// The stencil of a node along axis ``d``. The scaled coordinate u = n_d s_d +// of the fractional coordinate s_d reaches the points floor(u) - p + 1 + k, +// k = 0, ..., p - 1, with the weights value[k] = M_p(t + p - 1 - k), t the +// fractional part of u, and slope[k] their derivative with respect to u. The +// recursion M_n(t + j) = ((t + j) M_{n-1}(t + j) + (n - t - j) M_{n-1}(t + j - +// 1)) / (n - 1) builds a_j = M_n(t + j) from M_2; the slope follows from +// M_p'(x) = M_{p-1}(x) - M_{p-1}(x - 1). +template +__device__ void axis_stencil(const double* coord, + const FrameMesh& mesh, + int d, + AxisStencil& stencil) { + double s = coord[0] * mesh.inverse[d] + coord[1] * mesh.inverse[3 + d] + + coord[2] * mesh.inverse[6 + d]; + s -= floor(s); + const double u = s * mesh.n[d]; + const double base = floor(u); + const double t = u - base; + double a[kOrder]; + a[0] = t; + a[1] = 1.0 - t; +#pragma unroll + for (int n = 3; n <= kOrder; ++n) { + if (n == kOrder) { + // The weights of order p - 1 give the slopes. +#pragma unroll + for (int j = 0; j < kOrder; ++j) { + const double high = j < kOrder - 1 ? a[j] : 0.0; + const double low = j > 0 ? a[j - 1] : 0.0; + stencil.slope[kOrder - 1 - j] = static_cast(high - low); + } + } + const double scale = 1.0 / (n - 1); + a[n - 1] = scale * (1.0 - t) * a[n - 2]; +#pragma unroll + for (int j = n - 2; j > 0; --j) { + a[j] = scale * ((t + j) * a[j] + (n - t - j) * a[j - 1]); + } + a[0] = scale * t * a[0]; + } + if constexpr (kOrder == 2) { + stencil.slope[0] = -1.f; + stencil.slope[1] = 1.f; + } + const int n = mesh.n[d]; + const int stride = d == 0 ? mesh.n[1] * mesh.n[2] : d == 1 ? mesh.n[2] : 1; + // A stencil wider than the axis wraps more than once. + const int first = ((static_cast(base) - kOrder + 1) % n + n) % n; +#pragma unroll + for (int k = 0; k < kOrder; ++k) { + stencil.value[k] = static_cast(a[kOrder - 1 - k]); + stencil.offset[k] = (first + k) % n * stride; + } +} + +// Writes, one thread per node and axis, the stencils of the charged nodes +// along the three axes of their meshes, and, on the thread of the first axis, +// the charges of every node, shifted in the constrained channel and zero on +// the nodes that carry none. Every float64 warp instruction occupies the +// float64 pipe of the device whatever the number of its active lanes, so the +// float64 work of a node runs here with every lane busy rather than in the +// warps of the kernels that spread and interpolate. +template +__global__ void __launch_bounds__(kThreads) + node_stencils(const double* __restrict__ coord, + const float* __restrict__ charge, + const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const double* __restrict__ frame_stat, + int frame_count, + int64_t node_count, + int channels, + int channel, + AxisStencil* __restrict__ stencils, + float* __restrict__ node_charge) { + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < 3 * node_count; index += int64_t{gridDim.x} * blockDim.x) { + const int64_t node = index / 3; + const int axis = static_cast(index - 3 * node); + const int f = frame_of(frame_ptr, frame_count, node); + const bool charged = f >= 0 && node_mask[node]; + if (charged) { + axis_stencil(coord + 3 * node, frames[f], axis, stencils[index]); + } + if (axis == 0) { + const double shift = charged && channel >= 0 + ? frame_stat[f * kStatWidth + kStatShift] + : 0.0; + for (int c = 0; c < channels; ++c) { + const double raw = charged ? charge[node * channels + c] : 0.0; + node_charge[node * channels + c] = + static_cast(c == channel ? raw - shift : raw); + } + } + } +} + +// Rounds ``x`` to the nearest integer, ties to even, as __float2ll_rn does, +// in integer operations: the conversion instruction occupies the float64 pipe. +// A finite x is a 24-bit significand times 2^exponent; a negative exponent +// shifts the significand right with rounding. Magnitudes of 2^63 and beyond, +// infinities and NaN saturate. +__device__ __forceinline__ long long float_to_fixed(float x) { + const unsigned bits = __float_as_uint(x); + const int biased = static_cast((bits >> 23) & 0xffu); + // |x| < 1/2, zero and subnormal numbers included, rounds to zero. + if (biased < 126) { + return 0; + } + const int exponent = biased - 150; + if (exponent >= 40) { + return bits >> 31 ? std::numeric_limits::min() + : std::numeric_limits::max(); + } + const long long significand = (bits & 0x7fffffu) | 0x800000u; + long long magnitude; + if (exponent >= 0) { + magnitude = significand << exponent; + } else { + const int shift = -exponent; + const long long kept = significand >> shift; + const long long rest = significand & ((1ll << shift) - 1); + const long long half = 1ll << (shift - 1); + magnitude = kept + (rest > half || (rest == half && (kept & 1)) ? 1 : 0); + } + return bits >> 31 ? -magnitude : magnitude; +} + +// The assignment of the p^3 points of a stencil to the lanes of a warp: the +// lanes take the (k2, k3) pairs, k3 running fastest, so that a warp touches +// contiguous runs of the mesh, and every lane walks the k1 axis for its +// pairs; while the pairs are fewer than the lanes, groups of lanes split the +// k1 axis. +template +struct LaneSlots { + static constexpr int kPairs = kOrder * kOrder; + static constexpr int kGroups = kPairs >= kWarp ? 1 : kWarp / kPairs; + static constexpr int kSlots = kGroups * kPairs; + static constexpr int kRounds = (kSlots + kWarp - 1) / kWarp; +}; + +// Copies the stencils of a node along its three axes to the shared memory of +// its warp. +template +__device__ __forceinline__ void load_stencil( + const AxisStencil* __restrict__ stencils, + int64_t node, + int lane, + AxisStencil* shared) { + constexpr int kWords = 3 * sizeof(AxisStencil) / sizeof(int); + const int* source = reinterpret_cast(stencils + 3 * node); + int* target = reinterpret_cast(shared); + for (int word = lane; word < kWords; word += kWarp) { + target[word] = source[word]; + } + __syncwarp(); +} + +// Spreads the charges of the charged nodes to the meshes of their frames, one +// warp per node, one pass per channel. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + spread_charges(const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const float* __restrict__ node_charge, + int frame_count, + int64_t node_count, + int channels, + unsigned long long* __restrict__ mesh) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + const AxisStencil* axes = shared[warp]; + for (int64_t node = int64_t{blockIdx.x} * kStencilWarps + warp; + node < node_count; node += int64_t{gridDim.x} * kStencilWarps) { + const int f = frame_of(frame_ptr, frame_count, node); + if (f < 0 || !node_mask[node]) { + continue; + } + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + for (int c = 0; c < channels; ++c) { + const float q = node_charge[node * channels + c] * kFixedScaleFloat; + unsigned long long* channel_mesh = + mesh + grid.offset + c * mesh_points(grid); +#pragma unroll + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float weight = q * axes[1].value[k2] * axes[2].value[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; +#pragma unroll + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + atomicAdd(channel_mesh + axes[0].offset[k1] + column, + static_cast( + float_to_fixed(weight * axes[0].value[k1]))); + } + } + } + } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); + } +} + +// Converts the fixed-point mesh to the float32 mesh that the transforms read. +__global__ void fixed_to_float(int64_t count, + const unsigned long long* __restrict__ fixed, + float* __restrict__ mesh) { + constexpr float kFixedInverse = 1.f / kFixedScaleFloat; + for (int64_t index = blockIdx.x * int64_t{blockDim.x} + threadIdx.x; + index < count; index += int64_t{gridDim.x} * blockDim.x) { + mesh[index] = + __ll2float_rn(static_cast(fixed[index])) * kFixedInverse; + } +} + +// The weights M_p(1), ..., M_p(p - 1) of the assignment factor. +struct IntegerWeights { + double value[kMaxOrder - 1]; +}; + +// The inverse squared assignment factor |D_d(m)|^-2 along every axis of the +// half spectrum of one frame: n1 entries, then n2, then n3 / 2 + 1. +__global__ void assignment_factor(FrameMesh mesh, + int order, + IntegerWeights integer_weight, + float* __restrict__ factor) { + const int half = mesh.n[2] / 2 + 1; + const int total = mesh.n[0] + mesh.n[1] + half; + for (int index = blockIdx.x * blockDim.x + threadIdx.x; index < total; + index += gridDim.x * blockDim.x) { + int axis = 0; + int m = index; + if (m >= mesh.n[0]) { + m -= mesh.n[0]; + axis = 1; + if (m >= mesh.n[1]) { + m -= mesh.n[1]; + axis = 2; + } + } + const double theta = 2.0 * kPi * m / mesh.n[axis]; + double re = 0.0; + double im = 0.0; + for (int l = 0; l < order - 1; ++l) { + double s; + double c; + sincos(theta * l, &s, &c); + re += integer_weight.value[l] * c; + im += integer_weight.value[l] * s; + } + factor[index] = static_cast(1.0 / (re * re + im * im)); + } +} + +// The terms of the transform of the kernel, in float32, the precision of the +// transforms: 1 / (4 alpha^2) of the les kernel, or, for every Gaussian of the +// sog kernel, its weight a_g (2 pi)^(3/2) sigma_g^3 and half its squared +// width, from ``coefficient``, which holds alpha, or the amplitudes and then +// the rates c_g = log2(e) / (2 sigma_g^2). +struct KernelTerms { + float inv_four_alpha_sq; + float weight[kMaxGaussians]; + float half_width_sq[kMaxGaussians]; +}; + +__device__ void load_kernel_terms(int kernel, + const double* __restrict__ coefficient, + int gaussians, + KernelTerms& terms) { + if (kernel == kLes) { + if (threadIdx.x == 0) { + const double alpha = coefficient[0]; + terms.inv_four_alpha_sq = static_cast(0.25 / (alpha * alpha)); + } + return; + } + // (2 pi)^{3/2} + constexpr double kGaussianNorm = 15.749609945722419; + for (int g = threadIdx.x; g < gaussians; g += blockDim.x) { + // sigma^2 = log2(e) / (2 c) = 1 / (2 c ln 2) + const double width_sq = 0.72134752044448170 / coefficient[gaussians + g]; + terms.weight[g] = static_cast(coefficient[g] * kGaussianNorm * + width_sq * sqrt(width_sq)); + terms.half_width_sq[g] = static_cast(0.5 * width_sq); + } +} + +// The transform of the kernel at k^2 and its logarithmic derivative with +// respect to k^2. +__device__ __forceinline__ void kernel_transform(int kernel, + const KernelTerms& terms, + int gaussians, + float k_sq, + float& value, + float& log_slope) { + if (kernel == kLes) { + value = 4.f * static_cast(kPi) * + expf(-k_sq * terms.inv_four_alpha_sq) / k_sq; + log_slope = -terms.inv_four_alpha_sq - 1.f / k_sq; + return; + } + value = 0.f; + float slope = 0.f; + for (int g = 0; g < gaussians; ++g) { + const float term = terms.weight[g] * expf(-terms.half_width_sq[g] * k_sq); + value += term; + slope -= terms.half_width_sq[g] * term; + } + log_slope = value != 0.f ? slope / value : 0.f; +} + +// The influence function of one frame on its half spectrum: multiplies every +// mode of the spectra of the channels by G in place and sums, per block, the +// energy of the modes and their virial, xx yy zz xy xz yz. The modes are +// evaluated in float32, the precision of the transforms, summed in float32 +// over the modes of a row that a lane takes, and in float64 over the rows. +constexpr int kModeSums = 7; + +__global__ void __launch_bounds__(kThreads) + apply_influence(FrameMesh mesh, + int channels, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + const float* __restrict__ factor, + float2* __restrict__ spectrum, + double* __restrict__ block_sum) { + __shared__ KernelTerms terms; + load_kernel_terms(kernel, coefficient, gaussians, terms); + __syncthreads(); + const int n1 = mesh.n[0]; + const int n2 = mesh.n[1]; + const int n3 = mesh.n[2]; + const int half = n3 / 2 + 1; + const int modes = n1 * n2 * half; + // k_a = 2 pi sum_d (B^-1)_ad m_d + float reciprocal[9]; +#pragma unroll + for (int k = 0; k < 9; ++k) { + reciprocal[k] = static_cast(2.0 * kPi * mesh.inverse[k]); + } + const float scale = static_cast(prefactor / mesh.volume); + // Every warp takes rows (i1, i2) of the half spectrum, its lanes the modes + // of a row, which lie contiguously. + const int lane = threadIdx.x % kWarp; + const int warps = gridDim.x * (kThreads / kWarp); + double sum[kModeSums] = {}; + for (int row = blockIdx.x * (kThreads / kWarp) + threadIdx.x / kWarp; + row < n1 * n2; row += warps) { + float row_sum[kModeSums] = {}; + const int i1 = row / n2; + const int i2 = row - i1 * n2; + const int m1 = i1 <= n1 / 2 ? i1 : i1 - n1; + const int m2 = i2 <= n2 / 2 ? i2 : i2 - n2; + const bool row_weightless = 2 * m1 == n1 || 2 * m2 == n2; + float k_row[3]; +#pragma unroll + for (int a = 0; a < 3; ++a) { + k_row[a] = reciprocal[3 * a] * m1 + reciprocal[3 * a + 1] * m2; + } + const float row_factor = scale * factor[i1] * factor[n1 + i2]; + for (int i3 = lane; i3 < half; i3 += kWarp) { + const int m3 = i3; + const bool weightless = + row_weightless || 2 * m3 == n3 || (m1 == 0 && m2 == 0 && m3 == 0); + float influence = 0.f; + float k[3] = {0.f, 0.f, 0.f}; + float log_slope = 0.f; + if (!weightless) { +#pragma unroll + for (int a = 0; a < 3; ++a) { + k[a] = k_row[a] + reciprocal[3 * a + 2] * m3; + } + const float k_sq = k[0] * k[0] + k[1] * k[1] + k[2] * k[2]; + float transform; + kernel_transform(kernel, terms, gaussians, k_sq, transform, log_slope); + influence = row_factor * transform * factor[n1 + n2 + i3]; + } + // An interior plane of the half spectrum stands for its mirror as well. + const float multiplicity = m3 == 0 ? 1.f : 2.f; + const int index = row * half + i3; + float energy = 0.f; + for (int c = 0; c < channels; ++c) { + float2& value = spectrum[int64_t{c} * modes + index]; + energy += value.x * value.x + value.y * value.y; + value.x *= influence; + value.y *= influence; + } + energy *= 0.5f * multiplicity * influence; + const float twice = 2.f * log_slope * energy; + row_sum[0] += energy; + row_sum[1] += energy + twice * k[0] * k[0]; + row_sum[2] += energy + twice * k[1] * k[1]; + row_sum[3] += energy + twice * k[2] * k[2]; + row_sum[4] += twice * k[0] * k[1]; + row_sum[5] += twice * k[0] * k[2]; + row_sum[6] += twice * k[1] * k[2]; + } +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + sum[s] += row_sum[s]; + } + } + __shared__ double partial[kModeSums][kThreads]; + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = sum[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x < kModeSums) { + block_sum[blockIdx.x * kModeSums + threadIdx.x] = partial[threadIdx.x][0]; + } +} + +// Adds the block sums of apply_influence into the virial of the frame, +// row-major: every thread sums a fixed strided subset of the blocks, and the +// block adds the thread sums in a fixed order. +__global__ void __launch_bounds__(kThreads) + mode_virial(const double* __restrict__ block_sum, + int block_count, + double* __restrict__ virial) { + double value[kModeSums] = {}; + for (int block = threadIdx.x; block < block_count; block += kThreads) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + value[s] += block_sum[block * kModeSums + s]; + } + } + __shared__ double partial[kModeSums][kThreads]; +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] = value[s]; + } + __syncthreads(); + for (int stride = kThreads / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { +#pragma unroll + for (int s = 0; s < kModeSums; ++s) { + partial[s][threadIdx.x] += partial[s][threadIdx.x + stride]; + } + } + __syncthreads(); + } + // Components 1..6 are xx yy zz xy xz yz; the energy is not stored. + constexpr int kEntry[kModeSums][2] = {{-1, -1}, {0, 0}, {4, 4}, {8, 8}, + {1, 3}, {2, 6}, {5, 7}}; + if (threadIdx.x > 0 && threadIdx.x < kModeSums) { + virial[kEntry[threadIdx.x][0]] = partial[threadIdx.x][0]; + virial[kEntry[threadIdx.x][1]] = partial[threadIdx.x][0]; + } +} + +// The contact value K(0) of the kernel. +__device__ double contact_value(int kernel, + const double* __restrict__ coefficient, + int gaussians) { + if (kernel == kLes) { + return 1.1283791670955126 * coefficient[0]; + } + double value = 0.0; + for (int g = 0; g < gaussians; ++g) { + value += coefficient[g]; + } + return value; +} + +// Adds ``value`` over the lanes of a warp in a fixed butterfly order; every +// lane receives the total. +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int stride = kWarp / 2; stride > 0; stride /= 2) { + value += __shfl_xor_sync(0xffffffffu, value, stride); + } + return value; +} + +// Interpolates the potential meshes at every owner node, one warp per node, +// one pass per channel: the potential phi^rec_c of the mesh, which still holds +// that of the charge at itself, and the force at fixed charges. Owners that +// carry no charge get zeros. Every lane sums its points in float32, the +// precision of the meshes, and the warp adds the lane sums. +template +__global__ void __launch_bounds__(kStencilWarps* kWarp) + interpolate_owners(const bool* __restrict__ node_mask, + const int64_t* __restrict__ frame_ptr, + const FrameMesh* __restrict__ frames, + const AxisStencil* __restrict__ stencils, + const float* __restrict__ node_charge, + const float* __restrict__ potential_mesh, + int frame_count, + int owner_begin, + int owner_count, + int channels, + float* __restrict__ potential, + float* __restrict__ force) { + using Slots = LaneSlots; + __shared__ AxisStencil shared[kStencilWarps][3]; + const int warp = threadIdx.x / kWarp; + const int lane = threadIdx.x % kWarp; + const AxisStencil* axes = shared[warp]; + for (int owner = blockIdx.x * kStencilWarps + warp; owner < owner_count; + owner += gridDim.x * kStencilWarps) { + const int64_t node = int64_t{owner_begin} + owner; + const int f = frame_of(frame_ptr, frame_count, node); + if (f < 0 || !node_mask[node]) { + for (int c = lane; c < channels; c += kWarp) { + potential[int64_t{owner} * channels + c] = 0.f; + } + if (lane < 3) { + force[int64_t{owner} * 3 + lane] = 0.f; + } + continue; + } + load_stencil(stencils, node, lane, shared[warp]); + const FrameMesh& grid = frames[f]; + // The derivative of sum_c q_c phi_c with respect to u_d. + float gradient[3] = {0.f, 0.f, 0.f}; + for (int c = 0; c < channels; ++c) { + const float* channel_mesh = + potential_mesh + grid.offset + c * mesh_points(grid); + float phi = 0.f; + float slope[3] = {0.f, 0.f, 0.f}; +#pragma unroll + for (int round = 0; round < Slots::kRounds; ++round) { + const int slot = lane + kWarp * round; + if (slot < Slots::kSlots) { + const int group = slot / Slots::kPairs; + const int pair = slot - group * Slots::kPairs; + const int k2 = pair / kOrder; + const int k3 = pair - k2 * kOrder; + const float v2 = axes[1].value[k2]; + const float v3 = axes[2].value[k3]; + const float v23 = v2 * v3; + const float s2v3 = axes[1].slope[k2] * v3; + const float v2s3 = v2 * axes[2].slope[k3]; + const int column = axes[1].offset[k2] + axes[2].offset[k3]; +#pragma unroll + for (int k1 = group; k1 < kOrder; k1 += Slots::kGroups) { + const float value = channel_mesh[axes[0].offset[k1] + column]; + const float v1_value = axes[0].value[k1] * value; + phi += v1_value * v23; + slope[0] += axes[0].slope[k1] * value * v23; + slope[1] += v1_value * s2v3; + slope[2] += v1_value * v2s3; + } + } + } + phi = warp_sum(phi); + const float q = node_charge[node * channels + c]; +#pragma unroll + for (int d = 0; d < 3; ++d) { + gradient[d] += q * warp_sum(slope[d]); + } + if (lane == 0) { + potential[int64_t{owner} * channels + c] = phi; + } + } + if (lane < 3) { + // F_j = -sum_d dE/du_d n_d (B^-1)_jd + force[int64_t{owner} * 3 + lane] = + gradient[0] * grid.gradient_to_force[3 * lane] + + gradient[1] * grid.gradient_to_force[3 * lane + 1] + + gradient[2] * grid.gradient_to_force[3 * lane + 2]; + } + // The stencil of the next node of the warp overwrites this one. + __syncwarp(); + } +} + +// The potential prefactor * K(0) of a unit charge at itself, which the +// potential mesh holds and the potential of a node excludes, for every thread +// of the block. +__device__ double self_potential(int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor) { + __shared__ double self; + if (threadIdx.x == 0) { + self = prefactor * contact_value(kernel, coefficient, gaussians); + } + __syncthreads(); + return self; +} + +// Per-frame sums over one slice of the owners: the potential of the charged +// owners in the constrained channel, the number of charged owners and the +// nine components of sum_i F_i (x) r_i. +constexpr int kOwnerSums = 11; + +__global__ void __launch_bounds__(kFrameThreads) + owner_slices(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + int channels, + int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + double* __restrict__ slice_sum) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); + const NodeRange slice = slice_of_frame(owner_frame_ptr); + double value[kOwnerSums] = {}; + for (int64_t owner = slice.begin + threadIdx.x; owner < slice.end; + owner += kFrameThreads) { + if (!owner_mask[owner]) { + continue; + } + if (channel >= 0) { + value[0] += double{potential[owner * channels + channel]} - + self * owner_charge[owner * channels + channel]; + } + value[1] += 1.0; + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + value[2 + 3 * a + b] += + double{force[owner * 3 + a]} * owner_coord[owner * 3 + b]; + } + } + } + write_slice_sum(value, slice_sum); +} + +// Adds the slice sums of owner_slices of every frame in slice order. +__global__ void owner_totals(const double* __restrict__ slice_sum, + int slice_count, + double* __restrict__ total) { + if (threadIdx.x >= kOwnerSums) { + return; + } + double value = 0.0; + for (int slice = 0; slice < slice_count; ++slice) { + value += + slice_sum[(static_cast(blockIdx.x) * slice_count + slice) * + kOwnerSums + + threadIdx.x]; + } + total[blockIdx.x * kOwnerSums + threadIdx.x] = value; +} + +// Writes the outputs of every owner: the atomic energy q . phi / 2 with phi = +// phi^rec - prefactor K(0) q, the charge gradient w (phi - [c = ch] mean), and +// the virial F (x) r plus the share of the rest of the virial W of the frame +// over its charged nodes. The first threads write the share of the virial of +// every frame held by the owners, sum_i F_i (x) r_i + n (W - sum_j F_j (x) +// r_j) / N over the n charged owners of the N charged nodes of the frame. +// ``owner_local`` holds the owner sums of this domain, ``owner_total`` those +// of every domain. +__global__ void finish_owners(const float* __restrict__ potential, + const float* __restrict__ force, + const double* __restrict__ owner_coord, + const float* __restrict__ owner_charge, + const bool* __restrict__ owner_mask, + const int64_t* __restrict__ owner_frame_ptr, + const double* __restrict__ frame_virial, + const double* __restrict__ owner_local, + const double* __restrict__ owner_total, + const double* __restrict__ frame_stat, + int frame_count, + int owner_count, + int channels, + int channel, + int kernel, + const double* __restrict__ coefficient, + int gaussians, + double prefactor, + double* __restrict__ atom_energy, + float* __restrict__ charge_gradient, + float* __restrict__ atom_virial, + float* __restrict__ virial) { + const double self = self_potential(kernel, coefficient, gaussians, prefactor); + const int start = blockIdx.x * blockDim.x + threadIdx.x; + const int threads = gridDim.x * blockDim.x; + for (int index = start; index < 9 * frame_count; index += threads) { + const int f = index / 9; + const int k = index - 9 * f; + const double count = frame_stat[f * kStatWidth + kStatCount]; + const double* local = owner_local + f * kOwnerSums; + const double rest = + count > 0.0 + ? (frame_virial[index] - owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + virial[index] = static_cast(local[2 + k] + local[1] * rest); + } + for (int owner = start; owner < owner_count; owner += threads) { + const int f = frame_of(owner_frame_ptr, frame_count, owner); + const bool charged = f >= 0 && owner_mask[owner]; + // The charged nodes of the frame, of every domain. + const double count = + charged ? frame_stat[f * kStatWidth + kStatCount] : 1.0; + const double mean = + charged && channel >= 0 ? owner_total[f * kOwnerSums] / count : 0.0; + double energy = 0.0; + for (int c = 0; c < channels; ++c) { + const int64_t entry = int64_t{owner} * channels + c; + const double q = owner_charge[entry]; + const double phi = charged ? potential[entry] - self * q : 0.0; + energy += q * phi; + charge_gradient[entry] = + static_cast(c == channel ? phi - mean : phi); + } + atom_energy[owner] = 0.5 * energy; + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + const int k = 3 * a + b; + const double moment = double{force[int64_t{owner} * 3 + a]} * + owner_coord[int64_t{owner} * 3 + b]; + const double rest = charged ? (frame_virial[f * 9 + k] - + owner_total[f * kOwnerSums + 2 + k]) / + count + : 0.0; + atom_virial[int64_t{owner} * 9 + k] = static_cast(moment + rest); + } + } + } +} + +IntegerWeights integer_bspline_weights(int order) { + IntegerWeights weight{}; + for (int l = 0; l < order - 1; ++l) { + const double x = l + 1.0; + double value = 0.0; + double binomial = 1.0; + for (int j = 0; j <= order; ++j) { + if (j > 0) { + binomial = binomial * (order - j + 1) / j; + } + const double base = std::max(x - j, 0.0); + value += (j % 2 == 0 ? 1.0 : -1.0) * binomial * std::pow(base, order - 1); + } + double factorial = 1.0; + for (int k = 2; k < order; ++k) { + factorial *= k; + } + weight.value[l] = value / factorial; + } + return weight; +} + +void check_cufft(cufftResult result, const char* operation) { + TORCH_CHECK(result == CUFFT_SUCCESS, "latent_charge_sum: cuFFT ", operation, + " failed with error ", static_cast(result)); +} + +// The real-to-complex and complex-to-real transforms of the meshes of one +// shape, batched over the channels. The plans allocate no work area: every +// transform runs on a work area of the workspace of its call, set with the +// stream under ``mutex`` just before the transform is enqueued, so that +// concurrent calls never share one. +struct MeshPlans { + cufftHandle forward = 0; + cufftHandle inverse = 0; + size_t work_bytes = 0; + c10::DeviceIndex device = 0; + std::mutex mutex; + + MeshPlans() = default; + MeshPlans(const MeshPlans&) = delete; + MeshPlans& operator=(const MeshPlans&) = delete; + // A plan leaves the cache only past kPlanCacheSize shapes; its transforms + // may still be in flight, so the device finishes them first. The + // destructor reports no error. + ~MeshPlans() { + int current = 0; + cudaGetDevice(¤t); + cudaSetDevice(device); + cudaDeviceSynchronize(); + if (forward != 0) { + cufftDestroy(forward); + } + if (inverse != 0) { + cufftDestroy(inverse); + } + cudaSetDevice(current); + } +}; + +// The plans of the meshes of ``mesh`` with ``channels`` channels on the +// current device, kept across calls; beyond kPlanCacheSize shapes, the least +// recently used are released. The cache itself lives until the process ends +// and is never destroyed, since its plans would otherwise be released after +// the CUDA runtime shuts down. +constexpr size_t kPlanCacheSize = 16; + +std::shared_ptr mesh_plans(const FrameMesh& mesh, int channels) { + using Key = std::array; + using Cache = std::list>>; + static std::mutex cache_mutex; + static Cache& cache = *new Cache(); + const Key key = {c10::cuda::current_device(), mesh.n[0], mesh.n[1], mesh.n[2], + channels}; + // A plan that leaves the cache is released after the lock, since its + // release waits for the device. + std::shared_ptr evicted; + const std::lock_guard lock(cache_mutex); + for (auto entry = cache.begin(); entry != cache.end(); ++entry) { + if (entry->first == key) { + cache.splice(cache.begin(), cache, entry); + return entry->second; + } + } + auto plans = std::make_shared(); + plans->device = c10::cuda::current_device(); + long long shape[3] = {mesh.n[0], mesh.n[1], mesh.n[2]}; + const long long points = shape[0] * shape[1] * shape[2]; + const long long modes = shape[0] * shape[1] * (shape[2] / 2 + 1); + size_t forward_bytes = 0; + size_t inverse_bytes = 0; + check_cufft(cufftCreate(&plans->forward), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->forward, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->forward, 3, shape, nullptr, 1, points, nullptr, + 1, modes, CUFFT_R2C, channels, &forward_bytes), + "plan creation"); + check_cufft(cufftCreate(&plans->inverse), "plan creation"); + check_cufft(cufftSetAutoAllocation(plans->inverse, 0), "plan setup"); + check_cufft( + cufftMakePlanMany64(plans->inverse, 3, shape, nullptr, 1, modes, nullptr, + 1, points, CUFFT_C2R, channels, &inverse_bytes), + "plan creation"); + plans->work_bytes = std::max(forward_bytes, inverse_bytes); + cache.emplace_front(key, plans); + if (cache.size() > kPlanCacheSize) { + evicted = std::move(cache.back().second); + cache.pop_back(); + } + return plans; +} + +int blocks_for(int64_t work) { + return static_cast(std::max( + 1, std::min((work + kThreads - 1) / kThreads, 65535))); +} + +// Blocks of a stencil kernel over ``nodes`` nodes, one warp each. +int stencil_blocks_for(int64_t nodes) { + return static_cast(std::max( + 1, std::min((nodes + kStencilWarps - 1) / kStencilWarps, + int64_t{1} << 20))); +} + +// Calls ``launch`` with the B-spline order as a compile-time constant, for +// which the stencil kernels unroll their loops. +template +void dispatch_order(int64_t order, const Launch& launch) { + switch (order) { + case 2: + return launch(std::integral_constant{}); + case 3: + return launch(std::integral_constant{}); + case 4: + return launch(std::integral_constant{}); + case 5: + return launch(std::integral_constant{}); + case 6: + return launch(std::integral_constant{}); + case 7: + return launch(std::integral_constant{}); + case 8: + return launch(std::integral_constant{}); + case 9: + return launch(std::integral_constant{}); + case 10: + return launch(std::integral_constant{}); + case 11: + return launch(std::integral_constant{}); + case 12: + return launch(std::integral_constant{}); + default: + TORCH_CHECK(false, "latent_charge_sum: unsupported B-spline order ", + order); + } +} + +// The value of ``evaluate`` at the B-spline order as a compile-time constant. +template +auto dispatch_order_value(int64_t order, const Evaluate& evaluate) { + decltype(evaluate(std::integral_constant{})) value{}; + dispatch_order( + order, [&](auto order_constant) { value = evaluate(order_constant); }); + return value; +} + +} // namespace + +namespace deepmd { +namespace long_range { + +ChargeSumOutputs mesh_sum(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& box, + const torch::Tensor& coefficient, + int64_t kernel, + int64_t channel, + double prefactor, + int64_t order, + double spacing, + int owner_begin, + int owner_count, + const torch::Tensor* communicator) { + TORCH_CHECK(order >= 2 && order <= kMaxOrder, + "latent_charge_sum: the B-spline order must lie in [2, ", + kMaxOrder, "], got ", order); + TORCH_CHECK(std::isfinite(spacing) && spacing > 0.0, + "latent_charge_sum: the mesh spacing must be positive, got ", + spacing); + const int64_t node_count = coord.size(0); + TORCH_CHECK(node_count < std::numeric_limits::max(), + "latent_charge_sum: too many nodes"); + const int frame_count = static_cast(n_node_per_frame.size(0)); + const int channels = static_cast(charge.size(1)); + const int gaussians = + kernel == kSog ? static_cast(coefficient.numel() / 2) : 0; + const auto double_options = coord.options(); + const auto float_options = charge.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_sum: owners other than every node lie in one " + "frame"); + const auto stream = at::cuda::getCurrentCUDAStream(); + + // === Step 1. The meshes of the frames, sized on the host === + std::vector frames(frame_count); + std::vector> plans(frame_count); + const double* cells = box.data_ptr(); + int64_t mesh_total = 0; + int64_t spectrum_size = 0; + int64_t factor_size = 0; + size_t work_bytes = 0; + for (int f = 0; f < frame_count; ++f) { + frames[f] = frame_mesh(cells + 9 * f, spacing); + frames[f].offset = mesh_total; + mesh_total += channels * mesh_points(frames[f]); + const int half = frames[f].n[2] / 2 + 1; + spectrum_size = std::max(spectrum_size, int64_t{channels} * frames[f].n[0] * + frames[f].n[1] * half); + factor_size = + std::max(factor_size, int64_t{frames[f].n[0]} + frames[f].n[1] + half); + plans[f] = mesh_plans(frames[f], channels); + work_bytes = std::max(work_bytes, plans[f]->work_bytes); + } + const auto byte_options = torch::TensorOptions().dtype(torch::kUInt8); + auto host_frames = + torch::empty({static_cast(sizeof(FrameMesh)) * frame_count}, + byte_options.pinned_memory(true)); + std::memcpy(host_frames.data_ptr(), frames.data(), + sizeof(FrameMesh) * frame_count); + const IntegerWeights integer_weight = + integer_bspline_weights(static_cast(order)); + // A few blocks per multiprocessor fill the device; every thread of the + // influence function then walks several modes, which keeps the block sums + // few. + const int mode_blocks = + 4 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + + // === Step 2. The buffers of the call === + const int64_t average_nodes = (node_count + frame_count - 1) / frame_count; + const auto slices_for = [](int64_t nodes) { + return static_cast(std::min( + kMaxSlices, + std::max(1, (nodes + kSliceNodes - 1) / kSliceNodes))); + }; + const int slice_count = slices_for(average_nodes); + const int owner_slice_count = slices_for( + (int64_t{owner_count} + frame_count - 1) / std::max(frame_count, 1)); + const int64_t stencil_bytes = dispatch_order_value(order, [](auto constant) { + return static_cast(sizeof(AxisStencil)); + }); + Workspace workspace; + const int64_t frames_at = + workspace.reserve(sizeof(FrameMesh) * frame_count); + const int64_t frame_ptr_at = workspace.reserve(frame_count + 1); + const int64_t slice_charge_at = + workspace.reserve(int64_t{frame_count} * slice_count * 2); + const int64_t frame_stat_at = + workspace.reserve(int64_t{frame_count} * kStatWidth); + const int64_t stencils_at = + workspace.reserve(3 * node_count * stencil_bytes); + const int64_t node_charge_at = + workspace.reserve(node_count * channels); + const int64_t fixed_mesh_at = workspace.reserve(mesh_total); + // The potential mesh of a frame overwrites its charge mesh, which its + // forward transform has consumed. + const int64_t mesh_at = workspace.reserve(mesh_total); + const int64_t spectrum_at = workspace.reserve(spectrum_size); + const int64_t work_at = workspace.reserve(work_bytes); + const int64_t factor_at = workspace.reserve(factor_size); + const int64_t block_sum_at = + workspace.reserve(int64_t{mode_blocks} * kModeSums); + const int64_t frame_virial_at = + workspace.reserve(int64_t{frame_count} * 9); + const int64_t potential_at = + workspace.reserve(int64_t{owner_count} * channels); + const int64_t owner_slice_at = workspace.reserve( + int64_t{frame_count} * owner_slice_count * kOwnerSums); + const int64_t owner_local_at = + workspace.reserve(int64_t{frame_count} * kOwnerSums); + workspace.allocate(coord.device()); + auto* frames_ptr = workspace.at(frames_at); + auto* frame_ptr = workspace.at(frame_ptr_at); + auto* slice_charge = workspace.at(slice_charge_at); + auto* frame_stat = workspace.at(frame_stat_at); + auto* node_charge = workspace.at(node_charge_at); + auto* fixed_mesh = workspace.at(fixed_mesh_at); + auto* charge_mesh = workspace.at(mesh_at); + float* potential_mesh = charge_mesh; + auto* spectrum = workspace.at(spectrum_at); + void* work = workspace.at(work_at); + auto* factor = workspace.at(factor_at); + auto* block_sum = workspace.at(block_sum_at); + auto* frame_virial = workspace.at(frame_virial_at); + auto* potential = workspace.at(potential_at); + auto* owner_slice = workspace.at(owner_slice_at); + auto* owner_local = workspace.at(owner_local_at); + // The copy from the pinned staging buffer records its completion, before + // which the allocator does not hand the buffer out again. + torch::from_blob(frames_ptr, {host_frames.numel()}, + byte_options.device(coord.device())) + .copy_(host_frames, /*non_blocking=*/true); + C10_CUDA_CHECK( + cudaMemsetAsync(fixed_mesh, 0, sizeof(int64_t) * mesh_total, stream)); + + // === Step 3. The constraint, and the stencils and charges of the nodes === + const auto counts = n_node_per_frame.to(torch::kInt64).contiguous(); + frame_offsets<<<1, kOffsetThreads, 0, stream>>>(counts.data_ptr(), + frame_count, frame_ptr); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_slices<<>>( + charge.data_ptr(), node_mask.data_ptr(), frame_ptr, channels, + static_cast(std::max(channel, 0)), slice_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + charge_shift<<<(frame_count + kFrameThreads - 1) / kFrameThreads, + kFrameThreads, 0, stream>>>( + slice_charge, slice_count, frame_count, total_charge.data_ptr(), + frame_stat); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + auto* node_stencil = workspace.at>(stencils_at); + node_stencils<<>>( + coord.data_ptr(), charge.data_ptr(), + node_mask.data_ptr(), frame_ptr, frames_ptr, frame_stat, + frame_count, node_count, channels, static_cast(channel), + node_stencil, node_charge); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // === Step 4. The charge meshes === + spread_charges + <<>>( + node_mask.data_ptr(), frame_ptr, frames_ptr, node_stencil, + node_charge, frame_count, node_count, channels, fixed_mesh); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + fixed_to_float<<>>( + mesh_total, fixed_mesh, charge_mesh); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + // === Step 5. Influence function and potential meshes, frame by frame === + for (int f = 0; f < frame_count; ++f) { + const FrameMesh& grid = frames[f]; + MeshPlans& plan = *plans[f]; + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.forward, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.forward, work), "work area setup"); + check_cufft(cufftExecR2C(plan.forward, charge_mesh + grid.offset, + reinterpret_cast(spectrum)), + "transform"); + } + const int half = grid.n[2] / 2 + 1; + const int factor_count = grid.n[0] + grid.n[1] + half; + assignment_factor<<>>( + grid, static_cast(order), integer_weight, factor); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + apply_influence<<>>( + grid, channels, static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, factor, spectrum, + block_sum); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + mode_virial<<<1, kThreads, 0, stream>>>(block_sum, mode_blocks, + frame_virial + 9 * f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + { + const std::lock_guard lock(plan.mutex); + check_cufft(cufftSetStream(plan.inverse, stream), "stream setup"); + check_cufft(cufftSetWorkArea(plan.inverse, work), "work area setup"); + check_cufft( + cufftExecC2R(plan.inverse, reinterpret_cast(spectrum), + potential_mesh + grid.offset), + "transform"); + } + } + + // === Step 6. Potential, force and virial of the owners === + auto force = torch::empty({owner_count, 3}, float_options); + if (owner_count > 0) { + dispatch_order(order, [&](auto order_constant) { + constexpr int kOrder = decltype(order_constant)::value; + interpolate_owners<<>>( + node_mask.data_ptr(), frame_ptr, frames_ptr, + workspace.at>(stencils_at), node_charge, + potential_mesh, frame_count, owner_begin, owner_count, channels, + potential, force.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // The frames of the owners, on the axis of the owners: the frames + // themselves when every node is an owner, else the one frame they share. + torch::Tensor owner_frame_tensor; + const int64_t* owner_frame_ptr = frame_ptr; + if (!every_node) { + owner_frame_tensor = + torch::full({2}, int64_t{owner_count}, counts.options()); + owner_frame_tensor.select(0, 0).zero_(); + owner_frame_ptr = owner_frame_tensor.data_ptr(); + } + const bool* owner_mask = node_mask.data_ptr() + owner_begin; + const double* owner_coord = coord.data_ptr() + 3 * owner_begin; + const float* owner_charge = node_charge + int64_t{owner_begin} * channels; + owner_slices<<>>(potential, force.data_ptr(), owner_coord, + owner_charge, owner_mask, owner_frame_ptr, channels, + static_cast(channel), static_cast(kernel), + coefficient.data_ptr(), gaussians, prefactor, + owner_slice); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + owner_totals<<>>(owner_slice, owner_slice_count, + owner_local); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + // The mean potential of the constraint and the virial share run over the + // charged nodes of every domain. + torch::Tensor owner_sum; + const double* owner_total = owner_local; + if (communicator != nullptr) { + owner_sum = sum_over_domains( + *communicator, torch::from_blob(owner_local, {frame_count * kOwnerSums}, + double_options)); + owner_total = owner_sum.data_ptr(); + } + auto atom_energy = torch::empty({owner_count, 1}, double_options); + auto charge_gradient = torch::empty({owner_count, channels}, float_options); + auto atom_virial = torch::empty({owner_count, 9}, float_options); + auto virial = torch::empty({frame_count, 9}, float_options); + finish_owners<<(owner_count, 9 * frame_count)), + kThreads, 0, stream>>>( + potential, force.data_ptr(), owner_coord, owner_charge, owner_mask, + owner_frame_ptr, frame_virial, owner_local, owner_total, frame_stat, + frame_count, owner_count, channels, static_cast(channel), + static_cast(kernel), coefficient.data_ptr(), gaussians, + prefactor, atom_energy.data_ptr(), + charge_gradient.data_ptr(), atom_virial.data_ptr(), + virial.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {atom_energy, charge_gradient, force, atom_virial, virial}; +} + +} // 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..364fa6c2c2 --- /dev/null +++ b/source/op/pt/long_range/ops.cc @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// +// Operator schema of the latent-charge long-range energy. +// +// The schema is declared here, unconditionally, while the CUDA kernel +// registers itself only against a CUDA-enabled PyTorch (latent_charge_sum.cu), +// so that the Python front end can tell "library not loaded" from "this device +// has no kernel". +// +// For the nodes of the frames, 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 interactions 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. Without `box`, every frame is a complete +// non-periodic system, in which the potential of the charges is the direct +// sum +// +// phi_i = prefactor * sum_{j != i} q_j K(r_ij); +// +// with `box`, whose row f holds the row-major cell of frame f, the rows of the +// cell being its lattice vectors, phi_i sums over the images of the cell, by +// the smooth particle-mesh Ewald method (mesh_sum.cu), less the potential +// prefactor * K(0) q_i of every charge at itself. An all-zero row marks a +// non-periodic frame, which a call does not mix with periodic ones. Either way +// 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 = -dE / dr_i at fixed charges +// atom_virial_i (9 components, row-major) +// virial_f = sum_{i in f} atom_virial_i +// +// in float32 unless noted. In a non-periodic frame, with d_ij = r_i - r_j, +// +// 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; +// +// in a periodic frame the virial of the frame is that of the strain of the +// cell, W_f = -dE / d(strain), and W_i = F_i (x) r_i + (W_f - sum_{j in f} +// F_j (x) r_j) / n_f on the charged nodes, the force index first, so that the +// charged nodes of the frame sum to W_f. +// +// 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. +// `coefficient` is float64. `total_charge` holds Q_f in float64 and is read +// only with a constrained channel. `box` is a float64 CPU tensor of shape +// (F, 9). The mesh of a periodic frame holds along every lattice vector a_d +// the smallest even number of points without a prime factor above 7 of at +// least ceil(|a_d| / spacing), and spreads the charges with B-splines of +// order `order`. 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_sum(Tensor coord, Tensor charge, Tensor node_mask, " + "Tensor n_node_per_frame, Tensor total_charge, Tensor? box, " + "Tensor coefficient, Tensor communicator, int kernel, int channel, " + "float prefactor, int order, float spacing) -> " + "(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_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..79320de4a3 --- /dev/null +++ b/source/tests/common/dpmodel/test_long_range.py @@ -0,0 +1,890 @@ +# 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 ( + COULOMB_CONSTANT, + BaseInteractionKernel, + LatentChargeLongRange, + LESKernel, + SOGKernel, +) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + K_BLOCK, + build_k_grid, + build_pair_graph, + build_topologies, + enumerate_k_grid_bounds, + 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() + + +# --------------------------------------------------------------------------- +# Periodic (k-space) primitives +# --------------------------------------------------------------------------- + + +def _spherical_ft( + kernel: str, k: float, n_r: int = 20000, r_max: float = 60.0 +) -> float: + """Isotropic Fourier transform of the kernel by direct radial quadrature.""" + r = np.linspace(1e-9, r_max, n_r) + if kernel == "les": + k_r = 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_r = np.sum( + amplitude[None, :] * np.exp(-(r[:, None] ** 2) / (2 * width**2)), -1 + ) + return float(np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(k * r / np.pi), r)) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_kspace_matches_radial_ft(kernel: str) -> None: + """``kernel.kspace`` is the analytic Fourier transform of the real-space kernel.""" + k_sq = np.array([0.5, 1.0, 2.0, 5.0]) ** 2 + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + got = mod.kspace(k_sq) + for kk, value in zip(np.sqrt(k_sq), got, strict=True): + if kernel == "les": + # erf(alpha r)/r is not integrable; quadrature of the integrable + # complement erfc(alpha r)/r = 1/r - erf(alpha r)/r converges, and + # K_hat(k) = 4 pi / k^2 - FT(erfc(alpha r)/r). + r = np.linspace(1e-9, 60.0, 20000) + k_r = (1.0 - erf(0.8 * r)) / r # erfc(0.8 r)/r + expected = 4.0 * np.pi / kk**2 - float( + np.trapezoid(4.0 * np.pi * r**2 * k_r * np.sinc(kk * r / np.pi), r) + ) + else: + expected = _spherical_ft(kernel, kk) + np.testing.assert_allclose(value, expected, rtol=1e-3, atol=1e-5) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_kernel_contact_value_matches_real_space_limit(kernel: str) -> None: + """``kernel.contact_value`` is K(r) at contact.""" + if kernel == "les": + mod = LESKernel(alpha=0.8) + else: + mod = SOGKernel(amplitude=[0.7, -0.2, 0.1], width=[0.8, 1.7, 3.2]) + r_sq = np.array([1e-12]) + np.testing.assert_allclose( + mod.contact_value(r_sq), mod(r_sq)[0], rtol=1e-6, atol=1e-9 + ) + + +def test_build_k_grid_orthogonality_and_guards() -> None: + """The k points satisfy k . a_d = 2 pi n_d, every frame fills a row of its own.""" + box = np.array([[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]])[None] + grid = build_k_grid(box, 2.0, np.zeros((4, 3)), np.array([4])) + # the points fill whole blocks, the last one completed by guard points + assert grid.mask.shape == grid.vec.shape[:3] + assert grid.vec.shape[2] == K_BLOCK + flat_mask = grid.mask.reshape(-1) + count = int(flat_mask.sum()) + assert flat_mask[:count].all() and not flat_mask[count:].any() + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, 3]]) + vec = grid.vec.reshape(-1, 3)[:count] + n_integer = vec @ box[0].T / (2 * np.pi) + np.testing.assert_allclose(n_integer, np.round(n_integer), atol=1e-10) + k_sq = np.sum(vec**2, axis=-1) + assert (k_sq > 0).all() and (k_sq <= 4.0 + 1e-12).all() + np.testing.assert_allclose(grid.volume, np.abs(np.linalg.det(box[0]))) + # every integer point inside the sphere is present exactly once + ns = np.round(n_integer).astype(int) + assert len({tuple(row) for row in ns}) == ns.shape[0] + rec = 2 * np.pi * np.linalg.inv(box[0]) + n_d = np.ceil(2.0 * np.linalg.norm(box[0], axis=1) / (2 * np.pi)).astype(int) + count = 0 + for n0 in range(-n_d[0], n_d[0] + 1): + for n1 in range(-n_d[1], n_d[1] + 1): + for n2 in range(-n_d[2], n_d[2] + 1): + if (n0, n1, n2) != (0, 0, 0): + k = rec @ np.array([n0, n1, n2], dtype=float) + count += int(np.sum(k * k) <= 4.0) + assert count == ns.shape[0] + + +def test_build_topologies_periodic_branching() -> None: + """A cell switches the pair topology to guards and the k grid on, and back.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + ALL_PAIRS, + K_GRID, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + n_node = np.array([4]) + box = np.eye(3)[None] * 10.0 + names = (ALL_PAIRS, K_GRID) + params = {K_GRID: {"k_max": 2.0}} + with_box = build_topologies( + names, coord, n_node, box, atype=atype, topology_params=params + ) + assert with_box[ALL_PAIRS].index.shape[1] == 4 # guards only + assert not with_box[ALL_PAIRS].mask.any() + assert with_box[K_GRID].mask.sum() > 0 + without_box = build_topologies( + names, coord, n_node, None, atype=atype, topology_params=params + ) + assert without_box[ALL_PAIRS].mask.any() + assert not without_box[K_GRID].mask.any() # guards only + # a pair-only declaration still rejects a cell loudly + with pytest.raises(NotImplementedError, match="non-periodic"): + build_topologies((ALL_PAIRS,), coord, n_node, box, atype=atype) + # the k grid needs its build parameter + with pytest.raises(ValueError, match="k_max"): + build_topologies(names, coord, n_node, box, atype=atype) + # the reciprocal-space sum excludes no pair; a non-periodic batch still does + pair_excl = PairExcludeMask(2, [(0, 1)]) + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + build_topologies( + names, + coord, + n_node, + box, + atype=atype, + pair_excl=pair_excl, + topology_params=params, + ) + excluded = build_topologies( + names, coord, n_node, None, atype=atype, pair_excl=pair_excl + ) + assert 0 < excluded[ALL_PAIRS].mask.sum() < without_box[ALL_PAIRS].mask.sum() + # a cell without volume is no periodic system + with pytest.raises(ValueError, match="no volume"): + build_topologies( + names, coord, n_node, np.zeros((1, 9)), atype=atype, topology_params=params + ) + + +def test_build_k_grid_pads_the_rows_of_the_frames() -> None: + """Every frame keeps its own points and nodes; shorter rows end in masked guards.""" + box = np.stack([np.diag([8.0, 7.5, 9.0]), np.diag([5.0, 5.5, 6.0])]) + n_node = np.array([3, 5]) + grid = build_k_grid(box, 2.0, np.zeros((8, 3)), n_node) + alone = [ + build_k_grid(box[f : f + 1], 2.0, np.zeros((n, 3)), np.array([n])) + for f, n in enumerate(n_node) + ] + counts = [int(single.mask.sum()) for single in alone] + blocks = max(-(-max(counts) // K_BLOCK), 2) + assert grid.vec.shape == (2, blocks, K_BLOCK, 3) + for f, single in enumerate(alone): + mask = grid.mask[f].reshape(-1) + vec = grid.vec[f].reshape(-1, 3) + np.testing.assert_array_equal(mask, np.arange(blocks * K_BLOCK) < counts[f]) + np.testing.assert_allclose( + vec[: counts[f]], single.vec[0].reshape(-1, 3)[: counts[f]] + ) + assert np.all(vec[counts[f] :] == [1e-6, 0.0, 0.0]) + np.testing.assert_array_equal(grid.slot, [[0, 1, 2, -1, -1], [3, 4, 5, 6, 7]]) + + +@pytest.mark.parametrize("kernel", [LESKernel(alpha=0.6), SOGKernel(n_gaussian=3)]) +def test_kspace_energy_of_padded_frames_matches_frames_alone( + kernel: BaseInteractionKernel, +) -> None: + """Frames of unequal node counts, an empty one among them, sum as they do alone.""" + from deepmd.dpmodel.long_range.reciprocal import ( + kspace_energy, + ) + + box = np.stack( + [np.diag([8.0, 7.5, 9.0]), np.diag([6.0, 6.5, 7.0]), np.diag([5.0, 5.5, 6.0])] + ) + n_node = np.array([4, 0, 6]) + rng = np.random.default_rng(3) + coord = rng.uniform(0.0, 5.0, (10, 3)) + charge = rng.normal(0.0, 0.5, (10, 2)) + energy = kspace_energy(charge, build_k_grid(box, 2.5, coord, n_node), kernel) + start = 0 + for f, n in enumerate(n_node): + if n == 0: + continue + nodes = slice(start, start + n) + alone = build_k_grid(box[f : f + 1], 2.5, coord[nodes], np.array([n])) + np.testing.assert_allclose( + energy[nodes], kspace_energy(charge[nodes], alone, kernel), rtol=1e-12 + ) + start += n + + +# --------------------------------------------------------------------------- +# Fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _module_topologies( + module: LatentChargeLongRange, + coord: np.ndarray, + n_node: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, +) -> dict: + """The topologies of a module, built with its parameters.""" + return build_topologies( + module.required_topologies(), + coord, + n_node, + box, + atype=atype, + topology_params=module.topology_params(), + ) + + +def _periodic_long_range(**kwargs: object) -> LatentChargeLongRange: + """A periodic module, over the explicit wave numbers unless ``k_grid_mode`` says otherwise.""" + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + kwargs.setdefault("k_grid_mode", "ragged") + return LatentChargeLongRange( + charge_net, LESKernel(alpha=0.3), None, periodic=True, **kwargs + ) + + +def _cells() -> np.ndarray: + """Two cells of different sizes, the second the larger one.""" + return np.array( + [ + [[8.0, 0.6, 0.0], [0.0, 7.5, 0.4], [0.0, 0.0, 9.0]], + [[12.0, 0.5, 0.0], [0.2, 11.0, 0.3], [0.1, 0.4, 10.5]], + ] + ) + + +def test_enumerate_k_grid_bounds() -> None: + """The grid is the lexicographic meshgrid of the bounds, without the zero vector.""" + n_grid = enumerate_k_grid_bounds([2, 1, 0]) + assert n_grid.shape == (5 * 3 * 1 - 1, 3) + assert not np.any(np.all(n_grid == 0, axis=-1)) + expected = [ + (n0, n1, 0) for n0 in range(-2, 3) for n1 in range(-1, 2) if (n0, n1) != (0, 0) + ] + np.testing.assert_array_equal(n_grid, np.asarray(expected, dtype=np.int64)) + with pytest.raises(ValueError, match="bounds"): + enumerate_k_grid_bounds([1, -1, 0]) + + +def test_fixed_k_grid_matches_ragged_enumeration() -> None: + """The masked points of the fixed grid are exactly the ragged k points.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] + n_node = np.array([4, 5]) + fixed = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) + # a static layout: the blocks of the grid in every row + assert fixed.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + ragged = build_k_grid(box, k_max, np.zeros((9, 3)), n_node) + for frame in range(2): + expected = { + tuple(row) for row in np.round(ragged.vec[frame][ragged.mask[frame]], 8) + } + got = {tuple(row) for row in np.round(fixed.vec[frame][fixed.mask[frame]], 8)} + assert got == expected + # the smaller cell masks more points of the shared grid + assert fixed.mask[0].sum() < fixed.mask[1].sum() + np.testing.assert_allclose(fixed.volume, np.abs(np.linalg.det(box)), rtol=1e-12) + + +def test_fixed_k_grid_truncates_oversized_cells() -> None: + """A cell beyond the grid is served by the grid alone (truncated k-sum), with a warning.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_fixed_k_grid, + enumerate_k_grid_bounds, + k_grid_bounds_of_cells, + ) + + box = _cells() + k_max = 2.0 + bounds = np.max(k_grid_bounds_of_cells(box, k_max), axis=0) + nk = enumerate_k_grid_bounds(bounds).shape[0] + n_node = np.array([4, 5]) + with pytest.warns(UserWarning, match="truncated"): + bigger = build_fixed_k_grid(box * 2.0, bounds, k_max, np.zeros((9, 3)), n_node) + # the layout is untouched: static, every grid point present + assert bigger.vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + # every grid point is still a valid wave number of the larger cell; the + # mask keeps the whole k_max sphere the grid reaches + k2 = np.sum(bigger.vec**2, axis=-1) + real = (np.arange(bigger.mask[0].size) < nk).reshape(bigger.mask.shape[1:]) + np.testing.assert_array_equal(bigger.mask, real & (k2 <= k_max * k_max)) + # the larger cell activates more of the shared grid than the original did + original = build_fixed_k_grid(box, bounds, k_max, np.zeros((9, 3)), n_node) + assert bigger.mask.sum() > original.mask.sum() + # a larger wave-number cutoff (a drifting kernel) truncates the same way + with pytest.warns(UserWarning, match="truncated"): + build_fixed_k_grid(box, bounds, k_max * 2.0, np.zeros((9, 3)), n_node) + + +def test_fixed_k_grid_mode_configuration() -> None: + """The fixed mode requires periodicity, and the bounds the fixed mode.""" + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="fixed", + ) + with pytest.raises(ValueError, match="k_grid_bounds"): + _periodic_long_range(k_grid_bounds=[3, 3, 3]) + # explicit bounds build the grid at construction, no statistics needed + long_range = _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[2, 2, 2]) + np.testing.assert_array_equal(long_range.k_grid_bounds, [2, 2, 2]) + with pytest.raises(ValueError, match="non-negative"): + _periodic_long_range(k_grid_mode="fixed", k_grid_bounds=[0, 0, 0]) + + +def test_fixed_k_grid_stats_size_the_grid() -> None: + """The data statistics size the grid by the largest sampled cell.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + # the grid is not sized before the statistics run + coord = np.zeros((9, 3)) + atype = np.zeros(9, dtype=np.int64) + with pytest.raises(ValueError, match="not sized yet"): + _module_topologies(long_range, coord, np.array([4, 5]), _cells(), atype) + long_range.compute_input_stats( + [{"box": _cells()[:1].reshape(1, 9)}, {"box": _cells()}] + ) + k_max = long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + from deepmd.dpmodel.utils.neighbor_graph import ( + k_grid_bounds_of_cells, + ) + + expected = np.max(k_grid_bounds_of_cells(_cells(), k_max), axis=0) + np.testing.assert_array_equal(long_range.k_grid_bounds, expected) + topologies = _module_topologies( + long_range, coord, np.array([4, 5]), _cells(), atype + ) + nk = np.prod(2 * expected + 1) - 1 + assert topologies["k_grid"].vec.shape == (2, -(-nk // K_BLOCK), K_BLOCK, 3) + # statistics without a cell cannot size the grid + unsized = _periodic_long_range(k_grid_mode="fixed") + with pytest.raises(ValueError, match="no simulation cell"): + unsized.compute_input_stats([{"coord": np.zeros((1, 3))}]) + + +def test_fixed_k_grid_round_trip() -> None: + """The record keeps the grid mode and the sized grid.""" + long_range = _periodic_long_range(k_grid_mode="fixed") + long_range.compute_input_stats([{"box": _cells()}]) + record = long_range.serialize() + assert record["k_grid_mode"] == "fixed" + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "fixed" + np.testing.assert_array_equal(restored.k_grid_bounds, long_range.k_grid_bounds) + # the ragged mode serializes no grid bounds + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert plain.k_grid_mode == "ragged" + assert plain.k_grid_bounds is None + # a periodic module evaluates on the mesh unless told otherwise, and a + # periodic record without a mode is one of the explicit wave numbers + charge_net = LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1) + default = LatentChargeLongRange(charge_net, LESKernel(), None, periodic=True) + assert default.k_grid_mode == "pme" + assert default.required_topologies() == ("all_pairs", "pme_mesh") + record = default.serialize() + assert record["k_grid_mode"] == "pme" + record.pop("k_grid_mode") + assert LatentChargeLongRange.deserialize(record).k_grid_mode == "ragged" + + +# --------------------------------------------------------------------------- +# Decoupled k grid +# --------------------------------------------------------------------------- + + +def test_decoupled_k_grid_requires_periodic() -> None: + """Decoupling is only meaningful for a periodic module.""" + with pytest.raises(ValueError, match="decouple_k_grid"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + decouple_k_grid=True, + ) + + +def test_decoupled_k_grid_ignores_kernel_drift() -> None: + """A decoupled grid keeps the k points of the initial kernel however the + kernel parameters move; a coupled one re-derives them every batch. + """ + coord = np.random.default_rng(3).normal(size=(5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + decoupled = _periodic_long_range(decouple_k_grid=True) + coupled = _periodic_long_range() + # the frozen cutoff is the cutoff of the initial kernel + assert decoupled.frozen_k_max == pytest.approx( + LESKernel(alpha=0.3).kspace_cutoff(1e-8) + ) + initial = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + reference = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(initial.vec, reference.vec) + # a larger alpha moves the cutoff out, growing the coupled grid + for module in (decoupled, coupled): + module.kernel.log_alpha = np.log(1.5) + drifted = _module_topologies(decoupled, coord, n_node, box, atype)["k_grid"] + moved = _module_topologies(coupled, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(drifted.vec, initial.vec) + np.testing.assert_array_equal(drifted.mask, initial.mask) + assert np.sum(moved.mask) > np.sum(initial.mask) + + +def test_decoupled_fixed_grid_keeps_its_mask() -> None: + """The mask of a decoupled fixed grid ignores the kernel drift.""" + coord = np.zeros((5, 3)) + atype = np.zeros(5, dtype=np.int64) + n_node = np.array([5]) + box = np.eye(3)[None] * 8.0 + + long_range = _periodic_long_range( + k_grid_mode="fixed", k_grid_bounds=[5, 5, 5], decouple_k_grid=True + ) + before = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] + # a smaller alpha moves the cutoff in; a coupled mask would tighten + long_range.kernel.log_alpha = np.log(0.05) + after = _module_topologies(long_range, coord, n_node, box, atype)["k_grid"] + np.testing.assert_array_equal(after.mask, before.mask) + np.testing.assert_array_equal(after.vec, before.vec) + + +def test_decoupled_k_grid_round_trip() -> None: + """The frozen cutoff survives serialization, past the kernel drift.""" + long_range = _periodic_long_range(decouple_k_grid=True) + frozen = long_range.frozen_k_max + long_range.kernel.log_alpha = np.log(1.5) + record = long_range.serialize() + assert record["decouple_k_grid"] is True + assert record["frozen_k_max"] == pytest.approx(frozen) + restored = LatentChargeLongRange.deserialize(record) + assert restored.decouple_k_grid + assert restored.frozen_k_max == pytest.approx(frozen) + # ...even though the restored (drifted) kernel would derive a different one + assert restored.kernel.kspace_cutoff(1e-8) != pytest.approx(frozen) + # the default stays coupled and serializes nothing + plain = LatentChargeLongRange.deserialize(_periodic_long_range().serialize()) + assert not plain.decouple_k_grid + assert plain.frozen_k_max is None + + +# --------------------------------------------------------------------------- +# PME mesh mode +# --------------------------------------------------------------------------- + + +def _pme_long_range(**kwargs: object) -> LatentChargeLongRange: + kwargs.setdefault("k_grid_mode", "pme") + return _periodic_long_range(**kwargs) + + +def test_pme_mesh_grid_build_and_flatten() -> None: + """The mesh grid packs the cells and the mesh, flags the frames and flattens.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + PME_MESH, + build_topologies, + flatten_topologies, + mesh_dims_of_cells, + unflatten_topologies, + ) + + coord = np.random.default_rng(1).normal(size=(4, 3)) + atype = np.array([0, 1, 1, 0]) + box = _cells() + params = {PME_MESH: {"spacing": 1.0}} + tops = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + box, + atype=atype, + topology_params=params, + ) + mesh = tops[PME_MESH] + np.testing.assert_array_equal(mesh.box, box) + assert mesh.mask.shape == (2,) and mesh.mask.all() + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, 1.0) + # the modes of an axis are its integer frequencies in FFT order + n1 = dims[0] + expected = np.arange(n1, dtype=float) + np.testing.assert_array_equal( + mesh.mode_1, np.where(expected > n1 // 2, expected - n1, expected) + ) + assert tops["all_pairs"].index.shape[1] == 4 # guards only + flat = flatten_topologies(tops) + assert len(flat) == 3 + 8 + back = unflatten_topologies(("all_pairs", PME_MESH), flat) + np.testing.assert_array_equal(back[PME_MESH].box, box) + np.testing.assert_array_equal(back[PME_MESH].mask, mesh.mask) + np.testing.assert_array_equal(back[PME_MESH].mode_3, mesh.mode_3) + # a non-periodic batch: real pairs, masked-out mesh + tops_np = build_topologies( + ("all_pairs", PME_MESH), + coord, + np.array([2, 2]), + None, + atype=atype, + topology_params=params, + ) + assert tops_np["all_pairs"].mask.any() + assert not tops_np[PME_MESH].mask.any() + # missing build parameters are an error + with pytest.raises(ValueError, match="spacing"): + build_topologies((PME_MESH,), coord, np.array([4]), box, atype=atype) + + +def test_pme_mesh_dims_resolve_every_cell() -> None: + """The mesh is the smallest even 7-smooth one resolving every cell.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, + ) + + for spacing in (0.37, 0.5, 1.1): + dims = mesh_dims_of_cells(_cells(), spacing) + longest = np.max(np.linalg.norm(_cells(), axis=-1), axis=0) + for n, a in zip(dims, longest, strict=True): + assert n % 2 == 0 and n >= a / spacing + m = n + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + assert m == 1 + # no smaller even 7-smooth size reaches the spacing + for smaller in range(n - 2, 1, -2): + m = smaller + for p in (2, 3, 5, 7): + while m % p == 0: + m //= p + if m == 1: + assert smaller < a / spacing + break + # the larger cell alone sizes the mesh of the batch + assert mesh_dims_of_cells(_cells(), 0.5) == mesh_dims_of_cells(_cells()[1:], 0.5) + + +def _pme_energy(module, charge, coord, box) -> float: + from deepmd.dpmodel.long_range.reciprocal import ( + mesh_energy, + ) + from deepmd.dpmodel.utils.neighbor_graph import ( + build_mesh_grid, + ) + + mesh = build_mesh_grid(box[None], module.mesh_spacing(), coord) + frame_id = np.zeros(coord.shape[0], dtype=np.int64) + energy = mesh_energy(charge, mesh, frame_id, module.kernel, module.pme_order) + return float(np.sum(energy)) + + +def test_pme_energy_converges_to_direct_reference() -> None: + """The PME energy approaches the direct k-space sum as the mesh refines.""" + coord = 3.0 + np.random.default_rng(7).uniform(0.0, 1.4, (6, 3)) + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + charge = np.array([[1.0], [-0.5], [1.0], [-0.5], [-0.5], [1.0]]) + charge = charge - charge.mean(axis=0, keepdims=True) + + kernel = LESKernel(alpha=0.3) + # converged direct k-space reference (tol 1e-10) + k_max = kernel.kspace_cutoff(1e-10) + rec = 2.0 * np.pi * np.linalg.inv(box) + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid(*[np.arange(-n, n + 1) for n in n_d], indexing="ij") + n_all = np.stack(grids, -1).reshape(-1, 3) + k = n_all @ rec.T + k2 = (k * k).sum(-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(-1) + k_hat = kernel.kspace(k2) + volume = abs(np.linalg.det(box)) + k0 = 2.0 * 0.3 / np.sqrt(np.pi) + e_ref = COULOMB_CONSTANT * ( + (k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum() + ) + errors = [] + for spacing in (2.0, None): + module = LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + kernel, + None, + periodic=True, + k_grid_mode="pme", + pme_spacing=spacing, + ) + errors.append(abs(_pme_energy(module, charge, coord, box) - e_ref)) + assert errors[1] < errors[0] # refinement improves the energy + # the default mesh reaches the accuracy it is chosen for + assert errors[1] < 1e-6 * abs(e_ref) + + +def test_pme_mesh_follows_the_cells_of_the_batch() -> None: + """The mesh of a batch resolves its cells at the spacing, without statistics.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, + ) + + long_range = _pme_long_range() + coord = np.zeros((9, 3)) + atype = np.zeros(9, dtype=np.int64) + small = _module_topologies(long_range, coord, np.array([4, 5]), _cells()[:1], atype) + large = _module_topologies( + long_range, coord, np.array([4, 5]), 2.0 * _cells(), atype + ) + spacing = long_range.mesh_spacing() + assert spacing == pytest.approx( + 0.75 * np.pi / long_range.kernel.kspace_cutoff(long_range.kspace_tolerance) + ) + for topologies, box in ((small, _cells()[:1]), (large, 2.0 * _cells())): + mesh = topologies["pme_mesh"] + dims = tuple(m.shape[0] for m in (mesh.mode_1, mesh.mode_2, mesh.mode_3)) + assert dims == mesh_dims_of_cells(box, spacing) + assert mesh.mask.all() + + +def test_pme_mode_configuration() -> None: + """The pme mode requires periodicity, and so do the mesh parameters.""" + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + k_grid_mode="pme", + ) + for kwargs in ({"pme_spacing": 1.0}, {"pme_order": 6}): + with pytest.raises(ValueError, match="periodic"): + LatentChargeLongRange( + LatentChargeFitting("latent_charge", 2, 4, 1, neuron=[4], seed=1), + LESKernel(), + None, + **kwargs, + ) + # every periodic mode takes them: a frozen model evaluates on the mesh + assert _periodic_long_range(pme_spacing=1.0).mesh_spacing() == 1.0 + assert _periodic_long_range().pme_order == 8 + with pytest.raises(ValueError, match="pme_order"): + _pme_long_range(pme_order=1) + with pytest.raises(ValueError, match="pme_spacing"): + _pme_long_range(pme_spacing=-1.0) + with pytest.raises(ValueError, match="k_grid_mode"): + _periodic_long_range(k_grid_mode="dense") + + +def test_pme_round_trip() -> None: + """The record keeps the pme mode, the order and the spacing.""" + long_range = _pme_long_range(pme_spacing=1.0, pme_order=5) + record = long_range.serialize() + assert record["k_grid_mode"] == "pme" + assert record["pme_order"] == 5 + assert record["pme_spacing"] == 1.0 + restored = LatentChargeLongRange.deserialize(record) + assert restored.k_grid_mode == "pme" + assert restored.pme_order == 5 + assert restored.pme_spacing == 1.0 + # the default order is recorded too, so that it survives a change of default + assert _periodic_long_range().serialize()["pme_order"] == 8 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..8fa069eb05 --- /dev/null +++ b/source/tests/infer/gen_dpa4c_long_range.py @@ -0,0 +1,264 @@ +#!/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, and over the images of the cell of a periodic one when +it serves periodic systems, so the C++ interface hands the model what the sum +needs beyond the neighbor list: the graph lower of an uncompressed model +receives the pair list, or the mesh of the cell, that the interface builds +from the positions of the atoms, and the compact canonical lower of a +compressed model reads the positions and the cell themselves. One archive +covers each lower and each kind of module: + +- ``deeppot_dpa4c_long_range.pt2`` and + ``deeppot_dpa4c_long_range_periodic.pt2``: uncompressed, fp64, graph lower, + which compiles on the CPU; +- ``deeppot_dpa4c_long_range_compressed.pt2`` and + ``deeppot_dpa4c_long_range_periodic_compressed.pt2``: compressed, fp32, + compact canonical lower, whose operators are CUDA only; they are generated + for the CUDA variant of the C++ suite (``DP_VARIANT=cuda``). + +Every model conditions the descriptor on the charge state and constrains 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. For a +non-periodic module they are taken without a cell, at the default state and +at ``[1, 2]`` (sections ``default`` and ``explicit``); for a periodic module, +in the cell :data:`BOX` and without a cell, at the neutral default state +(sections ``periodic`` and ``cluster``). The six atoms form two groups about +9 Å apart, beyond the cutoff of the descriptor, which only the long-range +module couples; in the cell, the first group meets its images across the +boundary. +""" + +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) +# A triclinic cell, row-major with the lattice vectors as rows. +BOX = np.array( + [[13.0, 0.0, 0.0], [0.6, 12.5, 0.0], [-0.4, 0.8, 13.5]], dtype=np.float64 +).reshape(1, 9) + + +def _periodic(config: dict) -> dict: + """The configuration with a long-range module that serves periodic systems.""" + config = copy.deepcopy(config) + config["long_range"]["periodic"] = True + return config + + +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, periodic: bool = False) -> 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"] + expected = ["all_pairs", "pme_mesh"] if periodic else ["all_pairs"] + assert md["topologies"] == expected, md["topologies"] + if lower == "graph": + 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, periodic: bool = False) -> None: + """Evaluate the reference cases on the archive and write the C++ sidecar.""" + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + if periodic: + cases = (("periodic", BOX, None), ("cluster", None, None)) + else: + cases = (("default", None, None), ("explicit", None, _OTHER_CHG_SPIN)) + sections = {} + energies = {} + for name, box, charge_spin in cases: + condition = {} if charge_spin is None else {"charge_spin": charge_spin} + e, f, v, ae, av = dp.eval(_COORDS, box, _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 cases, which needs a real difference. + first, second = (name for name, _, _ in cases) + assert abs(energies[second] - energies[first]) > 1e-3, ( + "the two cases must differ, or a call that ignored the cell or 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, config: dict) -> 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(config)) + model.atomic_model.enable_compression(0.5) + data = { + "model": model.serialize(), + "model_def_script": 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__) + for suffix, periodic in (("", False), ("_periodic", True)): + config = _periodic(CONFIG) if periodic else CONFIG + stem = os.path.join(base_dir, f"deeppot_dpa4c_long_range{suffix}") + data = { + "model": _model_dict(config), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting to {stem}.pt2 (lower_kind='graph') ...") # noqa: T201 + deserialize_to_file( + f"{stem}.pt2", data, do_atomic_virial=True, lower_kind="graph" + ) + _check_metadata(f"{stem}.pt2", "graph", periodic) + _write_ref(f"{stem}.pt2", f"{stem}.expected", periodic) + + if os.environ.get("DP_VARIANT") == "cuda": + compressed = f"{stem}_compressed" + compressed_config = ( + _periodic(COMPRESSED_CONFIG) if periodic else COMPRESSED_CONFIG + ) + _export_compressed(f"{compressed}.pt2", compressed_config) + _check_metadata(f"{compressed}.pt2", "dpa4c_canonical", periodic) + _write_ref(f"{compressed}.pt2", f"{compressed}.expected", periodic) + + 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_sum.py b/source/tests/pt_expt/kernels/test_long_range_sum.py new file mode 100644 index 0000000000..b61c7390c9 --- /dev/null +++ b/source/tests/pt_expt/kernels/test_long_range_sum.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Correctness of the operator of the latent-charge long range. + +For non-periodic frames 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. + +For periodic frames the CUDA kernels evaluate the mesh sum in float32 on +B-spline weights formed in float64, and are held to the float64 mesh +evaluation of the training path, differentiated by autograd. +""" + +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.dpmodel.utils.neighbor_graph import ( + mesh_dims_of_cells, +) +from deepmd.pt_expt.kernels.long_range import ( + _reference, + kernel_operator_arguments, + latent_charge_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 long-range 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 _reference_of( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[torch.Tensor, ...]: + """The float64 reference of the operator.""" + coord, charge, node_mask, n_node, total_charge = system + code, coefficient = kernel_operator_arguments(kernel) + if box is None: + box = torch.zeros(n_node.shape[0], 9, dtype=torch.float64) + return _reference( + coord, + charge.double(), + node_mask, + n_node, + total_charge, + box, + coefficient, + torch.zeros(1, dtype=torch.int64), + code, + channel, + COULOMB_CONSTANT, + order, + spacing, + ) + + +def _evaluate( + system: tuple[torch.Tensor, ...], + kernel: Any, + channel: int, + box: torch.Tensor | None = None, + order: int = 8, + spacing: float = 1.0, +) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]: + actual = latent_charge_sum( + *system, + kernel, + channel, + COULOMB_CONSTANT, + box=box, + order=order, + spacing=spacing, + ) + return actual, _reference_of(system, kernel, channel, box, order, spacing) + + +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, ...], + tolerance: float = 2e-5, +) -> None: + """Every output agrees with the float64 reference, 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 <= tolerance * 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_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_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) + atom_energy, charge_gradient, force, _, _ = _reference_of( + (coord, charge, node_mask, n_node, total_charge), + KERNELS[kernel_name](), + channel, + ) + 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, with or without cells.""" + system = _system([3000], 1, 0, 5, 0.0) + for box in (None, _cells(1, 40.0)): + first = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + second = latent_charge_sum( + *system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box + ) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) + + +# --------------------------------------------------------------------------- +# Periodic frames +# --------------------------------------------------------------------------- + + +def _cells(frame_count: int, length: float, seed: int = 3) -> torch.Tensor: + """Triclinic cells of about ``length`` per side, row-major, one row per frame.""" + rng = np.random.default_rng(seed) + cells = [] + for _ in range(frame_count): + cell = np.diag(length * rng.uniform(0.8, 1.2, 3)) + cell[1, 0], cell[2, 0], cell[2, 1] = length * rng.uniform(-0.2, 0.2, 3) + cells.append(cell.reshape(9)) + return torch.tensor(np.asarray(cells), dtype=torch.float64) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +@pytest.mark.parametrize( + ("n_node", "channels", "padding", "channel", "order"), + [ + ([37, 1, 300], 1, 5, 0, 8), + ([37, 1, 300], 1, 5, -1, 4), + ([500], 3, 0, -1, 6), + ([5, 129], 8, 3, 2, 8), + ], +) +def test_mesh_matches_float64_reference( + kernel_name: str, + n_node: list[int], + channels: int, + padding: int, + channel: int, + order: int, +) -> None: + """Every output of a periodic frame agrees with the mesh evaluation of training. + + The positions extend beyond the cell, so that the fractional coordinates + wrap. + """ + system = _system(n_node, channels, padding, 7, -5.0) + box = _cells(len(n_node), 12.0) + actual, expected = _evaluate( + system, KERNELS[kernel_name](), channel, box, order, 0.6 + ) + _assert_close(actual, expected, 1e-6) + if padding: + for value in actual[:4]: + assert torch.all(value[-padding:] == 0) + + +@pytest.mark.parametrize("kernel_name", ["les", "sog"]) +def test_mesh_virial_is_the_strain_derivative(kernel_name: str) -> None: + """The virial of a periodic frame is the strain derivative of its energy. + + Straining the positions and the cell together by ``I + eps`` changes the + energy by ``-W : eps``; the atomic virial of the charged nodes sums to W. + The central difference takes the float64 mesh energy, whose rounding stays + far below that of the float32 outputs of the operator. + """ + system = _system([60], 2, 0, 13, 0.0) + coord, charge, node_mask, n_node, total_charge = system + box = _cells(1, 10.0) + kernel = KERNELS[kernel_name]() + _, _, _, atom_virial, virial = latent_charge_sum( + *system, kernel, 0, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + h = 1e-5 + finite = np.zeros((3, 3)) + for a in range(3): + for b in range(3): + energies = [] + for sign in (1.0, -1.0): + strain = torch.eye(3, dtype=torch.float64) + strain[a, b] += sign * h + strained = ( + coord @ strain.T.to(coord.device), + charge, + node_mask, + n_node, + total_charge, + ) + cell = box.reshape(3, 3) @ strain.T + # Both strained cells keep the mesh of the cell. + assert mesh_dims_of_cells(cell.reshape(1, 9).numpy(), 0.5) == ( + mesh_dims_of_cells(box.numpy(), 0.5) + ) + energies.append( + _reference_of(strained, kernel, 0, cell.reshape(1, 9), 8, 0.5)[ + 0 + ].sum() + ) + finite[a, b] = -(energies[0] - energies[1]).item() / (2 * h) + scale = np.abs(finite).max() + np.testing.assert_allclose( + virial[0].double().cpu().numpy().reshape(3, 3), finite, atol=2e-6 * scale + ) + np.testing.assert_allclose( + atom_virial.double().sum(0).cpu().numpy().reshape(3, 3), + finite, + atol=2e-6 * scale, + ) + + +@pytest.mark.parametrize("order", [2, 3, 5, 7, 9, 10, 11, 12]) +def test_mesh_orders(order: int) -> None: + """Every B-spline order agrees with the mesh evaluation of training. + + The atomic energies of these small cells are a fraction of the self + energy, to a few 1e-7 of which the float32 mesh resolves them. + """ + system = _system([40, 23], 2, 3, 5, -5.0) + box = _cells(2, 9.0) + actual, expected = _evaluate(system, KERNELS["sog"](), 1, box, order, 0.7) + _assert_close(actual, expected, 1e-5) + + +def test_mesh_stencil_wider_than_the_axis() -> None: + """A stencil wider than an axis of the mesh wraps around it more than once.""" + system = _system([12], 1, 0, 9, 0.0) + box = _cells(1, 3.0) + # Four points per axis, against eight of the stencil. + assert mesh_dims_of_cells(box.numpy(), 1.0) == (4, 4, 4) + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 1.0) + _assert_close(actual, expected, 1e-6) + + +@pytest.mark.parametrize("periodic", [False, True]) +def test_frame_without_charged_nodes(periodic: bool) -> None: + """A frame whose nodes carry no charge contributes zeros, the others are unaffected.""" + n_node = [40, 6, 50] + coord, charge, node_mask, n_node_t, total_charge = _system(n_node, 2, 0, 5, 0.0) + node_mask[40:46] = False + system = (coord, charge, node_mask, n_node_t, total_charge) + box = _cells(3, 10.0) if periodic else None + actual, expected = _evaluate(system, KERNELS["les"](), 0, box, 8, 0.6) + _assert_close(actual, expected, 1e-6) + for value in actual[:4]: + assert torch.all(value[40:46] == 0) + assert torch.all(actual[4][1] == 0) + + +def test_mesh_frames_are_independent() -> None: + """Every periodic frame is evaluated on its own mesh, bit for bit.""" + n_node = [64, 200] + system = _system(n_node, 2, 0, 11, 0.0) + coord, charge, node_mask, n_node_t, total_charge = system + box = _cells(2, 9.0) + kernel = KERNELS["sog"]() + together = latent_charge_sum( + *system, kernel, 1, COULOMB_CONSTANT, box=box, spacing=0.5 + ) + first = n_node[0] + alone = latent_charge_sum( + coord[:first], + charge[:first], + node_mask[:first], + n_node_t[:1], + total_charge[:1], + kernel, + 1, + COULOMB_CONSTANT, + box=box[:1], + spacing=0.5, + ) + for both, single in zip(together, alone, strict=True): + torch.testing.assert_close(both[: single.shape[0]], single, rtol=0, atol=0) + + +def test_mixed_frames_are_rejected() -> None: + """A call holds periodic frames or non-periodic ones, not both.""" + system = _system([10, 10], 1, 0, 1, 0.0) + box = _cells(2, 9.0) + box[1] = 0.0 + with pytest.raises(RuntimeError, match="periodic"): + latent_charge_sum(*system, KERNELS["les"](), 0, COULOMB_CONSTANT, box=box) 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..c48dc0ae4c --- /dev/null +++ b/source/tests/pt_expt/model/test_long_range.py @@ -0,0 +1,2002 @@ +# 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) + # a periodic module shares its composition with no other long-range module + periodic = copy.deepcopy(child) + periodic["long_range"]["periodic"] = True + for models in ([periodic, child], [periodic, periodic]): + config = copy.deepcopy(composition) + config.pop("pair_exclude_types") + config["models"] = copy.deepcopy(models) + with pytest.raises(NotImplementedError, match="serves periodic systems"): + get_model(config) + config["models"] = [copy.deepcopy(child), copy.deepcopy(child)] + get_model(config) + # a periodic module takes no pair exclusion, alone or in a composition + config = copy.deepcopy(composition) + config["models"] = [copy.deepcopy(periodic)] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + get_model(config) + config = copy.deepcopy(periodic) + config["type"] = "standard" + config["pair_exclude_types"] = [[0, 1]] + with pytest.raises(NotImplementedError, match="pair_exclude_types"): + 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("grid_mode", ["ragged", "pme"]) +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_checkpoint_inference_matches_direct_forward( + tmp_path, kernel: str, grid_mode: str +) -> None: + """DeepPot on a periodic checkpoint builds the periodic topology and matches.""" + if grid_mode == "pme": + model = _periodic_pme_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + else: + model = _periodic_model(kernel) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + checkpoint = tmp_path / f"periodic_{kernel}_{grid_mode}.pt" + wrapper = ModelWrapper(model, model_params=config) + torch.save({"model": wrapper.state_dict()}, checkpoint) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = model(coord_t, atype_t, box=box_t, do_atomic_virial=True) + dp = DeepPot(str(checkpoint), auto_batch_size=False) + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + actual = dict( + zip( + names, + dp.eval( + coord, + box, + atype, + atomic=True, + mixed_type=True, + ), + strict=True, + ) + ) + for name in names: + np.testing.assert_allclose( + actual[name], + expected[name].detach().cpu().numpy(), + rtol=1e-9, + atol=1e-9, + err_msg=name, + ) + + +@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": + # The pair topology follows the thirteen inputs of the graph lower. + signature = torch.export.load(str(frozen)).graph_signature + assert len(signature.user_inputs) == 13 + 3 + 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 long-range 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 + long-range 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 + + +# --------------------------------------------------------------------------- +# Periodic (k-space) evaluation +# --------------------------------------------------------------------------- + + +def _periodic_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model over the explicit wave numbers of every frame.""" + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + return get_model(config).to(env.DEVICE).eval() + + +def _periodic_case() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A compact cluster (all atoms within rcut of each other) in a large triclinic cell.""" + box = np.array([[9.0, 0.4, 0.0], [0.0, 8.5, 0.3], [0.0, 0.0, 9.5]]) + rng = np.random.default_rng(7) + coord = 3.0 + rng.uniform(0.0, 1.4, (6, 3)) + atype = np.array([0, 1, 0, 1, 1, 0], dtype=np.int64) + return coord[None], atype[None], box[None] + + +def _kspace_reference( + coord: np.ndarray, charge: np.ndarray, box: np.ndarray, kernel: str +) -> float: + """Independent k-space evaluation of the periodic long-range energy, in eV. + + E = k_e [ 1/(2V) sum_{k != 0} K_hat(k) |rho(k)|^2 - K(0)/2 sum_i |q_i|^2 ]. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) # columns g_d: k(n) = rec @ n + # the module's default kspace_tolerance=1e-8 sets the truncation + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + n_d = np.ceil(k_max * np.linalg.norm(box, axis=1) / (2 * np.pi)).astype(int) + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in n_d], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T # (nk, N) + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge # (nk, C) + s_sq = np.abs(rho) ** 2 + s_sq = s_sq.sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_energy_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the periodic energy is the documented k-space sum.""" + model = _periodic_model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose( + out["atom_energy"].sum().item(), expected, rtol=1e-9, atol=1e-9 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_force_virial_hessian_match_finite_differences(kernel: str) -> None: + """Force, strain virial and Hessian derive from the complete periodic energy.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + 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, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-6 + ) + + # The virial is -dE/d(strain) at zero strain, for r -> (I + eps) r and + # B -> B (I + eps)^T (the strain family of the k-space strain leaf). + 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).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-5 + ) + + hessian_model = _periodic_model(kernel) + hessian_model.load_state_dict(model.state_dict()) + hessian_model.enable_hessian() + hessian = ( + hessian_model(coord_t, atype_t, box=box_t)["hessian"].detach().cpu().numpy() + ) + hessian = hessian.reshape(18, 18) + hessian_fd = np.zeros((18, 18)) + for kk in range(18): + plus, minus = coord.copy(), coord.copy() + plus.reshape(-1)[kk] += h + minus.reshape(-1)[kk] -= h + f_plus = ( + model(_tensor(plus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + f_minus = ( + model(_tensor(minus), atype_t, box=box_t)["force"].detach().cpu().numpy() + ) + hessian_fd[:, kk] = -(f_plus - f_minus).reshape(-1) / (2 * h) + np.testing.assert_allclose(hessian, hessian_fd, atol=1e-4) + np.testing.assert_allclose(hessian, hessian.T, atol=1e-10) + + +@pytest.mark.parametrize("grid_mode", ["ragged", "pme"]) +def test_trace_samples_give_every_axis_its_own_length(grid_mode: str) -> None: + """No two dynamic axes of the trace-time topologies share a length, nor take 2 or 3.""" + from deepmd.pt_expt.utils.serialization import ( + build_synthetic_topologies, + ) + + model = ( + _periodic_model("les") if grid_mode == "ragged" else _periodic_pme_model("les") + ) + nframes, n_node, n_pair = 5, 7, 11 + topologies = build_synthetic_topologies( + model, n_node, n_pair, dtype=torch.float64, device=env.DEVICE, nframes=nframes + ) + lengths = {("", "frame"): nframes, ("", "node"): n_node} + for name, topology in topologies.items(): + for field, axes in type(topology).AXES.items(): + for axis, label in axes.items(): + key = ("", label) if label in ("frame", "node") else (name, label) + size = getattr(topology, field).shape[axis] + assert lengths.setdefault(key, size) == size, (name, field, label) + sizes = list(lengths.values()) + assert len(set(sizes)) == len(sizes), lengths + assert not {2, 3} & set(sizes), lengths + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the periodic model like the eager one.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_model_serves_non_periodic_batches(kernel: str) -> None: + """A periodic-enabled model keeps its non-periodic evaluation.""" + model = _periodic_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict()) + coord, atype = _batch() + periodic_enabled = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(periodic_enabled["energy"], plain_out["energy"]) + torch.testing.assert_close(periodic_enabled["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_frames_are_independent(kernel: str) -> None: + """Each frame's k grid and energy are independent of the rest of the batch.""" + model = _periodic_model(kernel) + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.1], axis=0) + batched = model(_tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2)) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_serialization_round_trip(kernel: str) -> None: + """The record keeps the periodic flag and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + model = _periodic_model(kernel) + assert model.atomic_model.required_topologies() == ("all_pairs", "k_grid") + record = model.serialize() + assert record["long_range"]["periodic"] is True + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.required_topologies() == ("all_pairs", "k_grid") + coord, atype, box = _periodic_case() + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common( + coord, + atype, + box, + neighbor_graph_method="dense", + ) + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +def _periodic_frames() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Two frames of the periodic case, the second strained and shifted.""" + coord, atype, box = _periodic_case() + strain = np.array([[1.04, 0.02, 0.0], [0.0, 0.97, 0.01], [0.03, 0.0, 1.02]]) + coord2 = np.concatenate([coord, coord @ strain.T + 0.3], axis=0) + box2 = np.concatenate([box, box @ strain.T], axis=0) + return coord2, np.concatenate([atype, atype], axis=0), box2 + + +@pytest.mark.parametrize( + ("kernel", "grid_mode", "suffix"), + [("les", "ragged", ".pt2"), ("sog", "pme", ".pte"), ("sog", "fixed", ".pt2")], +) +def test_periodic_frozen_model_matches_mesh_evaluation( + tmp_path, kernel: str, grid_mode: str, suffix: str +) -> None: + """A frozen periodic model evaluates every periodic frame on the mesh of its cell. + + Whatever its training mode, the frozen model reads the mesh of the cells + of a batch as inputs of dynamic size, and matches the model evaluated on + the mesh; a batch without a cell keeps the direct sum. + """ + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = grid_mode + if grid_mode == "fixed": + # Bounds that hold every wave number of the cells within the cutoff. + config["long_range"]["k_grid_bounds"] = [14, 14, 14] + model = get_model(config).to(env.DEVICE).eval() + frozen = tmp_path / f"periodic_{kernel}_{grid_mode}{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", "pme_mesh"] + mesh_model = copy.deepcopy(model) + mesh_model.atomic_model.long_range.use_mesh_evaluation() + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + coord, atype, box = _periodic_frames() + for frame in (slice(0, 1), slice(0, 2)): + expected = mesh_model( + _tensor(coord[frame]), + _tensor(atype[frame], torch.int64), + box=_tensor(box[frame]), + do_atomic_virial=True, + ) + actual = dp.eval( + coord[frame], box[frame], atype[frame], atomic=True, mixed_type=True + ) + 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, + ) + # The mesh approximates the sum of the training mode. + trained = model( + _tensor(coord[:1]), _tensor(atype[:1], torch.int64), box=_tensor(box[:1]) + ) + np.testing.assert_allclose( + dp.eval(coord[:1], box[:1], atype[:1], mixed_type=True)[0], + trained["energy"].detach().cpu().numpy(), + rtol=1e-5, + ) + cluster, cluster_type = _batch() + expected = model(cluster, cluster_type) + actual = dp.eval( + cluster.cpu().numpy(), None, cluster_type.cpu().numpy(), mixed_type=True + ) + np.testing.assert_allclose( + actual[0], expected["energy"].detach().cpu().numpy(), rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize(("kernel", "source"), [("les", "neutral"), ("sog", None)]) +def test_periodic_compressed_frozen_model_matches_uncompressed( + tmp_path, kernel: str, source: str | None +) -> None: + """A compressed periodic model keeps the mesh evaluation through the freeze. + + On a CUDA device the compact canonical lower evaluates every periodic + frame on the mesh of its cell with the long-range operator; on the CPU + the graph lower reads the mesh. Both match the uncompressed model on the + mesh, for a periodic frame and for a cluster. + """ + config = _config(kernel, constraint=None if source is None else {"source": source}) + config["descriptor"].update(n_radial=8, precision="float32") + 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} + config["long_range"]["periodic"] = True + 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) + reference.atomic_model.long_range.use_mesh_evaluation() + model.atomic_model.enable_compression(0.5) + frozen = tmp_path / f"periodic_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 + names = ("energy", "force", "virial", "atom_energy", "atom_virial") + rng = np.random.default_rng(3) + box = np.array([[[12.0, 0.0, 0.0], [0.8, 11.0, 0.0], [-0.5, 0.6, 13.0]]]) + periodic = rng.uniform(0.0, 1.0, (1, 40, 3)) @ box[0] + cluster = rng.uniform(0.0, 7.0, (1, 40, 3)) + 100.0 + atype = rng.integers(0, 2, (1, 40)) + for coord, cell in ((periodic, box), (cluster, None)): + expected = reference( + _tensor(coord), + _tensor(atype, torch.int64), + box=None if cell is None else _tensor(cell), + do_atomic_virial=True, + ) + actual = dp.eval(coord, cell, atype, atomic=True) + for name, value in zip(names, actual, strict=True): + target = expected[name].detach().cpu().numpy().reshape(value.shape) + scale = np.abs(target).max() + # The tables of the compressed descriptor approximate its + # networks, to a few 1e-6 of the scale of every output. + assert np.abs(value - target).max() <= 2e-5 * scale + 1e-5, ( + name, + cell is None, + np.abs(value - target).max(), + scale, + ) + # A cell without volume marks no system, on either lower: the graph lower + # meets it in the neighbor search, the compact lower at its input. + with pytest.raises((ValueError, RuntimeError), match=r"no volume|singular"): + dp.eval(periodic, np.zeros((1, 9)), atype) + + +# --------------------------------------------------------------------------- +# Periodic fixed k-grid mode +# --------------------------------------------------------------------------- + + +def _periodic_fixed_model( + kernel: str, boxes: np.ndarray | None = None, **kwargs: object +) -> torch.nn.Module: + """A periodic model in the fixed k-grid mode, sized by the given cells.""" + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + model = get_model(config).to(env.DEVICE).eval() + if boxes is not None: + boxes = torch.as_tensor(np.asarray(boxes, dtype=np.float64).reshape(-1, 9)) + model.atomic_model.long_range.compute_input_stats([{"box": boxes}]) + return model + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_matches_ragged(kernel: str) -> None: + """The fixed grid with its mask evaluates like the ragged enumeration.""" + coord, atype, box = _periodic_case() + ragged = _periodic_model(kernel) + fixed = _periodic_fixed_model(kernel, box) + # the state dicts of the two grid modes differ by the fixed grid buffer + fixed.load_state_dict(ragged.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + expected = ragged(coord_t, atype_t, box=box_t) + got = fixed(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-9, atol=1e-12) + # the gradient reaches the charge network and the kernel without NaN, + # though masked points carry a singular kernel transform nowhere + fixed.train() + out = fixed(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [ + parameter.grad + for parameter in fixed.atomic_model.long_range.parameters() + if parameter.grad is not None + ] + assert grads and all(torch.isfinite(grad).all() for grad in grads) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_mixed_cells(kernel: str) -> None: + """One shared grid serves a batch of different cells, each masked to its sphere.""" + coord, atype, box = _periodic_case() + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + atype2 = np.concatenate([atype, atype], axis=0) + box2 = np.concatenate([box, box * 1.4], axis=0) + fixed = _periodic_fixed_model(kernel, box2) + ragged = _periodic_model(kernel) + ragged.load_state_dict(fixed.state_dict(), strict=False) + coord_t, atype_t, box_t = ( + _tensor(coord2), + _tensor(atype2, torch.int64), + _tensor(box2), + ) + batched = fixed(coord_t, atype_t, box=box_t) + expected = ragged(coord_t, atype_t, box=box_t) + torch.testing.assert_close( + batched["energy"], expected["energy"], rtol=1e-9, atol=1e-12 + ) + torch.testing.assert_close( + batched["force"], expected["force"], rtol=1e-9, atol=1e-12 + ) + # the batched energy of a frame equals its standalone evaluation + for frame in range(2): + alone = fixed( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=1e-10, atol=1e-10 + ) + + +def _kspace_reference_bounded( + coord: np.ndarray, + charge: np.ndarray, + box: np.ndarray, + kernel: str, + bounds: np.ndarray, +) -> float: + """The k-space reference of :func:`_kspace_reference` restricted to the + integer indices within ``bounds`` — the truncated sum of a fixed grid. + """ + rec = 2.0 * np.pi * np.linalg.inv(box) + if kernel == "les": + k_max = 2.0 * 0.8 * np.sqrt(np.log(1e8)) + else: + k_max = np.sqrt(2.0 * np.log(1e8)) / 0.8 + grids = np.meshgrid( + *[np.arange(-n, n + 1, dtype=float) for n in bounds], indexing="ij" + ) + n_all = np.stack(grids, axis=-1).reshape(-1, 3) + k = n_all @ rec.T + k2 = np.sum(k * k, axis=-1) + keep = (k2 > 0) & (k2 <= k_max * k_max) + k, k2 = k[keep], k2[keep] + phase = k @ coord.T + rho = (np.cos(phase) + 1j * np.sin(phase)) @ charge + s_sq = (np.abs(rho) ** 2).sum(axis=-1) + if kernel == "les": + k_hat = 4.0 * np.pi * np.exp(-k2 / (4.0 * 0.8**2)) / k2 + k0 = 2.0 * 0.8 / np.sqrt(np.pi) + else: + amplitude = np.array([0.7, -0.2, 0.1]) + width = np.array([0.8, 1.7, 3.2]) + k_hat = np.sum( + amplitude[None, :] + * (2.0 * np.pi) ** 1.5 + * width[None, :] ** 3 + * np.exp(-0.5 * k2[:, None] * width[None, :] ** 2), + axis=-1, + ) + k0 = amplitude.sum() + volume = abs(np.linalg.det(box)) + return float( + COULOMB_CONSTANT + * ((k_hat * s_sq).sum() / (2 * volume) - 0.5 * k0 * (charge * charge).sum()) + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_truncates_oversized_cells(kernel: str) -> None: + """A cell beyond the grid is evaluated with the grid alone, not an error.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + big_box = box * 3.0 + with pytest.warns(UserWarning, match="truncated"): + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(big_box)) + bounds = model.atomic_model.long_range.k_grid_bounds.detach().cpu().numpy() + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference_bounded(coord[0], charge, big_box[0], kernel, bounds) + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_needs_statistics_or_bounds(kernel: str) -> None: + """Without statistics or explicit bounds the grid is not sized; bounds skip them.""" + coord, atype, box = _periodic_case() + unsized = _periodic_fixed_model(kernel) + with pytest.raises(ValueError, match="not sized yet"): + unsized(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "fixed" + config["long_range"]["k_grid_bounds"] = [12, 12, 12] + model = get_model(config).to(env.DEVICE).eval() + np.testing.assert_array_equal( + model.atomic_model.long_range.k_grid_bounds.cpu().numpy(), [12, 12, 12] + ) + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + reference = _periodic_model(kernel) + reference.load_state_dict(model.state_dict(), strict=False) + expected = reference(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + torch.testing.assert_close(out["energy"], expected["energy"], rtol=1e-9, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_restores_from_state_dict(kernel: str) -> None: + """The state dict of a sized grid loads into an unsized model, as a checkpoint does.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + restored = _periodic_fixed_model(kernel) + restored.load_state_dict(model.state_dict()) + args = (_tensor(coord), _tensor(atype, torch.int64)) + expected = model(*args, box=_tensor(box)) + got = restored(*args, box=_tensor(box)) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(got[key], expected[key], rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_fixed_grid_serialization_round_trip(kernel: str) -> None: + """The record keeps the sized grid and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "fixed" + assert "k_grid_bounds" in record["long_range"] + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "fixed" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_fixed_grid_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the fixed-grid model like the eager one.""" + coord, atype, box = _periodic_case() + model = _periodic_fixed_model(kernel, box) + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + + +# --------------------------------------------------------------------------- +# Periodic decoupled k grid +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_decoupled_grid_ignores_kernel_drift(kernel: str) -> None: + """A decoupled model selects the k points with the initial kernel but + evaluates them with the live one: the drift moves the energy, not the grid. + """ + config = _config(kernel) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "ragged" + config["long_range"]["decouple_k_grid"] = True + model = get_model(config).to(env.DEVICE).eval() + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + long_range = model.atomic_model.long_range + grid_args = ( + coord_t.reshape(-1, 3), + torch.tensor([6], dtype=torch.int64, device=env.DEVICE), + box_t, + ) + build = model.atomic_model.build_topologies + before = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + energy_before = model(coord_t, atype_t, box=box_t)["energy"] + # move the kernel far enough that a coupled grid would change + with torch.no_grad(): + if kernel == "les": + long_range.kernel.log_alpha.add_(1.0) + else: + long_range.kernel.log_width.sub_(1.0) + after = build(*grid_args, atype=atype_t.reshape(-1))["k_grid"] + torch.testing.assert_close(after.vec, before.vec) + torch.testing.assert_close(after.mask, before.mask) + # the kernel values still follow the live parameters + energy_after = model(coord_t, atype_t, box=box_t)["energy"] + assert not torch.allclose(energy_after, energy_before) + # and the drifted kernel still takes a finite gradient through the graph + model.train() + out = model(coord_t, atype_t, box=box_t) + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + kernel_grads = [parameter.grad for parameter in long_range.kernel.parameters()] + assert kernel_grads and all( + grad is not None and torch.isfinite(grad).all() for grad in kernel_grads + ) + + +# --------------------------------------------------------------------------- +# Periodic PME mesh mode +# --------------------------------------------------------------------------- + + +def _periodic_pme_model(kernel: str, **kwargs: object) -> torch.nn.Module: + """A periodic model in the PME mesh mode, ``pme_spacing`` and ``pme_order`` in ``kwargs``.""" + pme_kwargs = { + key: kwargs.pop(key) for key in ("pme_spacing", "pme_order") if key in kwargs + } + config = _config(kernel, **kwargs) + config["long_range"]["periodic"] = True + config["long_range"]["k_grid_mode"] = "pme" + config["long_range"].update(pme_kwargs) + return get_model(config).to(env.DEVICE).eval() + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_matches_kspace_reference(kernel: str) -> None: + """With fixed charges the PME energy approaches the converged k-space sum.""" + model = _periodic_pme_model(kernel, constraint=None) + _set_constant_charges(model, [1.0, -0.5]) + coord, atype, box = _periodic_case() + out = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + charge = np.zeros((6, CHANNELS[kernel])) + charge[:, 0] = [1.0, -0.5, 1.0, -0.5, -0.5, 1.0] + expected = _kspace_reference(coord[0], charge, box[0], kernel) + # The default mesh approximates the reciprocal sum within a few parts in a + # million; the energy, less the self term, is several times smaller. + np.testing.assert_allclose(out["energy"].item(), expected, rtol=1e-5) + np.testing.assert_allclose(out["atom_energy"].sum().item(), expected, rtol=1e-5) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_force_virial_match_finite_differences(kernel: str) -> None: + """Force and strain virial are the exact derivatives of the PME energy.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + out = model(coord_t, atype_t, box=box_t) + + def energy(c: np.ndarray, b: np.ndarray) -> float: + return model(_tensor(c), atype_t, box=_tensor(b))["energy"].item() + + h = 1e-5 + force_fd = np.zeros((6, 3)) + for i in range(6): + 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, box) - energy(minus, box)) / (2 * h) + np.testing.assert_allclose( + out["force"][0].detach().cpu().numpy(), force_fd, atol=1e-5 + ) + + # the strain family of the mesh strain leaf: r -> (I + eps) r and + # B -> B (I + eps)^T leave the fractional coordinates invariant + 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).T + minus = coord @ (np.eye(3) - strain).T + box_plus = box @ (np.eye(3) + strain).T + box_minus = box @ (np.eye(3) - strain).T + virial_fd[a, b] = -(energy(plus, box_plus) - energy(minus, box_minus)) / ( + 2 * h + ) + np.testing.assert_allclose( + out["virial"][0].detach().cpu().numpy().reshape(3, 3), virial_fd, atol=1e-4 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_frames_are_independent(kernel: str) -> None: + """Frames on one mesh are independent; a larger cell sharing the batch refines it. + + The frames of a batch share the mesh that resolves the largest cell, so a + smaller cell evaluated beside a larger one sees a finer mesh, which moves + its energy by less than the accuracy of the mesh. + """ + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + atype2 = np.concatenate([atype, atype], axis=0) + coord2 = np.concatenate([coord, coord + 0.1], axis=0) + for box2, rtol in ( + (np.concatenate([box, box], axis=0), 1e-10), + (np.concatenate([box, box * 1.3], axis=0), 1e-6), + ): + batched = model( + _tensor(coord2), _tensor(atype2, torch.int64), box=_tensor(box2) + ) + for frame in range(2): + alone = model( + _tensor(coord2[frame : frame + 1]), + _tensor(atype2[frame : frame + 1], torch.int64), + box=_tensor(box2[frame : frame + 1]), + ) + torch.testing.assert_close( + alone["energy"][0], batched["energy"][frame], rtol=rtol, atol=1e-10 + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serves_non_periodic_batches(kernel: str) -> None: + """A PME model keeps its non-periodic pair evaluation.""" + model = _periodic_pme_model(kernel) + plain = _model(kernel) + plain.load_state_dict(model.state_dict(), strict=False) + coord, atype = _batch() + mesh_out = model(coord, atype) + plain_out = plain(coord, atype) + torch.testing.assert_close(mesh_out["energy"], plain_out["energy"]) + torch.testing.assert_close(mesh_out["force"], plain_out["force"]) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +def test_periodic_pme_serialization_round_trip(kernel: str) -> None: + """The record keeps the pme mode and the order, and evaluates identically.""" + from deepmd.dpmodel.model.base_model import BaseModel as BaseModelDP + + coord, atype, box = _periodic_case() + model = _periodic_pme_model(kernel, pme_order=6) + record = model.serialize() + assert record["long_range"]["k_grid_mode"] == "pme" + assert record["long_range"]["pme_order"] == 6 + reference = BaseModelDP.deserialize(record) + assert reference.atomic_model.long_range.k_grid_mode == "pme" + expected = model(_tensor(coord), _tensor(atype, torch.int64), box=_tensor(box)) + ret = reference.call_common(coord, atype, box, neighbor_graph_method="dense") + np.testing.assert_allclose( + ret["energy_redu"].reshape(-1), + expected["energy"].detach().cpu().numpy().reshape(-1), + rtol=1e-9, + atol=1e-9, + ) + + +@pytest.mark.parametrize("kernel", ["les", "sog"]) +@REQUIRES_SUPPORTED_COMPILE +def test_periodic_pme_compiled_matches_eager(kernel: str) -> None: + """The compiled graph lower evaluates the PME model like the eager one.""" + model = _periodic_pme_model(kernel) + coord, atype, box = _periodic_case() + coord_t, atype_t, box_t = ( + _tensor(coord), + _tensor(atype, torch.int64), + _tensor(box), + ) + model.train() + runner = _CompiledModel(model, _get_model_structure_key(model)) + runner.train() + reference = model.eval()(coord_t, atype_t, box=box_t) + model.train() + out = runner(coord_t, atype_t, box=box_t) + for key in ("energy", "force", "virial"): + torch.testing.assert_close(out[key], reference[key], rtol=1e-6, atol=1e-8) + # the backward reaches the charge network and the kernel through the + # compiled graph + loss = out["energy"].square().sum() + out["force"].square().sum() + loss.backward() + grads = [parameter.grad for parameter in model.atomic_model.long_range.parameters()] + assert grads and all(grad is not None for grad in grads) + assert all(torch.isfinite(grad).all() for grad in grads) diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 719ada14b3..1e042b5fd3 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -541,6 +541,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]]) diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index 859e1df8c2..87d4fd9f58 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -868,6 +868,79 @@ def test_finetune_random_fitting_into_bridged(self) -> None: """A RANDOM fitting of a bridged target takes the plain descriptor.""" self._finetune_across_bridging(source_bridged=False, model_branch="RANDOM") + def test_finetune_short_range_into_long_range(self) -> None: + """A checkpoint without a long-range module fine-tunes into one with it. + + The descriptor and the fitting net come from the checkpoint, while the + new periodic module starts from its initialization and takes the + input statistics of the fine-tuning data, here the bounds of its fixed + k grid, before the bias adjustment evaluates the complete model. + """ + plain = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [16, 16], "seed": 1}, + "data_stat_nbatch": 1, + } + target = deepcopy(plain) + target["long_range"] = { + "type": "latent_charge", + "charge_net": {"neuron": [8], "seed": 1}, + "kernel": {"type": "les", "alpha": 1.0}, + "charge_constraint": {"source": "neutral"}, + "periodic": True, + "k_grid_mode": "fixed", + } + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_long_range_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = normalize( + update_deepmd_input( + _make_config(self.data_dir, plain, numb_steps=1), warning=False + ) + ) + ckpt_path = self._train_pretrained(config, tmpdir) + ft_config = normalize( + update_deepmd_input( + _make_config(self.data_dir, target, numb_steps=1), warning=False + ) + ) + ft_config["model"], finetune_links = get_finetune_rules( + ckpt_path, ft_config["model"] + ) + trainer_ft = get_trainer( + ft_config, + finetune_model=ckpt_path, + finetune_links=finetune_links, + ) + pretrained_state = torch.load( + ckpt_path, map_location=DEVICE, weights_only=True + )["model"] + ft_state = trainer_ft.wrapper.state_dict() + learned = [ + key + for key in ft_state + if ".descriptor." in key + or (".fitting_net." in key and "bias_atom_e" not in key) + ] + self.assertTrue(learned) + for key in learned: + torch.testing.assert_close(ft_state[key], pretrained_state[key]) + long_range = trainer_ft.wrapper.model["Default"].atomic_model.long_range + self.assertTrue(torch.all(long_range.k_grid_bounds > 0)) + trainer_ft.run() + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + def test_finetune_from_pte(self) -> None: """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint.""" from deepmd.pt_expt.entrypoints.main import (