Python ShipGeoConfig.ConfigRegistry类代码示例

 def test_true(self):
     c = ConfigRegistry.loadpys(self.config, MU_SHIELD_ENABLED=True)
     self.assertTrue("muShield" in c)
     assert len(ConfigRegistry.keys()) == 1, ConfigRegistry.keys()
     assert self.key in ConfigRegistry.keys()
     assert ConfigRegistry[self.key].Bfield.max  == 1.5*u.kilogauss
     assert ConfigRegistry[self.key].muShield.dZ1 == 2.5*u.m

    def setUp(self):
        self.key = "basic"
        config = """
import shipunit as u
from ShipGeoConfig import AttrDict, ConfigRegistry

with ConfigRegistry.register_config("basic") as c:
    c.vetoStation = AttrDict(z=-2390.*u.cm)
    c.TrackStation1 = AttrDict(z=1510.*u.cm)
    c.TrackStation2 = AttrDict(z=1710.*u.cm)
    c.TrackStation3 = AttrDict(z=2150.*u.cm)
    c.TrackStation4 = AttrDict(z=2370.*u.cm)

    c.z = c.TrackStation2.z + 0.5 * (c.TrackStation3.z - c.TrackStation2.z)

    c.Bfield = AttrDict(z=c.z)
    c.Bfield.max = 1.5*u.kilogauss  # was 1.15 in EOI

    # target absorber muon shield setup
    c.decayVolume            =  AttrDict(z=0*u.cm)
    c.decayVolume.length     =   50*u.m

    c.muShield               =  AttrDict(z=0*u.cm)
    c.muShield.dZ1 = 2.5*u.m
    c.muShield.dZ2 = 3.5*u.m
    c.muShield.dZ3 = 3.0*u.m
    c.muShield.dZ4 = 3.0*u.m
    c.muShield.dZ5 = 2.5*u.m
    c.muShield.dZ6 = 2.5*u.m
    c.muShield.LE  = 5*u.m
"""
        ConfigRegistry.loadpys(config)

def main(arguments):
    logger.info("file: %s" % arguments.config_file)
    if arguments.params is not None:
        logger.info("paramters: %s" % arguments.params)
    ConfigRegistry.loadpy(arguments.config_file, **arguments.params)
    # ConfigRegistry.loadpy(arguments.config_file, muShieldDesign=2, targetOpt=5)
    for k, v in ConfigRegistry().iteritems():
        print "%s: %s" % (k, v)

 def InitTask(self):
# prepare container for fitted tracks
  self.comp  = ROOT.TEveCompound('Fitted tracks')
  evmgr.AddElement(self.comp)
  self.trackColors = {13:ROOT.kGreen,211:ROOT.kRed,11:ROOT.kOrange,321:ROOT.kMagenta}
  ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = int(float(dy)))
  self.bfield = ROOT.genfit.BellField(ShipGeo.Bfield.max ,ShipGeo.Bfield.z,2, ShipGeo.Yheight/2.*u.m)
  self.fM = ROOT.genfit.FieldManager.getInstance()
  self.fM.init(self.bfield)
  self.geoMat =  ROOT.genfit.TGeoMaterialInterface()
  ROOT.genfit.MaterialEffects.getInstance().init(self.geoMat)
  dv = top.GetNode('DecayVolume_1')
  ns = dv.GetNodes()
  T1Lid = ns.FindObject("T1Lid_1").GetMatrix()
  self.z_start = T1Lid.GetTranslation()[2]
  mv = top.GetNode('MuonDetector_1').GetMatrix()
  self.z_end = mv.GetTranslation()[2]
  mM = top.GetNode('MCoil_1').GetMatrix()
  self.z_mag = mM.GetTranslation()[2]
  mE = top.GetNode('Ecal_1').GetMatrix()
  self.z_ecal = mE.GetTranslation()[2]
  self.niter = 100
  self.dz = (self.z_end - self.z_start) / float(self.niter)
  self.parallelToZ = ROOT.TVector3(0., 0., 1.) 
  sc    = evmgr.GetScenes()
  self.evscene = sc.FindChild('Event scene')

def execute(f,ox,name='ShipGeo'):
    if type(ox) == type(''): ox = ConfigRegistry.register_config("basic")
    o = retrieveGitTags(ox)
    if type(f)==type("s"): fg = ROOT.TFile.Open(f,'update')
    else:                  fg = f 
    pkl = Pickler(fg)
    pkl.dump(o,name)    
    if type(f)==type("s"): fg.Close()

