Skip to content

Commit 1c807e9

Browse files
nachawatinachawati
committed
feat(core): extend AD/IA/McCormick tape to handle preserved array mode via element-wise unrolling
1 parent 48e3c7b commit 1c807e9

3 files changed

Lines changed: 114 additions & 4 deletions

File tree

‎packages/core/src/compiler/modelica/ad-codegen.ts‎

Lines changed: 71 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import {
2+
ModelicaArray,
23
ModelicaArrayEquation,
34
ModelicaBinaryExpression,
45
ModelicaBooleanLiteral,
@@ -8,6 +9,7 @@ import {
89
ModelicaIntegerLiteral,
910
ModelicaNameExpression,
1011
ModelicaRealLiteral,
12+
ModelicaSubscriptedExpression,
1113
ModelicaUnaryExpression,
1214
} from "./dae.js";
1315
import { ModelicaBinaryOperator, ModelicaUnaryOperator } from "./syntax.js";
@@ -126,10 +128,65 @@ export class StaticTapeBuilder {
126128
}
127129
}
128130

131+
// Array expressions: walk each flat element, return last index
132+
if (expr instanceof ModelicaArray) {
133+
let lastIdx = this.pushOp({ type: "const", val: 0 });
134+
for (const elem of expr.flatElements) {
135+
lastIdx = this.walk(elem);
136+
}
137+
return lastIdx;
138+
}
139+
140+
// Subscripted expression: x[i]
141+
if (expr instanceof ModelicaSubscriptedExpression) {
142+
// If subscript is a literal integer, resolve statically
143+
if (expr.subscripts.length === 1) {
144+
const sub = expr.subscripts[0];
145+
if (sub instanceof ModelicaIntegerLiteral) {
146+
const base = expr.base;
147+
if (base instanceof ModelicaArray) {
148+
const elem = base.getFlatElement(sub.value - 1); // Modelica is 1-indexed
149+
if (elem) return this.walk(elem);
150+
}
151+
// Named variable with subscript: x[3] → "x[3]"
152+
const baseName =
153+
expr.base instanceof ModelicaNameExpression
154+
? expr.base.name
155+
: expr.base && typeof expr.base === "object" && "name" in expr.base
156+
? (expr.base as { name: string }).name
157+
: null;
158+
if (baseName) return this.pushOp({ type: "var", name: `${baseName}[${sub.value}]` });
159+
}
160+
}
161+
// Fallback: treat base as a variable if it has a name
162+
const baseName =
163+
expr.base instanceof ModelicaNameExpression
164+
? expr.base.name
165+
: expr.base && typeof expr.base === "object" && "name" in expr.base
166+
? (expr.base as { name: string }).name
167+
: null;
168+
if (baseName) return this.pushOp({ type: "var", name: baseName });
169+
}
170+
129171
// Unrecognized or non-differentiable features default to constant 0 tape node.
130172
return this.pushOp({ type: "const", val: 0 });
131173
}
132174

