feat: Switch from time-rs to chrono and implement non-daily timesteps (…

…#115)
pywr · Feb 24, 2024 · 8db2928 · 8db2928
1 parent edb3679
commit 8db2928
Show file tree

Hide file tree

Showing 27 changed files with 645 additions and 211 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -35,7 +35,6 @@ opt-level = 3 # fast and small wasm
 serde = { version = "1", features = ["derive"] }
 serde_json = "1.0"
 thiserror = "1.0.25"
-time = { version = "0.3", features = ["serde", "serde-well-known", "serde-human-readable", "macros"] }
 num = "0.4.0"
 ndarray = "0.15.3"
 polars = { version = "0.37.0", features = ["lazy", "rows", "ndarray"] }
@@ -46,3 +45,4 @@ tracing = {  version ="0.1", features = ["log"] }
 csv = "1.1"
 hdf5 = { git="https://github.com/aldanor/hdf5-rust.git", package = "hdf5", features=["static", "zlib"] }
 pywr-v1-schema = { git = "https://github.com/pywr/pywr-schema/", tag="v0.9.0", package = "pywr-schema" }
+chrono = { workspace = true }
diff --git a/pywr-cli/Cargo.toml b/pywr-cli/Cargo.toml
@@ -19,7 +19,6 @@ tracing = { workspace = true }
 tracing-subscriber = { version ="0.3.17", features=["env-filter"] }
 rand = "0.8.5"
 rand_chacha = "0.3.1"
-time = {  workspace = true, features = ["serde", "serde-well-known", "serde-human-readable", "macros"] }
 serde = { workspace = true }
 serde_json = { workspace = true }
 pywr-v1-schema = { workspace = true }

diff --git a/pywr-core/Cargo.toml b/pywr-core/Cargo.toml
@@ -24,13 +24,14 @@ csv = {  workspace = true }
 clp-sys = { path = "../clp-sys" }
 ipm-ocl = { path = "../ipm-ocl", optional = true }
 ipm-simd = { path = "../ipm-simd", optional = true }
-time = { workspace = true, features = ["macros"] }
 tracing = {  workspace = true }
 highs-sys = {  git = "https://github.com/jetuk/highs-sys", branch="fix-build-libz-linking", optional = true }
 # highs-sys = { path = "../../highs-sys" }
 nalgebra = "0.32.3"
+chrono = { workspace = true }
+polars = { workspace = true }
 
-pyo3 = {  workspace = true }
+pyo3 = {  workspace = true, features = ["chrono"] }
 
 
 rayon = "1.6.1"

diff --git a/pywr-core/src/lib.rs b/pywr-core/src/lib.rs
@@ -157,6 +157,10 @@ pub enum PywrError {
     ParameterNoInitialValue,
     #[error("parameter state not found for parameter index {0}")]
     ParameterStateNotFound(ParameterIndex),
+    #[error("Could not create timestep range due to following error: {0}")]
+    TimestepRangeGenerationError(String),
+    #[error("Could not create timesteps for frequency '{0}'")]
+    TimestepGenerationError(String),
 }
 
 // Python errors

diff --git a/pywr-core/src/models/mod.rs b/pywr-core/src/models/mod.rs
@@ -3,6 +3,7 @@ mod simple;
 
 use crate::scenario::{ScenarioDomain, ScenarioGroupCollection};
 use crate::timestep::{TimeDomain, Timestepper};
+use crate::PywrError;
 pub use multi::{MultiNetworkModel, MultiNetworkTransferIndex};
 pub use simple::{Model, ModelState};
 
@@ -16,11 +17,11 @@ impl ModelDomain {
         Self { time, scenarios }
     }
 
-    pub fn from(timestepper: Timestepper, scenario_collection: ScenarioGroupCollection) -> Self {
-        Self {
-            time: timestepper.into(),
+    pub fn from(timestepper: Timestepper, scenario_collection: ScenarioGroupCollection) -> Result<Self, PywrError> {
+        Ok(Self {
+            time: TimeDomain::try_from(timestepper)?,
             scenarios: scenario_collection.into(),
-        }
+        })
     }
 
     pub fn time(&self) -> &TimeDomain {
@@ -36,12 +37,15 @@ impl ModelDomain {
     }
 }
 
-impl From<Timestepper> for ModelDomain {
-    fn from(value: Timestepper) -> Self {
-        Self {
-            time: value.into(),
+impl TryFrom<Timestepper> for ModelDomain {
+    type Error = PywrError;
+
+    fn try_from(value: Timestepper) -> Result<Self, Self::Error> {
+        let time = TimeDomain::try_from(value)?;
+        Ok(Self {
+            time,
             scenarios: ScenarioGroupCollection::default().into(),
-        }
+        })
     }
 }
 

diff --git a/pywr-core/src/models/multi.rs b/pywr-core/src/models/multi.rs
@@ -367,7 +367,7 @@ mod tests {
         let mut scenario_collection = ScenarioGroupCollection::default();
         scenario_collection.add_group("test-scenario", 2);
 
-        let mut multi_model = MultiNetworkModel::new(ModelDomain::from(timestepper, scenario_collection));
+        let mut multi_model = MultiNetworkModel::new(ModelDomain::from(timestepper, scenario_collection).unwrap());
 
         let test_scenario_group_idx = multi_model
             .domain()

diff --git a/pywr-core/src/parameters/discount_factor.rs b/pywr-core/src/parameters/discount_factor.rs
@@ -5,6 +5,7 @@ use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
 use crate::timestep::Timestep;
 use crate::PywrError;
+use chrono::Datelike;
 use std::any::Any;
 
 pub struct DiscountFactorParameter {

diff --git a/pywr-core/src/parameters/mod.rs b/pywr-core/src/parameters/mod.rs
@@ -388,13 +388,16 @@ pub trait VariableParameter<T> {
 #[cfg(test)]
 mod tests {
 
-    use crate::timestep::Timestepper;
-    use time::macros::date;
+    use crate::timestep::{TimestepDuration, Timestepper};
+    use chrono::NaiveDateTime;
 
     // TODO tests need re-enabling
     #[allow(dead_code)]
     fn default_timestepper() -> Timestepper {
-        Timestepper::new(date!(2020 - 01 - 01), date!(2020 - 01 - 15), 1)
+        let start = NaiveDateTime::parse_from_str("2020-01-01 00:00:00", "%Y-%m-%d %H:%M:%S").unwrap();
+        let end = NaiveDateTime::parse_from_str("2020-01-15 00:00:00", "%Y-%m-%d %H:%M:%S").unwrap();
+        let duration = TimestepDuration::Days(1);
+        Timestepper::new(start, end, duration)
     }
 
     // #[test]

diff --git a/pywr-core/src/parameters/profiles/daily.rs b/pywr-core/src/parameters/profiles/daily.rs
@@ -4,6 +4,7 @@ use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
 use crate::timestep::Timestep;
 use crate::PywrError;
+use chrono::Datelike;
 use std::any::Any;
 
 pub struct DailyProfileParameter {

diff --git a/pywr-core/src/parameters/profiles/monthly.rs b/pywr-core/src/parameters/profiles/monthly.rs
@@ -4,9 +4,8 @@ use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
 use crate::timestep::Timestep;
 use crate::PywrError;
+use chrono::{Datelike, NaiveDateTime};
 use std::any::Any;
-use time::util::days_in_year_month;
-use time::Date;
 
 #[derive(Copy, Clone)]
 pub enum MonthlyInterpDay {
@@ -30,10 +29,20 @@ impl MonthlyProfileParameter {
     }
 }
 
+fn days_in_year_month(datetime: &NaiveDateTime) -> u32 {
+    match datetime.month() {
+        1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
+        4 | 6 | 9 | 11 => 30,
+        2 if datetime.date().leap_year() => 29,
+        2 => 28,
+        _ => panic!("Invalid month"),
+    }
+}
+
 /// Interpolate between first_value and last value based on the day of the month. The last
 /// value is assumed to correspond to the first day of the next month.
-fn interpolate_first(date: &Date, first_value: f64, last_value: f64) -> f64 {
-    let days_in_month = days_in_year_month(date.year(), date.month());
+fn interpolate_first(date: &NaiveDateTime, first_value: f64, last_value: f64) -> f64 {
+    let days_in_month = days_in_year_month(date);
 
     if date.day() <= 1 {
         first_value
@@ -46,8 +55,8 @@ fn interpolate_first(date: &Date, first_value: f64, last_value: f64) -> f64 {
 
 /// Interpolate between first_value and last value based on the day of the month. The first
 /// value is assumed to correspond to the last day of the previous month.
-fn interpolate_last(date: &Date, first_value: f64, last_value: f64) -> f64 {
-    let days_in_month = days_in_year_month(date.year(), date.month());
+fn interpolate_last(date: &NaiveDateTime, first_value: f64, last_value: f64) -> f64 {
+    let days_in_month = days_in_year_month(date);
 
     if date.day() < 1 {
         first_value
@@ -76,13 +85,16 @@ impl Parameter for MonthlyProfileParameter {
         let v = match &self.interp_day {
             Some(interp_day) => match interp_day {
                 MonthlyInterpDay::First => {
-                    let first_value = self.values[timestep.date.month() as usize - 1];
-                    let last_value = self.values[timestep.date.month().next() as usize - 1];
+                    let next_month0 = (timestep.date.month0() + 1) % 12;
+                    let first_value = self.values[timestep.date.month0() as usize];
+                    let last_value = self.values[next_month0 as usize];
 
                     interpolate_first(&timestep.date, first_value, last_value)
                 }
                 MonthlyInterpDay::Last => {
-                    let first_value = self.values[timestep.date.month().previous() as usize - 1];
+                    let current_month = timestep.date.month();
+                    let last_month = if current_month == 1 { 12 } else { current_month - 1 };
+                    let first_value = self.values[last_month as usize - 1];
                     let last_value = self.values[timestep.date.month() as usize - 1];
 
                     interpolate_last(&timestep.date, first_value, last_value)

diff --git a/pywr-core/src/parameters/profiles/rbf.rs b/pywr-core/src/parameters/profiles/rbf.rs
@@ -7,6 +7,7 @@ use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
 use crate::timestep::Timestep;
 use crate::PywrError;
+use chrono::Datelike;
 use nalgebra::DMatrix;
 use std::any::Any;
 

diff --git a/pywr-core/src/parameters/profiles/uniform_drawdown.rs b/pywr-core/src/parameters/profiles/uniform_drawdown.rs
@@ -4,8 +4,8 @@ use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
 use crate::timestep::Timestep;
 use crate::PywrError;
+use chrono::{Datelike, NaiveDate};
 use std::any::Any;
-use time::{Date, Month};
 
 fn is_leap_year(year: i32) -> bool {
     (year % 4 == 0) & ((year % 100 != 0) | (year % 400 == 0))
@@ -18,11 +18,11 @@ pub struct UniformDrawdownProfileParameter {
 }
 
 impl UniformDrawdownProfileParameter {
-    pub fn new(name: &str, reset_day: u8, reset_month: Month, residual_days: u8) -> Self {
+    pub fn new(name: &str, reset_day: u32, reset_month: u32, residual_days: u8) -> Self {
         // Calculate the reset day of year in a known leap year.
-        let reset_doy = Date::from_calendar_date(2016, reset_month, reset_day)
+        let reset_doy = NaiveDate::from_ymd_opt(2016, reset_month, reset_day)
             .expect("Invalid reset day")
-            .ordinal();
+            .ordinal() as u16;
 
         Self {
             meta: ParameterMeta::new(name),

diff --git a/pywr-core/src/parameters/py.rs b/pywr-core/src/parameters/py.rs
@@ -4,8 +4,9 @@ use crate::network::Network;
 use crate::parameters::{downcast_internal_state_mut, MultiValueParameter};
 use crate::scenario::ScenarioIndex;
 use crate::state::{MultiValue, ParameterState, State};
+use chrono::Datelike;
 use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyDate, PyDict, PyFloat, PyLong, PyTuple};
+use pyo3::types::{IntoPyDict, PyDict, PyFloat, PyLong, PyTuple};
 use std::any::Any;
 use std::collections::HashMap;
 
@@ -116,19 +117,14 @@ impl Parameter for PyParameter {
         let internal = downcast_internal_state_mut::<Internal>(internal_state);
 
         let value: f64 = Python::with_gil(|py| {
-            let date = PyDate::new(
-                py,
-                timestep.date.year(),
-                timestep.date.month() as u8,
-                timestep.date.day(),
-            )?;
+            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, si.as_ref(py), metric_dict, index_dict]);
+            let args = PyTuple::new(py, [date.as_ref(py), si.as_ref(py), metric_dict, index_dict]);
 
             internal.user_obj.call_method1(py, "calc", args)?.extract(py)
         })
@@ -150,19 +146,14 @@ impl Parameter for PyParameter {
         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 = PyDate::new(
-                    py,
-                    timestep.date.year(),
-                    timestep.date.month() as u8,
-                    timestep.date.day(),
-                )?;
+                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, si.as_ref(py), metric_dict, index_dict]);
+                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)?;
             }
@@ -218,20 +209,14 @@ impl MultiValueParameter for PyParameter {
         let internal = downcast_internal_state_mut::<Internal>(internal_state);
 
         let value: MultiValue = Python::with_gil(|py| {
-            let date = PyDate::new(
-                py,
-                timestep.date.year(),
-                timestep.date.month() as u8,
-                timestep.date.day(),
-            )
-            .map_err(|e: PyErr| PywrError::PythonError(e.to_string()))?;
+            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, si.as_ref(py), metric_dict, index_dict]);
+            let args = PyTuple::new(py, [date.as_ref(py), si.as_ref(py), metric_dict, index_dict]);
 
             let py_values: HashMap<String, PyObject> = internal
                 .user_obj
@@ -282,19 +267,14 @@ impl MultiValueParameter for PyParameter {
         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 = PyDate::new(
-                    py,
-                    timestep.date.year(),
-                    timestep.date.month() as u8,
-                    timestep.date.day(),
-                )?;
+                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, si.as_ref(py), metric_dict, index_dict]);
+                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)?;
             }
@@ -344,7 +324,7 @@ class MyParameter:
 
         let param = PyParameter::new("my-parameter", class, args, kwargs, &HashMap::new(), &HashMap::new());
         let timestepper = default_timestepper();
-        let time: TimeDomain = timestepper.into();
+        let time: TimeDomain = TimeDomain::try_from(timestepper).unwrap();
         let timesteps = time.timesteps();
 
         let scenario_indices = [
@@ -413,7 +393,7 @@ class MyParameter:
 
         let param = PyParameter::new("my-parameter", class, args, kwargs, &HashMap::new(), &HashMap::new());
         let timestepper = default_timestepper();
-        let time: TimeDomain = timestepper.into();
+        let time: TimeDomain = TimeDomain::try_from(timestepper).unwrap();
         let timesteps = time.timesteps();
 
         let scenario_indices = [

diff --git a/pywr-core/src/parameters/rhai.rs b/pywr-core/src/parameters/rhai.rs
@@ -4,6 +4,7 @@ use crate::network::Network;
 use crate::parameters::downcast_internal_state_mut;
 use crate::scenario::ScenarioIndex;
 use crate::state::{ParameterState, State};
+use chrono::Datelike;
 use rhai::{Dynamic, Engine, Map, Scope, AST};
 use std::any::Any;
 use std::collections::HashMap;
@@ -150,7 +151,7 @@ mod tests {
         );
 
         let timestepper = default_timestepper();
-        let time: TimeDomain = timestepper.into();
+        let time: TimeDomain = TimeDomain::try_from(timestepper).unwrap();
         let timesteps = time.timesteps();
 
         let scenario_indices = [