Add support for nosetests multiprocessing plugin.
[sawsim.git] / pysawsim / histogram.py
1 # Copyright (C) 2009-2010  W. Trevor King <wking@drexel.edu>
2 #
3 # This program is free software: you can redistribute it and/or modify
4 # it under the terms of the GNU General Public License as published by
5 # the Free Software Foundation, either version 3 of the License, or
6 # (at your option) any later version.
7 #
8 # This program is distributed in the hope that it will be useful,
9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 # GNU General Public License for more details.
12 #
13 # You should have received a copy of the GNU General Public License
14 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
15 #
16 # The author may be contacted at <wking@drexel.edu> on the Internet, or
17 # write to Trevor King, Drexel University, Physics Dept., 3141 Chestnut St.,
18 # Philadelphia PA 19104, USA.
19
20 """Histogram generation and comparison.
21 """
22
23 import numpy
24
25 from . import log
26
27
28 _multiprocess_can_split_ = True
29 """Allow nosetests to split tests between processes.
30 """
31
32
33 class Histogram (object):
34     """A histogram with a flexible comparison method, `residual()`.
35
36     >>> h = Histogram()
37     """
38     def calculate_bin_edges(self, data, bin_width):
39         """
40         >>> h = Histogram()
41         >>> h.calculate_bin_edges(numpy.array([-7.5, 18.2, 4]), 10)
42         array([-10.,   0.,  10.,  20.])
43         >>> h.calculate_bin_edges(numpy.array([-7.5, 18.2, 4, 20]), 10)
44         array([-10.,   0.,  10.,  20.])
45         >>> h.calculate_bin_edges(numpy.array([0, 18.2, 4, 20]), 10)
46         array([  0.,  10.,  20.])
47         >>> h.calculate_bin_edges(numpy.array([18.2, 4, 20]), 10)
48         array([  0.,  10.,  20.])
49         >>> h.calculate_bin_edges(numpy.array([18.2, 20]), 10)
50         array([ 10.,  20.])
51         """
52         m = data.min()
53         M = data.max()
54         bin_start = m - m % bin_width
55         return numpy.arange(bin_start, M+bin_width, bin_width, dtype=data.dtype)
56
57     def from_data(self, data, bin_edges):
58         """Initialize from `data`.
59
60         All bins should be of equal width (so we can calculate which
61         bin a data point belongs to).
62
63         `data` should be a numpy array.
64         """
65         self.headings = None
66         self.bin_edges = bin_edges
67         bin_width = self.bin_edges[1] - self.bin_edges[0]
68
69         bin_is = numpy.floor((data - self.bin_edges[0])/bin_width)
70         self.counts = []
71         for i in range(len(self.bin_edges)-1):
72             self.counts.append(sum(bin_is == i))
73         self.total = float(len(data)) # some data might be outside the bins
74         self.mean = data.mean()
75         self.std_dev = data.std()
76         self.normalize()
77
78     def from_stream(self, stream):
79         """Initialize from `stream`.
80
81         File format:
82
83             # comment and blank lines ignored
84             <bin_edge><whitespace><count>
85             ...
86
87         `<bin_edge>` should mark the left-hand side of the bin, and
88         all bins should be of equal width (so we know where the last
89         one ends).
90
91         >>> import StringIO
92         >>> h = Histogram()
93         >>> h.from_stream(StringIO.StringIO('''# Force (N)\\tUnfolding events
94         ... 150e-12\\t10
95         ... 200e-12\\t40
96         ... 250e-12\\t5
97         ... '''))
98         >>> h.headings
99         ['Force (N)', 'Unfolding events']
100         >>> h.total
101         55.0
102         >>> h.counts
103         [10.0, 40.0, 5.0]
104         >>> h.bin_edges  # doctest: +ELLIPSIS
105         [1.5e-10, 2.000...e-10, 2.500...e-10, 3e-10]
106         >>> h.probabilities  # doctest: +ELLIPSIS
107         [0.181..., 0.727..., 0.0909...]
108         """
109         self.headings = None
110         self.bin_edges = []
111         self.counts = []
112         for line in stream.readlines():
113             line = line.strip()
114             if len(line) == 0 or line.startswith('#'):
115                 if self.headings == None and line.startswith('#'):
116                     line = line[len('#'):]
117                     self.headings = [x.strip() for x in line.split('\t')]
118                 continue # ignore blank lines and comments
119             try:
120                 bin_edge,count = line.split()
121             except ValueError:
122                 log().error('Unable to parse histogram line: "%s"' % line)
123                 raise
124             self.bin_edges.append(float(bin_edge))
125             self.counts.append(float(count))
126         bin_width = self.bin_edges[1] - self.bin_edges[0]
127         self.bin_edges.append(self.bin_edges[-1]+bin_width)
128         self.total = float(sum(self.counts))
129         self.mean = 0
130         for bin,count in zip(self.bin_edges, self.counts):
131             bin += bin_width / 2.0
132             self.mean += bin * count
133         self.mean /=  float(self.total)
134         variance = 0
135         for bin,count in zip(self.bin_edges, self.counts):
136             bin += bin_width / 2.0
137             variance += (bin - self.mean)**2 * count
138         self.std_dev = numpy.sqrt(variance)
139         self.normalize()
140
141     def to_stream(self, stream):
142         """Write to `stream` with the same format as `from_stream()`.
143
144         >>> import sys
145         >>> h = Histogram()
146         >>> h.headings = ['Force (N)', 'Unfolding events']
147         >>> h.bin_edges = [1.5e-10, 2e-10, 2.5e-10, 3e-10]
148         >>> h.counts = [10, 40, 5]
149         >>> h.to_stream(sys.stdout)
150         ... # doctest: +NORMALIZE_WHITESPACE, +REPORT_UDIFF
151         #Force (N)\tUnfolding events
152         1.5e-10\t10
153         2e-10\t40
154         2.5e-10\t5
155         """
156         if self.headings != None:
157             stream.write('#%s\n' % '\t'.join(self.headings))
158         for bin,count in zip(self.bin_edges, self.counts):
159             stream.write('%g\t%g\n' % (bin, count))
160
161     def normalize(self):
162         self.total = float(self.total)
163         self.probabilities = [count/self.total for count in self.counts]
164
165     def mean_residual(self, other):
166         return abs(other.mean - self.mean)
167
168     def std_dev_residual(self, other):
169         return abs(other.std_dev - self.std_dev)
170
171     def chi_squared_residual(self, other):
172         assert self.bin_edges == other.bin_edges
173         residual = 0
174         for probA,probB in zip(self.probabilities, other.probabilities):
175             residual += (probA-probB)**2 / probB
176         return residual
177
178     def jensen_shannon_residual(self, other):
179         assert self.bin_edges == other.bin_edges
180         def d_KL_term(p,q):
181             """
182             Kullback-Leibler divergence for a single bin, with the
183             exception that we return 0 if q==0, rather than
184             exploding like d_KL should.  We can do this because
185             the way d_KL_term is used in the Jensen-Shannon
186             divergence, if q (really m) == 0, then p also == 0,
187             and the Jensen-Shannon divergence for the entire term
188             should be zero.
189             """
190             if p==0 or q==0 or p==q:
191                 return 0.0
192             return p * numpy.log2(p/q)
193         residual = 0
194         for probA,probB in zip(self.probabilities, other.probabilities):
195             m = (probA+probB) / 2.0
196             residual += 0.5*(d_KL_term(probA,m) + d_KL_term(probB,m))
197         return residual
198
199     def _method_to_type(self, name):
200         return name[:-len('_residual')].replace('_', '-')
201
202     def _type_to_method(self, name):
203         return '%s_residual' % name.replace('-', '_')
204
205     def types(self):
206         """Return a list of supported residual types.
207         """
208         return sorted([self._method_to_type(x)
209                        for x in dir(self) if x.endswith('_residual')])
210
211     def residual(self, other, type='jensen-shannon'):
212         """Compare this histogram with `other`.
213
214         Supported comparison `type`\s may be found with `types()`:
215
216         >>> h = Histogram()
217         >>> h.types()
218         ['chi-squared', 'jensen-shannon', 'mean', 'std-dev']
219
220         Selecting an invalid `type` raises an exception.
221
222         >>> h.residual(other=None, type='invalid-type')
223         Traceback (most recent call last):
224           ...
225         AttributeError: 'Histogram' object has no attribute 'invalid_type_residual'
226
227         For residual types where there is a convention, this histogram
228         is treated as the experimental histogram and the other
229         histogram is treated as the theoretical one.
230         """
231         r_method = getattr(self, self._type_to_method(type))
232         return r_method(other)