From 9f6ab7c4c14d6a4d439bc4447189daaeb198bbef Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 12 Jul 2026 20:43:54 +0800 Subject: [PATCH 1/3] feat(modifier): add dpmodel dipole-charge adapters Add the shared Array API dipole-charge core and eager adapters for PyTorch exportable, JAX, and TensorFlow 2. Cover energy, force, virial, and Array API Strict consistency. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/modifier/__init__.py | 6 + deepmd/dpmodel/modifier/dipole_charge.py | 259 ++++++++++++++++++ deepmd/jax/modifier/__init__.py | 8 + deepmd/jax/modifier/dipole_charge.py | 128 +++++++++ deepmd/pt_expt/modifier/__init__.py | 8 + deepmd/pt_expt/modifier/dipole_charge.py | 157 +++++++++++ deepmd/tf2/modifier/__init__.py | 8 + deepmd/tf2/modifier/dipole_charge.py | 129 +++++++++ source/tests/consistent/modifier/__init__.py | 1 + .../consistent/modifier/test_dipole_charge.py | 148 ++++++++++ 10 files changed, 852 insertions(+) create mode 100644 deepmd/dpmodel/modifier/dipole_charge.py create mode 100644 deepmd/jax/modifier/__init__.py create mode 100644 deepmd/jax/modifier/dipole_charge.py create mode 100644 deepmd/pt_expt/modifier/__init__.py create mode 100644 deepmd/pt_expt/modifier/dipole_charge.py create mode 100644 deepmd/tf2/modifier/__init__.py create mode 100644 deepmd/tf2/modifier/dipole_charge.py create mode 100644 source/tests/consistent/modifier/__init__.py create mode 100644 source/tests/consistent/modifier/test_dipole_charge.py diff --git a/deepmd/dpmodel/modifier/__init__.py b/deepmd/dpmodel/modifier/__init__.py index d4e8ab56e3..fbbbb08f46 100644 --- a/deepmd/dpmodel/modifier/__init__.py +++ b/deepmd/dpmodel/modifier/__init__.py @@ -2,7 +2,13 @@ from .base_modifier import ( make_base_modifier, ) +from .dipole_charge import ( + DipoleChargeModifier, + DipoleChargeModifierBase, +) __all__ = [ + "DipoleChargeModifier", + "DipoleChargeModifierBase", "make_base_modifier", ] diff --git a/deepmd/dpmodel/modifier/dipole_charge.py b/deepmd/dpmodel/modifier/dipole_charge.py new file mode 100644 index 0000000000..34abef7ccf --- /dev/null +++ b/deepmd/dpmodel/modifier/dipole_charge.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent helpers for the dipole-charge modifier.""" + +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_modifier import ( + make_base_modifier, +) + +BaseModifier = make_base_modifier() + +ELECTROSTATIC_CONVERSION = 14.39964535475696995031 + + +def compute_ewald_grids(box: Any, spacing: float) -> tuple[tuple[int, int, int], ...]: + """Compute one fixed even reciprocal grid for every frame's cell.""" + box_array = to_numpy_array(box) + grids = [] + for frame in range(box_array.shape[0]): + grid = [] + for axis in range(3): + size = int(np.ceil(np.linalg.norm(box_array[frame, axis]) / spacing)) + grid.append(size + size % 2) + grids.append((grid[0], grid[1], grid[2])) + return tuple(grids) + + +def validate_charge_maps( + atype: Any, + sel_type: list[int], + model_charge_map: list[float], + sys_charge_map: list[float], +) -> None: + """Validate type coverage before entering a backend differentiation graph.""" + atype_array = to_numpy_array(atype) + real_types = atype_array[atype_array >= 0] + if real_types.size and np.max(real_types) >= len(sys_charge_map): + raise ValueError("sys_charge_map does not cover all real atom types") + if any(atom_type < 0 or atom_type >= len(sys_charge_map) for atom_type in sel_type): + raise ValueError("sel_type contains an atom type outside sys_charge_map") + if len(sel_type) != len(model_charge_map): + raise ValueError( + "model_charge_map must follow get_sel_type() order and have equal length" + ) + + +def extend_dplr_system( + coord: Any, + atype: Any, + dipole: Any, + sel_type: list[int], + model_charge_map: list[float], + sys_charge_map: list[float], +) -> tuple[Any, Any]: + """Append fixed-shape WC slots and construct real/WC charges. + + One WC slot is appended for every input atom. Slots belonging to unselected + or virtual atoms carry zero charge, avoiding variable-length boolean fancy + indexing while remaining exactly equivalent in the Ewald sum. + """ + xp = array_api_compat.array_namespace(coord, atype, dipole) + device = array_api_compat.device(coord) + real_mask = atype >= 0 + safe_atype = xp.where(real_mask, atype, xp.zeros_like(atype)) + sys_charge = xp.asarray(sys_charge_map, dtype=coord.dtype, device=device) + type_index = xp.arange(len(sys_charge_map), dtype=atype.dtype, device=device) + type_one_hot = xp.astype(safe_atype[..., None] == type_index, coord.dtype) + real_charge = xp.where( + real_mask, + xp.sum(type_one_hot * sys_charge, axis=-1), + xp.zeros_like(coord[..., 0]), + ) + + wc_charge_by_type = xp.zeros( + (len(sys_charge_map),), dtype=coord.dtype, device=device + ) + for index, atom_type in enumerate(sel_type): + type_mask = type_index == xp.asarray( + atom_type, dtype=atype.dtype, device=device + ) + wc_charge_by_type = xp.where( + type_mask, + xp.asarray(model_charge_map[index], dtype=coord.dtype, device=device), + wc_charge_by_type, + ) + wc_charge = xp.where( + real_mask, + xp.sum(type_one_hot * wc_charge_by_type, axis=-1), + xp.zeros_like(coord[..., 0]), + ) + selected_mask = wc_charge != 0 + wc_coord = coord + dipole * xp.astype(selected_mask[..., None], coord.dtype) + return ( + xp.concat((coord, wc_coord), axis=1), + xp.concat((real_charge, wc_charge), axis=1), + ) + + +def ewald_reciprocal_energy( + coord: Any, + charge: Any, + box: Any, + grids: tuple[tuple[int, int, int], ...], + beta: float, +) -> Any: + """Evaluate reciprocal Ewald energy using only Array API operations.""" + xp = array_api_compat.array_namespace(coord, charge, box) + device = array_api_compat.device(coord) + pi = xp.asarray(np.pi, dtype=coord.dtype, device=device) + conversion = xp.asarray(ELECTROSTATIC_CONVERSION, dtype=coord.dtype, device=device) + frame_energy = [] + for frame in range(coord.shape[0]): + axes = tuple( + xp.arange(-size // 2, size // 2 + 1, dtype=coord.dtype, device=device) + for size in grids[frame] + ) + mesh = xp.meshgrid(*axes, indexing="ij") + wave_index = xp.reshape(xp.stack(mesh, axis=-1), (-1, 3)) + nonzero = xp.any(wave_index != 0, axis=-1) + cell = box[frame, ...] + inverse_box = xp.linalg.inv(cell) + wave = wave_index @ xp.permute_dims(inverse_box, (1, 0)) + wave2 = xp.sum(wave * wave, axis=-1) + safe_wave2 = xp.where( + nonzero, + wave2, + xp.ones_like(wave2), + ) + fractional = coord[frame, ...] @ inverse_box + phase = 2 * pi * (fractional @ xp.permute_dims(wave_index, (1, 0))) + sqr = xp.sum(charge[frame, :, None] * xp.cos(phase), axis=0) + sqi = xp.sum(charge[frame, :, None] * xp.sin(phase), axis=0) + kernel = xp.exp(-(pi * pi) * safe_wave2 / (beta * beta)) / safe_wave2 + kernel = kernel * xp.astype(nonzero, coord.dtype) + volume = xp.abs(xp.linalg.det(cell)) + frame_energy.append( + xp.sum(kernel * (sqr * sqr + sqi * sqi)) * conversion / (2 * pi * volume) + ) + return xp.stack(frame_energy, axis=0)[:, None] + + +class DipoleChargeModifierBase: + """Store the dipole-charge schema and define its atom-selection contract. + + The modifier uses distinct masking concepts. ``real_atom_mask`` excludes + padding or externally supplied virtual atoms (negative atom types), while + ``selected_wc_mask`` identifies the real atoms that create Wannier charge + centers according to the dipole model's ``sel_type``. Neighbor-list masks + remain the responsibility of the embedded dipole model and must not be + reused as either of these masks. + + ``model_charge_map`` follows the exact order returned by the embedded + dipole model's ``get_sel_type()`` method, matching the established + TensorFlow modifier input contract. + """ + + modifier_type = "dipole_charge" + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1.0, + ewald_beta: float = 0.4, + ) -> None: + """Initialize the shared dipole-charge configuration.""" + if not model_name: + raise ValueError("model_name must identify a dipole model") + if not model_charge_map: + raise ValueError("model_charge_map must not be empty") + if not sys_charge_map: + raise ValueError("sys_charge_map must not be empty") + if ewald_h <= 0.0: + raise ValueError("ewald_h must be positive") + if ewald_beta <= 0.0: + raise ValueError("ewald_beta must be positive") + self.model_name = model_name + self.model_charge_map = [float(value) for value in model_charge_map] + self.sys_charge_map = [float(value) for value in sys_charge_map] + self.ewald_h = float(ewald_h) + self.ewald_beta = float(ewald_beta) + + def serialize(self) -> dict[str, Any]: + """Serialize the backend-neutral dipole-charge configuration.""" + return { + "@class": "Modifier", + "type": self.modifier_type, + "@version": 3, + "model_name": self.model_name, + "model_charge_map": self.model_charge_map, + "sys_charge_map": self.sys_charge_map, + "ewald_h": self.ewald_h, + "ewald_beta": self.ewald_beta, + } + + @classmethod + def deserialize(cls, data: dict[str, Any]) -> "DipoleChargeModifierBase": + """Deserialize a dipole-charge configuration with version validation.""" + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 3, 1) + data.pop("@class", None) + data.pop("type", None) + return cls(**data) + + @staticmethod + def get_real_atom_mask(atype: np.ndarray) -> np.ndarray: + """Return the mask of physical atoms accepted by the dipole model.""" + return np.asarray(atype) >= 0 + + @staticmethod + def get_selected_wc_mask(atype: np.ndarray, sel_type: list[int]) -> np.ndarray: + """Return real atoms whose dipole output creates a virtual WC site.""" + atype = np.asarray(atype) + return (atype >= 0) & np.isin(atype, np.asarray(sel_type, dtype=np.int64)) + + def make_charge_maps( + self, atype: np.ndarray, sel_type: list[int] + ) -> tuple[np.ndarray, np.ndarray]: + """Map real atom and selected WC types to their configured charges.""" + atype = np.asarray(atype, dtype=np.int64) + real_mask = self.get_real_atom_mask(atype) + if np.any(atype[real_mask] >= len(self.sys_charge_map)): + raise ValueError("sys_charge_map does not cover all real atom types") + if len(sel_type) != len(self.model_charge_map): + raise ValueError( + "model_charge_map length must match the dipole model sel_type length" + ) + + real_charge = np.zeros_like(atype, dtype=np.float64) + real_charge[real_mask] = np.asarray(self.sys_charge_map)[atype[real_mask]] + selected_mask = self.get_selected_wc_mask(atype, sel_type) + selected_type = atype[selected_mask] + wc_charge_by_type = { + atom_type: self.model_charge_map[index] + for index, atom_type in enumerate(sel_type) + } + wc_charge = np.asarray( + [wc_charge_by_type[int(atom_type)] for atom_type in selected_type], + dtype=np.float64, + ) + return real_charge, wc_charge + + +@BaseModifier.register("dipole_charge") +class DipoleChargeModifier(DipoleChargeModifierBase, BaseModifier): + """Backend-neutral serialized representation of dipole-charge.""" diff --git a/deepmd/jax/modifier/__init__.py b/deepmd/jax/modifier/__init__.py new file mode 100644 index 0000000000..6d84e8ba32 --- /dev/null +++ b/deepmd/jax/modifier/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Data modifiers for the JAX backend.""" + +from .dipole_charge import ( + DipoleChargeModifier, +) + +__all__ = ["DipoleChargeModifier"] diff --git a/deepmd/jax/modifier/dipole_charge.py b/deepmd/jax/modifier/dipole_charge.py new file mode 100644 index 0000000000..4dc75f761d --- /dev/null +++ b/deepmd/jax/modifier/dipole_charge.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""JAX implementation of the dipole-charge modifier.""" + +from typing import ( + Any, +) + +from deepmd.dpmodel.modifier.dipole_charge import ( + DipoleChargeModifierBase, + compute_ewald_grids, + ewald_reciprocal_energy, + extend_dplr_system, + validate_charge_maps, +) +from deepmd.dpmodel.utils.serialization import ( + load_dp_model, +) +from deepmd.jax.env import ( + jax, + jnp, +) +from deepmd.jax.model.base_model import ( + BaseModel, +) + + +class DipoleChargeModifier(DipoleChargeModifierBase): + """Apply dipole-charge corrections with JAX automatic differentiation.""" + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1.0, + ewald_beta: float = 0.4, + dipole_model: Any | None = None, + ) -> None: + """Load or attach the JAX dipole model used to create WC positions.""" + super().__init__( + model_name, model_charge_map, sys_charge_map, ewald_h, ewald_beta + ) + self.dipole_model = ( + BaseModel.deserialize(load_dp_model(model_name)["model"]) + if dipole_model is None + else dipole_model + ) + self.sel_type = [int(value) for value in self.dipole_model.get_sel_type()] + if len(self.sel_type) != len(self.model_charge_map): + raise ValueError( + "model_charge_map length must match the dipole model sel_type length" + ) + + def __call__( + self, + coord: jnp.ndarray, + atype: jnp.ndarray, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, + do_atomic_virial: bool = False, + charge_spin: jnp.ndarray | None = None, + ) -> dict[str, jnp.ndarray]: + """Compute dipole-charge energy, force, and virial corrections.""" + if box is None: + raise RuntimeError("dipole_charge does not support non-periodic systems") + if do_atomic_virial: + raise RuntimeError("dipole_charge does not provide atomic virial") + coord = jnp.asarray(coord) + atype = jnp.asarray(atype) + box = jnp.asarray(box) + validate_charge_maps( + atype, + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + grids = compute_ewald_grids(box, self.ewald_h) + + def energy_fn(force_coord: jnp.ndarray, strain: jnp.ndarray) -> jnp.ndarray: + """Return per-frame energy while retaining the full gradient path.""" + transform = jnp.eye(3, dtype=coord.dtype)[None, :, :] + strain + strained_coord = force_coord @ transform + strained_box = box @ transform + prediction = self.dipole_model( + strained_coord, + atype, + box=strained_box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=False, + charge_spin=charge_spin, + ) + all_coord, all_charge = extend_dplr_system( + strained_coord, + atype, + prediction["dipole"], + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + return ewald_reciprocal_energy( + all_coord, + all_charge, + strained_box, + grids, + self.ewald_beta, + ) + + strain = jnp.zeros((coord.shape[0], 3, 3), dtype=coord.dtype) + + def energy_with_aux( + force_coord: jnp.ndarray, cell_strain: jnp.ndarray + ) -> tuple[jnp.ndarray, jnp.ndarray]: + """Return the scalar differentiation target and per-frame energies.""" + energy_by_frame = energy_fn(force_coord, cell_strain) + return jnp.sum(energy_by_frame), energy_by_frame + + (_, energy_by_frame), gradients = jax.value_and_grad( + energy_with_aux, + argnums=(0, 1), + has_aux=True, + )(coord, strain) + return { + "energy": energy_by_frame, + "force": -gradients[0], + "virial": -jnp.swapaxes(gradients[1], -1, -2).reshape(coord.shape[0], 9), + } diff --git a/deepmd/pt_expt/modifier/__init__.py b/deepmd/pt_expt/modifier/__init__.py new file mode 100644 index 0000000000..c3d89f9f98 --- /dev/null +++ b/deepmd/pt_expt/modifier/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Data modifiers for the PyTorch exportable backend.""" + +from .dipole_charge import ( + DipoleChargeModifier, +) + +__all__ = ["DipoleChargeModifier"] diff --git a/deepmd/pt_expt/modifier/dipole_charge.py b/deepmd/pt_expt/modifier/dipole_charge.py new file mode 100644 index 0000000000..616d49d4f6 --- /dev/null +++ b/deepmd/pt_expt/modifier/dipole_charge.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Dipole-charge modifier for the PyTorch exportable backend.""" + +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.modifier.dipole_charge import ( + DipoleChargeModifierBase, + compute_ewald_grids, + ewald_reciprocal_energy, + extend_dplr_system, + validate_charge_maps, +) +from deepmd.dpmodel.utils.serialization import ( + load_dp_model, +) +from deepmd.pt_expt.model.model import ( + BaseModel, +) + + +class DipoleChargeModifier(torch.nn.Module, DipoleChargeModifierBase): + """Apply dipole-charge corrections to a pt_expt dipole model. + + A portable ``.dp`` file is accepted through ``model_name``. Tests and + embedding workflows may pass an already-deserialized ``dipole_model``. + """ + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1.0, + ewald_beta: float = 0.4, + use_cache: bool = True, + dipole_model: Any | None = None, + ) -> None: + """Load or attach the exportable dipole model used by the modifier.""" + torch.nn.Module.__init__(self) + DipoleChargeModifierBase.__init__( + self, + model_name, + model_charge_map, + sys_charge_map, + ewald_h, + ewald_beta, + ) + if dipole_model is None: + dipole_model = BaseModel.deserialize(load_dp_model(model_name)["model"]) + self.dipole_model = dipole_model + self.dipole_model.eval() + self.sel_type = [int(value) for value in self.dipole_model.get_sel_type()] + if len(self.sel_type) != len(self.model_charge_map): + raise ValueError( + "model_charge_map length must match the dipole model sel_type length" + ) + self.use_cache = use_cache + self.modifier_type = "dipole_charge" + + def train(self, mode: bool = True) -> "DipoleChargeModifier": + """Set modifier mode while keeping the embedded dipole model frozen.""" + super().train(mode) + self.dipole_model.eval() + return self + + def serialize(self) -> dict[str, Any]: + """Serialize the user-facing dipole-charge configuration.""" + data = DipoleChargeModifierBase.serialize(self) + data["use_cache"] = self.use_cache + return data + + def _energy_with_grid( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor, + grids: tuple[tuple[int, int, int], ...], + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + ) -> torch.Tensor: + """Evaluate the shared energy core on a precomputed reciprocal grid.""" + prediction = self.dipole_model( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=False, + charge_spin=charge_spin, + ) + all_coord, all_charge = extend_dplr_system( + coord, + atype, + prediction["dipole"], + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + return ewald_reciprocal_energy( + all_coord, all_charge, box, grids, self.ewald_beta + ) + + def forward( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + do_atomic_virial: bool = False, + charge_spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Compute dipole-charge outputs with the shared dpmodel core.""" + if box is None: + raise RuntimeError("dipole_charge does not support non-periodic systems") + if do_atomic_virial: + raise RuntimeError("dipole_charge does not provide atomic virial") + validate_charge_maps( + atype, + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + grids = compute_ewald_grids(box.detach(), self.ewald_h) + force_coord = coord + if not force_coord.requires_grad: + force_coord = force_coord.clone().requires_grad_(True) + strain = torch.zeros( + (coord.shape[0], 3, 3), dtype=coord.dtype, device=coord.device + ).requires_grad_(True) + transform = torch.eye(3, dtype=coord.dtype, device=coord.device)[None] + strain + energy = self._energy_with_grid( + force_coord @ transform, + atype, + box @ transform, + grids, + fparam, + aparam, + charge_spin, + ) + gradients = torch.autograd.grad( + energy, + (force_coord, strain), + grad_outputs=torch.ones_like(energy), + create_graph=self.training, + retain_graph=True, + ) + return { + "energy": energy, + "force": -gradients[0], + "virial": -gradients[1].transpose(-1, -2).reshape(coord.shape[0], 9), + } diff --git a/deepmd/tf2/modifier/__init__.py b/deepmd/tf2/modifier/__init__.py new file mode 100644 index 0000000000..5c92c6400d --- /dev/null +++ b/deepmd/tf2/modifier/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Data modifiers for the TensorFlow 2 backend.""" + +from .dipole_charge import ( + DipoleChargeModifier, +) + +__all__ = ["DipoleChargeModifier"] diff --git a/deepmd/tf2/modifier/dipole_charge.py b/deepmd/tf2/modifier/dipole_charge.py new file mode 100644 index 0000000000..bde70fe12f --- /dev/null +++ b/deepmd/tf2/modifier/dipole_charge.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""TensorFlow 2 implementation of the dipole-charge modifier.""" + +from typing import ( + Any, +) + +from deepmd.dpmodel.modifier.dipole_charge import ( + DipoleChargeModifierBase, + compute_ewald_grids, + ewald_reciprocal_energy, + extend_dplr_system, + validate_charge_maps, +) +from deepmd.dpmodel.utils.serialization import ( + load_dp_model, +) +from deepmd.tf2.common import ( + to_tensorflow_array, + to_tf_tensor, + unwrap_value, +) +from deepmd.tf2.env import ( + tf, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) + + +class DipoleChargeModifier(DipoleChargeModifierBase): + """Apply dipole-charge corrections with TensorFlow gradient tapes.""" + + def __init__( + self, + model_name: str, + model_charge_map: list[float], + sys_charge_map: list[float], + ewald_h: float = 1.0, + ewald_beta: float = 0.4, + dipole_model: Any | None = None, + ) -> None: + """Load or attach the TF2 dipole model used to create WC positions.""" + super().__init__( + model_name, model_charge_map, sys_charge_map, ewald_h, ewald_beta + ) + self.dipole_model = ( + BaseModel.deserialize(load_dp_model(model_name)["model"]) + if dipole_model is None + else dipole_model + ) + self.sel_type = [int(value) for value in self.dipole_model.get_sel_type()] + if len(self.sel_type) != len(self.model_charge_map): + raise ValueError( + "model_charge_map length must match the dipole model sel_type length" + ) + + def __call__( + self, + coord: Any, + atype: Any, + box: Any | None = None, + fparam: Any | None = None, + aparam: Any | None = None, + do_atomic_virial: bool = False, + charge_spin: Any | None = None, + ) -> dict[str, tf.Tensor]: + """Compute dipole-charge energy, force, and virial corrections.""" + if box is None: + raise RuntimeError("dipole_charge does not support non-periodic systems") + if do_atomic_virial: + raise RuntimeError("dipole_charge does not provide atomic virial") + coord = to_tf_tensor(to_tensorflow_array(coord)) + atype = to_tf_tensor(to_tensorflow_array(atype)) + box = to_tf_tensor(to_tensorflow_array(box)) + assert coord is not None + assert atype is not None + assert box is not None + validate_charge_maps( + atype, + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + grids = compute_ewald_grids(box, self.ewald_h) + with tf.GradientTape() as tape: + tape.watch(coord) + strain = tf.zeros((tf.shape(coord)[0], 3, 3), dtype=coord.dtype) + tape.watch(strain) + transform = ( + tf.eye(3, batch_shape=[tf.shape(coord)[0]], dtype=coord.dtype) + strain + ) + strained_coord = coord @ transform + strained_box = box @ transform + prediction = self.dipole_model( + strained_coord, + atype, + box=strained_box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=False, + charge_spin=charge_spin, + ) + all_coord, all_charge = extend_dplr_system( + to_tensorflow_array(strained_coord), + to_tensorflow_array(atype), + prediction["dipole"], + self.sel_type, + self.model_charge_map, + self.sys_charge_map, + ) + energy = unwrap_value( + ewald_reciprocal_energy( + all_coord, + all_charge, + to_tensorflow_array(strained_box), + grids, + self.ewald_beta, + ) + ) + total_energy = tf.reduce_sum(energy) + force_grad, strain_grad = tape.gradient(total_energy, [coord, strain]) + assert force_grad is not None + assert strain_grad is not None + return { + "energy": energy, + "force": -force_grad, + "virial": -tf.reshape(tf.transpose(strain_grad, (0, 2, 1)), (-1, 9)), + } diff --git a/source/tests/consistent/modifier/__init__.py b/source/tests/consistent/modifier/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/consistent/modifier/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/consistent/modifier/test_dipole_charge.py b/source/tests/consistent/modifier/test_dipole_charge.py new file mode 100644 index 0000000000..7f7b040863 --- /dev/null +++ b/source/tests/consistent/modifier/test_dipole_charge.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Cross-backend consistency tests for the dpmodel dipole-charge modifier.""" + +import unittest + +import numpy as np + +from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.modifier.dipole_charge import ( + ewald_reciprocal_energy, + extend_dplr_system, +) + +from ..common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT_EXPT, + INSTALLED_TF2, +) + +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict + +if INSTALLED_PT_EXPT: + import torch + +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.model.model import BaseModel as PTExptBaseModel + from deepmd.pt_expt.modifier.dipole_charge import ( + DipoleChargeModifier as PTExptModifier, + ) +if INSTALLED_JAX: + from deepmd.jax.model.base_model import BaseModel as JAXBaseModel + from deepmd.jax.modifier.dipole_charge import DipoleChargeModifier as JAXModifier +if INSTALLED_TF2: + from deepmd.tf2.model.base_model import BaseModel as TF2BaseModel + from deepmd.tf2.modifier.dipole_charge import DipoleChargeModifier as TF2Modifier + + +class TestDipoleChargeModifierConsistency(unittest.TestCase): + """Ensure dpmodel-driven adapters produce the same dipole-charge correction.""" + + @classmethod + def setUpClass(cls) -> None: + model_config = { + "type_map": ["O", "H"], + "atom_exclude_types": [1], + "descriptor": { + "type": "se_e2_a", + "sel": [8, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [4, 8], + "axis_neuron": 4, + "precision": "float64", + "seed": 2026, + }, + "fitting_net": { + "type": "dipole", + "neuron": [8, 8], + "precision": "float64", + "seed": 2027, + }, + } + cls.serialized_model = get_model_dp(model_config).serialize() + cls.coord = np.asarray( + [ + [ + [1.0, 1.0, 1.0], + [1.8, 1.0, 1.0], + [1.0, 1.8, 1.0], + [6.0, 6.0, 6.0], + [6.8, 6.0, 6.0], + [6.0, 6.8, 6.0], + ] + ], + dtype=np.float64, + ) + cls.atype = np.asarray([[0, 1, 1, 0, 1, 1]], dtype=np.int32) + cls.box = np.asarray([np.eye(3) * 10.0], dtype=np.float64) + cls.kwargs = { + "model_name": "embedded.dp", + "model_charge_map": [-8.0], + "sys_charge_map": [6.0, 1.0], + "ewald_h": 2.0, + "ewald_beta": 1.0, + } + + def _evaluate_backends(self) -> dict[str, dict[str, np.ndarray]]: + results = {} + if INSTALLED_PT_EXPT: + model = PTExptBaseModel.deserialize(self.serialized_model).double() + modifier = PTExptModifier(**self.kwargs, dipole_model=model).double().eval() + result = modifier( + torch.tensor(self.coord, dtype=torch.float64, device="cpu"), + torch.tensor(self.atype, dtype=torch.int64, device="cpu"), + torch.tensor(self.box, dtype=torch.float64, device="cpu"), + ) + results["pt_expt"] = { + key: value.detach().cpu().numpy() for key, value in result.items() + } + if INSTALLED_JAX: + model = JAXBaseModel.deserialize(self.serialized_model) + result = JAXModifier(**self.kwargs, dipole_model=model)( + self.coord, self.atype, self.box + ) + results["jax"] = {key: np.asarray(value) for key, value in result.items()} + if INSTALLED_TF2: + model = TF2BaseModel.deserialize(self.serialized_model) + result = TF2Modifier(**self.kwargs, dipole_model=model)( + self.coord, self.atype, self.box + ) + results["tf2"] = {key: np.asarray(value) for key, value in result.items()} + return results + + def test_energy_force_and_virial(self) -> None: + results = self._evaluate_backends() + if len(results) < 2: + self.skipTest("At least two dpmodel-driven backends are required") + reference_name, reference = next(iter(results.items())) + for backend_name, result in results.items(): + for key in ("energy", "force", "virial"): + np.testing.assert_allclose( + result[key], + reference[key], + rtol=1e-9, + atol=1e-9, + err_msg=f"{backend_name} does not match {reference_name} for {key}", + ) + + @unittest.skipUnless(INSTALLED_ARRAY_API_STRICT, "array_api_strict is required") + def test_array_api_strict_core(self) -> None: + """Ensure the shared numerical core uses only standard array operations.""" + xp = array_api_strict + coord = xp.asarray(self.coord, dtype=xp.float64) + atype = xp.asarray(self.atype, dtype=xp.int64) + box = xp.asarray(self.box, dtype=xp.float64) + all_coord, all_charge = extend_dplr_system( + coord, + atype, + xp.zeros_like(coord), + [0], + [-8.0], + [6.0, 1.0], + ) + energy = ewald_reciprocal_energy(all_coord, all_charge, box, ((6, 6, 6),), 1.0) + self.assertEqual(energy.shape, (1, 1)) + self.assertTrue(np.isfinite(np.asarray(energy)).all()) From 79496f4fc7239d339ac490b4268d91bf07941839 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 13 Jul 2026 18:35:15 +0800 Subject: [PATCH 2/3] test(modifier): address dipole charge review comments Remove unused NumPy-only charge-map helpers and extend consistency coverage to heterogeneous two-frame cells and negative-type padding in the shared Array API path. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/modifier/dipole_charge.py | 38 ------------------ .../consistent/modifier/test_dipole_charge.py | 40 +++++++++++++++---- 2 files changed, 33 insertions(+), 45 deletions(-) diff --git a/deepmd/dpmodel/modifier/dipole_charge.py b/deepmd/dpmodel/modifier/dipole_charge.py index 34abef7ccf..22f6ca8597 100644 --- a/deepmd/dpmodel/modifier/dipole_charge.py +++ b/deepmd/dpmodel/modifier/dipole_charge.py @@ -215,44 +215,6 @@ def deserialize(cls, data: dict[str, Any]) -> "DipoleChargeModifierBase": data.pop("type", None) return cls(**data) - @staticmethod - def get_real_atom_mask(atype: np.ndarray) -> np.ndarray: - """Return the mask of physical atoms accepted by the dipole model.""" - return np.asarray(atype) >= 0 - - @staticmethod - def get_selected_wc_mask(atype: np.ndarray, sel_type: list[int]) -> np.ndarray: - """Return real atoms whose dipole output creates a virtual WC site.""" - atype = np.asarray(atype) - return (atype >= 0) & np.isin(atype, np.asarray(sel_type, dtype=np.int64)) - - def make_charge_maps( - self, atype: np.ndarray, sel_type: list[int] - ) -> tuple[np.ndarray, np.ndarray]: - """Map real atom and selected WC types to their configured charges.""" - atype = np.asarray(atype, dtype=np.int64) - real_mask = self.get_real_atom_mask(atype) - if np.any(atype[real_mask] >= len(self.sys_charge_map)): - raise ValueError("sys_charge_map does not cover all real atom types") - if len(sel_type) != len(self.model_charge_map): - raise ValueError( - "model_charge_map length must match the dipole model sel_type length" - ) - - real_charge = np.zeros_like(atype, dtype=np.float64) - real_charge[real_mask] = np.asarray(self.sys_charge_map)[atype[real_mask]] - selected_mask = self.get_selected_wc_mask(atype, sel_type) - selected_type = atype[selected_mask] - wc_charge_by_type = { - atom_type: self.model_charge_map[index] - for index, atom_type in enumerate(sel_type) - } - wc_charge = np.asarray( - [wc_charge_by_type[int(atom_type)] for atom_type in selected_type], - dtype=np.float64, - ) - return real_charge, wc_charge - @BaseModifier.register("dipole_charge") class DipoleChargeModifier(DipoleChargeModifierBase, BaseModifier): diff --git a/source/tests/consistent/modifier/test_dipole_charge.py b/source/tests/consistent/modifier/test_dipole_charge.py index 7f7b040863..f94c04beff 100644 --- a/source/tests/consistent/modifier/test_dipole_charge.py +++ b/source/tests/consistent/modifier/test_dipole_charge.py @@ -7,6 +7,7 @@ from deepmd.dpmodel.model.model import get_model as get_model_dp from deepmd.dpmodel.modifier.dipole_charge import ( + compute_ewald_grids, ewald_reciprocal_energy, extend_dplr_system, ) @@ -72,12 +73,25 @@ def setUpClass(cls) -> None: [6.0, 6.0, 6.0], [6.8, 6.0, 6.0], [6.0, 6.8, 6.0], - ] + ], + [ + [1.2, 1.1, 1.0], + [2.0, 1.1, 1.0], + [1.2, 1.9, 1.0], + [7.0, 7.0, 7.0], + [7.8, 7.0, 7.0], + [7.0, 7.8, 7.0], + ], ], dtype=np.float64, ) - cls.atype = np.asarray([[0, 1, 1, 0, 1, 1]], dtype=np.int32) - cls.box = np.asarray([np.eye(3) * 10.0], dtype=np.float64) + cls.atype = np.asarray( + [[0, 1, 1, 0, 1, 1], [0, 1, 1, 0, 1, 1]], + dtype=np.int32, + ) + cls.box = np.asarray( + [np.eye(3) * 10.0, np.diag([12.0, 14.0, 16.0])], dtype=np.float64 + ) cls.kwargs = { "model_name": "embedded.dp", "model_charge_map": [-8.0], @@ -132,8 +146,14 @@ def test_energy_force_and_virial(self) -> None: def test_array_api_strict_core(self) -> None: """Ensure the shared numerical core uses only standard array operations.""" xp = array_api_strict - coord = xp.asarray(self.coord, dtype=xp.float64) - atype = xp.asarray(self.atype, dtype=xp.int64) + padding_coord = np.zeros((self.coord.shape[0], 1, 3), dtype=np.float64) + coord = xp.asarray( + np.concatenate((self.coord, padding_coord), axis=1), dtype=xp.float64 + ) + padding_atype = -np.ones((self.atype.shape[0], 1), dtype=np.int32) + atype = xp.asarray( + np.concatenate((self.atype, padding_atype), axis=1), dtype=xp.int64 + ) box = xp.asarray(self.box, dtype=xp.float64) all_coord, all_charge = extend_dplr_system( coord, @@ -143,6 +163,12 @@ def test_array_api_strict_core(self) -> None: [-8.0], [6.0, 1.0], ) - energy = ewald_reciprocal_energy(all_coord, all_charge, box, ((6, 6, 6),), 1.0) - self.assertEqual(energy.shape, (1, 1)) + # Padding atoms and their paired WC slots must both remain uncharged. + charge_array = np.asarray(all_charge) + np.testing.assert_array_equal(charge_array[:, 6], 0.0) + np.testing.assert_array_equal(charge_array[:, 13], 0.0) + grids = compute_ewald_grids(box, self.kwargs["ewald_h"]) + self.assertNotEqual(grids[0], grids[1]) + energy = ewald_reciprocal_energy(all_coord, all_charge, box, grids, 1.0) + self.assertEqual(energy.shape, (2, 1)) self.assertTrue(np.isfinite(np.asarray(energy)).all()) From a0b9fe54f122cfd74ae60ebc70614c8cb6236913 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 14 Jul 2026 11:18:46 +0800 Subject: [PATCH 3/3] test(modifier): use configured device in dipole charge test Run the pt_expt consistency inputs on the backend configured device so CUDA models do not receive CPU index tensors. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/consistent/modifier/test_dipole_charge.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/source/tests/consistent/modifier/test_dipole_charge.py b/source/tests/consistent/modifier/test_dipole_charge.py index f94c04beff..713f410643 100644 --- a/source/tests/consistent/modifier/test_dipole_charge.py +++ b/source/tests/consistent/modifier/test_dipole_charge.py @@ -25,6 +25,8 @@ if INSTALLED_PT_EXPT: import torch + from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE + if INSTALLED_PT_EXPT: from deepmd.pt_expt.model.model import BaseModel as PTExptBaseModel from deepmd.pt_expt.modifier.dipole_charge import ( @@ -106,9 +108,9 @@ def _evaluate_backends(self) -> dict[str, dict[str, np.ndarray]]: model = PTExptBaseModel.deserialize(self.serialized_model).double() modifier = PTExptModifier(**self.kwargs, dipole_model=model).double().eval() result = modifier( - torch.tensor(self.coord, dtype=torch.float64, device="cpu"), - torch.tensor(self.atype, dtype=torch.int64, device="cpu"), - torch.tensor(self.box, dtype=torch.float64, device="cpu"), + torch.tensor(self.coord, dtype=torch.float64, device=PT_EXPT_DEVICE), + torch.tensor(self.atype, dtype=torch.int64, device=PT_EXPT_DEVICE), + torch.tensor(self.box, dtype=torch.float64, device=PT_EXPT_DEVICE), ) results["pt_expt"] = { key: value.detach().cpu().numpy() for key, value in result.items()