175+
/**
176+
* Walk an array expression element-wise, returning one tape index per flat element.
177+
* For non-array expressions, returns a single-element array.
178+
*/
179+
public walkArray(expr: ModelicaExpression): number[] {
180+
if (expr instanceof ModelicaArray) {
181+
const indices: number[] = [];
182+
for (const elem of expr.flatElements) {
183+
indices.push(this.walk(elem));
184+
}
185+
return indices;
186+
}
187+
return [this.walk(expr)];
188+
}
189+
133190
/**
134191
* Emit the C-code evaluating the expressions step-by-step.
135192
* `varResolver` maps the Modelica variable name into a valid C-code getter string
@@ -479,11 +536,24 @@ export function generateModelEvaluateJacobian(id: string, dae: ModelicaDAE, vars
479536
// Gather target equations (der(x) = f(x,u))
480537
const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
481538
for (const eq of dae.equations) {
482-
if (eq instanceof ModelicaArrayEquation) continue; // Skip unsupported for NLP right now
483539
if (!("expression1" in eq && "expression2" in eq)) continue;
484540
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
485541
const ld = extractDer(se.expression1);
486542
const rd = extractDer(se.expression2);
543+
544+
if (eq instanceof ModelicaArrayEquation) {
545+
const baseName = ld || rd;
546+
if (!baseName) continue;
547+
const rhs = ld ? se.expression2 : se.expression1;
548+
const v = dae.variables.find((dv) => dv.name === baseName);
549+
const dims = v?.arrayDimensions ?? [];
550+
const size = dims.length > 0 ? dims.reduce((a: number, b: number) => a * b, 1) : 1;
551+
for (let i = 0; i < size; i++) {
552+
derEqs.push({ state: `${baseName}[${i + 1}]`, rhs });
553+
}
554+
continue;
555+
}
556+
487557
if (ld) derEqs.push({ state: ld, rhs: se.expression2 });
488558
else if (rd) derEqs.push({ state: rd, rhs: se.expression1 });
489559
}

‎packages/core/src/compiler/modelica/ad-jacobian.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,11 +167,25 @@ export function buildAdJacobian(dae: ModelicaDAE): ((t: number, y: number[]) =>
167167
// Gather derivative equations: der(x) = f(x, u)
168168
const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
169169
for (const eq of dae.equations) {
170-
if (eq instanceof ModelicaArrayEquation) continue;
171170
if (!("expression1" in eq && "expression2" in eq)) continue;
172171
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
173172
const ld = extractDer(se.expression1);
174173
const rd = extractDer(se.expression2);
174+
175+
if (eq instanceof ModelicaArrayEquation) {
176+
// Unroll array equation element-wise
177+
const baseName = ld || rd;
178+
if (!baseName) continue;
179+
const rhs = ld ? se.expression2 : se.expression1;
180+
const v = dae.variables.find((dv) => dv.name === baseName);
181+
const dims = v?.arrayDimensions ?? [];
182+
const size = dims.length > 0 ? dims.reduce((a: number, b: number) => a * b, 1) : 1;
183+
for (let i = 0; i < size; i++) {
184+
derEqs.push({ state: `${baseName}[${i + 1}]`, rhs });
185+
}
186+
continue;
187+
}
188+
175189
if (ld) derEqs.push({ state: ld, rhs: se.expression2 });
176190
else if (rd) derEqs.push({ state: rd, rhs: se.expression1 });
177191
}

‎packages/core/src/compiler/modelica/sparse.ts‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,11 +66,24 @@ export function buildCCS(rowsDeps: Set<string>[], columns: string[]): CCSMatrix
6666
export function computeJacobianSparsity(dae: ModelicaDAE): { ccs: CCSMatrix; states: string[] } {
6767
const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
6868
for (const eq of dae.equations) {
69-
if (eq instanceof ModelicaArrayEquation) continue;
7069
if (!("expression1" in eq && "expression2" in eq)) continue;
7170
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
7271
const ld = extractDer(se.expression1);
7372
const rd = extractDer(se.expression2);
73+
74+
if (eq instanceof ModelicaArrayEquation) {
75+
const baseName = ld || rd;
76+
if (!baseName) continue;
77+
const rhs = ld ? se.expression2 : se.expression1;
78+
const v = dae.variables.find((dv) => dv.name === baseName);
79+
const dims = v?.arrayDimensions ?? [];
80+
const size = dims.length > 0 ? dims.reduce((a: number, b: number) => a * b, 1) : 1;
81+
for (let i = 0; i < size; i++) {
82+
derEqs.push({ state: `${baseName}[${i + 1}]`, rhs });
83+
}
84+
continue;
85+
}
86+
7487
if (ld) derEqs.push({ state: ld, rhs: se.expression2 });
7588
else if (rd) derEqs.push({ state: rd, rhs: se.expression1 });
7689
}
@@ -110,11 +123,24 @@ export function computeJacobianSparsity(dae: ModelicaDAE): { ccs: CCSMatrix; sta
110123
export function computeHessianSparsity(dae: ModelicaDAE): { ccs: CCSMatrix; states: string[] } {
111124
const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
112125
for (const eq of dae.equations) {
113-
if (eq instanceof ModelicaArrayEquation) continue;
114126
if (!("expression1" in eq && "expression2" in eq)) continue;
115127
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
116128
const ld = extractDer(se.expression1);
117129
const rd = extractDer(se.expression2);
130+
131+
if (eq instanceof ModelicaArrayEquation) {
132+
const baseName = ld || rd;
133+
if (!baseName) continue;
134+
const rhs = ld ? se.expression2 : se.expression1;
135+
const v = dae.variables.find((dv) => dv.name === baseName);
136+
const dims = v?.arrayDimensions ?? [];
137+
const size = dims.length > 0 ? dims.reduce((a: number, b: number) => a * b, 1) : 1;
138+
for (let i = 0; i < size; i++) {
139+
derEqs.push({ state: `${baseName}[${i + 1}]`, rhs });
140+
}
141+
continue;
142+
}
143+
118144
if (ld) derEqs.push({ state: ld, rhs: se.expression2 });
119145
else if (rd) derEqs.push({ state: rd, rhs: se.expression1 });
120146
}

0 commit comments

Comments
 (0)