use cobre_core::{AffineBound, CoefficientRef, EntityId, LinearTerm};
pub(crate) enum SideTerm {
Variable(LinearTerm),
Constant(f64),
Param(f64, EntityId),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RelOp {
Le,
Ge,
Eq,
}
pub(crate) fn normalize(
lhs: Vec<SideTerm>,
rhs: Vec<SideTerm>,
op: RelOp,
) -> (Vec<LinearTerm>, Option<AffineBound>, Option<AffineBound>) {
let mut variable_terms: Vec<LinearTerm> = Vec::with_capacity(lhs.len() + rhs.len());
let mut constant = 0.0_f64;
let mut param_terms: Vec<(f64, EntityId)> = Vec::new();
for term in rhs {
match term {
SideTerm::Variable(lt) => variable_terms.push(negate_variable_term(lt)),
SideTerm::Constant(v) => constant += v,
SideTerm::Param(coef, id) => param_terms.push((coef, id)),
}
}
for term in lhs {
match term {
SideTerm::Variable(lt) => variable_terms.push(lt),
SideTerm::Constant(v) => constant -= v,
SideTerm::Param(coef, id) => param_terms.push((-coef, id)),
}
}
let variable_terms = fold_variable_terms(variable_terms);
let bound = AffineBound {
constant,
terms: param_terms,
};
let (lower, upper) = match op {
RelOp::Le => (None, Some(bound)),
RelOp::Ge => (Some(bound), None),
RelOp::Eq => (Some(bound.clone()), Some(bound)),
};
(variable_terms, lower, upper)
}
fn fold_variable_terms(terms: Vec<LinearTerm>) -> Vec<LinearTerm> {
let mut merged: Vec<LinearTerm> = Vec::with_capacity(terms.len());
for term in terms {
let slot = merged.iter_mut().find(|existing| {
existing.variable == term.variable
&& coefficient_kind_matches(&existing.coefficient, &term.coefficient)
});
match slot {
Some(existing) => merge_into(existing, &term),
None => merged.push(term),
}
}
merged.retain(|term| !is_zero_literal(&term.coefficient));
merged
}
fn negate_variable_term(mut lt: LinearTerm) -> LinearTerm {
match &mut lt.coefficient {
CoefficientRef::Literal(v) => *v = -*v,
CoefficientRef::Parameter(_) => lt.scale = -lt.scale,
}
lt
}
fn coefficient_kind_matches(a: &CoefficientRef, b: &CoefficientRef) -> bool {
match (a, b) {
(CoefficientRef::Literal(_), CoefficientRef::Literal(_)) => true,
(CoefficientRef::Parameter(id_a), CoefficientRef::Parameter(id_b)) => id_a == id_b,
(CoefficientRef::Literal(_), CoefficientRef::Parameter(_))
| (CoefficientRef::Parameter(_), CoefficientRef::Literal(_)) => false,
}
}
fn merge_into(existing: &mut LinearTerm, incoming: &LinearTerm) {
match incoming.coefficient {
CoefficientRef::Literal(b) => {
if let CoefficientRef::Literal(a) = &mut existing.coefficient {
*a = *a * existing.scale + b * incoming.scale;
existing.scale = 1.0;
}
}
CoefficientRef::Parameter(_) => existing.scale += incoming.scale,
}
}
fn is_zero_literal(coef: &CoefficientRef) -> bool {
matches!(coef, CoefficientRef::Literal(v) if v.to_bits() == 0.0_f64.to_bits())
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic)]
mod tests {
use super::*;
use cobre_core::VariableRef;
fn hg(id: i32) -> VariableRef {
VariableRef::HydroGeneration {
hydro_id: EntityId(id),
block_id: None,
bus_id: None,
}
}
fn tg(id: i32) -> VariableRef {
VariableRef::ThermalGeneration {
thermal_id: EntityId(id),
block_id: None,
}
}
fn effective(term: &LinearTerm) -> f64 {
match term.coefficient {
CoefficientRef::Literal(v) => v * term.scale,
CoefficientRef::Parameter(_) => panic!("expected a literal coefficient"),
}
}
#[test]
fn normalize_rhs_variable_sign_flips_into_lhs() {
let lhs = vec![SideTerm::Variable(LinearTerm::literal(1.0, tg(5)))];
let rhs = vec![SideTerm::Variable(LinearTerm::literal(0.87, hg(140)))];
let (terms, lower, upper) = normalize(lhs, rhs, RelOp::Le);
assert_eq!(terms.len(), 2);
assert_eq!(terms[0].variable, hg(140));
assert!((effective(&terms[0]) - (-0.87)).abs() < f64::EPSILON);
assert_eq!(terms[1].variable, tg(5));
assert!((effective(&terms[1]) - 1.0).abs() < f64::EPSILON);
assert_eq!(lower, None);
assert_eq!(
upper,
Some(AffineBound {
constant: 0.0,
terms: vec![]
})
);
}
#[test]
fn normalize_unmerged_rhs_literal_negates_coefficient_not_scale() {
let rhs = vec![SideTerm::Variable(LinearTerm::literal(0.87, hg(140)))];
let (terms, _, _) = normalize(vec![], rhs, RelOp::Le);
assert_eq!(terms.len(), 1);
assert_eq!(terms[0].coefficient, CoefficientRef::Literal(-0.87));
assert!((terms[0].scale - 1.0).abs() < f64::EPSILON);
}
#[test]
fn normalize_same_variable_repeat_merges_and_drops_at_zero() {
let lhs = vec![
SideTerm::Variable(LinearTerm::literal(1.0, hg(0))),
SideTerm::Variable(LinearTerm::literal(-1.0, hg(0))),
];
let rhs = vec![SideTerm::Constant(5.0)];
let (terms, lower, upper) = normalize(lhs, rhs, RelOp::Le);
assert!(terms.is_empty(), "net-zero column must be dropped");
assert_eq!(lower, None);
assert_eq!(
upper,
Some(AffineBound {
constant: 5.0,
terms: vec![]
})
);
}
#[test]
fn normalize_same_variable_repeat_merges_to_sum() {
let lhs = vec![
SideTerm::Variable(LinearTerm::literal(2.0, hg(0))),
SideTerm::Variable(LinearTerm::literal(3.0, hg(0))),
];
let (terms, _, _) = normalize(lhs, vec![], RelOp::Le);
assert_eq!(terms.len(), 1);
assert!((effective(&terms[0]) - 5.0).abs() < f64::EPSILON);
}
#[test]
fn normalize_parameter_and_literal_on_same_variable_stay_separate() {
let lhs = vec![
SideTerm::Variable(LinearTerm::literal(1.0, hg(0))),
SideTerm::Variable(LinearTerm::parameter(EntityId(9), 1.0, hg(0))),
];
let (terms, _, _) = normalize(lhs, vec![], RelOp::Le);
assert_eq!(terms.len(), 2, "distinct coefficient kinds must not merge");
}
#[test]
fn normalize_same_parameter_on_same_variable_merges_scale() {
let lhs = vec![
SideTerm::Variable(LinearTerm::parameter(EntityId(9), 2.0, hg(0))),
SideTerm::Variable(LinearTerm::parameter(EntityId(9), 3.0, hg(0))),
];
let (terms, _, _) = normalize(lhs, vec![], RelOp::Le);
assert_eq!(terms.len(), 1);
assert_eq!(terms[0].coefficient, CoefficientRef::Parameter(EntityId(9)));
assert!((terms[0].scale - 5.0).abs() < f64::EPSILON);
}
#[test]
fn normalize_rhs_param_scalar_assigns_lower_as_single() {
let lhs = vec![SideTerm::Variable(LinearTerm::literal(1.0, hg(3)))];
let rhs = vec![SideTerm::Param(1.0, EntityId(42))];
let (terms, lower, upper) = normalize(lhs, rhs, RelOp::Ge);
assert_eq!(terms.len(), 1);
assert_eq!(upper, None);
assert_eq!(lower, Some(AffineBound::single(EntityId(42))));
}
#[test]
fn normalize_lhs_non_variable_terms_negate_into_bound() {
let lhs = vec![SideTerm::Constant(10.0), SideTerm::Param(2.0, EntityId(7))];
let (terms, lower, upper) = normalize(lhs, vec![], RelOp::Le);
assert!(terms.is_empty());
assert_eq!(lower, None);
assert_eq!(
upper,
Some(AffineBound {
constant: -10.0,
terms: vec![(-2.0, EntityId(7))]
})
);
}
#[test]
fn normalize_eq_assigns_both_endpoints_equal() {
let rhs = vec![SideTerm::Constant(3.0)];
let (_, lower, upper) = normalize(vec![], rhs, RelOp::Eq);
let expected = Some(AffineBound {
constant: 3.0,
terms: vec![],
});
assert_eq!(lower, expected);
assert_eq!(upper, expected);
}
}