use crate::nl_quadratic::{
FactoredQuadratic, Quad2, QuadForm, QuadHessian, is_expanded_quadratic, is_trivially_zero,
quad_form_readout, recognize_expr, recognize_factored_quadratic,
};
use crate::nl_tape::{HybridTape, Tape, hybrid_supported};
use pounce_common::types::{Index, Number, lower_bound_present, upper_bound_present};
use pounce_nlp::constant_derivatives::{DerivativeProof, DerivativeProofs};
use pounce_nlp::quadratic::{QuadraticStructure, SquareTerm};
use pounce_nlp::tnlp::{
BoundsInfo, IDX_NAMES, IndexStyle, IpoptCq, IpoptData, Linearity, MetaData, NlpInfo,
ScalingRequest, Solution, SparsityRequest, StartingPoint, TNLP,
};
use std::cell::RefCell;
use std::collections::{BTreeMap, BTreeSet};
use std::path::Path;
use std::rc::Rc;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub enum Expr {
Const(Number),
Var(usize),
Binary(BinOp, Box<Expr>, Box<Expr>),
Unary(UnaryOp, Box<Expr>),
Sum(Vec<Expr>),
Cse(Arc<Expr>),
Funcall { id: usize, args: Vec<FuncallArg> },
Compare(CmpOp, Box<Expr>, Box<Expr>),
And(Box<Expr>, Box<Expr>),
Or(Box<Expr>, Box<Expr>),
Not(Box<Expr>),
Cond {
cond: Box<Expr>,
then_: Box<Expr>,
else_: Box<Expr>,
},
MinList(Vec<Expr>),
MaxList(Vec<Expr>),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CmpOp {
Lt,
Le,
Eq,
Ge,
Gt,
Ne,
}
#[derive(Debug, Clone)]
pub enum FuncallArg {
Real(Expr),
Str(String),
}
#[derive(Debug, Clone)]
pub struct ImportedFunc {
pub id: usize,
pub kind: usize,
pub nargs: i64,
pub name: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinOp {
Add,
Sub,
Mul,
Div,
Pow,
Atan2,
CEntropy,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UnaryOp {
Neg,
Sqrt,
Log,
Exp,
Abs,
Sin,
Cos,
Log10,
Tan,
Atan,
Acos,
Sinh,
Cosh,
Tanh,
Asin,
Acosh,
Asinh,
Atanh,
Erf,
XLogX,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NlCounts {
pub nl_cons: usize,
pub nl_objs: usize,
pub nl_vars_cons: usize,
pub nl_vars_objs: usize,
pub nl_vars_both: usize,
}
impl NlCounts {
pub fn nonlinear_vars(&self) -> usize {
self.nl_vars_cons
.saturating_add(self.nl_vars_objs)
.saturating_sub(self.nl_vars_both)
}
}
#[derive(Debug, Clone)]
pub enum NlBody {
Tree(Expr),
Quad(Box<QuadBody>),
}
#[derive(Debug, Clone)]
pub struct QuadBody {
pub form: Quad2,
pub vars: Vec<u32>,
pub src: std::ops::Range<usize>,
pub depth: u32,
}
impl NlBody {
pub fn is_trivially_zero(&self) -> bool {
match self {
NlBody::Tree(e) => matches!(e, Expr::Const(c) if *c == 0.0),
NlBody::Quad(_) => false,
}
}
pub fn quad(&self) -> Option<&Quad2> {
match self {
NlBody::Tree(_) => None,
NlBody::Quad(q) => Some(&q.form),
}
}
pub fn tree(&self) -> Option<&Expr> {
match self {
NlBody::Tree(e) => Some(e),
NlBody::Quad(_) => None,
}
}
pub fn analyze_quadratic(&self) -> Option<QuadHessian> {
self.analyze_quadratic_full().map(|(h, _, _)| h)
}
pub fn analyze_quadratic_full(&self) -> Option<QuadForm> {
match self {
NlBody::Tree(e) => {
let form = recognize_expr(e)?;
(!form.lost_terms()).then(|| quad_form_readout(&form))
}
NlBody::Quad(q) => (!q.form.lost_terms()).then(|| quad_form_readout(&q.form)),
}
}
pub fn quad_terms_dropped(&self) -> bool {
match self {
NlBody::Tree(e) => recognize_expr(e).is_some_and(|f| f.lost_terms()),
NlBody::Quad(q) => q.form.lost_terms(),
}
}
pub fn admitted_quad_form(&self) -> Option<QuadForm> {
match self {
NlBody::Tree(e) => {
if is_trivially_zero(e) || !is_expanded_quadratic(e) {
return None;
}
let form = recognize_expr(e)?;
(!form.lost_terms()).then(|| quad_form_readout(&form))
}
NlBody::Quad(q) => (!q.form.lost_terms()).then(|| quad_form_readout(&q.form)),
}
}
pub fn admitted_factored_form(&self) -> Option<FactoredQuadratic> {
match self {
NlBody::Tree(e) => {
if is_trivially_zero(e) || is_expanded_quadratic(e) {
return None;
}
recognize_factored_quadratic(e)
}
NlBody::Quad(_) => None,
}
}
pub fn provably_affine(&self) -> Option<bool> {
match self {
NlBody::Tree(e) => {
if is_trivially_zero(e) {
return Some(true);
}
affine_from_form(&recognize_expr(e)?)
}
NlBody::Quad(q) => affine_from_form(&q.form),
}
}
pub fn collect_vars(&self, out: &mut BTreeSet<usize>) {
match self {
NlBody::Tree(e) => collect_vars(e, out),
NlBody::Quad(q) => out.extend(q.vars.iter().map(|&v| v as usize)),
}
}
}
fn push_body_form(quad: &mut QuadraticStructure, body: &NlBody) -> Option<u32> {
if let Some((h, lin, c)) = body.admitted_quad_form() {
return Some(quad.push_form(&h, &lin, c));
}
let fq = body.admitted_factored_form()?;
let terms: Vec<SquareTerm<'_>> = fq
.squares
.iter()
.map(|t| SquareTerm {
weight: t.weight,
coefs: &t.coefs,
constant: t.constant,
})
.collect();
quad.push_factored_form(&terms, &fq.linear, fq.constant)
}
fn affine_from_form(q: &Quad2) -> Option<bool> {
if q.quadratic().is_empty() {
(!q.lost_terms()).then_some(true)
} else {
Some(false)
}
}
impl From<Expr> for NlBody {
fn from(e: Expr) -> Self {
NlBody::Tree(e)
}
}
#[derive(Debug, Clone)]
pub struct NlProblem {
pub n: usize,
pub m: usize,
pub num_obj: usize,
pub minimize: bool,
pub obj_nonlinear: NlBody,
pub obj_linear: Vec<(usize, Number)>,
pub obj_constant: Number,
pub con_nonlinear: Vec<NlBody>,
pub con_linear: Vec<Vec<(usize, Number)>>,
pub x_l: Vec<Number>,
pub x_u: Vec<Number>,
pub g_l: Vec<Number>,
pub g_u: Vec<Number>,
pub x0: Vec<Number>,
pub lambda0: Vec<Number>,
pub suffixes: NlSuffixes,
pub ampl_options: Vec<i64>,
pub nl_counts: Option<NlCounts>,
pub imported_funcs: Vec<ImportedFunc>,
pub var_names: Vec<String>,
pub con_names: Vec<String>,
pub src: Option<Arc<String>>,
pub cse_bodies: Vec<Arc<Expr>>,
}
#[derive(Debug, Clone)]
pub struct NlProblemParts {
pub minimize: bool,
pub objective: Expr,
pub obj_constant: Number,
pub constraints: Vec<Expr>,
pub x_l: Vec<Number>,
pub x_u: Vec<Number>,
pub x0: Vec<Number>,
pub g_l: Vec<Number>,
pub g_u: Vec<Number>,
pub var_names: Vec<String>,
pub con_names: Vec<String>,
}
impl NlProblem {
pub fn from_expressions(parts: NlProblemParts) -> Result<NlProblem, String> {
let NlProblemParts {
minimize,
objective,
obj_constant,
constraints,
x_l,
x_u,
x0,
g_l,
g_u,
var_names,
con_names,
} = parts;
let n = x_l.len();
let m = constraints.len();
let check = |name: &str, got: usize, want: usize| -> Result<(), String> {
if got == want {
Ok(())
} else {
Err(format!(
"from_expressions: {name} has length {got}, expected {want}"
))
}
};
check("x_u", x_u.len(), n)?;
check("x0", x0.len(), n)?;
check("g_l", g_l.len(), m)?;
check("g_u", g_u.len(), m)?;
if !var_names.is_empty() {
check("var_names", var_names.len(), n)?;
}
if !con_names.is_empty() {
check("con_names", con_names.len(), m)?;
}
let mut x_l = x_l;
let mut x_u = x_u;
let mut g_l = g_l;
let mut g_u = g_u;
for (i, v) in x_l.iter_mut().enumerate() {
*v = finite_bound_or_err(&format!("x_l[{i}]"), *v, true)
.map_err(|e| format!("from_expressions: {e}"))?;
}
for (i, v) in x_u.iter_mut().enumerate() {
*v = finite_bound_or_err(&format!("x_u[{i}]"), *v, false)
.map_err(|e| format!("from_expressions: {e}"))?;
}
for (i, v) in g_l.iter_mut().enumerate() {
*v = finite_bound_or_err(&format!("g_l[{i}]"), *v, true)
.map_err(|e| format!("from_expressions: {e}"))?;
}
for (i, v) in g_u.iter_mut().enumerate() {
*v = finite_bound_or_err(&format!("g_u[{i}]"), *v, false)
.map_err(|e| format!("from_expressions: {e}"))?;
}
for (i, v) in x0.iter().enumerate() {
finite_or_err(&format!("x0[{i}]"), *v).map_err(|e| format!("from_expressions: {e}"))?;
}
finite_or_err("obj_constant", obj_constant)
.map_err(|e| format!("from_expressions: {e}"))?;
let mut seen: std::collections::HashSet<*const Expr> = std::collections::HashSet::new();
validate_expr(&objective, n, &mut seen).map_err(|e| format!("objective {e}"))?;
for (i, c) in constraints.iter().enumerate() {
validate_expr(c, n, &mut seen).map_err(|e| format!("constraint {i} {e}"))?;
}
Ok(NlProblem {
n,
m,
num_obj: 1,
minimize,
obj_nonlinear: NlBody::Tree(objective),
obj_linear: Vec::new(),
obj_constant,
con_nonlinear: constraints.into_iter().map(NlBody::Tree).collect(),
con_linear: vec![Vec::new(); m],
x_l,
x_u,
g_l,
g_u,
x0,
lambda0: vec![0.0; m],
suffixes: NlSuffixes::default(),
imported_funcs: Vec::new(),
ampl_options: Vec::new(),
nl_counts: None,
var_names,
con_names,
src: None,
cse_bodies: Vec::new(),
})
}
pub fn obj_expr(&self) -> std::borrow::Cow<'_, Expr> {
self.body_expr(&self.obj_nonlinear, "objective")
}
pub fn con_expr(&self, k: usize) -> std::borrow::Cow<'_, Expr> {
self.body_expr(&self.con_nonlinear[k], "constraint")
}
fn body_expr<'a>(&'a self, body: &'a NlBody, what: &str) -> std::borrow::Cow<'a, Expr> {
match body {
NlBody::Tree(e) => std::borrow::Cow::Borrowed(e),
NlBody::Quad(q) => {
let src = self
.src
.as_deref()
.unwrap_or_else(|| panic!("{what} body was recognized but no source is kept"));
std::borrow::Cow::Owned(
parse_body_fragment(&src[q.src.clone()], self.n, &self.cse_bodies)
.unwrap_or_else(|e| panic!("re-parsing a recognized {what} body: {e}")),
)
}
}
}
}
fn validate_expr(
e: &Expr,
n: usize,
seen: &mut std::collections::HashSet<*const Expr>,
) -> Result<(), String> {
match e {
Expr::Const(_) => Ok(()),
Expr::Var(i) => {
if *i < n {
Ok(())
} else {
Err(format!("references Var({i}) but n = {n}"))
}
}
Expr::Binary(_, a, b) | Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
validate_expr(a, n, seen)?;
validate_expr(b, n, seen)
}
Expr::Unary(_, a) | Expr::Not(a) => validate_expr(a, n, seen),
Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
for a in args {
validate_expr(a, n, seen)?;
}
Ok(())
}
Expr::Cond { cond, then_, else_ } => {
validate_expr(cond, n, seen)?;
validate_expr(then_, n, seen)?;
validate_expr(else_, n, seen)
}
Expr::Cse(body) => {
if seen.insert(Arc::as_ptr(body)) {
validate_expr(body, n, seen)
} else {
Ok(())
}
}
Expr::Funcall { id, .. } => Err(format!(
"references AMPL imported function id {id}, which this path cannot \
resolve: a problem built from expressions has no F-segment \
declarations to bind it to. Load such a model with read_nl or \
parse_nl_text instead."
)),
}
}
#[derive(Debug, Clone, Default)]
pub struct NlSuffixes {
pub var_int: BTreeMap<String, Vec<Index>>,
pub con_int: BTreeMap<String, Vec<Index>>,
pub obj_int: BTreeMap<String, Vec<Index>>,
pub problem_int: BTreeMap<String, Index>,
pub var_real: BTreeMap<String, Vec<Number>>,
pub con_real: BTreeMap<String, Vec<Number>>,
pub obj_real: BTreeMap<String, Vec<Number>>,
pub problem_real: BTreeMap<String, Number>,
}
pub fn read_nl_file(path: &Path) -> Result<NlProblem, String> {
let resolved = if path.exists() {
path.to_path_buf()
} else {
let with_nl = append_extension(path, "nl");
if with_nl.exists() {
with_nl
} else {
path.to_path_buf()
}
};
let txt = std::fs::read_to_string(&resolved)
.map_err(|e| format!("could not read {}: {}", resolved.display(), e))?;
let mut prob = parse_nl_string(txt, std::env::var("POUNCE_DBG_NO_QUAD").is_err())?;
prob.var_names = read_name_file(&resolved.with_extension("col"), prob.n);
prob.con_names = read_name_file(&resolved.with_extension("row"), prob.m);
Ok(prob)
}
fn append_extension(path: &Path, ext: &str) -> std::path::PathBuf {
let mut name = path.as_os_str().to_os_string();
name.push(".");
name.push(ext);
std::path::PathBuf::from(name)
}
fn read_name_file(path: &Path, expected: usize) -> Vec<String> {
let Ok(txt) = std::fs::read_to_string(path) else {
return Vec::new();
};
let names: Vec<String> = txt.lines().take(expected).map(str::to_owned).collect();
if names.len() == expected {
names
} else {
Vec::new()
}
}
fn row_constant_value(e: &Expr) -> Option<Number> {
if let Expr::Const(c) = e {
return c.is_finite().then_some(*c);
}
let mut vars: BTreeSet<usize> = BTreeSet::new();
collect_vars(e, &mut vars);
if !vars.is_empty() {
return None;
}
let mut funcs: BTreeSet<usize> = BTreeSet::new();
crate::nl_external::collect_funcall_ids(e, &mut funcs);
if !funcs.is_empty() {
return None;
}
let v = eval_expr(e, &[]);
v.is_finite().then_some(v)
}
pub fn parse_nl_text(txt: &str) -> Result<NlProblem, String> {
parse_nl_text_with_quadratic(txt, std::env::var("POUNCE_DBG_NO_QUAD").is_err())
}
pub fn parse_nl_text_with_quadratic(txt: &str, use_quadratic: bool) -> Result<NlProblem, String> {
parse_nl_string(txt.to_string(), use_quadratic)
}
pub fn parse_nl_string(txt: String, use_quadratic: bool) -> Result<NlProblem, String> {
let src = Arc::new(txt);
let mut p = Parser::new(&src, use_quadratic);
p.parse_header()?;
let n = p.n;
let m = p.m;
let num_obj = p.num_obj;
let mut con_nonlinear: Vec<NlBody> = (0..m).map(|_| NlBody::Tree(Expr::Const(0.0))).collect();
let mut obj_nonlinear = NlBody::Tree(Expr::Const(0.0));
let mut minimize = true;
let mut obj_linear: Vec<(usize, Number)> = Vec::new();
let mut con_linear: Vec<Vec<(usize, Number)>> = vec![Vec::new(); m];
let mut x_l = vec![-1e19; n];
let mut x_u = vec![1e19; n];
let mut g_l = vec![-1e19; m];
let mut g_u = vec![1e19; m];
let mut x0 = vec![0.0; n];
let mut lambda0 = vec![0.0; m];
let mut suffixes = NlSuffixes::default();
let mut imported_funcs: Vec<ImportedFunc> = Vec::new();
let mut saw_r = false;
let mut saw_b = false;
while let Some(line) = p.peek_segment_line() {
let tag = line
.trim_start()
.chars()
.next()
.ok_or("unexpected blank segment header")?;
match tag {
'C' => {
let (_hdr, rest) = p.eat_segment_header()?;
let _ = rest;
let idx = parse_segment_index(_hdr, 'C')?;
if idx >= m {
return Err(format!("C{idx} out of range; m={m}"));
}
con_nonlinear[idx] = p.parse_body()?;
}
'O' => {
let (hdr, _rest) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.len() < 2 {
return Err(format!("malformed O-segment header: {hdr}"));
}
let idx = parse_segment_index(parts[0], 'O')?;
let kind: i32 = parts[1].parse().map_err(|e| format!("O kind: {e}"))?;
if idx == 0 {
minimize = kind == 0;
obj_nonlinear = p.parse_body()?;
} else {
let _ = p.parse_expr()?;
}
}
'r' => {
p.eat_segment_header()?;
saw_r = true;
for i in 0..m {
let line = p.next_data_line()?;
let (lo, hi) = parse_bound_line(line)?;
g_l[i] = lo;
g_u[i] = hi;
}
}
'b' => {
p.eat_segment_header()?;
saw_b = true;
for i in 0..n {
let line = p.next_data_line()?;
let (lo, hi) = parse_bound_line(line)?;
x_l[i] = lo;
x_u[i] = hi;
}
}
'k' => {
let (hdr, _) = p.eat_segment_header()?;
let declared = parse_segment_index(hdr, 'k')?;
let expected = if n == 0 { 0 } else { n - 1 };
if declared != expected {
return Err(format!(
"k-segment declares {declared} column-count lines but \
the standard count for n={n} variables is {expected}"
));
}
for _ in 0..declared {
p.next_data_line()?;
}
}
'J' => {
let (hdr, _) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.len() < 2 {
return Err(format!("malformed J-segment header: {hdr}"));
}
let row = parse_segment_index(parts[0], 'J')?;
let nz: usize = parts[1].parse().map_err(|e| format!("J nz: {e}"))?;
if row >= m {
return Err(format!("J{row} out of range"));
}
for _ in 0..nz {
let line = p.next_data_line()?;
let (var, coef) = parse_var_coef(line)?;
if var >= n {
return Err(format!(
"J{row} entry variable index {var} out of range (n={n})"
));
}
con_linear[row].push((var, coef));
}
}
'G' => {
let (hdr, _) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.len() < 2 {
return Err(format!("malformed G-segment header: {hdr}"));
}
let idx = parse_segment_index(parts[0], 'G')?;
let nz: usize = parts[1].parse().map_err(|e| format!("G nz: {e}"))?;
let mut acc = Vec::with_capacity(nz);
for _ in 0..nz {
let line = p.next_data_line()?;
let (var, coef) = parse_var_coef(line)?;
if var >= n {
return Err(format!(
"G{idx} entry variable index {var} out of range (n={n})"
));
}
acc.push((var, coef));
}
if idx == 0 {
obj_linear = acc;
}
}
'x' => {
let (hdr, _) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
let nx: usize = parts
.first()
.and_then(|s| s.trim_start_matches('x').parse().ok())
.ok_or_else(|| format!("malformed x-segment header: {hdr}"))?;
for _ in 0..nx {
let line = p.next_data_line()?;
let (idx, val) = parse_var_coef(line)?;
if idx >= n {
return Err(format!(
"x-segment variable index {idx} out of range (n={n})"
));
}
x0[idx] = val;
}
}
'd' => {
let (hdr, _) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
let nd: usize = parts
.first()
.and_then(|s| s.trim_start_matches('d').parse().ok())
.ok_or_else(|| format!("malformed d-segment header: {hdr}"))?;
for _ in 0..nd {
let line = p.next_data_line()?;
let (idx, val) = parse_var_coef(line)?;
if idx >= m {
return Err(format!(
"d-segment constraint index {idx} out of range (m={m})"
));
}
lambda0[idx] = val;
}
}
'V' => p.parse_v_segment()?,
'S' => {
parse_suffix_segment(&mut p, n, m, num_obj, &mut suffixes)?;
}
'F' => {
let (hdr, _rest) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.is_empty() {
return Err(format!("malformed F-segment header: '{hdr}'"));
}
let id = parse_segment_index(parts[0], 'F')?;
let kind: usize = parts.get(1).and_then(|s| s.parse().ok()).unwrap_or(0);
let nargs: i64 = parts.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
let name = parts.get(3).copied().unwrap_or("").to_string();
imported_funcs.push(ImportedFunc {
id,
kind,
nargs,
name,
});
}
other => return Err(format!("unknown .nl segment tag '{other}'")),
}
}
if m > 0 && !saw_r {
return Err(format!(
"missing `r` (constraint-bounds) segment for a model declaring {m} \
constraint(s): the .nl file is truncated or corrupt"
));
}
if n > 0 && !saw_b {
return Err(format!(
"missing `b` (variable-bounds) segment for a model declaring {n} \
variable(s): the .nl file is truncated or corrupt"
));
}
if let Some(declared) = p.declared_jac_nnz {
let parsed: usize = con_linear.iter().map(Vec::len).sum();
if parsed != declared {
return Err(format!(
"header declares {declared} Jacobian nonzero(s) but the J \
segments supply {parsed}: the .nl file is truncated or corrupt"
));
}
}
for i in 0..m {
let Some(tree) = con_nonlinear[i].tree() else {
continue;
};
let Some(c) = row_constant_value(tree) else {
continue;
};
if lower_bound_present(g_l[i]) {
g_l[i] -= c;
}
if upper_bound_present(g_u[i]) {
g_u[i] -= c;
}
con_nonlinear[i] = NlBody::Tree(Expr::Const(0.0));
}
let any_recognized =
obj_nonlinear.quad().is_some() || con_nonlinear.iter().any(|b| b.quad().is_some());
let kept_src = any_recognized.then(|| Arc::clone(&src));
Ok(NlProblem {
n,
m,
num_obj,
minimize,
obj_nonlinear,
obj_linear,
obj_constant: 0.0,
con_nonlinear,
con_linear,
x_l,
x_u,
g_l,
g_u,
x0,
lambda0,
suffixes,
ampl_options: p.ampl_options.clone(),
nl_counts: p.nl_counts,
imported_funcs,
var_names: Vec::new(),
con_names: Vec::new(),
src: kept_src,
cse_bodies: p.cses.clone(),
})
}
fn parse_suffix_segment(
p: &mut Parser,
n: usize,
m: usize,
num_obj: usize,
out: &mut NlSuffixes,
) -> Result<(), String> {
let (hdr, _) = p.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.len() < 3 {
return Err(format!(
"malformed S-segment header: '{hdr}' (expected `S<kind> <n> <name>`)"
));
}
let kind_str = parts[0].trim_start_matches('S');
let kind: u32 = kind_str
.parse()
.map_err(|e| format!("S kind '{kind_str}': {e}"))?;
let nentries: usize = parts[1].parse().map_err(|e| format!("S nentries: {e}"))?;
let name = parts[2].to_string();
let is_real = (kind & 0x4) != 0;
let target = kind & 0x3;
let target_dim = match target {
0 => n,
1 => m,
2 => num_obj,
3 => 0, _ => unreachable!("kind & 0x3 is in 0..=3"),
};
let mut int_buf: Vec<Index> = if !is_real && target != 3 {
vec![0; target_dim]
} else {
Vec::new()
};
let mut real_buf: Vec<Number> = if is_real && target != 3 {
vec![0.0; target_dim]
} else {
Vec::new()
};
let mut problem_int: Index = 0;
let mut problem_real: Number = 0.0;
for _ in 0..nentries {
let line = p.next_data_line()?;
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() < 2 {
return Err(format!(
"malformed S-segment entry '{line}' (expected `<idx> <value>`)"
));
}
let idx: usize = parts[0]
.parse()
.map_err(|e| format!("S entry idx '{}': {e}", parts[0]))?;
if target != 3 && idx >= target_dim {
return Err(format!(
"S-suffix '{name}' index {idx} out of range for target dim {target_dim}"
));
}
if is_real {
let v: Number = parts[1]
.parse()
.map_err(|e| format!("S real entry value '{}': {e}", parts[1]))?;
if target == 3 {
problem_real = v;
} else {
real_buf[idx] = v;
}
} else {
let v: Index = parts[1]
.parse()
.map_err(|e| format!("S int entry value '{}': {e}", parts[1]))?;
if target == 3 {
problem_int = v;
} else {
int_buf[idx] = v;
}
}
}
match (target, is_real) {
(0, false) => {
out.var_int.insert(name, int_buf);
}
(1, false) => {
out.con_int.insert(name, int_buf);
}
(2, false) => {
out.obj_int.insert(name, int_buf);
}
(3, false) => {
out.problem_int.insert(name, problem_int);
}
(0, true) => {
out.var_real.insert(name, real_buf);
}
(1, true) => {
out.con_real.insert(name, real_buf);
}
(2, true) => {
out.obj_real.insert(name, real_buf);
}
(3, true) => {
out.problem_real.insert(name, problem_real);
}
_ => unreachable!(),
}
Ok(())
}
fn parse_segment_index(s: &str, tag: char) -> Result<usize, String> {
let trimmed = s.trim_start_matches(tag);
trimmed
.parse()
.map_err(|e| format!("malformed {tag}-segment index '{s}': {e}"))
}
fn finite_or_err(what: &str, v: Number) -> Result<Number, String> {
if v.is_finite() {
Ok(v)
} else {
Err(format!(
"invalid number: {what} is {v}, which is not finite"
))
}
}
fn finite_bound_or_err(what: &str, v: Number, lower: bool) -> Result<Number, String> {
if v.is_finite() {
return Ok(v);
}
if lower && v == Number::NEG_INFINITY {
return Ok(-1e19);
}
if !lower && v == Number::INFINITY {
return Ok(1e19);
}
Err(format!(
"invalid number: {what} is {v}, which is not finite (a `.nl` file \
states an absent bound with a bound kind of 1, 2 or 3, not with a \
non-finite value)"
))
}
fn parse_bound_line(line: &str) -> Result<(Number, Number), String> {
let mut parts = line.split_whitespace();
let kind: i32 = parts
.next()
.ok_or("empty bound line")?
.parse()
.map_err(|e| format!("bound kind: {e}"))?;
let lo;
let hi;
match kind {
0 => {
let (l, h) = (parts.next(), parts.next());
let (Some(l), Some(h)) = (l, h) else {
return Err(format!("bound kind 0 needs 2 values: '{line}'"));
};
lo = finite_bound_or_err("lo", l.parse().map_err(|e| format!("lo: {e}"))?, true)?;
hi = finite_bound_or_err("hi", h.parse().map_err(|e| format!("hi: {e}"))?, false)?;
}
1 => {
let Some(h) = parts.next() else {
return Err(format!("bound kind 1 needs 1 value: '{line}'"));
};
lo = -1e19;
hi = finite_bound_or_err("hi", h.parse().map_err(|e| format!("hi: {e}"))?, false)?;
}
2 => {
let Some(l) = parts.next() else {
return Err(format!("bound kind 2 needs 1 value: '{line}'"));
};
lo = finite_bound_or_err("lo", l.parse().map_err(|e| format!("lo: {e}"))?, true)?;
hi = 1e19;
}
3 => {
lo = -1e19;
hi = 1e19;
}
4 => {
let Some(v) = parts.next() else {
return Err(format!("bound kind 4 needs 1 value: '{line}'"));
};
let v: Number = finite_or_err("eq bound", v.parse().map_err(|e| format!("eq: {e}"))?)?;
lo = v;
hi = v;
}
5 => return Err("complementarity (kind 5) bounds are not supported".into()),
other => return Err(format!("unknown bound kind {other}")),
}
Ok((lo, hi))
}
fn parse_var_coef(line: &str) -> Result<(usize, Number), String> {
let mut parts = line.split_whitespace();
let (Some(v), Some(c)) = (parts.next(), parts.next()) else {
return Err(format!("malformed var/coef line: '{line}'"));
};
let v: usize = v.parse().map_err(|e| format!("var idx: {e}"))?;
let c: Number = finite_or_err("coefficient", c.parse().map_err(|e| format!("coef: {e}"))?)?;
Ok((v, c))
}
fn parse_nl_counts(line3: &str, line5: &str) -> Option<NlCounts> {
let nums = |line: &str, want: usize| -> Option<Vec<usize>> {
let v: Vec<usize> = line
.split_whitespace()
.take(want)
.map(str::parse)
.collect::<Result<_, _>>()
.ok()?;
(v.len() == want).then_some(v)
};
let cons_objs = nums(line3, 2)?;
let vars = nums(line5, 3)?;
Some(NlCounts {
nl_cons: cons_objs[0],
nl_objs: cons_objs[1],
nl_vars_cons: vars[0],
nl_vars_objs: vars[1],
nl_vars_both: vars[2],
})
}
struct Parser<'a> {
lines: Vec<&'a str>,
txt_base: usize,
txt_len: usize,
pos: usize,
n: usize,
m: usize,
num_obj: usize,
n_funcs: usize,
nl_counts: Option<NlCounts>,
declared_jac_nnz: Option<usize>,
ampl_options: Vec<i64>,
cses: Vec<Arc<Expr>>,
quad_enabled: bool,
cse_quad: Vec<Option<Quad2>>,
cse_sum_ok: Vec<bool>,
cse_mono_ok: Vec<bool>,
cse_vars: Vec<Vec<u32>>,
cse_depth: Vec<u32>,
}
struct QFrame {
op: QOp,
remaining: usize,
mono: bool,
}
enum QOp {
Neg,
Add,
Sub,
Mul,
Div,
Pow,
Square,
Sum(usize),
}
impl<'a> Parser<'a> {
fn new(txt: &'a str, quad_enabled: bool) -> Self {
let lines: Vec<&str> = txt.lines().collect();
Self {
lines,
txt_base: txt.as_ptr() as usize,
txt_len: txt.len(),
pos: 0,
n: 0,
m: 0,
num_obj: 0,
n_funcs: 0,
nl_counts: None,
declared_jac_nnz: None,
ampl_options: Vec::new(),
cses: Vec::new(),
quad_enabled,
cse_quad: Vec::new(),
cse_sum_ok: Vec::new(),
cse_mono_ok: Vec::new(),
cse_vars: Vec::new(),
cse_depth: Vec::new(),
}
}
fn byte_at(&self, line: usize) -> usize {
match self.lines.get(line) {
Some(l) => l.as_ptr() as usize - self.txt_base,
None => self.txt_len,
}
}
fn next_line(&mut self) -> Option<&'a str> {
while self.pos < self.lines.len() {
let l = self.lines[self.pos];
self.pos += 1;
let trimmed = strip_comment(l).trim();
if !trimmed.is_empty() {
return Some(l);
}
}
None
}
fn next_data_line(&mut self) -> Result<&'a str, String> {
while self.pos < self.lines.len() {
let l = self.lines[self.pos];
self.pos += 1;
let trimmed = strip_comment(l).trim();
if !trimmed.is_empty() {
return Ok(trimmed);
}
}
Err("unexpected end of file in data line".to_string())
}
fn parse_header(&mut self) -> Result<(), String> {
let line0 = self.next_line().ok_or("empty .nl file")?;
let trimmed = strip_comment(line0).trim();
let first = trimmed.chars().next().ok_or("empty header line")?;
if first != 'g' {
return Err(format!(
"only ASCII (g-) .nl files supported; got header '{trimmed}'"
));
}
let mut words = trimmed.split_whitespace();
let n_opts: usize = words.next().and_then(|w| w[1..].parse().ok()).unwrap_or(0);
let opts: Vec<i64> = words.filter_map(|w| w.parse().ok()).collect();
if opts.len() >= n_opts {
self.ampl_options = opts[..n_opts].to_vec();
}
let l2 = self.next_data_line()?;
let nums: Vec<&str> = l2.split_whitespace().collect();
if nums.len() < 3 {
return Err(format!("malformed line 2: '{l2}'"));
}
self.n = nums[0].parse().map_err(|e| format!("n: {e}"))?;
self.m = nums[1].parse().map_err(|e| format!("m: {e}"))?;
self.num_obj = nums[2].parse().map_err(|e| format!("num_obj: {e}"))?;
let l3 = self.next_data_line()?;
let _l4_network = self.next_data_line()?;
let l5 = self.next_data_line()?;
self.nl_counts = parse_nl_counts(l3, l5);
let l6 = self.next_data_line()?;
let nums5: Vec<&str> = l6.split_whitespace().collect();
self.n_funcs = nums5.get(1).and_then(|s| s.parse().ok()).unwrap_or(0);
let _l7_discrete = self.next_data_line()?;
let l8 = self.next_data_line()?;
self.declared_jac_nnz = l8.split_whitespace().next().and_then(|s| s.parse().ok());
let _l9_name_lens = self.next_data_line()?;
let _l10_common_exprs = self.next_data_line()?;
Ok(())
}
fn peek_segment_line(&mut self) -> Option<&'a str> {
let saved = self.pos;
let l = self.next_line()?;
self.pos = saved;
Some(l)
}
fn eat_segment_header(&mut self) -> Result<(&'a str, &'a str), String> {
let raw = self
.next_line()
.ok_or_else(|| "expected segment header".to_string())?;
let (hdr, comment) = split_comment(raw);
Ok((hdr.trim(), comment.trim()))
}
fn parse_body(&mut self) -> Result<NlBody, String> {
if !self.quad_enabled {
return Ok(NlBody::Tree(self.parse_expr()?));
}
let saved = self.pos;
if let Some((form, vars, depth)) = self.parse_expr_quadratic() {
if !form.quadratic().is_empty() {
let src = self.byte_at(saved)..self.byte_at(self.pos);
return Ok(NlBody::Quad(Box::new(QuadBody {
form,
vars,
src,
depth,
})));
}
}
self.pos = saved;
Ok(NlBody::Tree(self.parse_expr()?))
}
fn parse_expr_quadratic(&mut self) -> Option<(Quad2, Vec<u32>, u32)> {
let mut frames: Vec<QFrame> = Vec::new();
let mut vals: Vec<Quad2> = Vec::new();
let mut vars: Vec<u32> = Vec::new();
let mut depth: u32 = 0;
let note_leaf = |frames: &[QFrame], depth: &mut u32, below: u32| {
let d = u32::try_from(frames.len()).unwrap_or(u32::MAX);
*depth = (*depth).max(d.saturating_add(1).saturating_add(below));
};
loop {
let mono = frames.last().is_some_and(|f| f.mono);
let pow_base = matches!(
frames.last(),
Some(QFrame {
op: QOp::Pow,
remaining: 2,
..
}) | Some(QFrame {
op: QOp::Square,
remaining: 1,
..
})
);
let raw = self.next_line()?;
let tok = strip_comment(raw).trim();
let first = tok.chars().next()?;
match first {
'n' => {
let v: Number = tok[1..]
.trim()
.parse()
.ok()
.filter(|v: &Number| v.is_finite())?;
note_leaf(&frames, &mut depth, 0);
vals.push(Quad2::of_constant(v));
}
'v' => {
let i: usize = tok[1..].trim().parse().ok()?;
if i < self.n {
note_leaf(&frames, &mut depth, 0);
vars.push(u32::try_from(i).ok()?);
vals.push(Quad2::of_var(i));
} else {
if pow_base {
return None;
}
let local = i.checked_sub(self.n)?;
let ok = if mono {
*self.cse_mono_ok.get(local)?
} else {
*self.cse_sum_ok.get(local)?
};
if !ok {
return None;
}
let form = self.cse_quad.get(local)?.clone()?;
note_leaf(&frames, &mut depth, *self.cse_depth.get(local)?);
vars.extend_from_slice(self.cse_vars.get(local)?);
vals.push(form);
}
}
'o' => {
let code: i32 = tok[1..].trim().parse().ok()?;
if pow_base && code != 16 {
return None;
}
let (op, arity, child_mono) = match code {
0 if !mono => (QOp::Add, 2, false),
1 if !mono => (QOp::Sub, 2, false),
16 => (QOp::Neg, 1, mono),
54 if !mono => {
let count_line = self.next_data_line().ok()?;
let count: usize =
count_line.split_whitespace().next()?.parse().ok()?;
(QOp::Sum(count), count, false)
}
2 => (QOp::Mul, 2, true),
3 => (QOp::Div, 2, true),
5 | 81 | 83 => (QOp::Pow, 2, true),
82 => (QOp::Square, 1, true),
_ => return None,
};
if arity == 0 {
note_leaf(&frames, &mut depth, 0);
vals.push(Quad2::default());
} else {
frames.push(QFrame {
op,
remaining: arity,
mono: child_mono,
});
continue;
}
}
_ => return None,
}
while let Some(f) = frames.last_mut() {
f.remaining -= 1;
if f.remaining > 0 {
break;
}
let f = frames.pop()?;
let combined = apply_quad_op(f.op, &mut vals)?;
vals.push(combined);
}
if frames.is_empty() {
break;
}
}
if vals.len() != 1 {
return None;
}
vars.sort_unstable();
vars.dedup();
Some((vals.pop()?, vars, depth))
}
fn parse_expr(&mut self) -> Result<Expr, String> {
let raw = self
.next_line()
.ok_or_else(|| "expected expression token".to_string())?;
let tok = strip_comment(raw).trim();
if tok.is_empty() {
return Err("empty expression token".into());
}
let first = tok.chars().next().ok_or("empty expression token")?;
match first {
'n' => {
let v: Number = tok[1..]
.trim()
.parse()
.map_err(|e| format!("n value: {e}"))?;
Ok(Expr::Const(finite_or_err("numeric literal", v)?))
}
'v' => {
let i: usize = tok[1..]
.trim()
.parse()
.map_err(|e| format!("v index: {e}"))?;
Ok(self.var_or_cse(i)?)
}
'o' => {
let code: i32 = tok[1..]
.trim()
.parse()
.map_err(|e| format!("opcode: {e}"))?;
self.parse_opcode(code)
}
'f' => {
let rest = &tok[1..];
let mut parts = rest.split_whitespace();
let id_str = parts
.next()
.ok_or_else(|| format!("missing function id in '{tok}'"))?;
let nargs_str = parts
.next()
.ok_or_else(|| format!("missing nargs in '{tok}'"))?;
let id: usize = id_str
.parse()
.map_err(|e| format!("bad function id '{id_str}': {e}"))?;
let nargs: usize = nargs_str
.parse()
.map_err(|e| format!("bad funcall nargs '{nargs_str}': {e}"))?;
let mut args: Vec<FuncallArg> = Vec::with_capacity(nargs);
for _ in 0..nargs {
args.push(self.parse_funcall_arg()?);
}
Ok(Expr::Funcall { id, args })
}
't' | 'u' => Err(format!("unsupported expression token '{tok}'")),
other => Err(format!(
"unexpected expression token start '{other}': '{tok}'"
)),
}
}
fn parse_funcall_arg(&mut self) -> Result<FuncallArg, String> {
let saved = self.pos;
let raw = self
.next_line()
.ok_or_else(|| "expected funcall argument".to_string())?;
let lead = raw.trim_start();
if let Some(after_h) = lead.strip_prefix('h') {
let colon = after_h
.find(':')
.ok_or_else(|| format!("malformed Hollerith string arg (no ':'): {lead:?}"))?;
let len: usize = after_h[..colon]
.trim()
.parse()
.map_err(|e| format!("Hollerith length in {lead:?}: {e}"))?;
let chars = &after_h[colon + 1..];
if chars.len() < len {
return Err(format!(
"Hollerith string shorter than declared length {len}: {chars:?}"
));
}
if !chars.is_char_boundary(len) {
return Err(format!(
"Hollerith length {len} splits a multibyte char in {chars:?}"
));
}
Ok(FuncallArg::Str(chars[..len].to_string()))
} else {
self.pos = saved;
Ok(FuncallArg::Real(self.parse_expr()?))
}
}
fn parse_opcode(&mut self, code: i32) -> Result<Expr, String> {
match code {
0 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Add, Box::new(a), Box::new(b)))
}
1 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b)))
}
2 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Mul, Box::new(a), Box::new(b)))
}
3 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Div, Box::new(a), Box::new(b)))
}
5 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Pow, Box::new(a), Box::new(b)))
}
15 => Ok(Expr::Unary(UnaryOp::Abs, Box::new(self.parse_expr()?))),
16 => Ok(Expr::Unary(UnaryOp::Neg, Box::new(self.parse_expr()?))),
39 => Ok(Expr::Unary(UnaryOp::Sqrt, Box::new(self.parse_expr()?))),
41 => Ok(Expr::Unary(UnaryOp::Sin, Box::new(self.parse_expr()?))),
42 => Ok(Expr::Unary(UnaryOp::Log10, Box::new(self.parse_expr()?))),
43 => Ok(Expr::Unary(UnaryOp::Log, Box::new(self.parse_expr()?))),
44 => Ok(Expr::Unary(UnaryOp::Exp, Box::new(self.parse_expr()?))),
46 => Ok(Expr::Unary(UnaryOp::Cos, Box::new(self.parse_expr()?))),
38 => Ok(Expr::Unary(UnaryOp::Tan, Box::new(self.parse_expr()?))),
49 => Ok(Expr::Unary(UnaryOp::Atan, Box::new(self.parse_expr()?))),
53 => Ok(Expr::Unary(UnaryOp::Acos, Box::new(self.parse_expr()?))),
40 => Ok(Expr::Unary(UnaryOp::Sinh, Box::new(self.parse_expr()?))),
45 => Ok(Expr::Unary(UnaryOp::Cosh, Box::new(self.parse_expr()?))),
37 => Ok(Expr::Unary(UnaryOp::Tanh, Box::new(self.parse_expr()?))),
51 => Ok(Expr::Unary(UnaryOp::Asin, Box::new(self.parse_expr()?))),
52 => Ok(Expr::Unary(UnaryOp::Acosh, Box::new(self.parse_expr()?))),
50 => Ok(Expr::Unary(UnaryOp::Asinh, Box::new(self.parse_expr()?))),
47 => Ok(Expr::Unary(UnaryOp::Atanh, Box::new(self.parse_expr()?))),
48 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Atan2, Box::new(a), Box::new(b)))
}
22 => self.parse_compare(CmpOp::Lt),
23 => self.parse_compare(CmpOp::Le),
24 => self.parse_compare(CmpOp::Eq),
28 => self.parse_compare(CmpOp::Ge),
29 => self.parse_compare(CmpOp::Gt),
30 => self.parse_compare(CmpOp::Ne),
20 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Or(Box::new(a), Box::new(b)))
}
21 => {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::And(Box::new(a), Box::new(b)))
}
34 => Ok(Expr::Not(Box::new(self.parse_expr()?))),
35 => {
let cond = self.parse_expr()?;
let then_ = self.parse_expr()?;
let else_ = self.parse_expr()?;
Ok(Expr::Cond {
cond: Box::new(cond),
then_: Box::new(then_),
else_: Box::new(else_),
})
}
54 => {
let count_line = self.next_data_line()?;
let count: usize = count_line
.split_whitespace()
.next()
.ok_or_else(|| "missing variadic count".to_string())?
.parse()
.map_err(|e| format!("variadic count: {e}"))?;
let mut args = Vec::with_capacity(count);
for _ in 0..count {
args.push(self.parse_expr()?);
}
Ok(Expr::Sum(args))
}
11 | 12 => {
let count_line = self.next_data_line()?;
let count: usize = count_line
.split_whitespace()
.next()
.ok_or_else(|| "missing min/max list count".to_string())?
.parse()
.map_err(|e| format!("min/max list count: {e}"))?;
let mut args = Vec::with_capacity(count);
for _ in 0..count {
args.push(self.parse_expr()?);
}
if code == 11 {
Ok(Expr::MinList(args))
} else {
Ok(Expr::MaxList(args))
}
}
81 => {
let base = self.parse_expr()?;
let exp = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Pow, Box::new(base), Box::new(exp)))
}
82 => {
let base = self.parse_expr()?;
Ok(Expr::Binary(
BinOp::Pow,
Box::new(base),
Box::new(Expr::Const(2.0)),
))
}
83 => {
let base = self.parse_expr()?;
let exp = self.parse_expr()?;
Ok(Expr::Binary(BinOp::Pow, Box::new(base), Box::new(exp)))
}
other => Err(format!("unsupported opcode o{other}")),
}
}
fn parse_compare(&mut self, op: CmpOp) -> Result<Expr, String> {
let a = self.parse_expr()?;
let b = self.parse_expr()?;
Ok(Expr::Compare(op, Box::new(a), Box::new(b)))
}
fn var_or_cse(&self, i: usize) -> Result<Expr, String> {
if i < self.n {
Ok(Expr::Var(i))
} else {
let local = i - self.n;
self.cses
.get(local)
.map(|rc| Expr::Cse(rc.clone()))
.ok_or_else(|| {
format!(
"v{i} references CSE {local} but only {} have been defined",
self.cses.len()
)
})
}
}
fn parse_v_segment(&mut self) -> Result<(), String> {
let (hdr, _) = self.eat_segment_header()?;
let parts: Vec<&str> = hdr.split_whitespace().collect();
if parts.len() < 2 {
return Err(format!("malformed V-segment header: {hdr}"));
}
let cse_idx = parse_segment_index(parts[0], 'V')?;
let nlin: usize = parts[1].parse().map_err(|e| format!("V nlin: {e}"))?;
let mut linear: Vec<(usize, Number)> = Vec::with_capacity(nlin);
for _ in 0..nlin {
let line = self.next_data_line()?;
let (var, coef) = parse_var_coef(line)?;
linear.push((var, coef));
}
let nonlin = self.parse_expr()?;
let mut combined = nonlin;
for (var, coef) in linear {
let v_expr = self.var_or_cse(var)?;
let term = if coef == 1.0 {
v_expr
} else {
Expr::Binary(BinOp::Mul, Box::new(Expr::Const(coef)), Box::new(v_expr))
};
combined = Expr::Binary(BinOp::Add, Box::new(combined), Box::new(term));
}
if cse_idx < self.n {
return Err(format!("V{cse_idx} below n={}", self.n));
}
let local = cse_idx - self.n;
if local != self.cses.len() {
return Err(format!(
"V-segment index V{cse_idx} out of order; expected V{}",
self.n + self.cses.len()
));
}
if self.quad_enabled {
self.cse_quad
.push(crate::nl_quadratic::recognize_expr(&combined));
self.cse_sum_ok
.push(crate::nl_quadratic::is_expanded_quadratic(&combined));
self.cse_mono_ok
.push(crate::nl_quadratic::is_monomial_expr(&combined));
let mut vars: BTreeSet<usize> = BTreeSet::new();
collect_vars(&combined, &mut vars);
self.cse_vars
.push(vars.into_iter().map(|v| v as u32).collect());
self.cse_depth.push(expr_tree_depth(&combined));
}
self.cses.push(Arc::new(combined));
Ok(())
}
}
fn expr_tree_depth(e: &Expr) -> u32 {
let deepest = |kids: &mut dyn Iterator<Item = &Expr>| {
kids.fold(0u32, |acc, k| acc.max(expr_tree_depth(k)))
};
1 + match e {
Expr::Const(_) | Expr::Var(_) => 0,
Expr::Binary(_, a, b) | Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
deepest(&mut [&**a, &**b].into_iter())
}
Expr::Unary(_, a) | Expr::Not(a) => expr_tree_depth(a),
Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => deepest(&mut args.iter()),
Expr::Cond { cond, then_, else_ } => {
deepest(&mut [&**cond, &**then_, &**else_].into_iter())
}
Expr::Funcall { args, .. } => deepest(&mut args.iter().filter_map(|a| match a {
FuncallArg::Real(inner) => Some(inner),
FuncallArg::Str(_) => None,
})),
Expr::Cse(body) => expr_tree_depth(body),
}
}
fn apply_quad_op(op: QOp, vals: &mut Vec<Quad2>) -> Option<Quad2> {
let pop2 = |vals: &mut Vec<Quad2>| -> Option<(Quad2, Quad2)> {
let b = vals.pop()?;
let a = vals.pop()?;
Some((a, b))
};
Some(match op {
QOp::Sum(n) => {
let at = vals.len().checked_sub(n)?;
let mut acc = Quad2::default();
for p in vals.drain(at..) {
acc = Quad2::add(acc, p);
}
acc
}
QOp::Neg => vals.pop()?.neg(),
QOp::Add => {
let (a, b) = pop2(vals)?;
Quad2::add(a, b)
}
QOp::Sub => {
let (a, b) = pop2(vals)?;
Quad2::add(a, b.neg())
}
QOp::Mul => {
let (a, b) = pop2(vals)?;
a.mul(&b)?
}
QOp::Div => {
let (a, b) = pop2(vals)?;
let d = b.as_constant()?;
if d == 0.0 {
return None;
}
let mut out = a.div_by_constant(d);
out.absorb_flags(&b);
out
}
QOp::Pow => {
let (a, b) = pop2(vals)?;
let exp = b.as_constant()?;
let mut out = if exp == 0.0 {
Quad2::of_constant(1.0)
} else if exp == 1.0 {
a
} else if exp == 2.0 {
a.mul(&a)?
} else {
return None;
};
out.absorb_flags(&b);
out
}
QOp::Square => {
let a = vals.pop()?;
a.mul(&a)?
}
})
}
fn parse_body_fragment(txt: &str, n: usize, cses: &[Arc<Expr>]) -> Result<Expr, String> {
let mut p = Parser::new(txt, false);
p.n = n;
p.cses = cses.to_vec();
p.parse_expr()
}
fn strip_comment(s: &str) -> &str {
match s.find('#') {
Some(i) => &s[..i],
None => s,
}
}
fn split_comment(s: &str) -> (&str, &str) {
match s.find('#') {
Some(i) => (&s[..i], &s[i + 1..]),
None => (s, ""),
}
}
pub fn eval_expr(e: &Expr, x: &[Number]) -> Number {
match e {
Expr::Const(c) => *c,
Expr::Var(i) => x[*i],
Expr::Binary(op, a, b) => {
let va = eval_expr(a, x);
let vb = eval_expr(b, x);
match op {
BinOp::Add => va + vb,
BinOp::Sub => va - vb,
BinOp::Mul => va * vb,
BinOp::Div => va / vb,
BinOp::Pow => va.powf(vb),
BinOp::Atan2 => va.atan2(vb),
BinOp::CEntropy => crate::nl_tape::centropy(va, vb),
}
}
Expr::Unary(op, a) => {
let va = eval_expr(a, x);
match op {
UnaryOp::Neg => -va,
UnaryOp::Sqrt => va.sqrt(),
UnaryOp::Log => va.ln(),
UnaryOp::Log10 => va.log10(),
UnaryOp::Exp => va.exp(),
UnaryOp::Abs => va.abs(),
UnaryOp::Sin => va.sin(),
UnaryOp::Cos => va.cos(),
UnaryOp::Tan => va.tan(),
UnaryOp::Atan => va.atan(),
UnaryOp::Acos => va.acos(),
UnaryOp::Sinh => va.sinh(),
UnaryOp::Cosh => va.cosh(),
UnaryOp::Tanh => va.tanh(),
UnaryOp::Asin => va.asin(),
UnaryOp::Acosh => va.acosh(),
UnaryOp::Asinh => va.asinh(),
UnaryOp::Atanh => va.atanh(),
UnaryOp::Erf => crate::nl_tape::erf(va),
UnaryOp::XLogX => crate::nl_tape::xlogx(va),
}
}
Expr::Sum(args) => args.iter().map(|a| eval_expr(a, x)).sum(),
Expr::MinList(args) => args
.iter()
.map(|a| eval_expr(a, x))
.fold(Number::INFINITY, Number::min),
Expr::MaxList(args) => args
.iter()
.map(|a| eval_expr(a, x))
.fold(Number::NEG_INFINITY, Number::max),
Expr::Compare(op, a, b) => {
let va = eval_expr(a, x);
let vb = eval_expr(b, x);
let truth = match op {
CmpOp::Lt => va < vb,
CmpOp::Le => va <= vb,
CmpOp::Eq => va == vb,
CmpOp::Ge => va >= vb,
CmpOp::Gt => va > vb,
CmpOp::Ne => va != vb,
};
if truth { 1.0 } else { 0.0 }
}
Expr::And(a, b) => {
if eval_expr(a, x) != 0.0 && eval_expr(b, x) != 0.0 {
1.0
} else {
0.0
}
}
Expr::Or(a, b) => {
if eval_expr(a, x) != 0.0 || eval_expr(b, x) != 0.0 {
1.0
} else {
0.0
}
}
Expr::Not(a) => {
if eval_expr(a, x) == 0.0 {
1.0
} else {
0.0
}
}
Expr::Cond { cond, then_, else_ } => {
if eval_expr(cond, x) != 0.0 {
eval_expr(then_, x)
} else {
eval_expr(else_, x)
}
}
Expr::Cse(body) => eval_expr(body, x),
Expr::Funcall { .. } => panic!(
"eval_expr: AMPL imported function called without an external resolver; \
evaluate through the tape AD path (Tape::build_with_externals) instead"
),
}
}
fn argmin_argmax(args: &[Expr], x: &[Number], want_min: bool) -> Option<usize> {
let mut best: Option<(usize, Number)> = None;
for (i, a) in args.iter().enumerate() {
let v = eval_expr(a, x);
match best {
None => best = Some((i, v)),
Some((_, bv)) => {
if (want_min && v < bv) || (!want_min && v > bv) {
best = Some((i, v));
}
}
}
}
best.map(|(i, _)| i)
}
pub fn grad_expr(e: &Expr, x: &[Number], seed: Number, grad: &mut [Number]) {
match e {
Expr::Const(_) => {}
Expr::Var(i) => grad[*i] += seed,
Expr::Binary(op, a, b) => {
let va = eval_expr(a, x);
let vb = eval_expr(b, x);
match op {
BinOp::Add => {
grad_expr(a, x, seed, grad);
grad_expr(b, x, seed, grad);
}
BinOp::Sub => {
grad_expr(a, x, seed, grad);
grad_expr(b, x, -seed, grad);
}
BinOp::Mul => {
grad_expr(a, x, seed * vb, grad);
grad_expr(b, x, seed * va, grad);
}
BinOp::Div => {
grad_expr(a, x, seed / vb, grad);
grad_expr(b, x, -seed * va / (vb * vb), grad);
}
BinOp::Pow => {
let dpa = vb * va.powf(vb - 1.0);
grad_expr(a, x, seed * dpa, grad);
if va > 0.0 {
let dpb = va.powf(vb) * va.ln();
grad_expr(b, x, seed * dpb, grad);
}
}
BinOp::Atan2 => {
let d = va * va + vb * vb;
grad_expr(a, x, seed * vb / d, grad);
grad_expr(b, x, -seed * va / d, grad);
}
BinOp::CEntropy => {
grad_expr(a, x, seed * crate::nl_tape::centropy_da(va, vb), grad);
grad_expr(b, x, seed * crate::nl_tape::centropy_db(va, vb), grad);
}
}
}
Expr::Unary(op, a) => {
let va = eval_expr(a, x);
let d = match op {
UnaryOp::Neg => -1.0,
UnaryOp::Sqrt => 0.5 / va.sqrt(),
UnaryOp::Log => 1.0 / va,
UnaryOp::Log10 => 1.0 / (va * std::f64::consts::LN_10),
UnaryOp::Exp => va.exp(),
UnaryOp::Abs => {
if va > 0.0 {
1.0
} else if va < 0.0 {
-1.0
} else {
0.0
}
}
UnaryOp::Sin => va.cos(),
UnaryOp::Cos => -va.sin(),
UnaryOp::Tan => {
let t = va.tan();
1.0 + t * t
}
UnaryOp::Atan => 1.0 / (1.0 + va * va),
UnaryOp::Acos => -1.0 / (1.0 - va * va).sqrt(),
UnaryOp::Sinh => va.cosh(),
UnaryOp::Cosh => va.sinh(),
UnaryOp::Tanh => {
let t = va.tanh();
1.0 - t * t
}
UnaryOp::Asin => 1.0 / (1.0 - va * va).sqrt(),
UnaryOp::Acosh => 1.0 / (va * va - 1.0).sqrt(),
UnaryOp::Asinh => 1.0 / (va * va + 1.0).sqrt(),
UnaryOp::Atanh => 1.0 / (1.0 - va * va),
UnaryOp::Erf => crate::nl_tape::erf_d1(va),
UnaryOp::XLogX => crate::nl_tape::xlogx_d1(va),
};
grad_expr(a, x, seed * d, grad);
}
Expr::Sum(args) => {
for arg in args {
grad_expr(arg, x, seed, grad);
}
}
Expr::MinList(args) => {
if let Some(k) = argmin_argmax(args, x, true) {
grad_expr(&args[k], x, seed, grad);
}
}
Expr::MaxList(args) => {
if let Some(k) = argmin_argmax(args, x, false) {
grad_expr(&args[k], x, seed, grad);
}
}
Expr::Compare(_, _, _) | Expr::And(_, _) | Expr::Or(_, _) | Expr::Not(_) => {}
Expr::Cond { cond, then_, else_ } => {
if eval_expr(cond, x) != 0.0 {
grad_expr(then_, x, seed, grad);
} else {
grad_expr(else_, x, seed, grad);
}
}
Expr::Cse(body) => grad_expr(body, x, seed, grad),
Expr::Funcall { .. } => {
panic!("grad_expr: AMPL imported function called without an external resolver")
}
}
}
pub fn collect_vars(e: &Expr, out: &mut BTreeSet<usize>) {
let mut seen: std::collections::HashSet<*const Expr> = std::collections::HashSet::new();
collect_vars_memo(e, out, &mut seen);
}
fn collect_vars_memo(
e: &Expr,
out: &mut BTreeSet<usize>,
seen: &mut std::collections::HashSet<*const Expr>,
) {
match e {
Expr::Const(_) => {}
Expr::Var(i) => {
out.insert(*i);
}
Expr::Binary(_, a, b) => {
collect_vars_memo(a, out, seen);
collect_vars_memo(b, out, seen);
}
Expr::Unary(_, a) => collect_vars_memo(a, out, seen),
Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
for a in args {
collect_vars_memo(a, out, seen);
}
}
Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
collect_vars_memo(a, out, seen);
collect_vars_memo(b, out, seen);
}
Expr::Not(a) => collect_vars_memo(a, out, seen),
Expr::Cond { cond, then_, else_ } => {
collect_vars_memo(cond, out, seen);
collect_vars_memo(then_, out, seen);
collect_vars_memo(else_, out, seen);
}
Expr::Cse(body) => {
if seen.insert(Arc::as_ptr(body)) {
collect_vars_memo(body, out, seen);
}
}
Expr::Funcall { args, .. } => {
for a in args {
if let FuncallArg::Real(e) = a {
collect_vars_memo(e, out, seen);
}
}
}
}
}
#[derive(Debug, Clone)]
struct ColorWrite {
row: u32,
hess_idx: u32,
}
#[derive(Debug, Clone)]
struct ConHybrid {
tape: HybridTape,
row_start: Vec<usize>,
prelude_vals: Vec<f64>,
local_vals: Vec<f64>,
local_adj: Vec<f64>,
prelude_adj: Vec<f64>,
use_for_jac: bool,
use_for_hess: bool,
local_vals_all: Vec<f64>,
local_off: Vec<usize>,
summand_row: Vec<u32>,
hess_color_summands: Vec<Vec<u32>>,
hess_color_reach: Vec<u32>,
hess_color_reach_off: Vec<usize>,
prelude_dot: Vec<f64>,
hess_prelude_adj: Vec<f64>,
prelude_adj_dot: Vec<f64>,
local_dot: Vec<f64>,
local_adj_dot: Vec<f64>,
}
const HYBRID_JAC_MIN_OP_RATIO: f64 = 4.0;
const HYBRID_HESS_MIN_OP_RATIO: f64 = 3.0;
#[derive(Debug, Clone)]
pub struct NlTnlp {
prob: NlProblem,
obj_tapes: Vec<Tape>,
con_tapes: Vec<Vec<Tape>>,
con_hybrid: Option<ConHybrid>,
quad: QuadraticStructure,
h_tape_mask: Vec<bool>,
h_irow: Vec<i32>,
h_jcol: Vec<i32>,
jac_cols: Vec<Vec<usize>>,
jac_nnz: usize,
seeds: Vec<Vec<f64>>,
decoding: Vec<Vec<ColorWrite>>,
obj_tape_colors: Vec<Vec<u32>>,
con_tape_colors: Vec<Vec<Vec<u32>>>,
var_color: Vec<u32>,
peeled_cols: Vec<u32>,
final_x: Option<Vec<Number>>,
final_obj: Number,
final_lambda: Option<Vec<Number>>,
final_z_l: Option<Vec<Number>>,
final_z_u: Option<Vec<Number>>,
scratch_row_grad: Vec<f64>,
vals_scratch: Vec<f64>,
dot_scratch: Vec<f64>,
adj_scratch: Vec<f64>,
adj_dot_scratch: Vec<f64>,
compressed: Vec<Vec<f64>>,
hvp_live: Vec<bool>,
curvature_scaling: Option<crate::nl_scaling::CurvatureScaling>,
}
const P_ADD: u8 = 10;
const P_MUL: u8 = 20;
const P_NEG: u8 = 30;
const P_POW: u8 = 40;
const P_ATOM: u8 = 100;
fn fmt_num(x: Number) -> String {
if x.is_finite() && x == x.trunc() && x.abs() < 1e15 {
format!("{}", x as i64)
} else {
format!("{x}")
}
}
fn var_label(i: usize, var_names: &[String]) -> String {
match var_names.get(i) {
Some(s) if !s.is_empty() => s.clone(),
_ => format!("x[{i}]"),
}
}
fn expr_prec(e: &Expr) -> u8 {
match e {
Expr::Binary(BinOp::Add, ..) | Expr::Binary(BinOp::Sub, ..) | Expr::Sum(_) => P_ADD,
Expr::Binary(BinOp::Mul, ..) | Expr::Binary(BinOp::Div, ..) => P_MUL,
Expr::Unary(UnaryOp::Neg, _) => P_NEG,
Expr::Binary(BinOp::Pow, ..) => P_POW,
Expr::Cse(inner) => expr_prec(inner),
_ => P_ATOM,
}
}
pub fn render_expression(e: &Expr, var_names: &[String]) -> String {
render_expr(e, var_names, &[])
}
fn render_prec(e: &Expr, min_prec: u8, vn: &[String], funcs: &[ImportedFunc]) -> String {
let s = render_expr(e, vn, funcs);
if expr_prec(e) < min_prec {
format!("({s})")
} else {
s
}
}
fn unary_name(op: UnaryOp) -> &'static str {
match op {
UnaryOp::Neg => "-",
UnaryOp::Sqrt => "sqrt",
UnaryOp::Log => "log",
UnaryOp::Exp => "exp",
UnaryOp::Abs => "abs",
UnaryOp::Sin => "sin",
UnaryOp::Cos => "cos",
UnaryOp::Log10 => "log10",
UnaryOp::Tan => "tan",
UnaryOp::Atan => "atan",
UnaryOp::Acos => "acos",
UnaryOp::Sinh => "sinh",
UnaryOp::Cosh => "cosh",
UnaryOp::Tanh => "tanh",
UnaryOp::Asin => "asin",
UnaryOp::Acosh => "acosh",
UnaryOp::Asinh => "asinh",
UnaryOp::Atanh => "atanh",
UnaryOp::Erf => "erf",
UnaryOp::XLogX => "xlogx",
}
}
fn cmp_sym(op: CmpOp) -> &'static str {
match op {
CmpOp::Lt => "<",
CmpOp::Le => "<=",
CmpOp::Eq => "==",
CmpOp::Ge => ">=",
CmpOp::Gt => ">",
CmpOp::Ne => "!=",
}
}
fn push_additive(out: &mut String, rendered: &str, first: bool) {
if first {
out.push_str(rendered);
} else if let Some(rest) = rendered.strip_prefix('-') {
out.push_str(" - ");
out.push_str(rest);
} else {
out.push_str(" + ");
out.push_str(rendered);
}
}
fn render_expr(e: &Expr, vn: &[String], funcs: &[ImportedFunc]) -> String {
match e {
Expr::Const(c) => fmt_num(*c),
Expr::Var(i) => var_label(*i, vn),
Expr::Binary(op, l, r) => match op {
BinOp::Add => {
let mut s = render_prec(l, P_ADD, vn, funcs);
push_additive(&mut s, &render_prec(r, P_ADD, vn, funcs), false);
s
}
BinOp::Sub => format!(
"{} - {}",
render_prec(l, P_ADD, vn, funcs),
render_prec(r, P_ADD + 1, vn, funcs)
),
BinOp::Mul => format!(
"{}*{}",
render_prec(l, P_MUL, vn, funcs),
render_prec(r, P_MUL, vn, funcs)
),
BinOp::Div => format!(
"{}/{}",
render_prec(l, P_MUL, vn, funcs),
render_prec(r, P_MUL + 1, vn, funcs)
),
BinOp::Pow => format!(
"{}^{}",
render_prec(l, P_POW + 1, vn, funcs),
render_prec(r, P_POW, vn, funcs)
),
BinOp::Atan2 => format!(
"atan2({}, {})",
render_expr(l, vn, funcs),
render_expr(r, vn, funcs)
),
BinOp::CEntropy => format!(
"centropy({}, {})",
render_expr(l, vn, funcs),
render_expr(r, vn, funcs)
),
},
Expr::Unary(UnaryOp::Neg, a) => format!("-{}", render_prec(a, P_NEG, vn, funcs)),
Expr::Unary(op, a) => format!("{}({})", unary_name(*op), render_expr(a, vn, funcs)),
Expr::Sum(xs) => {
if xs.is_empty() {
"0".to_string()
} else {
let mut s = String::new();
for (k, x) in xs.iter().enumerate() {
push_additive(&mut s, &render_prec(x, P_ADD, vn, funcs), k == 0);
}
s
}
}
Expr::Cse(inner) => render_expr(inner, vn, funcs),
Expr::Funcall { id, args } => {
let name = funcs
.iter()
.find(|f| f.id == *id)
.map(|f| f.name.clone())
.unwrap_or_else(|| format!("extern#{id}"));
let parts: Vec<String> = args
.iter()
.map(|a| match a {
FuncallArg::Real(x) => render_expr(x, vn, funcs),
FuncallArg::Str(s) => format!("{s:?}"),
})
.collect();
format!("{name}({})", parts.join(", "))
}
Expr::Compare(op, a, b) => format!(
"({} {} {})",
render_expr(a, vn, funcs),
cmp_sym(*op),
render_expr(b, vn, funcs)
),
Expr::And(a, b) => format!(
"({} && {})",
render_expr(a, vn, funcs),
render_expr(b, vn, funcs)
),
Expr::Or(a, b) => format!(
"({} || {})",
render_expr(a, vn, funcs),
render_expr(b, vn, funcs)
),
Expr::Not(a) => format!("!({})", render_expr(a, vn, funcs)),
Expr::Cond { cond, then_, else_ } => format!(
"if({}, {}, {})",
render_expr(cond, vn, funcs),
render_expr(then_, vn, funcs),
render_expr(else_, vn, funcs)
),
Expr::MinList(xs) => format!(
"min({})",
xs.iter()
.map(|x| render_expr(x, vn, funcs))
.collect::<Vec<_>>()
.join(", ")
),
Expr::MaxList(xs) => format!(
"max({})",
xs.iter()
.map(|x| render_expr(x, vn, funcs))
.collect::<Vec<_>>()
.join(", ")
),
}
}
fn render_linear(linear: &[(usize, Number)], vn: &[String]) -> String {
let mut out = String::new();
let mut first = true;
for (var, coef) in linear {
if *coef == 0.0 {
continue;
}
let neg = *coef < 0.0;
let mag = coef.abs();
let term = if mag == 1.0 {
var_label(*var, vn)
} else {
format!("{}*{}", fmt_num(mag), var_label(*var, vn))
};
if first {
if neg {
out.push('-');
}
out.push_str(&term);
first = false;
} else {
out.push_str(if neg { " - " } else { " + " });
out.push_str(&term);
}
}
out
}
fn render_body(linear: &[(usize, Number)], nonlinear: &Expr, prob: &NlProblem) -> String {
let mut s = render_linear(linear, &prob.var_names);
let nl_is_zero = matches!(nonlinear, Expr::Const(c) if *c == 0.0);
if !nl_is_zero {
let nl = render_prec(nonlinear, P_ADD, &prob.var_names, &prob.imported_funcs);
if s.is_empty() {
s = nl;
} else {
push_additive(&mut s, &nl, false);
}
}
if s.is_empty() {
s = "0".to_string();
}
s
}
pub fn render_constraint_equation(prob: &NlProblem, k: usize) -> String {
let body = render_body(&prob.con_linear[k], &prob.con_expr(k), prob);
let lo = prob.g_l[k];
let hi = prob.g_u[k];
const INF: Number = 1.0e19;
let has_lo = lo > -INF;
let has_hi = hi < INF;
match (has_lo, has_hi) {
(true, true) if lo == hi => format!("{body} = {}", fmt_num(lo)),
(true, true) => format!("{} <= {body} <= {}", fmt_num(lo), fmt_num(hi)),
(true, false) => format!("{body} >= {}", fmt_num(lo)),
(false, true) => format!("{body} <= {}", fmt_num(hi)),
(false, false) => format!("{body} (free)"),
}
}
pub fn render_all_constraint_equations(prob: &NlProblem) -> Vec<String> {
(0..prob.m)
.map(|k| render_constraint_equation(prob, k))
.collect()
}
pub fn constraint_jacobian_sparsity(prob: &NlProblem) -> (Vec<Index>, Vec<Index>) {
let mut irow: Vec<Index> = Vec::new();
let mut jcol: Vec<Index> = Vec::new();
let mut support: BTreeSet<usize> = BTreeSet::new();
for k in 0..prob.m {
support.clear();
for &(j, _coef) in &prob.con_linear[k] {
support.insert(j);
}
prob.con_nonlinear[k].collect_vars(&mut support);
for &j in &support {
irow.push(k as Index);
jcol.push(j as Index);
}
}
(irow, jcol)
}
fn split_top_sums(expr: &Expr) -> Vec<Expr> {
let mut out = Vec::new();
fn push_leaf(e: &Expr, factor: f64, out: &mut Vec<Expr>) {
if factor == 1.0 {
out.push(e.clone());
} else if factor == -1.0 {
out.push(Expr::Unary(UnaryOp::Neg, Box::new(e.clone())));
} else {
out.push(Expr::Binary(
BinOp::Mul,
Box::new(Expr::Const(factor)),
Box::new(e.clone()),
));
}
}
fn go(e: &Expr, factor: f64, out: &mut Vec<Expr>) {
match e {
Expr::Sum(terms) => {
for t in terms {
go(t, factor, out);
}
}
Expr::Binary(BinOp::Add, l, r) => {
go(l, factor, out);
go(r, factor, out);
}
Expr::Binary(BinOp::Sub, l, r) => {
go(l, factor, out);
go(r, -factor, out);
}
Expr::Unary(UnaryOp::Neg, x) => {
go(x, -factor, out);
}
Expr::Binary(BinOp::Mul, l, r) => match (l.as_ref(), r.as_ref()) {
(Expr::Const(c), _) => go(r, factor * c, out),
(_, Expr::Const(c)) => go(l, factor * c, out),
_ => push_leaf(e, factor, out),
},
Expr::Binary(BinOp::Div, l, r) => match r.as_ref() {
Expr::Const(c) if *c != 0.0 => go(l, factor / c, out),
_ => push_leaf(e, factor, out),
},
_ => push_leaf(e, factor, out),
}
}
go(expr, 1.0, &mut out);
if out.is_empty() {
out.push(Expr::Const(0.0));
}
out
}
const DENSE_COL_FACTOR: usize = 16;
const DENSE_COL_MIN: usize = 32;
const PEEL_MAX_REL_ERR: f64 = 1e-8;
const MAX_PEELED_COLS: usize = 256;
fn peel_color_bound(n: usize, lower_pairs: &[(usize, usize)], peeled: &[bool]) -> usize {
let mut deg = vec![0usize; n];
for &(i, j) in lower_pairs {
if peeled[i] || peeled[j] {
continue;
}
deg[j] += 1;
if i != j {
deg[i] += 1;
}
}
let n_peeled = peeled.iter().filter(|&&p| p).count();
n_peeled + deg.iter().copied().max().unwrap_or(0)
}
fn select_peeled_cols(
n: usize,
lower_pairs: &[(usize, usize)],
deg: &[usize],
mut candidates: Vec<usize>,
) -> Vec<usize> {
if candidates.is_empty() {
return candidates;
}
candidates.sort_unstable_by(|&a, &b| deg[b].cmp(°[a]).then(a.cmp(&b)));
candidates.truncate(MAX_PEELED_COLS);
let mut mask = vec![false; n];
let mut best_k = 0usize;
let mut best_bound = peel_color_bound(n, lower_pairs, &mask);
let mut marked = 0usize;
let mut k = 1usize;
loop {
let k_now = k.min(candidates.len());
for &j in &candidates[marked..k_now] {
mask[j] = true;
}
marked = k_now;
let bound = peel_color_bound(n, lower_pairs, &mask);
if bound < best_bound {
best_bound = bound;
best_k = k_now;
}
if k_now == candidates.len() {
break;
}
k *= 2;
}
candidates.truncate(best_k);
candidates
}
fn greedy_hessian_coloring(
n: usize,
lower_pairs: &[(usize, usize)],
peel_veto: &[bool],
) -> (Vec<u32>, usize, Vec<bool>) {
if n == 0 {
return (Vec::new(), 0, Vec::new());
}
let mut deg = vec![0usize; n];
for &(i, j) in lower_pairs {
deg[j] += 1;
if i != j {
deg[i] += 1;
}
}
let total: usize = deg.iter().sum();
let threshold = DENSE_COL_MIN.max(DENSE_COL_FACTOR.saturating_mul(total / n));
let mut peeled = vec![false; n];
let candidates: Vec<usize> = (0..n)
.filter(|&j| deg[j] > threshold && !peel_veto.get(j).copied().unwrap_or(false))
.collect();
let dense = select_peeled_cols(n, lower_pairs, °, candidates);
for &j in &dense {
peeled[j] = true;
}
let mut col_rows: Vec<Vec<u32>> = vec![Vec::new(); n];
let mut row_cols: Vec<Vec<u32>> = vec![Vec::new(); n];
for &(i, j) in lower_pairs {
if peeled[i] || peeled[j] {
continue;
}
col_rows[j].push(i as u32);
row_cols[i].push(j as u32);
if i != j {
col_rows[i].push(j as u32);
row_cols[j].push(i as u32);
}
}
let mut var_color = vec![u32::MAX; n];
let mut forbidden = vec![u32::MAX; n + 1];
let mut n_colors: u32 = 0;
for j in 0..n {
if peeled[j] || col_rows[j].is_empty() {
continue;
}
for &r in &col_rows[j] {
for &c in &row_cols[r as usize] {
if c as usize == j {
continue;
}
let cc = var_color[c as usize];
if cc != u32::MAX {
forbidden[cc as usize] = j as u32;
}
}
}
let mut chosen: u32 = 0;
while (chosen as usize) < forbidden.len() && forbidden[chosen as usize] == j as u32 {
chosen += 1;
}
var_color[j] = chosen;
if chosen + 1 > n_colors {
n_colors = chosen + 1;
}
}
for &j in &dense {
var_color[j] = n_colors;
n_colors += 1;
}
(var_color, n_colors as usize, peeled)
}
fn build_color_tables(
n: usize,
m: usize,
lower_pairs: &[(usize, usize)],
tape_mask: &[bool],
peel_veto: &[bool],
obj_tapes: &[Tape],
con_tapes: &[Vec<Tape>],
con_hybrid: Option<&mut ConHybrid>,
) -> ColorTables {
let colored_pairs: Vec<(usize, usize)>;
let color_input: &[(usize, usize)] = if tape_mask.is_empty() {
lower_pairs
} else {
colored_pairs = lower_pairs
.iter()
.zip(tape_mask)
.filter_map(|(p, &t)| t.then_some(*p))
.collect();
&colored_pairs
};
let (var_color, n_colors, peeled) = greedy_hessian_coloring(n, color_input, peel_veto);
let mut seeds: Vec<Vec<f64>> = vec![vec![0.0; n]; n_colors];
for (k, &c) in var_color.iter().enumerate() {
if c != u32::MAX {
seeds[c as usize][k] = 1.0;
}
}
let mut decoding: Vec<Vec<ColorWrite>> = vec![Vec::new(); n_colors];
for (idx, &(i, j)) in lower_pairs.iter().enumerate() {
if !tape_mask.is_empty() && !tape_mask[idx] {
continue;
}
let (c, row) = if peeled[i] && !peeled[j] {
(var_color[i], j)
} else {
(var_color[j], i)
};
debug_assert!(
c != u32::MAX,
"Hessian pair ({i}, {j}) at index {idx} has no color"
);
decoding[c as usize].push(ColorWrite {
row: row as u32,
hess_idx: idx as u32,
});
}
let tape_colors = |t: &Tape| -> Vec<u32> {
let mut s: Vec<u32> = t
.variables()
.into_iter()
.map(|v| var_color[v])
.filter(|&c| c != u32::MAX)
.collect();
s.sort_unstable();
s.dedup();
s
};
let obj_tape_colors: Vec<Vec<u32>> = obj_tapes.iter().map(tape_colors).collect();
let con_tape_colors: Vec<Vec<Vec<u32>>> = con_tapes
.iter()
.map(|row| row.iter().map(tape_colors).collect())
.collect();
if let Some(h) = con_hybrid {
let n_sum = h.tape.n_summands();
let mut local_off: Vec<usize> = Vec::with_capacity(n_sum + 1);
let mut acc = 0usize;
for s in &h.tape.summands {
local_off.push(acc);
acc += s.ops.len();
}
local_off.push(acc);
h.local_vals_all = vec![0.0; acc];
h.local_off = local_off;
let mut summand_row = vec![0u32; n_sum];
for i in 0..m {
for si in h.row_start[i]..h.row_start[i + 1] {
summand_row[si] = i as u32;
}
}
h.summand_row = summand_row;
let mut by_color: Vec<Vec<u32>> = vec![Vec::new(); n_colors];
for (si, s) in h.tape.summands.iter().enumerate() {
let mut cs: Vec<u32> = s
.all_vars
.iter()
.map(|&v| var_color[v])
.filter(|&c| c != u32::MAX)
.collect();
cs.sort_unstable();
cs.dedup();
for c in cs {
by_color[c as usize].push(si as u32);
}
}
let np = h.tape.n_prelude_ops();
let mut seen: Vec<u32> = vec![0; np];
let mut epoch: u32 = 0;
let mut reach: Vec<u32> = Vec::new();
let mut reach_off: Vec<usize> = Vec::with_capacity(n_colors + 1);
for list in &by_color {
reach_off.push(reach.len());
epoch += 1;
let start = reach.len();
for &si in list {
for &p in &h.tape.summands[si as usize].prelude_reach {
if seen[p] != epoch {
seen[p] = epoch;
reach.push(p as u32);
}
}
}
reach[start..].sort_unstable();
}
reach_off.push(reach.len());
h.hess_color_reach = reach;
h.hess_color_reach_off = reach_off;
h.hess_color_summands = by_color;
h.prelude_dot = vec![0.0; h.tape.n_prelude_ops()];
h.hess_prelude_adj = vec![0.0; h.tape.n_prelude_ops()];
h.prelude_adj_dot = vec![0.0; h.tape.n_prelude_ops()];
h.local_dot = vec![0.0; h.tape.max_summand_ops()];
h.local_adj_dot = vec![0.0; h.tape.max_summand_ops()];
}
ColorTables {
var_color,
n_colors,
peeled_cols: peeled
.iter()
.enumerate()
.filter(|(_, p)| **p)
.map(|(j, _)| j as u32)
.collect(),
seeds,
decoding,
obj_tape_colors,
con_tape_colors,
}
}
struct ColorTables {
var_color: Vec<u32>,
n_colors: usize,
peeled_cols: Vec<u32>,
seeds: Vec<Vec<f64>>,
decoding: Vec<Vec<ColorWrite>>,
obj_tape_colors: Vec<Vec<u32>>,
con_tape_colors: Vec<Vec<Vec<u32>>>,
}
impl NlTnlp {
pub fn new(prob: NlProblem) -> Self {
Self::try_new(prob)
.unwrap_or_else(|e| panic!("failed to resolve AMPL external functions: {e}"))
}
pub fn try_new(prob: NlProblem) -> Result<Self, String> {
Self::try_new_with_quadratic(prob, std::env::var("POUNCE_DBG_NO_QUAD").is_err())
}
pub fn try_new_with_quadratic(prob: NlProblem, use_quadratic: bool) -> Result<Self, String> {
let mut referenced: BTreeSet<usize> = BTreeSet::new();
for body in std::iter::once(&prob.obj_nonlinear).chain(prob.con_nonlinear.iter()) {
if let Some(e) = body.tree() {
super::nl_external::collect_funcall_ids(e, &mut referenced);
}
}
let resolver = if referenced.is_empty() {
super::nl_external::ExternalResolver::default()
} else {
super::nl_external::ExternalResolver::build_for_problem(
&prob.imported_funcs,
&referenced,
)?
};
let mut quad = QuadraticStructure::new(prob.m);
if use_quadratic {
if let Some(f) = push_body_form(&mut quad, &prob.obj_nonlinear) {
quad.assign_objective(f);
}
for k in 0..prob.m {
if let Some(f) = push_body_form(&mut quad, &prob.con_nonlinear[k]) {
quad.assign_row(k, f);
}
}
}
let obj_tapes: Vec<Tape> = if quad.objective_form().is_some() {
Vec::new()
} else {
split_top_sums(&prob.obj_expr())
.iter()
.map(|e| Tape::build_with_externals(e, &resolver))
.collect()
};
let mut con_tapes: Vec<Vec<Tape>> = Vec::with_capacity(prob.m);
let mut con_roots: Vec<Expr> = Vec::new();
let mut row_start: Vec<usize> = Vec::with_capacity(prob.m + 1);
for k in 0..prob.m {
row_start.push(con_roots.len());
if quad.row_form(k).is_some() {
con_tapes.push(Vec::new());
continue;
}
let summands = split_top_sums(&prob.con_expr(k));
con_tapes.push(
summands
.iter()
.map(|e| Tape::build_with_externals(e, &resolver))
.collect(),
);
con_roots.extend(summands);
}
row_start.push(con_roots.len());
let mut con_hybrid = if std::env::var("POUNCE_DBG_NO_HYBRID").is_ok() {
None
} else if hybrid_supported(&con_roots) {
let tape = HybridTape::build_multi(&con_roots);
(tape.n_prelude_ops() > 0).then(|| {
let flat_ops: usize = con_tapes.iter().flatten().map(|t| t.ops.len()).sum();
let shared_ops = tape.n_prelude_ops() + tape.total_local_ops();
let force_hess = std::env::var("POUNCE_DBG_FORCE_HYBRID_HESS").is_ok();
ConHybrid {
prelude_vals: vec![0.0; tape.n_prelude_ops()],
local_vals: vec![0.0; tape.max_summand_ops()],
local_adj: vec![0.0; tape.max_summand_ops()],
prelude_adj: vec![0.0; tape.n_prelude_ops()],
use_for_jac: flat_ops as f64
>= HYBRID_JAC_MIN_OP_RATIO * shared_ops.max(1) as f64,
use_for_hess: force_hess
|| flat_ops as f64 >= HYBRID_HESS_MIN_OP_RATIO * shared_ops.max(1) as f64,
local_vals_all: Vec::new(),
local_off: Vec::new(),
summand_row: Vec::new(),
hess_color_summands: Vec::new(),
hess_color_reach: Vec::new(),
hess_color_reach_off: Vec::new(),
prelude_dot: Vec::new(),
hess_prelude_adj: Vec::new(),
prelude_adj_dot: Vec::new(),
local_dot: Vec::new(),
local_adj_dot: Vec::new(),
row_start,
tape,
}
})
} else {
None
};
drop(con_roots);
let mut tape_pairs: Vec<(usize, usize)> = Vec::new();
for t in &obj_tapes {
tape_pairs.extend(t.hessian_sparsity());
}
for row in &con_tapes {
for t in row {
tape_pairs.extend(t.hessian_sparsity());
}
}
tape_pairs.sort_unstable();
tape_pairs.dedup();
let mut lower_pairs = tape_pairs.clone();
if !quad.is_empty() {
for f in quad
.objective_form()
.into_iter()
.chain((0..prob.m).filter_map(|i| quad.row_form(i)))
{
lower_pairs.extend(
quad.lower_triangle(f)
.map(|(r, c, _)| (r as usize, c as usize)),
);
}
lower_pairs.sort_unstable();
lower_pairs.dedup();
}
let h_tape_mask: Vec<bool> = if quad.is_empty() {
Vec::new()
} else {
let mut mask = vec![false; lower_pairs.len()];
let mut t = 0usize;
for (idx, pair) in lower_pairs.iter().enumerate() {
if t < tape_pairs.len() && tape_pairs[t] == *pair {
mask[idx] = true;
t += 1;
}
}
debug_assert_eq!(t, tape_pairs.len(), "every tape pair is in the union");
mask
};
drop(tape_pairs);
let mut h_irow = Vec::with_capacity(lower_pairs.len());
let mut h_jcol = Vec::with_capacity(lower_pairs.len());
for &(hi, lo) in &lower_pairs {
h_irow.push(hi as i32);
h_jcol.push(lo as i32);
}
if !quad.is_empty() {
quad.bind_slots(|r, c| {
lower_pairs
.binary_search(&(r as usize, c as usize))
.unwrap_or_else(|_| {
unreachable!("quadratic entry ({r}, {c}) missing from the union pattern")
})
});
}
let ColorTables {
var_color,
n_colors,
peeled_cols,
seeds,
decoding,
obj_tape_colors,
con_tape_colors,
} = build_color_tables(
prob.n,
prob.m,
&lower_pairs,
&h_tape_mask,
&vec![false; prob.n],
&obj_tapes,
&con_tapes,
con_hybrid.as_mut(),
);
let mut jac_cols: Vec<Vec<usize>> = Vec::with_capacity(prob.m);
let mut jac_nnz = 0;
for (i, row_tapes) in con_tapes.iter().enumerate() {
let mut cols: Vec<usize> = Vec::with_capacity(prob.con_linear[i].len());
for t in row_tapes {
cols.extend(t.variables());
}
if let Some(f) = quad.row_form(i) {
cols.extend(quad.gradient_support(f).iter().map(|&v| v as usize));
}
cols.extend(prob.con_linear[i].iter().map(|(v, _)| *v));
cols.sort_unstable();
cols.dedup();
cols.shrink_to_fit();
jac_nnz += cols.len();
jac_cols.push(cols);
}
let mut max_tape_n: usize = 0;
for t in &obj_tapes {
max_tape_n = max_tape_n.max(t.ops.len());
}
for row in &con_tapes {
for t in row {
max_tape_n = max_tape_n.max(t.ops.len());
}
}
if std::env::var("POUNCE_DBG_TAPE_STATS").is_ok() {
let n_obj = obj_tapes.len();
let n_con: usize = con_tapes.iter().map(|r| r.len()).sum();
let total = n_obj + n_con;
let mut sum_ops: usize = 0;
for t in &obj_tapes {
sum_ops += t.ops.len();
}
for row in &con_tapes {
for t in row {
sum_ops += t.ops.len();
}
}
let t = total.max(1);
let nnz_h = h_irow.len();
let avg_decode =
decoding.iter().map(|d| d.len()).sum::<usize>() as f64 / n_colors.max(1) as f64;
eprintln!(
"[tape stats] summands={total} (obj={n_obj} con={n_con}) \
total_ops={sum_ops} avg_ops={:.1} max_ops={max_tape_n} \
n_colors={n_colors} avg_decode_per_color={avg_decode:.1} nnz_h={nnz_h}",
sum_ops as f64 / t as f64,
);
match &con_hybrid {
Some(h) => {
let flat: usize = con_tapes.iter().flatten().map(|t| t.ops.len()).sum();
let prelude = h.tape.n_prelude_ops();
let local = h.tape.total_local_ops();
eprintln!(
"[hybrid stats] con flat_ops={flat} prelude_ops={prelude} \
local_ops={local} shared_total={} flat/shared={:.2}x \
jac_gate={} hess_gate={}",
prelude + local,
flat as f64 / (prelude + local).max(1) as f64,
if h.use_for_jac { "on" } else { "off" },
if h.use_for_hess { "on" } else { "off" },
);
}
None => eprintln!("[hybrid stats] con hybrid not built (no shared CSE bodies)"),
}
let quad_rows = (0..prob.m).filter(|&i| quad.row_form(i).is_some()).count();
let parsed = std::iter::once(&prob.obj_nonlinear)
.chain(prob.con_nonlinear.iter())
.filter(|b| b.quad().is_some())
.count();
eprintln!(
"[quad stats] forms={} (parse-time {parsed}) rows={quad_rows}/{} obj={} \
stored_h_entries={} colored_pairs={}/{}",
quad.len(),
prob.m,
if quad.objective_form().is_some() {
"quadratic"
} else {
"taped"
},
quad.stored_entries(),
if h_tape_mask.is_empty() {
lower_pairs.len()
} else {
h_tape_mask.iter().filter(|&&t| t).count()
},
lower_pairs.len(),
);
}
let compressed: Vec<Vec<f64>> = vec![vec![0.0; prob.n]; n_colors];
let mut me = Self {
prob,
obj_tapes,
con_tapes,
con_hybrid,
quad,
h_tape_mask,
h_irow,
h_jcol,
jac_cols,
jac_nnz,
seeds,
decoding,
obj_tape_colors,
con_tape_colors,
var_color,
peeled_cols,
final_x: None,
final_obj: 0.0,
final_lambda: None,
final_z_l: None,
final_z_u: None,
scratch_row_grad: Vec::new(),
vals_scratch: vec![0.0; max_tape_n],
dot_scratch: vec![0.0; max_tape_n],
adj_scratch: vec![0.0; max_tape_n],
adj_dot_scratch: vec![0.0; max_tape_n],
compressed,
hvp_live: Vec::new(),
curvature_scaling: None,
};
me.veto_ill_conditioned_peels();
Ok(me)
}
fn veto_ill_conditioned_peels(&mut self) {
if self.peeled_cols.is_empty() {
return;
}
let mut values = vec![0.0; self.h_irow.len()];
let lambda = vec![1.0; self.prob.m];
let x0 = self.prob.x0.clone();
if !self.eval_h(
Some(&x0),
true,
1.0,
Some(&lambda),
true,
SparsityRequest::Values {
values: &mut values,
},
) {
return;
}
let dbg = std::env::var("POUNCE_DBG_TAPE_STATS").is_ok();
let h_scale = self
.compressed
.iter()
.flat_map(|c| c.iter())
.fold(0.0f64, |a, &v| a.max(v.abs()));
let mut peel_veto = vec![false; self.prob.n];
let mut vetoed = 0usize;
for &d in &self.peeled_cols {
let c = self.var_color[d as usize];
if c == u32::MAX {
continue;
}
let pass = &self.compressed[c as usize];
let scale = pass.iter().fold(0.0f64, |a, &v| a.max(v.abs()));
let noise = f64::EPSILON * scale;
let smallest = self.decoding[c as usize]
.iter()
.map(|w| pass[w.row as usize].abs())
.filter(|v| *v > noise)
.fold(f64::INFINITY, f64::min);
if !smallest.is_finite() || smallest == 0.0 || scale == 0.0 {
continue;
}
if scale <= h_scale * f64::EPSILON {
continue;
}
let rel_err = f64::EPSILON * scale / smallest;
if dbg {
eprintln!(
"[peel probe] col={d} color={c} ||pass||={scale:.3e} \
min_entry={smallest:.3e} rel_err={rel_err:.3e}{}",
if rel_err > PEEL_MAX_REL_ERR {
" VETO"
} else {
""
}
);
}
if rel_err > PEEL_MAX_REL_ERR {
peel_veto[d as usize] = true;
vetoed += 1;
}
}
if vetoed == 0 {
return;
}
if dbg {
eprintln!(
"[peel probe] vetoing {vetoed}/{} peeled columns; re-coloring",
self.peeled_cols.len()
);
}
self.recolor(&peel_veto);
}
fn recolor(&mut self, peel_veto: &[bool]) {
let lower_pairs: Vec<(usize, usize)> = self
.h_irow
.iter()
.zip(&self.h_jcol)
.map(|(&i, &j)| (i as usize, j as usize))
.collect();
let ColorTables {
var_color,
n_colors,
peeled_cols,
seeds,
decoding,
obj_tape_colors,
con_tape_colors,
} = build_color_tables(
self.prob.n,
self.prob.m,
&lower_pairs,
&self.h_tape_mask,
peel_veto,
&self.obj_tapes,
&self.con_tapes,
self.con_hybrid.as_mut(),
);
self.var_color = var_color;
self.peeled_cols = peeled_cols;
self.seeds = seeds;
self.decoding = decoding;
self.obj_tape_colors = obj_tape_colors;
self.con_tape_colors = con_tape_colors;
self.compressed = vec![vec![0.0; self.prob.n]; n_colors];
}
pub fn final_x(&self) -> Option<&[Number]> {
self.final_x.as_deref()
}
pub fn final_obj(&self) -> Number {
self.final_obj
}
pub fn final_lambda(&self) -> Option<&[Number]> {
self.final_lambda.as_deref()
}
pub fn final_bound_multipliers(&self) -> Option<(&[Number], &[Number])> {
Some((self.final_z_l.as_deref()?, self.final_z_u.as_deref()?))
}
pub fn problem(&self) -> &NlProblem {
&self.prob
}
pub fn enable_curvature_scaling(&mut self) -> bool {
match crate::nl_scaling::curvature_scaling(&self.prob) {
Some(sc) => {
self.curvature_scaling = Some(sc);
true
}
None => false,
}
}
pub fn curvature_scaling_enabled(&self) -> bool {
self.curvature_scaling.is_some()
}
pub fn curvature_scaling_read_curvature(&self) -> bool {
self.curvature_scaling
.as_ref()
.is_some_and(|sc| sc.quadratic)
}
pub fn quadratic_row(&self, i: usize) -> bool {
self.quad.row_form(i).is_some()
}
pub fn quadratic_objective(&self) -> bool {
self.quad.objective_form().is_some()
}
fn nonlinear_var_set(&self) -> BTreeSet<usize> {
let mut nonlinear: BTreeSet<usize> = BTreeSet::new();
self.prob.obj_nonlinear.collect_vars(&mut nonlinear);
for row in &self.prob.con_nonlinear {
row.collect_vars(&mut nonlinear);
}
nonlinear
}
fn nonlinear_variables(&self) -> Vec<Index> {
if let Some(c) = self.prob.nl_counts
&& c.nonlinear_vars() >= self.prob.n
{
return (0..self.prob.n as Index).collect();
}
self.nonlinear_var_set()
.into_iter()
.map(|i| i as Index)
.collect()
}
pub fn problem_mut(&mut self) -> &mut NlProblem {
&mut self.prob
}
pub fn hessian_vector_product(
&mut self,
x: &[Number],
v: &[Number],
obj_factor: Number,
lambda: Option<&[Number]>,
out: &mut [Number],
) -> Result<(), String> {
self.hessian_vector_products(x, v, 1, obj_factor, lambda, out)
}
pub fn hessian_vector_products(
&mut self,
x: &[Number],
v: &[Number],
k: usize,
obj_factor: Number,
lambda: Option<&[Number]>,
out: &mut [Number],
) -> Result<(), String> {
let (n, m) = (self.prob.n, self.prob.m);
let check = |name: &str, got: usize, want: usize| -> Result<(), String> {
if got == want {
Ok(())
} else {
Err(format!(
"hessian_vector_product: {name} has length {got}, expected {want}"
))
}
};
check("x", x.len(), n)?;
check("v", v.len(), n * k)?;
check("out", out.len(), n * k)?;
if let Some(lam) = lambda {
check("lambda", lam.len(), m)?;
}
out.fill(0.0);
if k == 0 || n == 0 {
return Ok(());
}
self.hvp_live.clear();
self.hvp_live
.extend((0..k).map(|c| v[c * n..(c + 1) * n].iter().any(|&s| s != 0.0)));
if !self.hvp_live.iter().any(|&l| l) {
return Ok(());
}
let obj_seed = if self.prob.minimize {
obj_factor
} else {
-obj_factor
};
if !self.quad.is_empty() {
if obj_seed != 0.0 {
if let Some(f) = self.quad.objective_form() {
for (c, out_col) in out.chunks_mut(n).enumerate() {
if self.hvp_live[c] {
self.quad.add_hessian_vector(
f,
&v[c * n..(c + 1) * n],
obj_seed,
out_col,
);
}
}
}
}
if let Some(lam) = lambda {
for (i, &w) in lam.iter().enumerate() {
if w == 0.0 {
continue;
}
let Some(f) = self.quad.row_form(i) else {
continue;
};
for (c, out_col) in out.chunks_mut(n).enumerate() {
if self.hvp_live[c] {
self.quad
.add_hessian_vector(f, &v[c * n..(c + 1) * n], w, out_col);
}
}
}
}
}
if obj_seed != 0.0 {
for t in &self.obj_tapes {
if t.ops.is_empty() {
continue;
}
t.forward_into(x, &mut self.vals_scratch);
for (c, out_col) in out.chunks_mut(n).enumerate() {
if !self.hvp_live[c] {
continue;
}
t.hessian_directional(
&self.vals_scratch,
&v[c * n..(c + 1) * n],
obj_seed,
out_col,
&mut self.dot_scratch,
&mut self.adj_scratch,
&mut self.adj_dot_scratch,
);
}
}
}
if let Some(lam) = lambda {
for (i, &w) in lam.iter().enumerate() {
if w == 0.0 {
continue;
}
for t in &self.con_tapes[i] {
if t.ops.is_empty() {
continue;
}
t.forward_into(x, &mut self.vals_scratch);
for (c, out_col) in out.chunks_mut(n).enumerate() {
if !self.hvp_live[c] {
continue;
}
t.hessian_directional(
&self.vals_scratch,
&v[c * n..(c + 1) * n],
w,
out_col,
&mut self.dot_scratch,
&mut self.adj_scratch,
&mut self.adj_dot_scratch,
);
}
}
}
}
Ok(())
}
pub fn variant(&self, v: &NlVariation) -> Result<Self, String> {
let check = |name: &str, got: usize, want: usize| -> Result<(), String> {
if got == want {
Ok(())
} else {
Err(format!(
"NlVariation.{name} has length {got}, expected {want}"
))
}
};
let mut out = self.clone();
out.final_x = None;
out.final_obj = 0.0;
out.final_lambda = None;
out.final_z_l = None;
out.final_z_u = None;
if let Some(x0) = &v.x0 {
check("x0", x0.len(), self.prob.n)?;
out.prob.x0.clone_from(x0);
}
if let Some(x_l) = &v.x_l {
check("x_l", x_l.len(), self.prob.n)?;
out.prob.x_l.clone_from(x_l);
}
if let Some(x_u) = &v.x_u {
check("x_u", x_u.len(), self.prob.n)?;
out.prob.x_u.clone_from(x_u);
}
if let Some(g_l) = &v.g_l {
check("g_l", g_l.len(), self.prob.m)?;
out.prob.g_l.clone_from(g_l);
}
if let Some(g_u) = &v.g_u {
check("g_u", g_u.len(), self.prob.m)?;
out.prob.g_u.clone_from(g_u);
}
Ok(out)
}
pub fn variants(&self, vs: &[NlVariation]) -> Result<Vec<Self>, String> {
vs.iter().map(|v| self.variant(v)).collect()
}
}
#[derive(Debug, Clone, Default)]
pub struct NlVariation {
pub x0: Option<Vec<Number>>,
pub x_l: Option<Vec<Number>>,
pub x_u: Option<Vec<Number>>,
pub g_l: Option<Vec<Number>>,
pub g_u: Option<Vec<Number>>,
}
impl pounce_nlp::expression_provider::ExpressionProvider for NlTnlp {
fn constraint_expression(&self, i: usize) -> Option<pounce_nlp::FbbtTape> {
if i >= self.prob.con_nonlinear.len() {
return None;
}
let nonlinear = self.prob.con_expr(i);
let linear = self
.prob
.con_linear
.get(i)
.map(|v| v.as_slice())
.unwrap_or(&[]);
crate::nl_fbbt_translate::translate_constraint(&nonlinear, linear)
}
fn variable_name(&self, i: usize) -> Option<&str> {
self.prob.var_names.get(i).map(String::as_str)
}
fn constraint_name(&self, i: usize) -> Option<&str> {
self.prob.con_names.get(i).map(String::as_str)
}
}
impl TNLP for NlTnlp {
fn get_nlp_info(&mut self) -> Option<NlpInfo> {
Some(NlpInfo {
n: self.prob.n as Index,
m: self.prob.m as Index,
nnz_jac_g: self.jac_nnz as Index,
nnz_h_lag: self.h_irow.len() as Index,
index_style: IndexStyle::C,
})
}
fn get_bounds_info(&mut self, b: BoundsInfo<'_>) -> bool {
b.x_l.copy_from_slice(&self.prob.x_l);
b.x_u.copy_from_slice(&self.prob.x_u);
if !self.prob.g_l.is_empty() {
b.g_l.copy_from_slice(&self.prob.g_l);
b.g_u.copy_from_slice(&self.prob.g_u);
}
true
}
fn get_starting_point(&mut self, sp: StartingPoint<'_>) -> bool {
sp.x.copy_from_slice(&self.prob.x0);
if sp.init_lambda {
sp.lambda.copy_from_slice(&self.prob.lambda0);
}
true
}
fn get_scaling_parameters(&mut self, req: ScalingRequest<'_>) -> bool {
const NAME: &str = "scaling_factor";
let sfx = &self.prob.suffixes;
let obj = sfx.obj_real.get(NAME);
let var = sfx.var_real.get(NAME);
let con = sfx.con_real.get(NAME);
let computed = self.curvature_scaling.as_ref();
if obj.is_none() && var.is_none() && con.is_none() && computed.is_none() {
return false;
}
if let Some(sc) = computed {
debug_assert_eq!(
sc.x.len(),
req.x_scaling.len(),
"curvature x-scaling sized {} for a {}-variable request",
sc.x.len(),
req.x_scaling.len()
);
debug_assert_eq!(
sc.g.len(),
req.g_scaling.len(),
"curvature g-scaling sized {} for a {}-row request",
sc.g.len(),
req.g_scaling.len()
);
if sc.x.len() == req.x_scaling.len() {
req.x_scaling.copy_from_slice(&sc.x);
*req.use_x_scaling = true;
}
if sc.g.len() == req.g_scaling.len() {
req.g_scaling.copy_from_slice(&sc.g);
*req.use_g_scaling = true;
}
}
*req.obj_scaling = obj
.and_then(|v| v.first().copied())
.filter(|&s| s != 0.0)
.unwrap_or(1.0);
if let Some(v) = var.filter(|v| v.len() == req.x_scaling.len()) {
for (slot, &s) in req.x_scaling.iter_mut().zip(v) {
if s != 0.0 {
*slot = s;
} else if computed.is_none() {
*slot = 1.0;
}
}
*req.use_x_scaling = true;
} else if computed.is_none() {
*req.use_x_scaling = false;
}
if let Some(g) = con.filter(|g| g.len() == req.g_scaling.len()) {
for (slot, &s) in req.g_scaling.iter_mut().zip(g) {
if s != 0.0 {
*slot = s;
} else if computed.is_none() {
*slot = 1.0;
}
}
*req.use_g_scaling = true;
} else if computed.is_none() {
*req.use_g_scaling = false;
}
true
}
fn eval_f(&mut self, x: &[Number], _new_x: bool) -> Option<Number> {
let (obj_tapes, vals) = (&self.obj_tapes, &mut self.vals_scratch);
let mut nl: Number = 0.0;
for t in obj_tapes {
nl += t.eval_into(x, vals);
}
if let Some(f) = self.quad.objective_form() {
nl += self.quad.value(f, x);
}
let lin: Number = self.prob.obj_linear.iter().map(|(i, c)| c * x[*i]).sum();
let v = self.prob.obj_constant + nl + lin;
let signed = if self.prob.minimize { v } else { -v };
Some(signed)
}
fn eval_grad_f(&mut self, x: &[Number], _new_x: bool, grad: &mut [Number]) -> bool {
grad.fill(0.0);
for t in &self.obj_tapes {
t.gradient_seed_into(x, 1.0, grad, &mut self.vals_scratch, &mut self.adj_scratch);
}
if let Some(f) = self.quad.objective_form() {
self.quad.add_gradient(f, x, 1.0, grad);
}
for (i, c) in &self.prob.obj_linear {
grad[*i] += c;
}
if !self.prob.minimize {
for g in grad.iter_mut() {
*g = -*g;
}
}
true
}
fn eval_g(&mut self, x: &[Number], _new_x: bool, g: &mut [Number]) -> bool {
let m = self.prob.m;
let con_linear = &self.prob.con_linear;
let quad = &self.quad;
if let Some(h) = &mut self.con_hybrid {
let ConHybrid {
tape,
row_start,
prelude_vals,
local_vals,
..
} = h;
tape.forward_prelude(x, prelude_vals);
for i in 0..m {
let mut nl: Number = 0.0;
for s in &tape.summands[row_start[i]..row_start[i + 1]] {
tape.forward_summand(s, x, prelude_vals, local_vals);
nl += tape.root_value(s, local_vals);
}
if let Some(f) = quad.row_form(i) {
nl += quad.value(f, x);
}
let lin: Number = con_linear[i].iter().map(|(j, c)| c * x[*j]).sum();
g[i] = nl + lin;
}
return true;
}
let (con_tapes, vals) = (&self.con_tapes, &mut self.vals_scratch);
for i in 0..m {
let mut nl: Number = 0.0;
for t in &con_tapes[i] {
nl += t.eval_into(x, vals);
}
if let Some(f) = quad.row_form(i) {
nl += quad.value(f, x);
}
let lin: Number = con_linear[i].iter().map(|(j, c)| c * x[*j]).sum();
g[i] = nl + lin;
}
true
}
fn eval_jac_g(
&mut self,
x: Option<&[Number]>,
_new_x: bool,
mode: SparsityRequest<'_>,
) -> bool {
match mode {
SparsityRequest::Structure { irow, jcol } => {
let mut k = 0;
for i in 0..self.prob.m {
for &j in &self.jac_cols[i] {
irow[k] = i as Index;
jcol[k] = j as Index;
k += 1;
}
}
true
}
SparsityRequest::Values { values } => {
let n = self.prob.n;
if self.scratch_row_grad.len() < n {
self.scratch_row_grad.resize(n, 0.0);
}
let Self {
prob,
con_tapes,
con_hybrid,
quad,
jac_cols,
scratch_row_grad,
vals_scratch,
adj_scratch,
..
} = self;
let xs = x.unwrap_or(&prob.x0);
let mut k = 0;
if let Some(h) = con_hybrid.as_mut().filter(|h| h.use_for_jac) {
let ConHybrid {
tape,
row_start,
prelude_vals,
local_vals,
local_adj,
prelude_adj,
..
} = h;
tape.forward_prelude(xs, prelude_vals);
for i in 0..prob.m {
for &j in &jac_cols[i] {
scratch_row_grad[j] = 0.0;
}
for s in &tape.summands[row_start[i]..row_start[i + 1]] {
tape.forward_summand(s, xs, prelude_vals, local_vals);
tape.gradient_summand(
s,
prelude_vals,
local_vals,
1.0,
scratch_row_grad,
local_adj,
prelude_adj,
);
}
if let Some(f) = quad.row_form(i) {
quad.add_gradient(f, xs, 1.0, scratch_row_grad);
}
for &(v, c) in &prob.con_linear[i] {
scratch_row_grad[v] += c;
}
for &j in &jac_cols[i] {
values[k] = scratch_row_grad[j];
k += 1;
}
}
return true;
}
for i in 0..prob.m {
for &j in &jac_cols[i] {
scratch_row_grad[j] = 0.0;
}
for t in &con_tapes[i] {
t.gradient_seed_into(xs, 1.0, scratch_row_grad, vals_scratch, adj_scratch);
}
if let Some(f) = quad.row_form(i) {
quad.add_gradient(f, xs, 1.0, scratch_row_grad);
}
for &(v, c) in &prob.con_linear[i] {
scratch_row_grad[v] += c;
}
for &j in &jac_cols[i] {
values[k] = scratch_row_grad[j];
k += 1;
}
}
true
}
}
}
fn eval_h(
&mut self,
x: Option<&[Number]>,
_new_x: bool,
obj_factor: Number,
lambda: Option<&[Number]>,
_new_lambda: bool,
mode: SparsityRequest<'_>,
) -> bool {
match mode {
SparsityRequest::Structure { irow, jcol } => {
irow.copy_from_slice(&self.h_irow);
jcol.copy_from_slice(&self.h_jcol);
true
}
SparsityRequest::Values { values } => {
let x = x.unwrap_or(&self.prob.x0);
values.fill(0.0);
let obj_seed = if self.prob.minimize {
obj_factor
} else {
-obj_factor
};
for buf in &mut self.compressed {
buf.fill(0.0);
}
if !self.quad.is_empty() {
if obj_seed != 0.0 {
if let Some(f) = self.quad.objective_form() {
self.quad.accumulate_hessian(f, obj_seed, values);
}
}
if let Some(lam) = lambda {
for i in 0..self.prob.m {
let w = lam[i];
if w == 0.0 {
continue;
}
if let Some(f) = self.quad.row_form(i) {
self.quad.accumulate_hessian(f, w, values);
}
}
}
}
if obj_seed != 0.0 {
for (ti, t) in self.obj_tapes.iter().enumerate() {
if t.ops.is_empty() {
continue;
}
t.forward_into(x, &mut self.vals_scratch);
for &c in &self.obj_tape_colors[ti] {
t.hessian_directional(
&self.vals_scratch,
&self.seeds[c as usize],
obj_seed,
&mut self.compressed[c as usize],
&mut self.dot_scratch,
&mut self.adj_scratch,
&mut self.adj_dot_scratch,
);
}
}
}
match (lambda, self.con_hybrid.as_mut()) {
(Some(lam), Some(h)) if h.use_for_hess => {
let ConHybrid {
tape,
prelude_vals,
local_vals_all,
local_off,
summand_row,
hess_color_summands,
hess_color_reach,
hess_color_reach_off,
prelude_dot,
hess_prelude_adj,
prelude_adj_dot,
local_dot,
local_adj,
local_adj_dot,
..
} = h;
tape.forward_prelude(x, prelude_vals);
for (si, s) in tape.summands.iter().enumerate() {
if lam[summand_row[si] as usize] == 0.0 {
continue;
}
tape.forward_summand(
s,
x,
prelude_vals,
&mut local_vals_all[local_off[si]..local_off[si + 1]],
);
}
for (c, list) in hess_color_summands.iter().enumerate() {
if !list
.iter()
.any(|&si| lam[summand_row[si as usize] as usize] != 0.0)
{
continue;
}
let seed = &self.seeds[c];
let out = &mut self.compressed[c];
let creach = &hess_color_reach
[hess_color_reach_off[c]..hess_color_reach_off[c + 1]];
tape.prelude_tangent(prelude_vals, seed, creach, prelude_dot);
for &si in list {
let si = si as usize;
let w = lam[summand_row[si] as usize];
if w == 0.0 {
continue;
}
tape.hessian_summand_directional(
&tape.summands[si],
&local_vals_all[local_off[si]..local_off[si + 1]],
prelude_dot,
seed,
w,
out,
local_dot,
local_adj,
local_adj_dot,
hess_prelude_adj,
prelude_adj_dot,
);
}
tape.prelude_reverse_directional(
prelude_vals,
prelude_dot,
creach,
out,
hess_prelude_adj,
prelude_adj_dot,
);
}
}
(Some(lam), _) => {
for k in 0..self.prob.m {
let w = lam[k];
if w == 0.0 {
continue;
}
for (ti, t) in self.con_tapes[k].iter().enumerate() {
if t.ops.is_empty() {
continue;
}
t.forward_into(x, &mut self.vals_scratch);
for &c in &self.con_tape_colors[k][ti] {
t.hessian_directional(
&self.vals_scratch,
&self.seeds[c as usize],
w,
&mut self.compressed[c as usize],
&mut self.dot_scratch,
&mut self.adj_scratch,
&mut self.adj_dot_scratch,
);
}
}
}
}
(None, _) => {}
}
for (c, table) in self.decoding.iter().enumerate() {
let comp = &self.compressed[c];
for w in table {
values[w.hess_idx as usize] += comp[w.row as usize];
}
}
true
}
}
}
fn finalize_solution(&mut self, sol: Solution<'_>, _d: &IpoptData, _q: &IpoptCq) {
self.final_x = Some(sol.x.to_vec());
self.final_obj = sol.obj_value;
self.final_lambda = Some(sol.lambda.to_vec());
self.final_z_l = Some(sol.z_l.to_vec());
self.final_z_u = Some(sol.z_u.to_vec());
}
fn get_var_con_metadata(&mut self, var: &mut MetaData, con: &mut MetaData) -> bool {
let mut any = false;
if !self.prob.var_names.is_empty() {
var.strings
.insert(IDX_NAMES.to_string(), self.prob.var_names.clone());
any = true;
}
if !self.prob.con_names.is_empty() {
con.strings
.insert(IDX_NAMES.to_string(), self.prob.con_names.clone());
any = true;
}
any
}
fn get_constraints_linearity(&mut self, types: &mut [Linearity]) -> bool {
for (i, t) in types.iter_mut().enumerate() {
*t = if self.prob.con_nonlinear[i].is_trivially_zero() {
Linearity::Linear
} else {
Linearity::NonLinear
};
}
true
}
fn get_variables_linearity(&mut self, types: &mut [Linearity]) -> bool {
let nonlinear = self.nonlinear_var_set();
for (i, t) in types.iter_mut().enumerate() {
*t = if nonlinear.contains(&i) {
Linearity::NonLinear
} else {
Linearity::Linear
};
}
true
}
fn get_objective_variables_linearity(&mut self, types: &mut [Linearity]) -> bool {
let mut nonlinear: BTreeSet<usize> = BTreeSet::new();
self.prob.obj_nonlinear.collect_vars(&mut nonlinear);
for (i, t) in types.iter_mut().enumerate() {
*t = if nonlinear.contains(&i) {
Linearity::NonLinear
} else {
Linearity::Linear
};
}
true
}
fn get_number_of_nonlinear_variables(&mut self) -> Index {
self.nonlinear_variables().len() as Index
}
fn get_list_of_nonlinear_variables(&mut self, pos_nonlin_vars: &mut [Index]) -> bool {
let list = self.nonlinear_variables();
if pos_nonlin_vars.len() < list.len() {
return false;
}
pos_nonlin_vars[..list.len()].copy_from_slice(&list);
true
}
fn derivative_proofs(&mut self) -> DerivativeProofs {
fn proof(affine: Option<bool>) -> DerivativeProof {
match affine {
Some(true) => DerivativeProof::Constant,
Some(false) => DerivativeProof::Varying,
None => DerivativeProof::Unknown,
}
}
let obj_affine = self.prob.obj_nonlinear.provably_affine();
let jac: Vec<DerivativeProof> = self
.prob
.con_nonlinear
.iter()
.map(|b| proof(b.provably_affine()))
.collect();
let hessian = if jac.iter().any(|&p| p == DerivativeProof::Varying) {
DerivativeProof::Varying
} else if obj_affine.is_some() && jac.iter().all(|&p| p == DerivativeProof::Constant) {
DerivativeProof::Constant
} else {
DerivativeProof::Unknown
};
DerivativeProofs {
grad_f: proof(obj_affine),
hessian,
jac,
}
}
}
pub fn load_nl_as_tnlp(path: &Path) -> Result<Rc<RefCell<dyn TNLP>>, String> {
let prob = read_nl_file(path)?;
Ok(Rc::new(RefCell::new(NlTnlp::new(prob))))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn nl_problem_and_tnlp_are_send() {
fn assert_send<T: Send>() {}
assert_send::<NlProblem>();
assert_send::<NlTnlp>();
assert_send::<Expr>();
}
#[test]
fn variant_overrides_bounds_and_x0() {
let p = parse_nl_text(SIMPLE).expect("parse");
let base = NlTnlp::new(p);
let var = base
.variant(&NlVariation {
x0: Some(vec![3.0, 4.0]),
x_l: Some(vec![-1.0, -2.0]),
x_u: Some(vec![5.0, 6.0]),
..Default::default()
})
.expect("variant");
let mut var = var;
let (mut x_l, mut x_u) = ([0.0; 2], [0.0; 2]);
let (mut g_l, mut g_u) = ([0.0; 0], [0.0; 0]);
assert!(var.get_bounds_info(BoundsInfo {
x_l: &mut x_l,
x_u: &mut x_u,
g_l: &mut g_l,
g_u: &mut g_u,
}));
assert_eq!(x_l, [-1.0, -2.0]);
assert_eq!(x_u, [5.0, 6.0]);
let mut x = [0.0; 2];
let (mut zl, mut zu, mut lam) = ([0.0; 2], [0.0; 2], [0.0; 0]);
assert!(var.get_starting_point(StartingPoint {
init_x: true,
x: &mut x,
init_z: false,
z_l: &mut zl,
z_u: &mut zu,
init_lambda: false,
lambda: &mut lam,
}));
assert_eq!(x, [3.0, 4.0]);
assert!(base.problem().x_l[0] < -1.0e18);
assert!(
base.variant(&NlVariation {
x0: Some(vec![1.0]),
..Default::default()
})
.is_err()
);
}
const SIMPLE: &str = "g3 0 1 0
2 0 1 0 0
0 1
0 0
0 2 0
0 0 0 1
0 0 0 0 0
0 0
0 0
0 0 0 0 0
O0 0
o0
o5
o1
v0
n1
n2
o5
o1
v1
n2
n2
b
3
3
";
#[test]
fn parses_simple_quadratic() {
let p = parse_nl_text(SIMPLE).expect("parse");
assert_eq!(p.n, 2);
assert_eq!(p.m, 0);
assert_eq!(p.num_obj, 1);
let f = eval_expr(&p.obj_expr(), &[0.0, 0.0]);
assert!((f - 5.0).abs() < 1e-12);
let f = eval_expr(&p.obj_expr(), &[1.0, 2.0]);
assert!(f.abs() < 1e-12);
}
#[test]
fn gradient_matches_analytic() {
let p = parse_nl_text(SIMPLE).expect("parse");
let x = [0.5, 1.0];
let mut g = [0.0_f64; 2];
grad_expr(&p.obj_expr(), &x, 1.0, &mut g);
assert!((g[0] - (-1.0)).abs() < 1e-12);
assert!((g[1] - (-2.0)).abs() < 1e-12);
}
#[test]
fn variables_linearity_tags_obj_nonlinear_vs_linear_vars() {
let obj_nl = Expr::Binary(
BinOp::Pow,
Box::new(Expr::Binary(
BinOp::Sub,
Box::new(Expr::Var(0)),
Box::new(Expr::Const(1.0)),
)),
Box::new(Expr::Const(2.0)),
);
let prob = NlProblem {
src: None,
cse_bodies: Vec::new(),
n: 2,
m: 0,
num_obj: 1,
minimize: true,
obj_nonlinear: NlBody::Tree(obj_nl),
obj_linear: vec![(1, 3.0)],
obj_constant: 0.0,
con_nonlinear: vec![],
con_linear: vec![],
x_l: vec![f64::NEG_INFINITY; 2],
x_u: vec![f64::INFINITY; 2],
g_l: vec![],
g_u: vec![],
x0: vec![0.0; 2],
lambda0: vec![],
suffixes: NlSuffixes::default(),
imported_funcs: vec![],
ampl_options: vec![],
nl_counts: None,
var_names: vec![],
con_names: vec![],
};
let mut tnlp = NlTnlp::new(prob);
let mut types = vec![Linearity::Linear; 2];
let ok = tnlp.get_variables_linearity(&mut types);
assert!(
ok,
"get_variables_linearity must report it filled the slice"
);
assert!(
matches!(types[0], Linearity::NonLinear),
"x0 is nonlinear in the objective"
);
assert!(
matches!(types[1], Linearity::Linear),
"x1 appears only in the linear part"
);
}
#[test]
fn objective_variables_linearity_ignores_constraint_nonlinearity() {
let con_nl = Expr::Binary(
BinOp::Pow,
Box::new(Expr::Var(0)),
Box::new(Expr::Const(2.0)),
);
let prob = NlProblem {
src: None,
cse_bodies: Vec::new(),
n: 2,
m: 1,
num_obj: 1,
minimize: true,
obj_nonlinear: NlBody::Tree(Expr::Const(0.0)),
obj_linear: vec![(1, 3.0)],
obj_constant: 0.0,
con_nonlinear: vec![NlBody::Tree(con_nl)],
con_linear: vec![vec![]],
x_l: vec![f64::NEG_INFINITY; 2],
x_u: vec![f64::INFINITY; 2],
g_l: vec![4.0],
g_u: vec![4.0],
x0: vec![0.0; 2],
lambda0: vec![0.0],
suffixes: NlSuffixes::default(),
imported_funcs: vec![],
ampl_options: vec![],
nl_counts: None,
var_names: vec![],
con_names: vec![],
};
let mut tnlp = NlTnlp::new(prob);
let mut global = vec![Linearity::Linear; 2];
assert!(tnlp.get_variables_linearity(&mut global));
assert!(
matches!(global[0], Linearity::NonLinear),
"global tags see x0's constraint nonlinearity"
);
let mut obj = vec![Linearity::NonLinear; 2];
assert!(tnlp.get_objective_variables_linearity(&mut obj));
assert!(
matches!(obj[0], Linearity::Linear),
"x0 is linear w.r.t. the objective despite the nonlinear constraint"
);
assert!(
matches!(obj[1], Linearity::Linear),
"x1 is linear everywhere"
);
}
#[test]
fn header_census_is_parsed() {
let p = parse_nl_text(SIMPLE).expect("parse");
let c = p.nl_counts.expect("SIMPLE has a well-formed header");
assert_eq!(c.nl_cons, 0);
assert_eq!(c.nl_objs, 1);
assert_eq!((c.nl_vars_cons, c.nl_vars_objs, c.nl_vars_both), (0, 2, 0));
assert_eq!(c.nonlinear_vars(), 2);
}
#[test]
fn nonlinear_var_total_uses_inclusion_exclusion() {
let c = |vc, vo, vb| NlCounts {
nl_cons: 0,
nl_objs: 0,
nl_vars_cons: vc,
nl_vars_objs: vo,
nl_vars_both: vb,
};
assert_eq!(c(5, 5, 5).nonlinear_vars(), 5);
assert_eq!(c(1, 1, 0).nonlinear_vars(), 2);
assert_eq!(c(4, 3, 2).nonlinear_vars(), 5);
assert_eq!(c(1, 1, 9).nonlinear_vars(), 0);
}
#[test]
fn short_header_line_records_no_census() {
let txt = SIMPLE.replacen("0 2 0\n", "0 2\n", 1);
assert_ne!(txt, SIMPLE, "the substitution must have applied");
let p = parse_nl_text(&txt).expect("parse");
assert!(p.nl_counts.is_none());
}
#[test]
fn nonlinear_variable_list_excludes_linear_columns() {
let obj_nl = Expr::Binary(
BinOp::Pow,
Box::new(Expr::Binary(
BinOp::Sub,
Box::new(Expr::Var(0)),
Box::new(Expr::Const(1.0)),
)),
Box::new(Expr::Const(2.0)),
);
let parts = NlProblemParts {
minimize: true,
objective: obj_nl,
obj_constant: 0.0,
constraints: vec![],
x_l: vec![-1e19; 2],
x_u: vec![1e19; 2],
x0: vec![0.0; 2],
g_l: vec![],
g_u: vec![],
var_names: vec![],
con_names: vec![],
};
let prob = NlProblem::from_expressions(parts).expect("build");
assert!(
prob.nl_counts.is_none(),
"a model built in memory has no header to read"
);
let mut tnlp = NlTnlp::new(prob);
assert_eq!(tnlp.get_number_of_nonlinear_variables(), 1);
let mut list = [-1 as Index; 2];
assert!(tnlp.get_list_of_nonlinear_variables(&mut list));
assert_eq!(list[0], 0);
}
#[test]
fn all_nonlinear_header_short_circuits_to_n() {
let p = parse_nl_text(SIMPLE).expect("parse");
assert_eq!(p.nl_counts.expect("census").nonlinear_vars(), p.n);
let mut tnlp = NlTnlp::new(p);
assert_eq!(tnlp.get_number_of_nonlinear_variables(), 2);
let mut list = [-1 as Index; 2];
assert!(tnlp.get_list_of_nonlinear_variables(&mut list));
assert_eq!(list, [0, 1]);
}
#[test]
fn nonlinear_variable_list_declines_a_short_slice() {
let p = parse_nl_text(SIMPLE).expect("parse");
let mut tnlp = NlTnlp::new(p);
let mut list = [-1 as Index; 1];
assert!(!tnlp.get_list_of_nonlinear_variables(&mut list));
assert_eq!(list, [-1]);
}
const EQ_LIN: &str = "g3 0 1 0
2 1 1 0 0
0 1
0 0
0 2 0
0 0 0 1
0 0 0 0 0
2 0
0 0
0 0 0 0 0
C0
n0
O0 0
o0
o5
v0
n2
o5
v1
n2
r
4 1
b
3
3
k1
2
J0 2
0 1
1 1
";
#[test]
fn parses_constrained_problem() {
let p = parse_nl_text(EQ_LIN).expect("parse");
assert_eq!(p.n, 2);
assert_eq!(p.m, 1);
assert!((p.g_l[0] - 1.0).abs() < 1e-12);
assert!((p.g_u[0] - 1.0).abs() < 1e-12);
assert_eq!(p.con_linear[0], vec![(0, 1.0), (1, 1.0)]);
}
#[test]
fn malformed_j_variable_index_is_parse_error_not_panic() {
let bad = EQ_LIN.replace("J0 2\n0 1\n1 1\n", "J0 2\n0 1\n5 1\n");
assert_ne!(bad, EQ_LIN, "fixture substitution must apply");
let err = parse_nl_text(&bad).expect_err("out-of-range J var must error");
assert!(err.contains("out of range"), "unexpected error: {err}");
}
#[test]
fn out_of_range_x_segment_index_is_parse_error() {
let bad = format!("{EQ_LIN}x1\n5 0.5\n");
let err = parse_nl_text(&bad).expect_err("out-of-range x index must error");
assert!(err.contains("out of range"), "unexpected error: {err}");
}
#[test]
fn a_constant_row_body_folds_into_both_bounds_and_reads_linear() {
let nl = EQ_LIN.replace("C0\nn0\n", "C0\nn3\n");
assert_ne!(nl, EQ_LIN, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse");
assert!(
p.con_nonlinear[0].is_trivially_zero(),
"the constant body must be replaced by the identity zero, got {:?}",
p.con_nonlinear[0]
);
assert!((p.g_l[0] - (-2.0)).abs() < 1e-12, "g_l = {}", p.g_l[0]);
assert!((p.g_u[0] - (-2.0)).abs() < 1e-12, "g_u = {}", p.g_u[0]);
let mut lin = [Linearity::NonLinear];
let mut t = NlTnlp::new(p);
assert!(t.get_constraints_linearity(&mut lin));
assert_eq!(lin[0], Linearity::Linear);
}
#[test]
fn folding_a_row_constant_gives_the_hand_folded_problem() {
let offset = EQ_LIN.replace("C0\nn0\n", "C0\nn3\n");
let folded = EQ_LIN.replace("r\n4 1\n", "r\n4 -2\n");
assert_ne!(offset, EQ_LIN);
assert_ne!(folded, EQ_LIN);
let a = parse_nl_text(&offset).expect("parse offset form");
let b = parse_nl_text(&folded).expect("parse folded form");
assert_eq!(a.g_l, b.g_l);
assert_eq!(a.g_u, b.g_u);
assert_eq!(a.con_linear, b.con_linear);
let mut ga = [0.0];
let mut gb = [0.0];
let x = [0.75, -1.25];
assert!(NlTnlp::new(a).eval_g(&x, true, &mut ga));
assert!(NlTnlp::new(b).eval_g(&x, true, &mut gb));
assert!((ga[0] - gb[0]).abs() < 1e-12, "{ga:?} vs {gb:?}");
}
#[test]
fn a_row_body_that_evaluates_to_a_constant_folds_too() {
let nl = EQ_LIN.replace("C0\nn0\n", "C0\no0\nn1\nn2\n");
assert_ne!(nl, EQ_LIN, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse");
assert!(p.con_nonlinear[0].is_trivially_zero());
assert!((p.g_l[0] - (-2.0)).abs() < 1e-12, "g_l = {}", p.g_l[0]);
assert!((p.g_u[0] - (-2.0)).abs() < 1e-12, "g_u = {}", p.g_u[0]);
}
#[test]
fn folding_a_row_constant_leaves_the_absent_bound_sentinel_alone() {
let nl = EQ_LIN
.replace("C0\nn0\n", "C0\nn-1e18\n")
.replace("r\n4 1\n", "r\n1 1\n");
let p = parse_nl_text(&nl).expect("parse");
assert!((p.g_u[0] - 1.0e18).abs() < 1024.0, "g_u = {}", p.g_u[0]);
assert!(
!lower_bound_present(p.g_l[0]),
"the absent-lower sentinel became a real bound: {}",
p.g_l[0]
);
let nl = EQ_LIN
.replace("C0\nn0\n", "C0\nn1e18\n")
.replace("r\n4 1\n", "r\n2 1\n");
let p = parse_nl_text(&nl).expect("parse");
assert!((p.g_l[0] + 1.0e18).abs() < 1024.0, "g_l = {}", p.g_l[0]);
assert!(
!upper_bound_present(p.g_u[0]),
"the absent-upper sentinel became a real bound: {}",
p.g_u[0]
);
}
#[test]
fn a_row_body_with_a_variable_is_not_folded() {
let nl = EQ_LIN.replace("C0\nn0\n", "C0\no5\nv0\nn2\n"); assert_ne!(nl, EQ_LIN, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse");
assert!(
!matches!(p.con_nonlinear[0].tree(), Some(Expr::Const(_))),
"a row in x0 was folded away: {:?}",
p.con_nonlinear[0]
);
assert!((p.g_l[0] - 1.0).abs() < 1e-12, "bounds moved: {}", p.g_l[0]);
assert!((p.g_u[0] - 1.0).abs() < 1e-12, "bounds moved: {}", p.g_u[0]);
}
#[test]
fn a_non_finite_constant_row_body_is_not_folded() {
let nl = EQ_LIN.replace("C0\nn0\n", "C0\no43\nn-1\n");
assert_ne!(nl, EQ_LIN, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse");
assert!(
!p.con_nonlinear[0].is_trivially_zero(),
"a NaN body was folded into the bounds"
);
assert!(p.g_l[0].is_finite() && p.g_u[0].is_finite());
assert!((p.g_l[0] - 1.0).abs() < 1e-12);
}
#[test]
fn a_constant_argument_funcall_row_body_is_not_folded() {
let nl = EQ_LIN.replace("C0\nn0\n", "F0 1 1 myfunc\nC0\nf0 1\nn2.0\n");
assert_ne!(nl, EQ_LIN, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse");
assert!(
matches!(p.con_nonlinear[0].tree(), Some(Expr::Funcall { .. })),
"expected the funcall to survive the fold, got {:?}",
p.con_nonlinear[0]
);
assert!((p.g_l[0] - 1.0).abs() < 1e-12, "bounds moved: {}", p.g_l[0]);
}
#[test]
fn k_segment_nonstandard_count_is_parse_error_at_source() {
let bad = EQ_LIN.replace("k1\n2\n", "k0\n");
assert_ne!(bad, EQ_LIN, "fixture substitution must apply");
let err = parse_nl_text(&bad).expect_err("nonstandard k count must error");
assert!(
err.contains("k-segment declares"),
"expected a clear k-segment count error, got: {err}"
);
}
#[test]
fn get_starting_point_returns_nl_initial_duals() {
let nl = format!("{EQ_LIN}\nd1\n0 2.5\n");
let p = parse_nl_text(&nl).expect("parse");
assert_eq!(p.lambda0, vec![2.5], "the `d` segment fills lambda0");
let mut t = NlTnlp::new(p);
let info = t.get_nlp_info().unwrap();
let (n, m) = (info.n as usize, info.m as usize);
let mut x = vec![0.0; n];
let mut z_l = vec![0.0; n];
let mut z_u = vec![0.0; n];
let mut lambda = vec![0.0; m];
assert!(t.get_starting_point(StartingPoint {
init_x: true,
x: &mut x,
init_z: false,
z_l: &mut z_l,
z_u: &mut z_u,
init_lambda: true,
lambda: &mut lambda,
}));
assert_eq!(
lambda,
vec![2.5],
"a warm start must use the `.nl` initial duals, not zero"
);
let mut lambda_untouched = vec![7.0; m];
assert!(t.get_starting_point(StartingPoint {
init_x: true,
x: &mut x,
init_z: false,
z_l: &mut z_l,
z_u: &mut z_u,
init_lambda: false,
lambda: &mut lambda_untouched,
}));
assert_eq!(
lambda_untouched,
vec![7.0],
"without init_lambda the multiplier buffer must be untouched"
);
}
fn dense_row_objective_nl(n: usize) -> String {
let mut s = String::new();
s.push_str("g3 1 1 0\n");
s.push_str(&format!(" {n} 0 1 0 0 0\n"));
s.push_str(" 0 1\n 0 0\n");
s.push_str(&format!(" {n} {n} {n}\n"));
s.push_str(" 0 0 0 1\n 0 0 0 0 0\n");
s.push_str(&format!(" 0 {n}\n"));
s.push_str(" 0 0\n 0 0 0 0 0\n");
s.push_str("O0 0\n");
s.push_str(&format!("o54\n{}\n", n + 1));
for j in 0..n {
s.push_str(&format!("o5\no1\nv{j}\nn1.0\nn2\n"));
}
s.push_str(&format!("o2\nv0\no54\n{n}\n"));
for j in 0..n {
s.push_str(&format!("v{j}\n"));
}
s.push_str(&format!("x{n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.5\n"));
}
s.push_str("b\n");
for _ in 0..n {
s.push_str("3\n");
}
s.push_str(&format!("G0 {n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.0\n"));
}
s
}
fn weighted_dense_row_objective_nl(n: usize, span: f64) -> String {
let w = |j: usize| 10_f64.powf(span / 2.0 - span * j as f64 / (n - 1) as f64);
let mut s = String::new();
s.push_str("g3 1 1 0\n");
s.push_str(&format!(" {n} 0 1 0 0 0\n"));
s.push_str(" 0 1\n 0 0\n");
s.push_str(&format!(" {n} {n} {n}\n"));
s.push_str(" 0 0 0 1\n 0 0 0 0 0\n");
s.push_str(&format!(" 0 {n}\n"));
s.push_str(" 0 0\n 0 0 0 0 0\n");
s.push_str("O0 0\n");
s.push_str(&format!("o54\n{}\n", n + 1));
for j in 0..n {
s.push_str(&format!("o5\no1\nv{j}\nn1.0\nn2\n"));
}
s.push_str(&format!("o2\nv0\no54\n{n}\n"));
for j in 0..n {
s.push_str(&format!("o2\nn{:.17e}\nv{j}\n", w(j)));
}
s.push_str(&format!("x{n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.5\n"));
}
s.push_str("b\n");
for _ in 0..n {
s.push_str("3\n");
}
s.push_str(&format!("G0 {n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.0\n"));
}
s
}
fn bench_model(rel: &str) -> Option<std::path::PathBuf> {
let root = std::env::var("POUNCE_BENCH_DATA").ok()?;
let p = std::path::PathBuf::from(root).join(rel);
p.is_file().then_some(p)
}
#[test]
fn cho_parmest_decodes_to_its_unpeeled_reference() {
let Some(path) = bench_model("cho/nl_export_results/cho_parmest.nl") else {
eprintln!("POUNCE_BENCH_DATA/cho not present — skipping corpus check");
return;
};
let p = read_nl_file(&path).expect("read cho_parmest");
let n = p.n;
let mut t = NlTnlp::new(p);
assert!(t.peeled_cols.is_empty());
let info = t.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
let (mut irow, mut jcol) = (vec![0_i32; nnz], vec![0_i32; nnz]);
assert!(t.eval_h(
None,
true,
1.0,
None,
true,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let x: Vec<f64> = t
.prob
.x0
.iter()
.enumerate()
.map(|(i, v)| v + 0.01 * (i % 7) as f64 + 0.001)
.collect();
let lambda: Vec<f64> = (0..t.prob.m).map(|i| 0.5 + 0.01 * (i % 5) as f64).collect();
let mut got = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
Some(&lambda),
true,
SparsityRequest::Values { values: &mut got }
));
t.recolor(&vec![true; n]);
assert!(t.peeled_cols.is_empty());
let mut want = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
Some(&lambda),
true,
SparsityRequest::Values { values: &mut want }
));
let scale = want.iter().fold(0.0_f64, |a, &v| a.max(v.abs()));
let mut worst = 0.0_f64;
let mut at = 0usize;
for k in 0..nnz {
let rel = (got[k] - want[k]).abs() / want[k].abs().max(f64::MIN_POSITIVE);
if rel > worst {
worst = rel;
at = k;
}
}
assert!(
worst <= 1e-13,
"H[{},{}] decoded {:e}, unpeeled reference {:e} — relative error \
{worst:e} (||H||inf = {scale:e}). A peeled column is being read \
out of a pass it cannot be read out of.",
irow[at],
jcol[at],
got[at],
want[at]
);
t.recolor(&vec![false; n]);
assert!(
!t.peeled_cols.is_empty(),
"restoring the unvetoed coloring must peel again"
);
let mut unguarded = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
Some(&lambda),
true,
SparsityRequest::Values {
values: &mut unguarded
}
));
let mut bad = 0usize;
let mut worst_unguarded = 0.0_f64;
for k in 0..nnz {
let rel = (unguarded[k] - want[k]).abs() / want[k].abs().max(f64::MIN_POSITIVE);
if rel > 1e-13 {
bad += 1;
}
worst_unguarded = worst_unguarded.max(rel);
}
assert!(
bad >= 50 && worst_unguarded > 1e-11,
"peeling cho_parmest unguarded should damage the entries the guard \
exists to protect (measured: 88 entries past 1e-13 relative, \
worst 9.9e-11); got {bad} entries, worst {worst_unguarded:e}. If \
this fires, the model or the corpus changed and the guard's \
calibration should be re-derived rather than the bound relaxed."
);
eprintln!("unguarded peeling damages {bad}/{nnz} entries, worst {worst_unguarded:e}");
}
fn weighted_dense_last_col_nl(n: usize, span: f64) -> String {
let w = |j: usize| 10_f64.powf(span / 2.0 - span * j as f64 / (n - 2) as f64);
let d = n - 1;
let mut s = String::new();
s.push_str("g3 1 1 0\n");
s.push_str(&format!(" {n} 0 1 0 0 0\n"));
s.push_str(" 0 1\n 0 0\n");
s.push_str(&format!(" {n} {n} {n}\n"));
s.push_str(" 0 0 0 1\n 0 0 0 0 0\n");
s.push_str(&format!(" 0 {n}\n"));
s.push_str(" 0 0\n 0 0 0 0 0\n");
s.push_str("O0 0\n");
s.push_str(&format!("o54\n{}\n", n + 1));
for j in 0..n {
s.push_str(&format!("o5\no1\nv{j}\nn1.0\nn2\n"));
}
s.push_str(&format!("o2\nv{d}\no54\n{}\n", n - 1));
for j in 0..d {
s.push_str(&format!("o2\nn{:.17e}\nv{j}\n", w(j)));
}
s.push_str(&format!("x{n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.5\n"));
}
s.push_str("b\n");
for _ in 0..n {
s.push_str("3\n");
}
s.push_str(&format!("G0 {n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.0\n"));
}
s
}
#[test]
fn a_well_scaled_dense_column_is_still_peeled() {
let n = 200;
let p = parse_nl_text(&weighted_dense_row_objective_nl(n, 0.0)).expect("parse");
let t = NlTnlp::new(p);
assert_eq!(
t.peeled_cols,
vec![0],
"a dense column that costs no accuracy must stay peeled"
);
assert!(
t.seeds.len() <= 4,
"peeling should keep the color count at O(1), got {}",
t.seeds.len()
);
}
#[test]
fn an_ill_scaled_dense_column_is_not_peeled_and_stays_exact() {
let n = 200;
let span = 12.0;
let w = |j: usize| 10_f64.powf(span / 2.0 - span * j as f64 / (n - 1) as f64);
let p = parse_nl_text(&weighted_dense_row_objective_nl(n, span)).expect("parse");
let mut t = NlTnlp::new(p);
assert!(
t.peeled_cols.is_empty(),
"a dense column spanning {span} orders must not be peeled; got {:?}",
t.peeled_cols
);
let info = t.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
let (mut irow, mut jcol) = (vec![0_i32; nnz], vec![0_i32; nnz]);
assert!(t.eval_h(
None,
true,
1.0,
None,
true,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let x: Vec<f64> = (0..n).map(|j| 0.1 * j as f64).collect();
let mut vals = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
None,
true,
SparsityRequest::Values { values: &mut vals }
));
let mut checked = 0;
for k in 0..nnz {
let (i, j) = (irow[k] as usize, jcol[k] as usize);
if i == j {
continue;
}
assert_eq!(j, 0, "unexpected off-diagonal ({i}, {j})");
checked += 1;
let want = w(i);
assert!(
(vals[k] - want).abs() <= 1e-13 * want,
"H[{i},0] = {:e}, want {want:e} (relative error {:e})",
vals[k],
(vals[k] - want).abs() / want
);
}
assert_eq!(checked, n - 1);
}
#[test]
fn a_peeled_decode_matches_an_unpeeled_reference() {
let cases = [
(200, 0.0, false),
(200, 2.0, false),
(200, 0.0, true),
(200, 1.0, true),
(200, 2.0, true),
(400, 3.0, true),
(600, 0.0, true),
];
let mut peeled_any = false;
let mut worst = 0.0_f64;
for &(n, span, last) in &cases {
let text = if last {
weighted_dense_last_col_nl(n, span)
} else {
weighted_dense_row_objective_nl(n, span)
};
let p = parse_nl_text(&text).expect("parse");
let mut t = NlTnlp::new(p);
let peeled = t.peeled_cols.clone();
peeled_any |= !peeled.is_empty();
let info = t.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
let (mut irow, mut jcol) = (vec![0_i32; nnz], vec![0_i32; nnz]);
assert!(t.eval_h(
None,
true,
1.0,
None,
true,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let x: Vec<f64> = (0..n).map(|j| 0.25 + 0.05 * (j % 13) as f64).collect();
let mut got = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
None,
true,
SparsityRequest::Values { values: &mut got }
));
t.recolor(&vec![true; n]);
assert!(
t.peeled_cols.is_empty(),
"a fully vetoed model must peel nothing"
);
let mut want = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
None,
true,
SparsityRequest::Values { values: &mut want }
));
for k in 0..nnz {
let scale = want[k].abs().max(f64::MIN_POSITIVE);
let rel = (got[k] - want[k]).abs() / scale;
worst = worst.max(rel);
assert!(
rel <= 1e-13,
"n={n} span={span} last={last} peeled={peeled:?}: H[{},{}] decoded {:e}, \
unpeeled reference {:e} (relative error {rel:e})",
irow[k],
jcol[k],
got[k],
want[k]
);
}
}
assert!(
peeled_any,
"no case peeled anything, so this test proved nothing about the \
decode path — the exact way the fixture suite missed the bug"
);
assert!(worst < 1e-13, "worst relative disagreement {worst:e}");
}
#[test]
fn a_vetoed_column_is_colored_the_ordinary_way() {
let n = 300;
let mut pairs: Vec<(usize, usize)> = (0..n).map(|j| (j, j)).collect();
pairs.extend((1..n).map(|i| (i, 0)));
pairs.sort_unstable();
let (_, colors_peeled, peeled) = greedy_hessian_coloring(n, &pairs, &vec![false; n]);
assert!(peeled[0], "the dense column should peel by default");
assert!(
colors_peeled <= 4,
"peeling should collapse the count, got {colors_peeled}"
);
let mut veto = vec![false; n];
veto[0] = true;
let (_, colors_vetoed, peeled) = greedy_hessian_coloring(n, &pairs, &veto);
assert!(!peeled[0], "a vetoed column must not be peeled");
assert!(
colors_vetoed > colors_peeled,
"un-peeling restores row 0's conflicts, so colors must rise: \
{colors_vetoed} vs {colors_peeled}"
);
}
#[test]
fn a_dense_hessian_row_does_not_explode_the_coloring() {
let n = 200;
let p = parse_nl_text(&dense_row_objective_nl(n)).expect("parse");
let t = NlTnlp::new(p);
assert_eq!(t.h_irow.len(), 2 * n - 1);
assert!(
t.seeds.len() <= 4,
"one dense row should cost one extra color, not n; got {} colors for n={n}",
t.seeds.len()
);
assert_eq!(t.seeds.len(), t.compressed.len());
}
#[test]
fn peeled_dense_column_still_recovers_the_exact_hessian() {
let n = 200;
let p = parse_nl_text(&dense_row_objective_nl(n)).expect("parse");
let mut t = NlTnlp::new(p);
let info = t.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
let mut irow = vec![0_i32; nnz];
let mut jcol = vec![0_i32; nnz];
assert!(t.eval_h(
None,
true,
1.0,
None,
true,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let x: Vec<f64> = (0..n).map(|j| 0.1 * j as f64).collect();
let obj_factor = 2.5;
let mut vals = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
obj_factor,
None,
true,
SparsityRequest::Values { values: &mut vals }
));
let mut seen_diag = 0;
let mut seen_coupling = 0;
for k in 0..nnz {
let (i, j) = (irow[k] as usize, jcol[k] as usize);
let want = if i == 0 && j == 0 {
4.0
} else if i == j {
seen_diag += 1;
2.0
} else {
assert_eq!(j, 0, "unexpected off-diagonal ({i}, {j})");
seen_coupling += 1;
1.0
} * obj_factor;
assert!(
(vals[k] - want).abs() < 1e-12,
"H[{i},{j}] = {}, want {want}",
vals[k]
);
}
assert_eq!(seen_diag, n - 1);
assert_eq!(seen_coupling, n - 1);
}
#[test]
fn a_sparse_hessian_is_colored_by_its_bandwidth_not_peeled() {
let n = 400;
let pairs: Vec<(usize, usize)> = (0..n)
.flat_map(|j| {
let mut v = vec![(j, j)];
if j + 1 < n {
v.push((j + 1, j));
}
v
})
.collect();
let (var_color, n_colors, peeled) = greedy_hessian_coloring(n, &pairs, &vec![false; n]);
assert!(!peeled.iter().any(|&p| p), "nothing in a band is dense");
assert!(
n_colors <= 3,
"a tridiagonal Hessian needs a handful of colors, got {n_colors}"
);
assert!(var_color.iter().all(|&c| c != u32::MAX));
}
fn disjoint_blocks(n: usize, blocks: usize, size: usize) -> Vec<(usize, usize)> {
let mut pairs: Vec<(usize, usize)> = (0..n).map(|j| (j, j)).collect();
let stride = n / blocks;
for b in 0..blocks {
let base = b * stride;
for i in 0..size {
for j in 0..=i {
if i != j {
pairs.push((base + i, base + j));
}
}
}
}
pairs
}
#[test]
fn thousands_of_medium_degree_cols_are_not_peeled() {
for (n, blocks, size) in [(200_000, 100, 50), (200_000, 300, 34)] {
let pairs = disjoint_blocks(n, blocks, size);
let (_, n_colors, peeled) = greedy_hessian_coloring(n, &pairs, &vec![false; n]);
let n_peeled = peeled.iter().filter(|&&p| p).count();
assert_eq!(
n_peeled, 0,
"degree-{size} columns do not pay for a singleton color \
(n={n}, blocks={blocks}), yet {n_peeled} were peeled"
);
assert!(
n_colors <= size + 1,
"disjoint {size}x{size} blocks color by block size regardless \
of block count; got {n_colors} for n={n}, blocks={blocks}"
);
}
}
#[test]
fn a_few_truly_dense_rows_are_still_peeled() {
let n = 5_000;
let dense_rows = 4;
let mut pairs: Vec<(usize, usize)> = (0..n).map(|j| (j, j)).collect();
for d in 0..dense_rows {
for j in 0..n {
if j != d {
pairs.push((j.max(d), j.min(d)));
}
}
}
let (_, n_colors, peeled) = greedy_hessian_coloring(n, &pairs, &vec![false; n]);
let n_peeled = peeled.iter().filter(|&&p| p).count();
assert_eq!(n_peeled, dense_rows, "every full row should peel");
assert!(
n_colors <= dense_rows + 2,
"peeling {dense_rows} full rows should leave a diagonal remainder, \
got {n_colors} colors"
);
}
#[test]
fn peeling_is_capped_and_keeps_the_worst_offenders() {
let n = 20_000;
let dense_rows = 400;
let mut pairs: Vec<(usize, usize)> = (0..n).map(|j| (j, j)).collect();
for d in 0..dense_rows {
for j in 0..(n - d) {
if j != d {
pairs.push((j.max(d), j.min(d)));
}
}
}
let (_, _, peeled) = greedy_hessian_coloring(n, &pairs, &vec![false; n]);
let n_peeled = peeled.iter().filter(|&&p| p).count();
assert_eq!(n_peeled, MAX_PEELED_COLS, "cap binds at {MAX_PEELED_COLS}");
assert!(
(0..MAX_PEELED_COLS).all(|d| peeled[d]),
"the {MAX_PEELED_COLS} densest rows are the ones kept"
);
}
#[test]
fn header_option_words_are_kept_verbatim() {
let p = parse_nl_text(EQ_LIN).expect("parse");
assert_eq!(p.ampl_options, vec![0, 1, 0]);
}
#[test]
fn a_truncated_option_list_is_dropped_not_padded() {
let text = EQ_LIN.replacen("g3 0 1 0", "g9 0 1 0", 1);
let p = parse_nl_text(&text).expect("parse");
assert!(
p.ampl_options.is_empty(),
"9 declared but 3 present: {:?}",
p.ampl_options
);
}
#[test]
fn constrained_tnlp_eval_g_jac_h() {
let p = parse_nl_text(EQ_LIN).expect("parse");
let mut t = NlTnlp::new(p);
let info = t.get_nlp_info().unwrap();
assert_eq!(info.m, 1);
assert_eq!(info.nnz_jac_g, 2);
let mut g = [0.0_f64; 1];
assert!(t.eval_g(&[0.3, 0.4], true, &mut g));
assert!((g[0] - 0.7).abs() < 1e-12);
let mut irow = [0_i32; 2];
let mut jcol = [0_i32; 2];
assert!(t.eval_jac_g(
None,
true,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
assert_eq!(irow, [0, 0]);
assert_eq!(jcol, [0, 1]);
let mut vals = [0.0_f64; 2];
assert!(t.eval_jac_g(
Some(&[0.3, 0.4]),
true,
SparsityRequest::Values { values: &mut vals }
));
assert!((vals[0] - 1.0).abs() < 1e-12);
assert!((vals[1] - 1.0).abs() < 1e-12);
assert_eq!(info.nnz_h_lag, 2);
let mut hirow = [0_i32; 2];
let mut hjcol = [0_i32; 2];
assert!(t.eval_h(
None,
true,
1.0,
None,
true,
SparsityRequest::Structure {
irow: &mut hirow,
jcol: &mut hjcol
}
));
assert_eq!(hirow, [0, 1]);
assert_eq!(hjcol, [0, 1]);
let mut hvals = [0.0_f64; 2];
assert!(t.eval_h(
Some(&[0.3, 0.4]),
true,
1.0,
Some(&[0.5]),
true,
SparsityRequest::Values { values: &mut hvals }
));
assert!((hvals[0] - 2.0).abs() < 1e-12);
assert!((hvals[1] - 2.0).abs() < 1e-12);
}
const CSE_OBJ: &str = "g3 0 1 0
2 0 1 0 0
0 1
0 0
0 2 0
0 0 0 1
0 0 0 0 0
0 0
0 0
0 1 0 0 0
V2 0 0
o0
v0
v1
O0 0
o0
o5
v2
n2
v2
b
3
3
";
#[test]
fn parses_v_segment_cse() {
let p = parse_nl_text(CSE_OBJ).expect("parse");
assert_eq!(p.n, 2);
let f = eval_expr(&p.obj_expr(), &[1.0, 2.0]);
assert!((f - 12.0).abs() < 1e-12, "got {f}");
let mut g = [0.0_f64; 2];
grad_expr(&p.obj_expr(), &[1.0, 2.0], 1.0, &mut g);
assert!((g[0] - 7.0).abs() < 1e-12, "g[0]={}", g[0]);
assert!((g[1] - 7.0).abs() < 1e-12, "g[1]={}", g[1]);
let mut vs = BTreeSet::new();
p.obj_nonlinear.collect_vars(&mut vs);
assert_eq!(vs.into_iter().collect::<Vec<_>>(), vec![0, 1]);
}
const WITH_SUFFIXES: &str = "g3 0 1 0
1 0 1 0 0
0 1
0 0
0 1 0
0 0 0 1
0 0 0 0 0
0 0
0 0
0 0 0 0 0
O0 0
o5
o1
v0
n1
n2
b
3
S0 1 sens_state_1
0 7
S4 1 sens_state_value_1
0 4.5
";
#[test]
fn parses_var_int_and_var_real_suffixes() {
let p = parse_nl_text(WITH_SUFFIXES).expect("parse");
let v = p.suffixes.var_int.get("sens_state_1").expect("var_int");
assert_eq!(v.as_slice(), &[7]);
let r = p
.suffixes
.var_real
.get("sens_state_value_1")
.expect("var_real");
assert_eq!(r.len(), 1);
assert!((r[0] - 4.5).abs() < 1e-12);
assert!(p.suffixes.con_int.is_empty());
assert!(p.suffixes.con_real.is_empty());
}
const WITH_CON_SUFFIX: &str = "g3 0 1 0
2 2 1 0 0
0 0
0 0
0 2 0
0 0 0 1
0 0 0 0 0
4 0
0 0
0 0 0 0 0 0
C0
n0
C1
n0
O0 0
n0
r
4 0.0
4 0.0
b
3
3
k1
0
J0 2
0 1
1 1
J1 2
0 1
1 -1
S1 2 sens_init_constr
0 1
1 2
";
#[test]
fn parses_con_int_suffix() {
let p = parse_nl_text(WITH_CON_SUFFIX).expect("parse");
let s = p.suffixes.con_int.get("sens_init_constr").expect("con_int");
assert_eq!(s.as_slice(), &[1, 2]);
}
fn truncate_before(txt: &str, at: &str) -> String {
let cut = txt
.find(at)
.unwrap_or_else(|| panic!("fixture has no {at:?} segment"));
txt[..cut].to_string()
}
#[test]
fn truncation_before_the_row_bounds_is_a_parse_error() {
let err = parse_nl_text(&truncate_before(WITH_CON_SUFFIX, "\nr\n"))
.expect_err("truncated file must not parse");
assert!(
err.contains("`r` (constraint-bounds) segment"),
"error should name the missing segment: {err}"
);
}
#[test]
fn truncation_before_the_variable_bounds_is_a_parse_error() {
let err = parse_nl_text(&truncate_before(WITH_CON_SUFFIX, "\nb\n"))
.expect_err("truncated file must not parse");
assert!(
err.contains("`b` (variable-bounds) segment"),
"error should name the missing segment: {err}"
);
}
#[test]
fn truncation_before_the_jacobian_is_a_parse_error() {
let err = parse_nl_text(&truncate_before(WITH_CON_SUFFIX, "\nk1\n"))
.expect_err("truncated file must not parse");
assert!(
err.contains("declares 4 Jacobian nonzero(s) but the J segments supply 0"),
"error should report the mismatch: {err}"
);
}
#[test]
fn more_jacobian_entries_than_declared_is_also_a_parse_error() {
let extra = WITH_CON_SUFFIX.replace("J1 2\n0 1\n1 -1\n", "J1 2\n0 1\n1 -1\nJ0 1\n0 5\n");
let err = parse_nl_text(&extra).expect_err("over-full file must not parse");
assert!(
err.contains("declares 4 Jacobian nonzero(s) but the J segments supply 5"),
"error should report the mismatch: {err}"
);
}
fn scaling_of(tnlp: &mut NlTnlp) -> (bool, Number, bool, Vec<Number>, bool, Vec<Number>) {
let mut obj = 1.0;
let mut use_x = false;
let mut x = vec![0.0; 2];
let mut use_g = false;
let mut g = vec![0.0; 2];
let ok = tnlp.get_scaling_parameters(ScalingRequest {
obj_scaling: &mut obj,
use_x_scaling: &mut use_x,
x_scaling: &mut x,
use_g_scaling: &mut use_g,
g_scaling: &mut g,
});
(ok, obj, use_x, x, use_g, g)
}
#[test]
fn scaling_factor_suffix_feeds_obj_and_constraint_scaling() {
let nl = WITH_CON_SUFFIX.to_string()
+ "S5 1 scaling_factor\n0 10.0\nS6 1 scaling_factor\n0 100.0\n";
let p = parse_nl_text(&nl).expect("parse");
let mut tnlp = NlTnlp::new(p);
let (ok, obj, use_x, _x, use_g, g) = scaling_of(&mut tnlp);
assert!(ok, "a tagged model must supply scaling");
assert!((obj - 100.0).abs() < 1e-12, "obj_scaling={obj}");
assert!(use_g);
assert_eq!(g, vec![10.0, 1.0]);
assert!(!use_x, "no variable suffix was declared");
}
#[test]
fn scaling_factor_suffix_forwards_variable_factors() {
let nl = WITH_CON_SUFFIX.to_string() + "S4 1 scaling_factor\n1 3.0\n";
let p = parse_nl_text(&nl).expect("parse");
let mut tnlp = NlTnlp::new(p);
let (ok, _obj, use_x, x, _use_g, _g) = scaling_of(&mut tnlp);
assert!(ok);
assert!(use_x, "variable factors must reach the engine");
assert_eq!(x, vec![1.0, 3.0]);
}
#[test]
fn a_user_suffix_overrides_the_computed_factor_component_wise() {
let baseline = {
let p = parse_nl_text(WITH_CON_SUFFIX).expect("parse");
let mut t = NlTnlp::new(p);
assert!(t.enable_curvature_scaling(), "an LP is degree ≤ 2");
scaling_of(&mut t).5
};
assert!(
baseline.iter().all(|v| *v > 0.0),
"curvature scaling should produce usable row factors, got {baseline:?}"
);
let nl = WITH_CON_SUFFIX.to_string() + "S5 1 scaling_factor\n0 7.0\n";
let p = parse_nl_text(&nl).expect("parse");
let mut t = NlTnlp::new(p);
assert!(t.enable_curvature_scaling());
let (ok, _obj, use_x, _x, use_g, g) = scaling_of(&mut t);
assert!(ok);
assert!(use_g && use_x);
assert_eq!(g[0], 7.0, "the tagged row takes the user's factor");
assert_eq!(
g[1], baseline[1],
"the untagged row keeps the computed one, not 1.0"
);
}
#[test]
fn no_scaling_factor_suffix_declines() {
let p = parse_nl_text(WITH_CON_SUFFIX).expect("parse");
let mut tnlp = NlTnlp::new(p);
let (ok, ..) = scaling_of(&mut tnlp);
assert!(!ok);
}
#[test]
fn rejects_suffix_with_out_of_range_index() {
let bad = WITH_CON_SUFFIX.replace("1 2\n", "5 2\n"); let err = parse_nl_text(&bad).expect_err("must reject");
assert!(
err.contains("out of range"),
"expected out-of-range error, got: {err}"
);
}
#[test]
fn tnlp_round_trip_solves() {
let p = parse_nl_text(SIMPLE).expect("parse");
let mut tnlp = NlTnlp::new(p);
let info = tnlp.get_nlp_info().unwrap();
assert_eq!(info.n, 2);
assert_eq!(info.m, 0);
let f0 = tnlp.eval_f(&[0.0, 0.0], true).unwrap();
assert!((f0 - 5.0).abs() < 1e-12);
let mut g = [0.0_f64; 2];
tnlp.eval_grad_f(&[0.0, 0.0], true, &mut g);
assert!((g[0] - (-2.0)).abs() < 1e-12);
assert!((g[1] - (-4.0)).abs() < 1e-12);
}
use pounce_nlp::expression_provider::ExpressionProvider;
use std::sync::atomic::{AtomicUsize, Ordering};
fn scratch_dir(tag: &str) -> std::path::PathBuf {
static N: AtomicUsize = AtomicUsize::new(0);
let seq = N.fetch_add(1, Ordering::Relaxed);
let dir = std::env::temp_dir().join(format!(
"pounce_nlnames_{}_{}_{}",
std::process::id(),
tag,
seq
));
std::fs::create_dir_all(&dir).expect("create scratch dir");
dir
}
#[test]
fn read_name_file_reads_in_order() {
let dir = scratch_dir("col_order");
let p = dir.join("m.col");
std::fs::write(&p, "x_in\nT_reactor\nflow\n").unwrap();
assert_eq!(read_name_file(&p, 3), vec!["x_in", "T_reactor", "flow"]);
}
#[test]
fn read_name_file_truncates_extra_lines() {
let dir = scratch_dir("row_obj");
let p = dir.join("m.row");
std::fs::write(&p, "mass_balance\nenergy_balance\nobj\n").unwrap();
assert_eq!(
read_name_file(&p, 2),
vec!["mass_balance", "energy_balance"]
);
}
#[test]
fn read_name_file_empty_on_short_or_missing() {
let dir = scratch_dir("short");
let short = dir.join("m.col");
std::fs::write(&short, "only_one\n").unwrap();
assert!(read_name_file(&short, 3).is_empty());
assert!(read_name_file(&dir.join("absent.col"), 2).is_empty());
}
#[test]
fn read_nl_file_captures_sibling_names() {
let dir = scratch_dir("sibling");
let nl = dir.join("m.nl");
std::fs::write(&nl, SIMPLE).unwrap();
std::fs::write(dir.join("m.col"), "alpha\nbeta\n").unwrap();
let prob = read_nl_file(&nl).expect("parse + name capture");
assert_eq!(prob.var_names, vec!["alpha", "beta"]);
assert!(prob.con_names.is_empty());
let tnlp = NlTnlp::new(prob);
assert_eq!(tnlp.variable_name(0), Some("alpha"));
assert_eq!(tnlp.variable_name(1), Some("beta"));
assert_eq!(tnlp.variable_name(2), None); }
#[test]
fn read_nl_file_without_names_yields_empty() {
let dir = scratch_dir("noname");
let nl = dir.join("m.nl");
std::fs::write(&nl, SIMPLE).unwrap();
let prob = read_nl_file(&nl).expect("parse");
assert!(prob.var_names.is_empty());
assert!(prob.con_names.is_empty());
let tnlp = NlTnlp::new(prob);
assert_eq!(tnlp.variable_name(0), None);
}
#[test]
fn read_nl_file_resolves_extensionless_ampl_stub() {
let dir = scratch_dir("stub");
std::fs::write(dir.join("mystub.nl"), SIMPLE).unwrap();
let stub = dir.join("mystub");
assert!(!stub.exists(), "stub must be extensionless / absent");
let prob = read_nl_file(&stub).expect("stub should resolve to mystub.nl");
assert_eq!(prob.n, 2);
assert_eq!(prob.m, 0);
std::fs::write(dir.join("mystub.col"), "alpha\nbeta\n").unwrap();
let prob = read_nl_file(&stub).expect("stub resolves, names ride along");
assert_eq!(prob.var_names, vec!["alpha", "beta"]);
}
#[test]
fn read_nl_file_prefers_exact_path_over_nl_sibling() {
let dir = scratch_dir("exact");
std::fs::write(dir.join("data"), SIMPLE).unwrap();
std::fs::write(dir.join("data.nl"), "not an nl file").unwrap();
let prob = read_nl_file(&dir.join("data")).expect("exact path wins");
assert_eq!(prob.n, 2);
}
#[test]
fn append_extension_appends_rather_than_replaces() {
use std::path::Path;
assert_eq!(
append_extension(Path::new("mystub"), "nl"),
Path::new("mystub.nl")
);
assert_eq!(
append_extension(Path::new("my.model"), "nl"),
Path::new("my.model.nl")
);
}
fn names(v: &[&str]) -> Vec<String> {
v.iter().map(|s| s.to_string()).collect()
}
#[test]
fn render_uses_variable_names_when_present() {
let e = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
assert_eq!(render_expr(&e, &names(&["T", "flow"]), &[]), "T*flow");
assert_eq!(render_expr(&e, &[], &[]), "x[0]*x[1]");
}
#[test]
fn render_parenthesizes_by_precedence() {
let sum = Expr::Binary(BinOp::Add, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
let e = Expr::Binary(BinOp::Mul, Box::new(sum), Box::new(Expr::Var(2)));
assert_eq!(render_expr(&e, &[], &[]), "(x[0] + x[1])*x[2]");
let mul = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(1)), Box::new(Expr::Var(2)));
let e2 = Expr::Binary(BinOp::Add, Box::new(Expr::Var(0)), Box::new(mul));
assert_eq!(render_expr(&e2, &[], &[]), "x[0] + x[1]*x[2]");
}
#[test]
fn render_subtraction_right_assoc_parens() {
let inner = Expr::Binary(BinOp::Sub, Box::new(Expr::Var(1)), Box::new(Expr::Var(2)));
let e = Expr::Binary(BinOp::Sub, Box::new(Expr::Var(0)), Box::new(inner));
assert_eq!(render_expr(&e, &[], &[]), "x[0] - (x[1] - x[2])");
}
#[test]
fn render_functions_and_pow() {
let sq = Expr::Binary(
BinOp::Pow,
Box::new(Expr::Var(0)),
Box::new(Expr::Const(2.0)),
);
let e = Expr::Unary(UnaryOp::Exp, Box::new(sq));
assert_eq!(render_expr(&e, &names(&["q"]), &[]), "exp(q^2)");
}
#[test]
fn render_linear_signs_are_tidy() {
let lin = vec![(0usize, 1.0), (1, -2.0), (2, 1.0)];
assert_eq!(render_linear(&lin, &names(&["a", "b", "c"])), "a - 2*b + c");
}
#[test]
fn render_linear_skips_zero_coefficients() {
let lin = vec![(0usize, 1.0), (1, 0.0), (2, -3.0)];
assert_eq!(render_linear(&lin, &names(&["a", "b", "c"])), "a - 3*c");
let lin = vec![(0usize, 0.0), (1, 2.0)];
assert_eq!(render_linear(&lin, &names(&["a", "b"])), "2*b");
}
#[test]
fn render_sum_folds_negative_terms() {
let sq = |i| {
Expr::Binary(
BinOp::Pow,
Box::new(Expr::Var(i)),
Box::new(Expr::Const(2.0)),
)
};
let neg = |i| {
Expr::Binary(
BinOp::Mul,
Box::new(Expr::Const(-1.0)),
Box::new(Expr::Var(i)),
)
};
let e = Expr::Sum(vec![
sq(0),
neg(1),
Expr::Unary(UnaryOp::Neg, Box::new(Expr::Var(2))),
]);
assert_eq!(
render_expr(&e, &names(&["a", "b", "c"]), &[]),
"a^2 - 1*b - c"
);
}
#[test]
fn render_constraint_equation_forms() {
let mut prob = parse_nl_text(SIMPLE).unwrap();
prob.n = 2;
prob.m = 2;
prob.var_names = names(&["mass_in", "mass_out"]);
prob.con_names = names(&["balance", "window"]);
prob.con_linear = vec![
vec![(0, 1.0), (1, -1.0)], vec![(0, 1.0)], ];
prob.con_nonlinear = vec![
NlBody::Tree(Expr::Const(0.0)),
NlBody::Tree(Expr::Const(0.0)),
];
prob.g_l = vec![0.0, 0.0];
prob.g_u = vec![0.0, 500.0];
assert_eq!(
render_constraint_equation(&prob, 0),
"mass_in - mass_out = 0"
);
assert_eq!(render_constraint_equation(&prob, 1), "0 <= mass_in <= 500");
let all = render_all_constraint_equations(&prob);
assert_eq!(all.len(), 2);
assert_eq!(all[1], "0 <= mass_in <= 500");
}
#[test]
fn constraint_jacobian_sparsity_unions_linear_and_nonlinear() {
let mut prob = parse_nl_text(SIMPLE).unwrap();
prob.n = 3;
prob.m = 2;
prob.con_linear = vec![vec![(1, 4.0)], vec![(2, 1.0)]];
prob.con_nonlinear = vec![
NlBody::Tree(Expr::Binary(
BinOp::Mul,
Box::new(Expr::Var(0)),
Box::new(Expr::Var(2)),
)),
NlBody::Tree(Expr::Const(0.0)),
];
prob.g_l = vec![0.0, 0.0];
prob.g_u = vec![0.0, 0.0];
let (irow, jcol) = constraint_jacobian_sparsity(&prob);
assert_eq!(irow, vec![0, 0, 0, 1]);
assert_eq!(jcol, vec![0, 1, 2, 2]);
}
#[test]
fn funcall_string_arg_with_hash_is_not_truncated() {
let mut p = Parser::new("h3:a#b\n", false);
match p.parse_funcall_arg().expect("parse hollerith arg") {
FuncallArg::Str(s) => assert_eq!(s, "a#b"),
other => panic!("expected Str, got {other:?}"),
}
}
#[test]
fn funcall_string_arg_honors_declared_length() {
let mut p = Parser::new("h3:abc # trailing comment\n", false);
match p.parse_funcall_arg().expect("parse hollerith arg") {
FuncallArg::Str(s) => assert_eq!(s, "abc"),
other => panic!("expected Str, got {other:?}"),
}
}
fn parse_one_expr(n: usize, expr_src: &str) -> Expr {
let mut p = Parser::new(expr_src, false);
p.n = n;
p.parse_expr().expect("parse expression")
}
#[test]
fn opcode_o82_square_is_unary_pow_of_two() {
let e = parse_one_expr(1, "o82\nv0\n");
match &e {
Expr::Binary(BinOp::Pow, base, exp) => {
assert!(matches!(**base, Expr::Var(0)));
match **exp {
Expr::Const(c) => assert!((c - 2.0).abs() < 1e-12, "exp const = {c}"),
ref other => panic!("o82 exponent must be Const(2.0), got {other:?}"),
}
}
other => panic!("o82 must parse to Pow(base, 2), got {other:?}"),
}
assert!((eval_expr(&e, &[3.0]) - 9.0).abs() < 1e-12);
assert!((eval_expr(&e, &[-3.0]) - 9.0).abs() < 1e-12);
let mut g = [0.0_f64; 1];
grad_expr(&e, &[3.0], 1.0, &mut g);
assert!((g[0] - 6.0).abs() < 1e-9, "grad at 3 = {}", g[0]);
g[0] = 0.0;
grad_expr(&e, &[-3.0], 1.0, &mut g);
assert!((g[0] + 6.0).abs() < 1e-9, "grad at -3 = {}", g[0]);
}
#[test]
fn opcode_o81_const_exponent_is_base_pow_const() {
let e = parse_one_expr(1, "o81\nv0\nn3\n");
match &e {
Expr::Binary(BinOp::Pow, base, exp) => {
assert!(matches!(**base, Expr::Var(0)), "base must be the variable");
match **exp {
Expr::Const(c) => assert!((c - 3.0).abs() < 1e-12, "exp const = {c}"),
ref other => panic!("o81 exponent must be Const(3.0), got {other:?}"),
}
}
other => panic!("o81 must parse to Pow(var, const), got {other:?}"),
}
assert!((eval_expr(&e, &[2.0]) - 8.0).abs() < 1e-12);
assert!((eval_expr(&e, &[-2.0]) + 8.0).abs() < 1e-12);
let mut g = [0.0_f64; 1];
grad_expr(&e, &[2.0], 1.0, &mut g);
assert!((g[0] - 12.0).abs() < 1e-9, "grad at 2 = {}", g[0]);
}
#[test]
fn opcode_o83_const_base_is_const_pow_exp() {
let e = parse_one_expr(1, "o83\nn2\nv0\n");
match &e {
Expr::Binary(BinOp::Pow, base, exp) => {
match **base {
Expr::Const(c) => assert!((c - 2.0).abs() < 1e-12, "base const = {c}"),
ref other => panic!("o83 base must be Const(2.0), got {other:?}"),
}
assert!(
matches!(**exp, Expr::Var(0)),
"exponent must be the variable"
);
}
other => panic!("o83 must parse to Pow(const, var), got {other:?}"),
}
assert!((eval_expr(&e, &[3.0]) - 8.0).abs() < 1e-12);
assert!((eval_expr(&e, &[0.0]) - 1.0).abs() < 1e-12);
let mut g = [0.0_f64; 1];
grad_expr(&e, &[3.0], 1.0, &mut g);
assert!(
(g[0] - 8.0 * 2.0_f64.ln()).abs() < 1e-9,
"grad at 3 = {} (want {})",
g[0],
8.0 * 2.0_f64.ln()
);
}
#[test]
fn power_specializations_agree_with_general_o5() {
let o5_sq = parse_one_expr(1, "o5\nv0\nn2\n"); let o82 = parse_one_expr(1, "o82\nv0\n");
let o5_cube = parse_one_expr(1, "o5\nv0\nn3\n"); let o81 = parse_one_expr(1, "o81\nv0\nn3\n");
let o5_exp = parse_one_expr(1, "o5\nn2\nv0\n"); let o83 = parse_one_expr(1, "o83\nn2\nv0\n");
for &x in &[-2.0_f64, -0.5, 0.0, 1.0, 2.5, 4.0] {
assert!((eval_expr(&o82, &[x]) - eval_expr(&o5_sq, &[x])).abs() < 1e-12);
assert!((eval_expr(&o81, &[x]) - eval_expr(&o5_cube, &[x])).abs() < 1e-12);
assert!((eval_expr(&o83, &[x]) - eval_expr(&o5_exp, &[x])).abs() < 1e-12);
}
}
#[test]
fn power_opcodes_round_trip_through_parse_nl_text() {
let nl = SIMPLE.replace(
"o0\no5\no1\nv0\nn1\nn2\no5\no1\nv1\nn2\nn2\n",
"o0\no82\nv0\no82\nv1\n",
);
assert_ne!(nl, SIMPLE, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse o82 objective");
assert!((eval_expr(&p.obj_expr(), &[3.0, 4.0]) - 25.0).abs() < 1e-12);
assert!((eval_expr(&p.obj_expr(), &[-3.0, -4.0]) - 25.0).abs() < 1e-12);
}
#[test]
fn power_opcode_o81_evaluates_through_the_tape_at_negative_base() {
let nl = SIMPLE.replace(
"o0\no5\no1\nv0\nn1\nn2\no5\no1\nv1\nn2\nn2\n",
"o0\no81\nv0\nn3\no81\nv1\nn3\n",
);
assert_ne!(nl, SIMPLE, "fixture substitution must apply");
let p = parse_nl_text(&nl).expect("parse o81 objective");
let mut tnlp = NlTnlp::new(p);
tnlp.get_nlp_info().unwrap();
let f = tnlp.eval_f(&[-2.0, 1.0], true).unwrap();
assert!((f + 7.0).abs() < 1e-12, "f(-2,1) = {f}");
let mut g = [0.0_f64; 2];
assert!(tnlp.eval_grad_f(&[-2.0, 1.0], true, &mut g));
assert!((g[0] - 12.0).abs() < 1e-9, "df/dx0 = {}", g[0]);
assert!((g[1] - 3.0).abs() < 1e-9, "df/dx1 = {}", g[1]);
}
const SHARED_CSE: &str = "g3 1 1 0
3 3 1 0 0
3 0
0 0
3 0 0
0 0 0 1
0 0 0 0 0
8 3
0 0
0 1 0 0 0
V3 2 0
0 2.0
1 3.0
n0
C0
o5
v3
n2
C1
o0
o5
v3
n3
v2
C2
{BODY2}
O0 0
n0
r
2 0
2 0
2 0
b
3
3
3
k2
3
6
J0 2
0 0
1 0
J1 3
0 0
1 0
2 0
J2 3
0 0
1 0
2 0
G0 3
0 1.0
1 1.0
2 1.0
";
fn shared_cse_nl(body2: &str) -> String {
SHARED_CSE.replace("{BODY2}", body2)
}
#[test]
fn shared_cse_jacobian_matches_flat_tape_bit_for_bit() {
let nl = shared_cse_nl("o2\nv3\nv2");
let p = parse_nl_text(&nl).expect("parse shared-CSE model");
let mut hybrid = NlTnlp::new(p.clone());
let info = hybrid.get_nlp_info().unwrap();
let nnz = info.nnz_jac_g as usize;
hybrid
.con_hybrid
.as_mut()
.expect("CSE shared by 3 constraints must build the hybrid tape")
.use_for_jac = true;
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
for x in [[1.0, 1.0, 1.0], [-2.0, 0.5, 3.0], [0.0, -1.5, -0.25]] {
let mut jh = vec![0.0_f64; nnz];
let mut jf = vec![0.0_f64; nnz];
assert!(hybrid.eval_jac_g(Some(&x), true, SparsityRequest::Values { values: &mut jh }));
assert!(flat.eval_jac_g(Some(&x), true, SparsityRequest::Values { values: &mut jf }));
assert_eq!(
jh, jf,
"hybrid Jacobian differs from the flat tape at {x:?}"
);
let s = 2.0 * x[0] + 3.0 * x[1];
let want = [
4.0 * s,
6.0 * s,
6.0 * s * s,
9.0 * s * s,
1.0,
2.0 * x[2],
3.0 * x[2],
s,
];
assert_eq!(nnz, want.len());
for k in 0..nnz {
assert!(
(jh[k] - want[k]).abs() < 1e-9,
"entry {k} at {x:?}: got {}, want {}",
jh[k],
want[k]
);
}
}
}
fn shared_body_chain_nl(m: usize, body: &str) -> String {
let n = m + 2;
let nzc = 3 * m;
let mut s = String::new();
s.push_str("g3 1 1 0\n");
s.push_str(&format!(" {n} {m} 1 0 0 0\n"));
s.push_str(&format!(" {m} 0\n 0 0\n"));
s.push_str(&format!(" {n} 0 0\n"));
s.push_str(" 0 0 0 1\n 0 0 0 0 0\n");
s.push_str(&format!(" {nzc} {n}\n"));
s.push_str(" 0 0\n 0 1 0 0 0\n");
s.push_str(&format!("V{n} 0 0\n"));
s.push_str(body);
for i in 0..m {
s.push_str(&format!("C{i}\no2\nv{n}\nv{}\n", i + 2));
}
s.push_str("O0 0\nn0\n");
s.push_str(&format!("x{n}\n"));
for j in 0..n {
s.push_str(&format!("{j} {}\n", 0.3 + 0.05 * j as f64));
}
s.push_str("r\n");
for _ in 0..m {
s.push_str("2 0\n");
}
s.push_str("b\n");
for _ in 0..n {
s.push_str("3\n");
}
s.push_str(&format!("k{}\n", n - 1));
let mut acc = 0;
for j in 0..n - 1 {
acc += if j < 2 { m } else { 1 };
s.push_str(&format!("{acc}\n"));
}
for i in 0..m {
s.push_str(&format!("J{i} 3\n0 0.0\n1 0.0\n{} 0.0\n", i + 2));
}
s.push_str(&format!("G0 {n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.0\n"));
}
s
}
fn deep_shared_cse_nl(m: usize, depth: usize) -> String {
let mut body = String::new();
for _ in 0..depth {
body.push_str("o44\n");
}
body.push_str("o2\nn0.01\no0\nv0\nv1\n");
shared_body_chain_nl(m, &body)
}
fn bounded_deep_shared_cse_nl(m: usize, pairs: usize) -> String {
let mut body = String::new();
for _ in 0..pairs {
body.push_str("o43\no44\n");
}
body.push_str("o0\nn2\no2\nn0.01\no0\nv0\nv1\n");
shared_body_chain_nl(m, &body)
}
#[test]
fn a_deep_shared_body_turns_the_jacobian_gate_on_and_still_agrees() {
let p = parse_nl_text(&deep_shared_cse_nl(16, 40)).expect("parse");
let mut hybrid = NlTnlp::new(p.clone());
let info = hybrid.get_nlp_info().unwrap();
let nnz = info.nnz_jac_g as usize;
assert!(
hybrid
.con_hybrid
.as_ref()
.expect("shared CSE must build the hybrid tape")
.use_for_jac,
"a 40-deep body shared by 16 rows is well past the op-ratio gate"
);
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
for scale in [1.0_f64, -0.7, 2.5] {
let x: Vec<f64> = (0..info.n as usize)
.map(|j| scale * (0.2 + 0.03 * j as f64))
.collect();
let mut jh = vec![0.0_f64; nnz];
let mut jf = vec![0.0_f64; nnz];
assert!(hybrid.eval_jac_g(Some(&x), true, SparsityRequest::Values { values: &mut jh }));
assert!(flat.eval_jac_g(Some(&x), true, SparsityRequest::Values { values: &mut jf }));
assert_eq!(jh, jf, "gate-on Jacobian differs from the flat tape");
assert!(jh.iter().any(|v| *v != 0.0), "all-zero Jacobian is no test");
}
}
#[test]
fn a_small_shared_body_leaves_the_jacobian_on_the_flat_path() {
let p = parse_nl_text(&shared_cse_nl("o2\nv3\nv2")).expect("parse");
let mut t = NlTnlp::new(p);
t.get_nlp_info().unwrap();
let h = t.con_hybrid.as_ref().expect("hybrid built for eval_g");
assert!(
!h.use_for_jac,
"a 3-row model with a 2-term CSE is far below the op-ratio gate"
);
}
#[test]
fn shared_cse_hessian_matches_flat_tape_bit_for_bit() {
let nl = shared_cse_nl("o2\nv3\nv2");
let p = parse_nl_text(&nl).expect("parse shared-CSE model");
let mut hybrid = NlTnlp::new(p.clone());
let info = hybrid.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
hybrid
.con_hybrid
.as_mut()
.expect("CSE shared by 3 constraints must build the hybrid tape")
.use_for_hess = true;
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
let pairs: Vec<(usize, usize)> = hybrid
.h_irow
.iter()
.zip(&hybrid.h_jcol)
.map(|(&i, &j)| (i as usize, j as usize))
.collect();
for lam in [[0.5, -1.25, 2.0], [0.0, 1.0, 0.5]] {
for x in [[1.0, 1.0, 1.0], [-2.0, 0.5, 3.0], [0.0, -1.5, -0.25]] {
let mut hh = vec![0.0_f64; nnz];
let mut hf = vec![0.0_f64; nnz];
assert!(hybrid.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hh }
));
assert!(flat.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hf }
));
assert_eq!(
hh, hf,
"hybrid Hessian differs from the flat tape at {x:?}, λ = {lam:?}"
);
let s = 2.0 * x[0] + 3.0 * x[1];
let q = 2.0 * lam[0] + 6.0 * s * lam[1];
let dv = [2.0, 3.0, 0.0];
for (k, &(i, j)) in pairs.iter().enumerate() {
let mut want = q * dv[i] * dv[j];
if i == 2 {
want += lam[2] * dv[j];
}
if j == 2 {
want += lam[2] * dv[i];
}
assert!(
(hh[k] - want).abs() < 1e-9,
"entry ({i}, {j}) at {x:?}, λ = {lam:?}: got {}, want {want}",
hh[k]
);
}
}
}
}
#[test]
fn a_deep_shared_body_turns_the_hessian_gate_on_and_still_agrees() {
let m = 16;
let p = parse_nl_text(&bounded_deep_shared_cse_nl(m, 20)).expect("parse");
let mut hybrid = NlTnlp::new(p.clone());
let info = hybrid.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
assert!(
hybrid
.con_hybrid
.as_ref()
.expect("shared CSE must build the hybrid tape")
.use_for_hess,
"a 40-op body shared by 16 rows is well past the op-ratio gate"
);
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
let lam: Vec<f64> = (0..m).map(|k| 0.25 + 0.125 * k as f64).collect();
for scale in [1.0_f64, -0.7, 2.5] {
let x: Vec<f64> = (0..info.n as usize)
.map(|j| scale * (0.2 + 0.03 * j as f64))
.collect();
let mut jac = vec![0.0_f64; info.nnz_jac_g as usize];
assert!(hybrid.eval_jac_g(
Some(&x),
true,
SparsityRequest::Values { values: &mut jac }
));
let mut hh = vec![0.0_f64; nnz];
let mut hf = vec![0.0_f64; nnz];
assert!(hybrid.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hh }
));
assert!(flat.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hf }
));
for k in 0..nnz {
assert!(
hh[k].is_finite() && hf[k].is_finite(),
"non-finite Hessian entry {k} defeats the comparison"
);
let tol = 1e-12 * hf[k].abs().max(1.0);
assert!(
(hh[k] - hf[k]).abs() <= tol,
"gate-on Hessian entry {k} at scale {scale}: hybrid {} vs flat {}",
hh[k],
hf[k]
);
}
assert!(hh.iter().any(|v| *v != 0.0), "all-zero Hessian is no test");
}
}
fn two_block_shared_cse_nl(wide: usize, narrow: usize, rows: usize) -> String {
let nvars = wide + narrow;
let m = 2 * rows;
let n = nvars + m;
let nzc = rows * (wide + 1) + rows * (narrow + 1);
let mut s = String::new();
s.push_str("g3 1 1 0\n");
s.push_str(&format!(" {n} {m} 1 0 0 0\n"));
s.push_str(&format!(" {m} 0\n 0 0\n"));
s.push_str(&format!(" {n} 0 0\n"));
s.push_str(" 0 0 0 1\n 0 0 0 0 0\n");
s.push_str(&format!(" {nzc} {n}\n"));
s.push_str(" 0 0\n 0 2 0 0 0\n");
for (b, (base, count)) in [(0, wide), (wide, narrow)].iter().enumerate() {
s.push_str(&format!("V{} 0 0\n", n + b));
s.push_str("o41\n");
for _ in 0..count - 1 {
s.push_str("o0\n");
}
for j in 0..*count {
s.push_str(&format!("v{}\n", base + j));
}
}
for i in 0..m {
let b = i / rows;
s.push_str(&format!("C{i}\no2\nv{}\nv{}\n", n + b, nvars + i));
}
s.push_str("O0 0\nn0\n");
s.push_str(&format!("x{n}\n"));
for j in 0..n {
s.push_str(&format!("{j} {}\n", 0.2 + 0.01 * j as f64));
}
s.push_str("r\n");
for _ in 0..m {
s.push_str("2 0\n");
}
s.push_str("b\n");
for _ in 0..n {
s.push_str("3\n");
}
s.push_str(&format!("k{}\n", n - 1));
let mut acc = 0;
for j in 0..n - 1 {
acc += if j < nvars { rows } else { 1 };
s.push_str(&format!("{acc}\n"));
}
for i in 0..m {
let (base, count) = if i < rows { (0, wide) } else { (wide, narrow) };
s.push_str(&format!("J{i} {}\n", count + 1));
let mut cols: Vec<usize> = (base..base + count).collect();
cols.push(nvars + i);
cols.sort_unstable();
for c in cols {
s.push_str(&format!("{c} 0.0\n"));
}
}
s.push_str(&format!("G0 {n}\n"));
for j in 0..n {
s.push_str(&format!("{j} 0.0\n"));
}
s
}
#[test]
fn per_color_prelude_reach_skips_bodies_the_color_cannot_touch() {
let p = parse_nl_text(&two_block_shared_cse_nl(6, 2, 3)).expect("parse");
let mut hybrid = NlTnlp::new(p.clone());
let info = hybrid.get_nlp_info().unwrap();
let nnz = info.nnz_h_lag as usize;
let m = info.m as usize;
{
let h = hybrid
.con_hybrid
.as_mut()
.expect("two shared CSE bodies must build the hybrid tape");
h.use_for_hess = true;
let np = h.tape.n_prelude_ops();
let n_colors = h.hess_color_reach_off.len() - 1;
let total: usize = h.hess_color_reach.len();
assert!(np > 0 && n_colors > 1, "np={np} n_colors={n_colors}");
assert!(
total < n_colors * np,
"per-color reach must be strictly smaller than walking the whole \
prelude per color: Σ|reach_c| = {total}, n_colors × |prelude| = {}",
n_colors * np
);
for c in 0..n_colors {
let r =
&h.hess_color_reach[h.hess_color_reach_off[c]..h.hess_color_reach_off[c + 1]];
assert!(
r.windows(2).all(|w| w[0] < w[1]),
"color {c} reach is not strictly ascending"
);
let member: std::collections::HashSet<u32> = r.iter().copied().collect();
for &i in r {
let (a, b) = crate::nl_tape::op_operands(&h.tape.prelude[i as usize]);
for opnd in [a, b].into_iter().flatten() {
assert!(
member.contains(&(opnd as u32)),
"color {c}: slot {i}'s operand {opnd} is missing from its reach"
);
}
}
}
}
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
let lam: Vec<f64> = (0..m).map(|k| 0.3 + 0.2 * k as f64).collect();
for scale in [1.0_f64, -0.6] {
let x: Vec<f64> = (0..info.n as usize)
.map(|j| scale * (0.15 + 0.02 * j as f64))
.collect();
let mut hh = vec![0.0_f64; nnz];
let mut hf = vec![0.0_f64; nnz];
assert!(hybrid.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hh }
));
assert!(flat.eval_h(
Some(&x),
true,
1.0,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hf }
));
for k in 0..nnz {
let tol = 1e-12 * hf[k].abs().max(1.0);
assert!(
(hh[k] - hf[k]).abs() <= tol,
"narrowed-reach Hessian entry {k} at scale {scale}: \
hybrid {} vs flat {}",
hh[k],
hf[k]
);
}
assert!(hh.iter().any(|v| *v != 0.0), "all-zero Hessian is no test");
}
}
#[test]
fn a_small_shared_body_leaves_the_hessian_on_the_flat_path() {
let p = parse_nl_text(&shared_cse_nl("o2\nv3\nv2")).expect("parse");
let mut t = NlTnlp::new(p);
t.get_nlp_info().unwrap();
let h = t.con_hybrid.as_ref().expect("hybrid built for eval_g");
assert!(
!h.use_for_hess,
"a 3-row model with a 2-term CSE is below the Hessian op-ratio gate"
);
}
#[test]
fn shared_cse_constraint_tape_matches_flat_tape_bit_for_bit() {
let nl = shared_cse_nl("o2\nv3\nv2");
let p = parse_nl_text(&nl).expect("parse shared-CSE model");
let mut hybrid = NlTnlp::new(p.clone());
hybrid.get_nlp_info().unwrap();
let h = hybrid
.con_hybrid
.as_ref()
.expect("CSE shared by 3 constraints must take the hybrid path");
assert!(
h.tape.n_prelude_ops() > 0,
"shared CSE body must land in the prelude"
);
let mut flat = NlTnlp::new(p);
flat.get_nlp_info().unwrap();
flat.con_hybrid = None;
for x in [[1.0, 1.0, 1.0], [-2.0, 0.5, 3.0], [0.0, -1.5, -0.25]] {
let mut gh = [0.0_f64; 3];
let mut gf = [0.0_f64; 3];
assert!(hybrid.eval_g(&x, true, &mut gh));
assert!(flat.eval_g(&x, true, &mut gf));
let s = 2.0 * x[0] + 3.0 * x[1];
let want = [s * s, s * s * s + x[2], s * x[2]];
for i in 0..3 {
assert_eq!(gh[i], gf[i], "row {i} differs from the flat tape at {x:?}");
assert!(
(gh[i] - want[i]).abs() < 1e-9,
"row {i}: got {}, want {}",
gh[i],
want[i]
);
}
}
}
#[test]
fn unsupported_opcode_falls_back_to_the_flat_tape() {
let nl = shared_cse_nl("o11\n2\nv3\nv2");
let p = parse_nl_text(&nl).expect("parse min-list model");
let mut tnlp = NlTnlp::new(p);
tnlp.get_nlp_info().unwrap();
assert!(
tnlp.con_hybrid.is_none(),
"a min-list anywhere in the constraint block must disable the hybrid path"
);
let mut g = [0.0_f64; 3];
assert!(tnlp.eval_g(&[-2.0, 0.5, 3.0], true, &mut g));
let s = 2.0 * -2.0 + 3.0 * 0.5; assert!((g[0] - s * s).abs() < 1e-9);
assert!((g[2] - s.min(3.0)).abs() < 1e-9, "min(V3, x2) = {}", g[2]);
}
fn v(i: usize) -> Expr {
Expr::Var(i)
}
fn c(x: Number) -> Expr {
Expr::Const(x)
}
fn bin(op: BinOp, a: Expr, b: Expr) -> Expr {
Expr::Binary(op, Box::new(a), Box::new(b))
}
fn un(op: UnaryOp, a: Expr) -> Expr {
Expr::Unary(op, Box::new(a))
}
fn parts(n: usize, objective: Expr, constraints: Vec<Expr>) -> NlProblemParts {
let m = constraints.len();
NlProblemParts {
minimize: true,
objective,
obj_constant: 0.0,
constraints,
x_l: vec![-1e19; n],
x_u: vec![1e19; n],
x0: vec![0.0; n],
g_l: vec![-1e19; m],
g_u: vec![1e19; m],
var_names: Vec::new(),
con_names: Vec::new(),
}
}
#[test]
fn from_expressions_builds_evaluable_problem() {
let rosen = bin(
BinOp::Add,
bin(BinOp::Pow, bin(BinOp::Sub, c(1.0), v(0)), c(2.0)),
bin(
BinOp::Mul,
c(100.0),
bin(
BinOp::Pow,
bin(BinOp::Sub, v(1), bin(BinOp::Pow, v(0), c(2.0))),
c(2.0),
),
),
);
let circle = bin(
BinOp::Add,
bin(BinOp::Pow, v(0), c(2.0)),
bin(BinOp::Pow, v(1), c(2.0)),
);
let mut p = parts(2, rosen, vec![circle]);
p.g_l = vec![0.0];
p.g_u = vec![2.0];
p.x0 = vec![-1.2, 1.0];
p.var_names = names(&["x", "y"]);
p.con_names = names(&["circle"]);
let prob = NlProblem::from_expressions(p).expect("build");
assert_eq!((prob.n, prob.m), (2, 1));
assert_eq!(prob.var_names, names(&["x", "y"]));
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let f = t.eval_f(&[-1.2, 1.0], true).unwrap();
assert!((f - 24.2).abs() < 1e-10, "f = {f}");
let (x0, x1) = (-1.2, 1.0);
let want = [
-2.0 * (1.0 - x0) - 400.0 * x0 * (x1 - x0 * x0),
200.0 * (x1 - x0 * x0),
];
let mut g = [0.0_f64; 2];
assert!(t.eval_grad_f(&[x0, x1], true, &mut g));
for j in 0..2 {
assert!((g[j] - want[j]).abs() < 1e-8, "g[{j}] = {} ", g[j]);
}
let mut gv = [0.0_f64; 1];
assert!(t.eval_g(&[x0, x1], true, &mut gv));
assert!((gv[0] - 2.44).abs() < 1e-10, "g = {}", gv[0]);
}
#[test]
fn from_expressions_carries_ops_nl_cannot_express() {
let obj = Expr::Sum(vec![
bin(BinOp::Atan2, v(0), v(1)),
Expr::MinList(vec![v(0), v(1)]),
Expr::MaxList(vec![v(0), v(1)]),
un(UnaryOp::Erf, v(0)),
]);
let prob = NlProblem::from_expressions(parts(2, obj, Vec::new())).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let x: [Number; 2] = [0.8, 1.5];
let want = x[0].atan2(x[1]) + x[0] + x[1] + crate::nl_tape::erf(x[0]);
let f = t.eval_f(&x, true).unwrap();
assert!((f - want).abs() < 1e-12, "f = {f}, want {want}");
}
#[test]
fn from_expressions_rejects_out_of_range_var() {
let err = NlProblem::from_expressions(parts(2, v(5), Vec::new()))
.expect_err("Var(5) with n = 2 must be rejected");
assert!(err.contains("Var(5)"), "{err}");
let err = NlProblem::from_expressions(parts(2, c(0.0), vec![v(2)]))
.expect_err("constraint Var(2) with n = 2 must be rejected");
assert!(err.contains("constraint 0"), "{err}");
let mut p = parts(2, c(0.0), Vec::new());
p.x0 = vec![0.0; 3];
let err = NlProblem::from_expressions(p).expect_err("x0 length must be checked");
assert!(err.contains("x0"), "{err}");
}
#[test]
fn from_expressions_finds_out_of_range_vars_in_every_position() {
let inner_cse = Arc::new(v(7));
let cases: Vec<(&str, Expr)> = vec![
("bare", v(7)),
("cse", Expr::Cse(Arc::new(v(7)))),
("nested cse", Expr::Cse(Arc::new(Expr::Cse(inner_cse)))),
(
"cond branch",
Expr::Cond {
cond: Box::new(c(1.0)),
then_: Box::new(v(7)),
else_: Box::new(c(0.0)),
},
),
("min list", Expr::MinList(vec![c(0.0), v(7)])),
(
"sum",
Expr::Sum(vec![c(0.0), bin(BinOp::Mul, c(2.0), v(7))]),
),
];
for (label, e) in cases {
let err = NlProblem::from_expressions(parts(2, e, Vec::new()))
.err()
.unwrap_or_else(|| panic!("{label}: Var(7) with n = 2 should be rejected"));
assert!(err.contains("Var(7)"), "{label}: {err}");
}
}
fn share_dag(depth: usize) -> Expr {
let mut e = v(0);
for _ in 0..depth {
let shared = Arc::new(e);
e = bin(
BinOp::Add,
Expr::Cse(Arc::clone(&shared)),
Expr::Cse(shared),
);
}
e
}
#[test]
fn shared_dag_walks_are_memoized_not_exponential() {
const DEPTH: usize = 30;
let e = share_dag(DEPTH);
let mut vars = BTreeSet::new();
collect_vars(&e, &mut vars);
assert_eq!(vars.iter().copied().collect::<Vec<_>>(), vec![0]);
let mut ids = BTreeSet::new();
super::super::nl_external::collect_funcall_ids(&e, &mut ids);
assert!(ids.is_empty());
let prob = NlProblem::from_expressions(parts(1, e, Vec::new())).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let mut lin = vec![Linearity::Linear; 1];
assert!(t.get_variables_linearity(&mut lin));
}
#[test]
fn from_expressions_rejects_imported_function_calls() {
let call = Expr::Funcall {
id: 0,
args: vec![FuncallArg::Real(v(0))],
};
let err = NlProblem::from_expressions(parts(1, call.clone(), Vec::new()))
.expect_err("a Funcall must be rejected, not deferred to AMPLFUNC");
assert!(err.contains("imported function"), "{err}");
assert!(
err.contains("read_nl") || err.contains("parse_nl_text"),
"the error must point at the paths that do support externals: {err}"
);
let buried = Expr::Cse(Arc::new(Expr::Sum(vec![c(1.0), call])));
let err = NlProblem::from_expressions(parts(1, c(0.0), vec![buried]))
.expect_err("a buried Funcall must be rejected too");
assert!(err.contains("constraint 0"), "{err}");
}
#[test]
fn hessian_vector_product_matches_dense_hessian() {
let obj = Expr::Sum(vec![
bin(BinOp::Mul, v(0), bin(BinOp::Mul, v(1), v(2))),
un(UnaryOp::Exp, bin(BinOp::Mul, v(0), v(1))),
un(UnaryOp::Erf, v(2)),
]);
let cons = vec![
bin(
BinOp::Add,
bin(BinOp::Pow, v(0), c(2.0)),
un(UnaryOp::Sin, v(2)),
),
bin(BinOp::Mul, v(1), v(2)),
];
let prob = NlProblem::from_expressions(parts(3, obj, cons)).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
let info = t.get_nlp_info().unwrap();
let x = [0.3, -0.7, 1.1];
let lam = [0.5, -1.25];
let obj_factor = 2.0;
let nnz = info.nnz_h_lag as usize;
let (mut irow, mut jcol) = (vec![0_i32; nnz], vec![0_i32; nnz]);
assert!(t.eval_h(
None,
false,
1.0,
None,
false,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let mut hvals = vec![0.0_f64; nnz];
assert!(t.eval_h(
Some(&x),
true,
obj_factor,
Some(&lam),
true,
SparsityRequest::Values { values: &mut hvals }
));
let mut dense = [[0.0_f64; 3]; 3];
for k in 0..nnz {
let (i, j) = (irow[k] as usize, jcol[k] as usize);
dense[i][j] += hvals[k];
if i != j {
dense[j][i] += hvals[k];
}
}
let seeds: [[Number; 3]; 4] = [
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[0.4, -1.3, 2.0],
];
let mut out = vec![0.0; 3];
for s in &seeds {
t.hessian_vector_product(&x, s, obj_factor, Some(&lam), &mut out)
.expect("hvp");
for i in 0..3 {
let want: Number = (0..3).map(|j| dense[i][j] * s[j]).sum();
assert!(
(out[i] - want).abs() < 1e-9,
"seed {s:?} row {i}: hvp={:.9e} dense={want:.9e}",
out[i]
);
}
}
}
#[test]
fn hessian_vector_product_defaults_and_validation() {
let obj = bin(
BinOp::Add,
bin(BinOp::Pow, v(0), c(2.0)),
bin(BinOp::Mul, c(3.0), bin(BinOp::Mul, v(0), v(1))),
);
let prob = NlProblem::from_expressions(parts(2, obj, Vec::new())).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let mut out = vec![7.0, -7.0]; t.hessian_vector_product(&[0.5, 2.0], &[1.0, 1.0], 1.0, None, &mut out)
.expect("hvp");
assert!((out[0] - 5.0).abs() < 1e-12, "out = {out:?}");
assert!((out[1] - 3.0).abs() < 1e-12, "out = {out:?}");
t.hessian_vector_product(&[0.5, 2.0], &[1.0, 1.0], -2.0, None, &mut out)
.expect("hvp");
assert!((out[0] + 10.0).abs() < 1e-12, "out = {out:?}");
let mut short = vec![0.0; 1];
assert!(
t.hessian_vector_product(&[0.5, 2.0], &[1.0, 1.0], 1.0, None, &mut short)
.is_err()
);
assert!(
t.hessian_vector_product(&[0.5], &[1.0, 1.0], 1.0, None, &mut out)
.is_err()
);
assert!(
t.hessian_vector_product(&[0.5, 2.0], &[1.0], 1.0, None, &mut out)
.is_err()
);
}
#[test]
fn hessian_vector_products_on_a_sparse_hessian() {
const N: usize = 8;
let mut terms = Vec::new();
for i in 0..N - 1 {
terms.push(bin(BinOp::Pow, bin(BinOp::Mul, v(i), v(i + 1)), c(2.0)));
}
for i in 0..N {
terms.push(un(UnaryOp::Exp, v(i)));
}
let prob =
NlProblem::from_expressions(parts(N, Expr::Sum(terms), Vec::new())).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
let info = t.get_nlp_info().unwrap();
assert_eq!(
info.nnz_h_lag as usize,
2 * N - 1,
"chain objective should give a tridiagonal Hessian, not a dense one"
);
let x: Vec<Number> = (0..N).map(|i| 0.2 + 0.1 * i as Number).collect();
let nnz = info.nnz_h_lag as usize;
let (mut irow, mut jcol) = (vec![0_i32; nnz], vec![0_i32; nnz]);
assert!(t.eval_h(
None,
false,
1.0,
None,
false,
SparsityRequest::Structure {
irow: &mut irow,
jcol: &mut jcol
}
));
let mut hvals = vec![0.0; nnz];
assert!(t.eval_h(
Some(&x),
true,
1.0,
None,
true,
SparsityRequest::Values { values: &mut hvals }
));
let mut dense = vec![vec![0.0; N]; N];
for k in 0..nnz {
let (i, j) = (irow[k] as usize, jcol[k] as usize);
dense[i][j] += hvals[k];
if i != j {
dense[j][i] += hvals[k];
}
}
let mut seeds = vec![0.0; N * N];
for cc in 0..N {
seeds[cc * N + cc] = 1.0;
}
let mut out = vec![0.0; N * N];
t.hessian_vector_products(&x, &seeds, N, 1.0, None, &mut out)
.expect("block hvp");
for cc in 0..N {
for i in 0..N {
assert!(
(out[cc * N + i] - dense[i][cc]).abs() < 1e-9,
"H[{i},{cc}]: block={:.9e} sparse={:.9e}",
out[cc * N + i],
dense[i][cc]
);
}
}
}
#[test]
fn hessian_vector_products_match_repeated_single_calls() {
let obj = Expr::Sum(vec![
un(UnaryOp::Exp, bin(BinOp::Mul, v(0), v(1))),
bin(BinOp::Pow, v(2), c(4.0)),
bin(BinOp::Mul, v(0), v(2)),
]);
let cons = vec![bin(BinOp::Mul, v(1), v(2))];
let prob = NlProblem::from_expressions(parts(3, obj, cons)).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let x = [0.4, -0.6, 1.3];
let lam = [0.75];
let cols: [[Number; 3]; 4] = [
[1.0, 2.0, -3.0],
[0.0, 0.0, 0.0], [0.5, 0.0, 0.0],
[-1.0, 1.0, 1.0],
];
let mut block = vec![0.0; 3 * cols.len()];
let flat: Vec<Number> = cols.iter().flat_map(|c| c.iter().copied()).collect();
t.hessian_vector_products(&x, &flat, cols.len(), 1.0, Some(&lam), &mut block)
.expect("block hvp");
for (c, col) in cols.iter().enumerate() {
let mut single = vec![0.0; 3];
t.hessian_vector_product(&x, col, 1.0, Some(&lam), &mut single)
.expect("single hvp");
for i in 0..3 {
assert!(
(block[c * 3 + i] - single[i]).abs() < 1e-12,
"direction {c} row {i}: block={:.12e} single={:.12e}",
block[c * 3 + i],
single[i]
);
}
}
assert!(block[3..6].iter().all(|&z| z == 0.0), "{block:?}");
}
#[test]
fn hessian_vector_products_validate_block_shape() {
let prob = NlProblem::from_expressions(parts(2, bin(BinOp::Pow, v(0), c(2.0)), Vec::new()))
.expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let mut empty: Vec<Number> = Vec::new();
assert!(
t.hessian_vector_products(&[1.0, 1.0], &[], 0, 1.0, None, &mut empty)
.is_ok()
);
let mut out = vec![0.0; 4];
assert!(
t.hessian_vector_products(&[1.0, 1.0], &[1.0, 1.0], 2, 1.0, None, &mut out)
.is_err()
);
let mut short = vec![0.0; 2];
assert!(
t.hessian_vector_products(&[1.0, 1.0], &[1.0; 4], 2, 1.0, None, &mut short)
.is_err()
);
}
#[test]
fn hessian_vector_product_respects_maximize_sign() {
let obj = bin(BinOp::Pow, v(0), c(2.0));
let mut p = parts(1, obj, Vec::new());
p.minimize = false;
let prob = NlProblem::from_expressions(p).expect("build");
let mut t = NlTnlp::try_new(prob).expect("tnlp");
t.get_nlp_info().unwrap();
let mut out = vec![0.0; 1];
t.hessian_vector_product(&[1.0], &[1.0], 1.0, None, &mut out)
.expect("hvp");
assert!((out[0] + 2.0).abs() < 1e-12, "out = {out:?}");
}
}