Spaces:
Running
Running
| import type { Graph, GraphNode, Mode, Operation } from "./types"; | |
| // Rounded for reading | |
| export function formatNumber(value: number): string { | |
| return String(Number(value.toPrecision(5))); | |
| } | |
| // A multi-part argument, such as the constant "2 * 3", in brackets | |
| function atom(arg: string): string { | |
| return arg.includes(" ") ? `(${arg})` : arg; | |
| } | |
| // Each operation's derivative with respect to each argument; backend/src/logic.py's PARTIALS as text | |
| const PARTIALS: Record<Operation, ((a: string, b: string) => string)[]> = { | |
| add: [() => "1", () => "1"], | |
| sub: [() => "1", () => "−1"], | |
| mul: [(_, b) => b, (a) => a], | |
| div: [(_, b) => `1 / ${b}`, (a, b) => `−${a} / ${b}^2`], | |
| pow: [(a, b) => `${b} · ${a}^(${b} − 1)`, (a, b) => `${a}^${b} · log(${a})`], | |
| neg: [() => "−1"], | |
| sin: [(a) => `cos(${a})`], | |
| cos: [(a) => `−sin(${a})`], | |
| tan: [(a) => `1 / cos(${a})^2`], | |
| exp: [(a) => `exp(${a})`], | |
| log: [(a) => `1 / ${a}`], | |
| sqrt: [(a) => `1 / (2 · sqrt(${a}))`], | |
| }; | |
| export function partialOf(op: Operation, index: number, args: string[]): string { | |
| return PARTIALS[op][index](atom(args[0]), atom(args[1] ?? "")); | |
| } | |
| // An id with a bar over its first letter: the adjoint of that node | |
| export function adjoint(id: string): string { | |
| return `${id[0]}̄${id.slice(1)}`; | |
| } | |
| // An id with a dot over its first letter: the tangent of that node | |
| function tangent(id: string): string { | |
| return `${id[0]}\u0307${id.slice(1)}`; | |
| } | |
| export function derivativeSymbol(mode: Mode, id: string): string { | |
| return mode === "forward" ? tangent(id) : adjoint(id); | |
| } | |
| // A derivative times a local partial, leaving out a factor of 1 | |
| function term(symbol: string, partial: string): string { | |
| if (partial === "1") { | |
| return symbol; | |
| } | |
| if (partial === "−1") { | |
| return `−${symbol}`; | |
| } | |
| return `${symbol} · ${partial.includes(" ") ? `(${partial})` : partial}`; | |
| } | |
| function sumOf(terms: string[]): string { | |
| return terms.join(" + ").replaceAll("+ −", "− "); | |
| } | |
| // How reverse mode gets a node's adjoint: one term for each use of the node | |
| export function adjointFormula(graph: Graph, node: GraphNode): string { | |
| if (node.id === graph.output) { | |
| return `${adjoint(node.id)} = 1`; | |
| } | |
| const terms: string[] = []; | |
| for (const consumer of graph.nodes) { | |
| if (consumer.kind === "op") { | |
| consumer.args.forEach((arg, i) => { | |
| if (arg === node.id) { | |
| terms.push(term(adjoint(consumer.id), partialOf(consumer.op, i, consumer.args))); | |
| } | |
| }); | |
| } | |
| } | |
| return `${adjoint(node.id)} = ${sumOf(terms)}`; | |
| } | |
| // A node's arguments that have derivatives, in the order of a step's partials; constants don't | |
| export function nodeArgs(graph: Graph, node: GraphNode & { kind: "op" }): { arg: string; index: number }[] { | |
| return node.args.map((arg, index) => ({ arg, index })).filter(({ arg }) => !(arg in graph.constants)); | |
| } | |
| // How forward mode gets a node's tangent: one term for each of its arguments | |
| export function tangentFormula(graph: Graph, node: GraphNode, seed: string): string { | |
| if (node.kind === "input") { | |
| return `${tangent(node.id)} = ${node.name === seed ? 1 : 0}`; | |
| } | |
| const terms = nodeArgs(graph, node).map(({ arg, index }) => term(tangent(arg), partialOf(node.op, index, node.args))); | |
| return `${tangent(node.id)} = ${sumOf(terms)}`; | |
| } | |