deepmd 0.1.0

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

//! Configuration structs for DeePMD descriptors, fittings, and models.
//!
//! Mirrors the keyword arguments accepted by the Python constructors in
//! `deepmd/dpmodel/descriptor/se_e2_a.py`,
//! `deepmd/dpmodel/fitting/{ener_fitting,general_fitting}.py`, and
//! `deepmd/dpmodel/model/ener_model.py`.  Defaults match the Python
//! defaults so a JSON-serialized DeePMD model can deserialize directly.

use serde::{Deserialize, Serialize};

/// Configuration for the `se_e2_a` (DeepPot-SE) descriptor.
///
/// Field semantics follow `DescrptSeA.__init__` in
/// `deepmd/dpmodel/descriptor/se_e2_a.py`.  The only knobs this Rust port
/// supports today are the ones that affect the IR graph topology —
/// `compress`, `set_davg_zero`, `spin`, etc. that only touch host-side
/// statistics are deliberately omitted.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct SeAConfig {
    /// Cutoff radius `r_c`.
    pub rcut: f64,
    /// Inner smoothing radius `r_s` (≤ `rcut`).
    pub rcut_smth: f64,
    /// Per-type neighbor caps: `sel[t]` = max neighbors of type `t`.
    pub sel: Vec<usize>,
    /// Embedding-net hidden widths; output dim is `neuron[last]`.
    #[serde(default = "default_embedding_neuron")]
    pub neuron: Vec<usize>,
    /// `M₂` — number of axis-neuron columns of `Gᵢ_<`.
    #[serde(default = "default_axis_neuron")]
    pub axis_neuron: usize,
    /// Trainable `dt` resnet scaling in the embedding net.
    #[serde(default)]
    pub resnet_dt: bool,
    /// `true` ⇒ one embedding net per neighbor type (only the
    /// neighbor's type matters, not the centre's). This is the
    /// default in upstream DeePMD-kit. The current Rust port only
    /// implements this case; setting `false` will panic at graph
    /// build time.
    #[serde(default = "default_type_one_side")]
    pub type_one_side: bool,
    /// Activation function name. Only `"tanh"` is honored at present.
    #[serde(default = "default_activation")]
    pub activation_function: String,
    /// Optional `type_map` (informational, not used by the graph).
    #[serde(default)]
    pub type_map: Option<Vec<String>>,
}

fn default_embedding_neuron() -> Vec<usize> {
    vec![24, 48, 96]
}
fn default_axis_neuron() -> usize {
    8
}
fn default_type_one_side() -> bool {
    true
}
fn default_activation() -> String {
    "tanh".into()
}

impl SeAConfig {
    /// Number of atom types (= `sel.len()`).
    pub fn ntypes(&self) -> usize {
        self.sel.len()
    }

    /// Total neighbor slots per centre atom (= `sum(sel)`).
    pub fn nnei(&self) -> usize {
        self.sel.iter().sum()
    }

    /// Output dimension of the embedding net (`G`'s last axis).
    pub fn ng(&self) -> usize {
        *self
            .neuron
            .last()
            .expect("se_e2_a: neuron list must be non-empty")
    }

    /// Output dimension of the descriptor (`ng * axis_neuron`).
    pub fn dim_out(&self) -> usize {
        self.ng() * self.axis_neuron
    }

    /// Cumulative selection offsets: `[0, sel[0], sel[0]+sel[1], ...]`.
    pub fn sel_cumsum(&self) -> Vec<usize> {
        let mut acc = Vec::with_capacity(self.sel.len() + 1);
        let mut s = 0usize;
        acc.push(0);
        for &n in &self.sel {
            s += n;
            acc.push(s);
        }
        acc
    }
}

/// Configuration for the energy (and general invariant scalar) fitting net.
///
/// Mirrors the subset of `EnergyFittingNet` / `InvarFitting` /
/// `GeneralFitting` kwargs that this port supports.  Frame and atomic
/// parameter inputs (`numb_fparam`, `numb_aparam`) and case
/// embeddings are kept as config fields so the graph builder can
/// declare matching inputs, but the type-resolved (`mixed_types=false`)
/// path is not yet implemented and will panic at build time.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct EnerFittingConfig {
    /// Number of atom types.
    pub ntypes: usize,
    /// Input descriptor dimension (must equal `SeAConfig::dim_out`).
    pub dim_descrpt: usize,
    /// Output dimension of the fit; `1` for energy.
    #[serde(default = "default_dim_out")]
    pub dim_out: usize,
    /// Hidden layer widths of the fitting net.
    #[serde(default = "default_fitting_neuron")]
    pub neuron: Vec<usize>,
    /// `dt` resnet scaling on hidden layers.
    #[serde(default = "default_resnet_dt")]
    pub resnet_dt: bool,
    /// Number of frame parameters.
    #[serde(default)]
    pub numb_fparam: usize,
    /// Number of atomic parameters.
    #[serde(default)]
    pub numb_aparam: usize,
    /// Case-embedding dimension (DPA-2 / DPA-3 multitask).
    #[serde(default)]
    pub dim_case_embd: usize,
    /// Activation function name; only `"tanh"` is honored today.
    #[serde(default = "default_activation")]
    pub activation_function: String,
    /// Single shared fitting net across atom types.
    /// Only `true` is implemented in this Rust port.
    #[serde(default = "default_mixed_types_fitting")]
    pub mixed_types: bool,
}

fn default_dim_out() -> usize {
    1
}
fn default_fitting_neuron() -> Vec<usize> {
    vec![120, 120, 120]
}
fn default_resnet_dt() -> bool {
    true
}
fn default_mixed_types_fitting() -> bool {
    true
}

/// End-to-end DP energy model config: descriptor + fitting.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct DPModelConfig {
    pub descriptor: SeAConfig,
    pub fitting_net: EnerFittingConfig,
    /// Optional `type_map`. Informational.
    #[serde(default)]
    pub type_map: Option<Vec<String>>,
}

impl DPModelConfig {
    /// Build a default-shaped DP energy model from `(ntypes, sel)`.
    pub fn new(rcut: f64, rcut_smth: f64, sel: Vec<usize>) -> Self {
        let ntypes = sel.len();
        let descriptor = SeAConfig {
            rcut,
            rcut_smth,
            sel,
            neuron: default_embedding_neuron(),
            axis_neuron: default_axis_neuron(),
            resnet_dt: false,
            type_one_side: true,
            activation_function: default_activation(),
            type_map: None,
        };
        let fitting_net = EnerFittingConfig {
            ntypes,
            dim_descrpt: descriptor.dim_out(),
            dim_out: 1,
            neuron: default_fitting_neuron(),
            resnet_dt: true,
            numb_fparam: 0,
            numb_aparam: 0,
            dim_case_embd: 0,
            activation_function: default_activation(),
            mixed_types: true,
        };
        Self {
            descriptor,
            fitting_net,
            type_map: None,
        }
    }
}