1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

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

from __future__ import absolute_import 

from builtins import range 

import numpy as num 

 

from . import parstack_ext 

 

 

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

lengthout=-1, 

offsetout=0, 

result=None, 

nparallel=None, 

impl='openmp'): 

 

if nparallel is None: 

import multiprocessing 

nparallel = multiprocessing.cpu_count() 

 

narrays = offsets.size 

assert(len(arrays) == narrays) 

nshifts = shifts.size // narrays 

assert shifts.shape == (nshifts, narrays) 

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

assert weights.shape == (nshifts, narrays) 

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

 

if impl == 'openmp': 

parstack_impl = parstack_ext.parstack 

elif impl == 'numpy': 

parstack_impl = parstack_numpy 

 

result, offset = parstack_impl( 

arrays, offsets, shifts, weights, method, 

lengthout, offsetout, result, nparallel) 

 

if method == 0: 

nsamps = result.size // nshifts 

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

 

return result, offset 

 

 

def get_offset_and_length(arrays, offsets, shifts): 

narrays = offsets.size 

nshifts = shifts.size // narrays 

if shifts.ndim == 2: 

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

 

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

imin = offsets[0] + shifts[0] 

imax = imin + lengths[0] 

for iarray in range(len(arrays)): 

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

iends = istarts + lengths[iarray] 

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

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

 

return imin, imax - imin 

 

 

def parstack_numpy( 

arrays, 

offsets, 

shifts, 

weights, 

method, 

lengthout, 

offsetout, 

result, 

nparallel): 

 

# nparallel is ignored here 

 

narrays = offsets.size 

 

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

if lengthout < 0: 

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

else: 

nsamp = lengthout 

imin = offsetout 

 

nshifts = shifts.size // narrays 

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

 

for ishift in range(nshifts): 

for iarray in range(narrays): 

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

weight = weights[ishift*narrays + iarray] 

istart_r = ishift*nsamp + istart - imin 

jstart = max(0, imin - istart) 

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

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

arrays[iarray][jstart:jstop] * weight 

 

if method == 0: 

return result, imin 

elif method == 1: 

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

 

 

def argmax(a, nparallel=1): 

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

but more memory efficient and twice as fast.''' 

 

return parstack_ext.argmax(a, nparallel)