squiggle/src/distPlus/expressionTree/MathJsParser.re
2020-07-03 14:55:27 -07:00

369 lines
12 KiB
ReasonML

module MathJsonToMathJsAdt = {
type arg =
| Symbol(string)
| Value(float)
| Fn(fn)
| Array(array(arg))
| Object(Js.Dict.t(arg))
and fn = {
name: string,
args: array(arg),
};
let rec run = (j: Js.Json.t) =>
Json.Decode.(
switch (field("mathjs", string, j)) {
| "FunctionNode" =>
let args = j |> field("args", array(run));
Some(
Fn({
name: j |> field("fn", field("name", string)),
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)))
| n =>
Js.log3("Couldn't parse mathjs node", j, n);
None;
}
);
};
module MathAdtToDistDst = {
open MathJsonToMathJsAdt;
module MathAdtCleaner = {
let transformWithSymbol = (f: float, s: string) =>
switch (s) {
| "K"
| "k" => f *. 1000.
| "M"
| "m" => f *. 1000000.
| "B"
| "b" => f *. 1000000000.
| "T"
| "t" => f *. 1000000000000.
| _ => f
};
let rec run =
fun
| Fn({name: "multiply", args: [|Value(f), Symbol(s)|]}) =>
Value(transformWithSymbol(f, s))
| 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)
| Object(v) =>
Object(
v
|> Js.Dict.entries
|> E.A.fmap(((key, value)) => (key, run(value)))
|> Js.Dict.fromArray,
);
};
let normal:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(mean), Value(stdev)|] =>
Ok(`SymbolicDist(`Normal({mean, stdev})))
| _ => Error("Wrong number of variables in normal distribution");
let lognormal:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(mu), Value(sigma)|] =>
Ok(`SymbolicDist(`Lognormal({mu, sigma})))
| [|Object(o)|] => {
let g = Js.Dict.get(o);
switch (g("mean"), g("stdev"), g("mu"), g("sigma")) {
| (Some(Value(mean)), Some(Value(stdev)), _, _) =>
Ok(
`SymbolicDist(
SymbolicDist.Lognormal.fromMeanAndStdev(mean, stdev),
),
)
| (_, _, Some(Value(mu)), Some(Value(sigma))) =>
Ok(`SymbolicDist(`Lognormal({mu, sigma})))
| _ => Error("Lognormal distribution would need mean and stdev")
};
}
| _ => Error("Wrong number of variables in lognormal distribution");
let to_: array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(low), Value(high)|] when low <= 0.0 && low < high => {
Ok(`SymbolicDist(SymbolicDist.Normal.from90PercentCI(low, high)));
}
| [|Value(low), Value(high)|] when low < high => {
Ok(
`SymbolicDist(SymbolicDist.Lognormal.from90PercentCI(low, high)),
);
}
| [|Value(_), Value(_)|] =>
Error("Low value must be less than high value.")
| _ => Error("Wrong number of variables in lognormal distribution");
let uniform:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(low), Value(high)|] =>
Ok(`SymbolicDist(`Uniform({low, high})))
| _ => Error("Wrong number of variables in lognormal distribution");
let beta: array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(alpha), Value(beta)|] =>
Ok(`SymbolicDist(`Beta({alpha, beta})))
| _ => Error("Wrong number of variables in lognormal distribution");
let exponential:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(rate)|] => Ok(`SymbolicDist(`Exponential({rate: rate})))
| _ => Error("Wrong number of variables in Exponential distribution");
let cauchy:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(local), Value(scale)|] =>
Ok(`SymbolicDist(`Cauchy({local, scale})))
| _ => Error("Wrong number of variables in cauchy distribution");
let triangular:
array(arg) => result(ExpressionTypes.ExpressionTree.node, string) =
fun
| [|Value(low), Value(medium), Value(high)|] =>
Ok(`SymbolicDist(`Triangular({low, medium, high})))
| _ => Error("Wrong number of variables in triangle distribution");
let multiModal =
(
args: array(result(ExpressionTypes.ExpressionTree.node, string)),
weights: option(array(float)),
) => {
let weights = weights |> E.O.default([||]);
/*let dists: =
args
|> E.A.fmap(
fun
| Ok(a) => a
| Error(e) => Error(e)
);*/
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 arrayParser =
(args: array(arg))
: result(ExpressionTypes.ExpressionTree.node, string) => {
let samples =
args
|> E.A.fmap(
fun
| Value(n) => Some(n)
| _ => None,
)
|> E.A.O.concatSomes;
let outputs = Samples.T.fromSamples(samples);
let pdf =
outputs.shape |> E.O.bind(_, Distributions.Shape.T.toContinuous);
let shape =
pdf
|> E.O.fmap(pdf => {
let _pdf = Distributions.Continuous.T.normalize(pdf);
let cdf = Distributions.Continuous.T.integral(~cache=None, _pdf);
SymbolicDist.ContinuousShape.make(_pdf, cdf);
});
switch (shape) {
| Some(s) => Ok(`SymbolicDist(`ContinuousShape(s)))
| None => Error("Rendering did not work")
};
};
let operationParser =
(
name: string,
args: array(result(ExpressionTypes.ExpressionTree.node, string)),
) => {
let toOkAlgebraic = r => Ok(`AlgebraicCombination(r));
let toOkTrunctate = r => Ok(`Truncate(r));
switch (name, args) {
| ("add", [|Ok(l), Ok(r)|]) => toOkAlgebraic((`Add, l, r))
| ("add", _) => Error("Addition needs two operands")
| ("subtract", [|Ok(l), Ok(r)|]) => toOkAlgebraic((`Subtract, l, r))
| ("subtract", _) => Error("Subtraction needs two operands")
| ("multiply", [|Ok(l), Ok(r)|]) => toOkAlgebraic((`Multiply, l, r))
| ("multiply", _) => Error("Multiplication needs two operands")
| ("divide", [|Ok(l), Ok(r)|]) => toOkAlgebraic((`Divide, l, r))
| ("divide", _) => Error("Division needs two operands")
| ("pow", _) => Error("Exponentiation is not yet supported.")
| ("leftTruncate", [|Ok(d), Ok(`SymbolicDist(`Float(lc)))|]) =>
toOkTrunctate((Some(lc), None, d))
| ("leftTruncate", _) =>
Error("leftTruncate needs two arguments: the expression and the cutoff")
| ("rightTruncate", [|Ok(d), Ok(`SymbolicDist(`Float(rc)))|]) =>
toOkTrunctate((None, Some(rc), d))
| ("rightTruncate", _) =>
Error(
"rightTruncate needs two arguments: the expression and the cutoff",
)
| (
"truncate",
[|
Ok(d),
Ok(`SymbolicDist(`Float(lc))),
Ok(`SymbolicDist(`Float(rc))),
|],
) =>
toOkTrunctate((Some(lc), Some(rc), d))
| ("truncate", _) =>
Error("truncate needs three arguments: the expression and both cutoffs")
| _ => Error("This type not currently supported")
};
};
let functionParser = (nodeParser, name, args) => {
let parseArgs = () => args |> E.A.fmap(nodeParser);
switch (name) {
| "normal" => normal(args)
| "lognormal" => lognormal(args)
| "uniform" => uniform(args)
| "beta" => beta(args)
| "to" => to_(args)
| "exponential" => exponential(args)
| "cauchy" => cauchy(args)
| "triangular" => triangular(args)
| "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"
| "divide"
| "pow"
| "leftTruncate"
| "rightTruncate"
| "truncate" => operationParser(name, parseArgs())
| "mean" as n
| "inv" as n
| "sample" as n
| "pdf" as n
| n => Error(n ++ "(...) is not currently supported")
};
};
let rec nodeParser =
fun
| Value(f) => Ok(`SymbolicDist(`Float(f)))
| Fn({name, args}) => functionParser(nodeParser, name, args)
| _ => {
Error("This type not currently supported");
};
let topLevel =
fun
| Array(r) => arrayParser(r)
| Value(_) as r => nodeParser(r)
| Fn(_) as r => nodeParser(r)
| Symbol(_) => Error("Symbol not valid as top level")
| Object(_) => Error("Object not valid as top level");
let run = (r): result(ExpressionTypes.ExpressionTree.node, string) =>
r |> MathAdtCleaner.run |> topLevel;
};
let fromString = 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 = Mathjs.parseMath(str);
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;
};