donut code in python

1#CREDITS: https://gist.github.com/Denbergvanthijs/7f6936ca90a683d37216fd80f5750e9c
2
3import numpy as np
4
5screen_size = 40
6theta_spacing = 0.07
7phi_spacing = 0.02
8illumination = np.fromiter(".,-~:;=!*#$@", dtype="<U1")
9
10A = 1
11B = 1
12R1 = 1
13R2 = 2
14K2 = 5
15K1 = screen_size * K2 * 3 / (8 * (R1 + R2))
16
17
18def render_frame(A: float, B: float) -> np.ndarray:
19    """
20    Returns a frame of the spinning 3D donut.
21    Based on the pseudocode from: https://www.a1k0n.net/2011/07/20/donut-math.html
22    """
23    cos_A = np.cos(A)
24    sin_A = np.sin(A)
25    cos_B = np.cos(B)
26    sin_B = np.sin(B)
27
28    output = np.full((screen_size, screen_size), " ")  # (40, 40)
29    zbuffer = np.zeros((screen_size, screen_size))  # (40, 40)
30
31    cos_phi = np.cos(phi := np.arange(0, 2 * np.pi, phi_spacing))  # (315,)
32    sin_phi = np.sin(phi)  # (315,)
33    cos_theta = np.cos(theta := np.arange(0, 2 * np.pi, theta_spacing))  # (90,)
34    sin_theta = np.sin(theta)  # (90,)
35    circle_x = R2 + R1 * cos_theta  # (90,)
36    circle_y = R1 * sin_theta  # (90,)
37
38    x = (np.outer(cos_B * cos_phi + sin_A * sin_B * sin_phi, circle_x) - circle_y * cos_A * sin_B).T  # (90, 315)
39    y = (np.outer(sin_B * cos_phi - sin_A * cos_B * sin_phi, circle_x) + circle_y * cos_A * cos_B).T  # (90, 315)
40    z = ((K2 + cos_A * np.outer(sin_phi, circle_x)) + circle_y * sin_A).T  # (90, 315)
41    ooz = np.reciprocal(z)  # Calculates 1/z
42    xp = (screen_size / 2 + K1 * ooz * x).astype(int)  # (90, 315)
43    yp = (screen_size / 2 - K1 * ooz * y).astype(int)  # (90, 315)
44    L1 = (((np.outer(cos_phi, cos_theta) * sin_B) - cos_A * np.outer(sin_phi, cos_theta)) - sin_A * sin_theta)  # (315, 90)
45    L2 = cos_B * (cos_A * sin_theta - np.outer(sin_phi, cos_theta * sin_A))  # (315, 90)
46    L = np.around(((L1 + L2) * 8)).astype(int).T  # (90, 315)
47    mask_L = L >= 0  # (90, 315)
48    chars = illumination[L]  # (90, 315)
49
50    for i in range(90):
51        mask = mask_L[i] & (ooz[i] > zbuffer[xp[i], yp[i]])  # (315,)
52
53        zbuffer[xp[i], yp[i]] = np.where(mask, ooz[i], zbuffer[xp[i], yp[i]])
54        output[xp[i], yp[i]] = np.where(mask, chars[i], output[xp[i], yp[i]])
55
56    return output
57
58
59def pprint(array: np.ndarray) -> None:
60    """Pretty print the frame."""
61    print(*[" ".join(row) for row in array], sep="\n")
62
63
64if __name__ == "__main__":
65    for _ in range(screen_size * screen_size):
66        A += theta_spacing
67        B += phi_spacing
68        print("\x1b[H")
69        pprint(render_frame(A, B))