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>

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
-f'(mu). Colors encode mu, from mu -> -inf (blue) to mu -> +inf (red).
+Panels: the trajectory in the plane, the secular function f(mu), the distance ||x0 - x(mu)||
+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 quadric import TOL, Ellipsoid
+from quadproj.fun import Fun
+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'
-    active = np.abs(Q.to_standardized(x0)) > TOL
+    assert Q.is_ellipse, 'only 2D ellipses can be plotted'
+    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_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),
+        (ax_f, F.f, r'$f(\mu)$', (-1, 1), 5),
+        (ax_dist, lambda m: F.dist(F.x_not_std(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', (0, 2), 6),
+        (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])
-    ax_x.plot(*Q.boundary().T, color=ELLIPSE_COLOR, zorder=1)
+    xs = np.array([F.x_not_std(m) for m in mu])
+    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)
-        xd = Q.degenerate_point(x0, k)
+        point, direction = pole_line(Q, 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))
-        fig = plot_trajectories(Q, np.array(case['x0']))
+        Q = Quadric(np.asarray(case['A'], dtype=float), np.asarray(case.get('b', np.zeros(2)), dtype=float),
+                    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)