deepmd.jax.model.hlo

Attributes

OUTPUT_DEFS

Classes

HLO	Base class for final exported model that will be directly used for inference.

Module Contents

deepmd.jax.model.hlo.OUTPUT_DEFS

class deepmd.jax.model.hlo.HLO(stablehlo, stablehlo_atomic_virial, stablehlo_no_ghost, stablehlo_atomic_virial_no_ghost, model_def_script, type_map, rcut, dim_fparam, dim_aparam, sel_type, is_aparam_nall, model_output_type, mixed_types, min_nbor_dist, sel)

Bases: deepmd.jax.model.base_model.BaseModel

Base class for final exported model that will be directly used for inference.

The class defines some abstractmethods that will be directly called by the inference interface. If the final model class inherits some of those methods from other classes, BaseModel should be inherited as the last class to ensure the correct method resolution order.

This class is backend-indepedent.

See also

deepmd.dpmodel.model.base_model.BaseModel

BaseModel class for DPModel backend.

_call_lower

_call_lower_atomic_virial

_call_lower_no_ghost

_call_lower_atomic_virial_no_ghost

stablehlo

type_map

rcut

dim_fparam

dim_aparam

sel_type

_is_aparam_nall

_model_output_type

_mixed_types

min_nbor_dist

The minimum distance between two atoms. Used for model compression. None when skipping neighbor statistics.

sel

model_def_script

The model definition script.

__call__(coord: deepmd.jax.env.jnp.ndarray, atype: deepmd.jax.env.jnp.ndarray, box: 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) → Any

Return model prediction.

Parameters:

coord

The coordinates of the atoms. shape: nf x (nloc x 3)

atype

The type of atoms. shape: nf x nloc

box

The simulation box. shape: nf x 9

fparam

frame parameter. nf x ndf

aparam

atomic parameter. nf x nloc x nda

do_atomic_virial

If calculate the atomic virial.

Returns:

ret_dict

The result dict of type dict[str,np.ndarray]. The keys are defined by the ModelOutputDef.

call(coord: deepmd.jax.env.jnp.ndarray, atype: deepmd.jax.env.jnp.ndarray, box: 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)

Return model prediction.

Parameters:

coord

The coordinates of the atoms. shape: nf x (nloc x 3)

atype

The type of atoms. shape: nf x nloc

box

The simulation box. shape: nf x 9

fparam

frame parameter. nf x ndf

aparam

atomic parameter. nf x nloc x nda

do_atomic_virial

If calculate the atomic virial.

Returns:

ret_dict

The result dict of type dict[str,np.ndarray]. The keys are defined by the ModelOutputDef.

model_output_def()

call_lower(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)

get_type_map() → list[str]

Get the type map.

get_rcut()

Get the cut-off radius.

get_dim_fparam()

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

get_dim_aparam()

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).

model_output_type() → list[str]

Get the output type for the model.

abstract serialize() → dict

Serialize the model.

Returns:

dict

The serialized data

classmethod deserialize(data: dict) → deepmd.jax.model.base_model.BaseModel

Abstractmethod:

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BaseModel

The deserialized model

get_model_def_script() → str

Get the model definition script.

get_min_nbor_dist() → float | None

Get the minimum distance between two atoms.

get_nnei() → int

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

get_sel() → list[int]

get_nsel() → int

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

mixed_types() → bool

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Abstractmethod:

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statictics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

classmethod get_model(model_params: dict) → deepmd.jax.model.base_model.BaseModel

Abstractmethod:

Get the model by the parameters.

By default, all the parameters are directly passed to the constructor. If not, override this method.

Parameters:

model_paramsdict

The model parameters

Returns:

BaseBaseModel

The model