Merge pull request #904 from EnergySystemsModellingLab/save_graphs

Save commodity graphs as DOT files

Author: tsmbland

.gitignore

@@ -7,6 +7,7 @@ settings.toml
 
 # Simulation output files
 **/muse2_results
+**/muse2_graphs
 
 # Generated by Cargo
 # will have compiled files and executables

docs/user_guide.md

@@ -5,6 +5,20 @@
 Once you have installed MUSE2, you should be able to run it via the `muse2` command-line program.
 For details of the command-line interface, [see here](./command_line_help.md).
 
+### Visualising commodity graphs
+
+To visualise the structure of your model, you can use the [the `muse2 save-graphs` command] to
+create graphs of commodity/process relationships.
+This command will output a graph for each region/year in the simulation, where nodes are commodities
+and edges are processes.
+Graphs will be saved in [DOT format], which can be visualised locally with [Graphviz], or online
+with [Graphviz online].
+
+[the `muse2 save-graphs` command]: https://energysystemsmodellinglab.github.io/MUSE2/command_line_help.html#muse2-save-graphs
+[DOT format]: https://graphviz.org/doc/info/lang.html
+[Graphviz]: https://graphviz.org/
+[Graphviz online]: https://dreampuf.github.io/GraphvizOnline
+
 ## Modifying the program settings
 
 You can configure the behaviour of MUSE2 with a `settings.toml` file. To edit this file, run:

src/cli.rs

@@ -1,7 +1,8 @@
 //! The command line interface for the simulation.
-use crate::input::load_model;
+use crate::graph::save_commodity_graphs_for_model;
+use crate::input::{load_commodity_graphs, load_model};
 use crate::log;
-use crate::output::{create_output_directory, get_output_dir};
+use crate::output::{create_output_directory, get_graphs_dir, get_output_dir};
 use crate::settings::Settings;
 use ::log::{info, warn};
 use anyhow::{Context, Result};
@@ -26,7 +27,7 @@ struct Cli {
     markdown_help: bool,
 }
 
