Skip to content

Commit

Permalink
feat: Implement Parameter<usize> for PyParameter. (#137)
Browse files Browse the repository at this point in the history 
This supports using PyParameter as an index parameter with updated
schema to define the return type of the Python method. Some minor related
docs improvements.
  • Loading branch information
@jetuk
jetuk authored Mar 12, 2024
1 parent fb89773 commit 2133bc0
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 86 deletions.
135 changes: 81 additions & 54 deletions pywr-core/src/parameters/py.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,18 +66,8 @@ impl PyParameter {

Ok(index_values.into_py_dict(py))
}
}

impl Parameter<f64> for PyParameter {
fn meta(&self) -> &ParameterMeta {
&self.meta
}

fn setup(
&self,
_timesteps: &[Timestep],
_scenario_index: &ScenarioIndex,
) -> Result<Option<Box<dyn ParameterState>>, PywrError> {
fn setup(&self) -> Result<Option<Box<dyn ParameterState>>, PywrError> {
pyo3::prepare_freethreaded_python();

let user_obj: PyObject = Python::with_gil(|py| -> PyResult<PyObject> {
Expand All @@ -92,26 +82,20 @@ impl Parameter<f64> for PyParameter {
Ok(Some(internal.into_boxed_any()))
}

// fn before(&self, internal_state: &mut Option<Box<dyn ParameterState>>) -> Result<(), PywrError> {
// let internal = downcast_internal_state::<Internal>(internal_state);
//
// Python::with_gil(|py| internal.user_obj.call_method0(py, "before"))
// .map_err(|e| PywrError::PythonError(e.to_string()))?;
//
// Ok(())
// }

fn compute(
fn compute<T>(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
model: &Network,
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<f64, PywrError> {
) -> Result<T, PywrError>
where
T: for<'a> FromPyObject<'a>,
{
let internal = downcast_internal_state_mut::<Internal>(internal_state);

let value: f64 = Python::with_gil(|py| {
let value: T = Python::with_gil(|py| {
let date = timestep.date.into_py(py);

let si = scenario_index.index.into_py(py);
Expand Down Expand Up @@ -160,7 +144,8 @@ impl Parameter<f64> for PyParameter {
}
}

impl Parameter<MultiValue> for PyParameter {
impl Parameter<f64> for PyParameter {

fn meta(&self) -> &ParameterMeta {
&self.meta
}
Expand All @@ -170,18 +155,80 @@ impl Parameter<MultiValue> for PyParameter {
_timesteps: &[Timestep],
_scenario_index: &ScenarioIndex,
) -> Result<Option<Box<dyn ParameterState>>, PywrError> {
pyo3::prepare_freethreaded_python();
self.setup()
}

let user_obj: PyObject = Python::with_gil(|py| -> PyResult<PyObject> {
let args = self.args.as_ref(py);
let kwargs = self.kwargs.as_ref(py);
self.object.call(py, args, Some(kwargs))
})
.unwrap();
fn compute(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
model: &Network,
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<f64, PywrError> {
self.compute(timestep, scenario_index, model, state, internal_state)
}

let internal = Internal { user_obj };
fn after(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
model: &Network,
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<(), PywrError> {
self.after(timestep, scenario_index, model, state, internal_state)
}
}

Ok(Some(internal.into_boxed_any()))
impl Parameter<usize> for PyParameter {

fn meta(&self) -> &ParameterMeta {
&self.meta
}

fn setup(
&self,
_timesteps: &[Timestep],
_scenario_index: &ScenarioIndex,
) -> Result<Option<Box<dyn ParameterState>>, PywrError> {
self.setup()
}

fn compute(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
model: &Network,
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<usize, PywrError> {
self.compute(timestep, scenario_index, model, state, internal_state)
}

fn after(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
model: &Network,
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<(), PywrError> {
self.after(timestep, scenario_index, model, state, internal_state)
}
}

impl Parameter<MultiValue> for PyParameter {
fn meta(&self) -> &ParameterMeta {
&self.meta
}

fn setup(
&self,
_timesteps: &[Timestep],
_scenario_index: &ScenarioIndex,
) -> Result<Option<Box<dyn ParameterState>>, PywrError> {
self.setup()
}

// fn before(&self, internal_state: &mut Option<Box<dyn ParameterState>>) -> Result<(), PywrError> {
Expand Down Expand Up @@ -257,27 +304,7 @@ impl Parameter<MultiValue> for PyParameter {
state: &State,
internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<(), PywrError> {
let internal = downcast_internal_state_mut::<Internal>(internal_state);

Python::with_gil(|py| {
// Only do this if the object has an "after" method defined.
if internal.user_obj.getattr(py, "after").is_ok() {
let date = timestep.date.into_py(py);

let si = scenario_index.index.into_py(py);

let metric_dict = self.get_metrics_dict(model, state, py)?;
let index_dict = self.get_indices_dict(state, py)?;

let args = PyTuple::new(py, [date.as_ref(py), si.as_ref(py), metric_dict, index_dict]);

internal.user_obj.call_method1(py, "after", args)?;
}
Ok(())
})
.map_err(|e: PyErr| PywrError::PythonError(e.to_string()))?;

Ok(())
self.after(timestep, scenario_index, model, state, internal_state)
}
}

Expand Down
2 changes: 1 addition & 1 deletion pywr-schema/src/parameters/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ pub use super::parameters::profiles::{
DailyProfileParameter, MonthlyProfileParameter, RadialBasisFunction, RbfProfileParameter,
RbfProfileVariableSettings, UniformDrawdownProfileParameter, WeeklyProfileParameter,
};
pub use super::parameters::python::PythonParameter;
pub use super::parameters::python::{PythonModule, PythonParameter, PythonReturnType};
pub use super::parameters::tables::TablesArrayParameter;
pub use super::parameters::thresholds::ParameterThresholdParameter;
use crate::error::{ConversionError, SchemaError};
Expand Down
92 changes: 61 additions & 31 deletions pywr-schema/src/parameters/python.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,22 @@ pub enum PythonModule {
Path(PathBuf),
}

/// The expected return type of the Python parameter.
#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, Default)]
#[serde(rename_all = "lowercase")]
pub enum PythonReturnType {
#[default]
Float,
Int,
Dict,
}

/// A Parameter that uses a Python object for its calculations.
///
/// This struct defines a schema for loading a [`crate::parameters::PyParameter`] from external
/// This struct defines a schema for loading a [`PyParameter`] from external
/// sources. The user provides the name of an object in the given module. Typically, this object will be
/// a class the user has written. For more information on the expected format and signature of
/// this object please refer to the [`crate::parameters::PyParameter`] documentation. The object
/// this object please refer to the [`PyParameter`] documentation. The object
/// is initialised with user provided positional and/or keyword arguments that can be provided
/// here.
///
Expand Down Expand Up @@ -67,10 +77,10 @@ pub struct PythonParameter {
pub module: PythonModule,
/// The name of Python object from the module to use.
pub object: String,
/// Is this a multi-valued parameter or not. If true then the calculation method should
/// return a dictionary with string keys and either floats or ints as values.
/// The return type of the Python calculation. This is used to convert the Python return value
/// to the appropriate type for the Parameter.
#[serde(default)]
pub multi: bool,
pub return_type: PythonReturnType,
/// Position arguments to pass to the object during setup.
pub args: Vec<serde_json::Value>,
/// Keyword arguments to pass to the object during setup.
Expand Down Expand Up @@ -193,10 +203,11 @@ impl PythonParameter {
};

let p = PyParameter::new(&self.meta.name, object, args, kwargs, &metrics, &indices);
let pt = if self.multi {
ParameterType::Multi(network.add_multi_value_parameter(Box::new(p))?)
} else {
ParameterType::Parameter(network.add_parameter(Box::new(p))?)

let pt = match self.return_type {
PythonReturnType::Float => ParameterType::Parameter(network.add_parameter(Box::new(p))?),
PythonReturnType::Int => ParameterType::Index(network.add_index_parameter(Box::new(p))?),
PythonReturnType::Dict => ParameterType::Multi(network.add_multi_value_parameter(Box::new(p))?),
};

Ok(pt)
Expand All @@ -212,42 +223,59 @@ mod tests {
use pywr_core::network::Network;
use pywr_core::test_utils::default_time_domain;
use serde_json::json;
use std::fs::File;
use std::io::Write;
use tempfile::tempdir;
use std::path::PathBuf;

#[test]
fn test_python_parameter() {
let dir = tempdir().unwrap();

let file_path = dir.path().join("my_parameter.py");
fn test_python_float_parameter() {
let mut py_fn = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
py_fn.push("src/test_models/test_parameters.py");

let data = json!(
{
"name": "my-custom-calculation",
"name": "my-float-parameter",
"type": "Python",
"path": file_path,
"object": "MyParameter",
"path": py_fn,
"object": "FloatParameter",
"args": [0, ],
"kwargs": {},
}
)
.to_string();

let mut file = File::create(file_path).unwrap();
write!(
file,
r#"
class MyParameter:
def __init__(self, count, *args, **kwargs):
self.count = 0
// Init Python
pyo3::prepare_freethreaded_python();
// Load the schema ...
let param: PythonParameter = serde_json::from_str(data.as_str()).unwrap();
// ... add it to an empty network
// this should trigger loading the module and extracting the class
let domain: ModelDomain = default_time_domain().into();
let schema = PywrNetwork::default();
let mut network = Network::default();
let tables = LoadedTableCollection::from_schema(None, None).unwrap();
param
.add_to_model(&mut network, &schema, &domain, &tables, None, &[])
.unwrap();

assert!(network.get_parameter_by_name("my-float-parameter").is_ok());
}

#[test]
fn test_python_int_parameter() {
let mut py_fn = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
py_fn.push("src/test_models/test_parameters.py");

def calc(self, ts, si, p_values):
self.count += si
return float(self.count + ts.day)
"#
let data = json!(
{
"name": "my-int-parameter",
"type": "Python",
"path": py_fn,
"return_type": "int",
"object": "FloatParameter",
"args": [0, ],
"kwargs": {},
}
)
.unwrap();
.to_string();

// Init Python
pyo3::prepare_freethreaded_python();
Expand All @@ -262,5 +290,7 @@ class MyParameter:
param
.add_to_model(&mut network, &schema, &domain, &tables, None, &[])
.unwrap();

assert!(network.get_index_parameter_by_name("my-int-parameter").is_ok());
}
}
20 changes: 20 additions & 0 deletions pywr-schema/src/test_models/test_parameters.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
class FloatParameter:
"""A simple float parameter"""

def __init__(self, count, *args, **kwargs):
self.count = 0

def calc(self, ts, si, p_values) -> float:
self.count += si
return float(self.count + ts.day)


class IntParameter:
"""A simple int parameter"""

def __init__(self, count, *args, **kwargs):
self.count = 0

def calc(self, ts, si, p_values) -> int:
self.count += si
return self.count + ts.day

0 comments on commit 2133bc0

Please sign in to comment.