Skip to content

Commit 19040a5

Browse files
nachawatinachawati
committed
feat(core): cas engine with polynomial algebra, equation solver, linalg, symbolic integration, egraph integration
1 parent fe24eaa commit 19040a5

7 files changed

Lines changed: 1818 additions & 0 deletions

File tree

‎packages/core/src/compiler/modelica/symbolic/expand.ts‎

Lines changed: 348 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,348 @@
1+
// SPDX-License-Identifier: AGPL-3.0-or-later
2+
3+
/**
4+
* Polynomial expansion, term collection, and expression normalization.
5+
*
6+
* Operates on the native ModelicaExpression AST. Uses the E-Graph engine
7+
* for canonical form computation when normalization is requested.
8+
*/
9+
10+
import type { ModelicaExpression } from "../dae.js";
11+
import {
12+
ModelicaBinaryExpression,
13+
ModelicaFunctionCallExpression,
14+
ModelicaIntegerLiteral,
15+
ModelicaNameExpression,
16+
ModelicaRealLiteral,
17+
ModelicaUnaryExpression,
18+
} from "../dae.js";
19+
import { add, div, isOne, isZero, mul, pow, sub, ZERO } from "../symbolic-diff.js";
20+
import { ModelicaBinaryOperator, ModelicaUnaryOperator } from "../syntax.js";
21+
import { egraphSimplify } from "./egraph.js";
22+
23+
// ─────────────────────────────────────────────────────────────────────
24+
// Polynomial Expansion
25+
// ─────────────────────────────────────────────────────────────────────
26+
27+
/**
28+
* Recursively expand an expression by distributing multiplication over
29+
* addition and applying binomial expansion for integer powers.
30+
*
31+
* Examples:
32+
* (a + b) * c → a*c + b*c
33+
* (a + b)^2 → a^2 + 2*a*b + b^2
34+
*/
35+
export function expandExpr(expr: ModelicaExpression): ModelicaExpression {
36+
if (expr instanceof ModelicaRealLiteral || expr instanceof ModelicaIntegerLiteral) {
37+
return expr;
38+
}
39+
if (expr instanceof ModelicaNameExpression) {
40+
return expr;
41+
}
42+
43+
if (expr instanceof ModelicaUnaryExpression) {
44+
const expanded = expandExpr(expr.operand);
45+
if (expr.operator === ModelicaUnaryOperator.UNARY_MINUS) {
46+
return distributeNeg(expanded);
47+
}
48+
return new ModelicaUnaryExpression(expr.operator, expanded);
49+
}
50+
51+
if (expr instanceof ModelicaBinaryExpression) {
52+
const left = expandExpr(expr.operand1);
53+
const right = expandExpr(expr.operand2);
54+
55+
switch (expr.operator) {
56+
case ModelicaBinaryOperator.ADDITION:
57+
case ModelicaBinaryOperator.ELEMENTWISE_ADDITION:
58+
return add(left, right);
59+
60+
case ModelicaBinaryOperator.SUBTRACTION:
61+
case ModelicaBinaryOperator.ELEMENTWISE_SUBTRACTION:
62+
return sub(left, right);
63+
64+
case ModelicaBinaryOperator.MULTIPLICATION:
65+
case ModelicaBinaryOperator.ELEMENTWISE_MULTIPLICATION:
66+
return distributeMultiply(left, right);
67+
68+
case ModelicaBinaryOperator.DIVISION:
69+
case ModelicaBinaryOperator.ELEMENTWISE_DIVISION:
70+
return div(left, right);
71+
72+
case ModelicaBinaryOperator.EXPONENTIATION:
73+
case ModelicaBinaryOperator.ELEMENTWISE_EXPONENTIATION: {
74+
const n = getIntegerValue(right);
75+
if (n !== null && n >= 0 && n <= 10) {
76+
return expandPower(left, n);
77+
}
78+
return pow(left, right);
79+
}
80+
81+
default:
82+
return new ModelicaBinaryExpression(expr.operator, left, right);
83+
}
84+
}
85+
86+
if (expr instanceof ModelicaFunctionCallExpression) {
87+
const args = (expr.args as ModelicaExpression[]).map(expandExpr);
88+
return new ModelicaFunctionCallExpression(expr.functionName, args);
89+
}
90+
91+
return expr;
92+
}
93+
94+
/**
95+
* Distribute multiplication: (a+b)*c → a*c + b*c, a*(b+c) → a*b + a*c
96+
*/
97+
function distributeMultiply(left: ModelicaExpression, right: ModelicaExpression): ModelicaExpression {
98+
// (a + b) * right → a*right + b*right
99+
const leftSum = extractSum(left);
100+
if (leftSum) {
101+
return add(distributeMultiply(leftSum.a, right), distributeMultiply(leftSum.b, right));
102+
}
103+
104+
// left * (a + b) → left*a + left*b
105+
const rightSum = extractSum(right);
106+
if (rightSum) {
107+
return add(distributeMultiply(left, rightSum.a), distributeMultiply(left, rightSum.b));
108+
}
109+
110+
// (a - b) * right → a*right - b*right
111+
const leftDiff = extractDiff(left);
112+
if (leftDiff) {
113+
return sub(distributeMultiply(leftDiff.a, right), distributeMultiply(leftDiff.b, right));
114+
}
115+
116+
// left * (a - b) → left*a - left*b
117+
const rightDiff = extractDiff(right);
118+
if (rightDiff) {
119+
return sub(distributeMultiply(left, rightDiff.a), distributeMultiply(left, rightDiff.b));
120+
}
121+
122+
return mul(left, right);
123+
}
124+
125+
/**
126+
* Expand integer power by repeated multiplication.
127+
* x^0 → 1, x^1 → x, x^n → x * x^(n-1) (expanded)
128+
*/
129+
function expandPower(base: ModelicaExpression, n: number): ModelicaExpression {
130+
if (n === 0) return new ModelicaRealLiteral(1);
131+
if (n === 1) return base;
132+
// Binary exponentiation with expansion
133+
let result = base;
134+
for (let i = 1; i < n; i++) {
135+
result = distributeMultiply(result, base);
136+
}
137+
return result;
138+
}
139+
140+
/** Distribute negation into sums. */
141+
function distributeNeg(expr: ModelicaExpression): ModelicaExpression {
142+
const sum = extractSum(expr);
143+
if (sum) {
144+
return add(distributeNeg(sum.a), distributeNeg(sum.b));
145+
}
146+
const diff = extractDiff(expr);
147+
if (diff) {
148+
return sub(diff.b, diff.a);
149+
}
150+
if (isZero(expr)) return ZERO;
151+
return new ModelicaUnaryExpression(ModelicaUnaryOperator.UNARY_MINUS, expr);
152+
}
153+
154+
// ─────────────────────────────────────────────────────────────────────
155+
// Term Collection
156+
// ─────────────────────────────────────────────────────────────────────
157+
158+
/**
159+
* Collect terms by powers of a variable.
160+
*
161+
* Given an expression that is polynomial in `varName`, returns a map
162+
* from degree → coefficient expression (independent of varName).
163+
*
164+
* Example: 3*x^2 + 2*x + 1 → Map { 2 → 3, 1 → 2, 0 → 1 }
165+
*/
166+
export function collectTerms(expr: ModelicaExpression, varName: string): Map<number, ModelicaExpression> {
167+
const expanded = expandExpr(expr);
168+
const terms = new Map<number, ModelicaExpression>();
169+
170+
function addTerm(degree: number, coeff: ModelicaExpression): void {
171+
const existing = terms.get(degree);
172+
if (existing) {
173+
terms.set(degree, add(existing, coeff));
174+
} else {
175+
terms.set(degree, coeff);
176+
}
177+
}
178+
179+
function collect(e: ModelicaExpression): void {
180+
// Sum: collect each side
181+
const sum = extractSum(e);
182+
if (sum) {
183+
collect(sum.a);
184+
collect(sum.b);
185+
return;
186+
}
187+
188+
// Difference: collect left, negate right
189+
const diff = extractDiff(e);
190+
if (diff) {
191+
collect(diff.a);
192+
collect(new ModelicaUnaryExpression(ModelicaUnaryOperator.UNARY_MINUS, diff.b));
193+
return;
194+
}
195+
196+
// Determine degree and coefficient
197+
const { degree, coeff } = extractDegreeAndCoeff(e, varName);
198+
addTerm(degree, coeff);
199+
}
200+
201+
collect(expanded);
202+
return terms;
203+
}
204+
205+
/**
206+
* Extract the degree and coefficient of a single term with respect to varName.
207+
*/
208+
function extractDegreeAndCoeff(
209+
expr: ModelicaExpression,
210+
varName: string,
211+
): { degree: number; coeff: ModelicaExpression } {
212+
// Variable itself: degree 1, coefficient 1
213+
if (expr instanceof ModelicaNameExpression && expr.name === varName) {
214+
return { degree: 1, coeff: new ModelicaRealLiteral(1) };
215+
}
216+
217+
// Doesn't contain the variable: degree 0
218+
if (!containsVar(expr, varName)) {
219+
return { degree: 0, coeff: expr };
220+
}
221+
222+
// x^n
223+
if (expr instanceof ModelicaBinaryExpression) {
224+
if (
225+
expr.operator === ModelicaBinaryOperator.EXPONENTIATION ||
226+
expr.operator === ModelicaBinaryOperator.ELEMENTWISE_EXPONENTIATION
227+
) {
228+
if (expr.operand1 instanceof ModelicaNameExpression && expr.operand1.name === varName) {
229+
const n = getIntegerValue(expr.operand2);
230+
if (n !== null) return { degree: n, coeff: new ModelicaRealLiteral(1) };
231+
}
232+
}
233+
234+
// a * b: split based on which side contains the variable
235+
if (
236+
expr.operator === ModelicaBinaryOperator.MULTIPLICATION ||
237+
expr.operator === ModelicaBinaryOperator.ELEMENTWISE_MULTIPLICATION
238+
) {
239+
const leftHasVar = containsVar(expr.operand1, varName);
240+
const rightHasVar = containsVar(expr.operand2, varName);
241+
242+
if (leftHasVar && !rightHasVar) {
243+
const inner = extractDegreeAndCoeff(expr.operand1, varName);
244+
return { degree: inner.degree, coeff: mul(inner.coeff, expr.operand2) };
245+
}
246+
if (!leftHasVar && rightHasVar) {
247+
const inner = extractDegreeAndCoeff(expr.operand2, varName);
248+
return { degree: inner.degree, coeff: mul(expr.operand1, inner.coeff) };
249+
}
250+
// Both sides contain the variable — multiply degrees
251+
if (leftHasVar && rightHasVar) {
252+
const leftDC = extractDegreeAndCoeff(expr.operand1, varName);
253+
const rightDC = extractDegreeAndCoeff(expr.operand2, varName);
254+
return {
255+
degree: leftDC.degree + rightDC.degree,
256+
coeff: mul(leftDC.coeff, rightDC.coeff),
257+
};
258+
}
259+
}
260+
}
261+
262+
// Negation
263+
if (expr instanceof ModelicaUnaryExpression && expr.operator === ModelicaUnaryOperator.UNARY_MINUS) {
264+
const inner = extractDegreeAndCoeff(expr.operand, varName);
265+
return {
266+
degree: inner.degree,
267+
coeff: new ModelicaUnaryExpression(ModelicaUnaryOperator.UNARY_MINUS, inner.coeff),
268+
};
269+
}
270+
271+
// Default fallback: treat as opaque (degree 0 is wrong, but safe)
272+
return { degree: 0, coeff: expr };
273+
}
274+
275+
// ─────────────────────────────────────────────────────────────────────
276+
// Normalization (via E-Graph)
277+
// ─────────────────────────────────────────────────────────────────────
278+
279+
/**
280+
* Normalize an expression to a canonical form using the E-Graph engine.
281+
* This first expands, then runs equality saturation to find the simplest form.
282+
*/
283+
export function normalizeExpr(expr: ModelicaExpression): ModelicaExpression {
284+
const expanded = expandExpr(expr);
285+
return egraphSimplify(expanded);
286+
}
287+
288+
// ─────────────────────────────────────────────────────────────────────
289+
// Utilities
290+
// ─────────────────────────────────────────────────────────────────────
291+
292+
/** Extract integer value from a literal expression. */
293+
function getIntegerValue(expr: ModelicaExpression): number | null {
294+
if (expr instanceof ModelicaIntegerLiteral) return expr.value;
295+
if (expr instanceof ModelicaRealLiteral && Number.isInteger(expr.value)) return expr.value;
296+
return null;
297+
}
298+
299+
/** Check if expression contains a variable by name. */
300+
function containsVar(expr: ModelicaExpression, varName: string): boolean {
301+
if (expr instanceof ModelicaNameExpression) return expr.name === varName;
302+
if (expr instanceof ModelicaUnaryExpression) return containsVar(expr.operand, varName);
303+
if (expr instanceof ModelicaBinaryExpression) {
304+
return containsVar(expr.operand1, varName) || containsVar(expr.operand2, varName);
305+
}
306+
if (expr instanceof ModelicaFunctionCallExpression) {
307+
return (expr.args as ModelicaExpression[]).some((a) => containsVar(a, varName));
308+
}
309+
return false;
310+
}
311+
312+
/** Extract addition operands: a + b. */
313+
function extractSum(expr: ModelicaExpression): { a: ModelicaExpression; b: ModelicaExpression } | null {
314+
if (
315+
expr instanceof ModelicaBinaryExpression &&
316+
(expr.operator === ModelicaBinaryOperator.ADDITION || expr.operator === ModelicaBinaryOperator.ELEMENTWISE_ADDITION)
317+
) {
318+
return { a: expr.operand1, b: expr.operand2 };
319+
}
320+
return null;
321+
}
322+
323+
/** Extract subtraction operands: a - b. */
324+
function extractDiff(expr: ModelicaExpression): { a: ModelicaExpression; b: ModelicaExpression } | null {
325+
if (
326+
expr instanceof ModelicaBinaryExpression &&
327+
(expr.operator === ModelicaBinaryOperator.SUBTRACTION ||
328+
expr.operator === ModelicaBinaryOperator.ELEMENTWISE_SUBTRACTION)
329+
) {
330+
return { a: expr.operand1, b: expr.operand2 };
331+
}
332+
return null;
333+
}
334+
335+
/** Check if expression is a literal constant. */
336+
export function isLiteral(expr: ModelicaExpression): boolean {
337+
return expr instanceof ModelicaRealLiteral || expr instanceof ModelicaIntegerLiteral;
338+
}
339+
340+
/** Get numeric value of a literal. */
341+
export function getLiteralValue(expr: ModelicaExpression): number | null {
342+
if (expr instanceof ModelicaRealLiteral) return expr.value;
343+
if (expr instanceof ModelicaIntegerLiteral) return expr.value;
344+
return null;
345+
}
346+
347+
/** Check if expression represents constant one. */
348+
export { isOne, isZero };

0 commit comments

Comments
 (0)