Merge pull request #598 from quantified-uncertainty/samplesMap2

mapSamples2 and mapSamples3

packages/squiggle-lang/.prettierignore

@@ -3,6 +3,6 @@ lib
 *.bs.js
 *.gen.tsx
 .nyc_output/
-coverage/
+_coverage/
 .cache/
 Reducer_Peggy_GeneratedParser.js

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

@@ -281,7 +281,7 @@ module AlgebraicCombination = {
       let fn = Operation.Algebraic.toFn(arithmeticOperation)
       E.R.merge(toSampleSet(t1), toSampleSet(t2))
       ->E.R.bind(((t1, t2)) => {
-        SampleSetDist.map2(~fn, ~t1, ~t2)->E.R2.errMap(x => DistributionTypes.OperationError(x))
+        SampleSetDist.map2(~fn, ~t1, ~t2)->E.R2.errMap(x => DistributionTypes.SampleSetError(x))
       })
       ->E.R2.fmap(r => DistributionTypes.SampleSet(r))
     }

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

@@ -20,6 +20,14 @@ module Error = {
     }
 
   let fromOperationError = e => OperationError(e)
+
+  let toString = (err: sampleSetError) => {
+    switch err {
+    | TooFewSamples => "Too few samples when constructing sample set"
+    | NonNumericInput(err) => `Found a non-number in input: ${err}`
+    | OperationError(err) => Operation.Error.toString(err)
+    }
+  }
 }
 
 include Error
@@ -87,30 +95,27 @@ let sampleN = (t: t, n) => {
   }
 }
 
+let _fromSampleResultArray = (samples: array<result<float, QuriSquiggleLang.Operation.Error.t>>) =>
+  E.A.R.firstErrorOrOpen(samples)->E.R2.errMap(Error.fromOperationError) |> E.R2.bind(make)
+
 let samplesMap = (~fn: float => result<float, Operation.Error.t>, t: t): result<
   t,
   sampleSetError,
-> => {
-  let samples = T.get(t)->E.A2.fmap(fn)
-  E.A.R.firstErrorOrOpen(samples)->E.R2.errMap(Error.fromOperationError) |> E.R2.bind(make)
-}
+> => T.get(t)->E.A2.fmap(fn)->_fromSampleResultArray
 
 //TODO: Figure out what to do if distributions are different lengths. ``zip`` is kind of inelegant for this.
 let map2 = (~fn: (float, float) => result<float, Operation.Error.t>, ~t1: t, ~t2: t): result<
   t,
-  Operation.Error.t,
-> => {
-  let samples = Belt.Array.zip(get(t1), get(t2))->E.A2.fmap(((a, b)) => fn(a, b))
+  sampleSetError,
+> => E.A.zip(get(t1), get(t2))->E.A2.fmap(E.Tuple2.toFnCall(fn))->_fromSampleResultArray
 
-  // This assertion should never be reached. In order for it to be reached, one
-  // of the input parameters would need to be a sample set distribution with less
-  // than 6 samples. Which should be impossible due to the smart constructor.
-  // I could prove this to the type system (say, creating a {first: float, second: float, ..., fifth: float, rest: array<float>}
-  // But doing so would take too much time, so I'll leave it as an assertion
-  E.A.R.firstErrorOrOpen(samples)->E.R2.fmap(x =>
-    E.R.toExnFnString(Error.sampleSetErrorToString, make(x))
-  )
-}
+let map3 = (
+  ~fn: (float, float, float) => result<float, Operation.Error.t>,
+  ~t1: t,
+  ~t2: t,
+  ~t3: t,
+): result<t, sampleSetError> =>
+  E.A.zip3(get(t1), get(t2), get(t3))->E.A2.fmap(E.Tuple3.toFnCall(fn))->_fromSampleResultArray
 
 let mean = t => T.get(t)->E.A.Floats.mean
 let geomean = t => T.get(t)->E.A.Floats.geomean

packages/squiggle-lang/src/rescript/FunctionRegistry/FunctionRegistry_Helpers.res

@@ -75,7 +75,7 @@ module Process = {
         | Ok((t1, t2)) =>
           switch SampleSetDist.map2(~fn=altFn, ~t1, ~t2) {
           | Ok(r) => Ok(DistributionTypes.SampleSet(r))
-          | Error(r) => Error(Operation.Error.toString(r))
+          | Error(r) => Error(SampleSetDist.Error.toString(r))
           }
         | Error(r) => Error(DistributionTypes.Error.toString(r))
         }

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

@@ -101,16 +101,34 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
     rMappedList->Result.map(mappedList => mappedList->Belt.List.toArray->EvArray)
   }
 
