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 interested only on the flows direction, which are always the same for all years. So this
73/// function checks if the process can be operating in the target year and region and, if so, it
74/// returns its flows. It considers not only when the process can be commissioned, but also the
75/// lifetime of the process, since a process can be opperating many years after the commission time
76/// window is over. If the process cannot be opperating in the target year and region, None is
77/// returned.
78fn get_flow_for_year(
79    process: &Process,
80    target: (RegionID, u32),
81) -> Option<Rc<IndexMap<CommodityID, ProcessFlow>>> {
82    // If its already in the map, we return it
83    if process.flows.contains_key(&target) {
84        return process.flows.get(&target).cloned();
85    }
86
87    // Otherwise we try to find one that operates in the target year. It is assumed that
88    // parameters are defined in the same (region, year) combinations than flows, at least.
89    let (target_region, target_year) = target;
90    for ((region, year), value) in &process.flows {
91        if *region != target_region {
92            continue;
93        }
94        if year
95            + process
96                .parameters
97                .get(&(region.clone(), *year))
98                .unwrap()
99                .lifetime
100            >= target_year
101        {
102            return Some(value.clone());
103        }
104    }
105    None
106}
107
108/// Creates a directed graph of commodity flows for a given region and year.
109///
110/// The graph contains nodes for all commodities that may be consumed/produced by processes in the
111/// specified region/year. There will be an edge from commodity A to B if there exists a process
112/// that consumes A and produces B.
113///
114/// There are also special `Source` and `Sink` nodes, which are used for processes that have no
115/// inputs or outputs.
116///
117/// The graph does not take into account process availabilities or commodity demands, both of which
118/// can vary by time slice. See `prepare_commodities_graph_for_validation`.
119fn create_commodities_graph_for_region_year(
120    processes: &ProcessMap,
121    region_id: &RegionID,
122    year: u32,
123) -> CommoditiesGraph {
124    let mut graph = Graph::new();
125    let mut commodity_to_node_index = HashMap::new();
126
127    let key = (region_id.clone(), year);
128    for process in processes.values() {
129        let Some(flows) = get_flow_for_year(process, key.clone()) else {
130            // Process doesn't operate in this region/year
131            continue;
132        };
133
134        // Get output nodes for the process
135        let mut outputs: Vec<_> = flows
136            .values()
137            .filter(|flow| flow.direction() == FlowDirection::Output)
138            .map(|flow| GraphNode::Commodity(flow.commodity.id.clone()))
139            .collect();
140
141        // Get input nodes for the process
142        let mut inputs: Vec<_> = flows
143            .values()
144            .filter(|flow| flow.direction() == FlowDirection::Input)
145            .map(|flow| GraphNode::Commodity(flow.commodity.id.clone()))
146            .collect();
147
148        // Use `Source` node if no inputs, `Sink` node if no outputs
149        if inputs.is_empty() {
150            inputs.push(GraphNode::Source);
151        }
152        if outputs.is_empty() {
153            outputs.push(GraphNode::Sink);
154        }
155
156        // Get primary output for the process
157        let primary_output = &process.primary_output;
158
159        // Create edges from all inputs to all outputs
160        // We also create nodes the first time they are encountered
161        for (input, output) in iproduct!(inputs, outputs) {
162            let source_node_index = *commodity_to_node_index
163                .entry(input.clone())
164                .or_insert_with(|| graph.add_node(input.clone()));
165            let target_node_index = *commodity_to_node_index
166                .entry(output.clone())
167                .or_insert_with(|| graph.add_node(output.clone()));
168            let is_primary = match &output {
169                GraphNode::Commodity(commodity_id) => primary_output.as_ref() == Some(commodity_id),
170                _ => false,
171            };
172
173            graph.add_edge(
174                source_node_index,
175                target_node_index,
176                if is_primary {
177                    GraphEdge::Primary(process.id.clone())
178                } else {
179                    GraphEdge::Secondary(process.id.clone())
180                },
181            );
182        }
183    }
184
185    graph
186}
187
188/// Builds base commodity graphs for each region and year
189///
190/// These do not take into account demand and process availability
191pub fn build_commodity_graphs_for_model(
192    processes: &ProcessMap,
193    region_ids: &IndexSet<RegionID>,
194    years: &[u32],
195) -> IndexMap<(RegionID, u32), CommoditiesGraph> {
196    iproduct!(region_ids, years.iter())
197        .map(|(region_id, year)| {
198            let graph = create_commodities_graph_for_region_year(processes, region_id, *year);
199            ((region_id.clone(), *year), graph)
200        })
201        .collect()
202}
203
204/// Gets custom DOT attributes for edges in a commodity graph
205fn get_edge_attributes(_: &CommoditiesGraph, edge_ref: EdgeReference<GraphEdge>) -> String {
206    match edge_ref.weight() {
207        // Use dashed lines for secondary flows
208        GraphEdge::Secondary(_) => "style=dashed".to_string(),
209        // Other edges use default attributes
210        _ => String::new(),
211    }
212}
213
214/// Saves commodity graphs to file
215///
216/// The graphs are saved as DOT files to the specified output path
217pub fn save_commodity_graphs_for_model(
218    commodity_graphs: &IndexMap<(RegionID, u32), CommoditiesGraph>,
219    output_path: &Path,
220) -> Result<()> {
221    for ((region_id, year), graph) in commodity_graphs {
222        let dot = Dot::with_attr_getters(
223            graph,
224            &[],
225            &get_edge_attributes,  // Custom attributes for edges
226            &|_, _| String::new(), // Use default attributes for nodes
227        );
228        let mut file = File::create(output_path.join(format!("{region_id}_{year}.dot")))?;
229        write!(file, "{dot}")?;
230    }
231    Ok(())
232}
muse2/graph.rs

muse2/
graph.rs
