Use quadproj instead of the local quadric module

The quadric, x(mu), f, f', mu* and the degenerate points now come from
quadproj; the script keeps small helpers for what quadproj does not
provide: mu**, the asymptotes and the ellipse boundary. The tests check
these helpers against brute-force sampling and quadproj.project.

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

.gitignore

@@ -1,3 +1,4 @@
 __pycache__/
 .pytest_cache/
 figures/
+venv/

README.md

@@ -24,6 +24,8 @@ Text is typeset with LaTeX (`text.usetex`) when `latex` is on the `PATH`, with t
 
 ## Files
 
-- `quadric.py`: the `Ellipsoid` class (standardization, trajectory, extreme roots, degenerate points).
+- `plot_mu_trajectories.py`: the figure. The quadric, `x(mu)`, `f`, `f'`, `mu*` and the degenerate
-- `plot_mu_trajectories.py`: the figure.
+  points come from [quadproj](https://loicvh.gitlab.io/quadproj); the script only adds `mu**`,
-- `test_quadric.py`: checks against brute-force sampling of the ellipse.
+  the asymptotes and the ellipse boundary.
+- `test_plot_mu_trajectories.py`: checks these helpers against brute-force sampling of the ellipse
+  and against `quadproj.project`.

plot_mu_trajectories.py

@@ -1,8 +1,9 @@
 #!/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
+Panels: the trajectory in the plane, the secular function f(mu), the distance ||x0 - x(mu)||
-f'(mu). Colors encode mu, from mu -> -inf (blue) to mu -> +inf (red).
+and f'(mu). Colors encode mu, from mu -> -inf (blue) to mu -> +inf (red). The quadric, x(mu),
+f, f' and mu* come from quadproj.
 """
 
 import argparse
@@ -16,8 +17,10 @@ from matplotlib.cm import ScalarMappable
 from matplotlib.collections import LineCollection
 from matplotlib.colors import Normalize
 from matplotlib.lines import Line2D
-
+from quadproj.fun import Fun
-from quadric import TOL, Ellipsoid
+from quadproj.project import get_KKT_point_root, get_squared_radius_k, get_xd_k
+from quadproj.quadrics import Quadric
+from scipy.optimize import brentq
 
 # One figure per case, for the ellipse x'Ax + b'x + c = 0 (b = 0 and c = -1 by default).
 CASES = {
@@ -25,6 +28,8 @@ CASES = {
     'x_trajectories_mu_2': dict(A=np.diag([2.0, 0.75]), x0=[0.0, -0.6]),
 }
 
+TOL = 1e-6            # normalized components of x0 below TOL are considered zero
+
 N_PER_BRANCH = 266    # mu samples between consecutive poles
 N_TAIL = 100          # extra samples extending the outer branches geometrically up to |mu| ~ 1e4 / lambda_min
 EPS = 0.01            # gap left on each side of a pole
@@ -41,6 +46,63 @@ 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 kkt_roots(Q, x0):
+    """Return (mu*, mu**), or None if x0 is the centre of the ellipse Q.
+
+    mu* is quadproj's root of f right of the rightmost active pole and mu** the root left of the
+    leftmost active pole: x(mu*) and x(mu**) are the nearest and farthest non-degenerate KKT
+    points (a degenerate point may be nearer than x(mu*)).
+    """
+    x0_norm = Q.to_normalized(x0)
+    active = np.abs(x0_norm) > TOL
+    if not active.any():
+        return None
+    mu_star, _ = get_KKT_point_root(Q, x0)
+
+    # f restricted to the active components, which avoids 0/0 at the poles of the other ones.
+    eig, u0 = Q.eig[active], x0_norm[active]
+
+    def f(mu):
+        return np.sum(eig * (u0 / (1 + mu * eig))**2) - 1
+
+    pole = -1 / eig.min()
+    near = far = 1.0
+    while f(pole - near) <= 0:
+        near /= 2
+    while f(pole - far) >= 0:
+        far *= 2
+    return mu_star, brentq(f, pole - far, pole - near)
+
+
+def degenerate_point(Q, x0, k):
+    """Return quadproj's degenerate KKT point at mu = -1/eig_k, or None if it does not exist."""
+    x0_norm = Q.to_normalized(x0)
+    L_k = [i for i in range(Q.dim) if abs(Q.eig[i] - Q.eig[k]) < TOL]
+    if np.any(np.abs(x0_norm[L_k]) > TOL) or get_squared_radius_k(Q, x0_norm, L_k, Q.eig[k]) <= 0:
+        return None
+    return Q.to_non_normalized(get_xd_k(Q, x0_norm, L_k, Q.eig[k]))
+
+
+def pole_line(Q, x0, k):
+    """Return (point, direction) of the line {u : u_i = u0_i / (1 - eig_i/eig_k) for eig_i != eig_k}.
+
+    Here u are normalized coordinates. The line is the asymptote of x(mu) at mu = -1/eig_k when
+    u0_k != 0, and contains the degenerate points otherwise.
+    """
+    x0_norm = Q.to_normalized(x0)
+    on_pole = np.abs(Q.eig - Q.eig[k]) < TOL
+    u = np.zeros(Q.dim)
+    np.divide(x0_norm, 1 - Q.eig / Q.eig[k], out=u, where=~on_pole)
+    return Q.to_non_normalized(u), Q.V[:, k]
+
+
+def ellipse_boundary(Q, n=1000):
+    """Return n points of the ellipse Q as an (n, 2) array."""
+    t = np.linspace(-np.pi, np.pi, n)
+    u = np.column_stack((np.cos(t), np.sin(t))) / np.sqrt(Q.eig)
+    return Q.d + np.sqrt(abs(Q.gamma)) * u @ Q.V.T
+
+
 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))
