laddu_python/lib.rs
1#![warn(clippy::perf, clippy::style)]
2#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
3use laddu_core::ThreadPoolManager;
4use pyo3::{prelude::*, types::PyDict};
5
6/// Returns the number of CPUs (logical cores) available for use by ``laddu``.
7///
8#[pyfunction]
9pub fn available_parallelism() -> usize {
10 num_cpus::get()
11}
12
13/// Set the process-global default thread count for Python APIs that accept ``threads=``.
14///
15/// Parameters
16/// ----------
17/// n_threads : int, optional
18/// The default number of threads to use for omitted thread arguments and for ``threads=0``.
19/// Setting this to ``0`` or ``None`` resets the default to "all available CPUs".
20///
21/// Notes
22/// -----
23/// Explicit positive ``threads=`` arguments on individual calls override this default.
24///
25#[pyfunction]
26pub fn set_threads(n_threads: Option<usize>) {
27 ThreadPoolManager::set_global_thread_count(n_threads.unwrap_or(0));
28}
29
30/// Return the process-global default thread count used by omitted or zero-valued thread requests.
31///
32/// Returns ``0`` when the default is "all available CPUs".
33#[pyfunction]
34pub fn get_threads() -> usize {
35 ThreadPoolManager::global_thread_count().unwrap_or(0)
36}
37
38#[cfg_attr(coverage_nightly, coverage(off))]
39pub mod amplitudes;
40#[cfg_attr(coverage_nightly, coverage(off))]
41pub mod data;
42#[cfg_attr(coverage_nightly, coverage(off))]
43pub mod generation;
44#[cfg_attr(coverage_nightly, coverage(off))]
45pub mod math;
46#[cfg_attr(coverage_nightly, coverage(off))]
47pub mod quantum;
48#[cfg_attr(coverage_nightly, coverage(off))]
49pub mod variables;
50#[cfg_attr(coverage_nightly, coverage(off))]
51pub mod vectors;
52
53#[cfg_attr(coverage_nightly, coverage(off))]
54pub mod extensions;
55
56#[cfg_attr(coverage_nightly, coverage(off))]
57pub mod mpi {
58 #[cfg(not(feature = "mpi"))]
59 use pyo3::exceptions::PyModuleNotFoundError;
60
61 use super::*;
62 /// Check if ``laddu`` was compiled with MPI support (returns ``True`` if it was).
63 ///
64 /// Since ``laddu-mpi`` has the same namespace as ``laddu`` (they both are imported with
65 /// ``import laddu``), this method can be used to check if MPI capabilities are available
66 /// without actually running any MPI code. While functions in the ``laddu.mpi`` module will
67 /// raise an ``ModuleNotFoundError`` if MPI is not supported, its sometimes convenient to have
68 /// a simple boolean check rather than a try-catch block, and this method provides that.
69 ///
70 #[pyfunction]
71 pub fn is_mpi_available() -> bool {
72 #[cfg(feature = "mpi")]
73 return true;
74 #[cfg(not(feature = "mpi"))]
75 return false;
76 }
77 /// Use the Message Passing Interface (MPI) to run on a distributed system
78 ///
79 /// Parameters
80 /// ----------
81 /// trigger: bool, default=True
82 /// An optional parameter which allows MPI to only be used under some boolean
83 /// condition.
84 ///
85 /// Notes
86 /// -----
87 /// You must have MPI installed for this to work, and you must call the program with
88 /// ``mpirun <executable>``, or bad things will happen.
89 ///
90 /// MPI runs an identical program on each process, but gives the program an ID called its
91 /// "rank". Only the results of methods on the root process (rank 0) should be
92 /// considered valid, as other processes only contain portions of each dataset. To ensure
93 /// you don't save or print data at other ranks, use the provided ``laddu.mpi.is_root()``
94 /// method to check if the process is the root process.
95 ///
96 /// Once MPI is enabled, it cannot be disabled. If MPI could be toggled (which it can't),
97 /// the other processes will still run, but they will be independent of the root process
98 /// and will no longer communicate with it. The root process stores no data, so it would
99 /// be difficult (and convoluted) to get the results which were already processed via
100 /// MPI.
101 ///
102 /// Additionally, MPI must be enabled at the beginning of a script, at least before any
103 /// other ``laddu`` functions are called. For this reason, it is suggested that you use the
104 /// context manager ``laddu.mpi.MPI`` to ensure the MPI backend is used properly.
105 ///
106 /// If ``laddu.mpi.use_mpi()`` is called multiple times, the subsequent calls will have no
107 /// effect.
108 ///
109 /// You **must** call ``laddu.mpi.finalize_mpi()`` before your program exits for MPI to terminate
110 /// smoothly.
111 ///
112 /// See Also
113 /// --------
114 /// laddu.mpi.MPI
115 /// laddu.mpi.using_mpi
116 /// laddu.mpi.is_root
117 /// laddu.mpi.get_rank
118 /// laddu.mpi.get_size
119 /// laddu.mpi.finalize_mpi
120 ///
121 #[pyfunction]
122 #[allow(unused_variables)]
123 #[pyo3(signature = (*, trigger=true))]
124 pub fn use_mpi(trigger: bool) -> PyResult<()> {
125 #[cfg(feature = "mpi")]
126 {
127 laddu_core::mpi::use_mpi(trigger);
128 Ok(())
129 }
130 #[cfg(not(feature = "mpi"))]
131 return Err(PyModuleNotFoundError::new_err(
132 "`laddu` was not compiled with MPI support! Please use `laddu-mpi` instead.",
133 ));
134 }
135
136 /// Drop the MPI universe and finalize MPI at the end of a program
137 ///
138 /// This should only be called once and should be called at the end of all ``laddu``-related
139 /// function calls. This **must** be called at the end of any program which uses MPI.
140 ///
141 /// See Also
142 /// --------
143 /// laddu.mpi.use_mpi
144 ///
145 #[pyfunction]
146 pub fn finalize_mpi() -> PyResult<()> {
147 #[cfg(feature = "mpi")]
148 {
149 laddu_core::mpi::finalize_mpi();
150 Ok(())
151 }
152 #[cfg(not(feature = "mpi"))]
153 return Err(PyModuleNotFoundError::new_err(
154 "`laddu` was not compiled with MPI support! Please use `laddu-mpi` instead.",
155 ));
156 }
157
158 /// Check if MPI is enabled
159 ///
160 /// This can be combined with ``laddu.mpi.is_root()`` to ensure valid results are only
161 /// returned from the root rank process on the condition that MPI is enabled.
162 ///
163 /// See Also
164 /// --------
165 /// laddu.mpi.use_mpi
166 /// laddu.mpi.is_root
167 ///
168 #[pyfunction]
169 pub fn using_mpi() -> bool {
170 #[cfg(feature = "mpi")]
171 return laddu_core::mpi::using_mpi();
172 #[cfg(not(feature = "mpi"))]
173 return false;
174 }
175
176 /// Check if the current MPI process is the root process
177 ///
178 /// This can be combined with ``laddu.mpi.using_mpi()`` to ensure valid results are only
179 /// returned from the root rank process on the condition that MPI is enabled.
180 ///
181 /// See Also
182 /// --------
183 /// laddu.mpi.use_mpi
184 /// laddu.mpi.using_mpi
185 ///
186 #[pyfunction]
187 pub fn is_root() -> bool {
188 #[cfg(feature = "mpi")]
189 return laddu_core::mpi::is_root();
190 #[cfg(not(feature = "mpi"))]
191 return true;
192 }
193
194 /// Get the rank of the current MPI process
195 ///
196 /// Returns ``0`` if MPI is not enabled
197 ///
198 /// See Also
199 /// --------
200 /// laddu.mpi.use_mpi
201 ///
202 #[pyfunction]
203 pub fn get_rank() -> i32 {
204 #[cfg(feature = "mpi")]
205 return laddu_core::mpi::get_rank();
206 #[cfg(not(feature = "mpi"))]
207 return 0;
208 }
209
210 /// Get the total number of MPI processes (including the root process)
211 ///
212 /// Returns ``1`` if MPI is not enabled
213 ///
214 /// See Also
215 /// --------
216 /// laddu.mpi.use_mpi
217 ///
218 #[pyfunction]
219 pub fn get_size() -> i32 {
220 #[cfg(feature = "mpi")]
221 return laddu_core::mpi::get_size();
222 #[cfg(not(feature = "mpi"))]
223 return 1;
224 }
225}
226
227pub trait GetStrExtractObj {
228 fn get_extract<T>(&self, key: &str) -> PyResult<Option<T>>
229 where
230 T: for<'a, 'py> FromPyObject<'a, 'py, Error = PyErr>;
231}
232
233#[cfg_attr(coverage_nightly, coverage(off))]
234impl GetStrExtractObj for Bound<'_, PyDict> {
235 fn get_extract<T>(&self, key: &str) -> PyResult<Option<T>>
236 where
237 T: for<'a, 'py> FromPyObject<'a, 'py, Error = PyErr>,
238 {
239 self.get_item(key)?
240 .map(|value| value.extract::<T>())
241 .transpose()
242 }
243}