297 lines
9.4 KiB
ReasonML
297 lines
9.4 KiB
ReasonML
open ExpressionTypes;
|
|
open ExpressionTypes.ExpressionTree;
|
|
|
|
type t = node;
|
|
type tResult = node => result(node, string);
|
|
|
|
/* Given two random variables A and B, this returns the distribution
|
|
of a new variable that is the result of the operation on A and B.
|
|
For instance, normal(0, 1) + normal(1, 1) -> normal(1, 2).
|
|
In general, this is implemented via convolution. */
|
|
module AlgebraicCombination = {
|
|
let tryAnalyticalSimplification = (operation, t1: t, t2: t) =>
|
|
switch (operation, t1, t2) {
|
|
| (operation, `SymbolicDist(d1), `SymbolicDist(d2)) =>
|
|
switch (SymbolicDist.T.tryAnalyticalSimplification(d1, d2, operation)) {
|
|
| `AnalyticalSolution(symbolicDist) => Ok(`SymbolicDist(symbolicDist))
|
|
| `Error(er) => Error(er)
|
|
| `NoSolution => Ok(`AlgebraicCombination((operation, t1, t2)))
|
|
}
|
|
| _ => Ok(`AlgebraicCombination((operation, t1, t2)))
|
|
};
|
|
|
|
let tryCombination = (n, algebraicOp, t1: node, t2: node) => {
|
|
let sampleN =
|
|
mapRenderable(Shape.sampleNRendered(n), SymbolicDist.T.sampleN(n));
|
|
switch (sampleN(t1), sampleN(t2)) {
|
|
| (Some(a), Some(b)) =>
|
|
Some(
|
|
Belt.Array.zip(a, b)
|
|
|> E.A.fmap(((a, b)) => Operation.Algebraic.toFn(algebraicOp, a, b)),
|
|
)
|
|
| _ => None
|
|
};
|
|
};
|
|
|
|
let renderIfNotRendered = (params, t) =>
|
|
!renderable(t)
|
|
? switch (render(params, t)) {
|
|
| Ok(r) => Ok(r)
|
|
| Error(e) => Error(e)
|
|
}
|
|
: Ok(t);
|
|
|
|
let combineAsShapes =
|
|
(evaluationParams: evaluationParams, algebraicOp, t1: node, t2: node) => {
|
|
let i1 = renderIfNotRendered(evaluationParams, t1);
|
|
let i2 = renderIfNotRendered(evaluationParams, t2);
|
|
E.R.merge(i1, i2)
|
|
|> E.R.bind(
|
|
_,
|
|
((a, b)) => {
|
|
let samples =
|
|
tryCombination(
|
|
evaluationParams.samplingInputs.sampleCount,
|
|
algebraicOp,
|
|
a,
|
|
b,
|
|
);
|
|
let shape =
|
|
samples
|
|
|> E.O.fmap(
|
|
Samples.T.fromSamples(
|
|
~samplingInputs={
|
|
sampleCount:
|
|
Some(evaluationParams.samplingInputs.sampleCount),
|
|
outputXYPoints:
|
|
Some(evaluationParams.samplingInputs.outputXYPoints),
|
|
kernelWidth: evaluationParams.samplingInputs.kernelWidth,
|
|
},
|
|
),
|
|
)
|
|
|> E.O.bind(_, (r: RenderTypes.ShapeRenderer.Sampling.outputs) =>
|
|
r.shape
|
|
)
|
|
|> E.O.toResult("No response");
|
|
shape |> E.R.fmap(r => `Normalize(`RenderedDist(r)));
|
|
},
|
|
);
|
|
};
|
|
|
|
let operationToLeaf =
|
|
(
|
|
evaluationParams: evaluationParams,
|
|
algebraicOp: ExpressionTypes.algebraicOperation,
|
|
t1: t,
|
|
t2: t,
|
|
)
|
|
: result(node, string) =>
|
|
algebraicOp
|
|
|> tryAnalyticalSimplification(_, t1, t2)
|
|
|> E.R.bind(
|
|
_,
|
|
fun
|
|
| `SymbolicDist(d) as t => Ok(t)
|
|
| _ => combineAsShapes(evaluationParams, algebraicOp, t1, t2),
|
|
);
|
|
};
|
|
|
|
module VerticalScaling = {
|
|
let operationToLeaf =
|
|
(evaluationParams: evaluationParams, scaleOp, t, scaleBy) => {
|
|
// scaleBy has to be a single float, otherwise we'll return an error.
|
|
let fn = Operation.Scale.toFn(scaleOp);
|
|
let integralSumCacheFn = Operation.Scale.toIntegralSumCacheFn(scaleOp);
|
|
let integralCacheFn = Operation.Scale.toIntegralCacheFn(scaleOp);
|
|
let renderedShape = render(evaluationParams, t);
|
|
|
|
switch (renderedShape, scaleBy) {
|
|
| (Ok(`RenderedDist(rs)), `SymbolicDist(`Float(sm))) =>
|
|
Ok(
|
|
`RenderedDist(
|
|
Shape.T.mapY(
|
|
~integralSumCacheFn=integralSumCacheFn(sm),
|
|
~integralCacheFn=integralCacheFn(sm),
|
|
~fn=fn(sm),
|
|
rs,
|
|
),
|
|
),
|
|
)
|
|
| (Error(e1), _) => Error(e1)
|
|
| (_, _) => Error("Can only scale by float values.")
|
|
};
|
|
};
|
|
};
|
|
|
|
module PointwiseCombination = {
|
|
let pointwiseAdd = (evaluationParams: evaluationParams, t1: t, t2: t) => {
|
|
switch (render(evaluationParams, t1), render(evaluationParams, t2)) {
|
|
| (Ok(`RenderedDist(rs1)), Ok(`RenderedDist(rs2))) =>
|
|
Ok(
|
|
`RenderedDist(
|
|
Shape.combinePointwise(
|
|
~integralSumCachesFn=(a, b) => Some(a +. b),
|
|
~integralCachesFn=(a, b) => Some(Continuous.combinePointwise(~distributionType=`CDF, (+.), a, b)),
|
|
(+.),
|
|
rs1,
|
|
rs2,
|
|
),
|
|
),
|
|
)
|
|
| (Error(e1), _) => Error(e1)
|
|
| (_, Error(e2)) => Error(e2)
|
|
| _ => Error("Pointwise combination: rendering failed.")
|
|
};
|
|
};
|
|
|
|
let pointwiseMultiply = (evaluationParams: evaluationParams, t1: t, t2: t) => {
|
|
// TODO: construct a function that we can easily sample from, to construct
|
|
// a RenderedDist. Use the xMin and xMax of the rendered shapes to tell the sampling function where to look.
|
|
Error(
|
|
"Pointwise multiplication not yet supported.",
|
|
);
|
|
};
|
|
|
|
let operationToLeaf =
|
|
(evaluationParams: evaluationParams, pointwiseOp: pointwiseOperation, t1: t, t2: t) => {
|
|
switch (pointwiseOp) {
|
|
| `Add => pointwiseAdd(evaluationParams, t1, t2)
|
|
| `Multiply => pointwiseMultiply(evaluationParams, t1, t2)
|
|
};
|
|
};
|
|
};
|
|
|
|
module Truncate = {
|
|
let trySimplification = (leftCutoff, rightCutoff, t): simplificationResult => {
|
|
switch (leftCutoff, rightCutoff, t) {
|
|
| (None, None, t) => `Solution(t)
|
|
| (Some(lc), Some(rc), t) when lc > rc =>
|
|
`Error("Left truncation bound must be smaller than right bound.")
|
|
| (lc, rc, `SymbolicDist(`Uniform(u))) =>
|
|
`Solution(
|
|
`SymbolicDist(`Uniform(SymbolicDist.Uniform.truncate(lc, rc, u))),
|
|
)
|
|
| _ => `NoSolution
|
|
};
|
|
};
|
|
|
|
let truncateAsShape =
|
|
(evaluationParams: evaluationParams, leftCutoff, rightCutoff, t) => {
|
|
// TODO: use named args for xMin/xMax in renderToShape; if we're lucky we can at least get the tail
|
|
// of a distribution we otherwise wouldn't get at all
|
|
switch (render(evaluationParams, t)) {
|
|
| Ok(`RenderedDist(rs)) =>
|
|
Ok(`RenderedDist(Shape.T.truncate(leftCutoff, rightCutoff, rs)))
|
|
| Error(e) => Error(e)
|
|
| _ => Error("Could not truncate distribution.")
|
|
};
|
|
};
|
|
|
|
let operationToLeaf =
|
|
(
|
|
evaluationParams,
|
|
leftCutoff: option(float),
|
|
rightCutoff: option(float),
|
|
t: node,
|
|
)
|
|
: result(node, string) => {
|
|
t
|
|
|> trySimplification(leftCutoff, rightCutoff)
|
|
|> (
|
|
fun
|
|
| `Solution(t) => Ok(t)
|
|
| `Error(e) => Error(e)
|
|
| `NoSolution =>
|
|
truncateAsShape(evaluationParams, leftCutoff, rightCutoff, t)
|
|
);
|
|
};
|
|
};
|
|
|
|
module Normalize = {
|
|
let rec operationToLeaf = (evaluationParams, t: node): result(node, string) => {
|
|
switch (t) {
|
|
| `RenderedDist(s) => Ok(`RenderedDist(Shape.T.normalize(s)))
|
|
| `SymbolicDist(_) => Ok(t)
|
|
| _ => evaluateAndRetry(evaluationParams, operationToLeaf, t)
|
|
};
|
|
};
|
|
};
|
|
|
|
module FloatFromDist = {
|
|
let rec operationToLeaf =
|
|
(evaluationParams, distToFloatOp: distToFloatOperation, t: node)
|
|
: result(node, string) => {
|
|
switch (t) {
|
|
| `SymbolicDist(s) =>
|
|
SymbolicDist.T.operate(distToFloatOp, s)
|
|
|> E.R.bind(_, v => Ok(`SymbolicDist(`Float(v))))
|
|
| `RenderedDist(rs) =>
|
|
Shape.operate(distToFloatOp, rs)
|
|
|> (v => Ok(`SymbolicDist(`Float(v))))
|
|
| _ =>
|
|
t
|
|
|> evaluateAndRetry(evaluationParams, r =>
|
|
operationToLeaf(r, distToFloatOp)
|
|
)
|
|
};
|
|
};
|
|
};
|
|
|
|
module Render = {
|
|
let rec operationToLeaf =
|
|
(evaluationParams: evaluationParams, t: node): result(t, string) => {
|
|
switch (t) {
|
|
| `SymbolicDist(d) =>
|
|
Ok(
|
|
`RenderedDist(
|
|
SymbolicDist.T.toShape(evaluationParams.intendedShapeLength, d),
|
|
),
|
|
)
|
|
| `RenderedDist(_) as t => Ok(t) // already a rendered shape, we're done here
|
|
| _ => evaluateAndRetry(evaluationParams, operationToLeaf, t)
|
|
};
|
|
};
|
|
};
|
|
|
|
/* This function recursively goes through the nodes of the parse tree,
|
|
replacing each Operation node and its subtree with a Data node.
|
|
Whenever possible, the replacement produces a new Symbolic Data node,
|
|
but most often it will produce a RenderedDist.
|
|
This function is used mainly to turn a parse tree into a single RenderedDist
|
|
that can then be displayed to the user. */
|
|
let toLeaf =
|
|
(
|
|
evaluationParams: ExpressionTypes.ExpressionTree.evaluationParams,
|
|
node: t,
|
|
)
|
|
: result(t, string) => {
|
|
switch (node) {
|
|
// Leaf nodes just stay leaf nodes
|
|
| `SymbolicDist(_)
|
|
| `RenderedDist(_) => Ok(node)
|
|
// Operations nevaluationParamsd to be turned into leaves
|
|
| `AlgebraicCombination(algebraicOp, t1, t2) =>
|
|
AlgebraicCombination.operationToLeaf(
|
|
evaluationParams,
|
|
algebraicOp,
|
|
t1,
|
|
t2,
|
|
)
|
|
| `PointwiseCombination(pointwiseOp, t1, t2) =>
|
|
PointwiseCombination.operationToLeaf(
|
|
evaluationParams,
|
|
pointwiseOp,
|
|
t1,
|
|
t2,
|
|
)
|
|
| `VerticalScaling(scaleOp, t, scaleBy) =>
|
|
VerticalScaling.operationToLeaf(evaluationParams, scaleOp, t, scaleBy)
|
|
| `Truncate(leftCutoff, rightCutoff, t) =>
|
|
Truncate.operationToLeaf(evaluationParams, leftCutoff, rightCutoff, t)
|
|
| `FloatFromDist(distToFloatOp, t) =>
|
|
FloatFromDist.operationToLeaf(evaluationParams, distToFloatOp, t)
|
|
| `Normalize(t) => Normalize.operationToLeaf(evaluationParams, t)
|
|
| `Render(t) => Render.operationToLeaf(evaluationParams, t)
|
|
};
|
|
};
|