Merge pull request #785 from quantified-uncertainty/beta-from-mean-stdev

`beta({mean: 0.39, stdev: 0.1})` #497

packages/squiggle-lang/__tests__/Distributions/Symbolic_test.res

@@ -77,6 +77,23 @@ describe("(Symbolic) mean", () => {
     meanValue->unpackFloat->expect->ExpectJs.toBeFalsy
   })
 
+  testAll(
+    "of beta distributions from mean and standard dev",
+    list{(0.39, 0.1), (0.08, 0.1), (0.8, 0.3)},
+    tup => {
+      let (mean, stdev) = tup
+      let betaDistribution = SymbolicDist.Beta.fromMeanAndStdev(mean, stdev)
+      let meanValue =
+        betaDistribution->E.R2.fmap(d =>
+          run(FromDist(ToFloat(#Mean), d->DistributionTypes.Symbolic))
+        )
+      switch meanValue {
+      | Ok(value) => value->unpackFloat->expect->toBeCloseTo(mean)
+      | Error(err) => err->expect->toBe("shouldn't happen")
+      }
+    },
+  )
+
   testAll(
     "of lognormal distributions",
     list{(2.0, 4.0), (1e-7, 1e-2), (-1e6, 10.0), (1e3, -1e2), (-1e8, -1e4), (1e2, 1e-5)},

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

@@ -124,6 +124,26 @@ 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
+    make(alpha, beta)
+  }
+
+  let fromMeanAndStdev = (mean, stdev) => {
+    // https://en.wikipedia.org/wiki/Beta_distribution#Mean_and_variance
+    if !(0.0 < stdev && stdev <= 0.5) {
+      Error("Stdev must be in in between 0 and 0.5.")
+    } else if !(0.0 <= mean && mean <= 1.0) {
+      Error("Mean must be in between 0 and 1.0.")
+    } else {
+      let var = stdev *. stdev
+      let sampleSize = mean *. (1.0 -. mean) /. var -. 1.0
+      fromMeanAndSampleSize(mean, sampleSize)
+    }
+  }
 }
 
 module Lognormal = {

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)`,
+    ~examples=`beta(20, 25)
-    ~definitions=[TwoArgDist.make("beta", twoArgs(SymbolicDist.Beta.make))],
+beta({mean: 0.39, stdev: 0.1})`,
+    ~definitions=[
+      TwoArgDist.make("beta", twoArgs(SymbolicDist.Beta.make)),
+      TwoArgDist.makeRecordMeanStdev("beta", twoArgs(SymbolicDist.Beta.fromMeanAndStdev)),
+    ],
     (),
   ),
   Function.make(

packages/website/docs/Api/DistGeneric.mdx

@@ -70,12 +70,14 @@ uniform(10, 12);
 
 ```
 beta: (distribution|number, distribution|number) => distribution
+beta: (dict<{mean: distribution|number, stdev: distribution|number}>) => distribution
 ```
 
 **Examples**
 
 ```javascript
 beta(20, 25);
+beta({ mean: 0.39, stdev: 0.1 });
 ```
 
 ### cauchy

packages/website/docs/Guides/DistributionCreation.mdx

@@ -259,6 +259,7 @@ with values at 1 and 2. Therefore, this is the same as `mixture(pointMass(1),poi
 ## Beta
 
 `beta(alpha:number, beta:number)`
+`beta({mean: number, stdev: number})`
 
 Creates a [beta distribution](https://en.wikipedia.org/wiki/Beta_distribution) with the given `alpha` and `beta` values. For a good summary of the beta distribution, see [this explanation](https://stats.stackexchange.com/a/47782) on Stack Overflow.
 
@@ -278,6 +279,12 @@ Creates a [beta distribution](https://en.wikipedia.org/wiki/Beta_distribution) w
   <TabItem value="ex5" label="beta(0.8, 0.8)">
     <SquiggleEditor defaultCode="beta(0.8, 0.8)" />
   </TabItem>
+  <TabItem
+    value="from mean and standard deviation"
+    label="beta({mean: 0.39, stdev: 0.1})"
+  >
+    <SquiggleEditor initialSquiggleString="beta({mean: 0.39, stdev: 0.1})" />
+  </TabItem>
 </Tabs>
 
 ### Arguments

packages/website/docs/Introduction.md

@@ -10,7 +10,7 @@ Squiggle is an _estimation language_, and a syntax for _calculating and expressi
 - [Gallery](./Discussions/Gallery)
 - [Squiggle playground](/playground)
 - [Language basics](./Guides/Language)
-- [Squiggle functions source of truth](./docs/Guides/Functions)
+- [Squiggle functions source of truth](./Guides/Functions)
 - [Known bugs](./Discussions/Bugs)
 - [Original lesswrong sequence](https://www.lesswrong.com/s/rDe8QE5NvXcZYzgZ3)
 - [Author your squiggle models as Observable notebooks](https://observablehq.com/@hazelfire/squiggle)