Skip to main content

rscm_core/python/
mod.rs

1/// Python API to access core aspects of the RSCM library.
2///
3/// This module wraps the core functionality of the RSCM library for use in Python.
4/// This doesn't create an extension module that can be imported into python,
5/// this is done in the rscm crate, but uses the core pymodule that is exposed here.
6///
7/// ## Background
8///
9/// pyo3 provides an interface to be able to interact with the Python interpreter and PyObjects
10/// which are allocated on the Python heap.
11/// This allows us to write Rust code that can be called from Python,
12/// and also Python code that can be called from Rust.
13///
14/// Rust and Python differ in how they access memory.
15/// Rust has strong guarantees about reference checks with only one immutable reference
16/// or many mutable reference at a time.
17/// Python instead uses reference checks and the GIL to ensure that only one thread can
18/// access a PyObject at any time.
19///
20/// These differences can cause some difficulties when trying to use Rust-owned data in Python.
21/// One particular issue is that we can't return references (mutable or not) to Rust objects into
22/// Python because Python can't guarantee that Rust's memory safety rules are followed.
23/// Generally we return clones of the underlying Rust object which mean that changes to these
24/// objects in Python won't be reflected in Rust.
25/// Instead, we need to provide methods to update the Rust object from Python.
26///
27/// An example is the ['TimeseriesCollection'] which owns some ['Timeseries'] structs.
28/// In Rust we can return mutable references so we have more control about where data is modified.
29///
30/// ```rust
31/// use numpy::array;
32/// use rscm_core::spatial::ScalarRegion;
33/// use rscm_core::timeseries::Timeseries;
34/// use rscm_core::timeseries_collection::{TimeseriesCollection, VariableType};
35/// let mut collection = TimeseriesCollection::new();
36/// let timeseries = Timeseries::from_values(array![1.0, 2.0, 3.0], array![1.0, 2.0, 3.0]);
37///
38/// collection.add_timeseries("Surface Temperature".to_string(), timeseries, VariableType::Exogenous);
39/// // We can't access timeseries anymore because the collection has taken ownership
40/// // The line below generates a compiler error if uncommented
41/// // timeseries.at(0, ScalarRegion::Global)
42///
43/// {
44///     // Get mutable data which enables modifying the timeseries
45///     if let Some(data) = collection.get_data_mut("Surface Temperature") {
46///         if let Some(ts) = data.as_scalar_mut() {
47///             ts.set(0, ScalarRegion::Global, 2.0);
48///         }
49///     }
50/// }
51/// {
52///     // Get an immutable reference to the timeseries data
53///     let timeseries = collection.get_data("Surface Temperature")
54///         .and_then(|data| data.as_scalar())
55///         .unwrap();
56///     assert_eq!(timeseries.at(0, ScalarRegion::Global).unwrap(), 2.0);
57/// }
58/// ```
59///
60/// We don't have quite the same options for the Python interface which means that we need to
61/// add methods to enable modification.
62/// This results in a less flexible API, but is better than nothing!
63///
64/// ```py
65/// import numpy as np
66/// from rscm.core import TimeseriesCollection, Timeseries, InterpolationStrategy, VariableType
67///
68/// collection = TimeseriesCollection()
69/// timeseries = Timeseries(np.arange(2000.0, 2010.0), np.arange(2000.0, 2010.0), "units", InterpolationStrategy.Linear)
70/// collection.add_timeseries("Surface Temperature", timeseries, VariableType.Exogenous) # A clone is performed so that Rust can control ownership
71///
72/// # timeseries can still be used as it doesn't follow Rust's ownership model
73/// # Modifications to timeseries aren't reflected in the collection so we don't expose any methods to modify it
74///
75/// # Setting values in the collection requires a method call instead of being able to modify a timeseries
76/// collection.set_value("Surface Temperature", 0, 2.0)
77/// assert timeseries.at(0) != 2
78/// timeseries_2 = collection.get_timeseries_by_name("Surface Temperature")
79/// assert timeseries_2.at(0) == 2
80/// ```
81///
82/// There might be ways around this, but it would likely require deep integration of pyo3 throughout
83/// the library which results in code that is less readable.
84/// This tradeoff might be worth it if we need to expose a lot of functionality to Python,
85/// and it would also reduce the amount of boilerplate code.
86use crate::errors::RSCMError;
87use crate::schema::{AggregateDefinition, SchemaVariableDefinition, VariableSchema};
88use pyo3::exceptions::PyRuntimeError;
89use pyo3::prelude::*;
90use pyo3::{pymodule, wrap_pymodule, Bound, PyResult};
91
92mod component;
93mod example_component;
94mod model;
95pub mod spatial;
96pub mod state;
97pub mod timeseries;
98mod timeseries_collection;
99
100pub use component::PyRustComponent;
101pub use state::{
102    PyFourBoxSlice, PyFourBoxTimeseriesWindow, PyHemisphericSlice, PyHemisphericTimeseriesWindow,
103    PyTimeseriesWindow,
104};
105
106#[pymodule]
107pub fn core(m: &Bound<'_, PyModule>) -> PyResult<()> {
108    // Time types
109    m.add_class::<timeseries::PyTimeAxis>()?;
110    m.add_class::<timeseries::PyTimeseries>()?;
111    m.add_class::<timeseries::PyFourBoxTimeseries>()?;
112    m.add_class::<timeseries::PyHemisphericTimeseries>()?;
113    m.add_class::<timeseries::PyInterpolationStrategy>()?;
114
115    // State management
116    m.add_class::<timeseries_collection::PyTimeseriesCollection>()?;
117    m.add_class::<timeseries_collection::VariableType>()?;
118
119    // Component definitions
120    m.add_class::<component::PyPythonComponent>()?;
121    m.add_class::<component::RequirementDefinition>()?;
122    m.add_class::<component::RequirementType>()?;
123    m.add_class::<component::GridType>()?;
124
125    // Model orchestration
126    m.add_class::<model::PyModelBuilder>()?;
127    m.add_class::<model::PyModel>()?;
128
129    // Schema types (for variable definitions and aggregates)
130    m.add_class::<VariableSchema>()?;
131    m.add_class::<SchemaVariableDefinition>()?;
132    m.add_class::<AggregateDefinition>()?;
133
134    // Example component
135    m.add_class::<example_component::TestComponentBuilder>()?;
136
137    // Register submodules
138    m.add_wrapped(wrap_pymodule!(spatial::spatial))?;
139    m.add_wrapped(wrap_pymodule!(state::state))?;
140
141    Ok(())
142}
143
144impl From<RSCMError> for PyErr {
145    fn from(e: RSCMError) -> PyErr {
146        PyRuntimeError::new_err(e.to_string())
147    }
148}