Symbolic model serialization#
Import Python libraries
from pathlib import Path
from textwrap import shorten
import graphviz
import polarimetry
import sympy as sp
from ampform.io import aslatex
from ampform.sympy import unevaluated
from IPython.display import Markdown, Math
from polarimetry.amplitude import simplify_latex_rendering
from polarimetry.io import perform_cached_doit
from polarimetry.lhcb import load_model
from polarimetry.lhcb.particle import load_particles
from sympy.printing.mathml import MathMLPresentationPrinter
simplify_latex_rendering()
Expression trees#
SymPy expressions are built up from symbols and mathematical operations as follows:
x, y, z = sp.symbols("x y z")
expression = sp.sin(x * y) / 2 - x**2 + 1 / z
expression
In the back, SymPy represents these expressions as trees. There are a few ways to visualize this for this particular example:
sp.printing.tree.print_tree(expression, assumptions=False)
Add: -x**2 + sin(x*y)/2 + 1/z
+-Pow: 1/z
| +-Symbol: z
| +-NegativeOne: -1
+-Mul: sin(x*y)/2
| +-Half: 1/2
| +-sin: sin(x*y)
| +-Mul: x*y
| +-Symbol: x
| +-Symbol: y
+-Mul: -x**2
+-NegativeOne: -1
+-Pow: x**2
+-Symbol: x
+-Integer: 2
Show code cell source
src = sp.dotprint(
expression,
styles=[
(sp.Number, {"color": "grey", "fontcolor": "grey"}),
(sp.Symbol, {"color": "royalblue", "fontcolor": "royalblue"}),
],
)
graphviz.Source(src)
Expression trees are powerful, because we can use them as templates for any human-readable presentation we are interested in. In fact, the LaTeX representation that we saw when constructing the expression was generated by SymPy’s LaTeX printer.
src = sp.latex(expression)
Markdown(f"```latex\n{src}\n```")
- x^{2} + \frac{\sin{\left(x y \right)}}{2} + \frac{1}{z}
Hint
SymPy expressions can serve as a template for generating code!
Here’s a number of other representations:
Show code cell source
def to_mathml(expr: sp.Expr) -> str:
printer = MathMLPresentationPrinter()
xml = printer._print(expr)
return xml.toprettyxml().replace("\t", " ")
Markdown(
f"""
```python
# Python
{sp.pycode(expression)}
```
```cpp
// C++
{sp.cxxcode(expression, standard="c++17")}
```
```fortran
! Fortran
{sp.fcode(expression).strip()}
```
```matlab
% Matlab / Octave
{sp.octave_code(expression)}
```
```julia
# Julia
{sp.julia_code(expression)}
```
```rust
// Rust
{sp.rust_code(expression)}
```
```xml
<!-- MathML -->
{to_mathml(expression)}
```
"""
)
# Python
-x**2 + (1/2)*math.sin(x*y) + 1/z
// C++
-std::pow(x, 2) + (1.0/2.0)*std::sin(x*y) + 1.0/z
! Fortran
-x**2 + (1.0d0/2.0d0)*sin(x*y) + 1d0/z
% Matlab / Octave
-x.^2 + sin(x.*y)/2 + 1./z
# Julia
-x .^ 2 + sin(x .* y) / 2 + 1 ./ z
// Rust
-x.powi(2) + (1_f64/2.0)*(x*y).sin() + z.recip()
<!-- MathML -->
<mrow>
<mrow>
<mo>-</mo>
<msup>
<mi>x</mi>
<mn>2</mn>
</msup>
</mrow>
<mo>+</mo>
<mrow>
<mfrac>
<mrow>
<mi>sin</mi>
<mfenced>
<mrow>
<mi>x</mi>
<mo>⁢</mo>
<mi>y</mi>
</mrow>
</mfenced>
</mrow>
<mn>2</mn>
</mfrac>
</mrow>
<mo>+</mo>
<mfrac>
<mn>1</mn>
<mi>z</mi>
</mfrac>
</mrow>
Foldable expressions#
The previous example is quite simple, but SymPy works just as well with huge expressions, as we will see in Large expressions. Before, though, let’s have a look how to define these larger expressions in such a way that we can still read them. A nice solution is to define sp.Expr classes with the @unevaluated decorator (see ComPWA/ampform#364). Here, we define a Chew-Mandelstam function \(\rho^\text{CM}\) for \(S\)-waves. This function requires the definition of a break-up momentum \(q\).
@unevaluated(real=False)
class PhspFactorSWave(sp.Expr):
s: sp.Symbol
m1: sp.Symbol
m2: sp.Symbol
_latex_repr_ = R"\rho^\text{{CM}}\left({s}\right)"
def evaluate(self) -> sp.Expr:
s, m1, m2 = self.args
q = BreakupMomentum(s, m1, m2)
cm = (
(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)
) / (16 * sp.pi**2)
return 16 * sp.pi * sp.I * cm
@unevaluated(real=False)
class BreakupMomentum(sp.Expr):
s: sp.Symbol
m1: sp.Symbol
m2: sp.Symbol
_latex_repr_ = R"q\left({s}\right)"
def evaluate(self) -> sp.Expr:
s, m1, m2 = self.args
return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))
We now have a very clean mathematical representation of how the \(\rho^\text{CM}\) function is defined in terms of \(q\):
Now, let’s build up a more complicated expression that contains this phase space factor. Here, we use SymPy to derive a Breit-Wigner using a single-channel \(K\) matrix [Chung et al., 1995]:
I = sp.Identity(n=1)
K = sp.MatrixSymbol("K", m=1, n=1)
ρ = sp.MatrixSymbol("rho", m=1, n=1)
T = (I - sp.I * K * ρ).inv() * K
T
T.as_explicit()[0, 0]
Here we need to provide definitions for the matrix elements of \(K\) and \(\rho\). A suitable choice is our phase space factor for \(S\) waves we defined above:
m0, Γ0, γ0 = sp.symbols("m0 Gamma0 gamma0")
K_expr = (γ0**2 * m0 * Γ0) / (s - m0**2)
And there we have it! After some algebraic simplifications, we get a Breit-Wigner with Chew-Mandelstam phase space factor for \(S\) waves:
T_expr = T.as_explicit().xreplace(substitutions)
BW_expr = T_expr[0, 0].simplify(doit=False)
BW_expr
The expression tree now has a node that is ‘folded’:
Show code cell source
dot_style = [
(sp.Basic, {"style": "filled", "fillcolor": "white"}),
(sp.Atom, {"color": "gray", "style": "filled", "fillcolor": "white"}),
(sp.Symbol, {"color": "dodgerblue1"}),
(PhspFactorSWave, {"color": "indianred2"}),
]
dot = sp.dotprint(BW_expr, bgcolor=None, styles=dot_style)
graphviz.Source(dot)
After unfolding, we get the full expression tree of fundamental mathematical operations:
Show code cell source
dot = sp.dotprint(BW_expr.doit(), bgcolor=None, styles=dot_style)
graphviz.Source(dot)
Large expressions#
Here, we import the large symbolic intensity expression that was used for 228.svg){kind=icon}
and see how well SymPy serialization performs on a much more complicated model.
The model contains 43,198 mathematical operations. See ComPWA/polarimetry#319 for the origin of this investigation.
Serialization with srepr#
SymPy expressions can directly be serialized to Python code as well, with the function srepr(). For the full intensity expression, we can do so with:
%%time
eval_str = sp.srepr(unfolded_intensity_expr)
CPU times: user 966 ms, sys: 0 ns, total: 966 ms
Wall time: 966 ms
Show code cell source
n_nodes = sp.count_ops(unfolded_intensity_expr)
byt = len(eval_str.encode("utf-8"))
mb = f"{1e-6 * byt:.2f}"
rendering = shorten(eval_str, placeholder=" ...", width=85)
src = f"""
This serializes the intensity expression of {n_nodes:,d} nodes
to a string of **{mb} MB**.
```python
{rendering} {")" * (rendering.count("(") - rendering.count(")"))}
```
"""
Markdown(src)
This serializes the intensity expression of 43,198 nodes to a string of 1.04 MB.
Add(Pow(Abs(Add(Mul(Add(Mul(Integer(-1), Pow(Add(Mul(Integer(-1), I, ... ))))))))))
It is up to the user, however, to import the classes of each exported node before the string can be unparsed with eval() (see this comment).
imported_intensity_expr = eval(eval_str)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[21], line 1
----> 1 imported_intensity_expr = eval(eval_str)
File <string>:1
NameError: name 'Add' is not defined
In the case of this intensity expression, it is sufficient to import all definition from the main sympy module and the Str class. Optionally, the required import statements can be embedded into the string:
exec_str = f"""\
from sympy import *
from sympy.core.symbol import Str
def get_intensity_function() -> Expr:
return {eval_str}
"""
exec_filename = Path("../_static/exported_intensity_model.py")
with open(exec_filename, "w") as f:
f.write(exec_str)
See exported_intensity_model.py for the exported model.
The parsing is then done with exec() instead of the eval() function:
%%time
exec(exec_str)
imported_intensity_expr = get_intensity_function()
CPU times: user 411 ms, sys: 88.1 ms, total: 499 ms
Wall time: 498 ms
Notice how the imported expression is exactly the same as the serialized one, including assumptions:
Common sub-expressions#
A problem is that the expression exported generated with srepr() is not human-readable in practice for large expressions. One way out may be to extract common components of the main expression with Foldable expressions. Another may be to use SymPy to detect and collect common sub-expressions.
sub_exprs, common_expr = sp.cse(unfolded_intensity_expr, order="none")
Show code cell source
Math(sp.multiline_latex(sp.Symbol("I"), common_expr[0], environment="eqnarray"))
This already works quite well with sp.lambdify (without cse=True, this would takes minutes):
%%time
args = sorted(unfolded_intensity_expr.free_symbols, key=str)
_ = sp.lambdify(args, unfolded_intensity_expr, cse=True, dummify=True)
CPU times: user 1.5 s, sys: 3.97 ms, total: 1.5 s
Wall time: 1.5 s
Still, as can be seen above, there are many sub-expressions that have exactly the same form. It would be better to find those expressions that have a similar structure, so that we can serialize them to functions or custom sub-definitions.
In SymPy, the equivalence between the expressions can be determined by the match() method using Wild symbols. We therefore first have to make all symbols in the common sub-expressions ‘wild’. In addition, in the case of this intensity expression, some of symbols are indexed and need to be replaced first.
pure_symbol_expr = unfolded_intensity_expr.replace(
query=lambda z: isinstance(z, sp.Indexed),
value=lambda z: sp.Symbol(sp.latex(z), **z.assumptions0),
)
sub_exprs, common_expr = sp.cse(pure_symbol_expr, order="none")
Note that for example the following two common sub-expressions are equivalent:
Wild symbols now allow us to find how these expressions relate to each other.
is_symbol = lambda z: isinstance(z, sp.Symbol)
make_wild = lambda z: sp.Wild(z.name)
X = [x.replace(is_symbol, make_wild) for _, x in sub_exprs]
Math(aslatex(X[5].match(X[8])))
Hint
This can be used to define functions for larger, common expression blocks.