Merge pull request #860 from quantified-uncertainty/sampleN

`mapSamplesN`

packages/squiggle-lang/__tests__/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn_test.res

@@ -21,6 +21,10 @@ describe("builtin", () => {
     "addOne(t)=t+1; toList(mapSamples(fromSamples([1,2,3,4,5,6]), addOne))",
     "Ok([2,3,4,5,6,7])",
   )
+  testEval(
+    "toList(mapSamplesN([fromSamples([1,2,3,4,5,6]), fromSamples([6, 5, 4, 3, 2, 1])], {|x| x[0] > x[1] ? x[0] : x[1]}))",
+    "Ok([6,5,4,4,5,6])",
+  )
 })
 
 describe("builtin exception", () => {

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

@@ -117,6 +117,11 @@ let map3 = (
 ): result<t, sampleSetError> =>
   E.A.zip3(get(t1), get(t2), get(t3))->E.A2.fmap(E.Tuple3.toFnCall(fn))->_fromSampleResultArray
 
+let mapN = (~fn: array<float> => result<float, Operation.Error.t>, ~t1: array<t>): result<
+  t,
+  sampleSetError,
+> => E.A.transpose(E.A.fmap(get, t1))->E.A2.fmap(fn)->_fromSampleResultArray
+
 let mean = t => T.get(t)->E.A.Floats.mean
 let geomean = t => T.get(t)->E.A.Floats.geomean
 let mode = t => T.get(t)->E.A.Floats.mode

packages/squiggle-lang/src/rescript/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res

@@ -149,6 +149,27 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
         doLambdaCall(aLambdaValue, list{IEvNumber(a), IEvNumber(b), IEvNumber(c)})
       SampleSetDist.map3(~fn, ~t1, ~t2, ~t3)->toType
     }
+
+    let parseSampleSetArray = (arr: array<internalExpressionValue>): option<
+      array<SampleSetDist.t>,
+    > => {
+      let parseSampleSet = (value: internalExpressionValue): option<SampleSetDist.t> =>
+        switch value {
+        | IEvDistribution(SampleSet(dist)) => Some(dist)
+        | _ => None
+        }
+      E.A.O.openIfAllSome(E.A.fmap(parseSampleSet, arr))
+    }
+
+    let mapN = (aValueArray: array<internalExpressionValue>, aLambdaValue) => {
+      switch parseSampleSetArray(aValueArray) {
+      | Some(t1) =>
+        let fn = a => doLambdaCall(aLambdaValue, list{IEvArray(E.A.fmap(x => IEvNumber(x), a))})
+        SampleSetDist.mapN(~fn, ~t1)->toType
+      | None =>
+        Error(REFunctionNotFound(call->functionCallToCallSignature->functionCallSignatureToString))
+      }
+    }
   }
 
   let doReduceArray = (aValueArray, initialValue, aLambdaValue) => {
@@ -230,6 +251,8 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
       ],
     ) =>
     SampleMap.map3(dist1, dist2, dist3, aLambdaValue)
+  | ("mapSamplesN", [IEvArray(aValueArray), IEvLambda(aLambdaValue)]) =>
+    SampleMap.mapN(aValueArray, aLambdaValue)
   | ("reduce", [IEvArray(aValueArray), initialValue, IEvLambda(aLambdaValue)]) =>
     doReduceArray(aValueArray, initialValue, aLambdaValue)
   | ("reduceReverse", [IEvArray(aValueArray), initialValue, IEvLambda(aLambdaValue)]) =>
@@ -246,7 +269,6 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
     Error(REFunctionNotFound(call->functionCallToCallSignature->functionCallSignatureToString)) // Report full type signature as error
   }
 }
-
 /*
   Reducer uses Result monad while reducing expressions
 */

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

@@ -631,6 +631,19 @@ module A = {
     )
   let filter = Js.Array.filter
   let joinWith = Js.Array.joinWith
+  let transpose = (xs: array<array<'a>>): array<array<'a>> => {
+    let arr: array<array<'a>> = []
+    for i in 0 to length(xs) - 1 {
+      for j in 0 to length(xs[i]) - 1 {
+        if Js.Array.length(arr) <= j {
+          ignore(Js.Array.push([xs[i][j]], arr))
+        } else {
+          ignore(Js.Array.push(xs[i][j], arr[j]))
+        }
+      }
+    }
+    arr
+  }
 
   let all = (p: 'a => bool, xs: array<'a>): bool => length(filter(p, xs)) == length(xs)
   let any = (p: 'a => bool, xs: array<'a>): bool => length(filter(p, xs)) > 0