deepmd.pt.model.atomic_model

The atomic model provides the prediction of some property on each atom. All the atomic models are not supposed to be directly accessed by users, but it provides a convenient interface for the implementation of models.

Taking the energy models for example, the developeres only needs to implement the atomic energy prediction via an atomic model, and the model can be automatically made by the deepmd.dpmodel.make_model method. The DPModel is made by ` DPModel = make_model(DPAtomicModel) `

Submodules

• deepmd.pt.model.atomic_model.base_atomic_model
• deepmd.pt.model.atomic_model.dipole_atomic_model
• deepmd.pt.model.atomic_model.dos_atomic_model
• deepmd.pt.model.atomic_model.dp_atomic_model
• deepmd.pt.model.atomic_model.energy_atomic_model
• deepmd.pt.model.atomic_model.linear_atomic_model
• deepmd.pt.model.atomic_model.pairtab_atomic_model
• deepmd.pt.model.atomic_model.polar_atomic_model

Package Contents

Classes

BaseAtomicModel	The base of atomic model.
DPDipoleAtomicModel	Model give atomic prediction of some physical property.
DPDOSAtomicModel	Model give atomic prediction of some physical property.
DPAtomicModel	Model give atomic prediction of some physical property.
DPEnergyAtomicModel	Model give atomic prediction of some physical property.
DPZBLLinearEnergyAtomicModel	Model linearly combine a list of AtomicModels.
LinearEnergyAtomicModel	Linear model make linear combinations of several existing models.
PairTabAtomicModel	Pairwise tabulation energy model.
DPPolarAtomicModel	Model give atomic prediction of some physical property.

class deepmd.pt.model.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, torch.Tensor] | None = None)

Bases: torch.nn.Module, BaseAtomicModel_

The base of atomic model.

Parameters:

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

atom_exclude_types

Exclude the atomic contribution of the given types

pair_exclude_types

Exclude the pair of atoms of the given types from computing the output of the atomic model. Implemented by removing the pairs from the nlist.

rcondfloat, optional

The condition number for the regression of atomic energy.

preset_out_biasDict[str, List[Optional[torch.Tensor]]], optional

Specifying atomic energy contribution in vacuum. Given by key:value pairs. The value is a list specifying the bias. the elements can be None or np.array of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The set_davg_zero key in the descrptor should be set.

init_out_stat()

Initialize the output bias.

__setitem__(key, value)

__getitem__(key)

get_type_map() → List[str]

Get the type map.

reinit_atom_exclude(exclude_types: List[int] = [])

reinit_pair_exclude(exclude_types: List[Tuple[int, int]] = [])

make_atom_mask(atype: torch.Tensor) → torch.Tensor

The atoms with type < 0 are treated as virutal atoms, which serves as place-holders for multi-frame calculations with different number of atoms in different frames.

Parameters:

atype

Atom types. >= 0 for real atoms <0 for virtual atoms.

Returns:

mask

True for real atoms and False for virutal atoms.

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.

forward_common_atomic(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None) → Dict[str, torch.Tensor]

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 coodinates, 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

comm_dict

The data needed for communication for parallel inference.

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.

serialize() → dict

classmethod deserialize(data: dict) → BaseAtomicModel

abstract compute_or_load_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

compute_or_load_out_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

apply_out_stat(ret: Dict[str, torch.Tensor], atype: torch.Tensor)

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

change_out_bias(sample_merged, stat_file_path: deepmd.utils.path.DPPath | None = None, bias_adjust_mode='change-by-statistic') → None

Change the output bias according to the input data and the pretrained model.

Parameters:

sample_mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

bias_adjust_modestr

The mode for changing output bias : [‘change-by-statistic’, ‘set-by-statistic’] ‘change-by-statistic’ : perform predictions on labels of target dataset,

and do least square on the errors to obtain the target shift as bias.

‘set-by-statistic’ : directly use the statistic output bias in the target dataset.

stat_file_pathOptional[DPPath]

The path to the stat file.

_get_forward_wrapper_func() → Callable[Ellipsis, torch.Tensor]

Get a forward wrapper of the atomic model for output bias calculation.

_default_bias()

_default_std()

_varsize(shape: List[int]) → int

_get_bias_index(kk: str) → int

_store_out_stat(out_bias: Dict[str, torch.Tensor], out_std: Dict[str, torch.Tensor], add: bool = False)

_fetch_out_stat(keys: List[str]) → Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]

