Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/parstack.py: 100%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Fast weight-delay-and-sum implementation for seismic array processing and 

8migration. 

9''' 

10 

11import numpy as num 

12 

13from . import parstack_ext 

14 

15try: 

16 range = xrange 

17except NameError: 

18 pass 

19 

20import multiprocessing 

21 

22nparallel = multiprocessing.cpu_count() 

23 

24 

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

26 lengthout=-1, 

27 offsetout=0, 

28 result=None, 

29 nparallel=nparallel, 

30 dtype=num.float64, 

31 impl='openmp'): 

32 narrays = offsets.size 

33 assert len(arrays) == narrays 

34 nshifts = shifts.size // narrays 

35 assert shifts.shape == (nshifts, narrays) 

36 shifts = num.reshape(shifts, (nshifts*narrays)) 

37 assert weights.shape == (nshifts, narrays) 

38 weights = num.reshape(weights, (nshifts*narrays)) 

39 

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

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

42 

43 if impl == 'openmp': 

44 parstack_impl = parstack_ext.parstack 

45 elif impl == 'numpy': 

46 parstack_impl = parstack_numpy 

47 

48 result, offset = parstack_impl( 

49 arrays, offsets, shifts, weights, method, 

50 lengthout, offsetout, result, nparallel) 

51 

52 if method == 0: 

53 nsamps = result.size // nshifts 

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

55 

56 return result, offset 

57 

58 

59def get_offset_and_length(arrays, offsets, shifts): 

60 narrays = offsets.size 

61 nshifts = shifts.size // narrays 

62 if shifts.ndim == 2: 

63 shifts = num.reshape(shifts, (nshifts*narrays)) 

64 

65 lengths = num.array([a.size for a in arrays], dtype=int) 

66 imin = offsets[0] + shifts[0] 

67 imax = imin + lengths[0] 

68 for iarray in range(len(arrays)): 

69 istarts = offsets[iarray] + shifts[iarray::narrays] 

70 iends = istarts + lengths[iarray] 

71 imin = min(imin, num.amin(istarts)) 

72 imax = max(imax, num.amax(iends)) 

73 

74 return imin, imax - imin 

75 

76 

77def parstack_numpy( 

78 arrays, 

79 offsets, 

80 shifts, 

81 weights, 

82 method, 

83 lengthout, 

84 offsetout, 

85 result, 

86 nparallel): 

87 

88 # nparallel is ignored here 

89 

90 narrays = offsets.size 

91 

92 lengths = num.array([a.size for a in arrays], dtype=int) 

93 if lengthout < 0: 

94 imin, nsamp = get_offset_and_length(arrays, offsets, shifts) 

95 else: 

96 nsamp = lengthout 

97 imin = offsetout 

98 

99 nshifts = shifts.size // narrays 

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

101 

102 for ishift in range(nshifts): 

103 for iarray in range(narrays): 

104 istart = offsets[iarray] + shifts[ishift*narrays + iarray] 

105 weight = weights[ishift*narrays + iarray] 

106 istart_r = ishift*nsamp + istart - imin 

107 jstart = max(0, imin - istart) 

108 jstop = max(jstart, min(nsamp - istart + imin, lengths[iarray])) 

109 result[istart_r+jstart:istart_r+jstop] += \ 

110 arrays[iarray][jstart:jstop] * weight 

111 

112 if method == 0: 

113 return result, imin 

114 elif method == 1: 

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

116 

117 

118def argmax(a, nparallel=1): 

119 ''' 

120 Same as numpys' argmax for 2 dimensional arrays along axis 0 

121 but more memory efficient and twice as fast. 

122 ''' 

123 

124 return parstack_ext.argmax(a, nparallel)