<jupyter_start><jupyter_text>GravNet for top-taggingThe data we will use for these examples are the constituents of jets, large collections particles which have been clustered together by some combination algorithm (such at [anti-kt](https://arxiv.org/abs/0802.1189)). Jets can be produced in particle collisions due to multiple reasons, such as the emission of an energetic light-quark or gluon (QCD) and the decay of some heavy particle (e.g. a *top* quark). Our data contains simulations for jets produced by both QCD interactions and hadronic top-quark decays, and our task is to learn to classify the production mechanism of each jet (either QCD - target 0, or top - target 1).For each jet we have up to the first 200 jet constituents (our objects), which have been ordered by decreasing transverse momentum (*pt*), and we are provided with the Cartesian 4-momenta for each constituent: *px*, *py*, *pz*, and *E*.This dataset is available from here https://desycloud.desy.de/index.php/s/llbX3zpLhazgPJ6, or by running the cell below. Warning, the dataset is quite large ~1.6 GB. A discussion of the dataset is documented here: https://docs.google.com/document/d/1Hcuc6LBxZNX16zjEGeq16DAzspkDC4nDTyjMp1bWHRo/edit, and the reference is "Deep-learned Top Tagging with a Lorentz Layer" by <NAME>, <NAME>, <NAME> <NAME> ([arXiv: 1707.08966](https://arxiv.org/abs/1707.08966))For the purpose of this tutorial, to save time and bandwidth (24MB versus >1GB), I've shared a preprocessed subsample of this dataset. The link to which becomes invalid on 18/07/21 to encourage the use of the official source.<jupyter_code>import os
if 'google.colab' in str(get_ipython()):
    os.system('pip install lumin==0.8.0')<jupyter_output><empty_output><jupyter_text>Basic Modules<jupyter_code>from typing import *
from functools import partial
from fastcore.all import Path
PATH = Path('./data')
PATH.mkdir(parents=True, exist_ok=True)
SCRATCH = False  # Whether to import data from scratch<jupyter_output><empty_output><jupyter_text>Importing data<jupyter_code>if SCRATCH:
    import h5py
    import pandas as pd
    import numpy as np

    os.system("wget -O data.zip https://desycloud.desy.de/index.php/s/llbX3zpLhazgPJ6/download?path=%2F&files=train.h5")
    os.system("unzip data.zip")
    os.system("rm data.zip")
    LOAD = Path('./data')

#     LOAD = Path('../lumin/examples/data/top/')
    
    trn_store = pd.HDFStore(LOAD/'train.h5')

    def load_df(store:pd.io.pytables.HDFStore, n_evts:Optional[int]=None, n_const:Optional[int]=None) -> pd.DataFrame:
        df = trn_store.select("table",stop=n_evts).drop(columns=['truthE','truthPX','truthPY','truthPZ','ttv'])
        if n_const is not None: df.drop(columns=[f for f in df.columns if f != 'is_signal_new' and int(f[f.find('_')+1:]) >= n_const], inplace=True)
        rename = {f:f'{f[f.find("_")+1:]}_{f[:f.find("_")].lower()}' for f in df.columns if f != 'is_signal_new'}
        rename['is_signal_new'] = 'gen_target'
        df.rename(columns=rename, inplace=True)
        return df

    df = load_df(trn_store, 100000, 15)
    train_feats = [f for f in df.columns if f != 'gen_target']

    from lumin.data_processing.pre_proc import fit_input_pipe
    from lumin.data_processing.hep_proc import to_pt_eta_phi, proc_event, get_vecs

    def proc_df(df:pd.DataFrame, train_feats:List[str], preproc:bool=True) -> None:
        print('Moving to float32')
        df[train_feats] = df[train_feats].values.astype('float32')
        print('Replacing zeros')
        df[train_feats] = df[train_feats].replace(0.0, np.nan)  # Don't bias preprocessing with non-existent tracks
        print('Converting vectors')
        for v in get_vecs(train_feats): to_pt_eta_phi(df, v, True)    
        print('Processing event')
        proc_event(df, fix_phi=True, fix_y=True, fix_z=True, ref_vec_0='0', ref_vec_1='1', use_cartesian=True)
        print('Readding 0_py')
        df['0_py'] = 0.
        if preproc:
            print('Preprocessing features')
            print('Fitting preproc pipe')
            input_pipe = fit_input_pipe(df, train_feats, PATH/'input_pipe')
            print('Transforming features')
            df[train_feats] = input_pipe.transform(df[train_feats])  # Rescale and shift track momenta values
        df[train_feats] = df[train_feats].replace(np.nan, 0.0) # Matrix dat amust not contain NaNs

    proc_df(df, train_feats)

    from lumin.data_processing.file_proc import df2foldfile

    def df2fy(df:pd.DataFrame, train_feats:List[str], n_folds:int=10) -> None:
        df2foldfile(df=df, n_folds=n_folds,
                    cont_feats=train_feats, cat_feats=[], targ_feats='gen_target',
                    savename=PATH/'train', targ_type='int', strat_key='gen_target')

    df2fy(df, train_feats)<jupyter_output><empty_output><jupyter_text>Loading Data<jupyter_code>if not SCRATCH:
    PATH.mkdir(parents=True, exist_ok=True)
    if not (PATH/"train.hdf5").exists():
        os.system(f"wget -O {str(PATH)}/data.tar https://cernbox.cern.ch/index.php/s/YsKrkmIM6rBcnfG/download")  # Link expires on 18/07/21
        os.system(f"tar -C {str(PATH)} -xvf {str(PATH)}/data.tar" )
        os.system(f"rm {str(PATH)}/data.tar")
from lumin.nn.data.fold_yielder import FoldYielder

train_fy = FoldYielder(PATH/'train.hdf5')
train_fy[0]['inputs'].shape
train_fy.get_df(inc_inputs=True, suppress_warn=True)<jupyter_output>100000 datapoints loaded<jupyter_text>Model BuildingWe'll be trying out several different models, but in each case we'll mostly be changing the `head` of the model, so let's write a function to return a suitable `ModelBuilder` from a few arguments.<jupyter_code>from lumin.nn.models.model_builder import ModelBuilder
from lumin.nn.models.model import Model
from lumin.nn.models.blocks.body import FullyConnected
from lumin.nn.models.helpers import CatEmbedder
from lumin.nn.models.blocks.head import AbsHead, CatEmbHead
from typing import *
from torch import nn
bs = 256
objective = 'classifier'
n_out = 1

def get_model_builder(fy:FoldYielder, depth:int, width:int, head:AbsHead=CatEmbHead) -> ModelBuilder:
    opt_args = {'opt':'adam', 'eps':1e-08}
    cat_embedder = CatEmbedder.from_fy(fy)
    body = partial(FullyConnected, depth=depth, width=width, act='swish')
    model_builder = ModelBuilder(objective, cont_feats=fy.cont_feats, n_out=n_out, cat_embedder=cat_embedder, 
                                 opt_args=opt_args, body=body, head=head)
    print(Model(model_builder))
    return model_builder<jupyter_output><empty_output><jupyter_text>DNNThe basic approach to inputting the data is to keep the data flat and feed every single feature into the model simultaneously. Using the first 15 constituents, we have 60 input features and as we increase the number of constituents, the number of inputs would grow considerably. This means that the number of parameters in the input layer of our DNN would also grow.The aim will be to try and replace this fully-connected input layer with a more efficient layer. Nonetheless, the DNN allows all constituents to be considered simultaneously, so should theoretically offer the best performance, provided it can be trained.For our tests, we'll use a DNN with 4 hidden layers of 50 neurons, and for later examples, drop a hidden layer in favour of a different input head.<jupyter_code>model_builder = get_model_builder(train_fy, depth=4, width=50)<jupyter_output>Inputs:
60 Continuous: ['0_e', '0_px', '0_py', '0_pz', '1_e', '1_px', '1_py', '1_pz', '2_e', '2_px', '2_py', '2_pz', '3_e', '3_px', '3_py', '3_pz', '4_e', '4_px', '4_py', '4_pz', '5_e', '5_px', '5_py', '5_pz', '6_e', '6_px', '6_py', '6_pz', '7_e', '7_px', '7_py', '7_pz', '8_e', '8_px', '8_py', '8_pz', '9_e', '9_px', '9_py', '9_pz', '10_e', '10_px', '10_py', '10_pz', '11_e', '11_px', '11_py', '11_pz', '12_e', '12_px', '12_py', '12_pz', '13_e', '13_px', '13_py', '13_pz', '14_e', '14_px', '14_py', '14_pz']
                   
0  Categorical: []
                   
0  Matrix elements: []
                   

Model:
<bound method Module.parameters of Sequential(
  (0): CatEmbHead()
  (1): FullyConnected(
    (layers): Sequential(
      (0): Sequential(
        (0): Linear(in_features=60, out_features=50, bias=True)
        (1): Swish()
      )
      (1): Sequential(
        (0): Linear(in_features=50, out_features=50, bias=True)
        (1): Swish()
      )
      (2): Sequential(
        (0[...]<jupyter_text>So, our baseline model uses 10751 parameters. Let's train it and check the score. We'll use the area under the ROC curve as a metric<jupyter_code>from lumin.optimisation.hyper_param import lr_find
lr_finder = lr_find(train_fy, model_builder, bs, lr_bounds=[1e-6,1e0])
from lumin.nn.callbacks.cyclic_callbacks import *
from lumin.nn.metrics.class_eval import RocAucScore, BinaryAccuracy
from functools import partial

cb_partials = [partial(OneCycle, lengths=[5, 10],lr_range=[1e-4, 1e-2], mom_range=[0.85, 0.95], interp='cosine')]
metric_partials = [RocAucScore, BinaryAccuracy]
from lumin.nn.training.train import train_models
_ = train_models(train_fy, n_models=1,
                 model_builder=model_builder,
                 bs=bs,
                 cb_partials=cb_partials,
                 metric_partials=metric_partials,
                 n_epochs=15)<jupyter_output><empty_output><jupyter_text>Our DNN achieves a validation AUC of 0.936  Gravitational NetWe'll construct a graph for each jet by using the sub-jets as nodes, and their 4-momenta as features per node. Our data is still a flat table of features per jet, but it can automatically be reshaped into a tensor of (batch x node x feature) by defining what prefixes to consider as nodes (vectors) and what suffixes consider as features per node.E.g. given a feature name of "0_px", "0" should be considered as a node identifier, and "px" should considered as a feature of the "0" node.Note, the flat data will be reshaped during the forward pass, however we could have saved the data already processed into the (batch x node x feature) tensor.<jupyter_code>from lumin.data_processing.hep_proc import get_vecs

fpv = ['px', 'py', 'pz', 'e']
vecs = [str(i) for i in sorted([int(f) for f in get_vecs(train_fy.cont_feats)])]
fpv, vecs<jupyter_output><empty_output><jupyter_text>GNNs in LUMIN consist of two parts:- extractor: GNN layers used to compute new features per node- collapser: aggregates/flattens node features into flat vector. Can also apply a few fully connected layers to nodes prior to collapsing`GNNHead` class wraps the extractor and collapser in a single object. Forward method:```pythondef forward(self, x:Union[Tensor,Tuple[Tensor,Tensor]]) -> Tensor:    r'''    Passes input through the GravNet head and returns a flat tensor.    Arguments:        x: If a tuple, the second element is assumed to the be the matrix data. If a flat tensor, will convert the data to a matrix    Returns:        Resulting tensor    '''    x = self._process_input(x)   features per vtx per datapoint    if self.use_in_bn: x = self.bn(x)    if self.cat_means:        if self.row_wise: x = torch.cat([x,x.mean(1).unsqueeze(2).repeat_interleave(repeats=x.shape[1],dim=2).transpose(1,2)],dim=2)        else:             x = torch.cat([x,x.mean(2).unsqueeze(1).repeat_interleave(repeats=x.shape[2],dim=1).transpose(1,2)],dim=1)    x = self.extractor(x)   new features per vtx per datapoint    if not self.row_wise: x = x.transpose(1,2)    x = self.collapser(x)   features per datapoint    return x```<jupyter_code>from lumin.nn.models.blocks.head import GNNHead<jupyter_output><empty_output><jupyter_text>GravNet layers can be stacked together. `GravNet` class passes node features through a set number of layers, and concatenates the output of each one to the initial node features. Forward method```pythondef forward(self, x:Tensor) -> Tensor:    r'''    Passes input through the GravNet head.    Arguments:        x: row-wise tensor (batch x vertices x features)    Returns:        Resulting tensor row-wise tensor (batch x vertices x new features)    '''    outs = [x]    for l in self.grav_layers: outs.append(l(outs[-1]))    return torch.cat(outs[1:], dim=-1)```Each layer is a `GravNetLayer` class. Breaking down the forwards pass step by step:```pythondef forward(self, x:Tensor) -> Tensor:    r'''    Pass batch of vertices through GravNet layer and return new features per vertex    Arguments:        x: Incoming data (batch x vertices x features)    Returns:        Data with new features per vertex (batch x vertices x new features)    '''     Concat means    if self.cat_means:        x = torch.cat([x,x.mean(1).unsqueeze(2).repeat_interleave(repeats=x.shape[1],dim=-1).transpose(1,2)],dim=2)```The features per node can be optionally extended to include the mean values of each feature across the graph.```python     Compute spatial and vertex features    slr = self.f_slr(x)    s,lr = slr[:,:,:self.n_s],slr[:,:,self.n_s:]```An initial DNN `f_slr` is used to compute: `s` spatial coordinates, and `lr` features per node. In PyTorch, DNNs automatically broadcast across all tensor dimensions except the last one, so `f_slr(x)` applies the same DNN to every node in every graph in the batch.```python     kNN    d_jk = torch.norm(s[:,:,None]-s[:,None], dim=-1)    idxs = self._knn(d_jk)```For every node, compute the Euclidean distance between it and all nodes in the same graph (including itself) using `s` as coordinates in latent space. For every node, select the *k* nearest neighbours (including itself). *N.B.* "for" is used here to describe the process, but the above lines are vectorised and compute all distances in a single function.```python    def _knn(self, dists:Tensor) -> Tuple[Tensor,Tensor,Tensor]:        idxs = dists.argsort()        i = np.arange(dists.shape[0])[:,None,None]        j = np.arange(dists.shape[1])[None,:,None]        return i,j,idxs[:,:,:self.k]```This just orders the distances, and returns indices to look up the *k* nearest neighbours.```python    d_jk = d_jk[idxs]     lr = lr[:,None].expand(lr.shape[0],lr.shape[1],lr.shape[1],lr.shape[2])[idxs]```The distances `d_jk` are already a (batch x node x node) tensor, and are indexed to return a (batch x node x *k*-nearest nodes) tensor.  The `lr` features are expanded and indexed into a (batch x node x *k*-nearest nodes x features).```python    v_jk = self.potential(d_jk)    ft_ijk = lr*v_jk.unsqueeze(-1)    if self.use_sa: ft_ijk = self.sa_agg(ft_ijk)```The `d_jk` distances are used to weight feature through a potential function, e.g. `torch.exp(-10*(d_jk**2)`. I.e. close-by nodes' features have a higher weight than those of more distance nodes. Optionally, and unique to LUMIN, [self-attention](https://arxiv.org/abs/1706.03762) can be used to here to also provide a learnable feature weighting (still experimental, and sometimes reduces performance).```python    fp = [x]    for agg in self.agg_methods: fp.append(agg(ft_ijk))    fp = torch.cat(fp,dim=-1)```Neighbourhhood features are constructed by aggregating the distance-weighted features of the *k* nearest nodes, e.g. taking the mean and maximum. These are then concatenated with the original features of the node.```python     Output    return self.f_out(fp)```Finally, the output features per node per graph are computed using a second DNN based on the node features + the neighbourhood features.<jupyter_code>from lumin.nn.models.blocks.gnn_blocks import  GravNet<jupyter_output><empty_output><jupyter_text>Finally, we pass the (batch x node x features) tensor through the `GraphCollapser` to transform it into a (batch x feature) tensor:```pythondef forward(self, x:Tensor) -> Tensor:    r'''    Collapses features per vertex down to features    Arguemnts:        x: incoming data (batch x vertices x features)    Returns:        Flattened data (batch x flat features)    '''    if self.global_feat_vec and self.gfv_pos == 'pre-initial':        x = torch.cat([x,x.mean(1).unsqueeze(2).repeat_interleave(repeats=x.shape[1],dim=2).transpose(1,2)],dim=2)     x = self.f_inital(x)    if self.n_sa_layers > 0:        outs = []        for i, sa in enumerate(self.sa_layers): outs.append(sa(outs[-1] if i > 0 else x))        x = torch.cat(outs, dim=-1)    if self.global_feat_vec and self.gfv_pos == 'pre-final':        x = torch.cat([x,x.mean(1).unsqueeze(2).repeat_interleave(repeats=x.shape[1],dim=2).transpose(1,2)],dim=2)    x = self.f_final(x)    return self._agg(x)```This is quite configurable:- DNNs (`f_inital` and `f_final`) can be used to reduce the number of features per node per graph and two different stages (or they can be `lambda x: x`).- A "global feature vertex" can be computed, which concatenates the mean values of features across all nodes per graph (like the `cat_means` in the gravnet layers). This can be done before either of the two DNNs.- Multiple self-attention blocks can be used to augment the features, a la [point-cloud transformer](https://link.springer.com/article/10.1007/s41095-021-0229-5) (see also [Point Cloud Transformers applied to Collider Physics](https://iopscience.iop.org/article/10.1088/2632-2153/ac07f6))- Finally the tensor is either reshaped or aggregated:```pythondef _agg(self, x:Tensor) -> Tensor:    if self.flatten: return x.reshape((len(x),-1))    else:            return torch.cat([agg(x) for agg in self.agg_methods],-1)```where `agg_methods` are the mean and maximum of the features across the nodes.<jupyter_code>from lumin.nn.models.blocks.gnn_blocks import  GravNet, GraphCollapser

head = partial(GNNHead, vecs=vecs, feats_per_vec=fpv,
               extractor=partial(GravNet,  # Use GravNet for feature extraction
                                 cat_means=True,  # Concatenate the means of features as new features
                                 f_slr_depth=3,  # Use 3 hidden layers to compute the coordinates and features in latent space
                                 n_s=4,  # Use 4 dimensions for latent space
                                 n_lr=22,  # Compute 22 features per vertex for the latent representation
                                 k=10,  # Each vertex should consider the 10 nearest vertices (self included) when computing neighbourhood features
                                 f_out_depth=1,  # Use one hidden layer to compute the new features per vertex
                                 n_out=[48,48,48,48],  # Use 4 GravNet layers, each computing 48 new features per vertex
                                 act='swish'),  # Use swish activation functions
               collapser=partial(GraphCollapser,  # Use default graph collapser
                                 flatten=True,  # Reshape the data into (batch x features) without aggregation. Vertices are ordered by pT
                                 f_final_outs=[40],  # Compute 40 features per vertex based on the GravNet layer outputs
                                 act='swish'))  # Use swish activation functions
model_builder = get_model_builder(train_fy, depth=3, width=50, head=head)
lr_finder = lr_find(train_fy, model_builder, bs, lr_bounds=[1e-6,1e-2])
cb_partials[0].keywords['lr_range'] = [1e-4, 1e-2]
_ = train_models(train_fy, n_models=1,
                 model_builder=model_builder,
                 bs=bs,
                 cb_partials=cb_partials,
                 metric_partials=metric_partials,
                 n_epochs=15)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11166150670794633, 'fraction_numerical': 0.021001031991744065, 'mean_word_length': 3.822343866633491, 'pattern_counts': {'":': 0, '<': 48, '<?xml version=': 0, '>': 61, 'https://': 9, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '576970', 'license': 'Apache-2.0', 'n_tokens_mistral': 6508, 'n_tokens_neox': 6116, 'n_words': 1960, 'path': 'GravNet_for_top_tagging.ipynb', 'repo_name': 'GilesStrong/workshop_LIP_GNN', 'size': '19389'}