Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/scipy/fft/_backend.py : 59%

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
5class _ScipyBackend:
6 """The default backend for fft calculations
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"
17 @staticmethod
18 def __ua_function__(method, args, kwargs):
19 fn = getattr(_pocketfft, method.__name__, None)
21 if fn is None:
22 return NotImplemented
23 return fn(*args, **kwargs)
26_named_backends = {
27 'scipy': _ScipyBackend,
28}
31def _backend_from_arg(backend):
32 """Maps strings to known backends and validates the backend"""
34 if isinstance(backend, str):
35 try:
36 backend = _named_backends[backend]
37 except KeyError:
38 raise ValueError('Unknown backend {}'.format(backend))
40 if backend.__ua_domain__ != 'numpy.scipy.fft':
41 raise ValueError('Backend does not implement "numpy.scipy.fft"')
43 return backend
46def set_global_backend(backend):
47 """Sets the global fft backend
49 The global backend has higher priority than registered backends, but lower
50 priority than context-specific backends set with `set_backend`.
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.
59 Raises
60 ------
61 ValueError: If the backend does not implement ``numpy.scipy.fft``.
63 Notes
64 -----
65 This will overwrite the previously set global backend, which, by default, is
66 the SciPy implementation.
68 Examples
69 --------
70 We can set the global fft backend:
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)
81def register_backend(backend):
82 """
83 Register a backend for permanent use.
85 Registered backends have the lowest priority and will be tried after the
86 global backend.
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.
95 Raises
96 ------
97 ValueError: If the backend does not implement ``numpy.scipy.fft``.
99 Examples
100 --------
101 We can register a new fft backend:
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
114 """
115 backend = _backend_from_arg(backend)
116 ua.register_backend(backend)
119def set_backend(backend, coerce=False, only=False):
120 """Context manager to set the backend within a fixed scope.
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.
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.
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)
151def skip_backend(backend):
152 """Context manager to skip a backend within a fixed scope.
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.
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.
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)
180set_global_backend('scipy')