Add scripts to generator events from Foresee and a sampler to read these in to...

Add scripts to generator events from Foresee and a sampler to read these in to the DIF + config for this.  Add ability to read HepMC input events.  Add Validation alg

Generators/FaserParticleGun/python/FaserParticleGunConfig.py

@@ -140,6 +140,51 @@ 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),
+        )
+    
+    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 +194,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 )

Generators/ForeseeGenerator/CMakeLists.txt

+################################################################################
+# 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 )
+

Generators/ForeseeGenerator/python/ForeseeSampler.py

+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):
+
+        self.path = 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
+
+        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}")
+
+        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 ...
+
+    import os
+    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")
+    
+
+
+        

Generators/ForeseeGenerator/python/Validate.py

+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, 12]
+        pmin, pmax, pnum = [ 0, 5, 5]
+        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("Masses", 100, 0, 0.01)
+        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}", 10, 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)
+
+        # Mother
+        self.hists.add("P_M", 100, 0, 10000) 
+        self.hists.add("Pt_M", 100, 0, 1)       
+        self.hists.add("Theta_M", 10, 0, 0.001)
+        self.hists.add("Phi_M", 16, -3.2, 3.2)
+        self.hists.add("Mass_M", 100, 0, 1)
+        self.hists.add("ThetaVsP_M", arrayX = tbins, arrayY = pbins)
+        
+        return StatusCode.Success
+
+
+    def fillKin(self, label, p, mass = False, 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["Masses"].Fill(p.momentum().m()/1000, 1)
+        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", 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())

Generators/ForeseeGenerator/share/generate_forsee_events.py

+import os
+
+import numpy as np
+
+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):
+
+        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
+
+        #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}.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)
+
+        # Return energy (converting to MeV), theta and weights
+        return [[e*1000 for e in flatten(energies)], flatten(thetas), flatten(weights)] 
+
+    def write(self):
+        # Write LLP results to a file
+        
+        energies, thetas, weights = self.data
+
+        dirname = f"files/models/{self.modelname}/events"
+        if not os.path.exists(dirname):
+            os.mkdir(dirname)
+
+        if len(self.couplings) == 1:
+            filename = f"{dirname}/events_{self.energy}TeV_m{self.mass}GeV_c{self.couplings[0]}to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
+        else:
+            filename = f"{dirname}/events_{self.energy}TeV_m{self.mass}GeV_to_{self.daughter1_pid}_{self.daughter2_pid}.npy"
+
+        print(f"save data to file: {filename}")
+        np.save(filename,[energies,thetas, weights])
+
+        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, 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("--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()
+    
+        
+

Generators/HEPMCReader/CMakeLists.txt

+################################################################################
+# 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} )

Generators/HEPMCReader/python/HepMCReaderConfig.py

+#!/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

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
 #