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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216

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

//! `se_t` (DeepPot-SE T, three-body) descriptor graph builder.
//!
//! Translated from `DescrptSeT.call` in
//! `deepmd/dpmodel/descriptor/se_t.py`.
//!
//! Topology (per unordered pair of neighbor types `(ti, tj)` with
//! `ti ≤ tj`):
//!
//! ```text
//!     R    = (R_raw - davg[atype]) / dstd[atype]                # [nf, nloc, nnei, 4]
//!     R   *= exclude_mask
//!     R_i  = R[:, :, sec_ti:sec_ti+nt_i, 1:4]                   # xyz only
//!     R_j  = R[:, :, sec_tj:sec_tj+nt_j, 1:4]
//!     env  = R_i · R_jᵀ                                         # [nf, nloc, nt_i, nt_j]
//!     gg   = N_{ti,tj}(env.unsqueeze(-1))                       # [nf, nloc, nt_i, nt_j, ng]
//!     res += (env.unsqueeze(-1) * gg).sum(dims=[2,3]) / (nt_i·nt_j)
//! ```

use anyhow::{bail, Result};
use rlx_ir::infer::GraphExt;
use rlx_ir::op::ReduceOp;
use rlx_ir::{DType, Graph, NodeId, Shape};
use serde::{Deserialize, Serialize};

use crate::nn::{embedding_mlp, scalar_const, ActivationKind, MlpSpec};

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SeTConfig {
    pub rcut: f64,
    pub rcut_smth: f64,
    pub sel: Vec<usize>,
    #[serde(default = "default_neuron")]
    pub neuron: Vec<usize>,
    #[serde(default)]
    pub resnet_dt: bool,
    #[serde(default = "default_activation")]
    pub activation_function: String,
    #[serde(default)]
    pub type_map: Option<Vec<String>>,
}

fn default_neuron() -> Vec<usize> {
    vec![24, 48, 96]
}
fn default_activation() -> String {
    "tanh".into()
}

impl SeTConfig {
    pub fn ntypes(&self) -> usize {
        self.sel.len()
    }
    pub fn nnei(&self) -> usize {
        self.sel.iter().sum()
    }
    pub fn ng(&self) -> usize {
        *self.neuron.last().expect("se_t: empty neuron list")
    }
    pub fn dim_out(&self) -> usize {
        self.ng()
    }
    pub fn sel_cumsum(&self) -> Vec<usize> {
        let mut acc = Vec::with_capacity(self.sel.len() + 1);
        acc.push(0);
        let mut s = 0;
        for &n in &self.sel {
            s += n;
            acc.push(s);
        }
        acc
    }
}

pub struct SeTDescriptor {
    /// Descriptor `D` node, shape `[nf, nloc, ng]`.
    pub descriptor: NodeId,
    pub dim_out: usize,
}

