Jit compile `to_awkward` with numba if possible

[ikrommyd]

It would be really nice if we can @numba.njit the to_awkward function. Awkward's array builder is numba-jittable but the input iterable of LHEEvents isn't. We would need to properly type it and tell numba how to compile it I think.

[eduardo-rodrigues]

That's a good idea. It would be helpful for anyone going to do this to have a pointer to similar examples of code, for help/inspiration ... Let me also tag @Saransh-cpp as I believe he has done similar work in the past.

[Saransh-cpp]

Thanks for the tag! Unfortunately, I have no experience working with LHE files and I won't have any time in the near future to look at this :(

[eduardo-rodrigues]

Hi @Saransh-cpp, I tagged you just in case you have experience with making code numba friendly and could point to useful examples. No need to have experience with LHE files: I was not expecting that :). Thanks anyway, was worth the ping.

[APN-Pucky]

I can take a look at optimizing this, however my perception of this package was never that it aims at being optimized since it is python after all. Given that since version 1.0.0 dataclasses are typed that makes it probably easy to @numba.njit the LHEEvent

[ikrommyd]

I can take a look at optimizing this, however my perception of this package was never that it aims at being optimized since it is python after all. Given that since version 1.0.0 dataclasses are typed that makes it probably easy to @numba.njit the LHEEvent

It's not critical or anything, I just happen to load a 100K event lhe file into awkward pretty often and I thought "hmm, it would be nice if that was faster" and awkward's array builder is jit-compilable. I think the only thing that numba doesn't know how to compile in the to_awkward function is the input.

[APN-Pucky]

Reading about numba's njit. I am not sure if numba.njit would really help if the bottle neck is the reading/parsing from the pylhe Generator (Iterable[LHEEvents] input). That opens the question of performance. I'll try to make a flamegraph of where the code spends its time, before looking into adding njit:

[ikrommyd]

Reading about numba's njit. I am not sure if numba.njit would really help if the bottle neck is the reading/parsing from the pylhe Generator (Iterable[LHEEvents] input). That opens the question of performance. I'll try to make a flamegraph of where the code spends its time, before looking into adding njit:

Oh yeah, I was assuming that the python for loop is what's slow. It's true that if it's lhe parsing, it doesn't help.

[APN-Pucky]

So the awkward builder

pylhe/src/pylhe/awkward.py

Line 45 in 36c3cce

builder.field(label).real(w)

is 21% while the IO/_generator is 23%

[APN-Pucky]

I'll give it a shot, however we also need to use Numba's List somewhere in between probably then since the events are not of fixed particle number (nor does numba support dicts).

EDIT: Or by ChatGPT prompt

spec = [
    ('coords', float64[:]),  # 1D array of floats
]

@jitclass(spec)
class PointArray:
    def __init__(self, coords):
        self.coords = coords

    def magnitude(self):
        return np.sqrt(np.sum(self.coords**2))

but that would again not be a generator then, but a list of fixed size.

[APN-Pucky]

I think I can speed up the to_awkward code just by massaging it a bit. Currently, I get

$ hyperfine "python3 debug/my_script.py"  -w 1
Benchmark 1: python3 debug/my_script.py
  Time (mean ± σ):      6.573 s ±  0.135 s    [User: 6.311 s, System: 0.255 s]
  Range (min … max):    6.452 s …  6.909 s    10 runs

The new one is faster

$ hyperfine "python3 debug/my_script.py"  -w 1
Benchmark 1: python3 debug/my_script.py
  Time (mean ± σ):      5.182 s ±  0.086 s    [User: 4.966 s, System: 0.211 s]
  Range (min … max):    5.067 s …  5.352 s    10 runs

and has correct integers instead of reals for some parameters. >21% faster (including import overhead) by the power of -funroll-loops:

    builder = ak.ArrayBuilder()
    for event in event_iterable:
        with builder.record(name="Event"):
            builder.field("eventinfo")
            with builder.record(name="EventInfo"):
                builder.field("nparticles").integer(event.eventinfo.nparticles)
                builder.field("pid").integer(event.eventinfo.pid)
                builder.field("weight").real(event.eventinfo.weight)
                builder.field("scale").real(event.eventinfo.scale)
                builder.field("aqed").real(event.eventinfo.aqed)
                builder.field("aqcd").real(event.eventinfo.aqcd)
            if event.weights:
                builder.field("weights")
                with builder.record(name="Weights"):
                    for label, w in event.weights.items():
                        builder.field(label).real(w)
            builder.field("particles")
            with builder.list():
                for particle in event.particles:
                    with builder.record(name="Particle"):
                        builder.field("vector")
                        with builder.record(name="Momentum4D"):
                            builder.field("px").real(particle.px)
                            builder.field("py").real(particle.py)
                            builder.field("pz").real(particle.pz)
                            builder.field("e").real(particle.e)
                        builder.field("id").integer(particle.id)
                        builder.field("status").integer(particle.status)
                        builder.field("mother1").integer(particle.mother1)
                        builder.field("mother2").integer(particle.mother2)
                        builder.field("color1").integer(particle.color1)
                        builder.field("color2").integer(particle.color2)
                        builder.field("m").real(particle.m)
                        builder.field("lifetime").real(particle.lifetime)
                        builder.field("spin").real(particle.spin)
    return builder.snapshot()  # awkward array

[APN-Pucky]

Instead of using a ArrayBuilder we discussed during the hands on workshop that one can use ak.to_buffers ak.from_buffers which is more low-level, but maybe faster, since it avoids much CPP-Python interoperation.