7-spheres-3d

diff --git a/model.py b/model.py
@@ -28,6 +28,9 @@ class Model:
                 for i in range(x.shape[2]) for j in range(x.shape[1]))
 
     def write_obj(self, path: str):
+        """
+        Exports the model to the an OBJ file at `path`
+        """
         with open(path, "w") as f:
             vs = set()       # vertices
             eff_vis = list() # effective indixes
@@ -79,16 +82,17 @@ def normalized_p(z1: ComplexGrid, z2: ComplexGrid) -> Patch:
 
     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 riemman_sphere_p(z1: ComplexGrid, z2: ComplexGrid) -> Patch:
+    "Riemman-sphere transform + projection on the first 3 coordinates"
+    assert z1.shape == z2.shape
+
+    norm_sqrd = z1.real ** 2 + z1.imag ** 2 + z2.real ** 2 + z2.imag ** 2
+
+    u0 = 2 * (norm_sqrd/4)
+    u1 = z1.real
+    u2 = z1.imag
+
+    return np.vstack([u0, u1, u2])/(1 + norm_sqrd/4)
 
 def s(k: int, n: int) -> np.complex128:
     return np.exp(2 * np.pi * 1j * k / n)

diff --git /dev/null b/test.py
@@ -0,0 +1,51 @@
+import numpy as np
+import numpy.linalg as la
+
+from main import s
+
+EPS = 0.1
+XI_MAX = 15
+
+#
+#
+#patch = []
+#
+#for theta in thetas:
+#    u1 = (np.exp(xis + 1j * theta)+np.exp(-xis - 1j * theta))/2
+#    z1 = s(k1, n) * (u1**(2/n))
+#    u2 = (np.exp(xis + 1j * theta)-np.exp(-xis - 1j * theta))/2j
+#    z2 = s(k2, n) * (u2**(2/n))
+#
+#    row = np.vstack([z1, z2])
+#
+#    patch.append(row)
+#
+#m.append_patch(patch)
+
+#thetas = 
+#xis = 
+
+n = 5
+k1, k2 = 1, 1
+
+theta, xi = np.meshgrid(np.arange(0, np.pi / 2 + EPS, EPS),
+                        np.arange(-XI_MAX, XI_MAX + EPS, EPS))
+
+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)
+w = z/la.norm(z, axis=0)
+v = np.vstack([w.real, w.imag])
+
+x = v[:3, :, :]/(1 - v[3, :, :])
+
+print(x.shape,
+      len(np.arange(0, np.pi / 2 + EPS, EPS)),
+      len(np.arange(-XI_MAX, XI_MAX + EPS, EPS)))
+
+d, cols, rows = x.shape
+assert d == 3
+