Fix unit tests

Author: tsmbland

src/commodity.rs

@@ -4,11 +4,9 @@ use crate::region::RegionID;
 use crate::time_slice::{TimeSliceID, TimeSliceLevel, TimeSliceSelection};
 use crate::units::{Flow, MoneyPerFlow};
 use indexmap::IndexMap;
-use itertools::Itertools;
 use serde::Deserialize;
 use serde_string_enum::DeserializeLabeledStringEnum;
 use std::collections::HashMap;
-use std::fmt::Display;
 use std::rc::Rc;
 
 define_id_type! {CommodityID}
@@ -54,33 +52,6 @@ pub struct Commodity {
 }
 define_id_getter! {Commodity, CommodityID}
 
-/// A set of commodities that can be invested in together
-pub enum InvestmentSet {
-    /// Single commodity
-    Single(CommodityID),
-    /// A set of commodities that form a cycle
-    Cycle(Vec<CommodityID>),
-}
-
-impl InvestmentSet {
-    /// Returns an iterator over the commodity IDs in this investment set
-    pub fn iter(&self) -> impl Iterator<Item = &CommodityID> {
-        match self {
-            InvestmentSet::Single(id) => std::slice::from_ref(id).iter(),
-            InvestmentSet::Cycle(ids) => ids.iter(),
-        }
-    }
-}
-
-impl Display for InvestmentSet {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            InvestmentSet::Single(id) => write!(f, "{id}"),
-            InvestmentSet::Cycle(ids) => write!(f, "[{}]", ids.iter().join(", ")),
-        }
-    }
-}
-
 /// Type of balance for application of cost
 #[derive(PartialEq, Clone, Debug, DeserializeLabeledStringEnum)]
 pub enum BalanceType {

src/graph.rs

@@ -1,7 +1,8 @@
 //! Module for creating and analysing commodity graphs
-use crate::commodity::{CommodityID, CommodityMap, CommodityType, InvestmentSet};
+use crate::commodity::{CommodityID, CommodityMap, CommodityType};
 use crate::process::{ProcessID, ProcessMap};
 use crate::region::RegionID;
+use crate::simulation::investment::InvestmentSet;
 use crate::time_slice::{TimeSliceInfo, TimeSliceLevel, TimeSliceSelection};
 use crate::units::{Dimensionless, Flow};
 use anyhow::{Context, Result, ensure};
@@ -338,7 +339,7 @@ fn solve_investment_order(
         .iter()
         .rev()
         .filter_map(|node_idx| {
-            // Get set of commodity IDs for the node
+            // Get set of commodity ID(s) for the node, referring back to `condensed_graph`
             let commodities: Vec<CommodityID> = condensed_graph
                 .node_weight(*node_idx)
                 .unwrap()
@@ -495,7 +496,10 @@ mod tests {
     use std::rc::Rc;
 
     #[rstest]
-    fn test_topo_sort_linear_graph(sed_commodity: Commodity, svd_commodity: Commodity) {
+    fn test_solve_investment_order_linear_graph(
+        sed_commodity: Commodity,
+        svd_commodity: Commodity,
+    ) {
         // Create a simple linear graph: A -> B -> C
         let mut graph = Graph::new();
 
@@ -513,17 +517,18 @@ mod tests {
         commodities.insert("B".into(), Rc::new(sed_commodity));
         commodities.insert("C".into(), Rc::new(svd_commodity));
 
-        let result = solve_investment_order(&graph, &commodities).unwrap();
+        let result = solve_investment_order(&graph, &commodities);
 
         // Expected order: C, B, A (leaf nodes first)
+        // No cycles, so all investment sets should be `Single`
         assert_eq!(result.len(), 3);
-        assert_eq!(result[0], "C".into());
-        assert_eq!(result[1], "B".into());
-        assert_eq!(result[2], "A".into());
+        assert_eq!(result[0], InvestmentSet::Single("C".into()));
+        assert_eq!(result[1], InvestmentSet::Single("B".into()));
+        assert_eq!(result[2], InvestmentSet::Single("A".into()));
     }
 
     #[rstest]
-    fn test_topo_sort_cyclic_graph(sed_commodity: Commodity) {
+    fn test_solve_investment_order_cyclic_graph(sed_commodity: Commodity) {
         // Create a simple cyclic graph: A -> B -> A
         let mut graph = Graph::new();
 
@@ -539,11 +544,14 @@ mod tests {
         commodities.insert("A".into(), Rc::new(sed_commodity.clone()));
         commodities.insert("B".into(), Rc::new(sed_commodity));
 
-        // This should return an error due to the cycle
-        // The error message should flag commodity B
-        // Note: A is also involved in the cycle, but B is flagged as it is encountered first
         let result = solve_investment_order(&graph, &commodities);
-        assert_error!(result, "Cycle detected in commodity graph for commodity B");
+
+        // Should be a single `Cycle` investment set containing both commodities
+        assert_eq!(result.len(), 1);
+        assert_eq!(
+            result[0],
+            InvestmentSet::Cycle(vec!["A".into(), "B".into()])
+        );
     }
 
     #[rstest]
@@ -570,8 +578,7 @@ mod tests {
         graph.add_edge(node_c, node_d, GraphEdge::Demand);
 
         // Validate the graph at DayNight level
-        let result = validate_commodities_graph(&graph, &commodities, TimeSliceLevel::Annual);
-        assert!(result.is_ok());
+        assert!(validate_commodities_graph(&graph, &commodities, TimeSliceLevel::Annual).is_ok());
     }
 
     #[rstest]
@@ -596,8 +603,10 @@ mod tests {
         graph.add_edge(node_a, node_b, GraphEdge::Primary("process2".into()));
 
         // Validate the graph at DayNight level
-        let result = validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight);
-        assert_error!(result, "SVD commodity A cannot be an input to a process");
+        assert_error!(
+            validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight),
+            "SVD commodity A cannot be an input to a process"
+        );
     }
 
     #[rstest]
@@ -614,8 +623,10 @@ mod tests {
         graph.add_edge(node_a, node_b, GraphEdge::Demand);
 
         // Validate the graph at DayNight level
-        let result = validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight);
-        assert_error!(result, "SVD commodity A is demanded but has no producers");
+        assert_error!(
+            validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight),
+            "SVD commodity A is demanded but has no producers"
+        );
     }
 
     #[rstest]
@@ -633,9 +644,8 @@ mod tests {
         graph.add_edge(node_b, node_a, GraphEdge::Primary("process1".into()));
 
         // Validate the graph at DayNight level
-        let result = validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight);
         assert_error!(
-            result,
+            validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight),
             "SED commodity B may be consumed but has no producers"
         );
     }
@@ -661,9 +671,8 @@ mod tests {
         graph.add_edge(node_a, node_c, GraphEdge::Primary("process2".into()));
 
         // Validate the graph at DayNight level
-        let result = validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight);
         assert_error!(
-            result,
+            validate_commodities_graph(&graph, &commodities, TimeSliceLevel::DayNight),
             "OTH commodity A cannot have both producers and consumers"
         );
     }

src/model.rs

@@ -1,8 +1,9 @@
 //! The model represents the static input data provided by the user.
 use crate::agent::AgentMap;
-use crate::commodity::{CommodityMap, InvestmentSet};
+use crate::commodity::CommodityMap;
 use crate::process::ProcessMap;
 use crate::region::{Region, RegionID, RegionMap};
+use crate::simulation::investment::InvestmentSet;
 use crate::time_slice::TimeSliceInfo;
 use std::collections::HashMap;
 use std::path::PathBuf;

src/simulation/investment.rs

@@ -2,7 +2,7 @@
 use super::optimisation::{DispatchRun, FlowMap};
 use crate::agent::Agent;
 use crate::asset::{Asset, AssetIterator, AssetRef, AssetState};
-use crate::commodity::{Commodity, CommodityID, CommodityMap, InvestmentSet};
+use crate::commodity::{Commodity, CommodityID, CommodityMap};
 use crate::model::Model;
 use crate::output::DataWriter;
 use crate::region::RegionID;
@@ -11,9 +11,10 @@ use crate::time_slice::{TimeSliceID, TimeSliceInfo};
 use crate::units::{Capacity, Dimensionless, Flow, FlowPerCapacity};
 use anyhow::{Result, bail, ensure};
 use indexmap::IndexMap;
-use itertools::{chain, iproduct};
+use itertools::{Itertools, chain, iproduct};
 use log::debug;
 use std::collections::HashMap;
+use std::fmt::Display;
 
 pub mod appraisal;
 use appraisal::coefficients::calculate_coefficients_for_assets;
@@ -25,6 +26,34 @@ type DemandMap = IndexMap<TimeSliceID, Flow>;
 /// Demand for a given combination of commodity, region and time slice
 type AllDemandMap = IndexMap<(CommodityID, RegionID, TimeSliceID), Flow>;
 
+/// Represents a set of commodities which are invested in together.
+#[derive(PartialEq, Debug)]
+pub enum InvestmentSet {
+    /// Assets are selected for a single commodity using `select_assets_for_commodity`
+    Single(CommodityID),
+    /// Assets are selected for a group of commodities which forms a cycle. NOT YET IMPLEMENTED.
+    Cycle(Vec<CommodityID>),
+}
+
+impl InvestmentSet {
+    /// Returns an iterator over the commodity IDs in this investment set
+    pub fn iter(&self) -> impl Iterator<Item = &CommodityID> {
+        match self {
+            InvestmentSet::Single(id) => std::slice::from_ref(id).iter(),
+            InvestmentSet::Cycle(ids) => ids.iter(),
+        }
+    }
+}
+
+impl Display for InvestmentSet {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            InvestmentSet::Single(id) => write!(f, "{id}"),
+            InvestmentSet::Cycle(ids) => write!(f, "[{}]", ids.iter().join(", ")),
+        }
+    }
+}
+
 /// Perform agent investment to determine capacity investment of new assets for next milestone year.
 ///
 /// # Arguments
