Extend PyPI package code example

[redeboer]

It would be nice to extend the basic usage example of the polarimetry-lc2pkpi package so that it becomes clearer to the user how to:

• Get the intensity expression and (optionally) replace its symbols by parameter values
• Convert the expression to a numerical function
• Evaluate it over a data sample (i.e., what is the expected input and how to compute the kinematic variables from input Mandelstam variables)

@mmikhasenko any other requirements?

[mmikhasenko]

Thanks @redeboer. Here is what I'm thinking of

Let's say we have a tree with two branches, msq_piK and msq_Kp. The example should evaluate the unpolarized intensity on this kinematic point.

The test file test_tree.root produced as follows:

TTree *t = new TTree("Lc2ppiK", "Lc2ppiK");

double msq_piK = 0.64;
double msq_Kp = 4.03;

t->Branch("msq_piK", &msq_piK);
t->Branch("msq_Kp", &msq_Kp);
t->Fill();

msq_piK = 0.68;
msq_Kp = 4.01;
t->Fill();

t->SaveAs("test_tree.root");

[mmikhasenko]

test_tree.zip

[mmikhasenko]

That is how to do it in julia

using UnROOT
using Lc2ppiKSemileptonicModelLHCb
using Lc2ppiKSemileptonicModelLHCb.ThreeBodyDecays

# get model
const model = published_model("Default amplitude model")
const ms = first(model.chains).tbs.ms

# get data
t = LazyTree("test_tree.root", "Lc2ppiK")

# event loop
weights = map(zip(t.msq_Kp, t.msq_piK)) do (msq_Kp, msq_piK)
    σs = Invariants(ms; σ1=msq_piK, σ2=msq_Kp)
    unpolarizedintensity(model, σs)
end 

[mmikhasenko]

Checked again that Julia code works with Lc2ppiKSemileptonicModelLHCb with v0.4.0

[redeboer]

That is how to do it in julia

using UnROOT
using Lc2ppiKSemileptonicModelLHCb
using Lc2ppiKSemileptonicModelLHCb.ThreeBodyDecays

# get model
const model = published_model("Default amplitude model")
const ms = first(model.chains).tbs.ms

# get data
t = LazyTree("test_tree.root", "Lc2ppiK")

# event loop
weights = map(zip(t.msq_Kp, t.msq_piK)) do (msq_Kp, msq_piK)
    σs = Invariants(ms; σ1=msq_piK, σ2=msq_Kp)
    unpolarizedintensity(model, σs)
end 

What is missing on the Python side is a convenient function that gives you the third Mandelstam variable in terms of the other two, i.e. compute_third_mandelstam(). The user can be told how to do this specifically, but there is room for error there, because the substitution to one expression with only two (Mandelstam) variables requires some understanding about symbolics.

import polarimetry
from polarimetry.io import perform_cached_doit

model = polarimetry.published_model("Default amplitude model")
intensity_expr = perform_cached_doit(model.full_expression)
fixed_intensity_expr = intensity_expr.xreplace(model.parameter_defaults)

Now replace the angle definitions algebraically:

fixed_intensity_expr = fixed_intensity_expr.xreplace(model.variables)
free_symbols = sorted(fixed_intensity_expr.free_symbols, key=str)
free_symbols
# [m0, m1, m2, m3, sigma1, sigma2, sigma3]

...express the third Mandelstam variable in terms of the others and substitute...

from ampform.kinematics.phasespace import compute_third_mandelstam

m0, m1, m2, m3, σ1, σ2, σ3 = free_symbols
σ3_expr = compute_third_mandelstam(σ1, σ2, m0, m1, m2, m3)
σ3_expr
# m0**2 + m1**2 + m2**2 + m3**2 - sigma1 - sigma2

...substitute the $\sigma_3$ expression and the masses using the model parameters:

fixed_intensity_expr = fixed_intensity_expr.xreplace({σ3: σ3_expr})
fixed_intensity_expr = fixed_intensity_expr.xreplace(model.parameter_defaults)
fixed_intensity_expr.free_symbols
# {sigma1, sigma2}

Finally you have a numeric function that can evaluate over two Mandelstam variables:

from polarimetry.io import perform_cached_lambdify

intensity_func = perform_cached_lambdify(fixed_intensity_expr.doit(), backend="jax")

import numpy as np  # jax.numpy for large data

data = {  # this can be imported with uproot
    "sigma1": np.array([0.64, 0.68]),  # piK
    "sigma2": np.array([4.03, 4.01]),  # Kp
}
intensity_func(data)
# Array([4507.65484864, 4992.72822529], dtype=float64)

[!NOTE] Algebraic substitutions are just one way. You could also create numeric data transformers (SympyDataTransformer) that numerically transforms your sigma1, sigma2 to the kinematic variables first.

[mmikhasenko]

I was thinking of the numeric transformation, would not it faster, at least for the one-time evaluation?

[redeboer]

With JAX probably comparable, but indeed the algebraic substitutions take a longer time. The benefit of direct lambdification would be that one can later tweak the parameters.

