Skip to content
Snippets Groups Projects
Forked from faser / calypso
773 commits behind the upstream repository.
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
FaserParticleGunConfig.py 8.27 KiB
#!/usr/bin/env python

# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration

# import sys
from AthenaConfiguration.MainServicesConfig import AthSequencer
import ParticleGun as PG

from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
# from AthenaCommon.Constants import VERBOSE, INFO
from AthenaCommon.SystemOfUnits import TeV
from AthenaCommon.PhysicalConstants import pi

### add radial pos sampler ### with gaussian beam implemented
from FaserParticleGun.RadialPosSampler import RadialPosSampler



def FaserParticleGunCommonCfg(ConfigFlags, **kwargs) :
    cfg = ComponentAccumulator(AthSequencer("AthBeginSeq", Sequential = True))


    # have to do this in a clunky way, since ParticleGun does not recognize all keywords we might get passed
    pg = PG.ParticleGun( name = kwargs.setdefault("name", "ParticleGun"),
                         randomSeed = kwargs.setdefault("randomSeed", 12345),
                         randomSvcName = kwargs.setdefault("randomSvcName", "AtRndmGenSvc"),
                         randomStream = kwargs.setdefault("randomStream", "ParticleGun") )

    pg.McEventKey = kwargs.setdefault("McEventKey", "BeamTruthEvent")
    cfg.addEventAlgo(pg, sequenceName = "AthBeginSeq", primary = True) # to run *before* G4

    return cfg


def FaserParticleGunSingleParticleCfg(ConfigFlags, **kwargs) :
    cfg = FaserParticleGunCommonCfg(ConfigFlags, **kwargs)

    pg = cfg.getPrimary()

    pg.sampler.n   = kwargs.setdefault("n", 1)
    pg.sampler.pid = kwargs.setdefault("pid", -13)
    pg.sampler.mom = PG.EThetaMPhiSampler(energy = kwargs.setdefault("energy", 1*TeV),
                                          theta = kwargs.setdefault("theta", [0, pi/20]), 
                                          phi = kwargs.setdefault("phi", [0, 2*pi]), 
                                          mass = kwargs.setdefault("mass", 0.0) )
    pg.sampler.pos = PG.PosSampler(x = kwargs.setdefault("x", [-5, 5]), 
                                   y = kwargs.setdefault("y", [-5, 5]), 
                                   z = kwargs.setdefault("z", -3750.0), 
                                   t = kwargs.setdefault("t", 0.0) )

    return cfg

def FaserParticleGunSingleEcalParticleCfg(ConfigFlags, **kwargs) :
    cfg = FaserParticleGunCommonCfg(ConfigFlags, **kwargs)

    pg = cfg.getPrimary()

    pg.sampler.pid = kwargs.setdefault("pid", 13)
    pg.sampler.mom = PG.EThetaMPhiSampler(energy = kwargs.setdefault("energy", 1*TeV),
                                          theta = kwargs.setdefault("theta", 0.0), 
                                          phi = kwargs.setdefault("phi", 0.0), 
                                          mass = kwargs.setdefault("mass", 105.7) )    

    if "radius" in kwargs:
        pg.sampler.pos =  RadialPosSampler(x = kwargs.setdefault("x", 0.0), 
                                   y = kwargs.setdefault("y", 0.0), 
                                   z = kwargs.setdefault("z", 0.0),
                                   r = kwargs.setdefault("radius", 1.0),
                                   t = kwargs.setdefault("t", 0.0) )
    else:
        pg.sampler.pos = PG.PosSampler(x = kwargs.setdefault("x", 0.0), 
                                   y = kwargs.setdefault("y", 0.0), 
                                   z = kwargs.setdefault("z", 0.0),
                                   t = kwargs.setdefault("t", 0.0) )

    return cfg

'''
def FaserParticleGunSingleEcalParticleCfg(ConfigFlags, **kwargs) :
    cfg = FaserParticleGunCommonCfg(ConfigFlags, **kwargs)

    pg = cfg.getPrimary()

    pg.sampler.pid = kwargs.setdefault("pid", 11)
    pg.sampler.mom = PG.EThetaMPhiSampler(energy = kwargs.setdefault("energy", 1*TeV),
                                          theta = kwargs.setdefault("theta", [0, pi/200]), 
                                          phi = kwargs.setdefault("phi", [0, 2*pi]), 
                                          mass = kwargs.setdefault("mass", 0.0) )
    pg.sampler.pos = PG.PosSampler(x = kwargs.setdefault("x", [-5, 5]), 
                                   y = kwargs.setdefault("y", [-5, 5]), 
                                   z = kwargs.setdefault("z", 2730.0), 
                                   t = kwargs.setdefault("t", 0.0) )

    return cfg
'''
def FaserParticleGunCosmicsCfg(ConfigFlags, **kwargs) :
    # Supported keyword arguments:
    #
    # chargeRatio       (default: 1.27)
    # maxMuonEnergyGeV  (default: 10000)
    # thetaMinDegree    (default: 0)
    # thetaMaxDegree    (default: 80)
    # targetDepthMeters (default: 85)
    # targetDxMm        (default: 600)  - width of target slab
    # targetDyMm        (default: 600)  - height of target slab
    # targetDzMm        (default: 7000) - length of target slab
    # x0Mm              (default: 0)    - X center of target slab in FASER coordinate system
    # y0Mm              (default: 0)    - Y center of target slab in FASER coordinate system
    # z0Mm              (default: 150)  - Z center of target slab in FASER coordinate system
    #
    
    cfg = FaserParticleGunCommonCfg(ConfigFlags, **kwargs)

    pg = cfg.getPrimary()

    from FaserCosmicGenerator import CosmicSampler
    pg.sampler = CosmicSampler(**kwargs)

    return cfg

def FaserParticleGunDecayInFlightCfg(ConfigFlags, **kwargs) :
    # Supported keyword arguments:
    #
    # daughter1_pid (default:  11)
    # daughter2_pid (default: -11)
    # mother_pid    (default: none)
    # mother_mom    (default: PG.EThetaMPhiSampler(sqrt((1*TeV)**2 + (10*MeV)**2),0,10*MeV,0))
    # 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 import DIFSampler
    pg.sampler = DIFSampler(**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
    if generator == "SingleEcalParticle" :
        return FaserParticleGunSingleEcalParticleCfg(ConfigFlags, **kwargs)
    elif generator == "Cosmics" :
        return FaserParticleGunCosmicsCfg(ConfigFlags, **kwargs)
    elif generator == "DecayInFlight" :
        return FaserParticleGunDecayInFlightCfg(ConfigFlags, **kwargs)
    elif generator == "Foresee" :
        return FaserParticleGunForeseeCfg(ConfigFlags, **kwargs)    
    else :
        return FaserParticleGunSingleParticleCfg(ConfigFlags, **kwargs )