Simplify construction of foldable expression with PEP 681

[redeboer]

Currently, AmpForm provides the following mechanism for defining 'foldable' expression classes (UnevaluatedExpression) that make it easier to read large expressions that consist of these folded definitions (see also {doc}/adr/002). (example taken from ComPWA/compwa.github.io#204, see also for example compwa.github.io/polarimetry/appendix/dynamics.html):

import sympy as sp
from ampform.sympy import UnevaluatedExpression, create_expression, implement_doit_method


@implement_doit_method
class ChewMandelstamm(UnevaluatedExpression):
    is_commutative = True
    is_real = False

    def __new__(cls, s, m1, m2, **hints) -> ChewMandelstamm:
        return create_expression(cls, s, m1, m2, *hints)

    def evaluate(self) -> sp.Expr:
        s, m1, m2 = self.args
        q = BreakupMomentum(s, m1, m2)
        return (
            1
            / (16 * (sp.pi) ** 2)
            * (
                (2 * q / sp.sqrt(s))
                * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))
                - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)
            )
        )

    def _latex(self, printer, *args) -> str:
        s = printer._print(self.args[0])
        return Rf"\rho\left({s}\right)"


@implement_doit_method
class BreakupMomentum(UnevaluatedExpression):
    is_commutative = True

    def __new__(cls, s, m_a, m2, **hints) -> BreakupMomentum:
        return create_expression(cls, s, m_a, m2, **hints)

    def evaluate(self) -> sp.Expr:
        s, m_a, m2 = self.args
        return sp.sqrt((s - (m_a + m2) ** 2) * (s - (m_a - m2) ** 2) / (s * 4))

    def _latex(self, printer, *args) -> str:
        s = printer._print(self.args[0])
        return Rf"q\left({s}\right)"

from ampform.io import aslatex
from IPython.display import Math

s, m1, m2 = sp.symbols("s m1 m2")
expressions = [
    ChewMandelstamm(s, m1, m2),
    BreakupMomentum(s, m1, m2),
]
Math(aslatex({expr: expr.doit(deep=False) for expr in expressions}))

The main problem is that there is way to much boilerplate code here. Ideally, the above example should be rewritten into something like this, where unevaluated_expression behaves just like attrs.define():

from ampform.sympy import unevaluated_expression


@unevaluated_expression(commutative=True, real=True)
class ChewMandelstamm:
    s: sp.Basic
    m1: sp.Basic
    m2: sp.Basic

    def _implementation_(self) -> sp.Expr:
        s, m1, m2 = self.args
        q = BreakupMomentum(s, m1, m2)
        return (
            1
            / (16 * (sp.pi) ** 2)
            * (
                (2 * q / sp.sqrt(s))
                * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))
                - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)
            )
        )

    def _latex_repr_(self) -> str:
        s = self.latex_args[0]
        return Rf"\rho\left({s}\right)"


@unevaluated_expression(commutative=True)
class BreakupMomentum:
    s: sp.Basic
    m1: sp.Basic
    m2: sp.Basic

    def _implementation_(self) -> sp.Expr:
        s, m_a, m2 = self.args
        return sp.sqrt((s - (m_a + m2) ** 2) * (s - (m_a - m2) ** 2) / (s * 4))

    def _latex_repr_(self) -> str:
        s = self.latex_args[0]
        return Rf"q\left({s}\right)"

Of course, even nicer would be something like this:

@unevaluated_expression(commutative=True)
class BreakupMomentum:
    s: sp.Basic
    m1: sp.Basic
    m2: sp.Basic
    _latex_repr_ = R"q\left({s}\right)"

    def _implementation_(self) -> sp.Expr:
        return sp.sqrt((s - (m_a + m2) ** 2) * (s - (m_a - m2) ** 2) / (s * 4))

but that would be hard, if not impossible, to implement.

We probably have to use typing.dataclass_transform, because we want the unevaluated_expression decorator to give the class 'dataclass'-like behavior (see PEP 681).

Tasks

Preview Give feedback

1. Investivate typing.dataclass_transform
2. Implement __new__ through class decorator that gets arguments from class annotations
3. Implement assumptions setter arguments for unevaluated_expression decorator
4. Generate _latex method from _latex_repr_ string
5. Support typing.Unpack in Sphinx builds policy#243
   🐛 Bug 🖱️ DX +2
   
6. Implement __slots__ mechanism for phsp_factor etc
7. Implement self.latex_args for simpler _latex_repr_ method implementation
8. Investigate abstract base class for typing of unevaluated_expression-dectorated classes
9. Replace UnevaluatedExpr in documentation with unevaluated_expression decorator
10. Implement name mechanism from UnevaluatedClass
11. Add deprecation warnings for UnevaluatedExpr etc