update model to 0612

app.py

@@ -46,14 +46,13 @@ z_events = ak.from_parquet(os.path.join(base_path, "data/z_test_subset.parquet")
 qq_events = ak.from_parquet(os.path.join(base_path, "data/qq_test_subset.parquet"))
 
 def run_full_inference(events):
-    from mltau.tools.io.ParT_dataloader import ParticleTransformerDataset
+    from mltau.tools.io.ParT_dataloader import build_tensors
     from mltau.tools.io.scaling import apply_saved_input_scaling_from_cfg
     
     # Ensure output_dir is defined for scaler path resolution if needed
     if "output_dir" not in cfg:
         cfg.output_dir = base_path
         
-    ds = ParticleTransformerDataset([], cfg, batch_size=256)
     all_preds = {}
     all_features = []
     all_masks = []
@@ -73,7 +72,7 @@ def run_full_inference(events):
         for i in range(0, len(events), 256):
             captured.clear()
             chunk = events[i:i+256]
-            tensors = ds.build_tensors(chunk)
+            tensors = build_tensors(chunk, cfg)
             
             # Apply scaling (it will check if enabled in cfg)
             tensors = apply_saved_input_scaling_from_cfg(tensors, cfg)
@@ -99,6 +98,11 @@ print("Pre-computing inference results...")
 z_results, z_feat, z_mask_feat, z_emb = run_full_inference(z_events)
 qq_results, qq_feat, qq_mask_feat, qq_emb = run_full_inference(qq_events)
 
+# Capture decay modes for the signal events
+z_dm_raw_all = ak.to_numpy(z_events.gen_jet_tau_decaymode)
+z_dm_reduced_all = g.get_reduced_decaymodes(z_dm_raw_all)
+z_dm_idx_all = g.prepare_one_hot_encoding(z_dm_reduced_all)
+
 # Compute t-SNE for a subset of jets
 print("Computing t-SNE visualization...")
 n_z = min(1000, len(z_emb))
@@ -108,10 +112,8 @@ tsne_results = TSNE(n_components=2, random_state=42).fit_transform(tsne_emb)
 z_tsne = tsne_results[:n_z]
 qq_tsne = tsne_results[n_z:]
 
-# Capture decay modes for the signal subset
-z_dm_raw = ak.to_numpy(z_events.gen_jet_tau_decaymode[:n_z])
-z_dm_reduced = g.get_reduced_decaymodes(z_dm_raw)
-z_dm_idx = g.prepare_one_hot_encoding(z_dm_reduced)
+# Use the pre-computed indices for t-SNE visualization mapping
+z_dm_idx = z_dm_idx_all[:n_z]
 
 FEATURE_NAMES = [
     "Δη (relative to jet)", "Δφ (relative to jet)", "log(pT)", "log(energy)", 
@@ -178,8 +180,7 @@ def update_scientific_plots(min_pt, max_pt, min_tot, max_tot, min_ch, max_ch, mi
         cm = confusion_matrix(true_dm_idx, pred_dm, labels=range(6), normalize='true')
         
         # Use descriptive labels for the 6 classes
-        class_names = ["1p 0pi0", "1p 1pi0", "1p >=2pi0", "3p 0pi0", "3p >=1pi0", "Rare"]
-        fig_cm = go.Figure(data=go.Heatmap(z=cm, x=class_names, y=class_names, colorscale='Viridis'))
+        fig_cm = go.Figure(data=go.Heatmap(z=cm, x=CLASS_NAMES, y=CLASS_NAMES, colorscale='Viridis'))
         fig_cm.update_layout(
             title="Decay mode confusion matrix",
             template="plotly_white",
@@ -275,16 +276,37 @@ def plot_histogram(dataset_name, prop_key, min_val, max_val, title, xlabel):
     fig.update_layout(title=title, xaxis_title=xlabel, yaxis_title="Count", template="plotly_white", margin=dict(l=10, r=10, t=40, b=10), height=300, showlegend=False, yaxis2=dict(overlaying='y', visible=False, range=[0, 1]))
     return fig
 
+CLASS_NAMES = ["1p 0pi0", "1p 1pi0", "1p >=2pi0", "3p 0pi0", "3p >=1pi0", "Rare"]
+
 def filter_events_handler(*filters):
     z_mask = get_mask(z_props, *filters)
     qq_mask = get_mask(qq_props, *filters)
-    z_indices = np.where(z_mask)[0].tolist()
-    qq_indices = np.where(qq_mask)[0].tolist()
+    z_indices = np.where(z_mask)[0]
+    qq_indices = np.where(qq_mask)[0]
     
     status_text = f"**Event selection summary:**\n- Signal (Z): {len(z_indices)} / {len(z_events)}\n- Background (QQ): {len(qq_indices)} / {len(qq_events)}"
     
-    z_update = gr.update(choices=z_indices, value=z_indices[0] if z_indices else None)
-    qq_update = gr.update(choices=qq_indices, value=qq_indices[0] if qq_indices else None)
+    # Create descriptive labels and group signal by DM
+    z_choices = []
+    for idx in z_indices:
+        dm_name = CLASS_NAMES[z_dm_idx_all[idx]]
+        n_ch, n_nh, n_ph = z_props["n_ch"][idx], z_props["n_nh"][idx], z_props["n_photon"][idx]
+        label = f"[{dm_name}] {n_ch}ch, {n_nh}nh, {n_ph}ph | ID:{idx}"
+        z_choices.append((label, int(idx)))
+    # Sort by DM index, then constituent counts
+    z_choices.sort(key=lambda x: (z_dm_idx_all[x[1]], z_props["n_ch"][x[1]], z_props["n_photon"][x[1]]))
+
+    # Create descriptive labels and group background by constituent counts
+    qq_choices = []
+    for idx in qq_indices:
+        n_ch, n_nh, n_ph = qq_props["n_ch"][idx], qq_props["n_nh"][idx], qq_props["n_photon"][idx]
+        label = f"{n_ch}ch, {n_nh}nh, {n_ph}ph | ID:{idx}"
+        qq_choices.append((label, int(idx)))
+    # Sort by constituent counts
+    qq_choices.sort(key=lambda x: (qq_props["n_ch"][x[1]], qq_props["n_photon"][x[1]], qq_props["n_nh"][x[1]]))
+
+    z_update = gr.update(choices=z_choices, value=z_choices[0][1] if z_choices else None)
+    qq_update = gr.update(choices=qq_choices, value=qq_choices[0][1] if qq_choices else None)
     return z_update, qq_update, status_text
 
 def plot_2d_jet(idx, dataset_name):
@@ -386,10 +408,9 @@ def plot_tsne():
     ))
     
     # Signal by decay mode
-    class_names = ["1p 0pi0", "1p 1pi0", "1p >=2pi0", "3p 0pi0", "3p >=1pi0", "Rare"]
     symbols = ['circle', 'diamond', 'square', 'cross', 'x', 'triangle-up']
     
-    for i, name in enumerate(class_names):
+    for i, name in enumerate(CLASS_NAMES):
         mask = z_dm_idx == i
         if np.any(mask):
             fig.add_trace(go.Scatter(
@@ -403,7 +424,7 @@ def plot_tsne():
         xaxis_title="t-SNE dimension 1",
         yaxis_title="t-SNE dimension 2",
         template="plotly_white",
-        height=600, width=800,
+        height=800, width=800,
         legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99)
     )
     return fig

hf_space/models/model.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b025e95007b99e95ccca77d106f67172edb2c82444b5b6707810eb27f545cb35
+oid sha256:647985cd3ad9ca838ee3610218042f39a3de19c2488509d37184e4131a207f2a
 size 14485529