#[cfg(feature = "solve")]
use super::tape::TapeProgram;
use super::*;
#[cfg_attr(not(feature = "solve"), allow(unused_imports))]
use crate::simulate::SimulateError;
#[cfg(feature = "solve")]
use crate::simulate::{
Alg, Progress, ProgressFn, ReturnCode, Solution, SolutionMetadata, SolveOptions, SolveStats,
};
#[cfg(feature = "solve")]
use diffsol::{Bdf, FaerLU, FaerMat, NewtonNonlinearSolver, OdeBuilder, Sdirk, VectorHost};
use std::collections::HashSet;
#[cfg(feature = "solve")]
struct TapedObserveds {
scratch: RhsScratch,
dy: Vec<f64>,
}
#[cfg(feature = "solve")]
impl TapedObserveds {
fn new(compiled: &ArrayCompiled, tape: &std::sync::Arc<TapeProgram>) -> Self {
let mut scratch = RhsScratch::new(&compiled.var_shapes);
scratch.set_const_arrays(std::sync::Arc::clone(&compiled.const_scope));
scratch.install_tape(std::sync::Arc::clone(tape));
scratch.set_exports_active(true);
TapedObserveds {
scratch,
dy: vec![0.0f64; compiled.n_states],
}
}
fn scratch_observeds(&self) -> &ArrMap {
self.scratch
.taped_observeds()
.expect("the harvester installs a tape on its own scratch")
}
fn at(&mut self, compiled: &ArrayCompiled, state: &[f64], params: &[f64], t: f64) -> &ArrMap {
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &compiled.rhs_rules,
observed_rules: &compiled.observed_rules,
var_shapes: &compiled.var_shapes,
param_names: &compiled.param_names,
state,
params,
forcing: &compiled.forcing,
t,
declared: &compiled.declared_names,
rule_set: compiled.rule_set,
},
&mut self.dy,
false,
&mut RhsStats::default(),
&mut self.scratch,
);
self.scratch
.taped_observeds()
.expect("the harvester installs a tape on its own scratch")
}
}
#[cfg(feature = "solve")]
fn xla_not_installed() -> SimulateError {
SimulateError::CompilerUnavailable {
compiler: "xla",
details: "this Problem names `xla` but carries no compiled program; the build did not \
install one, and the tape is NOT run in its place — a compiler that cannot \
run says so (esm-libraries-spec §2.5.10)"
.to_string(),
}
}
#[cfg(all(feature = "solve", feature = "xla"))]
fn run_xla_rhs(
program: &crate::xla_runtime::CompiledRhs,
fault: &Rc<RefCell<Option<String>>>,
state: &[f64],
params: &[f64],
t: f64,
out: &mut [f64],
) -> bool {
match program.eval_into(state, params, t, out) {
Ok(()) => true,
Err(e) => {
let mut slot = fault.borrow_mut();
if slot.is_none() {
*slot = Some(e.to_string());
}
for v in out.iter_mut() {
*v = f64::NAN;
}
false
}
}
}
#[cfg(all(feature = "solve", any(test, feature = "xla")))]
struct FdJvp {
base_y: Vec<f64>,
base_p: Vec<f64>,
base_t: f64,
f_y: Vec<f64>,
base_valid: bool,
y_perturbed: Vec<f64>,
f_yp: Vec<f64>,
}
#[cfg(all(feature = "solve", any(test, feature = "xla")))]
impl FdJvp {
fn new(n_states: usize) -> Self {
FdJvp {
base_y: vec![0.0; n_states],
base_p: Vec::new(),
base_t: 0.0,
f_y: vec![0.0; n_states],
base_valid: false,
y_perturbed: vec![0.0; n_states],
f_yp: vec![0.0; n_states],
}
}
fn apply(
&mut self,
y: &[f64],
p: &[f64],
t: f64,
v: &[f64],
jv: &mut [f64],
mut eval: impl FnMut(&[f64], &[f64], f64, &mut [f64]) -> bool,
) {
let n = y.len();
let mut y_norm = 0.0f64;
for &yi in y {
y_norm += yi * yi;
}
let y_norm = y_norm.sqrt().max(1.0);
let eps = f64::EPSILON.sqrt() * y_norm;
let same_point = self.base_valid
&& self.base_t.to_bits() == t.to_bits()
&& same_bits(&self.base_y, y)
&& same_bits(&self.base_p, p);
if !same_point {
self.base_y.clear();
self.base_y.extend_from_slice(y);
self.base_p.clear();
self.base_p.extend_from_slice(p);
self.base_t = t;
self.f_y.resize(n, 0.0);
self.base_valid = eval(y, p, t, &mut self.f_y);
}
self.y_perturbed.clear();
self.y_perturbed
.extend(y.iter().zip(v).map(|(&yi, &vi)| yi + eps * vi));
self.f_yp.resize(n, 0.0);
eval(&self.y_perturbed, p, t, &mut self.f_yp);
for ((out, &fp), &f0) in jv.iter_mut().zip(&self.f_yp).zip(&self.f_y) {
*out = (fp - f0) / eps;
}
}
}
#[cfg(all(feature = "solve", any(test, feature = "xla")))]
fn same_bits(a: &[f64], b: &[f64]) -> bool {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
}
#[cfg(feature = "solve")]
#[derive(Clone, Debug)]
pub(crate) struct FieldIcRecord {
pub name: String,
pub kind: &'static str,
pub per_cell: bool,
}
#[cfg(feature = "solve")]
pub(crate) struct FieldIcMemo {
params: Vec<u64>,
forcing_generation: u64,
slots: HashMap<usize, f64>,
}
#[cfg(all(test, feature = "solve"))]
thread_local! {
static FIELD_IC_RESOLUTIONS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}
#[cfg(all(test, feature = "solve"))]
pub(crate) fn field_ic_resolutions() -> u64 {
FIELD_IC_RESOLUTIONS.with(std::cell::Cell::get)
}
#[cfg(feature = "solve")]
fn param_bits(param_vec: &[f64]) -> Vec<u64> {
param_vec.iter().map(|x| x.to_bits()).collect()
}
impl ArrayCompiled {
pub(crate) fn tape_serves_passes(&self) -> bool {
matches!(self.runtime_mode, RuntimeMode::Native | RuntimeMode::Xla)
}
pub(crate) fn is_xla(&self) -> bool {
self.runtime_mode == RuntimeMode::Xla
}
#[cfg(feature = "xla")]
pub(crate) fn install_xla_rhs(&self) -> Result<(), crate::xla_runtime::CompileRhsError> {
use crate::xla_runtime::{CompileRhsError, CompiledRhs};
if self.xla_rhs.get().is_some() {
return Ok(());
}
let program = CompiledRhs::compile(self)?;
if program.n_states() != self.n_states {
return Err(CompileRhsError::Runtime(format!(
"the emitted program returns {} state slots, the model has {}",
program.n_states(),
self.n_states
)));
}
if program.params_len() < self.param_names.len() {
return Err(CompileRhsError::Runtime(format!(
"the emitted program takes {} parameters, the model has {}",
program.params_len(),
self.param_names.len()
)));
}
let _ = self.xla_rhs.set(Rc::new(program));
Ok(())
}
#[cfg(feature = "xla")]
pub(crate) fn xla_program(&self) -> Option<Rc<crate::xla_runtime::CompiledRhs>> {
if self.is_xla() {
self.xla_rhs.get().cloned()
} else {
None
}
}
pub(crate) fn is_interpreter(&self) -> bool {
self.runtime_mode == RuntimeMode::Interpreter
}
pub(crate) fn merged_renames(&self) -> &HashMap<String, String> {
&self.merged_renames
}
pub(crate) fn unvalued_shaped_params(&self) -> &[String] {
&self.unvalued_shaped_params
}
pub(crate) fn forcing_names(&self) -> impl Iterator<Item = &String> {
self.forcing_decls.keys()
}
pub(crate) fn forcing_default(&self, name: &str) -> Option<Option<ArrayD<f64>>> {
let entry = self.forcing_defaults.get(name)?;
Some(entry.as_ref().and_then(|(shape, values)| {
ArrayD::from_shape_vec(ndarray::IxDyn(shape), values.clone()).ok()
}))
}
pub fn forcing_handle(&self) -> Rc<RefCell<HashMap<String, ArrayD<f64>>>> {
self.forcing_generation
.set(self.forcing_generation.get().wrapping_add(1));
Rc::clone(&self.forcing)
}
pub(crate) fn forcing_buffer(&self) -> Rc<RefCell<HashMap<String, ArrayD<f64>>>> {
Rc::clone(&self.forcing)
}
pub(crate) fn data_fed(&self) -> &[(String, String)] {
&self.data_fed
}
pub(crate) fn qualify_data_fed(&mut self, namespace: &str) {
for (name, _) in self.data_fed.iter_mut() {
*name = format!("{namespace}.{name}");
}
}
pub fn has_differential_equations(&self) -> bool {
!self.rhs_rules.is_empty()
}
pub(crate) fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
pub fn observed_variable_names(&self) -> Vec<String> {
self.observed_rules
.iter()
.map(|r| observed_rule_var(r).clone())
.collect()
}
pub fn state_variable_names(&self) -> &[String] {
self.state_names
.get_or_init(|| super::layout::slot_names(&self.var_shapes))
}
pub(crate) fn qualified_state_names(&self) -> &[String] {
let Some(ns) = self.namespace.as_deref() else {
return self.state_variable_names();
};
self.qualified_state_names.get_or_init(|| {
super::layout::slot_names_with(&self.var_shapes, |v| crate::problem::qualify(ns, v))
})
}
pub fn n_states(&self) -> usize {
self.n_states
}
pub fn parameter_names(&self) -> &[String] {
&self.param_names
}
#[doc(hidden)]
pub fn debug_cadence_partition(
&self,
discrete_forcing: &[String],
) -> (Vec<String>, Vec<String>, Vec<String>) {
let df: HashSet<String> = discrete_forcing.iter().cloned().collect();
let const_set = self.classify_static_observeds(&df);
let leq_discrete = self.classify_segment_invariant_observeds(&df, true);
let mut const_: Vec<String> = const_set.iter().cloned().collect();
let mut discrete: Vec<String> = leq_discrete
.iter()
.filter(|n| !const_set.contains(*n))
.cloned()
.collect();
let mut continuous: Vec<String> = self
.observed_rules
.iter()
.map(observed_rule_var)
.filter(|n| !leq_discrete.contains(*n))
.cloned()
.collect();
const_.sort();
discrete.sort();
continuous.sort();
(const_, discrete, continuous)
}
#[doc(hidden)]
pub fn debug_eval_rhs(
&self,
state: &[f64],
t: f64,
params: &HashMap<String, f64>,
force_scalar: bool,
) -> (Vec<f64>, RhsStats) {
let param_vec = self.debug_resolve_params(params);
let mut dy = vec![0.0f64; self.n_states];
let mut stats = RhsStats::default();
let mut scratch = RhsScratch::new(&self.var_shapes);
scratch.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &self.rhs_rules,
observed_rules: &self.observed_rules,
var_shapes: &self.var_shapes,
param_names: &self.param_names,
state,
params: ¶m_vec,
forcing: &self.forcing,
t,
declared: &self.declared_names,
rule_set: self.rule_set,
},
&mut dy,
force_scalar,
&mut stats,
&mut scratch,
);
(dy, stats)
}
#[doc(hidden)]
pub fn debug_new_scratch(&self) -> RhsScratch {
let mut s = RhsScratch::new(&self.var_shapes);
s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
s
}
#[doc(hidden)]
pub fn debug_new_scratch_taped(&self) -> Result<RhsScratch, SimulateError> {
let mut s = RhsScratch::new(&self.var_shapes);
s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
s.install_tape(self.strict_tape(&HashSet::new())?);
Ok(s)
}
#[doc(hidden)]
pub fn debug_precision_env(&self) -> crate::precision::Env {
self.precision.clone()
}
#[doc(hidden)]
pub fn debug_resolve_params(&self, params: &HashMap<String, f64>) -> Vec<f64> {
let mut param_vec = vec![0.0f64; self.param_names.len()];
for (i, name) in self.param_names.iter().enumerate() {
if let Some(&v) = params.get(name) {
param_vec[i] = v;
} else if let Some(d) = self.param_defaults[i] {
param_vec[i] = d;
}
}
param_vec
}
#[doc(hidden)]
pub fn debug_eval_rhs_into(
&self,
state: &[f64],
t: f64,
param_vec: &[f64],
dy: &mut [f64],
scratch: &mut RhsScratch,
stats: &mut RhsStats,
) {
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &self.rhs_rules,
observed_rules: &self.observed_rules,
var_shapes: &self.var_shapes,
param_names: &self.param_names,
state,
params: param_vec,
forcing: &self.forcing,
t,
declared: &self.declared_names,
rule_set: self.rule_set,
},
dy,
false,
stats,
scratch,
);
}
#[cfg(all(not(target_arch = "wasm32"), feature = "solve"))]
pub(crate) fn observeds_are_scalar(&self) -> bool {
self.observed_rules.iter().all(|rule| match rule {
AlgebraicRule::Scalar { body, .. } => check_scalar_evaluable(body).is_ok(),
_ => false,
})
}
#[cfg(all(not(target_arch = "wasm32"), feature = "solve"))]
pub(crate) fn evaluate_stateless_observeds(
&self,
params: &HashMap<String, f64>,
times: &[f64],
) -> Result<Vec<(String, Vec<f64>)>, SimulateError> {
let mut out: Vec<(String, Vec<f64>)> = Vec::new();
self.stateless_pass(params, times, |k, obs| {
if k == 0 {
out = self
.observed_rules
.iter()
.map(observed_rule_var)
.filter(|name| obs.get(*name).is_some_and(|a| a.len() == 1))
.map(|name| (name.clone(), Vec::with_capacity(times.len())))
.collect();
}
for (name, values) in &mut out {
let v = obs.get(name).and_then(|a| a.first().copied());
values.push(v.unwrap_or(f64::NAN));
}
})?;
Ok(out)
}
pub(crate) fn evaluate_stateless_fields(
&self,
params: &HashMap<String, f64>,
t: f64,
) -> Result<Vec<(String, ArrayD<f64>)>, SimulateError> {
let mut out: Vec<(String, ArrayD<f64>)> = Vec::new();
self.stateless_pass(params, &[t], |_, obs| {
out = self
.observed_rules
.iter()
.map(observed_rule_var)
.filter_map(|name| obs.get(name).map(|a| (name.clone(), a.clone())))
.collect();
})?;
Ok(out)
}
fn stateless_pass(
&self,
params: &HashMap<String, f64>,
times: &[f64],
mut visit: impl FnMut(usize, &ArrMap),
) -> Result<(), SimulateError> {
if self.n_states != 0 {
return Err(SimulateError::Compile(
crate::compile_error::CompileError::InterpreterBuildError {
details: format!(
"a stateless evaluation was asked of a model with {} state slot(s)",
self.n_states
),
},
));
}
let _precision_guard = self.precision.enter();
let param_vec = self.build_param_vec(params)?;
let mut scratch = RhsScratch::new(&self.var_shapes);
scratch.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
if self.tape_serves_passes() {
scratch.install_tape(self.strict_tape(&HashSet::new())?);
scratch.set_exports_active(true);
}
let mut dy: Vec<f64> = Vec::new();
crate::simulate_array::take_const_array_oob();
for (k, &t) in times.iter().enumerate() {
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &self.rhs_rules,
observed_rules: &self.observed_rules,
var_shapes: &self.var_shapes,
param_names: &self.param_names,
state: &[],
params: ¶m_vec,
forcing: &self.forcing,
t,
declared: &self.declared_names,
rule_set: self.rule_set,
},
&mut dy,
self.is_interpreter(),
&mut RhsStats::default(),
&mut scratch,
);
if let Some(details) = crate::simulate_array::take_const_array_oob() {
return Err(
crate::compile_error::CompileError::InterpreterBuildError { details }.into(),
);
}
let obs = scratch
.taped_observeds()
.unwrap_or_else(|| scratch.observed_arrays());
visit(k, obs);
}
Ok(())
}
#[cfg(feature = "solve")]
fn resolve_field_ics(
&self,
params: &HashMap<String, f64>,
mut records: Option<&mut Vec<FieldIcRecord>>,
) -> Result<HashMap<usize, f64>, SimulateError> {
let mut out: HashMap<usize, f64> = HashMap::new();
if self.field_ics.is_empty() {
return Ok(out);
}
#[cfg(test)]
FIELD_IC_RESOLUTIONS.with(|c| c.set(c.get() + 1));
let walks = crate::simulate_array::per_cell_walks;
let _overlay = OverlayGuard::armed(self.is_interpreter());
let borrowed = self.forcing.borrow();
let mut scope: Option<HashMap<String, ArrayD<f64>>> = None;
let mut scope_params: Option<HashMap<String, f64>> = None;
let mut skipped: HashMap<&str, String> = self
.ic_scope
.excluded
.iter()
.map(|(name, reason)| (name.as_str(), reason.clone()))
.collect();
if !self.ic_scope.defs.is_empty() {
let by_name: HashMap<&str, &super::compile::IcScopeDef> = self
.ic_scope
.defs
.iter()
.map(|d| (d.name.as_str(), d))
.collect();
let mut live: HashSet<&str> = HashSet::new();
let mut stack: Vec<&str> = self.ic_scope.roots.iter().map(String::as_str).collect();
while let Some(name) = stack.pop() {
if borrowed.contains_key(name) || !live.insert(name) {
continue;
}
if let Some(def) = by_name.get(name) {
stack.extend(def.deps.iter().map(String::as_str));
}
}
let mut built: HashMap<String, ArrayD<f64>> = borrowed.clone();
let mut bound: HashMap<String, f64> = params.clone();
for def in &self.ic_scope.defs {
if !live.contains(def.name.as_str()) {
continue;
}
if let Some(missing) = def
.needs
.iter()
.find(|n| !built.contains_key(*n) && !bound.contains_key(*n))
{
let why = match skipped.get(missing.as_str()) {
Some(inner) => format!("reads '{missing}', which {inner}"),
None => format!("reads '{missing}', which is not in the build-time scope"),
};
skipped.insert(&def.name, why);
continue;
}
let before = walks();
let value =
eval_buildtime_field_in_scope(&def.body, &self.index_sets, &bound, &built);
if per_cell_walk_refused() {
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: def.name.clone(),
kind: "initial-condition scope",
per_cell: true,
});
}
return Ok(out);
}
match value {
Ok(Value::Array(arr)) => {
built.insert(def.name.clone(), *arr);
}
Ok(Value::Scalar(v)) => {
bound.insert(def.name.clone(), v);
}
Err(e) => {
skipped.insert(&def.name, format!("does not evaluate at build time: {e}"));
continue;
}
}
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: def.name.clone(),
kind: "initial-condition scope",
per_cell: walks() != before,
});
}
}
scope = Some(built);
scope_params = Some(bound);
}
let forcing: &HashMap<String, ArrayD<f64>> = scope.as_ref().unwrap_or(&borrowed);
let params: &HashMap<String, f64> = scope_params.as_ref().unwrap_or(params);
for (target, rhs) in &self.field_ics {
let vs = self.var_shapes.get(target).ok_or_else(|| {
SimulateError::InvalidFieldInitialCondition {
name: target.clone(),
details: "scoped-reference target is not a state variable of the flattened \
system"
.to_string(),
}
})?;
let total = vs.shape.iter().copied().product::<usize>().max(1);
let mut cached_field: Option<Value> = None;
let before = walks();
for flat in 0..total {
let multi = flat_to_multi_col_major(flat, &vs.shape);
let slot = vs.flat_offset + flat;
let value = resolve_field_ic_cell(
target,
rhs,
&multi,
forcing,
&self.index_sets,
params,
&mut cached_field,
);
if per_cell_walk_refused() {
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: target.clone(),
kind: "initial condition",
per_cell: true,
});
}
return Ok(out);
}
out.insert(
slot,
value.map_err(|e| explain_ic_scope_miss(e, rhs, &skipped))?,
);
}
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: target.clone(),
kind: "initial condition",
per_cell: walks() != before,
});
}
}
Ok(out)
}
#[cfg(feature = "solve")]
fn seeded_defaults(&self, field_ic_map: &HashMap<usize, f64>) -> (Vec<f64>, Option<Vec<bool>>) {
let mut ic_vec = vec![0.0f64; self.n_states];
let mut unset: Option<Vec<bool>> = None;
for ((_, vs), default) in self.var_shapes.iter().zip(&self.state_defaults) {
let n = vs.shape.iter().copied().product::<usize>().max(1);
let range = vs.flat_offset..vs.flat_offset + n;
if !super::layout::write_state_default(vs, default, &mut ic_vec[range.clone()]) {
unset.get_or_insert_with(|| vec![false; self.n_states])[range].fill(true);
}
}
for (&slot, &v) in field_ic_map {
ic_vec[slot] = v;
if let Some(u) = unset.as_mut() {
u[slot] = false;
}
}
(ic_vec, unset)
}
#[cfg(feature = "solve")]
pub(crate) fn seed_initialization_faqs(
&self,
ic_vec: &mut [f64],
params: &HashMap<String, f64>,
keep: &HashSet<usize>,
mut records: Option<&mut Vec<FieldIcRecord>>,
) -> Result<Vec<usize>, SimulateError> {
let mut assigned = Vec::new();
if self.init_faqs.is_empty() {
return Ok(assigned);
}
let walks = crate::simulate_array::per_cell_walks;
let _overlay = OverlayGuard::armed(self.is_interpreter());
for (target, rhs) in &self.init_faqs {
let rule = format!("init({target})");
let invalid = |details: String| SimulateError::InvalidFieldInitialCondition {
name: rule.clone(),
details,
};
let vs = self.var_shapes.get(target).ok_or_else(|| {
invalid(format!(
"the left-hand side `{target}` of a `faq` initialization equation is not a \
state variable"
))
})?;
let mut scope: HashMap<String, ArrayD<f64>> = HashMap::new();
for (name, s) in &self.var_shapes {
let n = s.shape.iter().copied().product::<usize>().max(1);
let flat = &ic_vec[s.flat_offset..s.flat_offset + n];
scope.insert(name.clone(), col_major_to_arrayd(flat, &s.shape));
}
for (name, arr) in self.forcing.borrow().iter() {
scope.entry(name.clone()).or_insert_with(|| arr.clone());
}
let mut resolved = rhs.clone();
crate::faq::resolve_expr_ranges(&mut resolved, &self.index_sets)
.map_err(|e| invalid(e.to_string()))?;
let Expr::Operator(node) = &resolved else {
unreachable!("an initialization faq is an operator");
};
let names = node.output_idx.clone().unwrap_or_default();
let mut ranges: Vec<(i64, i64)> = Vec::with_capacity(names.len());
for n in &names {
let b = node
.ranges
.as_ref()
.and_then(|r| r.get(n))
.and_then(|s| s.bounds());
let Some([lo, hi]) = b else {
return Err(invalid(format!("output index `{n}` has no static range")));
};
ranges.push((lo, hi));
}
if ranges.len() != vs.shape.len() {
return Err(invalid(format!(
"{} output indices for a {}-D state",
ranges.len(),
vs.shape.len()
)));
}
let before = walks();
let value = eval_buildtime_field_in_scope(&resolved, &self.index_sets, params, &scope);
if per_cell_walk_refused() {
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: rule,
kind: "initialization equation",
per_cell: true,
});
}
return Ok(assigned);
}
if let Some(sink) = records.as_deref_mut() {
sink.push(FieldIcRecord {
name: rule.clone(),
kind: "initialization equation",
per_cell: walks() != before,
});
}
let arr = match value {
Ok(Value::Array(a)) if a.ndim() == ranges.len() => a,
Ok(_) => return Err(invalid("the `faq` did not evaluate to its box".to_string())),
Err(e) => return Err(invalid(e.to_string())),
};
if let Some(details) = crate::simulate_array::take_const_array_oob() {
return Err(invalid(details));
}
let total: usize = ranges
.iter()
.map(|(lo, hi)| (hi - lo + 1).max(0) as usize)
.product();
let mut cell: Vec<i64> = ranges.iter().map(|(lo, _)| *lo).collect();
for _ in 0..total {
if cell
.iter()
.zip(&vs.origin)
.zip(&vs.shape)
.any(|((&c, &o), &n)| c < o || c >= o + n as i64)
{
return Err(invalid(format!("cell {cell:?} is outside `{target}`")));
}
let slot = vs.flat_offset + multi_to_flat_col_major(&cell, &vs.shape, &vs.origin);
if !keep.contains(&slot) {
let at: Vec<usize> = cell
.iter()
.zip(&ranges)
.map(|(c, (lo, _))| (c - lo) as usize)
.collect();
ic_vec[slot] = arr[IxDyn(&at)];
assigned.push(slot);
}
let mut d = cell.len();
while d > 0 {
d -= 1;
cell[d] += 1;
if cell[d] <= ranges[d].1 {
break;
}
cell[d] = ranges[d].0;
}
}
}
Ok(assigned)
}
#[cfg(feature = "solve")]
pub(crate) fn field_ic_records(
&self,
params: &HashMap<String, f64>,
strict: bool,
) -> Vec<FieldIcRecord> {
if self.field_ics.is_empty() && self.init_faqs.is_empty() {
return Vec::new();
}
let Ok(param_vec) = self.build_param_vec(params) else {
return Vec::new();
};
let resolved: HashMap<String, f64> = self
.param_names
.iter()
.cloned()
.zip(param_vec.iter().copied())
.collect();
let _precision_guard = self.precision.enter();
let mut records = Vec::new();
let stop = strict.then(crate::simulate_array::StopAtFirstCell::arm);
let result = self.resolve_field_ics(&resolved, Some(&mut records));
let stopped = stop.as_ref().is_some_and(|s| s.stopped());
if let Ok(slots) = &result
&& !stopped
&& !self.init_faqs.is_empty()
{
let (mut ic_vec, _) = self.seeded_defaults(slots);
let _ = self.seed_initialization_faqs(
&mut ic_vec,
&resolved,
&HashSet::new(),
Some(&mut records),
);
}
if let Ok(slots) = result
&& !stopped
{
*self.field_ic_memo.borrow_mut() = Some(FieldIcMemo {
params: param_bits(¶m_vec),
forcing_generation: self.forcing_generation.get(),
slots,
});
}
records
}
fn override_namespaces(&self) -> std::collections::HashSet<String> {
crate::simulate::namespace_scope(
self.param_names
.iter()
.chain(self.var_shapes.keys())
.map(String::as_str),
self.namespace.as_deref(),
)
}
pub(crate) fn unsupplied_params(&self, params: &HashMap<String, f64>) -> Vec<String> {
let Ok(params) = crate::simulate::canonicalize_override_keys(
&self.param_index,
&self.override_namespaces(),
params,
&self.merged_renames,
) else {
return Vec::new();
};
self.param_names
.iter()
.zip(&self.param_defaults)
.filter(|(name, d)| d.is_none() && !params.contains_key(*name))
.map(|(name, _)| name.clone())
.collect()
}
#[cfg(feature = "solve")]
pub(crate) fn unset_initial_slots(
&self,
initial_conditions: &HashMap<String, f64>,
) -> Vec<String> {
let Ok(initial_conditions) = crate::simulate::canonicalize_override_keys(
&super::layout::SlotNames(&self.var_shapes),
&self.override_namespaces(),
initial_conditions,
&self.merged_renames,
) else {
return Vec::new();
};
let memo_slots;
let empty = HashMap::new();
let field_ic_map = if self.field_ics.is_empty() {
&empty
} else {
match self.field_ic_memo.borrow().as_ref() {
Some(m) => {
memo_slots = m.slots.clone();
&memo_slots
}
None => return Vec::new(),
}
};
let (_, unset) = self.seeded_defaults(field_ic_map);
let Some(mut unset) = unset else {
return Vec::new();
};
for name in initial_conditions.keys() {
if let Some(slot) = super::layout::lookup_slot(&self.var_shapes, name) {
unset[slot] = false;
}
}
for target in self
.init_faqs
.iter()
.map(|(t, _)| t)
.chain(&self.held_at_ic)
{
if let Some(vs) = self.var_shapes.get(target) {
let n = vs.shape.iter().copied().product::<usize>().max(1);
unset[vs.flat_offset..vs.flat_offset + n].fill(false);
}
}
self.var_shapes
.iter()
.filter(|(_, vs)| {
let n = vs.shape.iter().copied().product::<usize>().max(1);
unset[vs.flat_offset..vs.flat_offset + n].iter().all(|&u| u)
})
.map(|(name, _)| name.clone())
.collect()
}
fn build_param_vec(&self, params: &HashMap<String, f64>) -> Result<Vec<f64>, SimulateError> {
let params = crate::simulate::canonicalize_override_keys(
&self.param_index,
&self.override_namespaces(),
params,
&self.merged_renames,
)
.map_err(crate::simulate::param_key_error)?;
let mut param_vec = vec![0.0f64; self.param_names.len()];
for (i, name) in self.param_names.iter().enumerate() {
if let Some(&v) = params.get(name) {
param_vec[i] = v;
} else if let Some(d) = self.param_defaults[i] {
param_vec[i] = d;
} else {
return Err(SimulateError::InvalidParameter { name: name.clone() });
}
}
Ok(param_vec)
}
#[cfg(feature = "solve")]
fn build_initial_state(
&self,
initial_conditions: &HashMap<String, f64>,
param_vec: &[f64],
) -> Result<Vec<f64>, SimulateError> {
let initial_conditions = crate::simulate::canonicalize_override_keys(
&super::layout::SlotNames(&self.var_shapes),
&self.override_namespaces(),
initial_conditions,
&self.merged_renames,
)
.map_err(crate::simulate::ic_key_error)?;
let memo = self.field_ic_memo.borrow_mut().take().filter(|m| {
m.params == param_bits(param_vec)
&& m.forcing_generation == self.forcing_generation.get()
});
let resolved_params: HashMap<String, f64> = self
.param_names
.iter()
.cloned()
.zip(param_vec.iter().copied())
.collect();
let field_ic_map = match memo {
Some(m) => m.slots,
None => self.resolve_field_ics(&resolved_params, None)?,
};
let (mut ic_vec, mut unset) = self.seeded_defaults(&field_ic_map);
let mut overridden: HashSet<usize> = HashSet::new();
for (name, &v) in &initial_conditions {
let slot = super::layout::lookup_slot(&self.var_shapes, name)
.expect("a canonical initial-condition key names a slot");
ic_vec[slot] = v;
overridden.insert(slot);
if let Some(u) = unset.as_mut() {
u[slot] = false;
}
}
for slot in
self.seed_initialization_faqs(&mut ic_vec, &resolved_params, &overridden, None)?
{
if let Some(u) = unset.as_mut() {
u[slot] = false;
}
}
if let Some(slot) = unset.and_then(|u| u.iter().position(|&x| x)) {
let name = super::layout::slot_name(&self.var_shapes, slot).unwrap_or_default();
return Err(SimulateError::InvalidInitialCondition { name });
}
Ok(ic_vec)
}
#[cfg(feature = "solve")]
pub fn solve(
&self,
tspan: (f64, f64),
params: &HashMap<String, f64>,
initial_conditions: &HashMap<String, f64>,
opts: &SolveOptions,
) -> Result<Solution, SimulateError> {
self.solve_inspect(tspan, params, initial_conditions, opts, None)
}
#[cfg(feature = "solve")]
pub fn solve_inspect(
&self,
tspan: (f64, f64),
params: &HashMap<String, f64>,
initial_conditions: &HashMap<String, f64>,
opts: &SolveOptions,
inspect: Option<&mut BuildInspection>,
) -> Result<Solution, SimulateError> {
let _precision_guard = self.precision.enter();
self.solve_core(
tspan,
params,
initial_conditions,
opts,
inspect,
&HashSet::new(),
&[],
|_t| Ok(()),
)
}
#[allow(clippy::too_many_arguments)]
#[cfg(all(feature = "solve", not(target_arch = "wasm32")))]
pub(crate) fn solve_with_refresh_inspect(
&self,
tspan: (f64, f64),
params: &HashMap<String, f64>,
initial_conditions: &HashMap<String, f64>,
opts: &SolveOptions,
inspect: Option<&mut BuildInspection>,
discrete_forcing: &HashSet<String>,
boundaries: &[f64],
refresh_fn: impl FnMut(f64) -> Result<(), SimulateError>,
) -> Result<Solution, SimulateError> {
self.solve_core(
tspan,
params,
initial_conditions,
opts,
inspect,
discrete_forcing,
boundaries,
refresh_fn,
)
}
#[allow(clippy::too_many_arguments)]
#[cfg(feature = "solve")]
fn solve_core(
&self,
tspan: (f64, f64),
params: &HashMap<String, f64>,
initial_conditions: &HashMap<String, f64>,
opts: &SolveOptions,
inspect: Option<&mut BuildInspection>,
discrete_forcing: &HashSet<String>,
boundaries: &[f64],
mut refresh_fn: impl FnMut(f64) -> Result<(), SimulateError>,
) -> Result<Solution, SimulateError> {
let params_owned = params;
let ics_owned = initial_conditions;
let (t0, t_end) = tspan;
crate::simulate::reject_nonfinite_span(t0, t_end)?;
crate::simulate_array::take_const_array_oob();
let param_vec = self.build_param_vec(params_owned)?;
let ic_vec = self.build_initial_state(ics_owned, ¶m_vec)?;
refresh_fn(t0)?;
let tape = self.build_solve_tape(discrete_forcing)?;
let cadence = self.partition_observed_cadence(discrete_forcing);
let setup = self.hoist_static_observeds(cadence, &ic_vec, ¶m_vec, t0, tape.as_ref());
let scratches = match &tape {
Some(prog) if self.tape_serves_passes() => Some(SolveScratches::new(self, prog)),
_ => None,
};
if let Some(insp) = inspect {
self.fill_solve_inspection(
insp,
&setup,
&ic_vec,
¶m_vec,
t0,
boundaries,
tape.as_ref(),
);
}
let solver_name = match opts.alg {
Alg::Bdf => "Bdf",
Alg::Sdirk => "Sdirk",
Alg::Erk => "Erk",
};
if boundaries.is_empty() || opts.saveat.is_none() {
let (time, state, stats, retcode) = self
.run_one_segment(
t0,
t_end,
&ic_vec,
¶m_vec,
&setup.static_obs,
&setup.cadence.segment_static_rules,
&setup.cadence.continuous_rules,
opts,
tape.as_ref(),
scratches.as_ref(),
)
.map_err(Self::const_oob_first)?;
return self.assemble_solution(
time,
state,
retcode,
solution_metadata(
solver_name,
&stats,
self.merged_renames.clone(),
self.namespace.clone(),
),
¶m_vec,
&setup,
&opts.output_observed,
tape.as_ref(),
);
}
let run = SegmentedRun {
t0,
t_end,
boundaries,
ic_vec: &ic_vec,
param_vec: ¶m_vec,
setup: &setup,
opts,
tape: tape.as_ref(),
scratches: scratches.as_ref(),
};
let (time, state, stats, retcode) = self.run_segmented(&run, &mut refresh_fn)?;
self.assemble_solution(
time,
state,
retcode,
solution_metadata(
solver_name,
&stats,
self.merged_renames.clone(),
self.namespace.clone(),
),
¶m_vec,
&setup,
&opts.output_observed,
tape.as_ref(),
)
}
#[cfg(feature = "solve")]
fn partition_observed_cadence(&self, discrete_forcing: &HashSet<String>) -> ObservedCadence {
let static_names = self.classify_static_observeds(discrete_forcing);
let seg_invariant_names = self.classify_segment_invariant_observeds(discrete_forcing, true);
let static_rules: Vec<AlgebraicRule> = self
.observed_rules
.iter()
.filter(|r| static_names.contains(observed_rule_var(r)))
.cloned()
.collect();
let segment_static_rules: Vec<AlgebraicRule> = self
.observed_rules
.iter()
.filter(|r| {
seg_invariant_names.contains(observed_rule_var(r))
&& !static_names.contains(observed_rule_var(r))
})
.cloned()
.collect();
let continuous_rules: Vec<AlgebraicRule> = self
.observed_rules
.iter()
.filter(|r| !seg_invariant_names.contains(observed_rule_var(r)))
.cloned()
.collect();
let varying_rules: Vec<AlgebraicRule> = self
.observed_rules
.iter()
.filter(|r| !static_names.contains(observed_rule_var(r)))
.cloned()
.collect();
ObservedCadence {
static_names,
static_rules,
segment_static_rules,
continuous_rules,
varying_rules,
}
}
#[cfg(feature = "solve")]
fn hoist_static_observeds(
&self,
cadence: ObservedCadence,
ic_vec: &[f64],
param_vec: &[f64],
t0: f64,
tape: Option<&std::sync::Arc<TapeProgram>>,
) -> SolveSetup {
let sa0 = build_state_arrays(&self.var_shapes, ic_vec);
if self.tape_serves_passes() && tape.is_some() {
return SolveSetup {
cadence,
sa0,
static_obs: ArrMap::default(),
};
}
let static_rings_cell: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
let mut static_obs = ArrMap::default();
let env = EvalEnv {
state_arrays: &sa0,
params: param_vec,
param_names: &self.param_names,
t: t0,
derived_rings: &static_rings_cell,
derived_extents: empty_derived_extents(),
forcing: &self.forcing,
cse: None,
const_lits: None,
const_arrays: &self.const_scope,
declared: &self.declared_names,
};
materialize_observeds_pass(
&mut static_obs,
&cadence.static_rules,
&ObsPass {
env,
force_scalar: self.is_interpreter(),
},
&mut RhsStats::default(),
);
drop(static_rings_cell);
SolveSetup {
cadence,
sa0,
static_obs,
}
}
#[cfg(feature = "solve")]
#[allow(clippy::too_many_arguments)]
fn fill_solve_inspection(
&self,
insp: &mut BuildInspection,
setup: &SolveSetup,
ic_vec: &[f64],
param_vec: &[f64],
t0: f64,
boundaries: &[f64],
tape: Option<&std::sync::Arc<TapeProgram>>,
) {
if self.tape_serves_passes()
&& let Some(tape) = tape
{
let mut harvest = TapedObserveds::new(self, tape);
let obs = harvest.at(self, ic_vec, param_vec, t0);
self.fill_inspection(insp, obs, &setup.cadence.static_names, param_vec);
if !boundaries.is_empty() {
for rule in &setup.cadence.varying_rules {
let name = observed_rule_var(rule);
if let Some(a) = obs.get(name) {
insp.setup_arrays.insert(name.clone(), a.clone());
}
}
}
return;
}
self.fill_inspection(
insp,
&setup.static_obs,
&setup.cadence.static_names,
param_vec,
);
if !boundaries.is_empty() {
let dr: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
let mut snapshot = setup.static_obs.clone();
materialize_observeds_pass(
&mut snapshot,
&setup.cadence.varying_rules,
&ObsPass {
env: EvalEnv {
state_arrays: &setup.sa0,
params: param_vec,
param_names: &self.param_names,
t: t0,
derived_rings: &dr,
derived_extents: empty_derived_extents(),
forcing: &self.forcing,
cse: None,
const_lits: None,
const_arrays: &self.const_scope,
declared: &self.declared_names,
},
force_scalar: self.is_interpreter(),
},
&mut RhsStats::default(),
);
for rule in &setup.cadence.varying_rules {
let name = observed_rule_var(rule);
if let Some(a) = snapshot.get(name) {
insp.setup_arrays.insert(name.clone(), a.clone());
}
}
}
}
#[cfg(feature = "solve")]
fn build_solve_tape(
&self,
discrete_forcing: &HashSet<String>,
) -> Result<SolveTape, SimulateError> {
if self.is_interpreter() {
return Ok(None);
}
Ok(Some(self.strict_tape(discrete_forcing)?))
}
#[cfg(feature = "solve")]
fn run_segmented(
&self,
run: &SegmentedRun<'_>,
refresh_fn: &mut impl FnMut(f64) -> Result<(), SimulateError>,
) -> Result<(Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode), SimulateError> {
let &SegmentedRun {
t0,
t_end,
boundaries,
ic_vec,
param_vec,
setup,
opts,
tape,
scratches,
} = run;
let n_states = self.n_states;
let mut endpoints: Vec<f64> = vec![t0];
for &b in boundaries {
if b > t0 && b < t_end && *endpoints.last().unwrap() < b {
endpoints.push(b);
}
}
if *endpoints.last().unwrap() < t_end {
endpoints.push(t_end);
}
let global_out = opts.saveat.clone().expect("output grid checked Some");
let mut u0 = ic_vec.to_vec();
let mut time: Vec<f64> = Vec::new();
let mut state: Vec<Vec<f64>> = vec![Vec::new(); n_states];
let mut stats = SolveStats::default();
let mut retcode = ReturnCode::Success;
for w in endpoints.windows(2) {
let (a, b) = (w[0], w[1]);
if a != t0 {
refresh_fn(a)?;
if let Some(s) = scratches {
s.forcing_refreshed();
}
}
let requested: Vec<f64> = global_out
.iter()
.copied()
.filter(|&g| (if a == t0 { g >= a } else { g > a }) && g <= b)
.collect();
let mut grid = requested.clone();
if grid.last() != Some(&b) {
grid.push(b);
}
let seg_progress: Option<ProgressFn> = opts.progress.as_ref().map(|user| {
let user = user.clone();
let steps_before = stats.n_accepted_steps;
let wrapped: ProgressFn = std::sync::Arc::new(move |p: &Progress<'_>| {
user(&Progress {
t0,
t: p.t,
t_end,
step: steps_before + p.step,
maxiters: p.maxiters,
u: p.u,
})
});
wrapped
});
let seg_opts = SolveOptions {
saveat: Some(grid),
progress: seg_progress,
..opts.clone()
};
let (seg_time, seg_state, seg_stats, seg_retcode) = self
.run_one_segment(
a,
b,
&u0,
param_vec,
&setup.static_obs,
&setup.cadence.segment_static_rules,
&setup.cadence.continuous_rules,
&seg_opts,
tape,
scratches,
)
.map_err(Self::const_oob_first)?;
stats += seg_stats;
if !seg_retcode.is_success() {
retcode = seg_retcode;
}
for (i, &t) in seg_time.iter().enumerate() {
if t == b {
u0 = (0..n_states).map(|r| seg_state[r][i]).collect();
}
if requested.contains(&t) {
time.push(t);
for r in 0..n_states {
state[r].push(seg_state[r][i]);
}
}
}
if !retcode.is_success() {
break;
}
}
Ok((time, state, stats, retcode))
}
#[cfg(feature = "solve")]
#[allow(clippy::too_many_arguments)]
fn assemble_solution(
&self,
time: Vec<f64>,
mut state: Vec<Vec<f64>>,
retcode: ReturnCode,
metadata: SolutionMetadata,
param_vec: &[f64],
setup: &SolveSetup,
output_observed: &[String],
tape: Option<&std::sync::Arc<TapeProgram>>,
) -> Result<Solution, SimulateError> {
let mut state_variable_names = self.state_variable_names().to_vec();
self.append_observed_trajectories(
&time,
&mut state,
&mut state_variable_names,
param_vec,
&setup.static_obs,
&setup.cadence.static_names,
&setup.cadence.varying_rules,
output_observed,
tape,
);
if let Some(details) = crate::simulate_array::take_const_array_oob() {
return Err(
crate::compile_error::CompileError::InterpreterBuildError { details }.into(),
);
}
Ok(Solution {
time,
state,
state_variable_names,
retcode,
metadata,
})
}
#[cfg(feature = "solve")]
fn const_oob_first(err: SimulateError) -> SimulateError {
match crate::simulate_array::take_const_array_oob() {
Some(details) => {
crate::compile_error::CompileError::InterpreterBuildError { details }.into()
}
None => err,
}
}
#[allow(clippy::too_many_arguments)]
#[cfg(feature = "solve")]
fn run_one_segment(
&self,
t0: f64,
t_end: f64,
u0: &[f64],
param_vec: &[f64],
static_obs: &ArrMap,
segment_static_rules: &[AlgebraicRule],
continuous_rules: &[AlgebraicRule],
opts: &SolveOptions,
tape: Option<&std::sync::Arc<TapeProgram>>,
scratches: Option<&SolveScratches>,
) -> Result<(Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode), SimulateError> {
if let Some((time, state, retcode)) =
crate::simulate::nonadvancing_trajectory(t0, t_end, u0, opts)
{
return Ok((time, state, SolveStats::default(), retcode));
}
let n_states = self.n_states;
let rhs_rules = self.rhs_rules.clone();
let var_shapes = self.var_shapes.clone();
let param_names = self.param_names.clone();
let rhs_rules_jac = rhs_rules.clone();
let varying_rules_rhs = continuous_rules.to_vec();
let varying_rules_jac = continuous_rules.to_vec();
let var_shapes_jac = var_shapes.clone();
let param_names_jac = param_names.clone();
let declared = self.declared_names.clone();
let declared_jac = declared.clone();
let rule_set_rhs = next_rule_set_id();
let rule_set_jac = next_rule_set_id();
let seg_seed: ArrMap = if segment_static_rules.is_empty() || self.tape_serves_passes() {
static_obs.clone()
} else {
let sa_seg = build_state_arrays(&self.var_shapes, u0);
let seg_rings: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
let mut seed = static_obs.clone();
materialize_observeds_pass(
&mut seed,
segment_static_rules,
&ObsPass {
env: EvalEnv {
state_arrays: &sa_seg,
params: param_vec,
param_names: &self.param_names,
t: t0,
derived_rings: &seg_rings,
derived_extents: empty_derived_extents(),
forcing: &self.forcing,
cse: None,
const_lits: None,
const_arrays: &self.const_scope,
declared: &self.declared_names,
},
force_scalar: false,
},
&mut RhsStats::default(),
);
seed
};
#[cfg(feature = "xla")]
let xla_rhs = match (self.is_xla(), self.xla_program()) {
(false, _) => None,
(true, Some(program)) => Some(program),
(true, None) => return Err(xla_not_installed()),
};
#[cfg(not(feature = "xla"))]
if self.is_xla() {
return Err(xla_not_installed());
}
#[cfg(feature = "xla")]
let tape_rhs = xla_rhs.is_none();
#[cfg(not(feature = "xla"))]
let tape_rhs = true;
let seg_seed = Rc::new(seg_seed);
let (rhs_scratch, jac_scratch) = match scratches {
Some(s) => (Rc::clone(&s.rhs), Rc::clone(&s.jac)),
None => (Rc::new(RefCell::new(None)), Rc::new(RefCell::new(None))),
};
let program_rhs = scratches.and_then(|s| s.program.clone());
let program_jac = program_rhs.clone();
if tape_rhs && program_rhs.is_none() && rhs_scratch.borrow().is_none() {
let mut s = RhsScratch::new(&var_shapes);
s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
s.set_static((*seg_seed).clone());
if let Some(prog) = tape {
s.install_tape(std::sync::Arc::clone(prog));
}
*rhs_scratch.borrow_mut() = Some(s);
}
let jac_seed = Rc::clone(&seg_seed);
let const_scope_jac = std::sync::Arc::clone(&self.const_scope);
let tape_jac: Option<std::sync::Arc<TapeProgram>> = match (self.tape_serves_passes(), tape)
{
(true, Some(prog)) => Some(std::sync::Arc::clone(prog)),
_ => None,
};
let force_scalar = self.is_interpreter();
#[cfg(feature = "xla")]
let xla_jac: Option<(Rc<crate::xla_runtime::CompiledRhs>, RefCell<Option<FdJvp>>)> =
xla_rhs.clone().map(|program| (program, RefCell::new(None)));
#[cfg(feature = "xla")]
let xla_fault: Rc<RefCell<Option<String>>> = Rc::new(RefCell::new(None));
#[cfg(feature = "xla")]
let xla_fault_rhs = Rc::clone(&xla_fault);
#[cfg(feature = "xla")]
let xla_fault_jac = Rc::clone(&xla_fault);
let forcing_rhs = Rc::clone(&self.forcing);
let forcing_jac = Rc::clone(&self.forcing);
let rhs_closure = move |y: &diffsol::FaerVec<f64>,
p: &diffsol::FaerVec<f64>,
t: f64,
dy: &mut diffsol::FaerVec<f64>| {
let y_s = y.as_slice();
let p_s = p.as_slice();
let dy_s = dy.as_mut_slice();
#[cfg(feature = "xla")]
if let Some(program) = &xla_rhs {
dy_s.fill(0.0);
run_xla_rhs(program, &xla_fault_rhs, y_s, p_s, t, dy_s);
return;
}
if let Some(sp) = &program_rhs {
let mut ws = sp.rhs.borrow_mut();
let ws = ws.get_or_insert_with(|| sp.program.solve_workspace());
sp.program
.eval_solve(y_s, p_s, t, &forcing_rhs.borrow(), dy_s, ws);
return;
}
let mut scratch_slot = rhs_scratch.borrow_mut();
let scratch = scratch_slot
.as_mut()
.expect("the tape scratch is built whenever no XLA program serves the rhs");
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &rhs_rules,
observed_rules: &varying_rules_rhs,
var_shapes: &var_shapes,
param_names: ¶m_names,
state: y_s,
params: p_s,
forcing: &forcing_rhs,
t,
declared: &declared,
rule_set: rule_set_rhs,
},
dy_s,
force_scalar,
&mut RhsStats::default(),
scratch,
);
};
let jac_closure = move |y: &diffsol::FaerVec<f64>,
p: &diffsol::FaerVec<f64>,
t: f64,
v: &diffsol::FaerVec<f64>,
jv: &mut diffsol::FaerVec<f64>| {
#[cfg(feature = "xla")]
if let Some((program, jvp)) = &xla_jac {
let mut jvp = jvp.borrow_mut();
let jvp = jvp.get_or_insert_with(|| FdJvp::new(n_states));
jvp.apply(
y.as_slice(),
p.as_slice(),
t,
v.as_slice(),
jv.as_mut_slice(),
|state, params, t, out| {
run_xla_rhs(program, &xla_fault_jac, state, params, t, out)
},
);
return;
}
let n = y.as_slice().len();
let v_s = v.as_slice();
let p_s = p.as_slice();
let y_s = y.as_slice();
let mut y_norm = 0.0f64;
for &yi in y_s {
y_norm += yi * yi;
}
let y_norm = y_norm.sqrt().max(1.0);
let eps = f64::EPSILON.sqrt() * y_norm;
let mut y_perturbed = vec![0.0f64; n];
for i in 0..n {
y_perturbed[i] = y_s[i] + eps * v_s[i];
}
let mut f_y = vec![0.0f64; n];
let mut f_yp = vec![0.0f64; n];
if let Some(sp) = &program_jac {
let mut ws = sp.jac.borrow_mut();
let ws = ws.get_or_insert_with(|| sp.program.solve_workspace());
let forcing = forcing_jac.borrow();
sp.program.eval_solve(y_s, p_s, t, &forcing, &mut f_y, ws);
sp.program
.eval_solve(&y_perturbed, p_s, t, &forcing, &mut f_yp, ws);
let jv_s = jv.as_mut_slice();
for i in 0..n {
jv_s[i] = (f_yp[i] - f_y[i]) / eps;
}
return;
}
let mut scratch_slot = jac_scratch.borrow_mut();
let scratch = scratch_slot.get_or_insert_with(|| {
let mut s = RhsScratch::new(&var_shapes_jac);
s.set_const_arrays(std::sync::Arc::clone(&const_scope_jac));
s.set_static((*jac_seed).clone());
if let Some(prog) = &tape_jac {
s.install_tape(std::sync::Arc::clone(prog));
}
s
});
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &rhs_rules_jac,
observed_rules: &varying_rules_jac,
var_shapes: &var_shapes_jac,
param_names: ¶m_names_jac,
state: y_s,
params: p_s,
forcing: &forcing_jac,
t,
declared: &declared_jac,
rule_set: rule_set_jac,
},
&mut f_y,
force_scalar,
&mut RhsStats::default(),
scratch,
);
evaluate_rhs_with_scratch(
&RhsCall {
rhs_rules: &rhs_rules_jac,
observed_rules: &varying_rules_jac,
var_shapes: &var_shapes_jac,
param_names: ¶m_names_jac,
state: &y_perturbed,
params: p_s,
forcing: &forcing_jac,
t,
declared: &declared_jac,
rule_set: rule_set_jac,
},
&mut f_yp,
force_scalar,
&mut RhsStats::default(),
scratch,
);
let jv_s = jv.as_mut_slice();
for i in 0..n {
jv_s[i] = (f_yp[i] - f_y[i]) / eps;
}
};
let abstol = opts.abstol_or_default();
let reltol = opts.reltol_or_default();
let ic_for_init = u0.to_vec();
let builder = OdeBuilder::<FaerMat<f64>>::new()
.t0(t0)
.rtol(reltol)
.atol(vec![abstol; n_states])
.p(param_vec.to_vec())
.rhs_implicit(rhs_closure, jac_closure)
.init(
move |_p: &diffsol::FaerVec<f64>, _t: f64, y: &mut diffsol::FaerVec<f64>| {
let y_s = y.as_mut_slice();
for (i, &v) in ic_for_init.iter().enumerate() {
y_s[i] = v;
}
},
n_states,
);
let problem = builder.build().map_err(|e| SimulateError::DiffsolError {
details: e.to_string(),
})?;
type SolvedSegment = (Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode);
let solved = (|| -> Result<SolvedSegment, SimulateError> {
let out = match opts.alg {
Alg::Bdf => {
let mut solver: Bdf<'_, _, NewtonNonlinearSolver<_, FaerLU<f64>, _>> = problem
.bdf::<FaerLU<f64>>()
.map_err(|e| SimulateError::DiffsolError {
details: e.to_string(),
})?;
let (time, state, retcode) =
crate::simulate::run_solver(&mut solver, t_end, opts)?;
let bs = solver.get_statistics();
let stats = SolveStats::from_solver(
&solver,
bs.number_of_steps,
bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
);
(time, state, stats, retcode)
}
Alg::Sdirk => {
let mut solver: Sdirk<'_, _, FaerLU<f64>> = problem
.tr_bdf2::<FaerLU<f64>>()
.map_err(|e| SimulateError::DiffsolError {
details: e.to_string(),
})?;
let (time, state, retcode) =
crate::simulate::run_solver(&mut solver, t_end, opts)?;
let bs = solver.get_statistics();
let stats = SolveStats::from_solver(
&solver,
bs.number_of_steps,
bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
);
(time, state, stats, retcode)
}
Alg::Erk => {
let _seq = crate::simulate::SequentialFaer::enter();
let mut solver = problem.tsit45().map_err(|e| SimulateError::DiffsolError {
details: e.to_string(),
})?;
let (time, state, retcode) =
crate::simulate::run_solver(&mut solver, t_end, opts)?;
let bs = solver.get_statistics();
let stats = SolveStats::from_solver(
&solver,
bs.number_of_steps,
bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
);
(time, state, stats, retcode)
}
};
Ok(out)
})();
#[cfg(feature = "xla")]
if let Some(details) = xla_fault.borrow().clone() {
return Err(SimulateError::CompilerUnavailable {
compiler: "xla",
details: format!("the compiled right-hand side failed during the solve: {details}"),
});
}
let (time, state, stats, retcode) = solved?;
Ok((time, state, stats, retcode))
}
#[cfg(feature = "solve")]
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
fn append_observed_trajectories(
&self,
time: &[f64],
state: &mut Vec<Vec<f64>>,
state_variable_names: &mut Vec<String>,
param_vec: &[f64],
static_obs: &ArrMap,
static_names: &HashSet<String>,
varying_rules: &[AlgebraicRule],
requested: &[String],
tape: Option<&std::sync::Arc<TapeProgram>>,
) {
if self.observed_rules.is_empty() || time.is_empty() {
return;
}
if self.tape_serves_passes()
&& let Some(tape) = tape
{
self.append_observed_trajectories_taped(
time,
state,
state_variable_names,
param_vec,
static_names,
varying_rules,
requested,
tape,
);
return;
}
let wanted = self.resolve_requested_observeds(requested);
let nt = time.len();
let mut flat = vec![0.0f64; self.n_states];
for (i, slot) in flat.iter_mut().enumerate() {
*slot = state[i][0];
}
let mut sa = build_state_arrays(&self.var_shapes, &flat);
let static_keys: HashSet<String> = static_obs.keys().cloned().collect();
let mut obs: ArrMap = ArrMap::default();
for (name, arr) in static_obs {
obs.insert(name.clone(), arr.clone());
}
let dr: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
let cse = CseRt::default();
let bind_cse = |rules: &[AlgebraicRule]| {
cse.retarget(
(rules.as_ptr() as u64).rotate_left(17)
^ ((rules.len() as u64) << 16)
^ 0x0b5e_0b5e_0b5e_0b5e,
);
};
let prune = varying_rules
.iter()
.all(|r| !expr_blocks_output_pruning(observed_rule_body(r)));
let probe_owned: Vec<AlgebraicRule>;
let probe_rules: &[AlgebraicRule] = if prune {
let unknown: HashSet<String> = self
.observed_rules
.iter()
.filter(|r| {
!observed_rule_is_array_valued(r) || wanted.contains(observed_rule_var(r))
})
.map(|r| observed_rule_var(r).clone())
.collect();
match dependency_cone(varying_rules, &unknown) {
None => varying_rules,
Some(cone) => {
probe_owned = cone;
&probe_owned
}
}
} else {
varying_rules
};
if !probe_rules.is_empty() {
bind_cse(probe_rules);
materialize_observeds_pass(
&mut obs,
probe_rules,
&ObsPass {
env: EvalEnv {
state_arrays: &sa,
params: param_vec,
param_names: &self.param_names,
t: time[0],
derived_rings: &dr,
derived_extents: empty_derived_extents(),
forcing: &self.forcing,
cse: Some(&cse),
const_lits: None,
const_arrays: &self.const_scope,
declared: &self.declared_names,
},
force_scalar: self.is_interpreter(),
},
&mut RhsStats::default(),
);
}
let emit: Vec<ObservedRows> = self
.observed_rules
.iter()
.filter_map(|rule| {
let name = observed_rule_var(rule);
let arr = obs.get(name)?;
if arr.ndim() == 0 {
Some(ObservedRows::scalar(name.clone()))
} else if wanted.contains(name) {
Some(ObservedRows::gridded(name.clone(), arr.shape().to_vec()))
} else {
None
}
})
.collect();
if emit.is_empty() {
return;
}
let n_rows: usize = emit.iter().map(ObservedRows::n_rows).sum();
let value_owned: Vec<AlgebraicRule>;
let value_rules: &[AlgebraicRule] = if prune {
let seeds: HashSet<String> = emit.iter().map(|e| e.name.clone()).collect();
match dependency_cone(varying_rules, &seeds) {
None => varying_rules,
Some(cone) => {
value_owned = cone;
&value_owned
}
}
} else {
varying_rules
};
let mut rows: Vec<Vec<f64>> = vec![Vec::with_capacity(nt); n_rows];
let record = |obs: &ArrMap, rows: &mut Vec<Vec<f64>>| {
let mut j = 0usize;
for e in &emit {
match obs.get(&e.name) {
Some(a) if e.shape.is_empty() => {
rows[j].push(a.first().copied().unwrap_or(f64::NAN));
j += 1;
}
Some(a) if a.shape() == e.shape.as_slice() => {
for v in arrayd_to_col_major(a) {
rows[j].push(v);
j += 1;
}
}
_ => {
for _ in 0..e.n_rows() {
rows[j].push(f64::NAN);
j += 1;
}
}
}
}
};
record(&obs, &mut rows);
if value_rules.is_empty() {
for row in rows.iter_mut() {
let v = row[0];
row.resize(nt, v);
}
} else {
bind_cse(value_rules);
for k in 1..nt {
for (i, slot) in flat.iter_mut().enumerate() {
*slot = state[i][k];
}
refill_state_arrays(&mut sa, &self.var_shapes, &flat);
obs.retain(|name, _| static_keys.contains(name));
dr.borrow_mut().clear();
materialize_observeds_pass(
&mut obs,
value_rules,
&ObsPass {
env: EvalEnv {
state_arrays: &sa,
params: param_vec,
param_names: &self.param_names,
t: time[k],
derived_rings: &dr,
derived_extents: empty_derived_extents(),
forcing: &self.forcing,
cse: Some(&cse),
const_lits: None,
const_arrays: &self.const_scope,
declared: &self.declared_names,
},
force_scalar: self.is_interpreter(),
},
&mut RhsStats::default(),
);
record(&obs, &mut rows);
}
}
let mut rows = rows.into_iter();
for e in &emit {
for name in e.row_names() {
let Some(row) = rows.next() else { break };
state_variable_names.push(name);
state.push(row);
}
}
}
#[cfg(feature = "solve")]
#[allow(clippy::too_many_arguments)]
fn append_observed_trajectories_taped(
&self,
time: &[f64],
state: &mut Vec<Vec<f64>>,
state_variable_names: &mut Vec<String>,
param_vec: &[f64],
static_names: &HashSet<String>,
varying_rules: &[AlgebraicRule],
requested: &[String],
tape: &std::sync::Arc<TapeProgram>,
) {
let wanted = self.resolve_requested_observeds(requested);
let emitted = self.output_row_candidates(static_names, varying_rules, &wanted);
let nt = time.len();
let mut harvest = TapedObserveds::new(self, tape);
let mut flat = vec![0.0f64; self.n_states];
for (i, slot) in flat.iter_mut().enumerate() {
*slot = state[i][0];
}
let emit: Vec<ObservedRows> = {
let obs = harvest.at(self, &flat, param_vec, time[0]);
self.observed_rules
.iter()
.filter_map(|rule| {
let name = observed_rule_var(rule);
if !emitted.contains(name) {
return None;
}
let arr = obs.get(name)?;
if arr.ndim() == 0 {
Some(ObservedRows::scalar(name.clone()))
} else if wanted.contains(name) {
Some(ObservedRows::gridded(name.clone(), arr.shape().to_vec()))
} else {
None
}
})
.collect()
};
if emit.is_empty() {
return;
}
let n_rows: usize = emit.iter().map(ObservedRows::n_rows).sum();
let mut rows: Vec<Vec<f64>> = vec![Vec::with_capacity(nt); n_rows];
let record = |obs: &ArrMap, rows: &mut Vec<Vec<f64>>| {
let mut j = 0usize;
for e in &emit {
match obs.get(&e.name) {
Some(a) if e.shape.is_empty() => {
rows[j].push(a.first().copied().unwrap_or(f64::NAN));
j += 1;
}
Some(a) if a.shape() == e.shape.as_slice() => {
for v in arrayd_to_col_major(a) {
rows[j].push(v);
j += 1;
}
}
_ => {
for _ in 0..e.n_rows() {
rows[j].push(f64::NAN);
j += 1;
}
}
}
}
};
record(harvest.scratch_observeds(), &mut rows);
for k in 1..nt {
for (i, slot) in flat.iter_mut().enumerate() {
*slot = state[i][k];
}
let obs = harvest.at(self, &flat, param_vec, time[k]);
record(obs, &mut rows);
}
let mut rows = rows.into_iter();
for e in &emit {
for name in e.row_names() {
let Some(row) = rows.next() else { break };
state_variable_names.push(name);
state.push(row);
}
}
}
#[cfg(feature = "solve")]
fn output_row_candidates(
&self,
static_names: &HashSet<String>,
varying_rules: &[AlgebraicRule],
wanted: &HashSet<String>,
) -> HashSet<String> {
let mut out: HashSet<String> = static_names.clone();
let prune = varying_rules
.iter()
.all(|r| !expr_blocks_output_pruning(observed_rule_body(r)));
let cone: Option<Vec<AlgebraicRule>> = if prune {
let unknown: HashSet<String> = self
.observed_rules
.iter()
.filter(|r| {
!observed_rule_is_array_valued(r) || wanted.contains(observed_rule_var(r))
})
.map(|r| observed_rule_var(r).clone())
.collect();
dependency_cone(varying_rules, &unknown)
} else {
None
};
match cone {
Some(rules) => out.extend(rules.iter().map(|r| observed_rule_var(r).clone())),
None => out.extend(varying_rules.iter().map(|r| observed_rule_var(r).clone())),
}
out
}
#[cfg(feature = "solve")]
fn resolve_requested_observeds(&self, requested: &[String]) -> HashSet<String> {
if requested.is_empty() {
return HashSet::new();
}
let vars: Vec<&str> = self
.observed_rules
.iter()
.map(|r| observed_rule_var(r).as_str())
.collect();
requested
.iter()
.filter_map(
|r| match crate::data_output::match_output_request(r, &vars) {
crate::data_output::RequestMatch::Named(var) => Some(var.to_string()),
_ => None,
},
)
.collect()
}
pub(super) fn classify_static_observeds(
&self,
discrete_forcing: &HashSet<String>,
) -> HashSet<String> {
self.classify_segment_invariant_observeds(discrete_forcing, false)
}
pub(super) fn classify_segment_invariant_observeds(
&self,
discrete_forcing: &HashSet<String>,
allow_forcing: bool,
) -> HashSet<String> {
let observed_names: HashSet<&String> =
self.observed_rules.iter().map(observed_rule_var).collect();
let mut set: HashSet<String> = HashSet::new();
for rule in &self.observed_rules {
let mut refs = HashSet::new();
collect_expr_var_refs(observed_rule_body(rule), &mut refs);
if matches!(rule, AlgebraicRule::Recurrence { .. }) {
let self_name = observed_rule_var(rule);
refs.retain(|r| r != self_name);
}
let ok = refs.iter().all(|r| {
r != "t"
&& !self.var_shapes.contains_key(r)
&& (allow_forcing || !discrete_forcing.contains(r))
&& (!observed_names.contains(r) || set.contains(r))
});
if ok {
set.insert(observed_rule_var(rule).clone());
}
}
set
}
#[cfg(feature = "solve")]
fn fill_inspection(
&self,
insp: &mut BuildInspection,
static_obs: &ArrMap,
static_names: &HashSet<String>,
param_vec: &[f64],
) {
for (i, name) in self.param_names.iter().enumerate() {
insp.params.insert(name.clone(), param_vec[i]);
}
for rule in &self.observed_rules {
insp.observed_exprs.insert(
observed_rule_var(rule).clone(),
observed_rule_body(rule).clone(),
);
if let AlgebraicRule::Recurrence {
var,
output_idx_names,
axis,
max_lag,
lag_proven,
..
} = rule
{
insp.recurrences
.push(crate::simulate_array::RecurrenceInfo {
var: var.clone(),
axis: output_idx_names[*axis].clone(),
max_lag: *max_lag,
lag_proven: *lag_proven,
});
}
}
for name in static_names {
if let Some(a) = static_obs.get(name) {
insp.setup_arrays.insert(name.clone(), a.clone());
}
}
}
}
#[cfg(feature = "solve")]
struct ObservedRows {
name: String,
shape: Vec<usize>,
}
#[cfg(feature = "solve")]
impl ObservedRows {
fn scalar(name: String) -> Self {
ObservedRows {
name,
shape: Vec::new(),
}
}
fn gridded(name: String, shape: Vec<usize>) -> Self {
ObservedRows { name, shape }
}
fn n_rows(&self) -> usize {
self.shape.iter().product::<usize>().max(1)
}
fn row_names(&self) -> Vec<String> {
if self.shape.is_empty() {
return vec![self.name.clone()];
}
(0..self.n_rows())
.map(|flat| {
let idx = flat_to_multi_col_major(flat, &self.shape)
.iter()
.map(|i| (i + 1).to_string())
.collect::<Vec<_>>()
.join(",");
format!("{}[{}]", self.name, idx)
})
.collect()
}
}
pub(super) fn observed_rule_is_array_valued(rule: &AlgebraicRule) -> bool {
match rule {
AlgebraicRule::ArrayLoop { output_ranges, .. }
| AlgebraicRule::Recurrence { output_ranges, .. } => !output_ranges.is_empty(),
AlgebraicRule::Scalar { body, .. } => match &**body {
Expr::Operator(node) => {
node.op == "faq"
&& node.expr.is_some()
&& node.output_idx.as_ref().is_some_and(|ix| !ix.is_empty())
}
_ => false,
},
}
}
pub(super) fn expr_blocks_output_pruning(expr: &Expr) -> bool {
let Expr::Operator(node) = expr else {
return false;
};
if matches!(
node.op.as_str(),
"intersect_polygon" | "polygon_intersection_area" | "skolem" | "rank" | "distinct"
) {
return true;
}
if node.join.is_some() || node.distinct == Some(true) {
return true;
}
if let Some(ranges) = &node.ranges
&& ranges
.values()
.any(|r| r.ragged().is_some() || r.derived().is_some())
{
return true;
}
node.any_child(&mut expr_blocks_output_pruning)
}
#[cfg(feature = "solve")]
type SolveTape = Option<std::sync::Arc<TapeProgram>>;
#[cfg(feature = "solve")]
fn solution_metadata(
solver_name: &str,
stats: &SolveStats,
merged_variable_renames: HashMap<String, String>,
namespace: Option<String>,
) -> SolutionMetadata {
SolutionMetadata {
alg: solver_name.to_string(),
n_rhs_calls: stats.n_rhs_calls,
n_jacobian_calls: stats.n_jacobian_calls,
n_accepted_steps: stats.n_accepted_steps,
n_rejected_steps: stats.n_rejected_steps,
tape_fallbacks: Vec::new(),
merged_variable_renames,
namespace,
}
}
#[cfg(feature = "solve")]
struct ObservedCadence {
static_names: HashSet<String>,
static_rules: Vec<AlgebraicRule>,
segment_static_rules: Vec<AlgebraicRule>,
continuous_rules: Vec<AlgebraicRule>,
varying_rules: Vec<AlgebraicRule>,
}
#[cfg(feature = "solve")]
struct SolveSetup {
cadence: ObservedCadence,
sa0: ArrMap,
static_obs: ArrMap,
}
#[cfg(feature = "solve")]
struct SegmentedRun<'a> {
t0: f64,
t_end: f64,
boundaries: &'a [f64],
ic_vec: &'a [f64],
param_vec: &'a [f64],
setup: &'a SolveSetup,
opts: &'a SolveOptions,
tape: Option<&'a std::sync::Arc<TapeProgram>>,
scratches: Option<&'a SolveScratches>,
}
#[cfg(feature = "solve")]
pub(super) struct SolveScratches {
rhs: Rc<RefCell<Option<RhsScratch>>>,
jac: Rc<RefCell<Option<RhsScratch>>>,
program: Option<SolveProgram>,
}
#[cfg(feature = "solve")]
#[derive(Clone)]
struct SolveProgram {
program: super::range::RhsProgram,
rhs: Rc<RefCell<Option<super::range::RhsWorkspace>>>,
jac: Rc<RefCell<Option<super::range::RhsWorkspace>>>,
}
#[cfg(feature = "solve")]
impl SolveScratches {
fn new(compiled: &ArrayCompiled, tape: &std::sync::Arc<TapeProgram>) -> Self {
let program = (!compiled.is_xla()).then(|| SolveProgram {
program: super::range::RhsProgram::for_solve(compiled, std::sync::Arc::clone(tape)),
rhs: Rc::new(RefCell::new(None)),
jac: Rc::new(RefCell::new(None)),
});
SolveScratches {
rhs: Rc::new(RefCell::new(None)),
jac: Rc::new(RefCell::new(None)),
program,
}
}
fn forcing_refreshed(&self) {
for s in [&self.rhs, &self.jac] {
if let Some(scratch) = s.borrow_mut().as_mut() {
scratch.bump_forcing_epoch();
}
}
for w in self.program.iter().flat_map(|p| [&p.rhs, &p.jac]) {
if let Some(ws) = w.borrow_mut().as_mut() {
ws.forcing_refreshed();
}
}
}
}
pub(super) fn dependency_cone(
rules: &[AlgebraicRule],
seeds: &HashSet<String>,
) -> Option<Vec<AlgebraicRule>> {
let mut needed: HashSet<String> = seeds.clone();
let mut keep = vec![false; rules.len()];
let mut n_kept = 0usize;
for (i, rule) in rules.iter().enumerate().rev() {
if needed.contains(observed_rule_var(rule)) {
keep[i] = true;
n_kept += 1;
collect_expr_var_refs(observed_rule_body(rule), &mut needed);
}
}
if n_kept == rules.len() {
return None;
}
Some(
rules
.iter()
.zip(keep)
.filter(|(_, k)| *k)
.map(|(r, _)| r.clone())
.collect(),
)
}
#[cfg(feature = "solve")]
fn explain_ic_scope_miss(
err: SimulateError,
rhs: &Expr,
skipped: &HashMap<&str, String>,
) -> SimulateError {
let SimulateError::InvalidFieldInitialCondition { name, mut details } = err else {
return err;
};
if !skipped.is_empty() {
let mut reads: Vec<String> = super::compile::free_names(rhs).into_iter().collect();
reads.sort();
for read in reads {
if let Some(why) = skipped.get(read.as_str()) {
details.push_str(&format!(
"; the observed '{read}' is not in the build-time scope: it {why}"
));
}
}
}
SimulateError::InvalidFieldInitialCondition { name, details }
}
#[cfg(test)]
mod forcing_channel_tests {
use super::*;
use crate::parse::load_string;
fn arr1(v: &[f64]) -> ArrayD<f64> {
ArrayD::from_shape_vec(IxDyn(&[v.len()]), v.to_vec()).unwrap()
}
fn forced_model() -> ArrayCompiled {
let json = r#"
{
"esm": "1.1.0",
"metadata": {
"name": "forcing_channel"
},
"models": {
"Forced": {
"variables": {
"u": {
"type": "unknown",
"shape": [
"i"
],
"default": 0.0
}
},
"equations": [
{
"lhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"expr": {
"op": "D",
"args": [
{
"op": "index",
"args": [
"u",
"i"
]
}
],
"wrt": "t"
},
"ranges": {
"i": [
1,
3
]
}
},
"rhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"ranges": {
"i": [
1,
3
]
},
"expr": {
"op": "index",
"args": [
"w",
"i"
]
}
}
}
]
}
}
}
"#;
let file = load_string(json).expect("parse forcing model");
ArrayCompiled::from_file(&file).expect("compile forcing model")
}
#[test]
fn rhs_reads_forcing_array_and_reflects_mutation() {
let compiled = forced_model();
let forcing = compiled.forcing_handle();
let params = HashMap::new();
let state = vec![0.0, 0.0, 0.0];
forcing
.borrow_mut()
.insert("w".to_string(), arr1(&[10.0, 20.0, 30.0]));
let (dy1, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(
dy1,
vec![10.0, 20.0, 30.0],
"RHS must read the forcing array live from the buffer"
);
forcing
.borrow_mut()
.insert("w".to_string(), arr1(&[1.0, 2.0, 3.0]));
let (dy2, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(
dy2,
vec![1.0, 2.0, 3.0],
"a buffer mutation must change the RHS output"
);
assert_ne!(
dy1, dy2,
"the refreshed forcing must produce a different RHS"
);
let (dy_oracle, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, true);
assert_eq!(
dy_oracle,
vec![1.0, 2.0, 3.0],
"the oracle path resolves forcing identically"
);
}
#[test]
fn forcing_flows_through_the_production_solve() {
let compiled = forced_model();
compiled
.forcing_handle()
.borrow_mut()
.insert("w".to_string(), arr1(&[2.0, 4.0, 6.0]));
let params = HashMap::new();
let ics = HashMap::new(); let opts = SolveOptions::default();
let sol = compiled
.solve((0.0, 1.0), ¶ms, &ics, &opts)
.expect("solve with forcing");
for (i, want) in [2.0, 4.0, 6.0].iter().enumerate() {
let got = *sol.state[i].last().expect("trajectory non-empty");
assert!(
(got - want).abs() < 1e-6,
"forcing must drive the solve: state[{i}] got {got}, want {want}"
);
}
}
#[test]
fn empty_forcing_leaves_param_path_unaffected() {
let json = r#"
{
"esm": "1.1.0",
"metadata": {
"name": "param_path"
},
"models": {
"P": {
"variables": {
"u": {
"type": "unknown",
"shape": [
"i"
]
},
"k": {
"type": "parameter"
}
},
"equations": [
{
"lhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"expr": {
"op": "D",
"args": [
{
"op": "index",
"args": [
"u",
"i"
]
}
],
"wrt": "t"
},
"ranges": {
"i": [
1,
2
]
}
},
"rhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"ranges": {
"i": [
1,
2
]
},
"expr": {
"op": "*",
"args": [
"k",
{
"op": "index",
"args": [
"u",
"i"
]
}
]
}
}
}
]
}
}
}
"#;
let file = load_string(json).expect("parse param model");
let compiled = ArrayCompiled::from_file(&file).expect("compile param model");
let mut params = HashMap::new();
params.insert("k".to_string(), 2.0);
let state = vec![3.0, 5.0];
let (dy_no_forcing, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(
dy_no_forcing,
vec![6.0, 10.0],
"empty forcing leaves the parameter path identical (k·u)"
);
compiled
.forcing_handle()
.borrow_mut()
.insert("unrelated".to_string(), arr1(&[99.0]));
let (dy_with_junk, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(
dy_with_junk,
vec![6.0, 10.0],
"an unrelated forcing entry must not perturb the parameter path"
);
}
#[test]
fn fn_op_interp_linear_in_array_runtime() {
let json = r#"
{
"esm": "1.1.0",
"metadata": {
"name": "fn_array_path"
},
"models": {
"F": {
"variables": {
"u": {
"type": "unknown",
"shape": [
"i"
]
},
"code": {
"type": "parameter"
},
"looked_up": {
"type": "unknown"
}
},
"equations": [
{
"lhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"expr": {
"op": "D",
"args": [
{
"op": "index",
"args": [
"u",
"i"
]
}
],
"wrt": "t"
},
"ranges": {
"i": [
1,
2
]
}
},
"rhs": {
"op": "faq",
"args": [],
"output_idx": [
"i"
],
"ranges": {
"i": [
1,
2
]
},
"expr": "looked_up"
}
},
{
"lhs": "looked_up",
"rhs": {
"op": "fn",
"name": "interp.linear",
"args": [
{
"op": "const",
"value": [
10.0,
20.0,
40.0,
80.0,
160.0
],
"args": []
},
{
"op": "const",
"value": [
0.0,
1.0,
2.0,
3.0,
4.0
],
"args": []
},
"code"
]
}
}
]
}
}
}
"#;
let file = load_string(json).expect("parse fn model");
let compiled = ArrayCompiled::from_file(&file).expect("compile fn model");
let mut params = HashMap::new();
params.insert("code".to_string(), 2.0);
let state = vec![0.0, 0.0];
let (dy, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(
dy,
vec![40.0, 40.0],
"interp.linear(...,2.0)=40 must drive both cells (was NaN before the `fn` arm)"
);
params.insert("code".to_string(), 0.5);
let (dy2, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
assert_eq!(dy2, vec![15.0, 15.0], "interp.linear(...,0.5)=15");
}
}
#[cfg(all(test, feature = "solve"))]
mod fd_jvp_tests {
use super::*;
use std::cell::Cell;
fn f(y: &[f64], p: &[f64], t: f64, out: &mut [f64]) {
let n = y.len();
for (i, o) in out.iter_mut().enumerate() {
*o = p[0] * y[i] * y[(i + 1) % n] - (t * y[i]).sin() + 1.0 / y[i];
}
}
fn uncached(y: &[f64], p: &[f64], t: f64, v: &[f64]) -> Vec<f64> {
let n = y.len();
let mut y_norm = 0.0f64;
for &yi in y {
y_norm += yi * yi;
}
let y_norm = y_norm.sqrt().max(1.0);
let eps = f64::EPSILON.sqrt() * y_norm;
let yp: Vec<f64> = (0..n).map(|i| y[i] + eps * v[i]).collect();
let (mut fy, mut fyp) = (vec![0.0; n], vec![0.0; n]);
f(y, p, t, &mut fy);
f(&yp, p, t, &mut fyp);
(0..n).map(|i| (fyp[i] - fy[i]) / eps).collect()
}
fn unit(n: usize, j: usize) -> Vec<f64> {
(0..n).map(|i| if i == j { 1.0 } else { 0.0 }).collect()
}
fn bits(v: &[f64]) -> Vec<u64> {
v.iter().map(|x| x.to_bits()).collect()
}
fn apply(jvp: &mut FdJvp, y: &[f64], p: &[f64], t: f64, v: &[f64]) -> (Vec<f64>, usize) {
let calls = Cell::new(0usize);
let mut jv = vec![0.0; y.len()];
jvp.apply(y, p, t, v, &mut jv, |s, q, t, out| {
calls.set(calls.get() + 1);
f(s, q, t, out);
true
});
(jv, calls.get())
}
#[test]
fn a_jacobian_at_one_point_evaluates_the_base_once() {
let (y, p, t) = (vec![0.7, 1.3, 2.9, 0.4], vec![1.5], 0.25);
let n = y.len();
let mut jvp = FdJvp::new(n);
let mut total = 0;
for j in 0..n {
let v = unit(n, j);
let (jv, calls) = apply(&mut jvp, &y, &p, t, &v);
total += calls;
assert_eq!(bits(&jv), bits(&uncached(&y, &p, t, &v)), "column {j}");
}
assert_eq!(total, n + 1);
}
#[test]
fn a_new_point_evaluates_the_base_again() {
let (y, p, t) = (vec![0.7, 1.3, 0.0], vec![1.5], 0.25);
let v = vec![0.3, -1.1, 0.8];
let mut jvp = FdJvp::new(y.len());
assert_eq!(apply(&mut jvp, &y, &p, t, &v).1, 2);
assert_eq!(apply(&mut jvp, &y, &p, t, &v).1, 1, "same point");
let mut y2 = y.clone();
y2[1] = 1.3000000000000003;
let mut signed_zero = y.clone();
signed_zero[2] = -0.0;
let points: [(&[f64], &[f64], f64, &str); 4] = [
(&y2, &p, t, "one state"),
(&y, &[2.5], t, "the parameters"),
(&y, &p, 0.5, "the time"),
(&signed_zero, &p, t, "the sign of a zero"),
];
for (yk, pk, tk, what) in points {
apply(&mut jvp, &y, &p, t, &v);
let (jv, calls) = apply(&mut jvp, yk, pk, tk, &v);
assert_eq!(calls, 2, "{what}");
assert_eq!(bits(&jv), bits(&uncached(yk, pk, tk, &v)), "{what}");
}
}
#[test]
fn a_failed_base_evaluation_is_not_kept() {
let (y, p, t) = (vec![0.7, 1.3], vec![1.5], 0.25);
let v = vec![1.0, 0.0];
let mut jvp = FdJvp::new(y.len());
let mut jv = vec![0.0; 2];
let calls = Cell::new(0usize);
jvp.apply(&y, &p, t, &v, &mut jv, |s, q, t, out| {
calls.set(calls.get() + 1);
if calls.get() == 1 {
out.fill(f64::NAN);
return false;
}
f(s, q, t, out);
true
});
assert_eq!(calls.get(), 2);
assert!(jv.iter().all(|x| x.is_nan()));
let (jv, calls) = apply(&mut jvp, &y, &p, t, &v);
assert_eq!(calls, 2, "the base is evaluated again");
assert_eq!(bits(&jv), bits(&uncached(&y, &p, t, &v)));
}
}
#[cfg(all(test, feature = "solve"))]
mod field_ic_memo_tests {
use super::field_ic_resolutions;
use crate::compile_error::CompileError;
use crate::problem::{Compiler, ProblemOptions, Remake, esm_problem, remake, solve};
use crate::simulate::SimulateError;
use crate::simulate_array::per_cell_cells;
use crate::{SolveOptions, load_string};
use serde_json::{Value, json};
fn with_ic(n: usize, rhs_ic: Value) -> crate::EsmFile {
let doc = json!({
"esm": "1.1.0",
"metadata": {"name": "FieldIcMemo"},
"index_sets": {"x": {"kind": "interval", "size": n}},
"models": {"M": {
"variables": {
"u": {"type": "unknown", "units": "1", "shape": ["x"]},
"a": {"type": "parameter", "units": "1", "default": 0.5}
},
"equations": [
{"lhs": {"op": "ic", "args": ["u"]}, "rhs": rhs_ic},
{"lhs": {"op": "faq", "args": [], "output_idx": ["i"],
"ranges": {"i": {"from": "x"}},
"expr": {"op": "D", "args": [{"op": "index", "args": ["u", "i"]}],
"wrt": "t"}},
"rhs": {"op": "faq", "args": [], "output_idx": ["i"],
"ranges": {"i": {"from": "x"}},
"expr": {"op": "-", "args": [{"op": "index", "args": ["u", "i"]}]}}},
],
}},
});
load_string(&doc.to_string()).expect("loads")
}
fn scaled_index() -> Value {
json!({"op": "faq", "args": [], "output_idx": ["i"],
"ranges": {"i": {"from": "x"}},
"expr": {"op": "*", "args": ["a", "i"]}})
}
fn opts(compiler: Compiler) -> ProblemOptions {
ProblemOptions {
compiler: Some(compiler),
..Default::default()
}
}
fn first_column(sol: &crate::Solution) -> Vec<f64> {
sol.state.iter().map(|r| r[0]).collect()
}
fn bits(sol: &crate::Solution) -> Vec<Vec<u64>> {
sol.state
.iter()
.map(|r| r.iter().map(|x| x.to_bits()).collect())
.collect()
}
#[test]
fn the_first_solve_takes_the_resolution_construction_made() {
let file = with_ic(3, scaled_index());
for compiler in [Compiler::Native, Compiler::Interpreter] {
let before = field_ic_resolutions();
let prob = esm_problem(&file, (0.0, 1.0), opts(compiler))
.unwrap_or_else(|e| panic!("[{compiler}] {e}"));
assert_eq!(
field_ic_resolutions() - before,
1,
"[{compiler}] construction"
);
assert!(
prob.compiler_report()
.rules()
.iter()
.any(|r| r.kind == "initial condition"),
"[{compiler}] the report keeps its row"
);
let first = solve(&prob, &SolveOptions::default()).expect("solves");
assert_eq!(
field_ic_resolutions() - before,
1,
"[{compiler}] the first solve must not resolve the ics again"
);
assert_eq!(first_column(&first), [0.5, 1.0, 1.5], "[{compiler}]");
let second = solve(&prob, &SolveOptions::default()).expect("solves");
assert_eq!(field_ic_resolutions() - before, 2, "[{compiler}]");
assert_eq!(bits(&first), bits(&second), "[{compiler}]");
}
}
#[test]
fn a_remake_with_new_parameters_resolves_again() {
let file = with_ic(3, scaled_index());
let prob = esm_problem(&file, (0.0, 1.0), opts(Compiler::Native)).expect("builds");
let before = field_ic_resolutions();
let changed = remake(
&prob,
&Remake {
p: [("a".to_string(), 2.0)].into_iter().collect(),
..Default::default()
},
)
.expect("remakes");
let sol = solve(&changed, &SolveOptions::default()).expect("solves");
assert_eq!(
field_ic_resolutions() - before,
1,
"new parameters, new ics"
);
assert_eq!(first_column(&sol), [2.0, 4.0, 6.0]);
let sol = solve(&prob, &SolveOptions::default()).expect("solves");
assert_eq!(field_ic_resolutions() - before, 2);
assert_eq!(first_column(&sol), [0.5, 1.0, 1.5]);
}
#[test]
fn a_large_per_cell_initial_condition_is_refused_after_one_cell() {
let cumulative = json!({"op": "faq", "args": [], "output_idx": ["i"],
"ranges": {"i": {"from": "x"}, "j": {"from": "x"}},
"filter": {"op": "<=", "args": ["j", "i"]},
"expr": 1.0});
let refusal = |n: usize, ic: Value| {
let file = with_ic(n, ic);
let before = per_cell_cells();
let err = esm_problem(&file, (0.0, 1.0), opts(Compiler::Native)).expect_err("per cell");
let cells = per_cell_cells() - before;
match err {
SimulateError::Compile(CompileError::CompilerRefusedRule {
kind,
rule,
reason,
..
}) => ((kind, rule, reason), cells),
other => panic!("expected compiler_refused_rule, got {other:?}"),
}
};
let (small, _) = refusal(3, cumulative.clone());
let (large, cells) = refusal(100_000, cumulative.clone());
assert_eq!(large, small);
assert_eq!(large.0, "initial condition");
assert!(
large.2.ends_with(crate::simulate_array::ONE_CELL_NOTE),
"{}",
large.2
);
assert_eq!(cells, 1, "the walk must stop at its first cell");
let reciprocal = json!({"op": "/", "args": [1.0, cumulative]});
let (large, cells) = refusal(100_000, reciprocal);
assert_eq!(large, small);
assert_eq!(cells, 1);
}
struct Zeros;
impl crate::provider::CadenceProvider for Zeros {
fn materialize(
&mut self,
) -> Result<
std::collections::HashMap<String, crate::provider::NativeField>,
crate::provider::ProviderError,
> {
let zeros = ndarray::ArrayD::zeros(ndarray::IxDyn(&[3]));
Ok([("w".to_string(), crate::provider::NativeField::new(zeros))].into())
}
fn refresh(
&mut self,
_t: f64,
) -> Result<
Option<std::collections::HashMap<String, crate::provider::NativeField>>,
crate::provider::ProviderError,
> {
Ok(None)
}
fn refresh_times(&self) -> Vec<f64> {
Vec::new()
}
}
fn loaded(compiler: Compiler) -> crate::EsmProblem {
let file = with_ic(3, json!("w"));
let options = ProblemOptions {
providers: [(
"w".to_string(),
Box::new(Zeros) as Box<dyn crate::provider::CadenceProvider>,
)]
.into(),
..opts(compiler)
};
esm_problem(&file, (0.0, 1.0), options).unwrap_or_else(|e| panic!("[{compiler}] {e}"))
}
#[test]
fn a_forcing_write_before_the_first_solve_is_not_masked_by_the_memo() {
for compiler in [Compiler::Native, Compiler::Interpreter] {
let prob = loaded(compiler);
let before = field_ic_resolutions();
let sol = solve(&prob, &SolveOptions::default()).expect("solves");
assert_eq!(field_ic_resolutions() - before, 0, "[{compiler}]");
assert_eq!(first_column(&sol), [0.0, 0.0, 0.0], "[{compiler}]");
let prob = loaded(compiler);
let compiled = prob.debug_array_compiled().expect("an array model");
compiled.forcing_handle().borrow_mut().insert(
"w".to_string(),
ndarray::ArrayD::from_shape_vec(ndarray::IxDyn(&[3]), vec![1.0, 2.0, 3.0])
.expect("shape"),
);
let before = field_ic_resolutions();
let sol = solve(&prob, &SolveOptions::default()).expect("solves");
assert_eq!(field_ic_resolutions() - before, 1, "[{compiler}]");
assert_eq!(first_column(&sol), [1.0, 2.0, 3.0], "[{compiler}]");
}
}
}