zfit intensity pdf

[jonas-eschle]

Hi all,

after our last meeting, I started playing around with a proposed example.

Things are mostly quite straightforward. There is a little grinding with the concept of parameters: as far as I understand, in tensorvawes, there is no distinction between a parameter and the name of the parametrization? I guess it comes from the fact that these always coincide? For example a gaussian may has a mean that is defined by a parameter. But that parameter can be anything, it can also be shared in other distributions for example. Or how would a sum of two Gaussians be handled with a common mean?

Anyway, here is the snippet that is needed to make something into zfit (this is just a draft, no general converter, but it all seems pretty straightforward)

So while the whole point of zfit is to provide a stable API and base components so that it is easy to build stuff with it, nothing is in stone (yet) which is the main reason I would like to know better where things fail and how to improve the API or workflow (in case that is needed) to have it in "stone".

We can also make the PDF more general, there are normalization hooks, there is support for importance sampling by providing a sampler (such as phasespace).

The full example can be found here as a gist.

import tensorflow.experimental.numpy as znp
import zfit  # also at top to suppress tf warnings
from zfit import z

zfit.run.set_graph_mode(False)
zfit.run.set_autograd_mode(False)


class TensorWavePDF(zfit.pdf.BasePDF):
    def __init__(self, intensity, norm, obs, params=None, name="tensorwaves"):
        """tensorwave intensity normalized over the *norm* dataset."""
        super().__init__(obs, params, name)
        self.intensity = intensity
        norm = {ob: znp.array(ar) for ob, ar in zip(self.obs, z.unstack_x(norm))}
        self.norm_sample = norm

    def _pdf(self, x, norm_range):  # we can also use better mechanics, where it automatically normalizes or not
        # this here is rather to take full control, it is always possible

        # updating the parameters of the model. This seems not very TF compatible?
        self.intensity.update_parameters({p.name: float(p) for p in self.params.values()})

        # converting the data to a dict for tensorwaves
        data = {ob: znp.array(ar) for ob, ar in zip(self.obs, z.unstack_x(x))}
        unnormalized_pdf = self.intensity(data)
        if norm_range is False:  # I mean this is not really needed,
            # but useful for e.g. sampling with `pdf(..., norm_range=False)
            return unnormalized_pdf
        else:
            return unnormalized_pdf / znp.mean(self.intensity(self.norm_sample))


# this conversion is actually not needed for the rest, we could just convert the ones in initial_parameters
params = [zfit.param.convert_to_parameter(val, name, prefer_constant=False)
          for name, val in sympy_model.parameters.items()]

# this is the "ugliest" part: zfit defines observables with limits (usually) as they often come in handy
obs = [zfit.Space(ob, limits=(np.min(data_frame[ob]) - 1, np.max(data_frame[ob]) + 1))
       for ob in phsp_set.to_pandas()]
obsall = zfit.dimension.combine_spaces(*obs)

# data conversion
phsp_zfit = zfit.Data.from_pandas(pd.DataFrame(phsp_set.to_pandas()), obs=obsall)
data_zfit = zfit.Data.from_pandas(pd.DataFrame(data_set.to_pandas()), obs=obsall)

# filter the parameters that we want to use in the model. We can also use more and only minimize
# with respect to a subset
params_fit = [p for p in params if p.name in initial_parameters if p.independent]  # remove the complex param

# set the initial values (can be done differently). A tricky thing here are complex parameters:
# in zfit they are composed of two free parameters. A complex parameter cannot be set directly
# in the minimizer, it is (AFAIK) always decomposed.
[p.set_value(initial_parameters[p.name]) for p in params_fit if p.name in initial_parameters]

pdf = TensorWavePDF(obs=obsall,
                    intensity=intensity,
                    norm=phsp_zfit,
                    params={f'param{i}': p for i, p in enumerate(params_fit)})

loss = zfit.loss.UnbinnedNLL(pdf, data_zfit)

# ok, here I was caught up playing around :) Minuit seems to perform clearly the best though
minimizer = zfit.minimize.Minuit(verbosity=8, gradient=True)
# minimizer = zfit.minimize.ScipyLBFGSBV1(verbosity=8)
# minimizer = zfit.minimize.ScipyTrustKrylovV1(verbosity=8)
# minimizer = zfit.minimize.NLoptMMAV1(verbosity=8)
# minimizer = zfit.minimize.IpyoptV1(verbosity=8)
result = minimizer.minimize(loss)
print(result)
result.hesse()
print(result)

[redeboer]

Hey @mayou36, thanks a lot for providing this example and for your thoughts! It's nice to see that it's indeed straightforward to implement your own PDF :)

Once this example works well, I think it would be nice to host it in notebook form on the usage pages (where it CI runs it and it is easy to add references to zfits API). That would be a fast and concrete next step. Ideally, such a notebook would be placed under the expertsystem docs ('here is how to fit the produced amplitude model to some data sample'). Some data sample needs to be generated though, so for now it would have to be put under the tensorwaves docs. Of course, such a notebook doesn't really solve any design issues, but it's at least a fast way of illustrating how to do things with what's available.

