Add differential equation solvers

src/compute-engine/differential-equation-utils.ts

@@ -0,0 +1,47 @@
+import type { Expression, IComputeEngine } from './global-types';
+import { isFunction, isSymbol } from './boxed-expression/type-guards';
+
+export function symbolArg(
+  engine: IComputeEngine,
+  arg: Expression | undefined
+): Expression {
+  if (arg === undefined) return engine.error('missing');
+  if (!isSymbol(arg)) return engine.typeError('symbol', arg.type, arg);
+  return arg;
+}
+
+export function isDependentFunction(
+  expr: Expression,
+  dependentName: string,
+  independentName: string
+): boolean {
+  return (
+    isFunction(expr) &&
+    expr.operator === dependentName &&
+    expr.nops === 1 &&
+    isSymbol(expr.op1, independentName)
+  );
+}
+
+export function isDerivativeOfDependent(
+  expr: Expression,
+  dependentName: string,
+  independentName: string
+): boolean {
+  if (isFunction(expr, 'D')) {
+    return (
+      isDependentFunction(expr.op1, dependentName, independentName) &&
+      isSymbol(expr.op2, independentName)
+    );
+  }
+
+  if (isFunction(expr, 'Apply') && isFunction(expr.op1, 'Derivative')) {
+    return (
+      isSymbol(expr.op1.op1, dependentName) &&
+      expr.nops === 2 &&
+      isSymbol(expr.op2, independentName)
+    );
+  }
+
+  return false;
+}

src/compute-engine/library/calculus.ts

@@ -13,8 +13,11 @@ import {
 import { monteCarloEstimate } from '../numerics/monte-carlo';
 import { integrateSemiInfiniteOscillatory } from '../numerics/oscillatory-quadrature';
 import { centeredDiff8thOrder, limit } from '../numerics/numeric';
+import { nDSolve } from '../numerics/differential-equations';
 import { derivative, differentiate } from '../symbolic/derivative';
 import { antiderivative } from '../symbolic/antiderivative';
+import { dSolve } from '../symbolic/differential-equations';
+import { symbolArg } from '../differential-equation-utils';
 import { symbolicLimit } from '../symbolic/limit';
 import { residue } from '../symbolic/residue';
 import { canonicalLimits, canonicalLimitsSequence } from './utils';
@@ -465,6 +468,66 @@ volumes
       },
     },
 
+    DSolve: {
+      description: 'Symbolic differential equation solver.',
+      broadcastable: false,
+      lazy: true,
+      signature: '(expression, symbol, symbol) -> expression',
+      canonical: (ops, { engine }) => {
+        if (ops.length === 0)
+          return engine._fn('DSolve', [
+            engine.error('missing'),
+            engine.error('missing'),
+            engine.error('missing'),
+          ]);
+        if (ops.length === 1)
+          return engine._fn('DSolve', [
+            ops[0],
+            engine.error('missing'),
+            engine.error('missing'),
+          ]);
+        if (ops.length === 2)
+          return engine._fn('DSolve', [
+            ops[0],
+            symbolArg(engine, ops[1]),
+            engine.error('missing'),
+          ]);
+
+        return engine._fn('DSolve', [
+          ops[0],
+          symbolArg(engine, ops[1]),
+          symbolArg(engine, ops[2]),
+        ]);
+      },
+      evaluate: ([equation, dependent, independent]) =>
+        dSolve(equation, dependent, independent),
+    },
+
+    NDSolve: {
+      description: 'Numerical differential equation solver.',
+      broadcastable: false,
+      lazy: true,
+      signature:
+        '(expression, symbol, limits:(tuple|symbol), number, number?) -> list',
+      canonical: (ops, { engine }) => {
+        const missing = engine.error('missing');
+        const limits =
+          ops[2] && isFunction(ops[2])
+            ? canonicalLimits(ops[2].ops, { engine })
+            : canonicalLimits(ops[2] ? [ops[2]] : [], { engine });
+
+        return engine._fn('NDSolve', [
+          ops[0] ?? missing,
+          symbolArg(engine, ops[1]),
+          limits ?? missing,
+          ops[3]?.canonical ?? missing,
+          ...(ops[4] ? [ops[4].canonical] : []),
+        ]);
+      },
+      evaluate: ([equation, dependent, limits, initialValue, steps]) =>
+        nDSolve(equation, dependent, limits, initialValue, steps),
+    },
+
     // This is used to represent the indexing set/limits (i.e.
     // an index, lower and upper bounds) of a function
     // (not to be confused with Limit, which calculates the limit of a

src/compute-engine/numerics/differential-equations.ts

