use super::{CommoditiesGraph, GraphEdge, GraphNode};
use crate::commodity::{CommodityMap, CommodityType, PricingStrategy};
use crate::region::RegionID;
use crate::simulation::investment::InvestmentSet;
use anyhow::{Result, ensure};
use highs::{Col, HighsModelStatus, RowProblem, Sense};
use indexmap::IndexMap;
use log::warn;
use petgraph::algo::{condensation, toposort};
use petgraph::graph::Graph;
use petgraph::prelude::NodeIndex;
use petgraph::visit::EdgeRef;
use petgraph::{Directed, Direction};
use std::collections::HashMap;
type InvestmentGraph = Graph<InvestmentSet, GraphEdge, Directed>;
fn solve_investment_order_for_year(
graphs: &IndexMap<(RegionID, u32), CommoditiesGraph>,
commodities: &CommodityMap,
year: u32,
) -> Result<Vec<InvestmentSet>> {
let mut investment_graph = init_investment_graph_for_year(graphs, year, commodities);
investment_graph = compress_cycles(&investment_graph, commodities)?;
let order = toposort(&investment_graph, None).unwrap();
Ok(compute_layers(&investment_graph, &order))
}
fn init_investment_graph_for_year(
graphs: &IndexMap<(RegionID, u32), CommoditiesGraph>,
year: u32,
commodities: &CommodityMap,
) -> InvestmentGraph {
let mut combined = InvestmentGraph::new();
for ((region_id, _), graph) in graphs.iter().filter(|((_, y), _)| *y == year) {
let filtered = graph.filter_map(
|_, n| match n {
GraphNode::Commodity(cid) => {
let kind = &commodities[cid].kind;
matches!(
kind,
CommodityType::ServiceDemand | CommodityType::SupplyEqualsDemand
)
.then_some(GraphNode::Commodity(cid.clone()))
}
_ => None,
},
|_, e| matches!(e, GraphEdge::Primary(_)).then_some(e.clone()),
);
let node_map: HashMap<_, _> = filtered
.node_indices()
.map(|ni| {
let GraphNode::Commodity(cid) = filtered.node_weight(ni).unwrap() else {
unreachable!()
};
(
ni,
combined.add_node(InvestmentSet::Single((cid.clone(), region_id.clone()))),
)
})
.collect();
for e in filtered.edge_references() {
combined.add_edge(
node_map[&e.source()],
node_map[&e.target()],
e.weight().clone(),
);
}
}
combined
}
fn compress_cycles(graph: &InvestmentGraph, commodities: &CommodityMap) -> Result<InvestmentGraph> {
let mut condensed_graph = condensation(graph.clone(), true);
order_sccs(&mut condensed_graph, graph);
for node_weight in condensed_graph.node_weights() {
if node_weight.len() <= 1 {
continue;
}
let offenders: Vec<_> = node_weight
.iter()
.flat_map(|s| s.iter_markets())
.filter(|(cid, _)| {
matches!(
&commodities[cid].pricing_strategy,
PricingStrategy::MarginalCost
| PricingStrategy::MarginalCostAverage
| PricingStrategy::FullCost
| PricingStrategy::FullCostAverage
)
})
.map(|(cid, _)| cid.clone())
.collect();
ensure!(
offenders.is_empty(),
"Cannot use marginal, marginal_average, full or full_average pricing strategies for \
commodities with circular dependencies. Offending commodities: {offenders:?}"
);
}
let mapped = condensed_graph.map(
|_, node_weight| match node_weight.len() {
0 => unreachable!("Condensed graph node must have at least one member"),
1 => node_weight[0].clone(),
_ => InvestmentSet::Cycle(
node_weight
.iter()
.flat_map(|s| s.iter_markets())
.cloned()
.collect(),
),
},
|_, edge_weight| edge_weight.clone(),
);
Ok(mapped)
}
#[allow(clippy::too_many_lines)]
fn order_sccs(
condensed_graph: &mut Graph<Vec<InvestmentSet>, GraphEdge>,
original_graph: &InvestmentGraph,
) {
const EXTERNAL_BIAS: f64 = 0.1;
let node_lookup: HashMap<InvestmentSet, NodeIndex> = original_graph
.node_indices()
.map(|idx| (original_graph.node_weight(idx).unwrap().clone(), idx))
.collect();
for group in condensed_graph.node_indices() {
let scc = condensed_graph.node_weight_mut(group).unwrap();
let n = scc.len();
if n <= 1 {
continue;
}
let original_order = scc.clone();
let original_indices = original_order
.iter()
.map(|set| {
node_lookup
.get(set)
.copied()
.expect("Condensed SCC node must exist in the original graph")
})
.collect::<Vec<_>>();
let mut index_position = HashMap::new();
for (pos, idx) in original_indices.iter().copied().enumerate() {
index_position.insert(idx, pos);
}
let mut penalties = vec![vec![0.0f64; n]; n];
let mut has_external_outgoing = vec![false; n];
for (i, &idx) in original_indices.iter().enumerate() {
for edge in original_graph.edges_directed(idx, Direction::Outgoing) {
if let Some(&j) = index_position.get(&edge.target()) {
penalties[i][j] = 1.0;
} else {
has_external_outgoing[i] = true;
}
}
}
for (j, has_external) in has_external_outgoing.iter().enumerate() {
if *has_external {
for (row_idx, row) in penalties.iter_mut().enumerate() {
if row_idx != j {
row[j] += EXTERNAL_BIAS;
}
}
}
}
let mut problem = RowProblem::default();
let mut vars: Vec<Vec<Option<Col>>> = vec![vec![None; n]; n];
for (i, row) in vars.iter_mut().enumerate() {
for (j, slot) in row.iter_mut().enumerate() {
if i == j {
continue;
}
let cost = penalties[i][j];
*slot = Some(problem.add_integer_column(cost, 0..=1));
}
}
for (i, row) in vars.iter().enumerate() {
for (j, _) in row.iter().enumerate().skip(i + 1) {
let Some(x_ij) = vars[i][j] else { continue };
let Some(x_ji) = vars[j][i] else { continue };
problem.add_row(1.0..=1.0, [(x_ij, 1.0), (x_ji, 1.0)]);
}
}
for (i, row) in vars.iter().enumerate() {
for (j, _) in row.iter().enumerate() {
if i == j {
continue;
}
for (k, _) in vars.iter().enumerate() {
if i == k || j == k {
continue;
}
let Some(x_ij) = vars[i][j] else { continue };
let Some(x_jk) = vars[j][k] else { continue };
let Some(x_ki) = vars[k][i] else { continue };
problem.add_row(..=2.0, [(x_ij, 1.0), (x_jk, 1.0), (x_ki, 1.0)]);
}
}
}
let model = problem.optimise(Sense::Minimise);
let solved = match model.try_solve() {
Ok(solved) => solved,
Err(status) => {
warn!("HiGHS failed while ordering an SCC: {status:?}");
continue;
}
};
if solved.status() != HighsModelStatus::Optimal {
let status = solved.status();
warn!("HiGHS returned a non-optimal status while ordering an SCC: {status:?}");
continue;
}
let solution = solved.get_solution();
let mut wins = vec![0usize; n];
for (i, row) in vars.iter().enumerate() {
for (j, var) in row.iter().enumerate() {
if i == j {
continue;
}
if var.is_some_and(|col| solution[col] > 0.5) {
wins[i] += 1;
}
}
}
let mut order: Vec<usize> = (0..n).collect();
order.sort_by(|&a, &b| wins[b].cmp(&wins[a]).then_with(|| a.cmp(&b)));
*scc = order
.into_iter()
.map(|idx| original_order[idx].clone())
.collect();
}
}
fn compute_layers(graph: &InvestmentGraph, order: &[NodeIndex]) -> Vec<InvestmentSet> {
let mut ranks: HashMap<_, usize> = graph.node_indices().map(|n| (n, 0)).collect();
for &u in order {
let current_rank = ranks[&u];
for v in graph.neighbors_directed(u, Direction::Outgoing) {
if let Some(r) = ranks.get_mut(&v) {
*r = (*r).max(current_rank + 1);
}
}
}
let max_rank = ranks.values().copied().max().unwrap_or(0);
let mut groups: Vec<Vec<InvestmentSet>> = vec![Vec::new(); max_rank + 1];
for node_idx in order {
let rank = ranks[node_idx];
let w = graph.node_weight(*node_idx).unwrap().clone();
groups[rank].push(w);
}
let mut result = Vec::new();
for mut items in groups.into_iter().rev() {
if items.is_empty() {
unreachable!("Should be no gaps in the ranking")
}
if items.len() == 1 {
result.push(items.remove(0));
} else {
result.push(InvestmentSet::Layer(items));
}
}
result
}
pub fn solve_investment_order_for_model(
commodity_graphs: &IndexMap<(RegionID, u32), CommoditiesGraph>,
commodities: &CommodityMap,
years: &[u32],
) -> Result<HashMap<u32, Vec<InvestmentSet>>> {
let mut investment_orders = HashMap::new();
for year in years {
let order = solve_investment_order_for_year(commodity_graphs, commodities, *year)?;
investment_orders.insert(*year, order);
}
Ok(investment_orders)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::commodity::Commodity;
use crate::fixture::{sed_commodity, svd_commodity};
use petgraph::graph::Graph;
use rstest::rstest;
use std::rc::Rc;
#[test]
fn order_sccs_simple_cycle() {
let markets = ["A", "B", "C"].map(|id| InvestmentSet::Single((id.into(), "GBR".into())));
let mut original = InvestmentGraph::new();
let node_indices: Vec<_> = markets
.iter()
.map(|set| original.add_node(set.clone()))
.collect();
for &(src, dst) in &[(1, 0), (2, 1), (0, 2)] {
original.add_edge(
node_indices[src],
node_indices[dst],
GraphEdge::Primary("process1".into()),
);
}
let external = original.add_node(InvestmentSet::Single(("X".into(), "GBR".into())));
original.add_edge(
node_indices[2],
external,
GraphEdge::Primary("process2".into()),
);
let mut condensed: Graph<Vec<InvestmentSet>, GraphEdge> = Graph::new();
let component = condensed.add_node(markets.to_vec());
order_sccs(&mut condensed, &original);
let expected = ["C", "A", "B"]
.map(|id| InvestmentSet::Single((id.into(), "GBR".into())))
.to_vec();
assert_eq!(condensed.node_weight(component).unwrap(), &expected);
}
#[rstest]
fn solve_investment_order_linear_graph(sed_commodity: Commodity, svd_commodity: Commodity) {
let mut graph = Graph::new();
let node_a = graph.add_node(GraphNode::Commodity("A".into()));
let node_b = graph.add_node(GraphNode::Commodity("B".into()));
let node_c = graph.add_node(GraphNode::Commodity("C".into()));
graph.add_edge(node_a, node_b, GraphEdge::Primary("process1".into()));
graph.add_edge(node_b, node_c, GraphEdge::Primary("process2".into()));
let mut commodities = CommodityMap::new();
commodities.insert("A".into(), Rc::new(sed_commodity.clone()));
commodities.insert("B".into(), Rc::new(sed_commodity));
commodities.insert("C".into(), Rc::new(svd_commodity));
let graphs = IndexMap::from([(("GBR".into(), 2020), graph)]);
let result = solve_investment_order_for_year(&graphs, &commodities, 2020).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(result[0], InvestmentSet::Single(("C".into(), "GBR".into())));
assert_eq!(result[1], InvestmentSet::Single(("B".into(), "GBR".into())));
assert_eq!(result[2], InvestmentSet::Single(("A".into(), "GBR".into())));
}
#[rstest]
fn solve_investment_order_cyclic_graph(sed_commodity: Commodity) {
let mut graph = Graph::new();
let node_a = graph.add_node(GraphNode::Commodity("A".into()));
let node_b = graph.add_node(GraphNode::Commodity("B".into()));
graph.add_edge(node_a, node_b, GraphEdge::Primary("process1".into()));
graph.add_edge(node_b, node_a, GraphEdge::Primary("process2".into()));
let mut commodities = CommodityMap::new();
commodities.insert("A".into(), Rc::new(sed_commodity.clone()));
commodities.insert("B".into(), Rc::new(sed_commodity));
let graphs = IndexMap::from([(("GBR".into(), 2020), graph)]);
let result = solve_investment_order_for_year(&graphs, &commodities, 2020).unwrap();
assert_eq!(result.len(), 1);
assert_eq!(
result[0],
InvestmentSet::Cycle(vec![("A".into(), "GBR".into()), ("B".into(), "GBR".into())])
);
}
#[rstest]
fn solve_investment_order_layered_graph(sed_commodity: Commodity, svd_commodity: Commodity) {
let mut graph = Graph::new();
let node_a = graph.add_node(GraphNode::Commodity("A".into()));
let node_b = graph.add_node(GraphNode::Commodity("B".into()));
let node_c = graph.add_node(GraphNode::Commodity("C".into()));
let node_d = graph.add_node(GraphNode::Commodity("D".into()));
graph.add_edge(node_a, node_b, GraphEdge::Primary("process1".into()));
graph.add_edge(node_a, node_c, GraphEdge::Primary("process2".into()));
graph.add_edge(node_b, node_d, GraphEdge::Primary("process3".into()));
graph.add_edge(node_c, node_d, GraphEdge::Primary("process4".into()));
let mut commodities = CommodityMap::new();
commodities.insert("A".into(), Rc::new(sed_commodity.clone()));
commodities.insert("B".into(), Rc::new(sed_commodity.clone()));
commodities.insert("C".into(), Rc::new(sed_commodity));
commodities.insert("D".into(), Rc::new(svd_commodity));
let graphs = IndexMap::from([(("GBR".into(), 2020), graph)]);
let result = solve_investment_order_for_year(&graphs, &commodities, 2020).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(result[0], InvestmentSet::Single(("D".into(), "GBR".into())));
assert_eq!(
result[1],
InvestmentSet::Layer(vec![
InvestmentSet::Single(("B".into(), "GBR".into())),
InvestmentSet::Single(("C".into(), "GBR".into()))
])
);
assert_eq!(result[2], InvestmentSet::Single(("A".into(), "GBR".into())));
}
#[rstest]
fn solve_investment_order_multiple_regions(sed_commodity: Commodity, svd_commodity: Commodity) {
let mut graph = Graph::new();
let node_a = graph.add_node(GraphNode::Commodity("A".into()));
let node_b = graph.add_node(GraphNode::Commodity("B".into()));
let node_c = graph.add_node(GraphNode::Commodity("C".into()));
graph.add_edge(node_a, node_b, GraphEdge::Primary("process1".into()));
graph.add_edge(node_b, node_c, GraphEdge::Primary("process2".into()));
let mut commodities = CommodityMap::new();
commodities.insert("A".into(), Rc::new(sed_commodity.clone()));
commodities.insert("B".into(), Rc::new(sed_commodity));
commodities.insert("C".into(), Rc::new(svd_commodity));
let graphs = IndexMap::from([
(("GBR".into(), 2020), graph.clone()),
(("FRA".into(), 2020), graph),
]);
let result = solve_investment_order_for_year(&graphs, &commodities, 2020).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(
result[0],
InvestmentSet::Layer(vec![
InvestmentSet::Single(("C".into(), "GBR".into())),
InvestmentSet::Single(("C".into(), "FRA".into()))
])
);
assert_eq!(
result[1],
InvestmentSet::Layer(vec![
InvestmentSet::Single(("B".into(), "GBR".into())),
InvestmentSet::Single(("B".into(), "FRA".into()))
])
);
assert_eq!(
result[2],
InvestmentSet::Layer(vec![
InvestmentSet::Single(("A".into(), "GBR".into())),
InvestmentSet::Single(("A".into(), "FRA".into()))
])
);
}
}