Skip to content

Variable config #436

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jetuk wants to merge 14 commits into
base: main
Choose a base branch
from
Draft

Variable config #436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
0f4aacb
WIP on variable config.
jetuk Feb 11, 2025
067e07c
Merge branch 'main' into variable-config
jetuk Mar 17, 2025
a035d51
Merge branch 'main' into variable-config
jetuk Aug 12, 2025
0132617
Further tweaks and improvements.
jetuk Aug 13, 2025
f903e40
Fix pyo3 feature gating.
jetuk Aug 13, 2025
8905b48
Fix core feature gating.
jetuk Aug 13, 2025
9048d14
Merge branch 'main' into variable-config
jetuk Sep 19, 2025
85dbf72
Merge branch 'main' into variable-config
jetuk Oct 7, 2025
9153d76
Feature gate import.
jetuk Oct 7, 2025
fd08b52
Merge branch 'main' into variable-config
jetuk Dec 15, 2025
c39749e
Merge branch 'main' into variable-config
jetuk Feb 2, 2026
288fda2
Merge branch 'main' into variable-config
jetuk Feb 7, 2026
d29d3a8
Merge branch 'main' into variable-config
jetuk Feb 18, 2026
f08ffbd
Merge branch 'main' into variable-config
jetuk Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pywr-core/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ pub mod edge;
pub mod metric;
pub mod models;
pub mod network;
pub mod network_variable_config;
pub mod node;
pub mod parameters;
pub mod recorders;
Expand Down
133 changes: 89 additions & 44 deletions pywr-core/src/network.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::parameters::{
GeneralParameterIndex, GeneralParameterType, ParameterCalculationError, ParameterCollection,
ParameterCollectionConstCalculationError, ParameterCollectionError, ParameterCollectionSetupError,
ParameterCollectionSimpleCalculationError, ParameterIndex, ParameterName, ParameterStates, VariableConfig,
VariableParameterValues,
};
use crate::recorders::{
MetricSet, MetricSetIndex, MetricSetSaveError, MetricSetState, RecorderAggregationError, RecorderFinalResult,
Expand Down Expand Up @@ -2098,18 +2099,19 @@ impl Network {
Ok(edge_index)
}

/// Set the variable values on the parameter `parameter_index`.
/// Set the variable values on the parameter [`parameter_index`].
///
/// This will update the internal state of the parameter with the new values for all scenarios.
pub fn set_f64_parameter_variable_values(
&self,
parameter_index: ParameterIndex<f64>,
values: &[f64],
values_f64: &[f64],
values_u64: &[u64],
variable_config: &dyn VariableConfig,
state: &mut NetworkState,
) -> Result<(), NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_f64_variable() {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
// Iterate over all scenarios and set the variable values
for parameter_states in state.iter_parameter_states_mut() {
Expand All @@ -2120,7 +2122,7 @@ impl Network {
)?;

variable
.set_variables(values, variable_config, internal_state)
.set_variables(values_f64, values_u64, variable_config, internal_state)
.map_err(|source| NetworkError::VariableParameterError {
name: parameter.name().clone(),
source,
Expand All @@ -2137,33 +2139,35 @@ impl Network {
}
}

/// Set the variable values on the parameter `parameter_index` and scenario `scenario_index`.
/// Set the variable values on the parameter [`parameter_index`] and scenario [`scenario_index`].
///
/// Only the internal state of the parameter for the given scenario will be updated.
pub fn set_f64_parameter_variable_values_for_scenario(
&self,
parameter_index: ParameterIndex<f64>,
scenario_index: ScenarioIndex,
values: &[f64],
values_f64: &[f64],
values_u64: &[u64],
variable_config: &dyn VariableConfig,
state: &mut NetworkState,
) -> Result<(), NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_f64_variable() {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let internal_state = state
.parameter_states_mut(&scenario_index)
.get_mut_f64_state(parameter_index)
.ok_or(NetworkError::ParameterStateNotFound {
name: parameter.name().clone(),
})?;

variable
.set_variables(values, variable_config, internal_state)
.set_variables(values_f64, values_u64, variable_config, internal_state)
.map_err(|source| NetworkError::VariableParameterError {
name: parameter.name().clone(),
source,
})
})?;

Ok(())
}
None => Err(NetworkError::ParameterTypeNotVariable {
name: parameter.name().clone(),
Expand All @@ -2179,9 +2183,9 @@ impl Network {
parameter_index: ParameterIndex<f64>,
scenario_index: ScenarioIndex,
state: &NetworkState,
) -> Result<Option<Vec<f64>>, NetworkError> {
) -> Result<Option<VariableParameterValues>, NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_f64_variable() {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let internal_state = state
.parameter_states(&scenario_index)
Expand All @@ -2204,9 +2208,9 @@ impl Network {
&self,
parameter_index: ParameterIndex<f64>,
state: &NetworkState,
) -> Result<Vec<Option<Vec<f64>>>, NetworkError> {
) -> Result<Vec<Option<VariableParameterValues>>, NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_f64_variable() {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let values = state
.iter_parameter_states()
Expand All @@ -2231,29 +2235,30 @@ impl Network {
}
}

/// Set the variable values on the parameter `parameter_index`.
/// Set the variable values on the parameter [`parameter_index`].
///
/// This will update the internal state of the parameter with the new values for scenarios.
pub fn set_u32_parameter_variable_values(
pub fn set_u64_parameter_variable_values(
&self,
parameter_index: ParameterIndex<f64>,
values: &[u32],
parameter_index: ParameterIndex<u64>,
values_f64: &[f64],
values_u64: &[u64],
variable_config: &dyn VariableConfig,
state: &mut NetworkState,
) -> Result<(), NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_u32_variable() {
match self.parameters.get_u64(parameter_index) {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
// Iterate over all scenarios and set the variable values
for parameter_states in state.iter_parameter_states_mut() {
let internal_state = parameter_states.get_mut_f64_state(parameter_index).ok_or(
let internal_state = parameter_states.get_mut_u64_state(parameter_index).ok_or(
NetworkError::ParameterStateNotFound {
name: parameter.name().clone(),
},
)?;

variable
.set_variables(values, variable_config, internal_state)
.set_variables(values_f64, values_u64, variable_config, internal_state)
.map_err(|source| NetworkError::VariableParameterError {
name: parameter.name().clone(),
source,
Expand All @@ -2266,73 +2271,107 @@ impl Network {
name: parameter.name().clone(),
}),
},
None => Err(NetworkError::ParameterF64IndexNotFound(parameter_index)),
None => Err(NetworkError::ParameterU64IndexNotFound(parameter_index)),
}
}

/// Set the variable values on the parameter `parameter_index` and scenario `scenario_index`.
/// Set the variable values on the parameter [`parameter_index`] and scenario [`scenario_index`].
///
/// Only the internal state of the parameter for the given scenario will be updated.
pub fn set_u32_parameter_variable_values_for_scenario(
pub fn set_u64_parameter_variable_values_for_scenario(
&self,
parameter_index: ParameterIndex<f64>,
parameter_index: ParameterIndex<u64>,
scenario_index: ScenarioIndex,
values: &[u32],
values_f64: &[f64],
values_u64: &[u64],
variable_config: &dyn VariableConfig,
state: &mut NetworkState,
) -> Result<(), NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_u32_variable() {
match self.parameters.get_u64(parameter_index) {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let internal_state = state
.parameter_states_mut(&scenario_index)
.get_mut_f64_state(parameter_index)
.get_mut_u64_state(parameter_index)
.ok_or(NetworkError::ParameterStateNotFound {
name: parameter.name().clone(),
})?;

variable
.set_variables(values, variable_config, internal_state)
.set_variables(values_f64, values_u64, variable_config, internal_state)
.map_err(|source| NetworkError::VariableParameterError {
name: parameter.name().clone(),
source,
})
})?;

Ok(())
}
None => Err(NetworkError::ParameterTypeNotVariable {
name: parameter.name().clone(),
}),
},
None => Err(NetworkError::ParameterF64IndexNotFound(parameter_index)),
None => Err(NetworkError::ParameterU64IndexNotFound(parameter_index)),
}
}

/// Return a vector of the current values of active variable parameters.
pub fn get_u32_parameter_variable_values_for_scenario(
pub fn get_u64_parameter_variable_values_for_scenario(
&self,
parameter_index: ParameterIndex<f64>,
parameter_index: ParameterIndex<u64>,
scenario_index: ScenarioIndex,
state: &NetworkState,
) -> Result<Option<Vec<u32>>, NetworkError> {
match self.parameters.get_f64(parameter_index) {
Some(parameter) => match parameter.as_u32_variable() {
) -> Result<Option<VariableParameterValues>, NetworkError> {
match self.parameters.get_u64(parameter_index) {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let internal_state = state
.parameter_states(&scenario_index)
.get_f64_state(parameter_index)
.get_u64_state(parameter_index)
.ok_or(NetworkError::ParameterStateNotFound {
name: parameter.name().clone(),
})?;

Ok(variable.get_variables(internal_state))
}
None => Err(NetworkError::ParameterTypeNotVariable {
name: parameter.name().clone(),
}),
},
None => Err(NetworkError::ParameterF64IndexNotFound(parameter_index)),
None => Err(NetworkError::ParameterU64IndexNotFound(parameter_index)),
}
}
}

pub fn get_u64_parameter_variable_values(
&self,
parameter_index: ParameterIndex<u64>,
state: &NetworkState,
) -> Result<Vec<Option<VariableParameterValues>>, NetworkError> {
match self.parameters.get_u64(parameter_index) {
Some(parameter) => match parameter.as_variable() {
Some(variable) => {
let values = state
.iter_parameter_states()
.map(|parameter_states| {
let internal_state = parameter_states.get_u64_state(parameter_index).ok_or(
NetworkError::ParameterStateNotFound {
name: parameter.name().clone(),
},
)?;

Ok(variable.get_variables(internal_state))
})
.collect::<Result<_, NetworkError>>()?;

Ok(values)
}
None => Err(NetworkError::ParameterTypeNotVariable {
name: parameter.name().clone(),
}),
},
None => Err(NetworkError::ParameterU64IndexNotFound(parameter_index)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -2561,7 +2600,7 @@ mod tests {
let variable = ActivationFunction::Unit { min: 0.0, max: 10.0 };
let input_max_flow = parameters::ConstantParameter::new("my-constant".into(), 10.0);

assert!(input_max_flow.can_be_f64_variable());
assert!(input_max_flow.can_be_variable());

let input_max_flow_idx = model
.network_mut()
Expand All @@ -2585,7 +2624,7 @@ mod tests {
// Update the variable values
model
.network_mut()
.set_f64_parameter_variable_values(input_max_flow_idx, &[5.0], &variable, state.network_state_mut())
.set_f64_parameter_variable_values(input_max_flow_idx, &[5.0], &[], &variable, state.network_state_mut())
.unwrap();

// After update the variable value should match what was set
Expand All @@ -2594,6 +2633,12 @@ mod tests {
.get_f64_parameter_variable_values(input_max_flow_idx, state.network_state())
.unwrap();

assert_eq!(variable_values, vec![Some(vec![5.0])]);
assert_eq!(
variable_values,
vec![Some(VariableParameterValues {
f64: vec![5.0],
u64: vec![]
})]
);
}
}
Loading