Skip to content
Snippets Groups Projects
Commit 60911ba0 authored by Chao Peng's avatar Chao Peng
Browse files

modify clustering to use the latest components

parent a9d56427
No related branches found
No related tags found
1 merge request!107modify clustering to use the latest components
Hide whitespace changes
Inline Side-by-side
Showing
with 557 additions and 203 deletions
benchmarks/clustering/barrel_clusters.sh deleted 100644 → 0
+ 0
− 66
View file @ a9d56427
#!/bin/bash
print_env.sh
## To run the reconstruction, we need the following global variables:
## - JUGGLER_INSTALL_PREFIX: Install prefix for Juggler (simu/recon)
## - JUGGLER_DETECTOR: the detector package we want to use for this benchmark
## - JUGGLER_DETECTOR_VERSION: the detector package we want to use for this benchmark
## - DETECTOR_PATH: full path to the detector definitions
##
## You can ready options/env.sh for more in-depth explanations of the variables
## and how they can be controlled.
# deprecated
export JUGGLER_DETECTOR_PATH=${DETECTOR_PATH}
if [[ ! -n "${JUGGLER_N_EVENTS}" ]] ; then
export JUGGLER_N_EVENTS=1
fi
if [[ ! -n "${CB_EMCAL_SAMP_FRAC}" ]] ; then
export CB_EMCAL_SAMP_FRAC=1.0
fi
export JUGGLER_FILE_NAME_TAG="barrel_clusters"
export JUGGLER_GEN_FILE="${JUGGLER_FILE_NAME_TAG}.hepmc"
export JUGGLER_SIM_FILE="sim_${JUGGLER_FILE_NAME_TAG}.root"
export JUGGLER_REC_FILE="rec_${JUGGLER_FILE_NAME_TAG}.root"
echo "JUGGLER_N_EVENTS = ${JUGGLER_N_EVENTS}"
echo "JUGGLER_DETECTOR = ${JUGGLER_DETECTOR}"
root -b -q "benchmarks/clustering/scripts/gen_central_pions.cxx(${JUGGLER_N_EVENTS}, \"${JUGGLER_FILE_NAME_TAG}.hepmc\")"
### run geant4 simulations
npsim --runType batch \
--part.minimalKineticEnergy 1000*GeV \
-v WARNING \
--numberOfEvents ${JUGGLER_N_EVENTS} \
--compactFile ${DETECTOR_PATH}/${JUGGLER_DETECTOR}.xml \
--inputFiles ${JUGGLER_FILE_NAME_TAG}.hepmc \
--outputFile ${JUGGLER_SIM_FILE}
if [[ "$?" -ne "0" ]] ; then
echo "ERROR running npdet"
exit 1
fi
# Need to figure out how to pass file name to juggler from the commandline
gaudirun.py benchmarks/clustering/options/fullcalo_clustering.py
if [[ "$?" -ne "0" ]] ; then
echo "ERROR running juggler"
exit 1
fi
mkdir -p results/clustering
#root -b -q "benchmarks/clustering/scripts/barrel_clusters.cxx(\"${JUGGLER_REC_FILE}\")"
#root_filesize=$(stat --format=%s "${JUGGLER_DETECTOR}/${JUGGLER_REC_FILE}")
#if [[ "${JUGGLER_N_EVENTS}" -lt "500" ]] ; then
# # file must be less than 10 MB to upload
# if [[ "${root_filesize}" -lt "10000000" ]] ; then
# cp ${JUGGLER_REC_FILE} results/clustering/.
# fi
#fi
benchmarks/clustering/config.yml
+ 1
− 1
View file @ 60911ba0
......@@ -3,7 +3,7 @@ clustering:process :
stage: process
timeout: 8 hour
script:
- bash benchmarks/clustering/barrel_clusters.sh
- bash benchmarks/clustering/full_cal_clusters.sh -t fullcalo -n 1000 -p "pion-" --pmin 5 --pmax 5
clustering:results:
extends: .rec_benchmark
......
benchmarks/clustering/full_cal_clusters.sh 0 → 100644
+ 137
− 0
View file @ 60911ba0
#!/bin/bash
print_env.sh
function print_the_help {
echo "USAGE: ${0} -n <nevents> -t <nametag> -p <particle> "
echo " OPTIONS: "
echo " -n,--nevents Number of events"
echo " -t,--nametag name tag"
echo " -p,--particle particle type"
echo " --pmin minimum particle momentum (GeV)"
echo " --pmax maximum particle momentum (GeV)"
echo " allowed types: pion0, pion+, pion-, kaon0, kaon+, kaon-, proton, neutron, electron, positron, photon"
exit
}
POSITIONAL=()
while [[ $# -gt 0 ]]
do
key="$1"
case $key in
-h|--help)
shift # past argument
print_the_help
;;
-t|--nametag)
nametag="$2"
shift # past argument
shift # past value
;;
-p|--particle)
particle="$2"
shift # past argument
shift # past value
;;
-n|--nevents)
export FULL_CAL_NUMEV="$2"
shift # past argument
shift # past value
;;
--pmin)
export FULL_CAL_PMIN="$2"
shift
shift
;;
--pmax)
export FULL_CAL_PMAX="$2"
shift
shift
;;
*) # unknown option
#POSITIONAL+=("$1") # save it in an array for later
echo "unknown option $1"
print_the_help
shift # past argument
;;
esac
done
set -- "${POSITIONAL[@]}" # restore positional parameters
export FULL_CAL_COMPACT_PATH=${DETECTOR_PATH}/${JUGGLER_DETECTOR}.xml
if [[ ! -n "${FULL_CAL_NUMEV}" ]] ; then
export FULL_CAL_NUMEV=1000
fi
if [[ ! -n "${FULL_CAL_PMIN}" ]] ; then
export FULL_CAL_ENERGY=1.0
fi
if [[ ! -n "${FULL_CAL_PMAX}" ]] ; then
export FULL_CAL_ENERGY=10.0
fi
export FULL_CAL_NAME_TAG="${nametag}"
export FULL_CAL_GEN_FILE="${FULL_CAL_NAME_TAG}.hepmc"
export FULL_CAL_SIM_FILE="sim_${FULL_CAL_NAME_TAG}.root"
export FULL_CAL_REC_FILE="rec_${FULL_CAL_NAME_TAG}.root"
echo "FULL_CAL_NUMEV = ${FULL_CAL_NUMEV}"
echo "FULL_CAL_COMPACT_PATH = ${FULL_CAL_COMPACT_PATH}"
# Generate the input events
python benchmarks/imaging_ecal/scripts/gen_particles.py ${FULL_CAL_GEN_FILE} -n ${FULL_CAL_NUMEV}\
--angmin 5 --angmax 175 --phmin 0 --phmax 360 --pmin ${FULL_CAL_PMIN} --pmax ${FULL_CAL_PMAX}\
--particles="${particle}"
if [[ "$?" -ne "0" ]] ; then
echo "ERROR running script: generating input events"
exit 1
fi
ls -lh ${FULL_CAL_GEN_FILE}
# Run geant4 simulations
npsim --runType batch \
-v WARNING \
--part.minimalKineticEnergy "1*TeV" \
--numberOfEvents ${FULL_CAL_NUMEV} \
--compactFile ${FULL_CAL_COMPACT_PATH} \
--inputFiles ${FULL_CAL_GEN_FILE} \
--outputFile ${FULL_CAL_SIM_FILE}
if [[ "$?" -ne "0" ]] ; then
echo "ERROR running npdet"
exit 1
fi
rootls -t "${FULL_CAL_SIM_FILE}"
# Directory for results
mkdir -p results
FULL_CAL_OPTION_DIR=benchmarks/clustering/options
FULL_CAL_SCRIPT_DIR=benchmarks/clustering/scripts
# Run Juggler
xenv -x ${JUGGLER_INSTALL_PREFIX}/Juggler.xenv \
gaudirun.py ${FULL_CAL_OPTION_DIR}/full_cal_clustering.py
if [[ "$?" -ne "0" ]] ; then
echo "ERROR running juggler"
exit 1
fi
# run analysis scripts
python ${FULL_CAL_SCRIPT_DIR}/cluster_plots.py ${FULL_CAL_REC_FILE} -o results
root_filesize=$(stat --format=%s "${FULL_CAL_REC_FILE}")
if [[ "${FULL_CAL_NUMEV}" -lt "500" ]] ; then
# file must be less than 10 MB to upload
if [[ "${root_filesize}" -lt "10000000" ]] ; then
cp ${FULL_CAL_REC_FILE} results/.
fi
fi
benchmarks/clustering/options/full_cal_clustering.py 0 → 100644
+ 173
− 0
View file @ 60911ba0
from Gaudi.Configuration import *
import os
import ROOT
from GaudiKernel.DataObjectHandleBase import DataObjectHandleBase
from Configurables import ApplicationMgr, EICDataSvc, PodioInput, PodioOutput, GeoSvc
from GaudiKernel.SystemOfUnits import MeV, GeV, mm, cm, mrad
detector_name = str(os.environ.get("JUGGLER_DETECTOR", "athena"))
detector_path = str(os.environ.get("DETECTOR_PATH", "."))
compact_path = os.path.join(detector_path, detector_name)
# get sampling fractions from system environment variable, 1.0 by default
cb_ecal_sf = float(os.environ.get("CB_EMCAL_SAMP_FRAC", 1.0))
cb_hcal_sf = float(os.environ.get("CB_HCAL_SAMP_FRAC", 1.0))
# input and output
input_sims = [f.strip() for f in str.split(os.environ["FULL_CAL_SIM_FILE"], ",") if f.strip()]
output_rec = str(os.environ["FULL_CAL_REC_FILE"])
n_events = int(os.environ["FULL_CAL_NUMEV"])
# geometry service
geo_service = GeoSvc("GeoSvc", detectors=["{}.xml".format(compact_path)])
# data service
podioevent = EICDataSvc("EventDataSvc", inputs=input_sims)
# juggler components
from Configurables import Jug__Base__InputCopier_dd4pod__Geant4ParticleCollection_dd4pod__Geant4ParticleCollection_ as MCCopier
from Configurables import Jug__Digi__CalorimeterHitDigi as CalHitDigi
from Configurables import Jug__Reco__CalorimeterHitReco as CalHitReco
from Configurables import Jug__Reco__CalorimeterHitsMerger as CalHitsMerger
from Configurables import Jug__Reco__CalorimeterIslandCluster as IslandCluster
from Configurables import Jug__Reco__ImagingPixelReco as ImCalPixelReco
from Configurables import Jug__Reco__ImagingTopoCluster as ImagingCluster
from Configurables import Jug__Reco__ClusterRecoCoG as RecoCoG
from Configurables import Jug__Reco__ImagingClusterReco as ImagingClusterReco
# branches needed from simulation root file
sim_coll = [
"mcparticles",
"CrystalEcalHits",
"EcalBarrelHits",
"HcalBarrelHits",
]
# input and output
podin = PodioInput("PodioReader", collections=sim_coll)
podout = PodioOutput("out", filename=output_rec)
# copier needed to get around input --> output copy bug. So truth (mcparticles) can be saved in output file
copier = MCCopier("MCCopier", inputCollection="mcparticles", outputCollection="mcparticles2")
# Crystal Endcap Ecal
ce_ecal_digi = CalHitDigi("ce_ecal_digi",
inputHitCollection="CrystalEcalHits",
outputHitCollection="CrystalEcalHitsDigi",
energyResolutions=[0., 0.02, 0.],
dynamicRangeADC=5.*GeV, # digi settings must match with reco
capacityADC=32768,
pedestalMean=400,
pedestalSigma=3)
ce_ecal_reco = CalHitReco("ce_ecal_reco",
inputHitCollection="CrystalEcalHitsDigi",
outputHitCollection="CrystalEcalHitsReco",
dynamicRangeADC=5.*GeV,
capacityADC=32768,
pedestalMean=400,
pedestalSigma=3,
thresholdFactor=4, # 4 sigma
minimumHitEdep=1.0*MeV) # discard low energy hits
ce_ecal_cl = IslandCluster("ce_ecal_cl",
# OutputLevel=DEBUG,
inputHitCollection="CrystalEcalHitsReco",
outputClusterCollection="CrystalEcalClusters",
splitHitCollection="CrystalEcalHitsSplit",
splitCluster=False,
minClusterCenterEdep=30*MeV,
groupRanges=[2.2*cm, 2.2*cm])
ce_ecal_clreco = RecoCoG("ce_ecal_clreco",
clusterCollection="CrystalEcalClusters",
logWeightBase=4.6)
# Central Barrel Ecal (Imaging Cal.)
cb_ecal_digi = CalHitDigi("cb_ecal_digi",
inputHitCollection="EcalBarrelHits",
outputHitCollection="EcalBarrelHitsDigi",
energyResolutions=[0., 0.02, 0.], # 2% flat resolution
dynamicRangeADC=3*MeV,
capacityADC=8192,
pedestalMean=400,
pedestalSigma=20) # about 6 keV
cb_ecal_reco = ImCalPixelReco("cb_ecal_reco",
inputHitCollection="EcalBarrelHitsDigi",
outputHitCollection="EcalBarrelHitsReco",
dynamicRangeADC=3.*MeV,
capacityADC=8192,
pedestalMean=400,
pedestalSigma=20,
thresholdFactor=3, # about 20 keV
readoutClass="EcalBarrelHits", # readout class
layerField="layer", # field to get layer id
sectorField="module") # field to get sector id
cb_ecal_cl = ImagingCluster("cb_ecal_cl",
inputHitCollection="EcalBarrelHitsReco",
outputClusterCollection="EcalBarrelClusters",
localRanges=[2.*mm, 2*mm], # same layer
adjLayerRanges=[10*mrad, 10*mrad], # adjacent layer
adjLayerDiff=2, # id diff for adjacent layer
adjSectorDist=3.*cm) # different sector
cb_ecal_clreco = ImagingClusterReco("cb_ecal_clreco",
inputClusterCollection="EcalBarrelClusters",
outputLayerCollection="EcalBarrelLayers")
# Central Barrel Hcal
cb_hcal_digi = CalHitDigi("cb_hcal_digi",
inputHitCollection="HcalBarrelHits",
outputHitCollection="HcalBarrelHitsDigi",
dynamicRangeADC=50.*MeV,
capacityADC=32768,
pedestalMean=400,
pedestalSigma=10)
cb_hcal_reco = CalHitReco("cb_hcal_reco",
inputHitCollection="HcalBarrelHitsDigi",
outputHitCollection="HcalBarrelHitsReco",
dynamicRangeADC=50.*MeV,
capacityADC=32768,
pedestalMean=400,
pedestalSigma=10,
thresholdFactor=5.0)
# merge hits in different layer (projection to local x-y plane)
cb_hcal_merger = CalHitsMerger("cb_hcal_merger",
fields=["layer", "slice"],
fieldRefNumbers=[1, 0],
inputHitCollection="HcalBarrelHitsReco",
outputHitCollection="HcalBarrelHitsRecoXY")
cb_hcal_cl = IslandCluster("cb_hcal_cl",
inputHitCollection="HcalBarrelHitsRecoXY",
outputClusterCollection="HcalBarrelClusters",
splitHitCollection="HcalBarrelHitsSplit",
splitCluster=False,
minClusterCenterEdep=30.*MeV,
groupRanges=[15.*cm, 15.*cm])
cb_hcal_clreco = RecoCoG("cb_hcal_clreco",
clusterCollection="HcalBarrelClusters",
logWeightBase=6.2,
samplingFraction=cb_hcal_sf)
podout.outputCommands = ["keep *"]
ApplicationMgr(
TopAlg = [podin, copier,
ce_ecal_digi, ce_ecal_reco, ce_ecal_cl, ce_ecal_clreco,
cb_ecal_digi, cb_ecal_reco, cb_ecal_cl, cb_ecal_clreco,
cb_hcal_digi, cb_hcal_reco, cb_hcal_merger, cb_hcal_cl, cb_hcal_clreco,
podout],
EvtSel = 'NONE',
EvtMax = n_events,
ExtSvc = [podioevent],
OutputLevel=DEBUG
)
benchmarks/clustering/options/fullcalo_clustering.py deleted 100644 → 0
+ 0
− 135
View file @ a9d56427
from Gaudi.Configuration import *
import os
from GaudiKernel.DataObjectHandleBase import DataObjectHandleBase
from Configurables import ApplicationMgr, EICDataSvc, PodioOutput, GeoSvc
from GaudiKernel import SystemOfUnits as units
detector_name = "topside"
if "JUGGLER_DETECTOR" in os.environ :
detector_name = str(os.environ["JUGGLER_DETECTOR"])
if "JUGGLER_DETECTOR_PATH" in os.environ :
detector_name = str(os.environ["JUGGLER_DETECTOR_PATH"]) + "/" + detector_name
# get sampling fraction from system environment variable, 1.0 by default
sf = 1.0
if "CB_EMCAL_SAMP_FRAC" in os.environ :
sf = str(os.environ["CB_EMCAL_SAMP_FRAC"])
# todo add checks
input_sim_file = str(os.environ["JUGGLER_SIM_FILE"])
output_rec_file = str(os.environ["JUGGLER_REC_FILE"])
n_events = str(os.environ["JUGGLER_N_EVENTS"])
geo_service = GeoSvc("GeoSvc", detectors=["{}.xml".format(detector_name)])
podioevent = EICDataSvc("EventDataSvc", inputs=[input_sim_file])
from Configurables import PodioInput
from Configurables import Jug__Base__InputCopier_dd4pod__Geant4ParticleCollection_dd4pod__Geant4ParticleCollection_ as MCCopier
from Configurables import Jug__Base__InputCopier_dd4pod__CalorimeterHitCollection_dd4pod__CalorimeterHitCollection_ as CalCopier
from Configurables import Jug__Digi__HadronicCalDigi as HadCalorimeterDigi
from Configurables import Jug__Digi__EMCalorimeterDigi as EMCalorimeterDigi
from Configurables import Jug__Digi__CrystalEndcapsDigi as CrystalEndcapsDigi
from Configurables import Jug__Reco__CrystalEndcapsReco as CrystalEndcapsReco
from Configurables import Jug__Reco__EMCalReconstruction as EMCalReconstruction
from Configurables import Jug__Reco__HCalReconstruction as HCalReconstruction
from Configurables import Jug__Reco__SimpleClustering as SimpleClustering
from Configurables import Jug__Reco__CalorimeterIslandCluster as IslandCluster
from Configurables import Jug__Reco__ClusterRecoCoG as RecoCoG
podioinput = PodioInput("PodioReader", collections=["mcparticles","CrystalEcalHits","EcalBarrelHits","HcalBarrelHits"], OutputLevel=DEBUG)
## copiers to get around input --> output copy bug. Note the "2" appended to the output collection.
copier = MCCopier("MCCopier",
inputCollection="mcparticles",
outputCollection="mcparticles2")
calcopier = CalCopier("CalCopier",
inputCollection="CrystalEcalHits",
outputCollection="CrystalEcalHits2")
##raw hits
emcaldigi = CrystalEndcapsDigi("ecal_digi",
inputHitCollection="CrystalEcalHits",
outputHitCollection="RawDigiEcalHits")
ecdigi = EMCalorimeterDigi("ec_barrel_digi",
inputHitCollection="EcalBarrelHits",
outputHitCollection="RawEcalBarrelHits")
hcaldigi = HadCalorimeterDigi("hcal_barrel_digi",
inputHitCollection="HcalBarrelHits",
outputHitCollection="RawHcalBarrelHits")
#digitized hits
crystal_ec_reco = CrystalEndcapsReco("crystal_ec_reco",
inputHitCollection="RawDigiEcalHits",
outputHitCollection="RecoEcalHits",
minModuleEdep=1.0*units.MeV)
ecal_reco = EMCalReconstruction("ecal_reco",
inputHitCollection="RawEcalBarrelHits",
outputHitCollection="RecEcalBarrelHits",
samplingFraction=0.25,
minModuleEdep=0.0*units.MeV)
hcal_reco = HCalReconstruction("hcal_reco",
inputHitCollection="RawHcalBarrelHits",
outputHitCollection="RecHcalBarrelHits",
samplingFraction=0.25,
minModuleEdep=0.0*units.MeV)
#clusters
ec_barrel_cluster = IslandCluster("ec_barrel_cluster",
inputHitCollection="RecEcalBarrelHits",
outputClusterCollection="EcalBarrelClusters",
splitHitCollection="splitEcalBarrelHitCollection",
minClusterCenterEdep=1*units.MeV,
groupRange=2.0,
OutputLevel=DEBUG)
crystal_ec_cluster = IslandCluster("crystal_ec_cluster",
inputHitCollection="RecoEcalHits",
outputClusterCollection="EcalClusters",
minClusterCenterEdep=30*units.MeV,
groupRange=2.0,
OutputLevel=DEBUG)
simple_cluster = SimpleClustering("simple_cluster",
inputHitCollection="RecEcalBarrelHits",
outputClusters="SimpleClusters",
OutputLevel=DEBUG)
# finalizing clustering (center-of-gravity step)
ec_barrel_clusterreco = RecoCoG("ec_barrel_clusterreco",
clusterCollection="EcalBarrelClusters",
logWeightBase=6.2,
samplingFraction=sf)
clusterreco = RecoCoG("cluster_reco",
clusterCollection="EcalClusters",
logWeightBase=4.2,
moduleDimZName="CrystalBox_z_length",
samplingFraction=sf,
OutputLevel=DEBUG)
out = PodioOutput("out", filename=output_rec_file)
out.outputCommands = ["keep *"]
ApplicationMgr(
TopAlg = [podioinput, copier, calcopier,
ecdigi, emcaldigi,hcaldigi,
crystal_ec_reco, ecal_reco, hcal_reco,
#ec_barrel_cluster, crystal_ec_cluster, ec_barrel_clusterreco, clusterreco,
#simple_cluster,
out],
EvtSel = 'NONE',
EvtMax = n_events,
ExtSvc = [podioevent],
OutputLevel=DEBUG
)
benchmarks/clustering/scripts/barrel_clusters.cxx
+ 8
− 1
View file @ 60911ba0
......@@ -39,6 +39,13 @@ auto pt = [](std::vector<dd4pod::Geant4ParticleData> const& in) {
}
return result;
};
auto pid = [](ROOT::VecOps::RVec<dd4pod::Geant4ParticleData> const& in) {
std::vector<int> result;
for (size_t i = 0; i < in.size(); ++i) {
result.push_back(in[i].pdgID);
}
return result;
};
auto eta = [](ROOT::VecOps::RVec<dd4pod::Geant4ParticleData> const& in) {
std::vector<float> result;
ROOT::Math::PxPyPzMVector lv;
......@@ -79,7 +86,7 @@ auto delta_E = [](std::vector<ROOT::Math::PxPyPzMVector> const& thrown, const st
};
void barrel_clusters(const char* in_fname = "topside/rec_barrel_clusters.root")
int barrel_clusters(const char* in_fname = "topside/rec_barrel_clusters.root")
{
ROOT::EnableImplicitMT();
ROOT::RDataFrame df("events", in_fname);
......
benchmarks/clustering/scripts/cluster_plots.py 0 → 100644
+ 134
− 0
View file @ 60911ba0
'''
A simple analysis script to extract some basic info of Monte-Carlo particles and clusters
'''
import os
import ROOT
import pandas as pd
import numpy as np
import argparse
from matplotlib import pyplot as plt
import matplotlib.ticker as ticker
# load root macros, input is an argument string
def load_root_macros(arg_macros):
for path in arg_macros.split(','):
path = path.strip()
if os.path.exists(path):
gROOT.Macro(path)
print('Loading root macro: \'{}\' loaded.'.format(path))
else:
print('Loading root macros: \'{}\' does not exist, skip loading it.'.format(path))
def thrown_particles_figure(df, save):
# define truth particle info
df = df.Define('isThrown', 'mcparticles2.genStatus == 1')\
.Define('thrownParticles', 'mcparticles2[isThrown]')\
.Define('thrownP', 'fourvec(thrownParticles)')\
.Define('thrownPID', 'pid(thrownParticles)')\
.Define('thrownEta', 'eta(thrownParticles)')\
.Define('thrownTheta', 'theta(thrownP)')\
.Define('thrownMomentum', 'momentum(thrownP)')
# figure
fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
# pid info need some special treatment
pids = df.AsNumpy(['thrownPID'])['thrownPID']
# unpack vectors
pids = np.asarray([v for vec in pids for v in vec], dtype=int)
# build a dictionary for particle id and particle name
pdgbase = ROOT.TDatabasePDG()
pdict = {pid: pdgbase.GetParticle(int(pid)).GetName() for pid in np.unique(pids)}
pmap = {pid: i*2 for i, pid in enumerate(list(pdict))}
# convert pid to index in dictionary
pmaps = [pmap[i] for i in pids]
ax = axs.flat[0]
ax.hist(pmaps, bins=np.arange(-4, len(pmap)*2 + 4), align='left', ec='k')
ax.set_ylabel('Particle Counts', fontsize=24)
ax.tick_params(labelsize=22)
ax.set_axisbelow(True)
ax.grid(linestyle=':', which='both', axis='y')
ax.xaxis.set_major_locator(ticker.FixedLocator(list(pmap.values())))
ax.set_xticklabels(list(pdict.values()))
# kinematics
data = df.AsNumpy(['thrownMomentum', 'thrownTheta', 'thrownEta'])
labels = {
'thrownMomentum': (r'$p$ (GeV)', 'Counts'),
'thrownTheta': (r'$\theta$ (degree)', 'Counts'),
'thrownEta': (r'$\eta$', 'Counts'),
}
for ax, (name, vals) in zip(axs.flat[1:], data.items()):
hvals = [v for vec in vals for v in vec]
ax.hist(hvals, bins=50, ec='k')
ax.tick_params(labelsize=22, direction='in', which='both')
ax.grid(linestyle=':', which='both')
label = labels[name]
ax.set_axisbelow(True)
ax.set_xlabel(label[0], fontsize=24)
ax.set_ylabel(label[1], fontsize=24)
fig.text(0.5, 0.95, 'Thrown Particles', ha='center', fontsize=24)
fig.savefig(save)
plt.close(fig)
def general_clusters_figure(df, collection, save):
data = df.AsNumpy([
'{}.nhits'.format(collection),
'{}.energy'.format(collection),
'{}.polar.theta'.format(collection),
'{}.polar.phi'.format(collection),
])
# figure
fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
labels = [
('Number of Hits', 'Counts'),
('Energy (MeV)', 'Counts'),
('Theta (rad)', 'Counts'),
('Phi (rad)', 'Counts'),
]
for ax, label, (name, vals) in zip(axs.flat, labels, data.items()):
hvals = [v for vec in vals for v in vec]
ax.hist(hvals, bins=50, ec='k')
ax.tick_params(labelsize=22, direction='in', which='both')
ax.grid(linestyle=':', which='both')
ax.set_axisbelow(True)
ax.set_xlabel(label[0], fontsize=24)
ax.set_ylabel(label[1], fontsize=24)
fig.text(0.5, 0.95, collection, ha='center', fontsize=24)
fig.savefig(save)
plt.close(fig)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('file', help='Path to reconstruction file.')
parser.add_argument('-o', dest='outdir', default='.', help='Output directory.')
parser.add_argument('-m', '--macros', dest='macros', default='rootlogon.C',
help='Macro files to be loaded by root, separated by \",\".')
args = parser.parse_args()
# multi-threading for RDataFrame
nthreads = os.cpu_count()//2
ROOT.ROOT.EnableImplicitMT(nthreads)
print('CPU available {}, using {}'.format(os.cpu_count(), nthreads))
# declare some functions from c++ script, needed for data frame processing
script_dir = os.path.dirname(os.path.realpath(__file__))
fcode = open(os.path.join(script_dir, 'barrel_clusters.cxx')).read()
ROOT.gInterpreter.Declare(fcode)
os.makedirs(args.outdir, exist_ok=True)
# load macros (add libraries/headers root cannot automatically found here)
load_root_macros(args.macros)
rdf = ROOT.RDataFrame('events', args.file)
thrown_particles_figure(rdf, save=os.path.join(args.outdir, 'thrown_particles.png'))
general_clusters_figure(rdf, collection='CrystalEcalClusters', save=os.path.join(args.outdir, 'crystal_ecal_clusters.png'))
general_clusters_figure(rdf, collection='EcalBarrelClusters', save=os.path.join(args.outdir, 'ecal_barrel_clusters.png'))
general_clusters_figure(rdf, collection='HcalBarrelClusters', save=os.path.join(args.outdir, 'hcal_barrel_clusters.png'))
benchmarks/clustering/scripts/gen_particles.py 0 → 100644
+ 104
− 0
View file @ 60911ba0
import os
from pyHepMC3 import HepMC3 as hm
import numpy as np
import argparse
PARTICLES = {
"pion0": (111, 0.1349766), # pi0
"pion+": (211, 0.13957018), # pi+
"pion-": (-211, 0.13957018), # pi-
"kaon0": (311, 0.497648), # K0
"kaon+": (321, 0.493677), # K+
"kaon-": (-321, 0.493677), # K-
"proton": (2212, 0.938272), # proton
"neutron": (2112, 0.939565), # neutron
"electron": (11, 0.51099895e-3), # electron
"positron": (-11, 0.51099895e-3),# positron
"photon": (22, 0), # photon
}
def gen_event(p, theta, phi, pid, mass):
evt = hm.GenEvent(momentum_unit=hm.Units.MomentumUnit.GEV, length_unit=hm.Units.LengthUnit.MM)
# final state
state = 1
e0 = np.sqrt(p*p + mass*mass)
px = np.cos(phi)*np.sin(theta)
py = np.sin(phi)*np.sin(theta)
pz = np.cos(theta)
# beam
pbeam = hm.GenParticle(hm.FourVector(0, 0, 0, 0.938272), 2212, 4)
ebeam = hm.GenParticle(hm.FourVector(0, 0, e0, np.sqrt(e0*e0 + 0.511e-3*0.511e-3)), 11, 4)
# out particle
hout = hm.GenParticle(hm.FourVector(px*p, py*p, pz*p, e0), pid, state)
# vertex
vert = hm.GenVertex()
vert.add_particle_in(ebeam)
vert.add_particle_in(pbeam)
vert.add_particle_out(hout)
evt.add_vertex(vert)
return evt
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('output', help='path to the output file')
parser.add_argument('-n', type=int, default=1000, dest='nev', help='number of events to generate')
parser.add_argument('-s', type=int, default=-1, dest='seed', help='seed for random generator')
parser.add_argument('--parray', type=str, default="", dest='parray',
help='an array of momenta in GeV, separated by \",\"')
parser.add_argument('--pmin', type=float, default=8.0, dest='pmin', help='minimum momentum in GeV')
parser.add_argument('--pmax', type=float, default=100.0, dest='pmax', help='maximum momentum in GeV')
parser.add_argument('--angmin', type=float, default=0.0, dest='angmin', help='minimum angle in degree')
parser.add_argument('--angmax', type=float, default=20.0, dest='angmax', help='maximum angle in degree')
parser.add_argument('--phmin', type=float, default=0.0, dest='phmin', help='minimum angle in degree')
parser.add_argument('--phmax', type=float, default=360.0, dest='phmax', help='maximum angle in degree')
parser.add_argument('--particles', type=str, default='electron', dest='particles',
help='particle names, support {}'.format(list(PARTICLES.keys())))
args = parser.parse_args()
# random seed (< 0 will get it from enviroment variable 'SEED', or a system random number)
if args.seed < 0:
args.seed = os.environ.get('SEED', int.from_bytes(os.urandom(4), byteorder='big', signed=False))
print("Random seed is {}".format(args.seed))
np.random.seed(args.seed)
output = hm.WriterAscii(args.output);
if output.failed():
print("Cannot open file \"{}\"".format(args.output))
sys.exit(2)
# build particle info
parts = []
for pid in args.particles.split(','):
pid = pid.strip()
if pid not in PARTICLES.keys():
print('pid {:d} not found in dictionary, ignored.'.format(pid))
continue
parts.append(PARTICLES[pid])
# p values
pvals = np.random.uniform(args.pmin, args.pmax, args.nev) if not args.parray else \
np.random.choice([float(p.strip()) for p in args.parray.split(',')], args.nev)
thvals = np.random.uniform(args.angmin, args.angmax, args.nev)/180.*np.pi
phivals = np.random.uniform(args.phmin, args.phmax, args.nev)/180.*np.pi
partvals = [parts[i] for i in np.random.choice(len(parts), args.nev)]
count = 0
for p, theta, phi, (pid, mass) in zip(pvals, thvals, phivals, partvals):
if (count % 1000 == 0):
print("Generated {} events".format(count), end='\r')
evt = gen_event(p, theta, phi, pid, mass)
output.write_event(evt)
evt.clear()
count += 1
print("Generated {} events".format(args.nev))
output.close()
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment