deepmd.jax.model.dp_zbl_model#
Classes#
The unit operation of a native model. |
Module Contents#
- class deepmd.jax.model.dp_zbl_model.DPZBLModel(*args, **kwargs)[source]#
Bases:
deepmd.dpmodel.model.dp_zbl_model.DPZBLModel
The unit operation of a native model.
- 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, do_atomic_virial: bool = False)[source]#
- format_nlist(extended_coord: deepmd.jax.env.jnp.ndarray, extended_atype: deepmd.jax.env.jnp.ndarray, nlist: deepmd.jax.env.jnp.ndarray, extra_nlist_sort: bool = False)[source]#
Format the neighbor list.
1. If the number of neighbors in the nlist is equal to sum(self.sel), it does nothong
2. If the number of neighbors in the nlist is smaller than sum(self.sel), the nlist is pad with -1.
3. If the number of neighbors in the nlist is larger than sum(self.sel), the nearest sum(sel) neighbors will be preserved.
Known limitations:
In the case of not self.mixed_types, the nlist is always formatted. May have side effact on the efficiency.
- Parameters:
- extended_coord
coordinates in extended region. nf x nall x 3
- extended_atype
atomic type in extended region. nf x nall
- nlist
neighbor list. nf x nloc x nsel
- extra_nlist_sort
whether to forcibly sort the nlist.
- Returns:
formated_nlist
the formatted nlist.