deepmd.pt.utils package

class deepmd.pt.utils.AtomExcludeMask(ntypes: int, exclude_types: List[int] = [])

Bases: Module

Computes the type exclusion mask for atoms.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(atype)	Compute type exclusion mask for atoms.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit

forward(atype: Tensor) → Tensor

Compute type exclusion mask for atoms.

Parameters

atype

The extended aotm types. shape: nf x natom

Returns

mask

The type exclusion mask for atoms. shape: nf x natom Element [ff,ii] being 0 if type(ii) is excluded, otherwise being 1.

reinit(ntypes: int, exclude_types: List[int] = [])

training: bool

class deepmd.pt.utils.PairExcludeMask(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

Bases: Module

Computes the type exclusion mask for atom pairs.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(nlist, atype_ext)	Compute type exclusion mask.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit

forward(nlist: Tensor, atype_ext: Tensor) → Tensor

Compute type exclusion mask.

Parameters

nlist

The neighbor list. shape: nf x nloc x nnei

atype_ext

The extended aotm types. shape: nf x nall

Returns

mask

The type exclusion mask of shape: nf x nloc x nnei. Element [ff,ii,jj] being 0 if type(ii), type(nlist[ff,ii,jj]) is excluded, otherwise being 1.

reinit(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

training: bool

Submodules

deepmd.pt.utils.ase_calc module

deepmd.pt.utils.ase_calc.DPCalculator

alias of DP

deepmd.pt.utils.auto_batch_size module

class deepmd.pt.utils.auto_batch_size.AutoBatchSize(initial_batch_size: int = 1024, factor: float = 2.0)

Bases: AutoBatchSize

Methods

execute(callable, start_index, natoms)	Excuate a method with given batch size.
execute_all(callable, total_size, natoms, ...)	Excuate a method with all given data.
is_gpu_available()	Check if GPU is available.
is_oom_error(e)	Check if the exception is an OOM error.

is_gpu_available() → bool

Check if GPU is available.

Returns

bool

True if GPU is available

is_oom_error(e: Exception) → bool

Check if the exception is an OOM error.

Parameters

eException

Exception

deepmd.pt.utils.cache module

deepmd.pt.utils.cache.lru_cache(maxsize=16, typed=False, copy=False, deepcopy=False)

deepmd.pt.utils.dataloader module

class deepmd.pt.utils.dataloader.BackgroundConsumer(queue, source, max_len)

Bases: Thread

Attributes

daemon

A boolean value indicating whether this thread is a daemon thread.

ident

Thread identifier of this thread or None if it has not been started.

name

A string used for identification purposes only.

native_id

Native integral thread ID of this thread, or None if it has not been started.

Methods

is_alive()	Return whether the thread is alive.
join([timeout])	Wait until the thread terminates.
run()	Method representing the thread's activity.
start()	Start the thread's activity.

getName
isDaemon
setDaemon
setName

run()

Method representing the thread’s activity.

You may override this method in a subclass. The standard run() method invokes the callable object passed to the object’s constructor as the target argument, if any, with sequential and keyword arguments taken from the args and kwargs arguments, respectively.

class deepmd.pt.utils.dataloader.BufferedIterator(iterable)

Bases: object

class deepmd.pt.utils.dataloader.DpLoaderSet(systems, batch_size, model_params, seed=10, shuffle=True)

Bases: Dataset

A dataset for storing DataLoaders to multiple Systems.

Methods

add_data_requirement(data_requirement)

Add data requirement for each system in multiple systems.

print_summary
set_noise

add_data_requirement(data_requirement: List[DataRequirementItem])

Add data requirement for each system in multiple systems.

print_summary(name: str, prob: List[float])

set_noise(noise_settings)

deepmd.pt.utils.dataloader.collate_batch(batch)

deepmd.pt.utils.dataloader.get_weighted_sampler(training_data, prob_style, sys_prob=False)

deepmd.pt.utils.dataloader.setup_seed(seed)

deepmd.pt.utils.dataset module

class deepmd.pt.utils.dataset.DeepmdDataSetForLoader(system: str, type_map: str, shuffle=True)

Bases: Dataset

Methods

add_data_requirement(data_requirement)

Add data requirement for this data system.

add_data_requirement(data_requirement: List[DataRequirementItem])

Add data requirement for this data system.

deepmd.pt.utils.dp_random module

deepmd.pt.utils.dp_random.choice(a: Union[ndarray, int], size: Optional[Union[int, Tuple[int, ...]]] = None, replace: bool = True, p: Optional[ndarray] = None)

Generates a random sample from a given 1-D array.

Parameters

a1-D array_like or int

If an ndarray, a random sample is generated from its elements. If an int, the random sample is generated as if it were np.arange(a)

sizeint or tuple of ints, optional

Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn. Default is None, in which case a single value is returned.

replacebool, optional

Whether the sample is with or without replacement. Default is True, meaning that a value of a can be selected multiple times.

p1-D array_like, optional

The probabilities associated with each entry in a. If not given, the sample assumes a uniform distribution over all entries in a.

Returns

np.ndarray

arrays with results and their shapes

deepmd.pt.utils.dp_random.random(size=None)

Return random floats in the half-open interval [0.0, 1.0).

Parameters

size

Output shape.

Returns

np.ndarray

Arrays with results and their shapes.

deepmd.pt.utils.dp_random.seed(val: Optional[int] = None)

Seed the generator.

Parameters

valint

Seed.

deepmd.pt.utils.dp_random.shuffle(x: ndarray)

Modify a sequence in-place by shuffling its contents.

Parameters

xnp.ndarray

The array or list to be shuffled.

deepmd.pt.utils.env module

deepmd.pt.utils.env.GLOBAL_ENER_FLOAT_PRECISION

alias of float64

deepmd.pt.utils.env.GLOBAL_NP_FLOAT_PRECISION

alias of float64

deepmd.pt.utils.env_mat_stat module

class deepmd.pt.utils.env_mat_stat.EnvMatStat

Bases: EnvMatStat

Methods

compute_stat(env_mat)	Compute the statistics of the environment matrix for a single system.
compute_stats(data)	Compute the statistics of the environment matrix.
get_avg([default])	Get the average of the environment matrix.
get_std([default, protection])	Get the standard deviation of the environment matrix.
iter(data)	Get the iterator of the environment matrix.
load_or_compute_stats(data[, path])	Load the statistics of the environment matrix if it exists, otherwise compute and save it.
load_stats(path)	Load the statistics of the environment matrix.
save_stats(path)	Save the statistics of the environment matrix.

compute_stat(env_mat: Dict[str, Tensor]) → Dict[str, StatItem]

Compute the statistics of the environment matrix for a single system.

Parameters

env_mattorch.Tensor

The environment matrix.

Returns

Dict[str, StatItem]

The statistics of the environment matrix.

class deepmd.pt.utils.env_mat_stat.EnvMatStatSe(descriptor: DescriptorBlock)

Bases: EnvMatStat

Environmental matrix statistics for the se_a/se_r environemntal matrix.

Parameters

descriptorDescriptorBlock

The descriptor of the model.

Methods

__call__()	Call self as a function.
compute_stat(env_mat)	Compute the statistics of the environment matrix for a single system.
compute_stats(data)	Compute the statistics of the environment matrix.
get_avg([default])	Get the average of the environment matrix.
get_hash()	Get the hash of the environment matrix.
get_std([default, protection])	Get the standard deviation of the environment matrix.
iter(data)	Get the iterator of the environment matrix.
load_or_compute_stats(data[, path])	Load the statistics of the environment matrix if it exists, otherwise compute and save it.
load_stats(path)	Load the statistics of the environment matrix.
save_stats(path)	Save the statistics of the environment matrix.

get_hash() → str

Get the hash of the environment matrix.

Returns

str

The hash of the environment matrix.

iter(data: List[Dict[str, Tensor]]) → Iterator[Dict[str, StatItem]]

Get the iterator of the environment matrix.

Parameters

dataList[Dict[str, torch.Tensor]]

The data.

Yields

Dict[str, StatItem]

The statistics of the environment matrix.

deepmd.pt.utils.exclude_mask module

class deepmd.pt.utils.exclude_mask.AtomExcludeMask(ntypes: int, exclude_types: List[int] = [])

Bases: Module

Computes the type exclusion mask for atoms.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(atype)	Compute type exclusion mask for atoms.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit

forward(atype: Tensor) → Tensor

Compute type exclusion mask for atoms.

Parameters

atype

The extended aotm types. shape: nf x natom

Returns

mask

The type exclusion mask for atoms. shape: nf x natom Element [ff,ii] being 0 if type(ii) is excluded, otherwise being 1.

reinit(ntypes: int, exclude_types: List[int] = [])

training: bool

class deepmd.pt.utils.exclude_mask.PairExcludeMask(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

Bases: Module

Computes the type exclusion mask for atom pairs.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(nlist, atype_ext)	Compute type exclusion mask.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit

forward(nlist: Tensor, atype_ext: Tensor) → Tensor

Compute type exclusion mask.

Parameters

nlist

The neighbor list. shape: nf x nloc x nnei

atype_ext

The extended aotm types. shape: nf x nall

Returns

mask

The type exclusion mask of shape: nf x nloc x nnei. Element [ff,ii,jj] being 0 if type(ii), type(nlist[ff,ii,jj]) is excluded, otherwise being 1.

reinit(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

training: bool

deepmd.pt.utils.finetune module

deepmd.pt.utils.finetune.change_finetune_model_params(ckpt, finetune_model, model_config, multi_task=False, model_branch='')

Load model_params according to the pretrained one.

Args: - ckpt & finetune_model: origin model. - config: Read from json file.

deepmd.pt.utils.learning_rate module

class deepmd.pt.utils.learning_rate.LearningRateExp(start_lr, stop_lr, decay_steps, stop_steps, **kwargs)

Bases: object

Methods

value(step)

Get the learning rate at the given step.

value(step)

Get the learning rate at the given step.

deepmd.pt.utils.multi_task module

deepmd.pt.utils.multi_task.get_class_name(item_key, item_params)

deepmd.pt.utils.multi_task.preprocess_shared_params(model_config)

Preprocess the model params for multitask model, and generate the links dict for further sharing.

Parameters

model_config – Model params of multitask model.

Returns

model_config: Preprocessed model params of multitask model.

Those string names are replaced with real params in shared_dict of model params.

shared_links: Dict of link infos for further sharing.

Each item, whose key must be in shared_dict, is a dict with following keys: - “type”: The real class type of this item. - “links”: List of shared settings, each sub-item is a dict with following keys:

• “model_key”: Model key in the model_dict to share this item.

• “shared_type”: Type of this shard item.

• “shared_level”: Shared level (int) of this item in this model.

  Lower for more params to share, 0 means to share all params in this item.

This list are sorted by “shared_level”.

For example, if one has model_config like this:

“model”: {

“shared_dict”: {

“my_type_map”: [“foo”, “bar”], “my_des1”: {

“type”: “se_e2_a”, “neuron”: [10, 20, 40] },

}, “model_dict”: {

“model_1”: {

“type_map”: “my_type_map”, “descriptor”: “my_des1”, “fitting_net”: {

“neuron”: [100, 100, 100]

}

}, “model_2”: {

“type_map”: “my_type_map”, “descriptor”: “my_des1”, “fitting_net”: {

“neuron”: [100, 100, 100]

}

} “model_3”: {

“type_map”: “my_type_map”, “descriptor”: “my_des1:1”, “fitting_net”: {

“neuron”: [100, 100, 100]

}

}

}

}

The above config will init three model branches named model_1 and model_2 and model_3,

in which:

• model_2 and model_3 will have the same type_map as that in model_1.

• model_2 will share all the parameters of descriptor with model_1,

while model_3 will share part of parameters of descriptor with model_1 on human-defined share-level 1 (default is 0, meaning share all the parameters). - model_1, model_2 and model_3 have three different `fitting_net`s.

The returned model_config will automatically fulfill the input model_config as if there’s no sharing,

and the shared_links will keep all the sharing information with looking:

{

‘my_des1’: {

‘type’: ‘DescrptSeA’, ‘links’: [

{‘model_key’: ‘model_1’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 0}, {‘model_key’: ‘model_2’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 0}, {‘model_key’: ‘model_3’, ‘shared_type’: ‘descriptor’, ‘shared_level’: 1} ]

}

}

deepmd.pt.utils.neighbor_stat module

class deepmd.pt.utils.neighbor_stat.NeighborStat(ntypes: int, rcut: float, mixed_type: bool = False)

Bases: NeighborStat

Neighbor statistics using pure NumPy.

Parameters

ntypesint

The num of atom types

rcutfloat

The cut-off radius

mixed_typebool, optional, default=False

Treat all types as a single type.

Methods

get_stat(data)	Get the data statistics of the training data, including nearest nbor distance between atoms, max nbor size of atoms.
iterator(data)	Abstract method for producing data.

iterator(data: DeepmdDataSystem) → Iterator[Tuple[ndarray, float, str]]

Abstract method for producing data.

Yields

np.ndarray

The maximal number of neighbors

float

The squared minimal distance between two atoms

str

The directory of the data system

class deepmd.pt.utils.neighbor_stat.NeighborStatOP(ntypes: int, rcut: float, mixed_types: bool)

Bases: Module

Class for getting neighbor statics data information.

Parameters

ntypes

The num of atom types

rcut

The cut-off radius

mixed_typesbool, optional

If True, treat neighbors of all types as a single type.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(coord, atype, cell)	Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(coord: Tensor, atype: Tensor, cell: Optional[Tensor]) → Tuple[Tensor, Tensor]

Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix.

Parameters

coord

The coordinates of atoms.

atype

The atom types.

cell

The cell.

Returns

torch.Tensor

The minimal squared distance between two atoms, in the shape of (nframes,)

torch.Tensor

The maximal number of neighbors

training: bool

deepmd.pt.utils.nlist module

deepmd.pt.utils.nlist.build_multiple_neighbor_list(coord: Tensor, nlist: Tensor, rcuts: List[float], nsels: List[int]) → Dict[str, Tensor]

Input one neighbor list, and produce multiple neighbor lists with different cutoff radius and numbers of selection out of it. The required rcuts and nsels should be smaller or equal to the input nlist.

Parameters

coordtorch.Tensor

exptended coordinates of shape [batch_size, nall x 3]

nlisttorch.Tensor

Neighbor list of shape [batch_size, nloc, nsel], the neighbors should be stored in an ascending order.

rcutsList[float]

list of cut-off radius in ascending order.

nselsList[int]

maximal number of neighbors in ascending order.

Returns

nlist_dictDict[str, torch.Tensor]

A dict of nlists, key given by get_multiple_nlist_key(rc, nsel) value being the corresponding nlist.

deepmd.pt.utils.nlist.build_neighbor_list(coord1: Tensor, atype: Tensor, nloc: int, rcut: float, sel: Union[int, List[int]], distinguish_types: bool = True) → Tensor

Build neightbor list for a single frame. keeps nsel neighbors.

Parameters

coord1torch.Tensor

exptended coordinates of shape [batch_size, nall x 3]

atypetorch.Tensor

extended atomic types of shape [batch_size, nall]

nlocint

number of local atoms.

rcutfloat

cut-off radius

selint or List[int]

maximal number of neighbors (of each type). if distinguish_types==True, nsel should be list and the length of nsel should be equal to number of types.

distinguish_typesbool

distinguish different types.

Returns

neighbor_listtorch.Tensor

Neighbor list of shape [batch_size, nloc, nsel], the neighbors are stored in an ascending order. If the number of neighbors is less than nsel, the positions are masked with -1. The neighbor list of an atom looks like |------ nsel ------| xx xx xx xx -1 -1 -1 if distinguish_types==True and we have two types |---- nsel[0] -----| |---- nsel[1] -----| xx xx xx xx -1 -1 -1 xx xx xx -1 -1 -1 -1

deepmd.pt.utils.nlist.extend_coord_with_ghosts(coord: Tensor, atype: Tensor, cell: Optional[Tensor], rcut: float)

Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors within rcut are appended.

Parameters

coordtorch.Tensor

original coordinates of shape [-1, nloc*3].

atypetorch.Tensor

atom type of shape [-1, nloc].

celltorch.Tensor

simulation cell tensor of shape [-1, 9].

rcutfloat

the cutoff radius

Returns

extended_coord: torch.Tensor

extended coordinates of shape [-1, nall*3].

extended_atype: torch.Tensor

extended atom type of shape [-1, nall].

index_mapping: torch.Tensor

maping extended index to the local index

deepmd.pt.utils.nlist.extend_input_and_build_neighbor_list(coord, atype, rcut: float, sel: List[int], mixed_types: bool = False, box: Optional[Tensor] = None)

deepmd.pt.utils.nlist.get_multiple_nlist_key(rcut: float, nsel: int) → str

deepmd.pt.utils.nlist.nlist_distinguish_types(nlist: Tensor, atype: Tensor, sel: List[int])

Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types.

deepmd.pt.utils.plugin module

Base of plugin systems.

class deepmd.pt.utils.plugin.Plugin

Bases: object

A class to register and restore plugins.

Examples

>>> plugin = Plugin()
>>> @plugin.register("xx")
    def xxx():
        pass
>>> print(plugin.plugins["xx"])

Attributes

pluginsDict[str, object]

plugins

Methods

get_plugin(key)	Visit a plugin by key.
register(key)	Register a plugin.

get_plugin(key) → object

Visit a plugin by key.

Parameters

keystr

key of the plugin

Returns

object

the plugin

register(key: str) → Callable[[object], object]

Register a plugin.

Parameters

keystr

key of the plugin

Returns

Callable[[object], object]

decorator

class deepmd.pt.utils.plugin.PluginVariant(*args, **kwargs)

Bases: object

A class to remove type from input arguments.

class deepmd.pt.utils.plugin.VariantABCMeta(name, bases, namespace, **kwargs)

Bases: VariantMeta, ABCMeta

Methods

__call__(*args, **kwargs)	Remove type and keys that starts with underline.
mro(/)	Return a type's method resolution order.
register(subclass)	Register a virtual subclass of an ABC.

class deepmd.pt.utils.plugin.VariantMeta

Bases: object

Methods

__call__(*args, **kwargs)

Remove type and keys that starts with underline.

deepmd.pt.utils.preprocess module

class deepmd.pt.utils.preprocess.Region3D(boxt)

Bases: object

Methods

get_face_distance()	Return face distinces to each surface of YZ, ZX, XY.
inter2phys(coord)	Convert internal coordinates to physical ones.
phys2inter(coord)	Convert physical coordinates to internal ones.

get_face_distance()

Return face distinces to each surface of YZ, ZX, XY.

inter2phys(coord)

Convert internal coordinates to physical ones.

phys2inter(coord)

Convert physical coordinates to internal ones.

deepmd.pt.utils.preprocess.append_neighbors(coord, region: Region3D, atype, rcut: float)

Make ghost atoms who are valid neighbors.

Args: - coord: shape is [nloc*3] - atype: shape is [nloc]

deepmd.pt.utils.preprocess.build_inside_clist(coord, region: Region3D, ncell)

Build cell list on atoms inside region.

Args: - coord: shape is [nloc*3] - ncell: shape is [3]

deepmd.pt.utils.preprocess.build_neighbor_list(nloc: int, coord, atype, rcut: float, sec, mapping, type_split=True, min_check=False)

For each atom inside region, build its neighbor list.

Args: - coord: shape is [nall*3] - atype: shape is [nall]

deepmd.pt.utils.preprocess.compute_pbc_shift(cell_offset, ncell)

Tell shift count to move the atom into region.

deepmd.pt.utils.preprocess.compute_serial_cid(cell_offset, ncell)

Tell the sequential cell ID in its 3D space.

Args: - cell_offset: shape is [3] - ncell: shape is [3]

deepmd.pt.utils.preprocess.compute_smooth_weight(distance, rmin: float, rmax: float)

Compute smooth weight for descriptor elements.

deepmd.pt.utils.preprocess.make_env_mat(coord, atype, region, rcut: Union[float, list], sec, pbc=True, type_split=True, min_check=False)

Based on atom coordinates, return environment matrix.

Returns

nlist: nlist, [nloc, nnei]

merged_coord_shift: shift on nall atoms, [nall, 3]

merged_mapping: mapping from nall index to nloc index, [nall]

deepmd.pt.utils.preprocess.normalize_coord(coord, region: Region3D, nloc: int)

Move outer atoms into region by mirror.

Args: - coord: shape is [nloc*3]

deepmd.pt.utils.region module

deepmd.pt.utils.region.b_to_face_distance(cell)

deepmd.pt.utils.region.inter2phys(coord: Tensor, cell: Tensor) → Tensor

Convert internal(direct) coordinates to physical coordinates.

Parameters

coordtorch.Tensor

internal coordinates of shape [*, na, 3].

celltorch.Tensor

simulation cell tensor of shape [*, 3, 3].

Returns

phys_coord: torch.Tensor

the physical coordinates

deepmd.pt.utils.region.normalize_coord(coord: Tensor, cell: Tensor) → Tensor

Apply PBC according to the atomic coordinates.

Parameters

coordtorch.Tensor

orignal coordinates of shape [*, na, 3].

Returns

wrapped_coord: torch.Tensor

wrapped coordinates of shape [*, na, 3].

deepmd.pt.utils.region.phys2inter(coord: Tensor, cell: Tensor) → Tensor

Convert physical coordinates to internal(direct) coordinates.

Parameters

coordtorch.Tensor

physical coordinates of shape [*, na, 3].

celltorch.Tensor

simulation cell tensor of shape [*, 3, 3].

Returns

inter_coord: torch.Tensor

the internal coordinates

deepmd.pt.utils.region.to_face_distance(cell: Tensor) → Tensor

Compute the to-face-distance of the simulation cell.

Parameters

celltorch.Tensor

simulation cell tensor of shape [*, 3, 3].

Returns

dist: torch.Tensor

the to face distances of shape [*, 3]

deepmd.pt.utils.serialization module

deepmd.pt.utils.serialization.deserialize_to_file(model_file: str, data: dict) → None

Deserialize the dictionary to a model file.

Parameters

model_filestr

The model file to be saved.

datadict

The dictionary to be deserialized.

deepmd.pt.utils.serialization.serialize_from_file(model_file: str) → dict

Serialize the model file to a dictionary.

Parameters

model_filestr

The model file to be serialized.

Returns

dict

The serialized model data.

deepmd.pt.utils.stat module

deepmd.pt.utils.stat.make_stat_input(datasets, dataloaders, nbatches)

Pack data for statistics.

Args: - dataset: A list of dataset to analyze. - nbatches: Batch count for collecting stats.

Returns

• a list of dicts, each of which contains data from a system

deepmd.pt.utils.update_sel module

class deepmd.pt.utils.update_sel.UpdateSel

Bases: BaseUpdateSel

Attributes

neighbor_stat

Methods

get_min_nbor_dist
get_nbor_stat
get_rcut
get_sel
get_type_map
hook
parse_auto_sel
parse_auto_sel_ratio
update_one_sel
wrap_up_4

hook(min_nbor_dist, max_nbor_size)

property neighbor_stat: Type[NeighborStat]

deepmd.pt.utils.utils module

class deepmd.pt.utils.utils.ActivationFn(activation: Optional[str])

Bases: Module

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(x)	Returns the tensor after applying activation function corresponding to activation.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(x: Tensor) → Tensor

Returns the tensor after applying activation function corresponding to activation.

training: bool

deepmd.pt.utils.utils.dict_to_device(sample_dict)

deepmd.pt.utils.utils.get_activation_fn(activation: str) → Callable

Returns the activation function corresponding to activation.

deepmd.pt.utils.utils.to_numpy_array(xx: Tensor) → ndarray

deepmd.pt.utils.utils.to_numpy_array(xx: None) → None

deepmd.pt.utils.utils.to_torch_tensor(xx: ndarray) → Tensor

deepmd.pt.utils.utils.to_torch_tensor(xx: None) → None