earthsci_ast/simulate_array/driver.rs
1//! Solver plumbing on [`ArrayCompiled`]: `solve` / `solve_inspect`
2//! (diffsol problem build, RHS/Jacobian closures, observed trajectory
3//! exposure), the `debug_*` RHS entry points, and the external forcing-channel
4//! handle ([`ArrayCompiled::forcing_handle`]).
5
6#[cfg(feature = "solve")]
7use super::tape::TapeProgram;
8use super::*;
9#[cfg_attr(not(feature = "solve"), allow(unused_imports))]
10use crate::simulate::SimulateError;
11#[cfg(feature = "solve")]
12use crate::simulate::{
13 Alg, Progress, ProgressFn, ReturnCode, Solution, SolutionMetadata, SolveOptions, SolveStats,
14};
15// The solver is optional (esm-libraries-spec §2.5.9). Everything below that
16// touches `diffsol` is gated on the `solve` feature; the `debug_*` RHS entry
17// points and the forcing-channel handle are not, because a caller that only
18// BUILDS a EsmProblem still reaches them.
19#[cfg(feature = "solve")]
20use diffsol::{Bdf, FaerLU, FaerMat, NewtonNonlinearSolver, OdeBuilder, Sdirk, VectorHost};
21use std::collections::HashSet;
22
23/// A taped read-out of every observed, at one state and one time.
24///
25/// Under [`RuntimeMode::Native`] every rule is on the tape — construction
26/// refused the document otherwise — so the observeds a setup or output-time
27/// pass needs are already computed by the tape's own CONST / SEGMENT /
28/// CONTINUOUS sections. This runs the tape with its `Export` publishes forced
29/// on and hands back the published map.
30///
31/// It exists because the passes it replaces evaluated the SAME rules through
32/// the whole-array overlay, with the per-cell oracle beneath it and no entry in
33/// any decline report — the silent demotion `esm-libraries-spec.md` §2.5.10
34/// refuses, and one that a gate over the tape's declined rules alone cannot
35/// see, because the overlay declines on its own terms. Serving them from the tape
36/// makes the gate honest: there is no second evaluator left to disagree.
37///
38/// The production RHS scratch keeps its exports OFF (nothing inside a call
39/// reads them, and the publish is a memcpy per observed per call),
40/// so this carries its own scratch and its own slab.
41#[cfg(feature = "solve")]
42struct TapedObserveds {
43 scratch: RhsScratch,
44 /// The tape writes `dy` whether or not the caller wants it; a setup or
45 /// output-time pass does not, so it lands here and is discarded.
46 dy: Vec<f64>,
47}
48
49#[cfg(feature = "solve")]
50impl TapedObserveds {
51 fn new(compiled: &ArrayCompiled, tape: &std::sync::Arc<TapeProgram>) -> Self {
52 let mut scratch = RhsScratch::new(&compiled.var_shapes);
53 scratch.set_const_arrays(std::sync::Arc::clone(&compiled.const_scope));
54 scratch.install_tape(std::sync::Arc::clone(tape));
55 scratch.set_exports_active(true);
56 TapedObserveds {
57 scratch,
58 dy: vec![0.0f64; compiled.n_states],
59 }
60 }
61
62 /// The map the last [`Self::at`] published, without running the tape
63 /// again — so the node-0 values that decided the row layout are recorded
64 /// rather than recomputed.
65 fn scratch_observeds(&self) -> &ArrMap {
66 self.scratch
67 .taped_observeds()
68 .expect("the harvester installs a tape on its own scratch")
69 }
70
71 /// Every observed's value at `state` and `t`, published by the tape.
72 ///
73 /// The CONST and SEGMENT sections prime on the first call and are not
74 /// re-run while the parameter vector is unchanged, so a sweep over output
75 /// nodes pays for the CONTINUOUS section only — the same amortization the
76 /// static hoist existed to provide.
77 fn at(&mut self, compiled: &ArrayCompiled, state: &[f64], params: &[f64], t: f64) -> &ArrMap {
78 evaluate_rhs_with_scratch(
79 &RhsCall {
80 rhs_rules: &compiled.rhs_rules,
81 observed_rules: &compiled.observed_rules,
82 var_shapes: &compiled.var_shapes,
83 param_names: &compiled.param_names,
84 state,
85 params,
86 forcing: &compiled.forcing,
87 t,
88 declared: &compiled.declared_names,
89 rule_set: compiled.rule_set,
90 },
91 &mut self.dy,
92 false,
93 &mut RhsStats::default(),
94 &mut self.scratch,
95 );
96 self.scratch
97 .taped_observeds()
98 .expect("the harvester installs a tape on its own scratch")
99 }
100}
101
102/// The refusal a model that serves [`crate::Compiler::Xla`] gets when no
103/// compiled program was installed on it.
104///
105/// It is `compiler_unavailable` rather than a refused rule because nothing
106/// about the DOCUMENT is wrong: the build either never got as far as emitting
107/// (no `xla` feature) or handed the integrator a model it had not finished.
108#[cfg(feature = "solve")]
109fn xla_not_installed() -> SimulateError {
110 SimulateError::CompilerUnavailable {
111 compiler: "xla",
112 details: "this Problem names `xla` but carries no compiled program; the build did not \
113 install one, and the tape is NOT run in its place — a compiler that cannot \
114 run says so (esm-libraries-spec §2.5.10)"
115 .to_string(),
116 }
117}
118
119/// Run the compiled XLA right-hand side into `out`, routing a device failure
120/// into `fault` (FIRST failure wins) and leaving `out` NaN so the solver stops
121/// instead of integrating whatever was there before.
122///
123/// Shared by the production right-hand side and the finite-difference
124/// Jacobian, which under [`crate::Compiler::Xla`] evaluate exactly the same
125/// program; a second copy of the failure handling would be a second place to
126/// get it wrong.
127#[cfg(all(feature = "solve", feature = "xla"))]
128fn run_xla_rhs(
129 program: &crate::xla_runtime::CompiledRhs,
130 fault: &Rc<RefCell<Option<String>>>,
131 state: &[f64],
132 params: &[f64],
133 t: f64,
134 out: &mut [f64],
135) -> bool {
136 match program.eval_into(state, params, t, out) {
137 Ok(()) => true,
138 Err(e) => {
139 let mut slot = fault.borrow_mut();
140 if slot.is_none() {
141 *slot = Some(e.to_string());
142 }
143 for v in out.iter_mut() {
144 *v = f64::NAN;
145 }
146 false
147 }
148 }
149}
150
151/// The finite-difference Jacobian-vector product `J v ≈ (f(y + εv) − f(y)) / ε`
152/// the XLA arm of the integrator's Jacobian closure computes, with the
153/// evaluation at the base point shared across calls and every buffer
154/// allocated once.
155///
156/// diffsol assembles a Jacobian by calling the product once per column, all
157/// at ONE `(y, p, t)`. Only `f(y + εv)` depends on the column, so keeping
158/// `f(y)` makes a Jacobian of `n` states cost `n + 1` right-hand-side
159/// evaluations instead of `2n` — and under [`crate::Compiler::Xla`] each one
160/// is a host-to-device-to-host round trip. The kept `f(y)` is reused only
161/// when `y`, `p` and `t` are bit-for-bit the ones it was computed at, and the
162/// forcing buffer the right-hand side may read is fixed for the segment the
163/// closure lives in, so the product is bit-identical to recomputing both
164/// evaluations every call.
165#[cfg(all(feature = "solve", any(test, feature = "xla")))]
166struct FdJvp {
167 base_y: Vec<f64>,
168 base_p: Vec<f64>,
169 base_t: f64,
170 /// `f(base_y, base_p, base_t)`; meaningful only while `base_valid`.
171 f_y: Vec<f64>,
172 base_valid: bool,
173 y_perturbed: Vec<f64>,
174 f_yp: Vec<f64>,
175}
176
177#[cfg(all(feature = "solve", any(test, feature = "xla")))]
178impl FdJvp {
179 fn new(n_states: usize) -> Self {
180 FdJvp {
181 base_y: vec![0.0; n_states],
182 base_p: Vec::new(),
183 base_t: 0.0,
184 f_y: vec![0.0; n_states],
185 base_valid: false,
186 y_perturbed: vec![0.0; n_states],
187 f_yp: vec![0.0; n_states],
188 }
189 }
190
191 /// Write `J(y, p, t) v` into `jv`. `eval(state, params, t, out)` fills
192 /// `out` with the right-hand side and reports whether it succeeded; a
193 /// failed evaluation of `f(y)` is not kept, so the next call tries again.
194 fn apply(
195 &mut self,
196 y: &[f64],
197 p: &[f64],
198 t: f64,
199 v: &[f64],
200 jv: &mut [f64],
201 mut eval: impl FnMut(&[f64], &[f64], f64, &mut [f64]) -> bool,
202 ) {
203 let n = y.len();
204 let mut y_norm = 0.0f64;
205 for &yi in y {
206 y_norm += yi * yi;
207 }
208 let y_norm = y_norm.sqrt().max(1.0);
209 let eps = f64::EPSILON.sqrt() * y_norm;
210
211 let same_point = self.base_valid
212 && self.base_t.to_bits() == t.to_bits()
213 && same_bits(&self.base_y, y)
214 && same_bits(&self.base_p, p);
215 if !same_point {
216 self.base_y.clear();
217 self.base_y.extend_from_slice(y);
218 self.base_p.clear();
219 self.base_p.extend_from_slice(p);
220 self.base_t = t;
221 self.f_y.resize(n, 0.0);
222 self.base_valid = eval(y, p, t, &mut self.f_y);
223 }
224
225 self.y_perturbed.clear();
226 self.y_perturbed
227 .extend(y.iter().zip(v).map(|(&yi, &vi)| yi + eps * vi));
228 self.f_yp.resize(n, 0.0);
229 eval(&self.y_perturbed, p, t, &mut self.f_yp);
230 for ((out, &fp), &f0) in jv.iter_mut().zip(&self.f_yp).zip(&self.f_y) {
231 *out = (fp - f0) / eps;
232 }
233 }
234}
235
236/// Bitwise slice equality: `-0.0` and `0.0` differ, and a `NaN` equals the
237/// same `NaN`, which is what reusing a function value requires.
238#[cfg(all(feature = "solve", any(test, feature = "xla")))]
239fn same_bits(a: &[f64], b: &[f64]) -> bool {
240 a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
241}
242
243/// One evaluation a model's field initial conditions performed
244/// ([`ArrayCompiled::field_ic_records`]).
245#[cfg(feature = "solve")]
246#[derive(Clone, Debug)]
247pub(crate) struct FieldIcRecord {
248 /// The `ic` target, or the state-free observed materialized for its scope.
249 pub name: String,
250 /// `"initial condition"`, `"initial-condition scope"` or
251 /// `"initialization equation"`.
252 pub kind: &'static str,
253 /// Whether the reference evaluator walked a `faq` in it cell by cell.
254 pub per_cell: bool,
255}
256
257/// The field initial conditions construction resolved, held for the first
258/// solve ([`ArrayCompiled::field_ic_records`] fills it, the first
259/// `build_initial_state` takes it).
260///
261/// Construction has to evaluate every field `ic` — a strict compiler refuses
262/// one it would walk per cell, and the report names the route of each — and
263/// the first solve needs exactly those values, so resolving them twice was a
264/// second full evaluation of every `ic` field on every Problem. The resolution
265/// reads two inputs, and the memo answers only when neither has moved:
266///
267/// * the parameters, checked here bit for bit against the positional vector
268/// the solve resolved (a [`crate::remake`] with new `p` shares this model and
269/// misses);
270/// * the provider forcing buffer, which the crate's own drivers leave as
271/// construction had it until the initial state is built (the CONST providers
272/// are bound before the compiler gate runs, and a discrete refresh comes
273/// after `u0`). A host can write it in between through
274/// [`ArrayCompiled::forcing_handle`], so the memo records the buffer's
275/// handle generation, which every call to that accessor bumps, and answers
276/// only while it is unchanged. Taking the memo, rather than keeping it, is
277/// what holds the rest: only the first solve can read it, and every later
278/// one resolves from the buffer as it stands.
279///
280/// Initial-condition overrides are applied over the result, not read by it.
281#[cfg(feature = "solve")]
282pub(crate) struct FieldIcMemo {
283 params: Vec<u64>,
284 forcing_generation: u64,
285 slots: HashMap<usize, f64>,
286}
287
288#[cfg(all(test, feature = "solve"))]
289thread_local! {
290 static FIELD_IC_RESOLUTIONS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
291}
292
293/// Test hook: how many times this thread has resolved a model's field initial
294/// conditions.
295#[cfg(all(test, feature = "solve"))]
296pub(crate) fn field_ic_resolutions() -> u64 {
297 FIELD_IC_RESOLUTIONS.with(std::cell::Cell::get)
298}
299
300#[cfg(feature = "solve")]
301fn param_bits(param_vec: &[f64]) -> Vec<u64> {
302 param_vec.iter().map(|x| x.to_bits()).collect()
303}
304
305impl ArrayCompiled {
306 /// Whether the TAPE serves this Problem's passes other than the right-hand
307 /// side — the build-time materialization of constants and static
308 /// observeds, the per-segment seed, the inspection snapshot and the
309 /// observeds reported at output times (esm-libraries-spec §2.5.10).
310 ///
311 /// True under the two STRICT compilers, [`crate::Compiler::Native`] and
312 /// [`crate::Compiler::Xla`], and for the same reason in both: each refuses
313 /// the document unless every rule lowered to the tape, so there is no
314 /// second evaluator left for those passes to drift against. Under `xla`
315 /// the right-hand side itself runs on the emitted executable instead (see
316 /// `is_xla`) — the emitted program's only output is `du`, so the
317 /// observed passes stay on the tape the emitter was built from, which is
318 /// still ONE evaluator rather than two.
319 pub(crate) fn tape_serves_passes(&self) -> bool {
320 matches!(self.runtime_mode, RuntimeMode::Native | RuntimeMode::Xla)
321 }
322
323 /// Whether this model serves a [`crate::Compiler::Xla`] build: the
324 /// right-hand side (and the finite-difference Jacobian differenced out of
325 /// it) is the XLA executable rather than the tape's own interpreter.
326 pub(crate) fn is_xla(&self) -> bool {
327 self.runtime_mode == RuntimeMode::Xla
328 }
329
330 /// Emit this model's tape as an XLA computation, compile it for the
331 /// device, and install it as the right-hand side
332 /// [`crate::Compiler::Xla`] runs (API_SPEC §5.8).
333 ///
334 /// Called ONCE, at construction, by [`crate::problem::esm_problem`] — so a
335 /// model the emitter cannot lower is a BUILD refusal naming the rule
336 /// rather than a surprise on the first step, and so the XLA compilation
337 /// (seconds for a large program) is paid once per Problem rather than once
338 /// per integration segment. Idempotent: a second call keeps the first
339 /// executable.
340 #[cfg(feature = "xla")]
341 pub(crate) fn install_xla_rhs(&self) -> Result<(), crate::xla_runtime::CompileRhsError> {
342 use crate::xla_runtime::{CompileRhsError, CompiledRhs};
343 if self.xla_rhs.get().is_some() {
344 return Ok(());
345 }
346 let program = CompiledRhs::compile(self)?;
347 // The two length contracts `CompiledRhs::eval` would otherwise raise
348 // per CALL, checked once here instead. After this the only way an
349 // evaluation can fail is the device itself, which is what the run-time
350 // error channel in [`Self::run_one_segment`] carries.
351 if program.n_states() != self.n_states {
352 return Err(CompileRhsError::Runtime(format!(
353 "the emitted program returns {} state slots, the model has {}",
354 program.n_states(),
355 self.n_states
356 )));
357 }
358 if program.params_len() < self.param_names.len() {
359 return Err(CompileRhsError::Runtime(format!(
360 "the emitted program takes {} parameters, the model has {}",
361 program.params_len(),
362 self.param_names.len()
363 )));
364 }
365 let _ = self.xla_rhs.set(Rc::new(program));
366 Ok(())
367 }
368
369 /// The installed XLA executable, when this model serves
370 /// [`crate::Compiler::Xla`]. `None` under every other compiler, so a
371 /// caller can hand the answer straight to the closure builders.
372 #[cfg(feature = "xla")]
373 pub(crate) fn xla_program(&self) -> Option<Rc<crate::xla_runtime::CompiledRhs>> {
374 if self.is_xla() {
375 self.xla_rhs.get().cloned()
376 } else {
377 None
378 }
379 }
380
381 /// Whether this model serves a [`crate::Compiler::Interpreter`] build: no
382 /// tape, and the whole-array overlay off, so every rule is walked per cell.
383 pub(crate) fn is_interpreter(&self) -> bool {
384 self.runtime_mode == RuntimeMode::Interpreter
385 }
386
387 /// The flatten-time merge map (issue #230): every state spelling an
388 /// `operator_compose` renaming match DELETED, mapped onto the survivor.
389 pub(crate) fn merged_renames(&self) -> &HashMap<String, String> {
390 &self.merged_renames
391 }
392
393 /// The shaped parameters this build left with no value (see the field).
394 pub(crate) fn unvalued_shaped_params(&self) -> &[String] {
395 &self.unvalued_shaped_params
396 }
397
398 /// The names the forcing buffer serves: every parameter refreshed from
399 /// outside the model (a data source or a registered handler).
400 pub(crate) fn forcing_names(&self) -> impl Iterator<Item = &String> {
401 self.forcing_decls.keys()
402 }
403
404 /// The declared default of forcing parameter `name`: `None` when it
405 /// declares none, `Some(None)` when its shape does not resolve yet, and
406 /// otherwise the dense field it denotes over that shape.
407 pub(crate) fn forcing_default(&self, name: &str) -> Option<Option<ArrayD<f64>>> {
408 let entry = self.forcing_defaults.get(name)?;
409 Some(entry.as_ref().and_then(|(shape, values)| {
410 ArrayD::from_shape_vec(ndarray::IxDyn(shape), values.clone()).ok()
411 }))
412 }
413
414 /// A clonable handle to the external forcing buffer (PR-1, ess-14f.7). A
415 /// driver that integrates this model in discrete-cadence segments holds the
416 /// returned `Rc` and, at each cadence boundary, refreshes a loader-fed
417 /// field: `compiled.forcing_handle().borrow_mut().insert(var, regridded)`.
418 /// The captured RHS/Jacobian closures read the *same* buffer live on the
419 /// next `step()`, so the refresh is reflected without rebuilding the
420 /// problem. The buffer is shared (the handle and the closures clone one
421 /// `Rc`); mutate it only *between* segments, never inside a solver step, to
422 /// keep the RHS pure within a segment.
423 ///
424 /// Whoever holds the handle may write the buffer at any later point, so
425 /// taking one also retires the field initial conditions construction
426 /// resolved from it ([`FieldIcMemo`]): the next solve resolves them from
427 /// the buffer as it then stands.
428 pub fn forcing_handle(&self) -> Rc<RefCell<HashMap<String, ArrayD<f64>>>> {
429 self.forcing_generation
430 .set(self.forcing_generation.get().wrapping_add(1));
431 Rc::clone(&self.forcing)
432 }
433
434 /// The forcing buffer, for the crate's own drivers, which write it only
435 /// where [`FieldIcMemo`] allows: before construction records the memo, or
436 /// after the initial state is built. Unlike [`Self::forcing_handle`] it
437 /// leaves the memo standing.
438 pub(crate) fn forcing_buffer(&self) -> Rc<RefCell<HashMap<String, ArrayD<f64>>>> {
439 Rc::clone(&self.forcing)
440 }
441
442 /// Every DATA-FED parameter this build routed to the forcing channel, as
443 /// `(name, data_sources key)` — see [`crate::data_fed::refuse_unbound`],
444 /// which is the only reader.
445 pub(crate) fn data_fed(&self) -> &[(String, String)] {
446 &self.data_fed
447 }
448
449 /// Prefix each recorded data-fed name with the single model's namespace.
450 ///
451 /// The single-model build keeps variable names BARE (`k`), while the
452 /// coupled build namespaces them (`Forcing.k`) — and a refusal that named
453 /// `k` would name something the caller cannot find in a document with two
454 /// components. Applied where `namespace` itself is recorded, so the two
455 /// never disagree.
456 pub(crate) fn qualify_data_fed(&mut self, namespace: &str) {
457 for (name, _) in self.data_fed.iter_mut() {
458 *name = format!("{namespace}.{name}");
459 }
460 }
461
462 /// Whether this model has anything for the ODE solver to integrate — at
463 /// least one `D(var, t) = …` rule.
464 ///
465 /// An array model whose every equation is ALGEBRAIC compiles to observed
466 /// rules and no `RhsRule`, so it evaluates without the integrator ever
467 /// touching a value. That distinction is what
468 /// `problem::reject_f32_integration` needs: a Float32 declaration is
469 /// honourable for such a model and not for one with real dynamics.
470 pub fn has_differential_equations(&self) -> bool {
471 !self.rhs_rules.is_empty()
472 }
473
474 /// The single model's own namespace (the top-level `models` map key), or
475 /// `None` on the flattened path, whose names are already qualified.
476 ///
477 /// The single-model build names its slots BARE (`u[1]`, `k`), because the
478 /// raw `Model` it consumed carries no namespace; the flattened build
479 /// qualifies them (`M.u[1]`). Reported here so
480 /// [`crate::problem::EsmProblem`] can present ONE spelling whichever route
481 /// its document took.
482 pub(crate) fn namespace(&self) -> Option<&str> {
483 self.namespace.as_deref()
484 }
485
486 /// The OBSERVED variables this model declares, in dependency order.
487 ///
488 /// What [`crate::problem::observed_trajectories`] resolves a caller's name
489 /// against, with the §5.8 precedence.
490 pub fn observed_variable_names(&self) -> Vec<String> {
491 self.observed_rules
492 .iter()
493 .map(|r| observed_rule_var(r).clone())
494 .collect()
495 }
496
497 /// Every state slot's name (`u[2,3]`, or a 0-D state's bare name), in
498 /// flat state-vector order. Built from the slot layout the first time it
499 /// is asked for, then kept.
500 pub fn state_variable_names(&self) -> &[String] {
501 self.state_names
502 .get_or_init(|| super::layout::slot_names(&self.var_shapes))
503 }
504
505 /// [`Self::state_variable_names`] under this model's single-model
506 /// namespace (`M.u[2,3]`), the spelling the Problem reports (API_SPEC
507 /// §5.8's qualification rule, `crate::problem::qualify`, applied to each
508 /// variable's name; the cell suffix carries no `.`, so qualifying the
509 /// variable and qualifying the slot name agree). The same names as
510 /// [`Self::state_variable_names`] when the model has no namespace.
511 pub(crate) fn qualified_state_names(&self) -> &[String] {
512 let Some(ns) = self.namespace.as_deref() else {
513 return self.state_variable_names();
514 };
515 self.qualified_state_names.get_or_init(|| {
516 super::layout::slot_names_with(&self.var_shapes, |v| crate::problem::qualify(ns, v))
517 })
518 }
519
520 /// The length of the flat state vector.
521 pub fn n_states(&self) -> usize {
522 self.n_states
523 }
524 pub fn parameter_names(&self) -> &[String] {
525 &self.param_names
526 }
527
528 /// Cadence partition of the observed rules given the discrete-forcing names,
529 /// as `(const, discrete, continuous)` sorted name lists. Exposed for the
530 /// cadence-tier test: a forcing-derived but state-free / `t`-free observed
531 /// must land in the DISCRETE tier (materialized once per segment), not
532 /// CONTINUOUS (recomputed every step). Mirrors the split in `run_segmented`.
533 #[doc(hidden)]
534 pub fn debug_cadence_partition(
535 &self,
536 discrete_forcing: &[String],
537 ) -> (Vec<String>, Vec<String>, Vec<String>) {
538 let df: HashSet<String> = discrete_forcing.iter().cloned().collect();
539 let const_set = self.classify_static_observeds(&df);
540 let leq_discrete = self.classify_segment_invariant_observeds(&df, true);
541 let mut const_: Vec<String> = const_set.iter().cloned().collect();
542 let mut discrete: Vec<String> = leq_discrete
543 .iter()
544 .filter(|n| !const_set.contains(*n))
545 .cloned()
546 .collect();
547 let mut continuous: Vec<String> = self
548 .observed_rules
549 .iter()
550 .map(observed_rule_var)
551 .filter(|n| !leq_discrete.contains(*n))
552 .cloned()
553 .collect();
554 const_.sort();
555 discrete.sort();
556 continuous.sort();
557 (const_, discrete, continuous)
558 }
559
560 /// Evaluate the RHS `f(state, t)` once and return `(dy, stats)`. Exposed for
561 /// the no-scalarization verification (ess-bdm): callers compare the
562 /// vectorized path (`force_scalar = false`) against the per-cell oracle
563 /// (`force_scalar = true`) for bit-equivalence, and assert that the
564 /// vectorized [`RhsStats::kernel_ops`] is independent of the grid size N.
565 #[doc(hidden)]
566 pub fn debug_eval_rhs(
567 &self,
568 state: &[f64],
569 t: f64,
570 params: &HashMap<String, f64>,
571 force_scalar: bool,
572 ) -> (Vec<f64>, RhsStats) {
573 let param_vec = self.debug_resolve_params(params);
574 let mut dy = vec![0.0f64; self.n_states];
575 let mut stats = RhsStats::default();
576 let mut scratch = RhsScratch::new(&self.var_shapes);
577 scratch.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
578 evaluate_rhs_with_scratch(
579 &RhsCall {
580 rhs_rules: &self.rhs_rules,
581 observed_rules: &self.observed_rules,
582 var_shapes: &self.var_shapes,
583 param_names: &self.param_names,
584 state,
585 params: ¶m_vec,
586 forcing: &self.forcing,
587 t,
588 declared: &self.declared_names,
589 rule_set: self.rule_set,
590 },
591 &mut dy,
592 force_scalar,
593 &mut stats,
594 &mut scratch,
595 );
596 (dy, stats)
597 }
598
599 /// Build a persistent [`RhsScratch`] sized to this model. Exposed for the
600 /// zero-allocation verification (ess-mro): a counting-allocator test drives
601 /// [`Self::debug_eval_rhs_into`] with a reused scratch and asserts that the
602 /// steady-state vectorized RHS allocates nothing.
603 #[doc(hidden)]
604 pub fn debug_new_scratch(&self) -> RhsScratch {
605 let mut s = RhsScratch::new(&self.var_shapes);
606 s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
607 s
608 }
609
610 /// Build a scratch with the compiled tape installed (Step 3b) — what the
611 /// production RHS closure carries. Exposed for the fast-executor A/B,
612 /// allocation-steady-state and invalidation tests, driven through
613 /// [`Self::debug_eval_rhs_into`].
614 ///
615 /// A model whose tape declined a rule has no program to install, so this
616 /// is the same refusal a solve raises ([`Self::strict_tape`]): a scratch
617 /// without a tape would run the legacy evaluator, and a tape-versus-legacy
618 /// comparison would then compare the legacy evaluator with itself.
619 #[doc(hidden)]
620 pub fn debug_new_scratch_taped(&self) -> Result<RhsScratch, SimulateError> {
621 let mut s = RhsScratch::new(&self.var_shapes);
622 s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
623 s.install_tape(self.strict_tape(&HashSet::new())?);
624 Ok(s)
625 }
626
627 /// The precision environment this model was compiled under (the
628 /// document's element type and its per-variable ones), which a solve arms
629 /// for its whole run. The `debug_*` right-hand-side entries do not arm it;
630 /// a caller comparing them on a document that declares element types
631 /// holds this guard around the calls, as a solve would.
632 #[doc(hidden)]
633 pub fn debug_precision_env(&self) -> crate::precision::Env {
634 self.precision.clone()
635 }
636
637 /// Resolve a parameter map into the positional parameter vector once, so the
638 /// zero-allocation RHS test can pre-build it outside the measured loop.
639 #[doc(hidden)]
640 pub fn debug_resolve_params(&self, params: &HashMap<String, f64>) -> Vec<f64> {
641 let mut param_vec = vec![0.0f64; self.param_names.len()];
642 for (i, name) in self.param_names.iter().enumerate() {
643 if let Some(&v) = params.get(name) {
644 param_vec[i] = v;
645 } else if let Some(d) = self.param_defaults[i] {
646 param_vec[i] = d;
647 }
648 }
649 param_vec
650 }
651
652 /// Evaluate the vectorized RHS into a caller-owned `dy` using a caller-owned
653 /// scratch — the allocation-free entry point. With a warmed scratch and a
654 /// pre-resolved `param_vec`, this performs no heap allocation (ess-mro
655 /// acceptance criterion 1).
656 #[doc(hidden)]
657 pub fn debug_eval_rhs_into(
658 &self,
659 state: &[f64],
660 t: f64,
661 param_vec: &[f64],
662 dy: &mut [f64],
663 scratch: &mut RhsScratch,
664 stats: &mut RhsStats,
665 ) {
666 evaluate_rhs_with_scratch(
667 &RhsCall {
668 rhs_rules: &self.rhs_rules,
669 observed_rules: &self.observed_rules,
670 var_shapes: &self.var_shapes,
671 param_names: &self.param_names,
672 state,
673 params: param_vec,
674 forcing: &self.forcing,
675 t,
676 declared: &self.declared_names,
677 rule_set: self.rule_set,
678 },
679 dy,
680 false,
681 stats,
682 scratch,
683 );
684 }
685
686 /// Whether every observed rule is a plain scalar expression: a
687 /// [`AlgebraicRule::Scalar`] whose body [`check_scalar_evaluable`] admits
688 /// (no array, tensor or geometry op).
689 ///
690 /// The precondition under which a stateless evaluation
691 /// ([`Self::evaluate_stateless_observeds`]) is the whole answer for a
692 /// document with nothing to integrate. A SHAPED document answers
693 /// differently — from the fields its build materialized — so it is told
694 /// apart here rather than by what an evaluation happens to produce.
695 #[cfg(all(not(target_arch = "wasm32"), feature = "solve"))]
696 pub(crate) fn observeds_are_scalar(&self) -> bool {
697 self.observed_rules.iter().all(|rule| match rule {
698 AlgebraicRule::Scalar { body, .. } => check_scalar_evaluable(body).is_ok(),
699 _ => false,
700 })
701 }
702
703 /// Every 0-D observed of a model with NO state vector, evaluated at each of
704 /// `times` — `(name, values)` in dependency order, one value per time.
705 ///
706 /// The answer for a document with nothing to integrate: its observeds are
707 /// pure functions of the parameters and `t`, so one evaluation per time is
708 /// the whole of it, and no integrator is involved. Each evaluation is one
709 /// right-hand-side call over the empty state, on the evaluator this model
710 /// serves — the tape under `native` / `xla`, the per-cell oracle under
711 /// `interpreter` — so the values are the ones a solve would report.
712 ///
713 /// Names are the rules' own (flattened on the `from_flattened` path, bare
714 /// with [`Self::namespace`] on the single-model one). An observed with more
715 /// than one cell is not reported; it has no scalar value per time. One
716 /// declared with a single cell (`[1]`) is, and the caller decides its rank
717 /// from the declaration.
718 ///
719 /// # Errors
720 ///
721 /// A parameter override that designates nothing, a parameter with neither
722 /// an override nor a default, and any fault the evaluator latched (an
723 /// out-of-range const-array gather, an unbound name) — all as a solve would
724 /// report them. A model that HAS state is refused rather than read at a
725 /// state nobody supplied.
726 #[cfg(all(not(target_arch = "wasm32"), feature = "solve"))]
727 pub(crate) fn evaluate_stateless_observeds(
728 &self,
729 params: &HashMap<String, f64>,
730 times: &[f64],
731 ) -> Result<Vec<(String, Vec<f64>)>, SimulateError> {
732 let mut out: Vec<(String, Vec<f64>)> = Vec::new();
733 self.stateless_pass(params, times, |k, obs| {
734 if k == 0 {
735 out = self
736 .observed_rules
737 .iter()
738 .map(observed_rule_var)
739 .filter(|name| obs.get(*name).is_some_and(|a| a.len() == 1))
740 .map(|name| (name.clone(), Vec::with_capacity(times.len())))
741 .collect();
742 }
743 for (name, values) in &mut out {
744 let v = obs.get(name).and_then(|a| a.first().copied());
745 values.push(v.unwrap_or(f64::NAN));
746 }
747 })?;
748 Ok(out)
749 }
750
751 /// Every observed of a model with NO state vector, whatever its shape,
752 /// evaluated once at `t` — `(name, field)` in dependency order.
753 ///
754 /// [`Self::evaluate_stateless_observeds`]'s single-time sibling for the
755 /// build-time fields of a state-free document (`observed_field`): the same
756 /// one right-hand-side call over the empty state, on the evaluator this
757 /// model serves, so a shaped observed comes off the tape under `native`
758 /// exactly as a scalar one does. Errors as that method's.
759 pub(crate) fn evaluate_stateless_fields(
760 &self,
761 params: &HashMap<String, f64>,
762 t: f64,
763 ) -> Result<Vec<(String, ArrayD<f64>)>, SimulateError> {
764 let mut out: Vec<(String, ArrayD<f64>)> = Vec::new();
765 self.stateless_pass(params, &[t], |_, obs| {
766 out = self
767 .observed_rules
768 .iter()
769 .map(observed_rule_var)
770 .filter_map(|name| obs.get(name).map(|a| (name.clone(), a.clone())))
771 .collect();
772 })?;
773 Ok(out)
774 }
775
776 /// One right-hand-side call over the empty state per entry of `times`,
777 /// handing `visit` the observeds each call published.
778 fn stateless_pass(
779 &self,
780 params: &HashMap<String, f64>,
781 times: &[f64],
782 mut visit: impl FnMut(usize, &ArrMap),
783 ) -> Result<(), SimulateError> {
784 if self.n_states != 0 {
785 return Err(SimulateError::Compile(
786 crate::compile_error::CompileError::InterpreterBuildError {
787 details: format!(
788 "a stateless evaluation was asked of a model with {} state slot(s)",
789 self.n_states
790 ),
791 },
792 ));
793 }
794 let _precision_guard = self.precision.enter();
795 let param_vec = self.build_param_vec(params)?;
796 let mut scratch = RhsScratch::new(&self.var_shapes);
797 scratch.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
798 if self.tape_serves_passes() {
799 scratch.install_tape(self.strict_tape(&HashSet::new())?);
800 scratch.set_exports_active(true);
801 }
802 let mut dy: Vec<f64> = Vec::new();
803 crate::simulate_array::take_const_array_oob();
804 for (k, &t) in times.iter().enumerate() {
805 evaluate_rhs_with_scratch(
806 &RhsCall {
807 rhs_rules: &self.rhs_rules,
808 observed_rules: &self.observed_rules,
809 var_shapes: &self.var_shapes,
810 param_names: &self.param_names,
811 state: &[],
812 params: ¶m_vec,
813 forcing: &self.forcing,
814 t,
815 declared: &self.declared_names,
816 rule_set: self.rule_set,
817 },
818 &mut dy,
819 self.is_interpreter(),
820 &mut RhsStats::default(),
821 &mut scratch,
822 );
823 if let Some(details) = crate::simulate_array::take_const_array_oob() {
824 return Err(
825 crate::compile_error::CompileError::InterpreterBuildError { details }.into(),
826 );
827 }
828 let obs = scratch
829 .taped_observeds()
830 .unwrap_or_else(|| scratch.observed_arrays());
831 visit(k, obs);
832 }
833 Ok(())
834 }
835
836 /// Resolve the deferred scoped-reference / array `ic` equations
837 /// (esm-spec §11.4.1) into per-slot initial values keyed by flat state slot.
838 /// A loaded-field RHS (`InitialConditions.O3_init`) is read from the
839 /// provider-seeded forcing buffer and folded into the lifted grid state's cells
840 /// (column-major, matching the slot enumeration in [`Self::from_model`]); a
841 /// constant RHS broadcasts to every cell. Empty on the non-`ic` path.
842 ///
843 /// `records`, when given, receives one [`FieldIcRecord`] per evaluation
844 /// this performs — each state-free observed materialized for the `ic`
845 /// scope, then each target — saying whether it was walked per cell.
846 #[cfg(feature = "solve")]
847 fn resolve_field_ics(
848 &self,
849 params: &HashMap<String, f64>,
850 mut records: Option<&mut Vec<FieldIcRecord>>,
851 ) -> Result<HashMap<usize, f64>, SimulateError> {
852 let mut out: HashMap<usize, f64> = HashMap::new();
853 if self.field_ics.is_empty() {
854 return Ok(out);
855 }
856 #[cfg(test)]
857 FIELD_IC_RESOLUTIONS.with(|c| c.set(c.get() + 1));
858 let walks = crate::simulate_array::per_cell_walks;
859 // `interpreter` runs these on the per-cell oracle like every other
860 // evaluation it performs; the other compilers leave the overlay on.
861 let _overlay = OverlayGuard::armed(self.is_interpreter());
862 // esm-spec §6.6.5 build-time scope: materialize the STATE-FREE
863 // observeds an `ic` RHS reads (a `const` gather, a parameter-only
864 // expression) over the provider forcing buffer. A provider-served field
865 // of the same name WINS — loaded data beats a document-side definition,
866 // the same direction `vi_factor_arrays` takes — and the definitions
867 // only it read are not evaluated. An observed that does not evaluate is
868 // skipped rather than raised on: it simply is not in the ic scope, and
869 // the resolver's own diagnostic then names the unusable RHS, and why.
870 //
871 // Built ONLY when the document has such definitions: with none (every
872 // document before this), the resolver reads the forcing buffer straight
873 // through and no provider field is copied.
874 let borrowed = self.forcing.borrow();
875 let mut scope: Option<HashMap<String, ArrayD<f64>>> = None;
876 // A scalar-valued scope definition (`dz = H/N`) binds where a
877 // parameter does, for the definitions after it and for the targets.
878 let mut scope_params: Option<HashMap<String, f64>> = None;
879 let mut skipped: HashMap<&str, String> = self
880 .ic_scope
881 .excluded
882 .iter()
883 .map(|(name, reason)| (name.as_str(), reason.clone()))
884 .collect();
885 if !self.ic_scope.defs.is_empty() {
886 // The cone still to evaluate: from the definitions the `ic`s name,
887 // through what each reads, stopping at a provider-served field.
888 let by_name: HashMap<&str, &super::compile::IcScopeDef> = self
889 .ic_scope
890 .defs
891 .iter()
892 .map(|d| (d.name.as_str(), d))
893 .collect();
894 let mut live: HashSet<&str> = HashSet::new();
895 let mut stack: Vec<&str> = self.ic_scope.roots.iter().map(String::as_str).collect();
896 while let Some(name) = stack.pop() {
897 if borrowed.contains_key(name) || !live.insert(name) {
898 continue;
899 }
900 if let Some(def) = by_name.get(name) {
901 stack.extend(def.deps.iter().map(String::as_str));
902 }
903 }
904 let mut built: HashMap<String, ArrayD<f64>> = borrowed.clone();
905 let mut bound: HashMap<String, f64> = params.clone();
906 // In dependency order: everything a definition reads is built (or
907 // known not to be) by the time it is reached, so one pass is enough.
908 for def in &self.ic_scope.defs {
909 if !live.contains(def.name.as_str()) {
910 continue;
911 }
912 // Never evaluated with a read missing: the overlay cannot
913 // resolve an absent name, so it would fall to a per-cell walk —
914 // which a strict compiler refuses — only to fail anyway.
915 if let Some(missing) = def
916 .needs
917 .iter()
918 .find(|n| !built.contains_key(*n) && !bound.contains_key(*n))
919 {
920 let why = match skipped.get(missing.as_str()) {
921 Some(inner) => format!("reads '{missing}', which {inner}"),
922 None => format!("reads '{missing}', which is not in the build-time scope"),
923 };
924 skipped.insert(&def.name, why);
925 continue;
926 }
927 let before = walks();
928 let value =
929 eval_buildtime_field_in_scope(&def.body, &self.index_sets, &bound, &built);
930 if per_cell_walk_refused() {
931 // A strict caller's walk stopped at its first cell
932 // (`StopAtFirstCell`): this evaluation is the refusal,
933 // and nothing computed from its placeholder counts —
934 // not even a failure. Only construction arms the stop.
935 if let Some(sink) = records.as_deref_mut() {
936 sink.push(FieldIcRecord {
937 name: def.name.clone(),
938 kind: "initial-condition scope",
939 per_cell: true,
940 });
941 }
942 return Ok(out);
943 }
944 match value {
945 Ok(Value::Array(arr)) => {
946 built.insert(def.name.clone(), *arr);
947 }
948 Ok(Value::Scalar(v)) => {
949 bound.insert(def.name.clone(), v);
950 }
951 Err(e) => {
952 skipped.insert(&def.name, format!("does not evaluate at build time: {e}"));
953 continue;
954 }
955 }
956 if let Some(sink) = records.as_deref_mut() {
957 sink.push(FieldIcRecord {
958 name: def.name.clone(),
959 kind: "initial-condition scope",
960 per_cell: walks() != before,
961 });
962 }
963 }
964 scope = Some(built);
965 scope_params = Some(bound);
966 }
967 let forcing: &HashMap<String, ArrayD<f64>> = scope.as_ref().unwrap_or(&borrowed);
968 let params: &HashMap<String, f64> = scope_params.as_ref().unwrap_or(params);
969 for (target, rhs) in &self.field_ics {
970 let vs = self.var_shapes.get(target).ok_or_else(|| {
971 SimulateError::InvalidFieldInitialCondition {
972 name: target.clone(),
973 details: "scoped-reference target is not a state variable of the flattened \
974 system"
975 .to_string(),
976 }
977 })?;
978 let total = vs.shape.iter().copied().product::<usize>().max(1);
979 // Coordinate-expression ICs (case 3 in `resolve_field_ic_cell`)
980 // evaluate the WHOLE field with one `eval_buildtime_field` call and
981 // then read a single cell — so recomputing it per cell was O(cells)
982 // full-field evaluations for an O(cells)-sized result. Resolve it
983 // once per target and let every cell index the cached field. The
984 // cell-independent cases (1 loaded field / 2 constant) ignore it.
985 let mut cached_field: Option<Value> = None;
986 let before = walks();
987 for flat in 0..total {
988 let multi = flat_to_multi_col_major(flat, &vs.shape);
989 let slot = vs.flat_offset + flat;
990 let value = resolve_field_ic_cell(
991 target,
992 rhs,
993 &multi,
994 forcing,
995 &self.index_sets,
996 params,
997 &mut cached_field,
998 );
999 if per_cell_walk_refused() {
1000 // As for the scope above: the stopped walk is the refusal.
1001 if let Some(sink) = records.as_deref_mut() {
1002 sink.push(FieldIcRecord {
1003 name: target.clone(),
1004 kind: "initial condition",
1005 per_cell: true,
1006 });
1007 }
1008 return Ok(out);
1009 }
1010 out.insert(
1011 slot,
1012 value.map_err(|e| explain_ic_scope_miss(e, rhs, &skipped))?,
1013 );
1014 }
1015 if let Some(sink) = records.as_deref_mut() {
1016 sink.push(FieldIcRecord {
1017 name: target.clone(),
1018 kind: "initial condition",
1019 per_cell: walks() != before,
1020 });
1021 }
1022 }
1023 Ok(out)
1024 }
1025
1026 /// The initial state before any override: each variable's default (one
1027 /// fill, or one gather of its inline data), then the field `ic` slots.
1028 /// The mask, when present, flags the slots neither of them set.
1029 #[cfg(feature = "solve")]
1030 fn seeded_defaults(&self, field_ic_map: &HashMap<usize, f64>) -> (Vec<f64>, Option<Vec<bool>>) {
1031 let mut ic_vec = vec![0.0f64; self.n_states];
1032 let mut unset: Option<Vec<bool>> = None;
1033 for ((_, vs), default) in self.var_shapes.iter().zip(&self.state_defaults) {
1034 let n = vs.shape.iter().copied().product::<usize>().max(1);
1035 let range = vs.flat_offset..vs.flat_offset + n;
1036 if !super::layout::write_state_default(vs, default, &mut ic_vec[range.clone()]) {
1037 unset.get_or_insert_with(|| vec![false; self.n_states])[range].fill(true);
1038 }
1039 }
1040 for (&slot, &v) in field_ic_map {
1041 ic_vec[slot] = v;
1042 if let Some(u) = unset.as_mut() {
1043 u[slot] = false;
1044 }
1045 }
1046 (ic_vec, unset)
1047 }
1048
1049 /// Evaluate the `faq`-valued initialization equations (esm-spec §6.2:
1050 /// equations that hold at t = 0) into `ic_vec`, in document order: each
1051 /// assigns every cell of its ranges the value its body takes there, read
1052 /// against the initial state seeded so far (defaults, field `ic`s, the
1053 /// caller's overrides and the equations before it) and the parameters and
1054 /// forcing fields, at t = 0. A cell in `keep` — one the caller's `u0`
1055 /// names — keeps its value. Returns the slots assigned. The Julia
1056 /// reference's `_seed_faq_init_u0!`, whose order and precedence this
1057 /// follows.
1058 ///
1059 /// The whole `faq` is one build-time field evaluation, as a field `ic`'s
1060 /// coordinate expression is: the interpreter walks it cell by cell, the
1061 /// other compilers take the whole-array overlay, and a strict compiler's
1062 /// walk stops at its first cell (`records` then carries the refusal).
1063 #[cfg(feature = "solve")]
1064 pub(crate) fn seed_initialization_faqs(
1065 &self,
1066 ic_vec: &mut [f64],
1067 params: &HashMap<String, f64>,
1068 keep: &HashSet<usize>,
1069 mut records: Option<&mut Vec<FieldIcRecord>>,
1070 ) -> Result<Vec<usize>, SimulateError> {
1071 let mut assigned = Vec::new();
1072 if self.init_faqs.is_empty() {
1073 return Ok(assigned);
1074 }
1075 let walks = crate::simulate_array::per_cell_walks;
1076 let _overlay = OverlayGuard::armed(self.is_interpreter());
1077 for (target, rhs) in &self.init_faqs {
1078 let rule = format!("init({target})");
1079 let invalid = |details: String| SimulateError::InvalidFieldInitialCondition {
1080 name: rule.clone(),
1081 details,
1082 };
1083 let vs = self.var_shapes.get(target).ok_or_else(|| {
1084 invalid(format!(
1085 "the left-hand side `{target}` of a `faq` initialization equation is not a \
1086 state variable"
1087 ))
1088 })?;
1089 // The state as seeded so far, where a state array would bind, and
1090 // the forcing fields a state does not shadow.
1091 let mut scope: HashMap<String, ArrayD<f64>> = HashMap::new();
1092 for (name, s) in &self.var_shapes {
1093 let n = s.shape.iter().copied().product::<usize>().max(1);
1094 let flat = &ic_vec[s.flat_offset..s.flat_offset + n];
1095 scope.insert(name.clone(), col_major_to_arrayd(flat, &s.shape));
1096 }
1097 for (name, arr) in self.forcing.borrow().iter() {
1098 scope.entry(name.clone()).or_insert_with(|| arr.clone());
1099 }
1100 let mut resolved = rhs.clone();
1101 crate::faq::resolve_expr_ranges(&mut resolved, &self.index_sets)
1102 .map_err(|e| invalid(e.to_string()))?;
1103 let Expr::Operator(node) = &resolved else {
1104 unreachable!("an initialization faq is an operator");
1105 };
1106 let names = node.output_idx.clone().unwrap_or_default();
1107 let mut ranges: Vec<(i64, i64)> = Vec::with_capacity(names.len());
1108 for n in &names {
1109 let b = node
1110 .ranges
1111 .as_ref()
1112 .and_then(|r| r.get(n))
1113 .and_then(|s| s.bounds());
1114 let Some([lo, hi]) = b else {
1115 return Err(invalid(format!("output index `{n}` has no static range")));
1116 };
1117 ranges.push((lo, hi));
1118 }
1119 if ranges.len() != vs.shape.len() {
1120 return Err(invalid(format!(
1121 "{} output indices for a {}-D state",
1122 ranges.len(),
1123 vs.shape.len()
1124 )));
1125 }
1126 let before = walks();
1127 let value = eval_buildtime_field_in_scope(&resolved, &self.index_sets, params, &scope);
1128 if per_cell_walk_refused() {
1129 if let Some(sink) = records.as_deref_mut() {
1130 sink.push(FieldIcRecord {
1131 name: rule,
1132 kind: "initialization equation",
1133 per_cell: true,
1134 });
1135 }
1136 return Ok(assigned);
1137 }
1138 if let Some(sink) = records.as_deref_mut() {
1139 sink.push(FieldIcRecord {
1140 name: rule.clone(),
1141 kind: "initialization equation",
1142 per_cell: walks() != before,
1143 });
1144 }
1145 let arr = match value {
1146 Ok(Value::Array(a)) if a.ndim() == ranges.len() => a,
1147 Ok(_) => return Err(invalid("the `faq` did not evaluate to its box".to_string())),
1148 Err(e) => return Err(invalid(e.to_string())),
1149 };
1150 if let Some(details) = crate::simulate_array::take_const_array_oob() {
1151 return Err(invalid(details));
1152 }
1153 // Every cell of the ranges, the box `eval_faq` fills from `lo`.
1154 let total: usize = ranges
1155 .iter()
1156 .map(|(lo, hi)| (hi - lo + 1).max(0) as usize)
1157 .product();
1158 let mut cell: Vec<i64> = ranges.iter().map(|(lo, _)| *lo).collect();
1159 for _ in 0..total {
1160 if cell
1161 .iter()
1162 .zip(&vs.origin)
1163 .zip(&vs.shape)
1164 .any(|((&c, &o), &n)| c < o || c >= o + n as i64)
1165 {
1166 return Err(invalid(format!("cell {cell:?} is outside `{target}`")));
1167 }
1168 let slot = vs.flat_offset + multi_to_flat_col_major(&cell, &vs.shape, &vs.origin);
1169 if !keep.contains(&slot) {
1170 let at: Vec<usize> = cell
1171 .iter()
1172 .zip(&ranges)
1173 .map(|(c, (lo, _))| (c - lo) as usize)
1174 .collect();
1175 ic_vec[slot] = arr[IxDyn(&at)];
1176 assigned.push(slot);
1177 }
1178 let mut d = cell.len();
1179 while d > 0 {
1180 d -= 1;
1181 cell[d] += 1;
1182 if cell[d] <= ranges[d].1 {
1183 break;
1184 }
1185 cell[d] = ranges[d].0;
1186 }
1187 }
1188 }
1189 Ok(assigned)
1190 }
1191
1192 /// Every evaluation the field initial conditions perform, exercised at
1193 /// CONSTRUCTION against `params` so a strict compiler can refuse a per-cell
1194 /// one there (esm-libraries-spec §2.5.10 asks that every evaluation a
1195 /// compiler performs for the Problem be exercised at construction) and the
1196 /// compiler report can show the route.
1197 ///
1198 /// Empty for a model with no field `ic`, and — deliberately — for one whose
1199 /// initial conditions cannot be resolved yet: `solve` resolves them again
1200 /// and raises the diagnostic there, where it always has been.
1201 ///
1202 /// `strict` is the compiler's: the first per-cell walk then stops at its
1203 /// first cell (see `StopAtFirstCell`), because its record is a refusal and
1204 /// the rest of the walk would buy nothing. A resolution that walked no cell
1205 /// is kept for the first solve ([`FieldIcMemo`]).
1206 #[cfg(feature = "solve")]
1207 pub(crate) fn field_ic_records(
1208 &self,
1209 params: &HashMap<String, f64>,
1210 strict: bool,
1211 ) -> Vec<FieldIcRecord> {
1212 if self.field_ics.is_empty() && self.init_faqs.is_empty() {
1213 return Vec::new();
1214 }
1215 let Ok(param_vec) = self.build_param_vec(params) else {
1216 return Vec::new();
1217 };
1218 let resolved: HashMap<String, f64> = self
1219 .param_names
1220 .iter()
1221 .cloned()
1222 .zip(param_vec.iter().copied())
1223 .collect();
1224 let _precision_guard = self.precision.enter();
1225 let mut records = Vec::new();
1226 let stop = strict.then(crate::simulate_array::StopAtFirstCell::arm);
1227 let result = self.resolve_field_ics(&resolved, Some(&mut records));
1228 let stopped = stop.as_ref().is_some_and(|s| s.stopped());
1229 if let Ok(slots) = &result
1230 && !stopped
1231 && !self.init_faqs.is_empty()
1232 {
1233 // The initialization equations over the defaults and these fields:
1234 // the caller's `u0` is not known yet, and only the route each one
1235 // takes is recorded here. `build_initial_state` evaluates them.
1236 let (mut ic_vec, _) = self.seeded_defaults(slots);
1237 let _ = self.seed_initialization_faqs(
1238 &mut ic_vec,
1239 &resolved,
1240 &HashSet::new(),
1241 Some(&mut records),
1242 );
1243 }
1244 if let Ok(slots) = result
1245 && !stopped
1246 {
1247 *self.field_ic_memo.borrow_mut() = Some(FieldIcMemo {
1248 params: param_bits(¶m_vec),
1249 forcing_generation: self.forcing_generation.get(),
1250 slots,
1251 });
1252 }
1253 records
1254 }
1255
1256 /// The component / subsystem names a rule-2 override key may spell in its
1257 /// LEADING segments (esm-spec §6.6.2, §4.6).
1258 ///
1259 /// The namespace segments the build's own names carry cover a mounted
1260 /// subsystem (`sub` in `sub.g`) and, on the `from_flattened` path, every
1261 /// contributing component (`Left` in `Left.gain`). [`Self::namespace`]
1262 /// supplies the one namespace the names CANNOT show: the enclosing model's
1263 /// own, which the single-model path does not qualify its variables with —
1264 /// it is exactly what makes `P.sub.g` a legal spelling of `sub.g`.
1265 fn override_namespaces(&self) -> std::collections::HashSet<String> {
1266 crate::simulate::namespace_scope(
1267 // A slot name's cell suffix carries no `.`, so the state
1268 // VARIABLES' names carry every namespace their slots' names do.
1269 self.param_names
1270 .iter()
1271 .chain(self.var_shapes.keys())
1272 .map(String::as_str),
1273 self.namespace.as_deref(),
1274 )
1275 }
1276
1277 /// The scalar-table parameters neither `params` (the caller's `p`,
1278 /// canonicalized as [`Self::build_param_vec`] does) nor a declared default
1279 /// gives a value. Empty when a key does not canonicalize: that is the
1280 /// solve's diagnostic, not this one.
1281 pub(crate) fn unsupplied_params(&self, params: &HashMap<String, f64>) -> Vec<String> {
1282 let Ok(params) = crate::simulate::canonicalize_override_keys(
1283 &self.param_index,
1284 &self.override_namespaces(),
1285 params,
1286 &self.merged_renames,
1287 ) else {
1288 return Vec::new();
1289 };
1290 self.param_names
1291 .iter()
1292 .zip(&self.param_defaults)
1293 .filter(|(name, d)| d.is_none() && !params.contains_key(*name))
1294 .map(|(name, _)| name.clone())
1295 .collect()
1296 }
1297
1298 /// The states none of whose slots anything gives a starting value — no
1299 /// default, no field `ic`, no initialization equation, no entry of
1300 /// `initial_conditions` (the caller's `u0`) — by name. Read at construction, from the field `ic`s
1301 /// construction resolved ([`FieldIcMemo`], left in place for the first
1302 /// solve); empty when those could not be resolved there, or a key does not
1303 /// canonicalize, which leaves the diagnostic to the solve.
1304 #[cfg(feature = "solve")]
1305 pub(crate) fn unset_initial_slots(
1306 &self,
1307 initial_conditions: &HashMap<String, f64>,
1308 ) -> Vec<String> {
1309 let Ok(initial_conditions) = crate::simulate::canonicalize_override_keys(
1310 &super::layout::SlotNames(&self.var_shapes),
1311 &self.override_namespaces(),
1312 initial_conditions,
1313 &self.merged_renames,
1314 ) else {
1315 return Vec::new();
1316 };
1317 let memo_slots;
1318 let empty = HashMap::new();
1319 let field_ic_map = if self.field_ics.is_empty() {
1320 &empty
1321 } else {
1322 match self.field_ic_memo.borrow().as_ref() {
1323 Some(m) => {
1324 memo_slots = m.slots.clone();
1325 &memo_slots
1326 }
1327 None => return Vec::new(),
1328 }
1329 };
1330 let (_, unset) = self.seeded_defaults(field_ic_map);
1331 let Some(mut unset) = unset else {
1332 return Vec::new();
1333 };
1334 for name in initial_conditions.keys() {
1335 if let Some(slot) = super::layout::lookup_slot(&self.var_shapes, name) {
1336 unset[slot] = false;
1337 }
1338 }
1339 // An initialization equation assigns its whole target, and a held
1340 // state (no `D`, no definition) is not an ODE state.
1341 for target in self
1342 .init_faqs
1343 .iter()
1344 .map(|(t, _)| t)
1345 .chain(&self.held_at_ic)
1346 {
1347 if let Some(vs) = self.var_shapes.get(target) {
1348 let n = vs.shape.iter().copied().product::<usize>().max(1);
1349 unset[vs.flat_offset..vs.flat_offset + n].fill(false);
1350 }
1351 }
1352 // A state is refused only when NONE of its slots got a value: cells left
1353 // over in a partly seeded one are layout (a grid widened past the cells
1354 // the document names) and keep their placeholder.
1355 self.var_shapes
1356 .iter()
1357 .filter(|(_, vs)| {
1358 let n = vs.shape.iter().copied().product::<usize>().max(1);
1359 unset[vs.flat_offset..vs.flat_offset + n].iter().all(|&u| u)
1360 })
1361 .map(|(name, _)| name.clone())
1362 .collect()
1363 }
1364
1365 /// Run the simulation.
1366 /// Validate override parameter names and build the positional param
1367 /// vector (override > variable default; a parameter with neither is an
1368 /// [`SimulateError::InvalidParameter`]). The strict simulate-time
1369 /// counterpart of the lenient [`Self::debug_resolve_params`].
1370 fn build_param_vec(&self, params: &HashMap<String, f64>) -> Result<Vec<f64>, SimulateError> {
1371 // esm-spec §6.6.2 caller-key canonicalization (see
1372 // `crate::simulate::canonicalize_override_keys`). This subsumes the
1373 // former `normalize_override_keys`, which only stripped this model's
1374 // `<namespace>.` prefix: rule 2 resolves `M.A` against a bare-named
1375 // single-model system, and rule 3 resolves the LOCAL `A` against a
1376 // flattening-qualified `M.A` — which the prefix strip could not do.
1377 let params = crate::simulate::canonicalize_override_keys(
1378 &self.param_index,
1379 &self.override_namespaces(),
1380 params,
1381 &self.merged_renames,
1382 )
1383 .map_err(crate::simulate::param_key_error)?;
1384 let mut param_vec = vec![0.0f64; self.param_names.len()];
1385 for (i, name) in self.param_names.iter().enumerate() {
1386 if let Some(&v) = params.get(name) {
1387 param_vec[i] = v;
1388 } else if let Some(d) = self.param_defaults[i] {
1389 param_vec[i] = d;
1390 } else {
1391 return Err(SimulateError::InvalidParameter { name: name.clone() });
1392 }
1393 }
1394 Ok(param_vec)
1395 }
1396
1397 /// Validate override initial-condition names and build the initial state
1398 /// vector `u0`. Scoped-reference / array `ic` fields (esm-spec §11.4.1)
1399 /// are folded in from the provider-seeded forcing buffer (DESIGN
1400 /// pde_simulation_pipeline §2 R2); priority per slot: explicit
1401 /// `initial_conditions` override > loaded field ic > variable default (a
1402 /// slot with none of the three is an
1403 /// [`SimulateError::InvalidInitialCondition`]).
1404 #[cfg(feature = "solve")]
1405 fn build_initial_state(
1406 &self,
1407 initial_conditions: &HashMap<String, f64>,
1408 param_vec: &[f64],
1409 ) -> Result<Vec<f64>, SimulateError> {
1410 // Same §6.6.2 canonicalization as `build_param_vec`, on the state side.
1411 let initial_conditions = crate::simulate::canonicalize_override_keys(
1412 &super::layout::SlotNames(&self.var_shapes),
1413 &self.override_namespaces(),
1414 initial_conditions,
1415 &self.merged_renames,
1416 )
1417 .map_err(crate::simulate::ic_key_error)?;
1418 // What construction resolved, when it resolved it under these
1419 // parameters ([`FieldIcMemo`]).
1420 let memo = self.field_ic_memo.borrow_mut().take().filter(|m| {
1421 m.params == param_bits(param_vec)
1422 && m.forcing_generation == self.forcing_generation.get()
1423 });
1424 // Resolved scalar-parameter scope (load-time constants) for the ic
1425 // coordinate-expression path — a parameter-dependent grid-geometry
1426 // template (`x0 + (i − 1/2)·dx`) binds here — and for the
1427 // initialization equations.
1428 let resolved_params: HashMap<String, f64> = self
1429 .param_names
1430 .iter()
1431 .cloned()
1432 .zip(param_vec.iter().copied())
1433 .collect();
1434 let field_ic_map = match memo {
1435 Some(m) => m.slots,
1436 None => self.resolve_field_ics(&resolved_params, None)?,
1437 };
1438 // Per slot, the first of: an explicit override, an initialization
1439 // equation, a field `ic`, the variable's default. Written lowest
1440 // priority first; a slot left with none of them is flagged, and only a
1441 // variable that declares no default can leave one.
1442 let (mut ic_vec, mut unset) = self.seeded_defaults(&field_ic_map);
1443 let mut overridden: HashSet<usize> = HashSet::new();
1444 for (name, &v) in &initial_conditions {
1445 // Canonicalization answers only with names the layout resolves.
1446 let slot = super::layout::lookup_slot(&self.var_shapes, name)
1447 .expect("a canonical initial-condition key names a slot");
1448 ic_vec[slot] = v;
1449 overridden.insert(slot);
1450 if let Some(u) = unset.as_mut() {
1451 u[slot] = false;
1452 }
1453 }
1454 // After the overrides, so a body reading a state reads the caller's
1455 // value, and around them, so the caller's value stands.
1456 for slot in
1457 self.seed_initialization_faqs(&mut ic_vec, &resolved_params, &overridden, None)?
1458 {
1459 if let Some(u) = unset.as_mut() {
1460 u[slot] = false;
1461 }
1462 }
1463 if let Some(slot) = unset.and_then(|u| u.iter().position(|&x| x)) {
1464 let name = super::layout::slot_name(&self.var_shapes, slot).unwrap_or_default();
1465 return Err(SimulateError::InvalidInitialCondition { name });
1466 }
1467 Ok(ic_vec)
1468 }
1469
1470 #[cfg(feature = "solve")]
1471 pub fn solve(
1472 &self,
1473 tspan: (f64, f64),
1474 params: &HashMap<String, f64>,
1475 initial_conditions: &HashMap<String, f64>,
1476 opts: &SolveOptions,
1477 ) -> Result<Solution, SimulateError> {
1478 self.solve_inspect(tspan, params, initial_conditions, opts, None)
1479 }
1480
1481 /// [`Self::solve`] with an optional build-observability sink (see
1482 /// [`BuildInspection`]). When `inspect` is `Some`, the sink is filled —
1483 /// after the initial state vector is assembled and before the solver runs —
1484 /// with the state-free observed arrays materialized at `u0`/`t0` and every
1485 /// observed rule's resolved body expression. The integration itself is
1486 /// byte-identical with or without a sink.
1487 #[cfg(feature = "solve")]
1488 pub fn solve_inspect(
1489 &self,
1490 tspan: (f64, f64),
1491 params: &HashMap<String, f64>,
1492 initial_conditions: &HashMap<String, f64>,
1493 opts: &SolveOptions,
1494 inspect: Option<&mut BuildInspection>,
1495 ) -> Result<Solution, SimulateError> {
1496 // Re-arm the precision this model was compiled under
1497 // (`crate::precision`); a no-op for a Float64 model.
1498 let _precision_guard = self.precision.enter();
1499 // CONST / single-segment: no discrete forcing, no refresh boundaries.
1500 self.solve_core(
1501 tspan,
1502 params,
1503 initial_conditions,
1504 opts,
1505 inspect,
1506 &HashSet::new(),
1507 &[],
1508 |_t| Ok(()),
1509 )
1510 }
1511
1512 /// [`Self::solve_inspect`] with a DISCRETE-cadence forcing refresh. The
1513 /// integration is SEGMENTED on `boundaries` (solver-second refresh anchors);
1514 /// at each boundary `refresh_fn(t)` re-slices the live forcing buffer to that
1515 /// record. Observeds transitively reaching a `discrete_forcing` name are
1516 /// excluded from the build-once static hoist, so they recompute over the
1517 /// refreshed buffer per segment while the CONST terrain regrid stays hoisted.
1518 /// This is the Rust analog of the ESS-Julia live-field taint + segmented driver.
1519 #[allow(clippy::too_many_arguments)]
1520 #[cfg(all(feature = "solve", not(target_arch = "wasm32")))]
1521 pub(crate) fn solve_with_refresh_inspect(
1522 &self,
1523 tspan: (f64, f64),
1524 params: &HashMap<String, f64>,
1525 initial_conditions: &HashMap<String, f64>,
1526 opts: &SolveOptions,
1527 inspect: Option<&mut BuildInspection>,
1528 discrete_forcing: &HashSet<String>,
1529 boundaries: &[f64],
1530 refresh_fn: impl FnMut(f64) -> Result<(), SimulateError>,
1531 ) -> Result<Solution, SimulateError> {
1532 self.solve_core(
1533 tspan,
1534 params,
1535 initial_conditions,
1536 opts,
1537 inspect,
1538 discrete_forcing,
1539 boundaries,
1540 refresh_fn,
1541 )
1542 }
1543
1544 /// Shared setup + segmented integration (see the two entry points above).
1545 /// `boundaries` are the sorted solver-second refresh anchors strictly inside
1546 /// `(t0, t_end)`; an empty `boundaries` (the CONST path) runs one segment,
1547 /// byte-identical to the un-segmented driver.
1548 #[allow(clippy::too_many_arguments)]
1549 #[cfg(feature = "solve")]
1550 fn solve_core(
1551 &self,
1552 tspan: (f64, f64),
1553 params: &HashMap<String, f64>,
1554 initial_conditions: &HashMap<String, f64>,
1555 opts: &SolveOptions,
1556 inspect: Option<&mut BuildInspection>,
1557 discrete_forcing: &HashSet<String>,
1558 boundaries: &[f64],
1559 mut refresh_fn: impl FnMut(f64) -> Result<(), SimulateError>,
1560 ) -> Result<Solution, SimulateError> {
1561 let params_owned = params;
1562 let ics_owned = initial_conditions;
1563 let (t0, t_end) = tspan;
1564 // Validated here, at the entry point every array solve shares, so the
1565 // non-advancing shortcut in `run_one_segment` and the solver path below
1566 // it are never handed a span they would read differently.
1567 crate::simulate::reject_nonfinite_span(t0, t_end)?;
1568
1569 // Validate the override names and build the positional param vector
1570 // and the initial state vector `u0` (loaded-field / coordinate
1571 // `ic`s folded in — see [`Self::build_initial_state`]).
1572 // §5.5.5: the per-cell oracle LATCHES an out-of-range const-array gather
1573 // rather than raising (the tree walk returns a bare `Value`). Clear the
1574 // latch before the solve; `assemble_solution` drains it at every exit,
1575 // so a trajectory built on a const-array bug fails loudly instead of
1576 // carrying the `NaN` the gather substituted.
1577 crate::simulate_array::take_const_array_oob();
1578 let param_vec = self.build_param_vec(params_owned)?;
1579 let ic_vec = self.build_initial_state(ics_owned, ¶m_vec)?;
1580
1581 // Seed the forcing buffer at t0 BEFORE the static hoist reads it — a
1582 // no-op for the CONST/single-segment path; for DISCRETE it primes the
1583 // first record so the static regrid geometry sees a populated buffer.
1584 refresh_fn(t0)?;
1585
1586 // The tape is built BEFORE the static hoist, not after it. Under
1587 // either strict compiler the hoist is served FROM the tape (see
1588 // [`Self::hoist_static_observeds`]), so the order the two ran in was
1589 // itself the defect: the hoist evaluated every CONST-tier observed
1590 // through the whole-array overlay, once per solve, before the tape the
1591 // caller asked for had been built at all.
1592 let tape = self.build_solve_tape(discrete_forcing)?;
1593
1594 let cadence = self.partition_observed_cadence(discrete_forcing);
1595 let setup = self.hoist_static_observeds(cadence, &ic_vec, ¶m_vec, t0, tape.as_ref());
1596 let scratches = match &tape {
1597 Some(prog) if self.tape_serves_passes() => Some(SolveScratches::new(self, prog)),
1598 _ => None,
1599 };
1600
1601 if let Some(insp) = inspect {
1602 self.fill_solve_inspection(
1603 insp,
1604 &setup,
1605 &ic_vec,
1606 ¶m_vec,
1607 t0,
1608 boundaries,
1609 tape.as_ref(),
1610 );
1611 }
1612
1613 let solver_name = match opts.alg {
1614 Alg::Bdf => "Bdf",
1615 Alg::Sdirk => "Sdirk",
1616 Alg::Erk => "Erk",
1617 };
1618
1619 // CONST / single-segment (or no output grid to align segment samples on):
1620 // the original un-segmented run — byte-identical to the pre-segmentation
1621 // driver (one `run_one_segment` over the whole span with `opts` verbatim).
1622 if boundaries.is_empty() || opts.saveat.is_none() {
1623 let (time, state, stats, retcode) = self
1624 .run_one_segment(
1625 t0,
1626 t_end,
1627 &ic_vec,
1628 ¶m_vec,
1629 &setup.static_obs,
1630 &setup.cadence.segment_static_rules,
1631 &setup.cadence.continuous_rules,
1632 opts,
1633 tape.as_ref(),
1634 scratches.as_ref(),
1635 )
1636 .map_err(Self::const_oob_first)?;
1637 return self.assemble_solution(
1638 time,
1639 state,
1640 retcode,
1641 solution_metadata(
1642 solver_name,
1643 &stats,
1644 self.merged_renames.clone(),
1645 self.namespace.clone(),
1646 ),
1647 ¶m_vec,
1648 &setup,
1649 &opts.output_observed,
1650 tape.as_ref(),
1651 );
1652 }
1653
1654 let run = SegmentedRun {
1655 t0,
1656 t_end,
1657 boundaries,
1658 ic_vec: &ic_vec,
1659 param_vec: ¶m_vec,
1660 setup: &setup,
1661 opts,
1662 tape: tape.as_ref(),
1663 scratches: scratches.as_ref(),
1664 };
1665 let (time, state, stats, retcode) = self.run_segmented(&run, &mut refresh_fn)?;
1666
1667 // Note: `assemble_solution` re-evaluates the varying observeds against
1668 // the CURRENT (last-segment) forcing buffer, so an appended scalar
1669 // observed reading a discrete forcing reflects the final hour — the
1670 // array STATE trajectory (the fire front) is per-segment correct, which
1671 // is what the runner reads.
1672 self.assemble_solution(
1673 time,
1674 state,
1675 retcode,
1676 solution_metadata(
1677 solver_name,
1678 &stats,
1679 self.merged_renames.clone(),
1680 self.namespace.clone(),
1681 ),
1682 ¶m_vec,
1683 &setup,
1684 &opts.output_observed,
1685 tape.as_ref(),
1686 )
1687 }
1688
1689 /// Three-tier cadence split of the observed rules (cadence.rs lattice
1690 /// `CONST ⊏ DISCRETE ⊏ CONTINUOUS`):
1691 ///
1692 /// ```text
1693 /// * CONST (static_rules) — materialized ONCE at setup (see
1694 /// hoist_static_observeds).
1695 /// * DISCRETE (segment_static_rules) — state-free & t-free but reaches a
1696 /// refreshed forcing buffer; constant
1697 /// WITHIN a segment, so materialized
1698 /// once per segment in run_one_segment.
1699 /// * CONTINUOUS (continuous_rules) — reaches t or state; re-evaluated
1700 /// every RHS step.
1701 /// ```
1702 ///
1703 /// Collapsing DISCRETE into CONTINUOUS (the old two-tier split) recomputed
1704 /// the per-cell conservative regrid every step — the dominant cost of a
1705 /// coupled loader model. `varying_rules` (DISCRETE ∪ CONTINUOUS) is retained
1706 /// for the non-hot observed-trajectory output pass.
1707 #[cfg(feature = "solve")]
1708 fn partition_observed_cadence(&self, discrete_forcing: &HashSet<String>) -> ObservedCadence {
1709 let static_names = self.classify_static_observeds(discrete_forcing);
1710 let seg_invariant_names = self.classify_segment_invariant_observeds(discrete_forcing, true);
1711 let static_rules: Vec<AlgebraicRule> = self
1712 .observed_rules
1713 .iter()
1714 .filter(|r| static_names.contains(observed_rule_var(r)))
1715 .cloned()
1716 .collect();
1717 let segment_static_rules: Vec<AlgebraicRule> = self
1718 .observed_rules
1719 .iter()
1720 .filter(|r| {
1721 seg_invariant_names.contains(observed_rule_var(r))
1722 && !static_names.contains(observed_rule_var(r))
1723 })
1724 .cloned()
1725 .collect();
1726 let continuous_rules: Vec<AlgebraicRule> = self
1727 .observed_rules
1728 .iter()
1729 .filter(|r| !seg_invariant_names.contains(observed_rule_var(r)))
1730 .cloned()
1731 .collect();
1732 let varying_rules: Vec<AlgebraicRule> = self
1733 .observed_rules
1734 .iter()
1735 .filter(|r| !static_names.contains(observed_rule_var(r)))
1736 .cloned()
1737 .collect();
1738 ObservedCadence {
1739 static_names,
1740 static_rules,
1741 segment_static_rules,
1742 continuous_rules,
1743 varying_rules,
1744 }
1745 }
1746
1747 /// Hoist the STATE-FREE / `t`-free observeds (ess: static-observed hoist)
1748 /// out of the per-step RHS. Within a single `simulate` call the forcing
1749 /// buffer is constant (the free `simulate` never refreshes it between
1750 /// segments), so a rule whose transitive references reach no state
1751 /// variable and no `t` is CONSTANT across the whole solve: the
1752 /// conservative-regrid geometry (`intersect_polygon` over the src×tgt cell
1753 /// rings), the regridded terrain and its slopes, the Rothermel
1754 /// coefficients derived from the CONST forcing. Materialize them ONCE here
1755 /// and seed them into every RHS eval, rather than recomputing the
1756 /// (expensive) regrid on every step. A model with no such observeds hoists
1757 /// nothing and stays byte-identical to the un-hoisted path.
1758 #[cfg(feature = "solve")]
1759 fn hoist_static_observeds(
1760 &self,
1761 cadence: ObservedCadence,
1762 ic_vec: &[f64],
1763 param_vec: &[f64],
1764 t0: f64,
1765 tape: Option<&std::sync::Arc<TapeProgram>>,
1766 ) -> SolveSetup {
1767 let sa0 = build_state_arrays(&self.var_shapes, ic_vec);
1768 // The two STRICT compilers, `native` and `xla` (API_SPEC §5.8): the
1769 // tape's own CONST section computes exactly these rules, once per
1770 // solve, on the same schedule this hoist used to. Materializing them a
1771 // SECOND time through the whole-array overlay would be the off-tape
1772 // per-cell evaluation `esm-libraries-spec.md` §2.5.10 refuses — and it would be invisible,
1773 // since the overlay declines on its own terms and reports nothing.
1774 //
1775 // What consumed the hoisted map still gets it: the RHS and Jacobian
1776 // scratches read the tape's slots rather than a seeded observed map,
1777 // and the inspection snapshot and the output-node pass harvest their
1778 // values from the tape ([`TapedObserveds`]).
1779 if self.tape_serves_passes() && tape.is_some() {
1780 return SolveSetup {
1781 cadence,
1782 sa0,
1783 static_obs: ArrMap::default(),
1784 };
1785 }
1786 let static_rings_cell: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
1787 let mut static_obs = ArrMap::default();
1788 let env = EvalEnv {
1789 state_arrays: &sa0,
1790 params: param_vec,
1791 param_names: &self.param_names,
1792 t: t0,
1793 // The regrid's FAQ rings are produced AND consumed within this
1794 // one-time static pass (its ring-consuming aggregates are themselves
1795 // static), so the sink is discarded after — no varying rule reads a
1796 // static ring, and each RHS eval starts from empty `derived_rings`.
1797 derived_rings: &static_rings_cell,
1798 derived_extents: empty_derived_extents(),
1799 forcing: &self.forcing,
1800 // One-shot materialization: no CSE memo (nothing to amortize the
1801 // structural analysis over).
1802 cse: None,
1803 const_lits: None,
1804 const_arrays: &self.const_scope,
1805 declared: &self.declared_names,
1806 };
1807 materialize_observeds_pass(
1808 &mut static_obs,
1809 &cadence.static_rules,
1810 &ObsPass {
1811 env,
1812 // `interpreter` is the reference and carries no performance
1813 // promise: every fast tier off, including the whole-array
1814 // overlay, at setup as much as on the hot path.
1815 force_scalar: self.is_interpreter(),
1816 },
1817 &mut RhsStats::default(),
1818 );
1819 drop(static_rings_cell);
1820 SolveSetup {
1821 cadence,
1822 sa0,
1823 static_obs,
1824 }
1825 }
1826
1827 /// Build observability (see `BuildInspection`): the hoisted static
1828 /// observeds ARE the build-once products (regrid geometry, regridded
1829 /// terrain, slopes). Nothing downstream consults the sink, so the
1830 /// integration is unchanged.
1831 #[cfg(feature = "solve")]
1832 #[allow(clippy::too_many_arguments)]
1833 fn fill_solve_inspection(
1834 &self,
1835 insp: &mut BuildInspection,
1836 setup: &SolveSetup,
1837 ic_vec: &[f64],
1838 param_vec: &[f64],
1839 t0: f64,
1840 boundaries: &[f64],
1841 tape: Option<&std::sync::Arc<TapeProgram>>,
1842 ) {
1843 // Under either strict compiler: one taped read-out at t0 answers BOTH
1844 // halves of this sink — the static observeds (which the hoist no
1845 // longer materializes) and, on a segmented run, the varying ones. The
1846 // off-tape snapshot below is what §2.5.10 would otherwise leave
1847 // un-gated.
1848 if self.tape_serves_passes()
1849 && let Some(tape) = tape
1850 {
1851 let mut harvest = TapedObserveds::new(self, tape);
1852 let obs = harvest.at(self, ic_vec, param_vec, t0);
1853 self.fill_inspection(insp, obs, &setup.cadence.static_names, param_vec);
1854 if !boundaries.is_empty() {
1855 for rule in &setup.cadence.varying_rules {
1856 let name = observed_rule_var(rule);
1857 if let Some(a) = obs.get(name) {
1858 insp.setup_arrays.insert(name.clone(), a.clone());
1859 }
1860 }
1861 }
1862 return;
1863 }
1864 self.fill_inspection(
1865 insp,
1866 &setup.static_obs,
1867 &setup.cadence.static_names,
1868 param_vec,
1869 );
1870 // Segmented (DISCRETE) run: the time-varying regrid observeds (the
1871 // ERA5 t_xy/rh_xy/u_xy/v_xy over the first hour's slice) are NOT in
1872 // the static hoist, so ALSO snapshot them at t0 into `setup_arrays` —
1873 // a caller reading the build-time per-cell forcing (the runner's
1874 // forcing print) then still sees the ERA5 fields at their t=0 record.
1875 if !boundaries.is_empty() {
1876 let dr: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
1877 let mut snapshot = setup.static_obs.clone();
1878 materialize_observeds_pass(
1879 &mut snapshot,
1880 &setup.cadence.varying_rules,
1881 &ObsPass {
1882 env: EvalEnv {
1883 state_arrays: &setup.sa0,
1884 params: param_vec,
1885 param_names: &self.param_names,
1886 t: t0,
1887 derived_rings: &dr,
1888 derived_extents: empty_derived_extents(),
1889 forcing: &self.forcing,
1890 cse: None,
1891 const_lits: None,
1892 const_arrays: &self.const_scope,
1893 declared: &self.declared_names,
1894 },
1895 // Build-time t0 snapshot: vectorized overlay
1896 // (bit-identical), unless this is the `interpreter`
1897 // reference, which runs every tier off.
1898 force_scalar: self.is_interpreter(),
1899 },
1900 &mut RhsStats::default(),
1901 );
1902 for rule in &setup.cadence.varying_rules {
1903 let name = observed_rule_var(rule);
1904 if let Some(a) = snapshot.get(name) {
1905 insp.setup_arrays.insert(name.clone(), a.clone());
1906 }
1907 }
1908 }
1909 }
1910
1911 /// Step 3b: compile the tape program ONCE per solve and share it across
1912 /// every integration segment (each segment's fresh RHS scratch gets its
1913 /// own slab and re-runs the CONST/SEGMENT sections — the same cadence as
1914 /// the static-observed hoist above). [`crate::Compiler::Interpreter`]
1915 /// builds no tape at all — it IS the per-cell oracle.
1916 ///
1917 /// The tape is mandatory under every routing but the `interpreter`: a
1918 /// build that declined a rule is a refusal ([`Self::strict_tape`], naming
1919 /// the rule and the reason), never a solve of the whole model on another
1920 /// evaluator. That holds for the historical routing (`compile_array`) as
1921 /// much as for the strict compilers — its right-hand side is the tape,
1922 /// and a model whose tape declined would otherwise run every rule per
1923 /// cell with nothing but the runtime to show for it.
1924 #[cfg(feature = "solve")]
1925 fn build_solve_tape(
1926 &self,
1927 discrete_forcing: &HashSet<String>,
1928 ) -> Result<SolveTape, SimulateError> {
1929 // `interpreter` (API_SPEC §5.8) is the reference and nothing else: no
1930 // tape, and the whole-array overlay off under it, so every rule is
1931 // walked per cell. It is the compiler the caller named, and the only
1932 // way to reach the oracle.
1933 if self.is_interpreter() {
1934 return Ok(None);
1935 }
1936 Ok(Some(self.strict_tape(discrete_forcing)?))
1937 }
1938
1939 /// DISCRETE: integrate in segments split on the refresh boundaries. Segment
1940 /// endpoints = t0, each boundary strictly inside (t0, t_end) ascending, t_end.
1941 #[cfg(feature = "solve")]
1942 fn run_segmented(
1943 &self,
1944 run: &SegmentedRun<'_>,
1945 refresh_fn: &mut impl FnMut(f64) -> Result<(), SimulateError>,
1946 ) -> Result<(Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode), SimulateError> {
1947 let &SegmentedRun {
1948 t0,
1949 t_end,
1950 boundaries,
1951 ic_vec,
1952 param_vec,
1953 setup,
1954 opts,
1955 tape,
1956 scratches,
1957 } = run;
1958 let n_states = self.n_states;
1959
1960 let mut endpoints: Vec<f64> = vec![t0];
1961 for &b in boundaries {
1962 if b > t0 && b < t_end && *endpoints.last().unwrap() < b {
1963 endpoints.push(b);
1964 }
1965 }
1966 if *endpoints.last().unwrap() < t_end {
1967 endpoints.push(t_end);
1968 }
1969
1970 let global_out = opts.saveat.clone().expect("output grid checked Some");
1971 let mut u0 = ic_vec.to_vec();
1972 let mut time: Vec<f64> = Vec::new();
1973 let mut state: Vec<Vec<f64>> = vec![Vec::new(); n_states];
1974 // Each segment integrates with its own fresh diffsol solver, so the
1975 // reported step/eval counts are the sum across all segments.
1976 let mut stats = SolveStats::default();
1977 let mut retcode = ReturnCode::Success;
1978
1979 for w in endpoints.windows(2) {
1980 let (a, b) = (w[0], w[1]);
1981 // Refresh the live buffer at the START of every segment after the
1982 // first (t0 was already primed by `refresh_fn(t0)` in `solve_core`).
1983 if a != t0 {
1984 refresh_fn(a)?;
1985 if let Some(s) = scratches {
1986 s.forcing_refreshed();
1987 }
1988 }
1989 // Requested outputs falling in this segment: (a, b] — or [a, b] for
1990 // the first. Always run the solver's grid up to `b` (append if
1991 // absent) so the state at `b` seeds the next segment.
1992 let requested: Vec<f64> = global_out
1993 .iter()
1994 .copied()
1995 .filter(|&g| (if a == t0 { g >= a } else { g > a }) && g <= b)
1996 .collect();
1997 let mut grid = requested.clone();
1998 if grid.last() != Some(&b) {
1999 grid.push(b);
2000 }
2001 // Report progress against the GLOBAL interval, not the segment's.
2002 // Each segment runs a fresh solver over [a, b] with its own step
2003 // counter, so handing the caller's observer through unwrapped would
2004 // restart the bar from 0% at every refresh boundary.
2005 let seg_progress: Option<ProgressFn> = opts.progress.as_ref().map(|user| {
2006 let user = user.clone();
2007 let steps_before = stats.n_accepted_steps;
2008 let wrapped: ProgressFn = std::sync::Arc::new(move |p: &Progress<'_>| {
2009 user(&Progress {
2010 t0,
2011 t: p.t,
2012 t_end,
2013 step: steps_before + p.step,
2014 maxiters: p.maxiters,
2015 u: p.u,
2016 })
2017 });
2018 wrapped
2019 });
2020 let seg_opts = SolveOptions {
2021 saveat: Some(grid),
2022 progress: seg_progress,
2023 ..opts.clone()
2024 };
2025 let (seg_time, seg_state, seg_stats, seg_retcode) = self
2026 .run_one_segment(
2027 a,
2028 b,
2029 &u0,
2030 param_vec,
2031 &setup.static_obs,
2032 &setup.cadence.segment_static_rules,
2033 &setup.cadence.continuous_rules,
2034 &seg_opts,
2035 tape,
2036 scratches,
2037 )
2038 .map_err(Self::const_oob_first)?;
2039 stats += seg_stats;
2040 // A segment that stopped early ends the whole run: the state at its
2041 // right endpoint never materialised, so there is nothing to seed the
2042 // next segment with. Its reason becomes the run's `retcode`.
2043 if !seg_retcode.is_success() {
2044 retcode = seg_retcode;
2045 }
2046 // `run_solver` pushes the REQUESTED grid time verbatim, so a float
2047 // equality against `requested`/`b` is exact.
2048 for (i, &t) in seg_time.iter().enumerate() {
2049 if t == b {
2050 u0 = (0..n_states).map(|r| seg_state[r][i]).collect();
2051 }
2052 if requested.contains(&t) {
2053 time.push(t);
2054 for r in 0..n_states {
2055 state[r].push(seg_state[r][i]);
2056 }
2057 }
2058 }
2059 if !retcode.is_success() {
2060 break;
2061 }
2062 }
2063
2064 Ok((time, state, stats, retcode))
2065 }
2066
2067 /// Solution assembly, shared by the single-segment and segmented exits:
2068 /// expose observed trajectories alongside the states (see
2069 /// [`Self::append_observed_trajectories`]), then drain the §5.5.5
2070 /// const-array OOB latch so a trajectory built on a const-array bug fails
2071 /// loudly instead of carrying the `NaN` the gather substituted.
2072 #[cfg(feature = "solve")]
2073 #[allow(clippy::too_many_arguments)]
2074 fn assemble_solution(
2075 &self,
2076 time: Vec<f64>,
2077 mut state: Vec<Vec<f64>>,
2078 retcode: ReturnCode,
2079 metadata: SolutionMetadata,
2080 param_vec: &[f64],
2081 setup: &SolveSetup,
2082 output_observed: &[String],
2083 tape: Option<&std::sync::Arc<TapeProgram>>,
2084 ) -> Result<Solution, SimulateError> {
2085 let mut state_variable_names = self.state_variable_names().to_vec();
2086 self.append_observed_trajectories(
2087 &time,
2088 &mut state,
2089 &mut state_variable_names,
2090 param_vec,
2091 &setup.static_obs,
2092 &setup.cadence.static_names,
2093 &setup.cadence.varying_rules,
2094 output_observed,
2095 tape,
2096 );
2097 if let Some(details) = crate::simulate_array::take_const_array_oob() {
2098 return Err(
2099 crate::compile_error::CompileError::InterpreterBuildError { details }.into(),
2100 );
2101 }
2102 Ok(Solution {
2103 time,
2104 state,
2105 state_variable_names,
2106 retcode,
2107 metadata,
2108 })
2109 }
2110
2111 /// A solve that fails after the RHS latched an out-of-range const-array gather
2112 /// failed BECAUSE of it: the gather substituted `NaN`, and the integrator's
2113 /// step-size or error-test failure is the symptom. Report the latched
2114 /// `E_TREEWALK_CONSTARRAY_OOB` instead, as `assemble_solution` does on success.
2115 #[cfg(feature = "solve")]
2116 fn const_oob_first(err: SimulateError) -> SimulateError {
2117 match crate::simulate_array::take_const_array_oob() {
2118 Some(details) => {
2119 crate::compile_error::CompileError::InterpreterBuildError { details }.into()
2120 }
2121 None => err,
2122 }
2123 }
2124
2125 /// Integrate ONE segment `[t0, t_end]` from initial state `u0`, reading the
2126 /// live forcing buffer (`self.forcing`) — which a segmented driver refreshes
2127 /// between segments. Builds a fresh RHS/Jacobian closure pair (each scratch
2128 /// pre-seeded with the already-materialized `static_obs`) and a fresh diffsol
2129 /// problem, returning the states at `opts.saveat`.
2130 #[allow(clippy::too_many_arguments)]
2131 #[cfg(feature = "solve")]
2132 fn run_one_segment(
2133 &self,
2134 t0: f64,
2135 t_end: f64,
2136 u0: &[f64],
2137 param_vec: &[f64],
2138 static_obs: &ArrMap,
2139 segment_static_rules: &[AlgebraicRule],
2140 continuous_rules: &[AlgebraicRule],
2141 opts: &SolveOptions,
2142 // Step 3b: the solve-wide tape program + the FULL observed rule list
2143 // its fallback indices resolve against. `None` ⇒ legacy interpreter.
2144 tape: Option<&std::sync::Arc<TapeProgram>>,
2145 // The solve's kept scratches, or `None` for a pair built for this
2146 // segment alone (see [`SolveScratches`]).
2147 scratches: Option<&SolveScratches>,
2148 ) -> Result<(Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode), SimulateError> {
2149 // A segment that never advances is answered from its own initial state,
2150 // before any closure, scratch buffer or diffsol problem is built
2151 // (issue #438): building the solver materializes a dense Jacobian,
2152 // which the matrix-free finite-difference Jacobian below pays for in
2153 // full right-hand-side evaluations, one pair per state column, each one
2154 // materializing every varying observed — recurrence sweeps included —
2155 // for a trajectory that is the untouched initial state. See
2156 // [`crate::simulate::nonadvancing_trajectory`], which also makes the
2157 // step-0 progress report. [`crate::simulate::run_solver`] keeps the same
2158 // check as a backstop and produces the identical trajectory, so this is
2159 // a cost-only shortcut.
2160 if let Some((time, state, retcode)) =
2161 crate::simulate::nonadvancing_trajectory(t0, t_end, u0, opts)
2162 {
2163 return Ok((time, state, SolveStats::default(), retcode));
2164 }
2165 let n_states = self.n_states;
2166 let rhs_rules = self.rhs_rules.clone();
2167 let var_shapes = self.var_shapes.clone();
2168 let param_names = self.param_names.clone();
2169
2170 let rhs_rules_jac = rhs_rules.clone();
2171 let varying_rules_rhs = continuous_rules.to_vec();
2172 let varying_rules_jac = continuous_rules.to_vec();
2173 let var_shapes_jac = var_shapes.clone();
2174 let param_names_jac = param_names.clone();
2175 // See [`EvalCtx::declared`] (issue #181): the declared-name set is a
2176 // property of the compiled MODEL, so each RHS/Jacobian closure carries
2177 // its own clone rather than borrowing `self`.
2178 let declared = self.declared_names.clone();
2179 let declared_jac = declared.clone();
2180 // The closures own their rule clones, so each clone is its own rule set.
2181 let rule_set_rhs = next_rule_set_id();
2182 let rule_set_jac = next_rule_set_id();
2183
2184 // Materialize the DISCRETE (segment-invariant) observeds ONCE for this
2185 // segment, on top of the CONST `static_obs`. The caller refreshed the
2186 // forcing buffer at this segment's start (`refresh_fn(a)` / the `t0`
2187 // prime), and it stays fixed within the segment, so these forcing-derived
2188 // but state-free / `t`-free observeds (the per-cell conservative regrid)
2189 // are constant here — evaluating them once and seeding them as "static"
2190 // for the segment removes the per-step recompute that dominated the
2191 // coupled-loader profile. They are state-free, so `sa_seg` is only a
2192 // consistency placeholder; their FAQ rings are produced-and-consumed in
2193 // this one pass (own transient registry, discarded after).
2194 let seg_seed: ArrMap = if segment_static_rules.is_empty() || self.tape_serves_passes() {
2195 // Under either strict compiler: the tape's SEGMENT section
2196 // computes the segment-invariant observeds itself, on the same
2197 // once-per-segment schedule, so seeding them here would be the same rules evaluated
2198 // a second time off the tape (§2.5.10). `static_obs` is empty
2199 // under them for the same reason — see `hoist_static_observeds`.
2200 static_obs.clone()
2201 } else {
2202 let sa_seg = build_state_arrays(&self.var_shapes, u0);
2203 let seg_rings: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
2204 let mut seed = static_obs.clone();
2205 materialize_observeds_pass(
2206 &mut seed,
2207 segment_static_rules,
2208 &ObsPass {
2209 env: EvalEnv {
2210 state_arrays: &sa_seg,
2211 params: param_vec,
2212 param_names: &self.param_names,
2213 t: t0,
2214 derived_rings: &seg_rings,
2215 derived_extents: empty_derived_extents(),
2216 forcing: &self.forcing,
2217 cse: None,
2218 const_lits: None,
2219 const_arrays: &self.const_scope,
2220 declared: &self.declared_names,
2221 },
2222 force_scalar: false,
2223 },
2224 &mut RhsStats::default(),
2225 );
2226 seed
2227 };
2228
2229 // `xla` (API_SPEC §5.8): the right-hand side is the XLA executable the
2230 // build installed, not the tape's own interpreter. The FD Jacobian
2231 // goes through it too — it differences two right-hand-side evaluations
2232 // of the same rules, so leaving it on the tape would mean an implicit
2233 // solve ran the whole model on a DIFFERENT compiler from the one the
2234 // caller named, once per Jacobian call.
2235 //
2236 // Both closures hold the ONE executable (`Rc`); nothing recompiles per
2237 // segment or per call.
2238 //
2239 // A missing executable under `xla` is a REFUSAL, not a quiet return to
2240 // the tape. Without this the one way the two could come apart — a
2241 // model that reached the integrator with `RuntimeMode::Xla` and an
2242 // empty cell — would be answered by the tape under the name `xla`,
2243 // which is exactly the substitution §2.5.10 exists to prevent.
2244 #[cfg(feature = "xla")]
2245 let xla_rhs = match (self.is_xla(), self.xla_program()) {
2246 (false, _) => None,
2247 (true, Some(program)) => Some(program),
2248 (true, None) => return Err(xla_not_installed()),
2249 };
2250 #[cfg(not(feature = "xla"))]
2251 if self.is_xla() {
2252 // Unreachable: a build without the feature answers
2253 // `compiler_unavailable` for `xla` before a backend exists. Kept
2254 // so the no-fallback property holds by CODE rather than by that
2255 // argument.
2256 return Err(xla_not_installed());
2257 }
2258 // Per-closure reusable scratch (ess-mro), pre-seeded ONCE with the
2259 // CONST + per-segment DISCRETE observeds (retained in place across steps,
2260 // never re-cloned) so each RHS eval materializes only the CONTINUOUS
2261 // observeds. `RefCell` gives the interior mutability diffsol's `Fn` RHS
2262 // requires; the Jacobian closure carries its own so the two never alias.
2263 // Under `xla` the tape never evaluates the right-hand side, so this
2264 // scratch (state-array buffers plus a copy of the seeded observed map)
2265 // is not built at all.
2266 #[cfg(feature = "xla")]
2267 let tape_rhs = xla_rhs.is_none();
2268 #[cfg(not(feature = "xla"))]
2269 let tape_rhs = true;
2270 let seg_seed = Rc::new(seg_seed);
2271 let (rhs_scratch, jac_scratch) = match scratches {
2272 Some(s) => (Rc::clone(&s.rhs), Rc::clone(&s.jac)),
2273 None => (Rc::new(RefCell::new(None)), Rc::new(RefCell::new(None))),
2274 };
2275 // `native`: both closures evaluate through the solve's program.
2276 let program_rhs = scratches.and_then(|s| s.program.clone());
2277 let program_jac = program_rhs.clone();
2278 if tape_rhs && program_rhs.is_none() && rhs_scratch.borrow().is_none() {
2279 let mut s = RhsScratch::new(&var_shapes);
2280 s.set_const_arrays(std::sync::Arc::clone(&self.const_scope));
2281 s.set_static((*seg_seed).clone());
2282 // Step 3b: the production RHS closure's scratch gets the compiled
2283 // tape. Its CONST/SEGMENT sections prime on the first call, and a
2284 // kept scratch re-runs SEGMENT after each forcing refresh.
2285 if let Some(prog) = tape {
2286 s.install_tape(std::sync::Arc::clone(prog));
2287 }
2288 *rhs_scratch.borrow_mut() = Some(s);
2289 }
2290 // The Jacobian scratch is built LAZILY on the first Jacobian call:
2291 // diffsol's `rhs_implicit` builder demands a Jacobian closure even for
2292 // the explicit (ERK) solver, which then never invokes it — so an eager
2293 // scratch (state-array buffers + a full copy of the seeded observed map)
2294 // was memory spent on a closure that never runs. The implicit solvers
2295 // (BDF/SDIRK) build it on their first Jacobian evaluation instead;
2296 // construction is deterministic, so results are bit-identical either way.
2297 let jac_seed = Rc::clone(&seg_seed);
2298 let const_scope_jac = std::sync::Arc::clone(&self.const_scope);
2299 let tape_jac: Option<std::sync::Arc<TapeProgram>> = match (self.tape_serves_passes(), tape)
2300 {
2301 (true, Some(prog)) => Some(std::sync::Arc::clone(prog)),
2302 _ => None,
2303 };
2304 // `interpreter`: the per-cell oracle for the right-hand side too, not
2305 // just for the observeds. The strict compilers and the legacy routing
2306 // pass `false` and take the whole-array overlay where the tape is
2307 // absent.
2308 let force_scalar = self.is_interpreter();
2309
2310 // The XLA Jacobian's buffers and kept base-point evaluation (see
2311 // [`FdJvp`]), built on its first call for the same reason the tape's
2312 // Jacobian scratch is.
2313 #[cfg(feature = "xla")]
2314 let xla_jac: Option<(Rc<crate::xla_runtime::CompiledRhs>, RefCell<Option<FdJvp>>)> =
2315 xla_rhs.clone().map(|program| (program, RefCell::new(None)));
2316 // Where an XLA execution failure goes. diffsol's right-hand side is
2317 // `Fn(..) -> ()`, so a device failure has no return channel: it lands
2318 // here, the derivative is filled with NaN so the solver stops rather
2319 // than integrating stale numbers, and the first message is raised
2320 // after the run as `compiler_unavailable` — which is what a device
2321 // that stopped working mid-solve IS (esm-spec §9.6.6: this binding,
2322 // build or PROCESS cannot provide the compiler). The lengths that
2323 // `CompiledRhs::eval_into` would otherwise report here were already
2324 // checked at install time, so nothing about the MODEL can reach this
2325 // channel.
2326 #[cfg(feature = "xla")]
2327 let xla_fault: Rc<RefCell<Option<String>>> = Rc::new(RefCell::new(None));
2328 #[cfg(feature = "xla")]
2329 let xla_fault_rhs = Rc::clone(&xla_fault);
2330 #[cfg(feature = "xla")]
2331 let xla_fault_jac = Rc::clone(&xla_fault);
2332
2333 // External forcing channel (PR-1, ess-14f.7): clone the `Rc` handle into
2334 // each closure so both the RHS and the Jacobian read the *same*
2335 // model-lifetime buffer the caller refreshes between segments.
2336 let forcing_rhs = Rc::clone(&self.forcing);
2337 let forcing_jac = Rc::clone(&self.forcing);
2338
2339 let rhs_closure = move |y: &diffsol::FaerVec<f64>,
2340 p: &diffsol::FaerVec<f64>,
2341 t: f64,
2342 dy: &mut diffsol::FaerVec<f64>| {
2343 let y_s = y.as_slice();
2344 let p_s = p.as_slice();
2345 let dy_s = dy.as_mut_slice();
2346 #[cfg(feature = "xla")]
2347 if let Some(program) = &xla_rhs {
2348 dy_s.fill(0.0);
2349 run_xla_rhs(program, &xla_fault_rhs, y_s, p_s, t, dy_s);
2350 return;
2351 }
2352 if let Some(sp) = &program_rhs {
2353 let mut ws = sp.rhs.borrow_mut();
2354 let ws = ws.get_or_insert_with(|| sp.program.solve_workspace());
2355 sp.program
2356 .eval_solve(y_s, p_s, t, &forcing_rhs.borrow(), dy_s, ws);
2357 return;
2358 }
2359 let mut scratch_slot = rhs_scratch.borrow_mut();
2360 let scratch = scratch_slot
2361 .as_mut()
2362 .expect("the tape scratch is built whenever no XLA program serves the rhs");
2363 evaluate_rhs_with_scratch(
2364 &RhsCall {
2365 rhs_rules: &rhs_rules,
2366 observed_rules: &varying_rules_rhs,
2367 var_shapes: &var_shapes,
2368 param_names: ¶m_names,
2369 state: y_s,
2370 params: p_s,
2371 forcing: &forcing_rhs,
2372 t,
2373 declared: &declared,
2374 rule_set: rule_set_rhs,
2375 },
2376 dy_s,
2377 force_scalar,
2378 &mut RhsStats::default(),
2379 scratch,
2380 );
2381 };
2382
2383 let jac_closure = move |y: &diffsol::FaerVec<f64>,
2384 p: &diffsol::FaerVec<f64>,
2385 t: f64,
2386 v: &diffsol::FaerVec<f64>,
2387 jv: &mut diffsol::FaerVec<f64>| {
2388 #[cfg(feature = "xla")]
2389 if let Some((program, jvp)) = &xla_jac {
2390 let mut jvp = jvp.borrow_mut();
2391 let jvp = jvp.get_or_insert_with(|| FdJvp::new(n_states));
2392 jvp.apply(
2393 y.as_slice(),
2394 p.as_slice(),
2395 t,
2396 v.as_slice(),
2397 jv.as_mut_slice(),
2398 |state, params, t, out| {
2399 run_xla_rhs(program, &xla_fault_jac, state, params, t, out)
2400 },
2401 );
2402 return;
2403 }
2404 let n = y.as_slice().len();
2405 let v_s = v.as_slice();
2406 let p_s = p.as_slice();
2407 let y_s = y.as_slice();
2408 let mut y_norm = 0.0f64;
2409 for &yi in y_s {
2410 y_norm += yi * yi;
2411 }
2412 let y_norm = y_norm.sqrt().max(1.0);
2413 let eps = f64::EPSILON.sqrt() * y_norm;
2414
2415 let mut y_perturbed = vec![0.0f64; n];
2416 for i in 0..n {
2417 y_perturbed[i] = y_s[i] + eps * v_s[i];
2418 }
2419
2420 let mut f_y = vec![0.0f64; n];
2421 let mut f_yp = vec![0.0f64; n];
2422 if let Some(sp) = &program_jac {
2423 let mut ws = sp.jac.borrow_mut();
2424 let ws = ws.get_or_insert_with(|| sp.program.solve_workspace());
2425 let forcing = forcing_jac.borrow();
2426 sp.program.eval_solve(y_s, p_s, t, &forcing, &mut f_y, ws);
2427 sp.program
2428 .eval_solve(&y_perturbed, p_s, t, &forcing, &mut f_yp, ws);
2429 let jv_s = jv.as_mut_slice();
2430 for i in 0..n {
2431 jv_s[i] = (f_yp[i] - f_y[i]) / eps;
2432 }
2433 return;
2434 }
2435 let mut scratch_slot = jac_scratch.borrow_mut();
2436 let scratch = scratch_slot.get_or_insert_with(|| {
2437 let mut s = RhsScratch::new(&var_shapes_jac);
2438 s.set_const_arrays(std::sync::Arc::clone(&const_scope_jac));
2439 s.set_static((*jac_seed).clone());
2440 // Under `native` the FD Jacobian runs on the TAPE as well
2441 // (under `xla` it runs on the emitted program instead, which
2442 // the arm above took).
2443 // The two right-hand-side evaluations it differences are the
2444 // same rules the production closure runs, so leaving them on
2445 // the legacy interpreter meant an implicit solve evaluated the
2446 // whole model through the overlay — and, wherever the overlay
2447 // declined, per cell — at every Jacobian call: the largest
2448 // off-tape evaluation in the driver, and one §2.5.10 covers
2449 // ("every evaluation the compiler performs for the Problem").
2450 // The tape is bit-identical to the legacy path (that is what
2451 // the compiler-agreement tier asserts, CONFORMANCE_SPEC
2452 // §5.44), so the differenced Jacobian is unchanged.
2453 if let Some(prog) = &tape_jac {
2454 s.install_tape(std::sync::Arc::clone(prog));
2455 }
2456 s
2457 });
2458 evaluate_rhs_with_scratch(
2459 &RhsCall {
2460 rhs_rules: &rhs_rules_jac,
2461 observed_rules: &varying_rules_jac,
2462 var_shapes: &var_shapes_jac,
2463 param_names: ¶m_names_jac,
2464 state: y_s,
2465 params: p_s,
2466 forcing: &forcing_jac,
2467 t,
2468 declared: &declared_jac,
2469 rule_set: rule_set_jac,
2470 },
2471 &mut f_y,
2472 force_scalar,
2473 &mut RhsStats::default(),
2474 scratch,
2475 );
2476 evaluate_rhs_with_scratch(
2477 &RhsCall {
2478 rhs_rules: &rhs_rules_jac,
2479 observed_rules: &varying_rules_jac,
2480 var_shapes: &var_shapes_jac,
2481 param_names: ¶m_names_jac,
2482 state: &y_perturbed,
2483 params: p_s,
2484 forcing: &forcing_jac,
2485 t,
2486 declared: &declared_jac,
2487 rule_set: rule_set_jac,
2488 },
2489 &mut f_yp,
2490 force_scalar,
2491 &mut RhsStats::default(),
2492 scratch,
2493 );
2494 let jv_s = jv.as_mut_slice();
2495 for i in 0..n {
2496 jv_s[i] = (f_yp[i] - f_y[i]) / eps;
2497 }
2498 };
2499
2500 // Concrete values: `solve` has already resolved the esm-spec §2.2.2
2501 // chain into `opts`. The fallback covers a direct call that bypasses
2502 // `solve` and therefore has no document to consult.
2503 let abstol = opts.abstol_or_default();
2504 let reltol = opts.reltol_or_default();
2505 let ic_for_init = u0.to_vec();
2506
2507 let builder = OdeBuilder::<FaerMat<f64>>::new()
2508 .t0(t0)
2509 .rtol(reltol)
2510 .atol(vec![abstol; n_states])
2511 .p(param_vec.to_vec())
2512 .rhs_implicit(rhs_closure, jac_closure)
2513 .init(
2514 move |_p: &diffsol::FaerVec<f64>, _t: f64, y: &mut diffsol::FaerVec<f64>| {
2515 let y_s = y.as_mut_slice();
2516 for (i, &v) in ic_for_init.iter().enumerate() {
2517 y_s[i] = v;
2518 }
2519 },
2520 n_states,
2521 );
2522
2523 let problem = builder.build().map_err(|e| SimulateError::DiffsolError {
2524 details: e.to_string(),
2525 })?;
2526
2527 // Run the solver, then read the real step/eval counters out of diffsol
2528 // before the concrete solver is dropped (see [`SolveStats::from_solver`]).
2529 //
2530 // Wrapped in an immediately-invoked closure so a solver failure does
2531 // not leave the function before the XLA fault channel below is read: a
2532 // device that stopped working reaches the solver as a NaN derivative,
2533 // and reporting "tolerance not met" for it would name the wrong thing.
2534 type SolvedSegment = (Vec<f64>, Vec<Vec<f64>>, SolveStats, ReturnCode);
2535 let solved = (|| -> Result<SolvedSegment, SimulateError> {
2536 let out = match opts.alg {
2537 Alg::Bdf => {
2538 let mut solver: Bdf<'_, _, NewtonNonlinearSolver<_, FaerLU<f64>, _>> = problem
2539 .bdf::<FaerLU<f64>>()
2540 .map_err(|e| SimulateError::DiffsolError {
2541 details: e.to_string(),
2542 })?;
2543 let (time, state, retcode) =
2544 crate::simulate::run_solver(&mut solver, t_end, opts)?;
2545 let bs = solver.get_statistics();
2546 let stats = SolveStats::from_solver(
2547 &solver,
2548 bs.number_of_steps,
2549 bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
2550 );
2551 (time, state, stats, retcode)
2552 }
2553 Alg::Sdirk => {
2554 let mut solver: Sdirk<'_, _, FaerLU<f64>> = problem
2555 .tr_bdf2::<FaerLU<f64>>()
2556 .map_err(|e| SimulateError::DiffsolError {
2557 details: e.to_string(),
2558 })?;
2559 let (time, state, retcode) =
2560 crate::simulate::run_solver(&mut solver, t_end, opts)?;
2561 let bs = solver.get_statistics();
2562 let stats = SolveStats::from_solver(
2563 &solver,
2564 bs.number_of_steps,
2565 bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
2566 );
2567 (time, state, stats, retcode)
2568 }
2569 Alg::Erk => {
2570 let _seq = crate::simulate::SequentialFaer::enter();
2571 let mut solver = problem.tsit45().map_err(|e| SimulateError::DiffsolError {
2572 details: e.to_string(),
2573 })?;
2574 let (time, state, retcode) =
2575 crate::simulate::run_solver(&mut solver, t_end, opts)?;
2576 let bs = solver.get_statistics();
2577 let stats = SolveStats::from_solver(
2578 &solver,
2579 bs.number_of_steps,
2580 bs.number_of_error_test_failures + bs.number_of_nonlinear_solver_fails,
2581 );
2582 (time, state, stats, retcode)
2583 }
2584 };
2585 Ok(out)
2586 })();
2587 // §9.6.6 `compiler_unavailable`: the caller named `xla`, and the XLA
2588 // runtime stopped being able to provide it mid-solve. It is reported
2589 // in preference to whatever the solver made of the NaN derivative,
2590 // because that is the fact — and the first failure is reported, since
2591 // every one after it ran on the same broken device.
2592 #[cfg(feature = "xla")]
2593 if let Some(details) = xla_fault.borrow().clone() {
2594 return Err(SimulateError::CompilerUnavailable {
2595 compiler: "xla",
2596 details: format!("the compiled right-hand side failed during the solve: {details}"),
2597 });
2598 }
2599 let (time, state, stats, retcode) = solved?;
2600 Ok((time, state, stats, retcode))
2601 }
2602
2603 /// Expose observed trajectories (e.g. an `area` FAQ) alongside the states
2604 /// so inline conformance assertions can read algebraic quantities
2605 /// (RFC §8.1; CONFORMANCE_SPEC.md §5.8). The integrator carries only the
2606 /// state vector, so re-evaluate the (dependency-ordered, derived-ring-aware)
2607 /// observeds from the state trajectory at each output node and append them
2608 /// as extra rows (with matching entries in `state_variable_names`). Mirrors
2609 /// the Python `_simulate_with_numpy` output-observed exposure. A model with
2610 /// no observed rules (or an empty trajectory) is untouched.
2611 ///
2612 /// **Scalar observeds are always exposed; array-valued ones only on
2613 /// request.** `requested` is [`SolveOptions::output_observed`] — the
2614 /// caller-named subset of streaming-output-sinks RFC decision 8. An
2615 /// array-valued observed (the clip ring, the const polygons, a gridded
2616 /// emissions field) is `n_cells` rows, not one, so materializing every one
2617 /// of them unasked would multiply the trajectory's memory by the grid size;
2618 /// a requested one is flattened to one row per cell named exactly like a
2619 /// state cell — `name[i,j,…]`, 1-based and column-major, the spelling
2620 /// [`build_slot_tables`] gives the state and the one
2621 /// [`crate::derive_output_plan`] inverts back into a labeled output array.
2622 /// Names may be bare or `Model.`-qualified.
2623 #[cfg(feature = "solve")]
2624 #[allow(clippy::too_many_arguments)]
2625 #[allow(clippy::too_many_arguments)]
2626 fn append_observed_trajectories(
2627 &self,
2628 time: &[f64],
2629 state: &mut Vec<Vec<f64>>,
2630 state_variable_names: &mut Vec<String>,
2631 param_vec: &[f64],
2632 static_obs: &ArrMap,
2633 static_names: &HashSet<String>,
2634 varying_rules: &[AlgebraicRule],
2635 requested: &[String],
2636 tape: Option<&std::sync::Arc<TapeProgram>>,
2637 ) {
2638 if self.observed_rules.is_empty() || time.is_empty() {
2639 return;
2640 }
2641 // Under either strict compiler: the observeds reported at output
2642 // times are read off the TAPE, one call per saved node, instead of being re-derived through
2643 // the whole-array overlay. esm-libraries-spec §2.5.10 names this pass
2644 // explicitly — "the observeds reported at output times" are under the
2645 // refusal like the right-hand side is — and it is the one that runs
2646 // most often, once per saved time point.
2647 if self.tape_serves_passes()
2648 && let Some(tape) = tape
2649 {
2650 self.append_observed_trajectories_taped(
2651 time,
2652 state,
2653 state_variable_names,
2654 param_vec,
2655 static_names,
2656 varying_rules,
2657 requested,
2658 tape,
2659 );
2660 return;
2661 }
2662 let wanted = self.resolve_requested_observeds(requested);
2663 let nt = time.len();
2664
2665 // ---- per-node evaluation state, allocated ONCE ----------------------
2666 // The old shape of this pass built a fresh `ArrMap`, a fresh state-array
2667 // map and a deep `arr.clone()` of every hoisted static observed at every
2668 // one of the NT output nodes. All three are hoisted out of the node loop
2669 // here: the state arrays are refilled in place (the same
2670 // `refill_state_arrays` the RHS uses — identical column-major placement),
2671 // and the observed map keeps its static entries in place across nodes,
2672 // dropping only the varying ones with `retain`, exactly as `RhsScratch`
2673 // does between RHS calls.
2674 let mut flat = vec![0.0f64; self.n_states];
2675 for (i, slot) in flat.iter_mut().enumerate() {
2676 *slot = state[i][0];
2677 }
2678 let mut sa = build_state_arrays(&self.var_shapes, &flat);
2679 let static_keys: HashSet<String> = static_obs.keys().cloned().collect();
2680 let mut obs: ArrMap = ArrMap::default();
2681 for (name, arr) in static_obs {
2682 obs.insert(name.clone(), arr.clone());
2683 }
2684 let dr: RefCell<HashMap<String, ArrayD<f64>>> = RefCell::new(HashMap::new());
2685
2686 // ess-cse: this pass walks the SAME expanded discretization bodies the
2687 // RHS does (one `simpleclimate` advection observed is 45k+ operator
2688 // nodes) and used to run with `cse: None` — i.e. with no memoization at
2689 // all, re-evaluating every repeated subtree at every occurrence at every
2690 // output node. Give it its own runtime. The class table is keyed by AST
2691 // node address, so it is (re)bound to whichever rule slice is about to be
2692 // evaluated; see [`CseRt::retarget`].
2693 let cse = CseRt::default();
2694 let bind_cse = |rules: &[AlgebraicRule]| {
2695 cse.retarget(
2696 (rules.as_ptr() as u64).rotate_left(17)
2697 ^ ((rules.len() as u64) << 16)
2698 ^ 0x0b5e_0b5e_0b5e_0b5e,
2699 );
2700 };
2701
2702 // ---- dependency-cone pruning ---------------------------------------
2703 // Only the 0-D observeds become rows, but every varying rule — the whole
2704 // PPM stencil chain — was materialized at every node and then thrown
2705 // away. Restrict the work to the transitive dependency cone of what is
2706 // actually read.
2707 //
2708 // The cone is computed from `Expr::Variable` leaves, so it is only valid
2709 // for rules whose dependency edges are all expressible that way: a FAQ
2710 // ring producer/consumer pair and a ragged/derived contraction bound
2711 // couple two rules through the `derived_rings` registry / an `offsets`
2712 // factor name instead, and a `join` couples them through key COLUMN
2713 // names. A model using any of those keeps the old un-pruned behaviour.
2714 let prune = varying_rules
2715 .iter()
2716 .all(|r| !expr_blocks_output_pruning(observed_rule_body(r)));
2717
2718 // Probe: the rules that must actually run before 0-D-ness is known. A
2719 // rule whose value is provably an array needs no probe at all, and a
2720 // hoisted static observed's rank is already in `obs`.
2721 let probe_owned: Vec<AlgebraicRule>;
2722 let probe_rules: &[AlgebraicRule] = if prune {
2723 // A REQUESTED array-valued observed becomes rows as well, so it
2724 // must be in the probe cone even though its rank is already known.
2725 let unknown: HashSet<String> = self
2726 .observed_rules
2727 .iter()
2728 .filter(|r| {
2729 !observed_rule_is_array_valued(r) || wanted.contains(observed_rule_var(r))
2730 })
2731 .map(|r| observed_rule_var(r).clone())
2732 .collect();
2733 match dependency_cone(varying_rules, &unknown) {
2734 None => varying_rules,
2735 Some(cone) => {
2736 probe_owned = cone;
2737 &probe_owned
2738 }
2739 }
2740 } else {
2741 varying_rules
2742 };
2743 if !probe_rules.is_empty() {
2744 bind_cse(probe_rules);
2745 materialize_observeds_pass(
2746 &mut obs,
2747 probe_rules,
2748 &ObsPass {
2749 env: EvalEnv {
2750 state_arrays: &sa,
2751 params: param_vec,
2752 param_names: &self.param_names,
2753 t: time[0],
2754 derived_rings: &dr,
2755 derived_extents: empty_derived_extents(),
2756 forcing: &self.forcing,
2757 cse: Some(&cse),
2758 const_lits: None,
2759 const_arrays: &self.const_scope,
2760 declared: &self.declared_names,
2761 },
2762 // Output-node observed snapshot: vectorized overlay,
2763 // unless this is the `interpreter` reference.
2764 force_scalar: self.is_interpreter(),
2765 },
2766 &mut RhsStats::default(),
2767 );
2768 }
2769 // What becomes rows: every 0-D observed, as before, plus every
2770 // CALLER-REQUESTED array-valued one, at one row per cell.
2771 //
2772 // A name the probe did not materialize at all is skipped, and the
2773 // verdict is the same one the full materialization would have produced:
2774 // by construction such a name is either array-valued and unrequested —
2775 // so not a row either way — or outside the cone of anything that could
2776 // be 0-D. A requested array observed IS in the probe cone (see above),
2777 // so its rank and shape are known here; taking the shape from the probe
2778 // rather than per node is what keeps the row block from shifting if a
2779 // later node fails to materialize the name.
2780 let emit: Vec<ObservedRows> = self
2781 .observed_rules
2782 .iter()
2783 .filter_map(|rule| {
2784 let name = observed_rule_var(rule);
2785 let arr = obs.get(name)?;
2786 if arr.ndim() == 0 {
2787 Some(ObservedRows::scalar(name.clone()))
2788 } else if wanted.contains(name) {
2789 Some(ObservedRows::gridded(name.clone(), arr.shape().to_vec()))
2790 } else {
2791 None
2792 }
2793 })
2794 .collect();
2795 if emit.is_empty() {
2796 return;
2797 }
2798 let n_rows: usize = emit.iter().map(ObservedRows::n_rows).sum();
2799
2800 // Value pass: the cone of the rows we actually emit.
2801 let value_owned: Vec<AlgebraicRule>;
2802 let value_rules: &[AlgebraicRule] = if prune {
2803 let seeds: HashSet<String> = emit.iter().map(|e| e.name.clone()).collect();
2804 match dependency_cone(varying_rules, &seeds) {
2805 None => varying_rules,
2806 Some(cone) => {
2807 value_owned = cone;
2808 &value_owned
2809 }
2810 }
2811 } else {
2812 varying_rules
2813 };
2814
2815 let mut rows: Vec<Vec<f64>> = vec![Vec::with_capacity(nt); n_rows];
2816 let record = |obs: &ArrMap, rows: &mut Vec<Vec<f64>>| {
2817 let mut j = 0usize;
2818 for e in &emit {
2819 match obs.get(&e.name) {
2820 Some(a) if e.shape.is_empty() => {
2821 rows[j].push(a.first().copied().unwrap_or(f64::NAN));
2822 j += 1;
2823 }
2824 Some(a) if a.shape() == e.shape.as_slice() => {
2825 // Column-major, the order the cell keys were named in.
2826 for v in arrayd_to_col_major(a) {
2827 rows[j].push(v);
2828 j += 1;
2829 }
2830 }
2831 // Absent at this node, or a rank the probe did not see: NaN
2832 // across the block, so the fault shows in the values rather
2833 // than silently shifting every later row.
2834 _ => {
2835 for _ in 0..e.n_rows() {
2836 rows[j].push(f64::NAN);
2837 j += 1;
2838 }
2839 }
2840 }
2841 }
2842 };
2843 record(&obs, &mut rows);
2844 if value_rules.is_empty() {
2845 // Every 0-D observed sits in the hoisted static set, so its value is
2846 // node-independent: replicate node 0 instead of re-running the whole
2847 // varying rule set NT times for a number that cannot change.
2848 for row in rows.iter_mut() {
2849 let v = row[0];
2850 row.resize(nt, v);
2851 }
2852 } else {
2853 bind_cse(value_rules);
2854 for k in 1..nt {
2855 for (i, slot) in flat.iter_mut().enumerate() {
2856 *slot = state[i][k];
2857 }
2858 refill_state_arrays(&mut sa, &self.var_shapes, &flat);
2859 obs.retain(|name, _| static_keys.contains(name));
2860 dr.borrow_mut().clear();
2861 materialize_observeds_pass(
2862 &mut obs,
2863 value_rules,
2864 &ObsPass {
2865 env: EvalEnv {
2866 state_arrays: &sa,
2867 params: param_vec,
2868 param_names: &self.param_names,
2869 t: time[k],
2870 derived_rings: &dr,
2871 derived_extents: empty_derived_extents(),
2872 forcing: &self.forcing,
2873 cse: Some(&cse),
2874 const_lits: None,
2875 const_arrays: &self.const_scope,
2876 declared: &self.declared_names,
2877 },
2878 force_scalar: self.is_interpreter(),
2879 },
2880 &mut RhsStats::default(),
2881 );
2882 record(&obs, &mut rows);
2883 }
2884 }
2885 let mut rows = rows.into_iter();
2886 for e in &emit {
2887 for name in e.row_names() {
2888 let Some(row) = rows.next() else { break };
2889 state_variable_names.push(name);
2890 state.push(row);
2891 }
2892 }
2893 }
2894
2895 /// [`Self::append_observed_trajectories`] served from the tape — the path
2896 /// both strict compilers take.
2897 ///
2898 /// Same rows, same order, same cell-key spelling; the only difference is
2899 /// where the numbers come from. One taped call per saved time point
2900 /// publishes every observed (a strict build exports them all, because a
2901 /// harvest that covered only the probe cone would silently skip a
2902 /// caller-requested array observed and a static one), and the rows are read
2903 /// straight off that map.
2904 ///
2905 /// The dependency-cone pruning the overlay path needs has no counterpart
2906 /// here: the tape computes the whole CONTINUOUS section either way, and
2907 /// its CONST and SEGMENT sections prime once for the whole sweep rather
2908 /// than once per node, which is what the pruning was buying back.
2909 #[cfg(feature = "solve")]
2910 #[allow(clippy::too_many_arguments)]
2911 fn append_observed_trajectories_taped(
2912 &self,
2913 time: &[f64],
2914 state: &mut Vec<Vec<f64>>,
2915 state_variable_names: &mut Vec<String>,
2916 param_vec: &[f64],
2917 static_names: &HashSet<String>,
2918 varying_rules: &[AlgebraicRule],
2919 requested: &[String],
2920 tape: &std::sync::Arc<TapeProgram>,
2921 ) {
2922 let wanted = self.resolve_requested_observeds(requested);
2923 // WHICH observeds become rows is decided exactly as the overlay path
2924 // decides it, and deliberately not by what the tape happens to
2925 // publish: a strict build exports every observed so the harvest
2926 // cannot miss one, and emitting every one of them would put rows in
2927 // the solution that the same document did not carry before. The
2928 // candidates are the hoisted static observeds plus the probe cone —
2929 // the potentially-scalar rules and their transitive dependencies —
2930 // which is what `obs` holds at this point on the overlay path.
2931 let emitted = self.output_row_candidates(static_names, varying_rules, &wanted);
2932 let nt = time.len();
2933 let mut harvest = TapedObserveds::new(self, tape);
2934 let mut flat = vec![0.0f64; self.n_states];
2935
2936 for (i, slot) in flat.iter_mut().enumerate() {
2937 *slot = state[i][0];
2938 }
2939 // What becomes rows: every 0-D observed, plus every CALLER-REQUESTED
2940 // array-valued one at one row per cell. Decided at node 0 and held
2941 // fixed, so a later node cannot shift the row block.
2942 let emit: Vec<ObservedRows> = {
2943 let obs = harvest.at(self, &flat, param_vec, time[0]);
2944 self.observed_rules
2945 .iter()
2946 .filter_map(|rule| {
2947 let name = observed_rule_var(rule);
2948 if !emitted.contains(name) {
2949 return None;
2950 }
2951 let arr = obs.get(name)?;
2952 if arr.ndim() == 0 {
2953 Some(ObservedRows::scalar(name.clone()))
2954 } else if wanted.contains(name) {
2955 Some(ObservedRows::gridded(name.clone(), arr.shape().to_vec()))
2956 } else {
2957 None
2958 }
2959 })
2960 .collect()
2961 };
2962 if emit.is_empty() {
2963 return;
2964 }
2965 let n_rows: usize = emit.iter().map(ObservedRows::n_rows).sum();
2966 let mut rows: Vec<Vec<f64>> = vec![Vec::with_capacity(nt); n_rows];
2967 let record = |obs: &ArrMap, rows: &mut Vec<Vec<f64>>| {
2968 let mut j = 0usize;
2969 for e in &emit {
2970 match obs.get(&e.name) {
2971 Some(a) if e.shape.is_empty() => {
2972 rows[j].push(a.first().copied().unwrap_or(f64::NAN));
2973 j += 1;
2974 }
2975 Some(a) if a.shape() == e.shape.as_slice() => {
2976 for v in arrayd_to_col_major(a) {
2977 rows[j].push(v);
2978 j += 1;
2979 }
2980 }
2981 // Absent, or a rank node 0 did not see: NaN across the
2982 // block, so the fault shows in the values rather than
2983 // silently shifting every later row.
2984 _ => {
2985 for _ in 0..e.n_rows() {
2986 rows[j].push(f64::NAN);
2987 j += 1;
2988 }
2989 }
2990 }
2991 }
2992 };
2993 record(harvest.scratch_observeds(), &mut rows);
2994 for k in 1..nt {
2995 for (i, slot) in flat.iter_mut().enumerate() {
2996 *slot = state[i][k];
2997 }
2998 let obs = harvest.at(self, &flat, param_vec, time[k]);
2999 record(obs, &mut rows);
3000 }
3001
3002 let mut rows = rows.into_iter();
3003 for e in &emit {
3004 for name in e.row_names() {
3005 let Some(row) = rows.next() else { break };
3006 state_variable_names.push(name);
3007 state.push(row);
3008 }
3009 }
3010 }
3011
3012 /// The observeds that may become solution rows: the hoisted static ones
3013 /// plus the output-node probe cone.
3014 ///
3015 /// The overlay path's candidate set, named. It materializes the statics
3016 /// once and then the PROBE CONE — every potentially-scalar observed and
3017 /// its transitive dependencies — and an observed outside both is simply
3018 /// not in `obs` when the row set is decided, so it becomes no row. The
3019 /// taped path publishes every observed (that is what makes its harvest
3020 /// complete), so it has to be told the same rule rather than inferring it
3021 /// from what happens to be published.
3022 #[cfg(feature = "solve")]
3023 fn output_row_candidates(
3024 &self,
3025 static_names: &HashSet<String>,
3026 varying_rules: &[AlgebraicRule],
3027 wanted: &HashSet<String>,
3028 ) -> HashSet<String> {
3029 let mut out: HashSet<String> = static_names.clone();
3030 let prune = varying_rules
3031 .iter()
3032 .all(|r| !expr_blocks_output_pruning(observed_rule_body(r)));
3033 let cone: Option<Vec<AlgebraicRule>> = if prune {
3034 let unknown: HashSet<String> = self
3035 .observed_rules
3036 .iter()
3037 .filter(|r| {
3038 !observed_rule_is_array_valued(r) || wanted.contains(observed_rule_var(r))
3039 })
3040 .map(|r| observed_rule_var(r).clone())
3041 .collect();
3042 dependency_cone(varying_rules, &unknown)
3043 } else {
3044 None
3045 };
3046 match cone {
3047 Some(rules) => out.extend(rules.iter().map(|r| observed_rule_var(r).clone())),
3048 // No cone: the un-pruned behaviour materializes every varying rule.
3049 None => out.extend(varying_rules.iter().map(|r| observed_rule_var(r).clone())),
3050 }
3051 out
3052 }
3053
3054 /// The observed-rule names `requested` names.
3055 ///
3056 /// The rule [`crate::derive_output_plan`] applies to an output request
3057 /// (CONFORMANCE_SPEC §5.17.4), over the observed rules: the exact name, else
3058 /// the ONE rule whose last dotted segment equals the request's. A name that
3059 /// matches nothing, or whose last segment several rules share, selects
3060 /// nothing and is not an error here — the output plan diagnoses it with
3061 /// [`crate::OutputError::UnknownObserved`] or
3062 /// [`crate::OutputError::AmbiguousRequest`], and it also sees the state
3063 /// slots, so it tells the caller the whole truth.
3064 #[cfg(feature = "solve")]
3065 fn resolve_requested_observeds(&self, requested: &[String]) -> HashSet<String> {
3066 if requested.is_empty() {
3067 return HashSet::new();
3068 }
3069 let vars: Vec<&str> = self
3070 .observed_rules
3071 .iter()
3072 .map(|r| observed_rule_var(r).as_str())
3073 .collect();
3074 requested
3075 .iter()
3076 .filter_map(
3077 |r| match crate::data_output::match_output_request(r, &vars) {
3078 crate::data_output::RequestMatch::Named(var) => Some(var.to_string()),
3079 _ => None,
3080 },
3081 )
3082 .collect()
3083 }
3084
3085 /// Names of the observeds that are STATE-FREE and `t`-free: their transitive
3086 /// references reach no state variable and no `t` (each reference is a
3087 /// parameter, a loop index, an external forcing entry, or an already-static
3088 /// observed). Because `observed_rules` is dependency-ordered (Kahn sweep at
3089 /// build), one forward pass classifies each rule after its references; a
3090 /// cycle survivor's unplaced reference correctly disqualifies it.
3091 ///
3092 /// These are the observeds the RHS hoists out of the per-step loop and the
3093 /// build-once products a [`BuildInspection`] records. Unlike a strict
3094 /// "state-free" set, an external **forcing** reference is ALLOWED: a CONST
3095 /// loader field is constant within a `simulate` call (the free `simulate`
3096 /// never refreshes the forcing buffer mid-run), so an observed reaching only
3097 /// params + forcing + static observeds is constant across the whole solve.
3098 /// The regridded terrain (`elev_xy`) and its slopes — forcing-derived but
3099 /// state-free — thus hoist and land in `setup_arrays`, matching the Julia /
3100 /// Python `BuildInspection`.
3101 pub(super) fn classify_static_observeds(
3102 &self,
3103 discrete_forcing: &HashSet<String>,
3104 ) -> HashSet<String> {
3105 // CONST tier: state-free, `t`-free, AND forcing-free.
3106 self.classify_segment_invariant_observeds(discrete_forcing, false)
3107 }
3108
3109 /// Classify the observeds that are invariant *within a cadence segment* — the
3110 /// cadence lattice's `≤ DISCRETE` set (`CONST ⊏ DISCRETE`, `cadence.rs`): their
3111 /// transitive references reach neither `t` nor any state variable.
3112 ///
3113 /// * `allow_forcing = false` → the CONST set (also forcing-free): constant for
3114 /// the entire solve, materialized ONCE at setup (the hoisted `static_obs`).
3115 /// * `allow_forcing = true` → CONST ∪ DISCRETE: an observed may reach a
3116 /// discrete forcing buffer. Since the driver refreshes forcing only *between*
3117 /// segments (never inside a solver step), such an observed is constant
3118 /// *within* a segment — so the DISCRETE remainder is materialized ONCE per
3119 /// segment rather than re-evaluated every RHS step. This is the fix for a
3120 /// coupled model whose per-cell conservative-regrid observeds (state-free,
3121 /// `t`-free, forcing-fed) were previously recomputed every step.
3122 ///
3123 /// `observed_rules` are dependency-ordered, so the transitive
3124 /// `set.contains(r)` check resolves against already-classified predecessors.
3125 pub(super) fn classify_segment_invariant_observeds(
3126 &self,
3127 discrete_forcing: &HashSet<String>,
3128 allow_forcing: bool,
3129 ) -> HashSet<String> {
3130 let observed_names: HashSet<&String> =
3131 self.observed_rules.iter().map(observed_rule_var).collect();
3132 let mut set: HashSet<String> = HashSet::new();
3133 for rule in &self.observed_rules {
3134 let mut refs = HashSet::new();
3135 collect_expr_var_refs(observed_rule_body(rule), &mut refs);
3136 // A RECURRENCE rule's body reads the variable the rule DEFINES
3137 // (esm-spec §4.3.1.1). That self-edge is not a dependency on
3138 // another observed — it is an ordering within this one — so it is
3139 // dropped, exactly as `dependency_order_observed` drops it. Leaving
3140 // it in made the rule classify as non-static and never be hoisted,
3141 // which is why a state-free recurrence materialized nowhere at all
3142 // before this feature landed.
3143 //
3144 // Conditioned on the rule KIND, not applied to every rule: a
3145 // self-reference the recurrence lowering did not recognize is not
3146 // an ordering, it is a rule that reads a name nothing binds, and
3147 // hoisting it as a build-once constant would freeze that. The
3148 // exemption belongs to the construct that earns it.
3149 if matches!(rule, AlgebraicRule::Recurrence { .. }) {
3150 let self_name = observed_rule_var(rule);
3151 refs.retain(|r| r != self_name);
3152 }
3153 let ok = refs.iter().all(|r| {
3154 r != "t"
3155 && !self.var_shapes.contains_key(r)
3156 // A DISCRETE (hourly) forcing buffer is a LIVE field the driver
3157 // refreshes between segments. Excluded from the CONST tier (it
3158 // must not freeze at setup), but ALLOWED in the ≤DISCRETE tier
3159 // (constant *within* a segment, so hoistable to per-segment).
3160 && (allow_forcing || !discrete_forcing.contains(r))
3161 && (!observed_names.contains(r) || set.contains(r))
3162 });
3163 if ok {
3164 set.insert(observed_rule_var(rule).clone());
3165 }
3166 }
3167 set
3168 }
3169
3170 /// Fill a [`BuildInspection`] sink from the already-hoisted static observeds:
3171 /// record every rule's resolved body expression, the resolved scalar
3172 /// parameters, and the arrays of the static (state-free / `t`-free) subset
3173 /// (`static_obs`, materialized once by the caller). Read-only with respect to
3174 /// the run.
3175 #[cfg(feature = "solve")]
3176 fn fill_inspection(
3177 &self,
3178 insp: &mut BuildInspection,
3179 static_obs: &ArrMap,
3180 static_names: &HashSet<String>,
3181 param_vec: &[f64],
3182 ) {
3183 // Resolved scalar parameters (load-time constants) so the reference / ic
3184 // positions can bind them into a build-time cellwise evaluation.
3185 for (i, name) in self.param_names.iter().enumerate() {
3186 insp.params.insert(name.clone(), param_vec[i]);
3187 }
3188 for rule in &self.observed_rules {
3189 insp.observed_exprs.insert(
3190 observed_rule_var(rule).clone(),
3191 observed_rule_body(rule).clone(),
3192 );
3193 // esm-spec §4.3.1.1 / RFC §2.1: the recurrence a document's
3194 // self-read was interpreted as, reported rather than authored.
3195 if let AlgebraicRule::Recurrence {
3196 var,
3197 output_idx_names,
3198 axis,
3199 max_lag,
3200 lag_proven,
3201 ..
3202 } = rule
3203 {
3204 insp.recurrences
3205 .push(crate::simulate_array::RecurrenceInfo {
3206 var: var.clone(),
3207 axis: output_idx_names[*axis].clone(),
3208 max_lag: *max_lag,
3209 lag_proven: *lag_proven,
3210 });
3211 }
3212 }
3213 for name in static_names {
3214 if let Some(a) = static_obs.get(name) {
3215 insp.setup_arrays.insert(name.clone(), a.clone());
3216 }
3217 }
3218 }
3219}
3220
3221/// One observed's contribution to the appended trajectory rows: a 0-D observed
3222/// is a single row named after itself, an array-valued one is `n_cells` rows
3223/// named `base[i,j,…]` (1-based, column-major) — the identical cell-key scheme
3224/// [`build_slot_tables`] gives an array STATE, so the two are indistinguishable
3225/// to [`crate::derive_output_gridding`] and land on the same emergent grid.
3226#[cfg(feature = "solve")]
3227struct ObservedRows {
3228 /// The observed rule's variable name.
3229 name: String,
3230 /// Gridded shape, in cell-key axis order; EMPTY for a 0-D observed.
3231 shape: Vec<usize>,
3232}
3233
3234#[cfg(feature = "solve")]
3235impl ObservedRows {
3236 fn scalar(name: String) -> Self {
3237 ObservedRows {
3238 name,
3239 shape: Vec::new(),
3240 }
3241 }
3242
3243 fn gridded(name: String, shape: Vec<usize>) -> Self {
3244 ObservedRows { name, shape }
3245 }
3246
3247 /// How many trajectory rows this observed contributes (1 for a scalar).
3248 fn n_rows(&self) -> usize {
3249 self.shape.iter().product::<usize>().max(1)
3250 }
3251
3252 /// The row names, in the order [`ObservedRows::n_rows`] counts them.
3253 fn row_names(&self) -> Vec<String> {
3254 if self.shape.is_empty() {
3255 return vec![self.name.clone()];
3256 }
3257 (0..self.n_rows())
3258 .map(|flat| {
3259 let idx = flat_to_multi_col_major(flat, &self.shape)
3260 .iter()
3261 .map(|i| (i + 1).to_string())
3262 .collect::<Vec<_>>()
3263 .join(",");
3264 format!("{}[{}]", self.name, idx)
3265 })
3266 .collect()
3267 }
3268}
3269
3270/// Is this observed rule's value provably an ARRAY (`ndim ≥ 1`) *without*
3271/// evaluating it? Used only to keep a rule out of the 0-D probe in
3272/// [`ArrayCompiled::append_observed_trajectories`], so it must never
3273/// claim "array" for a rule that can materialize 0-D. Both arms read the rank
3274/// straight off the rule's own output box:
3275///
3276/// * an `ArrayLoop` writes a box of rank `output_ranges.len()` on BOTH the
3277/// vectorized fast path and the per-cell oracle (see
3278/// [`materialize_observeds_pass`]);
3279/// * a `Scalar` rule whose body is a `faq` carrying a body (`expr`) and
3280/// a non-empty `output_idx` is evaluated by [`eval_faq`], which returns
3281/// a `Value::Array` of rank `output_idx.len()`. (Without `expr` the oracle
3282/// reports the scalar `NaN` sentinel, hence the `expr.is_some()` guard.)
3283///
3284/// Anything else — a bare variable reference, an arithmetic body, a contraction
3285/// with no free index — is "unknown" and gets probed.
3286pub(super) fn observed_rule_is_array_valued(rule: &AlgebraicRule) -> bool {
3287 match rule {
3288 // A recurrence materializes the same padded box its frame declares
3289 // (esm-spec §4.3.1.1), so it is array-valued exactly when it has axes.
3290 AlgebraicRule::ArrayLoop { output_ranges, .. }
3291 | AlgebraicRule::Recurrence { output_ranges, .. } => !output_ranges.is_empty(),
3292 AlgebraicRule::Scalar { body, .. } => match &**body {
3293 Expr::Operator(node) => {
3294 node.op == "faq"
3295 && node.expr.is_some()
3296 && node.output_idx.as_ref().is_some_and(|ix| !ix.is_empty())
3297 }
3298 _ => false,
3299 },
3300 }
3301}
3302
3303/// Does `expr` carry a dependency edge that a `Expr::Variable` walk cannot see?
3304/// Such an edge makes the observed-trajectory dependency cone unsound, so the
3305/// pruning is disabled wholesale for a model containing one:
3306///
3307/// * `intersect_polygon` / `polygon_intersection_area` register a FAQ overlap
3308/// ring in `derived_rings` under their node id, and a `kind:"derived"`
3309/// contraction bound elsewhere consumes it by that id — not by name;
3310/// * a `distinct` aggregate likewise produces a data-dependent index set;
3311/// * a ragged range names its `offsets` FACTOR as a bare string;
3312/// * a `join` clause couples factors through key COLUMN names.
3313///
3314/// `simpleclimate.esm` (and any pure-stencil model) contains none of these, so
3315/// pruning is live there; a conservative-regrid / relational model keeps the
3316/// old un-pruned behaviour, bit-for-bit.
3317pub(super) fn expr_blocks_output_pruning(expr: &Expr) -> bool {
3318 let Expr::Operator(node) = expr else {
3319 return false;
3320 };
3321 if matches!(
3322 node.op.as_str(),
3323 "intersect_polygon" | "polygon_intersection_area" | "skolem" | "rank" | "distinct"
3324 ) {
3325 return true;
3326 }
3327 if node.join.is_some() || node.distinct == Some(true) {
3328 return true;
3329 }
3330 if let Some(ranges) = &node.ranges
3331 && ranges
3332 .values()
3333 .any(|r| r.ragged().is_some() || r.derived().is_some())
3334 {
3335 return true;
3336 }
3337 node.any_child(&mut expr_blocks_output_pruning)
3338}
3339
3340/// The solve-wide tape program (`None` only under the `interpreter`
3341/// compiler, which builds no tape).
3342#[cfg(feature = "solve")]
3343type SolveTape = Option<std::sync::Arc<TapeProgram>>;
3344
3345/// The run's [`SolutionMetadata`] — one construction for both driver exits.
3346#[cfg(feature = "solve")]
3347fn solution_metadata(
3348 solver_name: &str,
3349 stats: &SolveStats,
3350 merged_variable_renames: HashMap<String, String>,
3351 namespace: Option<String>,
3352) -> SolutionMetadata {
3353 SolutionMetadata {
3354 alg: solver_name.to_string(),
3355 n_rhs_calls: stats.n_rhs_calls,
3356 n_jacobian_calls: stats.n_jacobian_calls,
3357 n_accepted_steps: stats.n_accepted_steps,
3358 n_rejected_steps: stats.n_rejected_steps,
3359 // A rule the tape declined is a refusal before any solve
3360 // ([`ArrayCompiled::strict_tape`]), so a finished solve has none.
3361 tape_fallbacks: Vec::new(),
3362 // Rides to the caller so a name-keyed read of the result can resolve a
3363 // spelling the merge deleted (issue #230).
3364 merged_variable_renames,
3365 namespace,
3366 }
3367}
3368
3369/// The three-tier cadence partition of one solve's observed rules (see
3370/// [`ArrayCompiled::partition_observed_cadence`]).
3371#[cfg(feature = "solve")]
3372struct ObservedCadence {
3373 /// The CONST tier's names — consulted by the build inspection.
3374 static_names: HashSet<String>,
3375 /// CONST: materialized once at setup (the hoisted `static_obs`).
3376 static_rules: Vec<AlgebraicRule>,
3377 /// DISCRETE remainder: materialized once per segment.
3378 segment_static_rules: Vec<AlgebraicRule>,
3379 /// CONTINUOUS: re-evaluated every RHS step.
3380 continuous_rules: Vec<AlgebraicRule>,
3381 /// DISCRETE ∪ CONTINUOUS — the non-hot observed-trajectory output pass.
3382 varying_rules: Vec<AlgebraicRule>,
3383}
3384
3385/// The solve-wide products of the cadence partition and the static hoist,
3386/// shared by every segment and by both driver exits.
3387#[cfg(feature = "solve")]
3388struct SolveSetup {
3389 cadence: ObservedCadence,
3390 /// The t0 state arrays the hoist (and the inspection snapshot) evaluate over.
3391 sa0: ArrMap,
3392 /// The CONST observeds, materialized ONCE — seeded into every RHS eval.
3393 static_obs: ArrMap,
3394}
3395
3396/// One segmented (DISCRETE-cadence) integration's fixed inputs — bundled so
3397/// [`ArrayCompiled::run_segmented`] takes the run plus the refresh callback.
3398#[cfg(feature = "solve")]
3399struct SegmentedRun<'a> {
3400 t0: f64,
3401 t_end: f64,
3402 boundaries: &'a [f64],
3403 ic_vec: &'a [f64],
3404 param_vec: &'a [f64],
3405 setup: &'a SolveSetup,
3406 opts: &'a SolveOptions,
3407 tape: Option<&'a std::sync::Arc<TapeProgram>>,
3408 scratches: Option<&'a SolveScratches>,
3409}
3410
3411/// The right-hand-side and Jacobian scratches of a solve the TAPE serves,
3412/// kept across all of its integration segments.
3413///
3414/// Each segment builds a fresh solver and fresh closures, but not fresh
3415/// scratches: the tape's CONST section primes once per solve (the parameter
3416/// vector does not change between segments), and after the driver refreshes
3417/// the forcing buffer at a segment boundary it bumps each scratch's forcing
3418/// epoch, which re-runs the SEGMENT section, where the DISCRETE forcing loads
3419/// and everything computed from them live, on the next call.
3420///
3421/// Only under a strict compiler: there every rule is taped and the scratch
3422/// carries no per-segment seed, so a kept scratch is exactly the state a fresh
3423/// one would reach. Under the historical routing the scratch is seeded with
3424/// the segment's own materialized observeds, and each segment builds its own.
3425///
3426/// Under `native` the right-hand side and the Jacobian's evaluations run
3427/// through [`RhsProgram`], the public program's evaluation, each with a
3428/// workspace kept here the same way.
3429#[cfg(feature = "solve")]
3430pub(super) struct SolveScratches {
3431 rhs: Rc<RefCell<Option<RhsScratch>>>,
3432 jac: Rc<RefCell<Option<RhsScratch>>>,
3433 program: Option<SolveProgram>,
3434}
3435
3436/// The solve's program and its two kept workspaces.
3437#[cfg(feature = "solve")]
3438#[derive(Clone)]
3439struct SolveProgram {
3440 program: super::range::RhsProgram,
3441 rhs: Rc<RefCell<Option<super::range::RhsWorkspace>>>,
3442 jac: Rc<RefCell<Option<super::range::RhsWorkspace>>>,
3443}
3444
3445#[cfg(feature = "solve")]
3446impl SolveScratches {
3447 fn new(compiled: &ArrayCompiled, tape: &std::sync::Arc<TapeProgram>) -> Self {
3448 let program = (!compiled.is_xla()).then(|| SolveProgram {
3449 program: super::range::RhsProgram::for_solve(compiled, std::sync::Arc::clone(tape)),
3450 rhs: Rc::new(RefCell::new(None)),
3451 jac: Rc::new(RefCell::new(None)),
3452 });
3453 SolveScratches {
3454 rhs: Rc::new(RefCell::new(None)),
3455 jac: Rc::new(RefCell::new(None)),
3456 program,
3457 }
3458 }
3459
3460 /// The driver refreshed the forcing buffer: every scratch built so far
3461 /// re-runs its SEGMENT section on its next call.
3462 fn forcing_refreshed(&self) {
3463 for s in [&self.rhs, &self.jac] {
3464 if let Some(scratch) = s.borrow_mut().as_mut() {
3465 scratch.bump_forcing_epoch();
3466 }
3467 }
3468 for w in self.program.iter().flat_map(|p| [&p.rhs, &p.jac]) {
3469 if let Some(ws) = w.borrow_mut().as_mut() {
3470 ws.forcing_refreshed();
3471 }
3472 }
3473 }
3474}
3475
3476/// The transitive dependency cone of `seeds` within a dependency-ORDERED rule
3477/// slice, returned in the original order (which is therefore still a valid
3478/// evaluation order). One backward pass suffices: a rule's dependencies always
3479/// precede it, so by the time the sweep reaches a rule every consumer of it has
3480/// already been visited and has registered its name.
3481///
3482/// `None` means "the cone is the whole slice" — the caller then keeps using the
3483/// original slice rather than paying a deep AST clone of every rule for nothing.
3484pub(super) fn dependency_cone(
3485 rules: &[AlgebraicRule],
3486 seeds: &HashSet<String>,
3487) -> Option<Vec<AlgebraicRule>> {
3488 let mut needed: HashSet<String> = seeds.clone();
3489 let mut keep = vec![false; rules.len()];
3490 let mut n_kept = 0usize;
3491 for (i, rule) in rules.iter().enumerate().rev() {
3492 if needed.contains(observed_rule_var(rule)) {
3493 keep[i] = true;
3494 n_kept += 1;
3495 collect_expr_var_refs(observed_rule_body(rule), &mut needed);
3496 }
3497 }
3498 if n_kept == rules.len() {
3499 return None;
3500 }
3501 Some(
3502 rules
3503 .iter()
3504 .zip(keep)
3505 .filter(|(_, k)| *k)
3506 .map(|(r, _)| r.clone())
3507 .collect(),
3508 )
3509}
3510
3511/// A field `ic` that could not be resolved, told why the build-time scope
3512/// lacks the observeds its RHS reads, when that is the reason: an observed
3513/// left out of the scope is otherwise just an "unusable RHS".
3514#[cfg(feature = "solve")]
3515fn explain_ic_scope_miss(
3516 err: SimulateError,
3517 rhs: &Expr,
3518 skipped: &HashMap<&str, String>,
3519) -> SimulateError {
3520 let SimulateError::InvalidFieldInitialCondition { name, mut details } = err else {
3521 return err;
3522 };
3523 if !skipped.is_empty() {
3524 let mut reads: Vec<String> = super::compile::free_names(rhs).into_iter().collect();
3525 reads.sort();
3526 for read in reads {
3527 if let Some(why) = skipped.get(read.as_str()) {
3528 details.push_str(&format!(
3529 "; the observed '{read}' is not in the build-time scope: it {why}"
3530 ));
3531 }
3532 }
3533 }
3534 SimulateError::InvalidFieldInitialCondition { name, details }
3535}
3536
3537#[cfg(test)]
3538mod forcing_channel_tests {
3539 //! PR-1 (ess-14f.7): the external refreshable forcing-array channel into the
3540 //! diffsol array RHS. These tests are the bead's acceptance evidence:
3541 //! 1. the RHS reads a forcing array *live* from the buffer,
3542 //! 2. a buffer mutation (a driver refreshing between cadence segments) is
3543 //! reflected in the RHS output, and
3544 //! 3. the existing scalar-`p` / parameter path is unaffected.
3545 //!
3546 //! The forcing buffer is the runtime landing zone for a discrete-cadence
3547 //! loader's regridded field; here it is driven by hand (no I/O), exactly the
3548 //! "testable with a hand-built buffer" contract the plan (PR-1) specifies.
3549 use super::*;
3550 use crate::parse::load_string;
3551
3552 fn arr1(v: &[f64]) -> ArrayD<f64> {
3553 ArrayD::from_shape_vec(IxDyn(&[v.len()]), v.to_vec()).unwrap()
3554 }
3555
3556 /// A model whose state derivative reads an external forcing array `w`
3557 /// elementwise: `D(u[i]) = w[i]`, i ∈ [1,3]. `w` is declared in no variable
3558 /// block — it is a loader-fed field that resolves through the forcing buffer
3559 /// (the new lowest-precedence binding), precisely the channel PR-1 adds.
3560 fn forced_model() -> ArrayCompiled {
3561 let json = r#"
3562 {
3563 "esm": "1.1.0",
3564 "metadata": {
3565 "name": "forcing_channel"
3566 },
3567 "models": {
3568 "Forced": {
3569 "variables": {
3570 "u": {
3571 "type": "unknown",
3572 "shape": [
3573 "i"
3574 ],
3575 "default": 0.0
3576 }
3577 },
3578 "equations": [
3579 {
3580 "lhs": {
3581 "op": "faq",
3582 "args": [],
3583 "output_idx": [
3584 "i"
3585 ],
3586 "expr": {
3587 "op": "D",
3588 "args": [
3589 {
3590 "op": "index",
3591 "args": [
3592 "u",
3593 "i"
3594 ]
3595 }
3596 ],
3597 "wrt": "t"
3598 },
3599 "ranges": {
3600 "i": [
3601 1,
3602 3
3603 ]
3604 }
3605 },
3606 "rhs": {
3607 "op": "faq",
3608 "args": [],
3609 "output_idx": [
3610 "i"
3611 ],
3612 "ranges": {
3613 "i": [
3614 1,
3615 3
3616 ]
3617 },
3618 "expr": {
3619 "op": "index",
3620 "args": [
3621 "w",
3622 "i"
3623 ]
3624 }
3625 }
3626 }
3627 ]
3628 }
3629 }
3630 }
3631 "#;
3632 let file = load_string(json).expect("parse forcing model");
3633 ArrayCompiled::from_file(&file).expect("compile forcing model")
3634 }
3635
3636 #[test]
3637 fn rhs_reads_forcing_array_and_reflects_mutation() {
3638 let compiled = forced_model();
3639 let forcing = compiled.forcing_handle();
3640 let params = HashMap::new();
3641 let state = vec![0.0, 0.0, 0.0];
3642
3643 // Refresh #1 — the RHS reads the forcing array live from the buffer.
3644 forcing
3645 .borrow_mut()
3646 .insert("w".to_string(), arr1(&[10.0, 20.0, 30.0]));
3647 let (dy1, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3648 assert_eq!(
3649 dy1,
3650 vec![10.0, 20.0, 30.0],
3651 "RHS must read the forcing array live from the buffer"
3652 );
3653
3654 // Refresh #2 — a driver mutating the buffer between segments. The change
3655 // is reflected in the RHS output: the channel is live, not build-frozen.
3656 forcing
3657 .borrow_mut()
3658 .insert("w".to_string(), arr1(&[1.0, 2.0, 3.0]));
3659 let (dy2, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3660 assert_eq!(
3661 dy2,
3662 vec![1.0, 2.0, 3.0],
3663 "a buffer mutation must change the RHS output"
3664 );
3665 assert_ne!(
3666 dy1, dy2,
3667 "the refreshed forcing must produce a different RHS"
3668 );
3669
3670 // The per-cell oracle path (force_scalar = true) reads the same buffer —
3671 // the production vectorized path bails forcing reads to this oracle.
3672 let (dy_oracle, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, true);
3673 assert_eq!(
3674 dy_oracle,
3675 vec![1.0, 2.0, 3.0],
3676 "the oracle path resolves forcing identically"
3677 );
3678 }
3679
3680 #[test]
3681 fn forcing_flows_through_the_production_solve() {
3682 // The forcing buffer is captured (Rc clone) into the diffsol RHS closure,
3683 // so a constant forcing `D(u[i]) = w[i]` integrates to `u(t) = u0 + w·t`
3684 // through the real solver — proving the channel is wired into `simulate`,
3685 // not only the debug RHS entry point.
3686 let compiled = forced_model();
3687 compiled
3688 .forcing_handle()
3689 .borrow_mut()
3690 .insert("w".to_string(), arr1(&[2.0, 4.0, 6.0]));
3691 let params = HashMap::new();
3692 let ics = HashMap::new(); // states default to 0
3693 let opts = SolveOptions::default();
3694 let sol = compiled
3695 .solve((0.0, 1.0), ¶ms, &ics, &opts)
3696 .expect("solve with forcing");
3697 // Final state ≈ u0 + w·1 = [2, 4, 6].
3698 for (i, want) in [2.0, 4.0, 6.0].iter().enumerate() {
3699 let got = *sol.state[i].last().expect("trajectory non-empty");
3700 assert!(
3701 (got - want).abs() < 1e-6,
3702 "forcing must drive the solve: state[{i}] got {got}, want {want}"
3703 );
3704 }
3705 }
3706
3707 #[test]
3708 fn empty_forcing_leaves_param_path_unaffected() {
3709 // A parameter+state model `D(u[i]) = k·u[i]` with no forcing reference.
3710 // With an empty buffer the parameter/state path is byte-identical; and an
3711 // *unrelated* forcing entry does not perturb it, because forcing is
3712 // resolved last and only fills otherwise-unbound names.
3713 let json = r#"
3714 {
3715 "esm": "1.1.0",
3716 "metadata": {
3717 "name": "param_path"
3718 },
3719 "models": {
3720 "P": {
3721 "variables": {
3722 "u": {
3723 "type": "unknown",
3724 "shape": [
3725 "i"
3726 ]
3727 },
3728 "k": {
3729 "type": "parameter"
3730 }
3731 },
3732 "equations": [
3733 {
3734 "lhs": {
3735 "op": "faq",
3736 "args": [],
3737 "output_idx": [
3738 "i"
3739 ],
3740 "expr": {
3741 "op": "D",
3742 "args": [
3743 {
3744 "op": "index",
3745 "args": [
3746 "u",
3747 "i"
3748 ]
3749 }
3750 ],
3751 "wrt": "t"
3752 },
3753 "ranges": {
3754 "i": [
3755 1,
3756 2
3757 ]
3758 }
3759 },
3760 "rhs": {
3761 "op": "faq",
3762 "args": [],
3763 "output_idx": [
3764 "i"
3765 ],
3766 "ranges": {
3767 "i": [
3768 1,
3769 2
3770 ]
3771 },
3772 "expr": {
3773 "op": "*",
3774 "args": [
3775 "k",
3776 {
3777 "op": "index",
3778 "args": [
3779 "u",
3780 "i"
3781 ]
3782 }
3783 ]
3784 }
3785 }
3786 }
3787 ]
3788 }
3789 }
3790 }
3791 "#;
3792 let file = load_string(json).expect("parse param model");
3793 let compiled = ArrayCompiled::from_file(&file).expect("compile param model");
3794 let mut params = HashMap::new();
3795 params.insert("k".to_string(), 2.0);
3796 let state = vec![3.0, 5.0];
3797
3798 let (dy_no_forcing, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3799 assert_eq!(
3800 dy_no_forcing,
3801 vec![6.0, 10.0],
3802 "empty forcing leaves the parameter path identical (k·u)"
3803 );
3804
3805 // An unrelated forcing entry must not leak into the parameter path.
3806 compiled
3807 .forcing_handle()
3808 .borrow_mut()
3809 .insert("unrelated".to_string(), arr1(&[99.0]));
3810 let (dy_with_junk, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3811 assert_eq!(
3812 dy_with_junk,
3813 vec![6.0, 10.0],
3814 "an unrelated forcing entry must not perturb the parameter path"
3815 );
3816 }
3817
3818 #[test]
3819 fn fn_op_interp_linear_in_array_runtime() {
3820 // A scalar observed computed via `interp.linear` (a fuel-table lookup,
3821 // as in the coupled fire stack's FuelModelLookup) drives an array
3822 // state: D(u[i]) = looked_up. Before the `fn` arm existed, the observed
3823 // NaN-ed out and poisoned the whole RHS. At code = 2.0 the lookup is
3824 // the exact knot 40.0, so both cells' derivative must be 40.0.
3825 let json = r#"
3826 {
3827 "esm": "1.1.0",
3828 "metadata": {
3829 "name": "fn_array_path"
3830 },
3831 "models": {
3832 "F": {
3833 "variables": {
3834 "u": {
3835 "type": "unknown",
3836 "shape": [
3837 "i"
3838 ]
3839 },
3840 "code": {
3841 "type": "parameter"
3842 },
3843 "looked_up": {
3844 "type": "unknown"
3845 }
3846 },
3847 "equations": [
3848 {
3849 "lhs": {
3850 "op": "faq",
3851 "args": [],
3852 "output_idx": [
3853 "i"
3854 ],
3855 "expr": {
3856 "op": "D",
3857 "args": [
3858 {
3859 "op": "index",
3860 "args": [
3861 "u",
3862 "i"
3863 ]
3864 }
3865 ],
3866 "wrt": "t"
3867 },
3868 "ranges": {
3869 "i": [
3870 1,
3871 2
3872 ]
3873 }
3874 },
3875 "rhs": {
3876 "op": "faq",
3877 "args": [],
3878 "output_idx": [
3879 "i"
3880 ],
3881 "ranges": {
3882 "i": [
3883 1,
3884 2
3885 ]
3886 },
3887 "expr": "looked_up"
3888 }
3889 },
3890 {
3891 "lhs": "looked_up",
3892 "rhs": {
3893 "op": "fn",
3894 "name": "interp.linear",
3895 "args": [
3896 {
3897 "op": "const",
3898 "value": [
3899 10.0,
3900 20.0,
3901 40.0,
3902 80.0,
3903 160.0
3904 ],
3905 "args": []
3906 },
3907 {
3908 "op": "const",
3909 "value": [
3910 0.0,
3911 1.0,
3912 2.0,
3913 3.0,
3914 4.0
3915 ],
3916 "args": []
3917 },
3918 "code"
3919 ]
3920 }
3921 }
3922 ]
3923 }
3924 }
3925 }
3926 "#;
3927 let file = load_string(json).expect("parse fn model");
3928 let compiled = ArrayCompiled::from_file(&file).expect("compile fn model");
3929 let mut params = HashMap::new();
3930 params.insert("code".to_string(), 2.0);
3931 let state = vec![0.0, 0.0];
3932 let (dy, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3933 assert_eq!(
3934 dy,
3935 vec![40.0, 40.0],
3936 "interp.linear(...,2.0)=40 must drive both cells (was NaN before the `fn` arm)"
3937 );
3938
3939 // The blend (not just a knot): code = 0.5 -> 15.0.
3940 params.insert("code".to_string(), 0.5);
3941 let (dy2, _) = compiled.debug_eval_rhs(&state, 0.0, ¶ms, false);
3942 assert_eq!(dy2, vec![15.0, 15.0], "interp.linear(...,0.5)=15");
3943 }
3944}
3945
3946#[cfg(all(test, feature = "solve"))]
3947mod fd_jvp_tests {
3948 //! The finite-difference Jacobian-vector product the XLA arm of the
3949 //! integrator's Jacobian closure runs ([`FdJvp`]): how many right-hand-side
3950 //! evaluations a Jacobian costs, when the kept `f(y)` is thrown away, and
3951 //! that keeping it changes no bit of the product.
3952 use super::*;
3953 use std::cell::Cell;
3954
3955 /// A nonlinear right-hand side reading every argument, so a stale `f(y)`
3956 /// kept across a change of `y`, `p` or `t` would show in the product.
3957 fn f(y: &[f64], p: &[f64], t: f64, out: &mut [f64]) {
3958 let n = y.len();
3959 for (i, o) in out.iter_mut().enumerate() {
3960 *o = p[0] * y[i] * y[(i + 1) % n] - (t * y[i]).sin() + 1.0 / y[i];
3961 }
3962 }
3963
3964 /// The product with both evaluations made every call: the reference a
3965 /// kept `f(y)` must match bit for bit.
3966 fn uncached(y: &[f64], p: &[f64], t: f64, v: &[f64]) -> Vec<f64> {
3967 let n = y.len();
3968 let mut y_norm = 0.0f64;
3969 for &yi in y {
3970 y_norm += yi * yi;
3971 }
3972 let y_norm = y_norm.sqrt().max(1.0);
3973 let eps = f64::EPSILON.sqrt() * y_norm;
3974 let yp: Vec<f64> = (0..n).map(|i| y[i] + eps * v[i]).collect();
3975 let (mut fy, mut fyp) = (vec![0.0; n], vec![0.0; n]);
3976 f(y, p, t, &mut fy);
3977 f(&yp, p, t, &mut fyp);
3978 (0..n).map(|i| (fyp[i] - fy[i]) / eps).collect()
3979 }
3980
3981 fn unit(n: usize, j: usize) -> Vec<f64> {
3982 (0..n).map(|i| if i == j { 1.0 } else { 0.0 }).collect()
3983 }
3984
3985 fn bits(v: &[f64]) -> Vec<u64> {
3986 v.iter().map(|x| x.to_bits()).collect()
3987 }
3988
3989 /// Run one product through `jvp`, returning it and how many evaluations
3990 /// it took.
3991 fn apply(jvp: &mut FdJvp, y: &[f64], p: &[f64], t: f64, v: &[f64]) -> (Vec<f64>, usize) {
3992 let calls = Cell::new(0usize);
3993 let mut jv = vec![0.0; y.len()];
3994 jvp.apply(y, p, t, v, &mut jv, |s, q, t, out| {
3995 calls.set(calls.get() + 1);
3996 f(s, q, t, out);
3997 true
3998 });
3999 (jv, calls.get())
4000 }
4001
4002 /// A dense Jacobian of `n` states, one product per column at one point,
4003 /// costs `n + 1` evaluations rather than `2n`, and every column is
4004 /// bit-identical to evaluating both points every call.
4005 #[test]
4006 fn a_jacobian_at_one_point_evaluates_the_base_once() {
4007 let (y, p, t) = (vec![0.7, 1.3, 2.9, 0.4], vec![1.5], 0.25);
4008 let n = y.len();
4009 let mut jvp = FdJvp::new(n);
4010 let mut total = 0;
4011 for j in 0..n {
4012 let v = unit(n, j);
4013 let (jv, calls) = apply(&mut jvp, &y, &p, t, &v);
4014 total += calls;
4015 assert_eq!(bits(&jv), bits(&uncached(&y, &p, t, &v)), "column {j}");
4016 }
4017 assert_eq!(total, n + 1);
4018 }
4019
4020 /// Any change to the point — one state, the parameters, the time, or
4021 /// only the sign of a zero — evaluates the base again, and the product
4022 /// still matches the uncached one bit for bit.
4023 #[test]
4024 fn a_new_point_evaluates_the_base_again() {
4025 let (y, p, t) = (vec![0.7, 1.3, 0.0], vec![1.5], 0.25);
4026 let v = vec![0.3, -1.1, 0.8];
4027 let mut jvp = FdJvp::new(y.len());
4028 assert_eq!(apply(&mut jvp, &y, &p, t, &v).1, 2);
4029 assert_eq!(apply(&mut jvp, &y, &p, t, &v).1, 1, "same point");
4030
4031 let mut y2 = y.clone();
4032 y2[1] = 1.3000000000000003;
4033 let mut signed_zero = y.clone();
4034 signed_zero[2] = -0.0;
4035 let points: [(&[f64], &[f64], f64, &str); 4] = [
4036 (&y2, &p, t, "one state"),
4037 (&y, &[2.5], t, "the parameters"),
4038 (&y, &p, 0.5, "the time"),
4039 (&signed_zero, &p, t, "the sign of a zero"),
4040 ];
4041 for (yk, pk, tk, what) in points {
4042 // Start each from the original point so each change is the only one.
4043 apply(&mut jvp, &y, &p, t, &v);
4044 let (jv, calls) = apply(&mut jvp, yk, pk, tk, &v);
4045 assert_eq!(calls, 2, "{what}");
4046 assert_eq!(bits(&jv), bits(&uncached(yk, pk, tk, &v)), "{what}");
4047 }
4048 }
4049
4050 /// A failed evaluation of `f(y)` — a device error under `xla` — is not
4051 /// kept: the next product at the same point tries it again.
4052 #[test]
4053 fn a_failed_base_evaluation_is_not_kept() {
4054 let (y, p, t) = (vec![0.7, 1.3], vec![1.5], 0.25);
4055 let v = vec![1.0, 0.0];
4056 let mut jvp = FdJvp::new(y.len());
4057 let mut jv = vec![0.0; 2];
4058 let calls = Cell::new(0usize);
4059 jvp.apply(&y, &p, t, &v, &mut jv, |s, q, t, out| {
4060 calls.set(calls.get() + 1);
4061 if calls.get() == 1 {
4062 out.fill(f64::NAN);
4063 return false;
4064 }
4065 f(s, q, t, out);
4066 true
4067 });
4068 assert_eq!(calls.get(), 2);
4069 assert!(jv.iter().all(|x| x.is_nan()));
4070 let (jv, calls) = apply(&mut jvp, &y, &p, t, &v);
4071 assert_eq!(calls, 2, "the base is evaluated again");
4072 assert_eq!(bits(&jv), bits(&uncached(&y, &p, t, &v)));
4073 }
4074}
4075
4076#[cfg(all(test, feature = "solve"))]
4077mod field_ic_memo_tests {
4078 //! Construction resolves every field initial condition (the compiler gate
4079 //! and the report need it), and the first solve takes that resolution
4080 //! instead of computing it again. Counted with the resolver's test hook.
4081 use super::field_ic_resolutions;
4082 use crate::compile_error::CompileError;
4083 use crate::problem::{Compiler, ProblemOptions, Remake, esm_problem, remake, solve};
4084 use crate::simulate::SimulateError;
4085 use crate::simulate_array::per_cell_cells;
4086 use crate::{SolveOptions, load_string};
4087 use serde_json::{Value, json};
4088
4089 /// `u` over `n` cells with `D(u) = -u` and `ic(u) = rhs_ic`, and a
4090 /// parameter `a = 0.5` the `ic` may read.
4091 fn with_ic(n: usize, rhs_ic: Value) -> crate::EsmFile {
4092 let doc = json!({
4093 "esm": "1.1.0",
4094 "metadata": {"name": "FieldIcMemo"},
4095 "index_sets": {"x": {"kind": "interval", "size": n}},
4096 "models": {"M": {
4097 "variables": {
4098 "u": {"type": "unknown", "units": "1", "shape": ["x"]},
4099 "a": {"type": "parameter", "units": "1", "default": 0.5}
4100 },
4101 "equations": [
4102 {"lhs": {"op": "ic", "args": ["u"]}, "rhs": rhs_ic},
4103 {"lhs": {"op": "faq", "args": [], "output_idx": ["i"],
4104 "ranges": {"i": {"from": "x"}},
4105 "expr": {"op": "D", "args": [{"op": "index", "args": ["u", "i"]}],
4106 "wrt": "t"}},
4107 "rhs": {"op": "faq", "args": [], "output_idx": ["i"],
4108 "ranges": {"i": {"from": "x"}},
4109 "expr": {"op": "-", "args": [{"op": "index", "args": ["u", "i"]}]}}},
4110 ],
4111 }},
4112 });
4113 load_string(&doc.to_string()).expect("loads")
4114 }
4115
4116 /// `ic(u)[i] = a · i`: a coordinate expression over the grid.
4117 fn scaled_index() -> Value {
4118 json!({"op": "faq", "args": [], "output_idx": ["i"],
4119 "ranges": {"i": {"from": "x"}},
4120 "expr": {"op": "*", "args": ["a", "i"]}})
4121 }
4122
4123 fn opts(compiler: Compiler) -> ProblemOptions {
4124 ProblemOptions {
4125 compiler: Some(compiler),
4126 ..Default::default()
4127 }
4128 }
4129
4130 fn first_column(sol: &crate::Solution) -> Vec<f64> {
4131 sol.state.iter().map(|r| r[0]).collect()
4132 }
4133
4134 fn bits(sol: &crate::Solution) -> Vec<Vec<u64>> {
4135 sol.state
4136 .iter()
4137 .map(|r| r.iter().map(|x| x.to_bits()).collect())
4138 .collect()
4139 }
4140
4141 #[test]
4142 fn the_first_solve_takes_the_resolution_construction_made() {
4143 let file = with_ic(3, scaled_index());
4144 for compiler in [Compiler::Native, Compiler::Interpreter] {
4145 let before = field_ic_resolutions();
4146 let prob = esm_problem(&file, (0.0, 1.0), opts(compiler))
4147 .unwrap_or_else(|e| panic!("[{compiler}] {e}"));
4148 assert_eq!(
4149 field_ic_resolutions() - before,
4150 1,
4151 "[{compiler}] construction"
4152 );
4153 assert!(
4154 prob.compiler_report()
4155 .rules()
4156 .iter()
4157 .any(|r| r.kind == "initial condition"),
4158 "[{compiler}] the report keeps its row"
4159 );
4160
4161 let first = solve(&prob, &SolveOptions::default()).expect("solves");
4162 assert_eq!(
4163 field_ic_resolutions() - before,
4164 1,
4165 "[{compiler}] the first solve must not resolve the ics again"
4166 );
4167 assert_eq!(first_column(&first), [0.5, 1.0, 1.5], "[{compiler}]");
4168
4169 // A later solve resolves from the buffer as it stands, to the same
4170 // bits.
4171 let second = solve(&prob, &SolveOptions::default()).expect("solves");
4172 assert_eq!(field_ic_resolutions() - before, 2, "[{compiler}]");
4173 assert_eq!(bits(&first), bits(&second), "[{compiler}]");
4174 }
4175 }
4176
4177 #[test]
4178 fn a_remake_with_new_parameters_resolves_again() {
4179 let file = with_ic(3, scaled_index());
4180 let prob = esm_problem(&file, (0.0, 1.0), opts(Compiler::Native)).expect("builds");
4181 let before = field_ic_resolutions();
4182 let changed = remake(
4183 &prob,
4184 &Remake {
4185 p: [("a".to_string(), 2.0)].into_iter().collect(),
4186 ..Default::default()
4187 },
4188 )
4189 .expect("remakes");
4190 let sol = solve(&changed, &SolveOptions::default()).expect("solves");
4191 assert_eq!(
4192 field_ic_resolutions() - before,
4193 1,
4194 "new parameters, new ics"
4195 );
4196 assert_eq!(first_column(&sol), [2.0, 4.0, 6.0]);
4197 // The memo was keyed to the construction's parameters, and a miss
4198 // spends it: the original resolves too, to its own values.
4199 let sol = solve(&prob, &SolveOptions::default()).expect("solves");
4200 assert_eq!(field_ic_resolutions() - before, 2);
4201 assert_eq!(first_column(&sol), [0.5, 1.0, 1.5]);
4202 }
4203
4204 /// A strict compiler refuses an `ic` the reference evaluator walks per
4205 /// cell, and refuses it after one cell of the walk.
4206 #[test]
4207 fn a_large_per_cell_initial_condition_is_refused_after_one_cell() {
4208 let cumulative = json!({"op": "faq", "args": [], "output_idx": ["i"],
4209 "ranges": {"i": {"from": "x"}, "j": {"from": "x"}},
4210 "filter": {"op": "<=", "args": ["j", "i"]},
4211 "expr": 1.0});
4212 let refusal = |n: usize, ic: Value| {
4213 let file = with_ic(n, ic);
4214 let before = per_cell_cells();
4215 let err = esm_problem(&file, (0.0, 1.0), opts(Compiler::Native)).expect_err("per cell");
4216 let cells = per_cell_cells() - before;
4217 match err {
4218 SimulateError::Compile(CompileError::CompilerRefusedRule {
4219 kind,
4220 rule,
4221 reason,
4222 ..
4223 }) => ((kind, rule, reason), cells),
4224 other => panic!("expected compiler_refused_rule, got {other:?}"),
4225 }
4226 };
4227 let (small, _) = refusal(3, cumulative.clone());
4228 let (large, cells) = refusal(100_000, cumulative.clone());
4229 assert_eq!(large, small);
4230 assert_eq!(large.0, "initial condition");
4231 assert!(
4232 large.2.ends_with(crate::simulate_array::ONE_CELL_NOTE),
4233 "{}",
4234 large.2
4235 );
4236 assert_eq!(cells, 1, "the walk must stop at its first cell");
4237
4238 // What the rest of the expression makes of the stopped walk's
4239 // placeholder is not the document's either: `1 / Σ` over a walk that
4240 // stopped reads a zero, and the infinity that follows must not turn
4241 // the refusal into an unresolvable-ic failure.
4242 let reciprocal = json!({"op": "/", "args": [1.0, cumulative]});
4243 let (large, cells) = refusal(100_000, reciprocal);
4244 assert_eq!(large, small);
4245 assert_eq!(cells, 1);
4246 }
4247
4248 /// A CONST provider serving `w` as zeros over three cells.
4249 struct Zeros;
4250
4251 impl crate::provider::CadenceProvider for Zeros {
4252 fn materialize(
4253 &mut self,
4254 ) -> Result<
4255 std::collections::HashMap<String, crate::provider::NativeField>,
4256 crate::provider::ProviderError,
4257 > {
4258 let zeros = ndarray::ArrayD::zeros(ndarray::IxDyn(&[3]));
4259 Ok([("w".to_string(), crate::provider::NativeField::new(zeros))].into())
4260 }
4261 fn refresh(
4262 &mut self,
4263 _t: f64,
4264 ) -> Result<
4265 Option<std::collections::HashMap<String, crate::provider::NativeField>>,
4266 crate::provider::ProviderError,
4267 > {
4268 Ok(None)
4269 }
4270 fn refresh_times(&self) -> Vec<f64> {
4271 Vec::new()
4272 }
4273 }
4274
4275 fn loaded(compiler: Compiler) -> crate::EsmProblem {
4276 let file = with_ic(3, json!("w"));
4277 let options = ProblemOptions {
4278 providers: [(
4279 "w".to_string(),
4280 Box::new(Zeros) as Box<dyn crate::provider::CadenceProvider>,
4281 )]
4282 .into(),
4283 ..opts(compiler)
4284 };
4285 esm_problem(&file, (0.0, 1.0), options).unwrap_or_else(|e| panic!("[{compiler}] {e}"))
4286 }
4287
4288 /// The memo was resolved from the forcing buffer as construction left it.
4289 /// A host that writes the buffer through `forcing_handle` before the first
4290 /// solve must get initial conditions from what it wrote.
4291 #[test]
4292 fn a_forcing_write_before_the_first_solve_is_not_masked_by_the_memo() {
4293 for compiler in [Compiler::Native, Compiler::Interpreter] {
4294 // Untouched, the first solve takes the construction's resolution.
4295 let prob = loaded(compiler);
4296 let before = field_ic_resolutions();
4297 let sol = solve(&prob, &SolveOptions::default()).expect("solves");
4298 assert_eq!(field_ic_resolutions() - before, 0, "[{compiler}]");
4299 assert_eq!(first_column(&sol), [0.0, 0.0, 0.0], "[{compiler}]");
4300
4301 let prob = loaded(compiler);
4302 let compiled = prob.debug_array_compiled().expect("an array model");
4303 compiled.forcing_handle().borrow_mut().insert(
4304 "w".to_string(),
4305 ndarray::ArrayD::from_shape_vec(ndarray::IxDyn(&[3]), vec![1.0, 2.0, 3.0])
4306 .expect("shape"),
4307 );
4308 let before = field_ic_resolutions();
4309 let sol = solve(&prob, &SolveOptions::default()).expect("solves");
4310 assert_eq!(field_ic_resolutions() - before, 1, "[{compiler}]");
4311 assert_eq!(first_column(&sol), [1.0, 2.0, 3.0], "[{compiler}]");
4312 }
4313 }
4314}