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,18 +12,7 @@ type genericDist = [
   | #Symbolic(SymbolicDistTypes.symbolicDist)
 ]
 
-type outputType = [
+module OperationType = {
-  | #Dist(genericDist)
-  | #Error(error)
-  | #Float(float)
-]
-
-let fromResult = (r: result<outputType, error>): outputType =>
-  switch r {
-  | Ok(o) => o
-  | Error(e) => #Error(e)
-  }
-
   type direction = [
     | #Algebraic
     | #Pointwise
@@ -34,6 +24,7 @@ type combination = [
     | #Subtract
     | #Divide
     | #Exponentiate
+    | #Log
   ]
 
   let combinationToFn = (combination: combination) =>
@@ -43,6 +34,7 @@ let combinationToFn = (combination: combination) =>
     | #Subtract => \"-."
     | #Exponentiate => \"**"
     | #Divide => \"/."
+    | #Log => (a, b) => log(a) /. log(b)
     }
 
   type toFloat = [
@@ -63,31 +55,48 @@ type toFloatArray = [
     | #Sample(int)
   ]
 
-type operation = [
+  type scale = [
-  | #toFloat(toFloat)
+    | #Multiply
-  | #toDist(toDist)
+    | #Exponentiate
-  | #toDistCombination(direction, combination, genericDist)
+    | #Log
   ]
 
-type params = {
+  type t = [
-  sampleCount: int,
+    | #toFloat(toFloat)
-  xyPointLength: int,
+    | #toDist(toDist)
+    | #toDistCombination(direction, combination, [#Dist(genericDist) | #Float(float)])
+  ]
 }
 
-let genericParams = {
+type operation = OperationType.t
-  sampleCount: 1000,
+
-  xyPointLength: 1000,
+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 wrapped = (genericDist, params)
+  let toFloat = (toPointSet: toPointSetFn, fnName, t: genericDist) => {
-type wrappedOutput = (outputType, params)
+    switch t {
-
+    | #Symbolic(r) if Belt.Result.isOk(SymbolicDist.T.operate(fnName, r)) =>
-let wrapWithParams = (g: genericDist, f: params): wrapped => (g, f)
+      switch SymbolicDist.T.operate(fnName, r) {
-
+      | Ok(float) => Ok(float)
-let exampleDist: genericDist = #PointSet(
+      | Error(_) => Error(ImpossiblePath)
-  Discrete(Discrete.make(~integralSumCache=Some(1.0), {xs: [3.0], ys: [1.0]})),
+      }
-)
+    | _ =>
+      switch toPointSet(t) {
+      | Ok(r) => Ok(PointSetDist.operate(fnName, r))
+      | Error(r) => Error(r)
+      }
+    }
+  }
 
+  //TODO: Refactor this bit.
   let defaultSamplingInputs: SamplingInputs.samplingInputs = {
     sampleCount: 10000,
     outputXYPoints: 10000,
@@ -95,57 +104,10 @@ let defaultSamplingInputs: SamplingInputs.samplingInputs = {
     kernelWidth: None,
   }
 
-/* Given two random variables A and B, this returns the distribution
+  let toPointSet = (xyPointLength, t: t) => {
-   of a new variable that is the result of the operation on A and B.
+    switch t {
-   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 = (
-  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))
+    | #Symbolic(r) => Ok(SymbolicDist.T.toPointSetDist(xyPointLength, r))
     | #SampleSet(r) => {
         let response = SampleSet.toPointSetDist(
           ~samples=r,
@@ -160,60 +122,110 @@ let distToPointSet = (sampleCount, dist: genericDist) => {
     }
   }
 
-let rec applyFnInternal = (wrapped: wrapped, fnName: operation): wrappedOutput => {
+  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,
+        )
+      })
+    }
+  }
+}
+
+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
     }