Crate deepmd 

DeePMD-kit deep potential models as RLX IR graph builders.

This crate is a Rust translation of the inference-side graph construction in deepmd.dpmodel (NumPy / array-api reference backend of DeePMD-kit). It emits RLX Graph objects ready for compilation onto any RLX runtime backend.

Coverage

• config — descriptor and fitting configs (se_e2_a, DPA-1/2/3, …).
• descriptor and sibling modules — DeepPot-SE, DPA, RepFlow, RepFormer, hybrid, and attention descriptors.
• fitting — invariant, dipole, and polar fitting nets.
• model / model_dp — end-to-end energy, DOS, dipole, and polar graphs.
• env_mat — host-side environment-matrix helpers (switching, stats).
• weights — .pd / JSON weight loading and key translation.
• parity — golden-fixture helpers used by the integration test suite.

The environment matrix R (radial + 3 unit-vector components per neighbor) is taken as a graph input, not constructed inside the graph: neighbor lists and the switching function live outside the IR. Producing R on the host or in a fused pre-pass is left to the caller.

Re-exports

pub use config::DPModelConfig;

pub use config::EnerFittingConfig;

pub use config::SeAConfig;

pub use descriptor::build_se_a_descriptor;

pub use descriptor::SeADescriptor;

pub use descriptor::SeAExtraInputs;

pub use descriptor_se_r::build_se_r_descriptor;

pub use descriptor_se_r::SeRConfig;

pub use descriptor_se_r::SeRDescriptor;

pub use descriptor_se_t::build_se_t_descriptor;

pub use descriptor_se_t::SeTConfig;

pub use descriptor_se_t::SeTDescriptor;

pub use descriptor_dpa1::build_dpa1_descriptor;

pub use descriptor_dpa1::DPA1Config;

pub use descriptor_dpa1::DPA1Descriptor;

pub use descriptor_dpa1::DPA1Inputs;

pub use descriptor_dpa2::build_dpa2_descriptor;

pub use descriptor_dpa2::DPA2Config;

pub use descriptor_dpa2::DPA2Descriptor;

pub use descriptor_dpa2::DPA2Inputs;

pub use descriptor_dpa2::RepinitArgs;

pub use descriptor_dpa3::build_dpa3_descriptor;

pub use descriptor_dpa3::DPA3Config;

pub use descriptor_dpa3::DPA3Descriptor;

pub use descriptor_dpa3::DPA3Inputs;

pub use descriptor_repflows::build_repflows;

pub use descriptor_repflows::RepflowsConfig;

pub use descriptor_repflows::RepflowsInputs;

pub use descriptor_repflows::RepflowsOutputs;

pub use descriptor_repformers::build_repformers;

pub use descriptor_repformers::RepformersConfig;

pub use descriptor_repformers::RepformersInputs;

pub use descriptor_repformers::RepformersOutputs;

pub use descriptor_se_atten_v2::build_se_atten_v2_descriptor;

pub use descriptor_se_atten_v2::SeAttenV2Config;

pub use descriptor_hybrid::concat_descriptor_outputs;

pub use env_mat::apply_stats;

pub use env_mat::exp_sw;

pub use env_mat::make_env_mat;

pub use env_mat::smooth_weight;

pub use env_mat::EnvMatOutput;

pub use env_mat::EnvMatParams;

pub use descriptor_se_t_tebd::build_se_t_tebd_descriptor;

pub use descriptor_se_t_tebd::SeTTebdConfig;

pub use descriptor_se_t_tebd::SeTTebdDescriptor;

pub use descriptor_se_t_tebd::SeTTebdInputs;

pub use descriptor_se_t_tebd::TebdInputMode;

pub use type_embed::build_type_embedding;

pub use type_embed::gather_per_atom_embed;

pub use type_embed::TypeEmbedConfig;

pub use weights::load_model_file;

pub use weights::map_to_graph_keys;

pub use weights::parse_model_value;

pub use weights::WeightTable;

pub use weights::WeightTensor;

pub use zbl::build_pair_tab_energy;

pub use zbl::build_zbl_weight;

pub use zbl::combine_dp_zbl;

pub use zbl::PairTabConfig;

pub use zbl::PairTabInputs;

pub use fitting::build_dipole_fitting;

pub use fitting::build_ener_fitting;

pub use fitting::build_invar_fitting;

pub use fitting::build_polar_fitting;

pub use fitting::DipoleFitting;

pub use fitting::DipoleFittingConfig;

pub use fitting::EnerFitting;

pub use fitting::FittingInputs;

pub use fitting::InvarFitting;

pub use fitting::PolarFitting;

pub use fitting::PolarFittingConfig;

pub use model::build_dp_energy_graph;

pub use model::DPEnergyGraph;

pub use nn::ActivationKind;

pub use nn::DenseLayerSpec;

pub use nn::MlpSpec;

pub use parity::assert_allclose;

pub use parity::assert_shape;

pub use parity::GoldenFixture;

pub use parity::GoldenInput;

pub use parity::GoldenOutput;

pub use spin::build_spin_inputs;

pub use spin::split_spin_output;

pub use spin::SpinConfig;

pub use spin::SpinExpanded;

pub use transform_output::build_force_grad_graph;

pub use transform_output::build_total_energy;

pub use transform_output::build_virial_via_position;

pub use transform_output::reduce_atomic_mean;

pub use transform_output::reduce_atomic_sum;

Modules

config

Configuration structs for DeePMD descriptors, fittings, and models.

descriptor

se_e2_a (DeepPot-SE) descriptor graph builder.

descriptor_dpa1

DPA-1 descriptor — attention-augmented DeepPot-SE.

descriptor_dpa2

DPA-2 descriptor — repinit + g1 shape transform + repformers stack.

descriptor_dpa3

DPA-3 descriptor — type embedding + repflows stack.

descriptor_hybrid

Hybrid descriptor — concatenates outputs of multiple descriptors.

descriptor_repflows

Repflows descriptor block — DPA-3’s message-passing stack.

descriptor_repformers

Repformer descriptor block — DPA-2’s message-passing layer stack.

descriptor_se_atten_v2

se_atten_v2 descriptor — DPA-1 with tebd_input_mode="strip" and smooth type-embedding mixing forced on.

descriptor_se_r

se_r (radial-only DeepPot-SE) descriptor graph builder.

descriptor_se_t

se_t (DeepPot-SE T, three-body) descriptor graph builder.

descriptor_se_t_tebd

se_t_tebd descriptor — three-body with type embeddings.

env_mat

Host-side environment matrix construction.

fitting

Fitting-net graph builders.

model

DP energy model graph: descriptor + fitting + total-energy reduction.

model_dp

DP-style model graphs for the non-energy output kinds: dipole, polarizability, DOS, and generic invariant property.

nn

Shared neural-network building blocks: activations, dense layers, embedding/fitting MLPs. Translated from deepmd/dpmodel/utils/network.py.

parity

Parity-testing utilities — compare Rust graph output against Python-reference golden fixtures.

spin

Spin model graph helpers.

transform_output

Per-atom → global output transforms and gradient builders.

type_embed

Type embedding network — deepmd/dpmodel/utils/type_embed.py.

weights

Weight loading for serialized DeePMD-kit models.

zbl

ZBL composite model graph builders.