diff --git a/fredipy/covariance.py b/fredipy/covariance.py index a11a91f..a70af8b 100644 --- a/fredipy/covariance.py +++ b/fredipy/covariance.py @@ -20,6 +20,17 @@ class TwoSided: New rules for user-defined operators can be defined by creating a class that inherits from this class, adding the appropriate functions following the naming scheme, and passing an instance to the model constructor. + + Every combiner takes a ``derivative`` flag, which is threaded through + unchanged from :meth:`__call__`. It is ``False`` when the covariance matrix + itself is being assembled and ``True`` when ``kernel`` is a *derivative* of + the kernel with respect to one hyperparameter (see + ``GaussianProcess.log_likelihood_grad``). Combiners that build a block by + plain linear application of ``kernel`` can ignore it; only pieces that add a + term which is *not* linear in ``kernel`` -- e.g. the constant analytic UV + tail moment in ``GaussLegendre_1D_log._uv_tail_correction``, whose + hyperparameter derivative is zero rather than itself -- need to act on it. + Custom combiners must accept the flag. """ def __init__(self): pass @@ -27,7 +38,8 @@ def __init__(self): def __call__( self, kernel: Kernel, - constraints: List[LinearEquality] + constraints: List[LinearEquality], + derivative: bool = False ) -> np.ndarray: rows = [] for c1 in constraints: @@ -36,9 +48,9 @@ def __call__( combiner12 = getattr(self, f"_{type(c1.op).__name__}_{type(c2.op).__name__}", None) combiner21 = getattr(self, f"_{type(c2.op).__name__}_{type(c1.op).__name__}", None) if combiner12: - entry = combiner12(c1, kernel, c2) + entry = combiner12(c1, kernel, c2, derivative=derivative) elif combiner21: - entry = combiner21(c2, kernel, c1).T + entry = combiner21(c2, kernel, c1, derivative=derivative).T else: raise NotImplementedError( f"No rule found to combine operators of types \ @@ -47,28 +59,39 @@ def __call__( rows.append(np.concatenate(columns, axis=1)) return np.concatenate(rows) - def _Integral_Integral(self, c1, k, c2): + def _Integral_Integral(self, c1, k, c2, derivative: bool = False): + # NOTE: c1's integrator drives both sides of the cross block, so the two constraints + # must discretise omega compatibly. That is not checked here: integrators that apply an + # analytic UV tail (GaussLegendre_1D_log._uv_tail_correction) enforce matching nodes and + # weights downstream and raise NotImplementedError otherwise. Blocks (i, j) and (j, i) + # are built independently, each through its own left constraint's integrator, so any + # correction added downstream must be symmetric by construction. + # `derivative` is forwarded because that UV tail correction contains one term (T1 T2^T) + # that is independent of the kernel and must therefore vanish, not be reproduced, when + # k is dK/dtheta. if c1 == c2: - return c1.op.integrator.doubleIntegrationSymmetric(c1, k) + return c1.op.integrator.doubleIntegrationSymmetric(c1, k, derivative=derivative) else: - return c1.op.integrator.doubleIntegration(c1, k, c2) + return c1.op.integrator.doubleIntegration(c1, k, c2, derivative=derivative) - def _Identity_Identity(self, c1, k, c2): + def _Identity_Identity(self, c1, k, c2, derivative: bool = False): return k(c1.x, c2.x) - def _Derivative_Derivative(self, c1, k, c2): + def _Derivative_Derivative(self, c1, k, c2, derivative: bool = False): return k.d2K_dxdy(c1.x, c2.x) - def _Integral_Identity(self, c1, k, c2): + def _Integral_Identity(self, c1, k, c2, derivative: bool = False): + # No `derivative` handling needed: singleIntegration's analytic tail, T @ K(w_uv, x) / + # f_uv(w_uv), is linear in the kernel, so substituting dK/dtheta already differentiates it. return c1.op.integrator.singleIntegration(c1, k, c2.x) - def _Integral_Derivative(self, c1, k, c2): + def _Integral_Derivative(self, c1, k, c2, derivative: bool = False): return c1.op.integrator.singleIntegration(c1, k.dK_dy, c2.x) - def _Identity_Derivative(self, c1, k, c2): + def _Identity_Derivative(self, c1, k, c2, derivative: bool = False): return k.dK_dy(c1.x, c2.x) - def _ConstMatrix_ConstMatrix(self, c1, k, c2): + def _ConstMatrix_ConstMatrix(self, c1, k, c2, derivative: bool = False): return c1.op() @ k(c1.x, c2.x) @ c2.op().T diff --git a/fredipy/integrators.py b/fredipy/integrators.py index 8251736..382201d 100644 --- a/fredipy/integrators.py +++ b/fredipy/integrators.py @@ -21,7 +21,8 @@ def __init__( def doubleIntegrationSymmetric( self, constraint: LinearEquality, - kernel: Callable + kernel: Callable, + derivative: bool = False ) -> np.ndarray: r"""Symmetric double integration of a constraint C with a GP kernel K, i.e. @@ -31,6 +32,8 @@ def doubleIntegrationSymmetric( ---------- constraint : Some integral constraint kernel : Some Gaussian process kernel + derivative : True when ``kernel`` is dK/dtheta rather than K, see + :meth:`doubleIntegration`. Returns ------- @@ -42,7 +45,8 @@ def doubleIntegration( self, constraint1: LinearEquality, kernel: Callable, - constraint2: LinearEquality + constraint2: LinearEquality, + derivative: bool = False ) -> np.ndarray: r"""Double integration of two (different) constraints C1 and C2 with a GP kernel K, i.e. @@ -52,6 +56,16 @@ def doubleIntegration( ---------- constraint1, constraint2 : Some integral constraints kernel : Some Gaussian process kernel + derivative : Flag set by ``GaussianProcess.log_likelihood_grad`` (via + ``TwoSided``) to say that ``kernel`` is the derivative + dK/dtheta of the kernel with respect to one hyperparameter + rather than the kernel itself. Purely quadrature-based + implementations are linear in ``kernel`` and can ignore it; + any term that is *not* linear in ``kernel`` (the constant + analytic tail moment of + :meth:`GaussLegendre_1D_log._uv_tail_correction`) must be + dropped when it is set, since the derivative of a constant + is zero. Returns ------- @@ -81,6 +95,61 @@ def singleIntegration( """ raise NotImplementedError + def uv_tail_moment( + self, + constraint: LinearEquality + ) -> np.ndarray: + r"""Analytic UV tail moment ``T`` of this integrator for one constraint. + + .. math:: + T_i[m] = \int_{w_{uv}}^{\infty} C_i(p_m, w) f_{uv}(w)\, dw + + i.e. the constraint kernel weighted by the *bare* UV asymptotic shape + ``f_uv`` over the part of the ω axis that the quadrature grid does not + cover. It is the only ingredient (besides the bulk-UV overlap ``A``, + see :meth:`uv_anchor`) of the general two-constraint tail correction + + .. math:: + \Sigma_{12} = BB + A_1 T_2^T + T_1 A_2^T + T_1 T_2^T \qquad (**) + + implemented in :meth:`GaussLegendre_1D_log._uv_tail_correction`. + + ``T`` is per-constraint **and** per-row: two constraints whose kernels + have different UV falloffs have different tail moments, so a single + shared scalar is *not* valid for the cross block (see + ``docs/fredipy_uvtail_fix_plan.md`` in the reconstructions repo). + + Returns + ------- + 2D array of shape ``(len(constraint.x), 1)``. The default is all + zeros — integrators that carry no analytic tail, which lets + ``_uv_tail_correction`` short-circuit and stay bit-identical to the + uncorrected quadrature. + """ + return np.zeros((constraint.x.shape[0], 1)) + + def uv_anchor(self) -> tuple | None: + r"""Anchor point used to evaluate the bulk-UV overlap ``A`` of ``(**)``. + + .. math:: + A_i[m] = \frac{1}{f_{uv}(w_{uv})} + \sum_j W_j\, C_i(p_m, w_j)\, K(w_j, w_{uv}) + + which is *exact*, not approximate, whenever ``w_uv >> mu_uv``: the UV + branch of ``AsymptoticKernel`` factorises there as + ``K(w, w_uv) = θ_uv(w) f_uv(w) f_uv(w_uv)``, so dividing by the anchor + value ``f_uv(w_uv)`` recovers ``Σ_j W_j C_i(w_j) θ_uv(w_j) f_uv(w_j)`` + to all digits. + + Returns + ------- + ``None`` for tail-free integrators (the default). Tail-carrying + integrators return ``(w_uv, f_uv_anchor)`` with ``w_uv`` a ``(1, 1)`` + column vector holding the UV split point and ``f_uv_anchor`` the scalar + ``f_uv(w_uv)``. + """ + return None + class Riemann(Integrator): """Implementation of the Riemann integration in arbitrary dimensions and kernels. @@ -103,7 +172,8 @@ def __init__( def doubleIntegrationSymmetric( self, constraint: LinearEquality, - kernel: Callable + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: p = constraint.x @@ -161,15 +231,17 @@ def __init__( def doubleIntegrationSymmetric( self, constraint: LinearEquality, - kernel: Callable + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: - return self.doubleIntegration(constraint, kernel, constraint) + return self.doubleIntegration(constraint, kernel, constraint, derivative=derivative) def doubleIntegration( self, constraint1: LinearEquality, kernel: Callable, - constraint2: LinearEquality + constraint2: LinearEquality, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: return self.dw**2 * ( constraint1(make_row_vector(self.w), x=make_column_vector(constraint1.x)) @@ -190,6 +262,439 @@ def singleIntegration( ) +class Riemann_1D_log(Integrator): + """Riemann integration in 1D on a logarithmic (geomspace) grid. + + Uses midpoints equally spaced in log-space and a scalar step + ``dw = d(log ω) = log(w_max/w_min) / int_n``. + + The Jacobian of the substitution t = log ω is ω and is applied + internally, so plain kernels ``K(p, ω)`` can be passed directly. + """ + + def __init__( + self, + w_min: float = 0.01, + w_max: float = 10., + int_n: int = 1000 + ) -> None: + + w_edges = make_column_vector(np.geomspace(w_min, w_max, int_n + 1)) + self.dw = np.log(w_max / w_min) / int_n # uniform step in log-space + self.w = np.sqrt(w_edges[:-1] * w_edges[1:]) # geometric midpoints + self.jac = make_row_vector(self.w) # log-space Jacobian: ω_i, shape (1, n) + + def doubleIntegrationSymmetric( + self, + constraint: LinearEquality, + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` + ) -> np.ndarray: + return self.doubleIntegration(constraint, kernel, constraint, derivative=derivative) + + def doubleIntegration( + self, + constraint1: LinearEquality, + kernel: Callable, + constraint2: LinearEquality, + derivative: bool = False # unused: this rule is linear in `kernel` + ) -> np.ndarray: + return self.dw**2 * ( + self.jac * constraint1(make_row_vector(self.w), x=make_column_vector(constraint1.x)) + @ kernel(self.w, self.w) + @ (self.jac * constraint2(make_row_vector(self.w), x=make_column_vector(constraint2.x))).T + ) + + def singleIntegration( + self, + constraint: LinearEquality, + kernel: Callable, + w_pred: np.ndarray + ) -> np.ndarray: + return self.dw * ( + self.jac * constraint(make_row_vector(self.w), x=make_column_vector(constraint.x)) + @ kernel(self.w, w_pred) + ) + + +class GaussLegendre_1D_log(Integrator): + """Gauss-Legendre quadrature in 1D on a logarithmic grid. + + Maps the standard Gauss-Legendre nodes from [-1, 1] to log-space + [log(w_min), log(w_max)] and then to ω-space via exponentiation. + + Achieves exponential convergence in the number of quadrature points for + smooth integrands, requiring far fewer points than Riemann rules for the + same accuracy (typically 50–200 instead of 1000). + + The log-space Jacobian factor ω is absorbed into the quadrature weights, + so plain kernels ``K(p, ω)`` can be passed directly. + """ + + def __init__( + self, + w_min: float = 0.01, + w_max: float = 10., + int_n: int = 100 + ) -> None: + + # Grid definition kept as attributes purely for diagnostics: the node-compatibility + # guard in _uv_tail_correction names these in its error message. + self.w_min = w_min # lower end of the quadrature range in ω + self.w_max = w_max # upper end of the quadrature range in ω + self.int_n = int_n # number of Gauss-Legendre nodes + + # GL nodes (xi in [-1,1]) and weights + xi, wi = np.polynomial.legendre.leggauss(int_n) + + # Map nodes from [-1,1] to log-space [log(w_min), log(w_max)] + log_min = np.log(w_min) + log_max = np.log(w_max) + half_range = 0.5 * (log_max - log_min) + + t = half_range * xi + 0.5 * (log_max + log_min) # nodes in log-space + self.w = make_column_vector(np.exp(t)) # nodes in ω-space + # Weights absorb the half-range Jacobian and the log-space Jacobian ω + self.weights = make_row_vector(wi * half_range) * self.w.T # shape (1, int_n) + + def doubleIntegrationSymmetric( + self, + constraint: LinearEquality, + kernel: Callable, + derivative: bool = False + ) -> np.ndarray: + return self.doubleIntegration(constraint, kernel, constraint, derivative=derivative) + + def doubleIntegration( + self, + constraint1: LinearEquality, + kernel: Callable, + constraint2: LinearEquality, + derivative: bool = False + ) -> np.ndarray: + bulk = ( + self.weights * constraint1(make_row_vector(self.w), x=make_column_vector(constraint1.x)) + @ kernel(self.w, self.w) + @ (self.weights * constraint2(make_row_vector(self.w), x=make_column_vector(constraint2.x))).T + ) + # The analytic UV tail lives here (not as an override on the UVtail subclass) so that + # both orderings of a constraint pair go through the same code and the assembled + # covariance matrix is symmetric *by construction* — np.linalg.cholesky reads only the + # lower triangle and would silently accept an asymmetric matrix otherwise. + return bulk + self._uv_tail_correction( + constraint1, kernel, constraint2, derivative=derivative) + + def singleIntegration( + self, + constraint: LinearEquality, + kernel: Callable, + w_pred: np.ndarray + ) -> np.ndarray: + return ( + self.weights * constraint(make_row_vector(self.w), x=make_column_vector(constraint.x)) + @ kernel(self.w, w_pred) + ) + + @staticmethod + def _same_quadrature( + integrator1: Integrator, + integrator2: Integrator + ) -> bool: + """True if two integrators share (numerically) the same nodes and weights.""" + w1, w2 = getattr(integrator1, 'w', None), getattr(integrator2, 'w', None) + v1, v2 = getattr(integrator1, 'weights', None), getattr(integrator2, 'weights', None) + if w1 is None or w2 is None or v1 is None or v2 is None: + return False + if w1.shape != w2.shape or v1.shape != v2.shape: + return False + return bool(np.allclose(w1, w2) and np.allclose(v1, v2)) + + @staticmethod + def _describe_grid(integrator: Integrator) -> str: + """Human-readable grid summary used in the node-mismatch error message.""" + return (f"{type(integrator).__name__}(w_min={getattr(integrator, 'w_min', '?')}, " + f"w_max={getattr(integrator, 'w_max', '?')}, " + f"int_n={getattr(integrator, 'int_n', '?')})") + + def _uv_tail_correction( + self, + constraint1: LinearEquality, + kernel: Callable, + constraint2: LinearEquality, + derivative: bool = False + ) -> np.ndarray | float: + r"""Analytic UV tail correction to the (constraint1, constraint2) covariance block. + + Splitting *each* of the two integrals at the UV point ``w_uv`` into a bulk part + ``(w_min, w_uv)`` — which the quadrature grid covers — and a tail part + ``(w_uv, ∞)`` — which it does not — and using that ``AsymptoticKernel`` + factorises as ``K(w, w') = [θ_uv(w) f_uv(w)] f_uv(w')`` whenever ``w' > w_uv`` + with ``w_uv >> mu_uv`` (rank-1 in the tail, exact up to + ``O(exp(-(w_uv - mu_uv)/l_uv))``), the four pieces give + + .. math:: + \Sigma_{12} = BB + A_1 T_2^T + T_1 A_2^T + T_1 T_2^T \qquad (**) + + with ``BB`` the plain double quadrature computed by + :meth:`doubleIntegration`, ``A_i`` the bulk-UV overlap (see + :meth:`Integrator.uv_anchor`) and ``T_i`` the per-constraint, per-row tail + moment (see :meth:`Integrator.uv_tail_moment`). + + Setting ``C_1 = C_2`` and ``T`` row-independent collapses ``(**)`` to the + legacy symmetric closed form ``bulk + 2 A T + T²``, so the symmetric case is + a special case of this method and needs no separate override. + + Gradient mode (``derivative=True``) + ----------------------------------- + ``GaussianProcess.log_likelihood_grad`` reassembles the same block with + ``dK/dtheta`` substituted for ``K``. Under that substitution the two ``A`` + terms differentiate themselves correctly — ``A_i`` is linear in the kernel + (``A_i = W C_i K(ω, w_uv) / f_uv(w_uv)`` with a hyperparameter-independent + anchor value), so ``A_i[dK/dtheta] = dA_i/dtheta`` exactly. The term + ``T_1 T_2^T`` however contains no kernel at all: ``T`` is the *analytic* + tail moment, a function of ``w_uv`` and the anomalous dimension only, so + ``d(T_1 T_2^T)/dtheta = 0``. Adding it unchanged in gradient mode would + treat the derivative of a constant as the constant itself, which is why it + is dropped here. The likelihood *value* is unaffected either way; only the + gradient is. + """ + integrator1 = constraint1.op.integrator + integrator2 = constraint2.op.integrator + tail1 = integrator1.uv_tail_moment(constraint1) + tail2 = integrator2.uv_tail_moment(constraint2) + + # Fast path: no constraint carries an analytic tail, so (**) reduces to BB. + # This short circuit is load-bearing — it keeps every plain-quadrature model + # (three-gluon-vert, ghost-gluon-vert, fgvert-sd, all pre-existing fredipy + # tests) bit-identical to the uncorrected result. Do not drop it. + if not np.any(tail1) and not np.any(tail2): + return 0.0 + + # (**) mixes the two constraints on a *shared* set of quadrature nodes; it is + # not well defined if the two integrators discretise ω differently. + if integrator1 is not integrator2 and not self._same_quadrature(integrator1, integrator2): + raise NotImplementedError( + "UV tail correction requires both integral constraints to share the same " + "quadrature nodes and weights, but got " + f"{self._describe_grid(integrator1)} and {self._describe_grid(integrator2)}." + ) + + anchors = [a for a in (integrator1.uv_anchor(), integrator2.uv_anchor()) if a is not None] + if not anchors: + raise NotImplementedError( + "A non-zero UV tail moment was reported but no integrator supplied a UV " + f"anchor via uv_anchor(): {self._describe_grid(integrator1)} and " + f"{self._describe_grid(integrator2)}." + ) + w_uv, f_uv_anchor = anchors[0] + for other_w_uv, other_f_uv in anchors[1:]: + assert np.allclose(w_uv, other_w_uv) and np.isclose(f_uv_anchor, other_f_uv), \ + "Integral constraints report incompatible UV anchors." + + # A_i = (W C_i) @ K(ω, w_uv) / f_uv(w_uv), both evaluated on *this* integrator's + # nodes (guaranteed above to coincide with the other integrator's). + c1_row = constraint1(make_row_vector(self.w), x=make_column_vector(constraint1.x)) + c2_row = constraint2(make_row_vector(self.w), x=make_column_vector(constraint2.x)) + K_col = kernel(self.w, w_uv) # (N, 1) + A1 = (self.weights * c1_row @ K_col) / f_uv_anchor # (M1, 1) + A2 = (self.weights * c2_row @ K_col) / f_uv_anchor # (M2, 1) + + correction = A1 @ tail2.T + tail1 @ A2.T + if derivative: + # d(T_1 T_2^T)/dtheta = 0: T carries no hyperparameter dependence. See the + # "Gradient mode" section of the docstring. + return correction + return correction + tail1 @ tail2.T + + +class GaussLegendre_1D_log_UVtail(GaussLegendre_1D_log): + """GL log-space quadrature on (w_min, w_uv) plus an analytic UV tail. + + **Rank-1 tail assumption.** For ω' > w_uv with w_uv >> mu_uv the soft + thetas of ``AsymptoticKernel`` saturate (θ_uv(ω') = 1 to machine + precision), so the kernel becomes rank-1 in the tail: + + K(ω, ω') = [θ_uv(ω) f_uv(ω)] · f_uv(ω') (*) + + This is the *only* assumption made here; it holds up to + ``O(exp(-(w_uv - mu_uv)/l_uv))``. + + **General two-constraint correction.** Splitting each of the two integrals + of a covariance block at w_uv into bulk (w_min, w_uv) and tail (w_uv, ∞) and + applying (*) to whichever argument lies in the tail gives + + A_i[m] = ∫_bulk C_i(p_m, ω) θ_uv(ω) f_uv(ω) dω + T_i[m] = ∫_{w_uv}^∞ C_i(p_m, ω) f_uv(ω) dω + + Σ_12 = BB + A_1 T_2^T + T_1 A_2^T + T_1 T_2^T (**) + + where ``BB`` is the plain double quadrature of the parent class. Two points + that are easy to get wrong: + + * ``T`` is **per constraint and per row**, not a single shared scalar: two + constraints whose kernels have different UV falloffs (e.g. a sum-rule + kernel ``C(p, ω) = ω`` versus a Källén-Lehmann data kernel) have entirely + different tail moments. This is why ``tail_moment`` is a required + constructor argument rather than a hardcoded closed form. + * ``A`` obtained via the anchor, ``A = (W C @ K(ω, w_uv)) / f_uv(w_uv)``, is + **exact**, not approximate: by (*) the anchor column is + ``K(ω_i, w_uv) = θ_uv(ω_i) f_uv(ω_i) f_uv(w_uv)``, so the division + recovers the bulk-UV overlap to all digits. + + The symmetric case ``C_1 = C_2`` with row-independent ``T`` collapses (**) + to the legacy closed form ``bulk + 2 A T + T²``, so this class deliberately + does **not** override ``doubleIntegration``/``doubleIntegrationSymmetric``: + the correction is applied by the parent class for *both* orderings of a + constraint pair, which is what keeps the assembled covariance matrix + symmetric by construction (``np.linalg.cholesky`` reads only the lower + triangle and would silently accept an asymmetric matrix). + + As with the parent class, plain kernels ``K(p, ω)`` are passed directly; + the log-space Jacobian ω is absorbed into the quadrature weights. + + Parameters + ---------- + w_min : float + Lower end of the quadrature range in ω. + w_uv : float + UV split where the UV asymptotics fully apply, i.e. the upper end of the + quadrature range; everything above it is covered by ``tail_moment``. + Recommended: ≥ 1000, and in any case >> mu_uv of the kernel. + int_n : int + Number of Gauss-Legendre nodes on (w_min, w_uv). + uv_func : callable + The UV asymptotic function f_uv(ω), e.g. ``uv_asymptotics``. Used only + to evaluate the anchor value f_uv(w_uv). + tail_moment : callable + Maps ``constraint.x`` of shape (M, d) to the tail moments T of shape + (M, 1), i.e. ``T[m] = ∫_{w_uv}^∞ C(p_m, ω) f_uv(ω) dω`` for the + constraint this integrator is attached to. Required (no default): a + silent default is exactly what let a gluon-specific closed form be + applied to the ghost sum rule, off by a factor 10.4. + """ + + def __init__( + self, + w_min: float, + w_uv: float, + int_n: int, + uv_func: Callable, + tail_moment: Callable[[np.ndarray], np.ndarray] + ) -> None: + super().__init__(w_min, w_uv, int_n) + self.uv_func = uv_func # UV asymptotic shape f_uv(ω), used for the anchor value + self.tail_moment = tail_moment # constraint.x -> (M, 1) analytic tail moments T + # Anchor point at the UV split boundary + self.w_uv = make_column_vector(np.array([w_uv])) # (1, 1) + self.f_uv_anchor = float(uv_func(self.w_uv).flat[0]) # scalar: f_uv(w_uv) + + def uv_tail_moment( + self, + constraint: LinearEquality + ) -> np.ndarray: + tail = np.asarray(self.tail_moment(constraint.x), dtype=float) + if tail.ndim == 1: + tail = tail.reshape(-1, 1) + expected = (constraint.x.shape[0], 1) + if tail.shape != expected: + raise ValueError( + f"tail_moment must return an array of shape {expected} for this constraint, " + f"got {tail.shape}." + ) + return tail + + def uv_anchor(self) -> tuple: + return (self.w_uv, self.f_uv_anchor) + + def singleIntegration( + self, + constraint: LinearEquality, + kernel: Callable, + w_pred: np.ndarray + ) -> np.ndarray: + numerical = super().singleIntegration(constraint, kernel, w_pred) + # tail[m, n] = T[m] × K(w_uv, ω_pred[n]) / f_uv(w_uv) — an outer product, since T is + # per-row: broadcasting a single scalar T is only correct for a one-row constraint. + # K(w_uv, ω_pred) / f_uv(w_uv) ≈ θ_uv(ω_pred) f_uv(ω_pred) for w_uv >> ω_pred. + tail = self.uv_tail_moment(constraint) @ (kernel(self.w_uv, w_pred) / self.f_uv_anchor) + return numerical + tail + + +class GaussLegendre_1D_semiinf(Integrator): + """Gauss-Legendre quadrature on the full semi-infinite interval (0, ∞). + + Uses the rational change of variables + + ω = w_scale · (1 + u) / (1 − u), u ∈ (−1, 1) + + which maps (−1, 1) exactly onto (0, ∞). The log-space Jacobian + d(log ω)/du = 2/(1 − u²) and the ω factor are both absorbed into + the quadrature weights, so plain kernels ``K(p, ω)`` can be passed + directly. + + Parameters + ---------- + w_scale : float + The scale point: u = 0 maps to ω = w_scale. Choose it near the + centre of the integrand in log-space for the best node distribution. + int_n : int + Number of Gauss-Legendre nodes. + + Notes + ----- + For n = 100 nodes and w_scale = 1, the outermost nodes lie at + roughly ω ≈ 7×10⁻⁵ and ω ≈ 1.4×10⁴. The GL weights at those + extreme nodes automatically vanish proportionally to (1 − u²), + exactly cancelling the Jacobian divergence, so the product weight + W_i = w_i · 2/(1 − u_i²) is bounded O(1/n²) at all nodes. + """ + + def __init__( + self, + w_scale: float = 1.0, + int_n: int = 100 + ) -> None: + + xi, wi = np.polynomial.legendre.leggauss(int_n) + + # rational map (-1,1) → (0,∞) + self.w = make_column_vector(w_scale * (1.0 + xi) / (1.0 - xi)) + # weights: log-space d(log ω)/du = 2/(1-u²), plus ω Jacobian absorbed + self.weights = make_row_vector(wi * 2.0 / (1.0 - xi**2)) * self.w.T + + def doubleIntegrationSymmetric( + self, + constraint: LinearEquality, + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` + ) -> np.ndarray: + return self.doubleIntegration(constraint, kernel, constraint, derivative=derivative) + + def doubleIntegration( + self, + constraint1: LinearEquality, + kernel: Callable, + constraint2: LinearEquality, + derivative: bool = False # unused: this rule is linear in `kernel` + ) -> np.ndarray: + return ( + self.weights * constraint1(make_row_vector(self.w), x=make_column_vector(constraint1.x)) + @ kernel(self.w, self.w) + @ (self.weights * constraint2(make_row_vector(self.w), x=make_column_vector(constraint2.x))).T + ) + + def singleIntegration( + self, + constraint: LinearEquality, + kernel: Callable, + w_pred: np.ndarray + ) -> np.ndarray: + return ( + self.weights * constraint(make_row_vector(self.w), x=make_column_vector(constraint.x)) + @ kernel(self.w, w_pred) + ) + + class Simpson_1D(Integrator): """Implementation of Simpson's rule for 1D integration""" @@ -215,15 +720,17 @@ def __init__( def doubleIntegrationSymmetric( self, constraint: LinearEquality, - kernel: Callable + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: - return self.doubleIntegration(constraint, kernel, constraint) + return self.doubleIntegration(constraint, kernel, constraint, derivative=derivative) def doubleIntegration( self, constraint1: LinearEquality, kernel: Callable, - constraint2: LinearEquality + constraint2: LinearEquality, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: return self.dw**2 / 9 * ( @@ -270,7 +777,8 @@ def __init__( def doubleIntegrationSymmetric( self, constraint: LinearEquality, - kernel: Callable + kernel: Callable, + derivative: bool = False # unused: this rule is linear in `kernel` ) -> np.ndarray: p = constraint.x diff --git a/fredipy/kernels.py b/fredipy/kernels.py index 3a2c87c..638fc98 100644 --- a/fredipy/kernels.py +++ b/fredipy/kernels.py @@ -278,6 +278,12 @@ def __init__( self._K_asymp = None self.dim = kernel.dim + def _empty_cache(self) -> None: + """Clear this kernel's cache and propagate to the inner kernel.""" + self._K_asymp = None + self._x, self._y = np.array([None]), np.array([None]) + self.kernel._empty_cache() + def add_asymptotics( self, region: str, @@ -370,6 +376,136 @@ def set_params( self.mu_uv = asymp_params[0] self.l_uv = asymp_params[1] + def params_gradient(self) -> List[Callable]: + """Analytic gradient of the asymptotic kernel w.r.t. all hyperparameters. + + Returns a list of functions ``f_i(x, y)`` such that + ``dK_asymp / dtheta_i = f_i(x, y)``. + + Parameter ordering (identical to ``set_params`` / flat-vector convention): + [base_kernel_params..., mu_uv, l_uv (if UV), mu_ir, l_ir (if IR)] + + Derivation + ---------- + With the shorthands + σ_+(x) = softtheta(x, mu, l, +1) (sigmoid, 0 → 1) + σ_-(x) = softtheta(x, mu, l, −1) (complement, 1 → 0) + and the three kernel regions + A = σ_-(x;ir)·σ_-(y;ir)·f_ir(x)·f_ir(y) + B = σ_+(x;ir)·σ_+(y;ir)·K_rbf·σ_-(x;uv)·σ_-(y;uv) + C = σ_+(x;uv)·σ_+(y;uv)·f_uv(x)·f_uv(y) + the sigmoid derivative identity dσ_±/dmu = ∓(1/l)·σ_+·σ_− gives: + + dK/dmu_uv = (1/l_uv)·[B·(σ_+(x;uv)+σ_+(y;uv)) − C·(σ_-(x;uv)+σ_-(y;uv))] + dK/dl_uv = (1/l_uv²)·[B·((x−mu_uv)·σ_+(x;uv)+(y−mu_uv)·σ_+(y;uv)) + −C·((x−mu_uv)·σ_-(x;uv)+(y−mu_uv)·σ_-(y;uv))] + dK/dmu_ir = (1/l_ir)·[A·(σ_+(x;ir)+σ_+(y;ir)) − B·(σ_-(x;ir)+σ_-(y;ir))] + dK/dl_ir = (1/l_ir²)·[A·((x−mu_ir)·σ_+(x;ir)+(y−mu_ir)·σ_+(y;ir)) + −B·((x−mu_ir)·σ_-(x;ir)+(y−mu_ir)·σ_-(y;ir))] + + The base-kernel gradients are windowed by the RBF transition region: + dK/d(rbf_i) = σ_+(x;ir)·σ_+(y;ir)·(dK_rbf/dtheta_i)·σ_-(x;uv)·σ_-(y;uv) + + When only one asymptotic region is active the absent softthetas reduce to + 1 (sign=0 path in softtheta), so the formulas remain valid throughout. + """ + grads = [] + + # ---- base-kernel parameter gradients -------------------------------- + base_grads = self.kernel.params_gradient() + for bg in base_grads: + def _wrap(x, y, _bg=bg): + x1 = make_column_vector(x[:, 0]) + y1 = make_row_vector(y[:, 0]) + tau_x = softtheta(x1, self.mu_ir, self.l_ir, -self.ir) # σ_+(x;ir) + tau_y = softtheta(y1, self.mu_ir, self.l_ir, -self.ir) + sig_x = softtheta(x1, self.mu_uv, self.l_uv, -self.uv) # σ_-(x;uv) + sig_y = softtheta(y1, self.mu_uv, self.l_uv, -self.uv) + return tau_x * tau_y * _bg(x, y) * sig_x * sig_y + grads.append(_wrap) + + # ---- UV asymptotic parameter gradients: mu_uv, l_uv ---------------- + if self.uv: + def _dK_dmu_uv(x, y): + x1 = make_column_vector(x[:, 0]) + y1 = make_row_vector(y[:, 0]) + tau_x = softtheta(x1, self.mu_ir, self.l_ir, -self.ir) + tau_y = softtheta(y1, self.mu_ir, self.l_ir, -self.ir) + sp_x = softtheta(x1, self.mu_uv, self.l_uv, self.uv) # σ_+(x;uv) + sp_y = softtheta(y1, self.mu_uv, self.l_uv, self.uv) + sm_x = softtheta(x1, self.mu_uv, self.l_uv, -self.uv) # σ_-(x;uv) + sm_y = softtheta(y1, self.mu_uv, self.l_uv, -self.uv) + K_rbf = self.kernel(x, y) + f_x = self.uv_asymptotics(x1) + f_y = self.uv_asymptotics(y1) + return (1.0 / self.l_uv) * ( + tau_x * tau_y * K_rbf * sm_x * sm_y * (sp_x + sp_y) + - sp_x * sp_y * f_x * f_y * (sm_x + sm_y) + ) + + def _dK_dl_uv(x, y): + x1 = make_column_vector(x[:, 0]) + y1 = make_row_vector(y[:, 0]) + tau_x = softtheta(x1, self.mu_ir, self.l_ir, -self.ir) + tau_y = softtheta(y1, self.mu_ir, self.l_ir, -self.ir) + sp_x = softtheta(x1, self.mu_uv, self.l_uv, self.uv) + sp_y = softtheta(y1, self.mu_uv, self.l_uv, self.uv) + sm_x = softtheta(x1, self.mu_uv, self.l_uv, -self.uv) + sm_y = softtheta(y1, self.mu_uv, self.l_uv, -self.uv) + K_rbf = self.kernel(x, y) + f_x = self.uv_asymptotics(x1) + f_y = self.uv_asymptotics(y1) + dx = x1 - self.mu_uv + dy = y1 - self.mu_uv + return (1.0 / self.l_uv**2) * ( + tau_x * tau_y * K_rbf * sm_x * sm_y * (dx * sp_x + dy * sp_y) + - sp_x * sp_y * f_x * f_y * (dx * sm_x + dy * sm_y) + ) + + grads.extend([_dK_dmu_uv, _dK_dl_uv]) + + # ---- IR asymptotic parameter gradients: mu_ir, l_ir ---------------- + if self.ir: + def _dK_dmu_ir(x, y): + x1 = make_column_vector(x[:, 0]) + y1 = make_row_vector(y[:, 0]) + tp_x = softtheta(x1, self.mu_ir, self.l_ir, -self.ir) # σ_+(x;ir) + tp_y = softtheta(y1, self.mu_ir, self.l_ir, -self.ir) + tm_x = softtheta(x1, self.mu_ir, self.l_ir, self.ir) # σ_-(x;ir) + tm_y = softtheta(y1, self.mu_ir, self.l_ir, self.ir) + sm_x = softtheta(x1, self.mu_uv, self.l_uv, -self.uv) + sm_y = softtheta(y1, self.mu_uv, self.l_uv, -self.uv) + K_rbf = self.kernel(x, y) + f_x = self.ir_asymptotics(x1) + f_y = self.ir_asymptotics(y1) + return (1.0 / self.l_ir) * ( + tm_x * tm_y * f_x * f_y * (tp_x + tp_y) + - tp_x * tp_y * K_rbf * sm_x * sm_y * (tm_x + tm_y) + ) + + def _dK_dl_ir(x, y): + x1 = make_column_vector(x[:, 0]) + y1 = make_row_vector(y[:, 0]) + tp_x = softtheta(x1, self.mu_ir, self.l_ir, -self.ir) + tp_y = softtheta(y1, self.mu_ir, self.l_ir, -self.ir) + tm_x = softtheta(x1, self.mu_ir, self.l_ir, self.ir) + tm_y = softtheta(y1, self.mu_ir, self.l_ir, self.ir) + sm_x = softtheta(x1, self.mu_uv, self.l_uv, -self.uv) + sm_y = softtheta(y1, self.mu_uv, self.l_uv, -self.uv) + K_rbf = self.kernel(x, y) + f_x = self.ir_asymptotics(x1) + f_y = self.ir_asymptotics(y1) + dx = x1 - self.mu_ir + dy = y1 - self.mu_ir + return (1.0 / self.l_ir**2) * ( + tm_x * tm_y * f_x * f_y * (dx * tp_x + dy * tp_y) + - tp_x * tp_y * K_rbf * sm_x * sm_y * (dx * tm_x + dy * tm_y) + ) + + grads.extend([_dK_dmu_ir, _dK_dl_ir]) + + return grads + class Matern12(Kernel): """Matern-1/2 kernel diff --git a/fredipy/models.py b/fredipy/models.py index 12c04b6..7d1f9aa 100644 --- a/fredipy/models.py +++ b/fredipy/models.py @@ -106,7 +106,7 @@ def predict( if full_cov: return mu, C else: - return mu, np.sqrt(np.diag(C)) + return mu, np.sqrt(np.clip(np.diag(C), a_min=0.0, a_max=None)) def predict_data( self, @@ -142,7 +142,7 @@ def predict_data( if full_cov: return mu, C else: - return mu, np.sqrt(np.diag(C)) + return mu, np.sqrt(np.clip(np.diag(C), a_min=0.0, a_max=None)) def log_likelihood(self) -> float: """Returns the log-likelihood of the posterior GP.""" @@ -175,7 +175,10 @@ def log_likelihood_grad(self) -> List[float]: kernel_grad = self.kernel.params_gradient() for i in range(self.kernel.dim): - OpKer_grad = self.OpKerOp(kernel_grad[i], self.constraints) + # derivative=True tells the covariance machinery that kernel_grad[i] is dK/dtheta_i, + # not K: blocks assembled with terms that do not depend on the kernel (the constant + # analytic UV tail moment T) must drop those terms rather than reproduce them. + OpKer_grad = self.OpKerOp(kernel_grad[i], self.constraints, derivative=True) loglik_grad[i] = 0.5 * np.trace((alpha @ alpha.T - K_inv) @ OpKer_grad) return loglik_grad diff --git a/fredipy/util.py b/fredipy/util.py index c7c4d3a..addce83 100644 --- a/fredipy/util.py +++ b/fredipy/util.py @@ -53,7 +53,11 @@ def softtheta( return 1/(np.exp(-1*(x - mu0)/l0) + 1) elif sign < 0: - return np.exp(-1*(x - mu0)/l0)/(np.exp(-1*(x - mu0)/l0) + 1) + # Written as the exact algebraic complement of the sign>0 branch above, rather than as + # exp(...)/(exp(...)+1) directly: for x far below mu0 (deep cutoff region), the direct + # form divides inf by inf and returns nan, even though the correct limit is exactly 1. + # The complement form only ever evaluates the already-overflow-safe sign>0 branch. + return 1 - 1/(np.exp(-1*(x - mu0)/l0) + 1) else: return np.ones_like(x) diff --git a/tests/test_integrators_uvtail.py b/tests/test_integrators_uvtail.py new file mode 100644 index 0000000..74b9bb2 --- /dev/null +++ b/tests/test_integrators_uvtail.py @@ -0,0 +1,495 @@ +"""Coverage for ``GaussLegendre_1D_log``'s analytic UV tail correction. + +The correction implements + + Sigma_12 = BB + A_1 T_2^T + T_1 A_2^T + T_1 T_2^T (**) + +for a pair of integral constraints, where BB is the plain double quadrature, T_i are the +per-constraint, per-row analytic tail moments above the UV split point w_uv, and A_i are the +bulk-UV overlaps evaluated through the anchor K(w, w_uv) / f_uv(w_uv). + +The tests below pin down, in order: that (**) makes the cross block consistent with the +prediction operator (T1), that the assembled covariance stays symmetric (T2), that the legacy +symmetric closed form is unchanged (T3), that the sum-rule tail moment is the right analytic +value (T4), that the end-to-end posterior is self-consistent (T5), that plain quadrature models +are bit-identical to their pre-fix values (T6), and the two error/shape guards (T7, T8). +""" + +import numpy as np +import pytest + +from fredipy import constraints, integrators, kernels, models, operators +from fredipy.covariance import OneSided, TwoSided +from fredipy.util import make_column_vector, make_row_vector + + +# -------------------------------------------------------------------------------------------- +# Synthetic, deliberately well-conditioned model. +# +# Tolerances in an end-to-end test are limited by cond(OpKerOp + cov_y); the production gluon and +# ghost models run at data_cov = 1e-10 where cond ~ 7e12 and no identity can be checked below +# ~1e-5. These settings keep the model far from that regime so the tests measure the correction, +# not the conditioning. +# -------------------------------------------------------------------------------------------- + +W_MIN = 1e-3 # Lower end of the omega quadrature range. +W_UV = 1e5 # Upper end of the quadrature range == UV split point w_uv. +INT_N = 120 # Number of Gauss-Legendre nodes on (W_MIN, W_UV). +DATA_COV = 1e-6 # Diagonal data covariance; large enough to keep the model well conditioned. +SR_COV = 1e-8 # Sum-rule target covariance. +N_DATA = 15 # Number of synthetic momentum points. +P_MIN, P_MAX = 0.5, 20.0 # Momentum window of the synthetic data, in the same units as omega. + +RBF_VARIANCE = 2.0 # Base RBF variance of the synthetic kernel. +RBF_LENGTHSCALE = 1.5 # Base RBF lengthscale of the synthetic kernel. +MU_UV = 1.1 # Centre of the UV blend-in; must satisfy W_UV >> MU_UV for (*) to hold. +L_UV = 0.15 # Width of the UV blend-in. + +GAMMA_GLUON = 13.0 / 22.0 # Gluon anomalous dimension (the exponent formerly hardcoded). +GAMMA_GHOST = 9.0 / 44.0 # Ghost anomalous dimension; differs, which is why it is a parameter. + + +def propagator_uv_shape(w, gamma=GAMMA_GLUON): + """Propagator-type UV asymptotic shape f_uv(w) = w^-2 (2 log w)^-(1+gamma). + + Uses the same smooth-maximum regularisation as the reconstructions repo's ``uv_asymptotics`` + so that the log stays finite below w = 1; the regularisation is irrelevant at w >= w_uv. + """ + delta = 0.05 # Smoothing width of the regularised maximum. + w1 = 0.5 * (w + 1.15 + np.sqrt((w - 1.15) ** 2 + delta ** 2)) + return 1.0 / (w1 ** 2 * np.log(w1 ** 2) ** (1.0 + gamma)) + + +def ghost_uv_shape(w): + """Propagator-type UV shape at the ghost anomalous dimension.""" + return propagator_uv_shape(w, gamma=GAMMA_GHOST) + + +def kl_p2(p, w): + """Kallen-Lehmann data kernel multiplied by p^2.""" + return p ** 2 * w / (w ** 2 + p ** 2) / np.pi + + +def sr_kernel(p, w): + """Sum-rule kernel C(p, w) = w, broadcast over the (ignored) momentum axis.""" + return w * np.ones_like(p) + + +def analytic_sum_rule_tail(w_uv, gamma): + """Closed form (2 log w_uv)^-gamma / (2 gamma) of the sum-rule tail moment.""" + return (2.0 * np.log(w_uv)) ** (-gamma) / (2.0 * gamma) + + +def sum_rule_tail_moment(w_uv, gamma): + """Row-independent ``tail_moment`` callable for the sum-rule kernel C(p, w) = w.""" + value = analytic_sum_rule_tail(w_uv, gamma) + + def _tail_moment(x): + return np.full((np.atleast_2d(x).shape[0], 1), value) + + return _tail_moment + + +def make_kernel(uv_shape=propagator_uv_shape): + """AsymptoticKernel with a UV branch only (no IR branch) at the synthetic hyperparameters.""" + rbf = kernels.RadialBasisFunction(RBF_VARIANCE, RBF_LENGTHSCALE) + kernel = kernels.AsymptoticKernel(rbf) + kernel.add_asymptotics(region="UV", asymptotics=uv_shape) + kernel.set_params(asymp_params=[MU_UV, L_UV]) + return kernel + + +def synthetic_data(): + """Momenta and correlator values of the synthetic data set.""" + p = np.geomspace(P_MIN, P_MAX, N_DATA) + return p, p ** 2 / (p ** 2 + 1.0) + + +def build_model(with_tail=True, gamma=GAMMA_GLUON, uv_shape=propagator_uv_shape): + """Two-constraint model: a KL data constraint plus a one-row sum rule. + + ``with_tail=False`` gives both constraints a plain ``GaussLegendre_1D_log`` on the same grid, + which is what the tgvert/ghgvert/fgvertsd projects do and what T6 pins down. + """ + p, G = synthetic_data() + kernel = make_kernel(uv_shape) + + data_integrator = integrators.GaussLegendre_1D_log(W_MIN, W_UV, INT_N) + c_data = constraints.LinearEquality( + operators.Integral(kl_p2, data_integrator), + {"x": p, "y": G, "cov_y": DATA_COV * np.ones_like(G)}, + ) + + if with_tail: + sr_integrator = integrators.GaussLegendre_1D_log_UVtail( + W_MIN, W_UV, INT_N, uv_shape, tail_moment=sum_rule_tail_moment(W_UV, gamma) + ) + else: + sr_integrator = integrators.GaussLegendre_1D_log(W_MIN, W_UV, INT_N) + c_sr = constraints.LinearEquality( + operators.Integral(sr_kernel, sr_integrator), + {"x": np.array([0.0]), "y": np.array([0.0]), "cov_y": np.array([SR_COV])}, + ) + + model = models.GaussianProcess(kernel, [c_data, c_sr]) + return model, c_data, c_sr, data_integrator, sr_integrator + + +def bulk_double_integration(integrator, c1, kernel, c2): + """Uncorrected double quadrature BB, written out so the tests do not lean on the fix.""" + row1 = integrator.weights * c1(make_row_vector(integrator.w), x=make_column_vector(c1.x)) + row2 = integrator.weights * c2(make_row_vector(integrator.w), x=make_column_vector(c2.x)) + return row1 @ kernel(integrator.w, integrator.w) @ row2.T + + +# -------------------------------------------------------------------------------------------- +# T1 -- the sharpest test: pure linear algebra, no linear solve, so it must hold to machine +# precision. The cross block of OpKerOp is, by definition, the data-side reintegration of the +# sum-rule row of the prediction operator OpKer. That identity is exactly what the missing +# cross-block correction used to break. +# -------------------------------------------------------------------------------------------- + +def test_cross_block_equals_reintegrated_opker_row(): + model, c_data, c_sr, data_integrator, _ = build_model() + kernel = model.kernel + n_data = c_data.x.shape[0] + + opkerop = TwoSided()(kernel, [c_data, c_sr]) + cross_block = opkerop[:n_data, n_data:] # (M, 1) + + nodes = data_integrator.w + opker = OneSided()(kernel, [c_data, c_sr], nodes) # (M + 1, N) + sr_row = opker[n_data:, :] # (1, N) + + data_weighted = data_integrator.weights * kl_p2( + make_column_vector(c_data.x), make_row_vector(nodes) + ) + reintegrated = data_weighted @ sr_row.T # (M, 1) + + np.testing.assert_allclose(cross_block, reintegrated, rtol=1e-12) + + +# -------------------------------------------------------------------------------------------- +# T2 -- np.linalg.cholesky in models.py reads only the lower triangle, so an asymmetric OpKerOp +# would never raise. This is the only guard against that. +# -------------------------------------------------------------------------------------------- + +def test_covariance_matrix_is_symmetric(): + model, c_data, c_sr, _, _ = build_model() + kernel = model.kernel + n_data = c_data.x.shape[0] + + opkerop = TwoSided()(kernel, [c_data, c_sr]) + scale = np.abs(opkerop).max() + np.testing.assert_allclose(opkerop, opkerop.T, rtol=1e-14, atol=1e-14 * scale) + + # Building the same model with the constraint order reversed must give the permuted matrix, + # i.e. block (sr, data) must carry exactly the same correction as block (data, sr). + reversed_opkerop = TwoSided()(kernel, [c_sr, c_data]) + perm = np.concatenate([np.arange(n_data, n_data + 1), np.arange(n_data)]) + np.testing.assert_allclose( + reversed_opkerop, opkerop[np.ix_(perm, perm)], rtol=1e-14, atol=1e-14 * scale + ) + + +# -------------------------------------------------------------------------------------------- +# T3 -- the UVtail subclass no longer overrides doubleIntegrationSymmetric. This locks in that +# the inherited path reproduces the deleted override's closed form bulk + 2 A T + T^2 for a +# multi-row constraint with a row-independent tail moment. +# -------------------------------------------------------------------------------------------- + +def test_symmetric_block_matches_legacy_closed_form(): + kernel = make_kernel() + tail_value = analytic_sum_rule_tail(W_UV, GAMMA_GLUON) + integrator = integrators.GaussLegendre_1D_log_UVtail( + W_MIN, W_UV, INT_N, propagator_uv_shape, + tail_moment=sum_rule_tail_moment(W_UV, GAMMA_GLUON), + ) + p = np.array([0.7, 3.0, 11.0]) # Three rows, so a broadcast-vs-outer-product bug would show. + constraint = constraints.LinearEquality( + operators.Integral(kl_p2, integrator), + {"x": p, "y": np.zeros_like(p), "cov_y": DATA_COV * np.ones_like(p)}, + ) + + symmetric = integrator.doubleIntegrationSymmetric(constraint, kernel) + general = integrator.doubleIntegration(constraint, kernel, constraint) + np.testing.assert_allclose(symmetric, general, rtol=1e-12) + + bulk = bulk_double_integration(integrator, constraint, kernel, constraint) + c_row = constraint(make_row_vector(integrator.w), x=make_column_vector(constraint.x)) + A = (integrator.weights * c_row @ kernel(integrator.w, integrator.w_uv)) / integrator.f_uv_anchor + ones_col = np.ones((A.shape[0], 1)) + legacy = bulk + tail_value * (A @ ones_col.T + ones_col @ A.T) + tail_value ** 2 + + np.testing.assert_allclose(symmetric, legacy, rtol=1e-12) + + +# -------------------------------------------------------------------------------------------- +# T4 -- the analytic value itself, for both anomalous dimensions, plus a quadrature cross-check +# that extending the numerical range towards infinity really does converge onto +# bulk(w_uv) + T_sr. Convergence in log(w) is slow, hence the loose few-percent tolerance. +# -------------------------------------------------------------------------------------------- + +@pytest.mark.parametrize("gamma, uv_shape, max_truncation_error", [ + # max_truncation_error: how far the integral truncated at 10^20 still is from bulk + T. The + # tail decays only like (2 log w)^-gamma, so the smaller gamma the slower the convergence -- + # ~1% for the gluon, ~18% for the ghost. That slowness is exactly why the analytic tail is + # mandatory and extending the quadrature range is not a viable alternative. + (GAMMA_GLUON, propagator_uv_shape, 0.05), + (GAMMA_GHOST, ghost_uv_shape, 0.20), +]) +def test_tail_moment_matches_analytic_sum_rule_value(gamma, uv_shape, max_truncation_error): + integrator = integrators.GaussLegendre_1D_log_UVtail( + W_MIN, W_UV, INT_N, uv_shape, tail_moment=sum_rule_tail_moment(W_UV, gamma) + ) + constraint = constraints.LinearEquality( + operators.Integral(sr_kernel, integrator), + {"x": np.array([0.0]), "y": np.array([0.0]), "cov_y": np.array([SR_COV])}, + ) + tail = integrator.uv_tail_moment(constraint) + assert tail.shape == (1, 1) + np.testing.assert_allclose(tail.item(), analytic_sum_rule_tail(W_UV, gamma), rtol=1e-14) + + # Cross-check: sum_rule integral over (W_MIN, 10^k) -> integral over (W_MIN, w_uv) + T. + # n_nodes=2000: the two truncations use different grids, so their shared bulk parts only + # cancel to quadrature accuracy; 2000 nodes brings that below 4e-6 relative (800 leaves 2e-2). + def numeric_sum_rule_integral(w_max, n_nodes=2000): + integ = integrators.GaussLegendre_1D_log(W_MIN, w_max, n_nodes) + return float((integ.weights @ (integ.w * uv_shape(integ.w))).item()) + + target = numeric_sum_rule_integral(W_UV) + analytic_sum_rule_tail(W_UV, gamma) + shortfalls = [target - numeric_sum_rule_integral(10.0 ** k) for k in (6, 8, 12, 20)] + + # Whatever the truncated quadrature misses must be precisely the closed form's own remainder + # above the truncation point -- the direct check of the closed form against quadrature of the + # actual (regularised) f_uv, independent of how slowly the tail converges. + for k, shortfall in zip((6, 8, 12, 20), shortfalls): + np.testing.assert_allclose(shortfall, analytic_sum_rule_tail(10.0 ** k, gamma), rtol=1e-5) + + errors = [abs(s) / abs(target) for s in shortfalls] + assert errors == sorted(errors, reverse=True), \ + f"truncated integral should approach bulk + T monotonically, got {errors}" + assert errors[-1] < max_truncation_error, \ + f"k=20 still {errors[-1]:.3%} away from bulk + T" + + +# -------------------------------------------------------------------------------------------- +# T5 -- end-to-end self-consistency of the posterior. predict_data() and predict() must describe +# the same model: reintegrating the posterior mean over fredipy's own nodes and weights has to +# reproduce predict_data(). This is the identity the bug broke by ~1e0. +# -------------------------------------------------------------------------------------------- + +def test_predict_data_matches_reintegrated_posterior_mean(): + model, c_data, c_sr, data_integrator, sr_integrator = build_model() + n_data = c_data.x.shape[0] + + nodes = data_integrator.w + rho, _ = model.predict(nodes) # (N, 1) + mu, _ = model.predict_data() # (M + 1, 1) + + data_weighted = data_integrator.weights * kl_p2( + make_column_vector(c_data.x), make_row_vector(nodes) + ) + np.testing.assert_allclose(data_weighted @ rho, mu[:n_data], rtol=1e-9) + + # Sum-rule row: the reintegration only covers the bulk, so the analytic tail T * A_uv has to + # be added back, with A_uv = rho(w_uv) / f_uv(w_uv) the UV amplitude of the posterior mean. + sr_weighted = sr_integrator.weights * sr_kernel( + make_column_vector(c_sr.x), make_row_vector(nodes) + ) + rho_at_w_uv, _ = model.predict(sr_integrator.w_uv) + tail_value = sr_integrator.uv_tail_moment(c_sr).item() + reintegrated_sr = (sr_weighted @ rho).item() + reintegrated_sr += tail_value * rho_at_w_uv.item() / sr_integrator.f_uv_anchor + + np.testing.assert_allclose(reintegrated_sr, mu[n_data:].item(), rtol=1e-4) + + +# -------------------------------------------------------------------------------------------- +# T6 -- reference values captured with the pre-fix code, before the tail hook existed. A model +# whose constraints all carry plain GaussLegendre_1D_log integrators must stay bit-identical: +# this is what protects three-gluon-vert, ghost-gluon-vert and fgvert-sd from the change. +# -------------------------------------------------------------------------------------------- + +REFERENCE_LOG_LIKELIHOOD = -278560.62089715304 +REFERENCE_PREDICT_DATA = np.array([ + 0.11838105684728362, 0.2086624375442625, 0.35474911867640913, 0.5465432399942074, + 0.75084480180521496, 0.92431427465635352, 1.034895190037787, 1.0745688577007968, + 1.0557558063883334, 0.99917723311227746, 0.92389010410988703, 0.8432306828617584, + 0.76490387672674842, 0.69258282921509817, 0.62751788803143427, 0.0017991249333135784, +]) +REFERENCE_PREDICT_MEAN = np.array([ + 32.848601571109612, -11.554450084397104, -0.77481442881980911, + -0.00017162581324248194, -1.0987647334001352e-09, +]) +REFERENCE_PREDICT_GRID = np.array([0.01, 0.5, 2.0, 50.0, 1e4]) # Prediction points of the reference. +REFERENCE_OPKEROP_TRACE = 10.170139984656482 + + +def test_plain_integrators_are_unchanged_by_the_tail_hook(): + model, _, _, _, _ = build_model(with_tail=False) + + np.testing.assert_allclose(model.log_likelihood(), REFERENCE_LOG_LIKELIHOOD, rtol=1e-15) + mu, _ = model.predict_data() + np.testing.assert_allclose(mu.ravel(), REFERENCE_PREDICT_DATA, rtol=1e-15) + mean, _ = model.predict(REFERENCE_PREDICT_GRID) + np.testing.assert_allclose(mean.ravel(), REFERENCE_PREDICT_MEAN, rtol=1e-15) + trace = np.trace(model._posterior_cache["OpKerOp"]) + np.testing.assert_allclose(trace, REFERENCE_OPKEROP_TRACE, rtol=1e-15) + + +# -------------------------------------------------------------------------------------------- +# T7 -- (**) mixes both constraints on a shared quadrature grid, so it is undefined when the two +# integrators discretise omega differently. Fail loudly rather than silently mixing grids. +# -------------------------------------------------------------------------------------------- + +def test_mismatched_integrator_grids_raise(): + kernel = make_kernel() + p, G = synthetic_data() + + data_integrator = integrators.GaussLegendre_1D_log(W_MIN, 1e4, INT_N + 30) # Different grid. + c_data = constraints.LinearEquality( + operators.Integral(kl_p2, data_integrator), + {"x": p, "y": G, "cov_y": DATA_COV * np.ones_like(G)}, + ) + sr_integrator = integrators.GaussLegendre_1D_log_UVtail( + W_MIN, W_UV, INT_N, propagator_uv_shape, + tail_moment=sum_rule_tail_moment(W_UV, GAMMA_GLUON), + ) + c_sr = constraints.LinearEquality( + operators.Integral(sr_kernel, sr_integrator), + {"x": np.array([0.0]), "y": np.array([0.0]), "cov_y": np.array([SR_COV])}, + ) + + with pytest.raises(NotImplementedError) as excinfo: + TwoSided()(kernel, [c_data, c_sr]) + message = str(excinfo.value) + assert "10000.0" in message and "100000.0" in message, message + assert str(INT_N) in message and str(INT_N + 30) in message, message + + +# -------------------------------------------------------------------------------------------- +# T8 -- singleIntegration's tail is T[m] * K(w_uv, w_pred) / f_uv(w_uv), an outer product. The +# pre-fix code broadcast a single scalar, which happened to be right only because the sum-rule +# constraint has exactly one row. +# -------------------------------------------------------------------------------------------- + +def test_single_integration_tail_is_row_wise_outer_product(): + kernel = make_kernel() + row_moments = np.array([[0.25], [0.5], [1.75]]) # Deliberately row-dependent tail moments. + integrator = integrators.GaussLegendre_1D_log_UVtail( + W_MIN, W_UV, INT_N, propagator_uv_shape, + tail_moment=lambda x: row_moments.copy(), + ) + p = np.array([0.7, 3.0, 11.0]) + constraint = constraints.LinearEquality( + operators.Integral(kl_p2, integrator), + {"x": p, "y": np.zeros_like(p), "cov_y": DATA_COV * np.ones_like(p)}, + ) + w_pred = make_column_vector(np.geomspace(0.05, 500.0, 7)) + + full = integrator.singleIntegration(constraint, kernel, w_pred) + assert full.shape == (3, w_pred.shape[0]) + + bulk = ( + integrator.weights * constraint(make_row_vector(integrator.w), x=make_column_vector(p)) + @ kernel(integrator.w, w_pred) + ) + anchor_row = kernel(integrator.w_uv, w_pred) / integrator.f_uv_anchor # (1, N_pred) + for m in range(3): + expected = bulk[m] + row_moments[m, 0] * anchor_row.ravel() + np.testing.assert_allclose(full[m], expected, rtol=1e-12) + + +# -------------------------------------------------------------------------------------------- +# T9 -- the analytic NLL gradient against finite differences of the NLL itself. +# +# log_likelihood_grad() reassembles OpKerOp with dK/dtheta substituted for K. The A terms of +# (**) are linear in the kernel and differentiate themselves correctly, but T_1 T_2^T contains +# no kernel at all -- T is the analytic tail moment, a function of w_uv and gamma only -- so its +# hyperparameter derivative is zero. Adding it unchanged (which is what the code did before the +# derivative flag existed) put the gradient off by factors of 34 to 1675 on this very model, with +# the wrong sign on the RBF lengthscale. This test is the acceptance criterion for that fix: no +# other test in either suite can see a wrong gradient, since the likelihood *value* is correct +# either way. +# -------------------------------------------------------------------------------------------- + +FD_REL_STEP = 1e-4 # Relative central-difference step; near the optimum of truncation vs roundoff. +FD_RTOL = 1e-5 # Agreement demanded per component; the FD truncation floor here is ~5e-7. + +# Hyperparameter vector [rbf_variance, rbf_lengthscale, mu_uv, l_uv] at which the check is run. +FD_PARAMS = [RBF_VARIANCE, RBF_LENGTHSCALE, MU_UV, L_UV] + + +def _set_flat_params(model, params): + """Push a flat [rbf_variance, rbf_lengthscale, mu_uv, l_uv] vector into the model.""" + model.reset() + model.kernel.set_params(kernel_params=list(params[:2]), asymp_params=list(params[2:])) + + +def _finite_difference_gradient(model, params, rel_step=FD_REL_STEP): + """Central-difference gradient of log_likelihood() w.r.t. every entry of ``params``.""" + grad = [] + for i in range(len(params)): + h = rel_step * abs(params[i]) + plus = list(params) + plus[i] += h + _set_flat_params(model, plus) + f_plus = model.log_likelihood() + minus = list(params) + minus[i] -= h + _set_flat_params(model, minus) + f_minus = model.log_likelihood() + grad.append((f_plus - f_minus) / (2.0 * h)) + _set_flat_params(model, params) + return np.array(grad) + + +@pytest.mark.parametrize("with_tail", [True, False]) +def test_log_likelihood_grad_matches_finite_differences(with_tail): + """Every active hyperparameter, with and without the UVtail sum-rule constraint.""" + model, _, _, _, _ = build_model(with_tail=with_tail) + assert model.kernel.dim == len(FD_PARAMS) + + _set_flat_params(model, FD_PARAMS) + analytic = np.array(model.log_likelihood_grad(), dtype=float) + numeric = _finite_difference_gradient(model, FD_PARAMS) + + np.testing.assert_allclose(analytic, numeric, rtol=FD_RTOL) + + +def test_gradient_flag_is_a_noop_for_plain_integrators(): + """The derivative flag must not perturb the plain-quadrature gradient path at all. + + Plain ``GaussLegendre_1D_log`` blocks are linear in the kernel, so the assembled gradient + covariance has to be bit-identical with and without the flag. This is what guarantees that + three-gluon-vert, ghost-gluon-vert and fgvert-sd optimize against exactly the gradients they + did before the flag was introduced. + """ + model, _, _, _, _ = build_model(with_tail=False) + _set_flat_params(model, FD_PARAMS) + for kernel_grad in model.kernel.params_gradient(): + with_flag = TwoSided()(kernel_grad, model.constraints, derivative=True) + without_flag = TwoSided()(kernel_grad, model.constraints, derivative=False) + assert np.array_equal(with_flag, without_flag) + + +def test_gradient_flag_drops_only_the_constant_tail_term(): + """With a UVtail constraint the flag must remove exactly T_1 T_2^T and nothing else.""" + model, c_data, c_sr, _, sr_integrator = build_model(with_tail=True) + _set_flat_params(model, FD_PARAMS) + n_data = c_data.x.shape[0] + tail = sr_integrator.uv_tail_moment(c_sr) # (1, 1) + + for kernel_grad in model.kernel.params_gradient(): + with_flag = TwoSided()(kernel_grad, model.constraints, derivative=True) + without_flag = TwoSided()(kernel_grad, model.constraints, derivative=False) + difference = without_flag - with_flag + + # Only the sum-rule diagonal block differs, by exactly T^2: the data constraint carries a + # plain integrator, hence a zero tail moment, hence no constant term in any block it enters. + expected = np.zeros_like(difference) + expected[n_data:, n_data:] = tail @ tail.T + # rtol 1e-12: the sum-rule block is a difference of two O(1) sums, so it carries a few + # ulp of cancellation. atol 0: every other block must be bitwise unchanged by the flag. + np.testing.assert_allclose(difference, expected, rtol=1e-12, atol=0.0) diff --git a/tests/test_util.py b/tests/test_util.py index b7af1c9..c2d377c 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -27,6 +27,19 @@ def test_softtheta(): assert np.allclose(result, expected) +def test_softtheta_sign_negative_does_not_overflow_to_nan(): + # Regression test: x deep below mu0 relative to l0 (e.g. deep-IR kernel evaluations with a + # narrow l_ir) previously overflowed exp() in the sign<0 branch's direct exp/(exp+1) form, + # producing inf/inf = nan instead of the correct limit of 1.0. + x = np.array([0.0, 1e-3]) + mu0, l0 = 1.0, 1e-3 # (mu0 - x)/l0 >> 709 for both x values: exp() overflows + + result = softtheta(x, mu0, l0, -1) + + assert np.all(np.isfinite(result)) + assert np.allclose(result, 1.0) + + def test_allclose(): a = np.array([1, 2, 3]) b = np.array([1, 2, 3])