Skip to content
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.

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

feat: Replace DynamicIndexValue with IndexMetric #317

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
24 changes: 24 additions & 0 deletions pywr-core/src/metric.rs
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,18 @@ impl From<(ParameterIndex<MultiValue>, String)> for MetricF64 {
}
}

impl From<(ParameterIndex<MultiValue>, String)> for MetricUsize {
fn from((idx, key): (ParameterIndex<MultiValue>, String)) -> Self {
match idx {
ParameterIndex::General(idx) => Self::MultiParameterValue((idx, key)),
ParameterIndex::Simple(idx) => Self::Simple(SimpleMetricUsize::MultiParameterValue((idx, key))),
ParameterIndex::Const(idx) => Self::Simple(SimpleMetricUsize::Constant(
ConstantMetricUsize::MultiParameterValue((idx, key)),
)),
}
}
}

impl TryFrom<ParameterIndex<f64>> for SimpleMetricF64 {
type Error = PywrError;
fn try_from(idx: ParameterIndex<f64>) -> Result<Self, Self::Error> {
Expand All @@ -288,13 +300,17 @@ impl TryFrom<ParameterIndex<usize>> for SimpleMetricUsize {
#[derive(Clone, Debug, PartialEq)]
pub enum ConstantMetricUsize {
IndexParameterValue(ConstParameterIndex<usize>),
MultiParameterValue((ConstParameterIndex<MultiValue>, String)),
Constant(usize),
}

impl ConstantMetricUsize {
pub fn get_value(&self, values: &ConstParameterValues) -> Result<usize, PywrError> {
match self {
ConstantMetricUsize::IndexParameterValue(idx) => values.get_const_parameter_usize(*idx),
ConstantMetricUsize::MultiParameterValue((idx, key)) => {
Ok(values.get_const_multi_parameter_usize(*idx, key)?)
}
ConstantMetricUsize::Constant(v) => Ok(*v),
}
}
Expand All @@ -303,13 +319,17 @@ impl ConstantMetricUsize {
#[derive(Clone, Debug, PartialEq)]
pub enum SimpleMetricUsize {
IndexParameterValue(SimpleParameterIndex<usize>),
MultiParameterValue((SimpleParameterIndex<MultiValue>, String)),
Constant(ConstantMetricUsize),
}

impl SimpleMetricUsize {
pub fn get_value(&self, values: &SimpleParameterValues) -> Result<usize, PywrError> {
match self {
SimpleMetricUsize::IndexParameterValue(idx) => values.get_simple_parameter_usize(*idx),
SimpleMetricUsize::MultiParameterValue((idx, key)) => {
Ok(values.get_simple_multi_parameter_usize(*idx, key)?)
}
SimpleMetricUsize::Constant(m) => m.get_value(values.get_constant_values()),
}
}
Expand All @@ -319,13 +339,17 @@ impl SimpleMetricUsize {
pub enum MetricUsize {
IndexParameterValue(GeneralParameterIndex<usize>),
Simple(SimpleMetricUsize),
MultiParameterValue((GeneralParameterIndex<MultiValue>, String)),
InterNetworkTransfer(MultiNetworkTransferIndex),
}

impl MetricUsize {
pub fn get_value(&self, _network: &Network, state: &State) -> Result<usize, PywrError> {
match self {
Self::IndexParameterValue(idx) => state.get_parameter_index(*idx),
Self::MultiParameterValue((idx, key)) => Ok(state.get_multi_parameter_index(*idx, key)?),
Self::Simple(s) => s.get_value(&state.get_simple_parameter_values()),
Self::InterNetworkTransfer(_idx) => todo!("Support usize for inter-network transfers"),
}
}
}
Expand Down
8 changes: 8 additions & 0 deletions pywr-core/src/network.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1362,6 +1362,14 @@ impl Network {
self.parameters.add_simple_f64(parameter)
}

/// Add a [`parameters::SimpleParameter`] to the network
pub fn add_simple_index_parameter(
&mut self,
parameter: Box<dyn parameters::SimpleParameter<usize>>,
) -> Result<ParameterIndex<usize>, PywrError> {
self.parameters.add_simple_usize(parameter)
}

/// Add a [`parameters::ConstParameter`] to the network
pub fn add_const_parameter(
&mut self,
Expand Down
126 changes: 96 additions & 30 deletions pywr-core/src/parameters/array.rs
Original file line number Diff line number Diff line change
@@ -1,44 +1,55 @@
use crate::network::Network;
use crate::parameters::{GeneralParameter, Parameter, ParameterMeta, ParameterName, ParameterState};
use crate::parameters::{Parameter, ParameterMeta, ParameterName, ParameterState, SimpleParameter};
use crate::scenario::ScenarioIndex;
use crate::state::State;
use crate::timestep::Timestep;
use crate::state::SimpleParameterValues;
use crate::timestep::{Timestep, TimestepIndex};
use crate::PywrError;
use ndarray::{Array1, Array2, Axis};

pub struct Array1Parameter {
pub struct Array1Parameter<T> {
meta: ParameterMeta,
array: Array1<f64>,
array: Array1<T>,
timestep_offset: Option<i32>,
}

impl Array1Parameter {
pub fn new(name: ParameterName, array: Array1<f64>, timestep_offset: Option<i32>) -> Self {
impl<T> Array1Parameter<T> {
pub fn new(name: ParameterName, array: Array1<T>, timestep_offset: Option<i32>) -> Self {
Self {
meta: ParameterMeta::new(name),
array,
timestep_offset,
}
}

/// Compute the time-step index to use accounting for any defined offset.
///
/// The offset is applied to the time-step index and then clamped to the bounds of the array.
/// This ensures that the time-step index is always within the bounds of the array.
fn timestep_index(&self, timestep: &Timestep) -> TimestepIndex {
match self.timestep_offset {
None => timestep.index,
Some(offset) => (timestep.index as i32 + offset)
.max(0)
.min(self.array.len_of(Axis(0)) as i32 - 1) as usize,
}
}
}
impl Parameter for Array1Parameter {
impl<T> Parameter for Array1Parameter<T>
where
T: Send + Sync + Clone,
{
fn meta(&self) -> &ParameterMeta {
&self.meta
}
}
impl GeneralParameter<f64> for Array1Parameter {
impl SimpleParameter<f64> for Array1Parameter<f64> {
fn compute(
&self,
timestep: &Timestep,
_scenario_index: &ScenarioIndex,
_model: &Network,
_state: &State,
_values: &SimpleParameterValues,
_internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<f64, PywrError> {
let idx = match self.timestep_offset {
None => timestep.index,
Some(offset) => (timestep.index as i32 + offset).max(0).min(self.array.len() as i32 - 1) as usize,
};
let idx = self.timestep_index(timestep);
// This panics if out-of-bounds
let value = self.array[[idx]];
Ok(value)
Expand All @@ -52,17 +63,39 @@ impl GeneralParameter<f64> for Array1Parameter {
}
}

pub struct Array2Parameter {
impl SimpleParameter<usize> for Array1Parameter<u64> {
fn compute(
&self,
timestep: &Timestep,
_scenario_index: &ScenarioIndex,
_values: &SimpleParameterValues,
_internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<usize, PywrError> {
let idx = self.timestep_index(timestep);
// This panics if out-of-bounds
let value = self.array[[idx]];
Ok(value as usize)
}

fn as_parameter(&self) -> &dyn Parameter
where
Self: Sized,
{
self
}
}

pub struct Array2Parameter<T> {
meta: ParameterMeta,
array: Array2<f64>,
array: Array2<T>,
scenario_group_index: usize,
timestep_offset: Option<i32>,
}

impl Array2Parameter {
impl<T> Array2Parameter<T> {
pub fn new(
name: ParameterName,
array: Array2<f64>,
array: Array2<T>,
scenario_group_index: usize,
timestep_offset: Option<i32>,
) -> Self {
Expand All @@ -73,30 +106,40 @@ impl Array2Parameter {
timestep_offset,
}
}

/// Compute the time-step index to use accounting for any defined offset.
///
/// The offset is applied to the time-step index and then clamped to the bounds of the array.
/// This ensures that the time-step index is always within the bounds of the array.
fn timestep_index(&self, timestep: &Timestep) -> TimestepIndex {
match self.timestep_offset {
None => timestep.index,
Some(offset) => (timestep.index as i32 + offset)
.max(0)
.min(self.array.len_of(Axis(0)) as i32 - 1) as usize,
}
}
}

impl Parameter for Array2Parameter {
impl<T> Parameter for Array2Parameter<T>
where
T: Send + Sync + Clone,
{
fn meta(&self) -> &ParameterMeta {
&self.meta
}
}

impl GeneralParameter<f64> for Array2Parameter {
impl SimpleParameter<f64> for Array2Parameter<f64> {
fn compute(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
_model: &Network,
_state: &State,
_values: &SimpleParameterValues,
_internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<f64, PywrError> {
// This panics if out-of-bounds
let t_idx = match self.timestep_offset {
None => timestep.index,
Some(offset) => (timestep.index as i32 + offset)
.max(0)
.min(self.array.len_of(Axis(0)) as i32 - 1) as usize,
};
let t_idx = self.timestep_index(timestep);
let s_idx = scenario_index.indices[self.scenario_group_index];

Ok(self.array[[t_idx, s_idx]])
Expand All @@ -109,3 +152,26 @@ impl GeneralParameter<f64> for Array2Parameter {
self
}
}

impl SimpleParameter<usize> for Array2Parameter<u64> {
fn compute(
&self,
timestep: &Timestep,
scenario_index: &ScenarioIndex,
_values: &SimpleParameterValues,
_internal_state: &mut Option<Box<dyn ParameterState>>,
) -> Result<usize, PywrError> {
// This panics if out-of-bounds
let t_idx = self.timestep_index(timestep);
let s_idx = scenario_index.indices[self.scenario_group_index];

Ok(self.array[[t_idx, s_idx]] as usize)
}

fn as_parameter(&self) -> &dyn Parameter
where
Self: Sized,
{
self
}
}
2 changes: 1 addition & 1 deletion pywr-core/src/parameters/control_curves/piecewise.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ mod test {
// Create an artificial volume series to use for the interpolation test
let volume = Array1Parameter::new("test-x".into(), Array1::linspace(1.0, 0.0, 21), None);

let volume_idx = model.network_mut().add_parameter(Box::new(volume)).unwrap();
let volume_idx = model.network_mut().add_simple_parameter(Box::new(volume)).unwrap();

let parameter = PiecewiseInterpolatedParameter::new(
"test-parameter".into(),
Expand Down
2 changes: 1 addition & 1 deletion pywr-core/src/parameters/delay.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ mod test {
let volumes = Array1::linspace(1.0, 0.0, 21);
let volume = Array1Parameter::new("test-x".into(), volumes.clone(), None);

let volume_idx = model.network_mut().add_parameter(Box::new(volume)).unwrap();
let volume_idx = model.network_mut().add_simple_parameter(Box::new(volume)).unwrap();

const DELAY: usize = 3; // 3 time-step delay
let parameter = DelayParameter::new(
Expand Down
2 changes: 1 addition & 1 deletion pywr-core/src/parameters/discount_factor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ mod test {
let volumes = Array1::linspace(1.0, 0.0, 21);
let volume = Array1Parameter::new("test-x".into(), volumes.clone(), None);

let _volume_idx = network.add_parameter(Box::new(volume)).unwrap();
let _volume_idx = network.add_simple_parameter(Box::new(volume)).unwrap();

let parameter = DiscountFactorParameter::new(
"test-parameter".into(),
Expand Down
17 changes: 11 additions & 6 deletions pywr-core/src/parameters/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1002,7 +1002,7 @@ impl ParameterCollection {
}

match parameter.try_into_simple() {
Some(simple) => self.add_simple_usize(simple).map(|idx| idx.into()),
Some(simple) => self.add_simple_usize(simple),
None => {
let index = GeneralParameterIndex::new(self.general_usize.len());
self.general_usize.push(parameter);
Expand All @@ -1014,17 +1014,22 @@ impl ParameterCollection {
pub fn add_simple_usize(
&mut self,
parameter: Box<dyn SimpleParameter<usize>>,
) -> Result<SimpleParameterIndex<usize>, PywrError> {
) -> Result<ParameterIndex<usize>, PywrError> {
if self.has_name(parameter.name()) {
return Err(PywrError::ParameterNameAlreadyExists(parameter.meta().name.to_string()));
}

let index = SimpleParameterIndex::new(self.simple_usize.len());
match parameter.try_into_const() {
Some(constant) => self.add_const_usize(constant),
None => {
let index = SimpleParameterIndex::new(self.simple_f64.len());

self.simple_usize.push(parameter);
self.simple_resolve_order.push(SimpleParameterType::Index(index));
self.simple_usize.push(parameter);
self.simple_resolve_order.push(SimpleParameterType::Index(index));

Ok(index)
Ok(index.into())
}
}
}

pub fn add_const_usize(
Expand Down
26 changes: 26 additions & 0 deletions pywr-core/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -428,6 +428,10 @@ impl ParameterValuesRef<'_> {
fn get_multi_value(&self, idx: usize, key: &str) -> Option<&f64> {
self.multi_values.get(idx).and_then(|s| s.get_value(key))
}

fn get_multi_index(&self, idx: usize, key: &str) -> Option<&usize> {
self.multi_values.get(idx).and_then(|s| s.get_index(key))
}
}

pub struct SimpleParameterValues<'a> {
Expand Down Expand Up @@ -461,6 +465,17 @@ impl SimpleParameterValues<'_> {
.copied()
}

pub fn get_simple_multi_parameter_usize(
&self,
idx: SimpleParameterIndex<MultiValue>,
key: &str,
) -> Result<usize, PywrError> {
self.simple
.get_multi_index(*idx.deref(), key)
.ok_or(PywrError::SimpleMultiValueParameterIndexNotFound(idx))
.copied()
}

pub fn get_constant_values(&self) -> &ConstParameterValues {
&self.constant
}
Expand Down Expand Up @@ -495,6 +510,17 @@ impl ConstParameterValues<'_> {
.ok_or(PywrError::ConstMultiValueParameterIndexNotFound(idx))
.copied()
}

pub fn get_const_multi_parameter_usize(
&self,
idx: ConstParameterIndex<MultiValue>,
key: &str,
) -> Result<usize, PywrError> {
self.constant
.get_multi_index(*idx.deref(), key)
.ok_or(PywrError::ConstMultiValueParameterIndexNotFound(idx))
.copied()
}
}

// State of the nodes and edges
Expand Down
Loading
Loading