From 7403d71b7af030efaf5cffcbd4826b005361446a Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Sun, 12 May 2024 21:04:51 +0200
Subject: [PATCH] DOC: explain `_latex_repr_()` as method implementation

---
 docs/usage/sympy.ipynb          | 48 ++++++++++++++++++++++++---------
 src/ampform/sympy/_decorator.py | 20 +++++++++++++-
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/docs/usage/sympy.ipynb b/docs/usage/sympy.ipynb
index 3526e4b97..88224d310 100644
--- a/docs/usage/sympy.ipynb
+++ b/docs/usage/sympy.ipynb
@@ -102,33 +102,55 @@
     "\n",
     "\n",
     "@unevaluated(real=False)\n",
-    "class PhspFactorSWave(sp.Expr):\n",
+    "class BreakupMomentum(sp.Expr):\n",
     "    s: sp.Symbol\n",
     "    m1: sp.Symbol\n",
     "    m2: sp.Symbol\n",
-    "    _latex_repr_ = R\"\\rho^\\text{{CM}}\\left({s}\\right)\"\n",
+    "    _latex_repr_ = R\"q\\left({s}\\right)\"  # not an f-string!\n",
     "\n",
     "    def evaluate(self) -> sp.Expr:\n",
     "        s, m1, m2 = self.args\n",
-    "        q = BreakupMomentum(s, m1, m2)\n",
-    "        cm = (\n",
-    "            (2 * q / sp.sqrt(s))\n",
-    "            * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n",
-    "            - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n",
-    "        ) / (16 * sp.pi**2)\n",
-    "        return 16 * sp.pi * sp.I * cm\n",
+    "        return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sympy.printing.latex import LatexPrinter\n",
     "\n",
     "\n",
     "@unevaluated(real=False)\n",
-    "class BreakupMomentum(sp.Expr):\n",
+    "class PhspFactorSWave(sp.Expr):\n",
     "    s: sp.Symbol\n",
     "    m1: sp.Symbol\n",
     "    m2: sp.Symbol\n",
-    "    _latex_repr_ = R\"q\\left({s}\\right)\"\n",
     "\n",
     "    def evaluate(self) -> sp.Expr:\n",
     "        s, m1, m2 = self.args\n",
-    "        return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))"
+    "        q = BreakupMomentum(s, m1, m2)\n",
+    "        cm = (\n",
+    "            (2 * q / sp.sqrt(s))\n",
+    "            * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n",
+    "            - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n",
+    "        ) / (16 * sp.pi**2)\n",
+    "        return 16 * sp.pi * sp.I * cm\n",
+    "\n",
+    "    def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n",
+    "        s = printer._print(self.s)\n",
+    "        s, *_ = map(printer._print, self.args)  # or via args\n",
+    "        return Rf\"\\rho^\\text{{CM}}\\left({s}\\right)\"  # f-string here!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{note}\n",
+    "For illustrative purposes, the phase space factor defines `_latex_repr_()` [as a printer method](https://docs.sympy.org/latest/modules/printing.html#example-of-custom-printing-method). It is recommended to do so only if rendering the expression class as $\\LaTeX$ requires more logics. The disadvantage of defining `_latex_repr_()` as a method is that it requires more boilerplate code, such as explicitly converting the symbolic {attr}`~sympy.core.basic.Basic.args` of the expression class first. In this phase space factor, defining `_latex_repr_` as a {class}`str` would have been just fine.\n",
+    ":::"
    ]
   },
   {
@@ -537,7 +559,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,
diff --git a/src/ampform/sympy/_decorator.py b/src/ampform/sympy/_decorator.py
index 5ecd0ffd0..557e8e80f 100644
--- a/src/ampform/sympy/_decorator.py
+++ b/src/ampform/sympy/_decorator.py
@@ -143,7 +143,7 @@ def unevaluated(
     >>> @unevaluated
     ... class Function(sp.Expr):
     ...     x: sp.Symbol
-    ...     _latex_repr_ = R"f\left({x}\right)"
+    ...     _latex_repr_ = R"f\left({x}\right)"  # not an f-string!
     ...
     ...     def evaluate(self) -> sp.Expr:
     ...         return sp.sqrt(self.x)
@@ -154,6 +154,24 @@ def unevaluated(
     >>> expr.doit()
     y
 
+    Or, `as a method <https://docs.sympy.org/latest/modules/printing.html#example-of-custom-printing-method>`_:
+
+    >>> from sympy.printing.latex import LatexPrinter
+    >>> @unevaluated
+    ... class Function(sp.Expr):
+    ...     x: sp.Symbol
+    ...
+    ...     def evaluate(self) -> sp.Expr:
+    ...         return self.x**2
+    ...
+    ...     def _latex_repr_(self, printer: LatexPrinter, *args) -> str:
+    ...         x = printer._print(self.x)  # important to convert to string first
+    ...         x, *_ = map(printer._print, self.args)  # also possible via its args
+    ...         return Rf"g\left({x}\right)"  # this is an f-string
+    >>> expr = Function(y)
+    >>> sp.latex(expr)
+    'g\\left(y\\right)'
+
     Attributes to the class are fed to the `~object.__new__` constructor of the
     :class:`~sympy.core.expr.Expr` class and are therefore also called "arguments". Just
     like in the :class:`~sympy.core.expr.Expr` class, these arguments are automatically