fix: lambdify argument of sqrt

ComPWA · Jun 10, 2021 · 0559c94 · 0559c94
1 parent fc08ab8
commit 0559c94
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/docs/report/001.ipynb b/docs/report/001.ipynb
@@ -414,7 +414,9 @@
     "    _module = \"jax\"\n",
     "\n",
     "    def _print_ComplexSqrt(self, expr: sp.Expr) -> str:\n",
-    "        return \"select([less(x, 0), True], [1j * sqrt(-x), sqrt(x)], default=nan,)\""
+    "        arg = expr.args[0]\n",
+    "        x = self._print(arg)\n",
+    "        return f\"select([less({x}, 0), True], [1j * sqrt(-{x}), sqrt({x})], default=nan,)\""
    ]
   },
   {