class deepmd.pt.model.atomic_model.DPDipoleAtomicModel(descriptor, fitting, type_map, **kwargs)

Bases: deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor

Descriptor

fitting_net

Fitting net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

apply_out_stat(ret: Dict[str, torch.Tensor], atype: torch.Tensor)

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

class deepmd.pt.model.atomic_model.DPDOSAtomicModel(descriptor, fitting, type_map, **kwargs)

Bases: deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor

Descriptor

fitting_net

Fitting net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

class deepmd.pt.model.atomic_model.DPAtomicModel(descriptor, fitting, type_map: List[str], **kwargs)

Bases: deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor

Descriptor

fitting_net

Fitting net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

fitting_output_def() → deepmd.dpmodel.FittingOutputDef

Get the output def of the fitting net.

get_rcut() → float

Get the cut-off radius.

get_sel() → List[int]

Get the neighbor selection.

mixed_types() → bool

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

serialize() → dict

classmethod deserialize(data) → DPAtomicModel

forward_atomic(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None) → Dict[str, torch.Tensor]

Return atomic prediction.

Parameters:

extended_coord

coodinates in extended region

extended_atype

atomic type in extended region

nlist

neighbor list. nf x nloc x nsel

mapping

mapps the extended indices to local indices

fparam

frame parameter. nf x ndf

aparam

atomic parameter. nf x nloc x nda

Returns:

result_dict

the result dict, defined by the FittingOutputDef.

get_out_bias() → torch.Tensor

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters:

sampled_func

The lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → List[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

class deepmd.pt.model.atomic_model.DPEnergyAtomicModel(descriptor, fitting, type_map, **kwargs)

Bases: deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor

Descriptor

fitting_net

Fitting net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

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

Bases: LinearEnergyAtomicModel

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 swithed according to the distance of the nearest neighbor. This distance is calculated by softmin.

serialize() → dict

classmethod deserialize(data) → DPZBLLinearEnergyAtomicModel

_compute_weight(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlists_: List[torch.Tensor]) → List[torch.Tensor]

ZBL weight.

Returns:

List[torch.Tensor]

the atomic ZBL weight for interpolation. (nframes, nloc, 1)

class deepmd.pt.model.atomic_model.LinearEnergyAtomicModel(models: List[deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel], type_map: List[str], **kwargs)

Bases: deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel

Linear model make linear combinations of several existing models.

Parameters:

modelslist[DPAtomicModel or PairTabAtomicModel]

A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.

type_maplist[str]

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

mixed_types() → bool

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

get_out_bias() → torch.Tensor

get_rcut() → float

Get the cut-off radius.

get_type_map() → List[str]

Get the type map.

get_model_rcuts() → List[float]

Get the cut-off radius for each individual models.

get_sel() → List[int]

Returns the number of selected atoms for each type.

get_model_nsels() → List[int]

Get the processed sels for each individual models. Not distinguishing types.

get_model_sels() → List[List[int]]

Get the sels for each individual models.

_sort_rcuts_sels() → Tuple[List[float], List[int]]

forward_atomic(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None) → Dict[str, torch.Tensor]

Return atomic prediction.

Parameters:

extended_coord

coodinates in extended region, (nframes, nall * 3)

extended_atype

atomic type in extended region, (nframes, nall)

nlist

neighbor list, (nframes, nloc, nsel).

mapping

mapps the extended indices to local indices.

fparam

frame parameter. (nframes, ndf)

aparam

atomic parameter. (nframes, nloc, nda)

Returns:

result_dict

the result dict, defined by the fitting net output def.

apply_out_stat(ret: Dict[str, torch.Tensor], atype: torch.Tensor)

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

static remap_atype(ori_map: List[str], new_map: List[str]) → torch.Tensor

This method is used to map the atype from the common type_map to the original type_map of indivial AtomicModels. It creates a index mapping for the conversion.

Parameters:

ori_mapList[str]

The original type map of an AtomicModel.

new_mapList[str]

The common type map of the DPZBLLinearEnergyAtomicModel, created by the get_type_map method, must be a subset of the ori_map.

Returns:

torch.Tensor

fitting_output_def() → deepmd.dpmodel.FittingOutputDef

Get the output def of developer implemented atomic models.

serialize() → dict

classmethod deserialize(data: dict) → LinearEnergyAtomicModel

_compute_weight(extended_coord, extended_atype, nlists_) → List[torch.Tensor]

This should be a list of user defined weights that matches the number of models to be combined.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → List[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

compute_or_load_out_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters:

sampled_func

The lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

class deepmd.pt.model.atomic_model.PairTabAtomicModel(tab_file: str, rcut: float, sel: int | List[int], type_map: List[str], **kwargs)

Bases: deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel

Pairwise tabulation energy model.

This model can be used to tabulate the pairwise energy between atoms for either short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not be used alone, but rather as one submodel of a linear (sum) model, such as DP+D3.

Do not put the model on the first model of a linear model, since the linear model fetches the type map from the first model.

At this moment, the model does not smooth the energy at the cutoff radius, so one needs to make sure the energy has been smoothed to zero.

Parameters:

tab_filestr

The path to the tabulation file.

rcutfloat

The cutoff radius.

selint or list[int]

The maxmum number of atoms in the cut-off radius.

type_mapList[str]

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

rcondfloat, optional

The condition number for the regression of atomic energy.

atom_ener

Specifying atomic energy contribution in vacuum. The set_davg_zero key in the descrptor should be set.

_set_pairtab(tab_file: str, rcut: float) → deepmd.utils.pair_tab.PairTab

fitting_output_def() → deepmd.dpmodel.FittingOutputDef

Get the output def of developer implemented atomic models.

get_out_bias() → torch.Tensor

get_rcut() → float

Get the cut-off radius.

get_type_map() → List[str]

Get the type map.

get_sel() → List[int]

Returns the number of selected atoms for each type.

get_nsel() → int

Returns the total number of selected neighboring atoms in the cut-off radius.

mixed_types() → bool

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

serialize() → dict

classmethod deserialize(data) → PairTabAtomicModel

compute_or_load_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | None = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

forward_atomic(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: Dict[str, torch.Tensor] | None = None) → Dict[str, torch.Tensor]

_pair_tabulated_inter(nlist: torch.Tensor, i_type: torch.Tensor, j_type: torch.Tensor, rr: torch.Tensor) → torch.Tensor

Pairwise tabulated energy.

Parameters:

nlisttorch.Tensor

The unmasked neighbour list. (nframes, nloc)

i_typetorch.Tensor

The integer representation of atom type for all local atoms for all frames. (nframes, nloc)

j_typetorch.Tensor

The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)

rrtorch.Tensor

The salar distance vector between two atoms. (nframes, nloc, nnei)

Returns:

torch.Tensor

The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei)

