deepmd.jax.model

deepmd.jax.model#

Submodules#

Classes#

`DipoleModel`
`DOSModel`
`DPZBLLinearEnergyAtomicModel`	Model linearly combine a list of AtomicModels.
`EnergyModel`
`PolarModel`
`PropertyModel`

Package Contents#

class deepmd.jax.model.DipoleModel[source]#: Bases: make_jax_dp_model_from_dpmodel(DipoleModelDP, DPAtomicModelDipole)

class deepmd.jax.model.DOSModel[source]#: Bases: make_jax_dp_model_from_dpmodel(DOSModelDP, DPAtomicModelDOS)

class deepmd.jax.model.DPZBLLinearEnergyAtomicModel(dp_model: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel, zbl_model: deepmd.dpmodel.atomic_model.pairtab_atomic_model.PairTabAtomicModel, sw_rmin: float, sw_rmax: float, type_map: list[str], smin_alpha: float | None = 0.1, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.linear_atomic_model.DPZBLLinearEnergyAtomicModel

Model linearly combine a list of AtomicModels.

Parameters:

dp_model: The DPAtomicModel being combined.
zbl_model: The PairTable model being combined.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP.
sw_rmax: The upper boundary of the interpolation between short-range tabulated interaction and DP.
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
smin_alpha: The short-range tabulated interaction will be switched according to the distance of the nearest neighbor. This distance is calculated by softmin.

__setattr__(name: str, value: Any) → None[source]#

forward_common_atomic(extended_coord: deepmd.jax.env.jnp.ndarray, extended_atype: deepmd.jax.env.jnp.ndarray, nlist: deepmd.jax.env.jnp.ndarray, mapping: deepmd.jax.env.jnp.ndarray | None = None, fparam: deepmd.jax.env.jnp.ndarray | None = None, aparam: deepmd.jax.env.jnp.ndarray | None = None) → dict[str, deepmd.jax.env.jnp.ndarray][source]#

Common interface for atomic inference.

This method accept extended coordinates, extended atom typs, neighbor list, and predict the atomic contribution of the fit property.

Parameters:

extended_coord: extended coordinates, shape: nf x (nall x 3)
extended_atype: extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual.
nlist: neighbor list, shape: nf x nloc x nsel
mapping: extended to local index mapping, shape: nf x nall
fparam: frame parameters, shape: nf x dim_fparam
aparam: atomic parameter, shape: nf x nloc x dim_aparam

Returns:

ret_dict: dict of output atomic properties. should implement the definition of fitting_output_def. ret_dict[“mask”] of shape nf x nloc will be provided. ret_dict[“mask”][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real. ret_dict[“mask”][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual.

class deepmd.jax.model.EnergyModel[source]#: Bases: make_jax_dp_model_from_dpmodel(EnergyModelDP, DPAtomicModelEnergy)

class deepmd.jax.model.PolarModel[source]#: Bases: make_jax_dp_model_from_dpmodel(PolarModelDP, DPAtomicModelPolar)

class deepmd.jax.model.PropertyModel[source]#: Bases: make_jax_dp_model_from_dpmodel(PropertyModelDP, DPAtomicModelProperty)