testing demo

README.md

@@ -1,3 +1,3 @@
 ---
 license: mit
----
+---

nemo/configs/physicsnemo/physicsnemo.yaml

@@ -1,54 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd
-Training_Directory: trainings/physicsnemo/test
-Model:
-  module: models.meshgraphnet
-  class: MeshGraphNet
-  args:
-    input_dim_nodes: 7
-    input_dim_edges: 3
-    output_dim: 64
-Training:
-  epochs: 500
-  batch_size: 1024
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 1024
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

nemo/configs/stats_100K/ttH_CP_even_vs_odd.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd
-Training_Directory: trainings/stats_100K/ttH_CP_even_vs_odd
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 1024
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 1024
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

nemo/models/Edge_Network.py

@@ -1,163 +0,0 @@
-import dgl
-import dgl.nn as dglnn
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-import sys
-import os
-file_path = os.getcwd()
-sys.path.append(file_path)
-
-def Make_SLP(in_size, out_size, activation = nn.ReLU, dropout = 0):
-    layers = []
-    layers.append(nn.Linear(in_size, out_size))
-    layers.append(activation())
-    layers.append(nn.Dropout(dropout))
-    return layers
-
-def Make_MLP(in_size, hid_size, out_size, n_layers, activation = nn.ReLU, dropout = 0):
-    layers = []
-    if n_layers > 1:
-        layers += Make_SLP(in_size, hid_size, activation, dropout)
-        for i in range(n_layers-2):
-            layers += Make_SLP(hid_size, hid_size, activation, dropout)
-        layers += Make_SLP(hid_size, out_size, activation, dropout)
-    else:
-        layers += Make_SLP(in_size, out_size, activation, dropout)
-    layers.append(torch.nn.LayerNorm(out_size))
-    return nn.Sequential(*layers)
-
-class Edge_Network(nn.Module):
-    def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        if (len(sample_global) == 0):
-            self.has_global = False
-        else:
-            self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def forward(self, g, global_feats):
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-
-        batch_num_nodes = None
-        sum_weights = None
-        if "w" in g.ndata:
-            batch_indices = g.batch_num_nodes()
-            # Find non-zero rows (non-padded nodes)
-            non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1)
-            # Split the mask according to the batch indices
-            batch_num_nodes = []
-            start_idx = 0
-            for num_nodes in batch_indices:
-                end_idx = start_idx + num_nodes
-                non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item()
-                batch_num_nodes.append(non_padded_count)
-                start_idx = end_idx
-            batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, 64)
-            global_feats = batch_num_nodes[:, None].to(torch.float)
-
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            if "w" in g.ndata:
-                mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights
-                h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1))
-            else:
-                h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-    
-    def representation(self, g, global_feats):
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-
-        batch_num_nodes = None
-        sum_weights = None
-        if "w" in g.ndata:
-            batch_indices = g.batch_num_nodes()
-            # Find non-zero rows (non-padded nodes)
-            non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1)
-            # Split the mask according to the batch indices
-            batch_num_nodes = []
-            start_idx = 0
-            for num_nodes in batch_indices:
-                end_idx = start_idx + num_nodes
-                non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item()
-                batch_num_nodes.append(non_padded_count)
-                start_idx = end_idx
-            batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, 64)
-            global_feats = batch_num_nodes[:, None].to(torch.float)
-
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            if "w" in g.ndata:
-                mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights
-                h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1))
-            else:
-                h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        before_global_decoder = h_global
-        after_global_decoder = self.global_decoder(before_global_decoder)
-        after_classify = self.classify(after_global_decoder)
-        return before_global_decoder, after_global_decoder, after_classify
-    
-    def __str__(self):
-        layer_names = ["node_encoder", "edge_encoder", "global_encoder", 
-                  "node_update", "edge_update", "global_update", "global_decoder"]
-        
-        layers = [self.node_encoder, self.edge_encoder, self.global_encoder,
-                  self.node_update, self.edge_update, self.global_update, self.global_decoder]
-        
-        for i in range(len(layers)):
-            print(layer_names[i])
-            for layer in layers[i].children():
-                if isinstance(layer, nn.Linear):
-                    print(layer.state_dict())
-        
-        print("classify")
-        print(self.classify.weight)
-        return ""
-
-    def __name__():
-        return "Edge_Network"

nemo/models/GCN.py

@@ -1,1818 +0,0 @@
-import dgl
-import dgl.nn as dglnn
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-import sys
-import os
-file_path = os.getcwd()
-sys.path.append(file_path)
-
-import root_gnn_base.dataset as datasets
-from root_gnn_base import utils
-
-import gc
-
-def Make_SLP(in_size, out_size, activation = nn.ReLU, dropout = 0):
-    layers = []
-    layers.append(nn.Linear(in_size, out_size))
-    layers.append(activation())
-    layers.append(nn.Dropout(dropout))
-    return layers
-
-def Make_MLP(in_size, hid_size, out_size, n_layers, activation = nn.ReLU, dropout = 0):
-    layers = []
-    if n_layers > 1:
-        layers += Make_SLP(in_size, hid_size, activation, dropout)
-        for i in range(n_layers-2):
-            layers += Make_SLP(hid_size, hid_size, activation, dropout)
-        layers += Make_SLP(hid_size, out_size, activation, dropout)
-    else:
-        layers += Make_SLP(in_size, out_size, activation, dropout)
-    layers.append(torch.nn.LayerNorm(out_size))
-    return nn.Sequential(*layers)
-
-class MLP(nn.Module):
-    def __init__(self, in_size, hid_size, out_size, n_layers, activation = nn.ReLU, dropout = 0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating MLP: {kwargs}')
-        self.layers = Make_MLP(in_size, hid_size, hid_size, n_layers-1, activation, dropout)
-        self.linear = nn.Linear(hid_size, out_size)
-    
-    def forward(self, x):
-        return self.linear(self.layers(x))
-
-def broadcast_global_to_nodes(g, globals):
-    boundaries = g.batch_num_nodes()
-    return torch.repeat_interleave(globals, boundaries, dim=0)
-
-def broadcast_global_to_edges(g, globals):
-    boundaries = g.batch_num_edges()
-    return torch.repeat_interleave(globals, boundaries, dim=0)
-
-def copy_v(edges):
-    return {'m_v': edges.dst['h']}
-
-def partial_reset(model : nn.Module):
-    in_size = len(model.classify.weight[0])
-    out_size = len(model.classify.weight)
-    device = next(model.classify.parameters()).device
-    torch.manual_seed(2)
-    model.classify = nn.Linear(in_size, out_size)
-    model.classify.to(device)
-    print(model.classify.weight)
-
-def print_model(model: nn.Module):
-    print(model)
-
-def print_mlp(layer):
-    for l in layer.children():
-        if isinstance(l, nn.Linear):
-            print(l.state_dict())
-        else:
-            print(l)
-
-
-# This function returns a model with the whole GNN completely reset
-def full_reset(model : nn.Module):
-    mlp_list = [model.node_encoder, model.edge_encoder, model.global_encoder,
-                model.node_update, model.edge_update, model.global_update,
-                model.global_decoder]
-    
-    for mlp in mlp_list:
-        for layer in mlp.children():
-            if hasattr(layer, 'reset_parameters'):
-                layer.reset_parameters()
-    partial_reset(model)
-
-
-class GCN(nn.Module):
-    def __init__(self, in_size, hid_size, out_size, n_layers, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.layers = nn.ModuleList()
-
-        # two-layer GCN
-        self.layers.extend(
-            [nn.Linear(in_size, hid_size),] +
-            [nn.Linear(hid_size, hid_size) for i in range(n_layers)] +
-            [dglnn.GraphConv(hid_size, hid_size) for i in range(n_layers)] +
-            [nn.Linear(hid_size, hid_size) for i in range(n_layers)]
-        )
-        self.classify = nn.Linear(hid_size, out_size)
-        #self.dropout = nn.Dropout(0.05)
-
-    def forward(self, g):
-        h = g.ndata['features']
-        for i, layer in enumerate(self.layers):
-            if i >= self.n_layers + 1 and i < self.n_layers * 2 + 1:
-                h = layer(g, h)
-            else:
-                h = layer(h)
-            h = F.relu(h)
-        with g.local_scope():
-            g.ndata['h'] = h
-            # Calculate graph representation by average readout.
-            hg = dgl.mean_nodes(g, 'h')
-            return self.classify(hg)
-    
-class GCN_global(nn.Module):
-    def __init__(self, in_size, hid_size=4, out_size=1, n_layers=1, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-
-        #encoder
-        self.node_encoder = Make_MLP(in_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(1, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GCN
-        self.node_update = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.conv = dglnn.GraphConv(hid_size, hid_size)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def forward(self, g):
-        h = self.node_encoder(g.ndata['features'])
-        h_global = self.global_encoder(g.batch_num_nodes()[:, None].to(torch.float))
-        for i in range(self.n_layers):
-            h = self.node_update(h)
-            h = self.conv(g, h)
-            g.ndata['h'] = h
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-    
-class GCN_global_2way(nn.Module):
-    def __init__(self, in_size, hid_size=4, out_size=1, n_layers=1, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-
-        #encoder
-        self.node_encoder = Make_MLP(in_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(1, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GCN
-        self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.conv = dglnn.GraphConv(hid_size, hid_size)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def forward(self, g):
-        h = self.node_encoder(g.ndata['features'])
-        h_global = self.global_encoder(g.batch_num_nodes()[:, None].to(torch.float))
-        for i in range(self.n_layers):
-            h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h = self.conv(g, h)
-            g.ndata['h'] = h
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-
-class Transferred_Learning(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-
-        if (len(sample_global) == 0):
-            self.has_global = False
-        else:
-            self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        self.global_decoder = Make_MLP(pretraining_model['args']['hid_size'], hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def forward(self, g, global_feats):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return self.classify(self.global_decoder(h_global))
-
-class Transferred_Learning_Graph(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, additional_proc_steps=1, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-
-        if (len(sample_global) == 0):
-            self.has_global = False
-        else:
-            self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        self.additional_proc_steps = additional_proc_steps
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-
-    def forward(self, g, global_feats):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        for j in range(self.additional_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-
-class Transferred_Learning_Parallel(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size + pretraining_model['args']['hid_size'], out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return h_global
-
-    def forward(self, g, global_feats):
-        pretrained_global = self.Pretrained_Output(g.clone())
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-
-        return self.classify(torch.cat((pretrained_global, h_global), dim = 1))
-    
-class Transferred_Learning_Sequential(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        #encoder
-        self.mlp = Make_MLP(pretraining_model['args']['hid_size'], hid_size, hid_size, n_layers, dropout=dropout)
-    
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return h_global
-
-    def forward(self, g, global_feats):
-        pretrained_global = self.Pretrained_Output(g.clone())
-        global_features = self.mlp(pretrained_global)
-        return self.classify(global_features)
-
-
-class Transferred_Learning_Message_Passing(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        #encoder
-        self.mlp = Make_MLP(pretraining_model['args']['hid_size']*pretraining_model['args']['n_proc_steps'], hid_size, hid_size, n_layers, dropout=dropout)
-    
-        self.classify = nn.Linear(hid_size, out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        message_passing = None
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-            if (message_passing is None):
-                message_passing = h_global.clone()
-            else:
-                message_passing = torch.cat((message_passing, h_global.clone()), dim=1)
-        h_global = self.TL_global_decoder(h_global)
-        return message_passing
-
-    def forward(self, g, global_feats):
-        pretrained_global = self.Pretrained_Output(g.clone())
-        #print(f"message_passing layers have size = {pretrained_global.shape}")
-        #print(pretrained_global)
-        global_features = self.mlp(pretrained_global)
-        return self.classify(global_features)
-    
-class Transferred_Learning_Message_Passing_Parallel(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        # Freeze Weights
-        for param in self.pretrained_model.parameters():
-            param.requires_grad = False  # Freeze all layers
-
-        self.classify = nn.Linear(pretraining_model['args']['hid_size']*pretraining_model['args']['n_proc_steps'] + hid_size, out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        message_passing = None
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-            if (message_passing is None):
-                message_passing = h_global.clone()
-            else:
-                message_passing = torch.cat((message_passing, h_global.clone()), dim=1)
-        h_global = self.TL_global_decoder(h_global)
-        return message_passing
-
-    def forward(self, g, global_feats):
-        pretrained_message = self.Pretrained_Output(g.clone())
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(torch.cat((pretrained_message, h_global), dim = 1))
-    
-class Transferred_Learning_Finetuning(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=False, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-
-        if (len(sample_global) == 0):
-            self.has_global = False
-        else:
-            self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        print(f"Freeze Pretraining = {frozen_pretraining}")
-        if (frozen_pretraining):
-            for param in self.pretrained_model.parameters():
-                param.requires_grad = False  # Freeze all layers
-            for param in self.pretrained_model[7]:
-                param.requires_grad = True
-
-        torch.manual_seed(2)
-        self.classify = nn.Linear(pretraining_model['args']['hid_size'], out_size)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return h_global
-
-    def forward(self, g, global_feats):
-        h_global = self.Pretrained_Output(g.clone())
-        return self.classify(h_global)
-    
-    def representation(self, g, global_feats):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        
-        before_global_decoder = h_global
-        after_global_decoder = self.TL_global_decoder(before_global_decoder)
-        after_classify = self.classify(after_global_decoder)
-        return before_global_decoder, after_global_decoder, after_classify
-
-    def __str__(self):
-        layer_names = ["node_encoder", "edge_encoder", "global_encoder", 
-                  "node_update", "edge_update", "global_update", "global_decoder"]
-        
-        layers = [self.pretrained_model[1], self.pretrained_model[2], self.pretrained_model[3],
-                  self.pretrained_model[4], self.pretrained_model[5], self.pretrained_model[6], 
-                  self.pretrained_model[7]]
-        
-        for i in range(len(layers)):
-            print(layer_names[i])
-            for layer in layers[i].children():
-                if isinstance(layer, nn.Linear):
-                    print(layer.state_dict())
-        
-        print("classify")
-        print(self.classify.weight)
-        return ""
-    
-
-class Transferred_Learning_Parallel_Finetuning(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, learning_rate=0.0001, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-
-        self.learning_rate = learning_rate
-
-        self.parallel_params = []
-        self.finetuning_params = []
-
-
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        self.finetuning_params.append(self.pretrained_model)
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size + pretraining_model['args']['hid_size'], out_size)
-
-        self.parallel_params.append(self.node_encoder)
-        self.parallel_params.append(self.edge_encoder)
-        self.parallel_params.append(self.global_encoder)
-        self.parallel_params.append(self.node_update)
-        self.parallel_params.append(self.edge_update)
-        self.parallel_params.append(self.global_update)
-        self.parallel_params.append(self.global_decoder)
-        self.parallel_params.append(self.classify)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_edge_encoder(self, x):
-        for layer in self.pretrained_model[2]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_node_update(self, x):
-        for layer in self.pretrained_model[4]:
-            x = layer(x)
-        return x
-
-    def TL_edge_update(self, x):
-        for layer in self.pretrained_model[5]:
-            x = layer(x)
-        return x
-
-    def TL_global_update(self, x):
-        for layer in self.pretrained_model[6]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def Pretrained_Output(self, g):
-        h = self.TL_node_encoder(g.ndata['features'])
-        e = self.TL_edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return h_global
-
-    def forward(self, g, global_feats):
-        pretrained_global = self.Pretrained_Output(g.clone())
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-
-        return self.classify(torch.cat((pretrained_global, h_global), dim = 1))
-    
-    def parameters(self, recurse: bool = True):
-        params = []
-        for model_section in self.parallel_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["trainable_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["trainable_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        for model_section in self.finetuning_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["finetuning_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["finetuning_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        return params    
-    
-class Attention(nn.Module):
-    def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        self.hid_size = hid_size
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-        
-        #attention
-        self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True)
-        self.queries = nn.Linear(hid_size, hid_size)
-        self.keys = nn.Linear(hid_size, hid_size)
-        self.values = nn.Linear(hid_size, hid_size)
-
-    def forward(self, g, global_feats):
-        h = self.node_encoder(g.ndata['features'])
-        g.ndata['h'] = h
-
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-
-        batch_num_nodes = None
-        sum_weights = None
-        if "w" in g.ndata:
-            batch_indices = g.batch_num_nodes()
-            # Find non-zero rows (non-padded nodes)
-            non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1)
-            # Split the mask according to the batch indices
-            batch_num_nodes = []
-            start_idx = 0
-            for num_nodes in batch_indices:
-                end_idx = start_idx + num_nodes
-                non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item()
-                batch_num_nodes.append(non_padded_count)
-                start_idx = end_idx
-            batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
-            global_feats = batch_num_nodes[:, None].to(torch.float)
-
-        h_global = self.global_encoder(global_feats)
-
-        h_original_shape = h.shape
-        num_graphs = len(dgl.unbatch(g))
-        num_nodes = g.batch_num_nodes()[0].item()
-        padding_mask = g.ndata['padding_mask'] > 0
-        padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes))
-
-        h = g.ndata['h']
-        query = self.queries(h)
-        key = self.keys(h)
-        value = self.values(h)
-        query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1]))
-        key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1]))
-        value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1]))
-        h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask)
-        h = torch.reshape(h, h_original_shape)
-        
-        h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1))
-        g.ndata['h'] = h
-        mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights
-        h_global = self.global_update(torch.cat((h_global, mean_nodes), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-
-class Attention_Edge_Network(nn.Module):
-    def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-        
-
-        #attention
-        self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True)
-        self.queries = nn.Linear(hid_size, hid_size)
-        self.keys = nn.Linear(hid_size, hid_size)
-        self.values = nn.Linear(hid_size, hid_size)
-
-    def forward(self, g, global_feats):
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-
-        h = g.ndata['h']
-        h_original_shape = h.shape
-        num_graphs = len(dgl.unbatch(g))
-        num_nodes = g.batch_num_nodes()[0].item()
-        padding_mask = g.ndata['padding_mask'].bool()
-
-        padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes))
-
-        for i in range(self.n_proc_steps):
-            
-            h = g.ndata['h']
-            query = self.queries(h)
-            key = self.keys(h)
-            value = self.values(h)
-            query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1]))
-            key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1]))
-            value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1]))
-            h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask)
-            h = torch.reshape(h, h_original_shape)
-
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h', 'w'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-
-class Attention_Unbatched(nn.Module):
-    def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-        
-
-        #attention
-        self.multihead_attn = nn.MultiheadAttention(hid_size, 1, dropout=dropout)
-        self.queries = nn.Linear(hid_size, hid_size)
-        self.keys = nn.Linear(hid_size, hid_size)
-        self.values = nn.Linear(hid_size, hid_size)
-
-
-
-    def forward(self, g, global_feats):
-
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-
-        for i in range(self.n_proc_steps):
-
-            unbatched_g = dgl.unbatch(g)
-            for graph in unbatched_g:
-                h = graph.ndata['h']
-                h, _ = self.multihead_attn(self.queries(h), self.keys(h), self.values(h))
-                graph.ndata['h'] = h
-            g = dgl.batch(unbatched_g)
-
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return self.classify(h_global)
-    
-class Transferred_Learning_Attention(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, num_heads, dropout=0, learning_rate=0.0001, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        self.hid_size = hid_size
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.learning_rate = learning_rate
-
-        self.pretraining_params = []
-        self.attention_params = []
-        
-        self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-        checkpoint = torch.load(pretraining_path)
-        self.pretrained_model.load_state_dict(checkpoint['model_state_dict'])
-        pretrained_layers = list(self.pretrained_model.children())
-        pretrained_layers = pretrained_layers[:-1]
-        self.pretrained_model = nn.Sequential(*pretrained_layers)
-
-        self.pretraining_params.append(self.pretrained_model[1])
-        self.pretraining_params.append(self.pretrained_model[3])
-        self.pretraining_params.append(self.pretrained_model[7])
-
-        #attention
-        self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True)
-        self.queries = nn.Linear(hid_size, hid_size)
-        self.keys = nn.Linear(hid_size, hid_size)
-        self.values = nn.Linear(hid_size, hid_size)
-
-        self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        
-        self.classify = nn.Linear(pretraining_model['args']['hid_size'], out_size)
-
-        self.attention_params.append(self.multihead_attn)
-
-        self.attention_params.append(self.queries)
-        self.attention_params.append(self.keys)
-        self.attention_params.append(self.values)
-        self.attention_params.append(self.classify)
-        self.attention_params.append(self.node_update)
-        self.attention_params.append(self.global_update)
-
-    def TL_node_encoder(self, x):
-        for layer in self.pretrained_model[1]:
-            x = layer(x)
-        return x
-    
-    def TL_global_encoder(self, x):
-        for layer in self.pretrained_model[3]:
-            x = layer(x)
-        return x
-    
-    def TL_global_decoder(self, x):
-        for layer in self.pretrained_model[7]:
-            x = layer(x)
-        return x
-
-    def forward(self, g, global_feats):
-        h = self.TL_node_encoder(g.ndata['features'])
-        g.ndata['h'] = h
-
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-
-        batch_num_nodes = None
-        sum_weights = None
-        if "w" in g.ndata:
-            batch_indices = g.batch_num_nodes()
-            # Find non-zero rows (non-padded nodes)
-            non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1)
-            # Split the mask according to the batch indices
-            batch_num_nodes = []
-            start_idx = 0
-            for num_nodes in batch_indices:
-                end_idx = start_idx + num_nodes
-                non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item()
-                batch_num_nodes.append(non_padded_count)
-                start_idx = end_idx
-            batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
-            global_feats = batch_num_nodes[:, None].to(torch.float)
-
-        h_global = self.TL_global_encoder(global_feats)
-
-        h_original_shape = h.shape
-        num_graphs = len(dgl.unbatch(g))
-        num_nodes = g.batch_num_nodes()[0].item()
-        padding_mask = g.ndata['padding_mask'] > 0
-        padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes))
-
-        h = g.ndata['h']
-        query = self.queries(h)
-        key = self.keys(h)
-        value = self.values(h)
-        query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1]))
-        key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1]))
-        value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1]))
-        h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask)
-        h = torch.reshape(h, h_original_shape)
-        
-        h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1))
-        g.ndata['h'] = h
-        mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights
-        h_global = self.global_update(torch.cat((h_global, mean_nodes), dim = 1))
-        h_global = self.TL_global_decoder(h_global)
-        return self.classify(h_global)
-    
-    def parameters(self, recurse: bool = True):
-        params = []
-        for model_section in self.pretraining_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        for model_section in self.attention_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["attention_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["attention_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        return params
-    
-class Multimodel_Transferred_Learning(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=True, learning_rate=None, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.learning_rate = learning_rate
-        input_size = 0
-
-        self.pretraining_params = []
-        self.model_params = []
-        
-        self.pretrained_models = []
-        for model, path in zip(pretraining_model, pretraining_path):
-            input_size += model['args']['hid_size']
-            model = utils.buildFromConfig(model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-            checkpoint = torch.load(path)['model_state_dict']
-            new_state_dict = {}
-            for k, v in checkpoint.items():
-                new_key = k.replace('module.', '')
-                new_state_dict[new_key] = v
-            model.load_state_dict(new_state_dict)
-            pretrained_layers = list(model.children())
-            pretrained_layers = pretrained_layers[:-1]
-        
-            model = nn.Sequential(*pretrained_layers)
-
-            # Freeze Weights
-            print(f"Freeze Pretraining = {frozen_pretraining}")
-            if (frozen_pretraining):
-                for param in model.parameters():
-                    param.requires_grad = False  # Freeze all layers
-            self.pretraining_params.append(model)
-            self.pretrained_models.append(model)
-
-        print(f"len(pretrained_models) = {len(self.pretrained_models)}")
-        print(f"input size = {input_size}")
-
-        self.final_mlp = Make_MLP(input_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-        self.model_params.append(self.final_mlp)
-        self.model_params.append(self.classify)
-
-    def TL_node_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][1]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][1]:
-                x = layer(x)
-            return x
-    
-    def TL_edge_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][2]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][2]:
-                x = layer(x)
-            return x
-    
-    def TL_global_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][3]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][3]:
-                x = layer(x)
-            return x
-    
-    def TL_node_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][4]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][4]:
-                x = layer(x)
-            return x
-
-    def TL_edge_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][5]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][5]:
-                x = layer(x)
-            return x
-
-    def TL_global_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][6]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][6]:
-                x = layer(x)
-            return x
-    
-    def TL_global_decoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][7]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][7]:
-                x = layer(x)
-            return x
-
-    def Pretrained_Output(self, g, model_idx):
-        h = self.TL_node_encoder(g.ndata['features'], model_idx)
-        e = self.TL_edge_encoder(g.edata['features'], model_idx)
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats, model_idx)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1), model_idx)
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1), model_idx)
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1), model_idx)
-        # h_global = self.TL_global_decoder(h_global, model_idx)
-        return h_global
-
-    def forward(self, g, global_feats):
-        h_global = []
-        for i in range(len(self.pretrained_models)):
-            h_global.append(self.Pretrained_Output(g.clone(), i))
-        h_global = torch.concatenate(h_global, dim=1)
-        return self.classify(self.final_mlp(h_global))
-    
-    def to(self, device):
-        for i in range(len(self.pretrained_models)):
-            self.pretrained_models[i].to(device)
-        self.classify.to(device)
-        self.final_mlp.to(device)
-        return self
-    
-    def parameters(self, recurse: bool = True):
-        params = []
-        for model_section in self.pretraining_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.00001})
-        for model_section in self.model_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["model_lr"]): 
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["model_lr"]})
-            else:
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        return params    
-
-
-class MultiModel(nn.Module):
-    def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=True, learning_rate=None, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        self.learning_rate = learning_rate
-        input_size = 0
-
-        self.model_params = []
-        self.pretraining_params = []
-        
-        self.pretrained_models = []
-        for model, path in zip(pretraining_model, pretraining_path):
-            input_size += model['args']['hid_size']
-            model = utils.buildFromConfig(model, {'sample_graph': sample_graph, 'sample_global': sample_global})
-
-            checkpoint = torch.load(path)['model_state_dict']
-            new_state_dict = {}
-            for k, v in checkpoint.items():
-                new_key = k.replace('module.', '')
-                new_state_dict[new_key] = v
-            model.load_state_dict(new_state_dict)
-            pretrained_layers = list(model.children())
-            pretrained_layers = pretrained_layers[:-1]
-        
-            model = nn.Sequential(*pretrained_layers)
-
-            # Freeze Weights
-            print(f"Freeze Pretraining = {frozen_pretraining}")
-            if (frozen_pretraining):
-                for param in model.parameters():
-                    param.requires_grad = False  # Freeze all layers
-            self.pretraining_params.append(model)
-            self.pretrained_models.append(model)
-
-        print(f"len(pretrained_models) = {len(self.pretrained_models)}")
-        print(f"input size = {input_size}")
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        
-        self.final_mlp = Make_MLP(input_size + hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.classify = nn.Linear(hid_size, out_size)
-
-        self.model_params.append(self.final_mlp)
-        self.model_params.append(self.classify)
-
-    def TL_node_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][1]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][1]:
-                x = layer(x)
-            return x
-    
-    def TL_edge_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][2]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][2]:
-                x = layer(x)
-            return x
-    
-    def TL_global_encoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][3]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][3]:
-                x = layer(x)
-            return x
-    
-    def TL_node_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][4]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][4]:
-                x = layer(x)
-            return x
-
-    def TL_edge_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][5]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][5]:
-                x = layer(x)
-            return x
-
-    def TL_global_update(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][6]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][6]:
-                x = layer(x)
-            return x
-    
-    def TL_global_decoder(self, x, model_idx):
-        try:
-            for layer in self.pretrained_models[model_idx][7]:
-                x = layer(x)
-            return x
-        except (NotImplementedError, IndexError):
-            for layer in self.pretrained_models[model_idx][1][7]:
-                x = layer(x)
-            return x
-
-    def Pretrained_Output(self, g, model_idx):
-        h = self.TL_node_encoder(g.ndata['features'], model_idx)
-        e = self.TL_edge_encoder(g.edata['features'], model_idx)
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.TL_global_encoder(global_feats, model_idx)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1), model_idx)
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1), model_idx)
-            h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1), model_idx)
-        # h_global = self.TL_global_decoder(h_global, model_idx)
-        return h_global
-
-    def forward(self, g, global_feats):
-        h = self.node_encoder(g.ndata['features'])
-        e = self.edge_encoder(g.edata['features'])
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = [h_global]
-        for i in range(len(self.pretrained_models)):
-            h_global.append(self.Pretrained_Output(g.clone(), i))
-        h_global = torch.concatenate(h_global, dim=1)
-        return self.classify(self.final_mlp(h_global))
-    
-    def to(self, device):
-        for i in range(len(self.pretrained_models)):
-            self.pretrained_models[i].to(device)
-        self.classify.to(device)
-        self.final_mlp.to(device)
-        self.node_encoder.to(device)
-        self.edge_encoder.to(device)
-        self.global_encoder.to(device)
-
-        self.node_update.to(device)
-        self.edge_update.to(device)
-        self.global_update.to(device)
-        return self
-    
-    def parameters(self, recurse: bool = True):
-        params = []
-        for i, model_section in enumerate(self.pretraining_params):
-            if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): 
-                print(f"Pretraining LR = {self.learning_rate['pretraining_lr'][i]}")
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"][i]})
-            else:
-                print(f"Pretraining LR = 0.00001")
-                params.append({'params': model_section.parameters(), 'lr': 0.00001})
-        for model_section in self.model_params:
-            if (type(self.learning_rate) == dict and self.learning_rate["model_lr"]): 
-                print(f"Model LR = {self.learning_rate['model_lr']}")
-                params.append({'params': model_section.parameters(), 'lr': self.learning_rate["model_lr"]})
-            else:
-                print(f"Model LR = 0.0001")
-                params.append({'params': model_section.parameters(), 'lr': 0.0001})
-        return params    
-
-
-class Clustering(nn.Module):
-    def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs):
-        super().__init__()
-        print(f'Unused args while creating GCN: {kwargs}')
-        self.n_layers = n_layers
-        self.n_proc_steps = n_proc_steps
-        self.layers = nn.ModuleList()
-        self.hid_size = hid_size
-        if (len(sample_global) == 0):
-            self.has_global = False
-        else:
-            self.has_global = sample_global.shape[1] != 0
-        gl_size = sample_global.shape[1] if self.has_global else 1
-
-        #encoder
-        self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #GNN
-        self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-        self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout)
-
-        #decoder
-        self.global_decoder = Make_MLP(hid_size, hid_size, out_size, n_layers, dropout=dropout)
-
-    def model_forward(self, g, global_feats, features = 'features'):
-        h = self.node_encoder(g.ndata[features])
-        e = self.edge_encoder(g.edata[features])
-
-        g.ndata['h'] = h
-        g.edata['e'] = e
-        if not self.has_global:
-            global_feats = g.batch_num_nodes()[:, None].to(torch.float)
-
-        batch_num_nodes = None
-        sum_weights = None
-        if "w" in g.ndata:
-            batch_indices = g.batch_num_nodes()
-            # Find non-zero rows (non-padded nodes)
-            non_padded_nodes_mask = torch.any(g.ndata[features] != 0, dim=1)
-            # Split the mask according to the batch indices
-            batch_num_nodes = []
-            start_idx = 0
-            for num_nodes in batch_indices:
-                end_idx = start_idx + num_nodes
-                non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item()
-                batch_num_nodes.append(non_padded_count)
-                start_idx = end_idx
-            batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata[features].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
-            global_feats = batch_num_nodes[:, None].to(torch.float)
-
-        h_global = self.global_encoder(global_feats)
-        for i in range(self.n_proc_steps):
-            g.apply_edges(dgl.function.copy_u('h', 'm_u'))
-            g.apply_edges(copy_v)
-            g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1))
-            g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
-            g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1))
-            if "w" in g.ndata:
-                mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights
-                h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1))
-            else:
-                h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1))
-        h_global = self.global_decoder(h_global)
-        return h_global
-    
-    def forward(self, g, global_feats):
-        h_global = self.model_forward(g, global_feats, 'features')
-        h_global_augmented = self.model_forward(g, global_feats, 'augmented_features')
-        return torch.cat((h_global, h_global_augmented), dim=1)
-
-    def representation(self, g, global_feats):
-        h_global = self.model_forward(g, global_feats, 'features')
-        h_global_augmented = self.model_forward(g, global_feats, 'augmented_features')
-        return h_global, h_global_augmented, torch.cat((h_global, h_global_augmented), dim=1)
-    
-    def __str__(self):
-        layer_names = ["node_encoder", "edge_encoder", "global_encoder", 
-                  "node_update", "edge_update", "global_update", "global_decoder"]
-        
-        layers = [self.node_encoder, self.edge_encoder, self.global_encoder,
-                  self.node_update, self.edge_update, self.global_update, self.global_decoder]
-        
-        for i in range(len(layers)):
-            print(layer_names[i])
-            for layer in layers[i].children():
-                if isinstance(layer, nn.Linear):
-                    print(layer.state_dict())
-        
-        print("classify")
-        print(self.classify.weight)
-        return ""