Raises:

Exception

If the distance is beyond the table.

Notes

This function is used to calculate the pairwise energy between two atoms. It uses a table containing cubic spline coefficients calculated in PairTab.

static _get_pairwise_dist(coords: torch.Tensor, nlist: torch.Tensor) → torch.Tensor

Get pairwise distance dr.

Parameters:

coordstorch.Tensor

The coordinate of the atoms, shape of (nframes, nall, 3).

nlist

The masked nlist, shape of (nframes, nloc, nnei)

Returns:

torch.Tensor

The pairwise distance between the atoms (nframes, nloc, nnei).

static _extract_spline_coefficient(i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor, tab_data: torch.Tensor, nspline: int) → torch.Tensor

Extract the spline coefficient from the table.

Parameters:

i_typetorch.Tensor

The integer representation of atom type for all local atoms for all frames. (nframes, nloc)

j_typetorch.Tensor

The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)

idxtorch.Tensor

The index of the spline coefficient. (nframes, nloc, nnei)

tab_datatorch.Tensor

The table storing all the spline coefficient. (ntype, ntype, nspline, 4)

nsplineint

The number of splines in the table.

Returns:

torch.Tensor

The spline coefficient. (nframes, nloc, nnei, 4), shape may be squeezed.

static _calculate_ener(coef: torch.Tensor, uu: torch.Tensor) → torch.Tensor

Calculate energy using spline coeeficients.

Parameters:

coeftorch.Tensor

The spline coefficients. (nframes, nloc, nnei, 4)

uutorch.Tensor

The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei)

Returns:

torch.Tensor

The atomic energy for all local atoms for all frames. (nframes, nloc, nnei)

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → List[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

class deepmd.pt.model.atomic_model.DPPolarAtomicModel(descriptor, fitting, type_map, **kwargs)

Bases: deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor

Descriptor

fitting_net

Fitting net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

apply_out_stat(ret: Dict[str, torch.Tensor], atype: torch.Tensor)

Apply the stat to each atomic output.

Parameters:

ret

The returned dict by the forward_atomic method

atype

The atom types. nf x nloc