loicvhandClaude Opus 5
bdcf0e0b02
The poles, branches and mu range now follow the eigenvalues (one or two poles of either sign), so hyperbolas, parabolas, line pairs, single lines and points are supported, not only ellipses. - Find every root of f from sign changes, including across inactive poles where quadproj's f is 0/0; mu* is the root between the active poles surrounding 0 and the other roots get their own marker. - Mark the apex of conical quadrics as a degenerate KKT point. - Draw the quadric as the contour psi = 0 and frame the trajectory panel on x0, the KKT points and the vertices. - Use the correct asymptote direction for repeated eigenvalues. - Test every quadric type against brute-force sampling. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
367 lines
17 KiB
Python
367 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""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 the degenerate points 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_squared_radius_k, get_xd_k
|
|
from quadproj.quadrics import Quadric
|
|
from scipy.optimize import brentq
|
|
|
|
# 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 # 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 / min |lambda|
|
|
EPS = 0.01 # gap left on each side of a pole
|
|
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'
|
|
QUADRIC_COLOR = 'royalblue'
|
|
STAR_STYLE = dict(color='purple', marker='v', 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)
|
|
DIST_D_LABEL = r'$||\mathbf{x}^0 - \mathbf{x}^\mathrm{d}||_2$'
|
|
|
|
|
|
def f_poles(Q, x0):
|
|
"""Return the poles of f as (mu, eigenvalue indices, active) tuples, sorted by mu.
|
|
|
|
Equal eigenvalues share a pole, which is active when x0 has a normalized component along them.
|
|
"""
|
|
x0_norm = Q.to_normalized(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])
|
|
|
|
|
|
def sample_mu(Q, poles):
|
|
"""Return the mu samples of each segment between consecutive poles, and the range of the mu axis.
|
|
|
|
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:
|
|
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 through lim x(mu) at mu = -1/eig_k.
|
|
|
|
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
|
|
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):
|
|
"""Return y-limits framing the finite features of a curve plotted over xlim.
|
|
|
|
The features are the values at both ends of xlim, the local extrema of each branch and
|
|
`extra` (markers, reference lines). Divergences at the poles are not features: the curve
|
|
runs off the axes there.
|
|
"""
|
|
visible = (mu >= xlim[0]) & (mu <= xlim[1])
|
|
features = [values[visible][0], values[visible][-1], *extra]
|
|
for branch in branches:
|
|
v = values[branch[visible[branch]]]
|
|
slopes = np.diff(v)
|
|
features.extend(v[1:-1][slopes[:-1] * slopes[1:] < 0])
|
|
low, high = min(features), max(features)
|
|
padding = Y_PADDING * (high - low) if high > low else 1
|
|
low -= padding
|
|
if floor is not None:
|
|
low = max(low, floor)
|
|
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))
|
|
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):
|
|
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()
|
|
|
|
segments, xlim = sample_mu(Q, poles)
|
|
mu = np.concatenate(segments)
|
|
idx = np.arange(len(mu))
|
|
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)
|
|
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$']
|
|
|
|
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
|
|
|
|
# (axis, function, label, dotted points of the middle branches, extra values to show, lower bound)
|
|
panels = [
|
|
(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, 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):
|
|
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 pole_mus[pole_active]:
|
|
ax.axvline(pole, **POLE_STYLE)
|
|
for root, style in markers:
|
|
ax.scatter(root, fun(root), **style)
|
|
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([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 (_, 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 i in (pole_index - 1, pole_index):
|
|
ax_x.text(*xs[i], label, **TEXT_STYLE)
|
|
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)
|
|
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)
|
|
ax_dist.axhline(dist + 0.1, 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=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='best', fontsize=9)
|
|
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])
|
|
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()
|