squiggle/src/distPlus/expressionTree/MathJsParser.re
2020-07-31 12:31:39 +01:00

382 lines
12 KiB
ReasonML

module MathJsonToMathJsAdt = {
type arg =
| Symbol(string)
| Value(float)
| Fn(fn)
| Array(array(arg))
| Blocks(array(arg))
| Object(Js.Dict.t(arg))
| Assignment(arg, arg)
| FunctionAssignment(fnAssignment)
and fn = {
name: string,
args: array(arg),
}
and fnAssignment = {
name: string,
args: array(string),
expression: arg,
};
let rec run = (j: Js.Json.t) =>
Json.Decode.(
switch (field("mathjs", string, j)) {
| "FunctionNode" =>
let args = j |> field("args", array(run));
let name = j |> optional(field("fn", field("name", string)));
name |> E.O.fmap(name => Fn({name, args: args |> E.A.O.concatSomes}));
| "OperatorNode" =>
let args = j |> field("args", array(run));
Some(
Fn({
name: j |> field("fn", string),
args: args |> E.A.O.concatSomes,
}),
);
| "ConstantNode" =>
optional(field("value", Json.Decode.float), j)
|> E.O.fmap(r => Value(r))
| "ParenthesisNode" => j |> field("content", run)
| "ObjectNode" =>
let properties = j |> field("properties", dict(run));
Js.Dict.entries(properties)
|> E.A.fmap(((key, value)) => value |> E.O.fmap(v => (key, v)))
|> E.A.O.concatSomes
|> Js.Dict.fromArray
|> (r => Some(Object(r)));
| "ArrayNode" =>
let items = field("items", array(run), j);
Some(Array(items |> E.A.O.concatSomes));
| "SymbolNode" => Some(Symbol(field("name", string, j)))
| "AssignmentNode" =>
let object_ = j |> field("object", run);
let value_ = j |> field("value", run);
switch (object_, value_) {
| (Some(o), Some(v)) => Some(Assignment(o, v))
| _ => None
};
| "BlockNode" =>
let block = r => r |> field("node", run);
let args = j |> field("blocks", array(block)) |> E.A.O.concatSomes;
Some(Blocks(args));
| "FunctionAssignmentNode" =>
let name = j |> field("name", string);
let args = j |> field("params", array(field("name", string)));
let expression = j |> field("expr", run);
expression
|> E.O.fmap(expression =>
FunctionAssignment({name, args, expression})
);
| n =>
Js.log3("Couldn't parse mathjs node", j, n);
None;
}
);
};
module MathAdtToDistDst = {
open MathJsonToMathJsAdt;
let handleSymbol = sym => {
Ok(`Symbol(sym));
};
module MathAdtCleaner = {
let transformWithSymbol = (f: float, s: string) =>
switch (s) {
| "K"
| "k" => Some(f *. 1000.)
| "M"
| "m" => Some(f *. 1000000.)
| "B"
| "b" => Some(f *. 1000000000.)
| "T"
| "t" => Some(f *. 1000000000000.)
| _ => None
};
let rec run =
fun
| Fn({name: "multiply", args: [|Value(f), Symbol(s)|]}) as doNothing =>
transformWithSymbol(f, s)
|> E.O.fmap(r => Value(r))
|> E.O.default(doNothing)
| Fn({name: "unaryMinus", args: [|Value(f)|]}) => Value((-1.0) *. f)
| Fn({name, args}) => Fn({name, args: args |> E.A.fmap(run)})
| Array(args) => Array(args |> E.A.fmap(run))
| Symbol(s) => Symbol(s)
| Value(v) => Value(v)
| Blocks(args) => Blocks(args |> E.A.fmap(run))
| Assignment(a, b) => Assignment(a, run(b))
| FunctionAssignment(a) => FunctionAssignment(a)
| Object(v) =>
Object(
v
|> Js.Dict.entries
|> E.A.fmap(((key, value)) => (key, run(value)))
|> Js.Dict.fromArray,
);
};
let lognormal = (args, parseArgs, nodeParser) =>
switch (args) {
| [|Object(o)|] =>
let g = s =>
Js.Dict.get(o, s)
|> E.O.toResult("Variable was empty")
|> E.R.bind(_, nodeParser);
switch (g("mean"), g("stdev"), g("mu"), g("sigma")) {
| (Ok(mean), Ok(stdev), _, _) =>
Ok(
`FunctionCall(("lognormalFromMeanAndStdDev", [|mean, stdev|])),
)
| (_, _, Ok(mu), Ok(sigma)) =>
Ok(`FunctionCall(("lognormal", [|mu, sigma|])))
| _ =>
Error(
"Lognormal distribution needs either mean and stdev or mu and sigma",
)
};
| _ =>
parseArgs()
|> E.R.fmap((args: array(ExpressionTypes.ExpressionTree.node)) =>
`FunctionCall(("lognormal", args))
)
};
let multiModal =
(
args: array(result(ExpressionTypes.ExpressionTree.node, string)),
weights: option(array(float)),
) => {
let weights = weights |> E.O.default([||]);
let firstWithError = args |> Belt.Array.getBy(_, Belt.Result.isError);
let withoutErrors = args |> E.A.fmap(E.R.toOption) |> E.A.O.concatSomes;
switch (firstWithError) {
| Some(Error(e)) => Error(e)
| None when withoutErrors |> E.A.length == 0 =>
Error("Multimodals need at least one input")
| _ =>
let components =
withoutErrors
|> E.A.fmapi((index, t) => {
let w = weights |> E.A.get(_, index) |> E.O.default(1.0);
`VerticalScaling((`Multiply, t, `SymbolicDist(`Float(w))));
});
let pointwiseSum =
components
|> Js.Array.sliceFrom(1)
|> E.A.fold_left(
(acc, x) => {`PointwiseCombination((`Add, acc, x))},
E.A.unsafe_get(components, 0),
);
Ok(`Normalize(pointwiseSum));
};
};
let operationParser =
(
name: string,
args: result(array(ExpressionTypes.ExpressionTree.node), string),
) => {
let toOkAlgebraic = r => Ok(`AlgebraicCombination(r));
let toOkPointwise = r => Ok(`PointwiseCombination(r));
let toOkTruncate = r => Ok(`Truncate(r));
let toOkFloatFromDist = r => Ok(`FloatFromDist(r));
args
|> E.R.bind(_, args => {
switch (name, args) {
| ("add", [|l, r|]) => toOkAlgebraic((`Add, l, r))
| ("add", _) => Error("Addition needs two operands")
| ("subtract", [|l, r|]) => toOkAlgebraic((`Subtract, l, r))
| ("subtract", _) => Error("Subtraction needs two operands")
| ("multiply", [|l, r|]) => toOkAlgebraic((`Multiply, l, r))
| ("multiply", _) => Error("Multiplication needs two operands")
| ("dotMultiply", [|l, r|]) => toOkPointwise((`Multiply, l, r))
| ("dotMultiply", _) =>
Error("Dotwise multiplication needs two operands")
| ("rightLogShift", [|l, r|]) => toOkPointwise((`Add, l, r))
| ("rightLogShift", _) =>
Error("Dotwise addition needs two operands")
| ("divide", [|l, r|]) => toOkAlgebraic((`Divide, l, r))
| ("divide", _) => Error("Division needs two operands")
| ("pow", _) => Error("Exponentiation is not yet supported.")
| ("leftTruncate", [|d, `SymbolicDist(`Float(lc))|]) =>
toOkTruncate((Some(lc), None, d))
| ("leftTruncate", _) =>
Error(
"leftTruncate needs two arguments: the expression and the cutoff",
)
| ("rightTruncate", [|d, `SymbolicDist(`Float(rc))|]) =>
toOkTruncate((None, Some(rc), d))
| ("rightTruncate", _) =>
Error(
"rightTruncate needs two arguments: the expression and the cutoff",
)
| (
"truncate",
[|d, `SymbolicDist(`Float(lc)), `SymbolicDist(`Float(rc))|],
) =>
toOkTruncate((Some(lc), Some(rc), d))
| ("truncate", _) =>
Error(
"truncate needs three arguments: the expression and both cutoffs",
)
| ("pdf", [|d, `SymbolicDist(`Float(v))|]) =>
toOkFloatFromDist((`Pdf(v), d))
| ("cdf", [|d, `SymbolicDist(`Float(v))|]) =>
toOkFloatFromDist((`Cdf(v), d))
| ("inv", [|d, `SymbolicDist(`Float(v))|]) =>
toOkFloatFromDist((`Inv(v), d))
| ("mean", [|d|]) => toOkFloatFromDist((`Mean, d))
| ("sample", [|d|]) => toOkFloatFromDist((`Sample, d))
| _ => Error("This type not currently supported")
}
});
};
let functionParser = (nodeParser, name, args) => {
let parseArray = ags =>
ags |> E.A.fmap(nodeParser) |> E.A.R.firstErrorOrOpen;
let parseArgs = () => parseArray(args);
switch (name) {
| "normal"
| "uniform"
| "beta"
| "triangular"
| "to"
| "exponential"
| "cauchy" =>
parseArgs()
|> E.R.fmap((args: array(ExpressionTypes.ExpressionTree.node)) =>
`FunctionCall((name, args))
)
| "lognormal" => lognormal(args, parseArgs, nodeParser)
| "mm" =>
let weights =
args
|> E.A.last
|> E.O.bind(
_,
fun
| Array(values) => Some(values)
| _ => None,
)
|> E.O.fmap(o =>
o
|> E.A.fmap(
fun
| Value(r) => Some(r)
| _ => None,
)
|> E.A.O.concatSomes
);
let possibleDists =
E.O.isSome(weights)
? Belt.Array.slice(args, ~offset=0, ~len=E.A.length(args) - 1)
: args;
let dists = possibleDists |> E.A.fmap(nodeParser);
multiModal(dists, weights);
| "add"
| "subtract"
| "multiply"
| "dotMultiply"
| "rightLogShift"
| "divide"
| "pow"
| "leftTruncate"
| "rightTruncate"
| "truncate"
| "mean"
| "inv"
| "sample"
| "cdf"
| "pdf" => operationParser(name, parseArgs())
| name =>
parseArgs()
|> E.R.fmap((args: array(ExpressionTypes.ExpressionTree.node)) =>
`FunctionCall((name, args))
)
};
};
let rec nodeParser:
MathJsonToMathJsAdt.arg =>
result(ExpressionTypes.ExpressionTree.node, string) =
fun
| Value(f) => Ok(`SymbolicDist(`Float(f)))
| Symbol(sym) => Ok(`Symbol(sym))
| Fn({name, args}) => functionParser(nodeParser, name, args)
| _ => {
Error("This type not currently supported")
};
// | FunctionAssignment({name, args, expression}) => {
// let evaluatedExpression = run(expression);
// `Function(_ => Ok(evaluatedExpression));
// }
let rec topLevel = (r): result(ExpressionTypes.Program.program, string) =>
switch (r) {
| FunctionAssignment({name, args, expression}) =>
switch (nodeParser(expression)) {
| Ok(r) => Ok([|`Assignment((name, `Function((args, r))))|])
| Error(r) => Error(r)
}
| Value(_) as r => nodeParser(r) |> E.R.fmap(r => [|`Expression(r)|])
| Fn(_) as r => nodeParser(r) |> E.R.fmap(r => [|`Expression(r)|])
| Array(_) => Error("Array not valid as top level")
| Symbol(s) => handleSymbol(s) |> E.R.fmap(r => [|`Expression(r)|])
| Object(_) => Error("Object not valid as top level")
| Assignment(name, value) =>
switch (name) {
| Symbol(symbol) =>
nodeParser(value) |> E.R.fmap(r => [|`Assignment((symbol, r))|])
| _ => Error("Symbol not a string")
}
| Blocks(blocks) =>
blocks
|> E.A.fmap(b => topLevel(b))
|> E.A.R.firstErrorOrOpen
|> E.R.fmap(E.A.concatMany)
};
let run = (r): result(ExpressionTypes.Program.program, string) =>
r |> MathAdtCleaner.run |> topLevel;
};
/* The MathJs parser doesn't support '.+' syntax, but we want it because it
would make sense with '.*'. Our workaround is to change this to >>>, which is
logShift in mathJS. We don't expect to use logShift anytime soon, so this tradeoff
seems fine.
*/
let pointwiseToRightLogShift = Js.String.replaceByRe([%re "/\.\+/g"], ">>>");
let fromString2 = str => {
/* We feed the user-typed string into Mathjs.parseMath,
which returns a JSON with (hopefully) a single-element array.
This array element is the top-level node of a nested-object tree
representing the functions/arguments/values/etc. in the string.
The function MathJsonToMathJsAdt then recursively unpacks this JSON into a typed data structure we can use.
Inside of this function, MathAdtToDistDst is called whenever a distribution function is encountered.
*/
let mathJsToJson = str |> pointwiseToRightLogShift |> Mathjs.parseMath;
let mathJsParse =
E.R.bind(mathJsToJson, r => {
switch (MathJsonToMathJsAdt.run(r)) {
| Some(r) => Ok(r)
| None => Error("MathJsParse Error")
}
});
let value = E.R.bind(mathJsParse, MathAdtToDistDst.run);
value;
};
let fromString = str => {
fromString2(str);
};