245 lines
9.2 KiB
Plaintext
245 lines
9.2 KiB
Plaintext
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 combinationByRendering = (evaluationParams, algebraicOp, t1: node, t2: node): result<
|
|
node,
|
|
string,
|
|
> =>
|
|
E.R.merge(
|
|
Render.ensureIsRenderedAndGetShape(evaluationParams, t1),
|
|
Render.ensureIsRenderedAndGetShape(evaluationParams, t2),
|
|
) |> E.R.fmap(((a, b)) => #RenderedDist(Shape.combineAlgebraically(algebraicOp, a, b)))
|
|
|
|
let nodeScore: node => int = x =>
|
|
switch x {
|
|
| #SymbolicDist(#Float(_)) => 1
|
|
| #SymbolicDist(_) => 1000
|
|
| #RenderedDist(Discrete(m)) => m.xyShape |> XYShape.T.length
|
|
| #RenderedDist(Mixed(_)) => 1000
|
|
| #RenderedDist(Continuous(_)) => 1000
|
|
| _ => 1000
|
|
}
|
|
|
|
let choose = (t1: node, t2: node) =>
|
|
nodeScore(t1) * nodeScore(t2) > 10000 ? #Sampling : #Analytical
|
|
|
|
let combine = (evaluationParams, algebraicOp, t1: node, t2: node): result<node, string> =>
|
|
E.R.merge(
|
|
PTypes.SamplingDistribution.renderIfIsNotSamplingDistribution(evaluationParams, t1),
|
|
PTypes.SamplingDistribution.renderIfIsNotSamplingDistribution(evaluationParams, t2),
|
|
) |> E.R.bind(_, ((a, b)) =>
|
|
switch choose(a, b) {
|
|
| #Sampling =>
|
|
PTypes.SamplingDistribution.combineShapesUsingSampling(evaluationParams, algebraicOp, a, b)
|
|
| #Analytical => combinationByRendering(evaluationParams, algebraicOp, a, b)
|
|
}
|
|
)
|
|
|
|
let operationToLeaf = (
|
|
evaluationParams: evaluationParams,
|
|
algebraicOp: ExpressionTypes.algebraicOperation,
|
|
t1: t,
|
|
t2: t,
|
|
): result<node, string> =>
|
|
algebraicOp
|
|
|> tryAnalyticalSimplification(_, t1, t2)
|
|
|> E.R.bind(_, x =>
|
|
switch x {
|
|
| #SymbolicDist(_) as t => Ok(t)
|
|
| _ => combine(evaluationParams, algebraicOp, t1, t2)
|
|
}
|
|
)
|
|
}
|
|
|
|
module PointwiseCombination = {
|
|
let pointwiseAdd = (evaluationParams: evaluationParams, t1: t, t2: t) =>
|
|
switch (Render.render(evaluationParams, t1), Render.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 pointwiseCombine = (fn, evaluationParams: evaluationParams, t1: t, t2: t) =>
|
|
switch // 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.
|
|
// TODO: This should work for symbolic distributions too!
|
|
(Render.render(evaluationParams, t1), Render.render(evaluationParams, t2)) {
|
|
| (Ok(#RenderedDist(rs1)), Ok(#RenderedDist(rs2))) =>
|
|
Ok(#RenderedDist(Shape.combinePointwise(fn, rs1, rs2)))
|
|
| (Error(e1), _) => Error(e1)
|
|
| (_, Error(e2)) => Error(e2)
|
|
| _ => Error("Pointwise combination: rendering failed.")
|
|
}
|
|
|
|
let operationToLeaf = (
|
|
evaluationParams: evaluationParams,
|
|
pointwiseOp: pointwiseOperation,
|
|
t1: t,
|
|
t2: t,
|
|
) =>
|
|
switch pointwiseOp {
|
|
| #Add => pointwiseAdd(evaluationParams, t1, t2)
|
|
| #Multiply => pointwiseCombine(\"*.", evaluationParams, t1, t2)
|
|
| #Exponentiate => pointwiseCombine(\"**", evaluationParams, t1, t2)
|
|
}
|
|
}
|
|
|
|
module Truncate = {
|
|
let trySimplification = (leftCutoff, rightCutoff, t): simplificationResult =>
|
|
switch (leftCutoff, rightCutoff, t) {
|
|
| (None, None, t) => #Solution(t)
|
|
| (Some(lc), Some(rc), _) if lc > rc =>
|
|
#Error("Left truncation bound must be smaller than right truncation bound.")
|
|
| (lc, rc, #SymbolicDist(#Uniform(u))) =>
|
|
#Solution(#SymbolicDist(#Uniform(SymbolicDist.Uniform.truncate(lc, rc, u))))
|
|
| _ => #NoSolution
|
|
}
|
|
|
|
let truncateAsShape = (evaluationParams: evaluationParams, leftCutoff, rightCutoff, t) =>
|
|
switch // 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
|
|
Render.ensureIsRendered(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)
|
|
|> (
|
|
x =>
|
|
switch x {
|
|
| #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 FunctionCall = {
|
|
let _runHardcodedFunction = (name, evaluationParams, args) =>
|
|
TypeSystem.Function.Ts.findByNameAndRun(HardcodedFunctions.all, name, evaluationParams, args)
|
|
|
|
let _runLocalFunction = (name, evaluationParams: evaluationParams, args) =>
|
|
Environment.getFunction(evaluationParams.environment, name) |> E.R.bind(_, ((argNames, fn)) =>
|
|
PTypes.Function.run(evaluationParams, args, (argNames, fn))
|
|
)
|
|
|
|
let _runWithEvaluatedInputs = (
|
|
evaluationParams: ExpressionTypes.ExpressionTree.evaluationParams,
|
|
name,
|
|
args: array<ExpressionTypes.ExpressionTree.node>,
|
|
) =>
|
|
_runHardcodedFunction(name, evaluationParams, args) |> E.O.default(
|
|
_runLocalFunction(name, evaluationParams, args),
|
|
)
|
|
|
|
// TODO: This forces things to be floats
|
|
let run = (evaluationParams, name, args) =>
|
|
args
|
|
|> E.A.fmap(a => evaluationParams.evaluateNode(evaluationParams, a))
|
|
|> E.A.R.firstErrorOrOpen
|
|
|> E.R.bind(_, _runWithEvaluatedInputs(evaluationParams, name))
|
|
}
|
|
|
|
module Render = {
|
|
let rec operationToLeaf = (evaluationParams: evaluationParams, t: node): result<t, string> =>
|
|
switch t {
|
|
| #Function(_) => Error("Cannot render a function")
|
|
| #SymbolicDist(d) =>
|
|
Ok(#RenderedDist(SymbolicDist.T.toShape(evaluationParams.samplingInputs.shapeLength, 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 rec toLeaf = (
|
|
evaluationParams: ExpressionTypes.ExpressionTree.evaluationParams,
|
|
node: t,
|
|
): result<t, string> =>
|
|
switch node {
|
|
// Leaf nodes just stay leaf nodes
|
|
| #SymbolicDist(_)
|
|
| #Function(_)
|
|
| #RenderedDist(_) =>
|
|
Ok(node)
|
|
| #Array(args) =>
|
|
args |> E.A.fmap(toLeaf(evaluationParams)) |> E.A.R.firstErrorOrOpen |> E.R.fmap(r => #Array(r))
|
|
// 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)
|
|
| #Truncate(leftCutoff, rightCutoff, t) =>
|
|
Truncate.operationToLeaf(evaluationParams, leftCutoff, rightCutoff, t)
|
|
| #Normalize(t) => Normalize.operationToLeaf(evaluationParams, t)
|
|
| #Render(t) => Render.operationToLeaf(evaluationParams, t)
|
|
| #Hash(t) =>
|
|
t
|
|
|> E.A.fmap(((name: string, node: node)) =>
|
|
toLeaf(evaluationParams, node) |> E.R.fmap(r => (name, r))
|
|
)
|
|
|> E.A.R.firstErrorOrOpen
|
|
|> E.R.fmap(r => #Hash(r))
|
|
| #Symbol(r) =>
|
|
ExpressionTypes.ExpressionTree.Environment.get(evaluationParams.environment, r)
|
|
|> E.O.toResult("Undeclared variable " ++ r)
|
|
|> E.R.bind(_, toLeaf(evaluationParams))
|
|
| #FunctionCall(name, args) =>
|
|
FunctionCall.run(evaluationParams, name, args) |> E.R.bind(_, toLeaf(evaluationParams))
|
|
}
|