Merge pull request #425 from quantified-uncertainty/scalePower

Scale power

packages/squiggle-lang/src/rescript/Distributions/DistributionOperation/DistributionOperation.res

@@ -154,6 +154,16 @@ let rec run = (~env, functionCallInfo: functionCallInfo): outputType => {
       ->GenericDist.toPointSet(~xyPointLength, ~sampleCount, ())
       ->E.R2.fmap(r => Dist(PointSet(r)))
       ->OutputLocal.fromResult
+    | ToDist(Scale(#Logarithm, f)) => 
+      dist
+      ->GenericDist.pointwiseCombinationFloat(~toPointSetFn, ~algebraicCombination=#Logarithm, ~f)
+      ->E.R2.fmap(r => Dist(r))
+      ->OutputLocal.fromResult
+    | ToDist(Scale(#Power, f)) => 
+      dist
+      ->GenericDist.pointwiseCombinationFloat(~toPointSetFn, ~algebraicCombination=#Power, ~f)
+      ->E.R2.fmap(r => Dist(r))
+      ->OutputLocal.fromResult
     | ToDistCombination(Algebraic(_), _, #Float(_)) => GenDistError(NotYetImplemented)
     | ToDistCombination(Algebraic(strategy), arithmeticOperation, #Dist(t2)) =>
       dist

packages/squiggle-lang/src/rescript/Distributions/DistributionTypes.res

@@ -66,12 +66,19 @@ module DistributionOperation = {
     | #Pdf(float)
     | #Mean
     | #Sample
+    | #IntegralSum
+  ]
+
+  type toScaleFn = [
+    | #Power
+    | #Logarithm
   ]
 
   type toDist =
     | Normalize
     | ToPointSet
     | ToSampleSet(int)
+    | Scale(toScaleFn, float)
     | Truncate(option<float>, option<float>)
     | Inspect
 
@@ -112,6 +119,8 @@ module DistributionOperation = {
     | ToDist(ToSampleSet(r)) => `toSampleSet(${E.I.toString(r)})`
     | ToDist(Truncate(_, _)) => `truncate`
     | ToDist(Inspect) => `inspect`
+    | ToDist(Scale(#Power, r)) => `scalePower(${E.Float.toFixed(r)})`
+    | ToDist(Scale(#Logarithm, r)) => `scaleLog(${E.Float.toFixed(r)})`
     | ToString(ToString) => `toString`
     | ToString(ToSparkline(n)) => `toSparkline(${E.I.toString(n)})`
     | ToBool(IsNormalized) => `isNormalized`
@@ -143,6 +152,8 @@ module Constructors = {
     let fromSamples = (xs): t => FromSamples(xs)
     let truncate = (dist, left, right): t => FromDist(ToDist(Truncate(left, right)), dist)
     let inspect = (dist): t => FromDist(ToDist(Inspect), dist)
+    let scalePower = (dist, n): t => FromDist(ToDist(Scale(#Power, n)), dist)
+    let scaleLogarithm = (dist, n): t => FromDist(ToDist(Scale(#Logarithm, n)), dist)
     let toString = (dist): t => FromDist(ToString(ToString), dist)
     let toSparkline = (dist, n): t => FromDist(ToString(ToSparkline(n)), dist)
     let algebraicAdd = (dist1, dist2: genericDist): t => FromDist(

packages/squiggle-lang/src/rescript/Distributions/GenericDist/GenericDist.res

@@ -62,26 +62,31 @@ let isNormalized = (t: t): bool => Js.Math.abs_float(integralEndY(t) -. 1.0) < 1
 let toFloatOperation = (
   t,
   ~toPointSetFn: toPointSetFn,
-  ~distToFloatOperation: Operation.distToFloatOperation,
+  ~distToFloatOperation: DistributionTypes.DistributionOperation.toFloat,
 ) => {
-  let trySymbolicSolution = switch (t: t) {
-  | Symbolic(r) => SymbolicDist.T.operate(distToFloatOperation, r)->E.R.toOption
-  | _ => None
-  }
+  switch distToFloatOperation {
+  | #IntegralSum => Ok(integralEndY(t))
+  | (#Pdf(_) | #Cdf(_) | #Inv(_) | #Mean | #Sample) as op => {
+      let trySymbolicSolution = switch (t: t) {
+      | Symbolic(r) => SymbolicDist.T.operate(op, r)->E.R.toOption
+      | _ => None
+      }
 
-  let trySampleSetSolution = switch ((t: t), distToFloatOperation) {
-  | (SampleSet(sampleSet), #Mean) => SampleSetDist.mean(sampleSet)->Some
-  | (SampleSet(sampleSet), #Sample) => SampleSetDist.sample(sampleSet)->Some
-  | (SampleSet(sampleSet), #Inv(r)) => SampleSetDist.percentile(sampleSet, r)->Some
-  | _ => None
-  }
+      let trySampleSetSolution = switch ((t: t), distToFloatOperation) {
+      | (SampleSet(sampleSet), #Mean) => SampleSetDist.mean(sampleSet)->Some
+      | (SampleSet(sampleSet), #Sample) => SampleSetDist.sample(sampleSet)->Some
+      | (SampleSet(sampleSet), #Inv(r)) => SampleSetDist.percentile(sampleSet, r)->Some
+      | _ => None
+      }
 
-  switch trySymbolicSolution {
-  | Some(r) => Ok(r)
-  | None =>
-    switch trySampleSetSolution {
-    | Some(r) => Ok(r)
-    | None => toPointSetFn(t)->E.R2.fmap(PointSetDist.operate(distToFloatOperation))
+      switch trySymbolicSolution {
+      | Some(r) => Ok(r)
+      | None =>
+        switch trySampleSetSolution {
+        | Some(r) => Ok(r)
+        | None => toPointSetFn(t)->E.R2.fmap(PointSetDist.operate(op))
+        }
+      }
     }
   }
 }

packages/squiggle-lang/src/rescript/Distributions/GenericDist/GenericDist.resi

@@ -20,7 +20,7 @@ let isNormalized: t => bool
 let toFloatOperation: (
   t,
   ~toPointSetFn: toPointSetFn,
-  ~distToFloatOperation: Operation.distToFloatOperation,
+  ~distToFloatOperation: DistributionTypes.DistributionOperation.toFloat,
 ) => result<float, error>
 
 @genType

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/Continuous.res

@@ -156,8 +156,10 @@ let reduce = (
   ~integralSumCachesFn: (float, float) => option<float>=(_, _) => None,
   fn: (float, float) => result<float, 'e>,
   continuousShapes,
-): result<t, 'e> =>
-  continuousShapes |> E.A.R.foldM(combinePointwise(~integralSumCachesFn, fn), empty)
+): result<t, 'e> => {
+  let merge = combinePointwise(~integralSumCachesFn, fn)
+  continuousShapes |> E.A.R.foldM(merge, empty)
+}
 
 let mapYResult = (
   ~integralSumCacheFn=_ => None,

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/Discrete.res

@@ -34,9 +34,10 @@ let lastY = (t: t) => t |> getShape |> XYShape.T.lastY
 
 let combinePointwise = (
   ~integralSumCachesFn=(_, _) => None,
+  fn,
   t1: PointSetTypes.discreteShape,
   t2: PointSetTypes.discreteShape,
-): PointSetTypes.discreteShape => {
+): result<PointSetTypes.discreteShape, 'e> => {
   let combinedIntegralSum = Common.combineIntegralSums(
     integralSumCachesFn,
     t1.integralSumCache,
@@ -49,16 +50,22 @@ let combinePointwise = (
   make(
     ~integralSumCache=combinedIntegralSum,
     XYShape.PointwiseCombination.combine(
-      (a, b) => Ok(a +. b),
+      fn,
       XYShape.XtoY.discreteInterpolator,
       t1.xyShape,
       t2.xyShape,
     )->E.R.toExn("Addition operation should never fail", _),
-  )
+  )->Ok
 }
 
-let reduce = (~integralSumCachesFn=(_, _) => None, discreteShapes): PointSetTypes.discreteShape =>
-  discreteShapes |> E.A.fold_left(combinePointwise(~integralSumCachesFn), empty)
+let reduce = (
+  ~integralSumCachesFn=(_, _) => None,
+  fn: (float, float) => result<float, 'e>,
+  discreteShapes: array<PointSetTypes.discreteShape>,
+): result<t, 'e> => {
+  let merge = combinePointwise(~integralSumCachesFn, fn)
+  discreteShapes |> E.A.R.foldM(merge, empty)
+}
 
 let updateIntegralSumCache = (integralSumCache, t: t): t => {
   ...t,

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/Mixed.res

@@ -316,7 +316,10 @@ let combinePointwise = (
   t2: t,
 ): result<t, 'e> => {
   let reducedDiscrete =
-    [t1, t2] |> E.A.fmap(toDiscrete) |> E.A.O.concatSomes |> Discrete.reduce(~integralSumCachesFn)
+    [t1, t2]
+    |> E.A.fmap(toDiscrete)
+    |> E.A.O.concatSomes
+    |> Discrete.reduce(~integralSumCachesFn, fn)
 
   let reducedContinuous =
     [t1, t2]
@@ -335,11 +338,11 @@ let combinePointwise = (
     t1.integralCache,
     t2.integralCache,
   )
-  reducedContinuous->E.R2.fmap(continuous =>
+  E.R.merge(reducedContinuous, reducedDiscrete)->E.R2.fmap(((continuous, discrete)) =>
     make(
       ~integralSumCache=combinedIntegralSum,
       ~integralCache=combinedIntegral,
-      ~discrete=reducedDiscrete,
+      ~discrete,
       ~continuous,
     )
   )

packages/squiggle-lang/src/rescript/Distributions/PointSetDist/PointSetDist.res

@@ -84,7 +84,12 @@ let combinePointwise = (
       m2,
     )->E.R2.fmap(x => PointSetTypes.Continuous(x))
   | (Discrete(m1), Discrete(m2)) =>
-    Ok(PointSetTypes.Discrete(Discrete.combinePointwise(~integralSumCachesFn, m1, m2)))
+    Discrete.combinePointwise(
+      ~integralSumCachesFn,
+      fn,
+      m1,
+      m2,
+    )->E.R2.fmap(x => PointSetTypes.Discrete(x))
   | (m1, m2) =>
     Mixed.combinePointwise(
       ~integralSumCachesFn,

packages/squiggle-lang/src/rescript/ReducerInterface/ReducerInterface_GenericDistribution.res

@@ -196,6 +196,7 @@ let dispatchToGenericOutput = (call: ExpressionValue.functionCall): option<
     ->SymbolicConstructors.symbolicResultToOutput
   | ("sample", [EvDistribution(dist)]) => Helpers.toFloatFn(#Sample, dist)
   | ("mean", [EvDistribution(dist)]) => Helpers.toFloatFn(#Mean, dist)
+  | ("integralSum", [EvDistribution(dist)]) => Helpers.toFloatFn(#IntegralSum, dist)
   | ("toString", [EvDistribution(dist)]) => Helpers.toStringFn(ToString, dist)
   | ("toSparkline", [EvDistribution(dist)]) => Helpers.toStringFn(ToSparkline(20), dist)
   | ("toSparkline", [EvDistribution(dist), EvNumber(n)]) =>
@@ -211,6 +212,15 @@ let dispatchToGenericOutput = (call: ExpressionValue.functionCall): option<
   | ("normalize", [EvDistribution(dist)]) => Helpers.toDistFn(Normalize, dist)
   | ("isNormalized", [EvDistribution(dist)]) => Helpers.toBoolFn(IsNormalized, dist)
   | ("toPointSet", [EvDistribution(dist)]) => Helpers.toDistFn(ToPointSet, dist)
+  | ("scaleLog", [EvDistribution(dist)]) =>
+    Helpers.toDistFn(Scale(#Logarithm, MagicNumbers.Math.e), dist)
+  | ("scaleLog10", [EvDistribution(dist)]) => Helpers.toDistFn(Scale(#Logarithm, 10.0), dist)
+  | ("scaleLog", [EvDistribution(dist), EvNumber(float)]) =>
+    Helpers.toDistFn(Scale(#Logarithm, float), dist)
+  | ("scalePow", [EvDistribution(dist), EvNumber(float)]) =>
+    Helpers.toDistFn(Scale(#Power, float), dist)
+  | ("scaleExp", [EvDistribution(dist)]) =>
+    Helpers.toDistFn(Scale(#Power, MagicNumbers.Math.e), dist)
   | ("cdf", [EvDistribution(dist), EvNumber(float)]) => Helpers.toFloatFn(#Cdf(float), dist)
   | ("pdf", [EvDistribution(dist), EvNumber(float)]) => Helpers.toFloatFn(#Pdf(float), dist)
   | ("inv", [EvDistribution(dist), EvNumber(float)]) => Helpers.toFloatFn(#Inv(float), dist)