Scripts/Python/bfield_writing.py

0001 #!/usr/bin/env python3
0002 from pathlib import Path
0003
0004 import acts
0005 import acts.examples
0006 import acts.examples.root
0007
0008 u = acts.UnitConstants
0009
0010
0011 def RootBFieldWrite(bField, fileName, treeName="solenoid", level=acts.logging.VERBOSE):
0012     cfg = acts.examples.root.RootBFieldWriter.Config()
0013     cfg.bField = bField
0014     cfg.gridType = acts.examples.root.RootBFieldWriter.GridType.rz
0015     cfg.fileName = str(fileName)
0016     cfg.treeName = treeName
0017     acts.examples.root.RootBFieldWriter.run(cfg, level)
0018     return cfg
0019
0020
0021 def CsvBFieldWrite(bField, fileName, level=acts.logging.VERBOSE):
0022     cfg = acts.examples.CsvBFieldWriter.ConfigRzGrid()
0023     cfg.bField = bField
0024     cfg.fileName = str(fileName)
0025     acts.examples.CsvBFieldWriter.runRzGrid(cfg, level)
0026     return cfg
0027
0028
0029 def runBFieldWriting(outputDir: Path, rewrites: int = 0):
0030     solenoid = acts.SolenoidBField(
0031         radius=1200 * u.mm, length=6000 * u.mm, bMagCenter=2 * u.T, nCoils=1194
0032     )
0033     field = acts.solenoidFieldMap(
0034         rlim=(0, 1200 * u.mm),
0035         zlim=(-5000 * u.mm, 5000 * u.mm),
0036         nbins=(10, 10),
0037         field=solenoid,
0038     )
0039
0040     print("Solenoid ready")
0041
0042     cfg = RootBFieldWrite(field, outputDir / "solenoid.root")
0043     CsvBFieldWrite(field, outputDir / "solenoid.csv")
0044
0045     for i in range(rewrites):
0046         print(f"Now read back {cfg.fileName}")
0047
0048         field2 = acts.root.MagneticFieldMapRz(cfg.fileName, tree="solenoid")
0049         cfg2 = RootBFieldWrite(field2, outputDir / f"solenoid{i+2}.root")
0050         CsvBFieldWrite(field2, outputDir / f"solenoid{i+2}.csv")
0051         cfg = cfg2
0052
0053     print("Done")
0054
0055
0056 if "__main__" == __name__:
0057     runBFieldWriting(Path.cwd())