1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6import numpy as num 

7 

8from . import parstack_ext 

9 

10try: 

11 range = xrange 

12except NameError: 

13 pass 

14 

15 

16def parstack(arrays, offsets, shifts, weights, method, 

17 lengthout=-1, 

18 offsetout=0, 

19 result=None, 

20 nparallel=None, 

21 dtype=num.float64, 

22 impl='openmp'): 

23 

24 if nparallel is None: 

25 import multiprocessing 

26 nparallel = multiprocessing.cpu_count() 

27 

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)) 

35 

36 weights = weights.astype(dtype, copy=False) 

37 arrays = [arr.astype(dtype, copy=False) for arr in arrays] 

38 

39 if impl == 'openmp': 

40 parstack_impl = parstack_ext.parstack 

41 elif impl == 'numpy': 

42 parstack_impl = parstack_numpy 

43 

44 result, offset = parstack_impl( 

45 arrays, offsets, shifts, weights, method, 

46 lengthout, offsetout, result, nparallel) 

47 

48 if method == 0: 

49 nsamps = result.size // nshifts 

50 result = result.reshape((nshifts, nsamps)) 

51 

52 return result, offset 

53 

54 

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)) 

60 

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)) 

69 

70 return imin, imax - imin 

71 

72 

73def parstack_numpy( 

74 arrays, 

75 offsets, 

76 shifts, 

77 weights, 

78 method, 

79 lengthout, 

80 offsetout, 

81 result, 

82 nparallel): 

83 

84 # nparallel is ignored here 

85 

86 narrays = offsets.size 

87 

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 

94 

95 nshifts = shifts.size // narrays 

96 result = num.zeros(nsamp*nshifts, dtype=float) 

97 

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 

107 

108 if method == 0: 

109 return result, imin 

110 elif method == 1: 

111 return num.amax(result.reshape((nshifts, nsamp)), axis=1), imin 

112 

113 

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 ''' 

119 

120 return parstack_ext.argmax(a, nparallel)