Document how to create the virtual environment

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

README.md

@@ -1,17 +1,27 @@
 # mu-trajectories
 
-Visualize the KKT trajectory of the Euclidean projection of a point `x0` onto an ellipse
-`x'Ax + b'x + c = 0`:
+Visualize the KKT trajectory of the Euclidean projection of a point `x0` onto a 2D quadric
+`x'Ax + b'x + c = 0` (ellipse, hyperbola, parabola, two lines, a line or a point):
 
     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`.
+The figure shows the trajectory in the plane, the secular function `f(mu)`, the distance
+`||x0 - x(mu)||` and `f'(mu)`, together with the KKT points: `x(mu*)`, the other roots of `f`
+and, when `x0` has a zero normalized coordinate, the degenerate points `x^d`.
+
+## Setup
+
+Create a virtual environment in `venv/` (ignored by git) and install the dependencies
+(tested with Python 3.13):
+
+    python3 -m venv venv
+    source venv/bin/activate
+    pip install -r requirements.txt
+
+Run `source venv/bin/activate` again in each new shell, and `deactivate` to leave the environment.
 
 ## Usage
 
-    pip install -r requirements.txt
     python plot_mu_trajectories.py                         # every case, into figures/
     python plot_mu_trajectories.py x_trajectories_mu_2     # a single case
     pytest
@@ -24,8 +34,8 @@ Text is typeset with LaTeX (`text.usetex`) when `latex` is on the `PATH`, with t
 
 ## Files
 
