Store row offsets with keys

This makes things cleaner and safer.

Author: alexdewar

src/simulation/optimisation.rs

@@ -78,52 +78,41 @@ impl Solution<'_> {
             })
     }
 
-    /// Returns constraint keys along with their corresponding dual values
-    fn zip_duals<'a, I, T>(&'a self, keys: I, offset: usize) -> impl Iterator<Item = (T, f64)> + 'a
-    where
-        I: Iterator<Item = T> + 'a,
-    {
-        keys.zip(self.solution.dual_rows()[offset..].iter().copied())
-    }
-
     /// Keys and dual values for commodity balance constraints.
     pub fn iter_commodity_balance_duals(
         &self,
     ) -> impl Iterator<Item = (&CommodityID, &RegionID, &TimeSliceID, f64)> {
         // Each commodity balance constraint applies to a particular time slice
         // selection (depending on time slice level). Where this covers multiple timeslices,
         // we return the same dual for each individual timeslice.
-        self.zip_duals(
-            self.constraint_keys.commodity_balance_keys.iter(),
-            self.constraint_keys.commodity_balance_keys_offset(),
-        )
-        .flat_map(|((commodity_id, region_id, ts_selection), price)| {
-            self.time_slice_info
-                .iter_selection(ts_selection)
-                .map(move |(ts, _)| (commodity_id, region_id, ts, price))
-        })
+        self.constraint_keys
+            .commodity_balance_keys
+            .zip_duals(self.solution.dual_rows())
+            .flat_map(|((commodity_id, region_id, ts_selection), price)| {
+                self.time_slice_info
+                    .iter_selection(ts_selection)
+                    .map(move |(ts, _)| (commodity_id, region_id, ts, price))
+            })
     }
 
     /// Keys and dual values for capacity constraints.
     pub fn iter_capacity_duals(&self) -> impl Iterator<Item = (AssetID, &TimeSliceID, f64)> {
-        self.zip_duals(
-            self.constraint_keys.capacity_keys.iter(),
-            self.constraint_keys.capacity_keys_offset(),
-        )
-        .map(|((asset_id, time_slice), dual)| (*asset_id, time_slice, dual))
+        self.constraint_keys
+            .capacity_keys
+            .zip_duals(self.solution.dual_rows())
+            .map(|((asset_id, time_slice), dual)| (*asset_id, time_slice, dual))
     }
 
     /// Keys and dual values for fixed asset constraints.
     pub fn iter_fixed_asset_duals(
         &self,
     ) -> impl Iterator<Item = (AssetID, &CommodityID, &TimeSliceID, f64)> {
-        self.zip_duals(
-            self.constraint_keys.fixed_asset_keys.iter(),
-            self.constraint_keys.fixed_asset_keys_offset(),
-        )
-        .map(|((asset_id, commodity_id, time_slice), dual)| {
-            (*asset_id, commodity_id, time_slice, dual)
-        })
+        self.constraint_keys
+            .fixed_asset_keys
+            .zip_duals(self.solution.dual_rows())
+            .map(|((asset_id, commodity_id, time_slice), dual)| {
+                (*asset_id, commodity_id, time_slice, dual)
+            })
     }
 }
 

src/simulation/optimisation/constraints.rs

@@ -1,4 +1,5 @@
 //! Code for adding constraints to the dispatch optimisation problem.
+use super::VariableMap;
 use crate::asset::{AssetID, AssetPool};
 use crate::commodity::{CommodityID, CommodityType};
 use crate::model::Model;
@@ -7,42 +8,41 @@ use crate::time_slice::{TimeSliceID, TimeSliceInfo, TimeSliceSelection};
 use highs::RowProblem as Problem;
 use std::rc::Rc;
 
-use super::VariableMap;
+/// Corresponding variables for a constraint along with the row offset in the solution
+pub struct KeysWithOffset<T> {
+    offset: usize,
+    keys: Vec<T>,
+}
+
+impl<T> KeysWithOffset<T> {
+    /// Zip the keys with the corresponding dual values in the solution, accounting for the offset
+    pub fn zip_duals<'a>(&'a self, duals: &'a [f64]) -> impl Iterator<Item = (&'a T, f64)> {
+        assert!(
+            self.offset + self.keys.len() <= duals.len(),
+            "Bad constraint keys: dual rows out of range"
+        );
+
+        self.keys.iter().zip(duals[self.offset..].iter().copied())
+    }
+}
 
 /// Indicates the commodity ID and time slice selection covered by each commodity balance constraint
-pub type CommodityBalanceConstraintKeys = Vec<(CommodityID, RegionID, TimeSliceSelection)>;
+pub type CommodityBalanceKeys = KeysWithOffset<(CommodityID, RegionID, TimeSliceSelection)>;
 
 /// Indicates the asset ID and time slice covered by each capacity constraint
-pub type CapacityConstraintKeys = Vec<(AssetID, TimeSliceID)>;
+pub type CapacityKeys = KeysWithOffset<(AssetID, TimeSliceID)>;
 
 /// Indicates the asset ID, commodity ID and time slice for each fixed asset constraint
