Move `import with_statement` from sawsim_histogram to sawsim.
[sawsim.git] / pysawsim / sawsim_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, Drudge's University, Physics Dept., 3141 Chestnut St.,
18 # Philadelphia PA 19104, USA.
19
20 import numpy
21
22 from . import __version__, log
23 from .histogram import Histogram
24 from .manager import MANAGERS, get_manager
25 from .sawsim import SawsimRunner
26
27
28 def sawsim_histogram(sawsim_runner, param_string, N=400, bin_edges=None):
29         """Run `N` simulations and return a histogram with `bin_edges`.
30
31         If `bin_edges == None`, return a numpy array of all unfolding
32         forces.
33         """
34         events = []
35         for run in sawsim_runner(param_string=param_string, N=N):
36             events.extend([event.force for event in run
37                            if (event.initial_state == 'folded'
38                                and event.final_state == 'unfolded')])
39         events = numpy.array(events)
40         if bin_edges == None:
41             return events
42         h = Histogram()
43         h.from_data(events, bin_edges)
44         return h
45
46
47 def main(argv=None):
48     """
49     >>> main(['--sawsim', 'bin/sawsim', '-N', '2'])
50     ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
51     #Force (N)  Unfolding events
52     ...
53     """
54     from optparse import OptionParser
55     import sys
56
57     if argv == None:
58         argv = sys.argv[1:]
59
60     sr = SawsimRunner()
61
62     usage = '%prog [options]'
63     epilog = '\n'.join([
64             'Generate an unfolding force histogram from a series of `sawsim`',
65             'runs.',
66             ])
67     parser = OptionParser(usage, epilog=epilog)
68     for option in sr.optparse_options:
69         parser.add_option(option)
70     parser.add_option("-w", "--bin-width", dest="bin_width",
71                       metavar="FLOAT", type='float',
72                       help="Histogram bin width in newtons (%default).",
73                       default=10e-12)
74
75     options,args = parser.parse_args(argv)
76
77     sr_call_params = sr.initialize_from_options(options)
78     try:
79         events = sawsim_histogram(
80             sawsim_runner=sr, bin_edges=None, **sr_call_params)
81     finally:
82         sr.teardown()
83
84     if options.bin_width == None:
85         sys.stdout.write('# Unfolding force (N)\n')
86         events.tofile(sys.stdout, sep='\n')
87         sys.stdout.write('\n')
88     else:
89         h = Histogram()
90         h.from_data(events, bin_edges=h.calculate_bin_edges(
91                 events, options.bin_width))
92         h.headings = ['Force (N)', 'Unfolding events']
93         h.to_stream(sys.stdout)