Add chi square func.

refs #17562

Arvados-DCO-1.1-Signed-off-by: Tom Clegg <tom@curii.com>

diff --git a/chisquare.go b/chisquare.go
new file mode 100644 (file)
index 0000000..e1c8bff
--- /dev/null
+++ b/chisquare.go
@@ -0,0 +1,42 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
+
+import (
+       "golang.org/x/exp/rand"
+       "gonum.org/v1/gonum/stat/distuv"
+)
+
+var chisquared = distuv.ChiSquared{K: 1, Src: rand.NewSource(rand.Uint64())}
+
+func pvalue(a, b []bool) float64 {
+       //     !b        b
+       // !a  tab[0]    tab[1]
+       // a   tab[2]    tab[3]
+       tab := make([]int, 4)
+       for ai, aval := range []bool{false, true} {
+               for bi, bval := range []bool{false, true} {
+                       obs := 0
+                       for i := range a {
+                               if a[i] == aval && b[i] == bval {
+                                       obs++
+                               }
+                       }
+                       tab[ai*2+bi] = obs
+               }
+       }
+       var sum float64
+       for ai := 0; ai < 2; ai++ {
+               for bi := 0; bi < 2; bi++ {
+                       rowtotal := tab[ai*2] + tab[ai*2+1]
+                       coltotal := tab[bi] + tab[2+bi]
+                       exp := float64(rowtotal) * float64(coltotal) / float64(len(a))
+                       obs := tab[ai*2+bi]
+                       d := float64(obs) - exp
+                       sum += (d * d) / exp
+               }
+       }
+       return 1 - chisquared.CDF(sum)
+}

diff --git a/chisquare_test.go b/chisquare_test.go
new file mode 100644 (file)
index 0000000..3a59671
--- /dev/null
+++ b/chisquare_test.go
@@ -0,0 +1,35 @@
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
+
+import (
+       "fmt"
+
+       "gopkg.in/check.v1"
+)
+
+type pvalueSuite struct{}
+
+var _ = check.Suite(&pvalueSuite{})
+
+func (s *pvalueSuite) TestPvalue(c *check.C) {
+       a := make([]bool, 54)
+       b := make([]bool, 54)
+       for i := 0; i < 25; i++ {
+               a[i] = true
+               b[i] = true
+       }
+       for i := 25; i < 31; i++ {
+               a[i] = true
+       }
+       for i := 31; i < 39; i++ {
+               b[i] = true
+       }
+       c.Check(fmt.Sprintf("%.7f", pvalue(a, b)), check.Equals, "0.0006297")
+       for i := range a {
+               a[i] = !a[i]
+       }
+       c.Check(fmt.Sprintf("%.7f", pvalue(a, b)), check.Equals, "0.0006297")
+}

diff --git a/go.mod b/go.mod
index bc812fedcd2ea5dc084f86588588fbc501a40e48..bd7a6cd2dd256cb4cad11d1faf5608c1ab3d07ed 100644 (file)
--- a/go.mod
+++ b/go.mod
@@ -21,6 +21,7 @@ require (
        github.com/sirupsen/logrus v1.6.0
        github.com/spaolacci/murmur3 v1.1.0 // indirect
        golang.org/x/crypto v0.0.0-20200510223506-06a226fb4e37
+       golang.org/x/exp v0.0.0-20190125153040-c74c464bbbf2
        golang.org/x/net v0.0.0-20200520182314-0ba52f642ac2
        golang.org/x/sys v0.0.0-20200519105757-fe76b779f299 // indirect
        gonum.org/v1/gonum v0.8.1