#!/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 # 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]), } 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 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 # 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 standardized 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 = 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 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) 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('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 = Ellipsoid(case['A'], case.get('b', np.zeros(2)), case.get('c', -1.0)) fig = plot_trajectories(Q, np.array(case['x0'])) 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()