-- `plot_mu_trajectories.py`: the figure. The quadric, `x(mu)`, `f`, `f'`, `mu*` and the degenerate
-  points come from [quadproj](https://loicvh.gitlab.io/quadproj); the script only adds `mu**`,
-  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`: the figure. The quadric, `x(mu)`, `f`, `f'` and the degenerate
+  points come from [quadproj](https://loicvh.gitlab.io/quadproj); the script adds the roots of
+  `f`, the asymptotes and the plot limits.
+- `test_plot_mu_trajectories.py`: checks, for each type of quadric, that the KKT points contain
+  the projection found by brute-force sampling, and that every type can be plotted.

plot_mu_trajectories.py

@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
-"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto an ellipse.
+"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto a 2D quadric.
 
 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.
+f, f' and the degenerate points come from quadproj.
 """
 
 import argparse
@@ -18,65 +18,109 @@ from matplotlib.collections import LineCollection
 from matplotlib.colors import Normalize
 from matplotlib.lines import Line2D
 from quadproj.fun import Fun
-from quadproj.project import get_KKT_point_root, get_squared_radius_k, get_xd_k
+from quadproj.project import 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).
+# One figure per case, for the quadric x'Ax + b'x + c = 0 (b = 0 and c = -1 by default).
 CASES = {
     'x_trajectories_mu_1': dict(A=np.diag([2.0, 1.0]), x0=[0.0, -0.1]),
     'x_trajectories_mu_2': dict(A=np.diag([2.0, 0.75]), x0=[0.0, -0.6]),
+    'x_trajectories_mu_parabola': dict(A=np.diag([1.0, 0.0]), b=[0.0, -1.0], c=0.0, x0=[0.3, 0.2]),
 }
 
-TOL = 1e-6            # normalized components of x0 below TOL are considered zero
+TOL = 1e-6            # eigenvalues and 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
+N_TAIL = 100          # extra samples extending the outer branches geometrically up to |mu| ~ 1e4 / min |lambda|
 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
+MU_MARGIN = 2.5       # the mu axis extends MU_MARGIN beyond the poles and 0
 N_DOTTED_OUTER = 2    # points drawn dotted at the ends of the outer branches
 Y_PADDING = 0.1       # fraction of the y-range of the features added below and above them
+WINDOW_SCALE = 2      # the trajectory panel is WINDOW_SCALE times as wide as x0 and the KKT points
 
 COLORMAP = 'RdYlBu_r'
-ELLIPSE_COLOR = 'royalblue'
+QUADRIC_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)
+ROOT_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 kkt_roots(Q, x0):
-    """Return (mu*, mu**), or None if x0 is the centre of the ellipse Q.
+def f_poles(Q, x0):
+    """Return the poles of f as (mu, eigenvalue indices, active) tuples, sorted by mu.
 
-    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*)).
+    Equal eigenvalues share a pole, which is active when x0 has a normalized component along them.
     """
     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)
+    groups = []
+    for k in np.flatnonzero(np.abs(Q.eig) > TOL):
+        if groups and abs(Q.eig[k] - Q.eig[groups[-1][0]]) < TOL:
+            groups[-1].append(k)
+        else:
+            groups.append([k])
+    poles = [(-1 / Q.eig[ks[0]], ks, bool(np.any(np.abs(x0_norm[ks]) > TOL))) for ks in groups]
+    return sorted(poles, key=lambda pole: pole[0])
 
-    # 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
+def sample_mu(Q, poles):
+    """Return the mu samples of each segment between consecutive poles, and the range of the mu axis.
 
-    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)
+    x(mu) only reaches its limit as |mu| -> inf, so the outer segments continue geometrically
+    beyond the mu axis; these samples show up in the trajectory but not on the mu axes.
+    """
+    pole_mus = [mu for mu, _, _ in poles]
+    mu_min, mu_max = min(pole_mus + [0]) - MU_MARGIN, max(pole_mus + [0]) + MU_MARGIN
+    tail = np.geomspace(MU_MARGIN / N_PER_BRANCH, 1e4 / np.abs(Q.eig[np.abs(Q.eig) > TOL]).min(), N_TAIL)
+    starts = [mu_min] + [mu + EPS for mu in pole_mus]
+    stops = [mu - EPS for mu in pole_mus] + [mu_max]
+    segments = [np.linspace(start, stop, N_PER_BRANCH) for start, stop in zip(starts, stops)]
+    segments[0] = np.concatenate((mu_min - tail[::-1], segments[0]))
+    segments[-1] = np.concatenate((segments[-1], mu_max + tail))
+    return segments, (mu_min, mu_max)
+
+
+def off_pole(fun):
+    """Wrap a function of mu from quadproj, which is 0/0 exactly at an inactive pole, where it is continuous."""
+    def wrapped(mu):
+        with np.errstate(divide='ignore', invalid='ignore'):
+            value = fun(mu)
+        return fun(np.nextafter(mu, np.inf)) if np.any(np.isnan(value)) else value
+    return wrapped
+
+
+def kkt_roots(Q, x0):
+    """Return (mu*, the other roots of f), found from the sign changes of f between samples.
+
+    mu* is the root between the active poles surrounding 0, as in quadproj, or None.
+    """
+    poles = f_poles(Q, x0)
+    segments, _ = sample_mu(Q, poles)
+    f = off_pole(Fun(Q, Q.to_normalized(x0)).f)
+    # f is continuous across inactive poles, so roots may lie in the gaps left around them.
+    branches = [segments[0]]
+    for samples, (_, _, is_active) in zip(segments[1:], poles):
+        if is_active:
+            branches.append(samples)
+        else:
+            branches[-1] = np.concatenate((branches[-1], samples))
+    roots = []
+    for samples in branches:
+        signs = np.sign([f(mu) for mu in samples])
+        roots += [brentq(f, samples[i], samples[i + 1]) for i in np.flatnonzero(signs[:-1] * signs[1:] < 0)]
+
+    active = [mu for mu, _, is_active in poles if is_active]
+    e1 = max((mu for mu in active if mu < 0), default=-np.inf)
+    e2 = min((mu for mu in active if mu > 0), default=np.inf)
+    mu_star = min((mu for mu in roots if e1 < mu < e2), key=abs, default=None)
+    return mu_star, [mu for mu in roots if mu != mu_star]
 
 
 def degenerate_point(Q, x0, k):
     """Return quadproj's degenerate KKT point at mu = -1/eig_k, or None if it does not exist."""
+    if abs(Q.eig[k]) < TOL:
+        return None
     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:
@@ -85,23 +129,20 @@ def degenerate_point(Q, x0, 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}.
+    """Return (point, direction) of the line through lim x(mu) at mu = -1/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.
+    The limit is taken with the components of the eig_k eigenspace set to 0. When x0 has a
+    component in that eigenspace, x(mu) diverges along it and the line is the asymptote;
+    otherwise the line is along eigenvector k and contains the degenerate points.
     """
     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
+    with np.errstate(divide='ignore', invalid='ignore'):
+        u = Fun(Q, x0_norm).x_std(-1 / Q.eig[k])
+    u[on_pole] = 0
+    direction = np.where(on_pole, x0_norm, 0) if np.any(np.abs(x0_norm[on_pole]) > TOL) else np.eye(Q.dim)[k]
+    origin = Q.to_non_normalized(np.zeros(Q.dim))
+    return Q.to_non_normalized(u), Q.to_non_normalized(direction) - origin
 
 
 def auto_ylim(mu, values, branches, xlim, extra=(), floor=None):
@@ -125,6 +166,28 @@ def auto_ylim(mu, values, branches, xlim, extra=(), floor=None):
     return low, high + padding
 
 
+def square_window(points):
+    """Return (xlim, ylim) of a square centred on the points and WINDOW_SCALE times as wide (at least 2)."""
+    points = np.asarray(points)
+    low, high = points.min(axis=0), points.max(axis=0)
+    half = max(WINDOW_SCALE * (high - low).max() / 2, 1)
+    centre = (low + high) / 2
+    return (centre[0] - half, centre[0] + half), (centre[1] - half, centre[1] + half)
+
+
+def draw_quadric(ax, Q, xlim, ylim):
+    if Q.is_single_plane:  # psi >= 0 vanishes on a line, without sign change for the contour
+        k = np.flatnonzero(np.abs(Q.eig) <= TOL)[0]
+        ax.axline(Q.d, Q.d + Q.V[:, k], color=QUADRIC_COLOR, zorder=1)
+    elif Q.is_conical and Q.n == 0:  # psi >= 0 vanishes at a single point
+        ax.scatter(*Q.d, color=QUADRIC_COLOR, zorder=1)
+    else:
+        X, Y = np.meshgrid(np.linspace(*xlim, 801), np.linspace(*ylim, 801))
+        P = np.stack((X, Y), axis=-1)
+        psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c
+        ax.contour(X, Y, psi, levels=[0], colors=QUADRIC_COLOR, zorder=1)
+
+
 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))
