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

muse2/
graph.rs
