loicvhandClaude Opus 5
f1d5a73cc6
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>
265 lines
11 KiB
Python
265 lines
11 KiB
Python
#!/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, 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
|
|
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 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 = {
|
|
'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]),
|
|
}
|
|
|
|
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
|
|
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 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))
|
|
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.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.
|
|
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)
|
|
idx = np.arange(len(mu))
|
|
pole_indices = np.cumsum([len(samples) for samples in branch_samples[:-1]])
|
|
# 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$']
|
|
|
|
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, 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])
|
|
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([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(*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 = pole_line(Q, x0, k)
|
|
ax_x.axline(point, point + direction, linestyle='dotted', color='grey')
|
|
|
|
for k in np.flatnonzero(~active):
|
|
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
|
|
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('cases', nargs='*', metavar='CASE',
|
|
help=f'cases to plot, among {", ".join(CASES)} (default: all)')
|
|
parser.add_argument('--output-dir', type=Path, default=Path('figures'))
|
|
parser.add_argument('--show', action='store_true')
|
|
args = parser.parse_args()
|
|
unknown = set(args.cases) - set(CASES)
|
|
if unknown:
|
|
parser.error(f'unknown cases: {", ".join(sorted(unknown))}')
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
for name in args.cases or CASES:
|
|
case = CASES[name]
|
|
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)
|
|
if args.show:
|
|
plt.show()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|