FEAT: implement cached.simplify etc (#467)

* FEAT: implement cached simplify and trigsimp
* FEAT: implement `cached.subs`
* FIX: add back `cached.doit()` to API

Author: redeboer

docs/conf.py

@@ -107,6 +107,7 @@ def _get_dataclasses(module):
     "P": "typing.ParamSpec",
     "ParameterValue": ("obj", "ampform.helicity.ParameterValue"),
     "Particle": "qrules.particle.Particle",
+    "SympyObject": "typing.TypeVar",
     "ReactionInfo": "qrules.transition.ReactionInfo",
     "Slider": ("obj", "symplot.Slider"),
     "State": "qrules.transition.State",
@@ -323,10 +324,12 @@ def _get_dataclasses(module):
 nb_execution_timeout = -1
 nb_output_stderr = "remove"
 nitpick_ignore = [
+    ("py:class", "ampform.sympy._decorator.SymPyAssumptions"),
     ("py:class", "ArraySum"),
+    ("py:class", "BufferedReader"),
     ("py:class", "ExprClass"),
     ("py:class", "MatrixMultiplication"),
-    ("py:class", "ampform.sympy._decorator.SymPyAssumptions"),
+    ("py:class", "SupportsWrite"),
 ]
 nitpicky = True
 primary_domain = "py"

src/ampform/sympy/_cache.py

@@ -1,4 +1,11 @@
-"""Helper functions for :func:`.cached.doit` and related functions."""
+"""Helper functions for :func:`.cached.doit` and related functions.
+
+These methods are private, but can be imported from this module:
+
+.. code-block:: python
+
+   import ampform.sympy._cache
+"""
 
 from __future__ import annotations
 

src/ampform/sympy/cached.py

@@ -1,4 +1,7 @@
-"""Handy aliases for working with cached SymPy expressions."""
+"""Handy aliases for working with cached SymPy expressions.
+
+.. autofunction:: doit
+"""
 
 from __future__ import annotations
 
@@ -35,6 +38,37 @@ def doit(expr: SympyObject) -> SympyObject:
     return expr.doit()
 
 
+@cache
+@cache_to_disk
+def simplify(expr: sp.Expr, *args, **kwargs) -> sp.Expr:
+    """Perform :func:`~sympy.simplify.simplify.simplify` and cache the result to disk.
+
+    .. versionadded:: 0.15.7
+    """
+    return sp.simplify(expr, *args, **kwargs)
+
+
+@cache
+@cache_to_disk
+def trigsimp(expr: sp.Expr, *args, **kwargs) -> sp.Expr:
+    """Perform :func:`~sympy.simplify.trigsimp.trigsimp` and cache the result to disk.
+
+    .. versionadded:: 0.15.7
+    """
+    return sp.trigsimp(expr, *args, **kwargs)
+
+
+def subs(expr: sp.Expr, substitutions: Mapping[sp.Basic, sp.Basic]) -> sp.Expr:
+    """Call :meth:`~sympy.core.basic.Basic.subs` and cache the result to disk."""
+    return _subs_impl(expr, frozendict(substitutions))
+
+
+@cache
+@cache_to_disk(function_name="subs", dependencies=["sympy"])
+def _subs_impl(expr: sp.Expr, substitutions: frozendict[sp.Basic, sp.Basic]) -> sp.Expr:
+    return expr.xreplace(substitutions)
+
+
 def xreplace(expr: sp.Expr, substitutions: Mapping[sp.Basic, sp.Basic]) -> sp.Expr:
     """Call :meth:`~sympy.core.basic.Basic.xreplace` and cache the result to disk."""
     return _xreplace_impl(expr, frozendict(substitutions))

tests/sympy/test_cached.py

@@ -1,14 +1,14 @@
+# cspell:ignore pbarksigma
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
 import pytest
+import sympy as sp
 
 from ampform.sympy import cached
 
 if TYPE_CHECKING:
-    import sympy as sp
-
     from ampform.helicity import HelicityModel
 
 
@@ -23,19 +23,69 @@ def test_doit(amplitude_model: tuple[str, HelicityModel]):
     assert unfolded_expr_2 == expected_expr
 
 
+def test_simplify():
+    a, b, c, d, x, y, z = sp.symbols("a b c d x y z")
+    expr = (
+        (a * x + b * y + c * z + d) ** 2
+        - (a * x) ** 2
+        - (b * y) ** 2
+        - (c * z) ** 2
+        - 2 * a * b * x * y
+        - 2 * a * c * x * z
+        - 2 * b * c * y * z
+        - 2 * d * (a * x + b * y + c * z)
+    )
+
+    simplified = expr.simplify()
+    assert simplified != expr
+    assert simplified == d**2
+
+    cached_simplified = cached.simplify(expr)
+    assert simplified == cached_simplified
+
+
+@pytest.mark.parametrize("amplitude_idx", list(range(4)))
+def test_simplify_model(amplitude_model: tuple[str, HelicityModel], amplitude_idx: int):
+    _, model = amplitude_model
+    amplitudes = [model.amplitudes[k] for k in sorted(model.amplitudes, key=str)]
+    amplitude_expr = amplitudes[amplitude_idx]
+
+    simplified = amplitude_expr.simplify()
+    assert simplified != amplitude_expr
+
+    cached_simplified = cached.simplify(amplitude_expr)
+    assert simplified == cached_simplified
+
+
+def test_trigsimp():
+    x, y = sp.symbols("x y")
+    expr = (sp.sin(x) * sp.cos(y) + sp.cos(x) * sp.sin(y)) ** 2 + (
+        sp.cos(x) * sp.cos(y) - sp.sin(x) * sp.sin(y)
+    ) ** 2
+    simplified = sp.trigsimp(expr)
+    assert expr != simplified
+    assert simplified == 1
+    cached_simplified = cached.trigsimp(expr)
+    assert simplified == cached_simplified
+
+
+@pytest.mark.parametrize("substitute", ["subs", "xreplace"])
 @pytest.mark.parametrize(
     "substitution_name", ["parameter_defaults", "kinematic_variables"]
 )
-def test_xreplace(amplitude_model: tuple[str, HelicityModel], substitution_name: str):
+def test_xreplace(
+    amplitude_model: tuple[str, HelicityModel], substitute: str, substitution_name: str
+):
+    cached_func = getattr(cached, substitute)
     _, model = amplitude_model
     full_expression = model.expression.doit()
     substitutions: dict[sp.Symbol, sp.Basic] = getattr(model, substitution_name)
     expected_expr = full_expression.xreplace(substitutions)
     assert expected_expr != full_expression
 
-    substituted_expr_1 = cached.xreplace(full_expression, substitutions)
+    substituted_expr_1 = cached_func(full_expression, substitutions)
     assert substituted_expr_1 == expected_expr
-    substituted_expr_2 = cached.xreplace(full_expression, substitutions)
+    substituted_expr_2 = cached_func(full_expression, substitutions)
     assert substituted_expr_2 == expected_expr