Significant doit() slowdown since AmpForm v0.15.2

[redeboer]

Since #418, calling doit() on an AmplitudeModel.full_expression is much slower:

• ampform==0.15.1: 7m04s
• ampform==0.15.3: 1m43s

Benchmark script

ampform==0.15.3  # <-- 0.15.1
ampform-dpd==0.2.0
attrs==23.2.0
cloudpickle==3.0.0
hepunits==2.3.3
iminuit==2.25.2
jax==0.4.28
jaxlib==0.4.28
jsonschema==4.22.0
jsonschema-specifications==2023.12.1
ml-dtypes==0.4.0
mpmath==1.3.0
numpy==1.26.4
opt-einsum==3.3.0
particle==0.24.0
python-constraint2==2.0.0b5
pyyaml==6.0.1
qrules==0.10.2
referencing==0.35.1
rpds-py==0.18.1
scipy==1.13.0
sympy==1.12
tensorwaves==0.4.12
tqdm==4.66.4
typing-extensions==4.11.0

import logging
import timeit

import qrules

from ampform_dpd import DalitzPlotDecompositionBuilder
from ampform_dpd.adapter.qrules import normalize_state_ids, to_three_body_decay
from ampform_dpd.dynamics.builder import formulate_breit_wigner_with_form_factor


def timeit_decorator(func):
    def wrapper(*args, **kwargs):
        start_time = timeit.default_timer()
        result = func(*args, **kwargs)
        end_time = timeit.default_timer()
        elapsed_time = end_time - start_time
        minutes = int(elapsed_time // 60)
        seconds = int(elapsed_time % 60)
        print(f"Took {minutes}:{seconds:02d} minutes to execute.")  # noqa: T201
        return result

    return wrapper


@timeit_decorator
def main() -> None:
    logging.getLogger("ampform_dpd.io").setLevel(logging.ERROR)
    logging.getLogger("qrules.transition").setLevel(logging.ERROR)

    reaction = qrules.generate_transitions(
        initial_state="J/psi(1S)",
        final_state=["K0", "Sigma+", "p~"],
        allowed_interaction_types="strong",
        formalism="canonical-helicity",
        mass_conservation_factor=0.05,
    )
    reaction = normalize_state_ids(reaction)

    decay = to_three_body_decay(reaction.transitions, min_ls=True)

    model_builder = DalitzPlotDecompositionBuilder(decay, min_ls=False)
    for chain in model_builder.decay.chains:
        model_builder.dynamics_choices.register_builder(
            chain, formulate_breit_wigner_with_form_factor
        )
    model = model_builder.formulate(reference_subsystem=2)

    unfolded_intensity_expr = model.intensity.doit()
    unfolded_intensity_expr.xreplace(model.amplitudes).doit()


if __name__ == "__main__":
    main()

[redeboer]

Better benchmark, without ampform-dpd:

• ampform==0.15.1: 0.40s
• ampform==0.15.3: 2.44s

Requirements file

ampform==0.15.3  # <-- 0.15.1
attrs==23.2.0
hepunits==2.3.3
jsonschema==4.22.0
jsonschema-specifications==2023.12.1
mpmath==1.3.0
particle==0.24.0
python-constraint2==2.0.0b5
pyyaml==6.0.1
qrules==0.10.2
referencing==0.35.1
rpds-py==0.18.1
sympy==1.12
tqdm==4.66.4
typing-extensions==4.11.0

Script

import ampform
import timeit
import qrules
from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff

reaction = qrules.generate_transitions(
    initial_state=("psi(4160)", [-1, +1]),
    final_state=["D-", "D0", "pi+"],
    allowed_intermediate_particles=["D*(2007)0"],
    formalism="canonical-helicity",
)
builder = ampform.get_builder(reaction)
for particle in reaction.get_intermediate_particles():
    builder.dynamics.assign(particle.name, create_relativistic_breit_wigner_with_ff)
model = builder.formulate()

start_time = timeit.default_timer()
model.expression.doit()
end_time = timeit.default_timer()
elapsed_time = end_time - start_time
print(f"Took {elapsed_time} seconds to execute.")