[mmikhasenko]

Let me try to put things in one block, I would leave sigma3 in, in the spirit of Julia implementation.

import polarimetry
from polarimetry.io import perform_cached_doit
import numpy as np  # jax.numpy for large data
from ampform.kinematics.phasespace import compute_third_mandelstam

model = polarimetry.published_model("Default amplitude model")
intensity_expr = perform_cached_doit(model.full_expression)
from polarimetry.io import perform_cached_lambdify

fixed_intensity_expr = intensity_expr.xreplace(model.parameter_defaults)
fixed_intensity_expr = fixed_intensity_expr.xreplace(model.variables)
fixed_intensity_expr = fixed_intensity_expr.xreplace(model.parameter_defaults)

intensity_func = perform_cached_lambdify(fixed_intensity_expr.doit(), backend="jax")

Now, in order to handle the data,

data = {  # this can be imported with uproot
    "sigma1": np.array([0.64, 0.68]),  # piK
    "sigma2": np.array([4.03, 4.01]),  # Kp
}

we need the completion function, that would magically do the job

completed_data = mandelstam_variables(masses, data)
intensity_func(completed_data)

how mandelstam_variables can look like? - figuring out what variable is missing, sigma1, sigma2, or sigma3

[redeboer]

how mandelstam_variables can look like? - figuring out what variable is missing, sigma1, sigma2, or sigma3

mandelstam_variables would be a SympyDataTransformer. Same as in https://github.com/ComPWA/polarimetry/issues/336#issuecomment-1905795986, the user needs to find an expression for sigma3, using the same symbols as in the model.

import sympy as sp
from ampform.kinematics.phasespace import compute_third_mandelstam
from tensorwaves.data.transform import SympyDataTransformer

m0, m1, m2, m3, σ1, σ2, σ3 = sp.symbols("m:4 sigma1:4", nonnegative=True)
variables = model.variables
variables[σ3] = compute_third_mandelstam(σ1, σ2, m0, m1, m2, m3)
variables = {
    symbol: expr.xreplace(variables).xreplace(model.parameter_defaults)
    for symbol, expr in variables.items()
}
mandelstam_variables = SympyDataTransformer.from_sympy(variables, backend="jax")

import numpy as np

data = {  # this can be imported with uproot
    "sigma1": np.array([0.64, 0.68]),  # piK
    "sigma2": np.array([4.03, 4.01]),  # Kp
}
completed_data = mandelstam_variables(data)
completed_data.update(data)  # or use IdentityTransformer
intensity_func(completed_data)

[mmikhasenko]

Nice, I haven't appreciated the SympyDataTransformer before. It's cool stuff.

sp.symbols("m:4 sigma1:4", nonnegative=True) not ideal since reintroduces the variables, the constraint nonnegative is not evident. I think it would be better to have it as the model property

m0, m1, m2, m3 = model.masses()
σ1, σ2, σ3 = model.invariants()

[redeboer]

sp.symbols("m:4 sigma1:4", nonnegative=True) not ideal since reintroduces the variables, the constraint nonnegative is not evident. I think it would be better to have it as the model property

m0, m1, m2, m3 = model.masses()
σ1, σ2, σ3 = model.invariants()

That's a good idea. You could even make them dict, so that you immediately have access to their values and/or expressions.

[redeboer]

Just realized that there already was a function that does parts of the above workflow 🤦 https://compwa.github.io/polarimetry/api/polarimetry.data.html#polarimetry.data.create_data_transformer Only the transformation from sigma1+sigma2 to sigma3 is missing.

[mmikhasenko]

Ok, I think it remain kind-of manual, since there are three cases, sigma1+sigma2 to sigma3 sigma2+sigma3 to sigma1 sigma3+sigma1 to sigma2

[mmikhasenko]

How does the workflow look like, then?

[redeboer]

Ok, I think it remain kind-of manual, since there are three cases, sigma1+sigma2 to sigma3 sigma2+sigma3 to sigma1 sigma3+sigma1 to sigma2

Yeah indeed, that's what I like about this Invariants class in Julia https://github.com/ComPWA/polarimetry/issues/336#issuecomment-1904095067. I'm wondering whether a `DataTransformer can do something like that, that you feed it a datasample with only specific keys (here: sigma1 and sigma2) and it will then recursively perform a transformation until all variables are computed.

How does the workflow look like, then?

Gonna move this to a PR. Too much hacking in Markdown textboxes 😆

[redeboer]

Nice, I haven't appreciated the SympyDataTransformer before. It's cool stuff.

sp.symbols("m:4 sigma1:4", nonnegative=True) not ideal since reintroduces the variables, the constraint nonnegative is not evident. I think it would be better to have it as the model property

m0, m1, m2, m3 = model.masses()
σ1, σ2, σ3 = model.invariants()

339 added these new attributes. Their use is illustrated here:

https://compwa.github.io/polarimetry/appendix/model-serialization.html#python-package and will become available in polarimetry-lc2pkpi==0.0.11.