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 impl='openmp'): 

22 

23 if nparallel is None: 

24 import multiprocessing 

25 nparallel = multiprocessing.cpu_count() 

26 

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

34 

35 if impl == 'openmp': 

36 parstack_impl = parstack_ext.parstack 

37 elif impl == 'numpy': 

38 parstack_impl = parstack_numpy 

39 

40 result, offset = parstack_impl( 

41 arrays, offsets, shifts, weights, method, 

42 lengthout, offsetout, result, nparallel) 

43 

44 if method == 0: 

45 nsamps = result.size // nshifts 

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

47 

48 return result, offset 

49 

50 

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

56 

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

65 

66 return imin, imax - imin 

67 

68 

69def parstack_numpy( 

70 arrays, 

71 offsets, 

72 shifts, 

73 weights, 

74 method, 

75 lengthout, 

76 offsetout, 

77 result, 

78 nparallel): 

79 

80 # nparallel is ignored here 

81 

82 narrays = offsets.size 

83 

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 

90 

91 nshifts = shifts.size // narrays 

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

93 

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 

103 

104 if method == 0: 

105 return result, imin 

106 elif method == 1: 

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

108 

109 

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

115 

116 return parstack_ext.argmax(a, nparallel)