@@ -50,15 +79,18 @@ pub fn perform_agent_investment(
     let mut all_selected_assets = Vec::new();
 
     // External prices to be used in dispatch optimisation
-    // As investments are performed for each commodity/region combination, the system will become
-    // able to produce its own prices endogenously, so we'll gradually remove these external prices.
+    // Once investments are performed for a commodity, the dispatch system will be able to produce
+    // endogenous prices for that commodity, so we'll gradually remove these external prices.
     let mut external_prices = prices.clone();
 
     for region_id in model.iter_regions() {
-        let investment_order = &model.investment_order[&(region_id.clone(), year)];
-
-        // Iterate over investment sets in the investment order
+        // Keep track of the commodities that have been seen so far. This will be used to apply
+        // balance constraints in the dispatch optimisation - e only apply balance constraints for
+        // commodities that have been seen so far.
         let mut seen_commodities = Vec::new();
+
+        // Iterate over investment sets in the investment order for this region/year
+        let investment_order = &model.investment_order[&(region_id.clone(), year)];
         for investment_set in investment_order {
             // Select assets for the commodity(/ies) of interest
             let selected_assets = match investment_set {
@@ -98,7 +130,7 @@ pub fn perform_agent_investment(
                 external_prices.remove(commodity_id, region_id, time_slice);
             }
 
-            // If no assets have been selected for this region/commodity, skip dispatch optimisation
+            // If no assets have been selected, skip dispatch optimisation
             // **TODO**: this probably means there's no demand for the commodity, which we could
             // presumably preempt
             if selected_assets.is_empty() {