paramrf 0.36.2__tar.gz → 0.36.4__tar.gz

{paramrf-0.36.2 → paramrf-0.36.4}/CHANGELOG.md

@@ -117,6 +117,24 @@ longer recompiles a model's RF methods.
 
 ## Unreleased
 
+- **Breaking (#245):** a frequentist fit with neither `loss` nor `likelihood`
+  now minimises the mean squared error (`MSELoss`) instead of the RMSE. MSE is
+  smooth in the residuals, pools every output's residuals equally rather than
+  averaging per-output RMSEs, and matches the fixed-noise Gaussian likelihood
+  of the Bayesian path. Fits relying on the default stop at slightly different
+  points and report loss values on a squared scale; pass `loss=RMSELoss()` for
+  the old objective. SciPy's default `gtol`, and its `ftol` (absolute below an
+  objective of 1), can now stop a fit whose MSE is already tiny before it
+  moves: tighten them through `ScipyMinimize(options=...)`.
+- **Breaking (#247):** `SharedIndependentKernel` appends its shared axes as
+  size-1 trailing axes instead of materializing them, which fixes its routing
+  over a batched base kernel. Wrapping an `AutoCrossKernel` previously
+  prepended the shared axes, so half the real/imaginary entries got the wrong
+  kernel (S11 Im got cross, S21 Re got auto). Log-likelihoods of every model
+  using this composition change. `output_shape` is replaced by
+  `num_shared_axes` (default 1), since only its length was ever meaningful;
+  `gram` now returns e.g. `(2, 2, 1, N, N)` rather than `(2, 2, 2, N, N)`, and
+  the orthogonal discrepancy path factorizes the smaller batch.
 - Complete Tesche coaxial conductor physics across the low-frequency regime,
   and pass evaluated material properties to pure coaxial formulations.
 - Correct microstrip results by enabling Kirschning--Jansen modal dispersion by
@@ -124,3 +142,5 @@ longer recompiles a model's RF methods.
   pipeline explicitly.
 - Add the Hammerstad--Jensen microstrip formulation and selectable complex
   (ADS-like) and real (QUCS-like) permittivity conventions.
+- Add `CosineKernel` to `pmrf.covariance_kernels` for zero-mean oscillatory
+  (negative) correlation. Multiply it by a decaying kernel for a damped cosine.

{paramrf-0.36.2/paramrf.egg-info → paramrf-0.36.4}/PKG-INFO

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

{paramrf-0.36.2 → paramrf-0.36.4/paramrf.egg-info}/PKG-INFO

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

{paramrf-0.36.2 → paramrf-0.36.4}/paramrf.egg-info/SOURCES.txt

@@ -180,6 +180,7 @@ pmrf/utils/tree.py
 pmrf/utils/type.py
 pmrf/viz/__init__.py
 pmrf/viz/plots.py
+scripts/benchmark_gp_marginal_likelihood.py
 scripts/install-test-deps.sh
 tests/__init__.py
 tests/_jit.py
@@ -244,4 +245,5 @@ tests/test_models/test_sections.py
 tests/test_models/test_transformed.py
 tests/test_models/test_transformers.py
 tests/test_utils/test_compress.py
+tests/test_utils/test_debug.py
 tests/test_utils/test_tree.py

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/covariance_kernels.py

@@ -7,10 +7,11 @@ for more details.
 from abc import abstractmethod
 from collections.abc import Callable
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 
-from pmrf.utils import field
+from pmrf.utils import field, unwrap
 from pmrf.types import ArrayLike
 from pmrf.modules.base import Module
 
@@ -51,6 +52,12 @@ def gram(
         The Gram matrix, of shape ``(*batch_shape, N, N)``, where
         ``batch_shape`` is the shape returned by the kernel for a single pair.
     """
+    # Unwrap the hyperparameters to their values once, behind a barrier. Otherwise
+    # XLA can fuse each one's raw-to-value bijector chain into the elementwise loop
+    # that builds the matrix and recompute it per entry.
+    leaves, structure = eqx.partition(unwrap(kernel), eqx.is_array)
+    kernel = eqx.combine(jax.lax.optimization_barrier(leaves), structure)
+
     x = jnp.asarray(x)
     if x.ndim == 1:
         x_feat = x[:, None]
@@ -231,6 +238,40 @@ class PeriodicKernel(AbstractCovarianceKernel):
         return jnp.exp(-2.0 * (sin_term / self.lengthscale)**2)
 
 
+class CosineKernel(AbstractCovarianceKernel):
+    r"""
+    Cosine kernel.
+
+    Models zero-mean oscillatory correlation with a single period, which can be
+    negative. Unlike :class:`PeriodicKernel`, the correlation goes to $-1$ at
+    half a period.
+
+    $$k(x_1, x_2) = \cos\left(\frac{2\pi\, \mathbf{1}^\top (x_1 - x_2)}{p}\right)$$
+
+    where $p$ is the ``period``.
+
+    For $d$-dimensional inputs the oscillation runs along the all-ones direction:
+    the differences are summed, not combined into a Euclidean distance, since
+    $\cos(2\pi \lVert x_1 - x_2 \rVert / p)$ is not positive semi-definite
+    for $d > 1$.
+
+    Its Gram matrix has rank two, so it is singular on its own. Multiply by a
+    decaying kernel (e.g. :class:`Matern52Kernel`) for a damped oscillation,
+    and by a float for its variance.
+
+    Parameters
+    ----------
+    period : ArrayLike
+        The period of the oscillation. A period of shape ``(D,)`` gives ``D``
+        batched kernels, as for :class:`PeriodicKernel`.
+    """
+    #: The period.
+    period: ArrayLike
+
+    def __call__(self, x1, x2, key=None):
+        return jnp.cos(2 * jnp.pi * jnp.sum(x1 - x2) / self.period)
+
+
 class WhiteNoiseKernel(AbstractCovarianceKernel):
     """
     Kernel representing independent Gaussian noise.
@@ -353,32 +394,32 @@ class AutoCrossKernel(AbstractCovarianceKernel):
 
 class SharedIndependentKernel(AbstractCovarianceKernel):
     """
-    Evaluates a base kernel and broadcasts its output to represent 
-    multiple independent dimensions (e.g., real and imaginary parts) 
-    withed share hyperparameters.
+    Shares a base kernel across trailing independent batch axes (e.g. the real
+    and imaginary parts), so that every element on those axes uses the same
+    hyperparameters.
+
+    The shared axes are appended to the base kernel's output as size-1 axes
+    rather than materialized. A base output of shape ``batch_shape`` becomes
+    ``(*batch_shape, 1, ..., 1)``, which broadcasts to any size along the
+    shared axes, so a batched base kernel (e.g. :class:`AutoCrossKernel`)
+    keeps routing along its own leading axes.
 
     Parameters
     ----------
-    base_kernel : CovarianceKernel
+    base_kernel : AbstractCovarianceKernel
         The underlying kernel whose parameters are shared.
-    output_shape : tuple
-        The shape of the independent outputs to broadcast to.
+    num_shared_axes : int, default=1
+        The number of trailing batch axes the base kernel is shared across.
     """
     #: The base kernel.
     base_kernel: AbstractCovarianceKernel
 
-    #: The output shape.
-    output_shape: tuple = field(static=True)
+    #: The number of trailing shared axes.
+    num_shared_axes: int = field(default=1, static=True)
 
     def __call__(self, x1, x2, key=None):
-        # Evaluate the underlying shared kernel
-        val = self.base_kernel(x1, x2, key=key)
-        
-        # Broadcast the evaluation to the target shape.
-        # By appending the output_shape to the end of val.shape, 
-        # this safely handles both scalar evaluations and already-batched evaluations.
-        target_shape = val.shape + self.output_shape
-        return jnp.broadcast_to(val, target_shape)
+        val = jnp.asarray(self.base_kernel(x1, x2, key=key))
+        return val.reshape(val.shape + (1,) * self.num_shared_axes)
     
 
 class ZeroKernel(AbstractCovarianceKernel):

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/discrepancy_models.py

@@ -6,7 +6,10 @@ Useful for modeling discrepancy of RF models during fitting.
 from collections.abc import Callable
 from abc import abstractmethod
 
+import math
+
 import equinox as eqx
+import jax
 import jax.scipy as jsp
 import jax.numpy as jnp
 import parax.distributions as dist
@@ -49,6 +52,44 @@ class AbstractDiscrepancyModel(Module):
         """        
         raise NotImplementedError
 
+
+@jax.custom_vjp
+def _gaussian_log_prob(M: jnp.ndarray, R: jnp.ndarray) -> jnp.ndarray:
+    """Sum the zero-mean Gaussian log densities of the columns of ``R`` under ``M``.
+
+    ``M`` has shape ``(*batch, N, N)`` and ``R`` has shape ``(*batch, N, k)``: the
+    ``k`` residuals in ``R[b]`` share the covariance ``M[b]`` and are solved together.
+    """
+    return _gaussian_log_prob_fwd(M, R)[0]
+
+
+def _gaussian_log_prob_fwd(M, R):
+    L = jnp.linalg.cholesky(M)
+    A = jsp.linalg.cho_solve((L, True), R)
+    n, k = R.shape[-2:]
+    num_matrices = math.prod(M.shape[:-2])
+    logdet = 2 * jnp.sum(jnp.log(jnp.diagonal(L, axis1=-2, axis2=-1)))
+    value = -0.5 * (
+        jnp.sum(R * A) + k * logdet + k * num_matrices * n * jnp.log(2 * jnp.pi)
+    )
+    return value, (L, A)
+
+
+def _gaussian_log_prob_bwd(residuals, g):
+    # With alpha = M^-1 r, d/dM = (sum alpha alpha^T - k M^-1) / 2 and d/dr = -alpha.
+    # JAX has no potri, so M^-1 = L^-T L^-1 from one triangular solve and a matmul.
+    L, A = residuals
+    n, k = A.shape[-2:]
+    identity = jnp.broadcast_to(jnp.eye(n, dtype=L.dtype), L.shape)
+    L_inv = jsp.linalg.solve_triangular(L, identity, lower=True)
+    M_inv = jnp.swapaxes(L_inv, -1, -2) @ L_inv
+    dM = 0.5 * g * (A @ jnp.swapaxes(A, -1, -2) - k * M_inv)
+    return dM, -g * A
+
+
+_gaussian_log_prob.defvjp(_gaussian_log_prob_fwd, _gaussian_log_prob_bwd)
+
+
 class GaussianProcess(AbstractDiscrepancyModel):
     """
     Gaussian process discrepancy model with a covariance kernel.
@@ -82,6 +123,65 @@ class GaussianProcess(AbstractDiscrepancyModel):
     #: The added jitter.
     jitter: float = field(default=1e-10, static=True)
 
+    def log_prob(
+        self,
+        y_event: jnp.ndarray,
+        observed: jnp.ndarray,
+        x: jnp.ndarray,
+        noise_variance: jnp.ndarray,
+    ) -> jnp.ndarray:
+        r"""Evaluate the summed log density of ``observed`` under the GP plus Gaussian noise.
+
+        Each batch entry of ``observed`` is distributed as
+        $\mathcal{N}(y, K + \sigma^2 I)$. Equal to the log probability of the
+        distribution built by :meth:`__call__` and
+        :class:`pmrf.likelihoods.GaussianLikelihood`, summed over the batch.
+
+        ``M = K + sigma^2 I`` is formed and factorized at the broadcast shape of the
+        kernel's Gram batch and ``noise_variance``, rather than the full batch shape,
+        and the residuals sharing each ``M`` are solved together.
+
+        Parameters
+        ----------
+        y_event : jnp.ndarray
+            The model prediction in event space, with shape ``(*batch_shape, N)``.
+        observed : jnp.ndarray
+            The observation in event space, broadcastable to ``y_event``.
+        x : jnp.ndarray
+            The frequency points, with shape ``(N,)``.
+        noise_variance : jnp.ndarray
+            The noise variance, constant along the event axis and broadcastable to
+            ``batch_shape``.
+
+        Returns
+        -------
+        jnp.ndarray
+            The scalar log density, summed over the batch.
+        """
+        # Materialize K before forming M. Otherwise XLA on CPU fuses the Gram
+        # construction into the transpose to LAPACK's column-major layout, and that
+        # strided build nearly doubles the value's cost at N = 1000.
+        K = jax.lax.optimization_barrier(gram(self.kernel, x, jitter=self.jitter))
+        variance = jnp.asarray(noise_variance)
+        batch_shape = y_event.shape[:-1]
+        n = y_event.shape[-1]
+        M = K + variance[..., None, None] * jnp.eye(n, dtype=K.dtype)
+        if jnp.broadcast_shapes(M.shape[:-2], batch_shape) != batch_shape:
+            raise ValueError(
+                f"The kernel's batch shape {K.shape[:-2]} and noise variance shape "
+                f"{variance.shape} do not broadcast to the event batch shape {batch_shape}."
+            )
+        matrix_batch = (1,) * (len(batch_shape) - (M.ndim - 2)) + M.shape[:-2]
+        # Batch axes along which M is shared become extra right-hand sides.
+        matrix_axes = [i for i, size in enumerate(matrix_batch) if size == batch_shape[i]]
+        shared_axes = [i for i in range(len(batch_shape)) if i not in matrix_axes]
+        residual = jnp.broadcast_to(observed - y_event, y_event.shape)
+        R = jnp.transpose(residual, matrix_axes + [len(batch_shape)] + shared_axes)
+        matrix_shape = tuple(batch_shape[i] for i in matrix_axes)
+        R = R.reshape(matrix_shape + (n, -1))
+        M = M.reshape(matrix_shape + (n, n))
+        return _gaussian_log_prob(M, R)
+
     def orthogonal_log_prob(
         self,
         y_event: jnp.ndarray,

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/evaluators.py

@@ -255,6 +255,11 @@ class MarginalLogLikelihood(AbstractEvaluator):
         An optional discrepancy model to cater for model misspecification.
         Can be a function or a PyTree with optional parameters.
         See :class:`pmrf.discrepancy_models` for common discrepancy models.
+        A :class:`~pmrf.discrepancy_models.GaussianProcess` with a
+        :class:`~pmrf.likelihoods.GaussianLikelihood` whose noise is constant along
+        the event axis is evaluated in closed form
+        (:meth:`~pmrf.discrepancy_models.GaussianProcess.log_prob`), which factorizes
+        each distinct covariance matrix once.
     use_orthogonal_discrepancy
         Constrain a Gaussian-process discrepancy to the complement of the free-parameter
         tangent space. This retains the full-data likelihood; it is not REML.
@@ -403,17 +408,55 @@ class MarginalLogLikelihood(AbstractEvaluator):
             log_prob, event_transform = self._orthogonal_log_prob(
                 model, frequency, **kwargs
             )
-            if self.has_conditional_event_transform:
-                log_prob = log_prob + jnp.sum(
-                    event_transform.forward_log_det_jacobian(observed)
+        else:
+            closed_form = self._closed_form_log_prob(model, frequency, **kwargs)
+            if closed_form is not None:
+                log_prob, event_transform = closed_form
+            else:
+                log_prob, event_transform = self._distribution_log_prob(
+                    model, frequency, **kwargs
                 )
-            return log_prob
 
+        # A prediction-dependent change of variables carries a Jacobian determinant that
+        # varies with the model, so it must be included for the density to be normalized.
+        # For a static transform the term is a constant offset, so it is omitted.
+        if self.has_conditional_event_transform:
+            log_prob = log_prob + jnp.sum(event_transform.forward_log_det_jacobian(observed))
+
+        return log_prob
+
+    def _closed_form_log_prob(
+        self, model: PyTree, frequency: Frequency, **kwargs
+    ) -> tuple[jnp.ndarray, bij.AbstractBijector] | None:
+        """The log-likelihood of a GP with Gaussian noise constant over the event axis.
+
+        Returns ``None`` when it does not apply, in which case the distribution path
+        gives the same value.
+        """
+        if not isinstance(self.discrepancy, GaussianProcess) or not isinstance(
+            self.likelihood, GaussianLikelihood
+        ):
+            return None
+        pred_event, event_transform = self._event(model, frequency, **kwargs)
+        variance = self.likelihood._constant_variance(pred_event)
+        if variance is None:
+            return None
+        log_prob = self.discrepancy.log_prob(
+            pred_event,
+            event_transform.forward(self.observed),
+            frequency.f_scaled,
+            variance,
+        )
+        return log_prob, event_transform
+
+    def _distribution_log_prob(
+        self, model: PyTree, frequency: Frequency, **kwargs
+    ) -> tuple[jnp.ndarray, bij.AbstractBijector]:
         # Get the distribution over obs_event and the actual observed event.
         # The observation is mapped by the *same* resolved transform as the prediction,
         # which for a conditional transform depends on the model.
         obs_dist, event_transform = self._predictive(model, frequency, **kwargs)
-        obs_event = event_transform.forward(observed)
+        obs_event = event_transform.forward(self.observed)
         batch_ndims = obs_event.ndim - self.event_ndims
         
         # We evaluate the log prob `batch_ndims` many times
@@ -424,15 +467,7 @@ class MarginalLogLikelihood(AbstractEvaluator):
             mapped_log_prob = eqx.filter_vmap(mapped_log_prob)
         
         log_probs = mapped_log_prob(obs_dist, obs_event)
-        log_prob = jnp.sum(log_probs)
-
-        # A prediction-dependent change of variables carries a Jacobian determinant that
-        # varies with the model, so it must be included for the density to be normalized.
-        # For a static transform the term is a constant offset, so it is omitted.
-        if self.has_conditional_event_transform:
-            log_prob = log_prob + jnp.sum(event_transform.forward_log_det_jacobian(observed))
-
-        return log_prob
+        return jnp.sum(log_probs), event_transform
 
     def with_orthogonal_reference(
         self, model: PyTree, frequency: Frequency, **kwargs

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/fitting/minimize.py

@@ -17,7 +17,7 @@ from pmrf.problems import SummedTerms, PriorPenalized
 from pmrf.terms import as_terms
 from pmrf.fitting.targets import resolve_datasets, union_frequency
 from pmrf.likelihoods import GaussianLikelihood
-from pmrf.losses import RMSELoss
+from pmrf.losses import MSELoss
 from pmrf.parameters import Random
 from pmrf.distributions import Uniform
 
@@ -74,16 +74,20 @@ def fit_minimize(
         Can be a function or a callable PyTree with optional parameters.
         Used to internally create a :class:`pmrf.evaluators.TargetLoss` evaluator.
         Mutually exclusive with `likelihood`. If neither `loss` nor `likelihood` is passed,
-        :class:`pmrf.losses.RMSELoss` is used for `loss` if `inference` is 'frequentist',
+        :class:`pmrf.losses.MSELoss` is used for `loss` if `inference` is 'frequentist',
         otherwise :class:`pmrf.likelihoods.GaussianLikelihood` is used for `likelihood`.
         See :mod:`pmrf.losses` for common losses.
+
+        An MSE that starts very small (e.g. fitting a feature that barely moves) can
+        fall below SciPy's absolute `gtol` and `ftol` before the fit moves; tighten them
+        via ``ScipyMinimize(options=...)``.
     likelihood : str | Callable, optional
         A likelihood model representing the probability of observing the data.
         Can be a function or a callable PyTree with optional parameters.
         Used to internally create a :class:`pmrf.evaluators.MarginalLogLikelihood`,
         wrapped in a :class:`pmrf.evaluators.Negated` evaluator.
         Mutually exclusive with `loss`. If neither `loss` nor `likelihood` is passed,
-        :class:`pmrf.losses.RMSELoss` is used for `loss` if `inference` is 'frequentist',
+        :class:`pmrf.losses.MSELoss` is used for `loss` if `inference` is 'frequentist',
         otherwise :class:`pmrf.likelihoods.GaussianLikelihood` is used for `likelihood`.
         See :mod:`pmrf.losses` for common losses.
     noise : prf.Param | Callable[[jnp.ndarray], jnp.ndarray], optional
@@ -135,7 +139,7 @@ def fit_minimize(
     # Resolve defaults e.g. loss vs MLE vs MAP optimization
     if loss is None and likelihood is None:
         if inference == 'frequentist':
-            loss = RMSELoss()
+            loss = MSELoss()
         else:
             if noise is None:
                 # Lower bound kept strictly positive: at noise=0 the Gaussian likelihood

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/likelihoods.py

@@ -92,6 +92,43 @@ class GaussianLikelihood(AbstractLikelihood):
             )
         return var
 
+    def _constant_variance(self, y_event: jnp.ndarray) -> jnp.ndarray | None:
+        """Return the noise variance over the batch shape if it is constant along the event axis.
+
+        The variance is interpreted exactly as in :meth:`__call__`, and kept at its own
+        shape, which broadcasts to the non-event batch shape of ``y_event``.
+
+        Returns
+        -------
+        jnp.ndarray | None
+            The variance, or ``None`` if it varies along the event axis or does not
+            broadcast to ``y_event``.
+        """
+        var = self._event_variance(y_event)
+        try:
+            broadcast_shape = jnp.broadcast_shapes(var.shape, y_event.shape)
+        except ValueError:
+            return None
+        if broadcast_shape != y_event.shape:
+            return None
+        if var.ndim == 0:
+            return var
+        if var.shape[-1] != 1:
+            return None
+        return var[..., 0]
+
+    def _event_variance(self, y_event: jnp.ndarray) -> jnp.ndarray:
+        """The noise variance with the event as its last axis, broadcastable to ``y_event``.
+
+        A variance of exactly the batch shape is per batch entry; any other shape
+        aligns its trailing axis with the event axis.
+        """
+        var = self.noise(y_event) if callable(self.noise) else self.noise
+        var = jnp.asarray(var)
+        if var.shape == y_event.shape[:-1]:
+            var = var[..., None]
+        return var
+
     def __call__(self, y_event: jnp.ndarray | AbstractDistribution) -> AbstractDistribution:
         # If y_event is an array, the prediction is deterministic
         # and we can simply use a regular gaussian likelihood.
@@ -101,15 +138,7 @@ class GaussianLikelihood(AbstractLikelihood):
         is_dist = isinstance(y_event, AbstractDistribution)
         y_mean = y_event.mean() if is_dist else y_event
         
-        # Evaluate noise
-        var = self.noise(y_mean) if callable(self.noise) else self.noise
-        var = jnp.asarray(var)
-        
-        # Broadcast noise onto y_mean
-        batch_shape = y_mean.shape[:-1]
-        if var.shape == batch_shape:
-            var = var[..., None]
-        mapped_var = jnp.broadcast_to(var, y_mean.shape)
+        mapped_var = jnp.broadcast_to(self._event_variance(y_mean), y_mean.shape)
         num_batch_dims = y_mean.ndim - 1
 
         if not is_dist:

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/parameters.py

@@ -450,11 +450,14 @@ def _check_in_constraint(value: ArrayLike, constraint: AbstractConstraint) -> Ar
     The returned array must be used: under `jax.jit` the check lives only in its graph.
     """
     value_array = jnp.asarray(value)
+    # Escaped so that braces in the constraint's repr are not read as format fields.
+    constraint_text = str(constraint).replace("{", "{{").replace("}", "}}")
     return error_if(
         value_array,
         constraint.is_outside(value_array),
-        f"\n\nA parameter value falls outside the constraint ({value} is not in {constraint}). "
+        f"\n\nA parameter value falls outside the constraint ({{}} is not in {constraint_text}). "
         f"\nMake sure the values match the parameter and model constraints.",
+        value_array,
     )
 
 

paramrf-0.36.4/pmrf/utils/debug.py

@@ -0,0 +1,83 @@
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def error_if(x, pred, msg, *print_args, on_error="default", **print_kwargs):
+    """
+    Conditionally halts JAX execution and prints formatted debug values at runtime.
+
+    This function evaluates a boolean condition inside JIT-compiled code. If the 
+    condition is met (e.g., an unphysical parameter like negative impedance is 
+    detected), it outputs the formatted runtime arrays to standard output before 
+    raising. Whether to print is decided on the host, so nothing is printed when
+    the check passes, under any transform. Under `jax.vmap`, the message is printed
+    once for each failing batch element, with that element's values. Without print
+    arguments nothing is printed: the raised error carries `msg`.
+
+    Parameters
+    ----------
+    x : Any
+        The input data (JAX array or PyTree) to pass through. The check runs only
+        if the returned value is used.
+    pred : bool or jax.Array
+        A boolean condition or boolean array. If any element evaluates to `True`,
+        execution halts, printing the message if print arguments are given.
+    msg : str
+        The format string for the error message and console output. Uses standard 
+        Python `{}` formatting to dynamically inject JAX arrays.
+    *print_args : Any
+        Dynamic tensors or values to sequentially format into the `msg` string.
+    on_error : str, optional
+        The internal error handling mode. Default is "default".
+    **print_kwargs : Any
+        Dynamic tensors or values to format into the `msg` string via keywords.
+
+    Returns
+    -------
+    Any
+        The unmodified input `x`.
+
+    Raises
+    ------
+    equinox.EquinoxRuntimeError
+        At runtime, if `pred` contains any `True` elements.
+
+    Examples
+    --------
+    Catching an invalid characteristic impedance inside a JIT-compiled simulation:
+
+    >>> z0 = jnp.array(-50.0)
+    >>> is_invalid = z0 < 0
+    >>> z0 = error_if(
+    ...     z0, 
+    ...     is_invalid, 
+    ...     "Characteristic impedance must be positive, got Z0 = {} Ohms", 
+    ...     z0
+    ... )
+    """
+    error_msg = msg
+    if print_args or print_kwargs:
+        # Decided on the host: a `lax.cond` on a batched predicate lowers to `select`
+        # under `vmap`, which runs the print whether or not the check fails. A
+        # `pure_callback` rather than `jax.debug.callback`, so that `eqx.error_if` can
+        # consume its output: that data dependency makes the print precede the raise.
+        def print_if(host_pred, args, kwargs):
+            if np.any(host_pred):
+                print(msg.format(*args, **kwargs))
+            return host_pred
+
+        pred = jnp.asarray(pred)
+        args, kwargs = jax.lax.stop_gradient((print_args, print_kwargs))
+        pred = jax.pure_callback(
+            print_if,
+            jax.ShapeDtypeStruct(pred.shape, pred.dtype),
+            pred,
+            args,
+            kwargs,
+            vmap_method="sequential",
+        )
+        error_msg = f"{msg} (Check standard output for runtime values)"
+
+    return eqx.error_if(x, pred, error_msg, on_error=on_error)

{paramrf-0.36.2 → paramrf-0.36.4}/pmrf/utils/transforms.py

@@ -1,6 +1,9 @@
+import dataclasses
+
 import jax
 import jax.numpy as jnp
 import equinox as eqx
+import parax as prx
 from typing import Callable, Any, Tuple, TypeVarTuple
 
 from pmrf.utils.tree import unwrap, batch_axes
@@ -35,6 +38,11 @@ def derivative(eval_fn: Callable[..., Any], *args: *Ts, space: str = 'declared')
     inside the differentiated function, so the source's derivative includes the
     path through the tie. Raw JAX arrays are differentiated as they are.
 
+    In declared and physical space a parameter on a closed bound, such as a
+    ``NonNegative`` loss at 0, has a finite derivative. In raw space it has none:
+    its raw value there is infinite, so a derivative with respect to it has no
+    finite meaning.
+
     Safely handles models as inputs by filtering out non-differentiable
     static fields (like strings, booleans, integers, AND NumPy arrays).
 
@@ -82,15 +90,15 @@ def derivative(eval_fn: Callable[..., Any], *args: *Ts, space: str = 'declared')
     -1.164e-01
     """
     # Imported here: pmrf.parameters imports pmrf.utils.
-    from pmrf.parameters import _check_space, _read, _set_paths, _with_fixed, _write, tree_param_paths
+    from pmrf.parameters import _check_space, _read, _set_paths, _with_fixed, _write, is_param, tree_param_paths
     from pmrf.utils.tree import Pathgetter
 
     _check_space(space)
 
     # Each parameter is replaced by its value in `space`, and rebuilt from it inside
     # the differentiated function, so the derivative is taken with respect to that
-    # value and the chain rule carries it through scale, constraint and ties. Fixed
-    # is an optimisation state, not a zero sensitivity, so parameters are rebuilt free.
+    # value and the chain rule carries it through scale and ties. Fixed is an
+    # optimisation state, not a zero sensitivity, so parameters are rebuilt free.
     rebuilds = []
     surrogates = []
     for arg in args:
@@ -103,12 +111,23 @@ def derivative(eval_fn: Callable[..., Any], *args: *Ts, space: str = 'declared')
         rebuilds.append((paths, nodes))
         surrogates.append(_set_paths(arg, paths, [_read(node, space) for node in nodes]))
 
+    def _rebuild_node(node, value):
+        # In declared and physical space a parameter holds its declared value directly,
+        # as in box space (#226), not through its raw value: on a closed bound that is
+        # infinite, the constraint's bijector gives 0 * inf = NaN, and forward mode
+        # spreads the NaN to every parameter. The value is the one the parameter was
+        # built with, so its bounds need no checking again. A `Real` is free.
+        if not is_param(node) or space == 'raw':
+            return _with_fixed(_write(node, value, space), False)
+        declared = value / node._scale if space == 'physical' and node._scale != 1.0 else value
+        return dataclasses.replace(node, variable=prx.Real(declared))
+
     def _rebuild(surrogate_args):
         rebuilt = []
         for surrogate, (paths, nodes) in zip(surrogate_args, rebuilds):
             values = [Pathgetter(path)(surrogate) for path in paths]
             rebuilt.append(_set_paths(
-                surrogate, paths, [_with_fixed(_write(node, value, space), False) for node, value in zip(nodes, values)]
+                surrogate, paths, [_rebuild_node(node, value) for node, value in zip(nodes, values)]
             ))
         return unwrap(tuple(rebuilt))
 

{paramrf-0.36.2 → paramrf-0.36.4}/pyproject.toml

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