nemo/models/loss.py

@@ -1,311 +0,0 @@
-from torch import nn
-import torch
-from root_gnn_base import utils
-import numpy as np
-
-class MaskedLoss():
-    def __init__(self, mask = []):
-        self.mask = mask
-
-    def make_mask(self, targets):
-        mask = torch.ones_like(targets[:,0])
-        for m in self.mask:
-            if m['op'] == 'eq':
-                mask[targets[:,m['idx']] == m['val']] = 0
-            elif m['op'] == 'gt':
-                mask[targets[:,m['idx']] > m['val']] = 0
-            elif m['op'] == 'lt':
-                mask[targets[:,m['idx']] < m['val']] = 0
-            elif m['op'] == 'ge':
-                mask[targets[:,m['idx']] >= m['val']] = 0
-            elif m['op'] == 'le':
-                mask[targets[:,m['idx']] <= m['val']] = 0
-            elif m['op'] == 'ne':
-                mask[targets[:,m['idx']] != m['val']] = 0
-            else:
-                raise ValueError(f'Unknown mask op {m["op"]}')
-        return mask == 1
-
-class MaskedL1Loss(MaskedLoss):
-    def __init__(self, mask = [], index = 0):
-        super().__init__(mask)
-        self.index = index
-        self.loss = nn.L1Loss()
-
-    def __call__(self, logits, targets):
-        mask = self.make_mask(targets)
-        return self.loss(logits[mask], targets[mask][:,self.index])
-
-class BCEWithLogitsLoss():
-    def __init__(self, weight=None, reduction='mean'):
-        self.loss = nn.BCEWithLogitsLoss(weight=weight, reduction=reduction)
-    
-    def __call__(self, logits, targets):
-        return self.loss(logits[:,0], targets.float())
-
-class MultiScore():
-    def __init__(self, scores):
-        self. score_fcns = []
-        self.start_idx = []
-        self.end_idx = []
-        for score in scores:
-            self.score_fcns.append(utils.buildFromConfig(score))
-            self.start_idx.append(score['start_idx'])
-            self.end_idx.append(score['end_idx'])
-    
-    def __call__(self, last_layer):
-        scores = []
-        for i in range(len(self.score_fcns)):
-            scores.append(self.score_fcns[i](last_layer[:, self.start_idx[i]:self.end_idx[i]]))
-        return torch.cat(scores, dim=1)
-
-class MultiLoss():
-    def __init__(self, losses):
-        self.loss_fcns = []
-        self.label_start_idx = []
-        self.label_end_idx = []
-        self.output_start_idx = []
-        self.output_end_idx = []
-        self.weights = []
-        self.label_types = []
-        for loss in losses:
-            self.loss_fcns.append(utils.buildFromConfig(loss))
-            self.label_start_idx.append(loss['label_start_idx'])
-            self.label_end_idx.append(loss['label_end_idx'])
-            self.output_start_idx.append(loss['output_start_idx'])
-            self.output_end_idx.append(loss['output_end_idx'])
-            self.weights.append(loss.get('weight', 1.0))
-            self.label_types.append(loss.get('label_type', 'float'))
-
-    def __call__(self, logits, targets):
-        loss = 0
-        # print(logits.shape, targets.shape)
-        for i in range(len(self.loss_fcns)):
-            if self.label_types[i] == 'int':
-                # print('loss', i, self.label_start_idx[i], self.label_end_idx[i], self.output_start_idx[i], self.output_end_idx[i])
-                # print(logits[:, self.output_start_idx[i]:self.output_end_idx[i]].shape, targets[:, self.label_start_idx[i]].shape)
-                loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]].to(int))
-            elif self.label_end_idx[i] - self.label_start_idx[i] == 1:
-                loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]])
-            else:
-                # print('loos', i, self.label_start_idx[i], self.label_end_idx[i], self.output_start_idx[i], self.output_end_idx[i])
-                # print(logits[:, self.output_start_idx[i]:self.output_end_idx[i]].shape, targets[:, self.label_start_idx[i]:self.label_end_idx[i]].shape)
-                loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]:self.label_end_idx[i]])
-        return loss
-    
-class AdvLoss():
-    def __init__(self, loss, adv_loss, adv_weight=1.0):
-        self.loss_fcn = utils.buildFromConfig(loss)
-        self.adv_loss_fcn = utils.buildFromConfig(adv_loss)
-        self.adv_weight = adv_weight
-
-    def __call__(self, logits, targets):
-        mask = targets[:,0] == 0
-        loss = self.loss_fcn(logits[:,0], targets[:,0])
-        adv_loss = self.adv_loss_fcn(logits[mask][:,1], targets[mask])
-        return loss - self.adv_weight * adv_loss
-
-class MassWindowAdvLoss(AdvLoss):
-    def __call__(self, logits, targets):
-        mask = (targets[:,0] == 0) & (targets[:,1] > 5) & (targets[:,1] < 25)
-        print(mask, mask.shape, mask.sum())
-        loss = self.loss_fcn(logits[:,0], targets[:,0])
-        print(loss)
-        adv_loss = self.adv_loss_fcn(logits[mask][:,1], targets[mask][:,1])
-        print(adv_loss)
-        return loss - self.adv_weight * adv_loss
-
-class KDELoss(MaskedLoss):
-    def __init__(self, mask = [], index = 0):
-        self.index = index
-        super().__init__(mask)
-
-    def __call__(self, logits, targets):
-        mask = self.make_mask(targets)
-        logits = logits[mask]
-        targets = targets[mask][:,self.index]
-        N = logits.shape[0]
-        masses = targets / torch.sqrt(torch.mean(targets**2))
-        scores = logits[:,0]  / torch.sqrt(torch.mean(logits**2))
-
-        factor_2d = (1.0*N) ** (-2/6)
-        covs = (factor_2d * torch.var(masses), factor_2d * torch.var(scores))
-
-        m_diffs = torch.unsqueeze(masses, 1) - torch.unsqueeze(masses, 0)
-        s_diffs = torch.unsqueeze(scores, 1) - torch.unsqueeze(scores, 0)
-
-        ymm = torch.exp(- (m_diffs**2) / (4 * covs[0]))
-        yss = torch.exp(- (s_diffs**2) / (4 * covs[1]))
-
-        integral_rho_2d_rho_2d = torch.einsum('ij,ij->', ymm, yss)
-        integral_rho_1d_rho_1d = torch.einsum('ij,kl->', ymm, yss)
-        integral_rho_2d_rho_1d = torch.einsum('ij,ik->', ymm, yss)
-        raw_integral = integral_rho_2d_rho_2d - 2 * integral_rho_2d_rho_1d / N + integral_rho_1d_rho_1d / N**2
-        return raw_integral / (4 * torch.pi * N**2)
-
-class MultiLabelLoss():
-    def __init__(self, label_names, label_types, label_weights = None):
-        self.loss_fcn = []
-        if (label_weights):
-            self.weights = torch.tensor(label_weights)
-        else:
-            self.weights = torch.ones(len(label_types))
-        for type in label_types:
-            if (type == "r"):
-                self.loss_fcn.append(torch.nn.MSELoss(reduce=False))
-            elif (type == "c"):
-                self.loss_fcn.append(torch.nn.BCEWithLogitsLoss())
-        print(f"self.weights = {self.weights}")
-
-    def __call__(self, logits, targets):
-        targets = targets.float()
-        loss = torch.zeros(len(logits[:, 0]), device = logits.get_device())
-        for i in range(len(self.loss_fcn)):
-            loss += self.weights[i] * self.loss_fcn[i](logits[:, i], targets[:, i])
-        return torch.mean(loss)
-        
-    
-class MultiLabelFinish():
-    def __init__(self, label_names, label_types):
-        self.finish_fcn = []
-        for type in label_types:
-            if (type == "r"):
-                self.finish_fcn.append(None)
-            elif (type == "c"):
-                self.finish_fcn.append(torch.special.expit)
-
-    def __call__(self, logits):
-        for i in range(len(self.finish_fcn)):
-            if (self.finish_fcn[i]):
-                logits[:, i] = self.finish_fcn[i](logits[:, i].to(torch.long))
-        return logits
-
-class ContrastiveClusterLoss():
-    def __init__(self, k=10, temperature=1, alpha=1):
-        self.k = k
-        self.temperature = temperature
-        self.alpha = alpha
-
-    def __call__(self, logits, targets):
-        targets = targets.float()
-        logits_combined = logits.float()
-
-        hid_size = int(len(logits[0]) / 2)
-
-        logits = normalize_embeddings(logits_combined[:, :hid_size])
-        logits_augmented = normalize_embeddings(logits_combined[:, hid_size:])
-
-        contrastive = contrastive_loss(logits, logits_augmented, self.temperature)
-        clustering, _ = clustering_loss(logits, self.k)
-
-        variance_loss = variance_regularization(logits) + variance_regularization(logits_augmented)
-
-        return torch.mean(contrastive + clustering + self.alpha * variance_loss)
-    
-class ContrastiveClusterFinish():
-    def __init__(self, k = 10, temperature = 1, max_cluster_iterations = 10):
-        self.k = k
-        self.temperature = temperature
-        self.max_cluster_iterations = max_cluster_iterations
-
-        print(f"ContrastiveClusterFinish: k = {k}, temperature = {temperature}")
-
-    def __call__(self, logits):
-        logits_combined = logits.float()
-
-        hid_size = int(len(logits[0]) / 2)
-
-        logits = logits_combined[:, :hid_size]
-        logits_augmented = logits_combined[:, hid_size:]
-        
-        contrastive = contrastive_loss(logits, logits_augmented, self.temperature)
-        clustering, _ = clustering_loss(logits, self.k, self.max_cluster_iterations)
-        variance = variance_regularization(logits) + variance_regularization(logits_augmented)
-
-        return contrastive, clustering, variance
-    
-def s(z_i, z_j):
-    z_i = torch.tensor(z_i) if not isinstance(z_i, torch.Tensor) else z_i
-    z_j = torch.tensor(z_j) if not isinstance(z_j, torch.Tensor) else z_j
-    
-    return torch.cdist(z_i, z_j, p=2)
-    # dot_product = torch.dot(z_i, z_j)
-    # norm_i = torch.linalg.norm(z_i)
-    # norm_j = torch.linalg.norm(z_j)
-    
-    # return dot_product / (norm_i * norm_j)
-
-def contrastive_loss(logits, logits_augmented, temperature=1, margin=1.0):
-    logits = torch.tensor(logits) if not isinstance(logits, torch.Tensor) else logits
-    logits_augmented = torch.tensor(logits_augmented) if not isinstance(logits_augmented, torch.Tensor) else logits_augmented
-
-    z = torch.cat((logits, logits_augmented), dim=0)
-    similarity_matrix = torch.mm(z, z.t()) / temperature
-    norms = torch.linalg.norm(z, dim=1)
-    norm_matrix = torch.ger(norms, norms)
-    similarity_matrix = similarity_matrix / norm_matrix
-    mask = torch.eye(similarity_matrix.size(0), dtype=torch.bool)
-
-    loss = 0
-    for k in range(len(logits)):
-        numerator = torch.exp(similarity_matrix[k, k + len(logits)])
-        denominator = torch.sum(torch.exp(similarity_matrix[k, ~mask[k]]))
-        
-        loss += -torch.log(numerator / denominator)
-
-    return loss
-
-
-def clustering_loss(logits, k=10, max_iterations=10):
-    # Step 1: Initialize cluster means
-    indices = torch.randperm(logits.size(0))[:k]
-    cluster_means = logits[indices]
-
-    prev_assignments = None
-    assignment_history = []
-    iteration = 0
-
-    while iteration < max_iterations:
-        iteration += 1
-
-        # Step 2: Assign each data point to the nearest cluster mean
-        distances = torch.cdist(logits, cluster_means, p=2)  # Compute distances between logits and cluster means
-        cluster_assignments = torch.argmin(distances, dim=1)  # Assign each point to the nearest cluster mean
-
-        # Check for convergence: if assignments do not change, break the loop
-        if prev_assignments is not None and torch.equal(cluster_assignments, prev_assignments):
-            break
-
-        # Check for cycles: if assignments have been seen before, break the loop
-        if any(torch.equal(cluster_assignments, prev) for prev in assignment_history):
-            break
-
-        assignment_history.append(cluster_assignments.clone())
-        prev_assignments = cluster_assignments.clone()
-
-        # Step 3: Update cluster means based on assignments
-        new_cluster_means = torch.zeros_like(cluster_means)
-        for i in range(k):
-            assigned_points = logits[cluster_assignments == i]
-            if assigned_points.size(0) > 0:
-                new_cluster_means[i] = assigned_points.mean(dim=0)
-            else:
-                # If no points are assigned to the cluster, reinitialize the mean randomly
-                new_cluster_means[i] = logits[torch.randint(0, logits.size(0), (1,)).item()]
-        cluster_means = new_cluster_means
-
-    # Step 4: Compute the clustering loss
-    distances = torch.cdist(logits, cluster_means, p=2)
-    min_distances = torch.min(distances, dim=1)[0]
-    loss = torch.sum(min_distances ** 2)
-
-    return loss, cluster_means
-
-def normalize_embeddings(embeddings):
-    return embeddings / embeddings.norm(dim=1, keepdim=True)
-
-def variance_regularization(embeddings):
-    mean_embedding = embeddings.mean(dim=0)
-    variance = ((embeddings - mean_embedding) ** 2).mean()
-    return variance
-

nemo/models/meshgraphnet.py

@@ -1,33 +0,0 @@
-import torch
-import torch.nn as nn
-import dgl
-
-# Import the PhysicsNemo MeshGraphNet model
-from physicsnemo.models.meshgraphnet import MeshGraphNet as PhysicsNemoMeshGraphNet
-
-class MeshGraphNet(nn.Module):
-    def __init__(self, *args, out_dim=1, **kwargs):
-        super().__init__()
-        # Initialize the PhysicsNemo MeshGraphNet
-        self.base_gnn = PhysicsNemoMeshGraphNet(*args, **kwargs)
-        # Assume node_output_dim is known or infer from args/kwargs
-        node_output_dim = 64
-        self.mlp = nn.Linear(node_output_dim, out_dim)
-
-    def forward(self, node_feats, edge_feats, batched_graph):
-        """
-        Args:
-            batched_graph: DGLGraph, batched graphs
-            node_feats: Tensor [total_num_nodes, node_feat_dim]
-            edge_feats: Tensor [total_num_edges, edge_feat_dim]
-        Returns:
-            graph_pred: Tensor [num_graphs, out_dim]
-        """
-        # 1. Node-level prediction from PhysicsNemo GNN
-        node_pred = self.base_gnn(node_feats, edge_feats, batched_graph)
-        batched_graph.ndata['h'] = node_pred
-        graph_feat = dgl.readout_nodes(batched_graph, 'h', op='mean')  # [num_graphs, node_output_dim]
-
-        # 3. Final MLP for graph-level prediction
-        graph_pred = self.mlp(graph_feat)  # [num_graphs, out_dim]
-        return graph_pred

nemo/root_gnn_base/batched_dataset.py

@@ -1,191 +0,0 @@
-from dgl.dataloading import GraphDataLoader
-from torch.utils.data.sampler import SubsetRandomSampler
-from torch.utils.data.sampler import SequentialSampler
-from dgl.data import DGLDataset
-import torch
-import time
-import os
-import dgl
-from root_gnn_base import utils
-
-def GetBatchedLoader(dataset, batch_size, mask_fn = None, drop_last=True, **kwargs):
-    if mask_fn == None:
-        mask_fn = lambda x: torch.ones(len(x), dtype=torch.bool)
-    dloader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(torch.arange(len(dataset))[mask_fn(dataset)]), batch_size=batch_size, drop_last=drop_last, num_workers = 0)
-    return dloader
-
-#Dataset which contains prebatched shuffled graphs. Cannot be saved to disk, else batching info is lost.
-class PreBatchedDataset(DGLDataset):
-    def __init__(self, start_dataset, batch_size, mask_fn = None, drop_last=True, save_to_disk = True, suffix = '', chunks = 1, chunkno = -1, shuffle = True, padding_mode = 'NONE', hidden_size=64, **kwargs):
-        print(f'Unused kwargs: {kwargs}')
-        self.start_dataset = start_dataset
-        self.start_dataset.load()
-        
-        self.batch_size = batch_size
-        self.chunks = chunks
-        self.chunkno = chunkno
-        self.mask_fn = mask_fn
-        self.drop_last = drop_last
-        self.graphs = []
-        self.label = []
-        self.padding_mode = padding_mode
-        self.save_to_disk = save_to_disk
-        self.shuffle = shuffle
-        self.suffix = suffix
-        self.current_chunk = None
-        self.current_chunk_idx = -1
-        self.hid_size = hidden_size
-        super().__init__(name = start_dataset.name + '_prebatched_padded', save_dir=start_dataset.save_dir)
-
-    def process(self):
-        first = 0
-        last = len(self.start_dataset)
-        if self.chunks > 1 and self.chunkno >= 0:
-            first = int(self.chunkno / self.chunks * len(self.start_dataset))
-            last = int((self.chunkno + 1) / self.chunks * len(self.start_dataset))
-        print(f'Processing chunk {self.chunkno} of {self.chunks} from {first} to {last} of {len(self.start_dataset)}')
-        mask = torch.logical_and(torch.logical_and(self.mask_fn(self.start_dataset), torch.arange(len(self.start_dataset)) >= first), torch.arange(len(self.start_dataset)) < last)
-        if self.shuffle:
-            dloader = GraphDataLoader(self.start_dataset, sampler=SubsetRandomSampler(torch.arange(len(self.start_dataset))[mask]), batch_size=self.batch_size, drop_last=self.drop_last)
-        else: #Only don't shuffle if we're doing inference. Then we want all of the events anyways?
-            dloader = GraphDataLoader(self.start_dataset, sampler=SequentialSampler(self.start_dataset), batch_size=self.batch_size, drop_last=self.drop_last)
-        self.graphs = []
-        self.labels = []
-        self.tracking = []
-        self.globals = []
-        self.batch_num_nodes = []
-        self.batch_num_edges = []
-        max_edges = 0
-        max_nodes = 0
-        load_batch_start = time.time()
-        for batch, label, tracking, global_feat in dloader:
-            if batch.num_edges() > max_edges:
-                max_edges = batch.num_edges()
-            if batch.num_nodes() > max_nodes:
-                max_nodes = batch.num_nodes()
-            self.graphs.append(batch)
-            self.labels.append(label)
-            self.tracking.append(tracking)
-            self.globals.append(global_feat)
-        load_batch_end = time.time()
-        print(f'Loaded {len(self.graphs)} batches in {load_batch_end - load_batch_start} seconds')
-        if self.padding_mode == 'STEPS':
-            pad_node, pad_edge = utils.pad_size(self.batch_size, max_edges, max_nodes)
-        elif self.padding_mode == 'FIXED':
-            print('Padding to fixed size. This is currently hardcoded.')
-            pad_node = 16000
-            pad_edge = 104000
-        elif self.padding_mode == 'NONE':
-            pad_node = 0
-            pad_edge = 0
-        else:
-            pad_node = 0
-            pad_edge = 0
-        print(f'Max edges: {max_edges}, Max nodes: {max_nodes}, Padding to {pad_edge} edges and {pad_node} nodes')
-        pad_start = time.time()
-        if self.padding_mode == 'NODE':
-            for i in range(len(self.graphs)):
-                unbatched_g = dgl.unbatch(self.graphs[i])
-                max_num_nodes = max(g.number_of_nodes() for g in unbatched_g)
-                self.graphs[i] = utils.pad_batch_num_nodes(self.graphs[i], max_num_nodes, hid_size=self.hid_size)
-                self.batch_num_nodes.append(self.graphs[i].batch_num_nodes())
-                self.batch_num_edges.append(self.graphs[i].batch_num_edges())
-        else:
-            for i in range(len(self.graphs)):
-                self.graphs[i] = utils.pad_batch(self.graphs[i], pad_edge, pad_node)
-                self.batch_num_nodes.append(self.graphs[i].batch_num_nodes())
-                self.batch_num_edges.append(self.graphs[i].batch_num_edges())
-        pad_end = time.time()
-        print(f'Padded {len(self.graphs)} batches in {pad_end - pad_start} seconds')
-
-    def save(self):
-        if not self.save_to_disk:
-            return
-        graph_path = os.path.join(self.save_dir, f'{self.name}_{self.chunkno}_{self.suffix}.bin')
-        print(f'Saving dataset to {graph_path}')
-        if len(self.graphs) == 0:
-            return
-        dgl.save_graphs(str(graph_path), self.graphs, {'labels': torch.stack(self.labels), 'batch_num_nodes': torch.stack(self.batch_num_nodes), 'batch_num_edges': torch.stack(self.batch_num_edges), 'tracking': torch.stack(self.tracking), 'globals': torch.stack(self.globals)})
-
-    def has_cache(self):
-        if not self.save_to_disk:
-            return False
-        for ch in range(self.chunks):
-            graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin')
-            if not os.path.exists(graph_path):
-                print(f'Cache file {graph_path} does not exist, not loading from cache.')
-                return False
-        return True
-    
-    def load(self):
-        if not self.save_to_disk:
-            return
-        self.graphs = []
-        label_chunks = []
-        tracking_chunks = []
-        global_chunks = []
-        for ch in range(self.chunks):
-            graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin')
-            print(f'Loading dataset from {graph_path}')
-            graphs, label_dict = dgl.load_graphs(graph_path)
-            label_chunks.append(label_dict['labels'])
-            tracking_chunks.append(label_dict['tracking'])
-            global_chunks.append(label_dict['globals'])
-            for g, bnn, bne in zip(graphs, label_dict['batch_num_nodes'], label_dict['batch_num_edges']):
-                g.set_batch_num_nodes(bnn)
-                g.set_batch_num_edges(bne)
-            self.graphs.extend(graphs)
-        self.labels = torch.cat(label_chunks)
-        self.tracking = torch.cat(tracking_chunks)
-        self.globals = torch.cat(global_chunks)
-
-    def __getitem__(self, idx):
-        return self.graphs[idx], self.labels[idx], self.tracking[idx], self.globals[idx]
-    
-    def __len__(self):
-        return len(self.graphs)
-
-#Dataset which contains prebatched shuffled graphs. Cannot be saved to disk, else batching info is lost.
-class LazyPreBatchedDataset(PreBatchedDataset):
-    def __init__(self, **kwargs):
-        # print(f'Unused kwargs: {kwargs}')
-        self.current_chunk = None
-        self.current_chunk_idx = -10
-        self.label_chunks = []
-        super().__init__(**kwargs)
-    
-    def load(self):
-        if not self.save_to_disk:
-            return
-        self.label_chunks = []
-        for ch in range(self.chunks):
-            graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin')
-            print(f'Loading dataset from {graph_path}')
-            label_dict = dgl.data.graph_serialize.load_labels_v2(graph_path)
-            self.label_chunks.append(label_dict)            
-
-    def __getitem__(self, idx):
-        chunk_idx = -1
-        sum = 0
-        ev_idx = -999
-        for i in range(len(self.label_chunks)):
-            count = len(self.label_chunks[i]['labels'])
-            if idx < sum + count: 
-                chunk_idx = i
-                ev_idx = idx - sum
-                break
-            sum += count
-        if chunk_idx != self.current_chunk_idx:
-            # print(f"rank {self.rank} getting data from {self.name}_{chunk_idx}_{self.suffix}.bin")
-            self.current_chunk, _ = dgl.load_graphs(os.path.join(self.save_dir, f'{self.name}_{chunk_idx}_{self.suffix}.bin'))
-            self.current_chunk_idx = chunk_idx
-        g = self.current_chunk[ev_idx]
-        g.set_batch_num_nodes(self.label_chunks[chunk_idx]['batch_num_nodes'][ev_idx])
-        g.set_batch_num_edges(self.label_chunks[chunk_idx]['batch_num_edges'][ev_idx])
-        return g, self.label_chunks[chunk_idx]['labels'][ev_idx], self.label_chunks[chunk_idx]['tracking'][ev_idx], self.label_chunks[chunk_idx]['globals'][ev_idx]
-    
-    def __len__(self):
-        l = 0
-        for chunk in self.label_chunks:
-            l += len(chunk['labels'])
-        return l

