Input opt

mlpf/heptfds/clic_pf_edm4hep/ttbar.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import tensorflow as tf
+import tensorflow_datasets as tfds
 from utils_edm import (
     X_FEATURES_CL,
     X_FEATURES_TRK,
@@ -9,8 +10,6 @@
     split_sample,
 )
 
-import tensorflow_datasets as tfds
-
 _DESCRIPTION = """
 CLIC EDM4HEP dataset with ee -> ttbar at 380GeV.
   - X: reconstructed tracks and clusters, variable number N per event
@@ -26,7 +25,7 @@
 
 
 class ClicEdmTtbarPf(tfds.core.GeneratorBasedBuilder):
-    VERSION = tfds.core.Version("2.1.0")
+    VERSION = tfds.core.Version("2.2.0")
     RELEASE_NOTES = {
         "1.0.0": "Initial release.",
         "1.1.0": "update stats, move to 380 GeV",
@@ -36,6 +35,7 @@ class ClicEdmTtbarPf(tfds.core.GeneratorBasedBuilder):
         "1.5.0": "Regenerate with ARRAY_RECORD",
         "2.0.0": "Add ispu, genjets, genmet; disable genjet_idx; truth def not based on gp.status==1",
         "2.1.0": "Bump dataset size",
+        "2.2.0": "Additional cluster input features",
     }
     MANUAL_DOWNLOAD_INSTRUCTIONS = """
     For the raw input files in ROOT EDM4HEP format, please see the citation above.

mlpf/heptfds/clic_pf_edm4hep/utils_edm.py

@@ -1,6 +1,7 @@
+import random
+
 import awkward as ak
 import numpy as np
-import random
 
 # from fcc/postprocessing.py
 X_FEATURES_TRK = [
@@ -39,6 +40,16 @@
     "sigma_x",
     "sigma_y",
     "sigma_z",
+    # additional cluster input features
+    "energyError",
+    "sigma_energy",
+    "sigma_x_weighted",
+    "sigma_y_weighted",
+    "sigma_z_weighted",
+    "energy_weighted_width",
+    "pos_shower_max",
+    "width_shower_max",
+    "energy_shower_max",
 ]
 
 Y_FEATURES = ["PDG", "charge", "pt", "eta", "sin_phi", "cos_phi", "energy", "ispu"]

mlpf/pyg/mlpf.py

@@ -1,13 +1,13 @@
+import math
+
+import numpy as np
 import torch
 import torch.nn as nn
-
-from .gnn_lsh import CombinedGraphLayer
-
 from pyg.logger import _logger
-import math
-import numpy as np
 from torch.nn.attention import SDPBackend, sdpa_kernel
 
+from .gnn_lsh import CombinedGraphLayer
+
 
 def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
     # From https://github.com/rwightman/pytorch-image-models/blob/
@@ -57,6 +57,32 @@ def norm_cdf(x):
         return tensor
 
 
+def standardize_input(X, elemtypes_nonzero, standardization_dict):
+
+    for i, ielem in enumerate(elemtypes_nonzero):
+
+        Xfeat_normed_msked = X.clone()
+
+        # get mean/std of features of that elem
+        mean = torch.tensor(standardization_dict[f"PFelement{ielem}"]["mean"]).to(Xfeat_normed_msked.device)
+        std = torch.tensor(standardization_dict[f"PFelement{ielem}"]["std"]).to(Xfeat_normed_msked.device)
+
+        # standardize
+        Xfeat_normed_msked[..., 1:] = (Xfeat_normed_msked[..., 1:] - mean[..., 1:]) / std[..., 1:]
+
+        # msk other elements
+        msk = Xfeat_normed_msked[..., 0:1] == ielem
+        Xfeat_normed_msked = Xfeat_normed_msked * msk
+        Xfeat_normed_msked = torch.nan_to_num(Xfeat_normed_msked, nan=0.0)
+
+        if i == 0:
+            Xfeat_normed = Xfeat_normed_msked
+        else:
+            Xfeat_normed += Xfeat_normed_msked
+
+    return Xfeat_normed
+
+
 def get_activation(activation):
     if activation == "elu":
         act = nn.ELU
@@ -372,9 +398,12 @@ def __init__(
             self.final_norm_reg = torch.nn.LayerNorm(embed_dim)
 
     # @torch.compile
-    def forward(self, X_features, mask):
+    def forward(self, X_features, mask, standardization_dict=None):
         Xfeat_normed = X_features
 
+        if standardization_dict is not None:
+            Xfeat_normed = standardize_input(X_features, self.elemtypes_nonzero, standardization_dict)
+
         embeddings_id, embeddings_reg = [], []
         if self.num_convs != 0:
             if self.input_encoding == "joint":