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

1import scipy._lib.uarray as ua 

2from . import _pocketfft 

3 

4 

5class _ScipyBackend: 

6 """The default backend for fft calculations 

7 

8 Notes 

9 ----- 

10 We use the domain ``numpy.scipy`` rather than ``scipy`` because in the 

11 future, ``uarray`` will treat the domain as a hierarchy. This means the user 

12 can install a single backend for ``numpy`` and have it implement 

13 ``numpy.scipy.fft`` as well. 

14 """ 

15 __ua_domain__ = "numpy.scipy.fft" 

16 

17 @staticmethod 

18 def __ua_function__(method, args, kwargs): 

19 fn = getattr(_pocketfft, method.__name__, None) 

20 

21 if fn is None: 

22 return NotImplemented 

23 return fn(*args, **kwargs) 

24 

25 

26_named_backends = { 

27 'scipy': _ScipyBackend, 

28} 

29 

30 

31def _backend_from_arg(backend): 

32 """Maps strings to known backends and validates the backend""" 

33 

34 if isinstance(backend, str): 

35 try: 

36 backend = _named_backends[backend] 

37 except KeyError: 

38 raise ValueError('Unknown backend {}'.format(backend)) 

39 

40 if backend.__ua_domain__ != 'numpy.scipy.fft': 

41 raise ValueError('Backend does not implement "numpy.scipy.fft"') 

42 

43 return backend 

44 

45 

46def set_global_backend(backend): 

47 """Sets the global fft backend 

48 

49 The global backend has higher priority than registered backends, but lower 

50 priority than context-specific backends set with `set_backend`. 

51 

52 Parameters 

53 ---------- 

54 backend: {object, 'scipy'} 

55 The backend to use. 

56 Can either be a ``str`` containing the name of a known backend 

57 {'scipy'} or an object that implements the uarray protocol. 

58 

59 Raises 

60 ------ 

61 ValueError: If the backend does not implement ``numpy.scipy.fft``. 

62 

63 Notes 

64 ----- 

65 This will overwrite the previously set global backend, which, by default, is 

66 the SciPy implementation. 

67 

68 Examples 

69 -------- 

70 We can set the global fft backend: 

71 

72 >>> from scipy.fft import fft, set_global_backend 

73 >>> set_global_backend("scipy") # Sets global backend. "scipy" is the default backend. 

74 >>> fft([1]) # Calls the global backend 

75 array([1.+0.j]) 

76 """ 

77 backend = _backend_from_arg(backend) 

78 ua.set_global_backend(backend) 

79 

80 

81def register_backend(backend): 

82 """ 

83 Register a backend for permanent use. 

84 

85 Registered backends have the lowest priority and will be tried after the 

86 global backend. 

87 

88 Parameters 

89 ---------- 

90 backend: {object, 'scipy'} 

91 The backend to use. 

92 Can either be a ``str`` containing the name of a known backend 

93 {'scipy'} or an object that implements the uarray protocol. 

94 

95 Raises 

96 ------ 

97 ValueError: If the backend does not implement ``numpy.scipy.fft``. 

98 

99 Examples 

100 -------- 

101 We can register a new fft backend: 

102 

103 >>> from scipy.fft import fft, register_backend, set_global_backend 

104 >>> class NoopBackend: # Define an invalid Backend 

105 ... __ua_domain__ = "numpy.scipy.fft" 

106 ... def __ua_function__(self, func, args, kwargs): 

107 ... return NotImplemented 

108 >>> set_global_backend(NoopBackend()) # Set the invalid backend as global 

109 >>> register_backend("scipy") # Register a new backend 

110 >>> fft([1]) # The registered backend is called because the global backend returns `NotImplemented` 

111 array([1.+0.j]) 

112 >>> set_global_backend("scipy") # Restore global backend to default 

113 

114 """ 

115 backend = _backend_from_arg(backend) 

116 ua.register_backend(backend) 

117 

118 

119def set_backend(backend, coerce=False, only=False): 

120 """Context manager to set the backend within a fixed scope. 

121 

122 Upon entering the ``with`` statement, the given backend will be added to 

123 the list of available backends with the highest priority. Upon exit, the 

124 backend is reset to the state before entering the scope. 

125 

126 Parameters 

127 ---------- 

128 backend: {object, 'scipy'} 

129 The backend to use. 

130 Can either be a ``str`` containing the name of a known backend 

131 {'scipy'} or an object that implements the uarray protocol. 

132 coerce: bool, optional 

133 Whether to allow expensive conversions for the ``x`` parameter. e.g., 

134 copying a NumPy array to the GPU for a CuPy backend. Implies ``only``. 

135 only: bool, optional 

136 If only is ``True`` and this backend returns ``NotImplemented``, then a 

137 BackendNotImplemented error will be raised immediately. Ignoring any 

138 lower priority backends. 

139 

140 Examples 

141 -------- 

142 >>> import scipy.fft as fft 

143 >>> with fft.set_backend('scipy', only=True): 

144 ... fft.fft([1]) # Always calls the scipy implementation 

145 array([1.+0.j]) 

146 """ 

147 backend = _backend_from_arg(backend) 

148 return ua.set_backend(backend, coerce=coerce, only=only) 

149 

150 

151def skip_backend(backend): 

152 """Context manager to skip a backend within a fixed scope. 

153 

154 Within the context of a ``with`` statement, the given backend will not be 

155 called. This covers backends registered both locally and globally. Upon 

156 exit, the backend will again be considered. 

157 

158 Parameters 

159 ---------- 

160 backend: {object, 'scipy'} 

161 The backend to skip. 

162 Can either be a ``str`` containing the name of a known backend 

163 {'scipy'} or an object that implements the uarray protocol. 

164 

165 Examples 

166 -------- 

167 >>> import scipy.fft as fft 

168 >>> fft.fft([1]) # Calls default SciPy backend 

169 array([1.+0.j]) 

170 >>> with fft.skip_backend('scipy'): # We explicitly skip the SciPy backend 

171 ... fft.fft([1]) # leaving no implementation available 

172 Traceback (most recent call last): 

173 ... 

174 BackendNotImplementedError: No selected backends had an implementation ... 

175 """ 

176 backend = _backend_from_arg(backend) 

177 return ua.skip_backend(backend) 

178 

179 

180set_global_backend('scipy')