deepmd 0.1.0

DeePMD-kit deep potential models as RLX IR graph builders
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! 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`](rlx_ir::Graph)
//! objects ready for compilation onto any [RLX runtime](https://docs.rs/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.

pub mod config;
pub mod descriptor;
pub mod descriptor_se_r;
pub mod descriptor_dpa1;
pub mod descriptor_dpa2;
pub mod descriptor_dpa3;
pub mod descriptor_hybrid;
pub mod descriptor_repflows;
pub mod descriptor_repformers;
pub mod descriptor_se_atten_v2;
pub mod descriptor_se_t;
pub mod env_mat;
pub mod descriptor_se_t_tebd;
pub mod type_embed;
pub mod weights;
pub mod zbl;
pub mod fitting;
pub mod model;
pub mod model_dp;
pub mod nn;
pub mod parity;
pub mod spin;
pub mod transform_output;

pub use config::{DPModelConfig, EnerFittingConfig, SeAConfig};
pub use descriptor::{build_se_a_descriptor, SeADescriptor, SeAExtraInputs};
pub use descriptor_se_r::{build_se_r_descriptor, SeRConfig, SeRDescriptor};
pub use descriptor_se_t::{build_se_t_descriptor, SeTConfig, SeTDescriptor};
pub use descriptor_dpa1::{build_dpa1_descriptor, DPA1Config, DPA1Descriptor, DPA1Inputs};
pub use descriptor_dpa2::{
    build_dpa2_descriptor, DPA2Config, DPA2Descriptor, DPA2Inputs, RepinitArgs,
};
pub use descriptor_dpa3::{build_dpa3_descriptor, DPA3Config, DPA3Descriptor, DPA3Inputs};
pub use descriptor_repflows::{
    build_repflows, RepflowsConfig, RepflowsInputs, RepflowsOutputs,
};
pub use descriptor_repformers::{
    build_repformers, RepformersConfig, RepformersInputs, RepformersOutputs,
};
pub use descriptor_se_atten_v2::{build_se_atten_v2_descriptor, SeAttenV2Config};
pub use descriptor_hybrid::concat_descriptor_outputs;
pub use env_mat::{apply_stats, exp_sw, make_env_mat, smooth_weight, EnvMatOutput, EnvMatParams};
pub use descriptor_se_t_tebd::{
    build_se_t_tebd_descriptor, SeTTebdConfig, SeTTebdDescriptor, SeTTebdInputs, TebdInputMode,
};
pub use type_embed::{build_type_embedding, gather_per_atom_embed, TypeEmbedConfig};
pub use weights::{
    load_model_file, map_to_graph_keys, parse_model_value, WeightTable, WeightTensor,
};
pub use zbl::{
    build_pair_tab_energy, build_zbl_weight, combine_dp_zbl, PairTabConfig, PairTabInputs,
};
pub use fitting::{
    build_dipole_fitting, build_ener_fitting, build_invar_fitting, build_polar_fitting,
    DipoleFitting, DipoleFittingConfig, EnerFitting, FittingInputs, InvarFitting,
    PolarFitting, PolarFittingConfig,
};
pub use model::{build_dp_energy_graph, DPEnergyGraph};
pub use nn::{ActivationKind, DenseLayerSpec, MlpSpec};
pub use parity::{assert_allclose, assert_shape, GoldenFixture, GoldenInput, GoldenOutput};
pub use spin::{build_spin_inputs, split_spin_output, SpinConfig, SpinExpanded};
pub use transform_output::{
    build_force_grad_graph, build_total_energy, build_virial_via_position,
    reduce_atomic_mean, reduce_atomic_sum,
};