wdm-transform 0.1.1__py3-none-any.whl

wdm_transform/__init__.py

@@ -0,0 +1,13 @@
+"""Public package exports for the WDM transform scaffold."""
+
+from .backends import Backend, get_backend, register_backend
+from .datatypes import FrequencySeries, TimeSeries, WDM
+
+__all__ = [
+    "Backend",
+    "FrequencySeries",
+    "TimeSeries",
+    "WDM",
+    "get_backend",
+    "register_backend",
+]

wdm_transform/backends/__init__.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import os
+from importlib import import_module
+from typing import Any
+
+from .base import Backend
+from .numpy_backend import NUMPY_BACKEND
+
+_BACKENDS: dict[str, Backend] = {
+    NUMPY_BACKEND.name: NUMPY_BACKEND,
+}
+
+_BUILTIN_BACKEND_LOADERS = {
+    "jax": ("wdm_transform.backends.jax_backend", "load_jax_backend"),
+    "cupy": ("wdm_transform.backends.cupy_backend", "load_cupy_backend"),
+}
+
+
+def register_backend(name: str, xp: Any, fft: Any) -> Backend:
+    """Register a backend by name and return the resulting backend object."""
+    backend = Backend(name=name, xp=xp, fft=fft)
+    _BACKENDS[name] = backend
+    return backend
+
+
+def _load_builtin_backend(name: str) -> Backend:
+    try:
+        module_name, loader_name = _BUILTIN_BACKEND_LOADERS[name]
+    except KeyError as exc:
+        raise ValueError(f"Unknown backend {name!r}.") from exc
+
+    module = import_module(module_name)
+    loader = getattr(module, loader_name)
+    backend = loader()
+    _BACKENDS[name] = backend
+    return backend
+
+
+def get_backend(backend: str | Backend | None = None) -> Backend:
+    """Resolve a backend object from a name, instance, or environment default."""
+    if isinstance(backend, Backend):
+        return backend
+    if backend is None:
+        backend = os.environ.get("WDM_BACKEND")
+    if backend is None:
+        return NUMPY_BACKEND
+    try:
+        return _BACKENDS[backend]
+    except KeyError as exc:
+        try:
+            return _load_builtin_backend(backend)
+        except ValueError:
+            available = ", ".join(
+                sorted(set(_BACKENDS) | set(_BUILTIN_BACKEND_LOADERS))
+            )
+            raise ValueError(
+                f"Unknown backend {backend!r}. Available backends: {available}."
+            ) from exc
+
+
+__all__ = ["Backend", "NUMPY_BACKEND", "get_backend", "register_backend"]

wdm_transform/backends/base.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Backend:
+    """Minimal backend wrapper around an array namespace and FFT namespace."""
+
+    name: str
+    xp: Any
+    fft: Any
+
+    def asarray(self, value: Any, dtype: Any | None = None) -> Any:
+        """Convert a value into an array for this backend."""
+        return self.xp.asarray(value, dtype=dtype)

wdm_transform/backends/cupy_backend.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from .base import Backend
+
+
+def load_cupy_backend() -> Backend:
+    try:
+        import cupy as cp
+        import cupy.fft as cpfft
+    except ImportError as exc:
+        raise ImportError(
+            "CuPy backend requested, but cupy is not installed. Install a compatible cupy package first."
+        ) from exc
+
+    return Backend(name="cupy", xp=cp, fft=cpfft)

wdm_transform/backends/jax_backend.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from .base import Backend
+
+
+def load_jax_backend() -> Backend:
+    try:
+        import jax
+        jax.config.update("jax_enable_x64", True)
+        import jax.numpy as jnp
+        import jax.numpy.fft as jfft
+    except ImportError as exc:
+        raise ImportError(
+            "JAX backend requested, but jax is not installed. Install a compatible jax package first."
+        ) from exc
+
+    return Backend(name="jax", xp=jnp, fft=jfft)

wdm_transform/backends/numpy_backend.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import numpy as np
+
+from .base import Backend
+
+
+NUMPY_BACKEND = Backend(name="numpy", xp=np, fft=np.fft)

wdm_transform/benchmarking.py

