laddu_python/
lib.rs

1#![warn(clippy::perf, clippy::style)]
2#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
3use pyo3::prelude::*;
4use pyo3::types::PyDict;
5
6#[cfg_attr(coverage_nightly, coverage(off))]
7pub mod amplitudes;
8#[cfg_attr(coverage_nightly, coverage(off))]
9pub mod data;
10#[cfg_attr(coverage_nightly, coverage(off))]
11pub mod utils;
12
13#[cfg(feature = "mpi")]
14#[cfg_attr(coverage_nightly, coverage(off))]
15pub mod mpi {
16    use super::*;
17    /// Use the Message Passing Interface (MPI) to run on a distributed system
18    ///
19    /// Parameters
20    /// ----------
21    /// trigger: bool, default=True
22    ///     An optional parameter which allows MPI to only be used under some boolean
23    ///     condition.
24    ///
25    /// Notes
26    /// -----
27    /// You must have MPI installed for this to work, and you must call the program with
28    /// ``mpirun <executable>``, or bad things will happen.
29    ///
30    /// MPI runs an identical program on each process, but gives the program an ID called its
31    /// "rank". Only the results of methods on the root process (rank 0) should be
32    /// considered valid, as other processes only contain portions of each dataset. To ensure
33    /// you don't save or print data at other ranks, use the provided ``laddu.mpi.is_root()``
34    /// method to check if the process is the root process.
35    ///
36    /// Once MPI is enabled, it cannot be disabled. If MPI could be toggled (which it can't),
37    /// the other processes will still run, but they will be independent of the root process
38    /// and will no longer communicate with it. The root process stores no data, so it would
39    /// be difficult (and convoluted) to get the results which were already processed via
40    /// MPI.
41    ///
42    /// Additionally, MPI must be enabled at the beginning of a script, at least before any
43    /// other ``laddu`` functions are called. For this reason, it is suggested that you use the
44    /// context manager ``laddu.mpi.MPI`` to ensure the MPI backend is used properly.
45    ///
46    /// If ``laddu.mpi.use_mpi()`` is called multiple times, the subsequent calls will have no
47    /// effect.
48    ///
49    /// You **must** call ``laddu.mpi.finalize_mpi()`` before your program exits for MPI to terminate
50    /// smoothly.
51    ///
52    /// See Also
53    /// --------
54    /// laddu.mpi.MPI
55    /// laddu.mpi.using_mpi
56    /// laddu.mpi.is_root
57    /// laddu.mpi.get_rank
58    /// laddu.mpi.get_size
59    /// laddu.mpi.finalize_mpi
60    ///
61    #[pyfunction]
62    #[pyo3(signature = (*, trigger=true))]
63    pub fn use_mpi(trigger: bool) {
64        laddu_core::mpi::use_mpi(trigger);
65    }
66
67    /// Drop the MPI universe and finalize MPI at the end of a program
68    ///
69    /// This should only be called once and should be called at the end of all ``laddu``-related
70    /// function calls. This **must** be called at the end of any program which uses MPI.
71    ///
72    /// See Also
73    /// --------
74    /// laddu.mpi.use_mpi
75    ///
76    #[pyfunction]
77    pub fn finalize_mpi() {
78        laddu_core::mpi::finalize_mpi();
79    }
80
81    /// Check if MPI is enabled
82    ///
83    /// This can be combined with ``laddu.mpi.is_root()`` to ensure valid results are only
84    /// returned from the root rank process on the condition that MPI is enabled.
85    ///
86    /// See Also
87    /// --------
88    /// laddu.mpi.use_mpi
89    /// laddu.mpi.is_root
90    ///
91    #[pyfunction]
92    pub fn using_mpi() -> bool {
93        laddu_core::mpi::using_mpi()
94    }
95
96    /// Check if the current MPI process is the root process
97    ///
98    /// This can be combined with ``laddu.mpi.using_mpi()`` to ensure valid results are only
99    /// returned from the root rank process on the condition that MPI is enabled.
100    ///
101    /// See Also
102    /// --------
103    /// laddu.mpi.use_mpi
104    /// laddu.mpi.using_mpi
105    ///
106    #[pyfunction]
107    pub fn is_root() -> bool {
108        laddu_core::mpi::is_root()
109    }
110
111    /// Get the rank of the current MPI process
112    ///
113    /// Returns ``None`` if MPI is not enabled
114    ///
115    /// See Also
116    /// --------
117    /// laddu.mpi.use_mpi
118    ///
119    #[pyfunction]
120    pub fn get_rank() -> Option<i32> {
121        laddu_core::mpi::get_rank()
122    }
123
124    /// Get the total number of MPI processes (including the root process)
125    ///
126    /// Returns ``None`` if MPI is not enabled
127    ///
128    /// See Also
129    /// --------
130    /// laddu.mpi.use_mpi
131    ///
132    #[pyfunction]
133    pub fn get_size() -> Option<i32> {
134        laddu_core::mpi::get_size()
135    }
136}
137
138pub trait GetStrExtractObj {
139    fn get_extract<T>(&self, key: &str) -> PyResult<Option<T>>
140    where
141        T: for<'py> FromPyObject<'py>;
142}
143
144#[cfg_attr(coverage_nightly, coverage(off))]
145impl GetStrExtractObj for Bound<'_, PyDict> {
146    fn get_extract<T>(&self, key: &str) -> PyResult<Option<T>>
147    where
148        T: for<'py> FromPyObject<'py>,
149    {
150        self.get_item(key)?
151            .map(|value| value.extract::<T>())
152            .transpose()
153    }
154}
laddu_python/lib.rs

laddu_python/
lib.rs