def loadGeometry(geofile, dy):
    # init geometry and mag. field
    ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy )
    # -----Create geometry----------------------------------------------
    import shipDet_conf
    tgeom = ROOT.TGeoManager("Geometry", "Geane geometry")
    gMan  = tgeom.Import(geofile)
    fGeo = ROOT.gGeoManager
    return {'fGeo':fGeo,'gMan':gMan, 'ShipGeo':ShipGeo}

def configure(run,ship_geo):
# ---- for backward compatibility ----
 if not hasattr(ship_geo,"tankDesign"): ship_geo.tankDesign = 5
 if not hasattr(ship_geo,"muShieldGeo"): ship_geo.muShieldGeo = None
 if not hasattr(ship_geo.hcal,"File"):  ship_geo.hcal.File = "hcal.geo"
 if not hasattr(ship_geo.Bfield,'x') :  ship_geo.Bfield.x   = 3.*u.m
 if not hasattr(ship_geo,'cave') :       
   ship_geo.cave = AttrDict(z=0*u.cm)
   ship_geo.cave.floorHeightMuonShield = 5*u.m
   ship_geo.cave.floorHeightTankA   = 4.5*u.m
   ship_geo.cave.floorHeightTankB   = 2.*u.m
 if not hasattr(ship_geo,'NuTauTT') : ship_geo.NuTauTT= AttrDict(z=0*u.cm)
 if not hasattr(ship_geo.NuTauTT,'design') : ship_geo.NuTauTT.design = 0
 if not hasattr(ship_geo,'EcalOption'):     ship_geo.EcalOption = 1      
 latestShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py",Yheight = ship_geo.Yheight/u.m, tankDesign = ship_geo.tankDesign, muShieldDesign = ship_geo.muShieldDesign, nuTauTargetDesign = ship_geo.nuTauTargetDesign, muShieldGeo = ship_geo.muShieldGeo)
# -----Create media-------------------------------------------------
 run.SetMaterials("media.geo")  # Materials
# ------------------------------------------------------------------------
  
