Skip to main content

cp2k_rs/
lib.rs

1//! Rust bindings for CP2K
2//!
3//! This crate provides Rust bindings for the CP2K quantum chemistry package,
4//! with additional Python bindings through PyO3.
5
6mod ffi;
7#[cfg(feature = "extended")]
8mod ffi_extended;
9/// Raw FFI bindings for extended interface (for direct shm writes in worker).
10#[cfg(feature = "extended")]
11pub mod ffi_extended_raw {
12    pub use crate::ffi_extended::*;
13}
14mod force_env;
15
16/// POSIX shared memory helpers for zero-copy IPC of large arrays.
17pub mod shm;
18
19/// IPC protocol shared between the Python frontend and the MPI worker binary.
20pub mod worker_protocol;
21
22/// Pure-Rust worker management (spawning, IPC, lifecycle). GIL-free.
23pub mod worker;
24
25#[cfg(feature = "mpi")]
26use mpi::{ffi::MPI_Comm_c2f, raw::AsRaw};
27
28use std::sync::atomic::{AtomicBool, Ordering};
29
30#[cfg(feature = "extended")]
31pub use force_env::GridInfo;
32pub use force_env::{CP2KError, CP2KResult, ForceEnv};
33
34#[cfg(feature = "python")]
35pub mod python;
36
37/// Tracks whether CP2K was initialized with or without MPI support.
38/// Used to ensure finalize() calls the matching teardown function.
39static INIT_WITH_MPI: AtomicBool = AtomicBool::new(true);
40
41/// Initialize CP2K with MPI support
42///
43/// This function automatically detects if MPI is already initialized (e.g., by mpi4py
44/// in Python or mpi crate in Rust) and uses the existing MPI environment if available.
45/// Otherwise, it initializes MPI as part of CP2K initialization.
46///
47/// This is the recommended initialization function when using MPI, as it avoids
48/// double MPI initialization conflicts.
49pub fn init() -> CP2KResult<()> {
50    unsafe {
51        let result = ffi::cp2k_init_with_mpi_check();
52        if result != 0 {
53            return Err(CP2KError::InitializationError(format!(
54                "CP2K initialization failed with code {}",
55                result
56            )));
57        }
58    }
59    INIT_WITH_MPI.store(true, Ordering::SeqCst);
60    Ok(())
61}
62
63/// Initialize CP2K without MPI support
64pub fn init_without_mpi() -> CP2KResult<()> {
65    unsafe {
66        ffi::cp2k_init_without_mpi();
67    }
68    INIT_WITH_MPI.store(false, Ordering::SeqCst);
69    Ok(())
70}
71
72/// Finalize CP2K with MPI support
73///
74/// This function finalizes CP2K without finalizing MPI if MPI was already
75/// initialized by the host application (e.g., mpi4py or mpi crate).
76/// The host application is responsible for finalizing MPI.
77///
78/// # Important
79/// All [`ForceEnv`](crate::ForceEnv) instances must be dropped before calling
80/// this function.  CP2K's `f77_interface` will abort with "invalid env_id" if
81/// a live force environment is still registered.  Explicitly drop your
82/// environments before finalization:
83///
84/// ```ignore
85/// drop(force_env);
86/// cp2k_rs::finalize()?;
87/// ```
88pub fn finalize() -> CP2KResult<()> {
89    unsafe {
90        let result = if INIT_WITH_MPI.load(Ordering::SeqCst) {
91            ffi::cp2k_finalize_with_mpi_check()
92        } else {
93            ffi::cp2k_finalize_without_mpi();
94            0
95        };
96        if result != 0 {
97            return Err(CP2KError::FinalizationError(format!(
98                "CP2K finalization failed with code {}",
99                result
100            )));
101        }
102    }
103    Ok(())
104}
105
106/// Finalize CP2K without MPI support
107pub fn finalize_without_mpi() -> CP2KResult<()> {
108    unsafe {
109        ffi::cp2k_finalize_without_mpi();
110    }
111    Ok(())
112}
113
114/// Get the CP2K version string
115pub fn get_version() -> CP2KResult<String> {
116    const MAX_VERSION_LEN: usize = 256;
117    let mut version_buf = vec![0u8; MAX_VERSION_LEN];
118
119    unsafe {
120        ffi::cp2k_get_version(version_buf.as_mut_ptr() as *mut i8, MAX_VERSION_LEN as i32);
121    }
122
123    // Find the null terminator
124    let end = version_buf
125        .iter()
126        .position(|&c| c == 0)
127        .unwrap_or(version_buf.len());
128    version_buf.truncate(end);
129
130    // Convert to String
131    Ok(String::from_utf8(version_buf)?)
132}
133
134/// Run a CP2K input file
135pub fn run_input(input_file: &str, output_file: &str) -> CP2KResult<()> {
136    let input_c = std::ffi::CString::new(input_file)?;
137    let output_c = std::ffi::CString::new(output_file)?;
138
139    unsafe {
140        ffi::cp2k_run_input(input_c.as_ptr(), output_c.as_ptr());
141    }
142
143    Ok(())
144}
145
146/// Run a CP2K input file with custom MPI communicator
147#[cfg(feature = "mpi")]
148pub fn run_input_comm(
149    input_file: &str,
150    output_file: &str,
151    comm: &mpi::topology::SimpleCommunicator,
152) -> CP2KResult<()> {
153    let input_c = std::ffi::CString::new(input_file)?;
154    let output_c = std::ffi::CString::new(output_file)?;
155    let raw_comm = comm.as_raw();
156    let fortran_comm = unsafe { MPI_Comm_c2f(raw_comm) };
157
158    unsafe {
159        ffi::cp2k_run_input_comm(input_c.as_ptr(), output_c.as_ptr(), fortran_comm);
160    }
161
162    Ok(())
163}