use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::PowerFlowReport;
use crate::network::{
ShuntAdm, YBusSparse, build_ybus, build_ybus_with_outages, linear_initial_guess, stamp_shunts,
structural_component_count,
};
pub use crate::scenario::{BatchError, BusOverride, Scenario, uniform_load_scaling};
use crate::solver::{JacobianBackend, PersistentSolver};
use crate::types::{Bus, Line, Transformer};
pub struct BatchSolver {
backend: JacobianBackend,
pool: Option<rayon::ThreadPool>,
}
impl BatchSolver {
pub fn new(backend: JacobianBackend) -> Self {
Self {
backend,
pool: None,
}
}
pub fn with_threads(backend: JacobianBackend, threads: usize) -> Result<Self, BatchError> {
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.map_err(|e| BatchError::ThreadPool(e.to_string()))?;
Ok(Self {
backend,
pool: Some(pool),
})
}
pub fn threads(&self) -> usize {
match &self.pool {
Some(p) => p.current_num_threads(),
None => rayon::current_num_threads(),
}
}
pub fn solve(
&self,
buses_template: &[Bus],
ybus: &YBusSparse,
scenarios: &[Scenario],
tol: f64,
max_iter: usize,
) -> Result<Vec<PowerFlowReport>, BatchError> {
for (i, sc) in scenarios.iter().enumerate() {
if !sc.branch_outages.is_empty() {
return Err(BatchError::OutagesUnsupported { scenario: i });
}
for ov in &sc.bus_overrides {
if ov.bus >= buses_template.len() {
return Err(BatchError::BusOutOfRange {
scenario: i,
bus: ov.bus,
n_buses: buses_template.len(),
});
}
}
}
if scenarios.is_empty() {
return Ok(Vec::new());
}
let next = AtomicUsize::new(0);
let collected: Mutex<Vec<(usize, PowerFlowReport)>> =
Mutex::new(Vec::with_capacity(scenarios.len()));
let n_workers = self.threads().max(1).min(scenarios.len());
let backend = self.backend;
let run = || {
rayon::scope(|s| {
for _ in 0..n_workers {
s.spawn(|_| {
let mut solver = PersistentSolver::new(backend);
loop {
let i = next.fetch_add(1, Ordering::Relaxed);
if i >= scenarios.len() {
break;
}
let mut buses = buses_template.to_vec();
for ov in &scenarios[i].bus_overrides {
let b = &mut buses[ov.bus];
if let Some(p) = ov.p_spec {
b.p_spec = p;
}
if let Some(q) = ov.q_spec {
b.q_spec = q;
}
if let Some(vm) = ov.voltage_mag {
b.voltage_mag = vm;
}
}
linear_initial_guess(&mut buses, ybus);
let (islands, stats) =
solver.solve_with_stats(&mut buses, ybus, tol, max_iter);
collected
.lock()
.expect("batch result mutex poisoned by a panicking worker")
.push((
i,
PowerFlowReport {
buses,
islands,
stats,
dc: None,
linear: None,
outer: None,
},
));
}
});
}
});
};
match &self.pool {
Some(pool) => pool.install(run),
None => run(),
}
let mut out = collected.into_inner().expect("batch result mutex poisoned");
out.sort_by_key(|(i, _)| *i);
Ok(out.into_iter().map(|(_, r)| r).collect())
}
#[allow(clippy::too_many_arguments)]
pub fn solve_contingencies(
&self,
buses_template: &[Bus],
lines: &[Line],
transformers: &[Transformer],
shunts: &[ShuntAdm],
scenarios: &[Scenario],
tol: f64,
max_iter: usize,
) -> Result<Vec<PowerFlowReport>, BatchError> {
let n_branches = lines.len() + transformers.len();
for (i, sc) in scenarios.iter().enumerate() {
for &b in &sc.branch_outages {
if b >= n_branches {
return Err(BatchError::BranchOutOfRange {
scenario: i,
branch: b,
n_branches,
});
}
}
for ov in &sc.bus_overrides {
if ov.bus >= buses_template.len() {
return Err(BatchError::BusOutOfRange {
scenario: i,
bus: ov.bus,
n_buses: buses_template.len(),
});
}
}
}
if scenarios.is_empty() {
return Ok(Vec::new());
}
let n = buses_template.len();
let intact_components = structural_component_count(n, lines, transformers, &[]);
let plans: Vec<(Vec<bool>, bool)> = scenarios
.iter()
.map(|sc| {
let mut outaged = vec![false; n_branches];
for &b in &sc.branch_outages {
outaged[b] = true;
}
let splits = !sc.branch_outages.is_empty()
&& structural_component_count(n, lines, transformers, &outaged)
!= intact_components;
(outaged, splits)
})
.collect();
let next = AtomicUsize::new(0);
let collected: Mutex<Vec<(usize, PowerFlowReport)>> =
Mutex::new(Vec::with_capacity(scenarios.len()));
let n_workers = self.threads().max(1).min(scenarios.len());
let backend = self.backend;
let run = || {
rayon::scope(|s| {
for _ in 0..n_workers {
s.spawn(|_| {
let mut solver = PersistentSolver::new(backend);
loop {
let i = next.fetch_add(1, Ordering::Relaxed);
if i >= scenarios.len() {
break;
}
let (outaged, splits) = &plans[i];
let mut buses = buses_template.to_vec();
for ov in &scenarios[i].bus_overrides {
let b = &mut buses[ov.bus];
if let Some(p) = ov.p_spec {
b.p_spec = p;
}
if let Some(q) = ov.q_spec {
b.q_spec = q;
}
if let Some(vm) = ov.voltage_mag {
b.voltage_mag = vm;
}
}
let mut ybus = if *splits {
let keep_lines: Vec<Line> = lines
.iter()
.enumerate()
.filter(|(k, _)| !outaged[*k])
.map(|(_, l)| l.clone())
.collect();
let keep_transformers: Vec<Transformer> = transformers
.iter()
.enumerate()
.filter(|(k, _)| !outaged[lines.len() + k])
.map(|(_, t)| t.clone())
.collect();
solver.reset();
build_ybus(n, &keep_lines, &keep_transformers)
} else {
solver.invalidate_admittances();
build_ybus_with_outages(n, lines, transformers, outaged)
};
stamp_shunts(&mut ybus, shunts);
let ybus = ybus.finish();
linear_initial_guess(&mut buses, &ybus);
let (islands, stats) =
solver.solve_with_stats(&mut buses, &ybus, tol, max_iter);
if *splits {
solver.reset();
}
collected
.lock()
.expect("batch result mutex poisoned by a panicking worker")
.push((
i,
PowerFlowReport {
buses,
islands,
stats,
dc: None,
linear: None,
outer: None,
},
));
}
});
}
});
};
match &self.pool {
Some(pool) => pool.install(run),
None => run(),
}
let mut out = collected.into_inner().expect("batch result mutex poisoned");
out.sort_by_key(|(i, _)| *i);
Ok(out.into_iter().map(|(_, r)| r).collect())
}
}