@@ -0,0 +1,148 @@
+import { checkDeadline } from '../../common/interruptible';
+import type { Expression } from '../global-types';
+import { isFunction, sym } from '../boxed-expression/type-guards';
+import {
+  isDependentFunction,
+  isDerivativeOfDependent,
+} from '../differential-equation-utils';
+
+export type RK4Options = {
+  steps: number;
+  deadline?: number;
+};
+
+export type ODESample = readonly [x: number, y: number];
+
+/**
+ * Fixed-step classical fourth-order Runge-Kutta solver for scalar explicit
+ * initial value problems: y' = f(x, y), y(x0) = y0.
+ */
+export function rk4(
+  f: (x: number, y: number) => number,
+  x0: number,
+  y0: number,
+  x1: number,
+  options: RK4Options
+): ODESample[] | undefined {
+  const steps = Math.trunc(options.steps);
+  if (
+    !Number.isFinite(x0) ||
+    !Number.isFinite(y0) ||
+    !Number.isFinite(x1) ||
+    !Number.isInteger(steps) ||
+    steps <= 0
+  )
+    return undefined;
+
+  const h = (x1 - x0) / steps;
+  const samples: ODESample[] = [[x0, y0]];
+  let x = x0;
+  let y = y0;
+
+  for (let i = 0; i < steps; i++) {
+    if ((i & 0xff) === 0) checkDeadline(options.deadline);
+
+    const k1 = f(x, y);
+    const k2 = f(x + h / 2, y + (h * k1) / 2);
+    const k3 = f(x + h / 2, y + (h * k2) / 2);
+    const k4 = f(x + h, y + h * k3);
+    if (![k1, k2, k3, k4].every(Number.isFinite)) return undefined;
+
+    y += (h / 6) * (k1 + 2 * k2 + 2 * k3 + k4);
+    x = i === steps - 1 ? x1 : x + h;
+    if (!Number.isFinite(y)) return undefined;
+    samples.push([x, y]);
+  }
+
+  return samples;
+}
+
+function explicitRhs(
+  equation: Expression,
+  dependentName: string,
+  independentName: string
+): Expression | undefined {
+  if (!isFunction(equation, 'Equal')) return undefined;
+  if (isDerivativeOfDependent(equation.op1, dependentName, independentName))
+    return equation.op2;
+  if (isDerivativeOfDependent(equation.op2, dependentName, independentName))
+    return equation.op1;
+  return undefined;
+}
+
+function substituteDependentCall(
+  expr: Expression,
+  dependentName: string,
+  independentName: string,
+  stateName: string
+): Expression {
+  if (isDependentFunction(expr, dependentName, independentName))
+    return expr.engine.symbol(stateName);
+  if (!isFunction(expr)) return expr;
+  return expr.engine._fn(
+    expr.operator,
+    expr.ops.map((op) =>
+      substituteDependentCall(op, dependentName, independentName, stateName)
+    )
+  );
+}
+
+export function nDSolve(
+  equation: Expression,
+  dependent: Expression,
+  limits: Expression,
+  initialValue: Expression,
+  stepsExpr?: Expression
+): Expression | undefined {
+  const ce = equation.engine;
+  const dependentName = sym(dependent);
+  if (!dependentName) return undefined;
+
+  if (!isFunction(limits, 'Limits')) return undefined;
+  const independentName = sym(limits.op1);
+  if (!independentName) return undefined;
+
+  const [x0, x1, y0] = [
+    limits.op2.N().re,
+    limits.op3.N().re,
+    initialValue.N().re,
+  ];
+  if (![x0, x1, y0].every(Number.isFinite)) return undefined;
+
+  const steps = stepsExpr === undefined ? 100 : stepsExpr.N().re;
+  if (
+    !Number.isInteger(steps) ||
+    steps <= 0 ||
+    steps > ce.iterationLimit ||
+    steps + 1 > ce.maxCollectionSize
+  )
+    return undefined;
+
+  const rhs = explicitRhs(equation.structural, dependentName, independentName);
+  if (!rhs) return undefined;
+
+  const stateName = `ndsolve${dependentName}state`;
+  const compiledRhs = substituteDependentCall(
+    rhs,
+    dependentName,
+    independentName,
+    stateName
+  );
+  const compiled = ce._compile(compiledRhs, { realOnly: true });
+  if (!compiled.success) return undefined;
+  const run = compiled.run as (vars: Record<string, number>) => number;
+
+  const samples = rk4(
+    (x, y) => run({ [independentName]: x, [stateName]: y }),
+    x0,
+    y0,
+    x1,
+    { steps, deadline: ce._deadline }
+  );
+  if (!samples) return undefined;
+
+  return ce._fn(
+    'List',
+    samples.map(([x, y]) => ce._fn('List', [ce.number(x), ce.number(y)]))
+  );
+}