muse2/
graph.rs

1//! Module for creating and analysing commodity graphs
2use crate::commodity::CommodityID;
3use crate::process::{FlowDirection, Process, ProcessFlow, ProcessID, ProcessMap};
4use crate::region::RegionID;
5use anyhow::Result;
6use indexmap::{IndexMap, IndexSet};
7use itertools::iproduct;
8use petgraph::Directed;
9use petgraph::dot::Dot;
10use petgraph::graph::{EdgeReference, Graph};
11use std::collections::HashMap;
12use std::fs::File;
13use std::io::Write as IoWrite;
14use std::path::Path;
15use std::rc::Rc;
16
17pub mod investment;
18pub mod validate;
19
20/// A graph of commodity flows for a given region and year
21pub type CommoditiesGraph = Graph<GraphNode, GraphEdge, Directed>;
22
23#[derive(Eq, PartialEq, Clone, Hash, derive_more::Display)]
24/// A node in the commodity graph
25pub enum GraphNode {
26    /// A node representing a commodity
27    #[display("{_0}")]
28    Commodity(CommodityID),
29    /// A source node for processes that have no inputs
30    #[display("SOURCE")]
31    Source,
32    /// A sink node for processes that have no outputs
33    #[display("SINK")]
34    Sink,
35    /// A demand node for commodities with service demands
36    #[display("DEMAND")]
37    Demand,
38}
39
40#[derive(Eq, PartialEq, Clone, Hash, derive_more::Display)]
41/// An edge in the commodity graph
42pub enum GraphEdge {
43    /// An edge representing a primary flow of a process
44    #[display("{_0}")]
45    Primary(ProcessID),
46    /// An edge representing a secondary (non-primary) flow of a process
47    #[display("{_0}")]
48    Secondary(ProcessID),
49    /// An edge representing a service demand
50    #[display("DEMAND")]
51    Demand,
52}
53
54/// Helper function to return a possible flow operating in the requested year
55///
56/// We are only interested in the flow directions, which are constant across years. This
57/// function checks whether the process can be operating in the target region and year and, if so,
58/// returns its flows. It considers both the commission year and the process lifetime, since a
59/// process may operate for years after its commission window. If the process cannot be operating
60/// in the target region/year, `None` is returned.
61fn get_flow_for_year(
62    process: &Process,
63    target: (RegionID, u32),
64) -> Option<Rc<IndexMap<CommodityID, ProcessFlow>>> {
65    // If its already in the map, we return it
66    if process.flows.contains_key(&target) {
67        return process.flows.get(&target).cloned();
68    }
69
70    // Otherwise we try to find one that operates in the target year. It is assumed that
71    // parameters are defined for at least the same (region, year) combinations as flows.
72    let (target_region, target_year) = target;
73    for ((region, year), value) in &process.flows {
74        if *region != target_region {
75            continue;
76        }
77        if year + process.parameters[&(region.clone(), *year)].lifetime >= target_year {
78            return Some(value.clone());
79        }
80    }
81    None
82}
83
84/// Creates a directed graph of commodity flows for a given region and year.
85///
86/// The graph contains nodes for all commodities that may be consumed/produced by processes in the
87/// specified region/year. There will be an edge from commodity A to B if there exists a process
88/// that consumes A and produces B.
89///
90/// There are also special `Source` and `Sink` nodes, which are used for processes that have no
91/// inputs or outputs.
92///
93/// The graph does not take into account process availabilities or commodity demands, both of which
94/// can vary by time slice. See `prepare_commodities_graph_for_validation`.
95fn create_commodities_graph_for_region_year(
96    processes: &ProcessMap,
97    region_id: &RegionID,
98    year: u32,
99) -> CommoditiesGraph {
100    let mut graph = Graph::new();
101    let mut commodity_to_node_index = HashMap::new();
102
103    let key = (region_id.clone(), year);
104    for process in processes.values() {
105        let Some(flows) = get_flow_for_year(process, key.clone()) else {
106            // Process doesn't operate in this region/year
107            continue;
108        };
109
110        // Get output nodes for the process
111        let mut outputs: Vec<_> = flows
112            .values()
113            .filter(|flow| flow.direction() == FlowDirection::Output)
114            .map(|flow| GraphNode::Commodity(flow.commodity.id.clone()))
115            .collect();
116
117        // Get input nodes for the process
118        let mut inputs: Vec<_> = flows
119            .values()
120            .filter(|flow| flow.direction() == FlowDirection::Input)
121            .map(|flow| GraphNode::Commodity(flow.commodity.id.clone()))
122            .collect();
123
124        // Use `Source` node if no inputs, `Sink` node if no outputs
125        if inputs.is_empty() {
126            inputs.push(GraphNode::Source);
127        }
128        if outputs.is_empty() {
129            outputs.push(GraphNode::Sink);
130        }
131
132        // Get primary output for the process
133        let primary_output = &process.primary_output;
134
135        // Create edges from all inputs to all outputs
136        // We also create nodes the first time they are encountered
137        for (input, output) in iproduct!(inputs, outputs) {
138            let source_node_index = *commodity_to_node_index
139                .entry(input.clone())
140                .or_insert_with(|| graph.add_node(input.clone()));
141            let target_node_index = *commodity_to_node_index
142                .entry(output.clone())
143                .or_insert_with(|| graph.add_node(output.clone()));
144            let is_primary = match &output {
145                GraphNode::Commodity(commodity_id) => primary_output.as_ref() == Some(commodity_id),
146                _ => false,
147            };
148
149            graph.add_edge(
150                source_node_index,
151                target_node_index,
152                if is_primary {
153                    GraphEdge::Primary(process.id.clone())
154                } else {
155                    GraphEdge::Secondary(process.id.clone())
156                },
157            );
158        }
159    }
160
161    graph
162}
163
164/// Builds base commodity graphs for each region and year
165///
166/// These do not take into account demand and process availability
167pub fn build_commodity_graphs_for_model(
168    processes: &ProcessMap,
169    region_ids: &IndexSet<RegionID>,
170    years: &[u32],
171) -> IndexMap<(RegionID, u32), CommoditiesGraph> {
172    iproduct!(region_ids, years.iter())
173        .map(|(region_id, year)| {
174            let graph = create_commodities_graph_for_region_year(processes, region_id, *year);
175            ((region_id.clone(), *year), graph)
176        })
177        .collect()
178}
179
180/// Gets custom DOT attributes for edges in a commodity graph
181fn get_edge_attributes(_: &CommoditiesGraph, edge_ref: EdgeReference<GraphEdge>) -> String {
182    match edge_ref.weight() {
183        // Use dashed lines for secondary flows
184        GraphEdge::Secondary(_) => "style=dashed".to_string(),
185        // Other edges use default attributes
186        _ => String::new(),
187    }
188}
189
190/// Saves commodity graphs to file
191///
192/// The graphs are saved as DOT files to the specified output path
193pub fn save_commodity_graphs_for_model(
194    commodity_graphs: &IndexMap<(RegionID, u32), CommoditiesGraph>,
195    output_path: &Path,
196) -> Result<()> {
197    for ((region_id, year), graph) in commodity_graphs {
198        let dot = Dot::with_attr_getters(
199            graph,
200            &[],
201            &get_edge_attributes,  // Custom attributes for edges
202            &|_, _| String::new(), // Use default attributes for nodes
203        );
204        let mut file = File::create(output_path.join(format!("{region_id}_{year}.dot")))?;
205        write!(file, "{dot}")?;
206    }
207    Ok(())
208}
muse2/graph.rs

muse2/
graph.rs