pub fn build_se_t_descriptor(
    g: &mut Graph,
    cfg: &SeTConfig,
    env_mat_raw: NodeId,
    atype_loc: NodeId,
    nf: usize,
    nloc: usize,
    exclude_mask: Option<NodeId>,
) -> Result<SeTDescriptor> {
    let activation = ActivationKind::parse(&cfg.activation_function)?;
    let ntypes = cfg.ntypes();
    let nnei = cfg.nnei();
    let ng = cfg.ng();
    let sec = cfg.sel_cumsum();

    let rr_shape = g.shape(env_mat_raw).clone();
    if rr_shape.rank() != 4 {
        bail!(
            "se_t: env-matrix input must have rank 4 [nf, nloc, nnei, 4], got rank {}",
            rr_shape.rank()
        );
    }

    let davg = g.param(
        "descriptor.davg",
        Shape::new(&[ntypes, nnei, 4], DType::F32),
    );
    let dstd = g.param(
        "descriptor.dstd",
        Shape::new(&[ntypes, nnei, 4], DType::F32),
    );
    let davg_g = g.gather_(davg, atype_loc, 0);
    let dstd_g = g.gather_(dstd, atype_loc, 0);
    let mut rr = g.sub(env_mat_raw, davg_g);
    rr = g.div(rr, dstd_g);

    if let Some(mask) = exclude_mask {
        let mask_shape = Shape::new(&[nf, nloc, nnei, 1], DType::F32);
        let mask_4d = g.reshape(
            mask,
            vec![nf as i64, nloc as i64, nnei as i64, 1],
            mask_shape,
        );
        rr = g.mul(rr, mask_4d);
    }

    // Drop the radial column; keep the xyz unit-vector columns.
    let rr_xyz = g.narrow_(rr, 3, 1, 3); // [nf, nloc, nnei, 3]

    let mut acc: Option<NodeId> = None;
    for ti in 0..ntypes {
        for tj in ti..ntypes {
            let nt_i = sec[ti + 1] - sec[ti];
            let nt_j = sec[tj + 1] - sec[tj];
            if nt_i == 0 || nt_j == 0 {
                continue;
            }
            // [nf, nloc, nt_i, 3]
            let rr_i = g.narrow_(rr_xyz, 2, sec[ti], nt_i);
            // [nf, nloc, nt_j, 3]
            let rr_j = g.narrow_(rr_xyz, 2, sec[tj], nt_j);
            // env = rr_i · rr_jᵀ : [nf, nloc, nt_i, nt_j]
            let rr_j_t = g.transpose_(rr_j, vec![0, 1, 3, 2]);
            let env_ij = g.mm(rr_i, rr_j_t);

            // env_ij[..., None] → [nf, nloc, nt_i, nt_j, 1]
            let env_5d_shape = Shape::new(&[nf, nloc, nt_i, nt_j, 1], DType::F32);
            let env_5d = g.reshape(
                env_ij,
                vec![nf as i64, nloc as i64, nt_i as i64, nt_j as i64, 1],
                env_5d_shape,
            );

            // Embedding MLP over the trailing "1" axis → [nf, nloc, nt_i, nt_j, ng]
            let prefix = format!("descriptor.embedding.{ti}.{tj}");
            let mlp = MlpSpec {
                param_prefix: &prefix,
                in_dim: 1,
                neuron: &cfg.neuron,
                activation,
                resnet_dt: cfg.resnet_dt,
            };
            let gg = embedding_mlp(g, &mlp, env_5d);

            // weighted_sum = sum_{i,j}(env_ij[..., None] * gg) over axes 2,3
            // → [nf, nloc, ng]
            let weighted = g.mul(env_5d, gg);
            let sum_shape = Shape::new(&[nf, nloc, ng], DType::F32);
            let summed = g.reduce(weighted, ReduceOp::Sum, vec![2, 3], false, sum_shape);

            let scale = scalar_const(g, 1.0 / (nt_i as f32 * nt_j as f32));
            let scaled = g.mul(summed, scale);

            acc = Some(match acc {
                None => scaled,
                Some(prev) => g.add(prev, scaled),
            });
        }
    }

    let res = acc.ok_or_else(|| anyhow::anyhow!("se_t: empty selection"))?;
    Ok(SeTDescriptor {
        descriptor: res,
        dim_out: ng,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn se_t_descriptor_builds() {
        let cfg = SeTConfig {
            rcut: 6.0,
            rcut_smth: 0.5,
            sel: vec![20, 40],
            neuron: vec![8, 16, 32],
            resnet_dt: false,
            activation_function: "tanh".into(),
            type_map: None,
        };
        let mut g = Graph::new("se_t_test");
        let nf = 1;
        let nloc = 4;
        let nnei = cfg.nnei();
        let env_mat = g.input("env_mat", Shape::new(&[nf, nloc, nnei, 4], DType::F32));
        let atype = g.input("atype", Shape::new(&[nf, nloc], DType::I32));
        let out = build_se_t_descriptor(&mut g, &cfg, env_mat, atype, nf, nloc, None)
            .expect("build");
        assert_eq!(out.dim_out, cfg.dim_out());
        assert!(g.len() > 10);
    }
}