Python ConfigRegistry.loadpy方法代码示例

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
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)

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
 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')

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
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}

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
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: 

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
          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(theSeed)  # this should be propagated via ROOT to Pythia8 and Geant4VMC
shipRoot_conf.configure()      # load basic libraries, prepare atexit for python
#this is for the muon flux geometry
ship_geo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/charm-geometry_config.py", Setup = args.CharmdetSetup)

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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
  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):

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
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

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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
def run_track_pattern_recognition(input_file, geo_file, output_file, method):
    """
    Runs all steps of track pattern recognition.

    Parameters
    ----------
    input_file : string
        Path to an input .root file with events.
    geo_file : string
        Path to a file with SHiP geometry.
    output_file : string
        Path to an output .root file with quality plots.
    method : string
        Name of a track pattern recognition method.
    """


    ############################################# Load SHiP geometry ###################################################

    # Check geo file
    try:
        fgeo = ROOT.TFile(geo_file)
    except:
        print "An error with opening the ship geo file."
        raise

    sGeo = fgeo.FAIRGeom

    # Prepare ShipGeo dictionary
    if not fgeo.FindKey('ShipGeo'):

        if sGeo.GetVolume('EcalModule3') :
            ecalGeoFile = "ecal_ellipse6x12m2.geo"
        else:
            ecalGeoFile = "ecal_ellipse5x10m2.geo"

        if dy:
            ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy, EcalGeoFile = ecalGeoFile)
        else:
            ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", EcalGeoFile = ecalGeoFile)

    else:
        upkl    = Unpickler(fgeo)
        ShipGeo = upkl.load('ShipGeo')
    
    # Globals
    builtin.ShipGeo = ShipGeo

    ############################################# Load SHiP modules ####################################################

    run = ROOT.FairRunSim()
    modules = shipDet_conf.configure(run,ShipGeo)

    ############################################# Load inpur data file #################################################

    # Check input file
    try:
        fn = ROOT.TFile(input_file,'update')
    except:
        print "An error with opening the input data file."
        raise

    sTree = fn.cbmsim
    sTree.Write()
    
    ############################################# Create hists #########################################################
    
    h = init_book_hist()

    ########################################## Start Track Pattern Recognition #########################################
    import shipPatRec

    # Init book of hists for the quality measurements
    metrics = {'n_hits': [], 
               'reconstructible': 0, 
               'passed_y12': 0, 'passed_stereo12': 0, 'passed_12': 0, 
               'passed_y34': 0, 'passed_stereo34': 0, 'passed_34': 0, 
               'passed_combined': 0, 'reco_passed': 0, 'reco_passed_no_clones': 0,
               'frac_y12': [], 'frac_stereo12': [], 'frac_12': [], 
               'frac_y34': [], 'frac_stereo34': [], 'frac_34': [], 
               'reco_frac_tot': [], 
               'reco_mc_p': [], 'reco_mc_theta': [], 
               'fitted_p': [], 'fitted_pval': [], 'fitted_chi': [], 
               'fitted_x': [], 'fitted_y': [], 'fitted_z': [], 'fitted_mass': []}

    # Start event loop
    nEvents   = sTree.GetEntries()
   

    for iEvent in range(nEvents):

        if iEvent%1000 == 0:
            print 'Event ', iEvent

        ########################################### Select one event ###################################################

        rc = sTree.GetEvent(iEvent)

        ########################################### Reconstructible tracks #############################################

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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
# 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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
        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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
 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

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
        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()

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
if withOutput:
 outFile = inputFile.replace('.root','_rec.root') 
 os.system('cp '+inputFile+' '+outFile)

#-----prepare python exit-----------------------------------------------
def pyExit():
 global fitter
 del fitter
import atexit
atexit.register(pyExit)

from array import array
import shipunit as u
import rootUtils as ut
from ShipGeoConfig import ConfigRegistry
ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py",muShieldDesign=2,targetOpt=5)


