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 

20 

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

22 lengthout=-1, 

23 offsetout=0, 

24 result=None, 

25 nparallel=None, 

26 dtype=num.float64, 

27 impl='openmp'): 

28 

29 if nparallel is None: 

30 import multiprocessing 

31 nparallel = multiprocessing.cpu_count() 

32 

33 narrays = offsets.size 

34 assert len(arrays) == narrays 

35 nshifts = shifts.size // narrays 

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

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

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

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

40 

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

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

43 

44 if impl == 'openmp': 

45 parstack_impl = parstack_ext.parstack 

46 elif impl == 'numpy': 

47 parstack_impl = parstack_numpy 

48 

49 result, offset = parstack_impl( 

50 arrays, offsets, shifts, weights, method, 

51 lengthout, offsetout, result, nparallel) 

52 

53 if method == 0: 

54 nsamps = result.size // nshifts 

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

56 

57 return result, offset 

58 

59 

60def get_offset_and_length(arrays, offsets, shifts): 

61 narrays = offsets.size 

62 nshifts = shifts.size // narrays 

63 if shifts.ndim == 2: 

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

65 

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

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

68 imax = imin + lengths[0] 

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

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

71 iends = istarts + lengths[iarray] 

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

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

74 

75 return imin, imax - imin 

76 

77 

78def parstack_numpy( 

79 arrays, 

80 offsets, 

81 shifts, 

82 weights, 

83 method, 

84 lengthout, 

85 offsetout, 

86 result, 

87 nparallel): 

88 

89 # nparallel is ignored here 

90 

91 narrays = offsets.size 

92 

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

94 if lengthout < 0: 

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

96 else: 

97 nsamp = lengthout 

98 imin = offsetout 

99 

100 nshifts = shifts.size // narrays 

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

102 

103 for ishift in range(nshifts): 

104 for iarray in range(narrays): 

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

106 weight = weights[ishift*narrays + iarray] 

107 istart_r = ishift*nsamp + istart - imin 

108 jstart = max(0, imin - istart) 

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

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

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

112 

113 if method == 0: 

114 return result, imin 

115 elif method == 1: 

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

117 

118 

119def argmax(a, nparallel=1): 

120 ''' 

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

122 but more memory efficient and twice as fast. 

123 ''' 

124 

125 return parstack_ext.argmax(a, nparallel)