Support complex-Hermitian variables in L1 and SPD objective functions

src/admmsolver/__init__.py

@@ -2,4 +2,4 @@
 # SPDX-License-Identifier: MIT
 
 __license__ = "MIT"
-__version__ = "0.7.6"
+__version__ = "0.7.8"

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)

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