refactor and cupy

Author: JaredChilders

notebooks/benchmark.ipynb

@@ -0,0 +1,170 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "6194860c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Found 60 audio files\n",
+      "Total audio: 2880000 samples (180.0s at 16000Hz)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Simple benchmark comparing NumPy (CPU) vs CuPy (GPU) performance.\n",
+    "\n",
+    "import time\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import soundfile as sf\n",
+    "import pygaborstm as stm\n",
+    "\n",
+    "# Load all 60 ripple files\n",
+    "data_dir = Path(\"data/mvripfft\")\n",
+    "audio_files = sorted(data_dir.glob(\"*.wav\"))\n",
+    "print(f\"Found {len(audio_files)} audio files\")\n",
+    "\n",
+    "# Load and concatenate for a longer test signal\n",
+    "audios = []\n",
+    "for f in audio_files:\n",
+    "    audio, sr = sf.read(f)\n",
+    "    audios.append(audio)\n",
+    "\n",
+    "audio = np.concatenate(audios)\n",
+    "print(f\"Total audio: {len(audio)} samples ({len(audio)/sr:.1f}s at {sr}Hz)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "38c1ede8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU time: 169.165s\n"
+     ]
+    }
+   ],
+   "source": [
+    "# CPU Benchmark\n",
+    "config_cpu = stm.Config(use_gpu=False, resolution=\"low\")\n",
+    "model_cpu = stm.PyGaborSTM(config_cpu)\n",
+    "\n",
+    "# Warmup\n",
+    "_ = model_cpu.spectrogram(audios[0])\n",
+    "\n",
+    "# Benchmark\n",
+    "start = time.perf_counter()\n",
+    "spec_cpu = model_cpu.spectrogram(audio)\n",
+    "rsf_cpu = model_cpu.rsf(spec_cpu)\n",
+    "cpu_time = time.perf_counter() - start\n",
+    "\n",
+    "print(f\"CPU time: {cpu_time:.3f}s\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "86bb6abe",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "GPU time: 170.680s\n"
+     ]
+    }
+   ],
+   "source": [
+    "# GPU Benchmark\n",
+    "config_gpu = stm.Config(use_gpu=True, resolution=\"low\")\n",
+    "model_gpu = stm.PyGaborSTM(config_gpu)\n",
+    "\n",
+    "# Warmup (includes kernel compilation)\n",
+    "_ = model_gpu.spectrogram(audios[0])\n",
+    "_ = model_gpu.rsf(model_gpu.spectrogram(audios[0]))\n",
+    "\n",
+    "# Benchmark\n",
+    "start = time.perf_counter()\n",
+    "spec_gpu = model_gpu.spectrogram(audio)\n",
+    "rsf_gpu = model_gpu.rsf(spec_gpu)\n",
+    "gpu_time = time.perf_counter() - start\n",
+    "\n",
+    "print(f\"GPU time: {gpu_time:.3f}s\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "a3308e88",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "========================================\n",
+      "Audio length: 180.0s\n",
+      "Resolution:   low (60 filters)\n",
+      "========================================\n",
+      "CPU time: 169.165s\n",
+      "GPU time: 170.680s\n",
+      "Speedup:  1.0x\n",
+      "========================================\n",
+      "\n",
+      "Spectrogram match: True\n",
+      "RSF match: True\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Results\n",
+    "print(f\"\\n{'='*40}\")\n",
+    "print(f\"Audio length: {len(audio)/sr:.1f}s\")\n",
+    "print(f\"Resolution:   low (60 filters)\")\n",
+    "print(f\"{'='*40}\")\n",
+    "print(f\"CPU time: {cpu_time:.3f}s\")\n",
+    "print(f\"GPU time: {gpu_time:.3f}s\")\n",
+    "print(f\"Speedup:  {cpu_time/gpu_time:.1f}x\")\n",
+    "print(f\"{'='*40}\")\n",
+    "\n",
+    "# Verify outputs match\n",
+    "spec_match = np.allclose(spec_cpu.data, spec_gpu.data, rtol=1e-4)\n",
+    "rsf_match = np.allclose(rsf_cpu.data, rsf_gpu.data, rtol=1e-4)\n",
+    "print(f\"\\nSpectrogram match: {spec_match}\")\n",
+    "print(f\"RSF match: {rsf_match}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/chi2005_validation.ipynb

@@ -10,7 +10,7 @@
 
 # Defaults
 SR = 16000
-DURATION = 1.0
+DURATION = 3.0
 
 
 def generate_tone(freq, duration=DURATION, sr=SR, amplitude=0.5):

notebooks/example_usage.ipynb

@@ -1,60 +1,38 @@
-from .config import Config, SpectrogramConfig, GaborConfig
+"""
+PyGaborSTM: Spectro-temporal modulation analysis library.
+
+Usage:
+    import pygaborstm as stm
+
+    # Create model
+    model = stm.PyGaborSTM(config=stm.Config(use_gpu=True))
+    
+    # Compute
+    spec = model.spectrogram(audio)
+    rsf = model.rsf(spec)
+
+    # Plot
+    stm.plot.spectrogram(spec)
+    stm.plot.rsf(rsf)
+"""
+
+from .config import Config
 from .structs import Spectrogram, RSF
-from .spectrogram import AuditorySpectrogram, auditory_spectrogram
-from .gabor import GaborFilterbank, rsf
-from .core import load, compute_rsf
-from .plotting import (
-    plot_spectrogram,
-    plot_spectrogram_grid,
-    plot_rsf,
-    plot_rsf_grid,
-    plot_filterbank,
-)
-from .utils import (
-    generate_tone,
-    generate_three_tones,
-    generate_broadband_noise,
-    generate_harmonic_complex,
-    generate_moving_ripple,
-    generate_ripple_set,
-    save_three_tones,
-    save_noise,
-    save_harmonic_complexes,
-)
+from .core import PyGaborSTM
+from . import plot
+from . import structs
 
 __version__ = "0.1.0"
 
 __all__ = [
-    # Core functions
-    "load",
-    "compute_rsf",
-    # Functional API
-    "auditory_spectrogram",
-    "rsf",
-    # Classes
-    "AuditorySpectrogram",
-    "GaborFilterbank",
+    # Main class
+    "PyGaborSTM",
     # Config
     "Config",
-    "SpectrogramConfig",
-    "GaborConfig",
     # Data structures
     "Spectrogram",
     "RSF",
-    # Plotting
-    "plot_spectrogram",
-    "plot_spectrogram_grid",
-    "plot_rsf",
-    "plot_rsf_grid",
-    "plot_filterbank",
-    # Stimulus generation
-    "generate_tone",
-    "generate_three_tones",
-    "generate_broadband_noise",
-    "generate_harmonic_complex",
-    "generate_moving_ripple",
-    "generate_ripple_set",
-    "save_three_tones",
-    "save_noise",
-    "save_harmonic_complexes",
-]
+    # Namespaced modules
+    "plot",
+    "structs"
+]