@@ -59,8 +121,10 @@ def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0):
 
 
 def plot_trajectories(Q, x0):
-    assert Q.dim == 2, 'only 2D ellipses can be plotted'
+    assert Q.is_ellipse, 'only 2D ellipses can be plotted'
-    active = np.abs(Q.to_standardized(x0)) > TOL
+    x0_norm = Q.to_normalized(x0)
+    F = Fun(Q, x0_norm)
+    active = np.abs(x0_norm) > TOL
     degenerate = not active.all()
 
     # Q.eig is positive and decreasing, so reversing it sorts the poles in increasing order.
@@ -80,7 +144,7 @@ def plot_trajectories(Q, x0):
     mu = np.concatenate(branch_samples)
     idx = np.arange(len(mu))
     pole_indices = np.cumsum([len(samples) for samples in branch_samples[:-1]])
-    # A pole where x0 has a zero standardized component vanishes and joins its neighbouring branches.
+    # A pole where x0 has a zero normalized component vanishes and joins its neighbouring branches.
     branches = np.split(idx, pole_indices[pole_active])
 
     cmap = colormaps[COLORMAP]
@@ -89,16 +153,16 @@ def plot_trajectories(Q, x0):
     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)
+    roots = kkt_roots(Q, 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_f, F.f, 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_dist, lambda m: F.dist(F.x_not_std(m)), 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),
+        (ax_df, F.d_f, 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])
@@ -115,25 +179,25 @@ def plot_trajectories(Q, x0):
     for ax in (ax_f, ax_df):
         ax.yaxis.set_ticks_position('right')
 
-    xs = np.array([Q.x(m, x0) for m in mu])
+    xs = np.array([F.x_not_std(m) for m in mu])
-    ax_x.plot(*Q.boundary().T, color=ELLIPSE_COLOR, zorder=1)
+    ax_x.plot(*ellipse_boundary(Q).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)
+        ax_x.scatter(*F.x_not_std(root), **style)
 
     if not degenerate:
         for pole_index, label in zip(pole_indices, POLE_LABELS):
             for i in (pole_index - 1, pole_index):
                 ax_x.text(*xs[i], label, **TEXT_STYLE)
         for k in range(Q.dim):
