geomGeant4.py

Go to the documentation of this file.

# SPDX-License-Identifier: LGPL-3.0-or-later
# SPDX-FileCopyrightText: Copyright CERN for the benefit of the SHiP Collaboration
 
from array import array
 
import hepunit as G4Unit
import ROOT
import shipunit as u
 
ROOT.gROOT.ProcessLine('#include "Geant4/G4TransportationManager.hh"')
ROOT.gROOT.ProcessLine('#include "Geant4/G4FieldManager.hh"')
ROOT.gROOT.ProcessLine('#include "Geant4/G4UIterminal.hh"')
ROOT.gROOT.ProcessLine('#include "Geant4/G4RunManager.hh"')
ROOT.gROOT.ProcessLine('#include "TG4GeometryServices.h"')
ROOT.gROOT.ProcessLine('#include "TG4GeometryManager.h"')
 
 
def _collectVolumeNames(volume, listOfVolumes) -> None:
    """Recursively walk the geometry tree collecting volume names from nodes."""
    nodes = volume.GetNodes()
    if not nodes:
        return
    for i in range(nodes.GetEntries()):
        node = nodes.At(i)
        name = node.GetVolume().GetName()
        if name not in listOfVolumes:
            listOfVolumes.append(name)
        _collectVolumeNames(node.GetVolume(), listOfVolumes)
 
 
def check4OrphanVolumes(fGeo) -> None:
    # fill list with volumes from nodes and compare with list of volumes
    top = fGeo.GetTopVolume()
    listOfVolumes = [top.GetName()]
    _collectVolumeNames(top, listOfVolumes)
    orphan = []
    gIndex: dict[str, list] = {}
    for v in fGeo.GetListOfVolumes():
        name = v.GetName()
        if name not in listOfVolumes:
            orphan.append(name)
        if name not in gIndex:
            gIndex[name] = []
        gIndex[name].append(v.GetNumber())
    print("list of orphan volumes:", orphan)
    vSame = {}
    for x in gIndex:
        if len(gIndex[x]) > 1:
            vSame[x] = len(gIndex[x])
    print("list of volumes with same name", vSame)
 
 
def setMagnetField(flag=None) -> None:
    print("setMagnetField() called. Out of date, does not set field for tau neutrino detector!")
    fGeo = ROOT.gGeoManager
    vols = fGeo.GetListOfVolumes()
    # copy field by hand to geant4
    listOfFields = {}
    for v in vols:
        field = v.GetField()
        if not field:
            continue
        bx = field.GetFieldValue()[0] / u.tesla * G4Unit.tesla
        by = field.GetFieldValue()[1] / u.tesla * G4Unit.tesla
        bz = field.GetFieldValue()[2] / u.tesla * G4Unit.tesla
        magFieldIron = ROOT.G4UniformMagField(ROOT.G4ThreeVector(bx, by, bz))
        FieldIronMgr = ROOT.G4FieldManager(magFieldIron)
        FieldIronMgr.CreateChordFinder(magFieldIron)
        listOfFields[v.GetName()] = FieldIronMgr
    gt = ROOT.G4TransportationManager.GetTransportationManager()
    gn = gt.GetNavigatorForTracking()
    world = gn.GetWorldVolume().GetLogicalVolume()
    setField = {}
    for da in range(world.GetNoDaughters()):
        vl0 = world.GetDaughter(da)
        vln = vl0.GetName().c_str()
        lvl0 = vl0.GetLogicalVolume()
        if vln in listOfFields:
            setField[lvl0] = vln
        for dda in range(lvl0.GetNoDaughters()):
            vl = lvl0.GetDaughter(dda)
            vln = vl.GetName().c_str()
            lvl = vl.GetLogicalVolume()
            if vln in listOfFields:
                setField[lvl] = vln
    modified = False
    for lvl in setField:
        # check if field already exists
        fm = lvl.GetFieldManager()
        if not fm.DoesFieldExist():
            lvl.SetFieldManager(listOfFields[setField[lvl]], True)
            modified = True
            if flag == "dump":
                constField = listOfFields[setField[lvl]].GetDetectorField().GetConstantFieldValue()
                print("set field for ", setField[lvl], constField)
        else:
            if flag == "dump":
                print("field already set:", setField[lvl])
    if modified:
        g4Run = ROOT.G4RunManager.GetRunManager()
        g4Run.GeometryHasBeenModified(True)
 
 
