adding edge network

physicsnemo/configs/tHjb_CP_0_vs_90.yaml

@@ -31,13 +31,21 @@ performance:
   jit: False
 
 architecture:
-  processor_size: 8
-  hidden_dim_node_encoder: 128
-  hidden_dim_edge_encoder: 128
-  hidden_dim_processor: 128
-  hidden_dim_node_decoder: 128
-  global_emb_dim: 128
-  out_dim: 1
+  module: models.MeshGraphNet
+  class: MeshGraphNet
+  args:
+    base_gnn:
+      input_dim_nodes: 7
+      input_dim_edges: 3
+      output_dim: 128
+      processor_size: 8
+      hidden_dim_node_encoder: 128
+      hidden_dim_edge_encoder: 128
+      hidden_dim_processor: 128
+      hidden_dim_node_decoder: 128
+    global_emb_dim: 128
+    global_feat_dim: 1
+    out_dim: 1
 
 paths:
   data_dir: /global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/v02/preselection/merged_fixed/train/

physicsnemo/configs/tHjb_CP_0_vs_90_edge_network.yaml

@@ -0,0 +1,82 @@
+# ignore_header_test
+# Copyright 2023 Stanford University
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+random_seed: 2
+
+scheduler:
+  lr: 1.E-3
+  lr_decay: 1.E-3
+
+training:
+  epochs: 100
+  
+checkpoints:
+  ckpt_path: "checkpoints"
+  ckpt_name: "tHjb_CP_0_vs_90_edge_network"
+
+performance:
+  amp: False
+  jit: False
+
+architecture:
+  module: models.Edge_Network
+  class: Edge_Network
+  args:
+    input_dim_nodes: 7
+    input_dim_edges: 3
+    input_dim_globals: 1
+    hid_size: 64
+    n_layers: 4
+    n_proc_steps: 4
+    out_dim: 1
+
+paths:
+  data_dir: /global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/v02/preselection/merged_fixed/train/
+  save_dir: /pscratch/sd/j/joshuaho/physicsnemo/ttHCP/graphs/tHjb_CP_0_vs_90/
+  training_dir: ./tHjb_CP_0_vs_90_edge_network/
+
+datasets:
+  - name: tHjb_cp_0_had
+    load_path: ${paths.data_dir}/merged_aMCPy8_tHjb125_CP_0_AF3_had_scaled.root
+    label: 0
+  - name: tHjb_cp_0_lep
+    load_path: ${paths.data_dir}/merged_aMCPy8_tHjb125_CP_0_AF3_lep_scaled.root
+    label: 0
+  - name: tHjb_cp_90_had
+    load_path: ${paths.data_dir}/merged_aMCPy8_tHjb125_CP_90_AF3_had_scaled.root
+    label: 1
+  - name: tHjb_cp_90_lep
+    load_path: ${paths.data_dir}/merged_aMCPy8_tHjb125_CP_90_AF3_lep_scaled.root
+    label: 1
+
+root_dataset:
+  ttree: output
+  dtype: torch.bfloat16
+  features:
+    # pt, eta, phi, energy, btag, charge, node_type
+    jet: [m_jet_pt, m_jet_eta, m_jet_phi, CALC_E, m_jet_PCbtag, 0, 0]
+    electron: [m_el_pt, m_el_eta, m_el_phi, CALC_E, 0, m_el_charge, 1]
+    muon: [m_mu_pt, m_mu_eta, m_mu_phi, CALC_E, 0, m_mu_charge, 2]
+    photon: [ph_pt_myy, ph_eta, ph_phi, CALC_E, 0, 0, 3]
+    met: [m_met, 0, m_met_phi, CALC_E, 0, 0, 4]
+  globals: [NUM_NODES]
+  weights: 1
+  tracking: []
+  step_size: 16384
+  batch_size: 16384
+  train_val_test_split: [0.5, 0.25, 0.25]
+  prebatch: 
+    enabled: True
+    chunk_size: 512

