diff --git a/src/admmsolver/__init__.py b/src/admmsolver/__init__.py index c7e6e82..6a5b851 100644 --- a/src/admmsolver/__init__.py +++ b/src/admmsolver/__init__.py @@ -2,4 +2,4 @@ # SPDX-License-Identifier: MIT __license__ = "MIT" -__version__ = "0.7.6" +__version__ = "0.7.8" diff --git a/src/admmsolver/objectivefunc.py b/src/admmsolver/objectivefunc.py index ec5859d..9261f83 100644 --- a/src/admmsolver/objectivefunc.py +++ b/src/admmsolver/objectivefunc.py @@ -178,10 +178,10 @@ def solve( """ x <- argmin_x alpha * |x|_1 + h^+ x + x^+ h + mu x^+ x - This function works only if all the following conditions are met: - * h and x are real vectors - * mu is a diagonal matrix - Return a real vector. + ``h`` (and hence ``x``) may be real or complex; the proximal operator is + the magnitude soft-threshold, which shrinks |x| by the threshold while + keeping the phase. ``mu`` must be a diagonal matrix. The result is + float64 for real input and complex128 for complex input. """ _assert_types(h, [np.ndarray]) assert isinstance(mu, DiagonalMatrix) or isinstance(mu, ScaledIdentityMatrix) @@ -190,8 +190,9 @@ def solve( raise ValueError("h must not be None!") if mu is None: raise ValueError("mu must not be None!") - if np.iscomplexobj(h): - h = h.real + # The L1 proximal operator (soft-thresholding) works for both real and + # complex-valued x; _softmax below handles the complex case by shrinking + # the magnitude while keeping the phase. return _softmax(-(h/mu.diagonals), 0.5*self._alpha/mu.diagonals) @@ -210,7 +211,7 @@ def __init__(self, alpha: float, A: Union[np.ndarray, MatrixBase]): self._AcA = matmul(A.conjugate().T, A) # B = (A^+ A + mu)^{-1} - self._B_cache = (0, DenseMatrix(np.zeros((1,1)))) + self._B_cache = (0, DenseMatrix(np.zeros((1,1)))) # type: Tuple[int, MatrixBase] def __call__(self, x: np.ndarray): return self._alpha * np.linalg.norm(matmul(self._A, x))**2 @@ -220,7 +221,7 @@ def _get_B(self, mu: MatrixBase): if self._B_cache[0] != hash_: self._B_cache = ( hash_, - inv(add(self._alpha * self._AcA, mu)) + cast(MatrixBase, inv(add(self._alpha * self._AcA, mu))) ) return self._B_cache[1] @@ -277,6 +278,13 @@ class SemiPositiveDefinitePenalty(ObjectiveFunctionBase): 1) Reshape x into a three-way tensor 2) Along a given axis, we compute eigenvalues and remove negative ones (we assume hermition matrices). + + Real-symmetric and complex-Hermitian blocks are both supported. ``solve`` + returns the exact minimizer of ``h^+ x + x^+ h + x^+ mu x`` over the PSD cone + when the penalty ``mu`` is uniform within each block (a ScaledIdentityMatrix, + as produced by the ADMM optimizer). For a genuinely non-uniform diagonal + ``mu`` the returned value is Hermitian, PSD and feasible but only an + approximation of the weighted argmin. """ def __init__(self, shape: Union[Sequence,np.ndarray], axis: int): assert len(shape) == 3 @@ -312,18 +320,55 @@ def solve(self, h : Optional[np.ndarray] = None, mu: Optional[MatrixBase] = None if h is None: h = np.zeros(np.prod(self._shape)) - elif np.iscomplexobj(h): - h = h.real assert diagonals.size == h.size - x_ = (-(h/diagonals)).reshape(self._shape) - x_ = np.moveaxis(x_, self._axis, 0) - for i in range(x_.shape[0]): - evals, evecs = np.linalg.eigh(x_[i,:,:]) + # Minimize h^+ x + x^+ h + x^+ mu x over the (real-symmetric or + # complex-Hermitian) positive semi-definite cone, block by block. + # + # The Hermitian constraint couples the paired entries (i,j) and (j,i) + # (x_ji = conj(x_ij)). Under the diagonal metric mu with per-entry + # weights d, the metric projection onto the Hermitian subspace of the + # unconstrained centre x* = -h/d is the weighted average + # centre_ij = (d_ij x*_ij + d_ji conj(x*_ji)) / (d_ij + d_ji) + # = -(h_ij + conj(h_ji)) / (d_ij + d_ji). + # We divide numerator and denominator of each pair by a common per-pair + # scale before summing, so neither ``h_ij + conj(h_ji)`` nor + # ``d_ij + d_ji`` overflows for very large (or underflows for very small) + # finite operands. The scale cancels in the ratio. + # + # Eigenvalue clipping of this Hermitian centre is the *exact* PSD + # projection when the weights are uniform within each block -- which is + # the case for the scalar ADMM penalty (a ScaledIdentityMatrix mu). For + # a genuinely non-uniform diagonal metric the result is still Hermitian, + # PSD and feasible, but only an approximation of the weighted argmin. + hblk = np.asarray(h).reshape(self._shape) + dblk = np.asarray(diagonals, dtype=np.float64).reshape(self._shape) + hblk = np.moveaxis(hblk, self._axis, 0) + dblk = np.moveaxis(dblk, self._axis, 0) + hblkT = np.conjugate(np.swapaxes(hblk, 1, 2)) # conj(h_ji) + dblkT = np.swapaxes(dblk, 1, 2) # d_ji + # Per-pair scale factor (cancels in the ratio). Use max(|re|,|im|) for + # complex entries so it never overflows the way np.abs(complex) can when + # the components are large but finite. (A single shared scale assumes the + # paired operands are of comparable magnitude, which always holds for the + # O(1)-scaled ADMM penalties and data; it is not designed for adversarial + # blocks mixing sub-normal and near-DBL_MAX entries.) + def _comp_scale(a): + if np.iscomplexobj(a): + return np.maximum(np.abs(a.real), np.abs(a.imag)) + return np.abs(a) + scale = np.maximum.reduce([_comp_scale(hblk), _comp_scale(hblkT), dblk, dblkT]) + scale = np.where(scale > 0.0, scale, 1.0) + hbar = hblk / scale + hblkT / scale # (h_ij + conj(h_ji)) / scale + dbar = dblk / scale + dblkT / scale # (d_ij + d_ji) / scale + centre = -(hbar / dbar) + x_ = np.empty_like(centre) + for i in range(centre.shape[0]): + evals, evecs = np.linalg.eigh(centre[i, :, :]) idx = evals >= 0 U = evecs[:, idx] - x_[i,:,:] = U @ (evals[idx,None] * U.conjugate().T) + x_[i, :, :] = U @ (evals[idx, None] * U.conjugate().T) return np.moveaxis(x_, 0, self._axis).ravel() @@ -334,22 +379,58 @@ def _project_plus(x): def _softmax(y: np.ndarray, lambda_: np.ndarray): """ - Softmax function + Soft-thresholding (proximal operator of the L1 norm). + + Shrinks the magnitude of each component by ``lambda_`` while keeping its + phase, and clips it at zero: + + max(|y| - lambda_, 0) * y / |y| (y != 0) + 0 (y == 0) + + This is valid for both real and complex ``y``. For real ``y`` it reduces to - y - lambda_ (y > lambda_) - y + lambda_ (y < -lambda_) - 0 (otherwise) + y - lambda_ (y > lambda_) + y + lambda_ (y < -lambda_) + 0 (otherwise) + + The result is returned as float64 for real input and complex128 for complex + input (integer/float32 input is promoted, matching the historical float64 + real behavior and avoiding integer truncation of the threshold). """ assert isinstance(y, np.ndarray) and y.ndim == 1 assert isinstance(lambda_, np.ndarray) and lambda_.ndim == 1 assert (np.asarray(lambda_) > 0).all() assert y.size == lambda_.size - res = np.zeros(y.size) - - idx = y > lambda_ - res[idx] = y[idx] - lambda_[idx] - idx = y < -lambda_ - res[idx] = y[idx] + lambda_[idx] - - return res + if not np.iscomplexobj(y): + # Real path: literally the historical piecewise implementation, so the + # result is bit-for-bit identical for every real dtype (the arithmetic + # happens in the operands' dtype and is stored into a float64 array). + res = np.zeros(y.size) + idx = y > lambda_ + res[idx] = y[idx] - lambda_[idx] + idx = y < -lambda_ + res[idx] = y[idx] + lambda_[idx] + return res + + # Complex path: shrink the magnitude by lambda_ while keeping the phase, + # i.e. multiply by scale = max(1 - lambda/|y|, 0). + y = y.astype(np.complex128, copy=False) + lambda_ = np.asarray(lambda_, dtype=np.float64) + # Evaluate lambda/|y| with per-element component scaling so it stays finite + # even when |y| = sqrt(re**2 + im**2) itself overflows float64 (both np.abs + # and np.hypot would overflow for e.g. y = 1.3e308 + 1.3e308j). + # m = max(|re|, |im|), |y| = m * hypot(re/m, im/m), + # lambda/|y| = (lambda/m) / hypot(re/m, im/m). + m = np.maximum(np.abs(y.real), np.abs(y.imag)) + with np.errstate(divide="ignore", invalid="ignore"): + rm = np.where(m > 0.0, np.hypot(y.real / m, y.imag / m), 0.0) # |y|/m in [1, sqrt2] + ratio = np.where(m > 0.0, (lambda_ / m) / rm, 0.0) # lambda / |y| + scale = np.maximum(1.0 - ratio, 0.0) + res_c = y * scale + # Entries with an infinite component have infinite magnitude, so a finite + # threshold leaves them unchanged (and ``y * scale`` would give inf + nan*1j). + inf_mag = np.isinf(m) + if inf_mag.any(): + res_c = np.where(inf_mag, y, res_c) + return cast(np.ndarray, res_c) diff --git a/test/test_objectivefunc.py b/test/test_objectivefunc.py index f5ce0e5..d327505 100644 --- a/test/test_objectivefunc.py +++ b/test/test_objectivefunc.py @@ -184,3 +184,196 @@ def test_semi_positive_definite_penalty(): evals, evecs = np.linalg.eigh(x[:,:,k]) assert all(evals > -1e-10) + + +def _old_softmax(y, lambda_): + """Reference: original real-only piecewise soft-threshold.""" + res = np.zeros(y.size) + idx = y > lambda_ + res[idx] = y[idx] - lambda_[idx] + idx = y < -lambda_ + res[idx] = y[idx] + lambda_[idx] + return res + + +def test_softmax_real_equivalence_and_dtype(): + from admmsolver.objectivefunc import _softmax + rng = np.random.RandomState(0) + y = rng.randn(200) + lam = 0.3 * np.ones_like(y) + new = _softmax(y, lam) + old = _old_softmax(y, lam) + assert np.array_equal(new, old) # exact, not just close + assert new.dtype == np.float64 # real stays float64 + # exact boundaries: 0, +-lambda, and nextafter on both sides + lm = 1.5 + ys = np.array([0.0, lm, -lm, + np.nextafter(lm, np.inf), np.nextafter(lm, -np.inf), + np.nextafter(-lm, np.inf), np.nextafter(-lm, -np.inf), + 1e12, -1e12], dtype=np.float64) + lams = np.full(ys.size, lm) + assert np.array_equal(_softmax(ys, lams), _old_softmax(ys, lams)) + # float32 input is promoted to float64 (no threshold truncation for ints) + assert _softmax(np.array([5, -5, 0]), np.full(3, 1.5)).dtype == np.float64 + assert np.array_equal(_softmax(np.array([5, -5, 0]), np.full(3, 1.5)), [3.5, -3.5, 0.0]) + assert _softmax(np.array([1.0], np.float32), np.array([0.3])).dtype == np.float64 + # real +-inf stays +-inf (not NaN) + r = _softmax(np.array([np.inf, -np.inf]), np.array([1.0, 1.0])) + assert r[0] == np.inf and r[1] == -np.inf + + +def test_softmax_complex_and_inf(): + from admmsolver.objectivefunc import _softmax + # complex phase preserved, magnitude shrunk + y = np.array([3.0 + 4.0j, 0.6 - 0.8j, 0.0 + 0.0j]) # |y| = 5, 1, 0 + lam = np.array([1.0, 2.0, 1.0]) + out = _softmax(y, lam) + assert out.dtype == np.complex128 + assert np.allclose(out[0], (5.0 - 1.0) * (y[0] / 5.0)) # shrink 5 -> 4, keep phase + assert out[1] == 0 # |y|=1 < lambda=2 + assert out[2] == 0 # zero input + # complex entries with an infinite component must not become NaN + yi = np.array([complex(np.inf, 1.0), complex(1.0, np.inf), complex(np.inf, np.inf)]) + ri = _softmax(yi, np.ones(3)) + assert not np.isnan(ri).any() + assert ri[0] == complex(np.inf, 1.0) and ri[1] == complex(1.0, np.inf) + + +def test_L1_complex(): + """Complex L1 prox = magnitude soft-threshold keeping the phase.""" + N = 30 + rng = np.random.RandomState(1) + h = rng.randn(N) + 1j * rng.randn(N) + d = 0.5 + np.arange(N) * 0.1 # nonuniform positive weights + mu = DiagonalMatrix(d) + alpha = 0.7 + x = L1Regularizer(alpha, N).solve(h, mu) + assert x.dtype == np.complex128 + # closed form + center = -h / d + lam = 0.5 * alpha / d + mag = np.abs(center) + expect = np.where(mag > 0, np.maximum(mag - lam, 0.0) * center / np.where(mag > 0, mag, 1.0), 0.0) + assert np.allclose(x, expect, atol=1e-12) + # objective optimality (per component): f(x) <= f(x + perturbation) + for i in range(0, N, 7): + f = lambda z: alpha * np.abs(z) + 2 * np.real(np.conj(h[i]) * z) + d[i] * np.abs(z) ** 2 + for pert in [0.05, -0.05, 0.05j, -0.05j, 0.03 + 0.03j]: + assert f(x[i]) <= f(x[i] + pert) + 1e-12 + + +def test_spd_penalty_complex_correctness(): + """Complex-Hermitian SPD projection with a scalar (uniform) penalty must + equal the eigenvalue-clipped Hermitization of the unconstrained centre.""" + rng = np.random.RandomState(2) + K, N = 6, 3 + h = rng.randn(N * N * K) + 1j * rng.randn(N * N * K) + d = 2.0 + mu = PartialDiagonalMatrix(ScaledIdentityMatrix(N * N, d), (K,)) + res = SemiPositiveDefinitePenalty((N, N, K), axis=2).solve(h, mu).reshape((N, N, K)) + hb = (-h / d).reshape((N, N, K)) + for k in range(K): + H = 0.5 * (hb[:, :, k] + hb[:, :, k].conj().T) # Hermitize + ev, U = np.linalg.eigh(H) + keep = ev >= 0 + ref = U[:, keep] @ (ev[keep][:, None] * U[:, keep].conj().T) + assert np.allclose(res[:, :, k], ref, atol=1e-10) # exact projection + assert np.allclose(res[:, :, k], res[:, :, k].conj().T, atol=1e-12) # Hermitian + assert np.linalg.eigvalsh(res[:, :, k]).min() > -1e-10 # PSD + # imaginary part is genuinely retained (not dropped) + assert np.max(np.abs(res.imag)) > 1e-3 + + +def test_spd_penalty_nonuniform_center(): + """Verify the weighted paired-Hermitian center for a non-uniform diagonal + metric (independent of the known non-uniform PSD-projection limitation). + Uses a case whose center is already PSD, so eigenvalue clipping is a no-op + and the result must equal the center exactly.""" + N, K = 2, 1 + # per-entry weights d (Hermitian-pattern positive); build via DiagonalMatrix + d = np.array([1.0, 3.0, 0.5, 2.0]) # (0,0),(0,1),(1,0),(1,1) + # choose h so that centre = -(h_ij + conj(h_ji))/(d_ij + d_ji) is PSD + h = np.array([-2.0 + 0j, -(0.3 + 0.2j), -(0.3 - 0.2j), -3.0 + 0j]) + mu = DiagonalMatrix(d) + res = SemiPositiveDefinitePenalty((N, N, K), axis=2).solve(h, mu).reshape((N, N)) + hb = h.reshape((N, N)); db = d.reshape((N, N)) + centre = -(hb + hb.conj().T) / (db + db.T) + assert np.linalg.eigvalsh(centre).min() >= 0 # center already PSD + assert np.allclose(res, centre, atol=1e-12) # clip is a no-op -> equals center + assert np.allclose(res, res.conj().T) # Hermitian + + +def test_spd_penalty_uniform_diagonal_matches_scaled_identity(): + """A uniform full DiagonalMatrix must give the same result as the scalar + ScaledIdentityMatrix path.""" + rng = np.random.RandomState(3) + N, K = 3, 4 + h = rng.randn(N * N * K) + 1j * rng.randn(N * N * K) + c = 1.7 + p = SemiPositiveDefinitePenalty((N, N, K), axis=2) + r_scaled = p.solve(h, PartialDiagonalMatrix(ScaledIdentityMatrix(N * N, c), (K,))) + r_diag = p.solve(h, DiagonalMatrix(np.full(N * N * K, c))) + assert np.allclose(r_scaled, r_diag, atol=1e-10) + + +def test_softmax_low_precision_bit_exact(): + """Real path must be bit-for-bit identical to the historical implementation + for every real dtype (arithmetic in operand dtype, float64 result).""" + from admmsolver.objectivefunc import _softmax + rng = np.random.RandomState(4) + for dt in (np.float16, np.float32, np.float64): + y = (rng.randn(64) * 3).astype(dt) + lam = np.abs(rng.randn(64)).astype(dt) + dt(0.1) + assert np.array_equal(_softmax(y, lam), _old_softmax(y, lam)) + # integer input: no threshold truncation, float64 output + yi = np.array([5, -5, 0, 2, -2]) + li = np.full(yi.size, 1.5) + assert np.array_equal(_softmax(yi, li), _old_softmax(yi, li)) + assert _softmax(yi, li).dtype == np.float64 + + +def test_spd_center_no_overflow_large_values(): + """Paired centre must stay finite for near-float64-max operands.""" + big = np.finfo(np.float64).max + N, K = 2, 1 + # uniform h = d = big -> centre_ij = -(big+big)/(big+big) = -1 (must not be NaN) + h = np.array([big, big, big, big], dtype=np.float64) + d = np.array([big, big, big, big], dtype=np.float64) + p = SemiPositiveDefinitePenalty((N, N, K), axis=2) + res = p.solve(h, DiagonalMatrix(d)).reshape((N, N)) + assert np.isfinite(res).all() + centre = -np.ones((N, N)) # all entries -1 + ev, U = np.linalg.eigh(centre) + ref = U[:, ev >= 0] @ (ev[ev >= 0][:, None] * U[:, ev >= 0].conj().T) + assert np.allclose(res, ref, atol=1e-8) + + +def test_softmax_large_finite_complex_no_overflow(): + """Large-but-finite complex magnitude must be shrunk (not passed through), + even when |y| = sqrt(re^2+im^2) overflows float64.""" + from admmsolver.objectivefunc import _softmax + y = np.array([1.3e308 + 1.3e308j]) + lam = np.array([1.0e308]) + out = _softmax(y, lam) + assert np.isfinite(out).all() + # exact prox = (1 - lambda/|y|) * y, with lambda/|y| via component scaling so + # the reference itself does not overflow ( |y| ~ 1.84e308 > DBL_MAX ). + m = max(abs(y[0].real), abs(y[0].imag)) + ratio = (lam[0] / m) / np.hypot(y[0].real / m, y[0].imag / m) + expect = (1.0 - ratio) * y[0] + assert np.isclose(out[0].real, expect.real, rtol=1e-12) + assert np.isclose(out[0].imag, expect.imag, rtol=1e-12) + assert abs(out[0].real) < 6.0e307 # genuinely shrunk from 1.3e308 + + +def test_spd_center_large_finite_complex(): + """SPD paired centre must stay finite for near-max complex operands.""" + big = np.finfo(np.float64).max + N, K = 2, 1 + h = np.array([complex(-big, big)] * (N * N)) + d = np.array([big] * (N * N)) + res = SemiPositiveDefinitePenalty((N, N, K), axis=2).solve(h, DiagonalMatrix(d)).reshape((N, N)) + assert np.isfinite(res).all() + # centre_ij = -(h_ij + conj(h_ji))/(d_ij+d_ji) = -(-big+big j + (-big-big j))/(2 big) = 1 + ones = np.ones((N, N)) + assert np.allclose(res, ones, atol=1e-8) # ones is PSD -> projection is itself