422 lines
11 KiB
Go
422 lines
11 KiB
Go
// Package mathkit is a dependency-free math engine for the problems
|
|
// that actually come up when building and debugging ML code: linear
|
|
// algebra (the shape of every tensor op), calculus (gradients are how
|
|
// everything trains), probability/statistics (data, metrics, losses),
|
|
// numerical optimization (gradient descent, root finding), and the
|
|
// handful of ML-specific primitives (activations, loss functions,
|
|
// least-squares linear regression) that sit on top of those four
|
|
// pillars. It has no third-party dependency — see the comment on
|
|
// pkg/envkit/dotenv.go for why this module keeps go.mod minimal — so
|
|
// the whole engine is just the Go standard library's math/sort/fmt.
|
|
//
|
|
// Every sub-area (linalg.go, calculus.go, statistics.go, optimize.go,
|
|
// ml.go) is usable directly from Go code. solve.go adds one more
|
|
// layer on top: Solve(operation, params) dispatches a string
|
|
// operation name and a loosely-typed params map (exactly the shape
|
|
// a JSON tool-call argument object decodes to) to the right function
|
|
// and renders a human-readable report — that's the layer the agent
|
|
// service's solve_ml_math tool calls (see
|
|
// services/agent/internal/infrastructure/toolexec/mlmath.go), because
|
|
// an LLM tool call arrives as JSON, not as typed Go arguments.
|
|
package mathkit
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"strings"
|
|
"unicode"
|
|
)
|
|
|
|
// Expr is a parsed arithmetic expression that can be evaluated
|
|
// repeatedly against different variable bindings without re-parsing —
|
|
// exactly what Gradient and GradientDescent need, since they evaluate
|
|
// the same expression at many points.
|
|
type Expr struct {
|
|
root node
|
|
src string
|
|
}
|
|
|
|
// String returns the original source text the expression was parsed
|
|
// from, for error messages and report rendering.
|
|
func (e *Expr) String() string { return e.src }
|
|
|
|
// node is the parsed AST. A tiny tagged union kept as an interface
|
|
// rather than a struct-with-a-kind-field: each node type carries only
|
|
// the fields it needs, and eval is a single recursive method.
|
|
type node interface {
|
|
eval(vars map[string]float64) (float64, error)
|
|
}
|
|
|
|
type numNode float64
|
|
|
|
func (n numNode) eval(map[string]float64) (float64, error) { return float64(n), nil }
|
|
|
|
type varNode string
|
|
|
|
func (n varNode) eval(vars map[string]float64) (float64, error) {
|
|
v, ok := vars[string(n)]
|
|
if !ok {
|
|
return 0, fmt.Errorf("unbound variable %q (pass it in \"variables\" or \"at\")", string(n))
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
type binNode struct {
|
|
op byte // '+', '-', '*', '/', '^'
|
|
lhs, rhs node
|
|
}
|
|
|
|
func (n binNode) eval(vars map[string]float64) (float64, error) {
|
|
l, err := n.lhs.eval(vars)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
r, err := n.rhs.eval(vars)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
switch n.op {
|
|
case '+':
|
|
return l + r, nil
|
|
case '-':
|
|
return l - r, nil
|
|
case '*':
|
|
return l * r, nil
|
|
case '/':
|
|
if r == 0 {
|
|
return 0, fmt.Errorf("division by zero")
|
|
}
|
|
return l / r, nil
|
|
case '^':
|
|
return math.Pow(l, r), nil
|
|
}
|
|
return 0, fmt.Errorf("internal: unknown operator %q", n.op)
|
|
}
|
|
|
|
type negNode struct{ x node }
|
|
|
|
func (n negNode) eval(vars map[string]float64) (float64, error) {
|
|
v, err := n.x.eval(vars)
|
|
return -v, err
|
|
}
|
|
|
|
type callNode struct {
|
|
fn string
|
|
args []node
|
|
}
|
|
|
|
// unaryFuncs are functions of exactly one argument — the overwhelming
|
|
// majority of what shows up in ML math (activations, logs, norms).
|
|
var unaryFuncs = map[string]func(float64) float64{
|
|
"sin": math.Sin, "cos": math.Cos, "tan": math.Tan,
|
|
"asin": math.Asin, "acos": math.Acos, "atan": math.Atan,
|
|
"sinh": math.Sinh, "cosh": math.Cosh, "tanh": math.Tanh,
|
|
"exp": math.Exp, "sqrt": math.Sqrt, "abs": math.Abs,
|
|
"ln": math.Log, "log": math.Log, "log2": math.Log2, "log10": math.Log10,
|
|
"floor": math.Floor, "ceil": math.Ceil, "round": math.Round,
|
|
"sigmoid": Sigmoid, "relu": ReLU, "sign": func(x float64) float64 { return math.Copysign(1, x) * boolF(x != 0) },
|
|
}
|
|
|
|
func boolF(b bool) float64 {
|
|
if b {
|
|
return 1
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// binaryFuncs are functions of exactly two arguments.
|
|
var binaryFuncs = map[string]func(float64, float64) float64{
|
|
"max": math.Max, "min": math.Min, "pow": math.Pow,
|
|
"mod": math.Mod, "atan2": math.Atan2,
|
|
}
|
|
|
|
func (n callNode) eval(vars map[string]float64) (float64, error) {
|
|
switch len(n.args) {
|
|
case 1:
|
|
if f, ok := unaryFuncs[n.fn]; ok {
|
|
v, err := n.args[0].eval(vars)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return f(v), nil
|
|
}
|
|
case 2:
|
|
if f, ok := binaryFuncs[n.fn]; ok {
|
|
a, err := n.args[0].eval(vars)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
b, err := n.args[1].eval(vars)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return f(a, b), nil
|
|
}
|
|
}
|
|
return 0, fmt.Errorf("unknown function %q with %d argument(s)", n.fn, len(n.args))
|
|
}
|
|
|
|
// Eval evaluates the expression against a variable binding. Missing
|
|
// variables are a reported error (never silently treated as zero),
|
|
// since a silent zero would turn a typo into a wrong-but-plausible
|
|
// answer — exactly the failure mode numerical tools must not have.
|
|
func (e *Expr) Eval(vars map[string]float64) (float64, error) {
|
|
return e.root.eval(vars)
|
|
}
|
|
|
|
// ParseExpr parses a math expression written in ordinary infix
|
|
// notation: + - * / ^ (right-associative power), unary minus,
|
|
// parentheses, named variables, the constants pi/e, and function
|
|
// calls (sin(x), sigmoid(x), max(a,b), ...). See unaryFuncs and
|
|
// binaryFuncs for the full function list.
|
|
func ParseExpr(src string) (*Expr, error) {
|
|
p := &parser{toks: tokenize(src)}
|
|
n, err := p.parseExpr()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("parsing %q: %w", src, err)
|
|
}
|
|
if p.pos != len(p.toks) {
|
|
return nil, fmt.Errorf("parsing %q: unexpected trailing input near %q", src, p.toks[p.pos].text)
|
|
}
|
|
return &Expr{root: n, src: src}, nil
|
|
}
|
|
|
|
// --- tokenizer ---
|
|
|
|
type tokKind int
|
|
|
|
const (
|
|
tokNum tokKind = iota
|
|
tokIdent
|
|
tokOp
|
|
tokLParen
|
|
tokRParen
|
|
tokComma
|
|
)
|
|
|
|
type token struct {
|
|
kind tokKind
|
|
text string
|
|
num float64
|
|
}
|
|
|
|
func tokenize(src string) []token {
|
|
var toks []token
|
|
r := []rune(src)
|
|
i := 0
|
|
for i < len(r) {
|
|
c := r[i]
|
|
switch {
|
|
case unicode.IsSpace(c):
|
|
i++
|
|
case c == '(':
|
|
toks = append(toks, token{kind: tokLParen, text: "("})
|
|
i++
|
|
case c == ')':
|
|
toks = append(toks, token{kind: tokRParen, text: ")"})
|
|
i++
|
|
case c == ',':
|
|
toks = append(toks, token{kind: tokComma, text: ","})
|
|
i++
|
|
case strings.ContainsRune("+-*/^", c):
|
|
toks = append(toks, token{kind: tokOp, text: string(c)})
|
|
i++
|
|
case unicode.IsDigit(c) || c == '.':
|
|
j := i
|
|
for j < len(r) && (unicode.IsDigit(r[j]) || r[j] == '.') {
|
|
j++
|
|
}
|
|
// scientific notation, e.g. 1e-3
|
|
if j < len(r) && (r[j] == 'e' || r[j] == 'E') {
|
|
k := j + 1
|
|
if k < len(r) && (r[k] == '+' || r[k] == '-') {
|
|
k++
|
|
}
|
|
if k < len(r) && unicode.IsDigit(r[k]) {
|
|
j = k
|
|
for j < len(r) && unicode.IsDigit(r[j]) {
|
|
j++
|
|
}
|
|
}
|
|
}
|
|
text := string(r[i:j])
|
|
var f float64
|
|
fmt.Sscanf(text, "%g", &f)
|
|
toks = append(toks, token{kind: tokNum, text: text, num: f})
|
|
i = j
|
|
case unicode.IsLetter(c) || c == '_':
|
|
j := i
|
|
for j < len(r) && (unicode.IsLetter(r[j]) || unicode.IsDigit(r[j]) || r[j] == '_') {
|
|
j++
|
|
}
|
|
toks = append(toks, token{kind: tokIdent, text: string(r[i:j])})
|
|
i = j
|
|
default:
|
|
// Unknown character: emit as its own operator token so the
|
|
// parser reports a clear "unexpected token" error instead
|
|
// of the tokenizer silently dropping it.
|
|
toks = append(toks, token{kind: tokOp, text: string(c)})
|
|
i++
|
|
}
|
|
}
|
|
return toks
|
|
}
|
|
|
|
// --- recursive-descent parser ---
|
|
//
|
|
// expr := term (('+' | '-') term)*
|
|
// term := unary (('*' | '/') unary)*
|
|
// unary := ('-' | '+') unary | power
|
|
// power := primary ('^' unary)? -- right-associative; the
|
|
// exponent recurses through unary (not power) so the
|
|
// precedence matches ordinary math/Python convention:
|
|
// unary minus binds LOOSER than '^' (-2^2 == -(2^2) == -4)
|
|
// but the exponent itself may carry its own sign
|
|
// (2^-2 == 0.25).
|
|
// primary:= number | const | ident '(' args ')' | ident | '(' expr ')'
|
|
// args := expr (',' expr)*
|
|
|
|
type parser struct {
|
|
toks []token
|
|
pos int
|
|
}
|
|
|
|
func (p *parser) peek() (token, bool) {
|
|
if p.pos >= len(p.toks) {
|
|
return token{}, false
|
|
}
|
|
return p.toks[p.pos], true
|
|
}
|
|
|
|
func (p *parser) next() (token, bool) {
|
|
t, ok := p.peek()
|
|
if ok {
|
|
p.pos++
|
|
}
|
|
return t, ok
|
|
}
|
|
|
|
func (p *parser) parseExpr() (node, error) {
|
|
lhs, err := p.parseTerm()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for {
|
|
t, ok := p.peek()
|
|
if !ok || t.kind != tokOp || (t.text != "+" && t.text != "-") {
|
|
return lhs, nil
|
|
}
|
|
p.next()
|
|
rhs, err := p.parseTerm()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
lhs = binNode{op: t.text[0], lhs: lhs, rhs: rhs}
|
|
}
|
|
}
|
|
|
|
func (p *parser) parseTerm() (node, error) {
|
|
lhs, err := p.parseUnary()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for {
|
|
t, ok := p.peek()
|
|
if !ok || t.kind != tokOp || (t.text != "*" && t.text != "/") {
|
|
return lhs, nil
|
|
}
|
|
p.next()
|
|
rhs, err := p.parseUnary()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
lhs = binNode{op: t.text[0], lhs: lhs, rhs: rhs}
|
|
}
|
|
}
|
|
|
|
func (p *parser) parseUnary() (node, error) {
|
|
if t, ok := p.peek(); ok && t.kind == tokOp && (t.text == "-" || t.text == "+") {
|
|
p.next()
|
|
x, err := p.parseUnary()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if t.text == "-" {
|
|
return negNode{x: x}, nil
|
|
}
|
|
return x, nil
|
|
}
|
|
return p.parsePower()
|
|
}
|
|
|
|
func (p *parser) parsePower() (node, error) {
|
|
lhs, err := p.parsePrimary()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if t, ok := p.peek(); ok && t.kind == tokOp && t.text == "^" {
|
|
p.next()
|
|
rhs, err := p.parseUnary() // right-associative; lets the exponent carry its own sign
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return binNode{op: '^', lhs: lhs, rhs: rhs}, nil
|
|
}
|
|
return lhs, nil
|
|
}
|
|
|
|
func (p *parser) parsePrimary() (node, error) {
|
|
t, ok := p.next()
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected end of expression")
|
|
}
|
|
switch t.kind {
|
|
case tokNum:
|
|
return numNode(t.num), nil
|
|
case tokLParen:
|
|
n, err := p.parseExpr()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if c, ok := p.next(); !ok || c.kind != tokRParen {
|
|
return nil, fmt.Errorf("missing closing ')'")
|
|
}
|
|
return n, nil
|
|
case tokIdent:
|
|
switch strings.ToLower(t.text) {
|
|
case "pi":
|
|
return numNode(math.Pi), nil
|
|
case "e":
|
|
return numNode(math.E), nil
|
|
}
|
|
if nx, ok := p.peek(); ok && nx.kind == tokLParen {
|
|
p.next()
|
|
var args []node
|
|
if c, ok := p.peek(); !ok || c.kind != tokRParen {
|
|
for {
|
|
a, err := p.parseExpr()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
args = append(args, a)
|
|
c, ok := p.peek()
|
|
if !ok {
|
|
return nil, fmt.Errorf("missing closing ')' in call to %s", t.text)
|
|
}
|
|
if c.kind == tokComma {
|
|
p.next()
|
|
continue
|
|
}
|
|
break
|
|
}
|
|
}
|
|
if c, ok := p.next(); !ok || c.kind != tokRParen {
|
|
return nil, fmt.Errorf("missing closing ')' in call to %s", t.text)
|
|
}
|
|
return callNode{fn: strings.ToLower(t.text), args: args}, nil
|
|
}
|
|
return varNode(t.text), nil
|
|
}
|
|
return nil, fmt.Errorf("unexpected token %q", t.text)
|
|
}
|