Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/scipy/fft/_backend.py: 58%

1import scipy._lib.uarray as ua 

2from . import _fftlog 

3from . import _pocketfft 

4 

5 

6class _ScipyBackend: 

7 """The default backend for fft calculations 

8 

9 Notes 

10 ----- 

11 We use the domain ``numpy.scipy`` rather than ``scipy`` because ``uarray`` 

12 treats the domain as a hierarchy. This means the user can install a single 

13 backend for ``numpy`` and have it implement ``numpy.scipy.fft`` as well. 

14 """ 

15 __ua_domain__ = "numpy.scipy.fft" 

16 

17 @staticmethod 

18 def __ua_function__(method, args, kwargs): 

19 

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

21 if fn is None: 

22 fn = getattr(_fftlog, method.__name__, None) 

23 if fn is None: 

24 return NotImplemented 

25 return fn(*args, **kwargs) 

26 

27 

28_named_backends = { 

29 'scipy': _ScipyBackend, 

30} 

31 

32 

33def _backend_from_arg(backend): 

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

35 

36 if isinstance(backend, str): 

37 try: 

38 backend = _named_backends[backend] 

39 except KeyError as e: 

40 raise ValueError('Unknown backend {}'.format(backend)) from e 

41 

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

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

44 

45 return backend 

46 

47 

48def set_global_backend(backend, coerce=False, only=False, try_last=False): 

49 """Sets the global fft backend 

50 

51 This utility method replaces the default backend for permanent use. It 

52 will be tried in the list of backends automatically, unless the 

53 ``only`` flag is set on a backend. This will be the first tried 

54 backend outside the :obj:`set_backend` context manager. 

55 

56 Parameters 

57 ---------- 

58 backend : {object, 'scipy'} 

59 The backend to use. 

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

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

62 coerce : bool 

63 Whether to coerce input types when trying this backend. 

64 only : bool 

65 If ``True``, no more backends will be tried if this fails. 

66 Implied by ``coerce=True``. 

67 try_last : bool 

68 If ``True``, the global backend is tried after registered backends. 

69 

70 Raises 

71 ------ 

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

73 

74 Notes 

75 ----- 

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

77 the SciPy implementation. 

78 

79 Examples 

80 -------- 

81 We can set the global fft backend: 

82 

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

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

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

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

87 """ 

88 backend = _backend_from_arg(backend) 

89 ua.set_global_backend(backend, coerce=coerce, only=only, try_last=try_last) 

90 

91 

92def register_backend(backend): 

93 """ 

94 Register a backend for permanent use. 

95 

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

97 global backend. 

98 

99 Parameters 

100 ---------- 

101 backend : {object, 'scipy'} 

102 The backend to use. 

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

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

105 

106 Raises 

107 ------ 

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

109 

110 Examples 

111 -------- 

112 We can register a new fft backend: 

113 

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

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

116 ... __ua_domain__ = "numpy.scipy.fft" 

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

118 ... return NotImplemented 

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

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

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

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

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

124 

125 """ 

126 backend = _backend_from_arg(backend) 

127 ua.register_backend(backend) 

128 

129 

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

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

132 

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

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

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

136 

137 Parameters 

138 ---------- 

139 backend : {object, 'scipy'} 

140 The backend to use. 

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

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

143 coerce : bool, optional 

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

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

146 only : bool, optional 

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

148 BackendNotImplemented error will be raised immediately. Ignoring any 

149 lower priority backends. 

150 

151 Examples 

152 -------- 

153 >>> import scipy.fft as fft 

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

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

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

157 """ 

158 backend = _backend_from_arg(backend) 

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

160 

161 

162def skip_backend(backend): 

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

164 

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

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

167 exit, the backend will again be considered. 

168 

169 Parameters 

170 ---------- 

171 backend : {object, 'scipy'} 

172 The backend to skip. 

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

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

175 

176 Examples 

177 -------- 

178 >>> import scipy.fft as fft 

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

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

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

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

183 Traceback (most recent call last): 

184 ... 

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

186 """ 

187 backend = _backend_from_arg(backend) 

188 return ua.skip_backend(backend) 

189 

190 

191set_global_backend('scipy', try_last=True)