ecal_gaps: add profile to the plots (#19)

eic · Apr 28, 2024 · 03850ad · 03850ad
1 parent 16fd651
commit 03850ad
Showing 1 changed file with 50 additions and 21 deletions.
diff --git a/benchmarks/ecal_gaps/ecal_gaps.org b/benchmarks/ecal_gaps/ecal_gaps.org
@@ -65,6 +65,9 @@ def filter_name(name):
 
 import functools
 
+axis_eta = bh.axis.Regular(200, -4, 4)
+axis_eta_coarse = bh.axis.Regular(100, -4, 4)
+
 @functools.cache
 def get_events(particle="e-", energy="20GeV", num_files=1):
     events = uproot.dask(
@@ -81,8 +84,6 @@ def get_events(particle="e-", energy="20GeV", num_files=1):
     eta = -np.log(np.tan(theta / 2))
     p = np.hypot(pt, events["MCParticles.momentum.z"][:,0])
 
-    axis_eta = bh.axis.Regular(200, -4, 4)
-
     hist_norm = dh.factory(
         eta,
         axes=(
@@ -103,27 +104,38 @@ def get_events(particle="e-", energy="20GeV", num_files=1):
     events["eta_thrown"] = eta
     events["weights"] = weights
 
-    return events, axis_eta
+    return events
 #+end_src
 
 #+begin_src jupyter-python
 particle = "e-"
 
 for energy in ["500MeV", "5GeV", "20GeV"]:
-    events, axis_eta = get_events(particle=particle, energy=energy)
+    events = get_events(particle=particle, energy=energy)
 
     for calo_name in ["EcalEndcapN", "EcalBarrelScFi", "EcalBarrelImaging", "EcalEndcapP"]:
         cond = ak.sum(events[f"{calo_name}RecHits.energy"], axis=1) > 10e-3 # GeV
 
-        hist = dh.factory(
-            events["eta_thrown"][cond],
-            (ak.sum(events[f"{calo_name}RecHits.energy"], axis=1) / events["p_thrown"])[cond],
-            axes=(
-                axis_eta,
-                bh.axis.Regular(100, 0.1, 1.1),
+        hist, profile = client.gather(client.compute([
+            dh.factory(
+                events["eta_thrown"][cond],
+                (ak.sum(events[f"{calo_name}RecHits.energy"], axis=1) / events["p_thrown"])[cond],
+                axes=(
+                    axis_eta,
+                    bh.axis.Regular(100, 0.1, 1.1),
+                ),
+                weights=events["weights"][cond],
+            ),
+            dh.factory(
+                events["eta_thrown"][cond],
+                sample=(ak.sum(events[f"{calo_name}RecHits.energy"], axis=1) / events["p_thrown"])[cond],
+                storage=bh.storage.WeightedMean(),
+                axes=(
+                    axis_eta_coarse,
+                ),
+                weights=events["weights"][cond],
             ),
-            weights=events["weights"][cond],
-        ).compute()
+        ]))
 
         cmap = plt.get_cmap(name="rainbow", lut=None).copy()
         cmap.set_under("none")
@@ -138,6 +150,9 @@ for energy in ["500MeV", "5GeV", "20GeV"]:
             ),
         )
         plt.colorbar()
+        std = np.sqrt(profile.variances())
+        cond = profile.values() > std
+        plt.errorbar(profile.axes[0].centers[cond], profile.values()[cond], yerr=std[cond], marker=".", markersize=2, color="black", ls="none", lw=0.6, capsize=1.)
         plt.xlabel(r"$\eta_{thrown}$", loc="right")
         plt.ylabel(r"$\sum E_{\mathrm{dep.}} / p_{\mathrm{thrown}}$", loc="top")
         plt.title(f"{energy} {particle} in {calo_name}")
@@ -151,23 +166,34 @@ for energy in ["500MeV", "5GeV", "20GeV"]:
 particle = "e-"
 
 for energy in ["500MeV", "5GeV", "20GeV"]:
-    events, axis_eta = get_events(particle=particle, energy=energy)
+    events = get_events(particle=particle, energy=energy)
 
     calos = ["EcalEndcapN", "EcalBarrelScFi", "EcalEndcapP"]
     total_energy = sum([
         ak.sum(events[f"{calo_name}RecHits.energy"], axis=1)
         for calo_name in calos
     ])
 
-    hist = dh.factory(
-        events["eta_thrown"],
-        total_energy / events["p_thrown"],
-        axes=(
-            axis_eta,
-            bh.axis.Regular(100, 0.0, 1.1),
+    hist, profile = client.gather(client.compute([
+        dh.factory(
+            events["eta_thrown"],
+            total_energy / events["p_thrown"],
+            axes=(
+                axis_eta,
+                bh.axis.Regular(100, 0.1, 1.1),
+            ),
+            weights=events["weights"],
         ),
-        weights=events["weights"],
-    ).compute()
+        dh.factory(
+            events["eta_thrown"],
+            sample=total_energy / events["p_thrown"],
+            storage=bh.storage.WeightedMean(),
+            axes=(
+                axis_eta_coarse,
+            ),
+            weights=events["weights"],
+        ),
+    ]))
 
     cmap = plt.get_cmap(name="rainbow", lut=None).copy()
     cmap.set_under("none")
@@ -182,6 +208,9 @@ for energy in ["500MeV", "5GeV", "20GeV"]:
         ),
     )
     plt.colorbar()
+    std = np.sqrt(profile.variances())
+    cond = profile.values() > std
+    plt.errorbar(profile.axes[0].centers[cond], profile.values()[cond], yerr=std[cond], marker=".", markersize=2, color="black", ls="none", lw=0.6, capsize=1.)
     plt.xlabel(r"$\eta_{thrown}$", loc="right")
     plt.ylabel(r"$\sum E_{\mathrm{dep.}} / p_{\mathrm{thrown}}$", loc="top")
     plt.title(f"{energy} {particle}\n" + "+".join(calos))