-  let doMapSampleSetDist = (sampleSetDist: SampleSetDist.t, aLambdaValue) => {
-    let fn = r =>
-      switch Lambda.doLambdaCall(aLambdaValue, list{EvNumber(r)}, environment, reducer) {
+  module SampleMap = {
+    type t = SampleSetDist.t
+    let doLambdaCall = (aLambdaValue, list) =>
+      switch Lambda.doLambdaCall(aLambdaValue, list, environment, reducer) {
       | Ok(EvNumber(f)) => Ok(f)
       | _ => Error(Operation.SampleMapNeedsNtoNFunction)
       }
-    switch SampleSetDist.samplesMap(~fn, sampleSetDist) {
+
+    let toType = r =>
+      switch r {
       | Ok(r) => Ok(EvDistribution(SampleSet(r)))
       | Error(r) => Error(REDistributionError(SampleSetError(r)))
       }
+
+    let map1 = (sampleSetDist: t, aLambdaValue) => {
+      let fn = r => doLambdaCall(aLambdaValue, list{EvNumber(r)})
+      toType(SampleSetDist.samplesMap(~fn, sampleSetDist))
+    }
+
+    let map2 = (t1: t, t2: t, aLambdaValue) => {
+      let fn = (a, b) => doLambdaCall(aLambdaValue, list{EvNumber(a), EvNumber(b)})
+      SampleSetDist.map2(~fn, ~t1, ~t2)->toType
+    }
+
+    let map3 = (t1: t, t2: t, t3: t, aLambdaValue) => {
+      let fn = (a, b, c) => doLambdaCall(aLambdaValue, list{EvNumber(a), EvNumber(b), EvNumber(c)})
+      SampleSetDist.map3(~fn, ~t1, ~t2, ~t3)->toType
+    }
   }
 
   let doReduceArray = (aValueArray, initialValue, aLambdaValue) => {
@@ -143,7 +161,22 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
     doKeepArray(aValueArray, aLambdaValue)
   | ("map", [EvArray(aValueArray), EvLambda(aLambdaValue)]) => doMapArray(aValueArray, aLambdaValue)
   | ("mapSamples", [EvDistribution(SampleSet(dist)), EvLambda(aLambdaValue)]) =>
-    doMapSampleSetDist(dist, aLambdaValue)
+    SampleMap.map1(dist, aLambdaValue)
+  | (
+      "mapSamples2",
+      [EvDistribution(SampleSet(dist1)), EvDistribution(SampleSet(dist2)), EvLambda(aLambdaValue)],
+    ) =>
+    SampleMap.map2(dist1, dist2, aLambdaValue)
+  | (
+      "mapSamples3",
+      [
+        EvDistribution(SampleSet(dist1)),
+        EvDistribution(SampleSet(dist2)),
+        EvDistribution(SampleSet(dist3)),
+        EvLambda(aLambdaValue),
+      ],
+    ) =>
+    SampleMap.map3(dist1, dist2, dist3, aLambdaValue)
   | ("reduce", [EvArray(aValueArray), initialValue, EvLambda(aLambdaValue)]) =>
     doReduceArray(aValueArray, initialValue, aLambdaValue)
   | ("reduceReverse", [EvArray(aValueArray), initialValue, EvLambda(aLambdaValue)]) =>

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

@@ -55,6 +55,10 @@ module Tuple2 = {
   let toFnCall = (fn, (a1, a2)) => fn(a1, a2)
 }
 
+module Tuple3 = {
+  let toFnCall = (fn, (a1, a2, a3)) => fn(a1, a2, a3)
+}
+
 module O = {
   let dimap = (sFn, rFn, e) =>
     switch e {
@@ -567,6 +571,8 @@ module A = {
   let tail = Belt.Array.sliceToEnd(_, 1)
 
   let zip = Belt.Array.zip
+  let zip3 = (a, b, c) =>
+    Belt.Array.zip(a, b)->Belt.Array.zip(c)->Belt.Array.map((((v1, v2), v3)) => (v1, v2, v3))
   // This zips while taking the longest elements of each array.
   let zipMaxLength = (array1, array2) => {
     let maxLength = Int.max(length(array1), length(array2))

yarn.lock

@@ -4310,10 +4310,10 @@
   dependencies:
     "@types/react" "*"
 
-"@types/react@*", "@types/react@17.0.43", "@types/react@^18.0.9":
-  version "17.0.43"
-  resolved "https://registry.yarnpkg.com/@types/react/-/react-17.0.43.tgz#4adc142887dd4a2601ce730bc56c3436fdb07a55"
-  integrity sha512-8Q+LNpdxf057brvPu1lMtC5Vn7J119xrP1aq4qiaefNioQUYANF/CYeK4NsKorSZyUGJ66g0IM+4bbjwx45o2A==
+"@types/react@*", "@types/react@^18.0.1", "@types/react@^18.0.9":
+  version "18.0.9"
+  resolved "https://registry.yarnpkg.com/@types/react/-/react-18.0.9.tgz#d6712a38bd6cd83469603e7359511126f122e878"
+  integrity sha512-9bjbg1hJHUm4De19L1cHiW0Jvx3geel6Qczhjd0qY5VKVE2X5+x77YxAepuCwVh4vrgZJdgEJw48zrhRIeF4Nw==
   dependencies:
     "@types/prop-types" "*"
     "@types/scheduler" "*"
@@ -14557,7 +14557,7 @@ react-vega@^7.5.1:
     prop-types "^15.8.1"
     vega-embed "^6.5.1"
 
-react@^18.1.0:
+react@^18.0.0, react@^18.1.0:
   version "18.1.0"
   resolved "https://registry.yarnpkg.com/react/-/react-18.1.0.tgz#6f8620382decb17fdc5cc223a115e2adbf104890"
   integrity sha512-4oL8ivCz5ZEPyclFQXaNksK3adutVS8l2xzZU0cqEFrE9Sb7fC0EFK5uEk74wIreL1DERyjvsU915j1pcT2uEQ==