7-spheres-3d

diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,3 @@
 __pycache__
 site/
-models.blend
+models/

diff --git a/main.py /dev/null
@@ -1,89 +0,0 @@
-import numpy as np
-from model import Model, normalized_p, riemman_sphere_p
-
-def s(k: int, n: int) -> np.complex128:
-    return np.exp(2 * np.pi * 1j * k / n)
-
-EPS = 0.05
-XI_MAX = 15
-
-def brieskorn_surface(n: int,
-                      eps: float = EPS, xi_max: float = XI_MAX) -> Model:
-    """
-    Returns a model of the Brieskorn surface z_1^3 + z_2^(6 n - 1) = 1
-    """
-
-    if not (1 <= n <= 11):
-        raise ValueError(f"Parameter \"n\" should be in between 1 and 11: n = {n}")
-
-    m = Model(normalized_p)
-    n_ = 6*n - 1
-
-    thetas = np.arange(0, np.pi / 2 + eps, eps)
-    xis = np.arange(-xi_max, xi_max + eps, eps)
-
-    for k1 in range(0, 3):
-        for k2 in range(0, n_):
-            patch = []
-
-            for theta in thetas:
-                row = []
-                for xi in xis:
-                    u1 = (np.exp(xi + 1j * theta)+np.exp(-xi - 1j * theta))/2
-                    z1 = s(k1, 3) * (u1**(2/3))
-                    u2 = (np.exp(xi + 1j * theta)-np.exp(-xi - 1j * theta))/2j
-                    z2 = s(k2, n_) * (u2**(2/n_))
-
-                    row.append(np.array([z1, z2]))
-
-                patch.append(row)
-
-            m.append_patch(patch)
-
-    return m
-
-def fermat_surface(n: int, eps: float = EPS, xi_max: float = XI_MAX) -> Model:
-    """
-    Returns a model of the Fermat surface z_1^n + z_2^n = 1
-    """
-
-    if n <= 0:
-        raise ValueError(f"Parameter \"n\" should be positive: n = {n}")
-
-    m = Model(normalized_p)
-
-    thetas = np.arange(0, np.pi / 2 + eps, eps)
-    xis = np.arange(-xi_max, xi_max + eps, eps)
-
-    for k1 in range(0, n):
-        for k2 in range(0, n):
-            patch = []
-
-            for theta in thetas:
-                row = []
-                for xi in xis:
-                    u1 = (np.exp(xi + 1j * theta)+np.exp(-xi - 1j * theta))/2
-                    z1 = s(k1, n) * (u1**(2/n))
-                    u2 = (np.exp(xi + 1j * theta)-np.exp(-xi - 1j * theta))/2j
-                    z2 = s(k2, n) * (u2**(2/n))
-
-                    row.append(np.array([z1, z2]))
-
-                patch.append(row)
-
-            m.append_patch(patch)
-
-    return m
-
-def main():
-    # print("Creating Brieskorn model for n = 1...")
-    # m = brieskorn_surface(1, eps=0.01)
-    # m.write_obj("models/brieskorn1.obj")
-    # print("Done!")
-
-    print("Creating Brieskorn model for n = 6...")
-    m = brieskorn_surface(6, eps=0.01)
-    m.write_obj("models/brieskorn6.obj")
-    print("Done!")
-
-if __name__ == "__main__": main()

diff --git a/model.py b/model.py
@@ -1,70 +1,45 @@
-import numpy.linalg as la
 import numpy as np
+import numpy.linalg as la
 from typing import List, Tuple, Iterator, Callable
 
+Vector3     = np.ndarray[Tuple[3], np.dtype[np.float64]]
+Patch       = np.ndarray[Tuple[3, int, int], np.dtype[np.float64]]
+ComplexGrid = np.ndarray[Tuple[int, int], np.dtype[np.complex128]]
+Projection  = Callable[[ComplexGrid, ComplexGrid], Patch]
+
 class Model:
-    def __init__(self, p: Callable[[np.complex128, np.complex128], np.array]):
-        self.patches = []
-        self.bounding_x = (0, 0)
-        self.bounding_y = (0, 0)
-        self.bounding_z = (0, 0)
+    def __init__(self, p: Projection):
+        self.patches: List[Patch] = []
         self.projection = p
 