notebooks/mvripfft_validation.ipynb

@@ -0,0 +1,75 @@
+"""
+Backend module for numpy/cupy switching.
+
+Provides a unified interface for array operations that can use either
+numpy (CPU) or cupy (GPU) depending on configuration.
+"""
+
+import numpy as np
+
+# Try to import cupy
+try:
+    import cupy as cp
+    import cupyx.scipy.signal as cp_signal
+    CUPY_AVAILABLE = True
+except ImportError:
+    cp = None
+    cp_signal = None
+    CUPY_AVAILABLE = False
+
+
+def get_array_module(use_gpu: bool = False):
+    """
+    Get the appropriate array module (numpy or cupy).
+    
+    Args:
+        use_gpu: Whether to use GPU acceleration
+        
+    Returns:
+        Module (numpy or cupy)
+    """
+    if use_gpu:
+        if not CUPY_AVAILABLE:
+            import warnings
+            warnings.warn(
+                "CuPy not available. Falling back to NumPy. "
+                "Install cupy-cuda13x for GPU acceleration.",
+                UserWarning
+            )
+            return np
+        return cp
+    return np
+
+
+def get_signal_module(use_gpu: bool = False):
+    """
+    Get the appropriate signal processing module.
+    
+    Args:
+        use_gpu: Whether to use GPU acceleration
+        
+    Returns:
+        Module (scipy.signal or cupyx.scipy.signal)
+    """
+    if use_gpu:
+        if not CUPY_AVAILABLE:
+            from scipy import signal
+            return signal
+        return cp_signal
+    from scipy import signal
+    return signal
+
+
+def to_numpy(array) -> np.ndarray:
+    """
+    Convert array to numpy (transfers from GPU if needed).
+    
+    Args:
+        array: numpy or cupy array
+        
+    Returns:
+        numpy array
+    """
+    if CUPY_AVAILABLE and isinstance(array, cp.ndarray):
+        return cp.asnumpy(array)
+    return np.asarray(array)

pygaborstm/utils/__init__.py notebooks/nb_utils/__init__.pypygaborstm/utils/__init__.py renamed to notebooks/nb_utils/__init__.py

@@ -8,33 +8,21 @@
 
 
 @dataclass
-class SpectrogramConfig:
-    """Configuration for auditory spectrogram."""
-
+class Config:
+    """Configuration for PyGaborSTM."""
+    
+    # General
+    use_gpu: bool = False
     sample_rate: int = 16000
+    
+    # Spectrogram
     n_filters: int = 256
     f_min: float = 180.0
     octaves: float = 5.3
     tau_ms: float = 4.0
     frmlen_ms: float = 8.0
-    constant_Q: float = 16.0
-
-
-@dataclass
-class GaborConfig:
-    """Configuration for Gabor filterbank."""
-
-    sample_rate: int = 16000
-    n_freq_bins: int = 128
+    
+    # RSF / Gabor
     resolution: str = "low"
     rsf_frame_size_ms: int = 500
     rsf_frame_shift_ms: int = 10
-
-
-@dataclass
-class Config:
-    """Top-level configuration."""
-
-    use_gpu: bool = False
-    spectrogram: SpectrogramConfig = field(default_factory=SpectrogramConfig)
-    gabor: GaborConfig = field(default_factory=GaborConfig)