physicsnemo/models/Edge_Network.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn as nn
+import dgl
+
+from models import utils
+
+class Edge_Network(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        hid_size = cfg.hid_size
+        n_layers = cfg.n_layers
+        self.n_proc_steps = cfg.n_proc_steps
+
+        #encoder
+        self.node_encoder = utils.Make_MLP(cfg.input_dim_nodes, hid_size, hid_size, n_layers)
+        self.edge_encoder = utils.Make_MLP(cfg.input_dim_edges, hid_size, hid_size, n_layers)
+        self.global_encoder = utils.Make_MLP(cfg.input_dim_globals, hid_size, hid_size, n_layers)
+
+        #GNN
+        self.node_update = utils.Make_MLP(3*hid_size, hid_size, hid_size, n_layers)
+        self.edge_update = utils.Make_MLP(4*hid_size, hid_size, hid_size, n_layers)
+        self.global_update = utils.Make_MLP(3*hid_size, hid_size, hid_size, n_layers)
+
+        #decoder
+        self.global_decoder = utils.Make_MLP(hid_size, hid_size, hid_size, n_layers)
+        self.classify = nn.Linear(hid_size, cfg.out_dim)
+
+    def forward(self, node_feats, edge_feats, global_feats, batched_graph, metadata={}):
+        # encoders
+        batched_graph.ndata['h'] = self.node_encoder(node_feats)
+        batched_graph.edata['e'] = self.edge_encoder(edge_feats)
+
+        if global_feats.ndim == 3:
+            global_feats = global_feats.view(-1, global_feats.shape[-1])
+        h_global = self.global_encoder(global_feats)
+
+        # message passing
+        for _ in range(self.n_proc_steps):
+            batched_graph.apply_edges(dgl.function.copy_u('h', 'm_u'))
+            batched_graph.apply_edges(utils.copy_v)
+
+            # edge update
+            edge_inputs = torch.cat([
+                batched_graph.edata['e'],
+                batched_graph.edata['m_u'],
+                batched_graph.edata['m_v'],
+                utils.broadcast_global_to_edges(h_global, edge_split=metadata.get("batch_num_edges", None))
+            ], dim=1)
+            batched_graph.edata['e'] = self.edge_update(edge_inputs)
+
+            # node update
+            batched_graph.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e'))
+            node_inputs = torch.cat([
+                batched_graph.ndata['h'],
+                batched_graph.ndata['h_e'],
+                utils.broadcast_global_to_nodes(h_global, node_split=metadata.get("batch_num_nodes", None))
+            ], dim=1)
+            batched_graph.ndata['h'] = self.node_update(node_inputs)
+
+            # global update
+            graph_node_feat = utils.mean_nodes(
+                batched_graph, 'h', node_split=metadata.get("batch_num_nodes", None)
+            )
+            graph_edge_feat = utils.mean_edges(
+                batched_graph, 'e', edge_split=metadata.get("batch_num_edges", None)
+            )
+            h_global = self.global_update(torch.cat([h_global, graph_node_feat, graph_edge_feat], dim=1))
+
+        h_global = self.global_decoder(h_global)
+        out = self.classify(h_global)
+        return out
+

physicsnemo/models/MeshGraphNet.py

@@ -2,6 +2,8 @@ import torch
 import torch.nn as nn
 import dgl
 
+from models import utils
+
 # Import the PhysicsNemo MeshGraphNet model
 from physicsnemo.models.meshgraphnet import MeshGraphNet as PhysicsNemoMeshGraphNet
 
@@ -35,8 +37,8 @@ class MeshGraphNet(nn.Module):
         batched_graph.ndata['h'] = node_pred
         batched_graph.edata['e'] = edge_feats
 
-        graph_node_feat = mean_nodes(batched_graph, 'h', node_split=metadata.get("batch_num_nodes", None))
-        graph_edge_feat = mean_edges(batched_graph, 'e', edge_split=metadata.get("batch_num_edges", None))
+        graph_node_feat = utils.mean_nodes(batched_graph, 'h', node_split=metadata.get("batch_num_nodes", None))
+        graph_edge_feat = utils.mean_edges(batched_graph, 'e', edge_split=metadata.get("batch_num_edges", None))
 
         # Flatten global_feats if needed
         if global_feats.ndim == 3:
@@ -47,82 +49,3 @@ class MeshGraphNet(nn.Module):
         graph_pred = self.mlp(combined_feat)
         return graph_pred
     
-def mean_nodes(batched_graph, feat_key='h', op='mean', node_split=None):
-    """
-    Aggregates node features per disjoint graph in a batched DGLGraph.
-
-    Args:
-        batched_graph: DGLGraph
-        feat_key: str, node feature key
-        op: 'mean', 'sum', or 'max'
-        node_split: 1D tensor or list of ints (num nodes per graph)
-
-    Returns:
-        Tensor of shape [num_graphs, node_feat_dim]
-    """
-    h = batched_graph.ndata[feat_key]
-    if node_split is None or len(node_split) == 0:
-        if op == 'mean':
-            return dgl.mean_nodes(batched_graph, feat_key)
-        elif op == 'sum':
-            return dgl.sum_nodes(batched_graph, feat_key)
-        elif op == 'max':
-            return dgl.max_nodes(batched_graph, feat_key)
-        else:
-            raise ValueError(f"Unknown op: {op}")
-    else:
-        # Ensure node_split is a flat list of ints
-        if isinstance(node_split, torch.Tensor):
-            splits = node_split.view(-1).tolist()
-        else:
-            splits = [int(x) for x in node_split]
-        chunks = torch.split(h, splits, dim=0)
-        if op == 'mean':
-            out = torch.stack([chunk.mean(0) if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
-        elif op == 'sum':
-            out = torch.stack([chunk.sum(0) if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
-        elif op == 'max':
-            out = torch.stack([chunk.max(0).values if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
-        else:
-            raise ValueError(f"Unknown op: {op}")
-        return out
-    
-def mean_edges(batched_graph, feat_key='e', op='mean', edge_split=None):
-    """
-    Aggregates edge features per disjoint graph in a batched DGLGraph.
-
-    Args:
-        batched_graph: DGLGraph
-        feat_key: str, edge feature key
-        op: 'mean', 'sum', or 'max'
-        edge_split: 1D tensor or list of ints (num edges per graph)
-
-    Returns:
-        Tensor of shape [num_graphs, edge_feat_dim]
-    """
-    e = batched_graph.edata[feat_key]
-    if edge_split is None or len(edge_split) == 0:
-        if op == 'mean':
-            return dgl.mean_edges(batched_graph, feat_key)
-        elif op == 'sum':
-            return dgl.sum_edges(batched_graph, feat_key)
-        elif op == 'max':
-            return dgl.max_edges(batched_graph, feat_key)
-        else:
-            raise ValueError(f"Unknown op: {op}")
-    else:
-        # Ensure edge_split is a flat list of ints
-        if isinstance(edge_split, torch.Tensor):
-            splits = edge_split.view(-1).tolist()
-        else:
-            splits = [int(x) for x in edge_split]
-        chunks = torch.split(e, splits, dim=0)
-        if op == 'mean':
-            out = torch.stack([chunk.mean(0) if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
-        elif op == 'sum':
-            out = torch.stack([chunk.sum(0) if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
-        elif op == 'max':
-            out = torch.stack([chunk.max(0).values if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
-        else:
-            raise ValueError(f"Unknown op: {op}")
-        return out

physicsnemo/models/utils.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+import dgl
+
+def mean_nodes(batched_graph, feat_key='h', op='mean', node_split=None):
+    """
+    Aggregates node features per disjoint graph in a batched DGLGraph.
+
+    Args:
+        batched_graph: DGLGraph
+        feat_key: str, node feature key
+        op: 'mean', 'sum', or 'max'
+        node_split: 1D tensor or list of ints (num nodes per graph)
+
+    Returns:
+        Tensor of shape [num_graphs, node_feat_dim]
+    """
+    h = batched_graph.ndata[feat_key]
+    if node_split is None or len(node_split) == 0:
+        if op == 'mean':
+            return dgl.mean_nodes(batched_graph, feat_key)
+        elif op == 'sum':
+            return dgl.sum_nodes(batched_graph, feat_key)
+        elif op == 'max':
+            return dgl.max_nodes(batched_graph, feat_key)
+        else:
+            raise ValueError(f"Unknown op: {op}")
+    else:
+        # Ensure node_split is a flat list of ints
+        if isinstance(node_split, torch.Tensor):
+            splits = node_split.view(-1).tolist()
+        else:
+            splits = [int(x) for x in node_split]
+        chunks = torch.split(h, splits, dim=0)
+        if op == 'mean':
+            out = torch.stack([chunk.mean(0) if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
+        elif op == 'sum':
+            out = torch.stack([chunk.sum(0) if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
+        elif op == 'max':
+            out = torch.stack([chunk.max(0).values if chunk.shape[0] > 0 else torch.zeros_like(h[0]) for chunk in chunks])
+        else:
+            raise ValueError(f"Unknown op: {op}")
+        return out
+    
+def mean_edges(batched_graph, feat_key='e', op='mean', edge_split=None):
+    """
+    Aggregates edge features per disjoint graph in a batched DGLGraph.
+
+    Args:
+        batched_graph: DGLGraph
+        feat_key: str, edge feature key
+        op: 'mean', 'sum', or 'max'
+        edge_split: 1D tensor or list of ints (num edges per graph)
+
+    Returns:
+        Tensor of shape [num_graphs, edge_feat_dim]
+    """
+    e = batched_graph.edata[feat_key]
+    if edge_split is None or len(edge_split) == 0:
+        if op == 'mean':
+            return dgl.mean_edges(batched_graph, feat_key)
+        elif op == 'sum':
+            return dgl.sum_edges(batched_graph, feat_key)
+        elif op == 'max':
+            return dgl.max_edges(batched_graph, feat_key)
+        else:
+            raise ValueError(f"Unknown op: {op}")
+    else:
+        # Ensure edge_split is a flat list of ints
+        if isinstance(edge_split, torch.Tensor):
+            splits = edge_split.view(-1).tolist()
+        else:
+            splits = [int(x) for x in edge_split]
+        chunks = torch.split(e, splits, dim=0)
+        if op == 'mean':
+            out = torch.stack([chunk.mean(0) if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
+        elif op == 'sum':
+            out = torch.stack([chunk.sum(0) if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
+        elif op == 'max':
+            out = torch.stack([chunk.max(0).values if chunk.shape[0] > 0 else torch.zeros_like(e[0]) for chunk in chunks])
+        else:
+            raise ValueError(f"Unknown op: {op}")
+        return out
+    
+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)
+
+def broadcast_global_to_nodes(globals, node_split):
+    """
+    globals: [num_graphs, global_dim]
+    node_split: list/1D tensor of length num_graphs, number of nodes per graph
+    Returns: [total_num_nodes, global_dim]
+    """
+    if node_split is None:
+        raise ValueError("node_split must be provided")
+    if not torch.is_tensor(node_split):
+        node_split = torch.tensor(node_split, dtype=torch.long, device=globals.device)
+    else:
+        node_split = node_split.to(device=globals.device, dtype=torch.long)
+    node_split = node_split.flatten()
+    return torch.repeat_interleave(globals, node_split, dim=0)
+
+def broadcast_global_to_edges(globals, edge_split):
+    """
+    globals: [num_graphs, global_dim] (on CUDA or CPU)
+    edge_split: list/1D tensor of length num_graphs, number of edges per graph (CPU or CUDA)
+    Returns: [total_num_edges, global_dim]
+    """
+    if edge_split is None:
+        raise ValueError("edge_split must be provided")
+    if not torch.is_tensor(edge_split):
+        edge_split = torch.tensor(edge_split, dtype=torch.long, device=globals.device)
+    else:
+        edge_split = edge_split.to(device=globals.device, dtype=torch.long)
+    edge_split = edge_split.flatten()
+    return torch.repeat_interleave(globals, edge_split, dim=0)
+
+def copy_v(edges):
+    return {'m_v': edges.dst['h']}

physicsnemo/train.py

@@ -23,6 +23,7 @@ import models.MeshGraphNet as MeshGraphNet
 from dataset.Dataset import get_dataset
 import metrics
 
+import utils
 
 class MGNTrainer:
     def __init__(self, logger, cfg, dist):
@@ -30,8 +31,6 @@ class MGNTrainer:
         self.device = dist.device
         logger.info(f"Using {self.device} device")
 
-        params = {}
-
         start = time.time()
         self.trainloader, self.valloader, self.testloader = get_dataset(cfg, self.device)
         print(f"total time loading dataset: {time.time() - start:.2f} seconds")
@@ -42,20 +41,10 @@ class MGNTrainer:
         else:
             self.dtype = torch.float32
 
-        node_features = list(cfg.root_dataset.features.values())[0]
-        edge_features = ["dR", "deta", "dphi"]
-        global_features = ["num_nodes"]
-
-        params["infeat_nodes"] = len(node_features)
-        params["infeat_edges"] = len(edge_features)
-        params["infeat_globals"] = len(global_features)
-        params["out_dim"] = cfg.architecture.out_dim
-        params["node_features"] = list(node_features)
-        params["edge_features"] = edge_features
-        params["global_features"] = global_features
-
-        self.model = MeshGraphNet.MeshGraphNet(cfg.architecture)
+        self.model = utils.build_from_module(cfg.architecture)
         self.model = self.model.to(dtype=self.dtype, device=self.device)
+        # num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        # print(f"Number of trainable parameters: {num_params}")
 
         if cfg.performance.jit:
             self.model = torch.jit.script(self.model).to(self.device)
@@ -81,7 +70,6 @@ class MGNTrainer:
             device=self.device,
         )
 
-        self.params = params
         self.cfg = cfg
 
     def backward(self, loss):
@@ -244,9 +232,6 @@ def do_training(cfg: DictConfig):
         )
         start = time.time()
         trainer.scheduler.step()
-
-        with open(cfg.checkpoints.ckpt_path + "/parameters.json", "w") as outf:
-            json.dump(trainer.params, outf, indent=4)
     logger.info("Training completed!")
 
 

physicsnemo/utils.py

@@ -0,0 +1,11 @@
+import importlib
+from types import SimpleNamespace
+
+def build_from_module(cfg):
+    modname = cfg['module']
+    classname = cfg['class']
+    args = cfg['args']
+    module = importlib.import_module(modname)
+    model_cls = getattr(module, classname)
+    cfg_obj = SimpleNamespace(**args)
+    return model_cls(cfg_obj)