gpuarray 0.2.0__tar.gz → 0.4.0__tar.gz

{gpuarray-0.2.0 → gpuarray-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.2.0
+Version: 0.4.0
 Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
 Author: Arnav Kewalram
 License: MIT

{gpuarray-0.2.0 → gpuarray-0.4.0}/gpuarray/__init__.py

@@ -27,6 +27,10 @@ from .api import (
     concatenate,
     stack,
     random,
+    fft,
+    ifft,
+    fft_freq,
+    fft_magnitude,
 )
 
 __version__ = "0.1.0"
@@ -58,4 +62,8 @@ __all__ = [
     "concatenate",
     "stack",
     "random",
+    "fft",
+    "ifft",
+    "fft_freq",
+    "fft_magnitude",
 ]

{gpuarray-0.2.0 → gpuarray-0.4.0}/gpuarray/api.py

@@ -6,7 +6,7 @@ from typing import Sequence
 
 import numpy as np
 
-from .core import GPUArray, _from_numpy, matmul as _matmul
+from .core import GPUArray, _from_numpy, matmul as _matmul, _complex_magnitude
 from . import shaders
 
 
@@ -180,3 +180,37 @@ class _Random:
 
 
 random = _Random()
+
+
+# ----- FFT functions -----
+
+def fft(a) -> GPUArray:
+    """1D FFT. Input is real or complex (interleaved). Returns complex GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    return a.fft()
+
+
+def ifft(a) -> GPUArray:
+    """1D inverse FFT. Input is complex interleaved. Returns complex GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    return a.ifft()
+
+
+def fft_freq(n: int, d: float = 1.0) -> GPUArray:
+    """Return the DFT sample frequencies (like numpy.fft.fftfreq)."""
+    results = np.empty(n, dtype=np.float32)
+    N = (n - 1) // 2 + 1
+    results[:N] = np.arange(0, N, dtype=np.float32)
+    results[N:] = np.arange(-(n // 2), 0, dtype=np.float32)
+    results /= (n * d)
+    return _from_numpy(results)
+
+
+def fft_magnitude(a) -> GPUArray:
+    """Compute magnitude spectrum |FFT(x)|. Returns real GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    fft_result = a.fft()
+    return _complex_magnitude(fft_result)