use std::collections::HashMap;
use crate::expression::contains;
use crate::substitute::substitute;
use crate::types::{DaeInfo, Equation, EsmFile, Expr, Model};
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
#[error("{code}: {message}")]
pub struct DaeError {
pub code: String,
pub message: String,
}
#[derive(Debug, Clone, Copy)]
pub struct DiscretizeOptions {
pub dae_support: bool,
}
impl Default for DiscretizeOptions {
fn default() -> Self {
Self {
dae_support: default_dae_support(),
}
}
}
pub fn default_dae_support() -> bool {
match std::env::var("ESM_DAE_SUPPORT") {
Ok(v) => !matches!(
v.trim().to_ascii_lowercase().as_str(),
"0" | "false" | "no" | "off"
),
Err(_) => true,
}
}
pub fn apply_dae_contract(esm: &mut EsmFile, dae_support: bool) -> Result<DaeInfo, DaeError> {
let mut pre_factor_count = 0usize;
let mut first_path: Option<String> = None;
if let Some(models) = esm.models.as_ref() {
for (mname, model) in models.iter() {
for (i, eq) in model.equations.iter().enumerate() {
if is_algebraic(eq) {
pre_factor_count += 1;
if first_path.is_none() {
first_path = Some(format!("models.{mname}.equations[{i}]"));
}
}
}
}
}
if pre_factor_count == 0 {
let info = DaeInfo {
algebraic_equation_count: 0,
per_model: per_model_zero(esm),
factored_equation_count: Some(0),
};
stamp_metadata(esm, "ode", &info);
return Ok(info);
}
if !dae_support {
let where_ = first_path.as_deref().unwrap_or("(unknown)");
return Err(DaeError {
code: "E_NO_DAE_SUPPORT".into(),
message: format!(
"discretize() input contains {pre_factor_count} algebraic equation(s) \
(first at {where_}); DAE support is disabled \
(DiscretizeOptions::dae_support=false / ESM_DAE_SUPPORT=0). \
Enable DAE support or remove the algebraic constraint(s). See RFC §12."
),
});
}
let mut factored_total = 0usize;
let mut residual_paths: Vec<String> = Vec::new();
let mut per_model: HashMap<String, usize> = HashMap::new();
if let Some(models) = esm.models.as_mut() {
let mnames: Vec<String> = models.keys().cloned().collect();
for mname in mnames {
let model = models.get_mut(&mname).expect("model key just listed");
let (factored, residual) = factor_model(model);
factored_total += factored;
for idx in &residual {
residual_paths.push(format!("models.{mname}.equations[{idx}]"));
}
per_model.insert(mname, residual.len());
}
}
if !residual_paths.is_empty() {
let residual_count = residual_paths.len();
return Err(DaeError {
code: "E_NONTRIVIAL_DAE".into(),
message: format!(
"discretize() input contains {residual_count} non-trivial \
algebraic equation(s) that cannot be factored into the ODE \
system ({factored_total} were factored out): {paths}. \
The rust binding supports only trivially substitutable \
algebraic equations (form `var ~ expr` where `var` does \
not appear in `expr`, acyclic). For full DAE support use \
the Julia binding (EarthSciAST.jl). See RFC §12 \
and docs/rfcs/dae-binding-strategies.md.",
paths = residual_paths.join(", "),
),
});
}
let info = DaeInfo {
algebraic_equation_count: 0,
per_model,
factored_equation_count: Some(factored_total),
};
stamp_metadata(esm, "ode", &info);
Ok(info)
}
pub fn discretize(esm: &EsmFile, options: DiscretizeOptions) -> Result<EsmFile, DaeError> {
let mut out = esm.clone();
let input_name = esm.metadata.name.clone();
apply_dae_contract(&mut out, options.dae_support)?;
out.metadata.discretized_from = Some(crate::types::DiscretizedFrom { name: input_name });
Ok(out)
}
fn is_differential(eq: &Equation) -> bool {
match &eq.lhs {
Expr::Operator(node) => {
node.op == "D" && !crate::op_registry::is_rewrite_target_derivative(node)
}
_ => false,
}
}
fn is_algebraic(eq: &Equation) -> bool {
!is_differential(eq)
}
fn per_model_zero(esm: &EsmFile) -> HashMap<String, usize> {
let mut out = HashMap::new();
if let Some(models) = esm.models.as_ref() {
for name in models.keys() {
out.insert(name.clone(), 0);
}
}
out
}
fn stamp_metadata(esm: &mut EsmFile, system_class: &str, info: &DaeInfo) {
esm.metadata.system_class = Some(system_class.into());
esm.metadata.dae_info = Some(info.clone());
}
fn factor_model(model: &mut Model) -> (usize, Vec<usize>) {
let originals: Vec<(usize, bool, Equation)> = model
.equations
.drain(..)
.enumerate()
.map(|(i, eq)| {
let alg = is_algebraic(&eq);
(i, alg, eq)
})
.collect();
let mut alive: Vec<(usize, bool, Equation)> = originals;
let mut factored: usize = 0;
loop {
let pick = alive.iter().position(|(_, is_alg, eq)| {
if !*is_alg {
return false;
}
match &eq.lhs {
Expr::Variable(v) => !contains(&eq.rhs, v),
_ => false,
}
});
let Some(idx) = pick else { break };
let (_orig_idx, _, picked) = alive.remove(idx);
let var = match picked.lhs {
Expr::Variable(v) => v,
_ => unreachable!("pick guaranteed bare-variable LHS"),
};
let rhs = picked.rhs;
let mut subs: HashMap<String, Expr> = HashMap::new();
subs.insert(var, rhs);
for (_, _, eq) in alive.iter_mut() {
eq.lhs = substitute(&eq.lhs, &subs);
eq.rhs = substitute(&eq.rhs, &subs);
}
for var_def in model.variables.values_mut() {
var_def.for_each_expression_mut(&mut |expr| *expr = substitute(expr, &subs));
}
if let Some(init_eqs) = model.initialization_equations.as_mut() {
for eq in init_eqs.iter_mut() {
eq.lhs = substitute(&eq.lhs, &subs);
eq.rhs = substitute(&eq.rhs, &subs);
}
}
factored += 1;
}
let mut residual: Vec<usize> = Vec::new();
let mut out_eqs: Vec<Equation> = Vec::with_capacity(alive.len());
for (orig_idx, is_alg, eq) in alive {
if is_alg {
residual.push(orig_idx);
}
out_eqs.push(eq);
}
model.equations = out_eqs;
(factored, residual)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::{test_file, var as typed_var};
use crate::types::{Domain, ExpressionNode, ModelVariable, VariableType};
use indexmap::IndexMap;
fn var(name: &str) -> Expr {
Expr::Variable(name.into())
}
fn op(op: &str, args: Vec<Expr>) -> Expr {
Expr::operator(ExpressionNode {
op: op.into(),
args,
..Default::default()
})
}
fn op_wrt(op: &str, args: Vec<Expr>, wrt: &str) -> Expr {
Expr::operator(ExpressionNode {
op: op.into(),
args,
wrt: Some(wrt.into()),
..Default::default()
})
}
fn minimal_esm(model_name: &str, model: Model) -> EsmFile {
let mut models = IndexMap::new();
models.insert(model_name.into(), model);
EsmFile {
models: Some(models),
..test_file()
}
}
fn state_var() -> ModelVariable {
ModelVariable {
default: Some(1.0.into()),
..typed_var(VariableType::Unknown, None)
}
}
fn observed_var() -> ModelVariable {
typed_var(VariableType::Unknown, None)
}
fn param_var() -> ModelVariable {
ModelVariable {
default: Some(0.5.into()),
..typed_var(VariableType::Parameter, None)
}
}
fn empty_model() -> Model {
Model::default()
}
#[test]
fn pure_ode_stamps_ode_system_class() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("k".into(), param_var());
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: op("*", vec![op("-", vec![var("k")]), var("x")]),
});
let mut esm = minimal_esm("M", m);
let info = apply_dae_contract(&mut esm, true).expect("pure ODE");
assert_eq!(info.algebraic_equation_count, 0);
assert_eq!(info.factored_equation_count, Some(0));
assert_eq!(esm.metadata.system_class.as_deref(), Some("ode"));
assert_eq!(
esm.metadata
.dae_info
.as_ref()
.unwrap()
.algebraic_equation_count,
0
);
assert_eq!(esm.models.as_ref().unwrap()["M"].equations.len(), 1);
}
#[test]
fn trivial_observed_is_factored_out() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), observed_var());
m.equations.push(Equation {
comment: None,
lhs: var("y"),
rhs: op("^", vec![var("x"), Expr::Integer(2)]),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let info = apply_dae_contract(&mut esm, true).expect("trivially factored");
assert_eq!(info.algebraic_equation_count, 0);
assert_eq!(info.factored_equation_count, Some(1));
assert_eq!(esm.metadata.system_class.as_deref(), Some("ode"));
let eqs = &esm.models.as_ref().unwrap()["M"].equations;
assert_eq!(eqs.len(), 1, "algebraic eq removed after factoring");
match &eqs[0].rhs {
Expr::Operator(n) if n.op == "^" => {
assert_eq!(n.args.len(), 2);
assert!(matches!(n.args[0], Expr::Variable(ref v) if v == "x"));
assert!(matches!(n.args[1], Expr::Integer(2)));
}
other => panic!("expected x^2 after substitution, got {other:?}"),
}
}
#[test]
fn transitive_trivial_chain_factors_out() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), observed_var());
m.variables.insert("z".into(), observed_var());
m.equations.push(Equation {
comment: None,
lhs: var("z"),
rhs: op("+", vec![var("x"), Expr::Integer(1)]),
});
m.equations.push(Equation {
comment: None,
lhs: var("y"),
rhs: op("*", vec![var("z"), Expr::Integer(2)]),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let info = apply_dae_contract(&mut esm, true).expect("chain factors");
assert_eq!(info.factored_equation_count, Some(2));
assert_eq!(info.algebraic_equation_count, 0);
assert_eq!(esm.models.as_ref().unwrap()["M"].equations.len(), 1);
}
#[test]
fn cyclic_algebraic_errors_nontrivial() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), observed_var());
m.variables.insert("w".into(), observed_var());
m.equations.push(Equation {
comment: None,
lhs: var("y"),
rhs: op("+", vec![var("x"), var("w")]),
});
m.equations.push(Equation {
comment: None,
lhs: var("w"),
rhs: op("-", vec![var("y"), Expr::Integer(1)]),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let err = apply_dae_contract(&mut esm, true).expect_err("cycle");
assert_eq!(err.code, "E_NONTRIVIAL_DAE");
assert!(err.message.contains("models.M.equations"));
assert!(err.message.contains("Julia"));
assert!(err.message.contains("RFC §12"));
}
#[test]
fn implicit_constraint_errors_nontrivial() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), state_var());
m.equations.push(Equation {
comment: None,
lhs: op(
"+",
vec![
op("^", vec![var("x"), Expr::Integer(2)]),
op("^", vec![var("y"), Expr::Integer(2)]),
],
),
rhs: Expr::Integer(1),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let err = apply_dae_contract(&mut esm, true).expect_err("implicit");
assert_eq!(err.code, "E_NONTRIVIAL_DAE");
}
#[test]
fn self_referential_algebraic_errors_nontrivial() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), observed_var());
m.equations.push(Equation {
comment: None,
lhs: var("y"),
rhs: op("+", vec![var("y"), Expr::Integer(1)]),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let err = apply_dae_contract(&mut esm, true).expect_err("self ref");
assert_eq!(err.code, "E_NONTRIVIAL_DAE");
}
#[test]
fn dae_support_disabled_errors_on_algebraic() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.variables.insert("y".into(), observed_var());
m.equations.push(Equation {
comment: None,
lhs: var("y"),
rhs: op("^", vec![var("x"), Expr::Integer(2)]),
});
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: var("y"),
});
let mut esm = minimal_esm("M", m);
let err = apply_dae_contract(&mut esm, false).expect_err("no dae support");
assert_eq!(err.code, "E_NO_DAE_SUPPORT");
assert!(err.message.contains("models.M.equations"));
assert!(err.message.contains("ESM_DAE_SUPPORT"));
}
#[test]
fn dae_support_disabled_accepts_pure_ode() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: op("-", vec![var("x")]),
});
let mut esm = minimal_esm("M", m);
let info = apply_dae_contract(&mut esm, false).expect("pure ODE under disabled");
assert_eq!(info.algebraic_equation_count, 0);
assert_eq!(esm.metadata.system_class.as_deref(), Some("ode"));
}
#[test]
fn discretize_stamps_discretized_from() {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.equations.push(Equation {
comment: None,
lhs: op_wrt("D", vec![var("x")], "t"),
rhs: Expr::Integer(0),
});
let esm = minimal_esm("M", m);
let out = discretize(&esm, DiscretizeOptions { dae_support: true }).expect("ok");
assert_eq!(
out.metadata
.discretized_from
.as_ref()
.and_then(|d| d.name.as_deref()),
Some("test")
);
assert_eq!(out.metadata.system_class.as_deref(), Some("ode"));
}
#[test]
fn the_derivative_axis_is_the_literal_t_not_the_independent_variable() {
let with_wrt = |wrt: Option<&str>| {
let mut m = empty_model();
m.variables.insert("x".into(), state_var());
m.equations.push(Equation {
comment: None,
lhs: match wrt {
Some(w) => op_wrt("D", vec![var("x")], w),
None => op("D", vec![var("x")]),
},
rhs: Expr::Integer(0),
});
let mut esm = minimal_esm("M", m);
esm.domain = Some(Domain {
independent_variable: Some("s".into()),
temporal: None,
element_type: None,
array_type: None,
});
esm
};
for spelling in [Some("t"), None] {
let mut esm = with_wrt(spelling);
let info = apply_dae_contract(&mut esm, true)
.expect("the literal `t` is the structural axis whatever the domain declares");
assert_eq!(info.algebraic_equation_count, 0, "{spelling:?}");
assert_eq!(
esm.metadata.system_class.as_deref(),
Some("ode"),
"{spelling:?}"
);
let m = &esm.models.as_ref().unwrap()["M"];
assert_eq!(crate::ode_states(m), vec!["x".to_string()], "{spelling:?}");
}
let mut esm = with_wrt(Some("s"));
let err = apply_dae_contract(&mut esm, true).expect_err("a spatial `D` is not a tendency");
assert_eq!(err.code, "E_NONTRIVIAL_DAE");
let esm2 = with_wrt(Some("s"));
let m = &esm2.models.as_ref().unwrap()["M"];
assert!(crate::ode_states(m).is_empty());
}
}