@@ -0,0 +1,284 @@
+"""Benchmarking CLI for WDM transforms across different backends and input sizes.
+
+This module is excluded from pytest coverage and should not be imported
+when loading the main wdm_transform package.
+
+Run with: python -m wdm_transform.benchmarking [options]
+"""
+
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from . import backends, transforms
+
+
+def validate_backend_available(backend_name: str) -> bool:  # pragma: no cover
+    """Check if a backend is available."""
+    try:
+        backends.get_backend(backend_name)
+        return True
+    except ValueError:
+        return False
+
+
+def generate_benchmark_data(n: int) -> np.ndarray:  # pragma: no cover
+    """Generate random complex benchmark data."""
+    np.random.seed(42)
+    return np.random.randn(n) + 1j * np.random.randn(n)
+
+
+def find_factorization(n: int) -> tuple[int, int] | None:  # pragma: no cover
+    """Find a factorization n = nt * nf where both nt and nf are even.
+    
+    Prioritizes balanced factorizations (nt close to sqrt(n)).
+    Returns (nt, nf) or None if no valid factorization exists.
+    """
+    if n < 4:
+        return None
+    
+    # Try factorizations starting from sqrt(n) and working outward
+    sqrt_n = int(np.sqrt(n))
+    
+    for nt in range(sqrt_n, 1, -1):
+        if n % nt != 0:
+            continue
+        nf = n // nt
+        # Both must be even
+        if nt % 2 == 0 and nf % 2 == 0:
+            return (nt, nf)
+    
+    return None
+
+
+def benchmark_forward(
+    data: np.ndarray,
+    backend_name: str,
+    nt: int,
+    nf: int,
+    num_runs: int = 3,
+) -> float:  # pragma: no cover
+    """Benchmark the forward WDM transform.
+
+    Returns the mean runtime in seconds.
+    """
+    backend = backends.get_backend(backend_name)
+    fixed_params = {"a": 0.4, "d": 0.8, "dt": 1.0, "backend": backend}
+
+    # Convert data to backend array
+    backend_data = backend.asarray(data)
+
+    # Warmup run
+    transforms.forward_wdm(backend_data, nt=nt, nf=nf, **fixed_params)
+
+    # Timed runs
+    times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        transforms.forward_wdm(backend_data, nt=nt, nf=nf, **fixed_params)
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+
+    return np.mean(times)
+
+
+def benchmark_inverse(
+    coeffs: np.ndarray,
+    backend_name: str,
+    num_runs: int = 3,
+) -> float:  # pragma: no cover
+    """Benchmark the inverse WDM transform.
+
+    Returns the mean runtime in seconds.
+    """
+    backend = backends.get_backend(backend_name)
+    fixed_params = {"a": 0.4, "d": 0.8, "dt": 1.0, "backend": backend}
+
+    # Convert coeffs to backend array
+    backend_coeffs = backend.asarray(coeffs)
+
+    # Warmup run
+    transforms.inverse_wdm(backend_coeffs, **fixed_params)
+
+    # Timed runs
+    times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        transforms.inverse_wdm(backend_coeffs, **fixed_params)
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+
+    return np.mean(times)
+
+
+def run_benchmarks(  # pragma: no cover
+    backends_to_test: list[str],
+    n_values: list[int],
+    num_runs: int = 3,
+) -> dict[str, Any]:
+    """Run benchmark suite across backends and input sizes.
+
+    Returns results dictionary with structure:
+    {
+        'forward': {backend_name: {n: time_in_seconds}},
+        'inverse': {backend_name: {n: time_in_seconds}},
+    }
+    """
+    results = {"forward": {}, "inverse": {}}
+
+    # Validate backends
+    available_backends = [b for b in backends_to_test if validate_backend_available(b)]
+    if not available_backends:
+        print(f"ERROR: No backends available from: {backends_to_test}")
+        return results
+
+    if len(available_backends) < len(backends_to_test):
+        unavailable = set(backends_to_test) - set(available_backends)
+        print(f"WARNING: Skipping unavailable backends: {unavailable}")
+
+    # Forward transform benchmarks
+    print("\n" + "=" * 70)
+    print("FORWARD WDM TRANSFORM BENCHMARKS")
+    print("=" * 70)
+
+    for backend_name in available_backends:
+        results["forward"][backend_name] = {}
+        print(f"\nBackend: {backend_name}")
+        print("-" * 40)
+
+        for n in n_values:
+            # Find a valid factorization for this N
+            factorization = find_factorization(n)
+            if factorization is None:
+                print(f"  N={n:>8}: SKIPPED (no valid factorization)")
+                continue
+            
+            nt, nf = factorization
+            data = generate_benchmark_data(nt * nf)
+            try:
+                mean_time = benchmark_forward(data, backend_name, nt, nf, num_runs)
+                results["forward"][backend_name][n] = mean_time
+                print(f"  N={n:>8}: {mean_time*1e3:>10.4f} ms (nt={nt}, nf={nf})")
+            except Exception as e:
+                print(f"  N={n:>8}: FAILED ({type(e).__name__}: {e})")
+
+    # Inverse transform benchmarks
+    print("\n" + "=" * 70)
+    print("INVERSE WDM TRANSFORM BENCHMARKS")
+    print("=" * 70)
+
+    for backend_name in available_backends:
+        results["inverse"][backend_name] = {}
+        print(f"\nBackend: {backend_name}")
+        print("-" * 40)
+
+        for n in n_values:
+            # Find a valid factorization for this N
+            factorization = find_factorization(n)
+            if factorization is None:
+                print(f"  N={n:>8}: SKIPPED (no valid factorization)")
+                continue
+            
+            nt, nf = factorization
+            coeffs = generate_benchmark_data(nt * nf).reshape(nt, nf)
+            try:
+                mean_time = benchmark_inverse(coeffs, backend_name, num_runs)
+                results["inverse"][backend_name][n] = mean_time
+                print(f"  N={n:>8}: {mean_time*1e3:>10.4f} ms (nt={nt}, nf={nf})")
+            except Exception as e:
+                print(f"  N={n:>8}: FAILED ({type(e).__name__}: {e})")
+
+    return results
+
+
+def print_summary(results: dict[str, Any]) -> None:  # pragma: no cover
+    """Print a summary table of results."""
+    print("\n" + "=" * 70)
+    print("SUMMARY TABLE")
+    print("=" * 70)
+
+    for transform_type, backends_data in results.items():
+        if not backends_data:
+            continue
+
+        print(f"\n{transform_type.upper()} TRANSFORM:")
+        print("-" * 40)
+
+        # Get all N values
+        all_n = set()
+        for n_dict in backends_data.values():
+            all_n.update(n_dict.keys())
+        all_n = sorted(all_n)
+
+        # Header
+        header = f"{'N':>10} | " + " | ".join(f"{b:>12}" for b in backends_data.keys())
+        print(header)
+        print("-" * len(header))
+
+        # Rows
+        for n in all_n:
+            row = f"{n:>10} | "
+            times = [backends_data.get(b, {}).get(n) for b in backends_data.keys()]
+            time_strs = [
+                f"{t*1e3:>10.4f} ms" if t is not None else f"{'SKIPPED':>12}" for t in times
+            ]
+            row += " | ".join(time_strs)
+            print(row)
+
+
+def main() -> None:  # pragma: no cover
+    """CLI entry point for benchmarking."""
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Benchmark WDM transforms across backends and input sizes.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python -m wdm_transform.benchmarking
+  python -m wdm_transform.benchmarking --backends numpy jax
+  python -m wdm_transform.benchmarking --n 2048 4096 8192
+  python -m wdm_transform.benchmarking --runs 5
+        """,
+    )
+
+    parser.add_argument(
+        "--backends",
+        nargs="+",
+        default=["numpy"],
+        choices=["numpy", "jax", "cupy"],
+        help="Backends to benchmark (default: numpy)",
+    )
+
+    parser.add_argument(
+        "--n",
+        nargs="+",
+        type=int,
+        default=[2048, 4096, 8192, 16384],
+        help="Input sizes to test (default: 2048 4096 8192 16384)",
+    )
+
+    parser.add_argument(
+        "--runs",
+        type=int,
+        default=3,
+        help="Number of runs per benchmark (default: 3)",
+    )
+
+    args = parser.parse_args()
+
+    results = run_benchmarks(
+        backends_to_test=args.backends,
+        n_values=args.n,
+        num_runs=args.runs,
+    )
+
+    print_summary(results)
+    print("\n")
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()
+

wdm_transform/datatypes/__init__.py

@@ -0,0 +1,4 @@
+from .series import FrequencySeries, TimeSeries
+from .wdm import WDM
+
+__all__ = ["FrequencySeries", "TimeSeries", "WDM"]

wdm_transform/datatypes/series.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from ..backends import Backend, get_backend
+
+if TYPE_CHECKING:
+    from .wdm import WDM
+
+
+@dataclass(frozen=True)
+class TimeSeries:
+    """One-dimensional sampled time-domain data."""
+
+    data: Any
+    dt: float
+    backend: Backend = field(default_factory=get_backend)
+
+    def __post_init__(self) -> None:
+        backend = get_backend(self.backend)
+        data = backend.asarray(self.data)
+
+        if data.ndim != 1:
+            raise ValueError("TimeSeries data must be one-dimensional.")
+        if self.dt <= 0:
+            raise ValueError("dt must be positive.")
+
+        object.__setattr__(self, "backend", backend)
+        object.__setattr__(self, "data", data)
+
+    @property
+    def n(self) -> int:
+        return int(self.data.shape[0])
+
+    @property
+    def df(self) -> float:
+        return 1.0 / (self.n * self.dt)
+
+    @property
+    def times(self) -> Any:
+        return self.backend.xp.arange(self.n) * self.dt
+
+    def to_frequency_series(self) -> "FrequencySeries":
+        transformed = self.backend.fft.fft(self.data)
+        return FrequencySeries(transformed, df=self.df, backend=self.backend)
+
+    def to_wdm(
+        self,
+        *,
+        nt: int,
+        a: float = 1.0 / 3.0,
+        d: float = 1.0,
+        backend: str | Backend | None = None,
+    ) -> "WDM":
+        """Compute the WDM transform of this time series."""
+        from .wdm import WDM
+
+        return WDM.from_time_series(self, nt=nt, a=a, d=d, backend=backend)
+
+    def plot(self, **kwargs: Any) -> tuple[Any, Any]:
+        """Plot the time-domain samples using the shared plotting helper."""
+        from ..plotting import plot_time_series
+
+        return plot_time_series(self, **kwargs)
+
+
+@dataclass(frozen=True)
+class FrequencySeries:
+    """One-dimensional FFT-domain data with spacing metadata."""
+
+    data: Any
+    df: float
+    backend: Backend = field(default_factory=get_backend)
+
+    def __post_init__(self) -> None:
+        backend = get_backend(self.backend)
+        data = backend.asarray(self.data)
+
+        if data.ndim != 1:
+            raise ValueError("FrequencySeries data must be one-dimensional.")
+        if self.df <= 0:
+            raise ValueError("df must be positive.")
+
+        object.__setattr__(self, "backend", backend)
+        object.__setattr__(self, "data", data)
+
+    @property
+    def n(self) -> int:
+        return int(self.data.shape[0])
+
+    @property
+    def dt(self) -> float:
+        return 1.0 / (self.n * self.df)
+
+    @property
+    def freqs(self) -> Any:
+        return self.backend.fft.fftfreq(self.n, d=self.dt)
+
+    def to_time_series(self, *, real: bool = False) -> TimeSeries:
+        recovered = self.backend.fft.ifft(self.data)
+        if real:
+            recovered = self.backend.xp.real(recovered)
+        return TimeSeries(recovered, dt=self.dt, backend=self.backend)
+
+    def to_wdm(
+        self,
+        *,
+        nt: int,
+        a: float = 1.0 / 3.0,
+        d: float = 1.0,
+        backend: str | Backend | None = None,
+    ) -> "WDM":
+        """Compute the WDM transform of this frequency-domain series."""
+        from .wdm import WDM
+
+        return WDM.from_frequency_series(self, nt=nt, a=a, d=d, backend=backend)
+
+    def plot(self, **kwargs: Any) -> tuple[Any, Any]:
+        """Plot the spectrum using the shared plotting helper."""
+        from ..plotting import plot_frequency_series
+
+        return plot_frequency_series(self, **kwargs)