deepmd.dpmodel.atomic_model.base_atomic_model

Attributes

BaseAtomicModel_

Classes

BaseAtomicModel

Base Atomic Model provides the interfaces of an atomic model.

Module Contents

deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel_

class deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel(type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, preset_out_bias: dict[str, deepmd.dpmodel.array_api.Array] | None = None)

Bases: BaseAtomicModel_, deepmd.dpmodel.common.NativeOP

Base Atomic Model provides the interfaces of an atomic model.

type_map

rcond = None

preset_out_bias = None

init_out_stat() → None

Initialize the output bias.

__setitem__(key: str, value: deepmd.dpmodel.array_api.Array) → None

__getitem__(key: str) → deepmd.dpmodel.array_api.Array

get_type_map() → list[str]

Get the type map.

has_default_fparam() → bool

Check if the model has default frame parameters.

reinit_atom_exclude(exclude_types: list[int] = []) → None

reinit_pair_exclude(exclude_types: list[tuple[int, int]] = []) → None

atomic_output_def() → deepmd.dpmodel.output_def.FittingOutputDef

Get the output def of the atomic model.

By default it is the same as FittingOutputDef, but it allows model level wrapper of the output defined by the developer.

change_type_map(type_map: list[str], model_with_new_type_stat: Any | None = None) → None

Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

forward_common_atomic(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None) → dict[str, deepmd.dpmodel.array_api.Array]

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.

call(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None) → dict[str, deepmd.dpmodel.array_api.Array]

Forward pass in NumPy implementation.

serialize() → dict

classmethod deserialize(data: dict) → BaseAtomicModel

apply_out_stat(ret: dict[str, deepmd.dpmodel.array_api.Array], atype: deepmd.dpmodel.array_api.Array) → dict[str, deepmd.dpmodel.array_api.Array]

Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model.

Parameters:

ret

The returned dict by the forward_atomic method

atype

The atom types. nf x nloc

_varsize(shape: list[int]) → int

_get_bias_index(kk: str) → int

_fetch_out_stat(keys: list[str]) → tuple[dict[str, deepmd.dpmodel.array_api.Array], dict[str, deepmd.dpmodel.array_api.Array]]