Add mu-trajectory plot for the projection onto an ellipse

Cleaned-up extraction of test_mu_ellipse from proj_ellipsoid:
- quadric.py: minimal Ellipsoid class (standardization, KKT trajectory,
  extreme roots by exact bracketing, degenerate points, asymptotes).
- plot_mu_trajectories.py: the 4-panel figure, typeset with LaTeX.
- test_quadric.py: checks against brute-force sampling of the ellipse.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+.pytest_cache/
+figures/

README.md

@@ -0,0 +1,28 @@
+# mu-trajectories
+
+Visualize the KKT trajectory of the Euclidean projection of a point `x0` onto an ellipse
+`x'Ax + b'x + c = 0`:
+
+    x(mu) = (I + mu A)^{-1} (x0 - mu b / 2)
+
+The figure shows the trajectory in the plane, `f(mu) = psi(x(mu))`, the distance
+`||x0 - x(mu)||` and `f'(mu)`, together with the nearest/farthest KKT points `x(mu*)`,
+`x(mu**)` and, when `x0` has a zero standardized coordinate, the degenerate point `x^d`.
+
+## Usage
+
+    pip install -r requirements.txt
+    python plot_mu_trajectories.py                 # degenerate case x0 = (0, -0.1)
+    python plot_mu_trajectories.py --x0 0.3 0.2 --output figures/generic
+    pytest
+
+The ellipse is set by `A`, `B`, `C` at the top of `plot_mu_trajectories.py`.
+
+Text is typeset with LaTeX (`text.usetex`) when `latex` is on the `PATH`, with the `cm-super` and
+`type1cm` packages; otherwise matplotlib's Computer Modern math font is used.
+
+## Files
+
+- `quadric.py`: the `Ellipsoid` class (standardization, trajectory, extreme roots, degenerate points).
+- `plot_mu_trajectories.py`: the figure.
+- `test_quadric.py`: checks against brute-force sampling of the ellipse.

