diff --git a/README.md b/README.md index 42d8a40..6982469 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,13 @@ # mu-trajectories -Visualize the KKT trajectory of the Euclidean projection of a point `x0` onto an ellipse -`x'Ax + b'x + c = 0`: +Visualize the KKT trajectory of the Euclidean projection of a point `x0` onto a 2D quadric +`x'Ax + b'x + c = 0` (ellipse, hyperbola, parabola, two lines, a line or a point): x(mu) = (I + mu A)^{-1} (x0 - mu b / 2) -The figure shows the trajectory in the plane, `f(mu) = psi(x(mu))`, the distance -`||x0 - x(mu)||` and `f'(mu)`, together with the nearest/farthest KKT points `x(mu*)`, -`x(mu**)` and, when `x0` has a zero standardized coordinate, the degenerate point `x^d`. +The figure shows the trajectory in the plane, the secular function `f(mu)`, the distance +`||x0 - x(mu)||` and `f'(mu)`, together with the KKT points: `x(mu*)`, the other roots of `f` +and, when `x0` has a zero normalized coordinate, the degenerate points `x^d`. ## Usage @@ -24,8 +24,8 @@ Text is typeset with LaTeX (`text.usetex`) when `latex` is on the `PATH`, with t ## Files -- `plot_mu_trajectories.py`: the figure. The quadric, `x(mu)`, `f`, `f'`, `mu*` and the degenerate - points come from [quadproj](https://loicvh.gitlab.io/quadproj); the script only adds `mu**`, - the asymptotes and the ellipse boundary. -- `test_plot_mu_trajectories.py`: checks these helpers against brute-force sampling of the ellipse - and against `quadproj.project`. +- `plot_mu_trajectories.py`: the figure. The quadric, `x(mu)`, `f`, `f'` and the degenerate + points come from [quadproj](https://loicvh.gitlab.io/quadproj); the script adds the roots of + `f`, the asymptotes and the plot limits. +- `test_plot_mu_trajectories.py`: checks, for each type of quadric, that the KKT points contain + the projection found by brute-force sampling, and that every type can be plotted. diff --git a/plot_mu_trajectories.py b/plot_mu_trajectories.py index 1c61f55..679359d 100644 --- a/plot_mu_trajectories.py +++ b/plot_mu_trajectories.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 -"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto an ellipse. +"""Plot the KKT trajectory mu -> x(mu) of the projection of x0 onto a 2D quadric. 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. +f, f' and the degenerate points come from quadproj. """ import argparse @@ -18,65 +18,109 @@ 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.project import 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). +# One figure per case, for the quadric 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]), + 'x_trajectories_mu_parabola': dict(A=np.diag([1.0, 0.0]), b=[0.0, -1.0], c=0.0, x0=[0.3, 0.2]), } -TOL = 1e-6 # normalized components of x0 below TOL are considered zero +TOL = 1e-6 # eigenvalues and 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 +N_TAIL = 100 # extra samples extending the outer branches geometrically up to |mu| ~ 1e4 / min |lambda| 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 +MU_MARGIN = 2.5 # the mu axis extends MU_MARGIN beyond the poles and 0 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 +WINDOW_SCALE = 2 # the trajectory panel is WINDOW_SCALE times as wide as x0 and the KKT points COLORMAP = 'RdYlBu_r' -ELLIPSE_COLOR = 'royalblue' +QUADRIC_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) +ROOT_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. +def f_poles(Q, x0): + """Return the poles of f as (mu, eigenvalue indices, active) tuples, sorted by mu. - 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*)). + Equal eigenvalues share a pole, which is active when x0 has a normalized component along them. """ 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) + groups = [] + for k in np.flatnonzero(np.abs(Q.eig) > TOL): + if groups and abs(Q.eig[k] - Q.eig[groups[-1][0]]) < TOL: + groups[-1].append(k) + else: + groups.append([k]) + poles = [(-1 / Q.eig[ks[0]], ks, bool(np.any(np.abs(x0_norm[ks]) > TOL))) for ks in groups] + return sorted(poles, key=lambda pole: pole[0]) - # 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 +def sample_mu(Q, poles): + """Return the mu samples of each segment between consecutive poles, and the range of the mu axis. - 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) + x(mu) only reaches its limit as |mu| -> inf, so the outer segments continue geometrically + beyond the mu axis; these samples show up in the trajectory but not on the mu axes. + """ + pole_mus = [mu for mu, _, _ in poles] + mu_min, mu_max = min(pole_mus + [0]) - MU_MARGIN, max(pole_mus + [0]) + MU_MARGIN + tail = np.geomspace(MU_MARGIN / N_PER_BRANCH, 1e4 / np.abs(Q.eig[np.abs(Q.eig) > TOL]).min(), N_TAIL) + starts = [mu_min] + [mu + EPS for mu in pole_mus] + stops = [mu - EPS for mu in pole_mus] + [mu_max] + segments = [np.linspace(start, stop, N_PER_BRANCH) for start, stop in zip(starts, stops)] + segments[0] = np.concatenate((mu_min - tail[::-1], segments[0])) + segments[-1] = np.concatenate((segments[-1], mu_max + tail)) + return segments, (mu_min, mu_max) + + +def off_pole(fun): + """Wrap a function of mu from quadproj, which is 0/0 exactly at an inactive pole, where it is continuous.""" + def wrapped(mu): + with np.errstate(divide='ignore', invalid='ignore'): + value = fun(mu) + return fun(np.nextafter(mu, np.inf)) if np.any(np.isnan(value)) else value + return wrapped + + +def kkt_roots(Q, x0): + """Return (mu*, the other roots of f), found from the sign changes of f between samples. + + mu* is the root between the active poles surrounding 0, as in quadproj, or None. + """ + poles = f_poles(Q, x0) + segments, _ = sample_mu(Q, poles) + f = off_pole(Fun(Q, Q.to_normalized(x0)).f) + # f is continuous across inactive poles, so roots may lie in the gaps left around them. + branches = [segments[0]] + for samples, (_, _, is_active) in zip(segments[1:], poles): + if is_active: + branches.append(samples) + else: + branches[-1] = np.concatenate((branches[-1], samples)) + roots = [] + for samples in branches: + signs = np.sign([f(mu) for mu in samples]) + roots += [brentq(f, samples[i], samples[i + 1]) for i in np.flatnonzero(signs[:-1] * signs[1:] < 0)] + + active = [mu for mu, _, is_active in poles if is_active] + e1 = max((mu for mu in active if mu < 0), default=-np.inf) + e2 = min((mu for mu in active if mu > 0), default=np.inf) + mu_star = min((mu for mu in roots if e1 < mu < e2), key=abs, default=None) + return mu_star, [mu for mu in roots if mu != mu_star] def degenerate_point(Q, x0, k): """Return quadproj's degenerate KKT point at mu = -1/eig_k, or None if it does not exist.""" + if abs(Q.eig[k]) < TOL: + return None 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: @@ -85,23 +129,20 @@ def degenerate_point(Q, x0, 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}. + """Return (point, direction) of the line through lim x(mu) at mu = -1/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. + The limit is taken with the components of the eig_k eigenspace set to 0. When x0 has a + component in that eigenspace, x(mu) diverges along it and the line is the asymptote; + otherwise the line is along eigenvector k and contains the degenerate points. """ 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 + with np.errstate(divide='ignore', invalid='ignore'): + u = Fun(Q, x0_norm).x_std(-1 / Q.eig[k]) + u[on_pole] = 0 + direction = np.where(on_pole, x0_norm, 0) if np.any(np.abs(x0_norm[on_pole]) > TOL) else np.eye(Q.dim)[k] + origin = Q.to_non_normalized(np.zeros(Q.dim)) + return Q.to_non_normalized(u), Q.to_non_normalized(direction) - origin def auto_ylim(mu, values, branches, xlim, extra=(), floor=None): @@ -125,6 +166,28 @@ def auto_ylim(mu, values, branches, xlim, extra=(), floor=None): return low, high + padding +def square_window(points): + """Return (xlim, ylim) of a square centred on the points and WINDOW_SCALE times as wide (at least 2).""" + points = np.asarray(points) + low, high = points.min(axis=0), points.max(axis=0) + half = max(WINDOW_SCALE * (high - low).max() / 2, 1) + centre = (low + high) / 2 + return (centre[0] - half, centre[0] + half), (centre[1] - half, centre[1] + half) + + +def draw_quadric(ax, Q, xlim, ylim): + if Q.is_single_plane: # psi >= 0 vanishes on a line, without sign change for the contour + k = np.flatnonzero(np.abs(Q.eig) <= TOL)[0] + ax.axline(Q.d, Q.d + Q.V[:, k], color=QUADRIC_COLOR, zorder=1) + elif Q.is_conical and Q.n == 0: # psi >= 0 vanishes at a single point + ax.scatter(*Q.d, color=QUADRIC_COLOR, zorder=1) + else: + X, Y = np.meshgrid(np.linspace(*xlim, 801), np.linspace(*ylim, 801)) + P = np.stack((X, Y), axis=-1) + psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c + ax.contour(X, Y, psi, levels=[0], colors=QUADRIC_COLOR, zorder=1) + + 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)) @@ -143,109 +206,120 @@ def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0): 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() + if Q.dim != 2 or Q.is_empty: + raise ValueError('only non-empty 2D quadrics can be plotted') + F = Fun(Q, Q.to_normalized(x0)) + x, f, df = off_pole(F.x_not_std), off_pole(F.f), off_pole(F.d_f) + poles = f_poles(Q, x0) + pole_mus = np.array([mu for mu, _, _ in poles]) + pole_active = np.array([is_active for _, _, is_active in poles], dtype=bool) + degenerate = not pole_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) + segments, xlim = sample_mu(Q, poles) + mu = np.concatenate(segments) idx = np.arange(len(mu)) - pole_indices = np.cumsum([len(samples) for samples in branch_samples[:-1]]) + pole_indices = np.cumsum([len(samples) for samples in segments[:-1]], dtype=int) # 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$'] + pole_labels = [rf'$\frac{{-1}}{{\lambda_{ks[0] + 1}}}$' for _, ks, _ in poles] + ticks = [xlim[0], *pole_mus, 0, xlim[1]] + 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))) + mu_star, other_roots = kkt_roots(Q, x0) + markers = ([] if mu_star is None else [(mu_star, STAR_STYLE)]) + [(root, ROOT_STYLE) for root in other_roots] + + # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0 + inactive = [(*pole_line(Q, x0, ks[0]), degenerate_point(Q, x0, ks[0])) + for _, ks, is_active in poles if not is_active] + xds = [xd for _, _, xd in inactive if xd is not None] + if Q.is_conical and not Q.is_single_plane: # the apex is a KKT point + xds.append(Q.d) + xd_dists = [np.linalg.norm(xd - x0) for xd in xds] 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) + # (axis, function, label, dotted points of the middle branches, 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, + (ax_f, f, r'$f(\mu)$', 5, [0], None), + (ax_dist, lambda m: F.dist(x(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), + (ax_df, df, 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) + if degenerate: + n_dotted = 0 + else: + n_dotted = n_dotted_middle if 0 < i < len(branches) - 1 else N_DOTTED_OUTER colored_line(ax, mu[branch], values[branch], branch, cmap, norm, n_dotted) - for pole in poles[pole_active]: + for pole in pole_mus[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) + marker_values = [fun(root) for root, _ in markers if xlim[0] <= root <= xlim[1]] + ylim = auto_ylim(mu, values, branches, xlim, extra + marker_values + [fun(0)], floor) # x(0) = x0 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) + xs = np.array([x(m) for m in mu]) + kkt_points = [x(root) for root, _ in markers] + xds + has_centre = not (Q.is_parabolic or Q.is_cylindrical) # otherwise the centre is not unique + window_points = [x0, *kkt_points] + if has_centre: + window_points.append(Q.center) + if Q.is_central: # the vertices, so that the whole shape (e.g. both arcs of a hyperbola) is shown + window_points += [Q.to_non_normalized(sign * np.eye(Q.dim)[k] / np.sqrt(Q.eig[k])) + for k in np.flatnonzero(Q.eig > TOL) for sign in (-1, 1)] + window_xlim, window_ylim = square_window(window_points) + + draw_quadric(ax_x, Q, window_xlim, window_ylim) + if has_centre: + ax_x.scatter(*Q.center, color=QUADRIC_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) + for (_, style), point in zip(markers, kkt_points): + ax_x.scatter(*point, **style) if not degenerate: - for pole_index, label in zip(pole_indices, POLE_LABELS): + 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) + for _, ks, _ in poles: + point, direction = pole_line(Q, x0, ks[0]) 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) + else: + ax_x.plot(*np.column_stack((point, xd)), color='grey', zorder=0) + for xd, dist in zip(xds, xd_dists): 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') + ax_dist.text(xlim[1] + (pole_mus[-1] - xlim[1]) * 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 = [Line2D([], [], color=QUADRIC_COLOR, lw=2, label=r'$\mathcal{Q}(\Psi)$')] + if mu_star is not None: + legend.append(Line2D([], [], linestyle='none', marker='v', color='purple', markersize=6, + label=r'$\mathbf{x}(\mu^*)$')) + if other_roots: + legend.append(Line2D([], [], linestyle='none', marker='s', color='purple', markersize=6, + label=r'$\mathbf{x}(\mu)$, $f(\mu) = 0$')) + legend.append(Line2D([], [], linestyle='none', marker='x', color='red', markersize=8, + label=r'$\mathbf{x}^0 = \mathbf{x}(0)$')) + if xds: 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') + ax_x.legend(handles=legend, loc='best', fontsize=9) + ax_x.set(xlabel=r'$x_1$', ylabel=r'$x_2$', xlim=window_xlim, ylim=window_ylim, aspect='equal') fig.tight_layout(pad=2.1) colorbar = fig.colorbar(ScalarMappable(norm, cmap), ax=ax_x, ticks=[0, len(mu) - 1]) diff --git a/test_plot_mu_trajectories.py b/test_plot_mu_trajectories.py index a2c4a1d..43ce37c 100644 --- a/test_plot_mu_trajectories.py +++ b/test_plot_mu_trajectories.py @@ -1,44 +1,100 @@ +import contourpy +import matplotlib +import matplotlib.pyplot as plt import numpy as np import pytest from quadproj.fun import Fun from quadproj.project import project from quadproj.quadrics import Quadric -from plot_mu_trajectories import degenerate_point, ellipse_boundary, kkt_roots, pole_line +from plot_mu_trajectories import degenerate_point, kkt_roots, off_pole, plot_trajectories, pole_line -ROTATED_ELLIPSE = Quadric(np.array([[3.0, 1.0], [1.0, 2.0]]), np.array([1.0, -1.0]), -4.0) +matplotlib.use('Agg') + +GRID = np.linspace(-4, 4, 4001) + +# (A, b, c, x0): generic points, and points with a zero normalized component (degenerate) +QUADRICS = { + 'ellipse': ([[3, 1], [1, 2]], [1, -1], -4, [0.3, 0.2]), + 'ellipse_outside': ([[3, 1], [1, 2]], [1, -1], -4, [2.0, -1.5]), + 'ellipse_degenerate': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, -0.1]), + 'ellipse_degenerate_2': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, 0.5]), + 'hyperbola': ([[1, 2], [2, -1]], [1, 0], -1, [0.3, 0.2]), + 'hyperbola_degenerate': ([[1, 0], [0, -1]], [0, 0], -1, [0.0, 0.5]), + 'parabola': ([[1, 1], [1, 1]], [1, -1], 0, [0.3, 0.2]), + 'parabola_degenerate': ([[1, 0], [0, 0]], [0, -1], 0, [0.0, 0.5]), + 'crossing_lines': ([[1, 0], [0, -1]], [0, 0], 0, [0.3, 0.2]), + 'parallel_lines': ([[1, 0], [0, 0]], [0, 0], -1, [0.3, 0.2]), +} +# Quadrics without finite roots of f: the projection is the limit of x(mu) as |mu| -> inf. +LIMIT_QUADRICS = { + 'single_line': ([[1, 0], [0, 0]], [0, 0], 0, [0.3, 0.2]), + 'point': ([[1, 0], [0, 1]], [0, 0], 0, [0.3, 0.2]), +} + + +def make_quadric(A, b, c): + return Quadric(np.array(A, dtype=float), np.array(b, dtype=float), float(c)) + + +ROTATED_ELLIPSE = make_quadric(*QUADRICS['ellipse'][:3]) def trajectory(Q, x0): - return Fun(Q, Q.to_normalized(x0)).x_not_std + return off_pole(Fun(Q, Q.to_normalized(x0)).x_not_std) -def boundary_distances(Q, x0): - return np.linalg.norm(ellipse_boundary(Q, 200_000) - x0, axis=1) +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 distances_to_quadric(Q, x0): + """Distances from x0 to the points of the contour psi = 0 on a fine grid.""" + X, Y = np.meshgrid(GRID, GRID) + P = np.stack((X, Y), axis=-1) + psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c + return np.linalg.norm(np.vstack(contourpy.contour_generator(X, Y, psi).lines(0)) - x0, axis=1) + + +@pytest.mark.parametrize('name', QUADRICS) +def test_kkt_points_contain_the_projection(name): + *quadric, x0 = QUADRICS[name] + Q, x0 = make_quadric(*quadric), np.array(x0) + mu_star, other_roots = kkt_roots(Q, x0) + x = trajectory(Q, x0) + points = [x(mu) for mu in other_roots + ([] if mu_star is None else [mu_star])] + points += [xd for k in range(Q.dim) if (xd := degenerate_point(Q, x0, k)) is not None] + if Q.is_conical: + points.append(Q.d) + + assert all(abs(Q.evaluate_point(point)) < 1e-8 for point in points) + nearest = min(np.linalg.norm(point - x0) for point in points) + assert nearest == pytest.approx(distances_to_quadric(Q, x0).min(), abs=1e-4) + + +@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]]) +def test_ellipse_roots_give_nearest_and_farthest_points(x0): + Q, x0 = ROTATED_ELLIPSE, np.array(x0) + mu_star, other_roots = kkt_roots(Q, x0) + x = trajectory(Q, x0) + + distances = np.linalg.norm(ellipse_boundary(Q, 200_000) - x0, axis=1) + assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6) + assert max(np.linalg.norm(x(mu) - x0) for mu in other_roots) == pytest.approx(distances.max(), abs=1e-6) def test_boundary_is_on_the_ellipse(): assert np.all([abs(ROTATED_ELLIPSE.evaluate_point(x)) < 1e-10 for x in ellipse_boundary(ROTATED_ELLIPSE)]) -@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]]) -def test_kkt_roots_give_nearest_and_farthest_points(x0): - Q, x0 = ROTATED_ELLIPSE, np.array(x0) - mu_star, mu_star_star = kkt_roots(Q, x0) - x = trajectory(Q, x0) - - distances = boundary_distances(Q, x0) - assert abs(Q.evaluate_point(x(mu_star_star))) < 1e-8 - assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6) - assert np.linalg.norm(x(mu_star_star) - x0) == pytest.approx(distances.max(), abs=1e-6) - - @pytest.mark.parametrize('eig, x0', [([2.0, 1.0], [0.0, -0.1]), ([2.0, 0.75], [0.0, -0.6])]) def test_degenerate_point_is_the_projection(eig, x0): Q, x0 = Quadric(np.diag(eig), np.zeros(2), -1.0), np.array(x0) xd = degenerate_point(Q, x0, 0) np.testing.assert_allclose(xd, project(Q, x0)) - assert np.linalg.norm(xd - x0) == pytest.approx(boundary_distances(Q, x0).min(), abs=1e-6) assert degenerate_point(Q, x0, 1) is None @@ -46,5 +102,17 @@ def test_pole_line_is_the_asymptote(): Q, x0 = ROTATED_ELLIPSE, np.array([0.3, 0.2]) point, direction = pole_line(Q, x0, 1) x = trajectory(Q, x0)(-1 / Q.eig[1] + 1e-9) + direction = direction / np.linalg.norm(direction) residual = (x - point) - ((x - point) @ direction) * direction np.testing.assert_allclose(residual, 0, atol=1e-6) + + +@pytest.mark.parametrize('name', [*QUADRICS, *LIMIT_QUADRICS]) +def test_every_quadric_type_can_be_plotted(name): + *quadric, x0 = {**QUADRICS, **LIMIT_QUADRICS}[name] + plt.close(plot_trajectories(make_quadric(*quadric), np.array(x0, dtype=float))) + + +def test_empty_quadric_is_rejected(): + with pytest.raises(ValueError): + plot_trajectories(make_quadric([[1, 0], [0, 1]], [0, 0], 1), np.zeros(2))