@@ -143,109 +206,120 @@ def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0):
 
 
 def plot_trajectories(Q, x0):
-    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()
+    if Q.dim != 2 or Q.is_empty:
+        raise ValueError('only non-empty 2D quadrics can be plotted')
+    F = Fun(Q, Q.to_normalized(x0))
+    x, f, df = off_pole(F.x_not_std), off_pole(F.f), off_pole(F.d_f)
+    poles = f_poles(Q, x0)
+    pole_mus = np.array([mu for mu, _, _ in poles])
+    pole_active = np.array([is_active for _, _, is_active in poles], dtype=bool)
+    degenerate = not pole_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
-    # x(mu) only reaches the centre as |mu| -> inf, so the outer branches continue geometrically
-    # beyond [mu_min, mu_max]; these samples show up in the trajectory but not on the mu axes.
-    tail = np.geomspace(MU_HALF_WIDTH / N_PER_BRANCH, 1e4 / Q.eig.min(), N_TAIL)
-    branch_samples = [
-        np.concatenate((mu_min - tail[::-1], np.linspace(mu_min, poles[0] - EPS, N_PER_BRANCH))),
-        np.linspace(poles[0] + EPS, poles[1] - EPS, N_PER_BRANCH),
-        np.concatenate((np.linspace(poles[1] + EPS, mu_max, N_PER_BRANCH), mu_max + tail)),
-    ]
-    mu = np.concatenate(branch_samples)
+    segments, xlim = sample_mu(Q, poles)
+    mu = np.concatenate(segments)
     idx = np.arange(len(mu))
-    pole_indices = np.cumsum([len(samples) for samples in branch_samples[:-1]])
+    pole_indices = np.cumsum([len(samples) for samples in segments[:-1]], dtype=int)
     # 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]
     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$']
+    pole_labels = [rf'$\frac{{-1}}{{\lambda_{ks[0] + 1}}}$' for _, ks, _ in poles]
+    ticks = [xlim[0], *pole_mus, 0, xlim[1]]
+    tick_labels = [r'$\mu \to -\infty$', *pole_labels, r'$0$', r'$\mu \to \infty$']
 
-    roots = kkt_roots(Q, x0)
-    markers = [] if roots is None else list(zip(roots, (STAR_STYLE, STAR_STAR_STYLE)))
+    mu_star, other_roots = kkt_roots(Q, x0)
+    markers = ([] if mu_star is None else [(mu_star, STAR_STYLE)]) + [(root, ROOT_STYLE) for root in other_roots]
+
+    # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0
+    inactive = [(*pole_line(Q, x0, ks[0]), degenerate_point(Q, x0, ks[0]))
+                for _, ks, is_active in poles if not is_active]
+    xds = [xd for _, _, xd in inactive if xd is not None]
+    if Q.is_conical and not Q.is_single_plane:  # the apex is a KKT point
+        xds.append(Q.d)
+    xd_dists = [np.linalg.norm(xd - x0) for xd in xds]
 
     fig, axs = plt.subplots(2, 2, figsize=(10, 7))
     ax_x, ax_f, ax_dist, ax_df = axs.flat
 
