feat: allow timeseries with single data column to be used without spe…

…cifying column name
pywr · Jun 10, 2024 · 6de2b4b · 6de2b4b
1 parent 1b07603
commit 6de2b4b
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 16 deletions.
diff --git a/pywr-schema/src/metric.rs b/pywr-schema/src/metric.rs
@@ -63,14 +63,18 @@ impl Metric {
             }
             Self::Timeseries(ts_ref) => {
                 let param_idx = match &ts_ref.columns {
-                    TimeseriesColumns::Scenario(scenario) => {
+                    Some(TimeseriesColumns::Scenario(scenario)) => {
                         args.timeseries
                             .load_df(network, ts_ref.name.as_ref(), args.domain, scenario.as_str())?
                     }
-                    TimeseriesColumns::Column(col) => {
+                    Some(TimeseriesColumns::Column(col)) => {
                         args.timeseries
                             .load_column(network, ts_ref.name.as_ref(), col.as_str())?
                     }
+                    None => {
+                        args.timeseries
+                            .load_single_column(network, ts_ref.name.as_ref())?
+                    }
                 };
                 Ok(MetricF64::ParameterValue(param_idx))
             }
@@ -211,12 +215,12 @@ impl TryFromV1Parameter<ParameterValueV1> for Metric {
                         };
 
                         let cols = match (&t.column, &t.scenario) {
-                            (Some(col), None) => TimeseriesColumns::Column(col.clone()),
-                            (None, Some(scenario)) => TimeseriesColumns::Scenario(scenario.clone()),
+                            (Some(col), None) => Some(TimeseriesColumns::Column(col.clone())),
+                            (None, Some(scenario)) => Some(TimeseriesColumns::Scenario(scenario.clone())),
                             (Some(_), Some(_)) => {
                                 return Err(ConversionError::AmbiguousColumnAndScenario(name.clone()))
                             }
-                            (None, None) => return Err(ConversionError::MissingColumnOrScenario(name.clone())),
+                            (None, None) => None
                         };
 
                         Self::Timeseries(TimeseriesReference::new(name, cols))
@@ -238,11 +242,11 @@ pub enum TimeseriesColumns {
 #[derive(serde::Deserialize, serde::Serialize, Debug, Clone, JsonSchema)]
 pub struct TimeseriesReference {
     name: String,
-    columns: TimeseriesColumns,
+    columns: Option<TimeseriesColumns>,
 }
 
 impl TimeseriesReference {
-    pub fn new(name: String, columns: TimeseriesColumns) -> Self {
+    pub fn new(name: String, columns: Option<TimeseriesColumns>) -> Self {
         Self { name, columns }
     }
 

diff --git a/pywr-schema/src/model.rs b/pywr-schema/src/model.rs
@@ -312,10 +312,10 @@ impl PywrNetwork {
                     };
 
                     let cols = match (&ts_ref.column, &ts_ref.scenario) {
-                        (Some(col), None) => TimeseriesColumns::Column(col.clone()),
-                        (None, Some(scenario)) => TimeseriesColumns::Scenario(scenario.clone()),
+                        (Some(col), None) => Some(TimeseriesColumns::Column(col.clone())),
+                        (None, Some(scenario)) => Some(TimeseriesColumns::Scenario(scenario.clone())),
                         (Some(_), Some(_)) => return,
-                        (None, None) => return,
+                        (None, None) => None,
                     };
 
                     *m = Metric::Timeseries(TimeseriesReference::new(name, cols));

diff --git a/pywr-schema/src/test_models/timeseries.json b/pywr-schema/src/test_models/timeseries.json
@@ -73,11 +73,7 @@
         "metrics": [
           {
             "type": "Timeseries",
-            "name": "inflow",
-            "columns": {
-              "type": "Column",
-              "name": "inflow1"
-            }
+            "name": "inflow"
           },
           {
             "type": "Constant",

diff --git a/pywr-schema/src/timeseries/align_and_resample.rs b/pywr-schema/src/timeseries/align_and_resample.rs
@@ -78,12 +78,15 @@ pub fn align_and_resample(
         Ordering::Equal => df,
     };
 
-    let df = slice_end(df, time_col, domain)?;
+    let mut df = slice_end(df, time_col, domain)?;
 
     if df.height() != domain.time().timesteps().len() {
         return Err(TimeseriesError::DataFrameTimestepMismatch(name.to_string()));
     }
 
+    // time column is no longer needed
+    let _ = df.drop_in_place(time_col)?;
+
     Ok(df)
 }
 

diff --git a/pywr-schema/src/timeseries/mod.rs b/pywr-schema/src/timeseries/mod.rs
@@ -44,6 +44,10 @@ pub enum TimeseriesError {
     DataFrameTimestepMismatch(String),
     #[error("A timeseries dataframe with the name '{0}' already exists.")]
     TimeseriesDataframeAlreadyExists(String),
+    #[error("The timeseries dataset '{0}' has more than one column of data so a column or scenario name must be provided for any reference")]
+    TimeseriesColumnOrScenarioRequired(String),
+    #[error("The timeseries dataset '{0}' has no columns")]
+    TimeseriesDataframeHasNoColumns(String),
     #[error("Polars error: {0}")]
     #[cfg(feature = "core")]
     PolarsError(#[from] PolarsError),
@@ -149,6 +153,44 @@ impl LoadedTimeseriesCollection {
         }
     }
 
+    pub fn load_single_column(
+        &self,
+        network: &mut pywr_core::network::Network,
+        name: &str,
+    ) -> Result<ParameterIndex<f64>, TimeseriesError> {
+        let df = self
+            .timeseries
+            .get(name)
+            .ok_or(TimeseriesError::TimeseriesNotFound(name.to_string()))?;
+
+        let cols = df.get_column_names();
+
+        if cols.len() > 1 {
+            return Err(TimeseriesError::TimeseriesColumnOrScenarioRequired(name.to_string()));
+        };
+
+        let col = cols.first().ok_or(TimeseriesError::ColumnNotFound {
+            col: "".to_string(),
+            name: name.to_string(),
+        })?;
+
+        let series = df.column(col)?;
+
+        let array = series.cast(&Float64)?.f64()?.to_ndarray()?.to_owned();
+        let name = format!("{}_{}", name, col);
+
+        match network.get_parameter_index_by_name(&name) {
+            Ok(idx) => Ok(idx),
+            Err(e) => match e {
+                PywrError::ParameterNotFound(_) => {
+                    let p = Array1Parameter::new(&name, array, None);
+                    Ok(network.add_parameter(Box::new(p))?)
+                }
+                _ => Err(TimeseriesError::PywrCore(e)),
+            },
+        }
+    }
+
     pub fn load_df(
         &self,
         network: &mut pywr_core::network::Network,