-    def append_patch(self, patch: List[List[np.array]]):
+    def append_patch(self, z1: ComplexGrid, z2: ComplexGrid):
         """
         Appends a 2D array of points in ℂ² to the model
         and updates the bounding box
         """
-        x_min, x_max = self.bounding_x
-        y_min, y_max = self.bounding_y
-        z_min, z_max = self.bounding_z
-
-        patch_real = [[self.projection(z1, z2) for z1, z2 in row]
-                      for row in patch]
-
-        for row in patch_real:
-            for x, y, z in row:
-                x_min = min(x, x_min)
-                x_max = max(x, x_max)
-                y_min = min(y, y_min)
-                y_max = max(y, y_max)
-                z_min = min(z, z_min)
-                z_max = max(z, z_max)
-
-        self.patches.append(patch_real)
-        self.bounding_x = x_min, x_max 
-        self.bounding_y = y_min, y_max 
-        self.bounding_z = z_min, z_max 
-
-    def vertices(self) -> Iterator[Tuple[np.float64]]:
+        self.patches.append(self.projection(z1, z2))
+
+    def vertices(self) -> Iterator[Vector3]:
         """
-        Iterate through the vertices of the mesh in normalized position:
-        vertices are aligned to the center using the model's bounding box
+        Iterate through the vertices of the mesh
         """
-        x_min, x_max = self.bounding_x
-        y_min, y_max = self.bounding_y
-        z_min, z_max = self.bounding_z
 
-        x_off = -(x_min + x_max)/2
-        y_off = -(y_min + y_max)/2
-        z_off = -(z_min + z_max)/2
-
-        return ((x + x_off, y + y_off, z + z_off)
-                for patch in self.patches for row in patch for x, y, z in row)
+        return (x[:, j, i] for x in self.patches
+                for i in range(x.shape[2]) for j in range(x.shape[1]))
 
     def write_obj(self, path: str):
         with open(path, "w") as f:
             f.write("# vertices\n")
 
-            for x, y, z in self.vertices():
-                f.write(f"v {x:.3f} {y:.3f} {z:.3f}\n")
+            # TODO: omit repeated points
+            for x1, x2, x3 in self.vertices():
+                f.write(f"v {x1:.3f} {x2:.3f} {x3:.3f}\n")
 
             f.write("\n# faces\n")
 
             v_i = 1 # OBJ indices start at 1
             for patch in self.patches:
-                rows = len(patch)
-                cols = len(patch[0])
+                _, cols, rows = patch.shape
 
                 for theta_i in range(rows - 1):
                     for xi_i in range(cols - 1):
@@ -78,23 +53,108 @@ class Model:
 
                 v_i += rows * cols
 
-def normalized_p(z1: np.complex128, z2: np.complex128) -> np.array:
+def normalized_p(z1: ComplexGrid, z2: ComplexGrid) -> Patch:
     "Projection ℝ⁴ → 𝕊³ + inclusion ℝ³ → 𝕊³"
