baluev5.1.2

#!/usr/bin/python
import ROOT, Geant4, math, sys
from Geant4 import cm, cm3, mm, MeV, g, mole, m, deg
from Geant4 import G4Element, G4Material, G4Box
from Geant4 import G4LogicalVolume, G4ThreeVector
import g4py.EMSTDpl
import g4py.NISTmaterials
import g4py.ezgeom
import g4py.ParticleGun
g4py.EMSTDpl.Construct() # electron/gamma standard EM physics list
g4py.NISTmaterials.Construct()
g4py.ezgeom.Construct()
# fill the world with air
g4py.ezgeom.ResizeWorld(2.*m, 2.*m, 4.*m)
g4py.ezgeom.SetWorldMaterial(Geant4.G4Material.GetMaterial('G4_AIR'))



au = G4Material.GetMaterial("G4_Au")
Pb = G4Material.GetMaterial("G4_Pb")

maxEnGun = 10.0

# Create a detector
class ClrSD(Geant4.G4VSensitiveDetector):
    "My sensitive detector"
    def __init__(self):
        Geant4.G4VSensitiveDetector.__init__(self, "ClrSD")
        self.eventEnergy = 0.0
        self.vec = (0., 0., 0.)
        #self.hist2 = hist2
    def ProcessHits(self, step, rohist):
#       print "hist"
        energy = (step.GetTotalEnergyDeposit())/MeV
        self.eventEnergy += energy
        pos = step.GetPostStepPoint().GetPosition()
        #theta = step.GetPostStepPoint().GetPosition().getTheta()
    cofVec = 1. / 1*cm
        self.vec = ( cofVec*(self.vec[0] + pos.getX() * energy),
                cofVec*(self.vec[1] + pos.getY() * energy),
                cofVec*(self.vec[2] + pos.getZ() * energy))
        #theta = math.atan(pos.getRho() / pos.getZ())
        #hist2.Fill(theta, energy, 1)


class ClrEventAction(Geant4.G4UserEventAction):
    def __init__(self, sd, ad, hist, hist2, histab, hist2ab):
        Geant4.G4UserEventAction.__init__(self)
        self.sd = sd
        self.ad = ad
        self.hist = hist
        self.hist2 = hist2
        self.histab = histab
        self.hist2ab = hist2ab

        self.count = 0
        self.count_high_energy = 0

        self.countab = 0
        self.count_high_energyab = 0

    def BeginOfEventAction(self, event):
        self.count += 1
        #print '==== start of event #', self.count
        self.sd.eventEnergy = 0.
        self.sd.vec = (0., 0., 0.)

        self.ad.eventEnergy = 0.
        self.ad.vec = (0., 0., 0.)


    def EndOfEventAction(self, event):
        if self.sd.eventEnergy :
            cofVecsd = 1. / (self.sd.eventEnergy )
            self.hist.Fill(self.sd.eventEnergy)

            self.sd.vec = ( cofVecsd * self.sd.vec[0] , cofVecsd * self.sd.vec[1] , cofVecsd * self.sd.vec[2] )
            self.hist2.Fill(cofVecsd * self.sd.vec[0], cofVecsd * self.sd.vec[1])
            print cofVecsd * self.sd.vec[0], cofVecsd * self.sd.vec[1]
            print '== element ', 1 , ' end of event #', self.count, ', energy=', self.sd.eventEnergy
            # if self.sd.eventEnergy > 0.01 * MeV:
            #     self.count_high_energy += 1

        if self.ad.eventEnergy:
        cofVecad = 1. / (self.ad.eventEnergy * cm )
            self.histab.Fill(self.ad.eventEnergy)

            self.ad.vec = (cofVecad * self.ad.vec[0] , cofVecad * self.ad.vec[1] , cofVecad * self.ad.vec[2] )
            self.hist2ab.Fill(cofVecad * self.ad.vec[0], cofVecad * self.ad.vec[1])
            print '== element ', 2 , ' end of event #', self.countab, ', energy=', self.ad.eventEnergy
            # if self.sd.eventEnergy > 1 * MeV:
            #     self.count_high_energy += 1


calorimeter = g4py.ezgeom.G4EzVolume('Calorimeter')
calorimeter.CreateBoxVolume(au, 10.*cm, 10.*cm, 1.*mm)
calorimeter.PlaceIt(Geant4.G4ThreeVector(0.0, 0.0, 0.5*m-0.5*mm))

absorber = g4py.ezgeom.G4EzVolume('Absorber')
absorber.CreateBoxVolume(Pb, 1.*m, 1.*m, 1.*m)
absorber.PlaceIt(Geant4.G4ThreeVector(0.0, 0.0, -1*mm))


print "1"
histSlr = ROOT.TH1F("histSlr", "Energy", 100, 0., maxEnGun*1.1) # name, description, nbins, low, high
hist2Slr =ROOT.TH2F("hist2Slr", "histogram title", 100, -15., 15., 100, -15., 15.);

histAbs = ROOT.TH1F("histAbs", "Energy2", 100, 0., maxEnGun*1.1) # name, description, nbins, low, high
hist2Abs =ROOT.TH2F("hist2Abs", "histogram title2", 100, -15., 15., 100, -15., 15.);

detector = ClrSD()
detectorAbs = ClrSD()

calorimeter.SetSensitiveDetector(detector)
absorber.SetSensitiveDetector(detectorAbs)

uaction = ClrEventAction(detector, detectorAbs, histSlr, hist2Slr, histAbs, hist2Abs )
# absuaction = ClrEventAction(detectorAbs, histAbs, hist2Abs, 2)

Geant4.gRunManager.Initialize()

# ???
Geant4.gRunManager.SetUserAction(uaction)
# Geant4.gRunManager.SetUserAction(absuaction)

Geant4.gApplyUICommand("/control/execute vis.mac") # enable drawing
Geant4.gApplyUICommand("/vis/disable") # disable drawing


print "2"
pg = g4py.ParticleGun.Construct()
print "3"
pg.SetParticleByName("e-") # particle type:
pg.SetParticlePosition(Geant4.G4ThreeVector(0.0*m, 0.0*m+1*mm, 1.5*m)) # starting point
pg.SetParticleEnergy(maxEnGun*MeV) # energy
for i in range(10): #1000
    phi = ROOT.gRandom.Uniform(0.0, 2 * math.pi)
    theta = ROOT.gRandom.Gaus(0.0, 5.0*math.pi/180.0) #  5
    vx = math.sin(theta) * math.cos(phi)
    vy = math.sin(theta) * math.sin(phi)
    vz = -math.cos(theta)
    if i == 10:
        Geant4.gApplyUICommand("/vis/disable") # disable drawing
    pg.SetParticleMomentumDirection(Geant4.G4ThreeVector(vx, vy, vz))
    Geant4.gRunManager.BeamOn(1) # fire one particle
print "4"
# print 'High energy:', uaction.get_count_high_energy(),
# print 'ratio:', uaction.get_ratio()

c = ROOT.TCanvas('lll', 'nnn')
c.Divide(2, 2)
c.cd(1)
histSlr.Draw()
c.cd(2)
hist2Slr.Draw("lego")
c.cd(3)
histAbs.Draw()
c.cd(4)
hist2Abs.Draw("lego2")

# c2 = ROOT.TCanvas('lll2', 'nnn2')
# c2.Divide(1, 2)
# c2.cd(1)
# histAbs.Draw()

# c2.cd(2)
# hist2Abs.Draw("lego")



sys.stdin.readline()