use crate::types::{
AffectEquation, ContinuousEvent, CouplingEntry, DiscreteEvent, DiscreteEventTrigger, Domain,
Equation, EsmFile, Expr, ExpressionNode, IndexSet, InlineValue, JoinClause, Model,
ModelVariable, OverlapClause, RangeSpec, ReactionSystem, VariableMapTransform, VariableType,
};
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeSet, HashMap, HashSet};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum FlattenError {
#[error(
"Conflicting derivative for species {species:?}: explicit D(X, t) equation and reaction participation both present without an operator_compose rule to merge them"
)]
ConflictingDerivative { species: Vec<String> },
#[error("Dimension promotion failed: {message}")]
DimensionPromotion { message: String },
#[error(
"Unmapped domain: systems {systems:?} have different dimensionality but no Interface defines their dimension mapping; candidate target domains: {candidate_targets:?}"
)]
UnmappedDomain {
systems: Vec<String>,
candidate_targets: Vec<String>,
},
#[error(
"Unsupported mapping type '{mapping_type}' at Rust Core tier (supported: broadcast, identity). Reason: {reason}"
)]
UnsupportedMapping {
mapping_type: String,
reason: String,
},
#[error(
"unlowered_operator: rewrite-target operator '{op}' reached compilation without being \
lowered to a stencil by a rewrite rule (esm-spec §4.2 / §9.6.8). Discretization rules \
live in EarthSciDiscretizations, not this format."
)]
UnloweredOperator { op: String },
#[error(
"Domain unit mismatch on variable '{variable}': source units '{source_units}' vs target units '{target_units}'"
)]
DomainUnitMismatch {
variable: String,
source_units: String,
target_units: String,
},
#[error("Domain extent mismatch on independent variable '{variable}' under identity mapping")]
DomainExtentMismatch { variable: String },
#[error(
"Slice out of domain: slice coordinate '{coordinate}' = {value} lies outside the source domain extent"
)]
SliceOutOfDomain { coordinate: String, value: String },
#[error("Cyclic promotion detected involving variables {variables:?}")]
CyclicPromotion { variables: Vec<String> },
#[error(
"variable_map({from} -> {to}): an expression `transform` takes no `factor` (fold the scaling into the expression)"
)]
VariableMapFactorWithExpression { from: String, to: String },
#[error(
"variable_map({from} -> {to}): expression transform does not reference the entry's 'from' variable '{from}' (esm-spec §10.4)"
)]
VariableMapExpressionMissingFrom { from: String, to: String },
#[error(
"variable_map({from} -> {to}): the '{side}' endpoint '{endpoint}' resolves to no variable, \
parameter or observed in the flattened system (esm-spec §4.6, §10.4). A scoped reference \
walks EVERY dot-separated segment, so a subsystem endpoint is spelled \
'<Model>.<Subsystem>.<name>'."
)]
VariableMapUnresolvedEndpoint {
from: String,
to: String,
side: String,
endpoint: String,
},
#[error("Reaction lowering failed: {0}")]
Reaction(#[from] crate::reactions::DeriveError),
#[error("{0}")]
CouplingImport(#[from] crate::diagnostic::DiagnosticError),
#[error(
"Malformed connector equation in couple({systems}): '{side}' is absent or did not deserialize as an expression"
)]
MalformedConnectorEquation { systems: String, side: String },
#[error(
"couple_multiplicative_no_tendency: couple connector 'multiplicative' transform targets \
'{target}', which has no tendency (D({target})) to multiply (esm-spec §10.3). To scale a \
constant parameter by a factor, use a variable_map entry with an Expression transform \
(esm-spec §10.4) instead."
)]
CoupleMultiplicativeNoTendency { target: String },
#[error(
"contributions_out_of_step: FlattenedSystem.{field} has {entries} entries but the system \
has {equations} equations (esm-libraries-spec §4.7.5). Add or remove the entry alongside \
its equation, or clear the list to treat every right-hand side as one contribution of \
unknown origin."
)]
ContributionsOutOfStep {
field: &'static str,
entries: usize,
equations: usize,
},
#[error(
"operator_compose_no_merge: operator_compose({a} + {b}) merged NONE of the {authored} \
equations '{b}' authored; no equation of '{a}' matches: {unmatched}. The entry is \
indistinguishable from one that is not there — '{b}' would integrate decoupled from its \
own defaults and '{a}' would receive no contribution. If '{b}' really does contribute \
only states of its own, declare it with `require_match: false` on this entry."
)]
OperatorComposeNoMerge {
a: String,
b: String,
authored: usize,
unmatched: String,
},
#[error(
"operator_compose_require_match_unmatched: operator_compose({a} + {b}) declares \
`require_match` and merged {merged} of {authored} equations '{b}' authored; no equation \
of '{a}' matches: {unmatched}"
)]
OperatorComposeRequireMatchUnmatched {
a: String,
b: String,
merged: usize,
authored: usize,
unmatched: String,
},
#[error(
"operator_compose_ambiguous_bare_name: operator_compose({a} + {b}) matched '{b_dep}' to \
'{target}' on their shared local name alone, but BOTH are state variables and the merge \
keeps only one. Each carries its own initial condition, so the choice decides which the \
flattened system integrates from, and the document does not express it. Name the \
surviving spelling with a `translate` entry ({{'{target}': '{b_dep}'}}), or rename one \
of them."
)]
OperatorComposeAmbiguousBareName {
a: String,
b: String,
b_dep: String,
target: String,
},
#[error(
"Malformed data-loader subsystem '{subsystem}' in model '{system}': carries loader keys but did not deserialize as a DataSource: {reason}"
)]
MalformedLoaderSubsystem {
system: String,
subsystem: String,
reason: String,
},
#[error("{0}")]
ExpressionTemplate(crate::diagnostic::DiagnosticError),
#[error("No models or reaction systems to flatten")]
Empty,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct DimensionPromotionRecord {
pub variable: String,
pub source_domain: String,
pub target_domain: String,
pub mapping_type: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct FlattenMetadata {
pub source_systems: Vec<String>,
pub coupling_rules_applied: Vec<String>,
pub dimension_promotions_applied: Vec<DimensionPromotionRecord>,
pub implicit_interface_inferred: bool,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub merged_variable_renames: IndexMap<String, String>,
#[serde(default, skip_serializing_if = "BTreeSet::is_empty")]
pub coupling_rewritten_names: BTreeSet<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(untagged)]
pub enum TermOrigin {
Component { component: String },
Coupling { coupling: usize },
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Contribution {
pub expr: Expr,
pub origins: Vec<TermOrigin>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlattenedSystem {
pub independent_variables: Vec<String>,
pub state_variables: IndexMap<String, ModelVariable>,
pub parameters: IndexMap<String, ModelVariable>,
pub observed_variables: IndexMap<String, ModelVariable>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub algebraic_variables: IndexMap<String, ModelVariable>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub brownian_parameters: IndexMap<String, ModelVariable>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub discrete_parameters: IndexMap<String, ModelVariable>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub field_ics: Vec<(String, Expr)>,
pub equations: Vec<Equation>,
#[serde(default)]
pub equation_contributions: Vec<Vec<Contribution>>,
#[serde(default, skip_serializing_if = "none_lifted")]
pub equation_lifted: Vec<bool>,
pub continuous_events: Vec<ContinuousEvent>,
pub discrete_events: Vec<DiscreteEvent>,
pub domain: Option<Domain>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub index_sets: IndexMap<String, IndexSet>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub function_tables: IndexMap<String, crate::types::FunctionTable>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub template_registry: IndexMap<String, serde_json::Value>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub loader_fields: Vec<LoaderField>,
#[serde(default, skip_serializing_if = "IndexMap::is_empty")]
pub lifted_shapes: IndexMap<String, Vec<i64>>,
pub metadata: FlattenMetadata,
}
impl FlattenedSystem {
pub fn system_kind(&self) -> crate::classification::SystemKind {
let mut view: IndexMap<String, ModelVariable> = IndexMap::new();
for (name, var) in &self.state_variables {
view.insert(name.clone(), var.clone());
}
for (name, var) in &self.observed_variables {
view.entry(name.clone()).or_insert_with(|| var.clone());
}
for (name, var) in &self.parameters {
view.entry(name.clone()).or_insert_with(|| var.clone());
}
crate::classification::Classification::from_parts(&view, &self.equations).system_kind
}
pub(crate) fn map_rhs_per_contribution(
&mut self,
mut f: impl FnMut(&Expr) -> Expr,
) -> Result<(), FlattenError> {
let n = self.equations.len();
for (field, entries) in [
("equation_contributions", self.equation_contributions.len()),
("equation_lifted", self.equation_lifted.len()),
] {
if entries != 0 && entries != n {
return Err(FlattenError::ContributionsOutOfStep {
field,
entries,
equations: n,
});
}
}
if self.equation_contributions.is_empty() {
self.equation_contributions = vec![Vec::new(); n];
}
for (i, eq) in self.equations.iter_mut().enumerate() {
let lifted = self.equation_lifted.get(i).copied().unwrap_or(false);
let wrapper = lift_wrapper(&eq.rhs, lifted).cloned();
let body = wrapper
.as_ref()
.and_then(|node| node.expr.as_deref())
.unwrap_or(&eq.rhs);
let contributions = &mut self.equation_contributions[i];
if contributions.is_empty() || fold_contributions(contributions) != *body {
*contributions = vec![Contribution {
expr: body.clone(),
origins: union_origins(contributions, None),
}];
}
for c in contributions.iter_mut() {
c.expr = f(&c.expr);
}
eq.rhs = rebuild_body(wrapper, fold_contributions(contributions));
}
Ok(())
}
}
fn none_lifted(flags: &[bool]) -> bool {
!flags.contains(&true)
}
pub(crate) fn lift_wrapper(rhs: &Expr, lifted: bool) -> Option<&ExpressionNode> {
match rhs {
Expr::Operator(node) if lifted => Some(node),
_ => None,
}
}
pub(crate) fn rebuild_body(wrapper: Option<ExpressionNode>, body: Expr) -> Expr {
match wrapper {
Some(mut node) => {
node.expr = Some(Box::new(body));
Expr::operator(node)
}
None => body,
}
}
pub(crate) fn fold_contributions(contributions: &[Contribution]) -> Expr {
let mut iter = contributions.iter();
let Some(first) = iter.next() else {
return zero();
};
iter.fold(first.expr.clone(), |acc, c| sum_exprs(acc, c.expr.clone()))
}
fn union_origins(contributions: &[Contribution], extra: Option<TermOrigin>) -> Vec<TermOrigin> {
let mut out: Vec<TermOrigin> = Vec::new();
for origin in contributions
.iter()
.flat_map(|c| c.origins.iter().cloned())
.chain(extra)
{
if !out.contains(&origin) {
out.push(origin);
}
}
out
}
#[derive(Debug, Clone)]
struct FlatEq {
eq: Equation,
contributions: Vec<Contribution>,
lifted: bool,
}
impl FlatEq {
fn authored(eq: Equation, origin: TermOrigin) -> Self {
let contributions = vec![Contribution {
expr: eq.rhs.clone(),
origins: vec![origin],
}];
FlatEq {
eq,
contributions,
lifted: false,
}
}
fn body(&self) -> &Expr {
lift_wrapper(&self.eq.rhs, self.lifted)
.and_then(|node| node.expr.as_deref())
.unwrap_or(&self.eq.rhs)
}
fn set_contributions(&mut self, contributions: Vec<Contribution>) {
let body = fold_contributions(&contributions);
self.eq.rhs = rebuild_body(lift_wrapper(&self.eq.rhs, self.lifted).cloned(), body);
self.contributions = contributions;
}
fn map_contributions(&mut self, mut f: impl FnMut(&Expr) -> Expr) {
let mapped = self
.contributions
.iter()
.map(|c| Contribution {
expr: f(&c.expr),
origins: c.origins.clone(),
})
.collect();
self.set_contributions(mapped);
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LoaderField {
pub name: String,
pub owner: String,
pub source: String,
pub file_variable: String,
pub cadence: String,
}
pub fn flatten(file: &EsmFile) -> Result<FlattenedSystem, FlattenError> {
flatten_with_options(
file,
&crate::coupling_imports::CouplingImportOptions::default(),
)
}
pub fn flatten_with_options(
file: &EsmFile,
options: &crate::coupling_imports::CouplingImportOptions,
) -> Result<FlattenedSystem, FlattenError> {
if crate::coupling_imports::has_coupling_import(file) {
let expanded = crate::coupling_imports::expand_coupling_imports(file, options)?;
let mut cloned = file.clone();
cloned.coupling = expanded;
flatten_impl(&cloned)
} else {
flatten_impl(file)
}
}
fn flatten_impl(file: &EsmFile) -> Result<FlattenedSystem, FlattenError> {
let has_models = file.models.as_ref().is_some_and(|m| !m.is_empty());
let has_rs = file
.reaction_systems
.as_ref()
.is_some_and(|rs| !rs.is_empty());
if !has_models && !has_rs {
return Err(FlattenError::Empty);
}
check_variable_map_units(file)?;
let (source_systems, mut per_system) = collect_component_systems(file)?;
check_variable_map_endpoints(file, &per_system)?;
let CouplingOutcome {
rules_applied: coupling_rules_applied,
merged_variable_renames,
coupling_rewritten_names,
} = apply_coupling_entries(file, &mut per_system)?;
detect_conflicts(file, &per_system)?;
let mut parts = assemble_output(per_system);
let mut loaded_producers =
apply_variable_map_removals(file, &mut parts, &merged_variable_renames);
loaded_producers.extend(source_fed_producers(&parts));
maybe_apply_pointwise_lift(file, &mut parts, &loaded_producers)?;
let time_invariant: HashSet<String> = parts.parameters.keys().cloned().collect();
resolve_rhs_time_derivatives_per_contribution(&mut parts.equations, &time_invariant);
extract_field_ics(&mut parts);
let class = flattened_classification(&parts);
let algebraic_variables = in_document_order(
&class.algebraic_unknowns,
&[&parts.state_variables, &parts.observed_variables],
);
let brownian_parameters = in_document_order(&class.brownian_parameters, &[&parts.parameters]);
let discrete_parameters = in_document_order(&class.discrete_parameters, &[&parts.parameters]);
let loader_fields = collect_loader_fields(file, &parts.parameters);
let independent_variables = derive_independent_variables(&parts);
normalize_angle_arguments(&mut parts);
let AssembledParts {
state_variables,
parameters,
observed_variables,
field_ics,
equations,
continuous_events,
discrete_events,
lifted_shapes,
} = parts;
let template_registry = merged_template_registry(file);
check_registry_coupling_rewrites(&template_registry, &coupling_rewritten_names)?;
let equation_lifted: Vec<bool> = if equations.iter().any(|fe| fe.lifted) {
equations.iter().map(|fe| fe.lifted).collect()
} else {
Vec::new()
};
let (equations, equation_contributions): (Vec<Equation>, Vec<Vec<Contribution>>) = equations
.into_iter()
.map(|fe| (fe.eq, fe.contributions))
.unzip();
Ok(FlattenedSystem {
independent_variables,
state_variables,
parameters,
observed_variables,
algebraic_variables,
brownian_parameters,
discrete_parameters,
field_ics,
equations,
equation_contributions,
equation_lifted,
continuous_events,
discrete_events,
domain: file.domain.clone(),
index_sets: file
.index_sets
.as_ref()
.map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
.unwrap_or_default(),
function_tables: file
.function_tables
.as_ref()
.map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
.unwrap_or_default(),
template_registry,
loader_fields,
lifted_shapes,
metadata: FlattenMetadata {
source_systems,
coupling_rules_applied,
dimension_promotions_applied: Vec::new(),
implicit_interface_inferred: false,
merged_variable_renames,
coupling_rewritten_names,
},
})
}
fn normalize_angle_arguments(flat: &mut AssembledParts) {
let mut vars: IndexMap<String, ModelVariable> = IndexMap::new();
for m in [
&flat.state_variables,
&flat.parameters,
&flat.observed_variables,
] {
for (k, v) in m {
vars.entry(k.clone()).or_insert_with(|| v.clone());
}
}
let (env, _) = crate::units::build_unit_env(&vars);
if !env
.values()
.any(|u| crate::units::angle_normalization_factor(u).is_some())
{
return;
}
let rewrite = |e: &mut Expr| {
if let Some(next) = crate::units::normalize_angle_arguments(e, &env) {
*e = next;
}
};
for fe in &mut flat.equations {
let rewritten: Vec<Option<Expr>> = fe
.contributions
.iter()
.map(|c| crate::units::normalize_angle_arguments(&c.expr, &env))
.collect();
if rewritten.iter().all(Option::is_none) {
continue;
}
let contributions = std::mem::take(&mut fe.contributions)
.into_iter()
.zip(rewritten)
.map(|(c, next)| match next {
Some(expr) => Contribution { expr, ..c },
None => c,
})
.collect();
fe.set_contributions(contributions);
}
for (_, e) in &mut flat.field_ics {
rewrite(e);
}
for ev in &mut flat.continuous_events {
for c in &mut ev.conditions {
rewrite(c);
}
for a in &mut ev.affects {
rewrite(&mut a.rhs);
}
for a in ev.affect_neg.iter_mut().flatten() {
rewrite(&mut a.rhs);
}
}
for ev in &mut flat.discrete_events {
if let DiscreteEventTrigger::Condition { expression } = &mut ev.trigger {
rewrite(expression);
}
for a in ev.affects.iter_mut().flatten() {
rewrite(&mut a.rhs);
}
}
}
const TEMPLATE_BODY_NON_OPERAND_KEYS: &[&str] = &[
"op", "wrt", "dim", "fn", "table", "id", "manifold", "name", "semiring", "reduce",
];
fn template_body_var_names(body: &serde_json::Value) -> BTreeSet<String> {
fn walk(node: &serde_json::Value, operand: bool, out: &mut BTreeSet<String>) {
match node {
serde_json::Value::String(name) if operand => {
out.insert(name.clone());
}
serde_json::Value::Array(items) => {
for item in items {
walk(item, operand, out);
}
}
serde_json::Value::Object(map) => {
for (key, value) in map {
walk(
value,
!TEMPLATE_BODY_NON_OPERAND_KEYS.contains(&key.as_str()),
out,
);
}
}
_ => {}
}
}
let mut out = BTreeSet::new();
walk(body, false, &mut out);
out
}
fn check_registry_coupling_rewrites(
registry: &IndexMap<String, serde_json::Value>,
rewritten: &BTreeSet<String>,
) -> Result<(), FlattenError> {
if registry.is_empty() || rewritten.is_empty() {
return Ok(());
}
let mut names: Vec<&String> = registry.keys().collect();
names.sort_unstable();
for tname in names {
let Some(decl) = registry[tname.as_str()].as_object() else {
continue;
};
let Some(body) = decl.get("body") else {
continue;
};
let params: BTreeSet<&str> = decl
.get("params")
.and_then(|p| p.as_array())
.map(|entries| entries.iter().filter_map(|v| v.as_str()).collect())
.unwrap_or_default();
let hits: Vec<String> = template_body_var_names(body)
.into_iter()
.filter(|name| !params.contains(name.as_str()) && rewritten.contains(name))
.collect();
if !hits.is_empty() {
return Err(FlattenError::ExpressionTemplate(crate::diagnostic::err(
crate::diagnostic::codes::TEMPLATE_BODY_REFERENCES_COUPLING_REWRITTEN_VARIABLE,
format!(
"expression template '{tname}' body references '{}', which a coupling rule \
rewrote in the flattened equations; the registry body would expand to a \
stale name at the build boundary. Bind the value through the template's \
params, or expand the reference before coupling (esm-spec §9.6.4).",
hits.join("', '")
),
)));
}
}
Ok(())
}
pub fn merged_template_registry(file: &EsmFile) -> IndexMap<String, serde_json::Value> {
let Some(component_templates) = &file.component_templates else {
return IndexMap::new();
};
let mut ordered: Vec<String> = Vec::new();
if let Some(models) = &file.models {
ordered.extend(models.keys().map(|n| format!("models.{n}")));
}
if let Some(rs) = &file.reaction_systems {
ordered.extend(rs.keys().map(|n| format!("reaction_systems.{n}")));
}
for key in component_templates.keys() {
if !ordered.contains(key) {
ordered.push(key.clone());
}
}
let mut byname: Vec<(String, Vec<(String, serde_json::Value)>)> = Vec::new();
for compkey in &ordered {
let Some(block) = component_templates.get(compkey).and_then(|v| v.as_object()) else {
continue;
};
let (section, cname) = compkey.split_once('.').unwrap_or(("", compkey.as_str()));
let model = (section == "models")
.then(|| file.models.as_ref().and_then(|m| m.get(cname)))
.flatten();
for (tname, decl) in block {
if decl.get("match").is_some_and(|m| !m.is_null()) {
continue; }
let scoped = match (model, decl.get("body")) {
(Some(model), Some(body)) => {
let params: HashSet<String> = decl
.get("params")
.and_then(|p| p.as_array())
.map(|a| {
a.iter()
.filter_map(|v| v.as_str().map(str::to_string))
.collect()
})
.unwrap_or_default();
let mut locals: HashSet<String> = model.variables.keys().cloned().collect();
if let Some(subs) = &model.subsystems {
locals.extend(subs.keys().cloned());
}
for p in ¶ms {
locals.remove(p);
}
match serde_json::from_value::<Expr>(body.clone()) {
Ok(parsed) => {
let scoped_body =
scope_template_body(&parsed, cname, &locals, &HashSet::new());
match serde_json::to_value(&scoped_body) {
Ok(v) => {
let mut d = decl.clone();
if let Some(obj) = d.as_object_mut() {
obj.insert("body".to_string(), v);
}
d
}
Err(_) => decl.clone(),
}
}
Err(_) => decl.clone(),
}
}
_ => decl.clone(),
};
match byname.iter_mut().find(|(k, _)| k == tname) {
Some((_, occ)) => occ.push((cname.to_string(), scoped)),
None => byname.push((tname.clone(), vec![(cname.to_string(), scoped)])),
}
}
}
let collide = crate::lower_expression_templates::registry_collision_names(&byname);
let mut merged: IndexMap<String, serde_json::Value> = IndexMap::new();
let mut rename: HashMap<String, HashMap<String, String>> = HashMap::new();
for (name, occ) in &byname {
if collide.contains(name) {
for (path, decl) in occ {
let newname = format!("{path}.{name}");
merged.insert(newname.clone(), decl.clone());
rename
.entry(path.clone())
.or_default()
.insert(name.clone(), newname);
}
} else {
merged.insert(name.clone(), occ[0].1.clone()); }
}
for per_owner in rename.values() {
for new in per_owner.values() {
if let Some(decl) = merged.get_mut(new) {
crate::lower_expression_templates::rename_apply_refs(decl, per_owner);
}
}
}
merged
}
fn scope_template_body(
expr: &Expr,
prefix: &str,
locals: &HashSet<String>,
bound: &HashSet<String>,
) -> Expr {
match expr {
Expr::Number(n) => Expr::Number(*n),
Expr::Integer(n) => Expr::Integer(*n),
Expr::Variable(name) => {
if bound.contains(name) {
return Expr::Variable(name.clone());
}
let head = name.split('.').next().unwrap_or(name.as_str());
if locals.contains(head) {
Expr::Variable(format!("{prefix}.{name}"))
} else {
Expr::Variable(name.clone())
}
}
Expr::Operator(node) => {
let mut child_bound = bound.clone();
if let Some(output_idx) = &node.output_idx {
child_bound.extend(output_idx.iter().cloned());
}
if let Some(ranges) = &node.ranges {
child_bound.extend(ranges.keys().cloned());
}
let mut out =
node.map_children(&mut |c| scope_template_body(c, prefix, locals, &child_bound));
if let Some(join) = &node.join {
let mut binders: HashSet<&str> = HashSet::new();
if let Some(output_idx) = &node.output_idx {
binders.extend(output_idx.iter().map(String::as_str));
}
if let Some(ranges) = &node.ranges {
binders.extend(ranges.keys().map(String::as_str));
}
if let Some(ns) = namespace_join_names(join, &binders, prefix, locals) {
out.join = Some(ns);
}
}
Expr::operator(out)
}
}
}
fn derive_independent_variables(parts: &AssembledParts) -> Vec<String> {
let mut out = vec!["t".to_string()];
let mut scan = |e: &Expr| collect_spatial_dims(e, &mut out);
for fe in &parts.equations {
scan(&fe.eq.lhs);
scan(&fe.eq.rhs);
}
for (_, rhs) in &parts.field_ics {
scan(rhs);
}
out
}
fn collect_spatial_dims(expr: &Expr, out: &mut Vec<String>) {
let Expr::Operator(node) = expr else { return };
if let Some(dim) = &node.dim
&& !out.iter().any(|d| d == dim)
{
out.push(dim.clone());
}
if let Some(wrt) = &node.wrt
&& wrt != "t"
&& !out.iter().any(|d| d == wrt)
{
out.push(wrt.clone());
}
node.for_each_child(&mut |child| collect_spatial_dims(child, out));
}
pub fn reject_unlowered_operators(flat: &FlattenedSystem) -> Result<(), FlattenError> {
match first_unlowered_operator(flat) {
Some(op) => Err(FlattenError::UnloweredOperator { op }),
None => Ok(()),
}
}
pub fn first_unlowered_operator(flat: &FlattenedSystem) -> Option<String> {
for eq in &flat.equations {
for side in [&eq.lhs, &eq.rhs] {
if let Err(FlattenError::UnloweredOperator { op }) = reject_spatial_operators(side) {
return Some(op);
}
}
}
None
}
pub fn first_unresolved_rhs_time_derivative(flat: &FlattenedSystem) -> Option<String> {
first_unresolved_rhs_time_derivative_in(&flat.equations)
}
pub(crate) fn first_unresolved_rhs_time_derivative_in(equations: &[Equation]) -> Option<String> {
for eq in equations {
if let Some(target) = first_rhs_time_derivative(&eq.rhs) {
return Some(target);
}
}
None
}
fn first_rhs_time_derivative(expr: &Expr) -> Option<String> {
let Expr::Operator(node) = expr else {
return None;
};
if is_structural_time_derivative(node) {
return Some(match &node.args[0] {
Expr::Variable(name) => name.clone(),
other => format!("{other}"),
});
}
let mut found: Option<String> = None;
node.for_each_child(&mut |child| {
if found.is_none() {
found = first_rhs_time_derivative(child);
}
});
found
}
fn collect_loader_fields(
file: &EsmFile,
parameters: &IndexMap<String, ModelVariable>,
) -> Vec<LoaderField> {
let mut out = Vec::new();
for (name, var) in parameters {
let Some(spec) = &var.update else { continue };
let rules: &[crate::types::ParameterUpdate] = match spec {
crate::types::ParameterUpdateSpec::Single(rule) => std::slice::from_ref(rule),
crate::types::ParameterUpdateSpec::Several(rules) => rules,
};
for rule in rules {
let Some(source) = rule.data_source() else {
continue;
};
let Some(binding) = rule.value().and_then(|v| v.from.as_ref()) else {
continue;
};
let has_temporal = file
.data_sources
.as_ref()
.and_then(|ds| ds.get(source))
.is_some_and(|ds| ds.temporal.is_some());
out.push(LoaderField {
name: name.clone(),
owner: name
.rsplit_once('.')
.map(|(owner, _)| owner.to_string())
.unwrap_or_default(),
source: source.to_string(),
file_variable: binding.file_variable.clone(),
cadence: if has_temporal { "discrete" } else { "const" }.to_string(),
});
break; }
}
out
}
fn collect_component_systems(
file: &EsmFile,
) -> Result<(Vec<String>, Vec<SystemBlock>), FlattenError> {
let mut source_systems = Vec::new();
let mut per_system: Vec<SystemBlock> = Vec::new();
if let Some(models) = &file.models {
for (name, model) in models {
let block = build_model_block(name, model)?;
source_systems.push(name.clone());
per_system.push(block);
}
}
if let Some(rsystems) = &file.reaction_systems {
for (name, rs) in rsystems {
let block = build_reaction_block(name, rs)?;
source_systems.push(name.clone());
per_system.push(block);
}
}
Ok((source_systems, per_system))
}
struct CouplingOutcome {
rules_applied: Vec<String>,
merged_variable_renames: IndexMap<String, String>,
coupling_rewritten_names: BTreeSet<String>,
}
fn apply_coupling_entries(
file: &EsmFile,
per_system: &mut Vec<SystemBlock>,
) -> Result<CouplingOutcome, FlattenError> {
let mut coupling_rules_applied = Vec::new();
let mut coupling_rewritten_names: BTreeSet<String> = BTreeSet::new();
let mut merged_renames: IndexMap<String, String> = IndexMap::new();
let Some(entries) = &file.coupling else {
return Ok(CouplingOutcome {
rules_applied: coupling_rules_applied,
merged_variable_renames: merged_renames,
coupling_rewritten_names,
});
};
let kind_rank = |entry: &CouplingEntry| match entry {
CouplingEntry::OperatorCompose { .. } => 0,
CouplingEntry::Couple { .. } => 1,
_ => 2,
};
let mut descriptions: Vec<Option<String>> = vec![None; entries.len()];
for rank in 0..=2 {
for (i, entry) in entries.iter().enumerate() {
if kind_rank(entry) != rank {
continue;
}
let mut one = Vec::new();
apply_coupling_entry(
entry,
i,
per_system,
&mut one,
&mut merged_renames,
&mut coupling_rewritten_names,
)?;
descriptions[i] = one.into_iter().next();
}
}
coupling_rules_applied.extend(descriptions.into_iter().flatten());
coupling_rewritten_names.extend(merged_renames.keys().cloned());
Ok(CouplingOutcome {
rules_applied: coupling_rules_applied,
merged_variable_renames: merged_renames,
coupling_rewritten_names,
})
}
fn compose_renames(acc: &mut IndexMap<String, String>, next: IndexMap<String, String>) {
for survivor in acc.values_mut() {
if let Some(further) = next.get(survivor.as_str()) {
*survivor = further.clone();
}
}
acc.extend(next);
}
fn resolve_renamed<'a>(name: &'a str, renames: &'a IndexMap<String, String>) -> &'a str {
renames.get(name).map(String::as_str).unwrap_or(name)
}
fn retarget_connector<'a>(
connector: &'a serde_json::Value,
renames: &IndexMap<String, String>,
) -> std::borrow::Cow<'a, serde_json::Value> {
use std::borrow::Cow;
if renames.is_empty() {
return Cow::Borrowed(connector);
}
let Some(eqs) = connector.get("equations").and_then(|e| e.as_array()) else {
return Cow::Borrowed(connector);
};
let subs: HashMap<String, Expr> = renames
.iter()
.map(|(from, to)| (from.clone(), Expr::Variable(to.clone())))
.collect();
let mut changed = false;
let mut out_eqs = Vec::with_capacity(eqs.len());
for eq in eqs {
let mut eq = eq.clone();
if let Some(obj) = eq.as_object_mut() {
for side in ["from", "to"] {
let Some(name) = obj.get(side).and_then(|v| v.as_str()) else {
continue;
};
let survivor = resolve_renamed(name, renames);
if survivor != name {
let survivor = survivor.to_string();
obj.insert(side.to_string(), serde_json::Value::String(survivor));
changed = true;
}
}
if let Some(raw) = obj.get("expression")
&& let Ok(expr) = serde_json::from_value::<Expr>(raw.clone())
{
let rewritten = rename_names(&expr, &subs, renames);
if rewritten != expr
&& let Ok(v) = serde_json::to_value(&rewritten)
{
obj.insert("expression".to_string(), v);
changed = true;
}
}
}
out_eqs.push(eq);
}
if !changed {
return Cow::Borrowed(connector);
}
let mut out = connector.clone();
if let Some(obj) = out.as_object_mut() {
obj.insert("equations".to_string(), serde_json::Value::Array(out_eqs));
}
Cow::Owned(out)
}
fn detect_conflicts(file: &EsmFile, per_system: &[SystemBlock]) -> Result<(), FlattenError> {
let operator_compose_systems: Vec<Vec<String>> = file
.coupling
.as_ref()
.map(|entries| {
entries
.iter()
.filter_map(|e| match e {
CouplingEntry::OperatorCompose { systems, .. } => Some(systems.clone()),
_ => None,
})
.collect()
})
.unwrap_or_default();
let mut lhs_targets: IndexMap<String, Vec<String>> = IndexMap::new();
for block in per_system {
for FlatEq { eq, .. } in &block.equations {
if let Some(dep) = extract_ddt_dependent(&eq.lhs) {
lhs_targets.entry(dep).or_default().push(block.name.clone());
}
}
}
let mut conflicting_species: Vec<String> = Vec::new();
for (species, owning_systems) in &lhs_targets {
if owning_systems.len() < 2 {
continue;
}
let was_composed = operator_compose_systems
.iter()
.any(|compose_systems| owning_systems.iter().all(|s| compose_systems.contains(s)));
if !was_composed {
conflicting_species.push(species.clone());
}
}
if !conflicting_species.is_empty() {
conflicting_species.sort();
conflicting_species.dedup();
return Err(FlattenError::ConflictingDerivative {
species: conflicting_species,
});
}
Ok(())
}
struct AssembledParts {
state_variables: IndexMap<String, ModelVariable>,
parameters: IndexMap<String, ModelVariable>,
observed_variables: IndexMap<String, ModelVariable>,
field_ics: Vec<(String, Expr)>,
equations: Vec<FlatEq>,
continuous_events: Vec<ContinuousEvent>,
discrete_events: Vec<DiscreteEvent>,
lifted_shapes: IndexMap<String, Vec<i64>>,
}
fn assemble_output(per_system: Vec<SystemBlock>) -> AssembledParts {
let mut parts = AssembledParts {
state_variables: IndexMap::new(),
parameters: IndexMap::new(),
observed_variables: IndexMap::new(),
field_ics: Vec::new(),
equations: Vec::new(),
continuous_events: Vec::new(),
discrete_events: Vec::new(),
lifted_shapes: IndexMap::new(),
};
for block in per_system {
for (name, var) in block.state_vars {
parts.state_variables.insert(name, var);
}
for (name, var) in block.parameters {
parts.parameters.insert(name, var);
}
for (name, var) in block.observed_vars {
parts.observed_variables.insert(name, var);
}
parts.equations.extend(block.equations);
parts.continuous_events.extend(block.continuous_events);
parts.discrete_events.extend(block.discrete_events);
}
parts
}
fn extract_field_ics(parts: &mut AssembledParts) {
let mut kept = Vec::with_capacity(parts.equations.len());
for fe in std::mem::take(&mut parts.equations) {
if let Some(target) = extract_ic_target(&fe.eq.lhs) {
parts.field_ics.push((target, fe.eq.rhs));
} else {
kept.push(fe);
}
}
parts.equations = kept;
}
fn source_fed_producers(parts: &AssembledParts) -> HashMap<String, usize> {
let class = flattened_classification(parts);
parts
.parameters
.iter()
.filter(|(name, _)| class.is_discrete_parameter(name))
.filter_map(|(name, var)| {
var.shape
.as_ref()
.filter(|s| !s.is_empty())
.map(|s| (name.clone(), s.len()))
})
.collect()
}
fn flattened_classification(parts: &AssembledParts) -> crate::classification::Classification {
let mut view: IndexMap<String, ModelVariable> = IndexMap::new();
for (name, var) in &parts.state_variables {
view.insert(name.clone(), var.clone());
}
for (name, var) in &parts.observed_variables {
view.entry(name.clone()).or_insert_with(|| var.clone());
}
for (name, var) in &parts.parameters {
view.entry(name.clone()).or_insert_with(|| var.clone());
}
let equations: Vec<Equation> = parts.equations.iter().map(|fe| fe.eq.clone()).collect();
crate::classification::Classification::from_parts(&view, &equations)
}
fn in_document_order(
names: &[String],
maps: &[&IndexMap<String, ModelVariable>],
) -> IndexMap<String, ModelVariable> {
let wanted: HashSet<&str> = names.iter().map(String::as_str).collect();
let mut out = IndexMap::new();
for m in maps {
for (name, var) in *m {
if wanted.contains(name.as_str()) && !out.contains_key(name) {
out.insert(name.clone(), var.clone());
}
}
}
out
}
fn apply_variable_map_removals(
file: &EsmFile,
parts: &mut AssembledParts,
merged_renames: &IndexMap<String, String>,
) -> HashMap<String, usize> {
let loader_names: HashSet<String> = file
.data_sources
.as_ref()
.map(|dl| dl.keys().cloned().collect())
.unwrap_or_default();
let mut loaded_producers: HashMap<String, usize> = HashMap::new();
if let Some(entries) = &file.coupling {
for (index, entry) in entries.iter().enumerate() {
let CouplingEntry::VariableMap {
from,
to,
transform,
..
} = entry
else {
continue;
};
match transform {
VariableMapTransform::Named(name)
if matches!(name.as_str(), "param_to_var" | "conversion_factor") =>
{
let consumer_shape_rank = parts
.parameters
.get(to)
.and_then(|v| v.shape.as_ref())
.map(|s| s.len())
.filter(|r| *r > 0);
parts.parameters.shift_remove(to);
let from_owner = from.split('.').next().unwrap_or("");
if let Some(rank) = consumer_shape_rank
&& loader_names.contains(from_owner)
&& !parts.parameters.contains_key(from)
{
loaded_producers.insert(from.clone(), rank);
}
}
VariableMapTransform::Expression(node) => {
let removed = parts.parameters.shift_remove(to);
let (units, shape, description) = removed
.map(|p| (p.units, p.shape, p.description))
.unwrap_or((None, None, None));
parts.observed_variables.insert(
to.clone(),
ModelVariable {
var_type: VariableType::Unknown,
units,
description,
shape,
..Default::default()
},
);
let subs: HashMap<String, Expr> = merged_renames
.iter()
.map(|(gone, survivor)| (gone.clone(), Expr::Variable(survivor.clone())))
.collect();
let rhs = rename_names(&Expr::operator(node.clone()), &subs, merged_renames);
parts.equations.push(FlatEq::authored(
Equation {
comment: None,
lhs: Expr::Variable(to.clone()),
rhs,
},
TermOrigin::Coupling { coupling: index },
));
}
VariableMapTransform::Named(_) => {}
}
}
}
loaded_producers
}
fn maybe_apply_pointwise_lift(
file: &EsmFile,
parts: &mut AssembledParts,
loaded_producers: &HashMap<String, usize>,
) -> Result<(), FlattenError> {
let pointwise = file
.coupling
.as_ref()
.map(|entries| {
entries.iter().any(|e| {
matches!(e, CouplingEntry::OperatorCompose { lifting: Some(l), .. } if l == "pointwise")
})
})
.unwrap_or(false);
if pointwise {
apply_pointwise_lift(
&mut parts.equations,
&mut parts.lifted_shapes,
loaded_producers,
)?;
}
Ok(())
}
fn is_structural_time_derivative(node: &ExpressionNode) -> bool {
node.op == "D"
&& node.args.len() == 1
&& !crate::op_registry::is_rewrite_target_derivative(node)
}
pub(crate) fn resolve_rhs_time_derivatives(
equations: &mut [Equation],
time_invariant: &HashSet<String>,
) -> bool {
let Some((carriers, tables)) =
plan_time_derivative_resolution(equations.iter(), time_invariant)
else {
return false;
};
let Some(tables) = tables else {
return true;
};
for (eq, carries) in equations.iter_mut().zip(carriers) {
if carries {
let mut active = self_defined(&eq.lhs);
eq.rhs = tables.substitute(&eq.rhs, &mut active);
}
}
true
}
fn plan_time_derivative_resolution<'e, 'a>(
equations: impl Iterator<Item = &'e Equation> + Clone,
time_invariant: &'a HashSet<String>,
) -> Option<(Vec<bool>, Option<DerivTables<'a>>)> {
let carriers: Vec<bool> = equations
.clone()
.map(|eq| contains_time_derivative(&eq.rhs))
.collect();
if !carriers.contains(&true) {
return None;
}
Some((carriers, DerivTables::build(equations, time_invariant)))
}
fn resolve_rhs_time_derivatives_per_contribution(
equations: &mut [FlatEq],
time_invariant: &HashSet<String>,
) {
let Some((carriers, Some(tables))) =
plan_time_derivative_resolution(equations.iter().map(|fe| &fe.eq), time_invariant)
else {
return;
};
for (fe, carries) in equations.iter_mut().zip(carriers) {
if carries {
let mut active = self_defined(&fe.eq.lhs);
fe.map_contributions(|e| tables.substitute(e, &mut active));
}
}
}
fn self_defined(lhs: &Expr) -> Vec<String> {
match lhs {
Expr::Operator(node) if is_structural_time_derivative(node) => match node.args.first() {
Some(Expr::Variable(name)) => vec![name.clone()],
_ => Vec::new(),
},
Expr::Variable(name) => vec![name.clone()],
_ => Vec::new(),
}
}
struct DerivTables<'a> {
tendency: HashMap<String, Expr>,
definition: HashMap<String, Expr>,
time_invariant: &'a HashSet<String>,
}
impl<'a> DerivTables<'a> {
fn build<'e>(
equations: impl IntoIterator<Item = &'e Equation>,
time_invariant: &'a HashSet<String>,
) -> Option<Self> {
let mut tables = DerivTables {
tendency: HashMap::new(),
definition: HashMap::new(),
time_invariant,
};
for eq in equations {
match &eq.lhs {
Expr::Operator(node) if is_structural_time_derivative(node) => {
if let Some(Expr::Variable(name)) = node.args.first() {
tables.tendency.insert(name.clone(), eq.rhs.clone());
}
}
Expr::Variable(name) => {
tables.definition.insert(name.clone(), eq.rhs.clone());
}
_ => {}
}
}
(!tables.tendency.is_empty() || !tables.definition.is_empty()).then_some(tables)
}
fn substitute(&self, expr: &Expr, active: &mut Vec<String>) -> Expr {
let Expr::Operator(node) = expr else {
return expr.clone();
};
if is_structural_time_derivative(node) {
return match node.args.first().and_then(|a| self.deriv(a, active)) {
Some(out) => out,
None => expr.clone(),
};
}
Expr::operator(node.map_children(&mut |child| self.substitute(child, active)))
}
fn deriv(&self, expr: &Expr, active: &mut Vec<String>) -> Option<Expr> {
match expr {
Expr::Integer(_) | Expr::Number(_) => Some(zero()),
Expr::Variable(name) => {
if active.iter().any(|n| n == name) {
return None;
}
if let Some(f) = self.tendency.get(name) {
let f = f.clone();
active.push(name.clone());
let out = self.substitute(&f, active);
active.pop();
Some(out)
} else if let Some(g) = self.definition.get(name) {
let g = g.clone();
active.push(name.clone());
let out = self.deriv(&g, active);
active.pop();
out
} else if self.time_invariant.contains(name) {
Some(zero())
} else {
None
}
}
Expr::Operator(node) => self.deriv_op(node, active),
}
}
fn deriv_op(&self, node: &ExpressionNode, active: &mut Vec<String>) -> Option<Expr> {
let d = |a: &Expr, active: &mut Vec<String>| self.deriv(a, active);
match node.op.as_str() {
"+" => {
let mut terms = Vec::with_capacity(node.args.len());
for a in &node.args {
terms.push(d(a, active)?);
}
Some(sum(terms))
}
"-" => match node.args.len() {
1 => Some(negate(d(&node.args[0], active)?)),
2 => {
let a = d(&node.args[0], active)?;
let b = d(&node.args[1], active)?;
Some(difference(a, b))
}
_ => None,
},
"neg" if node.args.len() == 1 => Some(negate(d(&node.args[0], active)?)),
"*" => {
let mut terms = Vec::with_capacity(node.args.len());
for (i, a) in node.args.iter().enumerate() {
let da = d(a, active)?;
if is_zero(&da) {
continue;
}
let mut factors = vec![da];
for (j, b) in node.args.iter().enumerate() {
if i != j {
factors.push(b.clone());
}
}
terms.push(product(factors));
}
Some(sum(terms))
}
"/" if node.args.len() == 2 => {
let (u, v) = (&node.args[0], &node.args[1]);
let du = d(u, active)?;
let dv = d(v, active)?;
if is_zero(&dv) {
return Some(quotient(du, v.clone()));
}
let num = difference(product(vec![du, v.clone()]), product(vec![u.clone(), dv]));
Some(quotient(num, product(vec![v.clone(), v.clone()])))
}
_ => None,
}
}
}
fn contains_time_derivative(expr: &Expr) -> bool {
let Expr::Operator(node) = expr else {
return false;
};
is_structural_time_derivative(node) || node.any_child(&mut contains_time_derivative)
}
fn zero() -> Expr {
Expr::Number(0.0)
}
fn is_zero(e: &Expr) -> bool {
matches!(e, Expr::Number(v) if *v == 0.0) || matches!(e, Expr::Integer(0))
}
fn is_one(e: &Expr) -> bool {
matches!(e, Expr::Number(v) if *v == 1.0) || matches!(e, Expr::Integer(1))
}
fn node(op: &str, args: Vec<Expr>) -> Expr {
Expr::operator(ExpressionNode {
op: op.to_string(),
args,
..Default::default()
})
}
fn sum(terms: Vec<Expr>) -> Expr {
let mut kept: Vec<Expr> = terms.into_iter().filter(|t| !is_zero(t)).collect();
match kept.len() {
0 => zero(),
1 => kept.pop().expect("len 1"),
_ => node("+", kept),
}
}
fn difference(a: Expr, b: Expr) -> Expr {
if is_zero(&b) {
return a;
}
if is_zero(&a) {
return negate(b);
}
node("-", vec![a, b])
}
fn negate(a: Expr) -> Expr {
match &a {
_ if is_zero(&a) => zero(),
Expr::Number(v) => Expr::Number(-v),
Expr::Integer(v) => Expr::Integer(-v),
_ => node("-", vec![a]),
}
}
fn product(factors: Vec<Expr>) -> Expr {
if factors.iter().any(is_zero) {
return zero();
}
let mut kept: Vec<Expr> = factors.into_iter().filter(|f| !is_one(f)).collect();
match kept.len() {
0 => Expr::Number(1.0),
1 => kept.pop().expect("len 1"),
_ => node("*", kept),
}
}
fn quotient(a: Expr, b: Expr) -> Expr {
if is_zero(&a) {
return zero();
}
node("/", vec![a, b])
}
pub fn flatten_model(model: &Model) -> Result<FlattenedSystem, FlattenError> {
let system_name = model.name.clone().unwrap_or_else(|| "model".to_string());
let mut models = IndexMap::new();
models.insert(system_name, model.clone());
let file = EsmFile {
models: Some(models),
..Default::default()
};
flatten(&file)
}
struct SystemBlock {
name: String,
state_vars: IndexMap<String, ModelVariable>,
parameters: IndexMap<String, ModelVariable>,
observed_vars: IndexMap<String, ModelVariable>,
equations: Vec<FlatEq>,
continuous_events: Vec<ContinuousEvent>,
discrete_events: Vec<DiscreteEvent>,
}
fn build_model_block(system_name: &str, model: &Model) -> Result<SystemBlock, FlattenError> {
let mut state_vars = IndexMap::new();
let mut parameters = IndexMap::new();
let mut observed_vars = IndexMap::new();
let sub_keys: HashSet<String> = model
.subsystems
.as_ref()
.map(|subs| subs.keys().cloned().collect())
.unwrap_or_default();
let locals: HashSet<String> = model
.variables
.keys()
.cloned()
.chain(sub_keys.iter().cloned())
.collect();
let class = crate::classification::Classification::of(model);
for (var_name, var) in &model.variables {
let namespaced = format!("{system_name}.{var_name}");
let mut cloned = var.clone();
cloned.for_each_expression_mut(&mut |expr| {
*expr = namespace_expr(expr, system_name, &sub_keys, &locals);
});
match var.var_type {
VariableType::Unknown => {
if class.is_observed(var_name) {
if !class.is_inlined(var_name) {
state_vars.insert(namespaced.clone(), cloned.clone());
}
observed_vars.insert(namespaced, cloned);
} else {
state_vars.insert(namespaced, cloned);
}
}
VariableType::Parameter => {
parameters.insert(namespaced, cloned);
}
}
}
let equations: Vec<FlatEq> = model
.equations
.iter()
.map(|eq| {
FlatEq::authored(
Equation {
lhs: namespace_expr(&eq.lhs, system_name, &sub_keys, &locals),
rhs: namespace_expr(&eq.rhs, system_name, &sub_keys, &locals),
comment: eq.comment.clone(),
},
TermOrigin::Component {
component: system_name.to_string(),
},
)
})
.collect();
let continuous_events = model
.continuous_events
.clone()
.unwrap_or_default()
.into_iter()
.map(|e| namespace_continuous_event(e, system_name, &sub_keys, &locals))
.collect();
let discrete_events = model
.discrete_events
.clone()
.unwrap_or_default()
.into_iter()
.map(|e| namespace_discrete_event(e, system_name, &sub_keys, &locals))
.collect();
let mut block = SystemBlock {
name: system_name.to_string(),
state_vars,
parameters,
observed_vars,
equations,
continuous_events,
discrete_events,
};
if let Some(subs) = &model.subsystems {
for (sub_name, sub_value) in subs {
let Ok(sub_model) = serde_json::from_value::<Model>(sub_value.clone()) else {
continue;
};
let sub_block = build_model_block(&format!("{system_name}.{sub_name}"), &sub_model)?;
block.state_vars.extend(sub_block.state_vars);
block.parameters.extend(sub_block.parameters);
block.observed_vars.extend(sub_block.observed_vars);
block.equations.extend(sub_block.equations);
}
}
Ok(block)
}
fn build_reaction_block(
system_name: &str,
rs: &ReactionSystem,
) -> Result<SystemBlock, FlattenError> {
let mut state_vars = IndexMap::new();
let mut parameters = IndexMap::new();
for (species_name, species) in &rs.species {
let namespaced = format!("{system_name}.{species_name}");
let var = ModelVariable {
var_type: if species.constant == Some(true) {
VariableType::Parameter
} else {
VariableType::Unknown
},
units: species.units.clone(),
default: species.default.map(InlineValue::Scalar),
description: species.description.clone(),
..Default::default()
};
if species.constant == Some(true) {
parameters.insert(namespaced, var);
} else {
state_vars.insert(namespaced, var);
}
}
for (param_name, param) in &rs.parameters {
let namespaced = format!("{system_name}.{param_name}");
parameters.insert(
namespaced,
ModelVariable {
var_type: VariableType::Parameter,
units: param.units.clone(),
default: param.default.map(InlineValue::Scalar),
description: param.description.clone(),
default_units: param.default_units.clone(),
shape: param.shape.clone(),
distribution: param.distribution.clone(),
update: param.update.clone(),
..Default::default()
},
);
}
let locals: HashSet<String> = rs
.species
.keys()
.chain(rs.parameters.keys())
.cloned()
.collect();
let lowered = crate::reactions::lower_reactions_to_equations(&rs.reactions, &rs.species)?;
let equations = lowered
.into_iter()
.map(|eq| {
FlatEq::authored(
Equation {
lhs: namespace_expr(&eq.lhs, system_name, &HashSet::new(), &locals),
rhs: namespace_expr(&eq.rhs, system_name, &HashSet::new(), &locals),
comment: eq.comment,
},
TermOrigin::Component {
component: system_name.to_string(),
},
)
})
.collect();
let continuous_events = rs
.continuous_events
.clone()
.unwrap_or_default()
.into_iter()
.map(|e| namespace_continuous_event(e, system_name, &HashSet::new(), &locals))
.collect();
let discrete_events = rs
.discrete_events
.clone()
.unwrap_or_default()
.into_iter()
.map(|e| namespace_discrete_event(e, system_name, &HashSet::new(), &locals))
.collect();
Ok(SystemBlock {
name: system_name.to_string(),
state_vars,
parameters,
observed_vars: IndexMap::new(),
equations,
continuous_events,
discrete_events,
})
}
fn namespace_expr(
expr: &Expr,
system_name: &str,
subsys: &HashSet<String>,
locals: &HashSet<String>,
) -> Expr {
namespace_expr_scoped(expr, system_name, &HashSet::new(), subsys, locals)
}
fn namespace_join_names(
join: &[JoinClause],
binders: &HashSet<&str>,
system_name: &str,
locals: &HashSet<String>,
) -> Option<Vec<JoinClause>> {
let ns = |n: &String| -> String {
if binders.contains(n.as_str()) {
n.clone()
} else if let Some((head, _)) = n.split_once('.') {
if locals.contains(head) {
return format!("{system_name}.{n}");
}
n.clone()
} else if locals.contains(n) {
format!("{system_name}.{n}")
} else {
n.clone()
}
};
let out: Vec<JoinClause> = join
.iter()
.map(|c| JoinClause {
on: c.on.iter().map(|[l, r]| [ns(l), ns(r)]).collect(),
syms: c.syms.clone(),
overlap: c.overlap.as_ref().map(|ov| OverlapClause {
src_env: ov.src_env.iter().map(&ns).collect(),
tgt_env: ov.tgt_env.iter().map(&ns).collect(),
eps: ov.eps,
sym_src: ov.sym_src.clone(),
sym_tgt: ov.sym_tgt.clone(),
}),
on_gate: c.on_gate.as_ref().map(|g| g.map_column_names(ns)),
})
.collect();
(out != join).then_some(out)
}
fn namespace_expr_scoped(
expr: &Expr,
system_name: &str,
bound: &HashSet<String>,
subsys: &HashSet<String>,
locals: &HashSet<String>,
) -> Expr {
match expr {
Expr::Number(n) => Expr::Number(*n),
Expr::Integer(n) => Expr::Integer(*n),
Expr::Variable(name) => {
if name == "t" || name == VAR_PLACEHOLDER || bound.contains(name) {
Expr::Variable(name.clone())
} else if name.contains('.') {
let head = name.split('.').next().unwrap_or("");
if subsys.contains(head) {
Expr::Variable(format!("{system_name}.{name}"))
} else {
Expr::Variable(name.clone())
}
} else {
Expr::Variable(format!("{system_name}.{name}"))
}
}
Expr::Operator(node) => {
let mut child_bound = bound.clone();
if let Some(output_idx) = &node.output_idx {
child_bound.extend(output_idx.iter().cloned());
}
if let Some(ranges) = &node.ranges {
child_bound.extend(ranges.keys().cloned());
}
if let Some(int_var) = &node.int_var {
child_bound.insert(int_var.clone());
}
let mut out = node.map_children(&mut |c| {
namespace_expr_scoped(c, system_name, &child_bound, subsys, locals)
});
out.wrt = node.wrt.as_ref().map(|w| {
if w.contains('.')
|| w == "t"
|| w == VAR_PLACEHOLDER
|| child_bound.contains(w)
|| !locals.contains(w)
{
w.clone()
} else {
format!("{system_name}.{w}")
}
});
if let Some(join) = &node.join {
let mut binders: HashSet<&str> = HashSet::new();
if let Some(output_idx) = &node.output_idx {
binders.extend(output_idx.iter().map(String::as_str));
}
if let Some(ranges) = &node.ranges {
binders.extend(ranges.keys().map(String::as_str));
}
if let Some(ns) = namespace_join_names(join, &binders, system_name, locals) {
out.join = Some(ns);
}
}
Expr::operator(out)
}
}
}
fn namespace_continuous_event(
mut event: ContinuousEvent,
system_name: &str,
subsys: &HashSet<String>,
locals: &HashSet<String>,
) -> ContinuousEvent {
event.conditions = event
.conditions
.into_iter()
.map(|c| namespace_expr(&c, system_name, subsys, locals))
.collect();
event.affects = event
.affects
.into_iter()
.map(|mut a| {
a.lhs = namespace_plain(&a.lhs, system_name);
a.rhs = namespace_expr(&a.rhs, system_name, subsys, locals);
a
})
.collect();
if let Some(neg) = event.affect_neg.take() {
event.affect_neg = Some(
neg.into_iter()
.map(|mut a| {
a.lhs = namespace_plain(&a.lhs, system_name);
a.rhs = namespace_expr(&a.rhs, system_name, subsys, locals);
a
})
.collect(),
);
}
event
}
fn namespace_discrete_event(
mut event: DiscreteEvent,
system_name: &str,
subsys: &HashSet<String>,
locals: &HashSet<String>,
) -> DiscreteEvent {
use crate::types::DiscreteEventTrigger;
event.trigger = match event.trigger {
DiscreteEventTrigger::Condition { expression } => DiscreteEventTrigger::Condition {
expression: namespace_expr(&expression, system_name, subsys, locals),
},
other => other,
};
if let Some(affects) = event.affects.take() {
event.affects = Some(
affects
.into_iter()
.map(|mut a| {
a.lhs = namespace_plain(&a.lhs, system_name);
a.rhs = namespace_expr(&a.rhs, system_name, subsys, locals);
a
})
.collect(),
);
}
event
}
pub(crate) const VAR_PLACEHOLDER: &str = "_var";
fn namespace_plain(name: &str, system_name: &str) -> String {
if name.contains('.') || name == VAR_PLACEHOLDER {
name.to_string()
} else {
format!("{system_name}.{name}")
}
}
fn reject_spatial_operators(expr: &Expr) -> Result<(), FlattenError> {
match expr {
Expr::Number(_) | Expr::Integer(_) | Expr::Variable(_) => Ok(()),
Expr::Operator(node) => {
if let Err(crate::op_registry::OpError::Unlowered { op }) =
crate::op_registry::check_node(node)
{
return Err(FlattenError::UnloweredOperator { op });
}
node.try_for_each_child(&mut |child| reject_spatial_operators(child))
}
}
}
fn extract_ddt_dependent(lhs: &Expr) -> Option<String> {
let Expr::Operator(node) = lhs else {
return None;
};
if node.op != "D" {
return None;
}
if crate::op_registry::is_rewrite_target_derivative(node) {
return None;
}
if node.args.len() != 1 {
return None;
}
match &node.args[0] {
Expr::Variable(name) => Some(name.clone()),
_ => None,
}
}
fn apply_coupling_entry(
entry: &CouplingEntry,
index: usize,
per_system: &mut Vec<SystemBlock>,
coupling_rules_applied: &mut Vec<String>,
merged_renames: &mut IndexMap<String, String>,
coupling_rewritten_names: &mut BTreeSet<String>,
) -> Result<(), FlattenError> {
match entry {
CouplingEntry::OperatorCompose {
systems,
translate,
require_match,
description,
..
} => {
let made = apply_operator_compose(
systems,
translate.as_ref(),
*require_match,
per_system,
merged_renames,
)?;
compose_renames(merged_renames, made);
coupling_rules_applied.push(
description
.clone()
.unwrap_or_else(|| format!("operator_compose({})", systems.join(" + "))),
);
}
CouplingEntry::Couple {
systems,
connector,
lifting: _,
description,
} => {
let connector = retarget_connector(connector, merged_renames);
apply_couple(systems, connector.as_ref(), index, per_system)?;
coupling_rules_applied.push(
description
.clone()
.unwrap_or_else(|| format!("couple({})", systems.join(" <-> "))),
);
}
CouplingEntry::VariableMap {
from,
to,
transform,
factor,
lifting: _,
description,
} => {
match transform {
VariableMapTransform::Named(_) => {
let resolved_to = resolve_renamed(to, merged_renames).to_string();
apply_variable_map(
resolve_renamed(from, merged_renames),
&resolved_to,
*factor,
per_system,
);
coupling_rewritten_names.insert(resolved_to);
}
VariableMapTransform::Expression(node) => {
if factor.is_some() {
return Err(FlattenError::VariableMapFactorWithExpression {
from: from.clone(),
to: to.clone(),
});
}
if !node.any_child(&mut |e| crate::expression::contains(e, from)) {
return Err(FlattenError::VariableMapExpressionMissingFrom {
from: from.clone(),
to: to.clone(),
});
}
}
}
coupling_rules_applied.push(description.clone().unwrap_or_else(|| {
let factor_str = factor.map(|f| format!(" [factor={f}]")).unwrap_or_default();
format!("variable_map({from} -> {to}, {transform}){factor_str}")
}));
}
CouplingEntry::OperatorApply {
operator,
description,
} => {
coupling_rules_applied.push(
description
.clone()
.unwrap_or_else(|| format!("operator_apply({operator})")),
);
}
CouplingEntry::Callback {
callback_id,
description,
..
} => {
coupling_rules_applied.push(
description
.clone()
.unwrap_or_else(|| format!("callback({callback_id})")),
);
}
CouplingEntry::Event {
event_type,
name,
description,
..
} => {
coupling_rules_applied.push(description.clone().unwrap_or_else(|| {
format!(
"event({}: {})",
event_type,
name.as_deref().unwrap_or("unnamed")
)
}));
}
CouplingEntry::CouplingImport { .. } => {}
}
Ok(())
}
fn compose_dependent(lhs: &Expr) -> Option<String> {
match crate::classification::lhs_form(lhs) {
crate::classification::LhsForm::Derivative(name)
| crate::classification::LhsForm::Bare(name) => Some(name),
crate::classification::LhsForm::Expression => None,
}
}
fn build_translate_map(
translate: Option<&serde_json::Value>,
a_system: &str,
b_system: &str,
) -> HashMap<String, (String, f64)> {
let mut out: HashMap<String, (String, f64)> = HashMap::new();
let Some(obj) = translate.and_then(|v| v.as_object()) else {
return out;
};
for (a_name, value) in obj {
let a_q = qualify_translate_endpoint(a_name, a_system);
match value {
serde_json::Value::String(b_name) => {
out.insert(qualify_translate_endpoint(b_name, b_system), (a_q, 1.0));
}
serde_json::Value::Object(spec) => {
let b_name = ["var", "to", "target"]
.iter()
.find_map(|k| spec.get(*k).and_then(|v| v.as_str()));
let factor = spec.get("factor").and_then(|f| f.as_f64()).unwrap_or(1.0);
if let Some(b_name) = b_name {
out.insert(qualify_translate_endpoint(b_name, b_system), (a_q, factor));
}
}
_ => {}
}
}
out
}
fn qualify_translate_endpoint(name: &str, system: &str) -> String {
if name.is_empty()
|| name == VAR_PLACEHOLDER
|| name.ends_with(&format!(".{VAR_PLACEHOLDER}"))
|| name.contains('.')
|| system.is_empty()
{
return name.to_string();
}
format!("{system}.{name}")
}
fn expand_placeholder_equations(a_idx: usize, b_idx: usize, per_system: &mut [SystemBlock]) {
let a_states: Vec<String> = per_system[a_idx].state_vars.keys().cloned().collect();
if a_states.is_empty() {
return;
}
let b_equations = std::mem::take(&mut per_system[b_idx].equations);
let mut expanded: Vec<FlatEq> = Vec::with_capacity(b_equations.len());
for fe in b_equations {
if crate::expression::contains(&fe.eq.lhs, VAR_PLACEHOLDER)
|| crate::expression::contains(&fe.eq.rhs, VAR_PLACEHOLDER)
{
for state in &a_states {
let subs: HashMap<String, Expr> =
std::iter::once((VAR_PLACEHOLDER.to_string(), Expr::Variable(state.clone())))
.collect();
let mut clone = fe.clone();
clone.eq.lhs = crate::substitute::substitute(&fe.eq.lhs, &subs);
clone.map_contributions(|e| crate::substitute::substitute(e, &subs));
expanded.push(clone);
}
} else {
expanded.push(fe);
}
}
per_system[b_idx].equations = expanded;
}
fn apply_operator_compose(
systems: &[String],
translate: Option<&serde_json::Value>,
require_match: Option<bool>,
per_system: &mut [SystemBlock],
renames: &IndexMap<String, String>,
) -> Result<IndexMap<String, String>, FlattenError> {
if systems.len() < 2 {
return Ok(IndexMap::new());
}
let Some(a_idx) = per_system.iter().position(|b| b.name == systems[0]) else {
return Ok(IndexMap::new());
};
let Some(b_idx) = per_system.iter().position(|b| b.name == systems[1]) else {
return Ok(IndexMap::new());
};
if a_idx == b_idx {
return Ok(IndexMap::new());
}
expand_placeholder_equations(a_idx, b_idx, per_system);
let translate: HashMap<String, (String, f64)> =
build_translate_map(translate, &systems[0], &systems[1])
.into_iter()
.map(|(b_name, (a_name, factor))| {
(
resolve_renamed(&b_name, renames).to_string(),
(resolve_renamed(&a_name, renames).to_string(), factor),
)
})
.collect();
let mut a_index: IndexMap<String, usize> = IndexMap::new();
for (i, FlatEq { eq, .. }) in per_system[a_idx].equations.iter().enumerate() {
if let Some(dep) = compose_dependent(&eq.lhs) {
a_index.insert(dep, i);
}
}
let b_equations = std::mem::take(&mut per_system[b_idx].equations);
let mut surviving: Vec<FlatEq> = Vec::new();
let mut merged_away: IndexMap<String, String> = IndexMap::new();
let mut merged_positions: HashSet<usize> = HashSet::new();
let mut bare_matches: IndexMap<String, String> = IndexMap::new();
let mut authored = 0_usize;
let mut unmatched: Vec<String> = Vec::new();
for b_fe in b_equations {
let Some(b_dep) = compose_dependent(&b_fe.eq.lhs) else {
surviving.push(b_fe);
continue;
};
authored += 1;
let mut target = b_dep.clone();
let mut factor = 1.0_f64;
let mut is_bare_match = false;
if a_index.contains_key(&b_dep) {
} else if let Some((a_name, f)) = translate.get(&b_dep) {
target = a_name.clone();
factor = *f;
} else {
let short = b_dep.split_once('.').map(|(_, s)| s).unwrap_or(&b_dep);
let suffix = format!(".{short}");
if let Some(hit) = a_index.keys().find(|k| k.ends_with(&suffix)) {
target = hit.clone();
is_bare_match = true;
}
}
let Some(&i) = a_index.get(&target) else {
unmatched.push(b_dep);
surviving.push(b_fe);
continue;
};
let b_eq = b_fe.eq;
let mut b_contributions = b_fe.contributions;
if target != b_dep {
let subs: HashMap<String, Expr> =
std::iter::once((b_dep.clone(), Expr::Variable(target.clone()))).collect();
for c in &mut b_contributions {
c.expr = crate::substitute::substitute(&c.expr, &subs);
}
}
if factor != 1.0 {
b_contributions = vec![Contribution {
expr: Expr::operator(ExpressionNode {
op: "*".to_string(),
args: vec![Expr::Number(factor), fold_contributions(&b_contributions)],
..Default::default()
}),
origins: union_origins(&b_contributions, None),
}];
}
let a_fe = &mut per_system[a_idx].equations[i];
let mut merged = std::mem::take(&mut a_fe.contributions);
merged.extend(b_contributions);
a_fe.set_contributions(merged);
let a_lhs = &per_system[a_idx].equations[i].eq.lhs;
if !matches!(
crate::classification::lhs_form(a_lhs),
crate::classification::LhsForm::Derivative(_)
) && matches!(
crate::classification::lhs_form(&b_eq.lhs),
crate::classification::LhsForm::Derivative(_)
) {
let subs: HashMap<String, Expr> =
std::iter::once((b_dep.clone(), Expr::Variable(target.clone()))).collect();
per_system[a_idx].equations[i].eq.lhs = crate::substitute::substitute(&b_eq.lhs, &subs);
}
merged_positions.insert(i);
if target != b_dep {
if is_bare_match {
bare_matches.insert(b_dep.clone(), target.clone());
}
merged_away.insert(b_dep, target);
}
}
per_system[b_idx].equations = surviving;
report_operator_compose_merge(
&systems[0],
&systems[1],
require_match,
authored,
&unmatched,
)?;
let mut inverted: IndexMap<String, String> = IndexMap::new();
for (b_dep, target) in &bare_matches {
if bare_name_owner(per_system, &systems[0], &systems[1], b_dep, target)? == *b_dep {
inverted.insert(target.clone(), b_dep.clone());
merged_away.shift_remove(b_dep);
}
}
if !inverted.is_empty() {
retarget_merged_names(per_system, &inverted);
for gone in inverted.keys() {
let owner = gone.split_once('.').map(|(head, _)| head).unwrap_or(gone);
if let Some(block) = per_system.iter_mut().find(|b| b.name == owner) {
block.state_vars.shift_remove(gone);
block.observed_vars.shift_remove(gone);
}
}
}
reattribute_merged_equations(a_idx, &merged_positions, per_system);
if !merged_away.is_empty() {
retarget_merged_names(per_system, &merged_away);
for gone in merged_away.keys() {
per_system[b_idx].state_vars.shift_remove(gone);
per_system[b_idx].observed_vars.shift_remove(gone);
}
}
merged_away.extend(inverted);
Ok(merged_away)
}
thread_local! {
static COUPLING_DIAGNOSTIC_SINK: std::cell::RefCell<Option<Vec<String>>> =
const { std::cell::RefCell::new(None) };
}
pub fn capture_coupling_diagnostics<T>(f: impl FnOnce() -> T) -> (T, Vec<String>) {
COUPLING_DIAGNOSTIC_SINK.with(|sink| *sink.borrow_mut() = Some(Vec::new()));
let out = f();
let collected =
COUPLING_DIAGNOSTIC_SINK.with(|sink| sink.borrow_mut().take().unwrap_or_default());
(out, collected)
}
fn emit_coupling_warning(message: String) {
let captured = COUPLING_DIAGNOSTIC_SINK.with(|sink| {
if let Some(buf) = sink.borrow_mut().as_mut() {
buf.push(message.clone());
true
} else {
false
}
});
if !captured {
eprintln!("warning: {message}");
}
}
fn is_state_var(per_system: &[SystemBlock], dep: &str) -> bool {
let root = dep.split_once('.').map(|(h, _)| h).unwrap_or(dep);
per_system
.iter()
.find(|b| b.name == root)
.is_some_and(|b| b.state_vars.contains_key(dep))
}
fn bare_name_owner(
per_system: &[SystemBlock],
a: &str,
b: &str,
b_dep: &str,
target: &str,
) -> Result<String, FlattenError> {
let b_state = is_state_var(per_system, b_dep);
let a_state = is_state_var(per_system, target);
if b_state && a_state {
return Err(FlattenError::OperatorComposeAmbiguousBareName {
a: a.to_string(),
b: b.to_string(),
b_dep: b_dep.to_string(),
target: target.to_string(),
});
}
Ok(if b_state {
b_dep.to_string()
} else {
target.to_string()
})
}
fn report_operator_compose_merge(
a: &str,
b: &str,
require_match: Option<bool>,
authored: usize,
unmatched: &[String],
) -> Result<(), FlattenError> {
if authored == 0 || unmatched.is_empty() {
return Ok(());
}
if require_match == Some(false) {
return Ok(());
}
let merged = authored - unmatched.len();
let names = unmatched.join(", ");
if require_match == Some(true) {
return Err(FlattenError::OperatorComposeRequireMatchUnmatched {
a: a.to_string(),
b: b.to_string(),
merged,
authored,
unmatched: names,
});
}
if merged == 0 {
return Err(FlattenError::OperatorComposeNoMerge {
a: a.to_string(),
b: b.to_string(),
authored,
unmatched: names,
});
}
emit_coupling_warning(format!(
"operator_compose_partial_merge: operator_compose({a} + {b}) merged {merged} of \
{authored} equations '{b}' authored; unmatched dependent variable(s): {names}. The \
unmatched equations are preserved unchanged (esm-libraries-spec §4.7.1 step 5), so they \
integrate DECOUPLED from '{a}'. Set `require_match: true` if they were meant to be \
contributions, or `require_match: false` if they were not."
));
Ok(())
}
fn reattribute_merged_equations(
a_idx: usize,
merged_positions: &HashSet<usize>,
per_system: &mut [SystemBlock],
) {
if merged_positions.is_empty() {
return;
}
let a_name = per_system[a_idx].name.clone();
let mut kept: Vec<FlatEq> = Vec::new();
let mut moved: Vec<(String, FlatEq)> = Vec::new();
for (i, eq) in std::mem::take(&mut per_system[a_idx].equations)
.into_iter()
.enumerate()
{
let owner = if merged_positions.contains(&i) {
compose_dependent(&eq.eq.lhs).map(|dep| {
dep.split_once('.')
.map(|(root, _)| root.to_string())
.unwrap_or(dep)
})
} else {
None
};
match owner {
Some(owner) if owner != a_name && per_system.iter().any(|blk| blk.name == owner) => {
moved.push((owner, eq));
}
_ => kept.push(eq),
}
}
per_system[a_idx].equations = kept;
for (owner, eq) in moved {
if let Some(blk) = per_system.iter_mut().find(|blk| blk.name == owner) {
blk.equations.push(eq);
}
}
}
fn retarget_merged_names(per_system: &mut [SystemBlock], renames: &IndexMap<String, String>) {
if renames.is_empty() {
return;
}
let subs: HashMap<String, Expr> = renames
.iter()
.map(|(from, to)| (from.clone(), Expr::Variable(to.clone())))
.collect();
let ren = |expr: &Expr| rename_names(expr, &subs, renames);
for block in per_system.iter_mut() {
for fe in block.equations.iter_mut() {
fe.eq.lhs = ren(&fe.eq.lhs);
fe.map_contributions(&ren);
}
for var in block
.state_vars
.values_mut()
.chain(block.parameters.values_mut())
.chain(block.observed_vars.values_mut())
{
var.for_each_expression_mut(&mut |expr| *expr = ren(expr));
}
for event in block.continuous_events.iter_mut() {
rename_continuous_event_names(event, &subs, renames);
}
for event in block.discrete_events.iter_mut() {
rename_discrete_event_names(event, &subs, renames);
}
}
}
fn rename_names(
expr: &Expr,
subs: &HashMap<String, Expr>,
renames: &IndexMap<String, String>,
) -> Expr {
rename_join_names(&crate::substitute::substitute(expr, subs), renames)
}
fn rename_affect_names(
affects: &mut [AffectEquation],
subs: &HashMap<String, Expr>,
renames: &IndexMap<String, String>,
) {
for affect in affects {
if let Some(survivor) = renames.get(&affect.lhs) {
affect.lhs = survivor.clone();
}
affect.rhs = rename_names(&affect.rhs, subs, renames);
}
}
fn rename_continuous_event_names(
event: &mut ContinuousEvent,
subs: &HashMap<String, Expr>,
renames: &IndexMap<String, String>,
) {
for condition in event.conditions.iter_mut() {
*condition = rename_names(condition, subs, renames);
}
rename_affect_names(&mut event.affects, subs, renames);
if let Some(neg) = event.affect_neg.as_mut() {
rename_affect_names(neg, subs, renames);
}
}
fn rename_discrete_event_names(
event: &mut DiscreteEvent,
subs: &HashMap<String, Expr>,
renames: &IndexMap<String, String>,
) {
if let DiscreteEventTrigger::Condition { expression } = &mut event.trigger {
*expression = rename_names(expression, subs, renames);
}
if let Some(affects) = event.affects.as_mut() {
rename_affect_names(affects, subs, renames);
}
}
fn sum_exprs(a: Expr, b: Expr) -> Expr {
Expr::operator(ExpressionNode {
op: "+".to_string(),
args: vec![a, b],
wrt: None,
dim: None,
..Default::default()
})
}
fn apply_couple(
systems: &[String],
connector: &serde_json::Value,
index: usize,
per_system: &mut Vec<SystemBlock>,
) -> Result<(), FlattenError> {
let Some(eqs_json) = connector.get("equations").and_then(|e| e.as_array()) else {
return Ok(());
};
let block_name = format!("couple({})", systems.join(","));
let mut new_equations = Vec::new();
for eq_val in eqs_json {
if eq_val.get("lhs").is_some() || eq_val.get("rhs").is_some() {
let lhs = parse_connector_side(eq_val, "lhs", systems)?;
let rhs = parse_connector_side(eq_val, "rhs", systems)?;
new_equations.push(FlatEq::authored(
Equation {
comment: None,
lhs,
rhs,
},
TermOrigin::Coupling { coupling: index },
));
} else if eq_val.get("to").is_some() {
inject_connector_term(eq_val, systems, index, per_system)?;
} else {
return Err(FlattenError::MalformedConnectorEquation {
systems: systems.join(","),
side: "lhs".to_string(),
});
}
}
if !new_equations.is_empty() {
per_system.push(SystemBlock {
name: block_name,
state_vars: IndexMap::new(),
parameters: IndexMap::new(),
observed_vars: IndexMap::new(),
equations: new_equations,
continuous_events: Vec::new(),
discrete_events: Vec::new(),
});
}
Ok(())
}
fn last_defining_equation(
per_system: &[SystemBlock],
target: &str,
tendency_only: bool,
) -> Option<(usize, usize)> {
use crate::classification::{LhsForm, lhs_form};
let mut found = None;
for (b, block) in per_system.iter().enumerate() {
for (i, FlatEq { eq, .. }) in block.equations.iter().enumerate() {
let dep = match lhs_form(&eq.lhs) {
LhsForm::Derivative(name) => name,
LhsForm::Bare(name) if !tendency_only => name,
_ => continue,
};
if dep == target {
found = Some((b, i));
}
}
}
found
}
fn inject_connector_term(
eq_val: &serde_json::Value,
systems: &[String],
index: usize,
per_system: &mut [SystemBlock],
) -> Result<(), FlattenError> {
let Some(target) = eq_val.get("to").and_then(|v| v.as_str()) else {
return Ok(());
};
let transform = eq_val.get("transform").and_then(|v| v.as_str());
let multiplicative = transform == Some("multiplicative");
let found = last_defining_equation(per_system, target, multiplicative);
if multiplicative && found.is_none() {
return Err(FlattenError::CoupleMultiplicativeNoTendency {
target: target.to_string(),
});
}
let term = match eq_val.get("expression") {
Some(v) => serde_json::from_value::<Expr>(v.clone()).map_err(|_| {
FlattenError::MalformedConnectorEquation {
systems: systems.join(","),
side: "expression".to_string(),
}
})?,
None => match eq_val.get("from").and_then(|v| v.as_str()) {
Some(from) => Expr::Variable(from.to_string()),
None => {
return Err(FlattenError::MalformedConnectorEquation {
systems: systems.join(","),
side: "from".to_string(),
});
}
},
};
let authored = Contribution {
expr: term.clone(),
origins: vec![TermOrigin::Coupling { coupling: index }],
};
let Some((b, i)) = found else {
if !matches!(transform, None | Some("additive") | Some("")) {
return Ok(());
}
let root = target.split('.').next().unwrap_or(target);
let Some(owner) = per_system
.iter_mut()
.find(|blk| blk.name == root && blk.state_vars.contains_key(target))
else {
return Ok(());
};
let lhs = Expr::operator(ExpressionNode {
op: "D".to_string(),
args: vec![Expr::Variable(target.to_string())],
wrt: Some("t".to_string()),
..Default::default()
});
owner.equations.push(FlatEq {
eq: Equation {
comment: None,
lhs,
rhs: term,
},
contributions: vec![authored],
lifted: false,
});
return Ok(());
};
let fe = &mut per_system[b].equations[i];
let contributions = match transform {
Some("multiplicative") => vec![Contribution {
expr: Expr::operator(ExpressionNode {
op: "*".to_string(),
args: vec![fe.body().clone(), term],
..Default::default()
}),
origins: union_origins(
&fe.contributions,
Some(TermOrigin::Coupling { coupling: index }),
),
}],
Some("replacement") => vec![authored],
_ => {
let mut list = std::mem::take(&mut fe.contributions);
list.push(authored);
list
}
};
fe.set_contributions(contributions);
Ok(())
}
fn parse_connector_side(
eq_val: &serde_json::Value,
side: &str,
systems: &[String],
) -> Result<Expr, FlattenError> {
eq_val
.get(side)
.cloned()
.and_then(|v| serde_json::from_value::<Expr>(v).ok())
.ok_or_else(|| FlattenError::MalformedConnectorEquation {
systems: systems.join(","),
side: side.to_string(),
})
}
fn apply_variable_map(from: &str, to: &str, factor: Option<f64>, per_system: &mut [SystemBlock]) {
let replacement = match factor {
Some(f) if f != 1.0 => Expr::operator(ExpressionNode {
op: "*".to_string(),
args: vec![Expr::Number(f), Expr::Variable(from.to_string())],
wrt: None,
dim: None,
..Default::default()
}),
_ => Expr::Variable(from.to_string()),
};
let subs: HashMap<String, Expr> = std::iter::once((to.to_string(), replacement)).collect();
let join_renames: IndexMap<String, String> =
std::iter::once((to.to_string(), from.to_string())).collect();
for block in per_system.iter_mut() {
for fe in &mut block.equations {
fe.eq.lhs = crate::substitute::substitute(&fe.eq.lhs, &subs);
fe.map_contributions(|e| {
rename_join_names(&crate::substitute::substitute(e, &subs), &join_renames)
});
}
for var in block
.observed_vars
.values_mut()
.chain(block.parameters.values_mut())
{
var.for_each_expression_mut(&mut |expr| {
*expr =
rename_join_names(&crate::substitute::substitute(expr, &subs), &join_renames);
});
}
for ev in &mut block.continuous_events {
*ev = crate::substitute::substitute_in_continuous_event(ev, &subs);
}
for ev in &mut block.discrete_events {
*ev = crate::substitute::substitute_in_discrete_event(ev, &subs);
}
}
}
fn contains_join(expr: &Expr) -> bool {
match expr {
Expr::Operator(node) => {
node.join.as_ref().is_some_and(|j| !j.is_empty()) || node.any_child(&mut contains_join)
}
_ => false,
}
}
fn rename_join_names(expr: &Expr, renames: &IndexMap<String, String>) -> Expr {
if renames.is_empty() || !contains_join(expr) {
return expr.clone();
}
rename_join_names_in(expr, renames)
}
fn rename_join_names_in(expr: &Expr, renames: &IndexMap<String, String>) -> Expr {
let Expr::Operator(node) = expr else {
return expr.clone();
};
let mut out = node.map_children(&mut |c| rename_join_names_in(c, renames));
if let Some(join) = &node.join {
let ren = |n: &String| -> String { renames.get(n).unwrap_or(n).clone() };
out.join = Some(
join.iter()
.map(|c| JoinClause {
on: c.on.iter().map(|[l, r]| [ren(l), ren(r)]).collect(),
syms: c.syms.clone(),
overlap: c.overlap.as_ref().map(|ov| OverlapClause {
src_env: ov.src_env.iter().map(&ren).collect(),
tgt_env: ov.tgt_env.iter().map(&ren).collect(),
eps: ov.eps,
sym_src: ov.sym_src.clone(),
sym_tgt: ov.sym_tgt.clone(),
}),
on_gate: c.on_gate.as_ref().map(|g| g.map_column_names(ren)),
})
.collect(),
);
}
Expr::operator(out)
}
fn check_variable_map_endpoints(
file: &EsmFile,
per_system: &[SystemBlock],
) -> Result<(), FlattenError> {
let Some(entries) = &file.coupling else {
return Ok(());
};
let mut known: HashSet<&str> = HashSet::new();
for block in per_system {
known.extend(block.state_vars.keys().map(String::as_str));
known.extend(block.parameters.keys().map(String::as_str));
known.extend(block.observed_vars.keys().map(String::as_str));
}
let loader_names: HashSet<&str> = file
.data_sources
.as_ref()
.map(|ds| ds.keys().map(String::as_str).collect())
.unwrap_or_default();
for entry in entries {
let CouplingEntry::VariableMap { from, to, .. } = entry else {
continue;
};
let from_owner = from
.split_once('.')
.map_or(from.as_str(), |(owner, _)| owner);
let from_is_loaded = loader_names.contains(from_owner);
for (side, endpoint) in [("from", from), ("to", to)] {
if endpoint.is_empty()
|| known.contains(endpoint.as_str())
|| (side == "from" && from_is_loaded)
{
continue;
}
return Err(FlattenError::VariableMapUnresolvedEndpoint {
from: from.clone(),
to: to.clone(),
side: side.to_string(),
endpoint: endpoint.clone(),
});
}
}
Ok(())
}
fn check_variable_map_units(file: &EsmFile) -> Result<(), FlattenError> {
let Some(entries) = &file.coupling else {
return Ok(());
};
for entry in entries {
let CouplingEntry::VariableMap {
from,
to,
transform,
..
} = entry
else {
continue;
};
if transform.as_named() != Some("identity") {
continue;
}
let (Some(source_units), Some(target_units)) = (
lookup_variable_units(file, from),
lookup_variable_units(file, to),
) else {
continue;
};
if source_units.is_empty() || target_units.is_empty() {
continue;
}
if source_units != target_units {
return Err(FlattenError::DomainUnitMismatch {
variable: from.clone(),
source_units,
target_units,
});
}
}
Ok(())
}
pub(crate) fn lookup_variable_units(file: &EsmFile, qualified: &str) -> Option<String> {
let (root, tail) = qualified.split_once('.')?;
if let Some(models) = &file.models
&& let Some(model) = models.get(root)
{
return lookup_model_units(model, tail);
}
if let Some(rsystems) = &file.reaction_systems
&& let Some(rsys) = rsystems.get(root)
{
return lookup_rsys_units(rsys, tail);
}
None
}
fn lookup_model_units(model: &Model, name: &str) -> Option<String> {
if let Some(var) = model.variables.get(name) {
return var.units.clone();
}
if let Some((head, rest)) = name.split_once('.')
&& let Some(subs) = &model.subsystems
&& let Some(sub_val) = subs.get(head)
&& let Ok(sub_model) = serde_json::from_value::<Model>(sub_val.clone())
{
return lookup_model_units(&sub_model, rest);
}
None
}
fn lookup_rsys_units(rsys: &ReactionSystem, name: &str) -> Option<String> {
if let Some(sp) = rsys.species.get(name) {
return sp.units.clone();
}
if let Some(p) = rsys.parameters.get(name) {
return p.units.clone();
}
if let Some((head, rest)) = name.split_once('.')
&& let Some(subs) = &rsys.subsystems
&& let Some(sub_val) = subs.get(head)
&& let Ok(sub_rsys) = serde_json::from_value::<ReactionSystem>(sub_val.clone())
{
return lookup_rsys_units(&sub_rsys, rest);
}
None
}
pub(crate) fn extract_ic_target(lhs: &Expr) -> Option<String> {
let Expr::Operator(node) = lhs else {
return None;
};
if node.op != "ic" || node.args.len() != 1 {
return None;
}
match &node.args[0] {
Expr::Variable(v) => Some(v.clone()),
_ => None,
}
}
fn collect_makearrays<'a>(acc: &mut Vec<&'a ExpressionNode>, expr: &'a Expr) {
let Expr::Operator(node) = expr else {
return;
};
if node.op == "makearray" {
acc.push(node);
}
node.for_each_child(&mut |c| collect_makearrays(acc, c));
}
fn index_arg_loop(expr: &Expr) -> Option<String> {
match expr {
Expr::Variable(v) => Some(v.clone()),
Expr::Operator(node) => {
for a in &node.args {
if let Some(v) = index_arg_loop(a) {
return Some(v);
}
}
None
}
_ => None,
}
}
fn detect_lift_loops(
ma: &ExpressionNode,
lifted: &HashSet<String>,
rank: usize,
) -> Option<Vec<String>> {
fn walk(expr: &Expr, lifted: &HashSet<String>, rank: usize, out: &mut Option<Vec<String>>) {
if out.is_some() {
return;
}
let Expr::Operator(node) = expr else {
return;
};
if node.op == "index"
&& node.args.len() == rank + 1
&& let Some(Expr::Variable(name)) = node.args.first()
&& lifted.contains(name)
{
let mut loops = Vec::with_capacity(rank);
let mut ok = true;
for a in node.args.iter().skip(1) {
match index_arg_loop(a) {
Some(lv) => loops.push(lv),
None => {
ok = false;
break;
}
}
}
if ok {
*out = Some(loops);
return;
}
}
for a in &node.args {
walk(a, lifted, rank, out);
}
if let Some(e) = &node.expr {
walk(e, lifted, rank, out);
}
if let Some(vs) = &node.values {
for v in vs {
walk(v, lifted, rank, out);
}
}
}
let mut out = None;
for a in &ma.args {
walk(a, lifted, rank, &mut out);
}
if let Some(vs) = &ma.values {
for v in vs {
walk(v, lifted, rank, &mut out);
}
}
out
}
fn makearray_extents(ma: &ExpressionNode) -> Vec<i64> {
let Some(regions) = &ma.regions else {
return Vec::new();
};
let Some(first) = regions.first() else {
return Vec::new();
};
let rank = first.len();
let mut ext = vec![0i64; rank];
for region in regions {
if region.len() != rank {
continue;
}
for (d, r) in region.iter().enumerate() {
if let Some(hi) = r[1].as_i64() {
ext[d] = ext[d].max(hi);
}
}
}
ext
}
fn lift_rhs_to_cell(expr: &Expr, arrayvars: &HashSet<String>, loops: &[String]) -> Expr {
match expr {
Expr::Variable(name) if arrayvars.contains(name) => index_node(name, loops),
Expr::Variable(_) | Expr::Number(_) | Expr::Integer(_) => expr.clone(),
Expr::Operator(node) => {
if node.op == "makearray" {
return index_makearray(node, loops);
}
if matches!(node.op.as_str(), "index" | "faq") {
return expr.clone();
}
let mut out = ExpressionNode::clone(node);
out.args = node
.args
.iter()
.map(|a| lift_rhs_to_cell(a, arrayvars, loops))
.collect();
Expr::operator(out)
}
}
}
fn index_node(name: &str, loops: &[String]) -> Expr {
let mut args = Vec::with_capacity(loops.len() + 1);
args.push(Expr::Variable(name.to_string()));
for l in loops {
args.push(Expr::Variable(l.clone()));
}
Expr::operator(ExpressionNode {
op: "index".to_string(),
args,
..Default::default()
})
}
fn index_makearray(ma: &ExpressionNode, loops: &[String]) -> Expr {
let mut args = Vec::with_capacity(loops.len() + 1);
args.push(Expr::operator(ma.clone()));
for l in loops {
args.push(Expr::Variable(l.clone()));
}
Expr::operator(ExpressionNode {
op: "index".to_string(),
args,
..Default::default()
})
}
fn apply_pointwise_lift(
equations: &mut [FlatEq],
lifted_shapes: &mut IndexMap<String, Vec<i64>>,
loaded_producers: &HashMap<String, usize>,
) -> Result<(), FlattenError> {
let mut lifted: HashSet<String> = HashSet::new();
for FlatEq { eq, .. } in equations.iter() {
let Some(species) = extract_ddt_dependent(&eq.lhs) else {
continue;
};
let mut mas: Vec<&ExpressionNode> = Vec::new();
collect_makearrays(&mut mas, &eq.rhs);
if !mas.is_empty() {
lifted.insert(species);
}
}
if lifted.is_empty() {
return Ok(());
}
for fe in equations.iter_mut() {
let eq = &fe.eq;
let Some(species) = extract_ddt_dependent(&eq.lhs) else {
continue;
};
if !lifted.contains(&species) {
continue;
}
let mut mas: Vec<&ExpressionNode> = Vec::new();
collect_makearrays(&mut mas, &eq.rhs);
let Some(first_ma) = mas.first() else {
continue;
};
let regions = match &first_ma.regions {
Some(r) if !r.is_empty() => r,
_ => continue,
};
let rank = regions[0].len();
let mut loops: Option<Vec<String>> = None;
for ma in &mas {
loops = detect_lift_loops(ma, &lifted, rank);
if loops.is_some() {
break;
}
}
let loops = loops.ok_or_else(|| FlattenError::DimensionPromotion {
message: format!(
"pointwise lift: could not determine the spatial loop variables \
for species '{species}' from its operator makearray"
),
})?;
let extents = makearray_extents(first_ma);
lifted_shapes.insert(species.clone(), extents.clone());
let mut arrayvars: HashSet<String> = lifted.clone();
for (name, r) in loaded_producers {
if *r == rank {
arrayvars.insert(name.clone());
}
}
let mut ranges: HashMap<String, RangeSpec> = HashMap::new();
for (d, loop_name) in loops.iter().enumerate() {
ranges.insert(loop_name.clone(), RangeSpec::Interval([1, extents[d]]));
}
let idx_species = index_node(&species, &loops);
let d_body = Expr::operator(ExpressionNode {
op: "D".to_string(),
args: vec![idx_species],
wrt: Some("t".to_string()),
..Default::default()
});
let new_lhs = Expr::operator(ExpressionNode {
op: "faq".to_string(),
output_idx: Some(loops.clone()),
ranges: Some(ranges.clone()),
expr: Some(Box::new(d_body)),
..Default::default()
});
let wrapper = ExpressionNode {
op: "faq".to_string(),
output_idx: Some(loops.clone()),
ranges: Some(ranges),
..Default::default()
};
fe.eq.lhs = new_lhs;
fe.eq.rhs = Expr::operator(wrapper);
fe.lifted = true;
fe.map_contributions(|e| lift_rhs_to_cell(e, &arrayvars, &loops));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::{ddt, test_file, var};
use crate::types::{Model, ModelVariable, VariableType};
use std::collections::HashMap;
#[test]
fn test_flatten_empty_file_errors() {
let err = flatten(&EsmFile::default()).unwrap_err();
assert!(matches!(err, FlattenError::Empty));
}
#[test]
fn test_flatten_single_model_namespaces_variables() {
let mut vars = IndexMap::new();
vars.insert(
"x".to_string(),
ModelVariable {
default: Some(0.0.into()),
..var(VariableType::Unknown, Some("m"))
},
);
vars.insert(
"k".to_string(),
ModelVariable {
default: Some(1.0.into()),
..var(VariableType::Parameter, None)
},
);
let mut models = IndexMap::new();
models.insert(
"sys".to_string(),
Model {
name: Some("System".to_string()),
variables: vars,
equations: vec![ddt("x", Expr::Variable("k".to_string()))],
..Default::default()
},
);
let file = EsmFile {
models: Some(models),
..Default::default()
};
let flat = flatten(&file).unwrap();
assert_eq!(flat.independent_variables, vec!["t".to_string()]);
assert!(flat.state_variables.contains_key("sys.x"));
assert!(flat.parameters.contains_key("sys.k"));
assert_eq!(flat.equations.len(), 1);
assert_eq!(
extract_ddt_dependent(&flat.equations[0].lhs).unwrap(),
"sys.x"
);
assert_eq!(flat.equations[0].rhs, Expr::Variable("sys.k".to_string()));
assert_eq!(flat.metadata.source_systems, vec!["sys".to_string()]);
}
#[test]
fn test_namespace_expr_preserves_bare_t() {
let expr = Expr::operator(ExpressionNode {
op: "*".to_string(),
args: vec![
Expr::Variable("decay_rate".to_string()),
Expr::Variable("t".to_string()),
],
..Default::default()
});
let out = namespace_expr(&expr, "ExponentialDecay", &HashSet::new(), &HashSet::new());
match out {
Expr::Operator(node) => {
assert_eq!(
node.args[0],
Expr::Variable("ExponentialDecay.decay_rate".to_string())
);
assert_eq!(node.args[1], Expr::Variable("t".to_string()));
}
_ => panic!("expected operator node"),
}
}
fn make_model(vars: Vec<(&str, ModelVariable)>, equations: Vec<Equation>) -> Model {
let mut variables = IndexMap::new();
for (name, v) in vars {
variables.insert(name.to_string(), v);
}
Model {
variables,
equations,
..Default::default()
}
}
fn identity_map_file(from_units: Option<&str>, to_units: Option<&str>) -> EsmFile {
let src = make_model(
vec![("T", var(VariableType::Unknown, from_units))],
vec![ddt("T", Expr::Number(0.0))],
);
let dst = make_model(
vec![
("temp", var(VariableType::Parameter, to_units)),
("y", var(VariableType::Unknown, Some("K"))),
],
vec![ddt("y", Expr::Variable("temp".to_string()))],
);
let mut models = IndexMap::new();
models.insert("src".to_string(), src);
models.insert("dst".to_string(), dst);
EsmFile {
models: Some(models),
coupling: Some(vec![CouplingEntry::VariableMap {
from: "src.T".to_string(),
to: "dst.temp".to_string(),
transform: VariableMapTransform::Named("identity".to_string()),
factor: None,
lifting: None,
description: None,
}]),
..Default::default()
}
}
#[test]
fn test_variable_map_identity_unit_mismatch_errors() {
let file = identity_map_file(Some("K"), Some("degC"));
match flatten(&file).unwrap_err() {
FlattenError::DomainUnitMismatch {
variable,
source_units,
target_units,
} => {
assert_eq!(variable, "src.T");
assert_eq!(source_units, "K");
assert_eq!(target_units, "degC");
}
other => panic!("expected DomainUnitMismatch, got {other:?}"),
}
}
#[test]
fn test_variable_map_identity_matching_units_ok() {
let file = identity_map_file(Some("K"), Some("K"));
assert!(flatten(&file).is_ok());
}
#[test]
fn test_variable_map_identity_missing_unit_ok() {
assert!(flatten(&identity_map_file(None, Some("degC"))).is_ok());
assert!(flatten(&identity_map_file(Some("K"), None)).is_ok());
}
fn subsystem_map_file(to_endpoint: &str) -> EsmFile {
let src = make_model(
vec![("T", var(VariableType::Unknown, Some("K")))],
vec![ddt("T", Expr::Number(1.0))],
);
let inner = make_model(
vec![
("gain", var(VariableType::Parameter, Some("K"))),
("x", var(VariableType::Unknown, Some("K"))),
],
vec![ddt("x", Expr::Variable("gain".to_string()))],
);
let mut subsystems = IndexMap::new();
subsystems.insert(
"inner".to_string(),
serde_json::to_value(&inner).expect("model serializes"),
);
let wrap = Model {
subsystems: Some(subsystems),
..Default::default()
};
let mut models = IndexMap::new();
models.insert("Src".to_string(), src);
models.insert("Wrap".to_string(), wrap);
EsmFile {
models: Some(models),
coupling: Some(vec![CouplingEntry::VariableMap {
from: "Src.T".to_string(),
to: to_endpoint.to_string(),
transform: VariableMapTransform::Named("param_to_var".to_string()),
factor: None,
lifting: None,
description: None,
}]),
..test_file()
}
}
#[test]
fn test_variable_map_into_subsystem_parameter_resolves() {
let flat = flatten(&subsystem_map_file("Wrap.inner.gain")).expect("flattens");
assert!(
!flat.parameters.contains_key("Wrap.inner.gain"),
"the nested target must be promoted away, got {:?}",
flat.parameters.keys().collect::<Vec<_>>()
);
let rhs = &flat
.equations
.iter()
.find(|e| {
matches!(&e.lhs, Expr::Operator(n) if n.op == "D"
&& n.args.first() == Some(&Expr::Variable("Wrap.inner.x".to_string())))
})
.expect("the nested ODE survives")
.rhs;
assert_eq!(*rhs, Expr::Variable("Src.T".to_string()));
}
#[test]
fn test_variable_map_unresolved_to_endpoint_errors() {
match flatten(&subsystem_map_file("Wrap.gain")).unwrap_err() {
FlattenError::VariableMapUnresolvedEndpoint {
side, endpoint, to, ..
} => {
assert_eq!(side, "to");
assert_eq!(endpoint, "Wrap.gain");
assert_eq!(to, "Wrap.gain");
}
other => panic!("expected VariableMapUnresolvedEndpoint, got {other:?}"),
}
}
#[test]
fn test_variable_map_unresolved_from_endpoint_errors() {
let mut file = subsystem_map_file("Wrap.inner.gain");
if let Some(CouplingEntry::VariableMap { from, .. }) =
file.coupling.as_mut().and_then(|c| c.first_mut())
{
*from = "Src.Nope.T".to_string();
}
match flatten(&file).unwrap_err() {
FlattenError::VariableMapUnresolvedEndpoint { side, endpoint, .. } => {
assert_eq!(side, "from");
assert_eq!(endpoint, "Src.Nope.T");
}
other => panic!("expected VariableMapUnresolvedEndpoint, got {other:?}"),
}
}
#[test]
fn test_pointwise_lift_failure_yields_dimension_promotion() {
let makearray = Expr::operator(ExpressionNode {
op: "makearray".to_string(),
regions: Some(vec![vec![[
crate::types::RegionBound::Int(1),
crate::types::RegionBound::Int(3),
]]]),
args: vec![Expr::Number(0.0)],
..Default::default()
});
let mut equations = vec![FlatEq::authored(
ddt("C", makearray),
TermOrigin::Component {
component: "M".to_string(),
},
)];
let loaded_producers: HashMap<String, usize> = HashMap::new();
let mut lifted_shapes: IndexMap<String, Vec<i64>> = IndexMap::new();
let err = apply_pointwise_lift(&mut equations, &mut lifted_shapes, &loaded_producers)
.unwrap_err();
match err {
FlattenError::DimensionPromotion { message } => {
assert!(message.contains("pointwise lift"), "message: {message}");
assert!(message.contains("'C'"), "message: {message}");
}
other => panic!("expected DimensionPromotion, got {other:?}"),
}
}
#[test]
fn test_reject_spatial_operator_hidden_in_aggregate_body() {
let grad = Expr::operator(ExpressionNode {
op: "grad".to_string(),
args: vec![Expr::Variable("u".to_string())],
..Default::default()
});
let aggregate = Expr::operator(ExpressionNode {
op: "faq".to_string(),
args: vec![],
expr: Some(Box::new(grad)),
output_idx: Some(vec!["i".to_string()]),
reduce: Some("+".to_string()),
..Default::default()
});
let err = reject_spatial_operators(&aggregate).unwrap_err();
assert!(
matches!(&err, FlattenError::UnloweredOperator { op } if op == "grad"),
"expected UnloweredOperator{{grad}}, got {err:?}"
);
let spatial_d = Expr::operator(ExpressionNode {
op: "D".to_string(),
args: vec![Expr::Variable("u".to_string())],
wrt: Some("x".to_string()),
..Default::default()
});
let filtered = Expr::operator(ExpressionNode {
op: "faq".to_string(),
args: vec![Expr::Variable("w".to_string())],
filter: Some(Box::new(spatial_d)),
reduce: Some("+".to_string()),
..Default::default()
});
let err = reject_spatial_operators(&filtered).unwrap_err();
assert!(
matches!(&err, FlattenError::UnloweredOperator { op } if op == "D"),
"expected UnloweredOperator{{D}}, got {err:?}"
);
}
#[test]
fn test_namespace_expr_covers_aggregate_key() {
let key = Expr::operator(ExpressionNode {
op: "+".to_string(),
args: vec![
Expr::Variable("region".to_string()),
Expr::Variable("i".to_string()),
],
..Default::default()
});
let aggregate = Expr::operator(ExpressionNode {
op: "faq".to_string(),
args: vec![Expr::Variable("w".to_string())],
output_idx: Some(vec!["i".to_string()]),
key: Some(Box::new(key)),
reduce: Some("+".to_string()),
..Default::default()
});
let out = namespace_expr(&aggregate, "sys", &HashSet::new(), &HashSet::new());
match out {
Expr::Operator(node) => {
assert_eq!(node.args[0], Expr::Variable("sys.w".to_string()));
let key = node.key.clone().expect("aggregate key preserved");
match *key {
Expr::Operator(k) => {
assert_eq!(k.args[0], Expr::Variable("sys.region".to_string()));
assert_eq!(k.args[1], Expr::Variable("i".to_string()));
}
other => panic!("expected operator key, got {other:?}"),
}
}
other => panic!("expected operator node, got {other:?}"),
}
}
#[test]
fn map_rhs_per_contribution_keeps_the_lift_wrapper() {
let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../../tests/conformance/term_provenance/fixtures/pointwise_lift.esm");
let file = crate::load_path(&path).expect("loads");
let mut flat = flatten(&file).expect("flattens");
let i = flat
.equation_contributions
.iter()
.position(|cs| cs.len() == 3)
.expect("the lifted equation has three contributions");
let subs: HashMap<String, Expr> =
std::iter::once(("Src.e".to_string(), Expr::Variable("Src.e2".to_string()))).collect();
assert!(flat.equation_lifted[i], "the lift wrapped the equation");
flat.map_rhs_per_contribution(|e| crate::substitute::substitute(e, &subs))
.expect("the lists are in step");
let contributions = &flat.equation_contributions[i];
assert_eq!(contributions[1].expr, Expr::Variable("Src.e2".to_string()));
let Expr::Operator(node) = &flat.equations[i].rhs else {
panic!("the lifted right-hand side is an operator");
};
assert_eq!(node.op, "faq");
assert_eq!(
node.expr.as_deref(),
Some(&fold_contributions(contributions))
);
}
}