def printWF(vl, alreadyPrinted, onlyWithField: bool = True):
    magnetMass = 0
    vname = vl.GetName().data()
    if vname in alreadyPrinted:
        return magnetMass
    vln = vname + " " + str(vl.GetCopyNo())
    mvl = vl.GetMotherLogical().GetName().data()
    alreadyPrinted[vname] = mvl
    if mvl != "cave":
        vln = mvl + "/" + vln
    lvl = vl.GetLogicalVolume()
    cvol = lvl.GetSolid().GetCubicVolume() / G4Unit.m3
    M = lvl.GetMass() / G4Unit.kg
    fm = lvl.GetFieldManager()
    if not fm and onlyWithField:
        return magnetMass
    if M < 5000.0:
        print("%-35s volume = %5.2Fm3  mass = %5.2F kg" % (vln, cvol, M))
    else:
        print("%-35s volume = %5.2Fm3  mass = %5.2F t" % (vln, cvol, M / 1000.0))
    if fm:
        fi = fm.GetDetectorField()
        print("   Magnetic field:", fi.GetConstantFieldValue() / G4Unit.tesla)
    # if vl.GetName().c_str()[0:3]=='Mag': magnetMass =  M # only count volumes starting with Mag
    name = vl.GetName().c_str()
    if "_" in name and "Mag" in name.split("_")[1]:
        magnetMass = M  # only count volumes starting with Mag
    return magnetMass
 
 
def nextLevel(lv, magnetMass: int, onlyWithField, exclude, alreadyPrinted):
    tmp = 0
    for da in range(lv.GetNoDaughters()):
        lvn = lv.GetDaughter(da)
        name = lvn.GetName().c_str()
        if name in exclude:
            continue
        lvln = lvn.GetLogicalVolume()
        if lvln.GetNoDaughters() > 0:
            xtmp, _dummy = nextLevel(lvln, magnetMass, onlyWithField, exclude, alreadyPrinted)
            magnetMass += xtmp
        else:
            tmp += printWF(lvn, alreadyPrinted, onlyWithField)
    return tmp, magnetMass
 
 
def printWeightsandFields(onlyWithField: bool = True, exclude=None) -> None:
    if exclude is None:
        exclude = []
    if len(exclude) != 0:
        print("will not search in ", exclude)
    gt = ROOT.G4TransportationManager.GetTransportationManager()
    gn = gt.GetNavigatorForTracking()
    world = gn.GetWorldVolume().GetLogicalVolume()
    magnetMass = 0
    alreadyPrinted: dict[str, str] = {}
    _dummy, nM = nextLevel(world, magnetMass, onlyWithField, exclude, alreadyPrinted)
    print("total magnet mass", nM / 1000.0, "t")
    return
 
 