nemo/root_gnn_base/custom_scheduler.py

@@ -1,565 +0,0 @@
-import types
-import math
-import torch
-from torch import inf
-from functools import wraps, partial
-import warnings
-import weakref
-from collections import Counter
-from bisect import bisect_right
-
-from models import GCN
-
-
-
-
-### Code from: https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#ReduceLROnPlateau
-
-Optimizer = torch.optim.Optimizer
-
-__all__ = ['LambdaLR', 'MultiplicativeLR', 'StepLR', 'MultiStepLR', 'ConstantLR', 'LinearLR',
-           'ExponentialLR', 'SequentialLR', 'CosineAnnealingLR', 'ChainedScheduler', 'ReduceLROnPlateau',
-           'CyclicLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'PolynomialLR', 'LRScheduler']
-
-EPOCH_DEPRECATION_WARNING = (
-    "The epoch parameter in `scheduler.step()` was not necessary and is being "
-    "deprecated where possible. Please use `scheduler.step()` to step the "
-    "scheduler. During the deprecation, if epoch is different from None, the "
-    "closed form is used instead of the new chainable form, where available. "
-    "Please open an issue if you are unable to replicate your use case: "
-    "https://github.com/pytorch/pytorch/issues/new/choose."
-)
-
-
-def update_LR(opt, lr):
-    for param_group in opt.param_groups:
-        param_group['lr'] = lr
-
-def print_LR(opt):
-    for param_group in opt.param_groups:
-        print(f"LR = {param_group['lr']}")
-
-def _check_verbose_deprecated_warning(verbose):
-    """Raises a warning when verbose is not the default value."""
-    if verbose != "deprecated":
-        warnings.warn("The verbose parameter is deprecated. Please use get_last_lr() "
-                      "to access the learning rate.", UserWarning)
-        return verbose
-    return False
-
-class LRScheduler:
-
-    def __init__(self, optimizer, last_epoch=-1, verbose="deprecated"):
-
-        # Attach optimizer
-        if not isinstance(optimizer, Optimizer):
-            raise TypeError(f'{type(optimizer).__name__} is not an Optimizer')
-        self.optimizer = optimizer
-
-        # Initialize epoch and base learning rates
-        if last_epoch == -1:
-            for group in optimizer.param_groups:
-                group.setdefault('initial_lr', group['lr'])
-        else:
-            for i, group in enumerate(optimizer.param_groups):
-                if 'initial_lr' not in group:
-                    raise KeyError("param 'initial_lr' is not specified "
-                                   f"in param_groups[{i}] when resuming an optimizer")
-        self.base_lrs = [group['initial_lr'] for group in optimizer.param_groups]
-        self.last_epoch = last_epoch
-
-        # Following https://github.com/pytorch/pytorch/issues/20124
-        # We would like to ensure that `lr_scheduler.step()` is called after
-        # `optimizer.step()`
-        def with_counter(method):
-            if getattr(method, '_with_counter', False):
-                # `optimizer.step()` has already been replaced, return.
-                return method
-
-            # Keep a weak reference to the optimizer instance to prevent
-            # cyclic references.
-            instance_ref = weakref.ref(method.__self__)
-            # Get the unbound method for the same purpose.
-            func = method.__func__
-            cls = instance_ref().__class__
-            del method
-
-            @wraps(func)
-            def wrapper(*args, **kwargs):
-                instance = instance_ref()
-                instance._step_count += 1
-                wrapped = func.__get__(instance, cls)
-                return wrapped(*args, **kwargs)
-
-            # Note that the returned function here is no longer a bound method,
-            # so attributes like `__func__` and `__self__` no longer exist.
-            wrapper._with_counter = True
-            return wrapper
-
-        self.optimizer.step = with_counter(self.optimizer.step)
-        self.verbose = _check_verbose_deprecated_warning(verbose)
-
-        self._initial_step()
-
-    def _initial_step(self):
-        """Initialize step counts and performs a step"""
-        self.optimizer._step_count = 0
-        self._step_count = 0
-        self.step()
-
-    def state_dict(self):
-        """Returns the state of the scheduler as a :class:`dict`.
-
-        It contains an entry for every variable in self.__dict__ which
-        is not the optimizer.
-        """
-        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
-
-    def load_state_dict(self, state_dict):
-        """Loads the schedulers state.
-
-        Args:
-            state_dict (dict): scheduler state. Should be an object returned
-                from a call to :meth:`state_dict`.
-        """
-        self.__dict__.update(state_dict)
-
-    def get_last_lr(self):
-        """ Return last computed learning rate by current scheduler.
-        """
-        return self._last_lr
-
-    def get_lr(self):
-        # Compute learning rate using chainable form of the scheduler
-        raise NotImplementedError
-
-    def print_lr(self, is_verbose, group, lr, epoch=None):
-        """Display the current learning rate.
-        """
-        if is_verbose:
-            if epoch is None:
-                print(f'Adjusting learning rate of group {group} to {lr:.4e}.')
-            else:
-                epoch_str = ("%.2f" if isinstance(epoch, float) else
-                             "%.5d") % epoch
-                print(f'Epoch {epoch_str}: adjusting learning rate of group {group} to {lr:.4e}.')
-
-
-    def step(self, epoch=None):
-        # Raise a warning if old pattern is detected
-        # https://github.com/pytorch/pytorch/issues/20124
-        if self._step_count == 1:
-            if not hasattr(self.optimizer.step, "_with_counter"):
-                warnings.warn("Seems like `optimizer.step()` has been overridden after learning rate scheduler "
-                              "initialization. Please, make sure to call `optimizer.step()` before "
-                              "`lr_scheduler.step()`. See more details at "
-                              "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning)
-
-            # Just check if there were two first lr_scheduler.step() calls before optimizer.step()
-            elif self.optimizer._step_count < 1:
-                warnings.warn("Detected call of `lr_scheduler.step()` before `optimizer.step()`. "
-                              "In PyTorch 1.1.0 and later, you should call them in the opposite order: "
-                              "`optimizer.step()` before `lr_scheduler.step()`.  Failure to do this "
-                              "will result in PyTorch skipping the first value of the learning rate schedule. "
-                              "See more details at "
-                              "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning)
-        self._step_count += 1
-
-        with _enable_get_lr_call(self):
-            if epoch is None:
-                self.last_epoch += 1
-                values = self.get_lr()
-            else:
-                warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
-                self.last_epoch = epoch
-                if hasattr(self, "_get_closed_form_lr"):
-                    values = self._get_closed_form_lr()
-                else:
-                    values = self.get_lr()
-
-        for i, data in enumerate(zip(self.optimizer.param_groups, values)):
-            param_group, lr = data
-            param_group['lr'] = lr
-
-        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
-
-
-# Including _LRScheduler for backwards compatibility
-# Subclass instead of assign because we want __name__ of _LRScheduler to be _LRScheduler (assigning would make it LRScheduler).
-class _LRScheduler(LRScheduler):
-    pass
-
-
-class _enable_get_lr_call:
-
-    def __init__(self, o):
-        self.o = o
-
-    def __enter__(self):
-        self.o._get_lr_called_within_step = True
-        return self
-
-    def __exit__(self, type, value, traceback):
-        self.o._get_lr_called_within_step = False
-
-
-class Dynamic_LR(LRScheduler):
-    """Reduce learning rate when a metric has stopped improving.
-    Models often benefit from reducing the learning rate by a factor
-    of 2-10 once learning stagnates. This scheduler reads a metrics
-    quantity and if no improvement is seen for a 'patience' number
-    of epochs, the learning rate is reduced.
-
-    Args:
-        optimizer (Optimizer): Wrapped optimizer.
-        mode (str): One of `min`, `max`. In `min` mode, lr will
-            be reduced when the quantity monitored has stopped
-            decreasing; in `max` mode it will be reduced when the
-            quantity monitored has stopped increasing. Default: 'min'.
-        factor (float): Factor by which the learning rate will be
-            reduced. new_lr = lr * factor. Default: 0.1.
-        patience (int): Number of epochs with no improvement after
-            which learning rate will be reduced. For example, if
-            `patience = 2`, then we will ignore the first 2 epochs
-            with no improvement, and will only decrease the LR after the
-            3rd epoch if the loss still hasn't improved then.
-            Default: 10.
-        threshold (float): Threshold for measuring the new optimum,
-            to only focus on significant changes. Default: 1e-4.
-        threshold_mode (str): One of `rel`, `abs`. In `rel` mode,
-            dynamic_threshold = best * ( 1 + threshold ) in 'max'
-            mode or best * ( 1 - threshold ) in `min` mode.
-            In `abs` mode, dynamic_threshold = best + threshold in
-            `max` mode or best - threshold in `min` mode. Default: 'rel'.
-        cooldown (int): Number of epochs to wait before resuming
-            normal operation after lr has been reduced. Default: 0.
-        min_lr (float or list): A scalar or a list of scalars. A
-            lower bound on the learning rate of all param groups
-            or each group respectively. Default: 0.
-        eps (float): Minimal decay applied to lr. If the difference
-            between new and old lr is smaller than eps, the update is
-            ignored. Default: 1e-8.
-        verbose (bool): If ``True``, prints a message to stdout for
-            each update. Default: ``False``.
-
-            .. deprecated:: 2.2
-                ``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
-                learning rate.
-
-    Example:
-        >>> # xdoctest: +SKIP
-        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
-        >>> scheduler = ReduceLROnPlateau(optimizer, 'min')
-        >>> for epoch in range(10):
-        >>>     train(...)
-        >>>     val_loss = validate(...)
-        >>>     # Note that step should be called after validate()
-        >>>     scheduler.step(val_loss)
-    """
-
-    def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10,
-                 plateau_var = "test_auc", 
-                 threshold=1e-4, threshold_mode='rel', cooldown=0,
-                 min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False):
-        
-        """
-        if factor >= 1.0:
-            raise ValueError('Factor should be < 1.0.')
-        """
-        self.factor = factor
-
-        # Attach optimizer
-        if not isinstance(optimizer, Optimizer):
-            raise TypeError(f'{type(optimizer).__name__} is not an Optimizer')
-        self.optimizer = optimizer
-
-        if isinstance(min_lr, (list, tuple)):
-            if len(min_lr) != len(optimizer.param_groups):
-                raise ValueError(f"expected {len(optimizer.param_groups)} min_lrs, got {len(min_lr)}")
-            self.min_lrs = list(min_lr)
-            self.max_lrs = list(max_lr)
-        else:
-            self.min_lrs = [min_lr] * len(optimizer.param_groups)
-            self.max_lrs = [max_lr] * len(optimizer.param_groups)
-
-        self.patience = patience
-        self.plateau_var = plateau_var
-
-        self.verbose = verbose
-        self.cooldown = cooldown
-        self.cooldown_counter = 0
-        self.mode = mode
-        self.threshold = threshold
-        self.threshold_mode = threshold_mode
-        self.best = None
-        self.num_bad_epochs = None
-        self.mode_worse = None  # the worse value for the chosen mode
-        self.eps = eps
-        self.last_epoch = 0
-        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
-        self._init_is_better(mode=mode, threshold=threshold,
-                             threshold_mode=threshold_mode)
-        self._reset()
-
-    def _reset(self):
-        """Resets num_bad_epochs counter and cooldown counter."""
-        self.best = self.mode_worse
-        self.cooldown_counter = 0
-        self.num_bad_epochs = 0
-
-    def step(self, model, metrics, epoch=None):
-        # convert `metrics` to float, in case it's a zero-dim Tensor
-        current = float(metrics[self.plateau_var])
-        if epoch is None:
-            epoch = self.last_epoch + 1
-        else:
-            warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
-        self.last_epoch = epoch
-
-        if self.is_better(current, self.best):
-            if(self.verbose):
-                print("Model is improving!")
-            self.best = current
-            self.num_bad_epochs = 0
-        else:
-            if(self.verbose):
-                print(f"Model is not improving :( best = {self.best}, current = {current}")
-            self.num_bad_epochs += 1
-
-        if self.in_cooldown:
-            self.cooldown_counter -= 1
-            self.num_bad_epochs = 0  # ignore any bad epochs in cooldown
-
-        if self.num_bad_epochs > self.patience:
-            self._reduce_lr(epoch)
-            self.cooldown_counter = self.cooldown
-            self.num_bad_epochs = 0
-
-        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
-
-    def _reduce_lr(self, epoch):
-        print("Adjusting Learning Rate")
-        self._reset()
-        for i, param_group in enumerate(self.optimizer.param_groups):
-            old_lr = float(param_group['lr'])
-            new_lr = max(old_lr * self.factor, self.min_lrs[i])
-            new_lr = min(new_lr, self.max_lrs[i])
-            if abs(old_lr - new_lr) > self.eps:
-                param_group['lr'] = new_lr
-
-    def get_last_lr(self):
-        return self._last_lr
-    @property
-    def in_cooldown(self):
-        return self.cooldown_counter > 0
-
-    def is_better(self, a, best):
-        if self.mode == 'min' and self.threshold_mode == 'rel':
-            rel_epsilon = 1. - self.threshold
-            return a < best * rel_epsilon
-
-        elif self.mode == 'min' and self.threshold_mode == 'abs':
-            return a < best - self.threshold
-
-        elif self.mode == 'max' and self.threshold_mode == 'rel':
-            rel_epsilon = self.threshold + 1.
-            return a > best * rel_epsilon
-
-        else:  # mode == 'max' and epsilon_mode == 'abs':
-            return a > best + self.threshold
-
-    def _init_is_better(self, mode, threshold, threshold_mode):
-        if mode not in {'min', 'max'}:
-            raise ValueError('mode ' + mode + ' is unknown!')
-        if threshold_mode not in {'rel', 'abs'}:
-            raise ValueError('threshold mode ' + threshold_mode + ' is unknown!')
-
-        if mode == 'min':
-            self.mode_worse = inf
-        else:  # mode == 'max':
-            self.mode_worse = -inf
-
-        self.mode = mode
-        self.threshold = threshold
-        self.threshold_mode = threshold_mode
-
-    def state_dict(self):
-        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
-
-    def load_state_dict(self, state_dict):
-        self.__dict__.update(state_dict)
-        self._init_is_better(mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode)
-
-class Action_On_Plateau():
-
-    def __init__(self, mode = 'max', patience=10,
-                 plateau_var = "test_auc", 
-                 threshold=1e-4, threshold_mode='rel', cooldown=0,
-                 eps=1e-8, verbose=False):
-
-        self.patience = patience
-        self.plateau_var = plateau_var
-
-        self.verbose = verbose
-        self.cooldown = cooldown
-        self.cooldown_counter = 0
-        self.mode = mode
-        self.threshold = threshold
-        self.threshold_mode = threshold_mode
-        self.best = None
-        self.num_bad_epochs = None
-        self.mode_worse = None  # the worse value for the chosen mode
-        self.eps = eps
-        self.last_epoch = 0
-        self._init_is_better(mode=mode, threshold=threshold,
-                             threshold_mode=threshold_mode)
-        self._reset()
-
-    def _reset(self):
-        """Resets num_bad_epochs counter and cooldown counter."""
-        self.best = self.mode_worse
-        self.cooldown_counter = 0
-        self.num_bad_epochs = 0
-
-    def step(self, model, metrics, epoch=None):
-        # convert `metrics` to float, in case it's a zero-dim Tensor
-        current = float(metrics[self.plateau_var])
-        if epoch is None:
-            epoch = self.last_epoch + 1
-        else:
-            warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
-        self.last_epoch = epoch
-
-        if self.is_better(current, self.best):
-            if(self.verbose):
-                print("Model is improving!")
-            self.best = current
-            self.num_bad_epochs = 0
-        else:
-            if(self.verbose):
-                print(f"Model is not improving :( best = {self.best}, current = {current}")
-            self.num_bad_epochs += 1
-
-        if self.in_cooldown:
-            self.cooldown_counter -= 1
-            self.num_bad_epochs = 0  # ignore any bad epochs in cooldown
-
-        if self.num_bad_epochs > self.patience:
-            self.action(model, metrics, epoch)
-    
-    def action(self, model, metrics, epoch=None):
-        if(self.verbose):
-            print("Doing my action")
-    
-    @property
-    def in_cooldown(self):
-        return self.cooldown_counter > 0
-
-    def is_better(self, a, best):
-        if self.mode == 'min' and self.threshold_mode == 'rel':
-            rel_epsilon = 1. - self.threshold
-            return a < best * rel_epsilon
-
-        elif self.mode == 'min' and self.threshold_mode == 'abs':
-            return a < best - self.threshold
-
-        elif self.mode == 'max' and self.threshold_mode == 'rel':
-            rel_epsilon = self.threshold + 1.
-            return a > best * rel_epsilon
-
-        else:  # mode == 'max' and epsilon_mode == 'abs':
-            return a > best + self.threshold
-
-    def _init_is_better(self, mode, threshold, threshold_mode):
-        if mode not in {'min', 'max'}:
-            raise ValueError('mode ' + mode + ' is unknown!')
-        if threshold_mode not in {'rel', 'abs'}:
-            raise ValueError('threshold mode ' + threshold_mode + ' is unknown!')
-
-        if mode == 'min':
-            self.mode_worse = inf
-        else:  # mode == 'max':
-            self.mode_worse = -inf
-
-        self.mode = mode
-        self.threshold = threshold
-        self.threshold_mode = threshold_mode
-
-class Partial_Reset(Action_On_Plateau):
-
-    def __init__(self, mode='max', patience=10, plateau_var="test_auc", 
-                        threshold=0.0001, threshold_mode='rel', cooldown=0, 
-                        eps=1e-8, verbose=False):
-        
-        super().__init__(mode, patience, plateau_var, threshold, 
-                         threshold_mode, cooldown, eps, verbose)
-    
-    def action(self, model, metrics, epoch=None):
-        print("Partial Reset!!")
-        GCN.partial_reset(model)
-        self._reset()
-        self.cooldown_counter = self.cooldown
-        self.num_bad_epochs = 0
-    
-
-class Full_Reset(Action_On_Plateau):
-
-    def __init__(self, mode='max', patience=10, plateau_var="test_auc", 
-                        threshold=0.0001, threshold_mode='rel', cooldown=0, 
-                        eps=1e-8, verbose=False):
-        
-        super().__init__(mode, patience, plateau_var, threshold, 
-                         threshold_mode, cooldown, eps, verbose)
-    
-    def action(self, model, metrics, epoch=None):
-        print("Full Reset!!")
-        GCN.full_reset(model)
-        self._reset()
-        self.cooldown_counter = self.cooldown
-        self.num_bad_epochs = 0
-
-class Dynamic_LR_AND_Partial_Reset():
-    def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10,
-                plateau_var = "test_auc", reset_patience=None, reset_plateau_var=None,
-                threshold=1e-4, threshold_mode='rel', cooldown=0,
-                min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False):
-
-        if (reset_patience == None):
-            reset_patience = patience
-        if(reset_plateau_var == None):
-            reset_plateau_var = plateau_var
-
-        self.dynamic_lr = Dynamic_LR(optimizer, mode=mode, factor=factor, patience = patience, 
-                        plateau_var=plateau_var, threshold=threshold, threshold_mode =threshold_mode, 
-                        cooldown=cooldown, min_lr=min_lr, max_lr=max_lr, eps=eps, verbose=verbose)
-        
-        self.partial_reset = Partial_Reset(mode=mode, patience=reset_patience, plateau_var=reset_plateau_var,
-                        threshold=threshold, threshold_mode=threshold_mode, cooldown=cooldown,
-                        eps=eps)
-
-    def step(self, model, metrics, epoch=None):
-        self.dynamic_lr.step(model=model, metrics=metrics, epoch=epoch)
-        self.partial_reset.step(model=model, metrics=metrics, epoch=epoch)
-
-class Dynamic_LR_AND_Full_Reset():
-    def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10,
-                plateau_var = "test_auc", reset_patience=None, reset_plateau_var=None,
-                threshold=1e-4, threshold_mode='rel', cooldown=0,
-                min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False):
-
-        if (reset_patience == None):
-            reset_patience = patience
-        if(reset_plateau_var == None):
-            reset_plateau_var = plateau_var
-
-        self.dynamic_lr = Dynamic_LR(optimizer, mode=mode, factor=factor, patience = patience, 
-                        plateau_var=plateau_var, threshold=threshold, threshold_mode =threshold_mode, 
-                        cooldown=cooldown, min_lr=min_lr, max_lr=max_lr, eps=eps, verbose=verbose)
-        
-        self.full_reset = Full_Reset(mode=mode, patience=reset_patience, plateau_var=reset_plateau_var,
-                        threshold=threshold, threshold_mode=threshold_mode, cooldown=cooldown,
-                        eps=eps)
-
-    def step(self, model, metrics, epoch=None):
-        self.dynamic_lr.step(model=model, metrics=metrics, epoch=epoch)
-        self.full_reset.step(model=model, metrics=metrics, epoch=epoch)

nemo/root_gnn_base/dataset.py