-    # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0
-    inactive = [(*pole_line(Q, x0, k), degenerate_point(Q, x0, k)) for k in np.flatnonzero(~active)]
-    xd_dists = [np.linalg.norm(xd - x0) for _, _, xd in inactive if xd is not None]
-
-    # (axis, function, label, dotted points of the middle branch, extra values to show, lower bound)
+    # (axis, function, label, dotted points of the middle branches, extra values to show, lower bound)
     panels = [
-        (ax_f, F.f, r'$f(\mu)$', 5, [0], None),
-        (ax_dist, lambda m: F.dist(F.x_not_std(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', 6,
+        (ax_f, f, r'$f(\mu)$', 5, [0], None),
+        (ax_dist, lambda m: F.dist(x(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', 6,
          [dist + 0.3 for dist in xd_dists], 0),  # green line and its label above it
-        (ax_df, F.d_f, r'$f^\prime(\mu)$', 50, [], None),
+        (ax_df, df, r'$f^\prime(\mu)$', 50, [], None),
     ]
     for ax, fun, ylabel, n_dotted_middle, extra, floor 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)
+            if degenerate:
+                n_dotted = 0
+            else:
+                n_dotted = n_dotted_middle if 0 < i < len(branches) - 1 else N_DOTTED_OUTER
             colored_line(ax, mu[branch], values[branch], branch, cmap, norm, n_dotted)
-        for pole in poles[pole_active]:
+        for pole in pole_mus[pole_active]:
             ax.axvline(pole, **POLE_STYLE)
         for root, style in markers:
             ax.scatter(root, fun(root), **style)
-        root_values = [fun(root) for root, _ in markers]
-        ylim = auto_ylim(mu, values, branches, xlim, extra + root_values, floor)
+        marker_values = [fun(root) for root, _ in markers if xlim[0] <= root <= xlim[1]]
+        ylim = auto_ylim(mu, values, branches, xlim, extra + marker_values + [fun(0)], floor)  # x(0) = x0
         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([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)
+    xs = np.array([x(m) for m in mu])
+    kkt_points = [x(root) for root, _ in markers] + xds
+    has_centre = not (Q.is_parabolic or Q.is_cylindrical)  # otherwise the centre is not unique
+    window_points = [x0, *kkt_points]
+    if has_centre:
+        window_points.append(Q.center)
+    if Q.is_central:  # the vertices, so that the whole shape (e.g. both arcs of a hyperbola) is shown
+        window_points += [Q.to_non_normalized(sign * np.eye(Q.dim)[k] / np.sqrt(Q.eig[k]))
+                          for k in np.flatnonzero(Q.eig > TOL) for sign in (-1, 1)]
+    window_xlim, window_ylim = square_window(window_points)
+
+    draw_quadric(ax_x, Q, window_xlim, window_ylim)
+    if has_centre:
+        ax_x.scatter(*Q.center, color=QUADRIC_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(*F.x_not_std(root), **style)
+    for (_, style), point in zip(markers, kkt_points):
+        ax_x.scatter(*point, **style)
 
     if not degenerate:
-        for pole_index, label in zip(pole_indices, POLE_LABELS):
+        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 = pole_line(Q, x0, k)
+        for _, ks, _ in poles:
+            point, direction = pole_line(Q, x0, ks[0])
             ax_x.axline(point, point + direction, linestyle='dotted', color='grey')
 
     for point, direction, xd in inactive:
         if xd is None:
             ax_x.plot(*np.column_stack((point - direction, point + direction)), color='grey', zorder=0)
-            continue
+        else:
             ax_x.plot(*np.column_stack((point, xd)), color='grey', zorder=0)
+    for xd, dist in zip(xds, xd_dists):
         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')
+        ax_dist.text(xlim[1] + (pole_mus[-1] - xlim[1]) * 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 = [Line2D([], [], color=QUADRIC_COLOR, lw=2, label=r'$\mathcal{Q}(\Psi)$')]
+    if mu_star is not None:
+        legend.append(Line2D([], [], linestyle='none', marker='v', color='purple', markersize=6,
+                             label=r'$\mathbf{x}(\mu^*)$'))
+    if other_roots:
+        legend.append(Line2D([], [], linestyle='none', marker='s', color='purple', markersize=6,
+                             label=r'$\mathbf{x}(\mu)$, $f(\mu) = 0$'))
+    legend.append(Line2D([], [], linestyle='none', marker='x', color='red', markersize=8,
+                         label=r'$\mathbf{x}^0 = \mathbf{x}(0)$'))
+    if xds:
         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')
+    ax_x.legend(handles=legend, loc='lower left', fontsize=8, markerscale=0.8, handlelength=1.5, labelspacing=0.3)
+    ax_x.set(xlabel=r'$x_1$', ylabel=r'$x_2$', xlim=window_xlim, ylim=window_ylim, aspect='equal')
 
     fig.tight_layout(pad=2.1)
     colorbar = fig.colorbar(ScalarMappable(norm, cmap), ax=ax_x, ticks=[0, len(mu) - 1])

test_plot_mu_trajectories.py

@@ -1,44 +1,100 @@
+import contourpy
+import matplotlib
+import matplotlib.pyplot as plt
 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
+from plot_mu_trajectories import degenerate_point, kkt_roots, off_pole, plot_trajectories, pole_line
 
-ROTATED_ELLIPSE = Quadric(np.array([[3.0, 1.0], [1.0, 2.0]]), np.array([1.0, -1.0]), -4.0)
+matplotlib.use('Agg')
+
+GRID = np.linspace(-4, 4, 4001)
+
+# (A, b, c, x0): generic points, and points with a zero normalized component (degenerate)
+QUADRICS = {
+    'ellipse': ([[3, 1], [1, 2]], [1, -1], -4, [0.3, 0.2]),
+    'ellipse_outside': ([[3, 1], [1, 2]], [1, -1], -4, [2.0, -1.5]),
+    'ellipse_degenerate': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, -0.1]),
+    'ellipse_degenerate_2': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, 0.5]),
+    'hyperbola': ([[1, 2], [2, -1]], [1, 0], -1, [0.3, 0.2]),
+    'hyperbola_degenerate': ([[1, 0], [0, -1]], [0, 0], -1, [0.0, 0.5]),
+    'parabola': ([[1, 1], [1, 1]], [1, -1], 0, [0.3, 0.2]),
+    'parabola_degenerate': ([[1, 0], [0, 0]], [0, -1], 0, [0.0, 0.5]),
+    'crossing_lines': ([[1, 0], [0, -1]], [0, 0], 0, [0.3, 0.2]),
+    'parallel_lines': ([[1, 0], [0, 0]], [0, 0], -1, [0.3, 0.2]),
+}
+# Quadrics without finite roots of f: the projection is the limit of x(mu) as |mu| -> inf.
+LIMIT_QUADRICS = {
+    'single_line': ([[1, 0], [0, 0]], [0, 0], 0, [0.3, 0.2]),
+    'point': ([[1, 0], [0, 1]], [0, 0], 0, [0.3, 0.2]),
+}
+
+
+def make_quadric(A, b, c):
+    return Quadric(np.array(A, dtype=float), np.array(b, dtype=float), float(c))
+
+
+ROTATED_ELLIPSE = make_quadric(*QUADRICS['ellipse'][:3])
 
 
 def trajectory(Q, x0):
-    return Fun(Q, Q.to_normalized(x0)).x_not_std
+    return off_pole(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 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 distances_to_quadric(Q, x0):
+    """Distances from x0 to the points of the contour psi = 0 on a fine grid."""
+    X, Y = np.meshgrid(GRID, GRID)
+    P = np.stack((X, Y), axis=-1)
+    psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c
+    return np.linalg.norm(np.vstack(contourpy.contour_generator(X, Y, psi).lines(0)) - x0, axis=1)
+
+
+@pytest.mark.parametrize('name', QUADRICS)
+def test_kkt_points_contain_the_projection(name):
+    *quadric, x0 = QUADRICS[name]
+    Q, x0 = make_quadric(*quadric), np.array(x0)
+    mu_star, other_roots = kkt_roots(Q, x0)
+    x = trajectory(Q, x0)
+    points = [x(mu) for mu in other_roots + ([] if mu_star is None else [mu_star])]
+    points += [xd for k in range(Q.dim) if (xd := degenerate_point(Q, x0, k)) is not None]
+    if Q.is_conical:
+        points.append(Q.d)
+
+    assert all(abs(Q.evaluate_point(point)) < 1e-8 for point in points)
+    nearest = min(np.linalg.norm(point - x0) for point in points)
+    assert nearest == pytest.approx(distances_to_quadric(Q, x0).min(), abs=1e-4)
+
+
+@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]])
+def test_ellipse_roots_give_nearest_and_farthest_points(x0):
+    Q, x0 = ROTATED_ELLIPSE, np.array(x0)
+    mu_star, other_roots = kkt_roots(Q, x0)
+    x = trajectory(Q, x0)
+
+    distances = np.linalg.norm(ellipse_boundary(Q, 200_000) - x0, axis=1)
+    assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6)
+    assert max(np.linalg.norm(x(mu) - x0) for mu in other_roots) == pytest.approx(distances.max(), abs=1e-6)
 
 
 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
 
 
@@ -46,5 +102,17 @@ 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)
+    direction = direction / np.linalg.norm(direction)
     residual = (x - point) - ((x - point) @ direction) * direction
     np.testing.assert_allclose(residual, 0, atol=1e-6)
+
+
+@pytest.mark.parametrize('name', [*QUADRICS, *LIMIT_QUADRICS])
+def test_every_quadric_type_can_be_plotted(name):
+    *quadric, x0 = {**QUADRICS, **LIMIT_QUADRICS}[name]
+    plt.close(plot_trajectories(make_quadric(*quadric), np.array(x0, dtype=float)))
+
+
+def test_empty_quadric_is_rejected():
+    with pytest.raises(ValueError):
+        plot_trajectories(make_quadric([[1, 0], [0, 1]], [0, 0], 1), np.zeros(2))