diff --git a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Danger.res b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Danger.res index 59ae77cb..e8dbfbbc 100644 --- a/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Danger.res +++ b/packages/squiggle-lang/src/rescript/FunctionRegistry/Library/FR_Danger.res @@ -4,6 +4,22 @@ open FunctionRegistry_Helpers let nameSpace = "Danger" let requiresNamespace = true +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) + }, + (), + ) +} + module TwoNumbersToNumber = { let make = (name, fn) => FnDefinition.make( @@ -28,18 +44,52 @@ module ThreeNumbersToNumber = { ) } +module Internals = { + let factorial = Stdlib.Math.factorial + let choose = ((n, k)) => factorial(n) /. (factorial(n -. k) *. factorial(k)) + let pow = (base, exp) => Js.Math.pow_float(~base, ~exp) + let binomial = ((n, k, p)) => choose((n, k)) *. pow(p, k) *. pow(1.0 -. p, n -. k) +} + let library = [ Function.make( ~name="laplace", ~nameSpace, ~requiresNamespace, ~output=EvtNumber, - ~examples=[`laplace(1, 20)`], + ~examples=[`Danger.laplace(1, 20)`], ~definitions=[ TwoNumbersToNumber.make("laplace", ((successes, trials)) => - (successes +. 1.0) /. (trials +. 1.0) + (successes +. 1.0) /. (trials +. 2.0) ), ], (), ), + Function.make( + ~name="factorial", + ~nameSpace, + ~requiresNamespace, + ~output=EvtNumber, + ~examples=[`Danger.factorial(20)`], + ~definitions=[NumberToNumber.make("factorial", Internals.factorial)], + (), + ), + Function.make( + ~name="choose", + ~nameSpace, + ~requiresNamespace, + ~output=EvtNumber, + ~examples=[`Danger.choose(1, 20)`], + ~definitions=[TwoNumbersToNumber.make("choose", Internals.choose)], + (), + ), + Function.make( + ~name="binomial", + ~nameSpace, + ~requiresNamespace, + ~output=EvtNumber, + ~examples=[`Danger.binomial(1, 20, 0.5)`], + ~definitions=[ThreeNumbersToNumber.make("binomial", Internals.binomial)], + (), + ), ] diff --git a/packages/squiggle-lang/src/rescript/Utility/Stdlib.res b/packages/squiggle-lang/src/rescript/Utility/Stdlib.res index ee43e681..249b93d7 100644 --- a/packages/squiggle-lang/src/rescript/Utility/Stdlib.res +++ b/packages/squiggle-lang/src/rescript/Utility/Stdlib.res @@ -47,3 +47,7 @@ module Random = { @module external sample: (array, sampleArgs) => array = "@stdlib/random/sample" let sample = sample } + +module Math = { + @module external factorial: float => float = "@stdlib/math/base/special/factorial" +}