Added pointwiseCombinationFloat to sci.res

packages/squiggle-lang/src/rescript/sci.res

@@ -3,6 +3,7 @@ type error =
   | InputsNeedPointSetConversion
   | NotYetImplemented
   | ImpossiblePath
+  | DistributionVerticalShiftIsInvalid
   | Other(string)
 
 type genericDist = [
@@ -11,141 +12,102 @@ type genericDist = [
   | #Symbolic(SymbolicDistTypes.symbolicDist)
 ]
 
-type outputType = [
+module OperationType = {
-  | #Dist(genericDist)
+  type direction = [
-  | #Error(error)
-  | #Float(float)
-]
-
-let fromResult = (r: result<outputType, error>): outputType =>
-  switch r {
-  | Ok(o) => o
-  | Error(e) => #Error(e)
-  }
-
-type direction = [
     | #Algebraic
     | #Pointwise
-]
+  ]
 
-type combination = [
+  type combination = [
     | #Add
     | #Multiply
     | #Subtract
     | #Divide
     | #Exponentiate
-]
+    | #Log
+  ]
 
-let combinationToFn = (combination: combination) =>
+  let combinationToFn = (combination: combination) =>
     switch combination {
     | #Add => \"+."
     | #Multiply => \"*."
     | #Subtract => \"-."
     | #Exponentiate => \"**"
     | #Divide => \"/."
+    | #Log => (a, b) => log(a) /. log(b)
     }
 
-type toFloat = [
+  type toFloat = [
     | #Cdf(float)
     | #Inv(float)
     | #Mean
     | #Pdf(float)
     | #Sample
-]
+  ]
 
-type toDist = [
+  type toDist = [
     | #normalize
     | #toPointSet
     | #toSampleSet(int)
-]
+  ]
 
-type toFloatArray = [
+  type toFloatArray = [
     | #Sample(int)
-]
+  ]
 
-type operation = [
+  type scale = [
+    | #Multiply
+    | #Exponentiate
+    | #Log
+  ]
+
+  type t = [
     | #toFloat(toFloat)
     | #toDist(toDist)
-  | #toDistCombination(direction, combination, genericDist)
+    | #toDistCombination(direction, combination, [#Dist(genericDist) | #Float(float)])
-]
+  ]
-
-type params = {
-  sampleCount: int,
-  xyPointLength: int,
 }
 
-let genericParams = {
+type operation = OperationType.t
-  sampleCount: 1000,
-  xyPointLength: 1000,
-}
 
-type wrapped = (genericDist, params)
+module T = {
-type wrappedOutput = (outputType, params)
+  type t = genericDist
-
+  type toPointSetFn = genericDist => result<PointSetTypes.pointSetDist, error>
-let wrapWithParams = (g: genericDist, f: params): wrapped => (g, f)
+  let sampleN = (n, t: t) => {
-
+    switch t {
-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)
-      }
-    | _ => 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 = (
+  let toFloat = (toPointSet: toPointSetFn, fnName, t: genericDist) => {
-  toPointSet: genericDist => result<PointSetTypes.pointSetDist, error>,
+    switch t {
-  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) {
+      switch toPointSet(t) {
       | Ok(r) => Ok(PointSetDist.operate(fnName, r))
       | Error(r) => Error(r)
       }
     }
-}
+  }
 
-let distToPointSet = (sampleCount, dist: genericDist) => {
+  //TODO: Refactor this bit.
-  switch dist {
+  let defaultSamplingInputs: SamplingInputs.samplingInputs = {
+    sampleCount: 10000,
+    outputXYPoints: 10000,
+    pointSetDistLength: 1000,
+    kernelWidth: None,
+  }
+
+  let toPointSet = (xyPointLength, t: t) => {
+    switch t {
     | #PointSet(pointSet) => Ok(pointSet)
-  | #Symbolic(r) => Ok(SymbolicDist.T.toPointSetDist(sampleCount, r))
+    | #Symbolic(r) => Ok(SymbolicDist.T.toPointSetDist(xyPointLength, r))
     | #SampleSet(r) => {
         let response = SampleSet.toPointSetDist(
           ~samples=r,
@@ -158,62 +120,112 @@ let distToPointSet = (sampleCount, dist: genericDist) => {
         }
       }
     }
+  }
+
+  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,
+        )
+      })
+    }
+  }
 }
 
-let rec applyFnInternal = (wrapped: wrapped, fnName: operation): wrappedOutput => {
+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 reCallUnwrapped = (~value=value, ~extra=extra, ~fnName=fnName, ()) => {
-    let (value, _) = applyFnInternal((value, extra), fnName)
-    value
-  }
     let toPointSet = r => {
-    switch reCallUnwrapped(~value=r, ~fnName=#toDist(#toPointSet), ()) {
+      switch reCall(~value=r, ~fnName=#toDist(#toPointSet), ()) {
       | #Dist(#PointSet(p)) => Ok(p)
       | #Error(r) => Error(r)
       | _ => Error(Other("Impossible error"))
       }
     }
-  let toPointSetAndReCall = v =>
+    let toPointSetAndReCall = v => toPointSet(v) |> E.R.fmap(r => reCall(~value=#PointSet(r), ()))
-    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
+      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 |> distToPointSet(sampleCount) |> E.R.fmap(r => #Dist(#PointSet(r))) |> fromResult
+      value |> T.toPointSet(xyPointLength) |> E.R.fmap(r => #Dist(#PointSet(r))) |> fromResult
     | (#toDist(#toSampleSet(n)), _) =>
-    value |> sampleN(n) |> E.R.fmap(r => #Dist(#SampleSet(r))) |> fromResult
+      value |> T.sampleN(n) |> E.R.fmap(r => #Dist(#SampleSet(r))) |> fromResult
-  | (#toDistCombination(#Algebraic, operation, p2), p1) => {
+    | (#toDistCombination(#Algebraic, _, #Float(_)), _) => #Error(NotYetImplemented)
-      // TODO: This could be more complex, to get possible simplification and similar.
+    | (#toDistCombination(#Algebraic, operation, #Dist(p2)), p1) =>
-      let dist1 = sampleN(sampleCount, p1)
+      T.algebraicCombination(operation, sampleCount, p1, p2)
-      let dist2 = sampleN(sampleCount, p2)
+      |> E.R.fmap(r => #Dist(r))
-      let samples = E.R.merge(dist1, dist2) |> E.R.fmap(((d1, d2)) => {
+      |> fromResult
-        Belt.Array.zip(d1, d2) |> E.A.fmap(((a, b)) => Operation.Algebraic.toFn(operation, a, b))
+    | (#toDistCombination(#Pointwise, operation, #Dist(p2)), p1) =>
-      })
+      T.pointwiseCombination(toPointSet, operation, p1, p2) |> E.R.fmap(r => #Dist(r)) |> fromResult
-      switch samples {
+    | (#toDistCombination(#Pointwise, operation, #Float(f)), _) =>
-      | Ok(r) => #Dist(#SampleSet(r))
+      T.pointwiseCombinationFloat(toPointSet, operation, value, f)
-      | Error(e) => #Error(e)
+      |> E.R.fmap(r => #Dist(#PointSet(r)))
+      |> fromResult
     }
+    newVal
   }
-  | (#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"))
-    }
-  | _ => #Error(Other("No Match or not supported"))
-  }
-  (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))

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
     }