@@ -1,678 +0,0 @@
-from dgl.data import DGLDataset
-import dgl
-import uproot
-import awkward as ak
-import torch
-import os
-import glob
-import time
-import numpy as np
-from root_gnn_base import utils
-
-def node_features_from_tree(ch, node_branch_names, node_branch_types, node_feature_scales):
-    lengths = []
-    for branch, node_type in zip(node_branch_names[0], node_branch_types):
-        if node_type == 'single':
-            lengths.append(1)
-        elif node_type == 'vector':
-            lengths.append(len(ch[branch]))
-        else:
-            print('Unknown node branch type: {}'.format(node_type))
-    features = []
-    for node_feat in node_branch_names:
-        if node_feat == 'CALC_E':
-            features.append(features[0]*torch.cosh(features[1]))
-            continue
-        elif node_feat == 'NODE_TYPE':
-            feat = []
-            for i, length in enumerate(lengths):
-                feat.extend([i,]*length)
-            features.append(torch.tensor(feat))
-            continue
-        feat = []
-        itype = 0
-        for length, branch, node_type in zip(lengths, node_feat, node_branch_types):
-            if isinstance(branch, (int, float, complex)):
-                feat.extend([branch,]*length)
-            elif branch == 'CALC_E':
-                this_type_starts_at = sum(lengths[:itype])
-                this_type_ends_at = sum(lengths[:itype+1])
-                feat.extend(features[0][this_type_starts_at:this_type_ends_at]*torch.cosh(features[1][this_type_starts_at:this_type_ends_at]))
-            elif node_type == 'single':
-                feat.append(ch[branch])
-            elif node_type == 'vector':
-                feat.extend(ch[branch])
-            itype += 1
-        features.append(torch.tensor(feat))
-    return torch.stack(features, dim=1) * node_feature_scales, lengths
-
-def full_connected_graph(n_nodes, self_loops=True):
-    senders = np.arange(n_nodes*n_nodes) // n_nodes
-    receivers = np.arange(n_nodes*n_nodes) % n_nodes
-    if not self_loops and n_nodes > 1:
-        mask = senders != receivers
-        senders = senders[mask]
-        receivers = receivers[mask]
-    return dgl.graph((senders, receivers))
-
-def check_selection(ch, selection):
-    var, cut, op = selection
-    if op == '>':
-        return ch[var] > cut
-    elif op == '<':
-        return ch[var] < cut
-    elif op == '==':
-        return ch[var] == cut
-
-def check_selections(ch, selections):
-    for selection in selections:
-        if not check_selection(ch, selection):
-            return False
-    return True
-
-class RootDataset(DGLDataset):
-    def __init__(self, name=None, raw_dir=None, save_dir=None, label=1, file_names = '*.root', node_branch_names=None, node_branch_types=None, node_feature_scales=None, 
-                 selections=[], save=True, tree_name = 'nominal_Loose', fold_var = 'eventNumber', weight_var = None, chunks = 1, process_chunks = None, global_features = [], tracking_info = [], **kwargs):
-        print(f'Unused args while creating RootDataset: {kwargs}')
-        self.label = label
-        self.counts = []
-        self.selections = selections
-        self.save_to_disk = save
-        self.file_names = file_names
-        self.node_branch_names = node_branch_names
-        self.node_branch_types = node_branch_types
-        self.node_feature_scales = torch.tensor([float(sf) for sf in node_feature_scales])
-        self.tree_name = tree_name
-        self.fold_var = fold_var
-        self.tracking_info = tracking_info
-        self.tracking_info.insert(0, fold_var)
-        if weight_var is None:
-            weight_var = 1
-        self.tracking_info.insert(1, weight_var)
-        self.global_features = global_features
-        self.chunks = chunks
-        self.process_chunks = process_chunks
-        if self.process_chunks is None:
-            self.process_chunks = [i for i in range(self.chunks)]
-        self.times = [0, 0]
-        super().__init__(name=name, raw_dir=raw_dir, save_dir=save_dir)
-
-    def get_list_of_branches(self):
-        branches = []
-        for feat in self.node_branch_names:
-            if isinstance(feat, list):
-                for branch in feat:
-                    if branch == 'CALC_E':
-                        continue
-                    if isinstance(branch, str):
-                        branches.append(branch)
-        for feat in self.global_features:
-            if isinstance(feat, str):
-                branches.append(feat)
-        for feat in self.tracking_info:
-            if isinstance(feat, str):
-                branches.append(feat)
-        for selection in self.selections:
-            branches.append(selection[0])
-        return list(set(branches))  # Remove duplicates
-
-    def make_graph(self, ch):
-        t1 = time.time()
-        features, _ = node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales)
-        features = features[features[:,0] != 0]
-        t2 = time.time()
-        g = full_connected_graph(features.shape[0], self_loops=False)
-        g.ndata['features'] = features
-        t3 = time.time()
-        self.times[0] += t2 - t1
-        self.times[1] += t3 - t2
-        return g
-
-    def process(self):
-        times = [0, 0, 0]
-        oldtime = time.time()
-        if isinstance(self.file_names, str):
-            self.files = glob.glob(os.path.join(self.raw_dir, self.file_names))
-        else:
-            self.files = []
-            for file_name in self.file_names:
-                self.files.extend(glob.glob(os.path.join(self.raw_dir, file_name)))
-        branches = self.get_list_of_branches()
-
-        # Read all files and concatenate arrays
-        arrays = []
-        for file in self.files:
-            with uproot.open(file) as f:
-                arrays.append(f[self.tree_name].arrays(branches, library="ak"))
-        if len(arrays) == 0:
-            print('No files found in {}'.format(os.path.join(self.raw_dir, self.file_names)))
-            return
-        data = ak.concatenate(arrays, axis=0)
-        n_entries = len(data[branches[0]])
-        newtime = time.time()
-        times[0] += newtime - oldtime
-        chunks = np.array_split(np.arange(n_entries), self.chunks)
-        chunks = [chunk for i, chunk in enumerate(chunks) if i in self.process_chunks]
-
-        self.graph_chunks = []
-        self.label_chunks = []
-        self.tracking_chunks = []
-        self.global_chunks = []
-        chunk_id = -1
-        for chunk in chunks:
-            print('Processing chunk {}/{}'.format(chunk_id + 1, len(chunks)), flush=True)
-            chunk_id += 1
-            graphs = []
-            labels = []
-            tracking = []
-            globals = []
-            for ientry in chunk:
-                if (ientry % 10000 == 0):
-                    print('Processing event {}/{}'.format(ientry, n_entries), flush=True)
-                ch = {b: data[b][ientry] for b in branches}
-                passed = True
-                for selection in self.selections:
-                    if not check_selection(ch, selection):
-                        passed = False
-                        continue
-                oldtime = newtime
-                newtime = time.time()
-                times[1] += newtime - oldtime
-                if passed:
-                    graphs.append(self.make_graph(ch))
-                    labels.append(self.label)
-                    tracking.append(torch.zeros(len(self.tracking_info), dtype=torch.double))
-                    globals.append(torch.zeros(len(self.global_features)))
-                    for i_ti, tr_branch in enumerate(self.tracking_info):
-                        if isinstance(tr_branch, str):
-                            tracking[-1][i_ti] = ch[tr_branch]
-                        else:
-                            tracking[-1][i_ti] = tr_branch
-                    for i_gl, gl_branch in enumerate(self.global_features):
-                        globals[-1][i_gl] = ch[gl_branch]
-                oldtime = newtime
-                newtime = time.time()
-                times[2] += newtime - oldtime
-
-            labels = torch.tensor(labels)
-            tracking = torch.stack(tracking)
-            globals = torch.stack(globals)
-
-            self.graph_chunks.append(graphs)
-            self.label_chunks.append(labels)
-            self.tracking_chunks.append(tracking)
-            self.global_chunks.append(globals)
-            self.counts.append(len(graphs))
-            
-            if (self.chunks > 1):
-                self.save_chunk(chunk_id, graphs, labels, tracking, globals)
-            else:
-                self.labels = labels
-                self.tracking = tracking
-                self.global_features = globals
-                self.graphs = graphs
-                self.save()
-        return
-
-    def save(self):
-        if not self.save_to_disk:
-            return
-        graph_path = os.path.join(self.save_dir, self.name + '.bin')
-        if self.chunks == 1:
-            print(f'Saving dataset to {os.path.join(self.save_dir, self.name + ".bin")}')
-            dgl.save_graphs(str(graph_path), self.graphs, {'labels': torch.tensor(self.labels), 'tracking': torch.tensor(self.tracking), 'global': torch.tensor(self.global_features)})
-        else:
-            for i in range(len(self.process_chunks)):
-                print(f'Saving dataset to {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}')
-
-                dgl.save_graphs(str(graph_path).replace('.bin', f'_{self.process_chunks[i]}.bin'), self.graph_chunks[i], {'labels': self.label_chunks[i], 'tracking': self.tracking_chunks[i], 'global': self.global_chunks[i]})
-
-    def save_chunk(self, chunk_id, graphs, labels, tracking, globals):
-        if not self.save_to_disk:
-            return
-        graph_path = os.path.join(self.save_dir, self.name + '.bin')
-        print(f'Saving dataset to {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[chunk_id]}.bin")}')
-        dgl.save_graphs(str(graph_path).replace('.bin', f'_{self.process_chunks[chunk_id]}.bin'), graphs, {'labels': labels, 'tracking': tracking, 'global': globals})
-
-    def has_cache(self):
-        print(f'Checking for cache of {self.name}')
-        if not self.save_to_disk:
-            print('Skipping load.')
-            return False
-        if self.chunks == 1:
-            graph_path = os.path.join(self.save_dir, self.name + '.bin')
-            return os.path.exists(graph_path)
-        else:
-            for i in range(len(self.process_chunks)):
-                graph_path = os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin')
-                if not os.path.exists(graph_path):
-                    print(f'File {graph_path} does not exist, processing.')
-                    return False
-            return True
-
-    def load(self):
-        if self.chunks == 1:
-            print(f'Loading dataset from {os.path.join(self.save_dir, self.name + ".bin")}')
-            graphs, label_dict = dgl.load_graphs(os.path.join(self.save_dir, self.name + '.bin'))
-            self.graphs = graphs
-            self.labels = label_dict['labels']
-            self.tracking = label_dict['tracking']
-            self.global_features = label_dict['global']
-        else:
-            self.graphs = []
-            self.labels = []
-            self.tracking = []
-            self.global_features = []
-            for i in range(self.chunks):
-                try:
-                    print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}')
-                    graphs, label = dgl.load_graphs(os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin'))
-                    self.graphs.extend(graphs)
-                    self.labels.append(label['labels'])
-                    self.tracking.append(label['tracking'])
-                    self.global_features.append(label['global'])
-                except Exception as e:
-                    print(e)
-            self.labels = torch.cat(self.labels)
-            self.tracking = torch.cat(self.tracking)
-            self.global_features = torch.cat(self.global_features)
-
-    def __getitem__(self, idx):
-        return self.graphs[idx], self.labels[idx], self.tracking[idx], self.global_features[idx]
-    
-    def __len__(self):
-        return len(self.graphs)
-
-#Dataset with edge features added (deta, dphi, dR)
-class EdgeDataset(RootDataset):
-    def make_graph(self, ch):
-        g = super().make_graph(ch)
-        u, v = g.edges()
-        deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1]
-        dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2]
-        dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi)
-        dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi)
-        dR   = torch.sqrt(deta**2 + dphi**2)
-        g.edata['features'] = torch.stack([deta, dphi, dR], dim=1)
-        return g
-
-class tHbbEdgeDataset(RootDataset):
-    def __init__(self, exclude_branches=None, **kwargs):
-        self.exclude_branches = exclude_branches
-        super().__init__(**kwargs)
-
-    def get_list_of_branches(self):
-        br = super().get_list_of_branches()
-        for sector in self.exclude_branches:
-            if sector == None:
-                continue
-            for excl in sector:
-                if type(excl) == str:
-                    br.append(excl)
-        return br
-
-    def make_graph(self, ch):
-        features, lengths = node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales)
-
-        include_mask = torch.ones(features.shape[0], dtype=torch.bool)
-        node_idx = 0
-        for sector, length in zip(self.exclude_branches, lengths):
-            if sector == None:
-                node_idx += length
-                continue
-            for excl in sector:
-                if type(excl) == int:
-                    include_mask[excl + node_idx] = False
-                elif type(excl) == str:
-                    include_mask[getattr(self.chain, excl) + node_idx] = False
-        g = full_connected_graph(features[include_mask].shape[0], self_loops=False)
-        g.ndata['features'] = features[include_mask]
-
-        u, v = g.edges()
-        deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1]
-        dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2]
-        dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi)
-        dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi)
-        dR   = torch.sqrt(deta**2 + dphi**2)
-        g.edata['features'] = torch.stack([deta, dphi, dR], dim=1)
-        return g
-    
-class LazyDataset(EdgeDataset):
-    def __init__(self, buffer_size = 2, **kwargs):
-        self.buffer = [None,] * buffer_size
-        self.buffer_ptr = 0
-        self.get_item_calls = 0
-        self.buffer_indices = [-1,] * buffer_size
-        super().__init__(**kwargs)
-
-    def __getitem__(self, idx):
-        self.get_item_calls += 1
-        chunk_idx = -1
-        sum = 0
-        ev_idx = -999
-        for i, count in enumerate(self.counts):
-            sum += count
-            if idx < sum:
-                chunk_idx = i
-                ev_idx = idx - sum + count
-                break
-        buf_idx = self.buffer_get(chunk_idx)
-        if ev_idx >= len(self.buffer[buf_idx][0]):
-            print(f'Getting event {ev_idx} from chunk {chunk_idx} from buffer {buf_idx}. Calls: {self.get_item_calls}')
-            print(len(self.buffer))
-            print(self.counts)
-            print(len(self.buffer[buf_idx][0]))
-        return self.buffer[buf_idx][0][ev_idx], self.buffer[buf_idx][1]['labels'][ev_idx], self.buffer[buf_idx][1]['tracking'][ev_idx], self.buffer[buf_idx][1]['global'][ev_idx]
-
-    def buffer_get(self, buffer_idx):
-        if buffer_idx in self.buffer_indices:
-            for i in range(len(self.buffer)):
-                if self.buffer_indices[i] == buffer_idx:
-                    return i
-        else:
-            print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{buffer_idx}.bin")}', flush=True)
-            self.buffer_ptr = (self.buffer_ptr + 1) % len(self.buffer)
-            self.buffer[self.buffer_ptr] = dgl.load_graphs(os.path.join(self.save_dir, self.name + f'_{buffer_idx}.bin'))
-            self.buffer_indices[self.buffer_ptr] = buffer_idx
-            return self.buffer_ptr
-
-    def load(self):
-        self.counts = []
-        self.tracking = []
-        try:
-            for i in range(self.chunks):
-                print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}')
-                l = dgl.data.graph_serialize.load_labels_v2(os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin'))
-                self.counts.append(len(l['tracking']))
-                self.tracking.append(l['tracking'])
-            self.tracking = torch.cat(self.tracking)
-        except Exception as e:
-            print(e)
-
-    def __len__(self):
-        return sum(self.counts)
-
-class MultiLabelDataset(EdgeDataset):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def get_list_of_branches(self):
-        br =  super().get_list_of_branches()
-        for l in self.label:
-            if isinstance(l, str):
-                br.append(l)
-            if isinstance(l, dict):
-                br.append(l['branch'])
-        return br
-
-    def get_label(self, ch):
-        label = []
-        for l in self.label:
-            if isinstance(l, str):
-                label.append((getattr(ch, l)))
-            if isinstance(l, dict):
-                label.append(getattr(ch, l['branch'])*float(l['scale']))
-            if isinstance(l, float) or isinstance(l, int):
-                label.append(l)
-
-        return torch.tensor(label)
-
-    def process(self):
-        times = [0, 0, 0]
-        oldtime = time.time()
-        if isinstance(self.file_names, str):
-            self.files = glob.glob(os.path.join(self.raw_dir, self.file_names))
-        else:
-            self.files = []
-            for file_name in self.file_names:
-                self.files.extend(glob.glob(os.path.join(self.raw_dir, file_name)))
-        self.chain = ROOT.TChain(self.tree_name)
-        if len(self.files) == 0:
-            print('No files found in {}'.format(os.path.join(self.raw_dir, self.file_names)))
-        for file in self.files:
-            utils.set_timeout(60*2)
-            self.chain.Add(file)
-            utils.unset_timeout()
-        branches = self.get_list_of_branches()
-        self.chain.SetBranchStatus('*', 0)
-        for branch in branches:
-            self.chain.SetBranchStatus(branch, 1)
-        newtime = time.time()
-        times[0] += newtime - oldtime
-        chunks = np.array_split(np.arange(self.chain.GetEntries()), self.chunks)
-        chunks = [chunk for i, chunk in enumerate(chunks) if i in self.process_chunks]
-        self.graph_chunks = []
-        self.label_chunks = []
-        self.tracking_chunks = []
-        self.global_chunks = []
-        chunk_id = -1
-        for chunk in chunks:
-            chunk_id += 1
-            graphs = []
-            labels = []
-            tracking = []
-            globals = []
-            for ientry in chunk:
-                if (ientry % 10000 == 0):
-                    print('Processing event {}/{}'.format(ientry, self.chain.GetEntries()), flush=True)
-                self.chain.GetEntry(ientry)
-                passed = True
-                for selection in self.selections:
-                    if not check_selection(self.chain, selection):
-                        passed = False
-                        continue
-                oldtime = newtime
-                newtime = time.time()
-                times[1] += newtime - oldtime
-                if passed:
-                    graphs.append(self.make_graph(self.chain))
-                    labels.append(self.get_label(self.chain))
-                    tracking.append(torch.zeros(len(self.tracking_info), dtype=torch.double))
-                    globals.append(torch.zeros(len(self.global_features)))
-                    for i_ti, tr_branch in enumerate(self.tracking_info):
-                        if isinstance(tr_branch, str):
-                            tracking[-1][i_ti] = getattr(self.chain, tr_branch)
-                        else:
-                            tracking[-1][i_ti] = tr_branch
-                    for i_gl, gl_branch in enumerate(self.global_features):
-                        globals[-1][i_gl] = getattr(self.chain, gl_branch)
-                oldtime = newtime
-                newtime = time.time()
-                times[2] += newtime - oldtime
-
-            labels = torch.stack(labels)
-            self.save_chunk(chunk_id, graphs, labels, torch.stack(tracking), torch.stack(globals))
-            # self.graph_chunks.append(graphs)
-            # self.label_chunks.append(labels)
-            # self.tracking_chunks.append(torch.stack(tracking))
-            # self.global_chunks.append(torch.stack(globals))
-            # self.counts.append(len(graphs))
-        return
-        self.graphs = self.graph_chunks[0]
-        for chunk in self.graph_chunks[1:]:
-            self.graphs += chunk
-
-        self.labels = torch.cat(self.label_chunks)
-        self.tracking = torch.cat(self.tracking_chunks)
-        self.global_features = torch.cat(self.global_chunks)
-        print('Time spent: Creating TChain: {}s, Getting Entries and Selection: {}s, Graph Creation: {}s'.format(*times))
-        print('Time spent in node_features_from_tree: {}s, full_connected_graph: {}s'.format(*self.times))
-
-class LazyMultiLabelDataset(MultiLabelDataset, LazyDataset):
-    def __init__(self, buffer_size = 2, **kwargs):
-        LazyDataset.__init__(self, buffer_size=buffer_size, **kwargs)
-
-class MultiLabeltHbbDataset(MultiLabelDataset, tHbbEdgeDataset):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def get_list_of_branches(self):
-        br =  super().get_list_of_branches()
-        for sector in self.exclude_branches:
-            if sector == None:
-                continue
-            for excl in sector:
-                if type(excl) == str:
-                    br.append(excl)
-        return br
-
-
-class AugmentedDataset(RootDataset):
-
-    def __init__(self, seed = 2, feature_index = None, node_mapping = None, **kwargs):
-        self.seed = seed
-        np.random.seed(seed)
-        if(feature_index == None):
-            self.feature_index = {"pt": 0, "eta": 1, "phi": 2, "energy": 3, "btag": 4, "charge": 5, "node_type": 6}
-        if (node_mapping == None):
-            self.node_mapping = {"jet": 0, "ele": 1, "mu": 2, "ph": 3, "MET": 4}
-        super().__init__(**kwargs)
-
-    def detector_noise(self, node_features):
-        noise = np.zeros_like(node_features)
-        
-        node_types = node_features[:, self.feature_index["node_type"]]
-        pts = node_features[:, self.feature_index["pt"]]
-        etas = node_features[:, self.feature_index["eta"]]
-        energies = node_features[:, self.feature_index["energy"]]
-
-        # Noise calculation for jets
-        jet_mask = (node_types == self.node_mapping["jet"])
-        jet_pts = pts[jet_mask]
-        jet_etas = etas[jet_mask]
-
-        if (jet_mask.sum() > 0):
-            jet_resolutions = np.where(
-                jet_pts <= 0.1, 0.0,
-                np.where(
-                    np.abs(jet_etas) <= 0.5, np.sqrt(0.06**2 + jet_pts**2 * 1.3e-3**2),
-                    np.where(
-                        np.abs(jet_etas) <= 1.5, np.sqrt(0.10**2 + jet_pts**2 * 1.7e-3**2),
-                        np.where(
-                            np.abs(jet_etas) <= 2.5, np.sqrt(0.25**2 + jet_pts**2 * 3.1e-3**2),
-                            0.0
-                        )
-                    )
-                )
-            )
-            noise[jet_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=jet_resolutions)
-
-        # Noise calculation for electrons
-        ele_mask = (node_types == self.node_mapping["ele"])
-        ele_pts = pts[ele_mask]
-        ele_etas = etas[ele_mask]
-
-        if (ele_mask.sum() > 0):
-            ele_resolutions = np.where(
-                np.abs(ele_etas) <= 0.5, np.sqrt(0.03**2 + ele_pts**2 * 1.3e-3**2),
-                np.where(
-                    np.abs(ele_etas) <= 1.5, np.sqrt(0.05**2 + ele_pts**2 * 1.7e-3**2),
-                    np.where(
-                        np.abs(ele_etas) <= 2.5, np.sqrt(0.15**2 + ele_pts**2 * 3.1e-3**2),
-                        0.0
-                    )
-                )
-            )
-            noise[ele_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=ele_resolutions)
-
-        # Noise calculation for muons
-        mu_mask = (node_types == self.node_mapping["mu"])
-        mu_pts = pts[mu_mask]
-        mu_etas = etas[mu_mask]
-
-        if (mu_mask.sum() > 0):
-            mu_resolutions = np.where(
-                np.abs(mu_etas) <= 0.5, np.sqrt(0.01**2 + mu_pts**2 * 1.0e-4**2),
-                np.where(
-                    np.abs(mu_etas) <= 1.5, np.sqrt(0.015**2 + mu_pts**2 * 1.5e-4**2),
-                    np.where(
-                        np.abs(mu_etas) <= 2.5, np.sqrt(0.025**2 + mu_pts**2 * 3.5e-4**2),
-                        0.0
-                    )
-                )
-            )
-            noise[mu_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=mu_resolutions)
-
-        # Noise calculation for photons
-        ph_mask = (node_types == self.node_mapping["ph"])
-        ph_etas = etas[ph_mask]
-        ph_energies = energies[ph_mask]
-
-        if (ph_mask.sum() > 0):
-            ph_resolutions = np.where(
-                np.abs(ph_etas) <= 3.2, np.sqrt(ph_energies**2 * 0.0017**2 + ph_energies * 0.101**2),
-                np.where(
-                    np.abs(ph_etas) <= 4.9, np.sqrt(ph_energies**2 * 0.0350**2 + ph_energies * 0.285**2),
-                    0.0
-                )
-            )
-            noise[ph_mask, self.feature_index["energy"]] = np.random.normal(loc=0.0, scale=ph_resolutions)
-        return noise
-    
-    def make_graph(self, ch):
-        g = super().make_graph(ch)
-
-        g.ndata['augmented_features'] = g.ndata['features']
-
-        num_nodes = len(g.ndata['features'][:, 0])
-
-        # Rotations: phi -> phi + delta_phi
-        phi_index = self.feature_index["phi"]
-        # Generate a single delta_phi for all nodes
-        delta_phi = np.random.uniform(low=-np.pi, high=np.pi)
-        
-        # Apply the same delta_phi to all nodes
-        g.ndata['augmented_features'][:, phi_index] = (g.ndata['augmented_features'][:, phi_index] + delta_phi + np.pi) % (2 * np.pi) - np.pi
-
-        # Reflections: eta -> -1 * eta, phi -> -1 * phi
-        eta_index = self.feature_index["eta"]
-
-        eta_reflection = np.random.choice([-1, 1])
-        phi_reflection = np.random.choice([-1, 1])
-
-        g.ndata['augmented_features'][:, eta_index] = g.ndata['augmented_features'][:, eta_index] * eta_reflection
-        g.ndata['augmented_features'][:, phi_index] = g.ndata['augmented_features'][:, phi_index] * phi_reflection
-
-
-        # Detector Noise: pt -> pt + normal(pt, noise(pt))
-        noise = self.detector_noise(g.ndata['augmented_features'])
-        g.ndata['augmented_features'] = g.ndata['augmented_features'] + noise
-
-        pt_index = self.feature_index["pt"]
-        if (g.ndata['augmented_features'][-1][self.feature_index["node_type"]] == self.node_mapping["MET"]):
-            # Initialize sums of px and py
-            sum_px = 0
-            sum_py = 0
-
-            # Loop over all nodes except the last one (MET node)
-            for i in range(len(g.ndata['augmented_features']) - 1):
-                pt = g.ndata['augmented_features'][i][pt_index]
-                phi = g.ndata['augmented_features'][i][phi_index]
-
-                # Compute px and py
-                px = pt * np.cos(phi)
-                py = pt * np.sin(phi)
-
-                # Sum px and py
-                sum_px += px
-                sum_py += py
-
-            # Calculate MET
-            g.ndata['augmented_features'][-1][pt_index] = np.sqrt(sum_px**2 + sum_py**2)
-
-        u, v = g.edges()
-        deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1]
-        dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2]
-        dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi)
-        dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi)
-        dR   = torch.sqrt(deta**2 + dphi**2)
-        g.edata['features'] = torch.stack([deta, dphi, dR], dim=1)
-
-        deta = g.ndata['augmented_features'][u, 1] - g.ndata['augmented_features'][v, 1]
-        dphi = g.ndata['augmented_features'][u, 2] - g.ndata['augmented_features'][v, 2]
-        dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi)
-        dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi)
-        dR   = torch.sqrt(deta**2 + dphi**2)
-        g.edata['augmented_features'] = torch.stack([deta, dphi, dR], dim=1)
-
-        return g

nemo/root_gnn_base/photon_ID_dataset.py

@@ -1,44 +0,0 @@
-from root_gnn_base import dataset
-import dgl
-import torch
-import numpy as np
-
-def radius_graph(features, radii, self_loops=False):
-    senders = []
-    receivers = []
-    n_nodes = features.shape[0]
-    senders = np.arange(n_nodes*n_nodes) // n_nodes
-    receivers = np.arange(n_nodes*n_nodes) % n_nodes
-    if not self_loops and n_nodes > 1:
-        mask = senders != receivers
-        senders = senders[mask]
-        receivers = receivers[mask]
-    for k, r in radii.items():
-        d = features[senders, k] - features[receivers, k]
-        mask = np.abs(d) < r
-        senders = senders[mask]
-        receivers = receivers[mask]
-    return dgl.graph((senders, receivers))
-
-class PhotonIDDataset(dataset.LazyMultiLabelDataset):
-    def __init__(self, eta_radius, phi_radius, **kwargs):
-        self.eta_radius = eta_radius
-        self.phi_radius = phi_radius
-        super().__init__(**kwargs)
-    def make_graph(self, ch):
-        features, _ = dataset.node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales)
-        features = features[features[:,0] != 0]
-        #Delta Eta, Delta Phi, Adjacent Layer
-        g = radius_graph(features, {1: self.eta_radius, 2: self.phi_radius, 6: 1.1}, self_loops=True) #Self loops ensure last cell is included even if disconnected
-        g.ndata['features'] = features
-        u, v = g.edges()
-        deta = features[u, 1] - features[v, 1]
-        dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2]
-        dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi)
-        dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi)
-        dR   = torch.sqrt(deta**2 + dphi**2)
-        dx   = features[u, 3] - features[v, 3]
-        dy   = features[u, 4] - features[v, 4]
-        dz   = features[u, 5] - features[v, 5]
-        g.edata['features'] = torch.stack([deta, dphi, dR, dx, dy, dz], dim=1)
-        return g

