Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from distutils.version import LooseVersion 

2 

3import numpy as np 

4import scipy 

5 

6SCIPY_11 = (LooseVersion(scipy.__version__) < LooseVersion('1.2.0') and 

7 LooseVersion(scipy.__version__) >= LooseVersion('1.1.0')) 

8 

9SCIPY_GT_14 = LooseVersion(scipy.__version__) >= LooseVersion('1.5') 

10 

11 

12def _next_regular(target): 

13 """ 

14 Find the next regular number greater than or equal to target. 

15 Regular numbers are composites of the prime factors 2, 3, and 5. 

16 Also known as 5-smooth numbers or Hamming numbers, these are the optimal 

17 size for inputs to FFTPACK. 

18 

19 Target must be a positive integer. 

20 """ 

21 if target <= 6: 

22 return target 

23 

24 # Quickly check if it's already a power of 2 

25 if not (target & (target - 1)): 

26 return target 

27 

28 match = float('inf') # Anything found will be smaller 

29 p5 = 1 

30 while p5 < target: 

31 p35 = p5 

32 while p35 < target: 

33 # Ceiling integer division, avoiding conversion to float 

34 # (quotient = ceil(target / p35)) 

35 quotient = -(-target // p35) 

36 # Quickly find next power of 2 >= quotient 

37 p2 = 2 ** ((quotient - 1).bit_length()) 

38 

39 N = p2 * p35 

40 if N == target: 

41 return N 

42 elif N < match: 

43 match = N 

44 p35 *= 3 

45 if p35 == target: 

46 return p35 

47 if p35 < match: 

48 match = p35 

49 p5 *= 5 

50 if p5 == target: 

51 return p5 

52 if p5 < match: 

53 match = p5 

54 return match 

55 

56 

57def _valarray(shape, value=np.nan, typecode=None): 

58 """Return an array of all value. 

59 """ 

60 

61 out = np.ones(shape, dtype=bool) * value 

62 if typecode is not None: 

63 out = out.astype(typecode) 

64 if not isinstance(out, np.ndarray): 

65 out = np.asarray(out) 

66 return out 

67 

68 

69def _lazywhere(cond, arrays, f, fillvalue=None, f2=None): 

70 """ 

71 np.where(cond, x, fillvalue) always evaluates x even where cond is False. 

72 This one only evaluates f(arr1[cond], arr2[cond], ...). 

73 For example, 

74 >>> a, b = np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]) 

75 >>> def f(a, b): 

76 return a*b 

77 >>> _lazywhere(a > 2, (a, b), f, np.nan) 

78 array([ nan, nan, 21., 32.]) 

79 Notice it assumes that all `arrays` are of the same shape, or can be 

80 broadcasted together. 

81 """ 

82 if fillvalue is None: 

83 if f2 is None: 

84 raise ValueError("One of (fillvalue, f2) must be given.") 

85 else: 

86 fillvalue = np.nan 

87 else: 

88 if f2 is not None: 

89 raise ValueError("Only one of (fillvalue, f2) can be given.") 

90 

91 arrays = np.broadcast_arrays(*arrays) 

92 temp = tuple(np.extract(cond, arr) for arr in arrays) 

93 tcode = np.mintypecode([a.dtype.char for a in arrays]) 

94 out = _valarray(np.shape(arrays[0]), value=fillvalue, typecode=tcode) 

95 np.place(out, cond, f(*temp)) 

96 if f2 is not None: 

97 temp = tuple(np.extract(~cond, arr) for arr in arrays) 

98 np.place(out, ~cond, f2(*temp)) 

99 

100 return out