deepmd.jax.modifier

deepmd.jax.modifier#

Data modifiers for the JAX backend.

Submodules#

deepmd.jax.modifier.dipole_charge

Classes#

DipoleChargeModifier

Apply dipole-charge corrections with JAX automatic differentiation.

Package Contents#

class deepmd.jax.modifier.DipoleChargeModifier(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)[source]#

Bases: deepmd.dpmodel.modifier.dipole_charge.DipoleChargeModifierBase

Apply dipole-charge corrections with JAX automatic differentiation.

dipole_model#

sel_type#

__call__(coord: deepmd.jax.env.jnp.ndarray, atype: deepmd.jax.env.jnp.ndarray, box: deepmd.jax.env.jnp.ndarray | None = None, fparam: deepmd.jax.env.jnp.ndarray | None = None, aparam: deepmd.jax.env.jnp.ndarray | None = None, do_atomic_virial: bool = False, charge_spin: deepmd.jax.env.jnp.ndarray | None = None) → dict[str, deepmd.jax.env.jnp.ndarray][source]#: Compute dipole-charge energy, force, and virial corrections.