Speed up lambdifying¶
import logging
import ampform
import graphviz
import qrules
import sympy as sp
from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff
from IPython.display import HTML, SVG
from tensorwaves.model import LambdifiedFunction, SympyModel
from tensorwaves.model.sympy import optimized_lambdify, split_expression
logger = logging.getLogger()
logger.setLevel(logging.ERROR)
Split expression¶
Lambdifying a SymPy expression can take rather long when an expression is complicated. Fortunately, TensorWaves offers a way to speed up the lambdify process. The idea is to split up an an expression into sub-expressions, separate those separately, and then recombining them. Let’s illustrate that idea with the following simplified example:
x, y, z = sp.symbols("x:z")
expr = x ** z + 2 * y
expr
This expression can be represented in a tree of mathematical operations.
dot = sp.dotprint(expr)
graphviz.Source(dot)
The function split_expression()
can now be used to split up this expression tree into a ‘top expression’ plus definitions for each of the sub-expressions into which it was split:
top_expr, sub_expressions = split_expression(expr, max_complexity=3)
top_expr
sub_expressions
{f0: x**z, f1: 2*y}
The original expression can easily be reconstructed with subs()
or xreplace()
:
top_expr.xreplace(sub_expressions)
Each of the expression trees are now smaller than the original:
dot = sp.dotprint(top_expr)
graphviz.Source(dot)
for symbol, definition in sub_expressions.items():
dot = sp.dotprint(definition)
graph = graphviz.Source(dot)
graph.render(filename=f"sub_expr_{symbol.name}", format="svg")
html = "<table>\n"
html += " <tr>\n"
html += "".join(
' <th style="text-align:center;'
f' background-color:white">{symbol.name}</th>\n'
for symbol in sub_expressions
)
html += " </tr>\n"
html += " <tr>\n"
for symbol in sub_expressions:
svg = SVG(f"sub_expr_{symbol.name}.svg").data
html += f' <td style="background-color:white">{svg}</td>\n'
html += " </tr>\n"
html += "</table>"
HTML(html)
f0 | f1 |
---|---|
Optimized lambdify¶
Generally, the lambdify time scales exponentially with the size of an expression tree. With larger expression trees, it’s therefore much faster to lambdify these sub-expressions separately and to recombine them. TensorWaves offers a function that does this for you: optimized_lambdify()
. We’ll use an HelicityModel
to illustrate this:
reaction = qrules.generate_transitions(
initial_state=("J/psi(1S)", [+1]),
final_state=["gamma", "pi0", "pi0"],
allowed_intermediate_particles=["f(0)"],
)
model_builder = ampform.get_builder(reaction)
for name in reaction.get_intermediate_particles().names:
model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)
model = model_builder.formulate()
expression = model.expression.doit()
sorted_symbols = sorted(expression.free_symbols, key=lambda s: s.name)
%%time
lambdified_optimized = optimized_lambdify(
expression,
sorted_symbols,
max_complexity=100,
backend="numpy",
)
CPU times: user 601 ms, sys: 4.06 ms, total: 605 ms
Wall time: 605 ms
%%time
sp.lambdify(sorted_symbols, expression)
CPU times: user 10.5 s, sys: 64.1 ms, total: 10.6 s
Wall time: 10.4 s
<function _lambdifygenerated(Dummy_164, Dummy_163, Dummy_162, Dummy_161, Dummy_160, Dummy_159, Dummy_158, Dummy_157, Dummy_156, Dummy_155, Dummy_154, Dummy_153, Dummy_152, Dummy_151, Dummy_150, m_1, m_12, m_2, Dummy_149, Dummy_148, Dummy_147, Dummy_146, Dummy_145, Dummy_144, Dummy_143)>
Specifying complexity¶
In the usually workflow (see Usage), TensorWaves uses SymPy’s own lambdify()
by default. You can change this behavior with the max_complexity
argument of SympyModel
:
sympy_model = SympyModel(
expression=model.expression.doit(),
parameters=model.parameter_defaults,
max_complexity=100,
)
If max_complexity
is specified (i.e., is not None
), LambdifiedFunction
uses TensorWaves’s optimized_lambdify()
.
%%time
intensity = LambdifiedFunction(sympy_model, backend="jax")
CPU times: user 945 ms, sys: 23.9 ms, total: 969 ms
Wall time: 968 ms