diff --git a/dispersion-integral.ipynb b/dispersion-integral.ipynb
new file mode 100644
index 00000000..548583a4
--- /dev/null
+++ b/dispersion-integral.ipynb
@@ -0,0 +1,414 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Note: you may need to restart the kernel to use updated packages.\n"
+ ]
+ }
+ ],
+ "source": [
+ "%pip install -q ampform==0.14.* quadpy tensorwaves[jax]==0.4.*"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from functools import partial\n",
+ "\n",
+ "import numpy as np\n",
+ "import qrules\n",
+ "import quadpy\n",
+ "import sympy as sp\n",
+ "\n",
+ "from ampform.dynamics import (\n",
+ " BlattWeisskopfSquared,\n",
+ " BreakupMomentumSquared,\n",
+ " PhaseSpaceFactor,\n",
+ ")\n",
+ "\n",
+ "PDG = qrules.load_pdg()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "eb2f6bbe3d8f4a03b266d655ed81b2f2",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Propagating quantum numbers: 0%| | 0/48 [00:0049\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m---> 50\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtensorflow\u001b[39;00m \u001b[39mimport\u001b[39;00m float64\n\u001b[1;32m 51\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mImportError\u001b[39;00m: \u001b[39m# pragma: no cover\u001b[39;00m\n",
+ "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'tensorflow'",
+ "\nDuring handling of the above exception, another exception occurred:\n",
+ "\u001b[0;31mImportError\u001b[0m Traceback (most recent call last)",
+ "\u001b[1;32m/mnt/c/dev/ComPWA/compwa-org/dispersion-integral.ipynb Cell 3'\u001b[0m in \u001b[0;36m| \u001b[0;34m()\u001b[0m\n\u001b[1;32m 72\u001b[0m model_energy_dependent_with_ff \u001b[39m=\u001b[39m model_builder\u001b[39m.\u001b[39mformulate()\n\u001b[1;32m 74\u001b[0m \u001b[39m# Generate phase space sample\u001b[39;00m\n\u001b[0;32m---> 75\u001b[0m rng \u001b[39m=\u001b[39m TFUniformRealNumberGenerator(seed\u001b[39m=\u001b[39;49m\u001b[39m0\u001b[39;49m)\n\u001b[1;32m 76\u001b[0m phsp_generator \u001b[39m=\u001b[39m TFPhaseSpaceGenerator(\n\u001b[1;32m 77\u001b[0m initial_state_mass\u001b[39m=\u001b[39mreaction\u001b[39m.\u001b[39minitial_state[\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m]\u001b[39m.\u001b[39mmass,\n\u001b[1;32m 78\u001b[0m final_state_masses\u001b[39m=\u001b[39m{i: p\u001b[39m.\u001b[39mmass \u001b[39mfor\u001b[39;00m i, p \u001b[39min\u001b[39;00m reaction\u001b[39m.\u001b[39mfinal_state\u001b[39m.\u001b[39mitems()},\n\u001b[1;32m 79\u001b[0m )\n\u001b[1;32m 80\u001b[0m phsp_momenta \u001b[39m=\u001b[39m phsp_generator\u001b[39m.\u001b[39mgenerate(\u001b[39m1_000_000\u001b[39m, rng)\n",
+ "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/tensorwaves/data/rng.py:52\u001b[0m, in \u001b[0;36mTFUniformRealNumberGenerator.__init__\u001b[0;34m(self, seed)\u001b[0m\n\u001b[1;32m 50\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtensorflow\u001b[39;00m \u001b[39mimport\u001b[39;00m float64\n\u001b[1;32m 51\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mImportError\u001b[39;00m: \u001b[39m# pragma: no cover\u001b[39;00m\n\u001b[0;32m---> 52\u001b[0m raise_missing_module_error(\u001b[39m\"\u001b[39;49m\u001b[39mtensorflow\u001b[39;49m\u001b[39m\"\u001b[39;49m, extras_require\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mtf\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m 54\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mseed \u001b[39m=\u001b[39m seed\n\u001b[1;32m 55\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdtype \u001b[39m=\u001b[39m float64\n",
+ "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/tensorwaves/function/_backend.py:110\u001b[0m, in \u001b[0;36mraise_missing_module_error\u001b[0;34m(module_name, extras_require)\u001b[0m\n\u001b[1;32m 105\u001b[0m \u001b[39mif\u001b[39;00m extras_require:\n\u001b[1;32m 106\u001b[0m error_message \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m (\n\u001b[1;32m 107\u001b[0m \u001b[39m\"\u001b[39m\u001b[39m Reinstall tensorwaves with:\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39m\\n\u001b[39;00m\u001b[39m\"\u001b[39m\n\u001b[1;32m 108\u001b[0m \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m pip install tensorwaves[\u001b[39m\u001b[39m{\u001b[39;00mextras_require\u001b[39m}\u001b[39;00m\u001b[39m]\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39m\"\u001b[39m\n\u001b[1;32m 109\u001b[0m )\n\u001b[0;32m--> 110\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mImportError\u001b[39;00m(error_message)\n",
+ "\u001b[0;31mImportError\u001b[0m: Module tensorflow not installed. Reinstall tensorwaves with:\n\n pip install tensorwaves[tf]\n"
+ ]
+ }
+ ],
+ "source": [
+ "# %config InlineBackend.figure_formats = ['svg']\n",
+ "import ampform\n",
+ "import attrs\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import qrules\n",
+ "from ampform.dynamics import PhaseSpaceFactorSWave\n",
+ "from ampform.dynamics.builder import RelativisticBreitWignerBuilder\n",
+ "from matplotlib import cm\n",
+ "\n",
+ "from tensorwaves.data import (\n",
+ " SympyDataTransformer,\n",
+ " TFPhaseSpaceGenerator,\n",
+ " TFUniformRealNumberGenerator,\n",
+ ")\n",
+ "from tensorwaves.function.sympy import create_parametrized_function\n",
+ "\n",
+ "# Modify particle masses\n",
+ "PDG = qrules.load_pdg()\n",
+ "eta_mass = PDG[\"eta\"].mass\n",
+ "proton_mass = PDG[\"eta\"].mass\n",
+ "resonance_mass = 0.97 * (proton_mass + eta_mass)\n",
+ "\n",
+ "particle_set = set(PDG)\n",
+ "\n",
+ "particle = PDG[\"p\"]\n",
+ "modified_particle = attrs.evolve(particle, mass=proton_mass)\n",
+ "particle_set.remove(particle)\n",
+ "particle_set.add(modified_particle)\n",
+ "\n",
+ "particle = PDG[\"N(1440)+\"]\n",
+ "modified_particle = attrs.evolve(particle, mass=resonance_mass)\n",
+ "particle_set.remove(particle)\n",
+ "particle_set.add(modified_particle)\n",
+ "\n",
+ "PDG = qrules.ParticleCollection(particle_set)\n",
+ "\n",
+ "# Create amplitude model\n",
+ "reaction = qrules.generate_transitions(\n",
+ " initial_state=(\"J/psi(1S)\", [+1]),\n",
+ " final_state=[\"p\", \"p~\", \"eta\"],\n",
+ " allowed_intermediate_particles=[\"N(1440)+\"],\n",
+ " allowed_interaction_types=[\"strong\", \"EM\"],\n",
+ " particle_db=PDG,\n",
+ " mass_conservation_factor=1,\n",
+ ")\n",
+ "assert len(reaction.get_intermediate_particles()) == 1\n",
+ "\n",
+ "model_builder = ampform.get_builder(reaction)\n",
+ "model_builder.scalar_initial_state_mass = True\n",
+ "model_builder.stable_final_state_ids = [0, 1, 2]\n",
+ "breit_wigner_builder = RelativisticBreitWignerBuilder(\n",
+ " phsp_factor=PhaseSpaceFactorSWave # <-- fixes the problem\n",
+ ")\n",
+ "for name in reaction.get_intermediate_particles().names:\n",
+ " model_builder.set_dynamics(name, breit_wigner_builder)\n",
+ "\n",
+ "breit_wigner_builder.energy_dependent_width = False\n",
+ "breit_wigner_builder.form_factor = False\n",
+ "model_standard_bw = model_builder.formulate()\n",
+ "\n",
+ "breit_wigner_builder.energy_dependent_width = False\n",
+ "breit_wigner_builder.form_factor = True\n",
+ "model_standard_bw_with_ff = model_builder.formulate()\n",
+ "\n",
+ "breit_wigner_builder.energy_dependent_width = True\n",
+ "breit_wigner_builder.form_factor = False\n",
+ "model_energy_dependent = model_builder.formulate()\n",
+ "\n",
+ "breit_wigner_builder.energy_dependent_width = True\n",
+ "breit_wigner_builder.form_factor = True\n",
+ "model_energy_dependent_with_ff = model_builder.formulate()\n",
+ "\n",
+ "# Generate phase space sample\n",
+ "rng = TFUniformRealNumberGenerator(seed=0)\n",
+ "phsp_generator = TFPhaseSpaceGenerator(\n",
+ " initial_state_mass=reaction.initial_state[-1].mass,\n",
+ " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n",
+ ")\n",
+ "phsp_momenta = phsp_generator.generate(1_000_000, rng)\n",
+ "helicity_transformer = SympyDataTransformer.from_sympy(\n",
+ " model_standard_bw.kinematic_variables, backend=\"jax\"\n",
+ ")\n",
+ "phsp = helicity_transformer(phsp_momenta)\n",
+ "phsp = {k: v.real for k, v in phsp.items()} # important for sqrt of Blatt-Weisskopf!\n",
+ "\n",
+ "# Generate intensity distributions\n",
+ "func_standard_bw = create_parametrized_function(\n",
+ " expression=model_standard_bw.expression.doit(),\n",
+ " parameters=model_standard_bw.parameter_defaults,\n",
+ " backend=\"jax\",\n",
+ ")\n",
+ "func_standard_bw_with_ff = create_parametrized_function(\n",
+ " expression=model_standard_bw_with_ff.expression.doit(),\n",
+ " parameters=model_standard_bw_with_ff.parameter_defaults,\n",
+ " backend=\"jax\",\n",
+ ")\n",
+ "func_energy_dependent = create_parametrized_function(\n",
+ " expression=model_energy_dependent.expression.doit(),\n",
+ " parameters=model_energy_dependent.parameter_defaults,\n",
+ " backend=\"jax\",\n",
+ ")\n",
+ "func_energy_dependent_with_ff = create_parametrized_function(\n",
+ " expression=model_energy_dependent_with_ff.expression.doit(),\n",
+ " parameters=model_energy_dependent_with_ff.parameter_defaults,\n",
+ " backend=\"jax\",\n",
+ ")\n",
+ "\n",
+ "# Plot 'em!\n",
+ "fig, ax = plt.subplots(figsize=(9, 4))\n",
+ "\n",
+ "hist_kwargs = dict(\n",
+ " bins=200,\n",
+ " histtype=\"step\",\n",
+ " density=True,\n",
+ ")\n",
+ "ax.hist(\n",
+ " np.array(phsp[\"m_02\"].real),\n",
+ " weights=np.array(func_standard_bw(phsp)),\n",
+ " label=\"Standard Breit-Wigner\",\n",
+ " **hist_kwargs,\n",
+ ")\n",
+ "ax.hist(\n",
+ " np.array(phsp[\"m_02\"].real),\n",
+ " weights=np.array(func_standard_bw_with_ff(phsp)),\n",
+ " label=\"Standard Breit-Wigner with form factor\",\n",
+ " linestyle=\"dotted\",\n",
+ " **hist_kwargs,\n",
+ ")\n",
+ "ax.hist(\n",
+ " np.array(phsp[\"m_02\"].real),\n",
+ " weights=np.array(func_energy_dependent(phsp)),\n",
+ " label=\"Energy-dependent Breit-Wigner\",\n",
+ " **hist_kwargs,\n",
+ ")\n",
+ "ax.hist(\n",
+ " np.array(phsp[\"m_02\"].real),\n",
+ " weights=np.array(func_energy_dependent_with_ff(phsp)),\n",
+ " label=\"Energy-dependent Breit-Wigner with form factor\",\n",
+ " linestyle=\"dotted\",\n",
+ " **hist_kwargs,\n",
+ ")\n",
+ "\n",
+ "ax.set_yscale(\"log\")\n",
+ "ax.set_xlabel(R\"$m_{p\\eta}$ [GeV]\")\n",
+ "\n",
+ "resonances = sorted(\n",
+ " reaction.get_intermediate_particles(),\n",
+ " key=lambda p: p.mass,\n",
+ ")\n",
+ "evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n",
+ "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n",
+ "for p, color in zip(resonances, colors):\n",
+ " ax.axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=color)\n",
+ "ax.legend()\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/latex": [
+ "$\\displaystyle \\frac{B_{2}^2\\left(\\frac{q^2\\left(s^{\\prime}\\right)}{q_{0}^{2}}\\right) \\rho\\left(s^{\\prime}\\right)}{\\left(s^{\\prime} - \\left(m_{1} + m_{2}\\right)^{2}\\right) \\left(- i \\epsilon - s + s^{\\prime}\\right)}$"
+ ],
+ "text/plain": [
+ "BlattWeisskopfSquared(2, BreakupMomentumSquared(s^{\\prime}, m1, m2)/q0**2)*PhaseSpaceFactor(s^{\\prime}, m1, m2)/((s^{\\prime} - (m1 + m2)**2)*(-I*epsilon - s + s^{\\prime}))"
+ ]
+ },
+ "execution_count": 6,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "def na2(s, m1, m2, L, q0):\n",
+ " q_squared = BreakupMomentumSquared(s, m1, m2)\n",
+ " return BlattWeisskopfSquared(\n",
+ " z=q_squared / (q0**2),\n",
+ " angular_momentum=L,\n",
+ " )\n",
+ "\n",
+ "\n",
+ "L = 2\n",
+ "s, m1, m2 = sp.symbols(\"s m1 m2\", real=True)\n",
+ "epsilon = sp.Symbol(\"epsilon\", positive=True)\n",
+ "s_prime = sp.Symbol(R\"s^{\\prime}\", real=True)\n",
+ "\n",
+ "q0 = sp.Symbol(\"q0\", real=True)\n",
+ "s_thr = (m1 + m2) ** 2\n",
+ "integrand = (\n",
+ " PhaseSpaceFactor(s_prime, m1, m2) * na2(s_prime, m1, m2, L, q0)\n",
+ ") / ((s_prime - s_thr) * (s_prime - s - epsilon * sp.I))\n",
+ "integrand"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "integrand_func = sp.lambdify(\n",
+ " args=(s_prime, s, epsilon, m1, m2, q0),\n",
+ " expr=integrand.doit(),\n",
+ " modules=\"numpy\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Note: you may need to restart the kernel to use updated packages.\n"
+ ]
+ }
+ ],
+ "source": [
+ "%pip install -q quadpy"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "@np.vectorize\n",
+ "def vectorized_quad(func, a, b, **func_kwargs):\n",
+ " values, _ = quadpy.quad(partial(func, **func_kwargs), a, b)\n",
+ " return values"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "s_min, s_max = -0.15, 1.4\n",
+ "s_domain = np.linspace(s_min, s_max, num=50)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "array([0.02149997+2.34976267e-06j, 0.021575 +2.39410036e-06j,\n",
+ " 0.02165146+2.44072113e-06j, 0.02172944+2.48985230e-06j,\n",
+ " 0.02180901+2.54175752e-06j, 0.02189027+2.59674521e-06j,\n",
+ " 0.02197333+2.65517959e-06j, 0.0220583 +2.71749557e-06j,\n",
+ " 0.0221453 +2.78421888e-06j, 0.02223449+2.85599443e-06j,\n",
+ " 0.02232605+2.93362692e-06j, 0.02242016+3.01814111e-06j,\n",
+ " 0.02251708+3.11087473e-06j, 0.02261708+3.21362882e-06j,\n",
+ " 0.02272052+3.32892580e-06j, 0.02282785+3.46048858e-06j,\n",
+ " 0.02293968+3.61423036e-06j, 0.02305685+3.80064418e-06j,\n",
+ " 0.0231807 +4.04232011e-06j, 0.02331387+4.41505201e-06j,\n",
+ " 0.02346414+1.06665934e-05j, 0.02362069+3.51329782e-05j,\n",
+ " 0.02377729+7.19677549e-05j, 0.02393282+1.19096483e-04j,\n",
+ " 0.02408646+1.75249336e-04j, 0.02423758+2.39431771e-04j,\n",
+ " 0.02438575+3.10799126e-04j, 0.02453062+3.88614125e-04j,\n",
+ " 0.02467198+4.72220957e-04j, 0.02480966+5.61035555e-04j,\n",
+ " 0.02494356+6.54535110e-04j, 0.02507363+7.52251722e-04j,\n",
+ " 0.02519985+8.53763339e-04j, 0.02532221+9.58691708e-04j,\n",
+ " 0.02544075+1.06669479e-03j, 0.02555551+1.17746537e-03j,\n",
+ " 0.02566653+1.29072425e-03j, 0.02577388+1.40621857e-03j,\n",
+ " 0.02587761+1.52372060e-03j, 0.02597781+1.64302166e-03j,\n",
+ " 0.02607453+1.76393362e-03j, 0.02616787+1.88628265e-03j,\n",
+ " 0.02625788+2.00991358e-03j, 0.02634466+2.13468021e-03j,\n",
+ " 0.02642827+2.26045254e-03j, 0.02650878+2.38710849e-03j,\n",
+ " 0.02658628+2.51453853e-03j, 0.02666083+2.64264027e-03j,\n",
+ " 0.0267325 +2.77131861e-03j, 0.02680137+2.90048853e-03j])"
+ ]
+ },
+ "execution_count": 11,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "m1_val = PDG[\"pi0\"].mass\n",
+ "m2_val = PDG[\"eta\"].mass\n",
+ "s_thr_val = float(s_thr.subs({m1: m1_val, m2: m2_val}))\n",
+ "integral_value = vectorized_quad(\n",
+ " integrand_func,\n",
+ " a=s_thr_val,\n",
+ " b=np.inf,\n",
+ " s=s_domain,\n",
+ " epsilon=1e-3,\n",
+ " m1=m1_val,\n",
+ " m2=m2_val,\n",
+ " q0=1.0,\n",
+ ")\n",
+ "integral_value"
+ ]
+ }
+ ],
+ "metadata": {
+ "interpreter": {
+ "hash": "e69c11b2f7fb560b2f005360aad317e409fdc986d5dd403d061abaed85739f2a"
+ },
+ "kernelspec": {
+ "display_name": "Python 3.8.13 ('compwa-org')",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.8.13"
+ },
+ "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
|