<jupyter_start><jupyter_text>Histogram of the Dimuon Mass SpectrumThis implements the dimuon mass spectrum analysis, a "Hello, World!" example for data analysis in High Energy Physics. It is intended as a technology demonstrator for the use Apache Spark for High Energy Physics.The workload and data:  - The input data is a series of candidate muon events.   - The job output is a histogram of the dimuon mass spectrum, where several peaks (resonances) can be identified corresponding to well-know particles (e.g. the Z boson at 91 Gev).  - The computation is based on https://root.cern.ch/doc/master/df102__NanoAODDimuonAnalysis_8C.html and CERN open data from the CMS collaboration linked there.   - See also https://github.com/LucaCanali/Miscellaneous/Spark_Physics  Author and contact: <EMAIL>  January, 2022  Dimuon mass spectrum calculation with Spark UDF using MapInArrow<jupyter_code>import findspark
findspark.init("/home/luca/Spark/spark-3.3.0-SNAPSHOT-bin-20220214")

from pyspark.sql import SparkSession
spark = (SparkSession.builder
         .appName("dimuon mass")
         .master("local[4]")
         .config("spark.driver.memory", "2g")
         .config("spark.sql.orc.enableNestedColumnVectorizedReader", "true")
         .getOrCreate()
        )
# Read data with the muon candidate events
# download data as detailed at https://github.com/LucaCanali/Miscellaneous/Spark_Physics

path = "/home/luca/DoubleMuParked/"
df_muons = spark.read.orc(path + "Run2012BC_DoubleMuParked_Muons.orc")

df_muons.printSchema()
print(f"Number of events: {df_muons.count()}")
# Apply filters to the input data
# - select only events with 2 muons
# - select only events where the 2 muons have opposite charge

df_muons = df_muons.filter("nMuon==2").filter("Muon_charge[0] != Muon_charge[1]")

# Add the column Dimuon_mass to the DataFrame
# it will be used as output column, to work around the limitation that 
# mapInArrow needs to have the same input and output schema.

df_muons = df_muons.selectExpr("Muon_pt[0] pt0", "Muon_eta[0] eta0","Muon_phi[0] phi0", 
                               "Muon_pt[1] pt1", "Muon_eta[1] eta1","Muon_phi[1] phi1",
                               "float(0.0) as Dimuon_mass")
# this requires pip install awkward and pip install vector
import awkward as ak
import vector
from numpy import cosh, cos, sqrt

# This is where the dimuon mass is computed
# the candidate muon events and summed as 4-vectors in the Pt, eta, phi, m space
# and the invarian mass is computed
# we use awkward arrray, serializing directly with arrow

def UDF_dimuonmass(iterator):
    vector.register_awkward()
    for batch in iterator:
        chunk = ak.from_arrow(batch)
        dimuon_mass = sqrt( 2 * chunk.pt0 * chunk.pt1 *
                      ( cosh (chunk.eta0 - chunk.eta1)
                        - cos(chunk.phi0 - chunk.phi1) ))
        out = ak.zip({"pt0": None, "eta0": None, "phi0": None,
                      "pt1": None, "eta1": None, "phi1": None,
                      "Dimoun_mass": dimuon_mass }, depth_limit=1)
        yield from ak.to_arrow_table(out).to_batches()
# apply the UDF with arrow to compute the dimuon mass spectrum

df_with_dimuonmass = df_muons.mapInArrow(UDF_dimuonmass, df_muons.schema)
# Compute the histogram of dimuon mass values
# The Spark function "width_bucket" is used to generate the histogram bucket number
# a groupBy operation with count is used to fill the histogram
# The result is a histogram with bins value and counts foreach bin (N_events)

min_val = 0.25
max_val = 300
num_bins = 30000
step = (max_val - min_val) / num_bins

histogram_data = ( 
    df_with_dimuonmass
        .selectExpr(f"width_bucket(Dimuon_mass, {min_val}, {max_val}, {num_bins}) as bucket") 
        .groupBy("bucket")
        .count()
        .orderBy("bucket")
    )

# convert bucket number to the corresponding dimoun mass value
histogram_data = histogram_data.selectExpr(f"round({min_val} + (bucket - 1) * {step},2) as value", "count as N_events")
# The action toPandas() here triggers the computation.
# Histogram data is fetched into the driver as a Pandas Dataframe.

%time histogram_data_pandas=histogram_data.toPandas()

import matplotlib.pyplot as plt 
plt.style.use('seaborn-darkgrid')
plt.rcParams.update({'font.size': 20, 'figure.figsize': [14,10]})

f, ax = plt.subplots()

# cut the first and last bin
x = histogram_data_pandas.iloc[1:-1]["value"]
y = histogram_data_pandas.iloc[1:-1]["N_events"]

# line plot
ax.plot(x, y, '-')

# the plot is in log-log axis to better show the peaks
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlim(min_val, max_val)
ax.set_ylim(1, 6e5)

ax.set_xlabel('$m_{dimuon}$ (GeV)')
ax.set_ylabel('Number of Events')
ax.set_title("Distribution of the Dimuon Mass Spectrum")

# Label for the resonances spectrum peaks
txt_opts = {'horizontalalignment': 'center',
            'verticalalignment': 'center',
            'transform': ax.transAxes}

plt.text(0.85, 0.75, 'Z', **txt_opts)
plt.text(0.55, 0.77, r"$\Upsilon$(1,2,3S)", **txt_opts)
plt.text(0.37, 0.95, r"J/$\Psi$", **txt_opts)
plt.text(0.40, 0.77, r"$\Psi$'", **txt_opts)
plt.text(0.22, 0.80, r"$\phi$", **txt_opts)
plt.text(0.16, 0.83, r"$\rho,\omega$", **txt_opts)
plt.text(0.11, 0.78, r"$\eta$", **txt_opts);

plt.show()
spark.stop()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11826741996233521, 'fraction_numerical': 0.024858757062146894, 'mean_word_length': 3.815049864007253, 'pattern_counts': {'":': 7, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '186995', 'license': 'Apache-2.0', 'n_tokens_mistral': 1917, 'n_tokens_neox': 1797, 'n_words': 533, 'path': 'Spark_Physics/Dimuon_mass_spectrum/6a.Dimuon_mass_spectrum_histogram_Spark_UDF_MapInArrow_flattened_data_ORC_vectorized.ipynb', 'repo_name': 'Code360In/Miscellaneous', 'size': '5322'}