# -----Create geometry----------------------------------------------
 cave= ROOT.ShipCave("CAVE")
 if ship_geo.tankDesign < 5: cave.SetGeometryFileName("cave.geo")
 else: cave.SetGeometryFileName("caveWithAir.geo")
 detectorList.append(cave)

 if ship_geo.muShieldDesign in [6, 7, 8, 9]:  # magnetized hadron absorber defined in ShipMuonShield 
  TargetStation = ROOT.ShipTargetStation("TargetStation",ship_geo.target.length,
                                                        ship_geo.target.z,ship_geo.targetOpt,ship_geo.target.sl)
 else:
  TargetStation = ROOT.ShipTargetStation("TargetStation",ship_geo.target.length,ship_geo.hadronAbsorber.length,
                                                        ship_geo.target.z,ship_geo.hadronAbsorber.z,ship_geo.targetOpt,ship_geo.target.sl)
   
 if ship_geo.targetOpt>10:
  slices_length   = ROOT.std.vector('float')()
  slices_material = ROOT.std.vector('std::string')()
  for i in range(1,ship_geo.targetOpt+1):
   slices_length.push_back(  eval("ship_geo.target.L"+str(i)))
   slices_material.push_back(eval("ship_geo.target.M"+str(i)))
  TargetStation.SetLayerPosMat(ship_geo.target.xy,slices_length,slices_material)
 detectorList.append(TargetStation)

 if ship_geo.muShieldDesign==1:
  MuonShield = ROOT.ShipMuonShield("MuonShield",ship_geo.muShieldDesign,"ShipMuonShield",ship_geo.muShield.z,ship_geo.muShield.dZ0,ship_geo.muShield.length,\
                                   ship_geo.muShield.LE) 
 elif ship_geo.muShieldDesign==2:
  MuonShield = ROOT.ShipMuonShield("MuonShield",ship_geo.muShieldDesign,"ShipMuonShield",ship_geo.muShield.z,ship_geo.muShield.dZ0,ship_geo.muShield.dZ1,\
               ship_geo.muShield.dZ2,ship_geo.muShield.dZ3,ship_geo.muShield.dZ4,ship_geo.muShield.dZ5,ship_geo.muShield.dZ6,ship_geo.muShield.LE) 
 elif ship_geo.muShieldDesign in [3, 4, 5, 6, 7, 9]:
  if not hasattr(ship_geo.muShield,"Field"):
        MuonShield = ROOT.ShipMuonShield(
            "MuonShield", ship_geo.muShieldDesign, "ShipMuonShield",
            ship_geo.muShield.z, ship_geo.muShield.dZ0, ship_geo.muShield.dZ1,
            ship_geo.muShield.dZ2, ship_geo.muShield.dZ3,
            ship_geo.muShield.dZ4, ship_geo.muShield.dZ5,
            ship_geo.muShield.dZ6, ship_geo.muShield.dZ7,
            ship_geo.muShield.dZ8, ship_geo.muShield.dXgap,
            ship_geo.muShield.LE, ship_geo.Yheight * 4. / 10.,
            ship_geo.cave.floorHeightMuonShield)
  else:
        MuonShield = ROOT.ShipMuonShield(
            "MuonShield", ship_geo.muShieldDesign, "ShipMuonShield",
            ship_geo.muShield.z, ship_geo.muShield.dZ0, ship_geo.muShield.dZ1,
            ship_geo.muShield.dZ2, ship_geo.muShield.dZ3,
            ship_geo.muShield.dZ4, ship_geo.muShield.dZ5,
            ship_geo.muShield.dZ6, ship_geo.muShield.dZ7,
            ship_geo.muShield.dZ8, ship_geo.muShield.dXgap,
            ship_geo.muShield.LE, ship_geo.Yheight * 4. / 10.,
               ship_geo.cave.floorHeightMuonShield,ship_geo.muShield.Field) 
 elif ship_geo.muShieldDesign == 8:
  MuonShield = ROOT.ShipMuonShield(ship_geo.muShieldGeo)
 
 detectorList.append(MuonShield)

 if not hasattr(ship_geo,"magnetDesign"):
 # backward compatibility
  magnet_design = 2
  if ship_geo.tankDesign == 5: magnet_design = 3
  if ship_geo.tankDesign == 6: magnet_design = 4
  ship_geo.magnetDesign = magnet_design
  ship_geo.Bfield.YokeWidth = 200.*u.cm 
  ship_geo.Bfield.YokeDepth = 200.*u.cm
  ship_geo.Bfield.CoilThick = 25.*u.cm
# sanity check, 2018 layout ship_geo.tankDesign == 6 has to be together with ship_geo.nuTauTargetDesign == 3
 if (ship_geo.tankDesign == 6 and ship_geo.nuTauTargetDesign != 3) or (ship_geo.tankDesign != 6 and ship_geo.nuTauTargetDesign == 3):
   print "version of tankDesign and nuTauTargetDesign are not compatible, should be 6 and 3, it is ",ship_geo.tankDesign, ship_geo.nuTauTargetDesign 
   exit()
 if ship_geo.strawDesign > 1 : 
  if ship_geo.magnetDesign>3:
   B = ship_geo.Bfield
   magnet = ROOT.ShipMagnet("Magnet","SHiP Magnet",B.z, ship_geo.magnetDesign, B.x, B.y, ship_geo.cave.floorHeightTankB, B.YokeWidth, B.YokeDepth, B.CoilThick)
#                                                               xaperture,  yaperture 
  else: 
   magnet = ROOT.ShipMagnet("Magnet","SHiP Magnet",ship_geo.Bfield.z, ship_geo.magnetDesign, ship_geo.Bfield.x, ship_geo.Bfield.y, ship_geo.cave.floorHeightTankB)
 else: magnet = ROOT.ShipMagnet("Magnet","SHiP Magnet",ship_geo.Bfield.z)
 detectorList.append(magnet)
  
 Veto = ROOT.veto("Veto", ROOT.kTRUE)   # vacuum tank
 Veto.SetLiquidVeto(1)  # liquid scintillator
 Veto.SetPlasticVeto(1) # plastic scintillator

#.........这里部分代码省略.........

  print " for example -f  /eos/experiment/ship/data/muonDIS/muonDis_1.root"
  sys.exit()
if simEngine == "Nuage" and not inputFile:
 inputFile = 'Numucc.root'

print "FairShip setup for",simEngine,"to produce",nEvents,"events"
if (simEngine == "Ntuple" or simEngine == "MuonBack") and defaultInputFile :
  print 'input file required if simEngine = Ntuple or MuonBack'
  print " for example -f /eos/experiment/ship/data/Mbias/pythia8_Geant4-withCharm_onlyMuons_4magTarget.root"
  sys.exit()
ROOT.gRandom.SetSeed(theSeed)  # this should be propagated via ROOT to Pythia8 and Geant4VMC
shipRoot_conf.configure(DarkPhoton)      # load basic libraries, prepare atexit for python
# - muShieldDesign = 2  # 1=passive 5=active (default) 7=short design+magnetized hadron absorber
# - targetOpt      = 5  # 0=solid   >0 sliced, 5: 5 pieces of tungsten, 4 H20 slits, 17: Mo + W +H2O (default)
#   nuTauTargetDesign = 0 # 0 = TP, 1 = NEW with magnet, 2 = NEW without magnet, 3 = 2018 design
if charm == 0: ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy, tankDesign = dv, \
                                                muShieldDesign = ds, nuTauTargetDesign=nud, CaloDesign=caloDesign, strawDesign=strawDesign, muShieldGeo=geofile)
else: ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/charm-geometry_config.py")

# switch off magnetic field to measure muon flux
#ship_geo.muShield.Field = 0.
#ship_geo.EmuMagnet.B = 0.
#ship_geo.tauMudet.B = 0.


# Output file name, add dy to be able to setup geometry with ambiguities.
tag = simEngine+"-"+mcEngine
if charmonly: tag = simEngine+"CharmOnly-"+mcEngine
if eventDisplay: tag = tag+'_D'
if dv > 4 : tag = 'conical.'+tag
elif dy: tag = str(dy)+'.'+tag 
if not os.path.exists(outputDir):

s = sys.argv[1]
thickness = setup[s]['thickness']
material = setup[s]['material']
momentum = setup[s]['momentum']

checkOverlap = True

outFile = "gconv"+s+".root"
theSeed      = int(10000 * time.time() % 10000000)
ecut      = 0.0
                 
# -------------------------------------------------------------------
ROOT.gRandom.SetSeed(theSeed)  # this should be propagated via ROOT to Pythia8 and Geant4VMC
shipRoot_conf.configure()      # load basic libraries, prepare atexit for python
ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = 10, tankDesign = 5, muShieldDesign = 7, nuTauTargetDesign=1)

# -----Timer--------------------------------------------------------
timer = ROOT.TStopwatch()
timer.Start()

# -----Create simulation run----------------------------------------
run = ROOT.FairRunSim()
run.SetName(mcEngine)  # Transport engine
run.SetOutputFile(outFile)  # Output file
run.SetUserConfig("g4Config.C") # user configuration file default g4Config.C
rtdb = run.GetRuntimeDb() 

# -----Materials----------------------------------------------
run.SetMaterials("media.geo")  
# -----Create geometry----------------------------------------------

 def test_false(self):
     ConfigRegistry.loadpys(self.config, MU_SHIELD_ENABLED=False)
     assert len(ConfigRegistry.keys()) == 1, ConfigRegistry.keys()
     assert self.key in ConfigRegistry.keys()
     assert ConfigRegistry[self.key].Bfield.max  == 1.5*u.kilogauss
     self.assertTrue("muShield" not in ConfigRegistry[self.key])

        for f in files:
          os.unlink(os.path.join(root, f))
        for d in dirs:
          shutil.rmtree(os.path.join(root, d))
    else:
      logger.warn("...use '-f' option to overwrite it")
  else:
    os.makedirs(work_dir)
  return args

args = init()
os.chdir(work_dir)
# -------------------------------------------------------------------
ROOT.gRandom.SetSeed(args.seed)  # this should be propagated via ROOT to Pythia8 and Geant4VMC
shipRoot_conf.configure()      # load basic libraries, prepare atexit for python
ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy, tankDesign = dv, muShieldDesign = ds, nuTauTargetDesign=nud)

txt = 'pythia8_Geant4_'
if withEvtGen: txt = 'pythia8_evtgen_Geant4_'
outFile = outputDir+'/'+txt+str(runnr)+'_'+str(ecut)+'.root'
parFile = outputDir+'/ship.params.'+txt+str(runnr)+'_'+str(ecut)+'.root'

# -----Timer--------------------------------------------------------
timer = ROOT.TStopwatch()
timer.Start()

# -----Create simulation run----------------------------------------
run = ROOT.FairRunSim()
run.SetName(mcEngine)  # Transport engine
run.SetOutputFile(outFile)  # Output file
run.SetUserConfig("g4Config.C") # user configuration file default g4Config.C

# analyze muon background /media/Data/HNL/PythiaGeant4Production/pythia8_Geant4_total.root 
import os,ROOT
import rootUtils as ut
import shipunit as u
PDG = ROOT.TDatabasePDG.Instance()
from ShipGeoConfig import ConfigRegistry
# init geometry and mag. field
ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py")
tgeom   = ROOT.TGeoManager("Geometry", "Geane geometry")

rz_inter = -1.,0.
def origin(it):
 at = sTree.MCTrack[it]
 im = at.GetMotherId()
 if im>0: origin(im)
 if im<0: 
   print 'does not come from muon'
   rz_inter = -1.,0.
 if im==0: 
   #print 'origin z',at.GetStartZ()
   rz_inter = ROOT.TMath.Sqrt(at.GetStartX()**2+at.GetStartY()**2),at.GetStartZ()

inputFile = 'ship.MuonBack-TGeant4.root'

withChain =  9  # 9  # 1

if withChain == 1:
 inputFile = 'ship.MuonBack-TGeant4_D.root'

# 11-19 with QGSP_BERT_EMV instead of QGSP_BERT_HP_PEN
# 51-59 passive shielding

        if o in ("-f"):
            inputFile = a
        if o in ("-A"):
            inclusive = True
        if o in ("-F"):
            deepCopy = True

print "FairShip setup for",simEngine,"to produce",nEvents,"events"
if (simEngine == "Ntuple" or simEngine == "MuonBack") and not inputFile :
  print 'input file required if simEngine = Ntuple or MuonBack'
ROOT.gRandom.SetSeed(theSeed)  # this should be propagated via ROOT to Pythia8 and Geant4VMC
shipRoot_conf.configure()      # load basic libraries, prepare atexit for python
# - muShieldDesign = 2  # 1=passive 2=active
# - targetOpt      = 5  # 0=solid   >0 sliced, 5 pieces of tungsten, 4 air slits
# - strawDesign       = 1  # simplistic tracker design,  3=sophisticated straw tube design, horizontal wires
ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py",strawDesign=4,muShieldDesign=5,targetOpt=5)
# Output file name
tag = simEngine+"-"+mcEngine
if eventDisplay: tag = tag+'_D'
outFile ="ship."+tag+".root"  

# rm older files !!! 
os.system("rm *."+tag+".root")
# Parameter file name
parFile="ship.params."+tag+".root"

# In general, the following parts need not be touched
# ========================================================================

# -----Timer--------------------------------------------------------
timer = ROOT.TStopwatch()

if options.updateFile:
 f=ROOT.TFile(fname,'update')
 sTree=f.Get('cbmsim')
 if not sTree: 
   print "Problem with updateFile",f
   exit(-1)
else:
 sTree = ROOT.TChain('cbmsim')
 for f in fnames: 
  print "add ",f
  if options.onEOS: sTree.Add(os.environ['EOSSHIP']+f)
  else:             sTree.Add(f)

#-------------------------------geometry initialization
from ShipGeoConfig import ConfigRegistry
ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/charm-geometry_config.py", Setup = 1, cTarget = 3)
builtin.ShipGeo = ShipGeo
import charmDet_conf
run = ROOT.FairRunSim()
run.SetName("TGeant4")  # Transport engine
run.SetOutputFile(ROOT.TMemFile('output', 'recreate'))  # Output file
run.SetUserConfig("g4Config_basic.C") # geant4 transport not used, only needed for creating VMC field
rtdb = run.GetRuntimeDb()
modules = charmDet_conf.configure(run,ShipGeo)
# -----Create geometry and draw display----------------------------------------------
run.Init()
sGeo = ROOT.gGeoManager
nav = sGeo.GetCurrentNavigator()
top = sGeo.GetTopVolume()
top.SetVisibility(0)
if options.withDisplay: 

 def test_readDOS(self):
     c = ConfigRegistry.loadpy(self.filename, Yheight = 1)
     self.assertTrue("vetoStation" in c)
     assert ConfigRegistry[self.key].vetoStation.z == -2390*u.cm