types/mathsfunc.js (117473e) - rtensor · DKL9 GitList

Refactor towards a working QuickJS version

dkl9 commited 117473e at 2023-196 12:46:40

mathsfunc.js

Blame History Raw

"use strict";

import {fts} from "../rtensor_common.js";
import {BasicNumber} from "./basic.js";

// check if vectors (Arrays of numbers) a and b are completely equal
const vecEq = (a, b) => (a?.eq && b?.eq) ?
    (a.eq(b)) :
    // divide-by-1 to avoid undefined == undefined giving true
    ((a.length / 1 == b.length / 1) &&
        a.map((x, i) => (x == b[i])).reduce((p, q) => p && q, true));

// a collection of case -> expr mappings
// each case is a list of zero or more identifiers or constant values
// (in the form of ASTs)
// each expr is an expression (wow); more precisely, a function of the variables
export class MathsFunc extends BasicNumber {
    // cems is an Array of Array-pairs of Array of AST and functions
    // de is an evaluation input in accordance with AST.evaluate
    constructor(cems, de) {
        super();
        this.cases = [];
        for (let i = 0; i < cems.length; i++) {
            this.addCase(cems[i], de);
        }
    }
    
    // cem is an Array-pair of Array of AST and function
    // de is an evaluation input in accordance with AST.evaluate
    addCase(cem, de) {
        const nc = [];
        // for each argument ...
        for (let i = 0; i < cem[0].length; i++) {
            // if the argument is a simple identifier ...
            if (cem[0][i].ts == 1) {
                // it's an input variable
                nc.push(cem[0][i].v || i.toString());
            // else ...
            } else {
                // evaluate the expression to get the constant
                const eres = cem[0][i].evaluate(de);
                if (eres == null) {
                    // ah hecc
                } else {
                    nc.push(eres);
                }
            }
        }
        this.cases.unshift([nc, cem[1]]);
    }
    
    // inps is an Array of numbers or Arrays (vectors)
    // de is an evaluation input in accordance with AST.evaluate
    apply(self, [inps, de]) {
        let rv = null;
        for (let i = 0; i < self.cases.length; i++) {
            const cc = self.cases[i];
            if (cc[0].length == inps.length) {
                let isMatch = true;
                for (let j = 0; j < cc[0].length; j++) {
                    if (!((typeof cc[0][j]) == "string" || vecEq(cc[0][j], inps[j]))) {
                        isMatch = false;
                        break;
                    }
                }
                if (isMatch) {
                    const pinps = Object.fromEntries(inps.map(
                        (x, j) => [cc[0][j], x]).filter(
                        x => (typeof x[0]) == "string"));
                    rv = cc[1].apply(cc[1], [pinps, de]);
                    return rv;
                }
            }
        }
        throw `no suitable argument pattern to match ${inps.map(fts).join(", ")}`;
        return null;
    }
    
    toString() {
        return this.cases.map(
            x => "(" + x.toString() + ")").reduce(
            (a, b) => a + ", " + b);
    }
}

["neg", "recip", "exp", "ln", "sqrt", "sin", "cos", "asin", "acos", "mag", "floor"].forEach(k => MathsFunc.prototype[k] = function() {
    const il = this.cases[0][0].length;
    const f = this;
    return new MathsFunc([[(new Array(il)).fill(0).map(_x => IDENT_AST), function(args, de) {
        const a = f.apply(f, [Object.entries(args).map(([k, v]) => v), de]);
        return vecOp(a, 0, (x, _y) => NumericType.ify(x)[k]());
    }]], {});
});

["add", "mul", "pow"].forEach(k => MathsFunc.prototype[k] = function(g) {
    const il = this.cases[0][0].length;
    const f = this;
    return new MathsFunc([[(new Array(il)).fill(0).map(_x => IDENT_AST), function(args, de) {
        const a = f.apply(f, [Object.entries(args).map(([k, v]) => v), de]);
        const b = g.apply ? g.apply(g, [Object.entries(args).map(([k, v]) => v), de]) : g;
        return vecOp(a, b, (x, y) => NumericType.ify(x)[k](NumericType.ify(y)));
    }]], {});
});