# Copyright 2026 Pavlo Pelikh
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""Quadratic separable spectral regularizer for regularized SDPs.
For a primal matrix :math:`X \in \operatorname{dom}(\mathcal{A})`, the
regularizer contributes
.. math::
R_\varepsilon(X)
= \varepsilon \operatorname{Tr}[\varphi(X)],
\qquad
\varphi(t) = \frac{t^2}{2} + \iota_{[0,\infty)}(t).
Its Legendre transform is applied spectrally to the scaled dual slack
.. math::
\frac{\mathcal{A}^\dagger y - C}{\varepsilon},
\qquad y \in \operatorname{cod}(\mathcal{A}).
"""
from spacecore import DenseArray, jax_pytree_class
from ._base import NEG_EIG_TOL, Regularizer
[docs]
@jax_pytree_class
class QuadraticReg(Regularizer):
r"""Quadratic spectral regularizer on nonnegative spectra.
The scalar convex function is
.. math::
\varphi(t) = \frac{t^2}{2} + \iota_{[0,\infty)}(t).
The indicator term is zero for :math:`t \ge 0` and :math:`+\infty` for
:math:`t < 0`. Since the primal constraint is :math:`X \succeq 0`, this
defines the separable spectral penalty
.. math::
R_\varepsilon(X)
= \varepsilon \operatorname{Tr}[\varphi(X)]
= \varepsilon \sum_i \frac{\lambda_i(X)^2}{2}.
The Legendre transform of :math:`\varphi` restricted to
nonnegative primal eigenvalues is
.. math::
\psi(s) = \frac{\max\{s, 0\}^2}{2}.
If :math:`s_i` are the eigenvalues of
:math:`\mathcal{A}^\dagger y - C`, then ``primal_from_dual`` uses
.. math::
\lambda_i(X) = \psi'(s_i / \varepsilon)
= \max\{s_i / \varepsilon, 0\}.
In plain language: the quadratic regularizer clips negative scaled slack
eigenvalues to zero and keeps positive scaled slack eigenvalues linearly.
"""
[docs]
def phi(self, x: DenseArray) -> DenseArray:
r"""Return :math:`x^2 / 2 + \iota_{[0,\infty)}(x)` elementwise.
Round-off-negative eigenvalues (down to ``-NEG_EIG_TOL``) evaluate at
the limit :math:`\varphi(0)=0` rather than out of domain.
"""
ops = self.ctx.ops
safe = ops.maximum(x, 0.)
return ops.where(x >= -NEG_EIG_TOL, safe ** 2 / 2, float("inf"))
[docs]
def phi_star(self, x: DenseArray) -> DenseArray:
r"""Return :math:`\psi(x) = \max\{x, 0\}^2 / 2` elementwise."""
ops = self.ctx.ops
return ops.maximum(x, 0.) ** 2 / 2
[docs]
def phi_star_prime(self, x: DenseArray) -> DenseArray:
r"""Return :math:`\psi'(x) = \max\{x, 0\}` elementwise."""
ops = self.ctx.ops
return ops.maximum(x, 0.)
[docs]
def log_phi_star_prime(self, x: DenseArray) -> DenseArray:
r"""Return :math:`\log(\psi'(x))` elementwise.
For the quadratic regularizer this is :math:`\log(\max\{x, 0\})`,
with :math:`-\infty` on nonpositive entries.
"""
ops = self.ctx.ops
safe = ops.where(x > 0., x, 1.)
return ops.where(x > 0., ops.log(safe), float("-inf"))
# ---- fixed-trace conjugate ---------------------------------------------
#
# The trace multiplier enters additively, X = max((S - theta)/eps, 0) with
# sum_i max((s_i - theta)/eps, 0) = 1 -- the Euclidean projection of the
# spectrum onto the simplex, i.e. sparsemax. Unlike a general psi this needs
# no root find: the active set is a prefix of the descending spectrum, so
# one sort fixes theta in closed form. The base class's softmax rescaling
# would instead keep every ratio and every zero fixed, which is not the
# constrained maximizer, so it is overridden here to keep the value and the
# gradient generated by the *same* theta.
def _simplex_threshold(self, scaled: DenseArray) -> DenseArray:
r"""Return the scaled chemical potential :math:`\tau = \theta/\varepsilon`.
Solves :math:`\sum_i \max(u_i - \tau, 0) = 1` over the whole (structured)
spectrum ``u``. With ``u`` sorted descending and
:math:`\tau_k = (\sum_{i\le k} u_i - 1)/k`, the active set
:math:`\{k : u_k > \tau_k\}` is a prefix, so its size and its sum give
:math:`\tau` directly -- no bracketing and no data-dependent trip count,
which keeps this traceable.
"""
ops = self.ops
u = self.eigval_space.flatten(scaled)
n = int(ops.shape(u)[-1])
rev = ops.asarray(list(range(n - 1, -1, -1)))
desc = ops.take(ops.sort(u), rev) # descending spectrum
# ops has no cumsum; a lower-triangular ones matrix does it in one einsum.
tri = ops.asarray([[1.0 if j <= i else 0.0 for j in range(n)]
for i in range(n)])
css = ops.einsum("ij,j->i", tri, desc)
k_idx = ops.asarray([float(i + 1) for i in range(n)])
active = desc > (css - 1.0) / k_idx
size = ops.sum(ops.where(active, 1.0, 0.0))
total = ops.sum(ops.where(active, desc, 0.0))
return (total - 1.0) / size
def _normalized_legendre(self, scaled: DenseArray, val: float) -> DenseArray:
r"""Return :math:`\theta + \varepsilon\operatorname{Tr}\psi((S-\theta)/\varepsilon)`.
The exact fixed-trace conjugate. Its gradient is
:meth:`_grad_robust_normalization`, which is built from the same
:math:`\tau`, so :meth:`~sdplab.regularization.Regularizer.legendre_and_grad`
returns a true value/gradient pair.
"""
tau = self._simplex_threshold(scaled)
shifted = self.eigval_space.spectral_apply(
scaled, lambda ev: self.phi_star(ev - tau)
)
return (tau + self._trace(shifted)) * val
def _grad_robust_normalization(self, scaled: DenseArray) -> DenseArray:
r"""Unit-trace primal eigenvalues :math:`\max(u_i - \tau, 0)`.
The simplex projection, not a rescaling: entries below the threshold go
to exactly zero, so the recovered primal is genuinely sparse rather than
a full-rank spectrum divided by its trace.
"""
tau = self._simplex_threshold(scaled)
return self.eigval_space.spectral_apply(
scaled, lambda ev: self.phi_star_prime(ev - tau)
)