pygranger 2.0.0__py3-none-any.whl

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"
+]

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"
+        )

pygc/ar_model.py

@@ -0,0 +1,81 @@
+import numpy as np
+
+
+def ar_model_dhamala(N=5000, Trials=10, Fs=200, C=0.2, t_start=0, t_stop=None, cov=None):
+    """
+    Bivariate AR model from Dhamala et al. (2008).
+
+    Y drives X with coupling strength C in the interval [t_start, t_stop].
+    GC(Y->X) > 0, GC(X->Y) ≈ 0.
+
+    Parameters
+    ----------
+    N      : int   — number of time samples per trial.
+    Trials : int   — number of independent trials.
+    Fs     : float — sampling rate (Hz).
+    C      : float — coupling coefficient from Y to X.
+    t_start, t_stop : float — time interval (seconds) during which coupling is active.
+    cov    : ndarray, shape (2,2) — noise covariance matrix.
+
+    Returns
+    -------
+    Z : ndarray, shape (2, Trials, N)
+    """
+    if cov is None:
+        cov = np.eye(2)
+
+    T = N / Fs
+    time = np.linspace(0, T, N)
+
+    X = np.random.random([Trials, N])
+    Y = np.random.random([Trials, N])
+
+    def _active(t):
+        if t_stop is None:
+            return t >= t_start
+        return (t >= t_start) & (t <= t_stop)
+
+    for i in range(Trials):
+        E = np.random.multivariate_normal(np.zeros(2), cov, size=(N,))
+        for t in range(2, N):
+            X[i, t] = (0.55 * X[i, t - 1] - 0.8 * X[i, t - 2]
+                       + _active(time[t]) * C * Y[i, t - 1] + E[t, 0])
+            Y[i, t] = 0.55 * Y[i, t - 1] - 0.8 * Y[i, t - 2] + E[t, 1]
+
+    Z = np.zeros([2, Trials, N])
+    Z[0] = X
+    Z[1] = Y
+    return Z
+
+
+def ar_model_baccala(nvars=5, N=1000, ntrials=1):
+    """
+    Five-variable AR model from Baccalá & Sameshima (2001).
+
+    Connectivity: X1 → X2, X1 → X3, X1 → X4, X1 → X5 (via X4).
+
+    Parameters
+    ----------
+    nvars   : int — number of variables (must be 5).
+    N       : int — number of time samples per trial.
+    ntrials : int — number of trials.
+
+    Returns
+    -------
+    Y : ndarray, shape (nvars, N, ntrials)
+    """
+    Y = np.random.uniform(size=(nvars, N, ntrials))
+    w = np.random.normal(0, 1, size=(nvars, N, ntrials))
+
+    for i in range(ntrials):
+        for t in range(5, N):
+            Y[0, t, i] = (0.95 * np.sqrt(2.0) * Y[0, t - 1, i]
+                          - 0.9025 * Y[0, t - 2, i] + w[0, t, i])
+            Y[1, t, i] = 0.5 * Y[0, t - 2, i] + w[1, t, i]
+            Y[2, t, i] = -0.4 * Y[0, t - 3, i] + w[2, t, i]
+            Y[3, t, i] = (-0.5 * Y[0, t - 2, i]
+                          + 0.25 * np.sqrt(2.0) * Y[3, t - 1, i]
+                          + 0.25 * np.sqrt(2.0) * Y[4, t - 1, i] + w[3, t, i])
+            Y[4, t, i] = (-0.25 * np.sqrt(2.0) * Y[3, t - 1, i]
+                          + 0.25 * np.sqrt(2.0) * Y[4, t - 1, i] + w[4, t, i])
+    return Y