akt.py

# standard library imports
from __future__ import absolute_import, division, print_function

# standard numerical library imports
import numpy as np

# matplotlib is required for this example
import matplotlib.pyplot as plt

# matplotlib inline
plt.rcParams["figure.figsize"] = (4, 4)

#############################################################
# NOTE: install ffmpeg and point matplotlib to it
#############################################################
# on linux
# FIXME: MLB change this to your path
plt.rcParams["animation.ffmpeg_path"] = '/home/username/anaconda/envs/env_name/bin/ffmpeg'

import energyflow as ef
import ot
from matplotlib import animation, rc
from IPython.display import HTML

# helper function to compute the EMD between (potentially unbalanced) jets
def emde(ev0, ev1, R=1, return_flow=False):
    pTs0, pTs1 = ev0[:, 0], ev1[:, 0]
    thetas = ot.dist(ev0[:, 1:3], ev1[:, 1:3], metric="euclidean")

    # add a fictitious particle to the lower-energy event to balance the energy
    pT0, pT1 = pTs0.sum(), pTs1.sum()
    pTs0 = np.hstack((pTs0, 0 if pT0 > pT1 else pT1 - pT0))
    pTs1 = np.hstack((pTs1, 0 if pT1 > pT0 else pT0 - pT1))

    # make its distance R to all particles in the other event
    Thetas = R * np.ones((np.shape(thetas)[0] + 1, np.shape(thetas)[1] + 1))
    Thetas[:-1, :-1] = thetas

    return np.array(
        ot.emd2(pTs0, pTs1, Thetas) if not return_flow else ot.emd(pTs0, pTs1, Thetas)
    )


# helper function to interpolate between the optimal transport of two events
def merge(ev0, ev1, R=1, lamb=0.5):
    # Cast as arrays so that can do multi-index slicing in emde
    ev0 = np.asarray(ev0, order="c")
    ev1 = np.asarray(ev1, order="c")
    G = emde(ev0, ev1, R=R, return_flow=True)

    merged = []
    for i in range(len(ev0)):
        for j in range(len(ev1)):
            if G[i, j] > 0:
                merged.append(
                    [
                        G[i, j],
                        ((lamb) * ev0[i, 1] + (1 - lamb) * ev1[j, 1]),
                        ((lamb) * ev0[i, 2] + (1 - lamb) * ev1[j, 2]),
                    ]
                )

    for i in range(len(ev0)):
        if G[i, -1] > 0:
            merged.append([G[i, -1] * lamb, ev0[i, 1], ev0[i, 2]])

    for j in range(len(ev1)):
        if G[-1, j] > 0:
            merged.append([G[-1, j] * (1 - lamb), ev1[j, 1], ev1[j, 2]])

    return np.asarray(merged)


#############################################################
# ANIMATION OPTIONS
#############################################################
zf = 30  # size of points in scatter plot
lw = 1  # linewidth of flow lines
fps = 60  # frames per second
nframes = 10 * fps  # total number of frames


#############################################################
# LOAD IN JETS
#############################################################
specs = ["375 <= corr_jet_pts <= 425", "abs_jet_eta < 1.9", "quality >= 2"]
events = ef.mod.load(*specs, dataset="cms", amount=0.01)

# events go here as lists of particle [pT,y,phi]
event0 = events.particles[14929][:, :3]

# center the jets
event0[:, 0] /= 100
event0[:, 1] -= np.average(event0[:, 1], weights=event0[:, 0])
event0[:, 2] -= np.average(event0[:, 2], weights=event0[:, 0])

print(event0)

constits_to_cluster = np.zeros(
    (len(event0), 4), dtype=[("pt", "f8"), ("eta", "f8"), ("phi", "f8"), ("E", "f8")]
)
for idx, c in enumerate(event0):
    constits_to_cluster[idx, 0] = c[0]
    constits_to_cluster[idx, 1] = c[1]
    constits_to_cluster[idx, 2] = c[2]
    constits_to_cluster[idx, 3] = 0.0

print(constits_to_cluster)

from pyjet import cluster

sequence = cluster(constits_to_cluster, R=99, p=-1)
jets = sequence.inclusive_jets()

print(jets[0].constituents_array())

stages_of_animation = []

for i in range(1, len(constits_to_cluster)):
    # print(str(i)+" PROTOJETS:")
    keep_for_this_stage = []
    for pj in sequence.exclusive_jets(i):
        # print(pj.pt,pj.eta,pj.phi,pj.mass)
        if pj.pt > 0:
            keep_for_this_stage.append(np.array([pj.pt, pj.eta, pj.phi]))
    # Ensure no empty lists
    if keep_for_this_stage:
        stages_of_animation.append(keep_for_this_stage)

for idx, stage in enumerate(stages_of_animation):
    print("Step " + str(idx))
    print(stage)

ev0 = np.copy(event0)

print(ev0)
print(stages_of_animation[1])
print(stages_of_animation[2])

#############################################################
# MAKE ANIMATION
#############################################################

fig, ax = plt.subplots()

merged = merge(stages_of_animation[1], stages_of_animation[2], lamb=0)
pts, ys, phis = merged[:, 0], merged[:, 1], merged[:, 2]

scat = ax.scatter(ys, phis, color=(1, 0, 0), s=pts)

# animation function. This is called sequentially
def animate(i):

    print(i)

    ax.clear()

    nstages = len(stages_of_animation)-1

    phases = int(nframes / (nstages+1))-2

    phase_index = int((i/phases)%phases)

    print(len(stages_of_animation), nframes, nstages, phases, phase_index)
    
    # first kind of phase is a static image of stageN

    if (int(phase_index+2) % 2) == 0:
        print("TIRED")
        lamb = nstages * phase_index / (nframes - 1)
        ev0 = stages_of_animation[phase_index]
        ev1 = stages_of_animation[phase_index]
        color = (1, 0, 0)

    # second kind of phase is a transition from stageN to stageN+1
    else:
        print("WIRED")
        lamb = nstages * (phase_index - nframes / nstages) / (nframes - phase_index)
        ev0 = stages_of_animation[phase_index]
        ev1 = stages_of_animation[phase_index + 1]
        color = (1, 0, 0)

    merged = merge(stages_of_animation[phase_index], stages_of_animation[phase_index + 1], lamb=lamb)
    pts, ys, phis = merged[:, 0], merged[:, 1], merged[:, 2]
    scat = ax.scatter(ys, phis, color=color, s=zf * pts, lw=0)

    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_axis_off()

    fig.savefig(str(int(i))+".png")
    
    return (scat,)

anim = animation.FuncAnimation(fig, animate, frames=nframes, repeat=True)
anim.save("akt.mp4", fps=fps, dpi=200, codec=None)
HTML(anim.to_html5_video())