diff --git a/Generators/FaserParticleGun/python/FaserParticleGunConfig.py b/Generators/FaserParticleGun/python/FaserParticleGunConfig.py index 28bb6147a3276b510f3b4c3de83ad281a2997af4..086d3c6aa5e3b5827df893a727d7e6a59f858780 100644 --- a/Generators/FaserParticleGun/python/FaserParticleGunConfig.py +++ b/Generators/FaserParticleGun/python/FaserParticleGunConfig.py @@ -140,6 +140,52 @@ def FaserParticleGunDecayInFlightCfg(ConfigFlags, **kwargs) : return cfg + +def FaserParticleGunForeseeCfg(ConfigFlags, **kwargs) : + # Supported keyword arguments: + # model_path (detault: $PWD) + # model_name (default: DarkPhoton) + # mother_mass (default: 0.01 GeV) + # com_energy (default: 14 TeV) + # daughter1_pid (default: 11) + # daughter2_pid (default: -11) + # mother_pid (default: none) + # mother_pos (default: CylinderSampler([0, 100**2],[0, 2*pi],[-1500, 0],0)) + # + # Note that ALL of these can be samplers themselves - either the simple, "literal" variety or a sampler object configured by the caller + # + + cfg = FaserParticleGunCommonCfg(ConfigFlags, **kwargs) + + pg = cfg.getPrimary() + + from DIFGenerator.DIFSampler import CylinderSampler + + from DIFGenerator import DIFSampler + pg.sampler = DIFSampler( + daughter1_pid = kwargs.get("daughter1_pid", 11), + daughter2_pid = kwargs.get("daughter2_pid", -11), + mother_pid = kwargs.get("mother_pid", None), + mother_pos = kwargs.get("mother_pos", CylinderSampler([0, 100**2],[0, 2*pi],[-1500, 0],0)) + ) + + from ForeseeGenerator.ForeseeSampler import ForeseeNumpySampler + mother_mom = ForeseeNumpySampler( + model_path = kwargs.get("model_path", "."), + model_name = kwargs.get("model_name", "DarkPhoton"), + com_energy = kwargs.get("com_energy", "14"), + mother_mass = kwargs.get("mother_mass", 0.01), + daughter1_pid = kwargs.get("daughter1_pid", 11), + daughter2_pid = kwargs.get("daughter2_pid", -11), + randomSeed = kwargs.get("randomSeed", None) + ) + + pg.sampler.mother_sampler = PG.ParticleSampler(pid=kwargs.get("mother_pid", None),mom=mother_mom,n=1,pos=pg.sampler.mother_sampler.pos) + pg.sampler.mother_sampler.mass_override = False + + + return cfg + def FaserParticleGunCfg(ConfigFlags) : generator = ConfigFlags.Sim.Gun.setdefault("Generator", "SingleParticle") kwargs = ConfigFlags.Sim.Gun @@ -149,5 +195,7 @@ def FaserParticleGunCfg(ConfigFlags) : return FaserParticleGunCosmicsCfg(ConfigFlags, **kwargs) elif generator == "DecayInFlight" : return FaserParticleGunDecayInFlightCfg(ConfigFlags, **kwargs) + elif generator == "Foresee" : + return FaserParticleGunForeseeCfg(ConfigFlags, **kwargs) else : return FaserParticleGunSingleParticleCfg(ConfigFlags, **kwargs ) diff --git a/Generators/ForeseeGenerator/CMakeLists.txt b/Generators/ForeseeGenerator/CMakeLists.txt new file mode 100644 index 0000000000000000000000000000000000000000..e2e5137b8c50f144c2cef2528f0b6be5c5548270 --- /dev/null +++ b/Generators/ForeseeGenerator/CMakeLists.txt @@ -0,0 +1,12 @@ +################################################################################ +# Package: ForeseeGenerator +################################################################################ + +# Declare the package name: +atlas_subdir( ForeseeGenerator ) + +# Install files from the package: +atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} ) + +atlas_install_joboptions( share/*.py ) + diff --git a/Generators/ForeseeGenerator/data/events_14TeV_m0.1GeV_c1e-05to_11_-11.npy b/Generators/ForeseeGenerator/data/events_14TeV_m0.1GeV_c1e-05to_11_-11.npy new file mode 100644 index 0000000000000000000000000000000000000000..36320c3f1c9cced47c9307b48529b444a2384be0 Binary files /dev/null and b/Generators/ForeseeGenerator/data/events_14TeV_m0.1GeV_c1e-05to_11_-11.npy differ diff --git a/Generators/ForeseeGenerator/python/ForeseeSampler.py b/Generators/ForeseeGenerator/python/ForeseeSampler.py new file mode 100644 index 0000000000000000000000000000000000000000..d80be0aa3aa080d93eb9ec7510b61aebf6089464 --- /dev/null +++ b/Generators/ForeseeGenerator/python/ForeseeSampler.py @@ -0,0 +1,441 @@ +import os, sys +import random +import numpy as np + +import ParticleGun as PG + +class ForeseeNumpySampler(PG.MomSampler): + """ Sample from the output of Foresee generation in numpy format with columns E, theta, weight""" + def __init__(self, model_path = ".", model_name = "DarkPhoton", com_energy = "14", mother_mass = 0.01, + coupling = None, daughter1_pid = 11, daughter2_pid = -11, randomSeed = None): + + self.path = os.path.expanduser(os.path.expandvars(model_path)) + self.modelname = model_name + self.energy = com_energy + self._mass = mother_mass + self.coupling = coupling + self.daughter1_pid = daughter1_pid + self.daughter2_pid = daughter2_pid + + if randomSeed is not None: + print(f"Setting seed to {randomSeed}") + self.rng = np.random.default_rng(randomSeed) + else: + self.rng = np.random.default_rng() + + self.xs = 0 + + self.read() + + def read(self): + "Read data from numpy file in format E, theta, weight" + + if self.path.endswith(".npy"): + filename = self.path + elif self.coupling is None: + filename = f"{self.path}/files/models/{self.modelname}/events/events_{self.energy}TeV_m{self._mass}GeV_to_{self.daughter1_pid}_{self.daughter2_pid}.npy" + else: + filename = f"{self.path}/files/models/{self.modelname}/events/events_{self.energy}TeV_m{self._mass}GeV_c{self.coupling}to_{self.daughter1_pid}_{self.daughter2_pid}.npy" + + print(f"Reading data from file: {filename}") + self.data = np.load(filename) + + # Create probablity for each mode as weight / sum(weights) + self.prob = self.data[2]/np.sum(self.data[2]) + + return + + def mass(self): + "Mass converted from GeV to MeV" + return self._mass * 1000 + + def shoot(self): + "Choose a random item from the data, based on the probability" + + # TODO: what about reuse of events? + energy, theta, weight = self.rng.choice(self.data, axis = 1, p = self.prob) + + self.xs += weight + + # Create mother momentum, randomly sampling phi + mother_mom = PG.EThetaMPhiSampler(energy, theta, self.mass(), [0, 2*np.pi]) + + return mother_mom.shoot() + +class ForeseeSampler(PG.MomSampler): + """Create events from foresee directly on the fly + + Requires: + * foresee to be downloaded and in python path + + cd <PATH> + git clone https://github.com/KlingFelix/FORESEE.git + export PYTHONPATH=$PYTHONPATH:<PATH>/FORESEE/src/ + + * scikit-hep installed + + pip install scikit-hep --user + + * forsee files dir symlinked to the run dir + + ln -s <PATH>/FORESEE/files . + + """ + + + def __init__(self, modelname, energy, mass, couplings, daughter1_pid, daughter2_pid, mother_pid = None): + self.modelname = modelname + self.model = Model(self.modelname) + self.energy = energy + self._mass = mass + self.couplings = [couplings] if isinstance(couplings, str) else couplings + self.mother_pid = mother_pid + self.daughter1_pid = daughter1_pid + self.daughter2_pid = daughter2_pid + + self.rng = np.random.default_rng() + self.xs = 0 + + if not os.path.exists("files"): + os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files") + + self.pid_map = { + (11, 11) : "e_e", + (13, 13) : "mu_mu", + (22, 22) : "gamma_gamma", + } + + self.mode = self.pid_map.get((self.daughter1_pid, self.daughter2_pid), None) + if self.mode is None: + sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}") + + from foresee import Foresee, Model, Utility + self.foresee = Foresee() + self.foresee.set_detector(selection="np.sqrt(x.x**2 + x.y**2)< 0.1", channels=[self.mode], distance=480, length=1.5 , luminosity=150) + + + if self.modelname == "DarkPhoton": + self.data = self.darkphoton() + elif self.modelname == "ALP-W": + self.data = self.alps() + else: + sys.exit(f"Unknown model {self.modelname}") + + return + + + def mass(self): + return self._mass * 1000 + + def darkphoton(self): + + # Production modes + self.model.add_production_2bodydecay( + pid0 = "111", + pid1 = "22", + br = "2.*0.99 * coupling**2 * pow(1.-pow(mass/self.masses('111'),2),3)", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10) + + self.model.add_production_2bodydecay( + pid0 = "221", + pid1 = "22", + br = "2.*0.39 * coupling**2 * pow(1.-pow(mass/self.masses('221'),2),3)", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10) + + # Handwavey + self.model.add_production_mixing( + pid = "113", + mixing = "coupling * 0.3/5. * 0.77545**2/abs(mass**2-0.77545**2+0.77545*0.147*1j)", + generator = "EPOSLHC", + energy = self.energy, + ) + + # Question on validity as FASER gets larger + self.model.add_production_direct( + label = "Brem", + energy = self.energy, + condition = "p.pt<1", + coupling_ref=1, + ) + + self.model.add_production_direct( + label = "DY", + energy = self.energy, + coupling_ref=1, + massrange=[1.5, 10.] + ) + + return self.decays() + + + def alps(self): + + self.model.add_production_2bodydecay( + pid0 = "5", + pid1 = "321", + br = "2.2e4 * coupling**2 * np.sqrt((1-(mass+0.495)**2/5.279**2)*(1-(mass-0.495)**2/5.279**2))", + generator = "Pythia8", + energy = self.energy, + nsample = 20, # Vary over out of phi and theta + ) + + + self.model.add_production_2bodydecay( + pid0 = "-5", + pid1 = "321", + br = "2.2e4 * coupling**2 * np.sqrt((1-(mass+0.495)**2/5.279**2)*(1-(mass-0.495)**2/5.279**2))", + generator = "Pythia8", + energy = self.energy, + nsample = 20, + ) + + self.model.add_production_2bodydecay( + pid0 = "130", + pid1 = "111", + br = "4.5 * coupling**2 * np.sqrt((1-(mass+0.135)**2/0.495**2)*(1-(mass-0.135)**2/0.495**2))", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10, + ) + + self.model.add_production_2bodydecay( + pid0 = "321", + pid1 = "211", + br = "10.5 * coupling**2 * np.sqrt((1-(mass+0.135)**2/0.495**2)*(1-(mass-0.135)**2/0.495**2))", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10, + ) + + return self.decays() + + + def decays(self): + # Decays + self.model.set_ctau_1d( + filename=f"files/models/{self.modelname}/ctau.txt", + coupling_ref=1 + ) + + # TODO take into account BR + self.model.set_br_1d( + modes = [self.mode], + filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] + ) + + # LLP spectrum + self.foresee.set_model(model=self.model) + + plt = self.foresee.get_llp_spectrum(self._mass, coupling=1, do_plot=True) # This is just a reference coupling + plt.savefig(f"{self.modelname}.png") + + def flatten(l): + return [i for sublist in l for i in sublist] + + coups, ctaus, nsigs, energies, weights, thetas = self.foresee.get_events(mass=self._mass, energy=self.energy, couplings=self.couplings) + + return [flatten(thetas), flatten(energies), flatten(weights)] + + def shoot(self): + # Create probablity for each mode as weight / sum(weights) + prob = self.data[2]/np.sum(self.data[2]) + + # Choose a random item from the data, base on the probability + # TODO: what about reuse of events? + theta_mother, e_mother, w = self.rng.choice(self.data, axis = 1, p = prob) + + self.xs += w + + # Create other momentum + mother_mom = PG.EThetaMPhiSampler(e_mother*1000, theta_mother, self.mass(), [0,2*np.pi]) + + return mother_mom.shoot() + +if __name__ == "__main__": + + # Testing ... + + from math import sqrt, log10 + import matplotlib.pyplot as plt + import matplotlib + from DIFGenerator import DIFSampler + + + path = os.path.expandvars("$Calypso_DIR/../calypso/Generators/ForeseeGenerator/data/events_14TeV_m0.1GeV_c1e-05to_11_-11.npy") + + + modelname = "DarkPhoton" + mass = 0.1 + + theta = [] + mom = [] + + d0theta = [] + d0mom = [] + + d1theta = [] + d1mom = [] + + # Accounting for rounding + epsilon = 5 + + # Mother position sampler over Z + mother_pos = PG.PosSampler(x=0,y=0,z=[-1500,0],t=0) + + # Create mother momentum sampler reading data from foresee + mother_mom_sampler = ForeseeNumpySampler(model_path = path, model_name = modelname, com_energy = "14", mother_mass = 0.1, coupling = 1e-5, daughter1_pid = 11, daughter2_pid = -11) + + # Create decay-in-flight + d = DIFSampler(11, -11, None) + d.mother_sampler = PG.ParticleSampler(pid=None,mom=mother_mom_sampler,n=1,pos=mother_pos) + d.mother_sampler.mass_override = False + + # Loop over a range of events + for i in range(1000): + + # Shoot the decay in flight + daughters = d.shoot() + + # Get mother and sum of daugthers and check these make sense. + mother_mom = d.mother_mom + s = daughters[0].mom+daughters[1].mom + + try: + assert mother_mom.E() - epsilon <= s.E() <= mother_mom.E() + epsilon + assert mother_mom.P() - epsilon <= s.P()<= mother_mom.P() + epsilon + assert mother_mom.Px() - epsilon <= s.Px() <= mother_mom.Px() + epsilon + assert mother_mom.Py() - epsilon <= s.Py() <= mother_mom.Py() + epsilon + assert mother_mom.Pz() - epsilon <= s.Pz() <= mother_mom.Pz() + epsilon + assert daughters[0].pos.X() == daughters[1].pos.X() == d.mother_pos.X() + assert daughters[0].pos.Y() == daughters[1].pos.Y() == d.mother_pos.Y() + assert daughters[0].pos.Z() == daughters[1].pos.Z() == d.mother_pos.Z() + except AssertionError: + print("Error on run " + str(i)) + + print("mother particle:") + print(" E = " + str(mother_mom.E())) + print(" M = " + str(mother_mom.M())) + print(" P = " + str(mother_mom.P())) + print(" Px = " + str(mother_mom.Px())) + print(" Py = " + str(mother_mom.Py())) + print(" Pz = " + str(mother_mom.Pz())) + print(" theta = " + str(mother_mom.Theta())) + print(" phi = " + str(mother_mom.Phi())) + print(" x = " + str(d.mother_pos.X())) + print(" y = " + str(d.mother_pos.Y())) + print(" z = " + str(d.mother_pos.Z())) + + print("daughter 0 particle:") + print(" E = " + str(daughters[0].mom.E())) + print(" M = " + str(daughters[0].mom.M())) + print(" P = " + str(daughters[0].mom.P())) + print(" Px = " + str(daughters[0].mom.Px())) + print(" Py = " + str(daughters[0].mom.Py())) + print(" Pz = " + str(daughters[0].mom.Pz())) + print(" theta = " + str(daughters[0].mom.Theta())) + print(" phi = " + str(daughters[0].mom.Phi())) + print(" x = " + str(daughters[0].pos.X())) + print(" y = " + str(daughters[0].pos.Y())) + print(" z = " + str(daughters[0].pos.Z())) + + print("daughter 1 particle:") + print(" E = " + str(daughters[1].mom.E())) + print(" M = " + str(daughters[1].mom.M())) + print(" P = " + str(daughters[1].mom.P())) + print(" Px = " + str(daughters[1].mom.Px())) + print(" Py = " + str(daughters[1].mom.Py())) + print(" Pz = " + str(daughters[1].mom.Pz())) + print(" theta = " + str(daughters[1].mom.Theta())) + print(" phi = " + str(daughters[1].mom.Phi())) + print(" x = " + str(daughters[1].pos.X())) + print(" y = " + str(daughters[1].pos.Y())) + print(" z = " + str(daughters[1].pos.Z())) + + raise + + # Store mother info to plot + theta.append(log10(mother_mom.Theta())) + mom.append(log10(mother_mom.P()/1000.)) + + # Store mother info to plot + d0theta.append(log10(daughters[0].mom.Theta())) + d0mom.append(log10(daughters[0].mom.P()/1000.)) + d1theta.append(log10(daughters[1].mom.Theta())) + d1mom.append(log10(daughters[1].mom.P()/1000.)) + + + + # Plot mother from sampling events + prange=[[-6, 0, 120],[ 0, 5, 50]] + tmin, tmax, tnum = prange[0] + pmin, pmax, pnum = prange[1] + t_edges = np.logspace(tmin, tmax, num=tnum+1) + p_edges = np.logspace(pmin, pmax, num=pnum+1) + + ticks = np.array([[np.linspace(10**(j),10**(j+1),9)] for j in range(-7,6)]).flatten() + ticks = [np.log10(x) for x in ticks] + ticklabels = np.array([[r"$10^{"+str(j)+"}$","","","","","","","",""] for j in range(-7,6)]).flatten() + matplotlib.rcParams.update({'font.size': 15}) + + + fig = plt.figure(figsize=(8,5.5)) + ax = plt.subplot(1,1,1) + h=ax.hist2d(x=theta,y=mom, + bins=[tnum,pnum],range=[[tmin,tmax],[pmin,pmax]], + norm=matplotlib.colors.LogNorm(), cmap="hsv", + ) + fig.colorbar(h[3], ax=ax) + ax.set_xlabel(r"angle wrt. beam axis $\theta$ [rad]") + ax.set_ylabel(r"momentum $p$ [GeV]") + ax.set_xticks(ticks) + ax.set_xticklabels(ticklabels) + ax.set_yticks(ticks) + ax.set_yticklabels(ticklabels) + ax.set_xlim(tmin, tmax) + ax.set_ylim(pmin, pmax) + plt.savefig(f"{modelname}_PG_m{mass}.png") + + fig = plt.figure(figsize=(8,5.5)) + ax = plt.subplot(1,1,1) + h=ax.hist2d(x=d0theta,y=d0mom, + bins=[tnum,pnum],range=[[tmin,tmax],[pmin,pmax]], + norm=matplotlib.colors.LogNorm(), cmap="hsv", + ) + fig.colorbar(h[3], ax=ax) + ax.set_xlabel(r"angle wrt. beam axis $\theta$ [rad]") + ax.set_ylabel(r"momentum $p$ [GeV]") + ax.set_xticks(ticks) + ax.set_xticklabels(ticklabels) + ax.set_yticks(ticks) + ax.set_yticklabels(ticklabels) + ax.set_xlim(tmin, tmax) + ax.set_ylim(pmin, pmax) + plt.savefig(f"{modelname}_PG_d0_m{mass}.png") + + fig = plt.figure(figsize=(8,5.5)) + ax = plt.subplot(1,1,1) + h=ax.hist2d(x=d1theta,y=d1mom, + bins=[tnum,pnum],range=[[tmin,tmax],[pmin,pmax]], + norm=matplotlib.colors.LogNorm(), cmap="hsv", + ) + fig.colorbar(h[3], ax=ax) + ax.set_xlabel(r"angle wrt. beam axis $\theta$ [rad]") + ax.set_ylabel(r"momentum $p$ [GeV]") + ax.set_xticks(ticks) + ax.set_xticklabels(ticklabels) + ax.set_yticks(ticks) + ax.set_yticklabels(ticklabels) + ax.set_xlim(tmin, tmax) + ax.set_ylim(pmin, pmax) + plt.savefig(f"{modelname}_PG_d1_m{mass}.png") + + + print (f"x-sect = {mother_mom_sampler.xs} pb") + + + + diff --git a/Generators/ForeseeGenerator/python/Validate.py b/Generators/ForeseeGenerator/python/Validate.py new file mode 100644 index 0000000000000000000000000000000000000000..8b53c84caad9190e4a9bda75f05970ec5f524da8 --- /dev/null +++ b/Generators/ForeseeGenerator/python/Validate.py @@ -0,0 +1,198 @@ +from AthenaPython.PyAthena import StatusCode, McEventCollection, HepMC, CLHEP +from GeneratorModules.EvgenAnalysisAlg import EvgenAnalysisAlg + +import ROOT as R +import numpy as np +import os + +def fix(): + "Python Fixes for HepMC" + def add(self, other): + self.set(self.x() + other.x(), self.y() + other.y(), + self.z() + other.z(), self.t() + other.t()) + return self + + HepMC.FourVector.__iadd__ = add + del add + + return + +class HistSvc(object): + "Class to deal with histograms" + + def __init__(self): + self.hists = {} + + def add(self, name, nbinsX = None, loX = None, hiX = None, nbinsY = None, loY = None, hiY = None, title = None, arrayX = None, arrayY = None): + hname = os.path.basename(name) + + if title is None: title = hname + + if nbinsY is not None: + self.hists[name] = R.TH2F(hname, title, nbinsX, loX, hiX, nbinsY, loY, hiY) + elif arrayX is not None and arrayY is not None: + self.hists[name] = R.TH2F(hname, title, len(arrayX) - 1, arrayX, len(arrayY) - 1, arrayY) + elif arrayX is not None and arrayY is None and nbinsY is not None: + self.hists[name] = R.TH2F(hname, title, len(arrayX) - 1, arrayX, nbinsY, loY, hiY) + elif arrayX is None and arrayY is not None: + self.hists[name] = R.TH2F(hname, title, nbinsX, loX, hiX, len(arrayY) - 1, arrayY) + elif arrayX is not None: + self.hists[name] = R.TH1F(hname, title, len(arrayX) - 1, arrayX) + else: + self.hists[name] = R.TH1F(hname, title, nbinsX, loX, hiX) + + def __getitem__(self, name): + return self.hists[name] + + def write(self, name): + + f = R.TFile.Open(name, "RECREATE") + + for n, h in self.hists.items(): + path = os.path.dirname(n) + if path and not f.GetDirectory(path): + f.mkdir(path) + + f.cd(path) + h.Write() + + f.Close() + + return + +class EvgenValidation(EvgenAnalysisAlg): + "Gen-level validation" + + def __init__(self, name = "EvgenValidation"): + super(EvgenValidation, self).__init__(name=name) + self.hists = HistSvc() + self.ndaughter = 2 + self.outname = "validation.root" + + def binning(self): + "binning for theta vs phi plot" + tmin, tmax, tnum = [-6, 0, 24] + pmin, pmax, pnum = [ 0, 5, 10] + t_edges = np.logspace(tmin, tmax, num=tnum+1) + p_edges = np.logspace(pmin, pmax, num=pnum+1) + return t_edges, p_edges + + def initialize(self): + + # All daughters + self.hists.add("PIDs", 60, -30, 30) + + # Daughter i + tbins, pbins = self.binning() + for i in range(self.ndaughter): + self.hists.add(f"P_d{i}", 100, 0, 10000) + self.hists.add(f"Pt_d{i}", 100, 0, 1) + self.hists.add(f"Theta_d{i}", 20, 0, 0.001) + self.hists.add(f"Phi_d{i}", 16, -3.2, 3.2) + self.hists.add(f"ThetaVsP_d{i}", arrayX = tbins, arrayY = pbins) + self.hists.add(f"Mass_d{i}", 200, 0, 0.01) + + # Mother + self.hists.add("P_M", 100, 0, 10000) + self.hists.add("Pt_M", 100, 0, 1) + self.hists.add("Theta_M", 20, 0, 0.001) + self.hists.add("Phi_M", 16, -3.2, 3.2) + self.hists.add("Mass_M", 200, 0, 1) + self.hists.add("ThetaVsP_M", arrayX = tbins, arrayY = pbins) + + return StatusCode.Success + + + def fillKin(self, label, p, mass = True, twoD = True): + + self.hists[f"P_{label}"].Fill(p.rho()/1000, 1) + self.hists[f"Pt_{label}"].Fill(p.perp()/1000, 1) + self.hists[f"Theta_{label}"].Fill(p.theta(), 1) + self.hists[f"Phi_{label}"].Fill(p.phi(), 1) + + if mass: + self.hists[f"Mass_{label}"].Fill(p.m()/1000, 1) + + if twoD: + self.hists[f"ThetaVsP_{label}"].Fill(p.theta(), p.rho()/1000, 1) + + return + + def fillDaughter(self, p): + self.hists["PIDs"].Fill(p.pdg_id()) + return + + def execute(self): + evt = self.events()[0] + + # Loop over all particles in events (assuming mother not stored) + momenta = [] + mother = HepMC.FourVector(0,0,0,0) + for i, p in enumerate(evt.particles): + #p.print() + self.fillDaughter(p) + momenta.append(p.momentum()) + mother += p.momentum() + + # Fill daughter plots + for i in range(self.ndaughter): + self.fillKin(f"d{i}", momenta[i]) + + # Fill mother plots + self.fillKin("M", mother, mass = True) + + return StatusCode.Success + + def finalize(self): + self.hists.write(self.outname) + return StatusCode.Success + + +if __name__ == "__main__": + + import argparse, sys + parser = argparse.ArgumentParser(description="Run gen-level validation") + parser.add_argument("file", nargs="+", help = "full path to imput file") + parser.add_argument("--ndaugthers", "-d", default = 2, type = float, help = "Number of daugthers to plot") + parser.add_argument("--output", "-o", default = "validation.root", help = "Name of output file") + parser.add_argument("--mcEventKey", "-k", default = "BeamTruthEvent", help = "Name of MC collection") + parser.add_argument("--nevents", "-n", default = -1, type = int, help = "Number of events to process") + args = parser.parse_args() + + from AthenaCommon.Logging import log + from AthenaCommon.Constants import DEBUG + log.setLevel(DEBUG) + + from AthenaCommon.Configurable import Configurable + Configurable.configurableRun3Behavior = 1 + + from CalypsoConfiguration.AllConfigFlags import ConfigFlags + ConfigFlags.Input.isMC = True + ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-01" # Always needed; must match FaserVersion + ConfigFlags.GeoModel.FaserVersion = "FASER-01" # Default FASER geometry + ConfigFlags.Detector.EnableFaserSCT = True + ConfigFlags.Input.Files = args.file + ConfigFlags.lock() + + from CalypsoConfiguration.MainServicesConfig import MainServicesCfg + cfg = MainServicesCfg(ConfigFlags) + + from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg + cfg.merge(PoolReadCfg(ConfigFlags)) + + from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator + from AthenaConfiguration.ComponentFactory import CompFactory + + import McParticleEvent.Pythonizations + fix() + + acc = ComponentAccumulator() + valid = EvgenValidation() + valid.McEventKey = args.mcEventKey + valid.ndaughters = args.ndaugthers + valid.outname = args.output + acc.addEventAlgo(valid) + cfg.merge(acc) + + sc = cfg.run(maxEvents = args.nevents) + sys.exit(not sc.isSuccess()) diff --git a/Generators/ForeseeGenerator/share/generate_forsee_events.py b/Generators/ForeseeGenerator/share/generate_forsee_events.py new file mode 100644 index 0000000000000000000000000000000000000000..2a9e7a7fefd775c9d012fdfefe38c4ca8227b6d7 --- /dev/null +++ b/Generators/ForeseeGenerator/share/generate_forsee_events.py @@ -0,0 +1,304 @@ +import os + +import numpy as np +import matplotlib.pyplot as plt +import matplotlib + +from foresee import Foresee, Model, Utility + +class ForeseeGenerator(object): + """ + Generate LLP particles within FASER acceptance from FORESEE + """ + + def __init__(self, modelname, energy, mass, couplings, daughter1_pid, daughter2_pid, mother_pid = None, outdir = None): + + self.modelname = modelname + self.energy = energy + self.mass = mass + self.couplings = [couplings] if isinstance(couplings, (str, int, float)) else couplings + self.mother_pid = mother_pid + self.daughter1_pid = daughter1_pid + self.daughter2_pid = daughter2_pid + self.outdir = outdir + + #if not os.path.exists("files"): + # os.symlink(os.path.expandvars("$Calypso_DIR/../calypso/Generators/foresee/files"), "files") + + # Set decay mode ... + + self.pid_map = { + (11, 11) : "e_e", + (13, 13) : "mu_mu", + (22, 22) : "gamma_gamma", + } + + self.mode = self.pid_map.get((abs(self.daughter1_pid), abs(self.daughter2_pid)), None) + if self.mode is None: + sys.exit(f"Undefined decay to {self.daughter1_pid} + {self.daughter2_pid} for {self.modelname}") + + # Set detector ... + self.foresee = Foresee() + self.foresee.set_detector(selection="np.sqrt(x.x**2 + x.y**2)< 0.1", channels=[self.mode], distance=480, length=1.5 , luminosity=150) + + # Set model ... + self.model = Model(self.modelname) + + if self.modelname == "DarkPhoton": + self.data = self.darkphoton() + elif self.modelname == "ALP-W": + self.data = self.alps() + else: + sys.exit(f"Unknown model {self.modelname}") + + return + + def darkphoton(self): + + # Production modes + self.model.add_production_2bodydecay( + pid0 = "111", + pid1 = "22", + br = "2.*0.99 * coupling**2 * pow(1.-pow(mass/self.masses('111'),2),3)", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10) + + self.model.add_production_2bodydecay( + pid0 = "221", + pid1 = "22", + br = "2.*0.39 * coupling**2 * pow(1.-pow(mass/self.masses('221'),2),3)", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10) + + # This is a bit handwavey according to Felix + self.model.add_production_mixing( + pid = "113", + mixing = "coupling * 0.3/5. * 0.77545**2/abs(mass**2-0.77545**2+0.77545*0.147*1j)", + generator = "EPOSLHC", + energy = self.energy, + ) + + # Question on validity as FASER gets larger but OK for currrent size according to Felix + self.model.add_production_direct( + label = "Brem", + energy = self.energy, + condition = "p.pt<1", + coupling_ref=1, + ) + + self.model.add_production_direct( + label = "DY", + energy = self.energy, + coupling_ref=1, + massrange=[1.5, 10.] + ) + + return self.decays() + + + def alps(self): + + self.model.add_production_2bodydecay( + pid0 = "5", + pid1 = "321", + br = "2.2e4 * coupling**2 * np.sqrt((1-(mass+0.495)**2/5.279**2)*(1-(mass-0.495)**2/5.279**2))", + generator = "Pythia8", + energy = self.energy, + nsample = 20, # Vary over phi and theta + ) + + self.model.add_production_2bodydecay( + pid0 = "-5", + pid1 = "321", + br = "2.2e4 * coupling**2 * np.sqrt((1-(mass+0.495)**2/5.279**2)*(1-(mass-0.495)**2/5.279**2))", + generator = "Pythia8", + energy = self.energy, + nsample = 20, + ) + + self.model.add_production_2bodydecay( + pid0 = "130", + pid1 = "111", + br = "4.5 * coupling**2 * np.sqrt((1-(mass+0.135)**2/0.495**2)*(1-(mass-0.135)**2/0.495**2))", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10, + ) + + self.model.add_production_2bodydecay( + pid0 = "321", + pid1 = "211", + br = "10.5 * coupling**2 * np.sqrt((1-(mass+0.135)**2/0.495**2)*(1-(mass-0.135)**2/0.495**2))", + generator = "EPOSLHC", + energy = self.energy, + nsample = 10, + ) + + return self.decays() + + + def decays(self): + # Set up liftime and BRs + + self.model.set_ctau_1d( + filename=f"files/models/{self.modelname}/ctau.txt", + coupling_ref=1 + ) + + self.model.set_br_1d( + modes = [self.mode], + filenames=[f"files/models/{self.modelname}/br/{self.mode}.txt"] + ) + + # Get LLP spectrum + self.foresee.set_model(model=self.model) + plt = self.foresee.get_llp_spectrum(self.mass, coupling=1, do_plot=True) # This is just a reference coupling + plt.savefig(f"{self.modelname}_m{self.mass}.png") + + def flatten(l): + return [i for sublist in l for i in sublist] + + # Get list of events within detector + coups, ctaus, nsigs, energies, weights, thetas = self.foresee.get_events(mass=self.mass, energy=self.energy, couplings=self.couplings) + + self.plot(flatten(thetas), flatten(energies), flatten(weights)) + + # Return energy (converting to MeV), theta and weights + return [[e*1000 for e in flatten(energies)], flatten(thetas), flatten(weights)] + + def plot(self, thetas, energies, weights): + # Plot the results in Forsee format + + t = np.array(thetas) + p = np.sqrt(np.array(energies)**2 - self.mass**2) + + prange=[[-6, 0, 120],[ 0, 5, 50]] + tmin, tmax, tnum = prange[0] + pmin, pmax, pnum = prange[1] + t_edges = np.logspace(tmin, tmax, num=tnum+1) + p_edges = np.logspace(pmin, pmax, num=pnum+1) + + ticks = np.array([[np.linspace(10**(j),10**(j+1),9)] for j in range(-7,6)]).flatten() + ticks = [np.log10(x) for x in ticks] + ticklabels = np.array([[r"$10^{"+str(j)+"}$","","","","","","","",""] for j in range(-7,6)]).flatten() + matplotlib.rcParams.update({'font.size': 15}) + + fig = plt.figure(figsize=(8,5.5)) + ax = plt.subplot(1,1,1) + h=ax.hist2d(x=np.log10(t),y=np.log10(p),weights=weights, + bins=[tnum,pnum],range=[[tmin,tmax],[pmin,pmax]], + norm=matplotlib.colors.LogNorm(), cmap="hsv", + ) + + fig.colorbar(h[3], ax=ax) + ax.set_xlabel(r"angle wrt. beam axis $\theta$ [rad]") + ax.set_ylabel(r"momentum $p$ [GeV]") + ax.set_xticks(ticks) + ax.set_xticklabels(ticklabels) + ax.set_yticks(ticks) + ax.set_yticklabels(ticklabels) + ax.set_xlim(tmin, tmax) + ax.set_ylim(pmin, pmax) + plt.savefig(f"{self.modelname}_m{self.mass}_acc.png") + + def write(self): + # Write LLP results to a file + + energies, thetas, weights = self.data + + if self.outdir is None: + self.outdir = f"files/models/{self.modelname}/events" + if not os.path.exists(self.outdir): + os.mkdir(self.outdir) + + if len(self.couplings) == 1: + filename = f"{self.outdir}/events_{self.energy}TeV_m{self.mass}GeV_c{self.couplings[0]}to_{self.daughter1_pid}_{self.daughter2_pid}.npy" + else: + filename = f"{self.outdir}/events_{self.energy}TeV_m{self.mass}GeV_to_{self.daughter1_pid}_{self.daughter2_pid}.npy" + + print(f"Generated {len(thetas)} events") + print(f"save data to file: {filename}") + np.save(filename,[energies,thetas, weights]) + + cfgname = filename.replace(".npy", ".cfg") + print(f"save config to file: {cfgname}") + with open(cfgname, "w") as f: + f.write(" ".join(sys.argv)) + + return + +def setup_foresee(path): + + if path is None: + return + + # Add foresee to python path + path = os.path.expandvars(os.path.expanduser(path)) + os.sys.path.append(f"{path}/FORESEE/src") + + # Symlink foresee files dir to current dir + if not os.path.exists("files"): + os.symlink(os.path.expandvars(f"{path}/FORESEE/files"), "files") + + # Install scikit-hep if needed. + + try: + from skhep.math.vectors import LorentzVector, Vector3D + except ModuleNotFoundError: + os.system("pip install scikit-hep --user") + try: + from skhep.math.vectors import LorentzVector, Vector3D + except ModuleNotFoundError: + raise ModuleNotFoundError("Unable to find skhep. Please install the scikit-hep package") + + return + + +def parse_couplings(data): + + try: + couplings = [float(d) for d in data] + except ValueError: + try: + couplings = np.logspace(*eval(data[0])) + except: + sys.exit("Unable to parse couplings") + + return couplings + +if __name__ == "__main__": + + import argparse, sys + + parser = argparse.ArgumentParser(description="Run FORSEE generation") + parser.add_argument("model", help = "Name of foresee model") + parser.add_argument("--mass", "-m", required = True, type = float, help = "Mass of mother [GeV]") + parser.add_argument("--couplings", "-c", required = True, nargs = "+", help = "Couplings of mother (either single/mulitple values or tuple to pass to np.logspace)") + parser.add_argument("--pid1", required = True, type = int, help = "PID of daughter 1") + parser.add_argument("--pid2", default = None, type = int, help = "PID of daughter 2 (if not set then will be -PID1)") + parser.add_argument("--Ecom", default = "14", help = "Center of mass energy [TeV]") + parser.add_argument("--outdir", "-o", default = ".", help = "Output path") + parser.add_argument("--path", default = ".", help = "Path to foresee installation") + args = parser.parse_args() + + setup_foresee(args.path) + + + # Create PIDs + if args.pid2 is None: + args.pid2 = -args.pid1 + + couplings = parse_couplings(args.couplings) + + print(f"Generating {args.model} events at Ecom = {args.Ecom}") + print(f" mother mass = {args.mass} GeV") + print(f" decay = {args.pid1} {args.pid2}") + print(f" couplings = {couplings}") + + f = ForeseeGenerator(args.model, args.Ecom, args.mass, couplings, args.pid1, args.pid2) + f.write() + + + diff --git a/Generators/ForeseeGenerator/share/plot_validation.py b/Generators/ForeseeGenerator/share/plot_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..5b3d8b56f2ec4dc60f308be26374191ca47bfc2e --- /dev/null +++ b/Generators/ForeseeGenerator/share/plot_validation.py @@ -0,0 +1,112 @@ +import ROOT as R +from collections import namedtuple + +Hist = namedtuple("Hist", "name, xtitle, ytitle, xlo, xhi, ylo, yhi, r, d, logx, logy, ndiv", + defaults = [None, None, None, None, None, None, 1, "hist", False, False, None]) + +def plot(f, name, xtitle, ytitle, xlo = None, xhi = None, ylo = None, yhi = None, + r = 1, d = "hist", logx = False, logy = False, ndiv = None): + + h = f.Get(name) + + if xlo is not None and xhi is not None: + h.SetAxisRange(xlo, xhi) + + if ylo is not None and yhi is not None: + h.SetAxisRange(ylo, yhi, "Y") + + if r != 1: + h.Rebin(r) + + if xtitle is not None: + h.GetXaxis().SetTitle(xtitle) + + if ytitle is not None: + h.GetYaxis().SetTitle(ytitle) + + if logx: + R.gPad.SetLogx() + + if logy: + R.gPad.SetLogy() + + if ndiv is not None: + h.SetNdivisions(ndiv) + + h.SetLabelSize(0.05, "X") + h.SetTitleSize(0.05, "X") + h.SetLabelSize(0.05, "Y") + h.SetTitleSize(0.05, "Y") + + h.GetXaxis().SetTitleOffset(1.2) + + R.gPad.SetBottomMargin(0.15) + R.gPad.SetLeftMargin(0.12) + R.gPad.SetRightMargin(0.2) + + h.Draw(d) + return h + +def plotn(f, configs, x, y, outname = "valplot"): + + c = R.TCanvas() + c.Divide(x, y) + c._objs = [] + + if isinstance(configs, tuple): + configs = [configs] + + for i, cfg in enumerate(configs): + c.cd(i+1) + c._objs.append(plot(f, *cfg)) + + c.Print(f"{outname}.eps") + + return + +if __name__ == "__main__": + + R.gROOT.SetBatch(True) + R.gStyle.SetOptStat(0) + + fname = "validation.root" + f = R.TFile.Open(fname) + + config = [Hist("P_d0", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5), + Hist("Theta_d0", xtitle = "#theta [rad]", ndiv = -4), + Hist("Mass_d0", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.001, ndiv = 4), + Hist("Pt_d0", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5), + Hist("Phi_d0", xtitle = "#phi [rad]"), + Hist("ThetaVsP_d0", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz") + ] + + plotn(f, config, 3, 2, "daug0") + + config = [Hist("P_d1", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5), + Hist("Theta_d1", xtitle = "#theta [rad]", ndiv = -4), + Hist("Mass_d1", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.001, ndiv = 4), + Hist("Pt_d1", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5), + Hist("Phi_d1", xtitle = "#phi [rad]"), + Hist("ThetaVsP_d1", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz") + ] + + plotn(f, config, 3, 2, "daug1") + + config = [Hist("P_M", logy = True, xtitle = "p^{0} [GeV]", ndiv = 5, r = 5), + Hist("Theta_M", xtitle = "#theta [rad]", ndiv = -4), + Hist("Mass_M", xtitle = "m^{0} [GeV]", xlo = 0, xhi = 0.05), + Hist("Pt_M", logy = True, xtitle = "p_{T}^{0} [GeV]", ndiv = 10, r = 5), + Hist("Phi_M", xtitle = "#phi [rad]"), + Hist("ThetaVsP_M", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz") + ] + + plotn(f, config, 3, 2, "mother") + + plotn(f, Hist("PIDs", xtitle="PDG Id"), 1, 1, "pid") + + config = [Hist("ThetaVsP_M", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz"), + Hist("ThetaVsP_d0", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz"), + Hist("ThetaVsP_d1", xtitle = "p^{0} [GeV]", ytitle = "#theta [rad]", logx = True, logy = True, d = "colz") + ] + + plotn(f, config, 2, 2, "twod") diff --git a/Generators/HEPMCReader/CMakeLists.txt b/Generators/HEPMCReader/CMakeLists.txt new file mode 100644 index 0000000000000000000000000000000000000000..3076e7abc887d5cc1fabaedddc42d5db36286270 --- /dev/null +++ b/Generators/HEPMCReader/CMakeLists.txt @@ -0,0 +1,9 @@ +################################################################################ +# Package: HEPMCGenie +################################################################################ + +# Declare the package name: +atlas_subdir( HEPMCReader ) + +# Install files from the package: +atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} ) diff --git a/Generators/HEPMCReader/python/HepMCReaderConfig.py b/Generators/HEPMCReader/python/HepMCReaderConfig.py new file mode 100644 index 0000000000000000000000000000000000000000..72081a86f237a8a99154a61bf7d7713f12282d06 --- /dev/null +++ b/Generators/HEPMCReader/python/HepMCReaderConfig.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python + +# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration + +# import sys +from AthenaConfiguration.MainServicesConfig import AthSequencer +from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator +from AthenaConfiguration.ComponentFactory import CompFactory + +from TruthIO.TruthIOConf import HepMCReadFromFile + + +def HepMCReaderCfg(ConfigFlags, **kwargs) : + cfg = ComponentAccumulator(AthSequencer("AthBeginSeq", Sequential = True)) + + + from TruthIO.TruthIOConf import HepMCReadFromFile + hepmc = CompFactory.HepMCReadFromFile(name = kwargs.setdefault("name", "FASERHepMCReader")) + hepmc.InputFile = ConfigFlags.Input.Files[0] + hepmc.McEventKey = kwargs.setdefault("McEventKey", "BeamTruthEvent") + + cfg.addEventAlgo(hepmc, sequenceName = "AthBeginSeq", primary = True) # to run *before* G4 + + return cfg diff --git a/Simulation/G4Faser/G4FaserAlg/test/G4FaserAlgConfigNew_Test.py b/Simulation/G4Faser/G4FaserAlg/test/G4FaserAlgConfigNew_Test.py index 709797ae28bec863f565eba2aa662d8bcefc0e2e..d6e254c86c1c9f8207512c52414847663fd39a11 100644 --- a/Simulation/G4Faser/G4FaserAlg/test/G4FaserAlgConfigNew_Test.py +++ b/Simulation/G4Faser/G4FaserAlg/test/G4FaserAlgConfigNew_Test.py @@ -82,11 +82,23 @@ if __name__ == '__main__': # if ConfigFlags.Input.Files and ConfigFlags.Input.Files != ["_CALYPSO_GENERIC_INPUTFILE_NAME_"] : print("Input.Files = ",ConfigFlags.Input.Files) + +# +# If so, and only one file that ends in .events read as HepMC +# + if len(ConfigFlags.Input.Files) == 1 and ConfigFlags.Input.Files[0].endswith(".events"): + from HEPMCReader.HepMCReaderConfig import HepMCReaderCfg + cfg.merge(HepMCReaderCfg(ConfigFlags)) + + from McEventSelector.McEventSelectorConfig import McEventSelectorCfg + cfg.merge(McEventSelectorCfg(ConfigFlags)) + # -# If so, set up to read it +# Else, set up to read it as a pool.root file # - from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg - cfg.merge(PoolReadCfg(ConfigFlags)) + else: + from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg + cfg.merge(PoolReadCfg(ConfigFlags)) # # If not, configure the particle gun as requested, or using defaults #