plot_mu_trajectories/plot_mu_trajectories.py

#!/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
Y_PADDING = 0.1       # fraction of the y-range of the features added below and above them

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 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 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

    # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0
    inactive = [(*pole_line(Q, x0, k), degenerate_point(Q, x0, k)) for k in np.flatnonzero(~active)]
    xd_dists = [np.linalg.norm(xd - x0) for _, _, xd in inactive if xd is not None]

    # (axis, function, label, dotted points of the middle branch, extra values to show, lower bound)
    panels = [
        (ax_f, F.f, r'$f(\mu)$', 5, [0], None),
        (ax_dist, lambda m: F.dist(F.x_not_std(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, F.d_f, 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):
            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)
        root_values = [fun(root) for root, _ in markers]
        ylim = auto_ylim(mu, values, branches, xlim, extra + root_values, floor)
        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 point, direction, xd in inactive:
        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()