forked from burakbayramli/books
-
Notifications
You must be signed in to change notification settings - Fork 0
/
chisq2.py
61 lines (53 loc) · 1.79 KB
/
chisq2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import pylab, scipy.special, random, math
N1 = input('number of measurements from source 1 -> ')
N2 = input('number of measurements from source 2 -> ')
mu1 = input('rate of events from source 1 -> ')
mu2 = input('rate of events from source 2 -> ')
seed = input('random number seed -> ')
random.seed(seed)
# generate data
n1 = [0]*N1 # number of events for each measurement from source 1
n2 = [0]*N2 # number of events for each measurement from source 2
q1 = math.exp(-mu1)
q2 = math.exp(-mu2)
for i in range(N1):
p = random.random()
while p > q1:
p *= random.random()
n1[i] += 1
for i in range(N2):
p = random.random()
while p > q2:
p *= random.random()
n2[i] += 1
# compute observed and expected distributions
K = max(n1+n2)+1
R = [n1.count(k) for k in range(K)] # frequencies from source 1
S = [n2.count(k) for k in range(K)] # frequencies from source 1
counts = range(K)
# delete bins where R and S are both zero; note: need to go backwards
for k in reversed(range(K)):
if R[k] == 0 and S[k] == 0:
del R[k]
del S[k]
del counts[k]
# remaining number of bins
K = len(counts)
# compute chi-squared statistic
Q1 = (float(N2)/float(N1))**0.5
Q2 = 1.0/Q1
chisq = sum((Q1*R[k]-Q2*S[k])**2/(R[k]+S[k]) for k in range(K))
# compute significance
nu = K-1 # degrees of freedom
p = 1.0-scipy.special.gammainc(0.5*nu, 0.5*chisq)
# plot results
counts = pylab.array(counts)
pylab.bar(counts-0.3, R, color='b', width=0.3, label='source 1')
pylab.bar(counts, S, color='g', width=0.3, label='source 2')
pylab.ylabel('normalized frequency')
pylab.xlabel('counts')
pylab.xticks(range(min(counts), max(counts)+1))
pylab.xlim(xmin=min(counts)-1, xmax=max(counts)+1)
pylab.legend()
pylab.title('chi-squared = %f, p-value = %f%%'%(chisq, 100.0*p))
pylab.show()