plot_mu_trajectories.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto an ellipse.
+
+Panels: the trajectory in the plane, f(mu) = psi(x(mu)), the distance ||x0 - x(mu)|| and
+f'(mu). Colors encode mu, from mu -> -inf (blue) to mu -> +inf (red).
+"""
+
+import argparse
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import colormaps
+from matplotlib.cm import ScalarMappable
+from matplotlib.collections import LineCollection
+from matplotlib.colors import Normalize
+from matplotlib.lines import Line2D
+
+from quadric import TOL, Ellipsoid
+
+A = np.array([[2.0, 0.0], [0.0, 1.0]])
+B = np.zeros(2)
+C = -1.0
+
+N_PER_BRANCH = 266    # mu samples between consecutive poles
+EPS = 0.01            # gap left on each side of a pole
+MU_HALF_WIDTH = 3     # mu is sampled in [m - MU_HALF_WIDTH, m + MU_HALF_WIDTH], m the mean of the poles
+N_DOTTED_OUTER = 2    # points drawn dotted at the ends of the outer branches
+
+COLORMAP = 'RdYlBu_r'
+ELLIPSE_COLOR = 'royalblue'
+STAR_STYLE = dict(color='purple', marker='v', s=12, zorder=2)
+STAR_STAR_STYLE = dict(color='purple', marker='s', s=12, zorder=2)
+POLE_STYLE = dict(linestyle='dashed', color='grey')
+TEXT_STYLE = dict(ha='center', va='center', fontsize=12, clip_on=True)
+POLE_LABELS = [r'$\frac{-1}{\lambda_2}$', r'$\frac{-1}{\lambda_1}$']
+DIST_D_LABEL = r'$||\mathbf{x}^0 - \mathbf{x}^\mathrm{d}||_2$'
+
+
+def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0):
+    """Draw y(x) colored by idx, with the first and last n_dotted points joined by dotted segments."""
+    points = np.column_stack((x, y))
+    n = len(points)
+    if n_dotted == 0:
+        parts = [(0, n, 'solid')]
+    else:
+        parts = [(0, n_dotted, 'dotted'),
+                 (n_dotted - 1, n - n_dotted + 1, 'solid'),
+                 (n - n_dotted, n, 'dotted')]
+    for start, stop, linestyle in parts:
+        p = points[start:stop]
+        lc = LineCollection(np.stack((p[:-1], p[1:]), axis=1), cmap=cmap, norm=norm, linestyle=linestyle)
+        lc.set_array(idx[start:stop - 1])
+        ax.add_collection(lc)
+
+
+def plot_trajectories(Q, x0):
+    assert Q.dim == 2, 'only 2D ellipses can be plotted'
+    active = np.abs(Q.to_standardized(x0)) > TOL
+    degenerate = not active.all()
+
+    # Q.eig is positive and decreasing, so reversing it sorts the poles in increasing order.
+    poles = -1 / Q.eig[::-1]
+    pole_active = active[::-1]
+
+    mu_min = poles.mean() - MU_HALF_WIDTH
+    mu_max = poles.mean() + MU_HALF_WIDTH
+    mu = np.concatenate([np.linspace(lo, hi, N_PER_BRANCH) for lo, hi in
+                         [(mu_min, poles[0] - EPS), (poles[0] + EPS, poles[1] - EPS), (poles[1] + EPS, mu_max)]])
+    idx = np.arange(len(mu))
+    # A pole where x0 has a zero standardized component vanishes and joins its neighbouring branches.
+    branches = np.split(idx, [k * N_PER_BRANCH for k in (1, 2) if pole_active[k - 1]])
+
+    cmap = colormaps[COLORMAP]
+    norm = Normalize(0, len(mu) - 1)
+    xlim = (mu_min, mu_max)
+    ticks = [mu_min, poles[0], poles[1], 0, mu_max]
+    tick_labels = [r'$\mu \to -\infty$', *POLE_LABELS, r'$0$', r'$\mu \to \infty$']
+
+    roots = Q.extreme_roots(x0)
+    markers = [] if roots is None else list(zip(roots, (STAR_STYLE, STAR_STAR_STYLE)))
+
+    fig, axs = plt.subplots(2, 2, figsize=(10, 7))
+    ax_x, ax_f, ax_dist, ax_df = axs.flat
+
+    panels = [
+        (ax_f, lambda m: Q.f(m, x0), r'$f(\mu)$', (-1, 1), 5),
+        (ax_dist, lambda m: np.linalg.norm(Q.x(m, x0) - x0), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', (0, 2), 6),
+        (ax_df, lambda m: Q.df(m, x0), r'$f^\prime(\mu)$', (-2, 2), 50),
+    ]
+    for ax, fun, ylabel, ylim, n_dotted_middle in panels:
+        values = np.array([fun(m) for m in mu])
+        for i, branch in enumerate(branches):
+            n_dotted = 0 if degenerate else (n_dotted_middle if i == 1 else N_DOTTED_OUTER)
+            colored_line(ax, mu[branch], values[branch], branch, cmap, norm, n_dotted)
+        for pole in poles[pole_active]:
+            ax.axvline(pole, **POLE_STYLE)
+        for root, style in markers:
+            ax.scatter(root, fun(root), **style)
+        ax.set(xlabel=r'$\mu$', ylabel=ylabel, xlim=xlim, ylim=ylim)
+        ax.set_xticks(ticks, tick_labels)
+    ax_f.axhline(0, color='red', alpha=0.5, zorder=1)
+    for ax in (ax_f, ax_df):
+        ax.yaxis.set_ticks_position('right')
+
+    xs = np.array([Q.x(m, x0) for m in mu])
+    ax_x.plot(*Q.boundary().T, color=ELLIPSE_COLOR, zorder=1)
+    ax_x.scatter(*Q.d, color=ELLIPSE_COLOR)
+    ax_x.scatter(xs[:, 0], xs[:, 1], c=idx, cmap=cmap, norm=norm, s=4)
+    ax_x.scatter(*x0, color='red', marker='x')
+    for root, style in markers:
+        ax_x.scatter(*Q.x(root, x0), **style)
+
+    if not degenerate:
+        for k, label in zip((1, 2), POLE_LABELS):
+            for i in (k * N_PER_BRANCH - 1, k * N_PER_BRANCH):
+                ax_x.text(*xs[i], label, **TEXT_STYLE)
+        for k in range(Q.dim):
+            point, direction = Q.pole_line(x0, k)
+            ax_x.axline(point, point + direction, linestyle='dotted', color='grey')
+
+    for k in np.flatnonzero(~active):
+        point, direction = Q.pole_line(x0, k)
+        xd = Q.degenerate_point(x0, k)
+        if xd is None:
+            ax_x.plot(*np.column_stack((point - direction, point + direction)), color='grey', zorder=0)
+            continue
+        ax_x.plot(*np.column_stack((point, xd)), color='grey', zorder=0)
+        ax_x.scatter(*xd, color='green', zorder=2)
+        dist = np.linalg.norm(xd - x0)
+        ax_dist.axhline(dist + 0.1, color='green')
+        ax_dist.text(mu_max + (poles[1] - mu_max) * 2 / 3, dist + 0.2, DIST_D_LABEL, color='green')
+
+    legend = [
+        Line2D([], [], color=ELLIPSE_COLOR, lw=2, label=r'$\mathcal{Q}(\Psi)$'),
+        Line2D([], [], linestyle='none', marker='v', color='purple', markersize=6, label=r'$\mathbf{x}(\mu^*)$'),
+        Line2D([], [], linestyle='none', marker='s', color='purple', markersize=6, label=r'$\mathbf{x}(\mu^{**})$'),
+        Line2D([], [], linestyle='none', marker='x', color='red', markersize=8, label=r'$\mathbf{x}^0 = \mathbf{x}(0)$'),
+    ]
+    if degenerate:
+        legend.append(Line2D([], [], linestyle='none', marker='o', color='green', markersize=8,
+                             label=r'$\mathbf{x}^\mathrm{d}$'))
+    ax_x.legend(handles=legend, loc='lower left', fontsize=9)
+    ax_x.set(xlabel=r'$x_1$', ylabel=r'$x_2$', xlim=(-2, 2), ylim=(-2, 2), aspect='equal')
+
+    fig.tight_layout(pad=2.1)
+    colorbar = fig.colorbar(ScalarMappable(norm, cmap), ax=ax_x, ticks=[0, len(mu) - 1])
+    colorbar.ax.set_yticklabels([r'$\mu \to -\infty$', r'$\mu \to \infty$'])
+    return fig
+
+
+def use_latex_fonts():
+    """Typeset text with LaTeX when available, otherwise with matplotlib's Computer Modern."""
+    plt.rcParams.update({'font.family': 'serif', 'font.size': 12})
+    if shutil.which('latex'):
+        plt.rcParams.update({'text.usetex': True, 'text.latex.preamble': r'\usepackage{amsmath}'})
+    else:
+        plt.rcParams.update({'mathtext.fontset': 'cm', 'font.serif': ['cmr10'],
+                             'axes.formatter.use_mathtext': True})
+
+
+def main():
+    use_latex_fonts()
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument('--x0', type=float, nargs=2, default=[0.0, -0.1], metavar=('X1', 'X2'))
+    parser.add_argument('--output', type=Path, default=Path('figures/x_trajectories_mu'),
+                        help='output path, without extension (a .png and a .pdf are written)')
+    parser.add_argument('--show', action='store_true')
+    args = parser.parse_args()
+
+    fig = plot_trajectories(Ellipsoid(A, B, C), np.array(args.x0))
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    for extension in ('.png', '.pdf'):
+        fig.savefig(args.output.with_name(args.output.name + extension))
+    if args.show:
+        plt.show()
+
+
+if __name__ == '__main__':
+    main()

