Use WeeklyProfileSize to match weekly profile length

pywr · Nov 23, 2023 · 5a40dbd · 5a40dbd
1 parent 27ccfdb
commit 5a40dbd
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 20 deletions.
diff --git a/pywr-core/src/parameters/profiles/weekly.rs b/pywr-core/src/parameters/profiles/weekly.rs
@@ -6,13 +6,44 @@ use crate::timestep::Timestep;
 use crate::PywrError;
 use std::any::Any;
 
+// A weekly profile can be 52 or 53 week long
+enum WeeklyProfileSize {
+    FiftyTwo([f64; 52]),
+    FiftyThree([f64; 53]),
+}
+
+impl WeeklyProfileSize {
+    // Calculate the value for the current week
+    fn value(&self, date: &time::Date) -> f64 {
+        let current_day = date.ordinal() as usize;
+        let current_day_index = current_day - 1;
+
+        match self {
+            Self::FiftyTwo(values) => {
+                if current_day >= 364 {
+                    values[51]
+                } else {
+                    values[current_day_index / 7]
+                }
+            }
+            Self::FiftyThree(values) => {
+                if current_day >= 364 {
+                    values[52]
+                } else {
+                    values[current_day_index / 7]
+                }
+            }
+        }
+    }
+}
+
 pub struct WeeklyProfileParameter {
     meta: ParameterMeta,
-    values: Vec<f64>,
+    values: WeeklyProfileSize,
 }
 
 impl WeeklyProfileParameter {
-    pub fn new(name: &str, values: Vec<f64>) -> Self {
+    pub fn new(name: &str, values: WeeklyProfileSize) -> Self {
         if values.len() != 52 || values.len() != 53 {
             panic!("52 or 53 values must be given for the weekly profile named {name}");
         }
@@ -39,20 +70,6 @@ impl Parameter for WeeklyProfileParameter {
         _state: &State,
         _internal_state: &mut Option<Box<dyn Any + Send>>,
     ) -> Result<f64, PywrError> {
-        let current_day = timestep.date.ordinal();
-        let current_day_index = current_day - 1;
-
-        let week_index: u16;
-        if current_day >= 364 {
-            if self.values.len() == 52 {
-                week_index = 51
-            } else {
-                week_index = 52
-            }
-        } else {
-            week_index = current_day_index / 7
-        }
-
-        Ok(self.values[week_index as usize])
+        Ok(self.values.value(&timestep.date))
     }
 }
diff --git a/pywr-schema/src/parameters/profiles.rs b/pywr-schema/src/parameters/profiles.rs
@@ -5,7 +5,8 @@ use crate::parameters::{
 };
 use pywr_core::parameters::ParameterIndex;
 use pywr_v1_schema::parameters::{
-    DailyProfileParameter as DailyProfileParameterV1, MonthInterpDay as MonthInterpDayV1,
+    DailyProfileParameter as DailyProfileParameterV1,
+    MonthInterpDay as MonthInterpDayV1,
     MonthlyProfileParameter as MonthlyProfileParameterV1,
     // WeeklyProfileParameter as WeeklyProfileParameterV1,
     UniformDrawdownProfileParameter as UniformDrawdownProfileParameterV1,
@@ -243,7 +244,7 @@ impl WeeklyProfileParameter {
         let values = &self.values.load(tables)?;
         let p = pywr_core::parameters::WeeklyProfileParameter::new(
             &self.meta.name,
-            values.try_into().expect(""),
+            values.try_into().expect("")
         );
         Ok(model.add_parameter(Box::new(p))?)
     }
@@ -275,4 +276,4 @@ impl WeeklyProfileParameter {
 //         let p = Self { meta, values };
 //         Ok(p)
 //     }
-// }
+// }