def makePlots():
   ut.bookCanvas(h,key='fitresults',title='Fit Results',nx=1600,ny=1200,cx=2,cy=2)
   cv = h['fitresults'].cd(1)
   h['delPOverP'].Draw('box')
   cv = h['fitresults'].cd(2)
   cv.SetLogy(1)
   h['chi2'].Draw()
   cv = h['fitresults'].cd(3)
   h['delPOverP_proj'] = h['delPOverP'].ProjectionY()
   ROOT.gStyle.SetOptFit(11111)
   h['delPOverP_proj'].Draw()
   h['delPOverP_proj'].Fit('gaus')
   cv = h['fitresults'].cd(4)

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
def configure(run,ship_geo):
# ---- for backward compatibility ----
 if not hasattr(ship_geo,"tankDesign"): ship_geo.tankDesign = 4
 if not hasattr(ship_geo.hcal,"File"): ship_geo.hcal.File = "hcal.geo"
 latestShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py",Yheight = ship_geo.Yheight/u.m, tankDesign = ship_geo.tankDesign, muShieldDesign = ship_geo.muShieldDesign)
# -----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==6 or ship_geo.muShieldDesign==7: # 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:
  TargetStation.SetLayerPosMat(ship_geo.target.xy,ship_geo.target.L1,ship_geo.target.M1,ship_geo.target.L2,ship_geo.target.M2,\
  ship_geo.target.L3,ship_geo.target.M3,ship_geo.target.L4,ship_geo.target.M4,ship_geo.target.L5,ship_geo.target.M5,\
  ship_geo.target.L6,ship_geo.target.M6,ship_geo.target.L7,ship_geo.target.M7,ship_geo.target.L8,ship_geo.target.M8,\
  ship_geo.target.L9,ship_geo.target.M9,ship_geo.target.L10,ship_geo.target.M10,ship_geo.target.L11,ship_geo.target.M11,\
  ship_geo.target.L12,ship_geo.target.M12,ship_geo.target.L13,ship_geo.target.M13,ship_geo.target.L14,ship_geo.target.M14,\
  ship_geo.target.L15,ship_geo.target.M15,ship_geo.target.L16,ship_geo.target.M16,ship_geo.target.L17,ship_geo.target.M17)
 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==3 or ship_geo.muShieldDesign==4 or ship_geo.muShieldDesign==5 or ship_geo.muShieldDesign==6 or ship_geo.muShieldDesign==7 :
  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.) 

 detectorList.append(MuonShield)


 magnet_design = 2
 if ship_geo.tankDesign == 5: magnet_design = 3
 if ship_geo.strawDesign > 1 : 
   if not hasattr(ship_geo.Bfield,'x') :  ship_geo.Bfield.x   = 3.*u.m
   magnet = ROOT.ShipMagnet("Magnet","SHiP Magnet",ship_geo.Bfield.z, magnet_design, ship_geo.Bfield.x, ship_geo.Bfield.y)
 else: magnet = ROOT.ShipMagnet("Magnet","SHiP Magnet",ship_geo.Bfield.z)
 detectorList.append(magnet)
  
 Veto = ROOT.veto("Veto", ROOT.kTRUE)   # vacuum tank, liquid scintillator, simplistic tracking stations
 Veto.SetZpositions(ship_geo.vetoStation.z, ship_geo.TrackStation1.z, ship_geo.TrackStation2.z, \
                    ship_geo.TrackStation3.z, ship_geo.TrackStation4.z,ship_geo.tankDesign)
 Veto.SetTubZpositions(ship_geo.Chamber1.z,ship_geo.Chamber2.z,ship_geo.Chamber3.z,ship_geo.Chamber4.z,ship_geo.Chamber5.z,ship_geo.Chamber6.z);
 Veto.SetTublengths(ship_geo.chambers.Tub1length,ship_geo.chambers.Tub2length,ship_geo.chambers.Tub3length, \
                    ship_geo.chambers.Tub4length,ship_geo.chambers.Tub5length,ship_geo.chambers.Tub6length);
 Veto.SetB(ship_geo.Yheight/2.)
 if ship_geo.tankDesign == 5: 
    dzX = ship_geo.zFocusX+ship_geo.target.z0    
    x1  = ship_geo.xMax/(ship_geo.Chamber1.z -ship_geo.chambers.Tub1length-dzX)*(ship_geo.TrackStation4.z-dzX)
    dzY = ship_geo.zFocusY+ship_geo.target.z0    
    y1  = ship_geo.Yheight/(ship_geo.Chamber1.z -ship_geo.chambers.Tub1length-dzY)*(ship_geo.TrackStation4.z-dzY)
    Veto.SetXYstart(x1,dzX,y1,dzY)
    Veto.SetVesselStructure(ship_geo.Veto.innerSupport,ship_geo.Veto.sensitiveThickness,ship_geo.Veto.outerSupport,\
                            ship_geo.Veto.innerSupportMed,ship_geo.Veto.sensitiveMed,ship_geo.Veto.outerSupportMed,ship_geo.Veto.decayMed,\
                            ship_geo.Veto.rib,ship_geo.Veto.ribMed)
 detectorList.append(Veto)

 if ship_geo.muShieldDesign not in [2,3,4] and hasattr(ship_geo.tauMS,'Xtot'):
  taumagneticspectrometer = ROOT.MagneticSpectrometer("MagneticSpectrometer", ship_geo.tauMS.zMSC,  ROOT.kTRUE)
  taumagneticspectrometer.SetTotDimensions(ship_geo.tauMS.Xtot,ship_geo.tauMS.Ytot, ship_geo.tauMS.Ztot )
  taumagneticspectrometer.SetFeDimensions(ship_geo.tauMS.XFe,ship_geo.tauMS.YFe, ship_geo.tauMS.ZFe)
  taumagneticspectrometer.SetRpcDimensions(ship_geo.tauMS.XRpc,ship_geo.tauMS.YRpc, ship_geo.tauMS.ZRpc)
  taumagneticspectrometer.SetRpcGasDimensions(ship_geo.tauMS.XGas,ship_geo.tauMS.YGas, ship_geo.tauMS.ZGas)
  taumagneticspectrometer.SetRpcStripDimensions(ship_geo.tauMS.XStrip,ship_geo.tauMS.YStrip, ship_geo.tauMS.ZStrip)
  taumagneticspectrometer.SetRpcElectrodeDimensions(ship_geo.tauMS.XEle,ship_geo.tauMS.YEle, ship_geo.tauMS.ZEle)
  taumagneticspectrometer.SetRpcPETDimensions(ship_geo.tauMS.XPet,ship_geo.tauMS.YPet, ship_geo.tauMS.ZPet)
  taumagneticspectrometer.SetReturnYokeDimensions(ship_geo.tauMS.XRyoke,ship_geo.tauMS.YRyoke, ship_geo.tauMS.ZRyoke)
  taumagneticspectrometer.SetSmallerYokeDimensions(ship_geo.tauMS.XRyoke_s,ship_geo.tauMS.YRyoke_s, ship_geo.tauMS.ZRyoke_s)
  taumagneticspectrometer.SetZDimensionArm(ship_geo.tauMS.ZArm)
  taumagneticspectrometer.SetGapDownstream(ship_geo.tauMS.GapD)
  taumagneticspectrometer.SetGapMiddle(ship_geo.tauMS.GapM)
  taumagneticspectrometer.SetNFeInArm(ship_geo.tauMS.NFe)
  taumagneticspectrometer.SetNRpcInArm(ship_geo.tauMS.NRpc)
  taumagneticspectrometer.SetMagneticField(ship_geo.tauMS.B)
  taumagneticspectrometer.SetCoilParameters(ship_geo.tauMS.CoilH, ship_geo.tauMS.CoilW, ship_geo.tauMS.N, ship_geo.tauMS.CoilG)
  detectorList.append(taumagneticspectrometer)

  tauHpt = ROOT.Hpt("HighPrecisionTrackers",ship_geo.tauHPT.DX, ship_geo.tauHPT.DY, ship_geo.tauHPT.DZ, ROOT.kTRUE)
  tauHpt.SetZsize(ship_geo.tauMS.Ztot)
  detectorList.append(tauHpt)

  EmuMagnet = ROOT.EmulsionMagnet("EmuMagnet",ship_geo.EmuMagnet.zC,"EmulsionMagnet")
  EmuMagnet.SetTPDesign(ship_geo.EmuMagnet.TPDesign)
  EmuMagnet.SetGaps(ship_geo.EmuMagnet.GapUp, ship_geo.EmuMagnet.GapDown)
  EmuMagnet.SetMagneticField(ship_geo.EmuMagnet.B)
  EmuMagnet.SetMagnetSizes(ship_geo.EmuMagnet.X, ship_geo.EmuMagnet.Y, ship_geo.EmuMagnet.Z)
