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::sync::Arc;
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 in a specific region
27 #[display("{_0} | {_1}")]
28 Commodity(CommodityID, RegionID),
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 installed in the given region
44 #[display("{_0} | {_1}")]
45 Primary(ProcessID, RegionID),
46 /// An edge representing a secondary (non-primary) flow of a process installed in the given region
47 #[display("{_0} | {_1}")]
48 Secondary(ProcessID, RegionID),
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<Arc<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 year, covering all regions.
85///
86/// The graph contains nodes for all `(commodity, region)` pairs that may be consumed/produced
87/// by processes in the specified year across all regions. There will be an edge from node
88/// `(A, R)` to `(B, R)` if there exists a process in region R that consumes commodity A and
89/// produces commodity B. Processes currently only consume and produce commodities in their own
90/// region, so there are no edges between nodes in different regions.
91///
92/// There are also special `Source` and `Sink` nodes, which are used for processes that have no
93/// inputs or outputs.
94///
95/// The graph does not take into account process availabilities or commodity demands, both of which
96/// can vary by time slice. See `prepare_commodities_graph_for_validation`.
97fn create_commodities_graph_for_year(
98 processes: &ProcessMap,
99 region_ids: &IndexSet<RegionID>,
100 year: u32,
101) -> CommoditiesGraph {
102 let mut graph = Graph::new();
103 let mut node_index = HashMap::new();
104
105 for region_id in region_ids {
106 let key = (region_id.clone(), year);
107 for process in processes.values() {
108 let Some(flows) = get_flow_for_year(process, key.clone()) else {
109 // Process doesn't operate in this region/year
110 continue;
111 };
112
113 // Get output nodes for the process
114 // For now, we assume processes always output commodities in their own region
115 let mut outputs: Vec<_> = flows
116 .values()
117 .filter(|flow| flow.direction() == FlowDirection::Output)
118 .map(|flow| GraphNode::Commodity(flow.commodity.id.clone(), region_id.clone()))
119 .collect();
120
121 // Get input nodes for the process
122 // For now, we assume processes always input commodities in their own region
123 let mut inputs: Vec<_> = flows
124 .values()
125 .filter(|flow| flow.direction() == FlowDirection::Input)
126 .map(|flow| GraphNode::Commodity(flow.commodity.id.clone(), region_id.clone()))
127 .collect();
128
129 // Use `Source` node if no inputs, `Sink` node if no outputs
130 if inputs.is_empty() {
131 inputs.push(GraphNode::Source);
132 }
133 if outputs.is_empty() {
134 outputs.push(GraphNode::Sink);
135 }
136
137 // Get primary output for the process
138 let primary_output = &process.primary_output;
139
140 // Create edges from all inputs to all outputs.
141 // Edges carry the asset's installing region so the availability filter in
142 // validation can look up `(process_id, asset_region, year)`.
143 for (input, output) in iproduct!(inputs, outputs) {
144 let source_node_index = *node_index
145 .entry(input.clone())
146 .or_insert_with(|| graph.add_node(input.clone()));
147 let target_node_index = *node_index
148 .entry(output.clone())
149 .or_insert_with(|| graph.add_node(output.clone()));
150 let is_primary = match &output {
151 GraphNode::Commodity(commodity_id, _) => {
152 primary_output.as_ref() == Some(commodity_id)
153 }
154 _ => false,
155 };
156
157 graph.add_edge(
158 source_node_index,
159 target_node_index,
160 if is_primary {
161 GraphEdge::Primary(process.id.clone(), region_id.clone())
162 } else {
163 GraphEdge::Secondary(process.id.clone(), region_id.clone())
164 },
165 );
166 }
167 }
168 }
169
170 graph
171}
172
173/// Builds base commodity graphs for each year, covering all regions.
174///
175/// These do not take into account demand and process availability.
176pub fn build_commodity_graphs_for_model(
177 processes: &ProcessMap,
178 region_ids: &IndexSet<RegionID>,
179 years: &[u32],
180) -> IndexMap<u32, CommoditiesGraph> {
181 years
182 .iter()
183 .map(|&year| {
184 let graph = create_commodities_graph_for_year(processes, region_ids, year);
185 (year, graph)
186 })
187 .collect()
188}
189
190/// Gets custom DOT attributes for edges in a commodity graph
191fn get_edge_attributes(_: &CommoditiesGraph, edge_ref: EdgeReference<GraphEdge>) -> String {
192 match edge_ref.weight() {
193 // Use dashed lines for secondary flows
194 GraphEdge::Secondary(..) => "style=dashed".to_string(),
195 // Other edges use default attributes
196 _ => String::new(),
197 }
198}
199
200/// Saves commodity graphs to file
201///
202/// The graphs are saved as DOT files to the specified output path, one file per year.
203pub fn save_commodity_graphs_for_model(
204 commodity_graphs: &IndexMap<u32, CommoditiesGraph>,
205 output_path: &Path,
206) -> Result<()> {
207 for (year, graph) in commodity_graphs {
208 let dot = Dot::with_attr_getters(
209 graph,
210 &[],
211 &get_edge_attributes, // Custom attributes for edges
212 &|_, _| String::new(), // Use default attributes for nodes
213 );
214 let mut file = File::create(output_path.join(format!("{year}.dot")))?;
215 write!(file, "{dot}")?;
216 }
217 Ok(())
218}