quadric.py

@@ -0,0 +1,129 @@
+"""Ellipsoid x'Ax + b'x + c = 0 and the KKT trajectory of the projection onto it.
+
+The KKT conditions of min ||x - x0|| subject to x on the ellipsoid give
+
+    x(mu) = (I + mu A)^{-1} (x0 - mu b / 2),
+
+and the projections are the points x(mu) with f(mu) = psi(x(mu)) = 0.
+
+With d = -A^{-1} b / 2, gamma = c + b'd + d'Ad < 0 and A = V diag(eig) V', the standardized
+coordinates u = V'(x - d) / sqrt(-gamma) turn the ellipsoid into sum_i eig_i u_i^2 = 1 and the
+trajectory into u_i(mu) = u0_i / (1 + mu eig_i), which has a pole at mu = -1/eig_i when u0_i != 0.
+"""
+
+import numpy as np
+from scipy.optimize import brentq
+
+TOL = 1e-6
+
+
+class Ellipsoid:
+    def __init__(self, A, b, c):
+        A = np.asarray(A, dtype=float)
+        b = np.asarray(b, dtype=float)
+        d = np.linalg.solve(A, -b / 2)
+        gamma = c + b @ d + d @ A @ d
+        if gamma > 0:  # same set, written with gamma < 0
+            A, b, c, gamma = -A, -b, -c, -gamma
+
+        eig, V = np.linalg.eigh(A)
+        order = np.argsort(-eig)
+        self.eig = eig[order]
+        self.V = V[:, order]
+        if not np.all(self.eig > 0):
+            raise ValueError('the quadric is not a non-empty ellipsoid')
+
+        self.A, self.b, self.c, self.d = A, b, c, d
+        self.dim = len(b)
+        self.scale = np.sqrt(-gamma)
+
+    def psi(self, x):
+        return x @ self.A @ x + self.b @ x + self.c
+
+    def to_standardized(self, x):
+        return self.V.T @ (x - self.d) / self.scale
+
+    def to_non_standardized(self, u):
+        return self.d + self.scale * (self.V @ u)
+
+    def _solve(self, mu, y):
+        """Return (I + mu A)^{-1} y."""
+        return self.V @ ((self.V.T @ y) / (1 + mu * self.eig))
+
+    def x(self, mu, x0):
+        return self._solve(mu, x0 - mu * self.b / 2)
+
+    def f(self, mu, x0):
+        return self.psi(self.x(mu, x0))
+
+    def df(self, mu, x0):
+        x = self.x(mu, x0)
+        dx = -self._solve(mu, self.A @ x + self.b / 2)
+        return (2 * self.A @ x + self.b) @ dx
+
+    def extreme_roots(self, x0):
+        """Return (mu*, mu**), the roots of f right of the rightmost and left of the leftmost active pole.
+
+        x(mu*) and x(mu**) are the nearest and farthest non-degenerate KKT points; a degenerate
+        point (see degenerate_point) may be nearer than x(mu*). Return None if x0 is the centre.
+        """
+        u0 = self.to_standardized(x0)
+        active = np.abs(u0) > TOL
+        if not active.any():
+            return None
+        eig, u0 = self.eig[active], u0[active]
+
+        def secular(mu):
+            return np.sum(eig * (u0 / (1 + mu * eig))**2) - 1
+
+        mu_star = brentq(secular, *_bracket(secular, -1 / eig.max(), +1))
+        mu_star_star = brentq(secular, *_bracket(secular, -1 / eig.min(), -1))
+        return mu_star, mu_star_star
+
+    def _pole_point(self, u0, k):
+        """Standardized limit of u(mu) as mu -> -1/eig_k, with the eig_k components set to 0."""
+        on_pole = np.abs(self.eig - self.eig[k]) < TOL
+        u = np.zeros(self.dim)
+        np.divide(u0, 1 - self.eig / self.eig[k], out=u, where=~on_pole)
+        return u, on_pole
+
+    def pole_line(self, x0, k):
+        """Return (point, direction) of the line {x : u_i = u0_i / (1 - eig_i/eig_k) for eig_i != eig_k}.
+
+        It is the asymptote of x(mu) at mu = -1/eig_k when u0_k != 0, and the line containing
+        the degenerate points otherwise.
+        """
+        u, _ = self._pole_point(self.to_standardized(x0), k)
+        return self.to_non_standardized(u), self.V[:, k]
+
+    def degenerate_point(self, x0, k):
+        """Return the KKT point at mu = -1/eig_k when u0 vanishes on that eigenspace, or None."""
+        u0 = self.to_standardized(x0)
+        u, on_pole = self._pole_point(u0, k)
+        if np.any(np.abs(u0[on_pole]) > TOL):
+            return None
+        r = (1 - self.eig @ u**2) / self.eig[k]
+        if r <= 0:
+            return None
+        u[np.argmax(on_pole)] = np.sqrt(r)
+        return self.to_non_standardized(u)
+
+    def boundary(self, n=1000):
+        """Return n points of a 2D ellipse as an (n, 2) array."""
+        assert self.dim == 2, 'only 2D ellipses can be sampled'
+        t = np.linspace(-np.pi, np.pi, n)
+        u = np.column_stack((np.cos(t), np.sin(t))) / np.sqrt(self.eig)
+        return self.d + self.scale * u @ self.V.T
+
+
+def _bracket(f, pole, direction):
+    """Return an interval on the `direction` side of `pole` where f changes sign.
+
+    Assumes f -> +inf next to the pole and f -> -1 far from it.
+    """
+    near = far = 1.0
+    while f(pole + direction * near) <= 0:
+        near /= 2
+    while f(pole + direction * far) >= 0:
+        far *= 2
+    return sorted((pole + direction * near, pole + direction * far))

