Revision 4fd44c7b3c641895603c91ef130c3d3b10b80eef authored by Axel Naumann on 07 February 2014, 22:40:52 UTC, committed by Axel Naumann on 07 February 2014, 22:41:59 UTC
(cherry picked from commit f5b530ffcaba9cd95b1aafef268bcd5904ba3718)
1 parent 5e1543d
ROOTwriter.py
# @(#)root/gdml:$Id$
# Author: Witold Pokorski 05/06/2006
from math import *
import libPyROOT
import ROOT
import math
import re
# This class provides ROOT binding for the 'writer' class. It implements specific
# methods for all the supported TGeo classes which call the appropriate 'add-element'
# methods from the 'writer' class.
# The list of presently supported classes is the following:
# Materials:
# TGeoElement
# TGeoMaterial
# GeoMixture
# Solids:
# TGeoBBox
# TGeoArb8
# TGeoTubeSeg
# TGeoConeSeg
# TGeoCtub
# TGeoPcon
# TGeoTrap
# TGeoGtra
# TGeoTrd2
# TGeoSphere
# TGeoPara
# TGeoTorus
# TGeoHype
# TGeoPgon
# TGeoXtru
# TGeoEltu
# TGeoParaboloid
# TGeoCompositeShape (subtraction, union, intersection)
# Geometry:
# TGeoVolume
# In addition the class contains three methods 'dumpMaterials', 'dumpSolids' and 'examineVol'
# which retrieve from the memory the materials, the solids and the geometry tree
# respectively. The user should instanciate this class passing and instance of 'writer'
# class as argument. In order to export the geometry in the form of a GDML file,
# the three methods (dumpMaterials, dumpSolids and examineVol) should be called.
# The argument of 'dumpMaterials' method should be the list of materials,
# the argument of the 'dumpSolids' method should be the list of solids and
# the argument of the 'examineVol' method should be the top volume of
# the geometry tree.
# For any question or remarks concerning this code, please send an email to
# Witold.Pokorski@cern.ch.
class ROOTwriter(object):
def __init__(self, writer):
self.writer = writer
self.elements = []
self.volumeCount = 0
self.nodeCount = 0
self.shapesCount = 0
self.bvols = []
self.vols = []
self.volsUseCount = {}
self.sortedVols = []
self.nodes = []
self.bnodes = []
self.solList = []
self.geomgr = ROOT.gGeoManager
self.geomgr.SetAllIndex()
pass
def genName(self, name):
re.sub('$', '', name)
return name
def rotXYZ(self, r):
rad = 180/acos(-1)
cosb = math.sqrt( r[0]*r[0] + r[1]*r[1] )
if cosb > 0.00001 : #I didn't find a proper constant to use here, so I just put a value that works with all the examples on a linux machine (P4)
a = atan2( r[5], r[8] ) * rad
b = atan2( -r[2], cosb ) * rad
c = atan2( r[1], r[0] ) * rad
else:
a = atan2( -r[7], r[4] ) * rad
b = atan2( -r[2], cosb ) * rad
c = 0.
return (a, b, c)
def TGeoBBox(self, solid):
self.writer.addBox(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDX(), 2*solid.GetDY(), 2*solid.GetDZ())
def TGeoParaboloid(self, solid):
self.writer.addParaboloid(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRlo(), solid.GetRhi(), solid.GetDz())
def TGeoSphere(self, solid):
self.writer.addSphere(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRmin(), solid.GetRmax(),
solid.GetPhi1(), solid.GetPhi2() - solid.GetPhi1(),
solid.GetTheta1(), solid.GetTheta2() - solid.GetTheta1())
def TGeoArb8(self, solid):
self.writer.addArb8(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]),
solid.GetVertices()[0],
solid.GetVertices()[1],
solid.GetVertices()[2],
solid.GetVertices()[3],
solid.GetVertices()[4],
solid.GetVertices()[5],
solid.GetVertices()[6],
solid.GetVertices()[7],
solid.GetVertices()[8],
solid.GetVertices()[9],
solid.GetVertices()[10],
solid.GetVertices()[11],
solid.GetVertices()[12],
solid.GetVertices()[13],
solid.GetVertices()[14],
solid.GetVertices()[15],
solid.GetDz())
def TGeoConeSeg(self, solid):
self.writer.addCone(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDz(), solid.GetRmin1(), solid.GetRmin2(),
solid.GetRmax1(), solid.GetRmax2(), solid.GetPhi1(), solid.GetPhi2() - solid.GetPhi1())
def TGeoCone(self, solid):
self.writer.addCone(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDz(), solid.GetRmin1(), solid.GetRmin2(),
solid.GetRmax1(), solid.GetRmax2(), 0, 360)
def TGeoPara(self, solid):
self.writer.addPara(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetX(), solid.GetY(), solid.GetZ(),
solid.GetAlpha(), solid.GetTheta(), solid.GetPhi())
def TGeoTrap(self, solid):
self.writer.addTrap(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDz(), solid.GetTheta(), solid.GetPhi(),
2*solid.GetH1(), 2*solid.GetBl1(), 2*solid.GetTl1(), solid.GetAlpha1(),
2*solid.GetH2(), 2*solid.GetBl2(), 2*solid.GetTl2(), solid.GetAlpha2())
def TGeoGtra(self, solid):
self.writer.addTwistedTrap(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDz(), solid.GetTheta(), solid.GetPhi(),
2*solid.GetH1(), 2*solid.GetBl1(), 2*solid.GetTl1(), solid.GetAlpha1(),
2*solid.GetH2(), 2*solid.GetBl2(), 2*solid.GetTl2(), solid.GetAlpha2(), solid.GetTwistAngle())
def TGeoTrd1(self, solid):
self.writer.addTrd(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDx1(), 2*solid.GetDx2(), 2*solid.GetDy(),
2*solid.GetDy(), 2*solid.GetDz())
def TGeoTrd2(self, solid):
self.writer.addTrd(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), 2*solid.GetDx1(), 2*solid.GetDx2(), 2*solid.GetDy1(),
2*solid.GetDy2(), 2*solid.GetDz())
def TGeoTubeSeg(self, solid):
self.writer.addTube(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRmin(), solid.GetRmax(),
2*solid.GetDz(), solid.GetPhi1(), solid.GetPhi2()-solid.GetPhi1())
def TGeoCtub(self, solid):
self.writer.addCutTube(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRmin(), solid.GetRmax(),
2*solid.GetDz(), solid.GetPhi1(), solid.GetPhi2()-solid.GetPhi1(),
solid.GetNlow()[0],
solid.GetNlow()[1],
solid.GetNlow()[2],
solid.GetNhigh()[0],
solid.GetNhigh()[1],
solid.GetNhigh()[2])
def TGeoTube(self, solid):
self.writer.addTube(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRmin(), solid.GetRmax(),
2*solid.GetDz(), 0, 360)
def TGeoPcon(self, solid):
zplanes = []
for i in range(solid.GetNz()):
zplanes.append( (solid.GetZ(i), solid.GetRmin(i), solid.GetRmax(i)) )
self.writer.addPolycone(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetPhi1(), solid.GetDphi(), zplanes)
def TGeoTorus(self, solid):
self.writer.addTorus(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetR(), solid.GetRmin(), solid.GetRmax(),
solid.GetPhi1(), solid.GetDphi())
def TGeoPgon(self, solid):
zplanes = []
for i in range(solid.GetNz()):
zplanes.append( (solid.GetZ(i), solid.GetRmin(i), solid.GetRmax(i)) )
self.writer.addPolyhedra(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetPhi1(), solid.GetDphi(),
solid.GetNedges(), zplanes)
def TGeoXtru(self, solid):
vertices = []
sections = []
for i in range(solid.GetNvert()):
vertices.append( (solid.GetX(i), solid.GetY(i)) )
for i in range(solid.GetNz()):
sections.append( (i, solid.GetZ(i), solid.GetXOffset(i), solid.GetYOffset(i), solid.GetScale(i)) )
self.writer.addXtrusion(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), vertices, sections)
def TGeoEltu(self, solid):
self.writer.addEltube(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetA(), solid.GetB(), solid.GetDz())
def TGeoHype(self, solid):
self.writer.addHype(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]), solid.GetRmin(), solid.GetRmax(),
solid.GetStIn(), solid.GetStOut(), 2*solid.GetDz())
def TGeoUnion(self, solid):
lrot = self.rotXYZ(solid.GetBoolNode().GetLeftMatrix().Inverse().GetRotationMatrix())
rrot = self.rotXYZ(solid.GetBoolNode().GetRightMatrix().Inverse().GetRotationMatrix())
if ([solid.GetBoolNode().GetLeftShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetLeftShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetLeftShape().__class__.__name__)(solid.GetBoolNode().GetLeftShape())
self.shapesCount = self.shapesCount + 1
if ([solid.GetBoolNode().GetRightShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetRightShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetRightShape().__class__.__name__)(solid.GetBoolNode().GetRightShape())
self.shapesCount = self.shapesCount + 1
self.writer.addUnion(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]),
solid.GetBoolNode().GetLeftShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetLeftShape())[0]),
solid.GetBoolNode().GetLeftMatrix().GetTranslation(),
lrot,
solid.GetBoolNode().GetRightShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetRightShape())[0]),
solid.GetBoolNode().GetRightMatrix().GetTranslation(),
rrot)
def TGeoIntersection(self, solid):
lrot = self.rotXYZ(solid.GetBoolNode().GetLeftMatrix().Inverse().GetRotationMatrix())
rrot = self.rotXYZ(solid.GetBoolNode().GetRightMatrix().Inverse().GetRotationMatrix())
if ([solid.GetBoolNode().GetLeftShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetLeftShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetLeftShape().__class__.__name__)(solid.GetBoolNode().GetLeftShape())
self.shapesCount = self.shapesCount + 1
if ([solid.GetBoolNode().GetRightShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetRightShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetRightShape().__class__.__name__)(solid.GetBoolNode().GetRightShape())
self.shapesCount = self.shapesCount + 1
self.writer.addIntersection(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]),
solid.GetBoolNode().GetLeftShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetLeftShape())[0]),
solid.GetBoolNode().GetLeftMatrix().GetTranslation(),
lrot,
solid.GetBoolNode().GetRightShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetRightShape())[0]),
solid.GetBoolNode().GetRightMatrix().GetTranslation(),
rrot)
def TGeoSubtraction(self, solid):
lrot = self.rotXYZ(solid.GetBoolNode().GetLeftMatrix().Inverse().GetRotationMatrix())
rrot = self.rotXYZ(solid.GetBoolNode().GetRightMatrix().Inverse().GetRotationMatrix())
if ([solid.GetBoolNode().GetLeftShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetLeftShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetLeftShape().__class__.__name__)(solid.GetBoolNode().GetLeftShape())
self.shapesCount = self.shapesCount + 1
if ([solid.GetBoolNode().GetRightShape(), 0]) in self.solList:
self.solList[self.solList.index([solid.GetBoolNode().GetRightShape(), 0])][1] = 1
eval('self.'+solid.GetBoolNode().GetRightShape().__class__.__name__)(solid.GetBoolNode().GetRightShape())
self.shapesCount = self.shapesCount + 1
self.writer.addSubtraction(self.genName(solid.GetName())+'_'+str(libPyROOT.AddressOf(solid)[0]),
solid.GetBoolNode().GetLeftShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetLeftShape())[0]),
solid.GetBoolNode().GetLeftMatrix().GetTranslation(),
lrot,
solid.GetBoolNode().GetRightShape().GetName()+'_'+str(libPyROOT.AddressOf(solid.GetBoolNode().GetRightShape())[0]),
solid.GetBoolNode().GetRightMatrix().GetTranslation(),
rrot)
def TGeoCompositeShape(self, solid):
eval('self.'+solid.GetBoolNode().__class__.__name__)(solid)
def dumpMaterials(self, matlist):
print 'Info in <TPython::Exec>: Found ', matlist.GetSize(),' materials'
for mat in matlist:
if not mat.IsMixture():
self.writer.addMaterial(self.genName(mat.GetName()), mat.GetA(), mat.GetZ(), mat.GetDensity())
else:
elems = {}
for index in range(mat.GetNelements()):
elems[mat.GetElement(index).GetName()] = mat.GetWmixt()[index]
el = mat.GetElement(index)
if el not in self.elements:
self.elements.append(el)
self.writer.addElement(mat.GetElement(index).GetTitle(), mat.GetElement(index).GetName(), mat.GetZmixt()[index], mat.GetAmixt()[index])
self.writer.addMixture(self.genName(mat.GetName()), mat.GetDensity(), elems)
def dumpSolids(self, shapelist):
print 'Info in <TPython::Exec>: Found ', shapelist.GetEntries(), ' shapes'
for shape in shapelist:
self.solList.append([shape, 0])
for sol in self.solList:
if sol[1] == 0:
sol[1] = 1
#print eval('self.'+sol[0].__class__.__name__)(sol[0])
eval('self.'+sol[0].__class__.__name__)(sol[0])
self.shapesCount = self.shapesCount + 1
print 'Info in <TPython::Exec>: Dumped ', self.shapesCount, ' shapes'
def orderVolumes(self, volume):
index = str(volume.GetNumber())+"_"+str(libPyROOT.AddressOf(volume)[0])
daughters = volume.GetNodes()
if len(self.sortedVols)<len(self.vols) and self.volsUseCount[index]>0:
self.volsUseCount[index] = self.volsUseCount[index]-1
if self.volsUseCount[index]==0:
self.sortedVols.append(volume)
if daughters:
for node in daughters:
self.orderVolumes(node.GetVolume())
self.nodeCount = self.nodeCount+1
if self.nodeCount%10000==0:
print '[FIRST STAGE] Node count: ', self.nodeCount
elif len(self.sortedVols)<len(self.volsUseCount) and self.volsUseCount[index]==0:
self.sortedVols.append(volume)
if daughters:
for node in daughters:
self.orderVolumes(node.GetVolume())
self.nodeCount = self.nodeCount+1
if self.nodeCount%10000==0:
print '[FIRST STAGE] Node count: ', self.nodeCount
def getNodes(self, volume):
nd = volume.GetNdaughters()
if nd:
for i in range(nd):
currentNode = volume.GetNode(i)
nextVol = currentNode.GetVolume()
index = str(nextVol.GetNumber())+"_"+str(libPyROOT.AddressOf(nextVol)[0])
self.volsUseCount[index] = self.volsUseCount[index]+1
self.nodes.append(currentNode)
self.getNodes(nextVol)
self.nodeCount = self.nodeCount+1
if self.nodeCount%10000==0:
print '[ZEROTH STAGE] Analysing node: ', self.nodeCount
def examineVol2(self, volume): #use with geometries containing many volumes
print ''
print '[RETRIEVING VOLUME LIST]'
self.bvols = geomgr.GetListOfVolumes()
print ''
print '[INITIALISING VOLUME USE COUNT]'
for vol in self.bvols:
self.vols.append(vol)
self.volsUseCount[str(vol.GetNumber())+"_"+str(libPyROOT.AddressOf(vol)[0])]=0
print ''
print '[CALCULATING VOLUME USE COUNT]'
self.nodeCount = 0
self.getNodes(volume)
print ''
print '[ORDERING VOLUMES]'
self.nodeCount = 0
self.orderVolumes(volume)
print ''
print '[DUMPING GEOMETRY TREE]'
self.sortedVols.reverse()
self.nodeCount = 0
self.dumpGeoTree()
print ''
print '[FINISHED!]'
print ''
def examineVol(self, volume): #use with geometries containing very few volumes and many nodes
daughters = []
if volume.GetNodes():
for node in volume.GetNodes():
subvol = node.GetVolume()
#if bit not set, set and save primitive
if not subvol.TestAttBit(524288): #value referring to TGeoAtt::kSavePrimitiveAtt (1 << 19)
subvol.SetAttBit(524288)
self.vols.append(subvol)
self.examineVol(subvol)
name = node.GetName()+str(libPyROOT.AddressOf(subvol)[0])+'in'+volume.GetName()+str(libPyROOT.AddressOf(volume)[0])
pos = node.GetMatrix().GetTranslation()
self.writer.addPosition(name+'pos', pos[0], pos[1], pos[2])
r = self.rotXYZ(node.GetMatrix().GetRotationMatrix())
rotname = ''
if r[0]!=0.0 or r[1]!=0.0 or r[2]!=0.0:
self.writer.addRotation(name+'rot', r[0], r[1], r[2])
rotname = name+'rot'
reflection = node.GetMatrix().IsReflection()#check if this daughter has a reflection matrix
if reflection:
rotmat = node.GetMatrix().GetRotationMatrix()
#add new 'reflectedSolid' shape to solids
self.writer.addReflSolid('refl_'+node.GetVolume().GetShape().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume().GetShape())[0]), node.GetVolume().GetShape().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume().GetShape())[0]), 0, 0, 0, rotmat[0], rotmat[4], rotmat[8], 0, 0, 0)
#add new volume with correct solidref to the new reflectedSolid
emptyd = []
self.writer.addVolume('refl_'+node.GetVolume().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume())[0]), 'refl_'+node.GetVolume().GetShape().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume().GetShape())[0]), self.genName(node.GetVolume().GetMaterial().GetName()), emptyd)
#add new volume as volumeref to this physvol
daughters.append( ('refl_'+node.GetVolume().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume())[0]), name+'pos', rotname) )
else:
daughters.append( (node.GetVolume().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume())[0]), name+'pos', rotname) )
if volume.IsTopVolume():
if not volume.IsAssembly():
self.writer.addVolume(volume.GetName(), volume.GetShape().GetName()+'_'+str(libPyROOT.AddressOf(volume.GetShape())[0]), self.genName(volume.GetMaterial().GetName()), daughters)
else:
self.writer.addAssembly(volume.GetName(), daughters)
else:
if not volume.IsAssembly():
self.writer.addVolume(volume.GetName()+'_'+str(libPyROOT.AddressOf(volume)[0]), volume.GetShape().GetName()+'_'+str(libPyROOT.AddressOf(volume.GetShape())[0]), self.genName(volume.GetMaterial().GetName()), daughters)
else:
self.writer.addAssembly(volume.GetName()+'_'+str(libPyROOT.AddressOf(volume)[0]), daughters)
def dumpGeoTree(self):
for volume in self.sortedVols:
nd = volume.GetNdaughters()
daughters = []
if nd:
for i in range(nd):
node = volume.GetNode(i)
name = node.GetName()+'in'+volume.GetName()
pos = node.GetMatrix().GetTranslation()
self.writer.addPosition(name+'pos', pos[0], pos[1], pos[2])
r = self.rotXYZ(node.GetMatrix().GetRotationMatrix())
rotname = ''
if r[0]!=0.0 or r[1]!=0.0 or r[2]!=0.0:
self.writer.addRotation(name+'rot', r[0], r[1], r[2])
rotname = name+'rot'
daughters.append( (node.GetVolume().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume())[0]), name+'pos', rotname) )
self.nodeCount = self.nodeCount+1
if self.nodeCount%100==0:
print '[SECOND STAGE] Volume Count: ', self.nodeCount, node.GetVolume().GetName()+'_'+str(libPyROOT.AddressOf(node.GetVolume())[0])
if volume.IsTopVolume():
if not volume.IsAssembly():
self.writer.addVolume(volume.GetName(), volume.GetShape().GetName()+'_'+str(libPyROOT.AddressOf(volume.GetShape())[0]), self.genName(volume.GetMaterial().GetName()), daughters)
else:
self.writer.addAssembly(volume.GetName(), daughters)
else:
if not volume.IsAssembly():
self.writer.addVolume(volume.GetName()+'_'+str(libPyROOT.AddressOf(volume)[0]), volume.GetShape().GetName()+'_'+str(libPyROOT.AddressOf(volume.GetShape())[0]), self.genName(volume.GetMaterial().GetName()), daughters)
else:
self.writer.addAssembly(volume.GetName()+'_'+str(libPyROOT.AddressOf(volume)[0]), daughters)
Computing file changes ...
