Add bayes classes (#22)

Clean up and organization Python code, add a test suite.

Co-authored-by: Jeremy Dorn <jeremy@jeremydorn.com>

.github/workflows/ci.yml

@@ -48,6 +48,11 @@ jobs:
         with:
           node-version: 12.x
 
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
       - name: Get yarn cache directory
         id: yarn-cache
         run: |
@@ -63,6 +68,9 @@ jobs:
       - name: install dependencies
         run: |
           yarn
+          python -m pip install --upgrade pip
+          pip install numpy scipy pandas
+
         env:
           CI: true
 
@@ -76,10 +84,6 @@ jobs:
           # Run the test suite
           yarn test
 
-          # Cleanup
-          rm -rf packages/front-end/coverage
-          rm -rf packages/back-end/coverage
-
   # Build and publish the commit to docker
   # Runs on pushes to the main branch or when a release is made
   # Only runs if code changed on the front-end or back-end (e.g. skip if only docs changed)

packages/back-end/package.json

@@ -10,7 +10,9 @@
     "dev": "node-dev src/server.ts",
     "build": "yarn build:clean && yarn build:typescript && yarn build:python && yarn build:emails",
     "start": "node dist/server.js",
-    "test": "jest --forceExit --coverage --verbose --detectOpenHandles",
+    "test:node": "jest --forceExit --coverage --verbose --detectOpenHandles",
+    "test:python": "cd src/python && python3 -m unittest discover",
+    "test": "yarn test:python && yarn test:node",
     "type-check": "tsc --pretty --noEmit",
     "cron": "node dist/cron.js",
     "generate-dummy-data": "node --stack-size=8192 ./test/data-generator/data-generator.js",

packages/back-end/src/python/bayesian/dists.py

@@ -0,0 +1,140 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, Tuple
+from warnings import warn
+import numpy as np
+from numpy import ndarray, vectorize
+from scipy.stats import beta, norm, rv_continuous
+from scipy.special import digamma, polygamma, roots_hermitenorm  # , roots_sh_jacobi
+from .orthogonal import roots_sh_jacobi
+
+
+EPSILON = 1e-04
+
+
+class BayesABDist(ABC):
+    dist: rv_continuous
+
+    @staticmethod
+    @abstractmethod
+    def posterior(prior, data):
+        """
+        :type prior: Iterable
+        :type data: Iterable
+        :rtype: Tuple[ndarray, ndarray]
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def moments(par1, par2, log=False):
+        """
+        :type par1: float or ndarray
+        :type par2: float or ndarray
+        :type log: bool
+        :rtype: Tuple[float or ndarray, float or ndarray]
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def gq(n, par1, par2):
+        """
+        :type n: int
+        :type par1: float
+        :type par2: float
+        :rtype: Tuple[ndarray, ndarray]
+        """
+        raise NotImplementedError
+
+    # todo: @vectorize
+    @classmethod
+    def risk(cls, a_par1, a_par2, b_par1, b_par2, n=24):
+        """
+        :type a_par1: float
+        :type a_par2: float
+        :type b_par1: float
+        :type b_par2: float
+        :type n: int
+        :rtype: ndarray
+        """
+        a_nodes, a_weights = cls.gq(n, a_par1, a_par2)
+        b_nodes, b_weights = cls.gq(n, b_par1, b_par2)
+
+        gq = sum(a_nodes * cls.dist.cdf(a_nodes, b_par1, b_par2) * a_weights) + \
+            sum(b_nodes * cls.dist.cdf(b_nodes, a_par1, a_par2) * b_weights)
+        out = gq - cls.dist.mean((a_par1, b_par1), (a_par2, b_par2))
+
+        return out
+
+
+class Beta(BayesABDist):
+    dist = beta
+
+    @staticmethod
+    def posterior(prior, data):
+        a = prior[0] + data[0]
+        b = prior[1] + data[1] - data[0]
+        return a, b
+
+    @staticmethod
+    def moments(par1, par2, log=False):
+        if np.sum(par2 < 0) + np.sum(par1 < 0):
+            raise RuntimeError('params of beta distribution cannot be negative')
+
+        if log:
+            mean = digamma(par1) - digamma(par1 + par2)
+            var = polygamma(1, par1) - polygamma(1, par1 + par2)
+        else:
+            mean = par1 / (par1 + par2)
+            var = par1 * par2 / (np.power(par1 + par2, 2) * (par1 + par2 + 1))
+        return mean, var
+
+    @staticmethod
+    def gq(n, par1, par2):
+        x, w = roots_sh_jacobi(int(n), par1 + par2 - 1, par1, False)
+        return x, w
+
+
+class Norm(BayesABDist):
+    dist = norm
+
+    @staticmethod
+    def posterior(prior, data):
+        inv_var_0 = prior[2] / np.power(prior[1], 2)
+        inv_var_d = data[2] / np.power(data[1], 2)
+        var = 1 / (inv_var_0 + inv_var_d)
+
+        loc = var * (inv_var_0 * prior[0] + inv_var_d * data[0])
+        scale = np.sqrt(var)
+        return loc, scale
+
+    @staticmethod
+    def moments(par1, par2, log=False):
+        if np.sum(par2 < 0):
+            raise RuntimeError('got negative standard deviation.')
+
+        if log:
+            if np.sum(par1 <= 0):
+                raise RuntimeError('got mu <= 0. cannot use log approximation.')
+
+            max_prob = np.max(norm.cdf(0, par1, par2))
+            if max_prob > EPSILON:
+                warn(f'probability of being negative is higher than {EPSILON} (={max_prob}). '
+                     f'log approximation is in-exact', RuntimeWarning)
+
+            mean = np.log(par1)
+            var = np.power(par2 / par1, 2)
+        else:
+            mean = par1
+            var = np.power(par2, 2)
+        return mean, var
+
+    @staticmethod
+    def gq(n, par1, par2):
+        if par2 <= 0:
+            raise RuntimeError('got negative standard deviation.')
+
+        x, w, m = roots_hermitenorm(int(n), True)
+        x = par2 * x + par1
+        w /= m
+        return x, w

packages/back-end/src/python/bayesian/main.py

@@ -1,141 +1,92 @@
-# Medium article inspiration:
-#   https://towardsdatascience.com/how-to-do-bayesian-a-b-testing-fast-41ee00d55be8
-# Original code:
-#   https://github.com/itamarfaran/public-sandbox/tree/master/bayesian_blog
-
-
 import numpy as np
-from scipy.stats import norm, beta
-from scipy.special import digamma, polygamma, roots_hermitenorm
-from orthogonal import roots_sh_jacobi
-import sys
-import json
+from scipy.stats import norm
+from .dists import Beta, Norm
 
 
-def log_beta_mean(a, b): return digamma(a) - digamma(a + b)
-def var_beta_mean(a, b): return polygamma(1, a) - polygamma(1, a + b)
+"""
+Medium article inspiration: 
+    https://towardsdatascience.com/how-to-do-bayesian-a-b-testing-fast-41ee00d55be8
+
+Original code:
+    https://github.com/itamarfaran/public-sandbox/tree/master/bayesian_blog
+"""
 
 
-def beta_gq(n, a, b):
-    x, w, m = roots_sh_jacobi(n, a + b - 1, a, True)
-    w /= m
-    return x, w
+BETA_PRIOR = 1, 1
+NORM_PRIOR = 0, 1, 0
 
 
-def norm_gq(n, loc, scale):
-    x, w, m = roots_hermitenorm(n, True)
-    x = scale * x + loc
-    w /= m
-    return x, w
+def binomial_ab_test(x_a, n_a, x_b, n_b, ccr=.05):
+    alpha_a, beta_a = Beta.posterior(BETA_PRIOR, [x_a, n_a])
+    alpha_b, beta_b = Beta.posterior(BETA_PRIOR, [x_b, n_b])
+
+    mean_a, var_a = Beta.moments(alpha_a, beta_a, log=True)
+    mean_b, var_b = Beta.moments(alpha_b, beta_b, log=True)
+
+    mean_diff = mean_b - mean_a
+    std_diff = np.sqrt(var_a + var_b)
+
+    chance_to_win = norm.sf(0, mean_diff, std_diff)
+    expected = np.exp(mean_diff) - 1
+    ci = np.exp(norm.ppf([ccr / 2, 1 - ccr / 2], mean_diff, std_diff)) - 1
+    risk_beta = Beta.risk(alpha_a, beta_a, alpha_b, beta_b)
+
+    output = {'chance_to_win': chance_to_win,
+              'expected': expected,
+              'ci': ci.tolist(),
+              'uplift': {'dist': 'lognormal',
+                         'mean': mean_diff,
+                         'stddev': std_diff},
+              'risk': risk_beta.tolist()}
+
+    return output
 
 
-def binomial_ab_test(x_a, n_a, x_b, n_b):
-    # Uninformative prior
-    alpha_0, beta_0 = 1, 1
+def gaussian_ab_test(m_a, s_a, n_a, m_b, s_b, n_b, ccr=.05):
+    mu_a, sd_a = Norm.posterior(NORM_PRIOR, [m_a, s_a, n_a])
+    mu_b, sd_b = Norm.posterior(NORM_PRIOR, [m_b, s_b, n_b])
 
-    # Updating prior with data
-    alpha_a = alpha_0 + x_a
-    beta_a = beta_0 + n_a - x_a
+    mean_a, var_a = Norm.moments(mu_a, sd_a, log=True)
+    mean_b, var_b = Norm.moments(mu_b, sd_b, log=True)
 
-    alpha_b = alpha_0 + x_b
-    beta_b = beta_0 + n_b - x_b
+    mean_diff = mean_b - mean_a
+    std_diff = np.sqrt(var_a + var_b)
 
-    # Chance to win
-    d1_beta = norm(
-        loc=beta.mean(alpha_b, beta_b) - beta.mean(alpha_a, beta_a),
-        scale=np.sqrt(beta.var(alpha_b, beta_b) + beta.var(alpha_a, beta_a))
-    )
+    chance_to_win = norm.sf(0, mean_diff, std_diff)
+    expected = np.exp(mean_diff) - 1
+    ci = np.exp(norm.ppf([ccr / 2, 1 - ccr / 2], mean_diff, std_diff)) - 1
+    risk_norm = Norm.risk(mu_a, sd_a, mu_b, sd_b)
 
-    # Credible Interval
-    ci_mean = log_beta_mean(alpha_b, beta_b) - log_beta_mean(alpha_a, beta_a)
-    ci_std = np.sqrt(var_beta_mean(alpha_b, beta_b) +
-                     var_beta_mean(alpha_a, beta_a))
-    d2_beta = norm(
-        loc=ci_mean,
-        scale=ci_std
-    )
+    output = {'chance_to_win': chance_to_win,
+              'expected': expected,
+              'ci': ci.tolist(),
+              'uplift': {'dist': 'lognormal',
+                         'mean': mean_diff,
+                         'stddev': std_diff},
+              'risk': risk_norm.tolist()}
 
-    # Risk
-    nodes_a, weights_a = beta_gq(24, alpha_a, beta_a)
-    nodes_b, weights_b = beta_gq(24, alpha_b, beta_b)
-    gq = sum(nodes_a * beta.cdf(nodes_a, alpha_b, beta_b) * weights_a) + \
-        sum(nodes_b * beta.cdf(nodes_b, alpha_a, beta_a) * weights_b)
-    risk_beta = gq - beta.mean((alpha_a, alpha_b), (beta_a, beta_b))
-
-    return {
-        "chance_to_win": d1_beta.sf(0),
-        "expected": (np.exp(d2_beta.ppf(0.5)) - 1),
-        "ci": (np.exp(d2_beta.ppf((.025, .975))) - 1).tolist(),
-        "uplift": {
-            "dist": "lognormal",
-            "mean": ci_mean,
-            "stddev": ci_std,
-        },
-        "risk": risk_beta.tolist()
-    }
-
-
-def gaussian_ab_test(n_a, m_a, s_a, n_b, m_b, s_b):
-    # Uninformative prior
-    mu0, s0, n0 = 0, 1, 0
-
-    # Update the prior
-    inv_vars = n0 / np.power(s0, 2), n_a / np.power(s_a, 2)
-    mu_a = np.average((mu0, m_a), weights=inv_vars)
-    sd_a = 1 / np.sqrt(np.sum(inv_vars))
-
-    inv_vars = n0 / np.power(s0, 2), n_b / np.power(s_b, 2)
-    mu_b = np.average((mu0, m_b), weights=inv_vars)
-    sd_b = 1 / np.sqrt(np.sum(inv_vars))
-
-    # Chance to win
-    d1_norm = norm(loc=mu_b - mu_a, scale=np.sqrt(sd_a ** 2 + sd_b ** 2))
-
-    # Credible interval
-    ci_mean = np.log(mu_b) - np.log(mu_a)
-    ci_std = np.sqrt((sd_a / mu_a)**2 + (sd_b / mu_b)**2)
-    d2_norm = norm(loc=ci_mean, scale=ci_std)
-
-    # Risk
-    nodes_a, weights_a = norm_gq(24, mu_a, sd_a)
-    nodes_b, weights_b = norm_gq(24, mu_b, sd_b)
-
-    gq = sum(nodes_a * norm.cdf(nodes_a, mu_b, sd_b) * weights_a) + \
-        sum(nodes_b * norm.cdf(nodes_b, mu_a, sd_a) * weights_b)
-    risk_norm = gq - norm.mean((mu_a, mu_b))
-
-    return {
-        "chance_to_win": d1_norm.sf(0),
-        "expected": (np.exp(d2_norm.ppf(0.5)) - 1),
-        "ci": (np.exp(d2_norm.ppf((.025, .975))) - 1).tolist(),
-        "uplift": {
-            "dist": "lognormal",
-            "mean": ci_mean,
-            "stddev": ci_std,
-        },
-        "risk": risk_norm.tolist()
-    }
+    return output
 
 
 # python main.py binomial \
 #   '{"users":[1283,1321],"count":[254,289],"mean":[52.3,14.1],"stddev":[14.1,13.7]}'
 # python main.py normal \
 #   '{"users":[1283,1321],"count":[254,289],"mean":[52.3,14.1],"stddev":[14.1,13.7]}'
+
 if __name__ == '__main__':
+    import json
+    import sys
+
     metric = sys.argv[1]
     data = json.loads(sys.argv[2])
 
-    x_a, x_b = data["count"]
-    n_a, n_b = data["users"]
-    m_a, m_b = data["mean"]
-    s_a, s_b = data["stddev"]
+    xa, xb = data['count']
+    na, nb = data['users']
+    ma, mb = data['mean']
+    sa, sb = data['stddev']
 
     if metric == 'binomial':
-        print(json.dumps(binomial_ab_test(
-            x_a, n_a, x_b, n_b
-        )))
+        print(json.dumps(binomial_ab_test(x_a=xa, n_a=na, x_b=xb, n_b=nb)))
 
-    else:
-        print(json.dumps(gaussian_ab_test(
-            n_a, m_a, s_a, n_b, m_b, s_b
-        )))
+    else:  # todo: should be elif
+        print(json.dumps(gaussian_ab_test(m_a=ma, s_a=sa, n_a=na, m_b=mb, s_b=sb, n_b=nb)))

packages/back-end/src/python/test/test_dists.py

@@ -0,0 +1,143 @@
+from functools import partial
+from unittest import TestCase, main as unittest_main
+
+import numpy as np
+import pandas as pd
+from scipy.special import digamma
+from scipy.stats import beta, norm
+from bayesian.dists import Beta, Norm
+
+
+DECIMALS = 5
+round_ = partial(np.round, decimals=DECIMALS)
+
+
+def roundsum(x, decimals=DECIMALS):
+    return np.round(np.sum(x), decimals=decimals)
+
+
+class TestBeta(TestCase):
+    def test_posterior(self):
+        prior = 1, 1
+        data = 1, 2
+        result = Beta.posterior(prior, data)
+        outcome = (2, 2)
+
+        for res, out in zip(result, outcome):
+            self.assertEqual(res, out)
+
+        prior = 1, 1
+        data = pd.Series([1, 10]), pd.Series([2, 20])
+        result = Beta.posterior(prior, data)
+        outcome = pd.Series([2, 11]), pd.Series([2, 11])
+
+        for res, out in zip(result, outcome):
+            pd.testing.assert_series_equal(res, out)
+
+        prior = pd.Series([1, 2]), pd.Series([1, 3])
+        data = pd.Series([1, 10]), pd.Series([2, 20])
+        result = Beta.posterior(prior, data)
+        outcome = pd.Series([2, 12]), pd.Series([2, 13])
+
+        for res, out in zip(result, outcome):
+            pd.testing.assert_series_equal(res, out)
+
+    def test_moments(self):
+        pars = 12, 745
+        result = Beta.moments(*pars)
+        expected = beta.mean(*pars), beta.var(*pars)
+        for res, out in zip(result, expected):
+            self.assertEqual(round_(res), round_(out))
+
+        pars = 12, 745
+        result = Beta.moments(*pars, log=True)
+        mean = beta.expect(np.log, pars)
+        var = beta.expect(lambda x: np.log(x) ** 2, pars) - mean ** 2
+        expected = mean, var
+        for res, out in zip(result, expected):
+            self.assertEqual(round_(res), round_(out))
+
+        pars = np.array([12, 745]), np.array([745, 12])
+        result = Beta.moments(*pars)
+        expected = beta.mean(*pars), beta.var(*pars)
+        for res, out in zip(result, expected):
+            np.testing.assert_array_almost_equal(res, out)
+
+        self.assertRaises(RuntimeError, Beta.moments, 1, -1)
+        self.assertRaises(RuntimeError, Beta.moments, -1, 1)
+        self.assertRaises(RuntimeError, Beta.moments, -1, -1)
+
+    def test_gq(self):
+        test_cases = zip([10, 100, 500, 1000, 10000],
+                         [10000, 1000, 500, 100, 10])
+        for a, b in test_cases:
+            x, w = Beta.gq(24, a, b)
+            for p in range(8):
+                self.assertEqual(roundsum(x ** p * w), roundsum(beta.moment(p, a, b)))
+            self.assertEqual(roundsum(np.log(x) * w), roundsum(digamma(a) - digamma(a + b)))
+
+
+class TestNorm(TestCase):
+    def test_posterior(self):
+        prior = 0, 1, 10
+        data = 12, 1, 10
+        result = Norm.posterior(prior, data)
+        outcome = (6, np.sqrt(1/20))
+
+        for res, out in zip(result, outcome):
+            self.assertEqual(res, out)
+
+        prior = 0, 1, 10
+        data = pd.Series([12, 100]), pd.Series([1, np.sqrt(2)]), pd.Series([10, 20])
+        result = Norm.posterior(prior, data)
+        outcome = pd.Series([6., 50.]), pd.Series([np.sqrt(1/20), np.sqrt(1/20)])
+
+        for res, out in zip(result, outcome):
+            pd.testing.assert_series_equal(res, out)
+
+        prior = pd.Series([0, 100]), pd.Series([1, np.sqrt(2)]), pd.Series([10, 20])
+        data = pd.Series([12, 100]), pd.Series([1, np.sqrt(2)]), pd.Series([10, 20])
+        result = Norm.posterior(prior, data)
+        outcome = pd.Series([6., 100.]), pd.Series([np.sqrt(1/20), np.sqrt(1/20)])
+
+        for res, out in zip(result, outcome):
+            pd.testing.assert_series_equal(res, out)
+
+    def test_moments(self):
+        pars = 10, 100
+        result = Norm.moments(*pars)
+        expected = norm.mean(*pars), norm.var(*pars)
+        for res, out in zip(result, expected):
+            self.assertEqual(round_(res), round_(out))
+
+        pars = 100, 10
+        result = Norm.moments(*pars, log=True)
+        expected = np.log(100), (10 / 100) ** 2
+        for res, out in zip(result, expected):
+            self.assertEqual(round_(res), round_(out))
+
+        pars = np.array([10, 100]), np.array([100, 10])
+        result = Norm.moments(*pars)
+        expected = norm.mean(*pars), norm.var(*pars)
+        for res, out in zip(result, expected):
+            np.testing.assert_array_almost_equal(res, out)
+
+        self.assertWarns(RuntimeWarning, Norm.moments, 0.1, 1, log=True)
+        self.assertRaises(RuntimeError, Norm.moments, 0, 1, log=True)
+        self.assertRaises(RuntimeError, Norm.moments, -1, 1, log=True)
+        self.assertRaises(RuntimeError, Norm.moments, 1, -1)
+
+    def test_gq(self):
+        test_cases = zip([0, -2, 2, 10],
+                         [.01, 1, 4, .0001])
+        for loc, scale in test_cases:
+            x, w = Norm.gq(24, loc, scale)
+            for p in range(8):
+                self.assertEqual(roundsum(x ** p * w), roundsum(norm.moment(p, loc, scale)))
+
+        self.assertRaises(RuntimeError, Norm.gq, 24, 0, 0)
+        self.assertRaises(RuntimeError, Norm.gq, 24, 0, -1)
+
+
+if __name__ == '__main__':
+    unittest_main()

packages/back-end/src/python/test/test_main.py

@@ -0,0 +1,51 @@
+from functools import partial
+from unittest import TestCase, main as unittest_main
+import numpy as np
+from bayesian.main import binomial_ab_test, gaussian_ab_test
+
+DECIMALS = 5
+round_ = partial(np.round, decimals=DECIMALS)
+
+
+class TestBinom(TestCase):
+    def test_binomial_ab_test(self):
+        result = binomial_ab_test(49, 100, 51, 100)
+        expected = {'chance_to_win': 0.61052,
+                    'expected': 0.0404,
+                    'ci': None,
+                    'uplift': None,
+                    'risk': [0.03872, 0.01912]}
+
+        for key in expected.keys():
+            ex = expected[key]
+            res = result[key]
+
+            if ex is None:
+                continue
+            res = [round_(x) for x in res] if isinstance(res, list) else round_(res)
+            ex = [round_(x) for x in ex] if isinstance(ex, list) else round_(ex)
+            self.assertEqual(res, ex)
+
+
+class TestNorm(TestCase):
+    def test_gaussian_ab_test(self):
+        result = gaussian_ab_test(10, .5, 10, 10.5, 1, 10)
+        expected = {'chance_to_win': 0.92427,
+                    'expected': 0.05,
+                    'ci': None,
+                    'uplift': None,
+                    'risk': [0.51256, 0.01256]}
+
+        for key in expected.keys():
+            ex = expected[key]
+            res = result[key]
+
+            if ex is None:
+                continue
+            res = [round_(x) for x in res] if isinstance(res, list) else round_(res)
+            ex = [round_(x) for x in ex] if isinstance(ex, list) else round_(ex)
+            self.assertEqual(res, ex)
+
+
+if __name__ == '__main__':
+    unittest_main()

packages/back-end/src/services/stats.ts

@@ -43,21 +43,21 @@ export async function abtest(
     };
   }
 
-  const options = {
-    args: [
-      metric.type,
-      JSON.stringify({
-        users: [aUsers, bUsers],
-        count: [aStats.count, bStats.count],
-        mean: [aStats.mean, bStats.mean],
-        stddev: [aStats.stddev, bStats.stddev],
-      }),
-    ],
-  };
+  const args = [
+    metric.type,
+    JSON.stringify({
+      users: [aUsers, bUsers],
+      count: [aStats.count, bStats.count],
+      mean: [aStats.mean, bStats.mean],
+      stddev: [aStats.stddev, bStats.stddev],
+    }),
+  ];
 
-  const script = path.join(__dirname, "..", "python", "bayesian", "main.py");
-
-  const result = await promisify(PythonShell.run)(script, options);
+  const result = await promisify(PythonShell.run)("bayesian.main", {
+    cwd: path.join(__dirname, "..", "python"),
+    pythonOptions: ["-m"],
+    args,
+  });
   let parsed: {
     chance_to_win: number;
     expected: number;
@@ -72,7 +72,7 @@ export async function abtest(
   try {
     parsed = JSON.parse(result[0]);
   } catch (e) {
-    console.error("Failed to run stats model", options.args, result);
+    console.error("Failed to run stats model", args, result);
     throw e;
   }