Add geometric distribution functions (#75)

SamuelDixxon · web-flow · commit 2602ac18b14f · 2023-01-08T15:07:30.000+07:00
diff --git a/README.md b/README.md
@@ -129,6 +129,9 @@ func Trimean(input Float64Data) (float64, error) {}
 func VarP(input Float64Data) (sdev float64, err error) {}
 func VarS(input Float64Data) (sdev float64, err error) {}
 func Variance(input Float64Data) (sdev float64, err error) {}
+func ProbGeom(a int, b int, p float64) (prob float64, err error) {}
+func ExpGeom(p float64) (exp float64, err error) {}
+func VarGeom(p float64) (exp float64, err error) {}
 
 type Coordinate struct {
     X, Y float64
diff --git a/examples/main.go b/examples/main.go
@@ -169,4 +169,21 @@ func main() {
 	e, _ = stats.Entropy([]float64{1.1, 2.2, 3.3})
 	fmt.Println(e)
 	// Output: 1.0114042647073518
+
+	p := 0.5
+	begin := 1
+	end := 2
+	chance, _ := stats.ProbGeom(begin, end, p)
+	fmt.Println(chance)
+	// Output: 0.25
+
+	prob1 := 0.5
+	exp, _ := stats.ExpGeom(prob1)
+	fmt.Println(exp)
+	// Output: 
+	
+	prob2:= 0.5
+	vari, _ := stats.VarGeom(prob2)
+	fmt.Println(vari)
+	// Output: 2
 }
diff --git a/geometric_distribution.go b/geometric_distribution.go
@@ -0,0 +1,42 @@
+package stats
+
+import (
+	"math"
+)
+
+// ProbGeom generates the probability for a geometric random variable
+// with parameter p to achieve success in the interval of [a, b] trials
+// See https://en.wikipedia.org/wiki/Geometric_distribution for more information
+func ProbGeom(a int, b int, p float64) (prob float64, err error) {
+	if (a > b) || (a < 1) {
+		return math.NaN(), ErrBounds
+	}
+
+	prob = 0
+	q := 1 - p // probability of failure
+
+	for k := a + 1; k <= b; k++ {
+		prob = prob + p*math.Pow(q, float64(k-1))
+	}
+
+	return prob, nil
+}
+
+// ProbGeom generates the expectation or average number of trials
+// for a geometric random variable with parameter p
+func ExpGeom(p float64) (exp float64, err error) {
+	if (p > 1) || (p < 0) {
+		return math.NaN(), ErrNegative
+	}
+
+	return 1 / p, nil
+}
+
+// ProbGeom generates the variance for number for a
+// geometric random variable with parameter p
+func VarGeom(p float64) (exp float64, err error) {
+	if (p > 1) || (p < 0) {
+		return math.NaN(), ErrNegative
+	}
+	return (1 - p) / math.Pow(p, 2), nil
+}
diff --git a/geometric_distribution_test.go b/geometric_distribution_test.go
@@ -0,0 +1,92 @@
+package stats_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/montanaflynn/stats"
+)
+
+func ExampleProbGeom() {
+	p := 0.5
+	a := 1
+	b := 2
+	chance, _ := stats.ProbGeom(a, b, p)
+	fmt.Println(chance)
+	// Output: 0.25
+}
+
+func TestProbGeomLarge(t *testing.T) {
+	p := 0.5
+	a := 1
+	b := 10000
+	chance, _ := stats.ProbGeom(a, b, p)
+	fmt.Println(chance)
+	// Output: 1
+}
+
+func TestErrBoundsProbGeom(t *testing.T) {
+	p := 0.5
+	a := -1
+	b := 4
+	chance, err := stats.ProbGeom(a, b, p)
+	if err == nil {
+		t.Errorf("Did not return an error when expected")
+	}
+	fmt.Println(chance)
+	// Output: NaN
+}
+
+func ExampleExpGeom() {
+	p := 0.5
+	exp, _ := stats.ExpGeom(p)
+	fmt.Println(exp)
+	// Output: 2
+}
+
+func TestExpGeom(t *testing.T) {
+	p := 0.5
+	exp, err := stats.ExpGeom(p)
+	if err != nil {
+		t.Errorf("Returned an error when not expected")
+	}
+	fmt.Println(exp)
+	// Output: 2
+}
+
+func TestErrExpGeom(t *testing.T) {
+	p := -1.0
+	exp, err := stats.ExpGeom(p)
+	if err == nil {
+		t.Errorf("Expected Error")
+	}
+	fmt.Println(exp)
+	// Output: NaN
+}
+
+func ExampleVarGeom() {
+	p := 0.5
+	vari, _ := stats.VarGeom(p)
+	fmt.Println(vari)
+	// Output: 2
+}
+
+func TestVarGeom(t *testing.T) {
+	p := 0.25
+	vari, err := stats.VarGeom(p)
+	if err != nil {
+		t.Errorf("Returned an error when not expected")
+	}
+	fmt.Println(vari)
+	// Output: 12.0
+}
+
+func TestErrVarGeom(t *testing.T) {
+	p := -1.0
+	vari, err := stats.VarGeom(p)
+	if err == nil {
+		t.Errorf("Expected Erorr")
+	}
+	fmt.Println(vari)
+	// Output: NaN
+}
