cp2k-rs 0.1.3

//! IPC protocol types for the CP2K worker process.
//!
//! Serialization uses bincode for compact binary encoding.
//! All messages are length-prefixed: [u32 length][payload bytes].

use serde::{Deserialize, Serialize};

/// A command sent from the Python frontend to the worker.
#[derive(Debug, Serialize, Deserialize)]
pub enum Command {
    /// Initialize a new force environment.
    InitForceEnv { input: String, output: String },
    /// Initialise a new force environment and supply the starting geometry.
    ///
    /// The worker patches the input file by injecting `&CELL` and `&COORD`
    /// sections into `&SUBSYS` before passing the (temporary) file to CP2K.
    /// This avoids the requirement that the user's `.inp` already contains
    /// valid cell vectors and coordinates.
    ///
    /// Fields
    /// ------
    /// input          Path to the user's CP2K input file (.inp).
    /// output         Path for CP2K output.
    /// symbols        Element symbol for each atom, e.g. `["W", "Al", "O"]`.
    /// positions_angstrom Flat interleaved array [x0,y0,z0, x1,y1,z1, …] in Å.
    /// cell_angstrom      9-element row-major cell matrix in Å.
    ///                    cell_angstrom[i*3 + j] = j-th component of i-th lattice vector.
    /// periodic           Periodicity string passed to CP2K's PERIODIC keyword.
    ///                    Use `"XYZ"` for fully periodic systems (the common case),
    ///                    `"NONE"` for isolated molecules.  Any valid CP2K PERIODIC
    ///                    token is accepted.
    InitForceEnvWithGeometry {
        input: String,
        output: String,
        symbols: Vec<String>,
        positions_angstrom: Vec<f64>,
        cell_angstrom: Vec<f64>,
        periodic: String,
    },
    /// Calculate energy and forces for the current geometry.
    CalcEnergyForce,
    /// Calculate energy only.
    CalcEnergy,
    /// Query the number of atoms.
    GetNatom,
    /// Query the number of particles.
    GetNparticle,
    /// Retrieve particle positions as a flat f64 array.
    GetPositions,
    /// Retrieve forces as a flat f64 array.
    GetForces,
    /// Retrieve the potential energy.
    GetPotentialEnergy,
    /// Retrieve the 3x3 simulation cell (row-major flat f64).
    GetCell,
    /// Retrieve the 3x3 QMMM cell.
    GetQmmmCell,
    /// Set particle positions.
    SetPositions { data: Vec<f64> },
    /// Set particle velocities.
    SetVelocities { data: Vec<f64> },
    /// Set simulation cell (9 floats, row-major).
    SetCell { data: Vec<f64> },
    /// Get number of MOs in the active space.
    GetMoCount,
    // ----- extended interface -----
    #[cfg(feature = "extended")]
    IsQuickstep,
    #[cfg(feature = "extended")]
    GetStressTensor,
    #[cfg(feature = "extended")]
    GetVirialTensor,
    #[cfg(feature = "extended")]
    GetNmo { spin: i32 },
    #[cfg(feature = "extended")]
    GetEigenvalues { spin: i32 },
    #[cfg(feature = "extended")]
    GetOccupationNumbers { spin: i32 },
    #[cfg(feature = "extended")]
    GetHomoLumo { spin: i32 },
    #[cfg(feature = "extended")]
    GetMullikenCharges,
    #[cfg(feature = "extended")]
    GetHirshfeldCharges,
    #[cfg(feature = "extended")]
    GetDipoleMoment,
    #[cfg(feature = "extended")]
    GetScfInfo,
    #[cfg(feature = "extended")]
    GetEnergyComponents,
    #[cfg(feature = "extended")]
    GetNelectron,
    #[cfg(feature = "extended")]
    GetFermiEnergy,
    #[cfg(feature = "extended")]
    GetTotalSpin,
    /// Ask the worker to shut down cleanly.
    Shutdown,
    /// Get CP2K version string.
    GetVersion,
}

/// A request wrapper carrying a unique ID and a command.
#[derive(Debug, Serialize, Deserialize)]
pub struct Request {
    pub request_id: u64,
    pub command: Command,
}

/// Status of a response.
#[derive(Debug, Serialize, Deserialize)]
pub enum Status {
    Ok,
    Error(String),
}

/// The payload of a successful response.
#[derive(Debug, Serialize, Deserialize)]
pub enum Payload {
    /// No data to return (void operations).
    Empty,
    /// A single integer value.
    Int(i64),
    /// A single unsigned integer value.
    UInt(u64),
    /// A single float value.
    Float(f64),
    /// A boolean value.
    Bool(bool),
    /// A 1-D f64 array (flat).
    Array1(Vec<f64>),
    /// A 2-D f64 array (shape[0], shape[1], flat row-major data).
    Array2 {
        rows: usize,
        cols: usize,
        data: Vec<f64>,
    },
    /// A string value.
    String(String),
    /// Four values: two floats and two ints (HOMO/LUMO).
    HomoLumo {
        homo: f64,
        lumo: f64,
        homo_idx: i32,
        lumo_idx: i32,
    },
    /// Five floats (energy components).
    EnergyComponents {
        e_kin: f64,
        e_hartree: f64,
        e_xc: f64,
        e_core: f64,
        e_total: f64,
    },
    /// SCF info: steps, converged, energy_change.
    ScfInfo {
        nsteps: i32,
        converged: bool,
        energy_change: f64,
    },
}

/// A response from the worker back to the Python frontend.
#[derive(Debug, Serialize, Deserialize)]
pub struct Response {
    pub request_id: u64,
    pub status: Status,
    pub payload: Payload,
}

impl Response {
    pub fn ok(request_id: u64, payload: Payload) -> Self {
        Response {
            request_id,
            status: Status::Ok,
            payload,
        }
    }

    pub fn error(request_id: u64, msg: impl Into<String>) -> Self {
        Response {
            request_id,
            status: Status::Error(msg.into()),
            payload: Payload::Empty,
        }
    }
}