#!/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 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('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()