-    v = np.array([z1, z2])
-    v = v/la.norm(v)
+    assert z1.shape == z2.shape
+
+    z = np.stack((z1, z2), axis=0)
+    w = z/la.norm(z, axis=0)
+    v = np.vstack([w.real, w.imag])
+    x = v[:3, :, :]/(1 - v[3, :, :])
+
+    return x
+
+# TODO: vectorize this shit
+# def riemman_sphere_p(z1: np.complex128, z2: np.complex128) -> Vector3:
+#     "Riemman-sphere transform + projection on the first 3 coordinates"
+#     v = np.array([z1.real, z1.imag, z2.real, z2.imag])
+#     n_sqrd = la.norm(v)**2
+# 
+#     u1, u2, _, _ = v / (1 + n_sqrd/4)
+#     u0 = np.float64(2 * (n_sqrd/4)/(1 + n_sqrd/4))
+# 
+#     return (u0, u1, u2)
+
+def s(k: int, n: int) -> np.complex128:
+    return np.exp(2 * np.pi * 1j * k / n)
+
+EPS = 0.1
+XI_MAX = 15
+
+def fermat_surface(n: int, eps: float = EPS, xi_max: float = XI_MAX) -> Model:
+    """
+    Returns a model of the Fermat surface z1^n + z2^n = 1
+    """
+    if n <= 0:
+        raise ValueError(f"Parameter \"n\" should be positive: n = {n}")
+
+    m = Model(normalized_p)
+
+    theta, xi = np.meshgrid(np.arange(0, np.pi / 2 + eps, eps),
+                            np.arange(-xi_max, xi_max + eps, eps))
+
+    for k1 in range(n):
+        for k2 in range(n):
+            u1 = (np.exp(xi + 1j * theta)+np.exp(-xi - 1j * theta))/2
+            z1 = s(k1, n) * (u1**(2/n))
+            u2 = (np.exp(xi + 1j * theta)-np.exp(-xi - 1j * theta))/2j
+            z2 = s(k2, n) * (u2**(2/n))
+            z = np.stack((z1, z2), axis=0)
+
+            m.append_patch(z)
+
+    return m
+
+def brieskorn_surface(n: int, patches: Iterator[Tuple[int, int]] = None,
+                      eps: float = EPS, xi_max: float = XI_MAX) -> Model:
+    """
+    Returns the model of the Brieskorn surface z1^3 + z2^(6 n - 1) = 1
+    """
+    if not (1 <= n <= 11):
+        raise ValueError(f"Parameter \"n\" should be in between 1 and 11: n = {n}")
+
+    k0 = 6*n - 1
+
+    if patches is None:
+        patches = ((k1, k2) for k1 in range(3) for k2 in range(k0))
+
+    m = Model(normalized_p)
+
+    theta, xi = np.meshgrid(np.arange(0, np.pi / 2 + eps, eps),
+                            np.arange(-xi_max, xi_max + eps, eps))
+
+    for k1, k2 in patches:
+        if not (0 <= k1 < 3) or not (0 <= k2 < k0): continue
+
+        u1 = (np.exp(xi + 1j * theta)+np.exp(-xi - 1j * theta))/2
+        z1 = s(k1, 3) * (u1**(2/3))
+        u2 = (np.exp(xi + 1j * theta)-np.exp(-xi - 1j * theta))/2j
+        z2 = s(k2, k0) * (u2**(2/k0))
+
+        m.append_patch(z1, z2)
+
+    return m
 
-    v_1, v_2 = v
-    x, y = v_1.real, v_1.imag
-    z, w = v_2.real, v_2.imag
+def main():
+    print("==> creating Brieskorn model for n = 1")
+    for k1 in range(3):
+        for k2 in range(5):
+            eps = 0.1
+            if k2 == 1: eps = 0.008
 
-    return np.array([x, y, z]) / (1 - w)
+            print(f"=>  creating patch ({k1}, {k2:2})...", end=" ", flush=True)
+            m = brieskorn_surface(1, patches = [(k1, k2)], eps=eps)
+            print("done!")
 
-def riemman_sphere_p(z1: np.complex128, z2: np.complex128) -> np.array:
-    "Riemman-sphere transform + projection on the first 3 coordinates"
-    v = np.array([z1.real, z1.imag, z2.real, z2.imag])
-    n_sqrd = la.norm(v)**2
+            print("=>  rendering model to OBJ...", end=" ", flush=True)
+            m.write_obj(f"models/brieskorn1/{k1}.{k2}.obj")
+            print("done!")
 
-    u1, u2, _, _ = v / (1 + n_sqrd/4)
-    u0 = 2 * (n_sqrd/4)/(1 + n_sqrd/4)
+    # print("=> creating Brieskorn model for n = 6...", end=" ", flush=True)
+    # m = brieskorn_surface(6, eps=0.05)
+    # print("done!")
+    # print("=> rendering model to OBJ...", end=" ", flush=True)
+    # m.write_obj("models/brieskorn6.obj")
+    # print("done!")
 
-    return np.array([u0, u1, u2])
+if __name__ == "__main__": main()