#!/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, 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 mu* 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_KKT_point_root, get_squared_radius_k, get_xd_k from quadproj.quadrics import Quadric from scipy.optimize import brentq # 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]), } TOL = 1e-6 # 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 / 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 kkt_roots(Q, x0): """Return (mu*, mu**), or None if x0 is the centre of the ellipse Q. mu* is quadproj's root of f right of the rightmost active pole and mu** the root left of the leftmost active pole: x(mu*) and x(mu**) are the nearest and farthest non-degenerate KKT points (a degenerate point may be nearer than x(mu*)). """ x0_norm = Q.to_normalized(x0) active = np.abs(x0_norm) > TOL if not active.any(): return None mu_star, _ = get_KKT_point_root(Q, x0) # f restricted to the active components, which avoids 0/0 at the poles of the other ones. eig, u0 = Q.eig[active], x0_norm[active] def f(mu): return np.sum(eig * (u0 / (1 + mu * eig))**2) - 1 pole = -1 / eig.min() near = far = 1.0 while f(pole - near) <= 0: near /= 2 while f(pole - far) >= 0: far *= 2 return mu_star, brentq(f, pole - far, pole - near) def degenerate_point(Q, x0, k): """Return quadproj's degenerate KKT point at mu = -1/eig_k, or None if it does not exist.""" 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 {u : u_i = u0_i / (1 - eig_i/eig_k) for eig_i != eig_k}. Here u are normalized coordinates. The line is the asymptote of x(mu) at mu = -1/eig_k when u0_k != 0, and contains the degenerate points otherwise. """ x0_norm = Q.to_normalized(x0) on_pole = np.abs(Q.eig - Q.eig[k]) < TOL u = np.zeros(Q.dim) np.divide(x0_norm, 1 - Q.eig / Q.eig[k], out=u, where=~on_pole) return Q.to_non_normalized(u), Q.V[:, k] def ellipse_boundary(Q, n=1000): """Return n points of the ellipse Q as an (n, 2) array.""" t = np.linspace(-np.pi, np.pi, n) u = np.column_stack((np.cos(t), np.sin(t))) / np.sqrt(Q.eig) return Q.d + np.sqrt(abs(Q.gamma)) * u @ Q.V.T 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.is_ellipse, 'only 2D ellipses can be plotted' x0_norm = Q.to_normalized(x0) F = Fun(Q, x0_norm) active = np.abs(x0_norm) > 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 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) 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 = kkt_roots(Q, 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, F.f, r'$f(\mu)$', (-1, 1), 5), (ax_dist, lambda m: F.dist(F.x_not_std(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', (0, 2), 6), (ax_df, F.d_f, 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([F.x_not_std(m) for m in mu]) ax_x.plot(*ellipse_boundary(Q).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(*F.x_not_std(root), **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 = pole_line(Q, x0, k) ax_x.axline(point, point + direction, linestyle='dotted', color='grey') for k in np.flatnonzero(~active): point, direction = pole_line(Q, x0, k) xd = degenerate_point(Q, 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 = 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()