def addVMCFields(shipGeo, controlFile: str = "", verbose: bool = False, withVirtualMC: bool = True):
    """
    Define VMC B fields, e.g. global field, field maps, local or local+global fields
    """
    print("Calling addVMCFields")
 
    fieldMaker = ROOT.ShipFieldMaker(verbose)
 
    # Read the input control file. If this is empty then the only fields that are
    # defined (so far) are those within the C++ geometry classes
    if controlFile != "":
        fieldMaker.readInputFile(controlFile)
 
    # Set the main spectrometer field map as a global field
    if hasattr(shipGeo, "Bfield"):
        fieldsList = []
        fieldMaker.defineFieldMap("MainSpecMap", shipGeo.Bfield.fieldMap, ROOT.TVector3(0.0, 0.0, shipGeo.Bfield.z))
        fieldsList.append("MainSpecMap")
 
    if not shipGeo.muShield.WithConstField:
        offset = shipGeo.muShield.Entrance[0]
        quadSymm = True
        file_name = f"files/{shipGeo.shieldName}.root"
        fieldMaker.defineFieldMap(
            "muonShieldField", file_name, ROOT.TVector3(0.0, 0.0, offset), ROOT.TVector3(0.0, 0.0, 0.0), quadSymm
        )
        fieldsList.append("muonShieldField")
 
    elif not shipGeo.hadronAbsorber.WithConstField:
        fieldMaker.defineFieldMap(
            "HadronAbsorberMap",
            "files/FieldHadronStopper_raised_20190411.root",
            ROOT.TVector3(0.0, 0.0, shipGeo.hadronAbsorber.z),
        )
        fieldsList.append("HadronAbsorberMap")
 
    # Combine the fields to obtain the global field
    if len(fieldsList) > 1:
        fieldMaker.defineComposite("TotalField", *fieldsList)  # fieldsList MUST have length <=4
        fieldMaker.defineGlobalField("TotalField")
    else:
        fieldMaker.defineGlobalField("MainSpecMap")
    if withVirtualMC:
        # Force the VMC to update/reset the fields defined by the fieldMaker object.
        # Get the ROOT/Geant4 geometry manager
        geom = ROOT.TG4GeometryManager.Instance()
        # Let the geometry know about the fieldMaker object
        geom.SetUserPostDetConstruction(fieldMaker)
        # Update the fields via the overridden ShipFieldMaker::Construct() function
        geom.ConstructSDandField()
 
    # Return the fieldMaker object, otherwise it will "go out of scope" and its
    # content will be deleted
    return fieldMaker
 
 
def printVMCFields() -> None:
    """
    Method to print out information about VMC fields
    """
    print("Printing VMC fields and associated volumes")
 
    fGeo = ROOT.gGeoManager
    vols = fGeo.GetListOfVolumes()
 
    for v in vols:
        field = v.GetField()
        if field:
            print(f"Vol is {v.GetName()}, field is {field}")
        else:
            print(f"Vol is {v.GetName()}")
 
        if field:
            # Get the field value assuming the global coordinate origin.
            # This needs to be modified to use the local volume centre
            centre = array("d", [0.0, 0.0, 0.0])
            B = array("d", [0.0, 0.0, 0.0])
            field.Field(centre, B)
            print(f"Volume {v.GetName()} has B = ({B[0] / u.tesla}, {B[1] / u.tesla}, {B[2] / u.tesla}) T")
 
 
def getRunManager():
    return ROOT.G4RunManager.GetRunManager()
 
 
def startUI() -> None:
    session = ROOT.G4UIterminal()
    session.SessionStart()
 
 
def debug() -> None:
    gt = ROOT.G4TransportationManager.GetTransportationManager()
    gn = gt.GetNavigatorForTracking()
    world = gn.GetWorldVolume().GetLogicalVolume()
    vmap = {}
    for da in range(world.GetNoDaughters()):
        vl = world.GetDaughter(da)
        vmap[vl.GetName().c_str()] = vl
        print(da, vl.GetName())
    lvl = vmap["MagB"].GetLogicalVolume()
    print(lvl.GetMass() / G4Unit.kg, lvl.GetMaterial().GetName())
    print(lvl.GetFieldManager())
    #
    for da in range(world.GetNoDaughters()):
        vl = world.GetDaughter(da)
        vln = vl.GetName().c_str()
        lvl = vl.GetLogicalVolume()
        fm = lvl.GetFieldManager()
        if fm:
            v = fm.GetDetectorField().GetConstantFieldValue()
            print(vln, fm, v.getX(), v.getY())
    # FairROOT view
    fgeom = ROOT.gGeoManager
    magB = fgeom.GetVolume("MagB")
    fl = magB.GetField()
    print(fl.GetFieldValue()[0], fl.GetFieldValue()[1], fl.GetFieldValue()[2])