nemo/root_gnn_base/similarity.py

@@ -1,158 +0,0 @@
-import numpy as np
-import scipy
-from sklearn.decomposition import PCA
-from sklearn.metrics.pairwise import cosine_similarity
-from sklearn.metrics.pairwise import euclidean_distances
-from sklearn.preprocessing import StandardScaler
-
-from scipy.stats import wasserstein_distance
-
-def cka(rep_a, rep_b, size=None):
-    """
-    Computes the Centered Kernel Alignment (CKA) between two large representation matrices rep_a and rep_b.
-    If size is provided, it performs CKA on a randomly selected subset of the data.
-    
-    Parameters:
-    rep_a : np.ndarray
-        First representation matrix of size (n_samples, n_features_a).
-    rep_b : np.ndarray
-        Second representation matrix of size (n_samples, n_features_b).
-    size : int, optional
-        Number of samples to use for the CKA calculation. If None, use the full dataset.
-        
-    Returns:
-    float
-        CKA similarity between rep_a and rep_b.
-    """
-    
-    def gram_linear(x):
-        """Compute the Gram (kernel) matrix using a linear kernel."""
-        return x @ x.T
-    
-    def center_gram(gram):
-        """Center the Gram matrix."""
-        n = gram.shape[0]
-        identity = np.eye(n)
-        ones = np.ones((n, n)) / n
-        return gram - ones @ gram - gram @ ones + ones @ gram @ ones
-    
-    # If sample_size is specified, randomly sample a subset of the data
-    if size is not None and size < rep_a.shape[0]:
-        indices = np.random.choice(rep_a.shape[0], size, replace=False)
-        rep_a = rep_a[indices]
-        rep_b = rep_b[indices]
-    
-    # Compute the Gram matrices
-    gram_a = gram_linear(rep_a)
-    gram_b = gram_linear(rep_b)
-
-    # Center the Gram matrices
-    centered_gram_a = center_gram(gram_a)
-    centered_gram_b = center_gram(gram_b)
-    
-    # Compute the CKA similarity
-    numerator = np.sum(centered_gram_a * centered_gram_b)
-    denominator = np.sqrt(np.sum(centered_gram_a**2) * np.sum(centered_gram_b**2))
-    
-    return numerator / denominator if denominator != 0 else 0
-
-def cca(X, Y, size = None, num_components=10):
-    """
-    Perform Canonical Correlation Analysis (CCA) between two datasets.
-    
-    Parameters:
-    X : np.ndarray
-        First dataset, shape (n_samples, n_features_X).
-    Y : np.ndarray
-        Second dataset, shape (n_samples, n_features_Y).
-    num_components : int
-        Number of CCA components to return.
-        
-    Returns:
-    w_X : np.ndarray
-        Canonical weights for the first dataset, shape (n_features_X, num_components).
-    w_Y : np.ndarray
-        Canonical weights for the second dataset, shape (n_features_Y, num_components).
-    corrs : np.ndarray
-        Array of canonical correlations for each component.
-    """
-    
-    # If sample size is specified, randomly sample a subset of the data
-    if size is not None and size < X.shape[0]:
-        indices = np.random.choice(X.shape[0], size, replace=False)
-        X = X[indices]
-        Y = Y[indices]
-
-    # Standardize both datasets (mean = 0, variance = 1)
-    scaler_X = StandardScaler()
-    scaler_Y = StandardScaler()
-    
-    X = scaler_X.fit_transform(X)
-    Y = scaler_Y.fit_transform(Y)
-    
-    # Covariance matrices
-    C_XX = np.cov(X, rowvar=False)  # Covariance of X
-    C_YY = np.cov(Y, rowvar=False)  # Covariance of Y
-    C_XY = np.cov(X, Y, rowvar=False)[:X.shape[1], X.shape[1]:]  # Cross-covariance of X and Y
-
-    # Regularization term to avoid singular matrices
-    reg = 1e-6
-    inv_C_XX = np.linalg.inv(C_XX + reg * np.eye(C_XX.shape[0]))
-    inv_C_YY = np.linalg.inv(C_YY + reg * np.eye(C_YY.shape[0]))
-
-    # Solve the generalized eigenvalue problem for CCA
-    # (inv_C_XX @ C_XY @ inv_C_YY @ C_XY.T) and vice versa for Y
-    A = inv_C_XX @ C_XY @ inv_C_YY @ C_XY.T
-    B = inv_C_YY @ C_XY.T @ inv_C_XX @ C_XY
-
-    # Perform eigenvalue decomposition
-    eigvals_X, eigvecs_X = np.linalg.eigh(A)
-    eigvals_Y, eigvecs_Y = np.linalg.eigh(B)
-
-    # Sort the eigenvalues and eigenvectors in descending order
-    idx_X = np.argsort(eigvals_X)[::-1]
-    idx_Y = np.argsort(eigvals_Y)[::-1]
-    
-    eigvecs_X = eigvecs_X[:, idx_X]
-    eigvecs_Y = eigvecs_Y[:, idx_Y]
-
-    # Canonical weights (the first `num_components` components)
-    w_X = eigvecs_X[:, :num_components]
-    w_Y = eigvecs_Y[:, :num_components]
-    
-    # Canonical correlations (square root of the eigenvalues, constrained to [0,1])
-    corrs = np.sqrt(np.clip(eigvals_X[:num_components], 0, 1))
-
-    return np.mean(corrs)
-    return w_X, w_Y, corrs
-
-def pca(X, Y, size=1000, n_components=3, bins=30):
-
-    pca_X = PCA(n_components=n_components)
-    X_pca = pca_X.fit_transform(X)
-
-    pca_Y = PCA(n_components=n_components)
-    Y_pca = pca_Y.fit_transform(Y)
-
-    # Step 2: Determine common bin edges based on the range of PCA components
-    min_value = min(X_pca.min(), Y_pca.min())
-    max_value = max(X_pca.max(), Y_pca.max())
-    bin_edges = np.linspace(min_value, max_value, bins + 1)
-
-    # Step 3: Calculate histograms for each PCA component using the same bins
-    histograms_X = [np.histogram(X_pca[:, i], bins=bin_edges, density=True)[0] for i in range(n_components)]
-    histograms_Y = [np.histogram(Y_pca[:, i], bins=bin_edges, density=True)[0] for i in range(n_components)]
-
-    # Step 4: Calculate Wasserstein distance between corresponding histograms
-    total_distance = 0
-    for i in range(n_components):
-        total_distance += wasserstein_distance(histograms_X[i], histograms_Y[i])
-
-    # Step 5: Normalize the total distance for a similarity score
-    # Calculate the maximum possible distance (theoretical max could be based on histogram size)
-    # This could be replaced with a more complex calculation if necessary.
-    max_distance = 1.0  # Replace this with a suitable maximum based on your dataset properties.
-    
-    similarity_score = 1 - (total_distance / max_distance)
-
-    return max(0, min(1, similarity_score))  # Ensure the score stays in [0, 1]

nemo/root_gnn_base/uproot_dataset.py

@@ -1,54 +0,0 @@
-from root_gnn_base import dataset
-import torch
-import uproot
-import glob
-import os
-import awkward as ak
-import numpy as np
-import time
-
-def node_features_from_ak(ch, node_branch_names, node_branch_types, node_feature_scales):
-    node_types = []
-    n_types = len(node_branch_names[0])
-    for i in range(n_types):
-        features = []
-        branch_type = node_branch_types[i]
-        for j in range(len(node_branch_names)):
-            if node_branch_names[j] == 'CALC_E':
-                features.append(features[0] * np.cosh(features[1]))
-            elif node_branch_names[j] == 'NODE_TYPE':
-                features.append(ak.full_like(features[0], i))
-            elif isinstance(node_branch_names[j][i], str):
-                features.append(ch[node_branch_names[j][i]])
-            elif isinstance(node_branch_names[j][i], (int, float)):
-                features.append(ak.full_like(features[0], node_branch_names[j][i]))
-        if branch_type == 'single':
-            features = [f[:,np.newaxis] for f in features]
-        node_types.append(ak.Array(features))
-    node_features = ak.concatenate(node_types, axis=2) * node_feature_scales #axis order at this point is (feature, event, node)
-    return node_features
-
-class UprootDataset(dataset.RootDataset):
-    def process(self):
-        starttime = time.time()
-        self.files = glob.glob(os.path.join(self.raw_dir, self.file_names))
-        branches = self.get_list_of_branches()
-        self.chain = uproot.concatenate([f + ':' + self.tree_name for f in self.files], branches, num_workers=4)
-        node_features = node_features_from_ak(self.chain, self.node_branch_names, self.node_branch_types, self.node_feature_scales)
-        loadtime = time.time()
-        n_nodes = ak.num(node_features[0], axis=1) #number of nodes for each event
-        ftime = time.time()
-        self.graphs = [dataset.full_connected_graph(n, False) for n in n_nodes]
-        itime = time.time()
-        for i in range(len(self.graphs)):
-            if i % 10000 == 0:
-                print(f'Processing event {i}/{len(self.graphs)}')
-            self.graphs[i].ndata['features'] = torch.transpose(torch.tensor(node_features[:,i,:]),0,1).to(torch.float)
-        self.label = torch.stack([torch.full((len(self.graphs),),torch.tensor(self.label)), torch.tensor(ak.values_astype(self.chain[self.fold_var], np.int64))], dim=1)
-        gtime = time.time()
-        print()
-        print(f'load time: {loadtime - starttime} s')
-        print(f'feature time: {ftime - loadtime} s')
-        print(f'graph time: {itime - ftime} s')
-        print(f'graph data time: {gtime - itime} s')
-

nemo/root_gnn_base/utils.py

