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

// Verify that DPA-2's repinit (DPA-1 with attn_layer=0) matches Python.

use std::collections::HashMap;
use std::path::Path;

use deepmd::descriptor_dpa1::{build_dpa1_descriptor, DPA1Config, DPA1Inputs};
use deepmd::descriptor_se_t_tebd::TebdInputMode;
use rlx_ir::{DType, Graph, Shape};
use rlx_runtime::{Device, Session};
use serde::Deserialize;

#[derive(Deserialize)]
struct Tensor { name: String, #[allow(dead_code)] shape: Vec<usize>, data: Vec<f32> }

#[derive(Deserialize)]
struct Config {
    rcut_repinit: f64, rcut_smth_repinit: f64,
    nsel: usize, ntypes: usize, tebd_dim: usize,
    repinit_neuron: Vec<usize>, repinit_axis_neuron: usize,
    nf: usize, nloc: usize, nall: usize,
}

#[derive(Deserialize)]
struct Fixture {
    config: Config, params: HashMap<String, Tensor>,
    inputs: Vec<Tensor>, outputs: Vec<Tensor>,
}

fn find_named<'a>(list: &'a [Tensor], name: &str) -> &'a Tensor {
    list.iter().find(|t| t.name == name).expect("missing")
}

#[test]
fn dpa2_repinit_matches_python() {
    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent().unwrap().parent().unwrap()
        .join("tests/fixtures/dpa2_small.json");
    let fx: Fixture = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();

    let nf = fx.config.nf;
    let nloc = fx.config.nloc;
    let nsel = fx.config.nsel;
    let _ = fx.config.nall;

    let cfg = DPA1Config {
        rcut: fx.config.rcut_repinit, rcut_smth: fx.config.rcut_smth_repinit,
        sel: vec![nsel], ntypes: fx.config.ntypes,
        neuron: fx.config.repinit_neuron.clone(),
        axis_neuron: fx.config.repinit_axis_neuron,
        tebd_dim: fx.config.tebd_dim,
        tebd_input_mode: TebdInputMode::Concat,
        resnet_dt: false,
        attn: 16, attn_layer: 0, attn_dotr: false,
        num_heads: 1, normalize: false,
        scaling_factor: 1.0, ln_eps: 1e-5,
        smooth: true, type_one_side: true,
        activation_function: "tanh".into(),
        concat_output_tebd: false,
    };

    let mut g = Graph::new("repinit_debug");
    let env_mat = g.input("env_mat", Shape::new(&[nf, nloc, nsel, 4], DType::F32));
    let atype = g.input("atype", Shape::new(&[nf, nloc], DType::I32));
    let nei_atype = g.input("nei_atype", Shape::new(&[nf, nloc, nsel], DType::I32));
    let nlist_mask = g.input("nlist_mask", Shape::new(&[nf, nloc, nsel], DType::F32));
    let sw = g.input("sw", Shape::new(&[nf, nloc, nsel], DType::F32));
    let t_table = g.input("type_embed", Shape::new(&[cfg.ntypes + 1, cfg.tebd_dim], DType::F32));
    let out = build_dpa1_descriptor(
        &mut g, &cfg,
        DPA1Inputs {
            env_mat_raw: env_mat, atype_loc: atype, nei_atype,
            nlist_mask, sw, type_embedding: t_table,
        },
        nf, nloc,
    ).expect("build");
    // Also expose the normalized env-mat by re-doing the davg/dstd
    // gather here so we can compare against Python directly.
    use rlx_ir::infer::GraphExt;
    let davg = g.param("descriptor.davg", Shape::new(&[cfg.ntypes, nsel, 4], DType::F32));
    let dstd = g.param("descriptor.dstd", Shape::new(&[cfg.ntypes, nsel, 4], DType::F32));
    let davg_g = g.gather_(davg, atype, 0);
    let dstd_g = g.gather_(dstd, atype, 0);
    let mut rr = g.sub(env_mat, davg_g);
    rr = g.div(rr, dstd_g);
    g.set_outputs(vec![out.descriptor, rr]);

    let session = Session::new(Device::Cpu);
    let mut compiled = session.compile(g);
    for (name, t) in &fx.params {
        // Only bind keys that look like repinit (no "repformers." prefix, no "g1_shape_tranform").
        if name.starts_with("repformers.") || name == "g1_shape_tranform.w" {
            continue;
        }
        let bytes: Vec<u8> = t.data.iter().flat_map(|v| v.to_le_bytes()).collect();
        compiled.set_param_typed(name, &bytes, DType::F32);
    }
    let em = find_named(&fx.inputs, "env_mat_raw_repinit");
    let atype_ext = find_named(&fx.inputs, "atype_ext");
    let nei = find_named(&fx.inputs, "nei_atype");
    let nm = find_named(&fx.inputs, "nlist_mask_repinit");
    let sw_in = find_named(&fx.inputs, "sw_repinit");
    let tab = find_named(&fx.inputs, "type_embedding");
    let atype_loc_f32: Vec<f32> = atype_ext.data.iter().take(nf * nloc).copied().collect();

    let outputs = compiled.run(&[
        ("env_mat", em.data.as_slice()),
        ("atype", atype_loc_f32.as_slice()),
        ("nei_atype", nei.data.as_slice()),
        ("nlist_mask", nm.data.as_slice()),
        ("sw", sw_in.data.as_slice()),
        ("type_embed", tab.data.as_slice()),
    ]);

    // Check normed env first.
    let expected_env = find_named(&fx.outputs, "repinit_normed_env");
    let mut env_max_abs = 0f32;
    for (a, e) in outputs[1].iter().zip(expected_env.data.iter()) {
        env_max_abs = env_max_abs.max((a - e).abs());
    }
    eprintln!("normed_env: max_abs={env_max_abs:.3e}");

    let expected = find_named(&fx.outputs, "repinit_g1");
    let mut max_abs = 0f32;
    for (i, (a, e)) in outputs[0].iter().zip(expected.data.iter()).enumerate() {
        let d = (a - e).abs();
        if d > max_abs { max_abs = d; }
        if d > 1e-4 && i < 8 {
            eprintln!("  [{i}]: actual={a:.6e} expected={e:.6e} |Δ|={d:.3e}");
        }
    }
    eprintln!("repinit_g1: max_abs={max_abs:.3e}");
}