PyAETNA 4.0.0__py3-none-any.whl

PyAETNA/kernels/multichannel.py

@@ -0,0 +1,417 @@
+"""Multi-dimensional GP kernel (Rajpaul+2015 / Barragán+2022) for PyAETNA.
+
+This module implements the *joint* cross-output covariance that couples several
+observables to one latent stellar-activity process — as opposed to the
+block-diagonal, one-GP-per-instrument likelihood built by
+:mod:`PyAETNA.kernels.factory`.
+
+The model
+---------
+A single latent, **unit-variance** process :math:`G(t)` with covariance
+:math:`\\gamma(t_i, t_j)`.  Each observable channel :math:`c` is a linear
+projection of that process and its time derivative:
+
+.. math::
+
+    \\Delta_c(t) = A_c\\, G(t) + B_c\\, \\dot{G}(t)
+
+This is Barragán et al. (2022) Eq. 1, the N-channel generalisation of Rajpaul
+et al. (2015) Eqs. 14-16.  The paper's three observables are the special case
+``RV -> (V_c, V_r)``, ``BIS -> (B_c, B_r)``, ``log R'HK -> (L_c, 0)``.
+
+The covariance between channels :math:`c` and :math:`d` (Rajpaul Eq. 23) is
+
+.. math::
+
+    \\mathrm{cov}(\\Delta_c(t_i), \\Delta_d(t_j)) =
+        A_c A_d\\, \\gamma
+      + A_c B_d\\, \\partial\\gamma/\\partial t_j
+      + B_c A_d\\, \\partial\\gamma/\\partial t_i
+      + B_c B_d\\, \\partial^2\\gamma/\\partial t_i \\partial t_j
+
+The two middle terms are antisymmetric: they vanish at zero lag but not
+otherwise, and they encode the phase relationship between observables.  They
+are exactly what a block-diagonal likelihood discards.
+
+Implementation notes
+--------------------
+* The derivative blocks are obtained by **nested** :func:`jax.grad` on the
+  scalar latent function, so any twice-differentiable latent kernel works with
+  no hand-derived algebra.
+* The kernel operates on ``tinygp`` pytree inputs ``X = (t, channel_index)``.
+* The latent kernel carries **no amplitude**: :math:`\\gamma` is normalised to
+  unit variance and all amplitude information lives in :math:`A_c`/:math:`B_c`.
+  Otherwise ``(G, A, B) -> (kG, A/k, B/k)`` would be unidentifiable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from tinygp import kernels
+
+__all__ = [
+    "MULTIGP_PARAM_ORDERS",
+    "MULTIGP_SUPPORTED_KERNELS",
+    "ChannelAutoKernel",
+    "MultiChannelKernel",
+    "build_latent_kernel_fn",
+    "get_multigp_param_order",
+    "multigp_log_probability",
+    "multigp_log_probability_blocked",
+]
+
+# Latent kernels admissible for a multi-GP block.
+#
+# ``Gdot`` only exists if ``gamma`` is twice differentiable in the mean-square
+# sense.  Matern32 is deliberately EXCLUDED: the Matern-3/2 process is only
+# *once* mean-square differentiable, so its derivative process is not defined.
+# It cannot be left to fail at runtime — autodiff happily returns a finite,
+# plausible-looking second derivative (JAX takes the subgradient of |tau| at
+# zero) and the sampled matrix stays numerically positive semi-definite, so the
+# misspecification is completely silent.  Hence an explicit whitelist.
+#
+# The ``_quasisep`` kernels and SHO are excluded for a different reason: the
+# quasiseparable solver cannot represent the cross-derivative blocks, so they
+# would silently fall back to a dense solve anyway.
+MULTIGP_SUPPORTED_KERNELS = (
+    "ExpSquared",
+    "RBF",
+    "QuasiPeriodic",
+    "ExpSineSquared",
+    "Cosine",
+    "Matern52",
+    "RationalQuadratic",
+)
+
+# Hyperparameter names per latent kernel.  These mirror
+# ``PyAETNA.sampling.utils._get_default_param_order`` but **omit
+# ``log_gp_sigma``**: the latent process is unit-variance by construction (see
+# the module docstring).  The naming convention is unchanged — everything is
+# sampled in log scale except ``gp_freq``, and the kernel converts
+# ``period = 1 / freq`` internally.
+MULTIGP_PARAM_ORDERS: dict[str, list[str]] = {
+    "ExpSquared": ["log_gp_scale"],
+    "RBF": ["log_gp_scale"],
+    "Matern52": ["log_gp_scale"],
+    "RationalQuadratic": ["log_gp_scale", "log_gp_alpha"],
+    "Cosine": ["gp_freq"],
+    "ExpSineSquared": ["gp_freq", "log_gp_gamma"],
+    "QuasiPeriodic": ["log_gp_se_scale", "gp_freq", "log_gp_gamma"],
+}
+
+
+def get_multigp_param_order(kernel_type: str) -> list[str]:
+    """Return the latent hyperparameter names for a multi-GP kernel type.
+
+    Args:
+        kernel_type: Kernel identifier (case-insensitive).
+
+    Returns:
+        Ordered list of hyperparameter names, or ``[]`` if unknown.
+    """
+    for name, order in MULTIGP_PARAM_ORDERS.items():
+        if name.lower() == kernel_type.lower():
+            return list(order)
+    return []
+
+
+def build_latent_kernel_fn(
+    kernel_type: str,
+) -> Callable[[dict[str, Any], Any, Any], Any]:
+    """Build the scalar, unit-variance latent covariance ``gamma(params, t1, t2)``.
+
+    The hyperparameters are an *argument*, not a closure capture, so that
+    :class:`MultiChannelKernel` can hold the function as a **static** equinox
+    field while the hyperparameters stay traced — which is what lets gradients
+    reach them under NUTS.
+
+    The returned function is differentiable in ``t1`` and ``t2``, so the kernel
+    can form its derivative blocks with :func:`jax.grad`.
+
+    Args:
+        kernel_type: One of :data:`MULTIGP_SUPPORTED_KERNELS`.
+
+    Returns:
+        ``gamma(params, t1, t2)`` returning a scalar.  ``params`` holds
+        hyperparameters in **linear** scale (already exponentiated from their
+        log-sampled values), keyed by the names in
+        :data:`MULTIGP_PARAM_ORDERS` with any ``log_`` prefix stripped — e.g.
+        ``{"gp_scale": 12.0}``.
+
+    Raises:
+        ValueError: If ``kernel_type`` is not admissible for a multi-GP block.
+    """
+    kt = kernel_type.lower()
+
+    if kt not in {k.lower() for k in MULTIGP_SUPPORTED_KERNELS}:
+        raise ValueError(
+            f"Kernel '{kernel_type}' cannot be used as a multi-GP latent kernel. "
+            f"The latent process must be twice differentiable so that its "
+            f"derivative Gdot exists. Supported: "
+            f"{', '.join(MULTIGP_SUPPORTED_KERNELS)}."
+        )
+
+    if kt in ("expsquared", "rbf"):
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            return jnp.exp(-0.5 * ((t1 - t2) / p["gp_scale"]) ** 2)
+
+    elif kt == "matern52":
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            r = jnp.sqrt(5.0) * jnp.abs(t1 - t2) / p["gp_scale"]
+            return (1.0 + r + r**2 / 3.0) * jnp.exp(-r)
+
+    elif kt == "rationalquadratic":
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            alpha = p["gp_alpha"]
+            return (1.0 + 0.5 * ((t1 - t2) / p["gp_scale"]) ** 2 / alpha) ** (-alpha)
+
+    elif kt == "cosine":
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            return jnp.cos(2.0 * jnp.pi * (t1 - t2) * p["gp_freq"])
+
+    elif kt == "expsinesquared":
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            return jnp.exp(
+                -p["gp_gamma"] * jnp.sin(jnp.pi * (t1 - t2) * p["gp_freq"]) ** 2
+            )
+
+    else:  # quasiperiodic
+
+        def gamma(p: dict[str, Any], t1: Any, t2: Any) -> Any:
+            tau = t1 - t2
+            return jnp.exp(
+                -0.5 * (tau / p["gp_se_scale"]) ** 2
+                - p["gp_gamma"] * jnp.sin(jnp.pi * tau * p["gp_freq"]) ** 2
+            )
+
+    return gamma
+
+
+class MultiChannelKernel(kernels.Kernel):
+    """Rajpaul/Barragán cross-output covariance over ``X = (t, channel_index)``.
+
+    Attributes:
+        latent: Scalar unit-variance covariance ``gamma(params, t1, t2)``
+            from :func:`build_latent_kernel_fn`; a **static** field.
+        params: Latent hyperparameters in linear scale; a traced field, so
+            gradients reach them.
+        A: Shape ``(n_channels,)``. Projection onto ``G``; signed.
+        B: Shape ``(n_channels,)``. Projection onto ``Gdot``; signed.
+
+    ``A`` and ``B`` are deliberately **linear, signed** quantities rather than
+    log-sampled: in Rajpaul et al. the relative signs of ``(V_c, V_r)`` and
+    ``(B_c, B_r)`` are physically meaningful (BIS typically anti-correlates
+    with RV), and the sign structure is normally what the fit should measure.
+
+    The residual ``(G, A, B) -> (-G, -A, -B)`` sign degeneracy is broken by the
+    parser, which requires the reference channel's ``A`` prior to be bounded
+    below at zero.
+    """
+
+    latent: Callable[[dict[str, Any], Any, Any], Any] = eqx.field(static=True)
+    params: dict[str, Any]
+    A: jnp.ndarray
+    B: jnp.ndarray
+
+    def evaluate(self, X1: Any, X2: Any) -> Any:
+        """Covariance between one pair of (time, channel) points."""
+        t1, c1 = X1
+        t2, c2 = X2
+
+        # `latent` is a static field (a plain function) while `params` is a
+        # traced field, so gradients reach the latent hyperparameters.
+        def k(a: Any, b: Any) -> Any:
+            return self.latent(self.params, a, b)
+
+        dk_dt1 = jax.grad(k, argnums=0)
+        dk_dt2 = jax.grad(k, argnums=1)
+        d2k = jax.grad(dk_dt1, argnums=1)
+
+        a1, b1 = self.A[c1], self.B[c1]
+        a2, b2 = self.A[c2], self.B[c2]
+
+        return (
+            a1 * a2 * k(t1, t2)
+            + a1 * b2 * dk_dt2(t1, t2)
+            + b1 * a2 * dk_dt1(t1, t2)
+            + b1 * b2 * d2k(t1, t2)
+        )
+
+
+def multigp_log_probability(
+    kernel: MultiChannelKernel,
+    t: jnp.ndarray,
+    channel: jnp.ndarray,
+    res: jnp.ndarray,
+    diag: jnp.ndarray,
+) -> jnp.ndarray:
+    """Marginal log-probability of the full joint multi-channel covariance.
+
+    Args:
+        kernel: The :class:`MultiChannelKernel`.
+        t: Times of every point, all channels concatenated. Shape ``(N,)``.
+        channel: Channel index of every point, same order. Shape ``(N,)``.
+        res: Residuals (observable minus its mean model). Shape ``(N,)``.
+        diag: Per-point variance (measurement error squared plus jitter).
+            Shape ``(N,)``.
+
+    Returns:
+        Scalar log-probability.  Cost is ``O(N^3)`` with ``N`` the total across
+        **all** channels.
+    """
+    from tinygp import GaussianProcess
+
+    # tinygp types X as an Array, but a pytree input is the documented way to
+    # carry structured inputs (its "Derivative Observations & Pytree Data"
+    # pattern).  The tuple is what MultiChannelKernel.evaluate unpacks.
+    gp = GaussianProcess(kernel, (t, channel), diag=diag)  # type: ignore[arg-type]
+    return gp.log_probability(res)
+
+
+def multigp_log_probability_blocked(
+    kernel: MultiChannelKernel,
+    t: jnp.ndarray,
+    channel: jnp.ndarray,
+    res: jnp.ndarray,
+    diag: jnp.ndarray,
+    block_indices: list[Any],
+) -> jnp.ndarray:
+    """Block-diagonal version of :func:`multigp_log_probability`.
+
+    Each block holds **all channels** for the epochs it covers, so the
+    cross-channel information — the whole point of the multi-GP — is preserved
+    *within* a block.  Only the covariance *between* blocks is dropped.
+
+    .. math::
+
+        \\log p(\\mathbf{y} \\mid \\theta) = \\sum_b \\log p(\\mathbf{y}_b \\mid \\theta)
+
+    Cost drops from ``N^3`` to ``sum_b n_b^3``, i.e. roughly ``N^3 / n_blocks^2``.
+
+    .. warning::
+
+        Valid only when the between-block gap greatly exceeds the latent
+        kernel's decay scale, so the dropped off-block covariance really is
+        negligible.  As with
+        :func:`PyAETNA.gp.gp_log_probability_per_visit`, this is **never**
+        valid for a pure ``ExpSineSquared`` latent kernel, whose covariance
+        does not decay at all: ``gamma(nP) == 1`` for any integer ``n``.
+        For ``QuasiPeriodic`` it is conditionally valid, when the gap greatly
+        exceeds ``exp(log_gp_se_scale)``.
+
+    Padding follows the convention of
+    :func:`PyAETNA.gp.gp_log_probability_per_visit`: padded points get
+    ``diag = 1e30`` (effectively infinite noise, so they cannot affect the
+    real-data contribution), zero residuals, and times pushed far beyond the
+    last real point so their covariance with real data vanishes.
+
+    Args:
+        kernel: The :class:`MultiChannelKernel`.
+        t: Times of every point, all channels concatenated.
+        channel: Channel index of every point, same order.
+        res: Residuals.
+        diag: Per-point variance.
+        block_indices: One index array per block, giving positions within the
+            concatenated arrays.  Each must already contain every channel's
+            points for its epochs.
+
+    Returns:
+        Scalar log-probability summed over blocks.
+    """
+    n_blocks = len(block_indices)
+    if n_blocks == 0:
+        return jnp.zeros(())
+    if n_blocks == 1:
+        idx = block_indices[0]
+        return multigp_log_probability(
+            kernel, t[idx], channel[idx], res[idx], diag[idx]
+        )
+
+    max_len = max(len(idx) for idx in block_indices)
+
+    t_pad = jnp.zeros((n_blocks, max_len), dtype=t.dtype)
+    c_pad = jnp.zeros((n_blocks, max_len), dtype=channel.dtype)
+    r_pad = jnp.zeros((n_blocks, max_len), dtype=res.dtype)
+    d_pad = jnp.full((n_blocks, max_len), 1e30, dtype=diag.dtype)
+
+    for i, idx in enumerate(block_indices):
+        n_b = len(idx)  # static at trace time
+        t_pad = t_pad.at[i, :n_b].set(t[idx])
+        c_pad = c_pad.at[i, :n_b].set(channel[idx])
+        r_pad = r_pad.at[i, :n_b].set(res[idx])
+        d_pad = d_pad.at[i, :n_b].set(diag[idx])
+        if n_b < max_len:
+            n_extra = max_len - n_b
+            last_t = t[idx[-1]]
+            pad_start = last_t + jnp.asarray(1e4, dtype=t.dtype)
+            t_pad = t_pad.at[i, n_b:].set(
+                pad_start + jnp.arange(n_extra, dtype=t.dtype) * 0.001
+            )
+
+    def _block_logprob(
+        t_b: jnp.ndarray, c_b: jnp.ndarray, r_b: jnp.ndarray, d_b: jnp.ndarray
+    ) -> jnp.ndarray:
+        return multigp_log_probability(kernel, t_b, c_b, r_b, d_b)
+
+    return jnp.sum(jax.vmap(_block_logprob)(t_pad, c_pad, r_pad, d_pad))
+
+
+class ChannelAutoKernel(kernels.Kernel):
+    """One multi-GP channel's own covariance, as a function of scalar time.
+
+    :class:`MultiChannelKernel` needs ``(time, channel_index)`` pairs, which
+    makes it unusable with the report's PSD helpers — those call
+    ``kernel(t, t)`` on a plain time array.  This wrapper exposes a single
+    channel's auto-covariance in that ordinary form:
+
+    .. math::
+
+        k(t_1,t_2) = A^2\\gamma
+                   + A B\\,\\partial\\gamma/\\partial t_2
+                   + B A\\,\\partial\\gamma/\\partial t_1
+                   + B^2\\,\\partial^2\\gamma/\\partial t_1\\partial t_2
+
+    which is the diagonal block of the joint covariance — exactly the kernel
+    that channel would have had if it were fitted on its own.  It is what a
+    periodogram PSD overlay should be compared against, since the periodogram
+    is likewise computed from that channel alone.
+
+    Attributes:
+        latent: Scalar unit-variance covariance ``gamma(params, t1, t2)``;
+            a **static** field.
+        params: Latent hyperparameters in linear scale; a traced field.
+        A: Projection onto ``G`` for this channel.
+        B: Projection onto ``Gdot`` for this channel.
+    """
+
+    latent: Callable[[dict[str, Any], Any, Any], Any] = eqx.field(static=True)
+    params: dict[str, Any]
+    A: Any
+    B: Any
+
+    def evaluate(self, X1: Any, X2: Any) -> Any:
+        """Covariance between two scalar times."""
+
+        def k(a: Any, b: Any) -> Any:
+            return self.latent(self.params, a, b)
+
+        dk_dt1 = jax.grad(k, argnums=0)
+        dk_dt2 = jax.grad(k, argnums=1)
+        d2k = jax.grad(dk_dt1, argnums=1)
+
+        return (
+            self.A * self.A * k(X1, X2)
+            + self.A * self.B * dk_dt2(X1, X2)
+            + self.B * self.A * dk_dt1(X1, X2)
+            + self.B * self.B * d2k(X1, X2)
+        )

PyAETNA/logger.py

@@ -0,0 +1,73 @@
+"""
+Canonical repository-wide logging module for PyAETNA.
+
+Provides a unified logger `_log` supporting both implicit module-level log file
+destinations (configured via `set_log_file`) and explicit model/file target parameters.
+"""
+
+from __future__ import annotations
+
+import time as ttt
+from typing import Any
+
+_LOG_FILE: str | None = None
+_EIV_TRIGGERED = False
+
+
+def set_log_file(path: str | None) -> None:
+    """Set the module-wide default log file path.
+
+    All subsequent calls to `_log` without an explicit target will write to `path`.
+    Pass `None` to log to stdout by default.
+    """
+    global _LOG_FILE
+    _LOG_FILE = path
+
+
+def get_log_file() -> str | None:
+    """Return the current default log file path."""
+    return _LOG_FILE
+
+
+def _log(target_or_msg: Any, msg: str | None = None, level: str | None = None) -> None:
+    """Write a timestamped message to a log file or stdout.
+
+    Signatures supported:
+        _log("Message", level="WARNING")         # Default target (_LOG_FILE or stdout)
+        _log(plModel, "Message")                 # Explicit model or path target
+        _log(log_file_path, "Message")           # Explicit string path target
+
+    Args:
+        target_or_msg: Target object (with `.logFile` attribute), path (str), or message string.
+        msg (str, optional): Message string if target_or_msg is a target object/path.
+        level (str, optional): Severity level (e.g. 'WARNING', 'ERROR'). Standard 'INFO' level
+            is omitted from the output string as messages are informative by nature.
+    """
+    if msg is not None:
+        target = target_or_msg
+        message = msg
+    else:
+        target = None
+        message = target_or_msg
+
+    # Extract path from target object if present, otherwise use target path or module default
+    log_file = target.logFile if hasattr(target, "logFile") else (target or _LOG_FILE)
+
+    # Omit 'INFO' tag from output line; include non-INFO severity levels (WARNING, ERROR, etc.)
+    level_tag = f" {level.upper()}:" if level and level.upper() != "INFO" else ""
+    timestamp = ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime())
+    formatted = f"{timestamp}{level_tag} {message}"
+
+    if not log_file:
+        print(formatted)
+        return
+
+    try:
+        with open(log_file, "a") as f:
+            f.write(formatted + "\n")
+    except OSError as e:
+        print(f"{formatted} (warning: failed to write to log file {log_file!r}: {e})")
+
+
+# Default log file is stdout (None)
+set_log_file(None)