#.........这里部分代码省略.........

# 需要导入模块: from ShipGeoConfig import ConfigRegistry [as 别名]
# 或者: from ShipGeoConfig.ConfigRegistry import loadpy [as 别名]
def run_track_pattern_recognition(input_file, geo_file, output_file, method):
    """
    Runs all steps of track pattern recognition.
    Parameters
    ----------
    input_file : string
        Path to an input .root file with events.
    geo_file : string
        Path to a file with SHiP geometry.
    output_file : string
        Path to an output .root file with quality plots.
    method : string
        Name of a track pattern recognition method.
    """


    ############################################# Load SHiP geometry ###################################################

    # Check geo file
    try:
        fgeo = ROOT.TFile(geo_file)
    except:
        print "An error with opening the ship geo file."
        raise

    sGeo = fgeo.FAIRGeom

    # Prepare ShipGeo dictionary
    if not fgeo.FindKey('ShipGeo'):

        if sGeo.GetVolume('EcalModule3') :
            ecalGeoFile = "ecal_ellipse6x12m2.geo"
        else:
            ecalGeoFile = "ecal_ellipse5x10m2.geo"

        if dy:
            ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy, EcalGeoFile = ecalGeoFile)
        else:
            ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", EcalGeoFile = ecalGeoFile)

    else:
        upkl    = Unpickler(fgeo)
        ShipGeo = upkl.load('ShipGeo')
    
    # Globals
    builtin.ShipGeo = ShipGeo

    ############################################# Load SHiP modules ####################################################

    import shipDet_conf
    run = ROOT.FairRunSim()
    run.SetName("TGeant4")  # Transport engine
    run.SetOutputFile("dummy")  # Output file
    run.SetUserConfig("g4Config_basic.C") # geant4 transport not used, only needed for the mag field
    rtdb = run.GetRuntimeDb()

    modules = shipDet_conf.configure(run,ShipGeo)
    run.Init()

    #run = ROOT.FairRunSim()
    #modules = shipDet_conf.configure(run,ShipGeo)


    ######################################### Load SHiP magnetic field #################################################


    import geomGeant4
    if hasattr(ShipGeo.Bfield,"fieldMap"):
      fieldMaker = geomGeant4.addVMCFields(ShipGeo, '', True, withVirtualMC = False)
    else:
      print "no fieldmap given, geofile too old, not anymore support"
      exit(-1)
    sGeo   = fgeo.FAIRGeom
    geoMat =  ROOT.genfit.TGeoMaterialInterface()
    ROOT.genfit.MaterialEffects.getInstance().init(geoMat)
    bfield = ROOT.genfit.FairShipFields()
    bfield.setField(fieldMaker.getGlobalField())
    fM = ROOT.genfit.FieldManager.getInstance()
    fM.init(bfield)


    ############################################# Load inpur data file #################################################

    # Check input file
    try:
        fn = ROOT.TFile(input_file,'update')
    except:
        print "An error with opening the input data file."
        raise

    sTree = fn.cbmsim
    sTree.Write()
    
    ############################################# Create hists #########################################################
    
    h = init_book_hist()

    ########################################## Start Track Pattern Recognition #########################################
    import shipPatRec

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