loicvhandClaude Opus 5
f193596743
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='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])
|
|
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()
|