From d490af38f027151a2fa6d666cab24e03f043e292 Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Sat, 26 Mar 2022 13:25:47 -0400
Subject: [PATCH] Added pointwiseCombinationFloat to sci.res

---
 packages/squiggle-lang/src/rescript/sci.res   | 390 +++++++++---------
 .../src/rescript/utility/Operation.res        |   9 +-
 2 files changed, 211 insertions(+), 188 deletions(-)

diff --git a/packages/squiggle-lang/src/rescript/sci.res b/packages/squiggle-lang/src/rescript/sci.res
index 15a57f44..66a5605a 100644
--- a/packages/squiggle-lang/src/rescript/sci.res
+++ b/packages/squiggle-lang/src/rescript/sci.res
@@ -3,6 +3,7 @@ type error =
   | InputsNeedPointSetConversion
   | NotYetImplemented
   | ImpossiblePath
+  | DistributionVerticalShiftIsInvalid
   | Other(string)
 
 type genericDist = [
@@ -11,209 +12,220 @@ type genericDist = [
   | #Symbolic(SymbolicDistTypes.symbolicDist)
 ]
 
-type outputType = [
-  | #Dist(genericDist)
-  | #Error(error)
-  | #Float(float)
-]
+module OperationType = {
+  type direction = [
+    | #Algebraic
+    | #Pointwise
+  ]
 
-let fromResult = (r: result<outputType, error>): outputType =>
-  switch r {
-  | Ok(o) => o
-  | Error(e) => #Error(e)
+  type combination = [
+    | #Add
+    | #Multiply
+    | #Subtract
+    | #Divide
+    | #Exponentiate
+    | #Log
+  ]
+
+  let combinationToFn = (combination: combination) =>
+    switch combination {
+    | #Add => \"+."
+    | #Multiply => \"*."
+    | #Subtract => \"-."
+    | #Exponentiate => \"**"
+    | #Divide => \"/."
+    | #Log => (a, b) => log(a) /. log(b)
+    }
+
+  type toFloat = [
+    | #Cdf(float)
+    | #Inv(float)
+    | #Mean
+    | #Pdf(float)
+    | #Sample
+  ]
+
+  type toDist = [
+    | #normalize
+    | #toPointSet
+    | #toSampleSet(int)
+  ]
+
+  type toFloatArray = [
+    | #Sample(int)
+  ]
+
+  type scale = [
+    | #Multiply
+    | #Exponentiate
+    | #Log
+  ]
+
+  type t = [
+    | #toFloat(toFloat)
+    | #toDist(toDist)
+    | #toDistCombination(direction, combination, [#Dist(genericDist) | #Float(float)])
+  ]
+}
+
+type operation = OperationType.t
+
+module T = {
+  type t = genericDist
+  type toPointSetFn = genericDist => result<PointSetTypes.pointSetDist, error>
+  let sampleN = (n, t: t) => {
+    switch t {
+    | #PointSet(r) => Ok(PointSetDist.sampleNRendered(n, r))
+    | #Symbolic(r) => Ok(SymbolicDist.T.sampleN(n, r))
+    | #SampleSet(_) => Error(NotYetImplemented)
+    }
   }
 
-type direction = [
-  | #Algebraic
-  | #Pointwise
-]
-
-type combination = [
-  | #Add
-  | #Multiply
-  | #Subtract
-  | #Divide
-  | #Exponentiate
-]
-
-let combinationToFn = (combination: combination) =>
-  switch combination {
-  | #Add => \"+."
-  | #Multiply => \"*."
-  | #Subtract => \"-."
-  | #Exponentiate => \"**"
-  | #Divide => \"/."
-  }
-
-type toFloat = [
-  | #Cdf(float)
-  | #Inv(float)
-  | #Mean
-  | #Pdf(float)
-  | #Sample
-]
-
-type toDist = [
-  | #normalize
-  | #toPointSet
-  | #toSampleSet(int)
-]
-
-type toFloatArray = [
-  | #Sample(int)
-]
-
-type operation = [
-  | #toFloat(toFloat)
-  | #toDist(toDist)
-  | #toDistCombination(direction, combination, genericDist)
-]
-
-type params = {
-  sampleCount: int,
-  xyPointLength: int,
-}
-
-let genericParams = {
-  sampleCount: 1000,
-  xyPointLength: 1000,
-}
-
-type wrapped = (genericDist, params)
-type wrappedOutput = (outputType, params)
-
-let wrapWithParams = (g: genericDist, f: params): wrapped => (g, f)
-
-let exampleDist: genericDist = #PointSet(
-  Discrete(Discrete.make(~integralSumCache=Some(1.0), {xs: [3.0], ys: [1.0]})),
-)
-
-let defaultSamplingInputs: SamplingInputs.samplingInputs = {
-  sampleCount: 10000,
-  outputXYPoints: 10000,
-  pointSetDistLength: 1000,
-  kernelWidth: None,
-}
-
-/* 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: genericDist, t2: genericDist) =>
-    switch (operation, t1, t2) {
-    | (operation, #Symbolic(d1), #Symbolic(d2)) =>
-      switch SymbolicDist.T.tryAnalyticalSimplification(d1, d2, operation) {
-      | #AnalyticalSolution(symbolicDist) => Ok(#Symbolic(symbolicDist))
-      | #Error(er) => Error(er)
-      | #NoSolution => Ok(#NoSolution)
+  let toFloat = (toPointSet: toPointSetFn, fnName, t: genericDist) => {
+    switch t {
+    | #Symbolic(r) if Belt.Result.isOk(SymbolicDist.T.operate(fnName, r)) =>
+      switch SymbolicDist.T.operate(fnName, r) {
+      | Ok(float) => Ok(float)
+      | Error(_) => Error(ImpossiblePath)
       }
-    | _ => Ok(#NoSolution)
-    }
-}
-
-// let toSampleSet = (r)
-
-let sampleN = (n, genericDist) => {
-  switch genericDist {
-  | #PointSet(r) => Ok(PointSetDist.sampleNRendered(n, r))
-  | #Symbolic(r) => Ok(SymbolicDist.T.sampleN(n, r))
-  | #SampleSet(_) => Error(NotYetImplemented)
-  }
-}
-
-let toFloat = (
-  toPointSet: genericDist => result<PointSetTypes.pointSetDist, error>,
-  fnName,
-  value,
-) => {
-  switch value {
-  | #Symbolic(r) if Belt.Result.isOk(SymbolicDist.T.operate(fnName, r)) =>
-    switch SymbolicDist.T.operate(fnName, r) {
-    | Ok(float) => Ok(float)
-    | Error(_) => Error(ImpossiblePath)
-    }
-  | #PointSet(r) => Ok(PointSetDist.operate(fnName, r))
-  | _ =>
-    switch toPointSet(value) {
-    | Ok(r) => Ok(PointSetDist.operate(fnName, r))
-    | Error(r) => Error(r)
-    }
-  }
-}
-
-let distToPointSet = (sampleCount, dist: genericDist) => {
-  switch dist {
-  | #PointSet(pointSet) => Ok(pointSet)
-  | #Symbolic(r) => Ok(SymbolicDist.T.toPointSetDist(sampleCount, r))
-  | #SampleSet(r) => {
-      let response = SampleSet.toPointSetDist(
-        ~samples=r,
-        ~samplingInputs=defaultSamplingInputs,
-        (),
-      ).pointSetDist
-      switch response {
-      | Some(r) => Ok(r)
-      | None => Error(Other("Converting sampleSet to pointSet failed"))
+    | _ =>
+      switch toPointSet(t) {
+      | Ok(r) => Ok(PointSetDist.operate(fnName, r))
+      | Error(r) => Error(r)
       }
     }
   }
-}
 
-let rec applyFnInternal = (wrapped: wrapped, fnName: operation): wrappedOutput => {
-  let (value, {sampleCount, xyPointLength} as extra) = wrapped
-  let reCall = (~value=value, ~extra=extra, ~fnName=fnName, ()) => {
-    applyFnInternal((value, extra), fnName)
+  //TODO: Refactor this bit.
+  let defaultSamplingInputs: SamplingInputs.samplingInputs = {
+    sampleCount: 10000,
+    outputXYPoints: 10000,
+    pointSetDistLength: 1000,
+    kernelWidth: None,
   }
-  let reCallUnwrapped = (~value=value, ~extra=extra, ~fnName=fnName, ()) => {
-    let (value, _) = applyFnInternal((value, extra), fnName)
-    value
-  }
-  let toPointSet = r => {
-    switch reCallUnwrapped(~value=r, ~fnName=#toDist(#toPointSet), ()) {
-    | #Dist(#PointSet(p)) => Ok(p)
-    | #Error(r) => Error(r)
-    | _ => Error(Other("Impossible error"))
+
+  let toPointSet = (xyPointLength, t: t) => {
+    switch t {
+    | #PointSet(pointSet) => Ok(pointSet)
+    | #Symbolic(r) => Ok(SymbolicDist.T.toPointSetDist(xyPointLength, r))
+    | #SampleSet(r) => {
+        let response = SampleSet.toPointSetDist(
+          ~samples=r,
+          ~samplingInputs=defaultSamplingInputs,
+          (),
+        ).pointSetDist
+        switch response {
+        | Some(r) => Ok(r)
+        | None => Error(Other("Converting sampleSet to pointSet failed"))
+        }
+      }
     }
   }
-  let toPointSetAndReCall = v =>
-    toPointSet(v) |> E.R.fmap(r => reCallUnwrapped(~value=#PointSet(r), ()))
-  let newVal: outputType = switch (fnName, value) {
-  // | (#toFloat(n), v) => toFloat(toPointSet, v, n)
-  | (#toFloat(fnName), _) =>
-    toFloat(toPointSet, fnName, value) |> E.R.fmap(r => #Float(r)) |> fromResult
-  | (#toDist(#normalize), #PointSet(r)) => #Dist(#PointSet(PointSetDist.T.normalize(r)))
-  | (#toDist(#normalize), #Symbolic(_)) => #Dist(value)
-  | (#toDist(#normalize), #SampleSet(_)) => #Dist(value)
-  | (#toDist(#toPointSet), _) =>
-    value |> distToPointSet(sampleCount) |> E.R.fmap(r => #Dist(#PointSet(r))) |> fromResult
-  | (#toDist(#toSampleSet(n)), _) =>
-    value |> sampleN(n) |> E.R.fmap(r => #Dist(#SampleSet(r))) |> fromResult
-  | (#toDistCombination(#Algebraic, operation, p2), p1) => {
-      // TODO: This could be more complex, to get possible simplification and similar.
-      let dist1 = sampleN(sampleCount, p1)
-      let dist2 = sampleN(sampleCount, p2)
-      let samples = E.R.merge(dist1, dist2) |> E.R.fmap(((d1, d2)) => {
-        Belt.Array.zip(d1, d2) |> E.A.fmap(((a, b)) => Operation.Algebraic.toFn(operation, a, b))
+
+  let algebraicCombination = (operation, sampleCount, dist1: t, dist2: t) => {
+    let dist1 = sampleN(sampleCount, dist1)
+    let dist2 = sampleN(sampleCount, dist2)
+    let samples = E.R.merge(dist1, dist2) |> E.R.fmap(((d1, d2)) => {
+      Belt.Array.zip(d1, d2) |> E.A.fmap(((a, b)) => Operation.Algebraic.toFn(operation, a, b))
+    })
+    samples |> E.R.fmap(r => #SampleSet(r))
+  }
+
+  let pointwiseCombination = (toPointSet: toPointSetFn, operation, t1: t, t2: t) => {
+    E.R.merge(toPointSet(t1), toPointSet(t2))
+    |> E.R.fmap(((t1, t2)) =>
+      PointSetDist.combinePointwise(OperationType.combinationToFn(operation), t1, t2)
+    )
+    |> E.R.fmap(r => #PointSet(r))
+  }
+
+  let pointwiseCombinationFloat = (
+    toPointSet: toPointSetFn,
+    operation: OperationType.combination,
+    t: t,
+    f: float,
+  ) => {
+    switch operation {
+    | #Add | #Subtract => Error(DistributionVerticalShiftIsInvalid)
+    | (#Multiply | #Divide | #Exponentiate | #Log) as operation =>
+      toPointSet(t) |> E.R.fmap(t => {
+        let fn = (secondary, main) => Operation.Scale.toFn(operation, main, secondary)
+        let integralSumCacheFn = Operation.Scale.toIntegralSumCacheFn(operation)
+        let integralCacheFn = Operation.Scale.toIntegralCacheFn(operation)
+        PointSetDist.T.mapY(
+          ~integralSumCacheFn=integralSumCacheFn(f),
+          ~integralCacheFn=integralCacheFn(f),
+          ~fn=fn(f),
+          t,
+        )
       })
-      switch samples {
-      | Ok(r) => #Dist(#SampleSet(r))
-      | Error(e) => #Error(e)
+    }
+  }
+}
+
+module OmniRunner = {
+  type params = {
+    sampleCount: int,
+    xyPointLength: int,
+  }
+
+  let genericParams = {
+    sampleCount: 1000,
+    xyPointLength: 1000,
+  }
+  type wrapped = (genericDist, params)
+
+  let wrapWithParams = (g: genericDist, f: params): wrapped => (g, f)
+  type outputType = [
+    | #Dist(genericDist)
+    | #Error(error)
+    | #Float(float)
+  ]
+
+  let fromResult = (r: result<outputType, error>): outputType =>
+    switch r {
+    | Ok(o) => o
+    | Error(e) => #Error(e)
+    }
+
+  let rec applyFnInternal = (wrapped: wrapped, fnName: operation): outputType => {
+    let (value, {sampleCount, xyPointLength} as extra) = wrapped
+    let reCall = (~value=value, ~extra=extra, ~fnName=fnName, ()) => {
+      applyFnInternal((value, extra), fnName)
+    }
+    let toPointSet = r => {
+      switch reCall(~value=r, ~fnName=#toDist(#toPointSet), ()) {
+      | #Dist(#PointSet(p)) => Ok(p)
+      | #Error(r) => Error(r)
+      | _ => Error(Other("Impossible error"))
       }
     }
-  | (#toDistCombination(#Pointwise, operation, p2), p1) =>
-    switch (
-      toPointSet(p1),
-      toPointSet(p2)
-    ) {
-    | (Ok(p1), Ok(p2)) =>
-      // TODO: If the dist is symbolic, then it doesn't need to be converted into a pointSet
-      #Dist(#PointSet(PointSetDist.combinePointwise(combinationToFn(operation), p1, p2)))
-    | (_, _) => #Error(Other("No Match or not supported"))
+    let toPointSetAndReCall = v => toPointSet(v) |> E.R.fmap(r => reCall(~value=#PointSet(r), ()))
+    let newVal: outputType = switch (fnName, value) {
+    // | (#toFloat(n), v) => toFloat(toPointSet, v, n)
+    | (#toFloat(fnName), _) =>
+      T.toFloat(toPointSet, fnName, value) |> E.R.fmap(r => #Float(r)) |> fromResult
+    | (#toDist(#normalize), #PointSet(r)) => #Dist(#PointSet(PointSetDist.T.normalize(r)))
+    | (#toDist(#normalize), #Symbolic(_)) => #Dist(value)
+    | (#toDist(#normalize), #SampleSet(_)) => #Dist(value)
+    | (#toDist(#toPointSet), _) =>
+      value |> T.toPointSet(xyPointLength) |> E.R.fmap(r => #Dist(#PointSet(r))) |> fromResult
+    | (#toDist(#toSampleSet(n)), _) =>
+      value |> T.sampleN(n) |> E.R.fmap(r => #Dist(#SampleSet(r))) |> fromResult
+    | (#toDistCombination(#Algebraic, _, #Float(_)), _) => #Error(NotYetImplemented)
+    | (#toDistCombination(#Algebraic, operation, #Dist(p2)), p1) =>
+      T.algebraicCombination(operation, sampleCount, p1, p2)
+      |> E.R.fmap(r => #Dist(r))
+      |> fromResult
+    | (#toDistCombination(#Pointwise, operation, #Dist(p2)), p1) =>
+      T.pointwiseCombination(toPointSet, operation, p1, p2) |> E.R.fmap(r => #Dist(r)) |> fromResult
+    | (#toDistCombination(#Pointwise, operation, #Float(f)), _) =>
+      T.pointwiseCombinationFloat(toPointSet, operation, value, f)
+      |> E.R.fmap(r => #Dist(#PointSet(r)))
+      |> fromResult
     }
-  | _ => #Error(Other("No Match or not supported"))
+    newVal
   }
-  (newVal, {sampleCount: sampleCount, xyPointLength: xyPointLength})
 }
 
 // let applyFn = (wrapped, fnName): wrapped => {
@@ -237,4 +249,8 @@ let rec applyFnInternal = (wrapped: wrapped, fnName: operation): wrappedOutput =
 //   }
 // }
 
+// let exampleDist: genericDist = #PointSet(
+//   Discrete(Discrete.make(~integralSumCache=Some(1.0), {xs: [3.0], ys: [1.0]})),
+// )
+
 // let foo = exampleDist->wrapWithParams(genericParams)->applyFn(#toDist(#normalize))
diff --git a/packages/squiggle-lang/src/rescript/utility/Operation.res b/packages/squiggle-lang/src/rescript/utility/Operation.res
index 4eb2c3cd..540bd08c 100644
--- a/packages/squiggle-lang/src/rescript/utility/Operation.res
+++ b/packages/squiggle-lang/src/rescript/utility/Operation.res
@@ -7,10 +7,11 @@ type algebraicOperation = [
   | #Subtract
   | #Divide
   | #Exponentiate
+  | #Log
 ]
 @genType
 type pointwiseOperation = [#Add | #Multiply | #Exponentiate]
-type scaleOperation = [#Multiply | #Exponentiate | #Log]
+type scaleOperation = [#Multiply | #Exponentiate | #Log | #Divide]
 type distToFloatOperation = [
   | #Pdf(float)
   | #Cdf(float)
@@ -28,6 +29,7 @@ module Algebraic = {
     | #Multiply => \"*."
     | #Exponentiate => \"**"
     | #Divide => \"/."
+    | #Log => (a, b) => log(a) /. log(b)
     }
 
   let applyFn = (t, f1, f2) =>
@@ -43,6 +45,7 @@ module Algebraic = {
     | #Multiply => "*"
     | #Exponentiate => "**"
     | #Divide => "/"
+    | #Log => "log"
     }
 
   let format = (a, b, c) => b ++ (" " ++ (toString(a) ++ (" " ++ c)))
@@ -79,6 +82,7 @@ module Scale = {
   let toFn = x =>
     switch x {
     | #Multiply => \"*."
+    | #Divide => \"/."
     | #Exponentiate => \"**"
     | #Log => (a, b) => log(a) /. log(b)
     }
@@ -86,6 +90,7 @@ module Scale = {
   let format = (operation: t, value, scaleBy) =>
     switch operation {
     | #Multiply => j`verticalMultiply($value, $scaleBy) `
+    | #Divide => j`verticalDivide($value, $scaleBy) `
     | #Exponentiate => j`verticalExponentiate($value, $scaleBy) `
     | #Log => j`verticalLog($value, $scaleBy) `
     }
@@ -93,6 +98,7 @@ module Scale = {
   let toIntegralSumCacheFn = x =>
     switch x {
     | #Multiply => (a, b) => Some(a *. b)
+    | #Divide => (a, b) => Some(a /. b)
     | #Exponentiate => (_, _) => None
     | #Log => (_, _) => None
     }
@@ -100,6 +106,7 @@ module Scale = {
   let toIntegralCacheFn = x =>
     switch x {
     | #Multiply => (_, _) => None // TODO: this could probably just be multiplied out (using Continuous.scaleBy)
+    | #Divide => (_, _) => None
     | #Exponentiate => (_, _) => None
     | #Log => (_, _) => None
     }