1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
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 dtype=num.float32,
22 impl='openmp'):
24 if nparallel is None:
25 import multiprocessing
26 nparallel = multiprocessing.cpu_count()
28 narrays = offsets.size
29 assert len(arrays) == narrays
30 nshifts = shifts.size // narrays
31 assert shifts.shape == (nshifts, narrays)
32 shifts = num.reshape(shifts, (nshifts*narrays))
33 assert weights.shape == (nshifts, narrays)
34 weights = num.reshape(weights, (nshifts*narrays))
36 weights = weights.astype(dtype, copy=False)
37 arrays = [arr.astype(dtype, copy=False) for arr in arrays]
39 if impl == 'openmp':
40 parstack_impl = parstack_ext.parstack
41 elif impl == 'numpy':
42 parstack_impl = parstack_numpy
44 result, offset = parstack_impl(
45 arrays, offsets, shifts, weights, method,
46 lengthout, offsetout, result, nparallel)
48 if method == 0:
49 nsamps = result.size // nshifts
50 result = result.reshape((nshifts, nsamps))
52 return result, offset
55def get_offset_and_length(arrays, offsets, shifts):
56 narrays = offsets.size
57 nshifts = shifts.size // narrays
58 if shifts.ndim == 2:
59 shifts = num.reshape(shifts, (nshifts*narrays))
61 lengths = num.array([a.size for a in arrays], dtype=int)
62 imin = offsets[0] + shifts[0]
63 imax = imin + lengths[0]
64 for iarray in range(len(arrays)):
65 istarts = offsets[iarray] + shifts[iarray::narrays]
66 iends = istarts + lengths[iarray]
67 imin = min(imin, num.amin(istarts))
68 imax = max(imax, num.amax(iends))
70 return imin, imax - imin
73def parstack_numpy(
74 arrays,
75 offsets,
76 shifts,
77 weights,
78 method,
79 lengthout,
80 offsetout,
81 result,
82 nparallel):
84 # nparallel is ignored here
86 narrays = offsets.size
88 lengths = num.array([a.size for a in arrays], dtype=int)
89 if lengthout < 0:
90 imin, nsamp = get_offset_and_length(arrays, offsets, shifts)
91 else:
92 nsamp = lengthout
93 imin = offsetout
95 nshifts = shifts.size // narrays
96 result = num.zeros(nsamp*nshifts, dtype=float)
98 for ishift in range(nshifts):
99 for iarray in range(narrays):
100 istart = offsets[iarray] + shifts[ishift*narrays + iarray]
101 weight = weights[ishift*narrays + iarray]
102 istart_r = ishift*nsamp + istart - imin
103 jstart = max(0, imin - istart)
104 jstop = max(jstart, min(nsamp - istart + imin, lengths[iarray]))
105 result[istart_r+jstart:istart_r+jstop] += \
106 arrays[iarray][jstart:jstop] * weight
108 if method == 0:
109 return result, imin
110 elif method == 1:
111 return num.amax(result.reshape((nshifts, nsamp)), axis=1), imin
114def argmax(a, nparallel=1):
115 '''
116 Same as numpys' argmax for 2 dimensional arrays along axis 0
117 but more memory efficient and twice as fast.
118 '''
120 return parstack_ext.argmax(a, nparallel)