A few questions/remarks (no particular order, just while reading):

1. What do you mean by name of the parametrisation?
2. The sum of two gaussians with a common mean would indeed be something like:

   import sympy as sp
   def gaussian(x, mu, sigma):
       return sp.exp(-(((x - mu) / sigma) ** 2) / 2)
   a, b, x, mu, sigma1, sigma2 = sp.symbols("x, mu, sigma_(:2)")
   expr = a * gaussian(x, mu, sigma1) + b * gaussian(x, mu, sigma2)

   (for the rest see here). In TensorWaves, these sympy.Symbols then become function arguments and are as such only identifiable by their 'parameter' name.
3. update_parameters: With TF compatible, do you mean that the interface is different or that this may cause performance issues (casts and so)?
4. Yeah the complex parameters are indeed a bit bothersome :S They are there mainly because they are convenient in a helicity amplitude model. So within tensorwaves, they are again 'flattened' into a real and imaginary part. See also Use magnitude and phase in the user interface of a model instead of real and imaginary part #161.
5. I converted your example to a notebook on Google Colab here, to get an idea of performance. I was a bit disappointed by the performance there (several minutes for zfit with minuit, around 10 seconds for tensorwaves with jax; worse for CPU when I ran it on my own machine). Perhaps something needs to be done with these set_graph_mode and set_autograd_mode statements at the beginning? Or maybe the normalization is the problem?

Side note about the notebook idea: ideally, I want to move into this direction ComPWA/expertsystem#553. So, some quantum number conservation rule solver package, then a package that builds amplitude models (automated PWA formalisms), and finally, data + fitting. TensorWaves is then some steering package for PWA physicists that brings these packages together, while the code that tensorwaves is now, become some general optimizer package. That said, what the data and optimizer packages are, doesn't really matter, as long as they perform well and are sustainable.

[jonas-eschle]

Once this example works well, I think it would be nice to host it in notebook form on the usage pages (where it CI runs it and it is easy to add references to zfits API). That would be a fast and concrete next step. Ideally, such a notebook would be placed under the expertsystem docs ('here is how to fit the produced amplitude model to some data sample'). Some data sample needs to be generated though, so for now it would have to be put under the tensorwaves docs. Of course, such a notebook doesn't really solve any design issues, but it's at least a fast way of illustrating how to do things with what's available.

Yes, I agree so far, on the sampling, that can also be done in zfit in principle. But there seems to me also nothing wrong in using tensorwaves sampling, it's just to to deduplicate what we can easily.

1. What do you mean by name of the parametrisation?

So for examle in 2., sigma1 is the parameter while sigma is the name of the parametrization. What is now unclear to me: how does one update the parameters now? Here it seems functional, but the intensity has a method for that and does not take the parameters as arguments, right? So how would this now work with an intensity? Like making the example an intensity, which are the params I update, the sigma1 and sigma2 and mu? So an intensity is a "parameter holder of the inner functions"?

2. The sum of two gaussians with a common mean would indeed be something like:

   import sympy as sp
   def gaussian(x, mu, sigma):
       return sp.exp(-(((x - mu) / sigma) ** 2) / 2)
   a, b, x, mu, sigma1, sigma2 = sp.symbols("x, mu, sigma_(:2)")
   expr = a * gaussian(x, mu, sigma1) + b * gaussian(x, mu, sigma2)

I see, makes sense, so all nice in Sympy. While this makes cross working impossible (any PDF has to be written now in sp as well), I am a huge fan of the idea to do everything in Sympy, we played quite around with it as well in the beginning. Show stoppers were (back then) uncommon things such as conditionals, say just a simple where, or any kind of control flow (numerical integration for example), things that e.g. JAX or TF offer but sympy does not. Or what is your experience with this things?

In TensorWaves, these sympy.Symbols then become function arguments and are as such only identifiable by their 'parameter' name.

So basically, the symbols are then the parameters with a unique name? makes sense.

3. update_parameters: With TF compatible, do you mean that the interface is different or that this may cause performance issues (casts and so)?

I mean with the update of the variables:

• either the minimizer needs to be aware of this mechanism and knows how to use it, which zfit minimizers are not (and others neither). This is oc fine to have an API for updating in principle.
• or the PDF should be able to set it to synchronize it's parameters with the wrapped intensity. This is what the Wrapper currently does. However, since it is not possible with the current interface to update the parameters inside a jitted function (correct?), this possibility does not work either. Respectively it is the reason I've switched off the graph mode.

In short: The limitation is that the neither the minimizer in eager (since not aware of the method) nor the PDF inside the TF function can update the parameters.

What I've done now was to refactor the PDF to jit just the function and then have that run. This works as long as it's only the PDF but will fail if e.g. the PDF is composed and the sum is jitted. But it will work quite far of course, and my worries are already in a scope beyond what tensorwaves does now I think.

