pygranger 2.0.0__tar.gz

pygranger-2.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Vinícius Lima
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pygranger-2.0.0/PKG-INFO

@@ -0,0 +1,239 @@
+Metadata-Version: 2.4
+Name: pygranger
+Version: 2.0.0
+Summary: Granger Causality in the frequency domain — parametric and non-parametric estimation
+Author-email: Vinicius Lima <vinicius.lima.cordeiro@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/ViniciusLima94/pyGC
+Project-URL: Source, https://github.com/ViniciusLima94/pyGC
+Project-URL: Issues, https://github.com/ViniciusLima94/pyGC/issues
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: mne>=1.4
+Requires-Dist: joblib>=1.3
+Requires-Dist: h5py>=3.9
+Requires-Dist: xarray>=2025.4.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Requires-Dist: pytest-cov>=4.1; extra == "dev"
+Provides-Extra: notebooks
+Requires-Dist: jupyter>=1.0; extra == "notebooks"
+Requires-Dist: ipykernel>=6.25; extra == "notebooks"
+Requires-Dist: h5netcdf>=1.0; extra == "notebooks"
+Provides-Extra: jax
+Requires-Dist: jax>=0.4; extra == "jax"
+Dynamic: license-file
+
+# pyGC — Granger Causality in the Frequency Domain
+
+[![PyPI](https://img.shields.io/pypi/v/pygranger.svg)](https://pypi.org/project/pygranger/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/)
+
+A Python package for estimating Granger Causality (GC) in the frequency domain, supporting multiple spectral estimators (Fourier, Welch, Morlet, multitaper) and both NumPy and JAX/XLA backends.
+
+If you use this package, please cite:
+
+---
+
+## Installation
+
+The package is published on PyPI as `pygranger` and imported as `pygc`:
+
+```bash
+pip install pygranger
+```
+
+```python
+import pygc
+```
+
+Optional extras:
+
+```bash
+pip install "pygranger[jax]"       # JAX/XLA GPU-accelerated backend
+pip install "pygranger[notebooks]" # Jupyter notebooks
+```
+
+For development, install from a clone in editable mode:
+
+```bash
+git clone https://github.com/ViniciusLima94/pyGC.git
+cd pyGC
+pip install -e ".[dev]"            # pytest + coverage
+```
+
+---
+
+## Package structure
+
+```
+pygc/
+├── __init__.py               # Public API
+├── parametric.py             # Yule-Walker VAR fitting + transfer function
+├── non_parametric.py         # Wilson spectral factorization
+├── granger.py                # Bivariate GC, conditional GC (time + spectral)
+├── output.py                 # Packages results into labeled xarray.Dataset objects
+├── ar_model.py                # Synthetic AR benchmarks (Dhamala, Baccalá)
+├── _jax_backend.py           # Optional JAX/XLA Wilson factorization
+├── spectral_analysis/
+│   ├── fourier.py            # compute_freq, CSD, Morlet transforms (MNE)
+│   └── time_frequency.py     # Welch, multitaper, wavelet/Gabor transforms
+└── misc/
+    └── smooth_spectra.py     # smooth_spectra, downsample utilities
+```
+
+---
+
+## Usage
+
+All main functions accept raw signal data `X` and a sampling rate `fs`. The cross-spectral density is computed internally via the chosen `spectral_method`. Every GC function returns an `xarray.Dataset` (see `output.py`) rather than raw arrays, so results come back labeled and ready to index by name.
+
+> **Data shape convention:** the GC functions expect `X` as `(trials, nvars, N)` (or `(nvars, N)` for a single trial).
+
+### Bivariate GC
+
+```python
+import numpy as np
+from pygc import spectral_granger_causality
+from pygc import ar_model
+
+Fs = 200
+# Dhamala benchmark: Y drives X at 40 Hz
+data = ar_model.ar_model_dhamala(N=5000, Trials=20, Fs=Fs, C=0.25)
+# data shape: (2, Trials, N) -> transpose to (Trials, 2, N) for the GC functions
+data = np.transpose(data, (1, 0, 2))
+
+# Non-parametric (Fourier CSD + Wilson factorization) — default
+ds = spectral_granger_causality(data, fs=Fs, spectral_method='fourier', pairs=[(0, 1)])
+# ds.y2x peaks at ~40 Hz; ds.x2y is near zero
+peak_freq = ds.freq.values[ds.y2x.values[0].argmax()]
+```
+
+`spectral_granger_causality` returns a frequency-resolved `Dataset` with `x2y`, `y2x`, and `xy` variables (dims `pairs`, `freq`). For a single scalar per direction per pair instead (the broadband/time-domain summary), use `granger_causality`, which returns a `Dataset` with a single `F` variable (dims `source`, `target`):
+
+```python
+from pygc import granger_causality
+
+ds_f = granger_causality(data, fs=Fs, spectral_method='fourier')
+# ds_f.F[1, 0] -> broadband GC from channel 1 (Y) to channel 0 (X)
+```
+
+Available spectral methods: `'fourier'` (default), `'welch'`, `'morlet'`, `'multitaper'`.
+
+```python
+# Welch CSD
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='welch',
+    spectral_params={'nperseg': 512}
+)
+
+# Morlet CSD (requires explicit frequency axis)
+freqs = np.linspace(1, 80, 80)
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='morlet',
+    spectral_params={'freqs': freqs, 'n_cycles': 7.0}
+)
+
+# Multitaper CSD
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='multitaper',
+    spectral_params={'bandwidth': 4.0}
+)
+```
+
+### Parametric estimation (Yule-Walker)
+
+Use `YuleWalker` and `compute_transfer_function` to obtain the VAR-based cross-spectral matrix:
+
+```python
+import numpy as np
+from pygc import YuleWalker, compute_transfer_function
+from pygc.spectral_analysis import compute_freq
+from pygc import ar_model
+
+Fs = 200
+data = ar_model.ar_model_dhamala(N=5000, Trials=20, Fs=Fs, C=0.25)
+# data shape: (2, Trials, N) — YuleWalker takes (nvars, N) per trial, so no
+# transpose is needed here (unlike the GC functions above).
+
+f  = compute_freq(data.shape[2], Fs)
+m  = 2                          # VAR model order
+AR  = np.zeros((m, 2, 2))
+SIG = np.zeros((2, 2))
+
+for trial in range(data.shape[1]):
+    a, s = YuleWalker(data[:, trial, :], m)
+    AR  += a / data.shape[1]
+    SIG += s / data.shape[1]
+
+H, S = compute_transfer_function(AR, SIG, f, Fs)
+# H: (2, 2, n_freq) transfer function
+# S: (2, 2, n_freq) cross-spectral matrix
+```
+
+`YuleWalker_multitrial` is also available if you'd rather not average the per-trial fits by hand.
+
+### Conditional GC (multivariate, p ≥ 3)
+
+```python
+import numpy as np
+from pygc import conditional_granger_causality, spectral_conditional_granger_causality
+from pygc import ar_model
+
+Fs = 200
+data = ar_model.ar_model_baccala(nvars=5, N=3000, ntrials=20)
+# data shape: (nvars, N, trials) -> transpose to (trials, nvars, N)
+data = np.transpose(data, (2, 0, 1))
+
+# Time-domain conditional GC — Dataset with an (nvars, nvars) 'F' matrix
+ds_f = conditional_granger_causality(data, fs=Fs, n_jobs=-1)
+
+# Spectral conditional GC — Dataset with an (nvars, nvars, n_freq) 'GC' array
+ds_gc = spectral_conditional_granger_causality(data, fs=Fs, n_jobs=-1)
+```
+
+Both functions support the same `spectral_method` and `spectral_params` arguments as `spectral_granger_causality`.
+
+### JAX/XLA accelerated backend
+
+```python
+from pygc import spectral_granger_causality, JAX_AVAILABLE
+
+if JAX_AVAILABLE:
+    ds = spectral_granger_causality(data, fs=Fs, backend='jax')
+```
+
+### `ensure_stability` parameter
+
+The GC functions accept `ensure_stability=True` (default), which clips near-zero or negative diagonal entries of the noise covariance after Wilson factorization to improve numerical stability.
+
+---
+
+## Testing
+
+```bash
+pytest              # run all 38 tests
+pytest --cov=pygc   # with coverage report
+```
+
+---
+
+## Example notebooks
+
+See `notebooks/` for end-to-end worked examples:
+
+- `01_basic_granger_causality.ipynb` — bivariate GC on the Dhamala benchmark
+- `02_spectral_analysis.ipynb` — cross-spectral and time-frequency estimation (Fourier, Welch, multitaper, Morlet)
+- `03_conditional_granger_causality.ipynb` — time-domain and spectral conditional GC on the 5-variable Baccalá model
+- `04_benchmarks.ipynb` — performance benchmarks across backends and model sizes
+- `05_spectral_methods_comparison.ipynb` — comparison of Fourier, Welch, Morlet, and multitaper estimators
+- `06_example_with_eletrophysiological_data.ipynb` — GC applied to real electrophysiological recordings
+
+## Hardware used to run notebooks
+
+Apple M4 Pro with 48 GB RAM

pygranger-2.0.0/README.md

@@ -0,0 +1,209 @@
+# pyGC — Granger Causality in the Frequency Domain
+
+[![PyPI](https://img.shields.io/pypi/v/pygranger.svg)](https://pypi.org/project/pygranger/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/)
+
+A Python package for estimating Granger Causality (GC) in the frequency domain, supporting multiple spectral estimators (Fourier, Welch, Morlet, multitaper) and both NumPy and JAX/XLA backends.
+
+If you use this package, please cite:
+
+---
+
+## Installation
+
+The package is published on PyPI as `pygranger` and imported as `pygc`:
+
+```bash
+pip install pygranger
+```
+
+```python
+import pygc
+```
+
+Optional extras:
+
+```bash
+pip install "pygranger[jax]"       # JAX/XLA GPU-accelerated backend
+pip install "pygranger[notebooks]" # Jupyter notebooks
+```
+
+For development, install from a clone in editable mode:
+
+```bash
+git clone https://github.com/ViniciusLima94/pyGC.git
+cd pyGC
+pip install -e ".[dev]"            # pytest + coverage
+```
+
+---
+
+## Package structure
+
+```
+pygc/
+├── __init__.py               # Public API
+├── parametric.py             # Yule-Walker VAR fitting + transfer function
+├── non_parametric.py         # Wilson spectral factorization
+├── granger.py                # Bivariate GC, conditional GC (time + spectral)
+├── output.py                 # Packages results into labeled xarray.Dataset objects
+├── ar_model.py                # Synthetic AR benchmarks (Dhamala, Baccalá)
+├── _jax_backend.py           # Optional JAX/XLA Wilson factorization
+├── spectral_analysis/
+│   ├── fourier.py            # compute_freq, CSD, Morlet transforms (MNE)
+│   └── time_frequency.py     # Welch, multitaper, wavelet/Gabor transforms
+└── misc/
+    └── smooth_spectra.py     # smooth_spectra, downsample utilities
+```
+
+---
+
+## Usage
+
+All main functions accept raw signal data `X` and a sampling rate `fs`. The cross-spectral density is computed internally via the chosen `spectral_method`. Every GC function returns an `xarray.Dataset` (see `output.py`) rather than raw arrays, so results come back labeled and ready to index by name.
+
+> **Data shape convention:** the GC functions expect `X` as `(trials, nvars, N)` (or `(nvars, N)` for a single trial).
+
+### Bivariate GC
+
+```python
+import numpy as np
+from pygc import spectral_granger_causality
+from pygc import ar_model
+
+Fs = 200
+# Dhamala benchmark: Y drives X at 40 Hz
+data = ar_model.ar_model_dhamala(N=5000, Trials=20, Fs=Fs, C=0.25)
+# data shape: (2, Trials, N) -> transpose to (Trials, 2, N) for the GC functions
+data = np.transpose(data, (1, 0, 2))
+
+# Non-parametric (Fourier CSD + Wilson factorization) — default
+ds = spectral_granger_causality(data, fs=Fs, spectral_method='fourier', pairs=[(0, 1)])
+# ds.y2x peaks at ~40 Hz; ds.x2y is near zero
+peak_freq = ds.freq.values[ds.y2x.values[0].argmax()]
+```
+
+`spectral_granger_causality` returns a frequency-resolved `Dataset` with `x2y`, `y2x`, and `xy` variables (dims `pairs`, `freq`). For a single scalar per direction per pair instead (the broadband/time-domain summary), use `granger_causality`, which returns a `Dataset` with a single `F` variable (dims `source`, `target`):
+
+```python
+from pygc import granger_causality
+
+ds_f = granger_causality(data, fs=Fs, spectral_method='fourier')
+# ds_f.F[1, 0] -> broadband GC from channel 1 (Y) to channel 0 (X)
+```
+
+Available spectral methods: `'fourier'` (default), `'welch'`, `'morlet'`, `'multitaper'`.
+
+```python
+# Welch CSD
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='welch',
+    spectral_params={'nperseg': 512}
+)
+
+# Morlet CSD (requires explicit frequency axis)
+freqs = np.linspace(1, 80, 80)
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='morlet',
+    spectral_params={'freqs': freqs, 'n_cycles': 7.0}
+)
+
+# Multitaper CSD
+ds = spectral_granger_causality(
+    data, fs=Fs, spectral_method='multitaper',
+    spectral_params={'bandwidth': 4.0}
+)
+```
+
+### Parametric estimation (Yule-Walker)
+
+Use `YuleWalker` and `compute_transfer_function` to obtain the VAR-based cross-spectral matrix:
+
+```python
+import numpy as np
+from pygc import YuleWalker, compute_transfer_function
+from pygc.spectral_analysis import compute_freq
+from pygc import ar_model
+
+Fs = 200
+data = ar_model.ar_model_dhamala(N=5000, Trials=20, Fs=Fs, C=0.25)
+# data shape: (2, Trials, N) — YuleWalker takes (nvars, N) per trial, so no
+# transpose is needed here (unlike the GC functions above).
+
+f  = compute_freq(data.shape[2], Fs)
+m  = 2                          # VAR model order
+AR  = np.zeros((m, 2, 2))
+SIG = np.zeros((2, 2))
+
+for trial in range(data.shape[1]):
+    a, s = YuleWalker(data[:, trial, :], m)
+    AR  += a / data.shape[1]
+    SIG += s / data.shape[1]
+
+H, S = compute_transfer_function(AR, SIG, f, Fs)
+# H: (2, 2, n_freq) transfer function
+# S: (2, 2, n_freq) cross-spectral matrix
+```
+
+`YuleWalker_multitrial` is also available if you'd rather not average the per-trial fits by hand.
+
+### Conditional GC (multivariate, p ≥ 3)
+
+```python
+import numpy as np
+from pygc import conditional_granger_causality, spectral_conditional_granger_causality
+from pygc import ar_model
+
+Fs = 200
+data = ar_model.ar_model_baccala(nvars=5, N=3000, ntrials=20)
+# data shape: (nvars, N, trials) -> transpose to (trials, nvars, N)
+data = np.transpose(data, (2, 0, 1))
+
+# Time-domain conditional GC — Dataset with an (nvars, nvars) 'F' matrix
+ds_f = conditional_granger_causality(data, fs=Fs, n_jobs=-1)
+
+# Spectral conditional GC — Dataset with an (nvars, nvars, n_freq) 'GC' array
+ds_gc = spectral_conditional_granger_causality(data, fs=Fs, n_jobs=-1)
+```
+
+Both functions support the same `spectral_method` and `spectral_params` arguments as `spectral_granger_causality`.
+
+### JAX/XLA accelerated backend
+
+```python
+from pygc import spectral_granger_causality, JAX_AVAILABLE
+
+if JAX_AVAILABLE:
+    ds = spectral_granger_causality(data, fs=Fs, backend='jax')
+```
+
+### `ensure_stability` parameter
+
+The GC functions accept `ensure_stability=True` (default), which clips near-zero or negative diagonal entries of the noise covariance after Wilson factorization to improve numerical stability.
+
+---
+
+## Testing
+
+```bash
+pytest              # run all 38 tests
+pytest --cov=pygc   # with coverage report
+```
+
+---
+
+## Example notebooks
+
+See `notebooks/` for end-to-end worked examples:
+
+- `01_basic_granger_causality.ipynb` — bivariate GC on the Dhamala benchmark
+- `02_spectral_analysis.ipynb` — cross-spectral and time-frequency estimation (Fourier, Welch, multitaper, Morlet)
+- `03_conditional_granger_causality.ipynb` — time-domain and spectral conditional GC on the 5-variable Baccalá model
+- `04_benchmarks.ipynb` — performance benchmarks across backends and model sizes
+- `05_spectral_methods_comparison.ipynb` — comparison of Fourier, Welch, Morlet, and multitaper estimators
+- `06_example_with_eletrophysiological_data.ipynb` — GC applied to real electrophysiological recordings
+
+## Hardware used to run notebooks
+
+Apple M4 Pro with 48 GB RAM

pygranger-2.0.0/pygc/__init__.py

@@ -0,0 +1,24 @@
+from .parametric import YuleWalker, YuleWalker_multitrial, compute_transfer_function
+from .non_parametric import wilson_factorization
+from .granger import granger_causality, spectral_granger_causality, conditional_granger_causality, spectral_conditional_granger_causality #noqa
+from ._jax_backend import JAX_AVAILABLE, JAX_FLOAT64, wilson_factorization_jax
+from .output import build_granger_dataset, build_conditional_gc_dataset, build_conditional_spec_gc_dataset
+
+__version__ = "2.0.0"
+
+__all__ = [
+    "YuleWalker",
+    "YuleWalker_multitrial",
+    "compute_transfer_function",
+    "wilson_factorization",
+    "wilson_factorization_jax",
+    "JAX_AVAILABLE",
+    "JAX_FLOAT64",
+    "granger_causality",
+    "spectral_granger_causality",
+    "conditional_granger_causality",
+    "spectral_conditional_granger_causality",
+    "build_granger_dataset",
+    "build_conditional_gc_dataset",
+    "build_conditional_spec_gc_dataset"
+]

pygranger-2.0.0/pygc/_jax_backend.py

@@ -0,0 +1,200 @@
+"""
+Optional JAX backend for pygc.
+
+Provides a JIT-compiled Wilson spectral factorization that uses XLA kernels
+(CPU or GPU) for batched matrix operations.  All loops — including the
+convergence-checked iteration — are compiled with ``jax.lax.while_loop``
+so no Python overhead remains after the first (tracing) call.
+
+GPU support
+-----------
+* **Apple Silicon (Metal)**: ``pip install jax-metal``
+  Metal does not support float64, so the computation runs in complex64
+  (float32 precision).  Accuracy is lower but throughput is higher.
+* **NVIDIA (CUDA)**: install the appropriate ``jax[cuda]`` wheel.
+  Full float64 / complex128 precision is available.
+
+Usage
+-----
+from pygc._jax_backend import wilson_factorization_jax, JAX_AVAILABLE, JAX_FLOAT64
+"""
+from __future__ import annotations
+
+import functools
+from typing import Any
+import numpy as np
+
+# Pre-bind to None so the names are always defined; typed as Any so that
+# attribute access inside `if JAX_AVAILABLE:` blocks is not flagged.
+jax: Any = None
+jnp: Any = None
+lax: Any = None
+
+try:
+    import jax
+    jax.config.update("jax_enable_x64", True)   # no-op on Metal, active on CPU/CUDA
+    import jax.numpy as jnp
+    from jax import lax
+    JAX_AVAILABLE: bool = True
+
+    # Metal GPU supports only float32; CPU and CUDA support float64.
+    # Metal raises ValueError on float64 ops instead of silently downcasting.
+    try:
+        _probe = jnp.ones(1, dtype=jnp.float64)
+        JAX_FLOAT64: bool = bool(_probe.dtype == jnp.float64)
+        del _probe
+    except (ValueError, TypeError):
+        JAX_FLOAT64 = False
+
+    # Choose working dtypes based on device capability.
+    _NP_CDTYPE  = np.complex128  if JAX_FLOAT64 else np.complex64
+    _JNP_CDTYPE = jnp.complex128 if JAX_FLOAT64 else jnp.complex64
+
+except ImportError:           # pragma: no cover
+    JAX_AVAILABLE = False
+    JAX_FLOAT64   = False
+    _NP_CDTYPE    = None
+    _JNP_CDTYPE   = None
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers (only defined when JAX is present)
+# ---------------------------------------------------------------------------
+
+if JAX_AVAILABLE:
+
+    def _plus_op(g: Any, freq_len: int) -> Any:
+        """Plus-operator: batch IFFT → zero causal half → batch FFT."""
+        gam  = jnp.fft.ifft(g, axis=2)
+        gamp = (gam
+                .at[:, :, 0].set(jnp.triu(0.5 * gam[:, :, 0]))
+                .at[:, :, freq_len:].set(0.0 + 0.0j))
+        return jnp.fft.fft(gamp, axis=2)
+
+    @functools.partial(jax.jit, static_argnames=["freq_len", "Niterations"])
+    def _wilson_loop(
+        psi: Any,
+        Sarr: Any,
+        I: Any,
+        freq_len: int,
+        Niterations: int,
+        tol: float,
+    ) -> Any:
+        """Full Wilson iteration loop compiled to XLA via lax.while_loop."""
+
+        def body(carry: tuple) -> tuple:
+            psi, _, itr = carry
+            psi_T   = psi.transpose(2, 0, 1)                       # (2N, m, m)
+            psi_inv = jnp.linalg.inv(psi_T)
+            g_T     = psi_inv @ Sarr.transpose(2, 0, 1) @ jnp.conj(psi_inv).swapaxes(-1, -2) + I
+            g       = g_T.transpose(1, 2, 0)
+
+            gp      = _plus_op(g, freq_len)
+
+            new_psi_T = psi_T @ gp.transpose(2, 0, 1)
+            diff    = new_psi_T - psi_T
+            psierr  = jnp.abs(diff).sum(axis=-2).max(axis=-1).mean()
+
+            return new_psi_T.transpose(1, 2, 0), psierr, itr + 1
+
+        def cond(carry: tuple) -> Any:
+            _, psierr, itr = carry
+            return (psierr >= tol) & (itr < Niterations)
+
+        init = (psi, jnp.array(jnp.inf), jnp.int32(0))
+        psi_final, _, _ = lax.while_loop(cond, body, init)
+        return psi_final
+
+    def wilson_factorization_jax(
+        S: np.ndarray,
+        freq: np.ndarray,
+        fs: float,
+        Niterations: int = 100,
+        tol: float = 1e-12,
+        verbose: bool = False,
+        ensure_stability: bool = True,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Wilson spectral factorization using a JAX JIT-compiled loop.
+
+        The first call incurs XLA compilation time (~1–5 s); subsequent calls
+        with the same array shapes reuse the compiled kernel.
+
+        Precision is chosen automatically:
+        - **complex128** on CPU and NVIDIA GPU (float64 available).
+        - **complex64**  on Apple Metal GPU (float64 not supported).
+
+        Parameters
+        ----------
+        S           : ndarray (m, m, N+1) — cross-spectral matrix.
+        freq        : ndarray (N+1,)      — frequency axis.
+        fs          : float               — sampling rate (unused; kept for API compat).
+        Niterations : int.
+        tol         : float.
+        verbose     : bool                — print JAX device / precision info.
+
+        Returns
+        -------
+        Snew, Hnew : ndarray (m, m, N+1)
+        Znew       : ndarray (m, m)
+        """
+        m        = S.shape[0]
+        N        = freq.shape[0] - 1
+        freq_len = len(freq)
+
+        if verbose:
+            prec = "complex128 (float64)" if JAX_FLOAT64 else "complex64 (float32, Metal GPU)"
+            print(f'[JAX] devices: {jax.devices()}  precision: {prec}')
+
+        # Build Hermitian-extended Sarr (numpy; done once outside JIT)
+        Sarr_np = np.zeros([m, m, 2 * N], dtype=_NP_CDTYPE)
+        Sarr_np[:, :, :N + 1] = S
+        for k in range(1, N):
+            Sarr_np[:, :, 2 * N - k] = S[:, :, k].T
+
+        # Diagonal regularization fix
+        if ensure_stability:
+            eps = np.finfo(float).eps * np.abs(Sarr_np).max() * 100
+            for i in range(m):
+                Sarr_np[i, i, :] += eps        
+
+        # Initialise psi from Cholesky of gam0.
+        # A small relative regularisation guards against float32 rounding that
+        # can make gam0 appear numerically indefinite on Metal GPU.
+        gam0 = np.fft.ifft(Sarr_np, axis=2).real[:, :, 0]
+        eps  = np.finfo(gam0.dtype).eps * np.abs(np.diag(gam0)).max() * 10
+        gam0 += np.eye(m) * eps
+        h      = np.linalg.cholesky(gam0).T
+        psi_np = np.tile(h[:, :, np.newaxis], (1, 1, 2 * N)).astype(_NP_CDTYPE)
+
+        # Move to JAX arrays and run compiled loop
+        psi   = jnp.array(psi_np)
+        Sarr  = jnp.array(Sarr_np)
+        I_jax = jnp.eye(m, dtype=_JNP_CDTYPE)[jnp.newaxis]  # (1, m, m)
+
+        psi = _wilson_loop(psi, Sarr, I_jax, freq_len, Niterations, tol)
+        psi = np.array(psi)  # back to NumPy for output
+
+        # Compute outputs (NumPy; trivially fast)
+        psi_T  = psi.transpose(2, 0, 1)
+        Snew   = (psi_T[:N + 1] @ np.conj(psi_T[:N + 1]).swapaxes(-1, -2)).transpose(1, 2, 0)
+        gamtmp = np.fft.ifft(psi, axis=2).real
+        A0     = gamtmp[:, :, 0]
+        A0inv  = np.linalg.inv(A0)
+        Znew   = (A0 @ A0.T).real
+        Hnew   = (psi_T[:N + 1] @ A0inv[np.newaxis]).transpose(1, 2, 0)
+
+        return Snew, Hnew, Znew
+
+else:  # JAX not available — stub that raises a helpful error  # pragma: no cover
+
+    def wilson_factorization_jax(  # type: ignore[misc]
+        S: np.ndarray,
+        freq: np.ndarray,
+        fs: float,
+        Niterations: int = 100,
+        tol: float = 1e-12,
+        verbose: bool = False,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        raise ImportError(
+            "JAX is not installed. Install it with: pip install jax"
+        )