plot_dual_norms.py

"""
Dual p-norms illustrated.

For any norm |.|, the dual norm is defined as |y|_* = max{ <x, y> for |x| <= 1 }.

The figure shows the unit balls of the p-norm, for p = 1.5, 2, and 3.

We compute the dual norm at a dual vector y (short black arrow), rotating
uniformly around the origin over time.

The surface shows the contour lines of <y, x> for all x in the ball, and the
longer arrow is the optimal primal vector x, maximizing <y, x>.

P.S. the 2-norm is self dual, while 1.5 and 3 are dual to each other.
"""

# author: Vlad Niculae <vlad@vene.ro>
# license: MIT

import numpy as np
import matplotlib.pyplot as plt

from scipy.optimize import root_scalar

from tqdm import tqdm


def dual_norm_argmax_2d(y, p=1.5):

    # the maximizer x in the dual norm |y|_*
    # takes the form [x, (1-x^p)^1/p]

    # we can solve for x using root finding.

    if np.abs(y[0]) < 1e-15 or np.abs(y[1]) < 1e-15:
        return y / np.linalg.norm(y)

    sgn = np.sign(y)

    y = y * sgn  # make y nonnegative

    def f(x):
        return y[0] - y[1] * ((x ** (p-1)) * (1 - x**p) ** (1/p - 1))

    eps = 1e-15
    res = root_scalar(f, method='brentq', bracket=(0 + eps, 1 - eps))
    x0 = res.root
    x_best = np.array([x0, (1 - x0 ** p) ** (1/p)])
    return x_best * sgn



if __name__ == '__main__':
    # make a 2d mesh
    N = 300

    t = np.linspace(-1.1, 1.1, num=N)
    XX, YY = np.meshgrid(t, t)
    pts = np.column_stack([np.ravel(XX), np.ravel(YY)])

    ps = [1.5, 2, 3]

    for k, phi in enumerate(tqdm(np.linspace(0, 2 * np.pi, num=500))):

        # pick dual vector (rotated around origin)
        y = .33 * np.array([np.sin(phi), np.cos(phi)])

        fig, axes = plt.subplots(1, len(ps),
                                 constrained_layout=True,
                                 figsize=(8,3))

        for i in range(len(ps)):
            p = ps[i]
            norms = np.sum(np.abs(pts) ** p, axis=1) ** (1/p)
            norms = norms.reshape(N, N)

            # plot contour of primal norm ball
            axes[i].contour(XX, YY, norms, levels=[1])
            axes[i].set_aspect("equal")

            # get points inside primal norm ball
            pts_inside_ball = pts[(norms <= 1).ravel()]

            # compute objective value for each of them
            obj = pts_inside_ball @ y

            # now, we could approximately pick the best x from the mesh.
            # i.e.,
            # best_x = pts_inside_ball[np.argmax(obj)]
            # but this creates some rendering artifacts around
            # non-differentiable points. So instead,
            best_x = dual_norm_argmax_2d(y, p=p)

            # repackage obj values in an x/y mesh grid and plot contours
            objs = np.full_like(norms, -np.inf)
            objs[norms <= 1] = obj
            axes[i].contourf(XX, YY, objs,
                             levels=np.linspace(-.5, .5, 11))

            # plot optimal primal vector, and dual vector
            axes[i].quiver(0, 0, best_x[0], best_x[1], units='xy', scale=1, width=.05, color='C5')
            axes[i].quiver(0, 0, y[0], y[1], units='xy', scale=1, width=.05)
            axes[i].set_title(f"p={p:.2f}")

        plt.savefig(f"{k:04d}.png")
        plt.close(fig)

out.mp4