update bounds to work for all cases with z > 0

docs/notebooks/boundaries.ipynb

@@ -93,35 +93,34 @@
     "\n",
     "#     return np.degrees(phi), z\n",
     "\n",
-    "\n",
     "def phi_z(k, kp):\n",
     "    # return np.degrees(polar(k, kp)[0]), polar(k, kp)[1]\n",
     "    eq_phi, eq_z = eq(k, kp)\n",
     "    polar_phi, polar_z = polar(k, kp)\n",
-    "    mask = np.logical_or(eq_z <= 2 / 3, np.abs(kp) + np.abs(k) <= nside)\n",
-    "    z = np.where(eq_z <= 2 / 3, eq_z, polar_z)\n",
-    "    phi = np.where(eq_z <= 2 / 3, eq_phi, polar_phi)\n",
-    "    phi = phi[mask]\n",
-    "    z = z[mask]\n",
+    "    z = np.where(eq_z <= 2/3, eq_z, polar_z)\n",
+    "    phi = np.where(eq_z <= 2/3, eq_phi, polar_phi)\n",
     "    return np.degrees(phi), z\n",
     "\n",
-    "\n",
     "def eq(k, kp):\n",
-    "    z = 2/3 * (2 - (abs(kp) + abs(k)) / nside)\n",
-    "    phi = np.pi/4/nside * (nside - abs(kp) + abs(k))\n",
-    "    phi = np.where(k >= 0, phi, -phi)\n",
-    "    phi = np.where(kp >= 0, phi, np.pi - phi)\n",
+    "    z = 2/3 * (2 - (kp + k) / nside)\n",
+    "    phi = np.pi/4/nside * (nside - kp + k)\n",
     "    return phi, z\n",
     "\n",
-    "\n",
     "def polar(k, kp):\n",
     "    j = np.abs(k) - 0.5\n",
     "    i = np.abs(kp) + np.abs(k)\n",
     "\n",
-    "    z = 1 - (i / nside) ** 2 / 3\n",
+    "    z = 1 - (i / nside)**2 / 3\n",
     "    phi = 0.5 * np.pi * (j + 0.5) / i\n",
+    "\n",
+    "    eq_phi, eq_z = eq(np.abs(k), np.abs(kp))\n",
+    "    \n",
+    "    z = np.where(np.abs(kp) + np.abs(k) <= nside, z, eq_z)\n",
+    "    phi = np.where(np.abs(kp) + np.abs(k) <= nside, phi, eq_phi)\n",
+    "    \n",
     "    phi = np.where(kp >= 0, phi, np.pi - phi)\n",
     "    phi = np.where(k >= 0, phi, -phi)\n",
+    "    \n",
     "    return phi, z\n",
     "\n",
     "\n",