"""Prony series conversion transform.
Converts between time-domain and frequency-domain viscoelastic data using
Prony series decomposition:
- G(t) → G'(ω), G''(ω) via analytical Prony relations
- G'(ω), G''(ω) → G(t) via inverse Prony fitting
The Prony series representation is:
G(t) = G_e + Σ G_i exp(-t/τ_i)
G'(ω) = G_e + Σ G_i ω²τ_i² / (1 + ω²τ_i²)
G''(ω) = Σ G_i ωτ_i / (1 + ω²τ_i²)
"""
from __future__ import annotations
import warnings
from dataclasses import dataclass
from typing import Any
import numpy as np
from rheojax.core.base import BaseTransform
from rheojax.core.data import RheoData
from rheojax.core.jax_config import safe_import_jax
from rheojax.core.registry import TransformRegistry
from rheojax.logging import get_logger
from rheojax.utils.prony import compute_r_squared
jax, jnp = safe_import_jax()
logger = get_logger(__name__)
@dataclass
class PronyResult:
"""Result from Prony conversion."""
G_i: np.ndarray
tau_i: np.ndarray
G_e: float
n_modes: int
r_squared: float | None = None
[docs]
@TransformRegistry.register("prony_conversion", type="spectral")
class PronyConversion(BaseTransform):
"""Convert between time-domain and frequency-domain via Prony series.
Fitting and analytical conversion are self-contained (NNLS via
``scipy.optimize.nnls``); fit-quality scoring is delegated to
``rheojax.utils.prony.compute_r_squared``.
Args:
n_modes: Number of Prony modes (default: auto-select).
direction: ``"time_to_freq"`` or ``"freq_to_time"``.
omega_out: Target frequency array for time→freq conversion.
t_out: Target time array for freq→time conversion.
"""
[docs]
def __init__(
self,
n_modes: int | None = None,
direction: str = "time_to_freq",
omega_out: np.ndarray | None = None,
t_out: np.ndarray | None = None,
):
super().__init__()
self.n_modes = n_modes
self.direction = direction
self.omega_out = omega_out
self.t_out = t_out
self._prony_result: PronyResult | None = None
def _transform(self, data: RheoData) -> tuple[RheoData, dict[str, Any]]:
"""Apply Prony conversion.
Args:
data: Input RheoData. For ``time_to_freq``, expects relaxation
data (t, G(t)). For ``freq_to_time``, expects oscillation
data (ω, G* = G' + iG'').
Returns:
Tuple of (converted RheoData, metadata dict with PronyResult).
"""
if self.direction == "time_to_freq":
return self._time_to_freq(data)
elif self.direction == "freq_to_time":
return self._freq_to_time(data)
else:
raise ValueError(f"Invalid direction: {self.direction}")
def _time_to_freq(self, data: RheoData) -> tuple[RheoData, dict]:
"""G(t) → G'(ω), G''(ω) via Prony fit."""
t = np.asarray(data.x)
G_t = np.asarray(data.y)
# Determine number of modes
n_modes = (
self.n_modes if self.n_modes is not None else max(5, min(len(t) // 5, 20))
)
# Fit Prony series to G(t)
G_i, tau_i, G_e = _fit_prony_relaxation(t, G_t, n_modes)
r_squared = compute_r_squared(G_t, _prony_to_time(G_i, tau_i, G_e, t))
self._prony_result = PronyResult(
G_i=G_i, tau_i=tau_i, G_e=G_e, n_modes=len(G_i), r_squared=r_squared
)
# Convert to frequency domain
omega = self.omega_out
if omega is None:
# Use the data time range to set the valid frequency range:
# ω ∈ [1/t_max, 1/t_min]. The wider [0.1/τ_max, 10/τ_min] range
# extends into the extrapolation regime where the Prony fit is
# poorly constrained, inflating G'' relative errors.
t_positive = t[t > 0]
t_min = float(np.min(t_positive))
t_max = float(np.max(t_positive))
omega = np.logspace(
np.log10(1.0 / t_max),
np.log10(1.0 / t_min),
100,
)
G_prime, G_double_prime = _prony_to_frequency(G_i, tau_i, G_e, omega)
G_star = G_prime + 1j * G_double_prime
result_data = RheoData(
x=omega,
y=G_star,
x_units="rad/s",
y_units="Pa",
domain="frequency",
metadata={
"test_mode": "oscillation",
"source_transform": "prony_conversion",
"n_modes": len(G_i),
},
)
return result_data, {"prony_result": self._prony_result}
def _freq_to_time(self, data: RheoData) -> tuple[RheoData, dict]:
"""G'(ω), G''(ω) → G(t) via Prony fit to dynamic moduli."""
omega = np.asarray(data.x)
y = np.asarray(data.y)
if np.iscomplexobj(y):
G_prime = y.real
G_double_prime = y.imag
elif y.ndim == 2 and y.shape[1] == 2:
G_prime = y[:, 0]
G_double_prime = y[:, 1]
else:
raise ValueError(
"freq_to_time requires complex G* or (N, 2) array [G', G'']"
)
n_modes = (
self.n_modes
if self.n_modes is not None
else max(5, min(len(omega) // 5, 20))
)
G_i, tau_i, G_e = _fit_prony_oscillation(
omega, G_prime, G_double_prime, n_modes
)
G_prime_pred, G_double_prime_pred = _prony_to_frequency(
G_i, tau_i, G_e, omega
)
r_squared = compute_r_squared(
np.concatenate([G_prime, G_double_prime]),
np.concatenate([G_prime_pred, G_double_prime_pred]),
)
self._prony_result = PronyResult(
G_i=G_i, tau_i=tau_i, G_e=G_e, n_modes=len(G_i), r_squared=r_squared
)
t = self.t_out
if t is None:
# Mirror the time→freq convention: output t ∈ [1/omega_max, 1/omega_min],
# which is the data-informed range rather than the wider tau_i-based range.
# PC-001 (frequency-domain twin): filter to omega > 0, same root cause
# as the tau_i computation in _fit_prony_oscillation above.
omega_positive = omega[omega > 0]
t = np.logspace(
np.log10(1.0 / np.max(omega_positive)),
np.log10(1.0 / np.min(omega_positive)),
100,
)
G_t = _prony_to_time(G_i, tau_i, G_e, t)
result_data = RheoData(
x=t,
y=G_t,
x_units="s",
y_units="Pa",
domain="time",
metadata={
"test_mode": "relaxation",
"source_transform": "prony_conversion",
"n_modes": len(G_i),
},
)
return result_data, {"prony_result": self._prony_result}
# ---------------------------------------------------------------------------
# Analytical Prony conversions (pure functions, JIT-safe)
# ---------------------------------------------------------------------------
@jax.jit
def _prony_to_frequency_jax(
G_i: Any, tau_i: Any, G_e: float, omega: Any
) -> tuple[Any, Any]:
"""JIT-compiled Prony to frequency conversion."""
wt2 = (omega[:, None] * tau_i[None, :]) ** 2
G_prime = G_e + jnp.sum(G_i[None, :] * wt2 / (1.0 + wt2), axis=1)
G_double_prime = jnp.sum(
G_i[None, :] * omega[:, None] * tau_i[None, :] / (1.0 + wt2), axis=1
)
return G_prime, G_double_prime
def _prony_to_frequency(
G_i: np.ndarray, tau_i: np.ndarray, G_e: float, omega: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
"""Compute G'(ω) and G''(ω) from Prony parameters (JIT-accelerated)."""
G_p, G_dp = _prony_to_frequency_jax(
jnp.asarray(G_i, dtype=jnp.float64),
jnp.asarray(tau_i, dtype=jnp.float64),
G_e,
jnp.asarray(omega, dtype=jnp.float64),
)
return np.asarray(G_p), np.asarray(G_dp)
@jax.jit
def _prony_to_time_jax(G_i: Any, tau_i: Any, G_e: float, t: Any) -> Any:
"""JIT-compiled Prony to time conversion."""
return G_e + jnp.sum(G_i[None, :] * jnp.exp(-t[:, None] / tau_i[None, :]), axis=1)
def _prony_to_time(
G_i: np.ndarray, tau_i: np.ndarray, G_e: float, t: np.ndarray
) -> np.ndarray:
"""Compute G(t) from Prony parameters (JIT-accelerated)."""
return np.asarray(
_prony_to_time_jax(
jnp.asarray(G_i, dtype=jnp.float64),
jnp.asarray(tau_i, dtype=jnp.float64),
G_e,
jnp.asarray(t, dtype=jnp.float64),
)
)
def _fit_prony_relaxation(
t: np.ndarray, G_t: np.ndarray, n_modes: int
) -> tuple[np.ndarray, np.ndarray, float]:
"""Fit Prony series to relaxation modulus G(t)."""
from scipy.optimize import nnls
# Log-spaced relaxation times (use t[0] if positive, else t[1])
if len(t) < 2:
raise ValueError("Prony fitting requires at least 2 time points.")
t_positive = t[t > 0]
if len(t_positive) == 0:
raise ValueError("Prony fitting requires at least one positive time value.")
t_min = float(np.min(t_positive))
# PC-001: use np.max(t) not t[-1] — t may be unsorted so the last element
# is not necessarily the largest time. np.max guarantees t_max >= t_min,
# preventing a ValueError from log10 of a non-positive number.
t_max = float(np.max(t_positive))
# Use n+2 log-spaced values and drop the two edge modes. Modes at exactly
# tau = t_min / t_max are only "visible" in the outermost data points (high
# noise), so NNLS assigns them spurious weight that creates G''(ω) artifacts
# at ω ≈ 1/t_min. Dropping the edges keeps modes well inside the data
# range where they are reliably constrained.
tau_all = np.logspace(np.log10(t_min), np.log10(t_max), n_modes + 2)
tau_i = tau_all[1:-1]
# Equilibrium modulus estimate: G(t→∞) ≈ G at maximum time.
# Use the G value at argmax(t) to handle unsorted time arrays correctly.
G_e = max(float(G_t[np.argmax(t)]), 0.0)
# Build kernel matrix: A_ij = exp(-t_j / tau_i)
A = np.exp(-t[:, None] / tau_i[None, :])
# Solve non-negative least squares: G(t) - G_e ≈ A @ G_i
b = np.maximum(G_t - G_e, 0.0)
G_i, _ = nnls(A, b)
return G_i, tau_i, G_e
def _fit_prony_oscillation(
omega: np.ndarray,
G_prime: np.ndarray,
G_double_prime: np.ndarray,
n_modes: int,
) -> tuple[np.ndarray, np.ndarray, float]:
"""Fit Prony series to dynamic moduli G'(ω), G''(ω)."""
from scipy.optimize import nnls
if len(omega) < 2:
raise ValueError("Prony fitting requires at least 2 frequency points.")
# T-04: Ensure omega is sorted ascending for correct tau range computation.
# Must reorder G_prime and G_double_prime to stay aligned with sorted omega.
sort_idx = np.argsort(omega)
omega = omega[sort_idx]
G_prime = G_prime[sort_idx]
G_double_prime = G_double_prime[sort_idx]
# PC-001 (frequency-domain twin of _fit_prony_relaxation's t_positive
# filtering): drop non-positive omega before taking 1/omega. Otherwise an
# omega=0 entry makes 1/omega non-finite, poisoning tau_i with inf/NaN
# which propagates into an all-NaN NNLS kernel matrix below.
omega_positive = omega[omega > 0]
if len(omega_positive) == 0:
raise ValueError(
"Prony fitting requires at least one positive frequency value."
)
omega_min = float(np.min(omega_positive))
omega_max = float(np.max(omega_positive))
# Same n+2 / drop-edges approach as _fit_prony_relaxation: modes at exactly
# 1/omega_max and 1/omega_min are at the boundary of the informative range
# and pick up spurious weight that distorts G(t) predictions.
tau_all = np.logspace(
np.log10(1.0 / omega_max), np.log10(1.0 / omega_min), n_modes + 2
)
tau_i = tau_all[1:-1]
G_e = max(float(np.min(G_prime)), 0.0)
# Build kernel matrices
n = len(omega)
A = np.zeros((2 * n, n_modes))
wt2 = (omega[:, None] * tau_i[None, :]) ** 2
A[:n, :] = wt2 / (1.0 + wt2) # G' kernel
A[n:, :] = omega[:, None] * tau_i[None, :] / (1.0 + wt2) # G'' kernel
b = np.concatenate([G_prime - G_e, G_double_prime])
neg_mask = b < 0
if np.any(neg_mask):
n_neg = int(np.sum(neg_mask))
warnings.warn(
f"Prony fit: {n_neg} negative target values clipped to zero "
"(noisy data or overestimated G_e).",
stacklevel=2,
)
b = np.maximum(b, 0.0)
G_i, _ = nnls(A, b)
return G_i, tau_i, G_e