feat(core): extend AD, IA, and McCormick to support preserved array mode

nachawati · nachawati · commit 7fb4ba49a0ad · 2026-03-30T13:57:19.000+03:00
diff --git a/packages/core/src/compiler/modelica/ad-codegen.ts b/packages/core/src/compiler/modelica/ad-codegen.ts
@@ -670,7 +670,7 @@ export function generateModelEvaluateJacobian(id: string, dae: ModelicaDAE, vars
 
   // Gather target equations (der(x) = f(x,u))
   const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
-  for (const eq of dae.equations) {
+  for (const eq of dae.sortedEquations.length > 0 ? dae.sortedEquations : dae.equations) {
     if (!("expression1" in eq && "expression2" in eq)) continue;
     const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
     const ld = extractDer(se.expression1);
@@ -819,7 +819,7 @@ export function generateModelEvaluateHessian(id: string, dae: ModelicaDAE, vars:
   varMap.set("time", `inst->time`);
 
   const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
-  for (const eq of dae.equations) {
+  for (const eq of dae.sortedEquations.length > 0 ? dae.sortedEquations : dae.equations) {
     if (eq instanceof ModelicaArrayEquation) continue;
     if (!("expression1" in eq && "expression2" in eq)) continue;
     const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
diff --git a/packages/core/src/compiler/modelica/ad-jacobian.ts b/packages/core/src/compiler/modelica/ad-jacobian.ts
@@ -239,7 +239,7 @@ export function evaluateTapeReverse(ops: TapeOp[], t: Float64Array, outputIndex:
 export function buildAdJacobian(dae: ModelicaDAE): ((t: number, y: number[]) => number[][]) | null {
   // Gather derivative equations: der(x) = f(x, u)
   const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
-  for (const eq of dae.equations) {
+  for (const eq of dae.sortedEquations.length > 0 ? dae.sortedEquations : dae.equations) {
     if (!("expression1" in eq && "expression2" in eq)) continue;
     const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
     const ld = extractDer(se.expression1);
diff --git a/packages/core/src/compiler/modelica/dae.ts b/packages/core/src/compiler/modelica/dae.ts
@@ -18,6 +18,8 @@ export class ModelicaDAE {
   description: string | null;
   classKind = "class";
   equations: ModelicaEquation[] = [];
+  /** Equations sorted by BLT transformation for simulation. */
+  sortedEquations: ModelicaEquation[] = [];
   algorithms: ModelicaStatement[][] = [];
   /** Equations from `initial equation` sections. */
   initialEquations: ModelicaEquation[] = [];
diff --git a/packages/core/src/compiler/modelica/flattener.ts b/packages/core/src/compiler/modelica/flattener.ts
@@ -7907,7 +7907,7 @@ function tryFoldBuiltinFunction(functionName: string, args: ModelicaExpression[]
  */
 export function findAlgebraicLoops(dae: ModelicaDAE): void {
   const { sortedEquations, algebraicLoops } = performBltTransformation(dae);
-  dae.equations = sortedEquations;
+  dae.sortedEquations = sortedEquations;
 
   if (algebraicLoops.length > 0) {
     console.log(`[DAE] Found ${algebraicLoops.length} algebraic loop(s) in ${dae.name}`);
