First implementation of sampleSet mixed distribution

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

@@ -216,7 +216,7 @@ let rec run = (~env: env, functionCallInfo: functionCallInfo): outputType => {
   | FromFloat(subFnName, x) => reCall(~functionCallInfo=FromFloat(subFnName, x), ())
   | Mixture(dists) =>
     dists
-    ->GenericDist.mixture(~scaleMultiplyFn=scaleMultiply, ~pointwiseAddFn=pointwiseAdd)
+    ->GenericDist.mixture(~scaleMultiplyFn=scaleMultiply, ~pointwiseAddFn=pointwiseAdd, ~env)
     ->E.R2.fmap(r => Dist(r))
     ->OutputLocal.fromResult
   | FromSamples(xs) =>

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

@@ -491,15 +491,30 @@ let pointwiseCombinationFloat = (
   m->E.R2.fmap(r => DistributionTypes.PointSet(r))
 }
 
-//Note: The result should always cumulatively sum to 1. This would be good to test.
-//Note: If the inputs are not normalized, this will return poor results. The weights probably refer to the post-normalized forms. It would be good to apply a catch to this.
+//TODO: The result should always cumulatively sum to 1. This would be good to test.
+//TODO: If the inputs are not normalized, this will return poor results. The weights probably refer to the post-normalized forms. It would be good to apply a catch to this.
 let mixture = (
   values: array<(t, float)>,
   ~scaleMultiplyFn: scaleMultiplyFn,
   ~pointwiseAddFn: pointwiseAddFn,
+  ~env: env,
 ) => {
-  if E.A.length(values) == 0 {
+  let allValuesAreSampleSet = v => E.A.all(((t, _)) => isSampleSetSet(t), v)
+
+  if E.A.isEmpty(values) {
     Error(DistributionTypes.OtherError("Mixture error: mixture must have at least 1 element"))
+  } else if allValuesAreSampleSet(values) {
+    let withSampleSetValues = values->E.A2.fmap(((value, weight)) =>
+      switch value {
+      | SampleSet(sampleSet) => Ok((sampleSet, weight))
+      | _ => Error("Unreachable")
+      } |> E.R.toExn("Mixture coding error: SampleSet expected. This should be inaccessible.")
+    )
+    let sampleSetMixture = SampleSetDist.mixture(withSampleSetValues, env.sampleCount)
+    switch sampleSetMixture {
+    | Ok(sampleSet) => Ok(DistributionTypes.SampleSet(sampleSet))
+    | Error(err) => Error(DistributionTypes.Error.sampleErrorToDistErr(err))
+    }
   } else {
     let totalWeight = values->E.A2.fmap(E.Tuple2.second)->E.A.Floats.sum
     let properlyWeightedValues =

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

@@ -81,6 +81,7 @@ let mixture: (
   array<(t, float)>,
   ~scaleMultiplyFn: scaleMultiplyFn,
   ~pointwiseAddFn: pointwiseAddFn,
+  ~env: env,
 ) => result<t, error>
 
 let isSymbolic: t => bool

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

@@ -131,3 +131,15 @@ let max = t => T.get(t)->E.A.Floats.max
 let stdev = t => T.get(t)->E.A.Floats.stdev
 let variance = t => T.get(t)->E.A.Floats.variance
 let percentile = (t, f) => T.get(t)->E.A.Floats.percentile(f)
+
+let mixture = (values: array<(t, float)>, intendedLength: int) => {
+  let totalWeight = values->E.A2.fmap(E.Tuple2.second)->E.A.Floats.sum
+  values
+  ->E.A2.fmap(((dist, weight)) => {
+    let adjustedWeight = weight /. totalWeight
+    let samplesToGet = adjustedWeight *. E.I.toFloat(intendedLength) |> E.Float.toInt
+    sampleN(dist, samplesToGet)
+  })
+  ->E.A.concatMany
+  ->T.make
+}

packages/squiggle-lang/src/rescript/Utility/E.res

@@ -220,6 +220,7 @@ module I = {
   let increment = n => n + 1
   let decrement = n => n - 1
   let toString = Js.Int.toString
+  let toFloat = Js.Int.toFloat
 }
 
 exception Assertion(string)