# -*- coding: utf-8 -*-

#   Copyright (c) 2010-2014, MIT Probabilistic Computing Project
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import math

from bayeslite.util import float_sum

def arithmetic_mean(array):
    """Return the arithmetic mean of elements of ARRAY in floating-point."""
    return float_sum(array) / len(array)

def pearsonr(a0, a1):
    n = len(a0)
    assert n == len(a1)
    m0 = arithmetic_mean(a0)
    m1 = arithmetic_mean(a1)
    num = float_sum((x0 - m0)*(x1 - m1) for x0, x1 in zip(a0, a1))
    den0_root = float_sum((x0 - m0)**2 for x0 in a0)
    den1_root = float_sum((x1 - m1)**2 for x1 in a1)
    r = (num / math.sqrt(den0_root*den1_root))
    # Clamp r in [-1, +1] in case of floating-point error.
    r = min(r, +1.0)
    r = max(r, -1.0)
    return r

assert pearsonr([1,2,3], [2,4,6]) == +1.0
assert pearsonr([1,2,3], [-2,-4,-6]) == -1.0
assert pearsonr([1,2,3], [6,4,2]) == -1.0
assert pearsonr([1,2,3], [+1,-1,+1]) == 0.0

@@ -1474,8 +1474,1082 @@
  == 0.0%0A
+%0Adef signum(x):%0A    if x %3C 0:%0A        return -1%0A    elif 0 %3C x:%0A        return +1%0A    else:%0A        return 0%0A%0Adef relerr(expected, actual):%0A    assert abs((actual - expected)/expected)%0A%0Adef chi2_contingency(contingency, correction=None):%0A    if correction is None:%0A        correction is True%0A    assert 0 %3C len(contingency)%0A    assert all(all(isinstance(v, int) for v in row) for row in contingency)%0A    n = float(sum(sum(row) for row in contingency))%0A    n0 = len(contingency)%0A    n1 = len(contingency%5B0%5D)%0A    assert all(n1 == len(row) for row in contingency)%0A    p0 = %5Bfloat_sum(contingency%5Bi0%5D%5Bi1%5D/n for i1 in range(n1))%0A        for i0 in range(n0)%5D%0A    p1 = %5Bfloat_sum(contingency%5Bi0%5D%5Bi1%5D/n for i0 in range(n0))%0A        for i1 in range(n1)%5D%0A    def q(i0, i1):%0A        O = contingency%5Bi0%5D%5Bi1%5D%0A        E = n*p0%5Bi0%5D*p1%5Bi1%5D%0A        if correction:%0A            O += 0.5*signum(E - O)%0A        return ((O - E)**2)/E%0A    return float_sum(q(i0, i1) for i0 in range(n0) for i1 in range(n1))%0A%0Aassert relerr(7.66, chi2_contingency(%5B%5B4,2,3%5D,%5B3,16,2%5D%5D, correction=False)) %5C%0A    %3C 0.01%0A

Add chi2_contingency to bayeslite.stats.