secondary_vertexing_dis/analysis/invariant_mass.py

0001 #!/usr/bin/env python3
0002
0003 import uproot
0004 import numpy as np
0005 import matplotlib.pyplot as plt
0006 import argparse
0007
0008 def main():
0009     parser = argparse.ArgumentParser(description='Analyze invariant mass from secondary vertex data')
0010     parser.add_argument('--input', '-i', nargs='+', default=['podio_output.root'],
0011                        help='Input ROOT file(s) (default: podio_output.root)')
0012     parser.add_argument('--output', '-o', default='invariant_mass.png',
0013                        help='Output plot filename (default: invariant_mass.png)')
0014     parser.add_argument('--vertices-collection', default='SecondaryVerticesHelix',
0015                        help='Name of the vertices collection (default: SecondaryVerticesHelix)')
0016     args = parser.parse_args()
0017
0018     print(f"Loading data from {len(args.input)} file(s): {args.input}")
0019
0020     # Combine data from all input files
0021     all_sec_x_raw = []
0022     all_sec_y_raw = []
0023     all_sec_z_raw = []
0024     all_sec_chi2_raw = []
0025     all_sec_ndf_raw = []
0026     all_px_raw = []
0027     all_py_raw = []
0028     all_pz_raw = []
0029     all_energy_raw = []
0030     all_assoc_begin = []
0031     all_assoc_end = []
0032     all_assoc_indices = []
0033
0034     for input_file in args.input:
0035         print(f"  Loading {input_file}...")
0036         file = uproot.open(input_file)
0037         tree = file['events']
0038
0039         # Load data from this file
0040         vertices_collection = args.vertices_collection
0041         all_sec_x_raw.extend(tree[f'{vertices_collection}/{vertices_collection}.position.x'].array())
0042         all_sec_y_raw.extend(tree[f'{vertices_collection}/{vertices_collection}.position.y'].array())
0043         all_sec_z_raw.extend(tree[f'{vertices_collection}/{vertices_collection}.position.z'].array())
0044         all_sec_chi2_raw.extend(tree[f'{vertices_collection}/{vertices_collection}.chi2'].array())
0045         all_sec_ndf_raw.extend(tree[f'{vertices_collection}/{vertices_collection}.ndf'].array())
0046         all_px_raw.extend(tree['ReconstructedChargedParticles/ReconstructedChargedParticles.momentum.x'].array())
0047         all_py_raw.extend(tree['ReconstructedChargedParticles/ReconstructedChargedParticles.momentum.y'].array())
0048         all_pz_raw.extend(tree['ReconstructedChargedParticles/ReconstructedChargedParticles.momentum.z'].array())
0049         all_energy_raw.extend(tree['ReconstructedChargedParticles/ReconstructedChargedParticles.energy'].array())
0050         all_assoc_begin.extend(tree[f'{vertices_collection}/{vertices_collection}.associatedParticles_begin'].array())
0051         all_assoc_end.extend(tree[f'{vertices_collection}/{vertices_collection}.associatedParticles_end'].array())
0052         all_assoc_indices.extend(tree[f'_{vertices_collection}_associatedParticles/_{vertices_collection}_associatedParticles.index'].array())
0053
0054         file.close()
0055
0056     print(f"Loaded data from {len(all_sec_x_raw)} events")
0057
0058     # Calculate total statistics
0059     total_events = len(all_sec_x_raw)
0060     total_sec_vertices = sum(len(all_sec_x_raw[i]) for i in range(len(all_sec_x_raw)))
0061     total_particles = sum(len(all_px_raw[i]) for i in range(len(all_px_raw)))
0062
0063     print(f"\n=== INITIAL STATISTICS ===")
0064     print(f"Total events: {total_events}")
0065     print(f"Total secondary vertices: {total_sec_vertices}")
0066     print(f"Total reconstructed particles: {total_particles}")
0067     print(f"Total association entries: {sum(len(all_assoc_indices[i]) for i in range(len(all_assoc_indices)))}")
0068
0069     # Process events with correct associations
0070     invariant_masses = []
0071     decay_lengths_2d = []
0072     decay_lengths_3d = []
0073     vertex_quality = []
0074     n_tracks_per_vertex = []
0075
0076     # Cut flow counters
0077     events_processed = 0
0078     vertices_processed = 0
0079     two_track_vertices = 0
0080     valid_associations = 0
0081     valid_invariant_masses = 0
0082     passed_quality_cuts = 0
0083     kaon_candidates = 0
0084
0085     print("\n=== PROCESSING WITH CORRECT ASSOCIATIONS ===")
0086
0087     for event_idx in range(len(all_sec_x_raw)):
0088         # Get event data
0089         event_sec_x = all_sec_x_raw[event_idx]
0090         event_sec_y = all_sec_y_raw[event_idx]
0091         event_sec_z = all_sec_z_raw[event_idx]
0092         event_chi2 = all_sec_chi2_raw[event_idx]
0093         event_ndf = all_sec_ndf_raw[event_idx]
0094
0095         event_px = all_px_raw[event_idx]
0096         event_py = all_py_raw[event_idx]
0097         event_pz = all_pz_raw[event_idx]
0098         event_energy = all_energy_raw[event_idx]
0099
0100         event_assoc_begin = all_assoc_begin[event_idx]
0101         event_assoc_end = all_assoc_end[event_idx]
0102         event_assoc_indices = all_assoc_indices[event_idx]
0103
0104         # Skip if no particles or vertices
0105         if len(event_px) == 0 or len(event_sec_x) == 0:
0106             continue
0107
0108         events_processed += 1
0109
0110         # Process each secondary vertex in this event
0111         for vtx_idx in range(len(event_sec_x)):
0112             vertices_processed += 1
0113
0114             vtx_begin = event_assoc_begin[vtx_idx]
0115             vtx_end = event_assoc_end[vtx_idx]
0116
0117             # Get the track indices for this vertex using correct slicing
0118             if vtx_begin < len(event_assoc_indices) and vtx_end <= len(event_assoc_indices):
0119                 vertex_track_indices = event_assoc_indices[vtx_begin:vtx_end]
0120                 n_tracks = len(vertex_track_indices)
0121                 n_tracks_per_vertex.append(n_tracks)
0122
0123                 print(f"Event {event_idx}, Vertex {vtx_idx}: {n_tracks} tracks, indices: {vertex_track_indices}")
0124
0125                 # Calculate decay length for this vertex
0126                 decay_2d = np.sqrt(event_sec_x[vtx_idx]**2 + event_sec_y[vtx_idx]**2)
0127                 decay_3d = np.sqrt(event_sec_x[vtx_idx]**2 + event_sec_y[vtx_idx]**2 + event_sec_z[vtx_idx]**2)
0128                 chi2_ndf = event_chi2[vtx_idx] / (event_ndf[vtx_idx] + 1e-10)
0129
0130                 # Focus on 2-track vertices (typical for kaon decays)
0131                 if n_tracks == 2:
0132                     two_track_vertices += 1
0133
0134                     idx1, idx2 = vertex_track_indices[0], vertex_track_indices[1]
0135
0136                     # Check if indices are valid for this event
0137                     if idx1 < len(event_px) and idx2 < len(event_px):
0138                         valid_associations += 1
0139
0140                         # Calculate invariant mass
0141                         p1 = [event_px[idx1], event_py[idx1], event_pz[idx1], event_energy[idx1]]
0142                         p2 = [event_px[idx2], event_py[idx2], event_pz[idx2], event_energy[idx2]]
0143
0144                         total_E = p1[3] + p2[3]
0145                         total_px = p1[0] + p2[0]
0146                         total_py = p1[1] + p2[1]
0147                         total_pz = p1[2] + p2[2]
0148
0149                         inv_mass_sq = total_E**2 - (total_px**2 + total_py**2 + total_pz**2)
0150
0151                         if inv_mass_sq > 0:
0152                             valid_invariant_masses += 1
0153                             inv_mass = np.sqrt(inv_mass_sq)
0154
0155                             # Apply kaon selection cuts
0156                             passes_vertex_quality = chi2_ndf < 3.0  # Good vertex fit
0157                             passes_decay_length = 0.05 < decay_2d < 2.0  # Reasonable decay length for kaons
0158
0159                             # Only store candidates that pass quality cuts
0160                             if passes_vertex_quality and passes_decay_length:
0161                                 passed_quality_cuts += 1
0162                                 invariant_masses.append(inv_mass)
0163                                 decay_lengths_2d.append(decay_2d)
0164                                 decay_lengths_3d.append(decay_3d)
0165                                 vertex_quality.append(chi2_ndf)
0166
0167                                 # Check if it's a kaon candidate
0168                                 if 0.4 <= inv_mass <= 0.6:
0169                                     kaon_candidates += 1
0170                                     print(f"  -> KAON CANDIDATE: mass={inv_mass:.4f} GeV, decay_length={decay_2d:.3f} mm, χ²/ndf={chi2_ndf:.2f}")
0171                             else:
0172                                 print(f"  -> Failed cuts: χ²/ndf={chi2_ndf:.2f} (cut<3.0), decay_length={decay_2d:.3f}mm (cut: 0.05-2.0mm)")
0173                         else:
0174                             print(f"  -> Invalid mass² = {inv_mass_sq}")
0175                     else:
0176                         print(f"  -> Invalid track indices: {idx1}, {idx2} (max: {len(event_px)-1})")
0177                 elif n_tracks > 0:
0178                     print(f"  -> {n_tracks}-track vertex (not 2-track)")
0179
0180     # Convert to numpy arrays
0181     invariant_masses = np.array(invariant_masses)
0182     decay_lengths_2d = np.array(decay_lengths_2d)
0183     decay_lengths_3d = np.array(decay_lengths_3d)
0184     vertex_quality = np.array(vertex_quality)
0185     n_tracks_per_vertex = np.array(n_tracks_per_vertex)
0186
0187     print(f"\n=== CUT FLOW RESULTS ===")
0188     print(f"Events processed:                {events_processed:>8}")
0189     print(f"Vertices processed:              {vertices_processed:>8}")
0190     print(f"Two-track vertices:              {two_track_vertices:>8} ({100*two_track_vertices/vertices_processed:.1f}% of vertices)")
0191     print(f"Valid associations:              {valid_associations:>8} ({100*valid_associations/two_track_vertices:.1f}% of 2-track vertices)")
0192     print(f"Valid invariant masses:          {valid_invariant_masses:>8} ({100*valid_invariant_masses/valid_associations:.1f}% of valid associations)")
0193     print(f"Passed quality cuts:             {passed_quality_cuts:>8} ({100*passed_quality_cuts/valid_invariant_masses:.1f}% of valid masses)")
0194     print(f"Kaon candidates (0.4-0.6 GeV):   {kaon_candidates:>8} ({100*kaon_candidates/passed_quality_cuts:.1f}% of quality candidates)")
0195
0196     if len(invariant_masses) > 0:
0197         print(f"\n=== ANALYSIS RESULTS (AFTER CUTS) ===")
0198         print(f"Applied cuts: χ²/ndf < 3.0, 0.05 < decay_length < 2.0 mm")
0199         print(f"Mass range: {np.min(invariant_masses):.4f} - {np.max(invariant_masses):.4f} GeV")
0200         print(f"Mean mass: {np.mean(invariant_masses):.4f} ± {np.std(invariant_masses):.4f} GeV")
0201         print(f"Decay length range: {np.min(decay_lengths_2d):.3f} - {np.max(decay_lengths_2d):.3f} mm")
0202         print(f"Mean decay length: {np.mean(decay_lengths_2d):.3f} ± {np.std(decay_lengths_2d):.3f} mm")
0203         print(f"Vertex quality range: {np.min(vertex_quality):.2f} - {np.max(vertex_quality):.2f}")
0204
0205         # Track multiplicity analysis
0206         print(f"\n=== TRACK MULTIPLICITY ===")
0207         unique_multiplicities, counts = np.unique(n_tracks_per_vertex, return_counts=True)
0208         for mult, count in zip(unique_multiplicities, counts):
0209             percentage = 100 * count / len(n_tracks_per_vertex)
0210             print(f"{mult} tracks per vertex: {count:>8} ({percentage:>5.1f}%)")
0211
0212         # Decay length cuts analysis
0213         print(f"\n=== DECAY LENGTH CUTS ANALYSIS ===")
0214         decay_cuts = {
0215             'Very Short (<0.1mm)': (0, 0.1),
0216             'Short (0.1-0.5mm)': (0.1, 0.5),
0217             'Medium (0.5-1.0mm)': (0.5, 1.0),
0218             'Long (>1.0mm)': (1.0, np.inf)
0219         }
0220
0221         for cut_name, (min_cut, max_cut) in decay_cuts.items():
0222             mask = (decay_lengths_2d >= min_cut) & (decay_lengths_2d < max_cut)
0223             masses = invariant_masses[mask]
0224             kaon_mask = (masses >= 0.45) & (masses <= 0.55)
0225
0226             print(f"{cut_name:<25}: {len(masses):>6} pairs, {np.sum(kaon_mask):>3} kaon candidates")
0227
0228         # Create plot
0229         fig, ax = plt.subplots(1, 1, figsize=(10, 8))
0230
0231         # Mass vs decay length (2D histogram)
0232         h = ax.hist2d(invariant_masses, decay_lengths_2d,
0233                      bins=[np.linspace(0, 1, 21), 15],
0234                      cmap='viridis', alpha=0.8, range=[[0, 1], None])
0235         plt.colorbar(h[3], ax=ax, label='Count')
0236         ax.axvline(0.4976, color='red', linestyle='--', linewidth=2, alpha=0.9, label='K⁰_S')
0237         ax.axvline(0.4937, color='orange', linestyle=':', linewidth=2, alpha=0.9, label='K±')
0238         ax.axvline(0.7755, color='purple', linestyle='-.', linewidth=2, alpha=0.9, label='ρ⁰')
0239         ax.set_xlabel('Invariant Mass (GeV)')
0240         ax.set_ylabel('2D Decay Length (mm)')
0241         ax.set_title('Mass vs Decay Length (2D Histogram)\n(0-1 GeV Mass Range)')
0242         ax.set_xlim(0, 1)
0243         ax.legend()
0244
0245         plt.tight_layout()
0246         plt.savefig(args.output, dpi=150, bbox_inches='tight')
0247         print(f"Plot saved to {args.output}")
0248
0249     else:
0250         print("No valid invariant masses calculated!")
0251
0252 if __name__ == "__main__":
0253     main()