Merge branch 'master' of git@gitorious.org:/sdepy/sdepy
[sdepy:sdepy.git] / sde.py
1 import copy
2 import cPickle as pickle
3 import math
4 import operator
5 import os
6 import pwd
7 import re
8 import signal
9 import sys
10 import time
11 from collections import namedtuple, Iterable
12
13 from optparse import OptionGroup, OptionParser, OptionValueError, Values
14
15 import numpy
16 import sympy
17 from sympy import Symbol
18 from sympy.core import basic
19 from sympy.printing.ccode import CCodePrinter
20 from sympy.printing.precedence import precedence, PRECEDENCE
21
22 import pycuda.autoinit
23 import pycuda.driver as cuda
24 import pycuda.compiler
25 import pycuda.tools
26
27 from mako.template import Template
28 from mako.lookup import TemplateLookup
29
30 PeriodInfo = namedtuple('PeriodInfo', 'period freq')
31 OutputDecl = namedtuple('OutputDecl', 'func vars')
32
33 # Map RNG name to number of uint state variables.
34 RNG_STATE = {
35     'xs32': 1,
36     'kiss32': 4,
37     'nr32': 4,
38 }
39
40 def drift_velocity(sde, *args):
41     ret = []
42     for starting, final in args:
43         a = starting.astype(numpy.float64)
44         b = final.astype(numpy.float64)
45
46         ret.append((numpy.average(b) - numpy.average(a)) /
47                 (sde.sim_t - sde.start_t))
48
49     return ret
50
51 def abs_drift_velocity(sde, *args):
52     ret = []
53     for starting, final in args:
54         a = starting.astype(numpy.float64)
55         b = final.astype(numpy.float64)
56
57         ret.append(numpy.average(numpy.abs(b - a)) /
58                 (sde.sim_t - sde.start_t))
59
60     return ret
61
62 def diffusion_coefficient(sde, *args):
63     ret = []
64     for starting, final in args:
65         a = starting.astype(numpy.float64)
66         b = final.astype(numpy.float64)
67
68         deff1 = numpy.average(numpy.square(b)) - numpy.average(b)**2
69         deff2 = numpy.average(numpy.square(a)) - numpy.average(a)**2
70 #        ret.append((deff1 - deff2) / (2.0 * (sde.sim_t - sde.start_t)))
71         ret.append(deff1 / (2.0 * sde.sim_t))
72         ret.append(deff2 / (2.0 * sde.start_t))
73
74     return ret
75
76 def avg_moments(sde, *args):
77     ret = []
78
79     for arg in args:
80         ret.append(numpy.average(arg))
81         ret.append(numpy.average(numpy.square(arg)))
82
83     return ret
84
85 want_save = False
86 want_dump = False
87 want_exit = False
88
89 def _sighandler(signum, frame):
90     global want_dump, want_exit, want_save
91
92     if signum == signal.SIGUSR2:
93         want_save = True
94     else:
95         want_dump = True
96
97         if signum in [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]:
98             want_exit = True
99
100 def _convert_to_double(src):
101     import re
102     s = re.sub(r'float([^\.])', r'double\1', src)
103     s = s.replace('logf(', 'log(')
104     s = s.replace('expf(', 'exp(')
105     s = s.replace('powf(', 'pow(')
106     s = s.replace('sqrtf(', 'sqrt(')
107     s = s.replace('cosf(', 'cos(')
108     s = s.replace('sinf(', 'sin(')
109     s = s.replace('FLT_EPSILON', 'DBL_EPSILON')
110     return re.sub('([0-9]+\.[0-9]*)f', '\\1', s)
111
112 def _parse_range(option, opt_str, value, parser):
113     vals = value.split(':')
114
115     # Use higher precision (float64) below.   If single precision is requested, the values
116     # will be automatically degraded to float32 later on.
117     if len(vals) == 1:
118         setattr(parser.values, option.dest, numpy.float64(value))
119         parser.par_single.add(option.dest)
120     elif len(vals) == 3:
121         start = float(vals[0])
122         stop = float(vals[1])
123         size = float(vals[2])
124         setattr(parser.values, option.dest, numpy.linspace(start, stop, size, numpy.float64))
125         parser.par_multi.add(option.dest)
126     elif len(vals) == 4 and vals[0] == 'log':
127         start = float(vals[1])
128         stop = float(vals[2])
129         size = float(vals[3])
130         setattr(parser.values, option.dest, numpy.logspace(start, stop, size, numpy.float64))
131         parser.par_multi.add(option.dest)
132     else:
133         raise OptionValueError('"%s" has to be a single value or a range of the form \'start:stop:steps\'' % opt_str)
134
135 def _get_module_source(*files):
136     """Load multiple .cu files and join them into a string."""
137     src = ""
138     for fname in files:
139         fp = open(fname)
140         src += fp.read()
141         fp.close()
142     return src
143
144
145 class SolverGenerator(object):
146     @classmethod
147     def get_source(cls, sde, const_parameters, kernel_parameters, local_vars):
148         """
149         sde: a SDE object for which the code is to be generated
150         const_parameters: list of constant parameters for the SDE
151         kernel_parameters: list of parameters which will be passed as arguments to
152             the kernel
153         """
154         pass
155
156 class SRK2(SolverGenerator):
157     """Stochastic Runge Kutta method of the 2nd order."""
158
159     @classmethod
160     def get_source(cls, sde, const_parameters, kernel_parameters, local_vars):
161         ctx = {}
162         ctx['const_parameters'] = const_parameters
163         ctx['local_vars'] = dict([(k, KernelCodePrinter().doprint(v)) for k, v in local_vars.iteritems()])
164         ctx['par_cuda'] = kernel_parameters
165         ctx['rhs_vars'] = sde.num_vars
166         ctx['noise_strength_map'] = sde.noise_map
167         ctx['num_noises'] = sde.num_noises
168         ctx['sde_code'] = sde.code
169         ctx['rng_state_size'] = RNG_STATE[sde.options.rng]
170         ctx['rng'] = sde.options.rng
171
172         lookup = TemplateLookup(directories=sys.path,
173                 module_directory='/tmp/pysde_modules-%s' %
174                                 (pwd.getpwuid(os.getuid())[0]))
175         sde_template = lookup.get_template('sde.mako')
176         return sde_template.render(**ctx)
177
178 def int2float(t):
179     #return re.sub(r'([0-9]+)([^\.])', r'\1.0\2', str(t))
180     return str(t)
181
182 class KernelCodePrinter(CCodePrinter):
183
184     def _print_Pow(self, expr):
185         PREC = precedence(expr)
186         if expr.exp is basic.S.NegativeOne:
187             return '1.0/%s' % (self.parenthesize(expr.base, PREC))
188         # For the kernel code, it's better to calculate the power
189         # here explicitly by multiplication.
190         elif expr.exp == 2:
191             return '%s*%s' % (self.parenthesize(expr.base, PREC),
192                               self.parenthesize(expr.base, PREC))
193         else:
194             return int2float('powf(%s,%s)' % (self.parenthesize(expr.base, PREC),
195                                               self.parenthesize(expr.exp, PREC)))
196
197     def _print_Function(self, expr):
198         if expr.func.__name__ == 'log':
199             return 'logf(%s)' % self.stringify(expr.args, ', ')
200         else:
201             return super(KernelCodePrinter, self)._print_Function(expr)
202
203 class TextOutput(object):
204     def __init__(self, sde, subfiles):
205         self.sde = sde
206         self.subfiles = subfiles
207         self.out = {}
208
209         if sde.options.output is not None:
210             self.out['main'] = open(sde.options.output, 'w')
211
212             for sub in subfiles:
213                 if sub == 'main':
214                     continue
215                 self.out[sub] = open('%s_%s' % (sde.options.output, sub), 'w')
216         else:
217             if len(subfiles) > 1:
218                 raise ValueError('Output file name required so that auxiliary data stream can be saved.')
219
220             self.out['main'] = sys.stdout
221
222     def finish_block(self):
223         print >>self.out['main'], ''
224         self.out['main'].flush()
225
226     def data(self, **kwargs):
227         for name, val in kwargs.iteritems():
228             def my_rep(val):
229                 if isinstance(val, Iterable):
230                     return ' '.join(my_rep(x) for x in val)
231                 else:
232                     return str(val)
233
234             rep = [my_rep(x) for x in val]
235             print >>self.out[name], ' '.join(rep)
236
237     def header(self):
238         out = self.out['main']
239
240         print >>out, '# %s' % ' '.join(sys.argv)
241         if self.sde.options.seed is not None:
242             print >>out, '# seed = %d' % self.sde.options.seed
243         print >>out, '# sim periods = %d' % self.sde.options.simperiods
244         print >>out, '# transient periods = %d' % self.sde.options.transients
245         print >>out, '# samples = %d' % self.sde.options.samples
246         print >>out, '# paths = %d' % self.sde.options.paths
247         print >>out, '# spp = %d' % self.sde.options.spp
248         for par in self.sde.parser.par_single:
249             print >>out, '# %s = %f' % (par, self.sde.options.__dict__[par])
250         for par in self.sde.par_multi_ordered:
251             print >>out, '# %s = %s' % (par, ' '.join(str(x) for x in self.sde.options.__dict__[par]))
252         for par in self.sde.scan_vars:
253             print >>out, '# %s = %s' % (par, ' '.join(str(x) for x in self.sde.options.__dict__[par]))
254  
255     def close(self):
256         pass
257
258 class NpyOutput(object):
259     def __init__(self, sde, subfiles):
260         self.sde = sde
261         self.subfiles = subfiles
262
263         # This dictionary maps the subfile name to a list of lists.
264         # Every entry in the outer list represents one point in the
265         # parameter space.  The inner list represents the results for
266         # a particular set of parameters.
267         self.cache = {}
268
269     def finish_block(self):
270         global want_save
271
272         if want_save:
273             self.close()
274             want_save = False
275
276     def data(self, **kwargs):
277         for name, val in kwargs.iteritems():
278             self.cache.setdefault(name, []).append(val)
279
280     def header(self):
281         self.cmdline = '%s' % ' '.join(sys.argv)
282         self.scan_vars = self.sde.scan_vars
283         self.par_multi_ordered = self.sde.par_multi_ordered
284
285     def close(self):
286         out = {}
287         shape = []
288
289         for par in self.par_multi_ordered:
290             out[par] = getattr(self.sde.options, par)
291             shape.append(len(out[par]))
292
293         for sv in self.scan_vars:
294             out[sv] = getattr(self.sde.options, sv)
295             shape.append(len(out[sv]))
296
297         for name, val in self.cache.iteritems():
298             inner_len = max(len(x) for x in val)
299             out[name] = numpy.array(val, dtype=self.sde.float)
300
301             if shape and reduce(operator.mul, shape) * inner_len == reduce(operator.mul, out[name].shape):
302                 out[name] = numpy.reshape(out[name], shape + [inner_len])
303
304         numpy.savez(self.sde.options.output, cmdline=self.cmdline, scan_vars=self.scan_vars,
305                     par_multi=self.par_multi_ordered, options=self.sde.options, **out)
306
307 class StoreOutput(object):
308     def __init__(self, sde, subfiles):
309         self.sde = sde
310         self.subfiles = subfiles
311         self.cache = {}
312
313     def finish_block(self):
314         pass
315
316     def data(self, **kwargs):
317         for name, val in kwargs.iteritems():
318             self.cache.setdefault(name, []).append(val)
319
320     def header(self):
321         self.scan_vars = self.sde.scan_vars
322         self.par_multi_ordered = self.sde.par_multi_ordered
323
324     def close(self):
325         out = {}
326         shape = []
327
328         for par in self.par_multi_ordered:
329             out[par] = getattr(self.sde.options, par)
330             shape.append(len(out[par]))
331
332         for sv in self.scan_vars:
333             out[sv] = getattr(self.sde.options, sv)
334             shape.append(len(out[sv]))
335
336         for name, val in self.cache.iteritems():
337             inner_len = max(len(x) for x in val)
338             out[name] = numpy.array(val, dtype=self.sde.float)
339
340             if shape and reduce(operator.mul, shape) * inner_len == reduce(operator.mul, out[name].shape):
341                 out[name] = numpy.reshape(out[name], shape + [inner_len])
342
343         self.out = out
344
345 class S(object):
346     def __init__(self):
347         self.dt = Symbol('dt')
348         self.samples = Symbol('samples')
349
350 class SDE(object):
351     """A class representing a SDE equation to solve."""
352
353     format_cmd = r"indent -linux -sob -l120 {file} ; sed -i -e '/^$/{{N; s/\n\([\t ]*}}\)$/\1/}}' -e '/{{$/{{N; s/{{\n$/{{/}}' {file}"
354
355     def __init__(self, code, params, num_vars, num_noises,
356             noise_map, period_map=None, args=None, local_vars=None):
357         """
358         :param code: the code defining the Stochastic Differential Equation
359         :param params: dict of simulation parameters; keys are parameter names,
360             values are descriptions
361         :param num_vars: number of variables in the SDE
362         :param num_noises: number of independent, white Gaussian noises
363         :param noise_map: a dictionary, mapping the variable number to a list of
364             ``num_noises`` noise strengths, which can be either 0 or the name of
365             a variable containing the noise strength value
366         :param period_map: a dictionary, mapping the variable number to PeriodInfo
367             objects.  If a variable has a corresponding entry in this dictionary,
368             it will be assumed to be a periodic variable, such that only its
369             value modulo ``period`` is important for the evolution of the system.
370
371             Every ``frequency`` * ``samples`` steps, the value
372             of this variable will be folded back to the range of [0, period).
373             Its full (unfolded) value will however be retained in the output.
374
375             The folded values will result in faster CUDA code if trigonometric
376             functions are used and if the magnitude of their arguments always
377             remains below 48039.0f (see CUDA documentation).
378         """
379         self.parser = OptionParser()
380
381         group = OptionGroup(self.parser, 'SDE engine settings')
382         group.add_option('--spp', dest='spp', help='steps per period', metavar='DT', type='int', action='store', default=100)
383         group.add_option('--samples', dest='samples', help='sample the position every N steps', metavar='N', type='int', action='store', default=100)
384         group.add_option('--paths', dest='paths', help='number of paths to sample', type='int', action='store', default=256)
385         group.add_option('--transients', dest='transients', help='number of periods to ignore because of transients', type='int', action='store', default=200)
386         group.add_option('--simperiods', dest='simperiods', help='number of periods in the simulation', type='int', action='store', default=2000)
387         group.add_option('--seed', dest='seed', help='RNG seed', type='int', action='store', default=None)
388         group.add_option('--precision', dest='precision', help='precision of the floating-point numbers (single, double)', type='choice', choices=['single', 'double'], default='single')
389         group.add_option('--rng', dest='rng', help='PRNG to use', type='choice', choices=RNG_STATE.keys(), default='kiss32')
390         group.add_option('--no-fast-math', dest='fast_math',
391                 help='do not use faster intrinsic mathematical functions everywhere',
392                 action='store_false', default=True)
393         group.add_option('--deterministic', dest='deterministic', help='do not generate any noises',
394                 action='store_true', default=False)
395         self.parser.add_option_group(group)
396
397         group = OptionGroup(self.parser, 'Debug settings')
398         group.add_option('--save_src', dest='save_src', help='save the generated source to FILE', metavar='FILE',
399                                type='string', action='store', default=None)
400         group.add_option('--use_src', dest='use_src', help='use FILE instead of the automatically generated code',
401                 metavar='FILE', type='string', action='store', default=None)
402         group.add_option('--noformat_src', dest='format_src', help='do not format the generated source code', action='store_false', default=True)
403         self.parser.add_option_group(group)
404
405         group = OptionGroup(self.parser, 'Output settings')
406         group.add_option('--output_mode', dest='omode', help='output mode', type='choice', choices=['summary', 'path'], action='store', default='summary')
407         group.add_option('--output_format', dest='oformat', help='output file format', type='choice',
408                 choices=['npy', 'text', 'store'], action='store', default='npy')
409         group.add_option('--output', dest='output', help='base output filename', type='string', action='store', default=None)
410         group.add_option('--save_every', dest='save_every', help='save output every N seconds', type='int',
411                 action='store', default=0)
412         self.parser.add_option_group(group)
413
414         group = OptionGroup(self.parser, 'Checkpointing settings')
415         group.add_option('--dump_every', dest='dump_every', help='dump system state every N seconds', type='int',
416                 action='store', default=0)
417         group.add_option('--dump_state', dest='dump_filename', help='save state of the simulation to FILE after it is completed',
418                 metavar='FILE', type='string', action='store', default=None)
419         group.add_option('--restore_state', dest='restore_filename', help='restore state of the solver from FILE',
420                 metavar='FILE', type='string', action='store', default=None)
421         group.add_option('--resume', dest='resume', help='resume simulation from a saved checkpoint',
422                 action='store_true', default=False)
423         group.add_option('--continue', dest='continue_', help='continue a finished simulation',
424                 action='store_true', default=False)
425         self.parser.add_option_group(group)
426
427         self.parser.par_multi = set()
428         self.parser.par_single = set()
429
430         self.sim_params = params
431         self.num_vars = num_vars
432         self.num_noises = num_noises
433         self.noise_map = noise_map
434
435         if period_map is None:
436             self.period_map = {}
437         else:
438             self.period_map = period_map
439
440         self.make_symbols(local_vars)
441
442         # Local variables are defined as lambdas since they need access to the
443         # symbols.  Replace the lambdas with their values here now that the
444         # symbols are defined.
445         self.local_vars = {}
446         if local_vars is not None:
447             for k, v in local_vars.iteritems():
448                 self.local_vars[k] = v(self)
449
450         self.code = code
451
452         self.last_dump = time.time()
453         self.last_save = time.time()
454
455         # Current results.
456         self.state_results = []
457         # Results loaded from a dump file.
458         self.prev_state_results = []
459
460         self._sim_sym = {}
461         self._gpu_sym = {}
462
463         group = OptionGroup(self.parser, 'Simulation-specific settings')
464
465         for name, help_string in params.iteritems():
466             group.add_option('--%s' % name, dest=name, action='callback',
467                     callback=_parse_range, type='string', help=help_string,
468                     default=None)
469
470         self.parser.add_option_group(group)
471
472         for k, v in noise_map.iteritems():
473             if len(v) != num_noises:
474                 raise ValueError('The number of noise strengths for variable %s'
475                     ' has to be equal to %d.' % (k, num_noises))
476
477         self.parse_args(args)
478
479         if self.options.deterministic:
480             self.num_noises = 0
481             self.noise_map = []
482
483     def make_symbols(self, local_vars):
484         """Create a sympy Symbol for each simulation parameter."""
485         self.S = S()
486         for param in self.sim_params.iterkeys():
487             setattr(self.S, param, Symbol(param))
488         for param in local_vars.iterkeys():
489             setattr(self.S, param, Symbol(param))
490
491     def parse_args(self, args=None):
492         if args is None:
493             args = sys.argv
494
495         self.options = Values(self.parser.defaults)
496         self.parser.parse_args(args, self.options)
497
498         for name, hs in self.sim_params.iteritems():
499             if self.options.__dict__[name] == None:
500                 raise OptionValueError('Required option "%s" not specified.' % name)
501
502         if self.options.output is None and self.options.oformat != 'text':
503             raise OptionValueError('Required option "output"" not specified.')
504
505         if self.options.precision == 'single':
506             self.float = numpy.float32
507         else:
508             self.float = numpy.float64
509
510         if (self.options.resume or self.options.continue_) and self.options.restore_filename is None:
511             raise OptionValueError('The resume and continue modes require '
512                 'a state file to be specified with --restore_state.')
513
514         if self.options.restore_filename is not None:
515             self.load_state()
516
517     def _find_cuda_scan_par(self):
518         """Automatically determine which parameter is to be scanned over
519         in a single kernel run.
520         """
521         max_par = None
522         max = 0
523
524         # Look for a parameter with the highest number of values.
525         for par in self.parser.par_multi:
526             if len(self.options.__dict__[par]) > max and par != self.freq_var:
527                 max_par = par
528                 max = len(self.options.__dict__[par])
529
530         # No variable to scan over.
531         if max_par is None:
532             return None
533
534         self.parser.par_multi.remove(max_par)
535         return max_par
536
537     def prepare(self, algorithm, init_vectors, freq_var=None):
538         """Prepare a SDE simulation.
539
540         :param algorithm: the SDE solver to use, sublass of SDESolver
541         :param init_vectors: a callable that will be used to set the initial conditions
542         :param freq_var: name of the parameter which is to be interpreted as a
543             frequency (determines the step size 'dt').  If ``None``, the
544             system period will be assumed to be 1.0 and the time step size
545             will be set to 1.0/spp.
546         """
547         self.freq_var = freq_var
548
549         # Determine the scan variable(s).
550         if not (self.options.resume or self.options.continue_):
551             sv = self._find_cuda_scan_par()
552             if sv is None:
553                 self.scan_vars = []
554             else:
555                 self.scan_vars = [sv]
556
557             # Create an ordered copy of the list of parameters which have multiple
558             # values.
559             self.par_multi_ordered = list(self.parser.par_multi)
560
561             # Multi-valued parameters that are scanned over are no longer
562             # contained in par_multi at this point.
563             self.global_vars = self.parser.par_single | self.parser.par_multi
564             userdef_global_vars = self.global_vars.copy()
565             self.global_vars.add('dt')
566             self.global_vars.add('samples')
567             self.const_local_vars = {}
568
569             # If a local variable only depends on constant variables, make
570             # it a global constant.
571             for name, value in self.local_vars.iteritems():
572                 if set([str(x) for x in value.atoms(Symbol)]) <= self.global_vars:
573                     self.global_vars.add(name)
574                     userdef_global_vars.add(name)
575                     self.const_local_vars[name] = value
576
577             for name in self.const_local_vars.iterkeys():
578                 del self.local_vars[name]
579
580         # Prepare the CUDA source.  Load/save it from/to a file if requested.
581         if self.options.use_src:
582             with open(self.options.use_src, 'r') as file:
583                 kernel_source = file.read()
584         else:
585             kernel_source = algorithm.get_source(self, userdef_global_vars,
586                    self.scan_vars, self.local_vars)
587
588             if self.options.precision == 'double':
589                 kernel_source = _convert_to_double(kernel_source)
590
591         if self.options.save_src is not None:
592             with open(self.options.save_src, 'w') as file:
593                 print >>file, kernel_source
594
595             if self.options.format_src:
596                 os.system(self.format_cmd.format(file=self.options.save_src))
597
598         return self.cuda_prep(init_vectors, kernel_source)
599
600     @property
601     def scan_var_size(self):
602         ret = 1
603
604         for par in self.scan_vars:
605             ret *= len(getattr(self.options, par))
606
607         return ret
608
609     def get_param(self, name):
610         if name in self.scan_vars:
611             idx = self.scan_vars.index(name)
612             return self._sv[idx]
613         elif name in self._sim_sym:
614             return self._sim_sym[name]
615         else:
616             return getattr(self.options, name)
617
618     def cuda_prep(self, init_vectors, sources, sim_func='AdvanceSim'):
619         """Prepare a SDE simulation for execution using CUDA.
620
621         init_vectors: a function which takes an instance of this class and the
622             variable number as arguments and returns an initialized vector of
623             size num_threads
624         sources: list of source code files for the simulation
625         sim_func: name of the kernel advacing the simulation in time
626         """
627         if self.options.seed is not None and not (self.options.resume or self.options.continue_):
628             numpy.random.seed(self.options.seed)
629
630         self.init_vectors = init_vectors
631         self.num_threads = self.scan_var_size * self.options.paths
632         self._sim_prep_mod(sources, sim_func)
633         self._sim_prep_const()
634         self._sim_prep_var()
635
636     def _sim_prep_mod(self, sources, sim_func):
637         # The use of fast math below will result in certain mathematical functions
638         # being automaticallky replaced with their faster counterparts prefixed with
639         # __, e.g. __sinf().
640         if self.options.fast_math:
641             options=['--use_fast_math']
642         else:
643             options=[]
644
645         if type(sources) is str:
646             self.mod = pycuda.compiler.SourceModule(sources, options=options)
647         else:
648             self.mod = pycuda.compiler.SourceModule(_get_module_source(*sources), options=options)
649         self.advance_sim = self.mod.get_function(sim_func)
650
651     def _sim_prep_const(self):
652         for var in self.global_vars:
653             self._gpu_sym[var] = self.mod.get_global(var)[0]
654
655         samples = numpy.uint32(self.options.samples)
656         cuda.memcpy_htod(self._gpu_sym['samples'], samples)
657
658         # Single-valued system parameters
659         for par in self.parser.par_single:
660             self._sim_sym[par] = self.options.__dict__[par]
661             cuda.memcpy_htod(self._gpu_sym[par], self.float(self.options.__dict__[par]))
662
663     def _init_rng(self):
664         # Initialize the RNG seeds.
665         self._rng_state = numpy.fromstring(numpy.random.bytes(
666             self.num_threads * RNG_STATE[self.options.rng] * numpy.uint32().nbytes),
667             dtype=numpy.uint32)
668         self._gpu_rng_state = cuda.mem_alloc(self._rng_state.nbytes)
669         cuda.memcpy_htod(self._gpu_rng_state, self._rng_state)
670
671     def _sim_prep_var(self):
672         self.vec = []
673         self._gpu_vec = []
674
675         # Prepare device vectors.
676         for i in range(0, self.num_vars):
677             vt = numpy.zeros(self.num_threads).astype(self.float)
678             self.vec.append(vt)
679             self._gpu_vec.append(cuda.mem_alloc(vt.nbytes))
680
681         self._sv = []
682         self._gpu_sv = []
683
684         # Initialize the scan variables.  If we're in resume or continue mode, the
685         # scan variables are already initialized and nothing needs to be done.
686         if not self.options.resume and not self.options.continue_:
687             for sv in self.scan_vars:
688                 vals = numpy.ones(self.options.paths)
689                 for sv2 in reversed(self.scan_vars):
690                     if sv2 != sv:
691                         vals = numpy.kron(numpy.ones_like(getattr(self.options, sv2)), vals)
692                     else:
693                         vals = numpy.kron(getattr(self.options, sv2), vals)
694
695                 self._sv.append(vals.astype(self.float))
696
697         for sv in self._sv:
698             self._gpu_sv.append(cuda.mem_alloc(sv.nbytes))
699             cuda.memcpy_htod(self._gpu_sv[-1], sv)
700
701     def get_var(self, i, starting=False):
702         if starting:
703             vec = self.vec_start
704             nx = self.vec_start_nx
705         else:
706             vec = self.vec
707             nx = self.vec_nx
708
709         if i in nx:
710             return vec[i] + self.period_map[i].period * nx[i]
711         else:
712             return vec[i]
713
714     @property
715     def max_sim_iter(self):
716         return int(self.options.simperiods * self.options.spp / self.options.samples)+1
717
718     def iter_to_sim_time(self, iter_):
719         return iter_ * self.dt * self.options.samples
720
721     def sim_time_to_iter(self, time_):
722         return int(time_ / (self.dt * self.options.samples))
723
724     def simulate(self, req_output, block_size=64):
725         """Run a CUDA SDE simulation.
726
727         req_output: a dictionary mapping the the output mode to a list of
728             tuples of ``(callable, vars)``, where ``callable`` is a function
729             that will compute the values to be returned, and ``vars`` is a list
730             of variables that will be passed to this function
731         block_size: CUDA block size
732         """
733         self.req_output = req_output[self.options.omode]
734
735         if self.options.oformat == 'text':
736             self.output = TextOutput(self, self.req_output.keys())
737         elif self.options.oformat == 'npy':
738             self.output = NpyOutput(self, self.req_output.keys())
739         elif self.options.oformat == 'store':
740             self.output = StoreOutput(self, self.req_output.keys())
741
742         self.output.header()
743
744         # Determine which variables are necessary for the output.
745         self.req_vars = set([])
746         for k, v in self.req_output.iteritems():
747             for func, vars in v:
748                 self.req_vars |= set(vars)
749
750         self.block_size = block_size
751         arg_types = ['P'] + ['P']*self.num_vars + ['P'] * len(self.scan_vars) + [self.float]
752         self.advance_sim.prepare(arg_types, block=(block_size, 1, 1))
753
754         kern = self.advance_sim
755         ddata = pycuda.tools.DeviceData()
756         occ = pycuda.tools.OccupancyRecord(ddata, block_size, kern.shared_size_bytes, kern.num_regs)
757
758         print '# CUDA stats l:%d  s:%d  r:%d  occ:(%f tb:%d w:%d l:%s)' % (kern.local_size_bytes, kern.shared_size_bytes,
759                     kern.num_regs, occ.occupancy, occ.tb_per_mp, occ.warps_per_mp, occ.limited_by)
760
761         self._scan_iter = 0
762
763         if self.options.dump_filename is not None:
764             signal.signal(signal.SIGUSR1, _sighandler)
765             signal.signal(signal.SIGINT, _sighandler)
766             signal.signal(signal.SIGQUIT, _sighandler)
767             signal.signal(signal.SIGHUP, _sighandler)
768             signal.signal(signal.SIGTERM, _sighandler)
769
770         signal.signal(signal.SIGUSR2, _sighandler)
771
772         self._run_nested(self.par_multi_ordered)
773         self.output.close()
774
775         if self.options.dump_filename is not None:
776             self.dump_state()
777
778     def _run_nested(self, range_pars):
779         # No more parameters to loop over, time to actually run the kernel.
780         if not range_pars:
781             # Reinitialize the RNG here so that there is no interdependence
782             # between runs.  This also guarantees that the resume/continue
783             # modes can work correctly in the case of scan over 2+ parameters.
784             self._init_rng()
785
786             # Calculate period and step size.
787             if self.freq_var is not None:
788                 period = 2.0 * math.pi / self._sim_sym[self.freq_var]
789             else:
790                 period = 1.0
791             self.dt = self.float(period / self.options.spp)
792             cuda.memcpy_htod(self._gpu_sym['dt'], self.dt)
793
794             # Evaluate constant local vars.
795             subs = {self.S.dt: self.dt}
796             for k, v in self._sim_sym.iteritems():
797                 subs[Symbol(k)] = v
798
799             for name, value in self.const_local_vars.iteritems():
800                 cuda.memcpy_htod(self._gpu_sym[name], self.float(value.subs(subs)))
801
802             # In the resume mode, we skip all the computations that have already
803             # been completed and thus are saved in self.prev_state_results.
804             if not self.options.resume:
805                 self._run_kernel(period)
806             elif (self._scan_iter == len(self.prev_state_results)-1 and
807                     self.prev_state_results[-1][0] < self.iter_to_sim_time(self.max_sim_iter)):
808                 self._run_kernel(period)
809             elif self._scan_iter > len(self.prev_state_results)-1:
810                 self._run_kernel(period)
811
812             self._scan_iter += 1
813         else:
814             par = range_pars[0]
815
816             # Loop over all values of a specific parameter.
817             for val in self.options.__dict__[par]:
818                 self._sim_sym[par] = self.float(val)
819                 cuda.memcpy_htod(self._gpu_sym[par], self.float(val))
820                 self._run_nested(range_pars[1:])
821
822     def _run_kernel(self, period):
823         kernel_args = [self._gpu_rng_state] + self._gpu_vec + self._gpu_sv
824
825         # Prepare an array for initial value of the variables (after
826         # transients).
827         if self.options.omode == 'path':
828             transient = False
829             pathwise = True
830         else:
831             transient = True
832             pathwise = False
833
834         if (self.options.continue_ or
835                 (self.options.resume and self._scan_iter < len(self.prev_state_results))):
836             self.vec = self.prev_state_results[self._scan_iter][1]
837             self.vec_nx = self.prev_state_results[self._scan_iter][2]
838             self._rng_state = self.prev_state_results[self._scan_iter][3]
839             cuda.memcpy_htod(self._gpu_rng_state, self._rng_state)
840             for i in range(0, self.num_vars):
841                 cuda.memcpy_htod(self._gpu_vec[i], self.vec[i])
842
843             self.sim_t = self.prev_state_results[self._scan_iter][0]
844
845             if self.options.omode == 'summary':
846                 self.start_t = self.prev_state_results[self._scan_iter][4]
847                 self.vec_start = self.prev_state_results[self._scan_iter][5]
848                 self.vec_start_nx = self.prev_state_results[self._scan_iter][6]
849
850                 if self.sim_t >= self.options.transients * period:
851                     transient = False
852         else:
853             # Reinitialize the positions.
854             self.vec = []
855             for i in range(0, self.num_vars):
856                 vt = self.init_vectors(self, i).astype(self.float)
857                 self.vec.append(vt)
858                 cuda.memcpy_htod(self._gpu_vec[i], vt)
859
860             # Prepare an array for number of periods for periodic variables.
861             self.vec_nx = {}
862             for i, v in self.period_map.iteritems():
863                 self.vec_nx[i] = numpy.zeros_like(self.vec[i]).astype(numpy.int64)
864
865             if transient:
866                 self.vec_start = []
867                 for i in range(0, self.num_vars):
868                     self.vec_start.append(numpy.zeros_like(self.vec[i]))
869
870             self.sim_t = 0.0
871
872         def fold_variables(iter_, need_copy):
873             for i, (period, freq) in self.period_map.iteritems():
874                 if iter_ % freq == 0:
875                     if need_copy:
876                         cuda.memcpy_dtoh(self.vec[i], self._gpu_vec[i])
877
878                     self.vec_nx[i] = numpy.add(self.vec_nx[i],
879                             numpy.floor_divide(self.vec[i], period).astype(numpy.int64))
880                     self.vec[i] = numpy.remainder(self.vec[i], period)
881                     cuda.memcpy_htod(self._gpu_vec[i], self.vec[i])
882
883         init_iter = self.sim_time_to_iter(self.sim_t)
884
885         global want_dump, want_exit, want_save
886
887         # Actually run the simulation here.
888         for j in xrange(init_iter, self.max_sim_iter):
889             self.sim_t = self.iter_to_sim_time(j)
890
891             if transient and self.sim_t >= self.options.transients * period:
892                 for i in range(0, self.num_vars):
893                     cuda.memcpy_dtoh(self.vec_start[i], self._gpu_vec[i])
894                 transient = False
895                 self.start_t = self.sim_t
896                 self.vec_start_nx = self.vec_nx.copy()
897
898             # Fold the time variable passed to the kernel.
899             # NOTE: Here we implicitly assume that only the reduced value of
900             # time matters for the evolution of the system.
901             args = kernel_args + [self.float(math.fmod(self.sim_t, period))]
902             self.advance_sim.prepared_call((self.num_threads/self.block_size, 1), *args)
903             self.sim_t += self.options.samples * self.dt
904
905             if (self.options.save_every > 0 and
906                     (time.time() - self.last_save > self.options.save_every)):
907                 want_save = True
908                 self.last_save = time.time()
909
910             fold_variables(j, True)
911
912             if pathwise:
913                 self.output_current()
914                 if self.scan_vars:
915                     self.output.finish_block()
916
917             if (want_dump or
918                     (self.options.dump_every > 0 and
919                      time.time() - self.last_dump > self.options.dump_every)):
920
921                 if self.options.dump_filename is not None:
922                     self.save_block()
923                     self.dump_state()
924                     del self.state_results[-1]
925                 want_dump = False
926
927                 if want_exit:
928                     sys.exit(0)
929
930         if not pathwise:
931             self.output_summary()
932
933         self.output.finish_block()
934         self.save_block()
935
936     def output_current(self):
937         vars = {}
938
939         # Get required variables from the compute device and
940         # unfold them.
941         for i in self.req_vars:
942             cuda.memcpy_dtoh(self.vec[i], self._gpu_vec[i])
943             vars[i] = self.get_var(i)
944
945         self._output_results(vars, self.sim_t)
946
947     def output_summary(self):
948         vars = {}
949
950         # Get required variables from the compute device and
951         # unfold them.  For each variable, store both the reference
952         # (starting) value and the final one.
953         for i in self.req_vars:
954             cuda.memcpy_dtoh(self.vec[i], self._gpu_vec[i])
955             vars[i] = (self.get_var(i, True), self.get_var(i))
956
957         self._output_results(vars)
958
959     def _output_results(self, vars, *misc_pars):
960         # Iterate over all values of the scan parameters.  For each unique
961         # value, calculate the requested output(s).
962         for i in range(0, self.scan_var_size):
963             out = {}
964             for out_name, _ in self.req_output.iteritems():
965                 out[out_name] = []
966
967 #            for par in self.parser.par_multi:
968 #                out['main'].append(self._sim_sym[par])
969             out['main'].extend(misc_pars)
970
971 #            for sv in self._sv:
972 #                out['main'].append(sv[i*self.options.paths])
973
974             for out_name, out_decl in self.req_output.iteritems():
975                 for decl in out_decl:
976                     # Map variable numbers to their actual values.
977                     args = map(lambda x: vars[x], decl.vars)
978
979                     # Cut the part of the variables for the current value of the scan vars.
980                     if args and type(args[0]) is tuple:
981                         args = map(lambda x:
982                             (x[0][i*self.options.paths:(i+1)*self.options.paths],
983                              x[1][i*self.options.paths:(i+1)*self.options.paths]), args)
984                     else:
985                         args = map(lambda x:
986                                 x[i*self.options.paths:(i+1)*self.options.paths],
987                                 args)
988
989                     # Evaluate the requested function.
990                     out[out_name].extend(decl.func(self, *args))
991
992             self.output.data(**out)
993
994     @property
995     def state(self):
996         """A dictionary representing the current state of the solver."""
997
998         names = ['sim_params', 'num_vars', 'num_noises',
999                  'noise_map', 'period_map', 'code', 'options', 'float',
1000                  'scan_vars', 'local_vars' 'const_local_vars', 'global_vars',
1001                  'par_multi_ordered']
1002
1003         ret = {}
1004
1005         for name in names:
1006             ret[name] = getattr(self, name)
1007
1008         ret['par_single'] = self.parser.par_single
1009         ret['par_multi'] = self.parser.par_multi
1010
1011         return ret
1012
1013     def save_block(self):
1014         """Save the current block into the state of the solver if necessary."""
1015
1016         # If no dump file is specified, do not store any data in state_results.
1017         if self.options.dump_filename is None:
1018             return
1019
1020         cuda.memcpy_dtoh(self._rng_state, self._gpu_rng_state)
1021         for i in range(0, self.num_vars):
1022             cuda.memcpy_dtoh(self.vec[i], self._gpu_vec[i])
1023
1024         if self.options.omode == 'path':
1025             self.state_results.append((self.sim_t, copy.deepcopy(self.vec), self.vec_nx.copy(),
1026                 self._rng_state.copy()))
1027         else:
1028              self.state_results.append((self.sim_t, copy.deepcopy(self.vec), self.vec_nx.copy(),
1029                 self._rng_state.copy(), self.start_t, copy.deepcopy(self.vec_start),
1030                 self.vec_start_nx.copy()))
1031
1032     def dump_state(self):
1033         """Dump the current state of the solver to a file.
1034
1035         This makes it possible to later restart the calculations from the saved
1036         checkpoint using the :meth:`load_state` function.
1037         """
1038         state = self.state
1039         state['results'] = self.state_results
1040         state['numpy.random'] = numpy.random.get_state()
1041
1042         if self.scan_vars:
1043             state['sv'] = self._sv.copy()
1044
1045         with open(self.options.dump_filename, 'w') as f:
1046             pickle.dump(state, f, pickle.HIGHEST_PROTOCOL)
1047
1048         self.last_dump = time.time()
1049
1050     def load_state(self):
1051         """Restore saved state of the solver.
1052
1053         After the state is restored, the simulation can be continued the standard
1054         way, i.e. by calling :meth:`prepare` and :meth:`simulate`.
1055         """
1056         with open(self.options.restore_filename, 'r') as f:
1057             state = pickle.load(f)
1058
1059         # The current options object will be overriden by the one saved in the
1060         # checkpoint file.
1061         new_options = self.options
1062         for par, val in state.iteritems():
1063             if par not in ['par_single', 'par_multi', 'results', 'numpy.random', 'sv']:
1064                 setattr(self, par, val)
1065
1066         self.parser.par_single = state['par_single']
1067         self.parser.par_multi = state['par_multi']
1068         self.prev_state_results = state['results']
1069
1070         numpy.random.set_state(state['numpy.random'])
1071
1072         if 'sv' in state:
1073             self._sv = state['sv']
1074
1075         # Options overridable from the command line.
1076         overridable = ['resume', 'continue_', 'dump_filename']
1077
1078         # If this is a continuation of a previous simulation, make output-related
1079         # parameters overridable.
1080         if new_options.continue_:
1081             overridable.extend(['output', 'oformat', 'omode', 'simperiods'])
1082             # TODO: This could potentially cause problems with transients if the original
1083             # simulation was run in summary mode and the new one is in path mode.
1084         else:
1085             self.state_results = self.prev_state_results
1086
1087         for option in overridable:
1088             if hasattr(new_options, option) and getattr(new_options, option) is not None:
1089                 setattr(self.options, option, getattr(new_options, option))
1090