-/// Options for the run command
+/// Options for the `run` command
 #[derive(Args)]
 pub struct RunOpts {
     /// Directory for output files
@@ -40,6 +41,17 @@ pub struct RunOpts {
     pub debug_model: Option<bool>,
 }
 
+/// Options for the `graph` command
+#[derive(Args)]
+pub struct GraphOpts {
+    /// Directory for graph files
+    #[arg(short, long)]
+    pub output_dir: Option<PathBuf>,
+    /// Whether to overwrite the output directory if it already exists
+    #[arg(long)]
+    pub overwrite: bool,
+}
+
 /// The available commands.
 #[derive(Subcommand)]
 enum Commands {
@@ -62,6 +74,14 @@ enum Commands {
         /// The path to the model directory.
         model_dir: PathBuf,
     },
+    /// Build and output commodity flow graphs for a model.
+    SaveGraphs {
+        /// The path to the model directory.
+        model_dir: PathBuf,
+        /// Other options
+        #[command(flatten)]
+        opts: GraphOpts,
+    },
     /// Manage settings file.
     Settings {
         /// The subcommands for managing the settings file.
@@ -77,6 +97,9 @@ impl Commands {
             Self::Run { model_dir, opts } => handle_run_command(&model_dir, &opts, None),
             Self::Example { subcommand } => subcommand.execute(),
             Self::Validate { model_dir } => handle_validate_command(&model_dir, None),
+            Self::SaveGraphs { model_dir, opts } => {
+                handle_save_graphs_command(&model_dir, &opts, None)
+            }
             Self::Settings { subcommand } => subcommand.execute(),
         }
     }
@@ -178,3 +201,49 @@ pub fn handle_validate_command(model_path: &Path, settings: Option<Settings>) ->
 
     Ok(())
 }
+
+/// Handle the `save-graphs` command.
+pub fn handle_save_graphs_command(
+    model_path: &Path,
+    opts: &GraphOpts,
+    settings: Option<Settings>,
+) -> Result<()> {
+    // Load program settings, if not provided
+    let settings = if let Some(settings) = settings {
+        settings
+    } else {
+        Settings::load().context("Failed to load settings.")?
+    };
+
+    // Get path to output folder
+    let pathbuf: PathBuf;
+    let output_path = if let Some(p) = opts.output_dir.as_deref() {
+        p
+    } else {
+        pathbuf = get_graphs_dir(model_path)?;
+        &pathbuf
+    };
+
+    let overwrite =
+        create_output_directory(output_path, settings.overwrite).with_context(|| {
+            format!(
+                "Failed to create graphs directory: {}",
+                output_path.display()
+            )
+        })?;
+
+    // Initialise program logger (we won't save log files when running this command)
+    log::init(&settings.log_level, None).context("Failed to initialise logging.")?;
+
+    // NB: We have to wait until the logger is initialised to display this warning
+    if overwrite {
+        warn!("Graphs directory will be overwritten");
+    }
+
+    // Load commodity flow graphs and save to file
+    let commodity_graphs = load_commodity_graphs(model_path).context("Failed to build graphs.")?;
+    save_commodity_graphs_for_model(&commodity_graphs, output_path)?;
+    info!("Graphs saved to: {}", output_path.display());
+
+    Ok(())
+}

src/graph.rs

@@ -9,17 +9,21 @@ use indexmap::IndexSet;
 use itertools::{Itertools, iproduct};
 use petgraph::Directed;
 use petgraph::algo::toposort;
+use petgraph::dot::Dot;
 use petgraph::graph::Graph;
 use std::collections::HashMap;
 use std::fmt::Display;
+use std::fs::File;
+use std::io::Write as IoWrite;
+use std::path::Path;
 use strum::IntoEnumIterator;
 
 /// A graph of commodity flows for a given region and year
-type CommoditiesGraph = Graph<GraphNode, GraphEdge, Directed>;
+pub type CommoditiesGraph = Graph<GraphNode, GraphEdge, Directed>;
 
 #[derive(Eq, PartialEq, Clone, Hash)]
 /// A node in the commodity graph
-enum GraphNode {
+pub enum GraphNode {
     /// A node representing a commodity
     Commodity(CommodityID),
     /// A source node for processes that have no inputs
@@ -43,13 +47,22 @@ impl Display for GraphNode {
 
 #[derive(Eq, PartialEq, Clone, Hash)]
 /// An edge in the commodity graph
-enum GraphEdge {
+pub enum GraphEdge {
     /// An edge representing a process
     Process(ProcessID),
     /// An edge representing a service demand
     Demand,
 }
 
+impl Display for GraphEdge {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            GraphEdge::Process(id) => write!(f, "{id}"),
+            GraphEdge::Demand => write!(f, "DEMAND"),
+        }
+    }
+}
+
 /// Creates a directed graph of commodity flows for a given region and year.
 ///
 /// The graph contains nodes for all commodities that may be consumed/produced by processes in the
@@ -118,7 +131,7 @@ fn create_commodities_graph_for_region_year(
     graph
 }
 
-/// Prepares a graph for validation with `validate_commodities_graph`.
+/// Prepares a graph for validation with [`validate_commodities_graph`].
 ///
 /// It takes a base graph produced by `create_commodities_graph_for_region_year`, and modifies it to
 /// account for process availabilities and commodity demands within the given time slice selection,
@@ -196,7 +209,7 @@ fn prepare_commodities_graph_for_validation(
 /// The validation is only performed for commodities with the specified time slice level. For full
 /// validation of all commodities in the model, we therefore need to run this function for all time
 /// slice selections at all time slice levels. This is handled by
-/// `build_and_validate_commodity_graphs_for_model`.
+/// [`validate_commodity_graphs_for_model`].
 fn validate_commodities_graph(
     graph: &CommoditiesGraph,
     commodities: &CommodityMap,
@@ -307,7 +320,26 @@ fn topo_sort_commodities(
     Ok(order)
 }
 
-/// Builds and validates commodity graphs for the entire model.
+/// Builds base commodity graphs for each region and year
+///
+/// These do not take into account demand and process availability
+pub fn build_commodity_graphs_for_model(
+    processes: &ProcessMap,
+    region_ids: &IndexSet<RegionID>,
+    years: &[u32],
+) -> Result<HashMap<(RegionID, u32), CommoditiesGraph>> {
+    let commodity_graphs: HashMap<(RegionID, u32), CommoditiesGraph> =
+        iproduct!(region_ids, years.iter())
+            .map(|(region_id, year)| {
+                let graph = create_commodities_graph_for_region_year(processes, region_id, *year);
+                ((region_id.clone(), *year), graph)
+            })
+            .collect();
+
+    Ok(commodity_graphs)
+}
+
+/// Validates commodity graphs for the entire model.
 ///
 /// This function creates commodity flow graphs for each region/year combination in the model,
 /// validates the graph structure against commodity type rules, and determines the optimal
@@ -337,23 +369,12 @@ fn topo_sort_commodities(
 /// - Any commodity graph contains cycles
 /// - Commodity type rules are violated (e.g., SVD commodities being consumed)
 /// - Demand cannot be satisfied
-pub fn build_and_validate_commodity_graphs_for_model(
+pub fn validate_commodity_graphs_for_model(
+    commodity_graphs: &HashMap<(RegionID, u32), CommoditiesGraph>,
     processes: &ProcessMap,
     commodities: &CommodityMap,
-    region_ids: &IndexSet<RegionID>,
-    years: &[u32],
     time_slice_info: &TimeSliceInfo,
 ) -> Result<HashMap<(RegionID, u32), Vec<CommodityID>>> {
-    // Build base commodity graphs for each region and year
-    // These do not take into account demand and process availability
-    let commodity_graphs: HashMap<(RegionID, u32), CommoditiesGraph> =
-        iproduct!(region_ids, years.iter())
-            .map(|(region_id, year)| {
-                let graph = create_commodities_graph_for_region_year(processes, region_id, *year);
-                ((region_id.clone(), *year), graph)
-            })
-            .collect();
-
     // Determine commodity ordering for each region and year
     let commodity_order: HashMap<(RegionID, u32), Vec<CommodityID>> = commodity_graphs
         .iter()
@@ -366,7 +387,7 @@ pub fn build_and_validate_commodity_graphs_for_model(
         .try_collect()?;
 
     // Validate graphs at all time slice levels (taking into account process availability and demand)
-    for ((region_id, year), base_graph) in &commodity_graphs {
+    for ((region_id, year), base_graph) in commodity_graphs {
         for ts_level in TimeSliceLevel::iter() {
             for ts_selection in time_slice_info.iter_selections_at_level(ts_level) {
                 let graph = prepare_commodities_graph_for_validation(
@@ -392,6 +413,21 @@ pub fn build_and_validate_commodity_graphs_for_model(
     Ok(commodity_order)
 }
 
+/// Saves commodity graphs to file
+///
+/// The graphs are saved as DOT files to the specified output path
+pub fn save_commodity_graphs_for_model(
+    commodity_graphs: &HashMap<(RegionID, u32), CommoditiesGraph>,
+    output_path: &Path,
+) -> Result<()> {
+    for ((region_id, year), graph) in commodity_graphs {
+        let dot = Dot::new(&graph);
+        let mut file = File::create(output_path.join(format!("{region_id}_{year}.dot")))?;
+        write!(file, "{dot}")?;
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/input.rs

@@ -1,8 +1,11 @@
 //! Common routines for handling input data.
 use crate::asset::AssetPool;
-use crate::graph::build_and_validate_commodity_graphs_for_model;
+use crate::graph::{
+    CommoditiesGraph, build_commodity_graphs_for_model, validate_commodity_graphs_for_model,
+};
 use crate::id::{HasID, IDLike};
 use crate::model::{Model, ModelParameters};
+use crate::region::RegionID;
 use crate::units::UnitType;
 use anyhow::{Context, Result, bail, ensure};
 use float_cmp::approx_eq;
@@ -213,11 +216,11 @@ pub fn load_model<P: AsRef<Path>>(model_dir: P) -> Result<(Model, AssetPool)> {
 
     // Build and validate commodity graphs for all regions and years
     // This gives us the commodity order for each region/year which is passed to the model
-    let commodity_order = build_and_validate_commodity_graphs_for_model(
+    let commodity_graphs = build_commodity_graphs_for_model(&processes, &region_ids, years)?;
+    let commodity_order = validate_commodity_graphs_for_model(
+        &commodity_graphs,
         &processes,
         &commodities,
-        &region_ids,
-        years,
         &time_slice_info,
     )?;
 
@@ -238,6 +241,36 @@ pub fn load_model<P: AsRef<Path>>(model_dir: P) -> Result<(Model, AssetPool)> {
     Ok((model, AssetPool::new(assets)))
 }
 
+/// Load commodity flow graphs for a model.
+///
+/// Loads necessary input data and creates a graph of commodity flows for each region and year,
+/// where nodes are commodities and edges are processes.
+///
+/// Graphs validation is NOT performed. This ensures that graphs can be generated even when
+/// validation would fail, which may be helpful for debugging.
+pub fn load_commodity_graphs<P: AsRef<Path>>(
+    model_dir: P,
+) -> Result<HashMap<(RegionID, u32), CommoditiesGraph>> {
+    let model_params = ModelParameters::from_path(&model_dir)?;
+
+    let time_slice_info = read_time_slice_info(model_dir.as_ref())?;
+    let regions = read_regions(model_dir.as_ref())?;
+    let region_ids = regions.keys().cloned().collect();
+    let years = &model_params.milestone_years;
+
+    let commodities = read_commodities(model_dir.as_ref(), &region_ids, &time_slice_info, years)?;
+    let processes = read_processes(
+        model_dir.as_ref(),
+        &commodities,
+        &region_ids,
+        &time_slice_info,
+        years,
+    )?;
+
+    let commodity_graphs = build_commodity_graphs_for_model(&processes, &region_ids, years)?;
+    Ok(commodity_graphs)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/output.rs

@@ -47,7 +47,10 @@ const SOLVER_VALUES_FILE_NAME: &str = "debug_solver.csv";
 /// The output file name for appraisal results
 const APPRAISAL_RESULTS_FILE_NAME: &str = "debug_appraisal_results.csv";
 
-/// Get the model name from the specified directory path
+/// The root folder in which commodity flow graphs will be created
+const GRAPHS_DIRECTORY_ROOT: &str = "muse2_graphs";
+
+/// Get the default output directory for the model
 pub fn get_output_dir(model_dir: &Path) -> Result<PathBuf> {
     // Get the model name from the dir path. This ends up being convoluted because we need to check
     // for all possible errors. Ugh.
@@ -65,6 +68,19 @@ pub fn get_output_dir(model_dir: &Path) -> Result<PathBuf> {
     Ok([OUTPUT_DIRECTORY_ROOT, model_name].iter().collect())
 }
 
+/// Get the default output directory for commodity flow graphs for the model
+pub fn get_graphs_dir(model_dir: &Path) -> Result<PathBuf> {
+    let model_dir = model_dir
+        .canonicalize() // canonicalise in case the user has specified "."
+        .context("Could not resolve path to model")?;
+    let model_name = model_dir
+        .file_name()
+        .context("Model cannot be in root folder")?
+        .to_str()
+        .context("Invalid chars in model dir name")?;
+    Ok([GRAPHS_DIRECTORY_ROOT, model_name].iter().collect())
+}
+
 /// Create a new output directory for the model, optionally overwriting existing data
 ///
 /// # Arguments