Merge pull request #662 from quantified-uncertainty/utility-functions

Adds a few utility functions

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

@@ -10,8 +10,10 @@ type rec frType =
   | FRTypeDistOrNumber
   | FRTypeLambda
   | FRTypeRecord(frTypeRecord)
+  | FRTypeDict(frType)
   | FRTypeArray(frType)
   | FRTypeString
+  | FRTypeAny
   | FRTypeVariant(array<string>)
 and frTypeRecord = array<frTypeRecordParam>
 and frTypeRecordParam = (string, frType)
@@ -29,8 +31,11 @@ type rec frValue =
   | FRValueLambda(ReducerInterface_ExpressionValue.lambdaValue)
   | FRValueString(string)
   | FRValueVariant(string)
+  | FRValueAny(frValue)
+  | FRValueDict(Js.Dict.t<frValue>)
 and frValueRecord = array<frValueRecordParam>
 and frValueRecordParam = (string, frValue)
+and frValueDictParam = (string, frValue)
 and frValueDistOrNumber = FRValueNumber(float) | FRValueDist(DistributionTypes.genericDist)
 
 type fnDefinition = {
@@ -57,14 +62,34 @@ module FRType = {
         let input = ((name, frType): frTypeRecordParam) => `${name}: ${toString(frType)}`
         `record({${r->E.A2.fmap(input)->E.A2.joinWith(", ")}})`
       }
-    | FRTypeArray(r) => `record(${toString(r)})`
+    | FRTypeArray(r) => `list(${toString(r)})`
     | FRTypeLambda => `lambda`
     | FRTypeString => `string`
     | FRTypeVariant(_) => "variant"
+    | FRTypeDict(r) => `dict(${toString(r)})`
+    | FRTypeAny => `any`
+    }
+
+  let rec toFrValue = (r: expressionValue): option<frValue> =>
+    switch r {
+    | EvNumber(f) => Some(FRValueNumber(f))
+    | EvString(f) => Some(FRValueString(f))
+    | EvDistribution(f) => Some(FRValueDistOrNumber(FRValueDist(f)))
+    | EvLambda(f) => Some(FRValueLambda(f))
+    | EvArray(elements) =>
+      elements->E.A2.fmap(toFrValue)->E.A.O.openIfAllSome->E.O2.fmap(r => FRValueArray(r))
+    | EvRecord(record) =>
+      Js.Dict.entries(record)
+      ->E.A2.fmap(((key, item)) => item->toFrValue->E.O2.fmap(o => (key, o)))
+      ->E.A.O.openIfAllSome
+      ->E.O2.fmap(r => FRValueRecord(r))
+    | _ => None
     }
 
   let rec matchWithExpressionValue = (t: t, r: expressionValue): option<frValue> =>
     switch (t, r) {
+    | (FRTypeAny, f) => toFrValue(f)
+    | (FRTypeString, EvString(f)) => Some(FRValueString(f))
     | (FRTypeNumber, EvNumber(f)) => Some(FRValueNumber(f))
     | (FRTypeDistOrNumber, EvNumber(f)) => Some(FRValueDistOrNumber(FRValueNumber(f)))
     | (FRTypeDistOrNumber, EvDistribution(Symbolic(#Float(f)))) =>
@@ -77,6 +102,12 @@ module FRType = {
         let el = elements->E.A2.fmap(matchWithExpressionValue(intendedType))
         E.A.O.openIfAllSome(el)->E.O2.fmap(r => FRValueArray(r))
       }
+    | (FRTypeDict(r), EvRecord(record)) =>
+      record
+      ->Js.Dict.entries
+      ->E.A2.fmap(((key, item)) => matchWithExpressionValue(r, item)->E.O2.fmap(o => (key, o)))
+      ->E.A.O.openIfAllSome
+      ->E.O2.fmap(r => FRValueDict(Js.Dict.fromArray(r)))
     | (FRTypeRecord(recordParams), EvRecord(record)) => {
         let getAndMatch = (name, input) =>
           E.Dict.get(record, name)->E.O.bind(matchWithExpressionValue(input))
@@ -103,9 +134,18 @@ module FRType = {
           frValueRecord->E.A2.fmap(((name, value)) => (name, matchReverse(value)))->E.Dict.fromArray
         EvRecord(record)
       }
+    | FRValueDict(frValueRecord) => {
+        let record =
+          frValueRecord
+          ->Js.Dict.entries
+          ->E.A2.fmap(((name, value)) => (name, matchReverse(value)))
+          ->E.Dict.fromArray
+        EvRecord(record)
+      }
     | FRValueLambda(l) => EvLambda(l)
     | FRValueString(string) => EvString(string)
     | FRValueVariant(string) => EvString(string)
+    | FRValueAny(f) => matchReverse(f)
     }
 
   let matchWithExpressionValueArray = (inputs: array<t>, args: array<expressionValue>): option<

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

@@ -5,9 +5,15 @@ let impossibleError = "Wrong inputs / Logically impossible"
 module Wrappers = {
   let symbolic = r => DistributionTypes.Symbolic(r)
   let evDistribution = r => ReducerInterface_ExpressionValue.EvDistribution(r)
+  let evNumber = r => ReducerInterface_ExpressionValue.EvNumber(r)
+  let evArray = r => ReducerInterface_ExpressionValue.EvArray(r)
+  let evRecord = r => ReducerInterface_ExpressionValue.EvRecord(r)
+  let evString = r => ReducerInterface_ExpressionValue.EvString(r)
   let symbolicEvDistribution = r => r->DistributionTypes.Symbolic->evDistribution
 }
 
+let getOrError = (a, g) => E.A.get(a, g) |> E.O.toResult(impossibleError)
+
 module Prepare = {
   type t = frValue
   type ts = array<frValue>
@@ -34,6 +40,12 @@ module Prepare = {
         | FRValueArray(n) => Ok(n)
         | _ => Error(impossibleError)
         }
+
+      let arrayOfArrays = (inputs: t): result<array<ts>, err> =>
+        switch inputs {
+        | FRValueArray(n) => n->E.A2.fmap(openA)->E.A.R.firstErrorOrOpen
+        | _ => Error(impossibleError)
+        }
     }
   }
 
@@ -52,6 +64,13 @@ module Prepare = {
       }
     }
 
+    let threeNumbers = (values: ts): result<(float, float, float), err> => {
+      switch values {
+      | [FRValueNumber(a1), FRValueNumber(a2), FRValueNumber(a3)] => Ok(a1, a2, a3)
+      | _ => Error(impossibleError)
+      }
+    }
+
     let oneDistOrNumber = (values: ts): result<frValueDistOrNumber, err> => {
       switch values {
       | [FRValueDistOrNumber(a1)] => Ok(a1)
@@ -77,6 +96,33 @@ module Prepare = {
       pairs
     }
   }
+
+  let oneNumber = (values: t): result<float, err> => {
+    switch values {
+    | FRValueNumber(a1) => Ok(a1)
+    | _ => Error(impossibleError)
+    }
+  }
+
+  let oneDict = (values: t): result<Js.Dict.t<frValue>, err> => {
+    switch values {
+    | FRValueDict(a1) => Ok(a1)
+    | _ => Error(impossibleError)
+    }
+  }
+
+  module ToTypedArray = {
+    let numbers = (inputs: ts): result<array<float>, err> => {
+      let openNumbers = (elements: array<t>) =>
+        elements->E.A2.fmap(oneNumber)->E.A.R.firstErrorOrOpen
+      inputs->getOrError(0)->E.R.bind(ToValueArray.Array.openA)->E.R.bind(openNumbers)
+    }
+
+    let dicts = (inputs: ts): Belt.Result.t<array<Js.Dict.t<frValue>>, err> => {
+      let openDicts = (elements: array<t>) => elements->E.A2.fmap(oneDict)->E.A.R.firstErrorOrOpen
+      inputs->getOrError(0)->E.R.bind(ToValueArray.Array.openA)->E.R.bind(openDicts)
+    }
+  }
 }
 
 module Process = {
@@ -183,9 +229,36 @@ module OneArgDist = {
     ->E.R.bind(Process.DistOrNumberToDist.oneValueUsingSymbolicDist(~fn, ~value=_, ~env))
     ->E.R2.fmap(Wrappers.evDistribution)
 
-  let make = (name, fn) => {
+  let make = (name, fn) =>
     FnDefinition.make(~name, ~inputs=[FRTypeDistOrNumber], ~run=(inputs, env) =>
       inputs->Prepare.ToValueTuple.oneDistOrNumber->process(~fn, ~env)
     )
+}
+
+module ArrayNumberDist = {
+  let make = (name, fn) => {
+    FnDefinition.make(~name, ~inputs=[FRTypeArray(FRTypeNumber)], ~run=(inputs, _) =>
+      Prepare.ToTypedArray.numbers(inputs)
+      ->E.R.bind(r => E.A.length(r) === 0 ? Error("List is empty") : Ok(r))
+      ->E.R.bind(fn)
+    )
+  }
+  let make2 = (name, fn) => {
+    FnDefinition.make(~name, ~inputs=[FRTypeArray(FRTypeAny)], ~run=(inputs, _) =>
+      Prepare.ToTypedArray.numbers(inputs)
+      ->E.R.bind(r => E.A.length(r) === 0 ? Error("List is empty") : Ok(r))
+      ->E.R.bind(fn)
+    )
   }
 }
+
+module NumberToNumber = {
+  let make = (name, fn) =>
+    FnDefinition.make(~name, ~inputs=[FRTypeNumber], ~run=(inputs, _) => {
+      inputs
+      ->getOrError(0)
+      ->E.R.bind(Prepare.oneNumber)
+      ->E.R2.fmap(fn)
+      ->E.R2.fmap(Wrappers.evNumber)
+    })
+}

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

@@ -31,7 +31,7 @@ module Declaration = {
 }
 
 let inputsTodist = (inputs: array<FunctionRegistry_Core.frValue>, makeDist) => {
-  let array = inputs->E.A.unsafe_get(0)->Prepare.ToValueArray.Array.openA
+  let array = inputs->getOrError(0)->E.R.bind(Prepare.ToValueArray.Array.openA)
   let xyCoords =
     array->E.R.bind(xyCoords =>
       xyCoords
@@ -72,7 +72,7 @@ let registry = [
     ~name="Declaration",
     ~definitions=[
       FnDefinition.make(~name="declareFn", ~inputs=[Declaration.frType], ~run=(inputs, _) => {
-        inputs->E.A.unsafe_get(0)->Declaration.fromExpressionValue
+        inputs->getOrError(0)->E.R.bind(Declaration.fromExpressionValue)
       }),
     ],
   ),
@@ -134,4 +134,247 @@ let registry = [
     ~name="Bernoulli",
     ~definitions=[OneArgDist.make("bernoulli", SymbolicDist.Bernoulli.make)],
   ),
+  Function.make(~name="Floor", ~definitions=[NumberToNumber.make("floor", Js.Math.floor_float)]),
+  Function.make(~name="Ceiling", ~definitions=[NumberToNumber.make("ceil", Js.Math.ceil_float)]),
+  Function.make(
+    ~name="Absolute Value",
+    ~definitions=[NumberToNumber.make("abs", Js.Math.abs_float)],
+  ),
+  Function.make(~name="Exponent", ~definitions=[NumberToNumber.make("exp", Js.Math.exp)]),
+  Function.make(~name="Log", ~definitions=[NumberToNumber.make("log", Js.Math.log)]),
+  Function.make(~name="Log Base 10", ~definitions=[NumberToNumber.make("log10", Js.Math.log10)]),
+  Function.make(~name="Log Base 2", ~definitions=[NumberToNumber.make("log2", Js.Math.log2)]),
+  Function.make(~name="Round", ~definitions=[NumberToNumber.make("round", Js.Math.round)]),
+  Function.make(
+    ~name="Sum",
+    ~definitions=[ArrayNumberDist.make("sum", r => r->E.A.Floats.sum->Wrappers.evNumber->Ok)],
+  ),
+  Function.make(
+    ~name="Product",
+    ~definitions=[
+      ArrayNumberDist.make("product", r => r->E.A.Floats.product->Wrappers.evNumber->Ok),
+    ],
+  ),
+  Function.make(
+    ~name="Min",
+    ~definitions=[ArrayNumberDist.make("min", r => r->E.A.Floats.min->Wrappers.evNumber->Ok)],
+  ),
+  Function.make(
+    ~name="Max",
+    ~definitions=[ArrayNumberDist.make("max", r => r->E.A.Floats.max->Wrappers.evNumber->Ok)],
+  ),
+  Function.make(
+    ~name="Mean",
+    ~definitions=[ArrayNumberDist.make("mean", r => r->E.A.Floats.mean->Wrappers.evNumber->Ok)],
+  ),
+  Function.make(
+    ~name="Geometric Mean",
+    ~definitions=[
+      ArrayNumberDist.make("geomean", r => r->E.A.Floats.geomean->Wrappers.evNumber->Ok),
+    ],
+  ),
+  Function.make(
+    ~name="Standard Deviation",
+    ~definitions=[ArrayNumberDist.make("stdev", r => r->E.A.Floats.stdev->Wrappers.evNumber->Ok)],
+  ),
+  Function.make(
+    ~name="Variance",
+    ~definitions=[
+      ArrayNumberDist.make("variance", r => r->E.A.Floats.stdev->Wrappers.evNumber->Ok),
+    ],
+  ),
+  Function.make(
+    ~name="First",
+    ~definitions=[
+      ArrayNumberDist.make2("first", r =>
+        r->E.A.first |> E.O.toResult(impossibleError) |> E.R.fmap(Wrappers.evNumber)
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Last",
+    ~definitions=[
+      ArrayNumberDist.make2("last", r =>
+        r->E.A.last |> E.O.toResult(impossibleError) |> E.R.fmap(Wrappers.evNumber)
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Sort",
+    ~definitions=[
+      ArrayNumberDist.make("sort", r =>
+        r->E.A.Floats.sort->E.A2.fmap(Wrappers.evNumber)->Wrappers.evArray->Ok
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Reverse",
+    ~definitions=[
+      ArrayNumberDist.make("reverse", r =>
+        r->Belt_Array.reverse->E.A2.fmap(Wrappers.evNumber)->Wrappers.evArray->Ok
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Cumulative Sum",
+    ~definitions=[
+      ArrayNumberDist.make("cumsum", r =>
+        r->E.A.Floats.cumsum->E.A2.fmap(Wrappers.evNumber)->Wrappers.evArray->Ok
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Cumulative Prod",
+    ~definitions=[
+      ArrayNumberDist.make("cumprod", r =>
+        r->E.A.Floats.cumsum->E.A2.fmap(Wrappers.evNumber)->Wrappers.evArray->Ok
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Diff",
+    ~definitions=[
+      ArrayNumberDist.make("diff", r =>
+        r->E.A.Floats.diff->E.A2.fmap(Wrappers.evNumber)->Wrappers.evArray->Ok
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Dict.merge",
+    ~definitions=[
+      FnDefinition.make(
+        ~name="merge",
+        ~inputs=[FRTypeDict(FRTypeAny), FRTypeDict(FRTypeAny)],
+        ~run=(inputs, _) => {
+          switch inputs {
+          | [FRValueDict(d1), FRValueDict(d2)] => {
+              let newDict =
+                E.Dict.concat(d1, d2) |> Js.Dict.map((. r) =>
+                  FunctionRegistry_Core.FRType.matchReverse(r)
+                )
+              newDict->Wrappers.evRecord->Ok
+            }
+          | _ => Error(impossibleError)
+          }
+        },
+      ),
+    ],
+  ),
+  //TODO: Make sure that two functions can't have the same name. This causes chaos elsewhere.
+  Function.make(
+    ~name="Dict.mergeMany",
+    ~definitions=[
+      FnDefinition.make(~name="mergeMany", ~inputs=[FRTypeArray(FRTypeDict(FRTypeAny))], ~run=(
+        inputs,
+        _,
+      ) =>
+        inputs
+        ->Prepare.ToTypedArray.dicts
+        ->E.R2.fmap(E.Dict.concatMany)
+        ->E.R2.fmap(Js.Dict.map((. r) => FunctionRegistry_Core.FRType.matchReverse(r)))
+        ->E.R2.fmap(Wrappers.evRecord)
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Dict.keys",
+    ~definitions=[
+      FnDefinition.make(~name="keys", ~inputs=[FRTypeDict(FRTypeAny)], ~run=(inputs, _) =>
+        switch inputs {
+        | [FRValueDict(d1)] => Js.Dict.keys(d1)->E.A2.fmap(Wrappers.evString)->Wrappers.evArray->Ok
+        | _ => Error(impossibleError)
+        }
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Dict.values",
+    ~definitions=[
+      FnDefinition.make(~name="values", ~inputs=[FRTypeDict(FRTypeAny)], ~run=(inputs, _) =>
+        switch inputs {
+        | [FRValueDict(d1)] =>
+          Js.Dict.values(d1)
+          ->E.A2.fmap(FunctionRegistry_Core.FRType.matchReverse)
+          ->Wrappers.evArray
+          ->Ok
+        | _ => Error(impossibleError)
+        }
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Dict.toList",
+    ~definitions=[
+      FnDefinition.make(~name="dictToList", ~inputs=[FRTypeDict(FRTypeAny)], ~run=(inputs, _) =>
+        switch inputs {
+        | [FRValueDict(dict)] =>
+          dict
+          ->Js.Dict.entries
+          ->E.A2.fmap(((key, value)) =>
+            Wrappers.evArray([
+              Wrappers.evString(key),
+              FunctionRegistry_Core.FRType.matchReverse(value),
+            ])
+          )
+          ->Wrappers.evArray
+          ->Ok
+        | _ => Error(impossibleError)
+        }
+      ),
+    ],
+  ),
+  Function.make(
+    ~name="Dict.fromList",
+    ~definitions=[
+      FnDefinition.make(~name="dictFromList", ~inputs=[FRTypeArray(FRTypeArray(FRTypeAny))], ~run=(
+        inputs,
+        _,
+      ) => {
+        let convertInternalItems = items =>
+          items
+          ->E.A2.fmap(item => {
+            switch item {
+            | [FRValueString(string), value] =>
+              (string, FunctionRegistry_Core.FRType.matchReverse(value))->Ok
+            | _ => Error(impossibleError)
+            }
+          })
+          ->E.A.R.firstErrorOrOpen
+          ->E.R2.fmap(Js.Dict.fromArray)
+          ->E.R2.fmap(Wrappers.evRecord)
+        inputs->getOrError(0)->E.R.bind(Prepare.ToValueArray.Array.arrayOfArrays)
+          |> E.R2.bind(convertInternalItems)
+      }),
+    ],
+  ),
+  Function.make(
+    ~name="List.make",
+    ~definitions=[
+      //Todo: If the second item is a function with no args, it could be nice to run this function and return the result.
+      FnDefinition.make(~name="listMake", ~inputs=[FRTypeNumber, FRTypeAny], ~run=(inputs, _) => {
+        switch inputs {
+        | [FRValueNumber(number), value] =>
+          Belt.Array.make(E.Float.toInt(number), value)
+          ->E.A2.fmap(FunctionRegistry_Core.FRType.matchReverse)
+          ->Wrappers.evArray
+          ->Ok
+        | _ => Error(impossibleError)
+        }
+      }),
+    ],
+  ),
+  Function.make(
+    ~name="upTo",
+    ~definitions=[
+      FnDefinition.make(~name="upTo", ~inputs=[FRTypeNumber, FRTypeNumber], ~run=(inputs, _) =>
+        inputs
+        ->Prepare.ToValueTuple.twoNumbers
+        ->E.R2.fmap(((low, high)) =>
+          E.A.Floats.range(low, high, (high -. low +. 1.0)->E.Float.toInt)
+          ->E.A2.fmap(Wrappers.evNumber)
+          ->Wrappers.evArray
+        )
+      ),
+    ],
+  ),
 ]

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

@@ -207,6 +207,7 @@ module Float = {
   let toFixed = Js.Float.toFixed
   let toString = Js.Float.toString
   let isFinite = Js.Float.isFinite
+  let toInt = Belt.Float.toInt
 }
 
 module I = {
@@ -539,6 +540,7 @@ module A = {
   let hasBy = (r, fn) => Belt.Array.getBy(r, fn) |> O.isSome
   let fold_left = Array.fold_left
   let fold_right = Array.fold_right
+  let concat = Belt.Array.concat
   let concatMany = Belt.Array.concatMany
   let keepMap = Belt.Array.keepMap
   let slice = Belt.Array.slice
@@ -720,6 +722,7 @@ module A = {
     let variance = Jstat.variance
     let stdev = Jstat.stdev
     let sum = Jstat.sum
+    let product = Jstat.product
     let random = Js.Math.random_int
 
     let floatCompare: (float, float) => int = compare
@@ -748,6 +751,9 @@ module A = {
     let diff = (t: t): array<float> =>
       Belt.Array.zipBy(t, Belt.Array.sliceToEnd(t, 1), (left, right) => right -. left)
 
+    let cumsum = (t: t): array<float> => accumulate((a, b) => a +. b, t)
+    let cumProd = (t: t): array<float> => accumulate((a, b) => a *. b, t)
+
     exception RangeError(string)
     let range = (min: float, max: float, n: int): array<float> =>
       switch n {
@@ -872,4 +878,7 @@ module Dict = {
   let get = Js.Dict.get
   let keys = Js.Dict.keys
   let fromArray = Js.Dict.fromArray
+  let toArray = Js.Dict.entries
+  let concat = (a, b) => A.concat(toArray(a), toArray(b))->fromArray
+  let concatMany = ts => ts->A2.fmap(toArray)->A.concatMany->fromArray
 }