Add symbolic differential and recurrence solvers

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

@@ -11,6 +11,19 @@ export function symbolArg(
   return arg;
 }
 
+export function symbolOrListArg(
+  engine: IComputeEngine,
+  arg: Expression | undefined
+): Expression {
+  if (arg === undefined) return engine.error('missing');
+  if (isSymbol(arg)) return arg;
+  if (isFunction(arg, 'List')) {
+    const symbols = arg.ops.map((op) => symbolArg(engine, op));
+    return engine._fn('List', symbols);
+  }
+  return engine.typeError('symbol', arg.type, arg);
+}
+
 export function isDependentFunction(
   expr: Expression,
   dependentName: string,
@@ -78,7 +91,7 @@ export function derivativeOrderOfDependent(
   return undefined;
 }
 
-function explicitDerivativeRhs(
+export function explicitDerivativeRhs(
   equation: Expression,
   dependentName: string,
   independentName: string
@@ -99,6 +112,15 @@ function explicitDerivativeRhs(
   return undefined;
 }
 
+function dependentNames(dependent: Expression): string[] | undefined {
+  if (isSymbol(dependent)) return [dependent.symbol];
+  if (!isFunction(dependent, 'List')) return undefined;
+  const names = dependent.ops.map((op) => sym(op));
+  if (names.some((name) => name === undefined)) return undefined;
+  const result = names as string[];
+  return new Set(result).size === result.length ? result : undefined;
+}
+
 function substituteDependentCall(
   expr: Expression,
   dependentName: string,
@@ -116,6 +138,31 @@ function substituteDependentCall(
   );
 }
 
+function substituteSystemDependentCalls(
+  expr: Expression,
+  dependentNames: readonly string[],
+  independentName: string,
+  stateNames: readonly string[]
+): Expression {
+  for (let i = 0; i < dependentNames.length; i++) {
+    if (isDependentFunction(expr, dependentNames[i], independentName))
+      return expr.engine.symbol(stateNames[i]);
+  }
+
+  if (!isFunction(expr)) return expr;
+  return expr.engine._fn(
+    expr.operator,
+    expr.ops.map((op) =>
+      substituteSystemDependentCalls(
+        op,
+        dependentNames,
+        independentName,
+        stateNames
+      )
+    )
+  );
+}
+
 function substituteDependentState(
   expr: Expression,
   dependentName: string,
@@ -150,8 +197,8 @@ export function nDSolve(
   stepsExpr?: Expression
 ): Expression | undefined {
   const ce = equation.engine;
-  const dependentName = sym(dependent);
-  if (!dependentName) return undefined;
+  const names = dependentNames(dependent);
+  if (!names) return undefined;
 
   if (!isFunction(limits, 'Limits')) return undefined;
   const independentName = sym(limits.op1);
@@ -169,6 +216,72 @@ export function nDSolve(
   )
     return undefined;
 
+  if (names.length > 1 || isFunction(equation, 'List')) {
+    if (
+      names.length === 1 ||
+      !isFunction(equation, 'List') ||
+      !isFunction(initialValue, 'List') ||
+      equation.ops.length !== names.length ||
+      initialValue.ops.length !== names.length
+    )
+      return undefined;
+
+    const initialValues = initialValue.ops.map((op) => op.N().re);
+    if (!initialValues.every(Number.isFinite)) return undefined;
+
+    const stateNames = names.map((name, i) => `ndsolve${name}state${i}`);
+    const runs: ((vars: Record<string, number>) => number)[] = [];
+
+    for (let i = 0; i < equation.ops.length; i++) {
+      const rhsInfo = explicitDerivativeRhs(
+        equation.ops[i].structural,
+        names[i],
+        independentName
+      );
+      if (!rhsInfo || rhsInfo.order !== 1) return undefined;
+
+      const compiledRhs = substituteSystemDependentCalls(
+        rhsInfo.rhs,
+        names,
+        independentName,
+        stateNames
+      );
+      const compiled = ce._compile(compiledRhs, { realOnly: true });
+      if (!compiled.success) return undefined;
+      runs.push(compiled.run as (vars: Record<string, number>) => number);
+    }
+
+    const samples = rk4System(
+      (x, y) => {
+        const vars: Record<string, number> = { [independentName]: x };
+        stateNames.forEach((name, i) => {
+          vars[name] = y[i];
+        });
+        const values = runs.map((run) => run(vars));
+        return values.every(Number.isFinite) ? values : undefined;
+      },
+      x0,
+      initialValues,
+      x1,
+      { steps, deadline: ce._deadline }
+    );
+    if (!samples) return undefined;
+
+    return ce._fn(
+      'List',
+      samples.map(([x, y]) =>
+        ce._fn('List', [
+          ce.number(x),
+          ce._fn(
+            'List',
+            y.map((yi) => ce.number(yi))
+          ),
+        ])
+      )
+    );
+  }
+
+  const dependentName = names[0];
   const rhsInfo = explicitDerivativeRhs(
     equation.structural,
     dependentName,

src/compute-engine/library/calculus.ts

@@ -18,7 +18,12 @@ import { derivative, differentiate } from '../symbolic/derivative';
 import '../symbolic/explain-derivative';
 import { antiderivative } from '../symbolic/antiderivative';
 import { dSolve } from '../symbolic/differential-equations';
-import { nDSolve, symbolArg } from '../differential-equation-utils';
+import { rSolve } from '../symbolic/recurrences';
+import {
+  nDSolve,
+  symbolArg,
+  symbolOrListArg,
+} from '../differential-equation-utils';
 import { symbolicLimit } from '../symbolic/limit';
 import { residue } from '../symbolic/residue';
 import { computeSeries, normalStrip } from '../symbolic/series';
@@ -573,20 +578,55 @@ volumes
         if (ops.length === 2)
           return engine._fn('DSolve', [
             ops[0],
-            symbolArg(engine, ops[1]),
+            symbolOrListArg(engine, ops[1]),
             engine.error('missing'),
           ]);
 
         return engine._fn('DSolve', [
           ops[0],
-          symbolArg(engine, ops[1]),
+          symbolOrListArg(engine, ops[1]),
           symbolArg(engine, ops[2]),
         ]);
       },
       evaluate: ([equation, dependent, independent]) =>
         dSolve(equation, dependent, independent),
     },
 
+    RSolve: {
+      description: 'Symbolic recurrence equation solver.',
+      broadcastable: false,
+      lazy: true,
+      signature: '(expression, symbol, symbol) -> expression',
+      canonical: (ops, { engine }) => {
+        if (ops.length === 0)
+          return engine._fn('RSolve', [
+            engine.error('missing'),
+            engine.error('missing'),
+            engine.error('missing'),
+          ]);
+        if (ops.length === 1)
+          return engine._fn('RSolve', [
+            ops[0],
+            engine.error('missing'),
+            engine.error('missing'),
+          ]);
+        if (ops.length === 2)
+          return engine._fn('RSolve', [
+            ops[0],
+            symbolArg(engine, ops[1]),
+            engine.error('missing'),
+          ]);
+
+        return engine._fn('RSolve', [
+          ops[0],
+          symbolArg(engine, ops[1]),
+          symbolArg(engine, ops[2]),
+        ]);
+      },
+      evaluate: ([equation, dependent, index]) =>
+        rSolve(equation, dependent, index),
+    },
+
     NDSolve: {
       description: 'Numerical differential equation solver.',
       broadcastable: false,
@@ -602,7 +642,7 @@ volumes
 
         return engine._fn('NDSolve', [
           ops[0] ?? missing,
-          symbolArg(engine, ops[1]),
+          symbolOrListArg(engine, ops[1]),
           limits ?? missing,
           ops[3]?.canonical ?? missing,
           ...(ops[4] ? [ops[4].canonical] : []),