-pub type FixedAssetConstraintKeys = Vec<(AssetID, CommodityID, TimeSliceID)>;
+pub type FixedAssetKeys = KeysWithOffset<(AssetID, CommodityID, TimeSliceID)>;
 
 /// The keys for different constraints
 pub struct ConstraintKeys {
     /// Keys for commodity balance constraints
-    pub commodity_balance_keys: CommodityBalanceConstraintKeys,
+    pub commodity_balance_keys: CommodityBalanceKeys,
     /// Keys for capacity constraints
-    pub capacity_keys: CapacityConstraintKeys,
+    pub capacity_keys: CapacityKeys,
     /// Keys for fixed asset constraints
-    pub fixed_asset_keys: FixedAssetConstraintKeys,
-}
-
-impl ConstraintKeys {
-    /// Start offset for commodity balance constraints
-    pub fn commodity_balance_keys_offset(&self) -> usize {
-        0
-    }
-
-    /// Start offset for capacity constraints
-    pub fn capacity_keys_offset(&self) -> usize {
-        self.commodity_balance_keys.len()
-    }
-
-    /// Start offset for capacity constraints
-    pub fn fixed_asset_keys_offset(&self) -> usize {
-        self.capacity_keys_offset() + self.capacity_keys.len()
-    }
+    pub fixed_asset_keys: FixedAssetKeys,
 }
 
 /// Add asset-level constraints
@@ -72,27 +72,16 @@ pub fn add_asset_constraints(
     let commodity_balance_keys =
         add_commodity_balance_constraints(problem, variables, model, assets, year);
 
-    let capacity_keys = add_asset_capacity_constraints(
-        problem,
-        variables,
-        assets,
-        &model.time_slice_info,
-        &commodity_balance_keys,
-    );
+    let capacity_keys =
+        add_asset_capacity_constraints(problem, variables, assets, &model.time_slice_info);
 
     // **TODO**: Currently it's safe to assume all process flows are non-flexible, as we enforce
     // this when reading data in. Once we've added support for flexible process flows, we will
     // need to add different constraints for assets with flexible and non-flexible flows.
     //
     // See: https://github.com/EnergySystemsModellingLab/MUSE_2.0/issues/360
-    let fixed_asset_keys = add_fixed_asset_constraints(
-        problem,
-        variables,
-        assets,
-        &model.time_slice_info,
-        &commodity_balance_keys,
-        &capacity_keys,
-    );
+    let fixed_asset_keys =
+        add_fixed_asset_constraints(problem, variables, assets, &model.time_slice_info);
 
     // Return constraint keys
     ConstraintKeys {
@@ -115,16 +104,12 @@ fn add_commodity_balance_constraints(
     model: &Model,
     assets: &AssetPool,
     year: u32,
-) -> CommodityBalanceConstraintKeys {
-    // Sanity check: we rely on the first n values of the dual row values corresponding to the
-    // commodity constraints, so these must be the first rows
-    assert!(
-        problem.num_rows() == 0,
-        "Commodity balance constraints must be added before other constraints"
-    );
+) -> CommodityBalanceKeys {
+    // Row offset in problem. This line **must** come before we add more constraints.
+    let offset = problem.num_rows();
 
     let mut terms = Vec::new();
-    let mut keys = CommodityBalanceConstraintKeys::new();
+    let mut keys = Vec::new();
     for commodity in model.commodities.values() {
         if commodity.kind != CommodityType::SupplyEqualsDemand
             && commodity.kind != CommodityType::ServiceDemand
@@ -183,7 +168,7 @@ fn add_commodity_balance_constraints(
         }
     }
 
-    keys
+    CommodityBalanceKeys { offset, keys }
 }
 
 /// Add asset-level capacity and availability constraints.
@@ -200,17 +185,12 @@ fn add_asset_capacity_constraints(
     variables: &VariableMap,
     assets: &AssetPool,
     time_slice_info: &TimeSliceInfo,
-    commodity_balance_keys: &CommodityBalanceConstraintKeys,
-) -> CapacityConstraintKeys {
-    // Sanity check: we rely on the dual rows corresponding to the capacity constraints being
-    // immediately after the commodity balance constraints in `problem`.
-    assert!(
-        problem.num_rows() == commodity_balance_keys.len(),
-        "Capacity constraints must be added immediately after commodity constraints."
-    );
+) -> CapacityKeys {
+    // Row offset in problem. This line **must** come before we add more constraints.
+    let offset = problem.num_rows();
 
     let mut terms = Vec::new();
-    let mut keys = CapacityConstraintKeys::new();
+    let mut keys = Vec::new();
     for asset in assets.iter() {
         for time_slice in time_slice_info.iter_ids() {
             let mut is_input = false; // NB: there will be at least one PAC
@@ -234,7 +214,8 @@ fn add_asset_capacity_constraints(
             keys.push((asset.id, time_slice.clone()));
         }
     }
-    keys
+
+    CapacityKeys { offset, keys }
 }
 
 /// Add constraints for non-flexible assets.
@@ -249,17 +230,11 @@ fn add_fixed_asset_constraints(
     variables: &VariableMap,
     assets: &AssetPool,
     time_slice_info: &TimeSliceInfo,
-    commodity_balance_keys: &CommodityBalanceConstraintKeys,
-    capacity_keys: &CapacityConstraintKeys,
-) -> FixedAssetConstraintKeys {
-    // Sanity check: we rely on the dual rows corresponding to these constraints being
-    // immediately after the commodity balance and capacity constraints in `problem`.
-    assert!(
-        problem.num_rows() == commodity_balance_keys.len() + capacity_keys.len(),
-        "Fixed asset constraints must be added immediately after commodity constraints."
-    );
-
-    let mut keys = FixedAssetConstraintKeys::new();
+) -> FixedAssetKeys {
+    // Row offset in problem. This line **must** come before we add more constraints.
+    let offset = problem.num_rows();
+
+    let mut keys = Vec::new();
     for asset in assets.iter() {
         // Get first PAC. unwrap is safe because all processes have at least one PAC.
         let pac1 = asset.iter_pacs().next().unwrap();
@@ -283,5 +258,5 @@ fn add_fixed_asset_constraints(
         }
     }
 
-    keys
+    FixedAssetKeys { offset, keys }
 }