Skip to content

Commit

Permalink
feat: Initial implementation of virtual storage costs. (#71)
Browse files Browse the repository at this point in the history 
Virtual storage nodes now add a reference to their linked
flow nodes. Those nodes now include an function that aggregates
over the local cost and the costs of any linked virtual storage
nodes.

There's currently no method to change the default aggregation
function from "max" via the schema.
  • Loading branch information
@jetuk
jetuk authored Nov 28, 2023
1 parent d3dca15 commit 5fcef40
Show file tree
Hide file tree
Showing 7 changed files with 193 additions and 29 deletions.
2 changes: 2 additions & 0 deletions pywr-core/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,8 @@ pub enum PywrError {
FlowConstraintsUndefined,
#[error("storage constraints are undefined for this node")]
StorageConstraintsUndefined,
#[error("can not add virtual storage node to a storage node")]
NoVirtualStorageOnStorageNode,
#[error("timestep index out of range")]
TimestepIndexOutOfRange,
#[error("solver not initialised")]
Expand Down
15 changes: 12 additions & 3 deletions pywr-core/src/model.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1352,12 +1352,13 @@ impl Model {
min_volume: ConstraintValue,
max_volume: ConstraintValue,
reset: VirtualStorageReset,
cost: ConstraintValue,
) -> Result<VirtualStorageIndex, PywrError> {
if let Ok(_agg_node) = self.get_virtual_storage_node_by_name(name, sub_name) {
return Err(PywrError::NodeNameAlreadyExists(name.to_string()));
}

let node_index = self.virtual_storage_nodes.push_new(
let vs_node_index = self.virtual_storage_nodes.push_new(
name,
sub_name,
nodes,
Expand All @@ -1366,12 +1367,20 @@ impl Model {
min_volume,
max_volume,
reset,
cost,
);

// Link the virtual storage node to the nodes it is including
for node_idx in nodes {
let node = self.nodes.get_mut(node_idx)?;
node.add_virtual_storage(vs_node_index)?;
}

// Add to the resolve order.
self.resolve_order.push(ComponentType::VirtualStorageNode(node_index));
self.resolve_order
.push(ComponentType::VirtualStorageNode(vs_node_index));

Ok(node_index)
Ok(vs_node_index)
}

/// Add a `parameters::Parameter` to the model
Expand Down
137 changes: 113 additions & 24 deletions pywr-core/src/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use crate::metric::Metric;
use crate::model::Model;
use crate::state::{NodeState, State};
use crate::timestep::Timestep;
use crate::virtual_storage::VirtualStorageIndex;
use crate::PywrError;
use std::ops::{Deref, DerefMut};

Expand Down Expand Up @@ -123,6 +124,13 @@ impl From<Metric> for ConstraintValue {
}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CostAggFunc {
Sum,
Max,
Min,
}

impl Node {
/// Create a new input node
pub fn new_input(node_index: &NodeIndex, name: &str, sub_name: Option<&str>) -> Self {
Expand Down Expand Up @@ -286,6 +294,24 @@ impl Node {
}
}

pub fn add_virtual_storage(&mut self, virtual_storage_index: VirtualStorageIndex) -> Result<(), PywrError> {
match self {
Self::Input(n) => {
n.cost.virtual_storage_nodes.push(virtual_storage_index);
Ok(())
}
Self::Output(n) => {
n.cost.virtual_storage_nodes.push(virtual_storage_index);
Ok(())
}
Self::Link(n) => {
n.cost.virtual_storage_nodes.push(virtual_storage_index);
Ok(())
}
Self::Storage(_) => Err(PywrError::NoVirtualStorageOnStorageNode),
}
}

// /// Return a reference to a node's flow constraints if they exist.
// fn flow_constraints(&self) -> Option<&FlowConstraints> {
// match self {
Expand Down Expand Up @@ -502,6 +528,17 @@ impl Node {
}
}

pub fn set_cost_agg_func(&mut self, agg_func: CostAggFunc) -> Result<(), PywrError> {
match self {
Self::Input(n) => n.set_cost_agg_func(agg_func),
Self::Link(n) => n.set_cost_agg_func(agg_func),
Self::Output(n) => n.set_cost_agg_func(agg_func),
Self::Storage(_) => return Err(PywrError::NoVirtualStorageOnStorageNode),
};

Ok(())
}

pub fn get_outgoing_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
match self {
Self::Input(n) => n.get_cost(model, state),
Expand Down Expand Up @@ -628,10 +665,65 @@ impl StorageConstraints {
}
}

/// Generic cost data for a node.
#[derive(Debug, PartialEq)]
struct NodeCost {
local: ConstraintValue,
virtual_storage_nodes: Vec<VirtualStorageIndex>,
agg_func: CostAggFunc,
}

impl Default for NodeCost {
fn default() -> Self {
Self {
local: ConstraintValue::None,
virtual_storage_nodes: Vec::new(),
agg_func: CostAggFunc::Max,
}
}
}

impl NodeCost {
fn get_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
let local_cost = match &self.local {
ConstraintValue::None => Ok(0.0),
ConstraintValue::Scalar(v) => Ok(*v),
ConstraintValue::Metric(m) => m.get_value(model, state),
}?;

let vs_costs: Vec<f64> = self
.virtual_storage_nodes
.iter()
.map(|idx| {
let vs = model.get_virtual_storage_node(idx)?;
vs.get_cost(model, state)
})
.collect::<Result<_, _>>()?;

let cost = match self.agg_func {
CostAggFunc::Sum => local_cost + vs_costs.iter().sum::<f64>(),
CostAggFunc::Max => local_cost.max(
vs_costs
.into_iter()
.max_by(|a, b| a.total_cmp(b))
.unwrap_or(f64::NEG_INFINITY),
),
CostAggFunc::Min => local_cost.min(
vs_costs
.into_iter()
.min_by(|a, b| a.total_cmp(b))
.unwrap_or(f64::INFINITY),
),
};

Ok(cost)
}
}

#[derive(Debug, PartialEq)]
pub struct InputNode {
pub meta: NodeMeta<NodeIndex>,
pub cost: ConstraintValue,
cost: NodeCost,
pub flow_constraints: FlowConstraints,
pub outgoing_edges: Vec<EdgeIndex>,
}
Expand All @@ -640,20 +732,19 @@ impl InputNode {
fn new(index: &NodeIndex, name: &str, sub_name: Option<&str>) -> Self {
Self {
meta: NodeMeta::new(index, name, sub_name),
cost: ConstraintValue::None,
cost: NodeCost::default(),
flow_constraints: FlowConstraints::new(),
outgoing_edges: Vec::new(),
}
}
fn set_cost(&mut self, value: ConstraintValue) {
self.cost = value
self.cost.local = value
}
fn set_cost_agg_func(&mut self, agg_func: CostAggFunc) {
self.cost.agg_func = agg_func
}
fn get_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
match &self.cost {
ConstraintValue::None => Ok(0.0),
ConstraintValue::Scalar(v) => Ok(*v),
ConstraintValue::Metric(m) => m.get_value(model, state),
}
self.cost.get_cost(model, state)
}
fn set_min_flow(&mut self, value: ConstraintValue) {
self.flow_constraints.min_flow = value;
Expand All @@ -678,7 +769,7 @@ impl InputNode {
#[derive(Debug, PartialEq)]
pub struct OutputNode {
pub meta: NodeMeta<NodeIndex>,
pub cost: ConstraintValue,
cost: NodeCost,
pub flow_constraints: FlowConstraints,
pub incoming_edges: Vec<EdgeIndex>,
}
Expand All @@ -687,20 +778,19 @@ impl OutputNode {
fn new(index: &NodeIndex, name: &str, sub_name: Option<&str>) -> Self {
Self {
meta: NodeMeta::new(index, name, sub_name),
cost: ConstraintValue::None,
cost: NodeCost::default(),
flow_constraints: FlowConstraints::new(),
incoming_edges: Vec::new(),
}
}
fn set_cost(&mut self, value: ConstraintValue) {
self.cost = value
self.cost.local = value
}
fn get_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
match &self.cost {
ConstraintValue::None => Ok(0.0),
ConstraintValue::Scalar(v) => Ok(*v),
ConstraintValue::Metric(m) => m.get_value(model, state),
}
self.cost.get_cost(model, state)
}
fn set_cost_agg_func(&mut self, agg_func: CostAggFunc) {
self.cost.agg_func = agg_func
}
fn set_min_flow(&mut self, value: ConstraintValue) {
self.flow_constraints.min_flow = value;
Expand All @@ -725,7 +815,7 @@ impl OutputNode {
#[derive(Debug, PartialEq)]
pub struct LinkNode {
pub meta: NodeMeta<NodeIndex>,
pub cost: ConstraintValue,
cost: NodeCost,
pub flow_constraints: FlowConstraints,
pub incoming_edges: Vec<EdgeIndex>,
pub outgoing_edges: Vec<EdgeIndex>,
Expand All @@ -735,21 +825,20 @@ impl LinkNode {
fn new(index: &NodeIndex, name: &str, sub_name: Option<&str>) -> Self {
Self {
meta: NodeMeta::new(index, name, sub_name),
cost: ConstraintValue::None,
cost: NodeCost::default(),
flow_constraints: FlowConstraints::new(),
incoming_edges: Vec::new(),
outgoing_edges: Vec::new(),
}
}
fn set_cost(&mut self, value: ConstraintValue) {
self.cost = value
self.cost.local = value
}
fn set_cost_agg_func(&mut self, agg_func: CostAggFunc) {
self.cost.agg_func = agg_func
}
fn get_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
match &self.cost {
ConstraintValue::None => Ok(0.0),
ConstraintValue::Scalar(v) => Ok(*v),
ConstraintValue::Metric(m) => m.get_value(model, state),
}
self.cost.get_cost(model, state)
}
fn set_min_flow(&mut self, value: ConstraintValue) {
self.flow_constraints.min_flow = value;
Expand Down
50 changes: 48 additions & 2 deletions pywr-core/src/virtual_storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ impl VirtualStorageVec {
min_volume: ConstraintValue,
max_volume: ConstraintValue,
reset: VirtualStorageReset,
cost: ConstraintValue,
) -> VirtualStorageIndex {
let node_index = VirtualStorageIndex(self.nodes.len());
let node = VirtualStorage::new(
Expand All @@ -67,6 +68,7 @@ impl VirtualStorageVec {
min_volume,
max_volume,
reset,
cost,
);
self.nodes.push(node);
node_index
Expand All @@ -88,6 +90,7 @@ pub struct VirtualStorage {
pub initial_volume: StorageInitialVolume,
pub storage_constraints: StorageConstraints,
pub reset: VirtualStorageReset,
pub cost: ConstraintValue,
}

impl VirtualStorage {
Expand All @@ -101,6 +104,7 @@ impl VirtualStorage {
min_volume: ConstraintValue,
max_volume: ConstraintValue,
reset: VirtualStorageReset,
cost: ConstraintValue,
) -> Self {
Self {
meta: NodeMeta::new(index, name, sub_name),
Expand All @@ -110,6 +114,7 @@ impl VirtualStorage {
initial_volume,
storage_constraints: StorageConstraints::new(min_volume, max_volume),
reset,
cost,
}
}

Expand Down Expand Up @@ -139,6 +144,14 @@ impl VirtualStorage {
VirtualStorageState::new(0.0)
}

pub fn get_cost(&self, model: &Model, state: &State) -> Result<f64, PywrError> {
match &self.cost {
ConstraintValue::None => Ok(0.0),
ConstraintValue::Scalar(v) => Ok(*v),
ConstraintValue::Metric(m) => m.get_value(model, state),
}
}

pub fn before(&self, timestep: &Timestep, model: &Model, state: &mut State) -> Result<(), PywrError> {
let do_reset = if timestep.is_first() {
// Set the initial volume if it is the first timestep.
Expand Down Expand Up @@ -219,12 +232,13 @@ mod tests {
use crate::metric::Metric;
use crate::model::Model;
use crate::node::{ConstraintValue, StorageInitialVolume};
use crate::recorders::AssertionFnRecorder;
use crate::recorders::{AssertionFnRecorder, AssertionRecorder};
use crate::scenario::ScenarioIndex;
use crate::solvers::{ClpSolver, ClpSolverSettings};
use crate::test_utils::{default_timestepper, run_all_solvers};
use crate::test_utils::{default_timestepper, run_all_solvers, simple_model};
use crate::timestep::Timestep;
use crate::virtual_storage::{months_since_last_reset, VirtualStorageReset};
use ndarray::Array;
use time::macros::date;

/// Test the calculation of number of months since last reset
Expand Down Expand Up @@ -278,6 +292,7 @@ mod tests {
ConstraintValue::Scalar(0.0),
ConstraintValue::Scalar(100.0),
VirtualStorageReset::Never,
ConstraintValue::Scalar(0.0),
);

// Setup a demand on output-0 and output-1
Expand Down Expand Up @@ -319,4 +334,35 @@ mod tests {
// Test all solvers
run_all_solvers(&model, &timestepper);
}

#[test]
/// Test virtual storage node costs
fn test_virtual_storage_node_costs() {
let mut model = simple_model(1);
let timestepper = default_timestepper();

let nodes = vec![model.get_node_index_by_name("input", None).unwrap()];
// Virtual storage node cost is high enough to prevent any flow
model
.add_virtual_storage_node(
"vs",
None,
&nodes,
None,
StorageInitialVolume::Proportional(1.0),
ConstraintValue::Scalar(0.0),
ConstraintValue::Scalar(100.0),
VirtualStorageReset::Never,
ConstraintValue::Scalar(20.0),
)
.unwrap();

let expected = Array::zeros((366, 1));
let idx = model.get_node_by_name("output", None).unwrap().index();
let recorder = AssertionRecorder::new("output-flow", Metric::NodeInFlow(idx), expected, None, None);
model.add_recorder(Box::new(recorder)).unwrap();

// Test all solvers
run_all_solvers(&model, &timestepper);
}
}
Loading

0 comments on commit 5fcef40

Please sign in to comment.