Skip to content

Commit 373048c

Browse files
nachawatinachawati
committed
feat(core): exact AD jacobians in FMU algebraic loop solver and BDF integrator
1 parent 03f8ec6 commit 373048c

3 files changed

Lines changed: 319 additions & 53 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import {
1212
} from "./dae.js";
1313
import { ModelicaBinaryOperator, ModelicaUnaryOperator } from "./syntax.js";
1414

15-
type TapeOp =
15+
export type TapeOp =
1616
| { type: "const"; val: number }
1717
| { type: "var"; name: string }
1818
| { type: "add"; a: number; b: number }

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

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
1+
// SPDX-License-Identifier: AGPL-3.0-or-later
2+
3+
/**
4+
* Runtime Algorithmic Differentiation (AD) Jacobian evaluator.
5+
*
6+
* Builds a StaticTapeBuilder from DAE derivative equations at compile time,
7+
* then provides a closure that evaluates the exact Jacobian at runtime by
8+
* walking the tape operations in TypeScript (no C-code generation needed).
9+
*
10+
* This is used by the BDF integrator as `options.jacobian` to replace
11+
* finite-difference approximations with exact analytical derivatives.
12+
*/
13+
14+
import { StaticTapeBuilder, type TapeOp } from "./ad-codegen.js";
15+
import {
16+
ModelicaArrayEquation,
17+
type ModelicaDAE,
18+
type ModelicaExpression,
19+
ModelicaFunctionCallExpression,
20+
} from "./dae.js";
21+
22+
/** Extract derivative name from expression like der(x). */
23+
function extractDer(expr: ModelicaExpression): string | null {
24+
if (expr instanceof ModelicaFunctionCallExpression && expr.functionName === "der" && expr.args.length === 1) {
25+
const a = expr.args[0];
26+
if (a && typeof a === "object" && "name" in a) return (a as { name: string }).name;
27+
}
28+
return null;
29+
}
30+
31+
/**
32+
* Evaluate a tape forward pass at runtime, returning the value array.
33+
*/
34+
function evaluateTapeForward(ops: TapeOp[], varValues: Map<string, number>): Float64Array {
35+
const t = new Float64Array(ops.length);
36+
for (let i = 0; i < ops.length; i++) {
37+
const op = ops[i]!; // eslint-disable-line @typescript-eslint/no-non-null-assertion
38+
switch (op.type) {
39+
case "const":
40+
t[i] = op.val;
41+
break;
42+
case "var":
43+
t[i] = varValues.get(op.name) ?? 0;
44+
break;
45+
case "add":
46+
t[i] = (t[op.a] ?? 0) + (t[op.b] ?? 0);
47+
break;
48+
case "sub":
49+
t[i] = (t[op.a] ?? 0) - (t[op.b] ?? 0);
50+
break;
51+
case "mul":
52+
t[i] = (t[op.a] ?? 0) * (t[op.b] ?? 0);
53+
break;
54+
case "div":
55+
t[i] = (t[op.a] ?? 0) / (t[op.b] ?? 0);
56+
break;
57+
case "pow":
58+
t[i] = Math.pow(t[op.a] ?? 0, t[op.b] ?? 0);
59+
break;
60+
case "neg":
61+
t[i] = -(t[op.a] ?? 0);
62+
break;
63+
case "sin":
64+
t[i] = Math.sin(t[op.a] ?? 0);
65+
break;
66+
case "cos":
67+
t[i] = Math.cos(t[op.a] ?? 0);
68+
break;
69+
case "tan":
70+
t[i] = Math.tan(t[op.a] ?? 0);
71+
break;
72+
case "exp":
73+
t[i] = Math.exp(t[op.a] ?? 0);
74+
break;
75+
case "log":
76+
t[i] = Math.log(t[op.a] ?? 0);
77+
break;
78+
case "sqrt":
79+
t[i] = Math.sqrt(t[op.a] ?? 0);
80+
break;
81+
}
82+
}
83+
return t;
84+
}
85+
86+
/**
87+
* Evaluate the reverse-mode AD sweep on a tape, returning gradients for all variables.
88+
*/
89+
function evaluateTapeReverse(ops: TapeOp[], t: Float64Array, outputIndex: number): Map<string, number> {
90+
const dt = new Float64Array(ops.length);
91+
dt[outputIndex] = 1.0;
92+
93+
for (let i = ops.length - 1; i >= 0; i--) {
94+
if (dt[i] === 0) continue;
95+
const op = ops[i]!; // eslint-disable-line @typescript-eslint/no-non-null-assertion
96+
const dti = dt[i] ?? 0;
97+
98+
switch (op.type) {
99+
case "add":
100+
dt[op.a] = (dt[op.a] ?? 0) + dti;
101+
dt[op.b] = (dt[op.b] ?? 0) + dti;
102+
break;
103+
case "sub":
104+
dt[op.a] = (dt[op.a] ?? 0) + dti;
105+
dt[op.b] = (dt[op.b] ?? 0) - dti;
106+
break;
107+
case "mul":
108+
dt[op.a] = (dt[op.a] ?? 0) + dti * (t[op.b] ?? 0);
109+
dt[op.b] = (dt[op.b] ?? 0) + dti * (t[op.a] ?? 0);
110+
break;
111+
case "div":
112+
dt[op.a] = (dt[op.a] ?? 0) + dti / (t[op.b] ?? 1);
113+
dt[op.b] = (dt[op.b] ?? 0) - (dti * (t[op.a] ?? 0)) / ((t[op.b] ?? 1) * (t[op.b] ?? 1));
114+
break;
115+
case "pow": {
116+
const base = t[op.a] ?? 0;
117+
const exp = t[op.b] ?? 0;
118+
dt[op.a] = (dt[op.a] ?? 0) + dti * exp * Math.pow(base, exp - 1);
119+
dt[op.b] = (dt[op.b] ?? 0) + dti * (t[i] ?? 0) * Math.log(base);
120+
break;
121+
}
122+
case "neg":
123+
dt[op.a] = (dt[op.a] ?? 0) - dti;
124+
break;
125+
case "sin":
126+
dt[op.a] = (dt[op.a] ?? 0) + dti * Math.cos(t[op.a] ?? 0);
127+
break;
128+
case "cos":
129+
dt[op.a] = (dt[op.a] ?? 0) - dti * Math.sin(t[op.a] ?? 0);
130+
break;
131+
case "tan":
132+
dt[op.a] = (dt[op.a] ?? 0) + dti * (1 + (t[i] ?? 0) * (t[i] ?? 0));
133+
break;
134+
case "exp":
135+
dt[op.a] = (dt[op.a] ?? 0) + dti * (t[i] ?? 0);
136+
break;
137+
case "log":
138+
dt[op.a] = (dt[op.a] ?? 0) + dti / (t[op.a] ?? 1);
139+
break;
140+
case "sqrt":
141+
dt[op.a] = (dt[op.a] ?? 0) + dti / (2 * (t[i] ?? 1));
142+
break;
143+
}
144+
}
145+
146+
// Collect variable gradients
147+
const gradients = new Map<string, number>();
148+
for (let i = 0; i < ops.length; i++) {
149+
const op = ops[i]!; // eslint-disable-line @typescript-eslint/no-non-null-assertion
150+
if (op.type === "var") {
151+
gradients.set(op.name, (gradients.get(op.name) ?? 0) + (dt[i] ?? 0));
152+
}
153+
}
154+
return gradients;
155+
}
156+
157+
/**
158+
* Build a runtime AD Jacobian evaluator from a ModelicaDAE.
159+
*
160+
* Returns a function `(t: number, y: number[]) => number[][]` that computes
161+
* the exact Jacobian of the derivative equations w.r.t. the state variables.
162+
*
163+
* @param dae The flattened DAE
164+
* @returns Jacobian evaluator closure, or null if no derivative equations found
165+
*/
166+
export function buildAdJacobian(dae: ModelicaDAE): ((t: number, y: number[]) => number[][]) | null {
167+
// Gather derivative equations: der(x) = f(x, u)
168+
const derEqs: { state: string; rhs: ModelicaExpression }[] = [];
169+
for (const eq of dae.equations) {
170+
if (eq instanceof ModelicaArrayEquation) continue;
171+
if (!("expression1" in eq && "expression2" in eq)) continue;
172+
const se = eq as { expression1: ModelicaExpression; expression2: ModelicaExpression };
173+
const ld = extractDer(se.expression1);
174+
const rd = extractDer(se.expression2);
175+
if (ld) derEqs.push({ state: ld, rhs: se.expression2 });
176+
else if (rd) derEqs.push({ state: rd, rhs: se.expression1 });
177+
}
178+
179+
if (derEqs.length === 0) return null;
180+
181+
// State variable names (in order)
182+
const stateNames = derEqs.map((eq) => eq.state);
183+
const n = stateNames.length;
184+
185+
// Build tapes for each equation (done once at compile time)
186+
const tapeData: { ops: TapeOp[]; outputIndex: number }[] = [];
187+
for (const eq of derEqs) {
188+
const tape = new StaticTapeBuilder();
189+
const outIdx = tape.walk(eq.rhs);
190+
tapeData.push({ ops: [...tape.ops], outputIndex: outIdx });
191+
}
192+
193+
// Return the closure that evaluates J(t, y) at runtime
194+
return (time: number, y: number[]): number[][] => {
195+
// Build variable value map: state[i] -> y[i]
196+
const varValues = new Map<string, number>();
197+
varValues.set("time", time);
198+
for (let i = 0; i < n; i++) {
199+
const name = stateNames[i];
200+
if (name) varValues.set(name, y[i] ?? 0);
201+
}
202+
// Also set any other DAE variables from their current values
203+
for (const v of dae.variables) {
204+
if (!varValues.has(v.name) && v.expression) {
205+
// For non-state variables, use start value as approximation
206+
// (In a full implementation, these would be computed from the DAE)
207+
varValues.set(v.name, 0);
208+
}
209+
}
210+
211+
// Allocate Jacobian (n x n)
212+
const J: number[][] = [];
213+
for (let i = 0; i < n; i++) {
214+
J[i] = new Array(n).fill(0) as number[];
215+
}
216+
217+
// For each equation (row), compute forward pass then reverse pass
218+
for (let row = 0; row < n; row++) {
219+
const td = tapeData[row];
220+
if (!td) continue;
221+
222+
const t = evaluateTapeForward(td.ops, varValues);
223+
const grads = evaluateTapeReverse(td.ops, t, td.outputIndex);
224+
225+
// Fill Jacobian row
226+
const jRow = J[row];
227+
if (!jRow) continue;
228+
for (let col = 0; col < n; col++) {
229+
const stateName = stateNames[col];
230+
if (stateName) {
231+
jRow[col] = grads.get(stateName) ?? 0;
232+
}
233+
}
234+
}
235+
236+
return J;
237+
};
238+
}

0 commit comments

Comments
 (0)