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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! `se_atten_v2` descriptor — DPA-1 with `tebd_input_mode="strip"` and
//! smooth type-embedding mixing forced on.
//!
//! Translated from
//! `deepmd/dpmodel/descriptor/se_atten_v2.py::DescrptSeAttenV2`.
//! That class is implemented in Python as a thin subclass of
//! `DescrptDPA1` that pins those two flags; we mirror that exactly via
//! [`DPA1Config`] presets.

use anyhow::Result;
use rlx_ir::Graph;

use crate::descriptor_dpa1::{
    build_dpa1_descriptor, DPA1Config, DPA1Descriptor, DPA1Inputs,
};
use crate::descriptor_se_t_tebd::TebdInputMode;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SeAttenV2Config {
    pub rcut: f64,
    pub rcut_smth: f64,
    pub sel: Vec<usize>,
    pub ntypes: usize,
    #[serde(default = "default_neuron")]
    pub neuron: Vec<usize>,
    #[serde(default = "default_axis_neuron")]
    pub axis_neuron: usize,
    #[serde(default = "default_tebd_dim")]
    pub tebd_dim: usize,
    #[serde(default)]
    pub resnet_dt: bool,
    #[serde(default = "default_attn")]
    pub attn: usize,
    #[serde(default = "default_attn_layer")]
    pub attn_layer: usize,
    #[serde(default = "default_true")]
    pub attn_dotr: bool,
    #[serde(default = "default_num_heads")]
    pub num_heads: usize,
    #[serde(default = "default_true")]
    pub normalize: bool,
    #[serde(default = "default_scaling_factor")]
    pub scaling_factor: f32,
    #[serde(default = "default_ln_eps")]
    pub ln_eps: f32,
    #[serde(default)]
    pub type_one_side: bool,
    #[serde(default = "default_activation")]
    pub activation_function: String,
    #[serde(default = "default_true")]
    pub concat_output_tebd: bool,
}

fn default_neuron() -> Vec<usize> {
    vec![25, 50, 100]
}
fn default_axis_neuron() -> usize {
    8
}
fn default_tebd_dim() -> usize {
    8
}
fn default_attn() -> usize {
    128
}
fn default_attn_layer() -> usize {
    2
}
fn default_num_heads() -> usize {
    1
}
fn default_ln_eps() -> f32 {
    1e-5
}
fn default_true() -> bool {
    true
}
fn default_scaling_factor() -> f32 {
    1.0
}
fn default_activation() -> String {
    "tanh".into()
}

impl SeAttenV2Config {
    pub fn to_dpa1(&self) -> DPA1Config {
        DPA1Config {
            rcut: self.rcut,
            rcut_smth: self.rcut_smth,
            sel: self.sel.clone(),
            ntypes: self.ntypes,
            neuron: self.neuron.clone(),
            axis_neuron: self.axis_neuron,
            tebd_dim: self.tebd_dim,
            tebd_input_mode: TebdInputMode::Strip,
            resnet_dt: self.resnet_dt,
            attn: self.attn,
            attn_layer: self.attn_layer,
            attn_dotr: self.attn_dotr,
            num_heads: self.num_heads,
            normalize: self.normalize,
            scaling_factor: self.scaling_factor,
            ln_eps: self.ln_eps,
            smooth: true,
            type_one_side: self.type_one_side,
            activation_function: self.activation_function.clone(),
            concat_output_tebd: self.concat_output_tebd,
        }
    }
}

pub fn build_se_atten_v2_descriptor(
    g: &mut Graph,
    cfg: &SeAttenV2Config,
    inputs: DPA1Inputs,
    nf: usize,
    nloc: usize,
) -> Result<DPA1Descriptor> {
    let dpa1 = cfg.to_dpa1();
    build_dpa1_descriptor(g, &dpa1, inputs, nf, nloc)
}