diff --git a/workflow/src/legendsimflow/utils.py b/workflow/src/legendsimflow/utils.py index 2671f46..9a75010 100644 --- a/workflow/src/legendsimflow/utils.py +++ b/workflow/src/legendsimflow/utils.py @@ -15,11 +15,12 @@ from __future__ import annotations -import os +import fnmatch from datetime import datetime from pathlib import Path import legenddataflowscripts as ldfs +import pyg4ometry as pg4 from dbetto import AttrsDict from legendmeta import LegendMetadata @@ -81,6 +82,88 @@ def _make_path(d): ) +def _get_matching_volumes(volume_list: list, patterns: str | list) -> list[int]: + """Get the list of volumes from the GDML. The string can include wildcards.""" + + wildcard_list = [patterns] if isinstance(patterns, str) else patterns + + # find all volumes matching at least one pattern + matched_list = [] + for key in volume_list: + for name in wildcard_list: + if fnmatch.fnmatch(key, name): + matched_list.append(key) + return matched_list + + +def get_lar_minishroud_confine_commands( + reg: pg4.geant4.Registry, + pattern: str = "minishroud_side*", + inside: bool = True, + lar_name: str = "LAr", +) -> list[str]: + """Extract the commands for the LAr confinement inside/ outside the NMS from the GDML. + + Parameters + ---------- + reg + The registry describing the geometry. + pattern + The pattern used to search for physical volumes of minishrouds. + inside + Where to generate points only inside NMS to exclude them. + lar_name + The name of the physical volume of the LAr. + + Returns + ------- + a list of confinement commands for remage. + """ + string_list = _get_matching_volumes(list(reg.physicalVolumeDict.keys()), pattern) + # correct sampling mode + mode = "IntersectPhysicalWithGeometrical" if inside else "SubtractGeometrical" + + # physical volume sampling + lines = [f"/RMG/Generator/Confinement/SamplingMode {mode}"] + lines += [f"/RMG/Generator/Confinement/Physical/AddVolume {lar_name}"] + + for s in string_list: + vol = reg.physicalVolumeDict[s] + + center = vol.position.eval() + solid = vol.logicalVolume.solid + + outer_ms = solid.obj1 + r_max = outer_ms.pRMax + dz = outer_ms.pDz + + # type conversions from pyg4ometry types + if not isinstance(r_max, float): + r_max = r_max.eval() + + if not isinstance(dz, float): + dz = dz.eval() + + command = "AddSolid" if inside else "AddExcludeSolid" + lines.append(f"/RMG/Generator/Confinement/Geometrical/{command} Cylinder ") + + lines.append( + f"/RMG/Generator/Confinement/Geometrical/CenterPositionX {center[0]} mm" + ) + lines.append( + f"/RMG/Generator/Confinement/Geometrical/CenterPositionY {center[1]} mm" + ) + lines.append( + f"/RMG/Generator/Confinement/Geometrical/CenterPositionZ {center[2]} mm" + ) + lines.append( + f"/RMG/Generator/Confinement/Geometrical/Cylinder/OuterRadius {r_max} mm" + ) + lines.append(f"/RMG/Generator/Confinement/Geometrical/Cylinder/Height {dz} mm") + + return lines + + def setup_logdir_link(config: SimflowConfig, proctime): logdir = Path(config.paths.log) logdir.mkdir(parents=True, exist_ok=True)