@@ -1,393 +0,0 @@
-import importlib
-import yaml
-import os
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-import dgl
-import signal
-
-def buildFromConfig(conf, run_time_args = {}):
-    device = run_time_args.get('device', 'cpu')
-    if 'module' in conf:
-        module = importlib.import_module(conf['module'])
-        cls = getattr(module, conf['class'])
-        args = conf['args'].copy()
-        if 'weight' in args and isinstance(args['weight'], list):
-            args['weight'] = torch.tensor(args['weight'], dtype=torch.float, device=device)
-        # Remove device from run_time_args to not pass it to the class
-        run_time_args = {k: v for k, v in run_time_args.items() if k != 'device'}
-        return cls(**args, **run_time_args)
-    else:
-        print('No module specified in config. Returning None.')
-
-def cycler(iterable):
-    while True:
-        #print('Cycler is cycling...')
-        for i in iterable:
-            yield i
-
-def include_config(conf):
-    if 'include' in conf:
-        for i in conf['include']:
-            with open(i) as f:
-                conf.update(yaml.load(f, Loader=yaml.FullLoader))
-        del conf['include']
-
-def load_config(config_file):
-    with open(config_file) as f:
-        conf = yaml.load(f, Loader=yaml.FullLoader)
-    include_config(conf)
-    return conf
-
-#Timeout function from https://stackoverflow.com/questions/492519/timeout-on-a-function-call
-class TimeoutException(Exception):
-    pass
-
-def timeout_handler(signum, frame):
-    raise TimeoutException()
-
-def set_timeout(timeout):
-    signal.signal(signal.SIGALRM, timeout_handler)
-    signal.alarm(timeout)
-
-def unset_timeout():
-    signal.alarm(0)
-    signal.signal(signal.SIGALRM, signal.SIG_DFL) 
-
-def make_padding_graph(batch, pad_nodes, pad_edges):
-    senders = []
-    receivers = []
-    senders = torch.arange(0,pad_edges) // pad_nodes
-    receivers = torch.arange(1,pad_edges+1) % pad_nodes
-    if pad_nodes < 0 or pad_edges < 0 or pad_edges > pad_nodes * pad_nodes / 2:
-        print('Batch is larger than padding size or e > n^2/2. Repeating edges as necessary.')
-        print(f'Batch nodes: {batch.num_nodes()}, Batch edges: {batch.num_edges()}, Padding nodes: {pad_nodes}, Padding edges: {pad_edges}')
-        senders = senders % pad_nodes
-    padg = dgl.graph((senders[:pad_edges], receivers[:pad_edges]), num_nodes = pad_nodes)
-    for k in batch.ndata.keys():
-        padg.ndata[k] = torch.zeros( (pad_nodes, batch.ndata[k].shape[1]) )
-    for k in batch.edata.keys():
-        padg.edata[k] = torch.zeros( (pad_edges, batch.edata[k].shape[1]) )
-    return dgl.batch([batch, padg.to(batch.device)])
-
-def pad_size(graphs, edges, nodes, edge_per_graph=3, node_per_graph=14):
-    pad_nodes = ((nodes // (node_per_graph * graphs))+1) * graphs * node_per_graph
-    pad_edges = ((edges // (edge_per_graph * graphs))+1) * graphs * edge_per_graph
-    return pad_nodes, pad_edges
-
-def pad_batch_to_step_per_graph(batch, edge_per_graph=3, node_per_graph=14):
-    n_graphs = batch.batch_num_nodes().shape[0]
-    pad_nodes = (batch.num_nodes() + node_per_graph * n_graphs) % int(n_graphs * node_per_graph)
-    pad_edges = (batch.num_edges() + edge_per_graph * n_graphs) % int(n_graphs * edge_per_graph)
-    return make_padding_graph(batch, pad_nodes, pad_edges)
-
-def pad_batch(batch, edges = 104000, nodes = 16000):
-    if edges == 0 and nodes == 0:
-        return batch
-    pad_nodes = 0
-    pad_edges = 0
-    pad_nodes = nodes - batch.num_nodes()
-    pad_edges = edges - batch.num_edges()
-    return make_padding_graph(batch, pad_nodes, pad_edges)
-
-def pad_batch_num_nodes(batch, max_num_nodes, hid_size = 64):
-    print(f"Padding each graph to have {max_num_nodes} nodes. Using hidden size {hid_size}.")
-
-    unbatched = dgl.unbatch(batch)
-    for g in unbatched:
-        num_nodes_to_add = max_num_nodes - g.number_of_nodes()
-        if num_nodes_to_add > 0:
-            g.add_nodes(num_nodes_to_add)  # Add isolated nodes
-
-    batch = dgl.batch(unbatched)
-
-    padding_mask = torch.zeros((batch.ndata['features'].shape[0]), dtype=torch.bool)
-    global_update_weights = torch.ones((batch.ndata['features'].shape[0], hid_size))
-
-    for i in range(len(batch.ndata['features'])):
-        if (torch.count_nonzero(batch.ndata['features'][i]) == 0):
-            padding_mask[i] = True
-            global_update_weights[i] = 0
-
-    batch.ndata['w'] = global_update_weights
-    batch.ndata['padding_mask'] = padding_mask
-
-    return batch
-
-
-def fold_selection(fold_config, sample):
-    n_folds = fold_config['n_folds']
-    folds_opt = fold_config[sample]
-    folds = []
-    if type(folds_opt) == int:
-        return lambda x : x.tracking[:,0] % n_folds == folds_opt
-    elif type(folds_opt) == list:
-        print("fold type is list")
-        print(f"fold_config = {fold_config}")
-        print(f"folds_opt = {folds_opt}")
-        return lambda x : sum([x.tracking[:,0] % n_folds == f for f in folds_opt]) == 1
-    else:
-        raise ValueError("Invalid fold selection option with type {}".format(type(folds_opt)))
-    
-def fold_selection_name(fold_config, sample):
-    n_folds = fold_config['n_folds']
-    folds_opt = fold_config[sample]
-    if type(folds_opt) == int:
-        return f'n_{n_folds}_f_{folds_opt}'
-    elif type(folds_opt) == list:
-        return f'n_{n_folds}_f_{"_".join([str(f) for f in folds_opt])}'
-    else:
-        raise ValueError("Invalid fold selection option with type {}".format(type(folds_opt)))
-
-#Return the index and checkpoint of the last epoch.
-def get_last_epoch(config, max_ep = -1, device = None):
-    last_epoch = -1
-    checkpoint = None
-    if max_ep < 0:
-        max_ep = config['Training']['epochs']
-    for ep in range(max_ep):
-        if os.path.exists(os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')):
-            last_epoch = ep
-        else:
-            print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}')
-            print('File not found: ', os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt'))
-            break
-    if last_epoch >= 0:
-        checkpoint = torch.load(os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device)
-    return last_epoch, checkpoint
-
-#Return the index and checkpoint of the last epoch.
-def get_specific_epoch(config, target_epoch, device = None, from_ryan = False):
-    last_epoch = -1
-    checkpoint = None
-    for ep in range(target_epoch + 1):
-        if (from_ryan):
-            if os.path.exists(os.path.join('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + config['Training_Directory'], f'model_epoch_{ep}.pt')):
-                last_epoch = ep
-            else:
-                print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}')
-                print('File not found: ', os.path.join('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + config['Training_Directory'], f'model_epoch_{ep}.pt'))
-                break
-        else:
-            if os.path.exists(os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')):
-                last_epoch = ep
-            else:
-                print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}')
-                print('File not found: ', os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt'))
-                break
-    if last_epoch >= 0:
-        if (from_ryan):
-            checkpoint = torch.load('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device)
-        else:
-            checkpoint = torch.load(os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device)
-    return last_epoch, checkpoint
-
-#Return the index and checkpoint of the nest epoch.
-def get_best_epoch(config, var='Test_AUC', mode='max', device=None, from_ryan=False):
-    # Read the training log
-    log = read_log(config)
-    
-    # Ensure the specified variable exists in the log
-    if var not in log:
-        raise ValueError(f"Variable '{var}' not found in the training log.")
-    
-    # Determine the target epoch based on the mode ('max' or 'min')
-    if mode == 'max':
-        target_epoch = int(np.argmax(log[var]))
-        print(f"Best epoch based on '{var}' (max): {target_epoch} with value: {log[var][target_epoch]}")
-    elif mode == 'min':
-        target_epoch = int(np.argmin(log[var]))
-        print(f"Best epoch based on '{var}' (min): {target_epoch} with value: {log[var][target_epoch]}")
-    else:
-        raise ValueError(f"Invalid mode '{mode}'. Expected 'max' or 'min'.")
-    
-    # Initialize checkpoint retrieval variables
-    last_epoch = -1
-    checkpoint = None
-
-    # Iterate through epochs up to the target epoch to find the corresponding checkpoint
-    for ep in range(target_epoch + 1):
-        if from_ryan:
-            checkpoint_path = os.path.join(
-                '/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/',
-                config['Training_Directory'],
-                f'model_epoch_{ep}.pt'
-            )
-        else:
-            checkpoint_path = os.path.join(
-                config['Training_Directory'],
-                f'model_epoch_{ep}.pt'
-            )
-        
-        if os.path.exists(checkpoint_path):
-            last_epoch = ep
-        else:
-            print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}')
-            print('File not found: ', checkpoint_path)
-            break
-
-    # Load the checkpoint for the last valid epoch
-    if last_epoch >= 0:
-        if from_ryan:
-            checkpoint_path = os.path.join(
-                '/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/',
-                config['Training_Directory'],
-                f'model_epoch_{last_epoch}.pt'
-            )
-        else:
-            checkpoint_path = os.path.join(
-                config['Training_Directory'],
-                f'model_epoch_{last_epoch}.pt'
-            )
-        
-        checkpoint = torch.load(checkpoint_path, map_location=device)
-    
-    return last_epoch, checkpoint
-
-def read_log(config):
-    lines = []
-    with open(config['Training_Directory'] + '/training.log', 'r') as f:
-        lines = f.readlines()
-    lines = [l for l in lines if 'Epoch' in l]
-    
-    labels = []
-    for field in lines[0].split('|'):
-        labels.append(field.split()[0])
-    
-    # Initialize log as a dictionary with empty lists
-    log = {label: [] for label in labels}
-    
-    for line in lines:
-        valid_row = True  # Flag to check if the row is valid
-        temp_row = {}  # Temporary row to store values before adding to log
-        
-        for field in line.split('|'):
-            spl = field.split()
-            try:
-                temp_row[spl[0]] = float(spl[1])
-            except (ValueError, IndexError):
-                valid_row = False  # Mark row as invalid if conversion fails
-                break
-        
-        if valid_row:  # Only add the row if all fields are valid
-            for label in labels:
-                log[label].append(temp_row.get(label, np.nan))  # Handle missing labels gracefully
-    
-    # Convert lists to numpy arrays for consistency
-    for label in labels:
-        log[label] = np.array(log[label])
-    
-    return log
-
-#Plot training logs.
-def plot_log(log, output_file):
-    fig, ax = plt.subplots(2, 2, figsize=(10,10))
-    #Time
-
-    ax[0][0].plot(log['Epoch'], np.cumsum(log['Time']), label='Time')
-    ax[0][0].set_xlabel('Epoch')
-    ax[0][0].set_ylabel('Time (s)')
-    ax[0][0].legend()
-
-    """
-    ax[0][0].plot(log['Epoch'], log['LR'], label='Learning Rate')
-    ax[0][0].set_xlabel('Epoch')
-    ax[0][0].set_ylabel('Learning Rate')
-    ax[0][0].set_yscale('log')
-    ax[0][0].legend()
-    """
-
-    #Loss
-    ax[0][1].plot(log['Epoch'], log['Loss'], label='Train Loss')
-    ax[0][1].plot(log['Epoch'], log['Test_Loss'], label='Test Loss')
-    ax[0][1].set_xlabel('Epoch')
-    ax[0][1].set_ylabel('Loss')
-    ax[0][1].legend()
-
-    #Accuracy
-    ax[1][0].plot(log['Epoch'], log['Accuracy'], label='Test Accuracy')
-    ax[1][0].set_xlabel('Epoch')
-    ax[1][0].set_ylabel('Accuracy')
-    ax[1][0].set_ylim((0.44, 0.56))
-    ax[1][0].legend()
-
-    #AUC
-    ax[1][1].plot(log['Epoch'], log['Test_AUC'], label='Test AUC')
-    ax[1][1].set_xlabel('Epoch')
-    ax[1][1].set_ylabel('AUC')
-    ax[1][1].legend()
-
-    fig.savefig(output_file)
-
-class EarlyStop():
-    def __init__(self, patience=15, threshold=1e-8, mode='min'):
-        self.patience = patience
-        self.threshold = threshold
-        self.mode = mode
-        self.count = 0
-        self.current_best = np.inf if mode == 'min' else -np.inf
-        self.should_stop = False
-
-    def update(self, value):
-        if self.mode == 'min':  # Minimizing loss
-            if value < self.current_best - self.threshold:
-                self.current_best = value
-                self.count = 0
-            else:
-                self.count += 1
-        elif self.mode == 'max':  # Maximizing metric
-            if value > self.current_best + self.threshold:
-                self.current_best = value
-                self.count = 0
-            else:
-                self.count += 1
-
-        # Check if patience is exceeded
-        if self.count >= self.patience:
-            self.should_stop = True
-
-    def reset(self):
-        self.count = 0
-        self.current_best = np.inf if self.mode == 'min' else -np.inf
-        self.should_stop = False
-
-    def to_str(self):
-        status = (
-            f"EarlyStop Status:\n"
-            f"  Mode: {'Minimize' if self.mode == 'min' else 'Maximize'}\n"
-            f"  Patience: {self.patience}\n"
-            f"  Threshold: {self.threshold:.3e}\n"
-            f"  Current Best: {self.current_best:.6f}\n"
-            f"  Consecutive Epochs Without Improvement: {self.count}\n"
-            f"  Stopping Triggered: {'Yes' if self.should_stop else 'No'}"
-        )
-        return status
-
-    def to_dict(self):
-
-        return {
-            'patience': self.patience,
-            'threshold': self.threshold,
-            'mode': self.mode,
-            'count': self.count,
-            'current_best': self.current_best,
-            'should_stop': self.should_stop,
-        }
-    
-    @classmethod
-    def load_from_dict(cls, state_dict):
-        instance = cls(
-            patience=state_dict['patience'],
-            threshold=state_dict['threshold'],
-            mode=state_dict['mode']
-        )
-        instance.count = state_dict['count']
-        instance.current_best = state_dict['current_best']
-        instance.should_stop = state_dict['should_stop']
-        return instance
-    
-
-def graph_augmentation(graph):
-    print("Augmenting Graph")
-    return

nemo/scripts/check_dataset_files.py

@@ -1,130 +0,0 @@
-import yaml
-import os
-import subprocess
-import argparse
-
-def check_dataset_files(yaml_file, rerun=False):
-    """
-    Check if all required .bin files exist for each dataset in the YAML file.
-    """
-    try:
-        # Open and parse the YAML file
-        with open(yaml_file, 'r') as file:
-            config = yaml.safe_load(file)
-
-        # Check if 'Datasets' exists in the YAML file
-        if 'Datasets' not in config:
-            print(f"No 'Datasets' section found in {yaml_file}.")
-            return
-
-        datasets = config['Datasets']
-        all_files_exist = True
-
-        for dataset_name, dataset_config in datasets.items():
-            # Extract required information
-            save_dir = dataset_config['args']['save_dir']
-            chunks = dataset_config['args']['chunks']
-            folding = dataset_config.get('folding', {})
-            n_folds = folding.get('n_folds', 0)
-            test_folds = folding.get('test', [])
-            train_folds = folding.get('train', [])
-
-            print(f"\n== Checking dataset: {dataset_name} ==")
-            print(f"  save_dir: {save_dir}")
-            print(f"  chunks: {chunks}")
-            print(f"  n_folds: {n_folds}")
-            print(f"  test_folds: {test_folds}")
-            print(f"  train_folds: {train_folds}")
-
-            missing_files = []
-
-            # 1. Check for chunk files
-            for chunk in range(chunks):
-                chunk_file = os.path.join(save_dir, f"{dataset_name}_{chunk}.bin")
-                if not os.path.exists(chunk_file):
-                    missing_files.append(chunk_file)
-
-            # 2. Check for prebatched fold files (test and train)
-            #    Naming: dataset_name_prebatched_padded_{fold}_n_{n_folds}_f_{foldlist}.bin
-            fold_types = [('test', test_folds), ('train', train_folds)]
-            for fold_type, folds in fold_types:
-                if not folds:
-                    continue
-                foldlist_str = '_'.join(map(str, folds))
-                for i in range(chunks):
-                    prebatched_file = os.path.join(
-                        save_dir,
-                        f"{dataset_name}_prebatched_padded_{i}_n_{n_folds}_f_{foldlist_str}.bin"
-                    )
-                    if not os.path.exists(prebatched_file):
-                        missing_files.append(prebatched_file)
-
-            # Print results for the current dataset
-            if missing_files:
-                all_files_exist = False
-                print(f"  Missing files for dataset '{dataset_name}':")
-                for missing_file in missing_files:
-                    print(f"    - {missing_file}")
-
-                # Optionally rerun data prep
-                if rerun:
-                    print(f"  Reprocessing dataset '{dataset_name}' ...")
-                    prep_command = f"bash/prep_data.sh {yaml_file} {dataset_name} {chunks}"
-                    try:
-                        subprocess.run(prep_command, shell=True, check=True)
-                    except subprocess.CalledProcessError as e:
-                        print(f"  Could NOT reprocess '{dataset_name}': {e}")
-            else:
-                print(f"  All files exist for dataset '{dataset_name}'.")
-
-        # Final summary
-        if all_files_exist:
-            print("\nAll required files exist for all datasets.")
-        else:
-            print("\nSome files are missing.")
-
-    except Exception as e:
-        print(f"Error processing {yaml_file}: {e}")
-
-def main(pargs):
-    # Base directory containing the YAML files
-    base_directory = os.getcwd() + "/configs/"
-
-    if pargs.configs:
-        configs = [p.strip() for p in pargs.configs.split(',')]
-    else:
-        configs = [
-            "attention/ttH_CP_even_vs_odd.yaml",
-
-            "stats_100K/finetuning_ttH_CP_even_vs_odd.yaml",
-            "stats_100K/pretraining_multiclass.yaml",
-            "stats_100K/ttH_CP_even_vs_odd.yaml",
-
-            "stats_all/finetuning_ttH_CP_even_vs_odd.yaml",
-            "stats_all/pretraining_multiclass.yaml",
-            "stats_all/ttH_CP_even_vs_odd.yaml",
-        ]
-
-    for config in configs:
-        yaml_file = os.path.join(base_directory, config)
-        if os.path.exists(yaml_file):
-            print(f"\nProcessing file: {config}")
-            check_dataset_files(yaml_file, pargs.rerun)
-        else:
-            print(f"File not found: {yaml_file}")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Check YAML config files")
-    parser.add_argument(
-        "--configs", "-c",
-        type=str,
-        required=False,
-        help="Comma-separated list of YAML config paths relative to base directory"
-    )
-    parser.add_argument(
-        "--rerun", "-r",
-        action='store_true',   # Correct way for a boolean flag
-        help="Automatically re-run data processing to fix missing files"
-    )
-    args = parser.parse_args()
-    main(args)

nemo/scripts/find_free_port.py

@@ -1,12 +0,0 @@
-# find_free_port.py
-def find_free_port():
-    import socket
-    from contextlib import closing
-
-    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
-        s.bind(('', 0))
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return str(s.getsockname()[1])
-
-if __name__ == "__main__":
-    print(find_free_port())

nemo/scripts/inference.py

@@ -1,289 +0,0 @@
-import sys
-import os
-file_path = os.getcwd()
-sys.path.append(file_path)
-import os
-import argparse
-import yaml
-import gc
-
-import torch
-import dgl
-from dgl.data import DGLDataset
-from dgl.dataloading import GraphDataLoader
-from torch.utils.data import SubsetRandomSampler, SequentialSampler
-
-class CustomPreBatchedDataset(DGLDataset):
-    def __init__(self, start_dataset, batch_size, chunkno=0, chunks=1, mask_fn=None, drop_last=False, shuffle=False, **kwargs):
-        self.start_dataset = start_dataset
-        self.batch_size = batch_size
-        self.mask_fn = mask_fn or (lambda x: torch.ones(len(x), dtype=torch.bool))
-        self.drop_last = drop_last
-        self.shuffle = shuffle
-        self.chunkno = chunkno
-        self.chunks = chunks
-        super().__init__(name=start_dataset.name + '_custom_prebatched', save_dir=start_dataset.save_dir)
-
-    def process(self):
-        mask = self.mask_fn(self.start_dataset)
-        indices = torch.arange(len(self.start_dataset))[mask]
-        print(f"Number of elements after masking: {len(indices)}")  # Debugging print
-
-        # --- CHUNK SPLITTING ---
-        total = len(indices)
-        if self.chunks == 1:
-            chunk_indices = indices
-            print(f"Chunks=1, using all {total} indices.")
-        else:
-            chunk_size = (total + self.chunks - 1) // self.chunks
-            start = self.chunkno * chunk_size
-            end = min((self.chunkno + 1) * chunk_size, total)
-            chunk_indices = indices[start:end]
-            print(f"Working on chunk {self.chunkno}/{self.chunks}: indices {start}:{end} (total {len(chunk_indices)})")
-
-        if self.shuffle:
-            sampler = SubsetRandomSampler(chunk_indices)
-        else:
-            sampler = SequentialSampler(chunk_indices)
-
-        self.dataloader = GraphDataLoader(
-            self.start_dataset,
-            sampler=sampler,
-            batch_size=self.batch_size,
-            drop_last=self.drop_last
-        )
-
-    def __getitem__(self, idx):
-        if isinstance(idx, int):
-            idx = [idx]
-        sampler = SequentialSampler(idx)
-        dloader = GraphDataLoader(self.start_dataset, sampler=sampler, batch_size=self.batch_size, drop_last=False)
-        return next(iter(dloader))
-
-    def __len__(self):
-        mask = self.mask_fn(self.start_dataset)
-        indices = torch.arange(len(self.start_dataset))[mask]
-        total = len(indices)
-        if self.chunks == 1:
-            return total
-        chunk_size = (total + self.chunks - 1) // self.chunks
-        start = self.chunkno * chunk_size
-        end = min((self.chunkno + 1) * chunk_size, total)
-        return end - start
-
-def include_config(conf):
-    if 'include' in conf:
-        for i in conf['include']:
-            with open(i) as f:
-                conf.update(yaml.load(f, Loader=yaml.FullLoader))
-        del conf['include']
-
-def load_config(config_file):
-    with open(config_file) as f:
-        conf = yaml.load(f, Loader=yaml.FullLoader)
-    include_config(conf)
-    return conf
-
-def main():
-
-    parser = argparse.ArgumentParser()
-    add_arg = parser.add_argument
-    add_arg('--config', type=str, nargs='+', required=True, help="List of config files")
-    add_arg('--target', type=str, required=True)
-    add_arg('--destination', type=str, default='')
-    add_arg('--chunkno', type=int, default=0)
-    add_arg('--chunks', type=int, default=1)
-    add_arg('--write', action='store_true')
-    add_arg('--ckpt', type=int, default=-1)
-    add_arg('--var', type=str, default='Test_AUC')
-    add_arg('--mode', type=str, default='max')
-    add_arg('--clobber', action='store_true')
-    add_arg('--tree', type=str, default='')
-    add_arg('--branch_name', type=str, nargs='+', required=True, help="List of branch names corresponding to configs")
-    args = parser.parse_args()
-
-    if(len(args.config) != len(args.branch_name)):
-        print(f"configs and branch names do not match")
-        return
-
-    config = load_config(args.config[0])
-
-    # --- OUTPUT DESTINATION LOGIC ---
-    if args.destination == '':
-        base_dest = os.path.join(config['Training_Directory'], 'inference/', os.path.split(args.target)[1])
-    else:
-        base_dest = args.destination
-
-    base_dest = base_dest.replace('.root', '').replace('.npz', '')
-    if args.chunks > 1:
-        chunked_dest = f"{base_dest}_chunk{args.chunkno}"
-    else:
-        chunked_dest = base_dest
-    chunked_dest += '.root' if args.write else '.npz'
-    args.destination = chunked_dest
-
-    # --- FILE EXISTENCE CHECK ---
-    if os.path.exists(args.destination):
-        print(f'File {args.destination} already exists.')
-        if args.clobber:
-            print('Clobbering.')
-        else:
-            print('Exiting.')
-            return
-    else:
-        print(f'Writing to {args.destination}')
-
-    import time
-    start  = time.time()
-    import torch
-    from array import array
-    import numpy as np
-    from root_gnn_base import batched_dataset as dataset
-    from root_gnn_base import utils
-    end = time.time()
-    print('Imports finished in {:.2f} seconds'.format(end - start))
-
-    start = time.time()
-    dset_config = config['Datasets'][list(config['Datasets'].keys())[0]]
-    if dset_config['class'] == 'LazyDataset':
-        dset_config['class'] = 'EdgeDataset'
-    elif dset_config['class'] == 'LazyMultiLabelDataset':
-        dset_config['class'] = 'MultiLabelDataset'
-    elif dset_config['class'] == 'PhotonIDDataset':
-        dset_config['class'] = 'UnlazyPhotonIDDataset'
-    elif dset_config['class'] == 'kNNDataset':
-        dset_config['class'] = 'UnlazyKNNDataset'
-    dset_config['args']['raw_dir'] = os.path.split(args.target)[0]
-    dset_config['args']['file_names'] = os.path.split(args.target)[1]
-    dset_config['args']['save'] = False
-    dset_config['args']['chunks'] = args.chunks
-    dset_config['args']['process_chunks'] = [args.chunkno,]
-    dset_config['args']['selections'] = []
-
-    dset_config['args']['save_dir'] = os.path.dirname(args.destination)
-
-    if args.tree != '':
-        dset_config['args']['tree_name'] = args.tree
-
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-    dstart = time.time()
-    dset = utils.buildFromConfig(dset_config)
-    dend = time.time()
-    print('Dataset finished in {:.2f} seconds'.format(dend - dstart))
-
-    print(dset)
-
-    batch_size = config['Training']['batch_size']
-    lstart = time.time()
-    loader = CustomPreBatchedDataset(
-        dset,
-        batch_size,
-        chunkno=args.chunkno,
-        chunks=args.chunks
-    )
-    loader.process()
-    lend = time.time()
-    print('Loader finished in {:.2f} seconds'.format(lend - lstart))
-    sample_graph, _, _, global_sample = loader[0]
-
-    print('dset length =', len(dset))
-    print('loader length =', len(loader))
-
-    all_scores = {}
-    all_labels = {}
-    all_tracking = {}
-    with torch.no_grad():
-        for config_file, branch in zip(args.config, args.branch_name):
-            config = load_config(config_file)
-            model = utils.buildFromConfig(config['Model'], {'sample_graph' : sample_graph, 'sample_global': global_sample}).to(device)
-            if args.ckpt < 0:
-                ep, checkpoint = utils.get_best_epoch(config, var=args.var, mode='max', device=device)
-            else:
-                ep, checkpoint = utils.get_specific_epoch(config, args.ckpt, device=device)
-            # Remove distributed/compiled prefixes if present
-            mds_copy = {}
-            for key in checkpoint['model_state_dict'].keys():
-                newkey = key.replace('module.', '')
-                newkey = newkey.replace('_orig_mod.', '')
-                mds_copy[newkey] = checkpoint['model_state_dict'][key]
-            model.load_state_dict(mds_copy)
-            model.eval()
-
-            end = time.time()
-            print('Model and dataset finished in {:.2f} seconds'.format(end - start))
-            print('Starting inference')
-            start = time.time()
-
-            finish_fn = torch.nn.Sigmoid()
-            if 'Loss' in config:
-                finish_fn = utils.buildFromConfig(config['Loss']['finish'])
-
-            scores = []
-            labels = []
-            tracking_info = []
-            ibatch = 0
-
-            for batch, label, track, globals in loader.dataloader:
-                batch = batch.to(device)
-                pred = model(batch, globals.to(device))
-                ibatch += 1
-                if (finish_fn.__class__.__name__ == "ContrastiveClusterFinish"):
-                    scores.append(pred.detach().cpu().numpy())
-                else:
-                    scores.append(finish_fn(pred).detach().cpu().numpy())
-                labels.append(label.detach().cpu().numpy())
-                tracking_info.append(track.detach().cpu().numpy())
-
-            score_size = scores[0].shape[1] if len(scores[0].shape) > 1 else 1
-            scores = np.concatenate(scores)
-            labels = np.concatenate(labels)
-            tracking_info = np.concatenate(tracking_info)
-            end = time.time()
-
-            print('Inference finished in {:.2f} seconds'.format(end - start))
-            all_scores[branch] = scores
-            all_labels[branch] = labels
-            all_tracking[branch] = tracking_info
-
-
-    if args.write:
-        import uproot
-        import awkward as ak
-
-        # Open the original ROOT file and get the tree
-        infile = uproot.open(args.target)
-        tree = infile[dset_config['args']['tree_name']]
-
-        # Read the original tree as an awkward array
-        original_data = tree.arrays(library="ak")
-
-        # Prepare new branches as dicts of arrays
-        new_branches = {}
-        n_entries = len(original_data)
-        for branch, scores in all_scores.items():
-            # Ensure the scores array is the right length
-            scores = np.asarray(scores)
-            if scores.shape[0] != n_entries:
-                raise ValueError(f"Branch '{branch}' has {scores.shape[0]} entries, but tree has {n_entries}")
-            new_branches[branch] = scores
-
-        # Merge all arrays (original + new branches)
-        # Convert awkward to dict of numpy arrays for uproot
-        out_dict = {k: np.asarray(v) for k, v in ak.to_numpy(original_data).items()}
-        out_dict.update(new_branches)
-
-        # Write to new ROOT file
-        os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
-        with uproot.recreate(args.destination) as outfile:
-            outfile.mktree(dset_config['args']['tree_name'], {k: v.dtype for k, v in out_dict.items()})
-            outfile[dset_config['args']['tree_name']].extend(out_dict)
-
-        print(f"Wrote new ROOT file {args.destination} with new branches {list(new_branches.keys())}")
-
-    else:
-        os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
-        np.savez(args.destination, scores=all_scores, labels=all_labels, tracking_info=all_tracking)
-
-if __name__ == '__main__':
-    main()

nemo/scripts/prep_data.py

@@ -1,44 +0,0 @@
-import sys
-import os
-file_path = os.getcwd()
-sys.path.append(file_path)
-
-import root_gnn_base.utils as utils
-import argparse
-from root_gnn_base.batched_dataset import PreBatchedDataset
-from root_gnn_base.batched_dataset import LazyPreBatchedDataset
-
-def main():
-    parser = argparse.ArgumentParser()
-    add_arg = parser.add_argument
-    add_arg('--config', type=str, required=True)
-    add_arg('--dataset', type=str, required=True)
-    add_arg('--chunk', type=int, default=0)
-    add_arg('--shuffle_mode', action='store_true', help='Shuffle the dataset before training.')
-    add_arg('--drop_last', action='store_false', help='Set drop_last to False if the flag is provided. Defaults to True.')
-    args = parser.parse_args()
-
-    config = utils.load_config(args.config)
-    dset_config = config['Datasets'][args.dataset]
-    batch_size = config['Training']['batch_size']
-    if not args.shuffle_mode:
-        dset = utils.buildFromConfig(dset_config, {'process_chunks': [args.chunk,]})
-    else:
-        dset = utils.buildFromConfig(dset_config)
-        if 'batch_size' in dset_config:
-            batch_size = dset_config['batch_size']
-
-        shuffle_chunks = dset_config.get('shuffle_chunks', 10)
-        padding_mode = dset_config.get('padding_mode', 'STEPS')
-        fold_conf = dset_config["folding"]
-        print(f"shuffle_chunks = {shuffle_chunks}, args.chunk = {args.chunk}, padding_mode = {padding_mode}")
-        if dset_config["class"] == "LazyMultiLabelDataset":
-            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'] )
-            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'])
-
-        else:
-            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'])
-            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'] )
-
-if __name__ == "__main__":
-    main()

nemo/scripts/training_script.py

@@ -1,463 +0,0 @@
-import argparse
-import time
-import datetime
-import yaml
-import os
-import glob
-
-start_time = time.time()
-
-import dgl
-import torch
-import torch.nn as nn
-
-import sys
-file_path = os.getcwd()
-sys.path.append(file_path)
-import root_gnn_base.batched_dataset as datasets
-from root_gnn_base import utils
-import root_gnn_base.custom_scheduler as lr_utils
-from models import GCN
-
-import numpy as np
-from sklearn.metrics import roc_auc_score
-import resource
-import gc
-
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from torch.utils.data.distributed import DistributedSampler
-from torch.nn.parallel import DistributedDataParallel as DDP
-
-from physicsnemo.models.module import Module
-from physicsnemo.models.meta import ModelMetaData
-from dataclasses import dataclass
-
-print("import time: {:.4f} s".format(time.time() - start_time))
-
-def mem():
-    print(f'Current memory usage: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 / 1024:.2f} GB')
-
-def gpu_mem():
-    print()
-    print('GPU Memory Usage:')
-    print(f'Current GPU memory usage: {torch.cuda.memory_allocated() / 1024 / 1024 / 1024:.2f} GB')
-    print(f'Current GPU cache usage: {torch.cuda.memory_cached() / 1024 / 1024 / 1024:.2f} GB')
-    print(f'Current GPU max memory usage: {torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024:.2f} GB')
-    print(f'Current GPU max cache usage: {torch.cuda.memory_reserved() / 1024 / 1024 / 1024:.2f} GB')
-    mem()
-
-def train(train_loaders, test_loaders, model, device, config, args):
-    restart = args.restart
-
-    if ('Loss' in config):
-        loss_fcn = utils.buildFromConfig(config['Loss'], {'reduction':'none'})
-        finish_fn = utils.buildFromConfig(config['Loss']['finish'])
-    else:
-        loss_fcn = torch.nn.BCEWithLogitsLoss(reduction='none')
-        finish_fn = torch.nn.Sigmoid()
-    optimizer = torch.optim.Adam(model.parameters(), lr=config['Training']['learning_rate'])
-    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = config['Training']['gamma']) 
-
-    early_termination = utils.EarlyStop()
-    if 'early_termination' in config['Training']:
-        early_termination.patience = config['Training']['early_termination']['patience']
-        early_termination.threshold = config['Training']['early_termination']['threshold']
-        early_termination.mode = config['Training']['early_termination']['mode']
-
-    starting_epoch = 0
-    if not restart:
-        last_ep, checkpoint = utils.get_last_epoch(config)  
-        if (last_ep >= 0): 
-            model.load_state_dict(checkpoint['model_state_dict'])
-            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-            starting_epoch = checkpoint['epoch'] + 1
-            early_termination = utils.EarlyStop.load_from_dict(checkpoint['early_stop'])
-            if early_termination.should_stop:
-                print(f"Early Termination at Epoch {last_ep}")
-                return
-            print(f"Loaded epoch {checkpoint['epoch']} from checkpoint")
-        else:
-            print("Starting from scratch")
-        log = open(config['Training_Directory'] + '/training.log', 'a', buffering=1)
-    else:
-        # Remove all *.pt and *.npz files in the Training_Directory
-        for pattern in ('*.pt', '*.npz'):
-            for file in glob.glob(os.path.join(config['Training_Directory'], pattern)):
-                os.remove(file)
-        log = open(config['Training_Directory'] + '/training.log', 'w', buffering=1)
-
-    train_cyclers = []
-    for loader in train_loaders:
-        train_cyclers.append(utils.cycler((loader)))
-
-    log.write(f'Training {config["Training_Name"]} {datetime.datetime.now()} \n')
-    print(f"Starting training for {config['Training']['epochs']} epochs")
-
-    if hasattr(train_loaders[0].dataset, 'padding_mode'):
-        is_padded = train_loaders[0].dataset.padding_mode != 'NONE'
-        if (train_loaders[0].dataset.padding_mode == 'NODE'):
-            is_padded = False
-    else:
-        is_padded = False
-
-    lr_utils.print_LR(optimizer)
-
-    # training loop
-    for epoch in range(starting_epoch, config['Training']['epochs']):
-        start = time.time()
-        run = start
-        if (args.profile):
-            if (epoch == 0):
-                torch.cuda.cudart().cudaProfilerStart()
-            torch.cuda.nvtx.range_push("Epoch Start")
-
-        # training
-
-        model.train()
-
-        ibatch = 0
-        total_loss = 0
-        for batched_graph, labels, _, global_feats in train_loaders[0]:
-            batch_start = time.time()
-            logits = torch.tensor([])
-            tlabels = torch.tensor([])
-            weights = torch.tensor([])
-            batch_lengths = []
-            for cycler in train_cyclers:
-                graph, label, track, global_feats = next(cycler)
-                graph = graph.to(device)
-                label = label.to(device)
-                track = track.to(device)
-                global_feats = global_feats.to(device)
-                if is_padded: #Padding the globals to match padded graphs.
-                    global_feats = torch.concatenate((global_feats, torch.zeros(1, len(global_feats[0])).to(device)))
-                load = time.time()
-                if (args.profile):
-                    torch.cuda.nvtx.range_push("Model Forward")
-                if (len(logits) == 0):
-                    logits = model(graph.ndata['features'], graph.edata['features'], graph)
-                    tlabels = label
-                    weights = track[:,1]
-                else:
-                    logits = torch.concatenate((logits, model(graph.ndata['features'], graph.edata['features'], graph)), dim=0)
-                    tlabels = torch.concatenate((tlabels, label), dim=0)
-                    weights = torch.concatenate((weights, track[:,1]), dim=0)
-                batch_lengths.append(logits.shape[0] - 1)
-
-                if (args.profile):
-                    torch.cuda.nvtx.range_pop() # popping model forward
-
-            if is_padded:
-                keepmask = torch.full_like(logits[:,0], True, dtype=torch.bool)
-                keepmask[batch_lengths] = False
-                logits = logits[keepmask]
-            tlabels = tlabels.to(torch.float)
-            if logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'BCEWithLogitsLoss':
-                logits = logits[:,0]
-                tlabels = tlabels.to(torch.float)
-            if loss_fcn.__class__.__name__ == 'CrossEntropyLoss':
-                tlabels = tlabels.to(torch.long)
-
-            if args.abs:
-                weights = torch.abs(weights)
-
-            loss = loss_fcn(logits, tlabels.to(device))
-            # Normalize loss within each label
-            unique_labels = torch.unique(tlabels)  # Get unique labels
-            normalized_loss = 0.0
-            for label in unique_labels:
-                # Mask for samples belonging to the current label
-                label_mask = (tlabels == label)
-                # Extract weights and losses for the current label
-                label_weights = weights[label_mask]
-                label_losses = loss[label_mask]
-                # Compute normalized loss for the current label
-                label_loss = torch.sum(label_weights * label_losses) / torch.sum(label_weights)
-                # Add to the total normalized loss
-                normalized_loss += label_loss
-            loss = normalized_loss / len(unique_labels)
-
-            if (args.profile):
-                torch.cuda.nvtx.range_push("Model Backward")
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            total_loss += loss.detach().cpu().item()
-
-            if (args.profile):
-                torch.cuda.nvtx.range_pop() # pop model backward
-
-            ibatch += 1
-
-            if ibatch % 1000 == 0:
-                print(f'Batch {ibatch} out of {len(train_loaders[0])}', end='\r')
-                # gpu_mem()
-        else:
-            print("Epoch Done.")
-       
-       # validation
-
-        scores = []
-        labels = []
-        weights = []
-        model.eval()
-
-        if (args.profile):
-            torch.cuda.nvtx.range_push("Model Evaluation")
-
-        with torch.no_grad():
-            for loader in test_loaders:
-                for batch, label, track, global_feats in loader:
-                    #Don't use compiled model for testing since we can't control the batch size.
-                    #We could before, but it assumes each dataset has the same number of batches...
-                    if is_padded:
-                        global_feats = torch.cat([global_feats, torch.zeros(1, len(global_feats[0]))])
-
-                    # batch_scores = model(batch.to(device), global_feats.to(device))
-                    batch_scores = model(graph.ndata['features'].to(device), graph.edata['features'].to(device), graph)
-                    
-                    if is_padded:
-                        scores.append(batch_scores[:-1,:])
-                    else:
-                        scores.append(batch_scores)
-                    labels.append(label)
-                    weights.append(track[:,1])
-
-        if (args.profile):
-            torch.cuda.nvtx.range_pop() # pop evaluation
-
-        if scores == []: #If validation set is empty.
-            continue
-
-        logits = torch.concatenate(scores).to(device)
-        labels = torch.concatenate(labels).to(device)
-        weights = torch.concatenate(weights).to(device)
-
-        wgt_mask = weights > 0
-
-        if args.abs:
-            weights = torch.abs(weights)
-
-        print(f"Num batches trained = {ibatch}")
-
-        if (loss_fcn.__class__.__name__ == "ContrastiveClusterLoss"):
-            scores = logits
-            preds = scores
-            accuracy = 0
-            test_auc = 0
-            acc = 0
-            contrastive_cluster_loss = finish_fn(logits)
-
-        elif (loss_fcn.__class__.__name__ == "MultiLabelLoss"):
-            scores = finish_fn(logits)
-            preds = torch.round(scores)
-            multilabel_accuracy = []
-            threshold = 0.1
-            for i in range(len(labels[0])):
-                multilabel_accuracy.append(torch.sum(preds[:, i].to("cpu") == labels[:, i].to("cpu")) / len(labels))
-            test_auc = 0
-            acc = np.mean(multilabel_accuracy)
-
-        elif logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'BCEWithLogitsLoss':
-            test_auc = 0
-            acc = 0
-            logits = logits[:,0]
-            scores = finish_fn(logits)
-            labels =labels.to(torch.float)
-            preds = scores > 0.5
-            test_auc = roc_auc_score(labels[wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), sample_weight=weights[wgt_mask].to("cpu"))
-            acc = torch.sum(preds.to("cpu") == labels.to("cpu")) / len(labels)
-
-        elif logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'MSELoss':
-            logits = logits[:,0]
-            scores = finish_fn(logits)
-            labels = labels.to(torch.float)
-            acc = 0
-            test_auc = 0
-
-        else:
-            preds = torch.argmax(logits, dim=1)
-            scores = finish_fn(logits)
-            if labels.dim() == 1: #Multi-class
-                acc = torch.sum(preds.to("cpu") == labels.to("cpu")) / len(labels) #TODO: Make each class weighted equally?
-                
-                labels = labels.to("cpu")
-                weights = weights.to("cpu")
-                logits = logits.to("cpu")
-                wgt_mask = wgt_mask.to("cpu")
-
-                labels_onehot = np.zeros((len(labels), len(scores[0])))
-                labels_onehot[np.arange(len(labels)), labels] = 1
-
-                try:
-                    #test_auc = roc_auc_score(labels[wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu"))
-                    if (len(scores[0]) != config["Model"]["args"]["out_size"]):
-                        print("ERROR: The out_size and the number of class labels don't match! Please check config.")
-                    test_auc = roc_auc_score(labels_onehot[wgt_mask], scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu"))
-                except ValueError:
-                    test_auc = np.nan
-            else: #Multi-loss
-                acc = torch.sum(preds.to("cpu") == labels[:,0].to("cpu")) / len(labels)
-                try:
-                    test_auc = roc_auc_score(labels[:,0][wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu"))
-                except ValueError:
-                    test_auc = np.nan
-
-        if (loss_fcn.__class__.__name__ == "MultiLabelLoss"):
-            multilabel_log_str = "MultiLabel_Accuracy "
-            for accuracy in multilabel_accuracy:
-                multilabel_log_str += f" | {accuracy:.4f}"
-            log.write(multilabel_log_str + '\n')
-            print(multilabel_log_str, flush=True)
-        elif (loss_fcn.__class__.__name__ == "ContrastiveClusterLoss"):
-            contrastive_cluster_log_str = "ContrastiveClusterLoss "
-            contrastive_cluster_log_str += f"Contrastive Loss: {contrastive_cluster_loss[0]:.4f}, Clustering Loss: {contrastive_cluster_loss[1]:.4f}, Variance Loss: {contrastive_cluster_loss[2]:.4f}"
-            log.write(contrastive_cluster_log_str + '\n')
-            print(contrastive_cluster_log_str, flush=True)
-
-        test_loss = loss_fcn(logits, labels)
-        # Normalize loss within each label
-        unique_labels = torch.unique(labels)  # Get unique labels
-        normalized_loss = 0.0
-        for label in unique_labels:
-            # Mask for samples belonging to the current label
-            label_mask = (labels == label)
-            # Extract weights and losses for the current label
-            label_weights = weights[label_mask]
-            label_losses = test_loss[label_mask]
-            # Compute normalized loss for the current label
-            label_loss = torch.sum(label_weights * label_losses) / torch.sum(label_weights)
-            # Add to the total normalized loss
-            normalized_loss += label_loss
-        test_loss = normalized_loss / len(unique_labels)
-
-
-        end = time.time()
-        log_str = "Epoch {:05d} | LR {:.4e} | Loss {:.4f} | Accuracy {:.4f} | Test_Loss {:.4f} | Test_AUC {:.4f} | Time {:.4f} s".format(
-                epoch, optimizer.param_groups[0]['lr'], total_loss/ibatch, acc, test_loss, test_auc, end - start
-        )
-        log.write(log_str + '\n')
-        print(log_str, flush=True)
-
-        state_dict = model.state_dict()         
-
-        torch.save({
-                'epoch': epoch,
-                'model_state_dict': state_dict,
-                'optimizer_state_dict': optimizer.state_dict(),
-                'early_stop': early_termination.to_dict()
-                }, os.path.join(config['Training_Directory'], f"model_epoch_{epoch}.pt"))
-        np.savez(os.path.join(config['Training_Directory'], f'model_epoch_{epoch}.npz'), scores=scores.to("cpu"), labels=labels.to("cpu"))
-
-        early_termination.update(test_loss)
-        if early_termination.should_stop:
-            log_str = f"Early Termination at Epoch {epoch}"
-            log.write(log_str + "\n")
-            print(log_str)
-            log_str = early_termination.to_str()
-            log.write(log_str + "\n")
-            print(log_str)
-            break
-
-        scheduler.step()
-
-        if (args.profile):
-            torch.cuda.nvtx.range_pop() # pop epoch
-            
-    log.close()
-
-def main(args=None):
-    config = utils.load_config(args.config)
-
-    if not os.path.exists(config['Training_Directory']):
-        os.makedirs(config['Training_Directory'], exist_ok=True)
-    with open(config['Training_Directory'] + '/config.yaml', 'w') as f:
-        yaml.dump(config, f)
-    batch_size = config["Training"]["batch_size"]
-    
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    train_loaders = []
-    test_loaders = []
-    val_loaders = []
-    load_start = time.time()
-
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-    # ldr_type = datasets.LazyPreBatchedDataset if args.lazy else datasets.PreBatchedDataset
-    ldr_type = datasets.LazyPreBatchedDataset
-
-    for dset_conf in config["Datasets"]:
-        dset = utils.buildFromConfig(config["Datasets"][dset_conf])
-        if 'batch_size' in config["Datasets"][dset_conf]:
-            batch_size = config["Datasets"][dset_conf]['batch_size']
-        fold_conf = config["Datasets"][dset_conf]["folding"]
-        shuffle_chunks = config["Datasets"][dset_conf].get("shuffle_chunks", 10)
-        padding_mode = config["Datasets"][dset_conf].get("padding_mode", "STEPS")
-        mask_fn = utils.fold_selection(fold_conf, "train")
-        if args.preshuffle:
-            # ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode, use_ddp = args.multigpu, rank=rank, world_size=world_size)
-            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode, hidden_size = 128)
-            gsamp, _, _, global_samp = ldr[0]
-            sampler = None
-            train_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler = sampler))
-            
-            sampler = None
-            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, padding_mode = padding_mode, hidden_size=128)
-            test_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler=sampler))
-
-            if "validation" in fold_conf:
-                val_loaders.append(torch.utils.data.DataLoader((ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=utils.fold_selection(fold_conf, "validation"), suffix = utils.fold_selection_name(fold_conf, 'validation'), chunks = shuffle_chunks, hid_size=128,  padding_mode = padding_mode)), batch_size = None, num_workers = 0, sampler = sampler))
-            else:
-                print("No validation set for dataset ", dset_conf)
-        else:
-            train_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "train")))  
-            gsamp, _, _, global_samp = dset[0]
-            test_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "test")))
-            if "validation" in fold_conf:
-                val_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "validation")))
-            else:
-                print("No validation set for dataset ", dset_conf)
-    
-    load_end = time.time()
-    print("Load time: {:.4f} s".format(load_end - load_start))
-
-    # model = utils.buildFromConfig(config["Model"], {'sample_graph': gsamp, 'sample_global': global_samp, 'seed': args.seed}).to(device)
-    
-    model = utils.buildFromConfig(config["Model"]).to(device)
-    
-    # @dataclass
-    # class MetaData(ModelMetaData):
-    #     name: str = "Edge_Network"
-
-    # physicsnemo_model = Module.from_torch(model)
-    # print(f"physicsnemo_model = {physicsnemo_model}")
-    # model = physicsnemo_model()
-    
-
-    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    print(f"Number of trainable parameters = {pytorch_total_params}")
-
-    # model training
-    print("Training...")
-    gpu_mem()
-    train(train_loaders, test_loaders, model, device, config, args)
-
-if __name__ == "__main__":
-    #Handle CLI arguments
-    parser = argparse.ArgumentParser()
-    add_arg = parser.add_argument
-    add_arg("--config", type=str, help="Config file.", required=True)
-    add_arg("--restart", action="store_true", help="Restart training from scratch.")
-    add_arg("--preshuffle", action="store_true", help="Shuffle data before training.")
-    add_arg("--seed", type=int, default=2, help="Sets random seed")
-    add_arg("--abs", action="store_true", help="Use abs value of per-event weight")
-    add_arg("--profile", action="store_true", help="use nsight systems profiler")
-
-    pargs = parser.parse_args()
-    
-    main(pargs)
-        
-        

