From f1d5a73cc67603ad5334c666e3d78fa65a22ce46 Mon Sep 17 00:00:00 2001 From: loicvh Date: Tue, 15 Sep 2026 21:08:35 +0200 Subject: [PATCH] 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) --- .gitignore | 1 + README.md | 8 ++- plot_mu_trajectories.py | 103 ++++++++++++++++++++++------ quadric.py | 129 ----------------------------------- requirements.txt | 2 + test_plot_mu_trajectories.py | 50 ++++++++++++++ test_quadric.py | 58 ---------------- 7 files changed, 142 insertions(+), 209 deletions(-) delete mode 100644 quadric.py create mode 100644 test_plot_mu_trajectories.py delete mode 100644 test_quadric.py diff --git a/.gitignore b/.gitignore index 54363e3..c035281 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__/ .pytest_cache/ figures/ +venv/ diff --git a/README.md b/README.md index 47ea95b..42d8a40 100644 --- a/README.md +++ b/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. -- `test_quadric.py`: checks against brute-force sampling of the ellipse. +- `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`. diff --git a/plot_mu_trajectories.py b/plot_mu_trajectories.py index e6c4bbf..7cc3996 100644 --- a/plot_mu_trajectories.py +++ b/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) diff --git a/quadric.py b/quadric.py deleted file mode 100644 index 787fea8..0000000 --- a/quadric.py +++ /dev/null @@ -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)) diff --git a/requirements.txt b/requirements.txt index 27c52aa..3225342 100644 --- a/requirements.txt +++ b/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 diff --git a/test_plot_mu_trajectories.py b/test_plot_mu_trajectories.py new file mode 100644 index 0000000..a2c4a1d --- /dev/null +++ b/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) diff --git a/test_quadric.py b/test_quadric.py deleted file mode 100644 index 66999dc..0000000 --- a/test_quadric.py +++ /dev/null @@ -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))