FIX: create single branch with PosArg class (#14)

* MAINT: remove Colab's `cellView` metadata
ComPWA · Sep 13, 2023 · 32f5910 · 32f5910
1 parent 0eb0a8d
commit 32f5910
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 17 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -32,6 +32,7 @@ repos:
           - |
             cell.attachments
             cell.id
+            cell.metadata.cellView
             cell.metadata.code_folding
             cell.metadata.colab
             cell.metadata.editable

diff --git a/docs/lecture17.ipynb b/docs/lecture17.ipynb
@@ -294,31 +294,40 @@
     "    is_commutative = True\n",
     "    is_real = False\n",
     "\n",
-    "    def __new__(cls, z, sign, **hints) -> SignedSqrt:\n",
-    "        return create_expression(cls, z, sign, **hints)\n",
+    "    def __new__(cls, z, **hints) -> SignedSqrt:\n",
+    "        return create_expression(cls, z, **hints)\n",
     "\n",
     "    def evaluate(self) -> sp.Expr:\n",
-    "        z, sign = self.args\n",
-    "        return sp.sqrt(abs(z)) * sp.exp(sign * sp.I * sp.arg(z) / 2)\n",
+    "        z = self.args[0]\n",
+    "        return sp.sqrt(abs(z)) * sp.exp(sp.I * PosArg(z) / 2)\n",
     "\n",
     "    def _latex(self, printer, *args) -> str:\n",
     "        z = printer._print(self.args[0])\n",
-    "        sign = _render_sign(printer, self.args[1])\n",
-    "        return Rf\"\\sqrt[{sign}]{{{z}}}\"\n",
+    "        return Rf\"\\sqrt[+]{{{z}}}\"\n",
     "\n",
     "\n",
-    "def _render_sign(printer, sign) -> str:\n",
-    "    if sign == +1:\n",
-    "        return \"+\"\n",
-    "    if sign == -1:\n",
-    "        return \"-\"\n",
-    "    return printer._print(sign)\n",
+    "@implement_doit_method\n",
+    "class PosArg(UnevaluatedExpression):\n",
+    "    is_commutative = True\n",
+    "\n",
+    "    def __new__(cls, z, **hints) -> SignedSqrt:\n",
+    "        return create_expression(cls, z, **hints)\n",
+    "\n",
+    "    def evaluate(self) -> sp.Expr:\n",
+    "        z = self.args[0]\n",
+    "        arg = sp.arg(z)\n",
+    "        return sp.Piecewise(\n",
+    "            (arg + 2 * sp.pi, sp.im(z) < 0),\n",
+    "            (arg, True),\n",
+    "        )\n",
+    "\n",
+    "    def _latex(self, printer, *args) -> str:\n",
+    "        z = printer._print(self.args[0])\n",
+    "        return Rf\"\\arg^+\\left({z}\\right)\"\n",
     "\n",
     "\n",
     "z = sp.Symbol(\"z\", complex=True)\n",
-    "neg_sqrt = SignedSqrt(z, sign=-1)\n",
-    "pos_sqrt = SignedSqrt(z, sign=+1)\n",
-    "Math(aslatex({e: e.doit() for e in [neg_sqrt, pos_sqrt]}))"
+    "Math(aslatex({e: e.evaluate() for e in [SignedSqrt(z), PosArg(z)]}))"
    ]
   },
   {
@@ -337,8 +346,10 @@
    "outputs": [],
    "source": [
     "plot_riemann_surfaces(\n",
-    "    funcs=[sp.lambdify(z, pos_sqrt.doit())],\n",
+    "    funcs=[sp.lambdify(z, SignedSqrt(z).doit())],\n",
     "    func_unicode=\"⁺√z\",\n",
+    "    mask=lambda x, y: (np.abs(y) < 1e-5) & (x > 0),\n",
+    "    resolution=(30, 301),\n",
     ")"
    ]
   },
@@ -393,7 +404,7 @@
     "\n",
     "    def evaluate(self) -> sp.Expr:\n",
     "        s, m = self.args\n",
-    "        return sp.sqrt(1 - 4 * m**2 / s)\n",
+    "        return SignedSqrt(1 - 4 * m**2 / s)\n",
     "\n",
     "    def _latex(self, printer, *args) -> str:\n",
     "        s = printer._print(self.args[0])\n",