paramrf 0.37.2__tar.gz → 0.37.4__tar.gz

{paramrf-0.37.2 → paramrf-0.37.4}/CONTEXT.md

@@ -223,6 +223,12 @@ the prior precision $\Sigma_0^{-1}$ it gives the linearised posterior covariance
 $\Sigma_\text{post} = (F + \Sigma_0^{-1})^{-1}$, over parameter values in one space
 (declared by default). Fisher matrices of independent fits add.
 
+The **posterior covariance factor** $W$ is a square root of this covariance:
+$\Sigma_\text{post}=WW^\top$. It follows the same free-parameter order and space as
+the linearisation. For a stack $A$ of whitened Jacobians and a square root of the
+prior precision, $A=USV^\top$ gives $W=VS^{-1}$. Joint prediction uses $W$ to
+carry parameter uncertainty without first forming the Fisher matrix.
+
 ### Joint prediction
 
 The Gaussian over the free parameters and the discrepancy at new frequencies

{paramrf-0.37.2/paramrf.egg-info → paramrf-0.37.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: paramrf
-Version: 0.37.2
+Version: 0.37.4
 Summary: Parametric radio frequency modeling
 Author-email: Gary Allen <gvcallen@gmail.com>
 Project-URL: homepage, https://github.com/gvcallen/paramrf

{paramrf-0.37.2 → paramrf-0.37.4/paramrf.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: paramrf
-Version: 0.37.2
+Version: 0.37.4
 Summary: Parametric radio frequency modeling
 Author-email: Gary Allen <gvcallen@gmail.com>
 Project-URL: homepage, https://github.com/gvcallen/paramrf

{paramrf-0.37.2 → paramrf-0.37.4}/pmrf/models/components/lines/empirical.py

@@ -20,10 +20,14 @@ class PhysicalLine(AbstractImmittanceLine):
 
     The per-unit-length parameters are
     $$R = 2 z_n \alpha_c$$
-    $$L = \frac{z_n \sqrt{\varepsilon_r}}{c}$$
+    $$L = \frac{z_n \sqrt{\varepsilon_r}}{c} + \frac{R}{\omega}$$
     $$G = \frac{2 \alpha_d}{z_n}$$
     $$C = \frac{\sqrt{\varepsilon_r}}{z_n c}$$
 
+    The skin-effect term adds conductor reactance to the series impedance:
+    $$Z = R + j\omega z_n\sqrt{\varepsilon_r}/c + jR.$$
+    This contribution vanishes at DC.
+
     Example
     --------
     .. code-block:: python
@@ -93,7 +97,11 @@ class PhysicalLine(AbstractImmittanceLine):
         G = G_val * ones
         C = C_val * ones
         
-        return ImmittanceResult.from_rlgc(R, L, G, C, freq.w)
+        return ImmittanceResult(
+            Z=R + 1j * (freq.w * L + R),
+            Y=G + 1j * freq.w * C,
+            omega=freq.w,
+        )
     
 
 class DatasheetLine(AbstractImmittanceLine):
@@ -108,10 +116,14 @@ class DatasheetLine(AbstractImmittanceLine):
 
     and the per-unit-length parameters are
     $$R = 2 z_n \alpha_c$$
-    $$L = \frac{z_n}{v_f c}$$
+    $$L = \frac{z_n}{v_f c} + \frac{R}{\omega}$$
     $$G = \frac{2 \alpha_d}{z_n}$$
     $$C = \frac{1}{z_n v_f c}$$
 
+    The skin-effect term adds conductor reactance to the series impedance:
+    $$Z = R + j\omega z_n/(v_f c) + jR.$$
+    This contribution vanishes at DC.
+
     Example
     --------
     .. code-block:: python
@@ -186,5 +198,9 @@ class DatasheetLine(AbstractImmittanceLine):
         L = (zn / (vf * c)) * jnp.ones_like(w)
         C = (1.0 / (zn * vf * c)) * jnp.ones_like(w)
         
-        return ImmittanceResult.from_rlgc(R, L, G, C, w)
+        return ImmittanceResult(
+            Z=R + 1j * (w * L + R),
+            Y=G + 1j * w * C,
+            omega=w,
+        )
 

{paramrf-0.37.2 → paramrf-0.37.4}/pmrf/objectives/evaluators.py

@@ -22,7 +22,7 @@ from pmrf.objectives.losses import HingeLoss, RMSELoss
 from pmrf.stats.likelihoods import GaussianLikelihood
 from pmrf.stats.covariance_kernels import cross_gram, gram
 from pmrf.stats.discrepancy_models import GaussianProcess, _add_noise
-from pmrf.stats.linearization import Linearization, _linearize, posterior_covariance
+from pmrf.stats.linearization import Linearization, _linearize, posterior_covariance_factor
 from pmrf.modules.base import Module
 from pmrf.parameters import Space, values as _values, update
 from pmrf.utils import derivative, field, unwrap, unwrap_self
@@ -593,6 +593,8 @@ class MarginalLogLikelihood(AbstractEvaluator):
         frequency: Frequency,
         new_frequency: Frequency,
         covariance: jnp.ndarray | None = None,
+        *,
+        covariance_factor: jnp.ndarray | None = None,
         **kwargs,
     ) -> tuple[tuple[str, ...], dist.AbstractDistribution]:
         r"""Predict the free parameters and discrepancy together as a joint prior.
@@ -600,13 +602,14 @@ class MarginalLogLikelihood(AbstractEvaluator):
         **Mathematical Formulation**
 
         At the fitted model, let $A = K_{*h}\Sigma_D^{-1}$ and let
-        $\Sigma_\text{post}$ be the linearised parameter covariance. Then
+        $W W^\top = \Sigma_\text{post}$ factor the linearised parameter
+        covariance. Then
 
         $$\hat\delta = A r, \qquad
         \Sigma_\delta = K_{**} - A K_{*h}^\top
-            + A J\Sigma_\text{post}J^\top A^\top$$
+            + (AJW)(AJW)^\top$$
 
-        $$\Sigma_{\delta\theta} = -A J\Sigma_\text{post}$$
+        $$\Sigma_{\delta\theta} = -(AJW)W^\top$$
 
         The parameter term couples event blocks, giving one dense Gaussian over
         $[\theta;\delta]$. Kernel hyperparameters are held fixed, as in
@@ -632,6 +635,12 @@ class MarginalLogLikelihood(AbstractEvaluator):
             order. If omitted, computed from this fit and the model's prior with
             :func:`pmrf.stats.linearization.posterior_covariance`. Supply a covariance
             from combined linearisations when several fits inform the parameters.
+        covariance_factor : jnp.ndarray or None, default=None
+            A factor $W$ of the parameter covariance in declared space and
+            linearisation name order, with $W W^\top$ equal to the covariance.
+            Obtain one from :func:`pmrf.stats.linearization.posterior_covariance_factor`
+            when several fits inform the parameters. Cannot be combined with
+            covariance.
         **kwargs
             Passed to the predictor.
 
@@ -653,18 +662,29 @@ class MarginalLogLikelihood(AbstractEvaluator):
             If the discrepancy is not a Gaussian process or the likelihood is not
             Gaussian.
         ValueError
-            If orthogonal discrepancy is enabled or ``covariance`` has the wrong
-            shape for the free parameters.
+            If orthogonal discrepancy is enabled, an input has the wrong shape,
+            or both covariance arguments are supplied.
         """
         if not isinstance(self.discrepancy, GaussianProcess):
             raise TypeError("Joint prediction requires a `GaussianProcess` discrepancy.")
         linearization = self.linearize(model, frequency, **kwargs)
         p = linearization.J.shape[-1]
-        if covariance is None:
-            covariance = posterior_covariance([linearization], model)
-        covariance = jnp.asarray(covariance)
-        if covariance.shape != (p, p):
-            raise ValueError(f"Parameter covariance must have shape {(p, p)}.")
+        if covariance is not None and covariance_factor is not None:
+            raise ValueError("Supply either covariance or covariance_factor, not both.")
+        if covariance_factor is None:
+            if covariance is None:
+                covariance_factor = posterior_covariance_factor([linearization], model)
+            else:
+                covariance = jnp.asarray(covariance)
+                if covariance.shape != (p, p):
+                    raise ValueError(f"Parameter covariance must have shape {(p, p)}.")
+                covariance_factor = jnp.linalg.cholesky(covariance)
+        covariance_factor = jnp.asarray(covariance_factor)
+        if covariance_factor.shape != (p, p):
+            raise ValueError(f"Parameter covariance factor must have shape {(p, p)}.")
+        # QR gives a triangular scale without first forming an ill-conditioned covariance.
+        _, R = jnp.linalg.qr(covariance_factor.T)
+        parameter_scale = R.T
 
         batch_shape = linearization.residual.shape[:-1]
         n = linearization.residual.shape[-1]
@@ -682,17 +702,19 @@ class MarginalLogLikelihood(AbstractEvaluator):
         conditional = conditional - jnp.swapaxes(V, -1, -2) @ V
         conditional = conditional.reshape((-1, n_new, n_new))
         blocks = conditional.shape[0]
-        conditional = jnp.einsum('bij,bc->bicj', conditional, jnp.eye(blocks)).reshape(
-            (blocks * n_new, blocks * n_new)
-        )
+        conditional_scale = jnp.linalg.cholesky(conditional)
+        conditional_scale = jnp.einsum(
+            'bij,bc->bicj', conditional_scale, jnp.eye(blocks)
+        ).reshape((blocks * n_new, blocks * n_new))
         AJ = (A @ linearization.J).reshape((-1, p))
-        cross = -AJ @ covariance
-        sigma_delta = conditional + AJ @ covariance @ AJ.T
-        joint_covariance = jnp.block([[covariance, cross.T], [cross, sigma_delta]])
+        joint_scale = jnp.block([
+            [parameter_scale, jnp.zeros((p, blocks * n_new))],
+            [-AJ @ parameter_scale, conditional_scale],
+        ])
         named_values = _values(model, free_only=True, space='declared')
         theta = jnp.concatenate([jnp.ravel(named_values[name]) for name in linearization.names])
         mean = jnp.concatenate([theta, mean_delta])
-        return linearization.names, dist.MultivariateNormalFullCovariance(mean, joint_covariance)
+        return linearization.names, dist.MultivariateNormalTri(mean, joint_scale)
 
     @unwrap_self
     def linearize(

{paramrf-0.37.2 → paramrf-0.37.4}/pmrf/stats/__init__.py

@@ -28,9 +28,11 @@ if TYPE_CHECKING:
         AutoCrossKernel as AutoCrossKernel,
         ConstantKernel as ConstantKernel,
         CosineKernel as CosineKernel,
+        LogFrequencyKernel as LogFrequencyKernel,
         Matern32Kernel as Matern32Kernel,
         Matern52Kernel as Matern52Kernel,
         PeriodicKernel as PeriodicKernel,
+        PowerEnvelopeKernel as PowerEnvelopeKernel,
         ProductKernel as ProductKernel,
         RBFKernel as RBFKernel,
         SharedIndependentKernel as SharedIndependentKernel,
@@ -65,6 +67,7 @@ if TYPE_CHECKING:
     from .linearization import (
         Linearization as Linearization,
         posterior_covariance as posterior_covariance,
+        posterior_covariance_factor as posterior_covariance_factor,
     )
     from .noise_models import (
         AbstractNoiseModel as AbstractNoiseModel,
@@ -94,11 +97,13 @@ _EXPORTS = {
     'Joint': 'distributions',
     'Leafwise': 'bijectors',
     'Linearization': 'linearization',
+    'LogFrequencyKernel': 'covariance_kernels',
     'LogNormal': 'distributions',
     'Matern32Kernel': 'covariance_kernels',
     'Matern52Kernel': 'covariance_kernels',
     'Normal': 'distributions',
     'PeriodicKernel': 'covariance_kernels',
+    'PowerEnvelopeKernel': 'covariance_kernels',
     'Permute': 'bijectors',
     'ProductKernel': 'covariance_kernels',
     'R2ToComplex': 'bijectors',
@@ -122,6 +127,7 @@ _EXPORTS = {
     'cross_gram': 'covariance_kernels',
     'gram': 'covariance_kernels',
     'posterior_covariance': 'linearization',
+    'posterior_covariance_factor': 'linearization',
     'truncate': 'distributions',
 }
 _MODULES = {'bijectors', 'discrepancy_models', 'linearization', 'likelihoods', 'noise_models', 'distributions', 'covariance_kernels'}

{paramrf-0.37.2 → paramrf-0.37.4}/pmrf/stats/covariance_kernels.py

@@ -304,6 +304,74 @@ class CosineKernel(AbstractCovarianceKernel):
         return jnp.cos(2 * jnp.pi * jnp.sum(x1 - x2) / self.period)
 
 
+class LogFrequencyKernel(AbstractCovarianceKernel):
+    r"""
+    Evaluate a kernel on log frequency.
+
+    $$k(x_1, x_2) = k_\text{inner}(\ln x_1, \ln x_2)$$
+
+    A stationary inner kernel depends on the ratio $x_1/x_2$. For a
+    :class:`CosineKernel` with period $p$, one cycle spans a factor of $e^p$
+    in frequency. The warp needs no frequency unit: rescaling all inputs adds
+    the same constant to their logarithms, leaving stationary inner kernels
+    unchanged. Inputs must be strictly positive.
+
+    Parameters
+    ----------
+    kernel : AbstractCovarianceKernel
+        The kernel evaluated on log frequency. Its batching is preserved.
+    """
+    #: The inner kernel.
+    kernel: AbstractCovarianceKernel
+
+    def __call__(self, x1, x2, key=None):
+        x1, x2 = jnp.asarray(x1), jnp.asarray(x2)
+        x1 = eqx.error_if(x1, jnp.any(x1 <= 0),
+                          "LogFrequencyKernel inputs must be strictly positive")
+        x2 = eqx.error_if(x2, jnp.any(x2 <= 0),
+                          "LogFrequencyKernel inputs must be strictly positive")
+        return self.kernel(jnp.log(x1), jnp.log(x2), key=key)
+
+
+class PowerEnvelopeKernel(AbstractCovarianceKernel):
+    r"""
+    Power-law amplitude envelope for a one-dimensional frequency input.
+
+    $$k(x_1, x_2) = s(x_1)s(x_2), \qquad
+    s(x) = \left(\frac{x}{x_\text{ref}}\right)^{-\alpha}$$
+
+    This rank-one kernel is singular on its own. Multiply it by another
+    kernel to make that kernel's prior standard deviation scale as $s(x)$.
+    Inputs must be strictly positive and contain exactly one feature.
+    ``reference`` must use the same units as the GP input (``frequency.f_scaled``
+    in the fitting path). A unit mismatch silently rescales the prior variance
+    by $(\text{unit ratio})^{-2\alpha}$.
+
+    Parameters
+    ----------
+    reference : ArrayLike
+        Input at which the amplitude is one, in the same units as the GP input.
+    exponent : ArrayLike
+        Decay exponent. A shape ``(D,)`` gives ``D`` batched kernels.
+    """
+    #: The reference input.
+    reference: ArrayLike
+
+    #: The decay exponent.
+    exponent: ArrayLike
+
+    def _amplitude(self, x):
+        x = jnp.asarray(x)
+        if x.size != 1:
+            raise ValueError("PowerEnvelopeKernel requires one-dimensional input points")
+        x = eqx.error_if(x, jnp.any(x <= 0),
+                         "PowerEnvelopeKernel inputs must be strictly positive")
+        return (x.reshape(()) / self.reference) ** (-jnp.asarray(self.exponent))
+
+    def __call__(self, x1, x2, key=None):
+        return self._amplitude(x1) * self._amplitude(x2)
+
+
 class WhiteNoiseKernel(AbstractCovarianceKernel):
     """
     Kernel representing independent Gaussian noise.

{paramrf-0.37.2 → paramrf-0.37.4}/pmrf/stats/linearization.py

@@ -122,34 +122,45 @@ def _linearize(
     residual = residual_fn(model)
     J = -jax.jacfwd(lambda x: residual_fn(rebuild(x)))(x0)
 
+    W = _whitened_jacobian(J, chol)
+    F = jnp.einsum('...nps,...nqs->pq', W, W)
+    return Linearization(J, residual, chol, F, names, shapes, space)
+
+
+def _whitened_jacobian(J: Array, chol: Array) -> Array:
+    """Return $L^{-1}J$ grouped by shared covariance factors."""
+
     # Event blocks sharing a factor are solved together, with the parameter axis
     # carried as one more batch axis the factor is shared along.
     n, p = J.shape[-2], J.shape[-1]
     matrix_shape, R, _ = _group_by_matrix(chol.shape[:-2], jnp.moveaxis(J, -1, 0))
     L = jnp.broadcast_to(chol, chol.shape[:-2] + (n, n)).reshape(matrix_shape + (n, n))
     W = jsp.linalg.solve_triangular(L, R, lower=True)
-    W = W.reshape(W.shape[:-1] + (p, -1))
-    F = jnp.einsum('...nps,...nqs->pq', W, W)
-    return Linearization(J, residual, chol, F, names, shapes, space)
+    return W.reshape(W.shape[:-1] + (p, -1))
 
 
-def posterior_covariance(
+def posterior_covariance_factor(
     linearizations: Sequence[Linearization],
     model: PyTree,
     space: Space = 'declared',
 ) -> Array:
-    r"""The linearised posterior covariance of a model's free parameters.
+    r"""A square root of the linearised posterior covariance of free parameters.
 
     **Mathematical Formulation**
 
-    The Fisher matrices of independent fits add, and the prior enters through its
-    precision at `model`:
+    Stack the whitened Jacobians of independent fits with a square root of the
+    prior precision, then use the singular values of that stack:
 
-    $$\Sigma_\text{post} = \left(\sum_k F_k + \Sigma_0^{-1}\right)^{-1}, \qquad
-    \Sigma_0^{-1} = -\nabla^2 \log p(\theta)$$
+    $$A = \begin{bmatrix}L_1^{-1}J_1 \\ \cdots \\ L_k^{-1}J_k \\
+    \Sigma_0^{-1/2}\end{bmatrix}, \qquad
+    A = U S V^\top, \qquad W = V S^{-1}, \qquad
+    \Sigma_\text{post} = W W^\top,$$
 
-    with the Hessian of :func:`pmrf.log_prior` taken in `space`. A parameter without a
-    prior, or with a uniform prior inside its bounds, adds zero to $\Sigma_0^{-1}$.
+    where $\Sigma_0^{-1} = -\nabla^2\log p(\theta)$ in `space`. Its symmetric
+    eigendecomposition supplies the square root; negative eigenvalues are clipped to
+    zero. This local approximation is only meaningful where the prior is log-concave.
+    A parameter without a prior, or with a uniform prior inside its bounds, adds zero
+    prior precision.
 
     Parameters
     ----------
@@ -164,8 +175,8 @@ def posterior_covariance(
     Returns
     -------
     jax.Array
-        The covariance, with shape ``(p, p)``, over the free parameters in the
-        linearisations' :attr:`~Linearization.names` order.
+        A factor ``W`` with shape ``(p, p)`` such that ``W @ W.T`` is the
+        covariance over free parameters in :attr:`~Linearization.names` order.
 
     Raises
     ------
@@ -185,14 +196,42 @@ def posterior_covariance(
                 "A linearisation is over different free parameters from the model: "
                 f"{linearization.names} rather than {names}."
             )
-    fisher = sum(linearization.F for linearization in linearizations)
     prior_precision = -jax.hessian(lambda x: log_prior(rebuild(x), space=space))(x0)
-    precision = fisher + prior_precision
-    chol = jnp.linalg.cholesky(0.5 * (precision + precision.T))
-    return jsp.linalg.cho_solve((chol, True), jnp.eye(precision.shape[-1], dtype=precision.dtype))
+    eigenvalues, eigenvectors = jnp.linalg.eigh(0.5 * (prior_precision + prior_precision.T))
+    prior_root = jnp.sqrt(jnp.maximum(eigenvalues, 0.0))[:, None] * eigenvectors.T
+    p = x0.size
+    factors = [
+        jnp.moveaxis(_whitened_jacobian(lin.J, lin.chol), -2, -1).reshape(-1, p)
+        for lin in linearizations
+    ]
+    A = jnp.concatenate([*factors, prior_root], axis=0)
+    _, singular_values, vt = jnp.linalg.svd(A, full_matrices=False)
+    threshold = singular_values[0] * max(A.shape) * jnp.finfo(A.dtype).eps
+    singular_values = eqx.error_if(
+        singular_values, singular_values[-1] <= threshold,
+        "A parameter direction has neither data nor prior information.",
+    )
+    V = vt.T
+    return V / singular_values
+
+
+def posterior_covariance(
+    linearizations: Sequence[Linearization],
+    model: PyTree,
+    space: Space = 'declared',
+) -> Array:
+    r"""The linearised posterior covariance of a model's free parameters.
+
+    See :func:`posterior_covariance_factor` for the factorisation and its
+    assumptions. The result has shape ``(p, p)`` in the linearisations'
+    :attr:`~Linearization.names` order.
+    """
+    W = posterior_covariance_factor(linearizations, model, space)
+    return W @ W.T
 
 
 __all__ = [
     'Linearization',
     'posterior_covariance',
+    'posterior_covariance_factor',
 ]

{paramrf-0.37.2 → paramrf-0.37.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "paramrf"
-version = "0.37.2"
+version = "0.37.4"
 description = "Parametric radio frequency modeling"
 readme = "README.rst"
 authors = [

{paramrf-0.37.2 → paramrf-0.37.4}/tests/test_covariance_kernels.py

@@ -15,6 +15,8 @@ from pmrf.stats.covariance_kernels import (
     Matern52Kernel,
     ConstantKernel,
     CosineKernel,
+    LogFrequencyKernel,
+    PowerEnvelopeKernel,
     AutoCrossKernel,
     SharedIndependentKernel,
 )
@@ -38,6 +40,12 @@ def x_periods():
     return jnp.linspace(0.0, 50.0, 200)
 
 
+@pytest.fixture
+def x_log():
+    assert jax.config.jax_enable_x64
+    return jnp.geomspace(10e6, 500e6, 300)
+
+
 def _damped_cosine():
     return CosineKernel(period=7.3) * Matern52Kernel(lengthscale=20.0) * 2.0
 
@@ -210,6 +218,130 @@ def test_cosine_kernel_under_jit_and_vmap(x_periods):
     assert jnp.allclose(vmapped, CosineKernel(period=periods).gram(x_periods))
 
 
+def test_log_frequency_kernel_has_ratio_period():
+    """One cycle in log frequency spans a fixed frequency ratio."""
+    period = 1.7
+    start = jnp.array([10e6])
+    kernel = LogFrequencyKernel(CosineKernel(period))
+    assert jnp.allclose(kernel(start, start * jnp.exp(period)), 1.0, atol=1e-12)
+    assert jnp.allclose(kernel(start, start * jnp.exp(period / 2)), -1.0, atol=1e-12)
+
+
+def test_log_frequency_kernel_unit_invariance_and_psd(x_log):
+    """Changing frequency units leaves the damped covariance unchanged and PSD."""
+    kernel = LogFrequencyKernel(CosineKernel(1.7) * Matern52Kernel(1.5))
+    covariance = kernel.gram(x_log)
+    assert jnp.max(jnp.abs(covariance - kernel.gram(x_log / 1e6))) < 1e-12
+    eigenvalues = jnp.linalg.eigvalsh(covariance)
+    assert eigenvalues.min() >= -1e-8 * eigenvalues.max()
+
+
+def test_log_frequency_kernel_preserves_inner_batching(x_log):
+    """A vector of inner periods gives one Gram matrix per period."""
+    periods = jnp.array([1.0, 2.0])
+    covariance = LogFrequencyKernel(CosineKernel(periods)).gram(x_log)
+    assert covariance.shape == (2, 300, 300)
+    for index, period in enumerate(periods):
+        expected = LogFrequencyKernel(CosineKernel(period)).gram(x_log)
+        assert jnp.allclose(covariance[index], expected)
+
+
+@pytest.mark.parametrize("invalid", [0.0, -1.0])
+@pytest.mark.parametrize("side", [0, 1])
+def test_log_frequency_kernel_rejects_nonpositive_inputs(invalid, side):
+    """Both arguments require positive frequencies, eagerly and under jit."""
+    kernel = LogFrequencyKernel(CosineKernel(1.7))
+    inputs = [jnp.array([10e6]), jnp.array([10e6])]
+    inputs[side] = jnp.array([invalid])
+    with pytest.raises(Exception, match="strictly positive"):
+        kernel(*inputs)
+    with pytest.raises(Exception, match="strictly positive"):
+        jax.jit(kernel)(*inputs)
+
+
+def test_log_frequency_kernel_is_exported():
+    from pmrf.stats import LogFrequencyKernel as PublicLogFrequencyKernel
+    assert PublicLogFrequencyKernel is LogFrequencyKernel
+
+
+def test_power_envelope_kernel_value():
+    """The covariance is the outer product of two power-law amplitudes."""
+    kernel = PowerEnvelopeKernel(10e6, 0.5)
+    assert jnp.allclose(kernel(40e6, 90e6), 1 / 6, rtol=1e-14)
+
+
+def _damped_log_periodic_envelope(exponent=0.65, period=1.7, lengthscale=1.5):
+    return (PowerEnvelopeKernel(10e6, exponent)
+            * LogFrequencyKernel(CosineKernel(period) * Matern52Kernel(lengthscale))
+            * 2.5e-3 ** 2)
+
+
+def test_power_envelope_kernel_diagonal_and_psd(x_log):
+    """The envelope sets the prior variance and preserves positive semidefiniteness."""
+    kernel = _damped_log_periodic_envelope()
+    for index in [0, 150, 299]:
+        expected = 2.5e-3 ** 2 * (x_log[index] / 10e6) ** -1.3
+        assert jnp.allclose(kernel(x_log[index], x_log[index]), expected, rtol=1e-10)
+    eigenvalues = jnp.linalg.eigvalsh(kernel.gram(x_log))
+    assert eigenvalues.min() >= -1e-8 * eigenvalues.max()
+
+
+def test_power_envelope_kernel_batches_exponents(x_log):
+    """A vector of exponents gives the corresponding scalar Gram matrices."""
+    exponents = jnp.array([0.5, 1.0])
+    covariance = PowerEnvelopeKernel(10e6, exponents).gram(x_log)
+    assert covariance.shape == (2, 300, 300)
+    for index, exponent in enumerate(exponents):
+        expected = PowerEnvelopeKernel(10e6, exponent).gram(x_log)
+        assert jnp.allclose(covariance[index], expected)
+
+
+def test_power_envelope_kernel_composition_gradients_are_finite(x_log):
+    """Envelope, period, lengthscale, and coincident input gradients stay finite."""
+    def covariance_sum(parameters):
+        return _damped_log_periodic_envelope(*parameters).gram(x_log).sum()
+
+    gradients = jax.grad(covariance_sum)(jnp.array([0.65, 1.7, 1.5]))
+    assert jnp.all(jnp.isfinite(gradients))
+
+    kernel = _damped_log_periodic_envelope()
+    input_gradient = jax.grad(lambda x1: kernel(x1, jnp.array([10e6])))(jnp.array([10e6]))
+    assert jnp.all(jnp.isfinite(input_gradient))
+
+
+def test_power_envelope_kernel_composition_under_jit(x_log):
+    """Compilation preserves the composed covariance."""
+    kernel = _damped_log_periodic_envelope()
+    assert jnp.allclose(jax.jit(lambda x: kernel.gram(x))(x_log), kernel.gram(x_log))
+
+
+@pytest.mark.parametrize("invalid", [0.0, -1.0])
+@pytest.mark.parametrize("side", [0, 1])
+def test_power_envelope_kernel_rejects_nonpositive_inputs(invalid, side):
+    """Both arguments require positive frequencies, eagerly and under jit."""
+    kernel = PowerEnvelopeKernel(10e6, 0.65)
+    inputs = [jnp.array([10e6]), jnp.array([10e6])]
+    inputs[side] = jnp.array([invalid])
+    with pytest.raises(Exception, match="strictly positive"):
+        kernel(*inputs)
+    with pytest.raises(Exception, match="strictly positive"):
+        jax.jit(kernel)(*inputs)
+
+
+def test_power_envelope_kernel_rejects_multidimensional_inputs():
+    """A point with multiple features is rejected while tracing."""
+    kernel = PowerEnvelopeKernel(10e6, 0.65)
+    with pytest.raises(ValueError, match="one-dimensional"):
+        kernel(jnp.array([10e6, 20e6]), jnp.array([10e6]))
+    with pytest.raises(ValueError, match="one-dimensional"):
+        jax.jit(kernel)(jnp.array([10e6, 20e6]), jnp.array([10e6]))
+
+
+def test_power_envelope_kernel_is_exported():
+    from pmrf.stats import PowerEnvelopeKernel as PublicPowerEnvelopeKernel
+    assert PublicPowerEnvelopeKernel is PowerEnvelopeKernel
+
+
 _AUTO = RBFKernel(lengthscale=0.3)
 _CROSS = Matern32Kernel(lengthscale=2.0) * 0.5