1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
5from __future__ import absolute_import
6import numpy as num
8from . import parstack_ext
10try:
11 range = xrange
12except NameError:
13 pass
16def parstack(arrays, offsets, shifts, weights, method,
17 lengthout=-1,
18 offsetout=0,
19 result=None,
20 nparallel=None,
21 impl='openmp'):
23 if nparallel is None:
24 import multiprocessing
25 nparallel = multiprocessing.cpu_count()
27 narrays = offsets.size
28 assert(len(arrays) == narrays)
29 nshifts = shifts.size // narrays
30 assert shifts.shape == (nshifts, narrays)
31 shifts = num.reshape(shifts, (nshifts*narrays))
32 assert weights.shape == (nshifts, narrays)
33 weights = num.reshape(weights, (nshifts*narrays))
35 if impl == 'openmp':
36 parstack_impl = parstack_ext.parstack
37 elif impl == 'numpy':
38 parstack_impl = parstack_numpy
40 result, offset = parstack_impl(
41 arrays, offsets, shifts, weights, method,
42 lengthout, offsetout, result, nparallel)
44 if method == 0:
45 nsamps = result.size // nshifts
46 result = result.reshape((nshifts, nsamps))
48 return result, offset
51def get_offset_and_length(arrays, offsets, shifts):
52 narrays = offsets.size
53 nshifts = shifts.size // narrays
54 if shifts.ndim == 2:
55 shifts = num.reshape(shifts, (nshifts*narrays))
57 lengths = num.array([a.size for a in arrays], dtype=int)
58 imin = offsets[0] + shifts[0]
59 imax = imin + lengths[0]
60 for iarray in range(len(arrays)):
61 istarts = offsets[iarray] + shifts[iarray::narrays]
62 iends = istarts + lengths[iarray]
63 imin = min(imin, num.amin(istarts))
64 imax = max(imax, num.amax(iends))
66 return imin, imax - imin
69def parstack_numpy(
70 arrays,
71 offsets,
72 shifts,
73 weights,
74 method,
75 lengthout,
76 offsetout,
77 result,
78 nparallel):
80 # nparallel is ignored here
82 narrays = offsets.size
84 lengths = num.array([a.size for a in arrays], dtype=int)
85 if lengthout < 0:
86 imin, nsamp = get_offset_and_length(arrays, offsets, shifts)
87 else:
88 nsamp = lengthout
89 imin = offsetout
91 nshifts = shifts.size // narrays
92 result = num.zeros(nsamp*nshifts, dtype=float)
94 for ishift in range(nshifts):
95 for iarray in range(narrays):
96 istart = offsets[iarray] + shifts[ishift*narrays + iarray]
97 weight = weights[ishift*narrays + iarray]
98 istart_r = ishift*nsamp + istart - imin
99 jstart = max(0, imin - istart)
100 jstop = max(jstart, min(nsamp - istart + imin, lengths[iarray]))
101 result[istart_r+jstart:istart_r+jstop] += \
102 arrays[iarray][jstart:jstop] * weight
104 if method == 0:
105 return result, imin
106 elif method == 1:
107 return num.amax(result.reshape((nshifts, nsamp)), axis=1), imin
110def argmax(a, nparallel=1):
111 '''
112 Same as numpys' argmax for 2 dimensional arrays along axis 0
113 but more memory efficient and twice as fast.
114 '''
116 return parstack_ext.argmax(a, nparallel)