Matplotlib 3D Animated L1 Norm Plot

norms.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import matplotlib.animation as animation
from itertools import product, combinations

# Inspiration: https://stackoverflow.com/a/11156353

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.set_aspect("equal")
ax.autoscale(True)
ax.autoscale_view(None, False, False, False)

def sample_unit_circle(p=np.inf,samples=100,color='r',marker='^',size=3):
  '''
    plot some 2D vectors with p-norm < 1
  '''
  for i in range(samples):
    xyz = np.array(
    [[np.random.rand()*2-1],[np.random.rand()*2-1],[np.random.rand()*2-1]]
    )
    if np.linalg.norm(xyz, p) < 1:
      ax.scatter(*xyz, c=color,marker=marker,s=size,depthshade=True)

def remove_decoration():
    transparent = (1.0, 1.0, 1.0, 0.0)

    ax.w_xaxis.set_pane_color(transparent)
    ax.w_yaxis.set_pane_color(transparent)
    ax.w_zaxis.set_pane_color(transparent)

    ax.w_xaxis.line.set_color(transparent)
    ax.w_yaxis.line.set_color(transparent)
    ax.w_zaxis.line.set_color(transparent)

    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])


remove_decoration()

# ax.set_xlabel('X')
# ax.set_ylabel('Y')
# ax.set_zlabel('Z')

def draw_l1(color="r"):
    r = [-1, 0, 1]
    cartesian_prod = np.array(list(product(r, r, r)))
    line_segments = combinations(cartesian_prod, 2)
    for p1, p2 in line_segments:
        diff = p1 - p2
        distance = np.sqrt(np.sum(diff ** 2))
        if (
            distance == np.sqrt(2)
            and np.sum(np.abs(p1) + np.abs(p2)) == 2
            and np.sum(np.abs(p1)) == 1
            and np.sum(np.abs(p2)) == 1
        ):
            ax.plot3D(*zip(p1, p2), color=color)


def draw_l2(color="g"):
    # u, v = np.mgrid[0:2 * np.pi:20j, 0:np.pi:10j]
    # x = np.cos(u) * np.sin(v)
    # y = np.sin(u) * np.sin(v)
    # z = np.cos(v)

    u = np.linspace(0, 2 * np.pi, 9)
    v = np.linspace(0, np.pi, 9)
    x = np.outer(np.cos(u), np.sin(v))
    y = np.outer(np.sin(u), np.sin(v))
    z = np.outer(np.ones(np.size(u)), np.cos(v))
    ax.plot_wireframe(x, y, z, color=color)


def draw_inf(color="b"):
    r = [-1, 1]
    cartesian_prod = np.array(list(product(r, r, r)))
    line_segments = combinations(cartesian_prod, 2)
    for s, e in line_segments:
        if np.sum(np.abs(s - e)) == r[1] - r[0]:
            ax.plot3D(*zip(s, e), color=color)


def rotate(pitch=30):
    for angle in range(0, 360):
        ax.view_init(pitch, angle)
        plt.draw()
        plt.pause(.01)


draw_l1()
draw_l2()
draw_inf()
#sample_unit_circle(1,color='r')
#sample_unit_circle(2,color='g')
#sample_unit_circle(color='b')

from pynput import keyboard
import draugr

COMBINATIONS = [
  {keyboard.Key.shift, keyboard.Key.alt, keyboard.KeyCode(char='s')},
  {keyboard.Key.shift, keyboard.Key.alt, keyboard.KeyCode(char='S')},
  ]

CALLBACKS = []
current = set()


def add_early_stopping_key_combination(callback, key='ctrl+c'):
  CALLBACKS.append(callback)
  draugr.sprint(f'\n\nPress any of:\n{COMBINATIONS}\n for early stopping\n', color='red', bold=True,
                highlight=True)
  print('')
  return keyboard.Listener(on_press=on_press, on_release=on_release)


def on_press(key):
  if any([key in COMBO for COMBO in COMBINATIONS]):
    current.add(key)
    if any(all(k in current for k in COMBO) for COMBO in COMBINATIONS):
      for callback in CALLBACKS:
        callback()

def on_release(key):
  if any([key in COMBO for COMBO in COMBINATIONS]):
    current.remove(key)


listener = add_early_stopping_key_combination(exit)

listener.start()
try:
  rotate(10)
finally:
  listener.stop()