requirements.txt

@@ -0,0 +1,4 @@
+numpy
+scipy
+matplotlib>=3.6
+pytest

test_quadric.py

@@ -0,0 +1,58 @@
+import numpy as np
+import pytest
+
+from quadric import Ellipsoid
+
+AXIS_ELLIPSE = Ellipsoid([[2, 0], [0, 1]], [0, 0], -1)
+ROTATED_ELLIPSE = Ellipsoid([[3, 1], [1, 2]], [1, -1], -4)
+
+
+def boundary_distances(Q, x0):
+    return np.linalg.norm(Q.boundary(200_000) - x0, axis=1)
+
+
+@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]])
+def test_extreme_roots_give_nearest_and_farthest_points(x0):
+    Q, x0 = ROTATED_ELLIPSE, np.array(x0)
+    mu_star, mu_star_star = Q.extreme_roots(x0)
+    x_star, x_star_star = Q.x(mu_star, x0), Q.x(mu_star_star, x0)
+
+    assert abs(Q.psi(x_star)) < 1e-8
+    assert abs(Q.psi(x_star_star)) < 1e-8
+    distances = boundary_distances(Q, x0)
+    assert np.linalg.norm(x_star - x0) == pytest.approx(distances.min(), abs=1e-6)
+    assert np.linalg.norm(x_star_star - x0) == pytest.approx(distances.max(), abs=1e-6)
+
+
+def test_trajectory_satisfies_kkt_conditions():
+    Q, x0, mu = ROTATED_ELLIPSE, np.array([0.3, 0.2]), 0.7
+    x = Q.x(mu, x0)
+    np.testing.assert_allclose(x - x0 + mu * (Q.A @ x + Q.b / 2), 0, atol=1e-12)
+
+
+def test_df_matches_finite_difference():
+    Q, x0, mu, h = ROTATED_ELLIPSE, np.array([0.3, 0.2]), 0.7, 1e-6
+    finite_difference = (Q.f(mu + h, x0) - Q.f(mu - h, x0)) / (2 * h)
+    assert Q.df(mu, x0) == pytest.approx(finite_difference, rel=1e-6)
+
+
+def test_degenerate_point_is_the_projection():
+    Q, x0 = AXIS_ELLIPSE, np.array([0.0, -0.1])
+    xd = Q.degenerate_point(x0, 0)
+    assert abs(Q.psi(xd)) < 1e-12
+    assert np.linalg.norm(xd - x0) == pytest.approx(boundary_distances(Q, x0).min(), abs=1e-6)
+    assert Q.degenerate_point(x0, 1) is None
+
+
+def test_pole_line_is_the_asymptote():
+    Q, x0 = AXIS_ELLIPSE, np.array([0.3, 0.2])
+    point, direction = Q.pole_line(x0, 1)
+    x = Q.x(-1 / Q.eig[1] + 1e-9, x0)
+    residual = (x - point) - ((x - point) @ direction) * direction
+    np.testing.assert_allclose(residual, 0, atol=1e-6)
+
+
+def test_sign_convention_does_not_matter():
+    flipped = Ellipsoid(-ROTATED_ELLIPSE.A, -ROTATED_ELLIPSE.b, -ROTATED_ELLIPSE.c)
+    x0 = np.array([0.3, 0.2])
+    np.testing.assert_allclose(flipped.extreme_roots(x0), ROTATED_ELLIPSE.extreme_roots(x0))