wdm-transform 0.2.0__tar.gz → 0.3.0__tar.gz

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/CHANGELOG.md

@@ -1,3 +1,11 @@
+## v0.3.0 (2026-05-01)
+
+### fix
+
+- workflows for pypi
+- pytest fixes
+- add gpu and cpu tracking on benchmarking
+
 ## v0.2.0 (2026-04-29)
 
 ### feat

{wdm_transform-0.2.0/src/wdm_transform.egg-info → wdm_transform-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wdm-transform
-Version: 0.2.0
+Version: 0.3.0
 Summary: Object-oriented WDM transforms for sampled one-dimensional signals.
 Project-URL: Homepage, https://github.com/pywavelet/wdm_transform
 Project-URL: Repository, https://github.com/pywavelet/wdm_transform

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/mkdocs.yml

@@ -7,9 +7,7 @@ docs_dir: docs
 site_dir: site
 
 exclude_docs: |
-  studies/**
   manuscript/**
-  reference/api.md
   dev-notes.md
 
 theme:
@@ -102,4 +100,4 @@ nav:
   - API Reference:
       - API Overview: guide/api-overview.md
       - Package Layout: guide/package-layout.md
-      - Benchmarks: benchmarks.md
+      - Benchmarks: benchmarks.ipynb

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/src/wdm_transform/benchmarking.py

@@ -16,6 +16,8 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
+import shutil
 import time
 from pathlib import Path
 from typing import Any, Callable
@@ -36,6 +38,22 @@ DEFAULT_JSON_NAME = "benchmark_results.json"
 DEFAULT_PLOT_NAME = "benchmark_runtime.png"
 FIXED_PARAMS = {"a": 0.4, "d": 1, "dt": 1.0}
 
+def _resolve_device(requested: str) -> str:
+    """Resolve 'cpu'/'gpu'/'auto' and configure JAX platform env vars.
+
+    Must be called before any backend is imported.
+    """
+    if requested == "auto":
+        device = "gpu" if shutil.which("nvidia-smi") is not None else "cpu"
+    else:
+        device = requested
+
+    jax_platform = "cuda" if device == "gpu" else "cpu"
+    os.environ.setdefault("JAX_PLATFORMS", jax_platform)
+    os.environ.setdefault("JAX_PLATFORM_NAME", jax_platform)
+    return device
+
+
 OPERATION_TITLES = {
     "from_freq": "Forward Kernel: from_freq_to_wdm",
     "to_freq": "Inverse Kernel: from_wdm_to_freq",
@@ -494,6 +512,7 @@ def plot_results(  # pragma: no cover
     backend_styles = {
         "numpy": {"color": "tab:blue"},
         "jax": {"color": "tab:orange"},
+        "cupy": {"color": "tab:green"},
     }
     for column, operation in enumerate(operations):
         runtime_ax = axes[0, column]
@@ -598,12 +617,25 @@ def _build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="Benchmark low-level WDM kernels across backends and input sizes.",
     )
+    parser.add_argument(
+        "--device",
+        choices=["cpu", "gpu", "auto"],
+        default="auto",
+        help=(
+            "Target device for JAX (and CuPy). "
+            "'auto' detects GPU via nvidia-smi. "
+            "Forces JAX platform."
+        ),
+    )
     parser.add_argument(
         "--backends",
         nargs="+",
-        default=DEFAULT_BACKENDS,
+        default=None,
         choices=["numpy", "jax", "cupy"],
-        help="Backends to benchmark (default: numpy jax)",
+        help=(
+            "Backends to benchmark. "
+            "Default: numpy+jax on CPU, numpy+jax+cupy on GPU."
+        ),
     )
     parser.add_argument(
         "--pow2",
@@ -651,6 +683,15 @@ def _build_parser() -> argparse.ArgumentParser:
 def main(argv: list[str] | None = None) -> None:  # pragma: no cover
     """CLI entry point for benchmarking."""
     args = _build_parser().parse_args(argv)
+
+    # Resolve device and configure JAX platform BEFORE any backend is imported.
+    device = _resolve_device(args.device)
+    print(f"Device: {device.upper()}")
+
+    backends_to_test = args.backends or (
+        ["numpy", "jax", "cupy"] if device == "gpu" else DEFAULT_BACKENDS
+    )
+
     if args.batch_size < 1:
         raise SystemExit("--batch-size must be at least 1.")
     pow2_range = tuple(args.pow2) if args.pow2 is not None else None
@@ -659,18 +700,32 @@ def main(argv: list[str] | None = None) -> None:  # pragma: no cover
     except ValueError as exc:
         raise SystemExit(str(exc)) from exc
 
+    # Log which JAX devices are in use before running (confirms CPU vs GPU).
+    _jax = None
+    try:
+        import jax as _jax  # noqa: PLC0415
+
+        print(f"JAX devices: {_jax.devices()}")
+    except Exception:
+        pass
+
     results = run_benchmarks(
-        backends_to_test=args.backends,
+        backends_to_test=backends_to_test,
         n_values=n_values,
         num_runs=args.runs,
         batch_size=args.batch_size,
     )
+
+    if _jax is not None:
+        results["metadata"]["jax_devices"] = [str(d) for d in _jax.devices()]
+
     print_summary(results)
 
     output_json = args.output_json or (args.outdir / DEFAULT_JSON_NAME)
     output_plot = args.output_plot or (args.outdir / DEFAULT_PLOT_NAME)
+    plot_title = args.plot_title + f" ({device.upper()})"
     json_path = save_results(results, output_json)
-    plot_path = plot_results(results, output_plot, title=args.plot_title)
+    plot_path = plot_results(results, output_plot, title=plot_title)
 
     print(f"\nSaved JSON results to {json_path}")
     print(f"Saved benchmark plot to {plot_path}")

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/src/wdm_transform/transforms/__init__.py

@@ -154,7 +154,36 @@ def from_freq_to_wdm_subband(
     d: float = 1.0,
     backend: str | Backend | None = None,
 ) -> tuple[Any, int]:
-    """Compute the overlapping WDM sub-band from a compact Fourier span."""
+    """Compute the overlapping WDM sub-band from a compact Fourier span.
+
+    Use this when the input is already cropped in the one-sided Fourier
+    domain. The function evaluates only the WDM frequency channels whose
+    window support intersects the supplied Fourier span, rather than
+    materializing the full WDM grid.
+
+    Args:
+        data: One-dimensional one-sided Fourier samples for the compact span
+            ``[kmin, kmin + len(data))``.
+        df: Frequency spacing of the full one-sided Fourier grid.
+        nfreqs_fourier: Number of samples in the full one-sided Fourier grid,
+            including DC and Nyquist.
+        kmin: Starting Fourier-bin index of ``data`` in the full one-sided
+            grid.
+        nfreqs_wdm: Number of positive WDM frequency intervals in the full
+            transform. A full WDM coefficient grid has ``nfreqs_wdm + 1``
+            frequency channels.
+        ntimes_wdm: Number of WDM time bins.
+        a: Window roll-off parameter.
+        d: Window order parameter.
+        backend: Backend name or object. If omitted, the default backend is
+            used.
+
+    Returns:
+        A tuple ``(coeffs, mmin)``. ``coeffs`` has shape
+        ``(ntimes_wdm, nf_sub_wdm)`` and stores the touched WDM channels.
+        ``mmin`` is the first WDM frequency-channel index represented by
+        ``coeffs[:, 0]``.
+    """
     resolved_backend = get_backend(backend)
     module = _get_subband_transform_module(resolved_backend)
     return module.forward_wdm_subband(
@@ -190,6 +219,29 @@ def from_freq_to_wdm_band(
     minimal touched WDM span, then returns only the requested compact block
     ``mmin : mmin + nf_sub_wdm``. The requested block must lie inside the
     touched span of the supplied Fourier support.
+
+    Args:
+        data: One-dimensional one-sided Fourier samples for the compact span
+            ``[kmin, kmin + len(data))``.
+        df: Frequency spacing of the full one-sided Fourier grid.
+        nfreqs_fourier: Number of samples in the full one-sided Fourier grid,
+            including DC and Nyquist.
+        kmin: Starting Fourier-bin index of ``data`` in the full one-sided
+            grid.
+        nfreqs_wdm: Number of positive WDM frequency intervals in the full
+            transform. A full WDM coefficient grid has ``nfreqs_wdm + 1``
+            frequency channels.
+        ntimes_wdm: Number of WDM time bins.
+        mmin: First WDM frequency-channel index to return.
+        nf_sub_wdm: Number of WDM frequency channels to return.
+        a: Window roll-off parameter.
+        d: Window order parameter.
+        backend: Backend name or object. If omitted, the default backend is
+            used.
+
+    Returns:
+        The requested coefficient block with shape
+        ``(ntimes_wdm, nf_sub_wdm)``.
     """
     coeffs, touched_mmin = from_freq_to_wdm_subband(
         data,
@@ -225,7 +277,34 @@ def from_wdm_to_freq_subband(
     d: float = 1.0,
     backend: str | Backend | None = None,
 ) -> tuple[Any, int]:
-    """Reconstruct the touched Fourier span from a compact WDM block."""
+    """Reconstruct the touched Fourier span from a compact WDM block.
+
+    This is the inverse of the compact sub-band transform. It expands a block
+    of adjacent WDM frequency channels into the one-sided Fourier bins touched
+    by those channels' window support.
+
+    Args:
+        coeffs: Two-dimensional WDM coefficient block with shape
+            ``(ntimes_wdm, nf_sub_wdm)``.
+        df: Frequency spacing of the full one-sided Fourier grid.
+        nfreqs_fourier: Number of samples in the full one-sided Fourier grid,
+            including DC and Nyquist.
+        mmin: First WDM frequency-channel index represented by
+            ``coeffs[:, 0]``.
+        nfreqs_wdm: Number of positive WDM frequency intervals in the full
+            transform. A full WDM coefficient grid has ``nfreqs_wdm + 1``
+            frequency channels.
+        ntimes_wdm: Number of WDM time bins.
+        a: Window roll-off parameter.
+        d: Window order parameter.
+        backend: Backend name or object. If omitted, the default backend is
+            used.
+
+    Returns:
+        A tuple ``(spectrum, kmin)``. ``spectrum`` contains the reconstructed
+        one-sided Fourier span, and ``kmin`` is the starting Fourier-bin index
+        of ``spectrum[0]`` in the full one-sided grid.
+    """
     resolved_backend = get_backend(backend)
     module = _get_subband_transform_module(resolved_backend)
     return module.inverse_wdm_subband(

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/src/wdm_transform/transforms/_subband.py

@@ -93,7 +93,23 @@ def wdm_span_from_fourier_span(
     kmin: int,
     lendata: int,
 ) -> tuple[int, int]:
-    """Map a compact one-sided Fourier span onto the touched WDM channels."""
+    """Map a compact one-sided Fourier span onto the touched WDM channels.
+
+    Args:
+        nfreqs_fourier: Number of samples in the full one-sided Fourier grid,
+            including DC and Nyquist.
+        nfreqs_wdm: Number of positive WDM frequency intervals in the full
+            transform. A full WDM coefficient grid has ``nfreqs_wdm + 1``
+            frequency channels.
+        ntimes_wdm: Number of WDM time bins.
+        kmin: Starting Fourier-bin index of the compact Fourier span.
+        lendata: Number of Fourier bins in the compact span.
+
+    Returns:
+        A tuple ``(mmin, nf_sub_wdm)``. ``mmin`` is the first touched WDM
+        frequency-channel index, and ``nf_sub_wdm`` is the number of adjacent
+        WDM channels touched by the Fourier span.
+    """
     validate_subband_grid(
         nfreqs_fourier=nfreqs_fourier,
         nfreqs_wdm=nfreqs_wdm,
@@ -123,7 +139,24 @@ def fourier_span_from_wdm_span(
     mmin: int,
     nf_sub_wdm: int,
 ) -> tuple[int, int]:
-    """Map a compact WDM channel span onto the touched one-sided Fourier bins."""
+    """Map a compact WDM channel span onto the touched one-sided Fourier bins.
+
+    Args:
+        nfreqs_fourier: Number of samples in the full one-sided Fourier grid,
+            including DC and Nyquist.
+        nfreqs_wdm: Number of positive WDM frequency intervals in the full
+            transform. A full WDM coefficient grid has ``nfreqs_wdm + 1``
+            frequency channels.
+        ntimes_wdm: Number of WDM time bins.
+        mmin: First WDM frequency-channel index in the compact WDM span.
+        nf_sub_wdm: Number of adjacent WDM frequency channels in the compact
+            span.
+
+    Returns:
+        A tuple ``(kmin, lendata)``. ``kmin`` is the first touched one-sided
+        Fourier-bin index, and ``lendata`` is the number of touched Fourier
+        bins.
+    """
     validate_subband_grid(
         nfreqs_fourier=nfreqs_fourier,
         nfreqs_wdm=nfreqs_wdm,
@@ -143,4 +176,3 @@ def fourier_span_from_wdm_span(
     nyquist = int(nfreqs_wdm) * half
     kmax = nyquist if mmax == int(nfreqs_wdm) else (mmax + 1) * half - 1
     return kmin, kmax - kmin + 1
-

{wdm_transform-0.2.0 → wdm_transform-0.3.0/src/wdm_transform.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wdm-transform
-Version: 0.2.0
+Version: 0.3.0
 Summary: Object-oriented WDM transforms for sampled one-dimensional signals.
 Project-URL: Homepage, https://github.com/pywavelet/wdm_transform
 Project-URL: Repository, https://github.com/pywavelet/wdm_transform

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/tests/test_roundtrip.py

@@ -8,6 +8,7 @@ import pytest
 from wdm_transform import FrequencySeries, TimeSeries, WDM
 from wdm_transform.backends import Backend, get_backend
 from wdm_transform.plotting import plot_periodogram, plot_spectrogram
+import matplotlib.pyplot as plt
 
 
 def _chirplet(times: np.ndarray, duration: float) -> np.ndarray:
@@ -253,7 +254,7 @@ def test_batched_plot_methods_create_one_axis_per_batch_element() -> None:
 
     time_series = TimeSeries(np.zeros((2, 16)), dt=0.25)
     freq_series = FrequencySeries(np.zeros((2, 16), dtype=complex), df=0.25)
-    wdm = WDM(np.zeros((2, 4, 5)), dt=0.25, backend="jax")
+    wdm = WDM(np.zeros((2, 4, 5)), dt=0.25, backend="numpy")
 
     fig, axes = time_series.plot()
     assert len(np.ravel(axes)) == 2
@@ -294,9 +295,6 @@ def test_singleton_plot_methods_return_one_element_axes_arrays() -> None:
 
 
 def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
-    matplotlib = pytest.importorskip("matplotlib")
-    matplotlib.use("Agg", force=True)
-    plt = pytest.importorskip("matplotlib.pyplot")
 
     series, nt, _ = _example_series()
     coeffs = WDM.from_time_series(series, nt=nt, a=1.0 / 3.0)
@@ -305,15 +303,13 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     reconstructed = coeffs.to_time_series()
 
     residual = reconstructed.data - series.data
-    test_outdir = outdir / "test_roundtrip"
-    test_outdir.mkdir(parents=True, exist_ok=True)
 
     fig, ax = plt.subplots(figsize=(8, 3.5))
     series.plot(ax=ax, color="C0", label="input")
     ax.set_title("Original time series")
     assert ax.get_xlabel() == "Time [min]"
     assert ax.get_ylabel() == "Amplitude"
-    fig.savefig(test_outdir / "1_original.png", dpi=140)
+    fig.savefig(outdir / "1_original.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(8, 3.5))
@@ -327,7 +323,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     ax.set_title("Complex FFT coefficients")
     assert ax.get_xlabel() == "Frequency [Hz]"
     assert ax.get_ylabel() == "Value"
-    fig.savefig(test_outdir / "2_frequency.png", dpi=140)
+    fig.savefig(outdir / "2_frequency.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(7, 3.5))
@@ -335,7 +331,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     ax.set_title("Periodogram")
     assert ax.get_xlabel() == "Frequency [Hz]"
     assert ax.get_ylabel() == "Periodogram"
-    fig.savefig(test_outdir / "3_periodogram.png", dpi=140)
+    fig.savefig(outdir / "3_periodogram.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(8, 4.5))
@@ -348,7 +344,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     ax.set_title("Spectrogram")
     assert ax.get_xlabel() == "Time [min]"
     assert ax.get_ylabel() == "Frequency [Hz]"
-    fig.savefig(test_outdir / "4_spectrogram.png", dpi=140)
+    fig.savefig(outdir / "4_spectrogram.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(8, 4.5))
@@ -365,7 +361,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     assert "Time bins" in ax.get_xlabel()
     assert "Frequency bins" in ax.get_ylabel()
     assert any("coeffs" in text.get_text() for text in ax.texts)
-    fig.savefig(test_outdir / "5_wdm.png", dpi=140)
+    fig.savefig(outdir / "5_wdm.png", dpi=140)
     plt.close(fig)
 
     with pytest.warns(UserWarning, match="Falling back to default linear normalization"):
@@ -378,7 +374,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
             whiten_by=np.full((coeffs.nf + 1, coeffs.nt), np.nan),
         )
     assert len(fig.axes) == 1
-    fig.savefig(test_outdir / "6_wdm_fallback.png", dpi=140)
+    fig.savefig(outdir / "6_wdm_fallback.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(8, 3.5))
@@ -386,7 +382,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     ax.set_title("Reconstruction")
     assert ax.get_xlabel() == "Time [min]"
     assert ax.get_ylabel() == "Amplitude"
-    fig.savefig(test_outdir / "7_reconstruction.png", dpi=140)
+    fig.savefig(outdir / "7_reconstruction.png", dpi=140)
     plt.close(fig)
 
     fig, ax = plt.subplots(figsize=(7, 3.5))
@@ -395,7 +391,7 @@ def test_wdm_roundtrip_diagnostics_plots(outdir: Path) -> None:
     ax.set_xlabel("reconstruction - original")
     ax.set_ylabel("count")
     fig.tight_layout()
-    fig.savefig(test_outdir / "8_residual_hist.png", dpi=140)
+    fig.savefig(outdir / "8_residual_hist.png", dpi=140)
     plt.close(fig)
 
     np.testing.assert_allclose(reconstructed.data, series.data, atol=1e-10, rtol=1e-10)

{wdm_transform-0.2.0 → wdm_transform-0.3.0}/tests/test_subband.py

@@ -18,6 +18,7 @@ from wdm_transform.transforms import (
     wdm_span_from_fourier_span,
 )
 from wdm_transform.windows import cnm, phi_unit
+import matplotlib.pyplot as plt
 
 
 NT = 32
@@ -691,11 +692,6 @@ def test_chirping_binary_subband_roundtrip_jax() -> None:
 
 
 def test_chirping_binary_subband_diagnostics_plot(outdir) -> None:
-    matplotlib = pytest.importorskip("matplotlib")
-    matplotlib.use("Agg", force=True)
-    plt = pytest.importorskip("matplotlib.pyplot")
-    test_outdir = outdir / "test_subband"
-    test_outdir.mkdir(parents=True, exist_ok=True)
 
     signal, one_sided, kmin, lendata = _chirping_binary_subband_case()
     full_runtime_ms = _mean_runtime_ms(
@@ -871,5 +867,5 @@ def test_chirping_binary_subband_diagnostics_plot(outdir) -> None:
     axes[3].legend(loc="upper right")
     _annotate_metrics(axes[3], mf_snr=fourier_mf_snr)
 
-    fig.savefig(test_outdir / "subband_diagnostics_vertical.png", dpi=140)
+    fig.savefig(outdir / "subband_diagnostics_vertical.png", dpi=140)
     plt.close(fig)