diff --git a/docs/usage.ipynb b/docs/usage.ipynb index 43e642294..918ccd40a 100644 --- a/docs/usage.ipynb +++ b/docs/usage.ipynb @@ -243,9 +243,10 @@ }, "outputs": [], "source": [ + "from ampform.io import aslatex\n", + "\n", "(symbol, expr), *_ = model.amplitudes.items()\n", - "latex = sp.multiline_latex(symbol, expr)\n", - "Math(latex)" + "Math(aslatex({symbol: expr}, terms_per_line=1))" ] }, { @@ -351,8 +352,16 @@ "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "version": "3.8.17" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/amplitude.ipynb b/docs/usage/amplitude.ipynb index 00f4a15df..d022e41a6 100644 --- a/docs/usage/amplitude.ipynb +++ b/docs/usage/amplitude.ipynb @@ -70,6 +70,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] @@ -225,30 +228,18 @@ "This shows that the main intensity is an **incoherent** sum of the amplitude for each spin projection combination of the initial and final states. The expressions for each of these amplitudes are provided with the {attr}`~.amplitudes` attribute. This is an {class}`~collections.OrderedDict`, so we can inspect the first of these amplitudes as follows:" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "(symbol, expression), *_ = model_no_dynamics.amplitudes.items()" - ] - }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ - "remove-input", "full-width" ] }, "outputs": [], "source": [ - "latex = sp.multiline_latex(symbol, expression)\n", - "Math(latex)" + "(symbol, expression), *_ = model_no_dynamics.amplitudes.items()\n", + "Math(aslatex({symbol: expression}, terms_per_line=1))" ] }, { @@ -262,9 +253,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -619,6 +608,21 @@ "model = model_builder.formulate()" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "(symbol, expression), *_ = model.amplitudes.items()\n", + "Math(aslatex({symbol: expression}, terms_per_line=1))" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -1078,7 +1082,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/dynamics.ipynb b/docs/usage/dynamics.ipynb index 67b6a1440..2f3a957e7 100644 --- a/docs/usage/dynamics.ipynb +++ b/docs/usage/dynamics.ipynb @@ -107,6 +107,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-input" ] @@ -160,8 +163,29 @@ "\n", "L = sp.Symbol(\"L\", integer=True)\n", "z = sp.Symbol(\"z\", real=True)\n", - "ff2 = BlattWeisskopfSquared(z, L)\n", - "Math(sp.multiline_latex(ff2, ff2.doit(), environment=\"eqnarray\"))" + "ff2 = BlattWeisskopfSquared(z, L)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "editable": true, + "jupyter": { + "source_hidden": true + }, + "slideshow": { + "slide_type": "" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "from ampform.io import aslatex\n", + "\n", + "Math(aslatex({ff2: ff2.doit()}))" ] }, { @@ -430,7 +454,7 @@ " meson_radius=1,\n", " phsp_factor=PhaseSpaceFactorSWave,\n", ")\n", - "Math(sp.multiline_latex(width, width.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({width: width.evaluate()}))" ] }, { @@ -651,8 +675,16 @@ "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "version": "3.8.17" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/dynamics/analytic-continuation.ipynb b/docs/usage/dynamics/analytic-continuation.ipynb index 1a82cfe01..b9ed3a207 100644 --- a/docs/usage/dynamics/analytic-continuation.ipynb +++ b/docs/usage/dynamics/analytic-continuation.ipynb @@ -102,6 +102,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-input" ] @@ -115,6 +118,8 @@ "import sympy as sp\n", "from IPython.display import Math\n", "\n", + "from ampform.io import aslatex\n", + "\n", "warnings.filterwarnings(\"ignore\")" ] }, @@ -144,7 +149,7 @@ "\n", "s, m_a, m_b = sp.symbols(\"s, m_a, m_b\", nonnegative=True)\n", "q_squared = BreakupMomentumSquared(s, m_a, m_b)\n", - "Math(sp.multiline_latex(q_squared, q_squared.doit(), environment=\"eqnarray\"))" + "Math(aslatex({q_squared: q_squared.evaluate()}))" ] }, { @@ -170,7 +175,7 @@ "from ampform.dynamics import PhaseSpaceFactor\n", "\n", "rho = PhaseSpaceFactor(s, m_a, m_b)\n", - "Math(sp.multiline_latex(rho, rho.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({rho: rho.evaluate()}))" ] }, { @@ -196,7 +201,7 @@ "from ampform.dynamics import PhaseSpaceFactorComplex\n", "\n", "rho_c = PhaseSpaceFactorComplex(s, m_a, m_b)\n", - "Math(sp.multiline_latex(rho_c, rho_c.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({rho_c: rho_c.evaluate()}))" ] }, { @@ -224,7 +229,7 @@ "from ampform.dynamics import EqualMassPhaseSpaceFactor\n", "\n", "rho_ac = EqualMassPhaseSpaceFactor(s, m_a, m_b)\n", - "Math(sp.multiline_latex(rho_ac, rho_ac.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({rho_ac: rho_ac.evaluate()}))" ] }, { @@ -250,7 +255,7 @@ "from ampform.dynamics import PhaseSpaceFactorAbs\n", "\n", "rho_hat = PhaseSpaceFactorAbs(s, m_a, m_b)\n", - "Math(sp.multiline_latex(rho_hat, rho_hat.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({rho_hat: rho_hat.evaluate()}))" ] }, { @@ -289,7 +294,7 @@ "from ampform.dynamics import PhaseSpaceFactorSWave\n", "\n", "rho_cm = PhaseSpaceFactorSWave(s, m_a, m_b)\n", - "Math(sp.multiline_latex(rho_cm, rho_cm.evaluate(), environment=\"eqnarray\"))" + "Math(aslatex({rho_cm: rho_cm.evaluate()}))" ] }, { diff --git a/docs/usage/dynamics/custom.ipynb b/docs/usage/dynamics/custom.ipynb index 4c676d46e..e5cfe7d92 100644 --- a/docs/usage/dynamics/custom.ipynb +++ b/docs/usage/dynamics/custom.ipynb @@ -70,6 +70,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] @@ -83,7 +86,9 @@ "import graphviz\n", "import qrules\n", "import sympy as sp\n", - "from IPython.display import display" + "from IPython.display import Math\n", + "\n", + "from ampform.io import aslatex" ] }, { @@ -155,9 +160,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -273,7 +276,7 @@ "metadata": {}, "outputs": [], "source": [ - "display(*sorted(model.parameter_defaults, key=lambda p: p.name))" + "Math(aslatex(model.parameter_defaults))" ] }, { @@ -305,9 +308,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -321,7 +322,7 @@ " heurisch=None,\n", " manual=None,\n", ")\n", - "integrated_expr.n(1)" + "Math(aslatex(integrated_expr.n(1), terms_per_line=1))" ] }, { @@ -350,6 +351,18 @@ "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/dynamics/k-matrix.ipynb b/docs/usage/dynamics/k-matrix.ipynb index 134d1cc22..dc6e642fa 100644 --- a/docs/usage/dynamics/k-matrix.ipynb +++ b/docs/usage/dynamics/k-matrix.ipynb @@ -104,6 +104,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] diff --git a/docs/usage/helicity/formalism.ipynb b/docs/usage/helicity/formalism.ipynb index e9003a404..5efbd656c 100644 --- a/docs/usage/helicity/formalism.ipynb +++ b/docs/usage/helicity/formalism.ipynb @@ -70,6 +70,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] diff --git a/docs/usage/helicity/spin-alignment.ipynb b/docs/usage/helicity/spin-alignment.ipynb index e880f5d04..93068a556 100644 --- a/docs/usage/helicity/spin-alignment.ipynb +++ b/docs/usage/helicity/spin-alignment.ipynb @@ -72,6 +72,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] @@ -191,8 +194,8 @@ " for i, (symbol, expr) in enumerate(model.amplitudes.items())\n", " if i % 5 == 0\n", " }\n", - " src = aslatex(selected_amplitudes)\n", - " src = src.replace(R\"\\end{array}\", R\"& \\dots \\\\ \\end{array}\")\n", + " src = aslatex(selected_amplitudes, terms_per_line=1)\n", + " src = src.replace(R\"\\end{array}\", R\"\\dots \\\\ \\end{array}\")\n", " return Math(src)\n", "\n", "\n", @@ -237,7 +240,7 @@ "builder_123.config.scalar_initial_state_mass = True\n", "builder_123.config.stable_final_state_ids = [1, 2, 3]\n", "dpd_model = builder_123.formulate()\n", - "dpd_model.intensity" + "dpd_model.intensity.cleanup()" ] }, { @@ -369,12 +372,7 @@ }, "outputs": [], "source": [ - "latex = sp.multiline_latex(\n", - " sp.Symbol(\"I\"),\n", - " axisangle_model.intensity.evaluate(),\n", - " environment=\"eqnarray\",\n", - ")\n", - "Math(latex)" + "Math(aslatex({\"I\": axisangle_model.intensity.evaluate()}, terms_per_line=1))" ] }, { @@ -429,7 +427,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/kinematics.ipynb b/docs/usage/kinematics.ipynb index 4930467f6..4431d4fb8 100644 --- a/docs/usage/kinematics.ipynb +++ b/docs/usage/kinematics.ipynb @@ -72,6 +72,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] @@ -428,7 +431,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.11.8" } }, "nbformat": 4, diff --git a/docs/usage/modify.ipynb b/docs/usage/modify.ipynb index e201e6904..ad18b2724 100644 --- a/docs/usage/modify.ipynb +++ b/docs/usage/modify.ipynb @@ -70,6 +70,9 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Import Python libraries" + }, "tags": [ "hide-cell" ] @@ -149,9 +152,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -181,9 +182,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -208,9 +207,7 @@ "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "full-width" - ] + "tags": [] }, "outputs": [], "source": [ @@ -338,8 +335,7 @@ "outputs": [], "source": [ "new_model.components[\n", - " R\"A_{J/\\psi(1S)_{-1} \\to {f_{0}(980)}_{0} \\gamma_{-1}; {f_{0}(980)}_{0}\"\n", - " R\" \\to \\pi^{0}_{0} \\pi^{0}_{0}}\"\n", + " R\"A_{J/\\psi(1S)_{-1} \\to {f_{0}(980)}_{0} \\gamma_{-1}; {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}\"\n", "]" ] }, @@ -390,8 +386,16 @@ "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "version": "3.8.17" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" } }, "nbformat": 4, diff --git a/src/ampform/io/__init__.py b/src/ampform/io/__init__.py index 22166e819..8003d7a4c 100644 --- a/src/ampform/io/__init__.py +++ b/src/ampform/io/__init__.py @@ -17,18 +17,25 @@ from collections import abc from functools import singledispatch -from typing import Iterable, Mapping +from typing import Iterable, Mapping, Sequence import sympy as sp @singledispatch -def aslatex(obj, **kwargs) -> str: +def aslatex(obj, **kwargs) -> str: # noqa: D417 """Render objects as a LaTeX `str`. The resulting `str` can for instance be given to `IPython.display.Math`. .. versionadded:: 0.14.1 + + Args: + terms_per_line: If set to a non-zero, positive number, + `sp.Expr ` objects on the right-hand-side with multiple + terms are split over multiple lines. The terms are split at the addition. + + .. versionadded:: 0.15.2 """ return str(obj) @@ -47,23 +54,61 @@ def __downcast(obj: float, **kwargs) -> float | int: return obj +@aslatex.register(str) +def _(obj: str, **kwargs) -> str: + return obj + + @aslatex.register(sp.Basic) def _(obj: sp.Basic, **kwargs) -> str: return sp.latex(obj) +@aslatex.register(sp.Expr) +def _(obj: sp.Expr, *, terms_per_line: int = 0, **kwargs) -> str: + terms = obj.as_ordered_terms() + if terms_per_line > 0 and len(terms) > terms_per_line: + return _render_broken_expression(terms, terms_per_line, **kwargs) + return sp.latex(obj) + + +def _render_broken_expression( + terms: Sequence[sp.Basic], terms_per_line: int, **kwargs +) -> str: + n = terms_per_line + groups = [sp.Add(*terms[i : i + n]) for i in range(0, len(terms), n)] + latex = R"\begin{array}{l}" + "\n" + latex += Rf" {aslatex(groups[0], **kwargs)} \\" + "\n" + for term in groups[1:]: + latex += Rf" \; + \; {aslatex(term, **kwargs)} \\" + "\n" + latex += R"\end{array}" + return latex + + @aslatex.register(abc.Mapping) -def _(obj: Mapping, **kwargs) -> str: +def _(obj: Mapping, *, terms_per_line: int = 0, **kwargs) -> str: if len(obj) == 0: msg = "Need at least one dictionary item" raise ValueError(msg) latex = R"\begin{array}{rcl}" + "\n" for lhs, rhs in obj.items(): - latex += Rf" {aslatex(lhs)} &=& {aslatex(rhs)} \\" + "\n" + latex += _render_row(lhs, rhs, terms_per_line, **kwargs) latex += R"\end{array}" return latex +def _render_row(lhs, rhs, terms_per_line: int, **kwargs) -> str: + if terms_per_line > 0 and isinstance(rhs, sp.Expr): + n = terms_per_line + terms = rhs.as_ordered_terms() + terms = [sum(terms[i : i + n]) for i in range(0, len(terms), n)] + row = _render_row(lhs, terms[0], terms_per_line=False) + for term in terms[1:]: + row += Rf" &+& {aslatex(term, **kwargs)} \\" + "\n" + return row + return Rf" {aslatex(lhs)} &=& {aslatex(rhs, **kwargs)} \\" + "\n" + + @aslatex.register(abc.Iterable) def _(obj: Iterable, **kwargs) -> str: obj = list(obj) diff --git a/tests/io/test_latex.py b/tests/io/test_latex.py index e234ba721..41aaa9014 100644 --- a/tests/io/test_latex.py +++ b/tests/io/test_latex.py @@ -1,5 +1,6 @@ from textwrap import dedent +import pytest import sympy as sp from ampform.io import aslatex @@ -13,6 +14,31 @@ def test_complex(): assert aslatex(1 + 1j) == "1+1i" +def test_expr(): + x, y, z = sp.symbols("x:z") + expr = x + y + z + assert aslatex(expr) == "x + y + z" + assert aslatex(expr, terms_per_line=0) == "x + y + z" + assert aslatex(expr, terms_per_line=3) == "x + y + z" + + expected = dedent(R""" + \begin{array}{l} + x \\ + \; + \; y \\ + \; + \; z \\ + \end{array} + """) + + assert aslatex(expr, terms_per_line=1) == expected.strip() + expected = dedent(R""" + \begin{array}{l} + x + y \\ + \; + \; z \\ + \end{array} + """) + assert aslatex(expr, terms_per_line=2) == expected.strip() + + def test_iterable(): items = [ a * x**2 + b, @@ -31,13 +57,14 @@ def test_iterable(): assert latex == dedent(expected).strip() -def test_mapping(): +@pytest.mark.parametrize("terms_per_line", [0, 2]) +def test_mapping(terms_per_line: int): definitions = { y: a * x**2 + b, a: 3.0, b: 2 - 1.3j, } - latex = aslatex(definitions) + latex = aslatex(definitions, terms_per_line=terms_per_line) expected = R""" \begin{array}{rcl} y &=& a x^{2} + b \\ @@ -46,3 +73,14 @@ def test_mapping(): \end{array} """ assert latex == dedent(expected).strip() + + latex = aslatex(definitions, terms_per_line=1) + expected = R""" + \begin{array}{rcl} + y &=& a x^{2} \\ + &+& b \\ + a &=& 3.0 \\ + b &=& 2-1.3i \\ + \end{array} + """ + assert latex == dedent(expected).strip()