Add mu-trajectory plot for the projection onto an ellipse
Cleaned-up extraction of test_mu_ellipse from proj_ellipsoid: - quadric.py: minimal Ellipsoid class (standardization, KKT trajectory, extreme roots by exact bracketing, degenerate points, asymptotes). - plot_mu_trajectories.py: the 4-panel figure, typeset with LaTeX. - test_quadric.py: checks against brute-force sampling of the ellipse. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
loicvhandClaude Opus 5 committed
commit
b2b782394a
6 files changed
+403
No files matched your search
@@ -0,0 +1,181 @@
|
||||
#!/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).
|
||||
"""
|
||||
|
||||
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 quadric import TOL, Ellipsoid
|
||||
|
||||
A = np.array([[2.0, 0.0], [0.0, 1.0]])
|
||||
B = np.zeros(2)
|
||||
C = -1.0
|
||||
|
||||
N_PER_BRANCH = 266 # mu samples between consecutive poles
|
||||
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 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.dim == 2, 'only 2D ellipses can be plotted'
|
||||
active = np.abs(Q.to_standardized(x0)) > 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
|
||||
mu = np.concatenate([np.linspace(lo, hi, N_PER_BRANCH) for lo, hi in
|
||||
[(mu_min, poles[0] - EPS), (poles[0] + EPS, poles[1] - EPS), (poles[1] + EPS, mu_max)]])
|
||||
idx = np.arange(len(mu))
|
||||
# A pole where x0 has a zero standardized component vanishes and joins its neighbouring branches.
|
||||
branches = np.split(idx, [k * N_PER_BRANCH for k in (1, 2) if pole_active[k - 1]])
|
||||
|
||||
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 = Q.extreme_roots(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),
|
||||
]
|
||||
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([Q.x(m, x0) for m in mu])
|
||||
ax_x.plot(*Q.boundary().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)
|
||||
|
||||
if not degenerate:
|
||||
for k, label in zip((1, 2), POLE_LABELS):
|
||||
for i in (k * N_PER_BRANCH - 1, k * N_PER_BRANCH):
|
||||
ax_x.text(*xs[i], label, **TEXT_STYLE)
|
||||
for k in range(Q.dim):
|
||||
point, direction = Q.pole_line(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)
|
||||
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('--x0', type=float, nargs=2, default=[0.0, -0.1], metavar=('X1', 'X2'))
|
||||
parser.add_argument('--output', type=Path, default=Path('figures/x_trajectories_mu'),
|
||||
help='output path, without extension (a .png and a .pdf are written)')
|
||||
parser.add_argument('--show', action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
fig = plot_trajectories(Ellipsoid(A, B, C), np.array(args.x0))
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
for extension in ('.png', '.pdf'):
|
||||
fig.savefig(args.output.with_name(args.output.name + extension))
|
||||
if args.show:
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user