feat(modifier): add dpmodel dipole-charge adapters

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",
 ]

deepmd/dpmodel/modifier/dipole_charge.py

@@ -0,0 +1,221 @@
+# 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)
+
+
+@BaseModifier.register("dipole_charge")
+class DipoleChargeModifier(DipoleChargeModifierBase, BaseModifier):
+    """Backend-neutral serialized representation of dipole-charge."""

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"]

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),
+        }

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"]