Coverage for eminus/backend.py: 49.38%
81 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-26 09:44 +0000
1# SPDX-FileCopyrightText: 2025 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""Array backend handling.
5Includes helper and compatibility functions.
6For more information see https://wangenau.gitlab.io/eminus/backend.html.
7"""
9import pathlib
10import sys
12import numpy as np
13import scipy
15from . import config
17if "stubtest" not in pathlib.Path(sys.argv[0]).name:
18 # Do not overwrite getattr when stubtest is running
19 def __getattr__(name):
20 """Access modules and functions of array backends by their name."""
21 if config.backend == "torch":
22 from array_api_compat import torch as xp
23 elif config.backend == "cupy":
24 from array_api_compat import cupy as xp
25 else:
26 xp = np
27 return getattr(xp, name)
30# ### Helper functions ###
33def is_array(value):
34 """Check if the object is coming from NumPy, CuPy, or Torch.
36 Args:
37 value: Input array.
39 Returns:
40 If the value is an array supported by the available backends.
41 """
42 if isinstance(value, np.ndarray):
43 return True
44 if config.backend == "torch":
45 from array_api_compat import is_torch_array
47 return is_torch_array(value)
48 if config.backend == "cupy":
49 from array_api_compat import is_cupy_array
51 return is_cupy_array(value)
52 return False
55def to_np(arr):
56 """Copy the array from the current device to a CPU NumPy array.
58 Args:
59 arr: Input array.
61 Returns:
62 Copied array on the CPU.
63 """
64 if arr is None:
65 return None
66 try:
67 from array_api_compat import is_cupy_array, is_torch_array
69 if is_torch_array(arr):
70 return np.asarray(arr.resolve_conj().cpu())
71 # Wave functions are a list of arrays, handle them as well
72 if isinstance(arr, list) and len(arr) > 0 and is_torch_array(arr[0]):
73 return [np.asarray(arr.resolve_conj().cpu()) for arr in arr]
75 if is_cupy_array(arr):
76 return arr.get()
77 # Wave functions are a list of arrays, handle them as well
78 if isinstance(arr, list) and len(arr) > 0 and is_cupy_array(arr[0]):
79 return [arr.get() for arr in arr]
81 return np.asarray(arr)
82 except ImportError:
83 return np.asarray(arr)
86# ### Compatibility functions ###
89def delete(arr, obj, axis=None):
90 """Return a new array with sub-arrays along an axis deleted.
92 Ref: https://gist.github.com/velikodniy/6efef837e67aee2e7152eb5900eb0258
94 Args:
95 arr: Input array.
96 obj: Indicate indices of sub-arrays to remove along the specified axis.
98 Keyword Args:
99 axis: The axis along which to delete the subarray defined by obj. If `axis` is `None`, `obj`
100 is applied to the flattened array.
102 Returns:
103 A copy of `arr` with the elements specified by `obj` removed. If `axis` is `None`, `out` is
104 a flattened array.
105 """
106 if isinstance(arr, np.ndarray):
107 return np.delete(arr, obj, axis)
108 from array_api_compat import array_namespace, is_torch_array
110 if is_torch_array(arr):
111 if axis is None:
112 axis = 0
113 arr = arr.ravel()
114 if is_array(obj):
115 obj = to_np(obj)
116 else:
117 obj = np.asarray(obj)
118 skip = [i for i in range(arr.size(axis)) if i not in obj]
119 indices = tuple(slice(None) if i != axis else skip for i in range(arr.ndim))
120 return arr[indices]
122 xp = array_namespace(arr)
123 return xp.delete(arr, obj, axis)
126def indices(dimensions, dtype=int, *args, **kwargs):
127 """Return an array representing the indices of a grid.
129 Args:
130 dimensions: The shape of the grid.
131 dtype: Data type of the result.
132 args: Pass-through arguments.
134 Keyword Args:
135 **kwargs: Pass-through keyword arguments.
137 Returns:
138 Returns a tuple of arrays.
139 """
140 if config.backend == "torch":
141 from array_api_compat import torch as xp
143 return xp.asarray(np.indices(to_np(dimensions), dtype, *args, **kwargs))
145 if config.backend == "cupy":
146 from array_api_compat import cupy as xp
148 return xp.indices(tuple(to_np(dimensions)), dtype, *args, **kwargs)
150 return np.indices(dimensions, dtype, *args, **kwargs)
153def fftn(x, *args, **kwargs):
154 """Compute the N-D discrete Fourier Transform.
156 Use SciPy FFTs since they are faster, support parallelism, and are more accurate.
157 They will upcast some complex arrays to complex256. Using NumPy FFTs will fail some tests.
159 Args:
160 x: Input array, can be complex.
161 args: Pass-through arguments.
163 Keyword Args:
164 **kwargs: Pass-through keyword arguments.
166 Returns:
167 Value of the fftn function at x.
168 """
169 if isinstance(x, np.ndarray):
170 return scipy.fft.fftn(x, *args, **kwargs, workers=config.threads)
171 from array_api_compat import array_namespace
173 xp = array_namespace(x)
174 return xp.fft.fftn(x, *args, **kwargs)
177def ifftn(x, *args, **kwargs):
178 """Compute the N-D inverse discrete Fourier Transform.
180 Use SciPy FFTs since they are faster, support parallelism, and are more accurate.
181 They will upcast some complex arrays to complex256. Using NumPy FFTs will fail some tests.
183 Args:
184 x: Input array, can be complex.
185 args: Pass-through arguments.
187 Keyword Args:
188 **kwargs: Pass-through keyword arguments.
190 Returns:
191 Value of the ifftn function at x.
192 """
193 if isinstance(x, np.ndarray):
194 return scipy.fft.ifftn(x, *args, **kwargs, workers=config.threads)
195 from array_api_compat import array_namespace
197 xp = array_namespace(x)
198 return xp.fft.ifftn(x, *args, **kwargs)
201def sqrtm(A, *args, **kwargs):
202 """Matrix square root.
204 Args:
205 A: Matrix whose square root to evaluate.
206 args: Pass-through arguments.
208 Keyword Args:
209 **kwargs: Pass-through keyword arguments.
211 Returns:
212 Value of the sqrt function at A.
213 """
214 if isinstance(A, np.ndarray):
215 return np.asarray(scipy.linalg.sqrtm(A, *args, **kwargs), dtype=complex)
216 from array_api_compat import array_namespace
218 xp = array_namespace(A)
219 return xp.asarray(
220 np.asarray(scipy.linalg.sqrtm(to_np(A), *args, **kwargs), dtype=complex), dtype=complex
221 )