-            point, direction = Q.pole_line(x0, k)
+            point, direction = pole_line(Q, 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)
+        point, direction = pole_line(Q, x0, k)
-        xd = Q.degenerate_point(x0, k)
+        xd = degenerate_point(Q, x0, k)
         if xd is None:
             ax_x.plot(*np.column_stack((point - direction, point + direction)), color='grey', zorder=0)
             continue
@@ -186,8 +250,9 @@ def main():
     args.output_dir.mkdir(parents=True, exist_ok=True)
     for name in args.cases or CASES:
         case = CASES[name]
-        Q = Ellipsoid(case['A'], case.get('b', np.zeros(2)), case.get('c', -1.0))
+        Q = Quadric(np.asarray(case['A'], dtype=float), np.asarray(case.get('b', np.zeros(2)), dtype=float),
-        fig = plot_trajectories(Q, np.array(case['x0']))
+                    case.get('c', -1.0))
+        fig = plot_trajectories(Q, np.array(case['x0'], dtype=float))
         for extension in ('.png', '.pdf'):
             fig.savefig(args.output_dir / (name + extension))
         print('wrote', args.output_dir / name)

quadric.py

@@ -1,129 +0,0 @@
-"""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

@@ -1,4 +1,6 @@
 numpy
 scipy
 matplotlib>=3.6
+quadproj>=2.1
+imageio  # imported by quadproj but not declared as a dependency
 pytest

test_plot_mu_trajectories.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pytest
+from quadproj.fun import Fun
+from quadproj.project import project
+from quadproj.quadrics import Quadric
+
+from plot_mu_trajectories import degenerate_point, ellipse_boundary, kkt_roots, pole_line
+
+ROTATED_ELLIPSE = Quadric(np.array([[3.0, 1.0], [1.0, 2.0]]), np.array([1.0, -1.0]), -4.0)
+
+
+def trajectory(Q, x0):
+    return Fun(Q, Q.to_normalized(x0)).x_not_std
+
+
+def boundary_distances(Q, x0):
+    return np.linalg.norm(ellipse_boundary(Q, 200_000) - x0, axis=1)
+
+
+def test_boundary_is_on_the_ellipse():
+    assert np.all([abs(ROTATED_ELLIPSE.evaluate_point(x)) < 1e-10 for x in ellipse_boundary(ROTATED_ELLIPSE)])
+
+
+@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]])
+def test_kkt_roots_give_nearest_and_farthest_points(x0):
+    Q, x0 = ROTATED_ELLIPSE, np.array(x0)
+    mu_star, mu_star_star = kkt_roots(Q, x0)
+    x = trajectory(Q, x0)
+
+    distances = boundary_distances(Q, x0)
+    assert abs(Q.evaluate_point(x(mu_star_star))) < 1e-8
+    assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6)
+    assert np.linalg.norm(x(mu_star_star) - x0) == pytest.approx(distances.max(), abs=1e-6)
+
+
+@pytest.mark.parametrize('eig, x0', [([2.0, 1.0], [0.0, -0.1]), ([2.0, 0.75], [0.0, -0.6])])
+def test_degenerate_point_is_the_projection(eig, x0):
+    Q, x0 = Quadric(np.diag(eig), np.zeros(2), -1.0), np.array(x0)
+    xd = degenerate_point(Q, x0, 0)
+    np.testing.assert_allclose(xd, project(Q, x0))
+    assert np.linalg.norm(xd - x0) == pytest.approx(boundary_distances(Q, x0).min(), abs=1e-6)
+    assert degenerate_point(Q, x0, 1) is None
+
+
+def test_pole_line_is_the_asymptote():
+    Q, x0 = ROTATED_ELLIPSE, np.array([0.3, 0.2])
+    point, direction = pole_line(Q, x0, 1)
+    x = trajectory(Q, x0)(-1 / Q.eig[1] + 1e-9)
+    residual = (x - point) - ((x - point) @ direction) * direction
+    np.testing.assert_allclose(residual, 0, atol=1e-6)

test_quadric.py

@@ -1,58 +0,0 @@
-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))