deepmd 0.1.0

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

//! DP energy model graph: descriptor + fitting + total-energy reduction.
//!
//! Translated from `EnergyModel` / `make_model.forward_common` in
//! `deepmd/dpmodel/model/{ener_model,make_model}.py`.  Forces and
//! virials are reductive derivatives of `E`; they're produced at the
//! runtime layer by autodiff over the same graph (see `rlx-autodiff`)
//! and are not built explicitly here.

use anyhow::Result;
use rlx_ir::op::ReduceOp;
use rlx_ir::{DType, Graph, NodeId, Shape};

use crate::config::DPModelConfig;
use crate::descriptor::{build_se_a_descriptor, SeADescriptor};
use crate::fitting::{build_ener_fitting, EnerFitting};

/// Returned handle for the DP energy graph.
pub struct DPEnergyGraph {
    /// Total energy per frame, shape `[nf, 1]`.
    pub energy: NodeId,
    /// Per-atom energy, shape `[nf, nloc, 1]`.
    pub atom_energy: NodeId,
    /// Descriptor handle (descriptor + equivariant single-particle rep).
    pub descriptor: SeADescriptor,
    /// Input nodes declared on the graph, for caller convenience.
    pub inputs: DPEnergyInputs,
}

#[derive(Debug, Clone, Copy)]
pub struct DPEnergyInputs {
    /// Environment matrix `R`, shape `[nf, nloc, nnei, 4]`, f32.
    pub env_mat: NodeId,
    /// Atom-type indices, shape `[nf, nloc]`, integer (declared as i32).
    pub atype: NodeId,
}

/// Build the end-to-end DP energy graph for fixed `(nf, nloc)`.
///
/// Re-call with different `(nf, nloc)` to recompile for a different
/// system size — `rlx-flow`'s `BertFlow` follows the same convention.
pub fn build_dp_energy_graph(
    cfg: &DPModelConfig,
    nf: usize,
    nloc: usize,
) -> Result<(Graph, DPEnergyGraph)> {
    let mut g = Graph::new("dp_energy");

    let nnei = cfg.descriptor.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 descriptor = build_se_a_descriptor(
        &mut g,
        &cfg.descriptor,
        env_mat,
        atype,
        nf,
        nloc,
        crate::descriptor::SeAExtraInputs::default(),
    )?;

    assert_eq!(
        descriptor.dim_out, cfg.fitting_net.dim_descrpt,
        "descriptor.dim_out ({}) must match fitting_net.dim_descrpt ({})",
        descriptor.dim_out, cfg.fitting_net.dim_descrpt
    );

    let EnerFitting { atom_energy } = build_ener_fitting(
        &mut g,
        &cfg.fitting_net,
        nf,
        nloc,
        descriptor.descriptor,
        atype,
        None,
        None,
        None,
    )?;

    // E = Σ_i E_i over the nloc axis → [nf, dim_out].
    // We transpose to put nloc last and reduce the last axis; the rlx-cuda
    // and rlx-wgpu reduce kernels only support last-axis reductions, and
    // rlx-cpu / rlx-metal handle either form, so this is the portable shape.
    let dim_out = cfg.fitting_net.dim_out;
    let atom_e_t = {
        use rlx_ir::infer::GraphExt;
        g.transpose_(atom_energy, vec![0, 2, 1]) // [nf, dim_out, nloc]
    };
    let energy = g.reduce(
        atom_e_t,
        ReduceOp::Sum,
        vec![2],
        false,
        Shape::new(&[nf, dim_out], DType::F32),
    );

    Ok((
        g,
        DPEnergyGraph {
            energy,
            atom_energy,
            descriptor,
            inputs: DPEnergyInputs { env_mat, atype },
        },
    ))
}

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

    #[test]
    fn dp_energy_graph_builds_single_type() {
        // Minimal one-type DP model: rcut=6, sel=[46], one-atom-type system.
        let cfg = DPModelConfig::new(6.0, 0.5, vec![46]);
        let (g, out) = build_dp_energy_graph(&cfg, 1, 8).expect("build");
        // Smoke check: at least the descriptor MLP + final dense produce
        // a non-trivial number of nodes.
        assert!(g.len() > 10);
        assert_eq!(out.descriptor.dim_out, cfg.descriptor.dim_out());
    }

    #[test]
    fn dp_energy_graph_builds_multi_type() {
        // Water-like two-type system.
        let cfg = DPModelConfig::new(6.0, 0.5, vec![46, 92]);
        let (g, out) = build_dp_energy_graph(&cfg, 2, 192).expect("build");
        assert!(g.len() > 20);
        assert_eq!(out.descriptor.dim_out, cfg.descriptor.dim_out());
    }

    #[test]
    fn dp_energy_graph_type_two_side() {
        // Same shape as the multi-type test but with type_one_side=false
        // — exercises the ntypes² embedding-net branch and the
        // atype-based mask reduction.
        let mut cfg = DPModelConfig::new(6.0, 0.5, vec![46, 92]);
        cfg.descriptor.type_one_side = false;
        let (g, out) = build_dp_energy_graph(&cfg, 1, 16).expect("build");
        assert!(g.len() > 40);
        assert_eq!(out.descriptor.dim_out, cfg.descriptor.dim_out());
    }
}