diff --git a/packages/squiggle-lang/src/rescript/Distributions/SymbolicDist/SymbolicDist.res b/packages/squiggle-lang/src/rescript/Distributions/SymbolicDist/SymbolicDist.res index f33fc450..762e5fd5 100644 --- a/packages/squiggle-lang/src/rescript/Distributions/SymbolicDist/SymbolicDist.res +++ b/packages/squiggle-lang/src/rescript/Distributions/SymbolicDist/SymbolicDist.res @@ -124,6 +124,19 @@ module Beta = { let sample = (t: t) => Jstat.Beta.sample(t.alpha, t.beta) let mean = (t: t) => Ok(Jstat.Beta.mean(t.alpha, t.beta)) let toString = ({alpha, beta}: t) => j`Beta($alpha,$beta)` + + let fromMeanAndSampleSize = (mean, sampleSize) => { + // https://en.wikipedia.org/wiki/Beta_distribution#Mean_and_sample_size + let alpha = mean *. sampleSize + let beta = (1.0 -. mean) *. sampleSize + alpha->make(beta) + } + + let fromMeanAndStdev = (mean, stdev) => { + let var = stdev *. stdev + let sampleSize = mean *. (1.0 -. mean) /. var + mean->fromMeanAndSampleSize(sampleSize) + } } module Lognormal = { diff --git a/packages/squiggle-lang/src/rescript/FunctionRegistry/FunctionRegistry_Library.res b/packages/squiggle-lang/src/rescript/FunctionRegistry/FunctionRegistry_Library.res index 56315860..a37b8dc4 100644 --- a/packages/squiggle-lang/src/rescript/FunctionRegistry/FunctionRegistry_Library.res +++ b/packages/squiggle-lang/src/rescript/FunctionRegistry/FunctionRegistry_Library.res @@ -117,8 +117,12 @@ lognormal({mean: 5, stdev: 2})`, ), Function.make( ~name="Beta", - ~examples=`beta(20, 25)`, - ~definitions=[TwoArgDist.make("beta", twoArgs(SymbolicDist.Beta.make))], + ~examples=`beta(20, 25) +beta({mean: 0.39, stdev: 0.1})`, + ~definitions=[ + TwoArgDist.make("beta", twoArgs(SymbolicDist.Beta.make)), + TwoArgDist.makeRecordMeanStdev("beta", twoArgs(SymbolicDist.Beta.fromMeanAndStdev)), + ], (), ), Function.make(