4. Yeah the complex parameters are indeed a bit bothersome :S They are there mainly because they are convenient in a helicity amplitude model. So within tensorwaves, they are again 'flattened' into a real and imaginary part. See also Use magnitude and phase in the user interface of a model instead of real and imaginary part #161.

That is nice, but it seems to hardcode this way then. What we have in zfit is two constructors, one for cartesian, one for polar, so it is up to the user. Also, a parameter can in general be composed of anything and all the deps are retrieved with get_params, so no matter how a complex parameter is created or what it depends on (shift or whatever), we always get the params that the minimizer has to change. It's oc a way more general way of handling this but the complex parameter falls out nicely as a simple special case.

5. I converted your example to a notebook on Google Colab here, to get an idea of performance. I was a bit disappointed by the performance there (several minutes for zfit with minuit, around 10 seconds for tensorwaves with jax; worse for CPU when I ran it on my own machine). Perhaps something needs to be done with these set_graph_mode and set_autograd_mode statements at the beginning? Or maybe the normalization is the problem?

I've checked it: This is due to the numpy converted intensity, which will make a copy from the tf tensor every time I assume.

I've adjusted the notebook so that the backend can be chosen freely in the beginning and will be used for tensorwaves and zfit to have a good comparison.
The results are:
Tensorflow (without autograd for zfit):
zfit eager: 45 secs
zfit graph: 18 secs
tensorwaves: 100 sec (not graph I assume?)

JAX:
zfit 5 secs
tensorwaves 7-8 secs

The example can be found here

(I've removed the complex parameter for the moment being from the initial ones, just to simplify)

Side note about the notebook idea: ideally, I want to move into this direction ComPWA/expertsystem#553. So, some quantum number conservation rule solver package, then a package that builds amplitude models (automated PWA formalisms), and finally, data + fitting. TensorWaves is then some steering package for PWA physicists that brings these packages together, while the code that tensorwaves is now, become some general optimizer package. That said, what the data and optimizer packages are, doesn't really matter, as long as they perform well and are sustainable.

That sounds very good, separation of things is great! I am very eager to test the limits of zfit and see where it fails with tensorwaves. Anyhow, this should be followed by more discussions of course

[redeboer]

on the sampling, that can also be done in zfit in principle

Ah yes, good point, I saw some notes on toy studies on the zfit docs. So in the notebook, would it be easy to create some intensity sample based on a custom PDF over a phasespace as domain? If so, it could indeed be hosted standalone on (what is now) the expertsystem pages (as a follow-up to generating the amplitude model).

So for examle in 2., sigma1 is the parameter while sigma is the name of the parametrization. What is now unclear to me: how does one update the parameters now? Here it seems functional, but the intensity has a method for that and does not take the parameters as arguments, right? So how would this now work with an intensity? Like making the example an intensity, which are the params I update, the sigma1 and sigma2 and mu? So an intensity is a "parameter holder of the inner functions"?

Yes you're right, a Function carries a mapping of parameter names and its call takes a dict for a data sample, then it plugs in the values they point to as arguments for the function it carries around. Parameter values are updated with update_parameters. But indeed, better may be to make this purely functional, that its call method is supposed to get a dict with both parameter values and data. Then the distinction between parameter and data input comes only at the optimizer, and it might solve the problems you rightly point out with regard to the minimizing and eager mode.

I see, makes sense, so all nice in Sympy. While this makes cross working impossible (any PDF has to be written now in sp as well), I am a huge fan of the idea to do everything in Sympy, we played quite around with it as well in the beginning.

What @spflueger and I were looking for, was some sort of 'math language' (see expertsystem ADR-001), because the starting point for us is building arbitrary amplitude models, that can be easily tweaked and adjusted, or extended with other formalisms.

Sympy offered that (with CAS, latex representation, lambdifying, etc.), but it does have its flaws -- it seems to have become an overly large framework with a massive inheritance structure and it has thousands of open issues, which make me a bit worried about its maintainability. At core though, all we need, is a way to express mathematics in some expression tree with symbols as placeholders for variables, that can then be converted to some computational backend for fitting (where you have to start distinguishing between variables for input data and/or parameters that are to be optimized).

So this is the background for (1) using Sympy (while there is no better alternative for now) and (2) trying to follow a functional approach with function arguments only that become fit parameters only once we get to the minimizing. (The model that the expertsystem produces has 'parameters' but those are rather suggested values, like expected pole position.)

Show stoppers were (back then) uncommon things such as conditionals, say just a simple where, or any kind of control flow (numerical integration for example), things that e.g. JAX or TF offer but sympy does not. Or what is your experience with this things?

As for control flow in Sympy, we haven't had a case for that yet on the optimizer side (apart from BlattWeiskopf, but there L is fixed).

[redeboer]

@jonas-eschle, this may be a first step: ComPWA/compwa.github.io#150. I convert that amplitude analysis notebook to a "Technical Report" and strip the tensorwaves part. The resulting notebook can be found at TR-020: Amplitude analysis with zfit.