nemo/setup/Dockerfile

@@ -1,25 +0,0 @@
-FROM nvcr.io/nvidia/physicsnemo/physicsnemo:25.06
-
-WORKDIR /global/cfs/projectdirs/atlas/joshua/GNN4Colliders
-
-LABEL maintainer.name="Joshua Ho"
-LABEL maintainer.email="ho22joshua@berkeley.edu"
-
-ENV LANG=C.UTF-8
-
-# Install system dependencies: vim, OpenMPI, and build tools
-RUN apt-get update -qq \
- && apt-get install -y --no-install-recommends \
-    wget lsb-release gnupg software-properties-common \
-    vim \
-    g++-11 gcc-11 libstdc++-11-dev \
-    openmpi-bin openmpi-common libopenmpi-dev \
- && rm -rf /var/lib/apt/lists/*
-
-# Install Python packages: mpi4py and jupyter
-RUN pip install --no-cache-dir mpi4py jupyter uproot
-
-# (Optional) Expose Jupyter port
-EXPOSE 8888
-
-

nemo/setup/build_image.sh

@@ -1,4 +0,0 @@
-tag=$1
-echo $tag
-podman-hpc build -t joshuaho/nemo:$tag --platform linux/amd64 .
-podman-hpc migrate joshuaho/nemo:$tag

nemo/setup/environment.yml

@@ -1,391 +0,0 @@
-name: dgl
-channels:
-  - pytorch
-  - dglteam/label/cu118
-  - nvidia
-  - conda-forge
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1
-  - _openmp_mutex=4.5
-  - _sysroot_linux-64_curr_repodata_hack=3
-  - afterimage=1.21
-  - anyio=3.7.1
-  - appdirs=1.4.4
-  - argon2-cffi=21.3.0
-  - argon2-cffi-bindings=21.2.0
-  - arrow=1.2.3
-  - asttokens=2.2.1
-  - async-lru=2.0.4
-  - atk-1.0=2.38.0
-  - attrs=23.1.0
-  - awkward-pandas=2023.8.0
-  - aws-c-auth=0.7.0
-  - aws-c-cal=0.6.0
-  - aws-c-common=0.8.23
-  - aws-c-compression=0.2.17
-  - aws-c-event-stream=0.3.1
-  - aws-c-http=0.7.11
-  - aws-c-io=0.13.28
-  - aws-c-mqtt=0.8.14
-  - aws-c-s3=0.3.13
-  - aws-c-sdkutils=0.1.11
-  - aws-checksums=0.1.16
-  - aws-crt-cpp=0.20.3
-  - aws-sdk-cpp=1.10.57
-  - babel=2.12.1
-  - backcall=0.2.0
-  - backports=1.0
-  - backports.functools_lru_cache=1.6.5
-  - beautifulsoup4=4.12.2
-  - binutils=2.38
-  - binutils_impl_linux-64=2.38
-  - binutils_linux-64=2.38.0
-  - blas=1.0
-  - bleach=6.0.0
-  - brotlipy=0.7.0
-  - bzip2=1.0.8
-  - c-ares=1.19.1
-  - c-compiler=1.5.2
-  - ca-certificates=2025.4.26
-  - cached-property=1.5.2
-  - cached_property=1.5.2
-  - cairo=1.16.0
-  - certifi=2024.8.30
-  - cffi=1.15.1
-  - cfitsio=4.2.0
-  - charset-normalizer=2.0.4
-  - comm=0.1.4
-  - compilers=1.5.2
-  - cryptography=41.0.2
-  - cuda-cudart=11.8.89
-  - cuda-cupti=11.8.87
-  - cuda-libraries=11.8.0
-  - cuda-nvrtc=11.8.89
-  - cuda-nvtx=11.8.86
-  - cuda-runtime=11.8.0
-  - cxx-compiler=1.5.2
-  - davix=0.8.4
-  - debugpy=1.6.8
-  - decorator=5.1.1
-  - defusedxml=0.7.1
-  - dgl=1.1.1.cu118
-  - entrypoints=0.4
-  - exceptiongroup=1.1.3
-  - executing=1.2.0
-  - expat=2.5.0
-  - ffmpeg=4.3
-  - fftw=3.3.10
-  - filelock=3.9.0
-  - flit-core=3.9.0
-  - font-ttf-dejavu-sans-mono=2.37
-  - font-ttf-inconsolata=3.000
-  - font-ttf-source-code-pro=2.038
-  - font-ttf-ubuntu=0.83
-  - fontconfig=2.14.2
-  - fonts-conda-ecosystem=1
-  - fonts-conda-forge=1
-  - fortran-compiler=1.5.2
-  - fqdn=1.5.1
-  - freetype=2.12.1
-  - fribidi=1.0.10
-  - ftgl=2.4.0
-  - gcc=11.2.0
-  - gcc_impl_linux-64=11.2.0
-  - gcc_linux-64=11.2.0
-  - gdk-pixbuf=2.42.8
-  - gettext=0.21.1
-  - gflags=2.2.2
-  - gfortran=11.2.0
-  - gfortran_impl_linux-64=11.2.0
-  - gfortran_linux-64=11.2.0
-  - giflib=5.2.1
-  - gl2ps=1.4.2
-  - glew=2.1.0
-  - glog=0.6.0
-  - gmp=6.2.1
-  - gmpy2=2.1.2
-  - gnutls=3.6.15
-  - graphite2=1.3.13
-  - graphviz=6.0.2
-  - gsl=2.7
-  - gsoap=2.8.123
-  - gtk2=2.24.33
-  - gts=0.7.6
-  - gxx=11.2.0
-  - gxx_impl_linux-64=11.2.0
-  - gxx_linux-64=11.2.0
-  - harfbuzz=7.3.0
-  - icu=72.1
-  - idna=3.4
-  - importlib-metadata=6.8.0
-  - importlib-resources=6.0.1
-  - importlib_metadata=6.8.0
-  - importlib_resources=6.0.1
-  - intel-openmp=2023.1.0
-  - ipykernel=6.25.1
-  - ipyparallel=8.6.1
-  - ipython=8.12.2
-  - isoduration=20.11.0
-  - jedi=0.19.0
-  - jinja2=3.1.2
-  - jpeg=9e
-  - json5=0.9.14
-  - jsonpointer=2.0
-  - jsonschema=4.19.0
-  - jsonschema-specifications=2023.7.1
-  - jsonschema-with-format-nongpl=4.19.0
-  - jupyter-lsp=2.2.0
-  - jupyter_client=8.3.0
-  - jupyter_core=5.3.0
-  - jupyter_events=0.7.0
-  - jupyter_server=2.7.0
-  - jupyter_server_terminals=0.4.4
-  - jupyterlab=4.0.5
-  - jupyterlab_pygments=0.2.2
-  - jupyterlab_server=2.24.0
-  - kernel-headers_linux-64=3.10.0
-  - keyutils=1.6.1
-  - krb5=1.20.1
-  - lame=3.100
-  - lcms2=2.12
-  - ld_impl_linux-64=2.38
-  - lerc=3.0
-  - libabseil=20230125.3
-  - libarrow=12.0.1
-  - libblas=3.9.0
-  - libbrotlicommon=1.0.9
-  - libbrotlidec=1.0.9
-  - libbrotlienc=1.0.9
-  - libcblas=3.9.0
-  - libcrc32c=1.1.2
-  - libcublas=11.11.3.6
-  - libcufft=10.9.0.58
-  - libcufile=1.7.1.12
-  - libcurand=10.3.3.129
-  - libcurl=8.1.2
-  - libcusolver=11.4.1.48
-  - libcusparse=11.7.5.86
-  - libcxx=15.0.7
-  - libcxxabi=15.0.7
-  - libdeflate=1.12
-  - libedit=3.1.20191231
-  - libev=4.33
-  - libevent=2.1.12
-  - libexpat=2.5.0
-  - libffi=3.4.4
-  - libgcc-devel_linux-64=11.2.0
-  - libgcc-ng=13.1.0
-  - libgd=2.3.3
-  - libgfortran-ng=11.2.0
-  - libgfortran5=11.2.0
-  - libglib=2.76.4
-  - libglu=9.0.0
-  - libgomp=13.1.0
-  - libgoogle-cloud=2.12.0
-  - libgrpc=1.56.2
-  - libiconv=1.17
-  - libidn2=2.3.4
-  - libllvm13=13.0.1
-  - libllvm14=14.0.6
-  - libnghttp2=1.52.0
-  - libnpp=11.8.0.86
-  - libnsl=2.0.0
-  - libnuma=2.0.18
-  - libnvjpeg=11.9.0.86
-  - libpng=1.6.39
-  - libprotobuf=4.23.3
-  - librsvg=2.54.4
-  - libsodium=1.0.18
-  - libsqlite=3.42.0
-  - libssh2=1.11.0
-  - libstdcxx-devel_linux-64=11.2.0
-  - libstdcxx-ng=13.1.0
-  - libtasn1=4.19.0
-  - libthrift=0.18.1
-  - libtiff=4.4.0
-  - libtool=2.4.7
-  - libunistring=0.9.10
-  - libutf8proc=2.8.0
-  - libuuid=2.38.1
-  - libwebp=1.2.4
-  - libwebp-base=1.2.4
-  - libxcb=1.15
-  - libxml2=2.10.4
-  - libzlib=1.2.13
-  - llvmlite=0.40.1
-  - lz4-c=1.9.4
-  - markupsafe=2.1.1
-  - matplotlib-inline=0.1.6
-  - metakernel=0.29.5
-  - mistune=3.0.0
-  - mkl=2023.1.0
-  - mkl-service=2.4.0
-  - mkl_fft=1.3.6
-  - mkl_random=1.2.2
-  - mpc=1.1.0
-  - mpfr=4.0.2
-  - mpmath=1.3.0
-  - nbclient=0.8.0
-  - nbconvert-core=7.7.3
-  - nbformat=5.9.2
-  - ncurses=6.4
-  - nest-asyncio=1.5.6
-  - nettle=3.7.3
-  - networkx=3.1
-  - nlohmann_json=3.11.2
-  - notebook=7.0.2
-  - notebook-shim=0.2.3
-  - numba=0.57.1
-  - numpy=1.24.3
-  - numpy-base=1.24.3
-  - openh264=2.1.1
-  - openssl=3.3.1
-  - orc=1.9.0
-  - overrides=7.4.0
-  - packaging=23.0
-  - pandas=2.0.3
-  - pandocfilters=1.5.0
-  - pango=1.50.14
-  - parso=0.8.3
-  - pcre=8.45
-  - pcre2=10.40
-  - pexpect=4.8.0
-  - pickleshare=0.7.5
-  - pillow=9.4.0
-  - pip=23.2.1
-  - pixman=0.40.0
-  - pkgutil-resolve-name=1.3.10
-  - platformdirs=2.6.0
-  - pooch=1.4.0
-  - portalocker=2.7.0
-  - prometheus_client=0.17.1
-  - prompt-toolkit=3.0.39
-  - prompt_toolkit=3.0.39
-  - psutil=5.9.0
-  - pthread-stubs=0.4
-  - ptyprocess=0.7.0
-  - pure_eval=0.2.2
-  - pyarrow=12.0.1
-  - pycparser=2.21
-  - pygments=2.16.1
-  - pyopenssl=23.2.0
-  - pysocks=1.7.1
-  - pythia8=8.309
-  - python=3.8.17
-  - python-dateutil=2.8.2
-  - python-fastjsonschema=2.18.0
-  - python-json-logger=2.0.7
-  - python-tzdata=2024.2
-  - python_abi=3.8
-  - pytorch=2.0.1
-  - pytorch-cuda=11.8
-  - pytorch-mutex=1.0
-  - pytz=2023.3
-  - pyyaml=6.0
-  - pyzmq=25.1.1
-  - rdma-core=28.9
-  - re2=2023.03.02
-  - readline=8.2
-  - referencing=0.30.2
-  - requests=2.31.0
-  - rfc3339-validator=0.1.4
-  - rfc3986-validator=0.1.1
-  - root=6.28.0
-  - root_base=6.28.0
-  - rpds-py=0.9.2
-  - s2n=1.3.46
-  - scipy=1.10.1
-  - scitokens-cpp=0.7.3
-  - send2trash=1.8.2
-  - setuptools=68.0.0
-  - six=1.16.0
-  - snappy=1.1.10
-  - sniffio=1.3.0
-  - soupsieve=2.3.2.post1
-  - sqlite=3.41.2
-  - stack_data=0.6.2
-  - sympy=1.11.1
-  - sysroot_linux-64=2.17
-  - tbb=2021.8.0
-  - terminado=0.17.1
-  - tinycss2=1.2.1
-  - tk=8.6.12
-  - tomli=2.0.1
-  - torchaudio=2.0.2
-  - torchtriton=2.0.0
-  - torchvision=0.15.2
-  - tornado=6.3.2
-  - tqdm=4.65.0
-  - traitlets=5.9.0
-  - typing_extensions=4.12.2
-  - typing_utils=0.1.0
-  - ucx=1.14.1
-  - uri-template=1.3.0
-  - urllib3=1.26.16
-  - vdt=0.4.3
-  - vector-classes=1.4.3
-  - wcwidth=0.2.6
-  - webcolors=1.13
-  - webencodings=0.5.1
-  - websocket-client=1.6.1
-  - wheel=0.38.4
-  - xorg-fixesproto=5.0
-  - xorg-kbproto=1.0.7
-  - xorg-libice=1.1.1
-  - xorg-libsm=1.2.4
-  - xorg-libx11=1.8.6
-  - xorg-libxau=1.0.11
-  - xorg-libxcursor=1.2.0
-  - xorg-libxdmcp=1.1.3
-  - xorg-libxext=1.3.4
-  - xorg-libxfixes=5.0.3
-  - xorg-libxft=2.3.8
-  - xorg-libxpm=3.5.16
-  - xorg-libxrender=0.9.11
-  - xorg-libxt=1.3.0
-  - xorg-renderproto=0.11.1
-  - xorg-xextproto=7.3.0
-  - xorg-xproto=7.0.31
-  - xrootd=5.5.4
-  - xxhash=0.8.1
-  - xz=5.2.6
-  - yaml=0.2.5
-  - zeromq=4.3.4
-  - zipp=3.16.2
-  - zlib=1.2.13
-  - zstd=1.5.2
-  - pip:
-      - awkward==2.6.4
-      - awkward-cpp==33
-      - contourpy==1.1.0
-      - cramjam==2.8.3
-      - cycler==0.11.0
-      - fonttools==4.42.0
-      - fsspec==2024.3.1
-      - h5py==3.9.0
-      - pip-install==1.3.5
-      - joblib==1.3.2
-      - kiwisolver==1.4.4
-      - matplotlib==3.7.2
-      - nvidia-cublas-cu12==12.1.3.1
-      - nvidia-cuda-cupti-cu12==12.1.105
-      - nvidia-cuda-nvrtc-cu12==12.1.105
-      - nvidia-cuda-runtime-cu12==12.1.105
-      - nvidia-cudnn-cu12==8.9.2.26
-      - nvidia-cufft-cu12==11.0.2.54
-      - nvidia-curand-cu12==10.3.2.106
-      - nvidia-cusolver-cu12==11.4.5.107
-      - nvidia-cusparse-cu12==12.1.0.106
-      - nvidia-nccl-cu12==2.20.5
-      - nvidia-nvjitlink-cu12==12.4.127
-      - nvidia-nvtx-cu12==12.1.105
-      - pyparsing==3.0.9
-      - scikit-learn==1.3.0
-      - threadpoolctl==3.2.0
-      - torch==2.3.0
-      - triton==2.3.0
-      - typing-extensions==4.11.0
-      - tzdata==2024.1
-      - uproot==5.3.7
-prefix: /global/homes/j/joshuaho/.conda/envs/dgl

nemo/setup/setup/Dockerfile

@@ -1,29 +0,0 @@
-FROM nvcr.io/nvidia/dgl:25.05-py3
-
-WORKDIR /global/cfs/projectdirs/atlas/joshua/GNN4Colliders
-
-LABEL maintainer.name="Joshua Ho"
-LABEL maintainer.email="ho22joshua@berkeley.edu"
-
-ENV LANG=C.UTF-8
-
-# Install system dependencies: vim, OpenMPI, and build tools
-RUN apt-get update -qq \
- && apt-get install -y --no-install-recommends \
-    wget lsb-release gnupg software-properties-common \
-    vim \
-    g++-11 gcc-11 libstdc++-11-dev \
-    openmpi-bin openmpi-common libopenmpi-dev \
- && rm -rf /var/lib/apt/lists/*
-
-# Install Python packages: mpi4py and jupyter
-RUN pip install --no-cache-dir mpi4py jupyter uproot
-
-RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
-
-RUN pip install  dgl -f https://data.dgl.ai/wheels/torch-2.1/cu118/repo.html
-i
-# (Optional) Expose Jupyter port
-EXPOSE 8888
-
-

nemo/setup/setup/build_image.sh

@@ -1,4 +0,0 @@
-tag=$1
-echo $tag
-podman-hpc build -t joshuaho/pytorch:$tag --platform linux/amd64 .
-podman-hpc migrate joshuaho/pytorch:$tag

nemo/setup/setup/environment.yml

@@ -1,391 +0,0 @@
-name: dgl
-channels:
-  - pytorch
-  - dglteam/label/cu118
-  - nvidia
-  - conda-forge
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1
-  - _openmp_mutex=4.5
-  - _sysroot_linux-64_curr_repodata_hack=3
-  - afterimage=1.21
-  - anyio=3.7.1
-  - appdirs=1.4.4
-  - argon2-cffi=21.3.0
-  - argon2-cffi-bindings=21.2.0
-  - arrow=1.2.3
-  - asttokens=2.2.1
-  - async-lru=2.0.4
-  - atk-1.0=2.38.0
-  - attrs=23.1.0
-  - awkward-pandas=2023.8.0
-  - aws-c-auth=0.7.0
-  - aws-c-cal=0.6.0
-  - aws-c-common=0.8.23
-  - aws-c-compression=0.2.17
-  - aws-c-event-stream=0.3.1
-  - aws-c-http=0.7.11
-  - aws-c-io=0.13.28
-  - aws-c-mqtt=0.8.14
-  - aws-c-s3=0.3.13
-  - aws-c-sdkutils=0.1.11
-  - aws-checksums=0.1.16
-  - aws-crt-cpp=0.20.3
-  - aws-sdk-cpp=1.10.57
-  - babel=2.12.1
-  - backcall=0.2.0
-  - backports=1.0
-  - backports.functools_lru_cache=1.6.5
-  - beautifulsoup4=4.12.2
-  - binutils=2.38
-  - binutils_impl_linux-64=2.38
-  - binutils_linux-64=2.38.0
-  - blas=1.0
-  - bleach=6.0.0
-  - brotlipy=0.7.0
-  - bzip2=1.0.8
-  - c-ares=1.19.1
-  - c-compiler=1.5.2
-  - ca-certificates=2025.4.26
-  - cached-property=1.5.2
-  - cached_property=1.5.2
-  - cairo=1.16.0
-  - certifi=2024.8.30
-  - cffi=1.15.1
-  - cfitsio=4.2.0
-  - charset-normalizer=2.0.4
-  - comm=0.1.4
-  - compilers=1.5.2
-  - cryptography=41.0.2
-  - cuda-cudart=11.8.89
-  - cuda-cupti=11.8.87
-  - cuda-libraries=11.8.0
-  - cuda-nvrtc=11.8.89
-  - cuda-nvtx=11.8.86
-  - cuda-runtime=11.8.0
-  - cxx-compiler=1.5.2
-  - davix=0.8.4
-  - debugpy=1.6.8
-  - decorator=5.1.1
-  - defusedxml=0.7.1
-  - dgl=1.1.1.cu118
-  - entrypoints=0.4
-  - exceptiongroup=1.1.3
-  - executing=1.2.0
-  - expat=2.5.0
-  - ffmpeg=4.3
-  - fftw=3.3.10
-  - filelock=3.9.0
-  - flit-core=3.9.0
-  - font-ttf-dejavu-sans-mono=2.37
-  - font-ttf-inconsolata=3.000
-  - font-ttf-source-code-pro=2.038
-  - font-ttf-ubuntu=0.83
-  - fontconfig=2.14.2
-  - fonts-conda-ecosystem=1
-  - fonts-conda-forge=1
-  - fortran-compiler=1.5.2
-  - fqdn=1.5.1
-  - freetype=2.12.1
-  - fribidi=1.0.10
-  - ftgl=2.4.0
-  - gcc=11.2.0
-  - gcc_impl_linux-64=11.2.0
-  - gcc_linux-64=11.2.0
-  - gdk-pixbuf=2.42.8
-  - gettext=0.21.1
-  - gflags=2.2.2
-  - gfortran=11.2.0
-  - gfortran_impl_linux-64=11.2.0
-  - gfortran_linux-64=11.2.0
-  - giflib=5.2.1
-  - gl2ps=1.4.2
-  - glew=2.1.0
-  - glog=0.6.0
-  - gmp=6.2.1
-  - gmpy2=2.1.2
-  - gnutls=3.6.15
-  - graphite2=1.3.13
-  - graphviz=6.0.2
-  - gsl=2.7
-  - gsoap=2.8.123
-  - gtk2=2.24.33
-  - gts=0.7.6
-  - gxx=11.2.0
-  - gxx_impl_linux-64=11.2.0
-  - gxx_linux-64=11.2.0
-  - harfbuzz=7.3.0
-  - icu=72.1
-  - idna=3.4
-  - importlib-metadata=6.8.0
-  - importlib-resources=6.0.1
-  - importlib_metadata=6.8.0
-  - importlib_resources=6.0.1
-  - intel-openmp=2023.1.0
-  - ipykernel=6.25.1
-  - ipyparallel=8.6.1
-  - ipython=8.12.2
-  - isoduration=20.11.0
-  - jedi=0.19.0
-  - jinja2=3.1.2
-  - jpeg=9e
-  - json5=0.9.14
-  - jsonpointer=2.0
-  - jsonschema=4.19.0
-  - jsonschema-specifications=2023.7.1
-  - jsonschema-with-format-nongpl=4.19.0
-  - jupyter-lsp=2.2.0
-  - jupyter_client=8.3.0
-  - jupyter_core=5.3.0
-  - jupyter_events=0.7.0
-  - jupyter_server=2.7.0
-  - jupyter_server_terminals=0.4.4
-  - jupyterlab=4.0.5
-  - jupyterlab_pygments=0.2.2
-  - jupyterlab_server=2.24.0
-  - kernel-headers_linux-64=3.10.0
-  - keyutils=1.6.1
-  - krb5=1.20.1
-  - lame=3.100
-  - lcms2=2.12
-  - ld_impl_linux-64=2.38
-  - lerc=3.0
-  - libabseil=20230125.3
-  - libarrow=12.0.1
-  - libblas=3.9.0
-  - libbrotlicommon=1.0.9
-  - libbrotlidec=1.0.9
-  - libbrotlienc=1.0.9
-  - libcblas=3.9.0
-  - libcrc32c=1.1.2
-  - libcublas=11.11.3.6
-  - libcufft=10.9.0.58
-  - libcufile=1.7.1.12
-  - libcurand=10.3.3.129
-  - libcurl=8.1.2
-  - libcusolver=11.4.1.48
-  - libcusparse=11.7.5.86
-  - libcxx=15.0.7
-  - libcxxabi=15.0.7
-  - libdeflate=1.12
-  - libedit=3.1.20191231
-  - libev=4.33
-  - libevent=2.1.12
-  - libexpat=2.5.0
-  - libffi=3.4.4
-  - libgcc-devel_linux-64=11.2.0
-  - libgcc-ng=13.1.0
-  - libgd=2.3.3
-  - libgfortran-ng=11.2.0
-  - libgfortran5=11.2.0
-  - libglib=2.76.4
-  - libglu=9.0.0
-  - libgomp=13.1.0
-  - libgoogle-cloud=2.12.0
-  - libgrpc=1.56.2
-  - libiconv=1.17
-  - libidn2=2.3.4
-  - libllvm13=13.0.1
-  - libllvm14=14.0.6
-  - libnghttp2=1.52.0
-  - libnpp=11.8.0.86
-  - libnsl=2.0.0
-  - libnuma=2.0.18
-  - libnvjpeg=11.9.0.86
-  - libpng=1.6.39
-  - libprotobuf=4.23.3
-  - librsvg=2.54.4
-  - libsodium=1.0.18
-  - libsqlite=3.42.0
-  - libssh2=1.11.0
-  - libstdcxx-devel_linux-64=11.2.0
-  - libstdcxx-ng=13.1.0
-  - libtasn1=4.19.0
-  - libthrift=0.18.1
-  - libtiff=4.4.0
-  - libtool=2.4.7
-  - libunistring=0.9.10
-  - libutf8proc=2.8.0
-  - libuuid=2.38.1
-  - libwebp=1.2.4
-  - libwebp-base=1.2.4
-  - libxcb=1.15
-  - libxml2=2.10.4
-  - libzlib=1.2.13
-  - llvmlite=0.40.1
-  - lz4-c=1.9.4
-  - markupsafe=2.1.1
-  - matplotlib-inline=0.1.6
-  - metakernel=0.29.5
-  - mistune=3.0.0
-  - mkl=2023.1.0
-  - mkl-service=2.4.0
-  - mkl_fft=1.3.6
-  - mkl_random=1.2.2
-  - mpc=1.1.0
-  - mpfr=4.0.2
-  - mpmath=1.3.0
-  - nbclient=0.8.0
-  - nbconvert-core=7.7.3
-  - nbformat=5.9.2
-  - ncurses=6.4
-  - nest-asyncio=1.5.6
-  - nettle=3.7.3
-  - networkx=3.1
-  - nlohmann_json=3.11.2
-  - notebook=7.0.2
-  - notebook-shim=0.2.3
-  - numba=0.57.1
-  - numpy=1.24.3
-  - numpy-base=1.24.3
-  - openh264=2.1.1
-  - openssl=3.3.1
-  - orc=1.9.0
-  - overrides=7.4.0
-  - packaging=23.0
-  - pandas=2.0.3
-  - pandocfilters=1.5.0
-  - pango=1.50.14
-  - parso=0.8.3
-  - pcre=8.45
-  - pcre2=10.40
-  - pexpect=4.8.0
-  - pickleshare=0.7.5
-  - pillow=9.4.0
-  - pip=23.2.1
-  - pixman=0.40.0
-  - pkgutil-resolve-name=1.3.10
-  - platformdirs=2.6.0
-  - pooch=1.4.0
-  - portalocker=2.7.0
-  - prometheus_client=0.17.1
-  - prompt-toolkit=3.0.39
-  - prompt_toolkit=3.0.39
-  - psutil=5.9.0
-  - pthread-stubs=0.4
-  - ptyprocess=0.7.0
-  - pure_eval=0.2.2
-  - pyarrow=12.0.1
-  - pycparser=2.21
-  - pygments=2.16.1
-  - pyopenssl=23.2.0
-  - pysocks=1.7.1
-  - pythia8=8.309
-  - python=3.8.17
-  - python-dateutil=2.8.2
-  - python-fastjsonschema=2.18.0
-  - python-json-logger=2.0.7
-  - python-tzdata=2024.2
-  - python_abi=3.8
-  - pytorch=2.0.1
-  - pytorch-cuda=11.8
-  - pytorch-mutex=1.0
-  - pytz=2023.3
-  - pyyaml=6.0
-  - pyzmq=25.1.1
-  - rdma-core=28.9
-  - re2=2023.03.02
-  - readline=8.2
-  - referencing=0.30.2
-  - requests=2.31.0
-  - rfc3339-validator=0.1.4
-  - rfc3986-validator=0.1.1
-  - root=6.28.0
-  - root_base=6.28.0
-  - rpds-py=0.9.2
-  - s2n=1.3.46
-  - scipy=1.10.1
-  - scitokens-cpp=0.7.3
-  - send2trash=1.8.2
-  - setuptools=68.0.0
-  - six=1.16.0
-  - snappy=1.1.10
-  - sniffio=1.3.0
-  - soupsieve=2.3.2.post1
-  - sqlite=3.41.2
-  - stack_data=0.6.2
-  - sympy=1.11.1
-  - sysroot_linux-64=2.17
-  - tbb=2021.8.0
-  - terminado=0.17.1
-  - tinycss2=1.2.1
-  - tk=8.6.12
-  - tomli=2.0.1
-  - torchaudio=2.0.2
-  - torchtriton=2.0.0
-  - torchvision=0.15.2
-  - tornado=6.3.2
-  - tqdm=4.65.0
-  - traitlets=5.9.0
-  - typing_extensions=4.12.2
-  - typing_utils=0.1.0
-  - ucx=1.14.1
-  - uri-template=1.3.0
-  - urllib3=1.26.16
-  - vdt=0.4.3
-  - vector-classes=1.4.3
-  - wcwidth=0.2.6
-  - webcolors=1.13
-  - webencodings=0.5.1
-  - websocket-client=1.6.1
-  - wheel=0.38.4
-  - xorg-fixesproto=5.0
-  - xorg-kbproto=1.0.7
-  - xorg-libice=1.1.1
-  - xorg-libsm=1.2.4
-  - xorg-libx11=1.8.6
-  - xorg-libxau=1.0.11
-  - xorg-libxcursor=1.2.0
-  - xorg-libxdmcp=1.1.3
-  - xorg-libxext=1.3.4
-  - xorg-libxfixes=5.0.3
-  - xorg-libxft=2.3.8
-  - xorg-libxpm=3.5.16
-  - xorg-libxrender=0.9.11
-  - xorg-libxt=1.3.0
-  - xorg-renderproto=0.11.1
-  - xorg-xextproto=7.3.0
-  - xorg-xproto=7.0.31
-  - xrootd=5.5.4
-  - xxhash=0.8.1
-  - xz=5.2.6
-  - yaml=0.2.5
-  - zeromq=4.3.4
-  - zipp=3.16.2
-  - zlib=1.2.13
-  - zstd=1.5.2
-  - pip:
-      - awkward==2.6.4
-      - awkward-cpp==33
-      - contourpy==1.1.0
-      - cramjam==2.8.3
-      - cycler==0.11.0
-      - fonttools==4.42.0
-      - fsspec==2024.3.1
-      - h5py==3.9.0
-      - pip-install==1.3.5
-      - joblib==1.3.2
-      - kiwisolver==1.4.4
-      - matplotlib==3.7.2
-      - nvidia-cublas-cu12==12.1.3.1
-      - nvidia-cuda-cupti-cu12==12.1.105
-      - nvidia-cuda-nvrtc-cu12==12.1.105
-      - nvidia-cuda-runtime-cu12==12.1.105
-      - nvidia-cudnn-cu12==8.9.2.26
-      - nvidia-cufft-cu12==11.0.2.54
-      - nvidia-curand-cu12==10.3.2.106
-      - nvidia-cusolver-cu12==11.4.5.107
-      - nvidia-cusparse-cu12==12.1.0.106
-      - nvidia-nccl-cu12==2.20.5
-      - nvidia-nvjitlink-cu12==12.4.127
-      - nvidia-nvtx-cu12==12.1.105
-      - pyparsing==3.0.9
-      - scikit-learn==1.3.0
-      - threadpoolctl==3.2.0
-      - torch==2.3.0
-      - triton==2.3.0
-      - typing-extensions==4.11.0
-      - tzdata==2024.1
-      - uproot==5.3.7
-prefix: /global/homes/j/joshuaho/.conda/envs/dgl

nemo/setup/setup/test_setup.py

@@ -1,48 +0,0 @@
-import os
-import importlib.util
-import sys
-
-def test_imports(directories):
-    """
-    Test importing all Python files in the specified directories.
-    
-    Parameters:
-    - directories: List of directory paths to test.
-    """
-    print("Testing Conda environment...")
-
-    for directory in directories:
-        print(f"\nChecking directory: {directory}")
-        
-        # Check if the directory exists
-        if not os.path.isdir(directory):
-            print(f"Directory not found: {directory}")
-            continue
-        
-        # Iterate through all files in the directory
-        for filename in os.listdir(directory):
-            # Only consider Python files
-            if filename.endswith(".py"):
-                filepath = os.path.join(directory, filename)
-                module_name = os.path.splitext(filename)[0]  # Remove .py extension
-                
-                try:
-                    # Dynamically import the module
-                    spec = importlib.util.spec_from_file_location(module_name, filepath)
-                    module = importlib.util.module_from_spec(spec)
-                    spec.loader.exec_module(module)
-                    print(f"Successfully imported: {filepath}")
-                except Exception as e:
-                    # Print the file and the error message if import fails
-                    print(f"Failed to import: {filepath}")
-                    print(f"Error: {e}")
-
-if __name__ == "__main__":
-    # Automatically append the current directory to sys.path
-    current_directory = os.getcwd()
-    sys.path.append(current_directory)
-    print(f"Current directory added to sys.path: {current_directory}")
-
-    # List of directories to check
-    directories = ["scripts", "root_gnn_base", "models"]
-    test_imports(directories)

nemo/setup/test_setup.py

@@ -1,48 +0,0 @@
-import os
-import importlib.util
-import sys
-
-def test_imports(directories):
-    """
-    Test importing all Python files in the specified directories.
-    
-    Parameters:
-    - directories: List of directory paths to test.
-    """
-    print("Testing Conda environment...")
-
-    for directory in directories:
-        print(f"\nChecking directory: {directory}")
-        
-        # Check if the directory exists
-        if not os.path.isdir(directory):
-            print(f"Directory not found: {directory}")
-            continue
-        
-        # Iterate through all files in the directory
-        for filename in os.listdir(directory):
-            # Only consider Python files
-            if filename.endswith(".py"):
-                filepath = os.path.join(directory, filename)
-                module_name = os.path.splitext(filename)[0]  # Remove .py extension
-                
-                try:
-                    # Dynamically import the module
-                    spec = importlib.util.spec_from_file_location(module_name, filepath)
-                    module = importlib.util.module_from_spec(spec)
-                    spec.loader.exec_module(module)
-                    print(f"Successfully imported: {filepath}")
-                except Exception as e:
-                    # Print the file and the error message if import fails
-                    print(f"Failed to import: {filepath}")
-                    print(f"Error: {e}")
-
-if __name__ == "__main__":
-    # Automatically append the current directory to sys.path
-    current_directory = os.getcwd()
-    sys.path.append(current_directory)
-    print(f"Current directory added to sys.path: {current_directory}")
-
-    # List of directories to check
-    directories = ["scripts", "root_gnn_base", "models"]
-    test_imports(directories)

root_gnn_dgl/README.md

@@ -1,53 +1,62 @@
-# root_gnn_dgl
-
-## Data Directory (for Hackathon)
-`/global/cfs/projectdirs/trn007/lbl_atlas/data/`
 
-* `stats_all`: full statistics sample, ~10M events per process
-* `stats_100K`: reduced statistics sample, 100K events per process
-* `processed_graphs`: graphs that have already been processed
-* `scores`: a copy of the samples along with the GNN scores for each event
+# root_gnn_dgl
 
-## Environment Setup
+Pretrained DGL-based ROOT graph neural network.
 
-The environment dependencies for this project are listed in `setup/environment.yml`. Follow the steps below to set up the environment:
+## Overview
+- Stable release with pretrained model weights.
 
-### Step 1: Install Conda
-If you don’t already have Conda installed, install either Miniconda (lightweight) or Anaconda (full version):
+Pretrained model location: ``
 
-- **Miniconda**: Download and install from [https://docs.conda.io/en/latest/miniconda.html](https://docs.conda.io/en/latest/miniconda.html).
-- **Anaconda**: Download and install from [https://www.anaconda.com/products/distribution](https://www.anaconda.com/products/distribution).
+## Conda setup
 
-### Step 2: Clone the Repository
-Clone this repository to your local machine:
 ```bash
-git init
-git lfs install
-git clone https://huggingface.co/HWresearch/GNN4Colliders
+cd setup
+conda env create -f environment.yml
+conda activate pytorch
+cd ..
+python setup/test_setup.py 
 ```
-If you want to clone without large files - just their pointers
+
+##  Container Setup (Podman-HPC)
+
+- NERSC Perlmutter environment with `podman-hpc` available.
+- Access to `joshuaho/pytorch:1.0` on Docker Hub [https://hub.docker.com/r/joshuaho/pytorch](https://hub.docker.com/r/joshuaho/pytorch)
+
+### Pull the Prebuilt Image
+
 ```bash
-GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/HWresearch/GNN4Colliders
+podman-hpc pull docker.io/joshuaho/pytorch:1.0
 ```
 
-### Step 3: Create the Conda Environment
-Use the `environment.yml` file to create the Conda environment:
+Or, you can build your own container here:
+
 ```bash
-conda env create -f setup/environment.yml -n <environment_name>
+cd setup
+source build_image.sh
 ```
 
-### Step 4: Activate the Environment
-Activate the newly created environment:
+Run the image and mount the paths you need, replaceing `<source>` with source directory path and `<target>` with the path for when you are inside the container.
 ```bash
-conda activate <environment_name>
-```
-Replace <environment_name> with the name of the environment specified in Step 4.
 
-### Step 5: Test the Environment
+podman-hpc run \
+  -it \
+  --mount type=bind,source=<source>,target=<target> \
+  --rm \
+  --network host \
+  --gpu \
+  --userns keep-id \
+  --shm-size=32g \
+  joshuaho/pytorch:1.0
+  ```
+
+### Test the Environment
 Run the `setup/test_setup.py` script to confirm that all packages needed for training are properly set up.
 ```bash
 python setup/test_setup.py
 ```
+
+
 ## Running the Demo
 The demo training is an example of our ML workflow, consisting of training a pretrained model, then finetuning it for an analysis task, while also training a model for the analysis task from scratch. The config files for the demo are located in the directory `configs/stats_100K/`. The demo can be run on a login node on Perlmutter (if enough GPU memory is availble).
 

root_gnn_dgl/configs/attention/ttH_CP_even_vs_odd.yaml

@@ -1,58 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd
-Training_Directory: trainings/attention/ttH_CP_even_vs_odd
-Model:
-  module: models.GCN
-  class: Attention_Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-    num_heads: 2
-Training:
-  epochs: 500
-  batch_size: 1024
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 1024
-    padding_mode: NODE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/attention/ttH_CP_even_vs_odd/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml

@@ -41,8 +41,8 @@ Datasets:
       file_names: ttH_NLO.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_100K/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml

@@ -38,8 +38,8 @@ Datasets:
       file_names: ttH_NLO_inc.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/pretraining_multiclass/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_100K/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_100K/pretraining_multiclass/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml

@@ -31,8 +31,8 @@ Datasets:
       file_names: ttH_NLO.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_100K/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd_batch_size_2048.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_2048
-Training_Directory: trainings/stats_100K/ttH_CP_even_vs_odd_batch_size_2048
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 2048
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 2048
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd_batch_size_2048/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd_batch_size_4096.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_4096
-Training_Directory: trainings/stats_100K/ttH_CP_even_vs_odd_batch_size_4096
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 1024
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 4096
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd_batch_size_4096/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd_batch_size_8192.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_8192
-Training_Directory: trainings/stats_100K/ttH_CP_even_vs_odd_batch_size_8192
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 2048
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 3
-    batch_size: 2048
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 3
-      buffer_size: 2
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd_batch_size_8192/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml

@@ -41,8 +41,8 @@ Datasets:
       file_names: ttH_NLO.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_all/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_all/ttH_CP_even_vs_odd/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml

@@ -38,8 +38,8 @@ Datasets:
       file_names: ttH_NLO_inc.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/pretraining_multiclass/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_all/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_all/pretraining_multiclass/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml

@@ -31,8 +31,8 @@ Datasets:
       file_names: ttH_NLO.root
       tree_name: output
       fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd/
+      raw_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/stats_all/
+      save_dir: /global/cfs/projectdirs/atlas/joshua/gnn_data/processed_graphs/stats_all/ttH_CP_even_vs_odd/
       node_branch_names:
         - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
         - [jet_eta, ele_eta, mu_eta, ph_eta, 0]

root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd_batch_size_2048.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_2048
-Training_Directory: trainings/stats_all/ttH_CP_even_vs_odd_batch_size_2048
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 2048
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 10
-    batch_size: 2048
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 10
-      buffer_size: 3
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd_batch_size_2048/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd_batch_size_4096.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_4096
-Training_Directory: trainings/stats_all/ttH_CP_even_vs_odd_batch_size_4096
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 4096
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 10
-    batch_size: 4096
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 10
-      buffer_size: 3
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd_batch_size_4096/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd_batch_size_8192.yaml

@@ -1,57 +0,0 @@
-Training_Name: ttH_CP_even_vs_odd_batch_size_8192
-Training_Directory: trainings/stats_all/ttH_CP_even_vs_odd_batch_size_8192
-Model:
-  module: models.GCN
-  class: Edge_Network
-  args:
-    hid_size: 64
-    in_size: 7
-    out_size: 1
-    n_layers: 4
-    n_proc_steps: 4
-    dropout: 0
-Training:
-  epochs: 500
-  batch_size: 8192
-  learning_rate: 0.0001
-  gamma: 0.99
-Datasets:
-  ttH_CP_even: &dataset_defn
-    module: root_gnn_base.dataset
-    class: LazyDataset
-    shuffle_chunks: 10
-    batch_size: 8192
-    padding_mode: NONE #one of STEPS, FIXED, or NONE
-    args: &dataset_args
-      name: ttH_CP_even
-      label: 0
-      # weight_var: weight
-      chunks: 10
-      buffer_size: 3
-      file_names: ttH_NLO.root
-      tree_name: output
-      fold_var: Number
-      raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/
-      save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd_batch_size_8192/
-      node_branch_names:
-        - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met]
-        - [jet_eta, ele_eta, mu_eta, ph_eta, 0]
-        - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi]
-        - CALC_E
-        - [jet_btag, 0, 0, 0, 0]
-        - [0, ele_charge, mu_charge, 0, 0]
-        - NODE_TYPE
-      node_branch_types: [vector, vector, vector, vector, single]
-      node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1]
-    folding:
-      n_folds: 4
-      test: [0]
-      # validation: 1
-      train: [1, 2, 3]
-  ttH_CP_odd:
-    <<: *dataset_defn
-    args:
-      <<: *dataset_args
-      name: ttH_CP_odd
-      label: 1
-      file_names: ttH_CPodd.root

root_gnn_dgl/jobs/interactive.sh

@@ -1 +1 @@
-salloc --nodes 1 --qos shared_interactive --time 04:00:00 --constraint gpu --account=trn007 --gres=gpu:1
+salloc --nodes 1 --qos shared_interactive --time 04:00:00 --constraint gpu --account=atlas --gres=gpu:1

root_gnn_dgl/run_demo.sh

@@ -31,7 +31,7 @@ python scripts/training_script.py --config configs/stats_100K/ttH_CP_even_vs_odd
 
 python scripts/training_script.py --config configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml --preshuffle --nocompile --lazy
 
-# Inference
+# Inference: Writing GNN Scores for from-scratch training and finetuned training to root files
 files=(
     "ttH_NLO.root"
     "ttH_CPodd.root"
@@ -50,8 +50,8 @@ branch_name=(
 for ((j=0; j<${#files[@]}; j++))
 do
     python scripts/inference.py \
-        --target "/global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/${files[j]}" \
-        --destination "/global/cfs/projectdirs/trn007/lbl_atlas/data/scores/stats_100K/${files[j]}" \
+        --target "/global/cfs/projectdirs/atlas/joshua/gnn_data/stats_100K/${files[j]}" \
+        --destination "/global/cfs/projectdirs/atlas/joshua/gnn_data/scores/stats_100K/${files[j]}" \
         --config "${config[@]}" \
         --branch_name "${branch_name[@]}" \
         --chunks 1 \

root_gnn_dgl/setup/Dockerfile

@@ -1,6 +1,6 @@
 FROM nvcr.io/nvidia/dgl:25.05-py3
 
-WORKDIR /global/cfs/projectdirs/atlas/joshua/GNN4Colliders
+WORKDIR /workspace
 
 LABEL maintainer.name="Joshua Ho"
 LABEL maintainer.email="ho22joshua@berkeley.edu"

root_gnn_dgl/setup/build_image.sh

@@ -1,4 +1,2 @@
-tag=$1
-echo $tag
-podman-hpc build -t joshuaho/pytorch:$tag --platform linux/amd64 .
-podman-hpc migrate joshuaho/pytorch:$tag
+podman-hpc build -t joshuaho/pytorch:1.0 --platform linux/amd64 .
+podman-hpc migrate joshuaho/pytorch:1.0

root_gnn_dgl/setup/environment.yml

@@ -1,4 +1,4 @@
-name: dgl
+name: pytorch
 channels:
   - pytorch
   - dglteam/label/cu118
@@ -387,5 +387,4 @@ dependencies:
       - triton==2.3.0
       - typing-extensions==4.11.0
       - tzdata==2024.1
-      - uproot==5.3.7
-prefix: /global/homes/j/joshuaho/.conda/envs/dgl
+      - uproot==5.3.7

root_gnn_dgl/setup/launch_image.sh

@@ -0,0 +1,9 @@
+podman-hpc run \
+  -it \
+  --mount type=bind,source=/pscratch/sd/j/joshuaho/,target=/pscratch/sd/j/joshuaho/ \
+  --mount type=bind,source=/global/cfs/projectdirs/atlas/joshua/,target=/global/cfs/projectdirs/atlas/joshua/ \
+  --rm \
+  --network host \
+  --gpu \
+  --shm-size=32g \
+  joshuaho/pytorch:1.0