(resolution, mode="vector", dtype=None)
| 37 | |
| 38 | |
| 39 | def fftfreq(resolution, mode="vector", dtype=None): |
| 40 | assert mode in ("vector", "absolute", "square") |
| 41 | k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution], indexing="ij") |
| 42 | k = expand_dims(np.stack(k, -1), 0) |
| 43 | k = k.astype(float) |
| 44 | if mode == "vector": |
| 45 | return k |
| 46 | k = np.sum(k ** 2, axis=-1, keepdims=True) |
| 47 | if mode == "square": |
| 48 | return k |
| 49 | else: |
| 50 | return np.sqrt(k) |
| 51 | |
| 52 | |
| 53 | def expand_dims(a, axis=0, number=1): |
no test coverage detected