Scripts/Python/material_merge_rz_view.py

0001 #!/usr/bin/env python3
0002 """Draw an r-z view of a (Gen3) tracking geometry and highlight lossy material merges.
0003
0004 This is a standalone script that only depends on the **pure-core** ACTS Python
0005 bindings (the ``acts`` module, including ``acts.json``) plus matplotlib. It does
0006 *not* use the Examples framework.
0007
0008 It reads a tracking-geometry JSON file (as produced by
0009 ``acts.json.TrackingGeometryJsonConverter``), draws every cylinder volume as a
0010 rectangle in the r-z plane, and overlays -- as orange line segments -- any portal
0011 surface that carries a ``MergedMaterialMarker``. The marker is left behind by
0012 ``Portal::merge`` when it runs in "keep going" mode and has to discard surface
0013 material during container stacking, so this plot shows exactly *where* those
0014 problematic merges happened.
0015
0016 Example::
0017
0018     python material_merge_rz_view.py tracking-geometry.json -o detector_rz.svg
0019 """
0020
0021 import argparse
0022 from pathlib import Path
0023
0024 import acts
0025
0026
0027 class _RZCollector(acts.TrackingGeometryMutableVisitor):
0028     """Visitor collecting cylinder volume rectangles and marker line segments."""
0029
0030     def __init__(self, gctx: acts.GeometryContext):
0031         super().__init__()
0032         self._gctx = gctx
0033         # Each entry: (z_min, z_max, r_min, r_max, name)
0034         self.volumes = []
0035         # Each entry: (z0, r0, z1, r1)
0036         self.markers = []
0037         self._seen_markers = set()
0038
0039     def visitVolume(self, volume: acts.Volume):
0040         bounds = volume.volumeBounds
0041         if bounds.type() != acts.VolumeBoundsType.Cylinder:
0042             return
0043         # CylinderVolumeBounds: values()[0..2] = rMin, rMax, halfLengthZ
0044         values = bounds.values()
0045         r_min, r_max, half_z = values[0], values[1], values[2]
0046         z_center = volume.center(self._gctx)[2]
0047         name = ""
0048         if isinstance(volume, acts.TrackingVolume):
0049             name = volume.volumeName
0050         self.volumes.append((z_center - half_z, z_center + half_z, r_min, r_max, name))
0051
0052     def visitPortal(self, portal: acts.Portal):
0053         self._maybe_marker(portal.surface)
0054
0055     def visitSurface(self, surface: acts.Surface):
0056         self._maybe_marker(surface)
0057
0058     def _maybe_marker(self, surface: acts.Surface):
0059         material = surface.surfaceMaterial
0060         if not isinstance(material, acts.MergedMaterialMarker):
0061             return
0062         segment = self._surface_segment(surface)
0063         if segment is None:
0064             return
0065         # A merged portal surface is reachable both as a portal and (potentially)
0066         # as a volume surface, and is shared between the stacked volumes, so the
0067         # same segment is seen multiple times. Deduplicate by rounded geometry.
0068         key = tuple(round(v, 3) for v in segment)
0069         if key in self._seen_markers:
0070             return
0071         self._seen_markers.add(key)
0072         self.markers.append(segment)
0073
0074     def _surface_segment(self, surface: acts.Surface):
0075         """Return the (z0, r0, z1, r1) r-z footprint of a portal surface."""
0076         bounds = surface.bounds
0077         center = surface.center(self._gctx)
0078         z_center = center[2]
0079         if isinstance(bounds, acts.CylinderBounds):
0080             # Cylinder face: horizontal line at fixed radius spanning z.
0081             r = bounds.values()[int(acts.CylinderBoundsValue.R)]
0082             half_z = bounds.values()[int(acts.CylinderBoundsValue.HalfLengthZ)]
0083             return (z_center - half_z, r, z_center + half_z, r)
0084         if isinstance(bounds, acts.RadialBounds):
0085             # Disc face: vertical line at fixed z spanning r.
0086             r_min = bounds.values()[int(acts.RadialBoundsValue.MinR)]
0087             r_max = bounds.values()[int(acts.RadialBoundsValue.MaxR)]
0088             return (z_center, r_min, z_center, r_max)
0089         # Unsupported surface kind for the r-z view; skip it.
0090         return None
0091
0092
0093 def load_tracking_geometry(
0094     json_path: Path, gctx: acts.GeometryContext
0095 ) -> acts.TrackingGeometry:
0096     """Load a Gen3 tracking geometry from a JSON file."""
0097     from acts.json import TrackingGeometryJsonConverter
0098
0099     converter = TrackingGeometryJsonConverter()
0100     return converter.fromJson(gctx, Path(json_path).read_text())
0101
0102
0103 def collect_rz(tracking_geometry: acts.TrackingGeometry, gctx: acts.GeometryContext):
0104     """Collect cylinder volume rectangles and marker segments from the geometry.
0105
0106     Returns a tuple ``(volumes, markers)``.
0107     """
0108     collector = _RZCollector(gctx)
0109     tracking_geometry.apply(collector)
0110     return collector.volumes, collector.markers
0111
0112
0113 def render(volumes, markers, output_path: Path, title: str = "Material merge r-z view"):
0114     """Render the r-z view to an SVG (or any matplotlib-supported) file."""
0115     import matplotlib
0116
0117     matplotlib.use("Agg")
0118     import matplotlib.pyplot as plt
0119     from matplotlib.patches import Rectangle
0120
0121     fig, ax = plt.subplots(figsize=(12, 6))
0122
0123     for z_min, z_max, r_min, r_max, name in volumes:
0124         ax.add_patch(
0125             Rectangle(
0126                 (z_min, r_min),
0127                 z_max - z_min,
0128                 r_max - r_min,
0129                 facecolor="#dfe7f2",
0130                 edgecolor="#34507a",
0131                 linewidth=0.6,
0132                 zorder=1,
0133             )
0134         )
0135         if name:
0136             # ODD volume names are hierarchical (pipe-separated); show the leaf.
0137             leaf = name.split("|")[-1]
0138             # Rotate the label to run along the longer side of the rectangle.
0139             rotation = 90 if (r_max - r_min) > (z_max - z_min) else 0
0140             ax.text(
0141                 0.5 * (z_min + z_max),
0142                 0.5 * (r_min + r_max),
0143                 leaf,
0144                 ha="center",
0145                 va="center",
0146                 rotation=rotation,
0147                 fontsize=5,
0148                 color="#1b2a44",
0149                 clip_on=True,
0150                 zorder=3,
0151             )
0152
0153     marker_label_used = False
0154     for z0, r0, z1, r1 in markers:
0155         ax.plot(
0156             [z0, z1],
0157             [r0, r1],
0158             color="darkorange",
0159             linewidth=3.0,
0160             solid_capstyle="round",
0161             zorder=5,
0162             label=None if marker_label_used else "Discarded merge material",
0163         )
0164         marker_label_used = True
0165
0166     ax.set_xlabel("z [mm]")
0167     ax.set_ylabel("r [mm]")
0168     ax.set_title(title)
0169     ax.autoscale_view()
0170     ax.margins(0.02)
0171     ax.set_ylim(bottom=0)
0172     if marker_label_used:
0173         ax.legend(loc="upper right")
0174
0175     fig.tight_layout()
0176     fig.savefig(str(output_path))
0177     plt.close(fig)
0178     return output_path
0179
0180
0181 def run(json_path: Path, output_path: Path, title: str = "Material merge r-z view"):
0182     """Load a geometry JSON, build the r-z view, and write it to ``output_path``."""
0183     gctx = acts.GeometryContext.dangerouslyDefaultConstruct()
0184     tracking_geometry = load_tracking_geometry(json_path, gctx)
0185     volumes, markers = collect_rz(tracking_geometry, gctx)
0186     render(volumes, markers, output_path, title=title)
0187     return volumes, markers
0188
0189
0190 def main():
0191     parser = argparse.ArgumentParser(description=__doc__)
0192     parser.add_argument(
0193         "geometry",
0194         type=Path,
0195         help="Path to a Gen3 tracking-geometry JSON file",
0196     )
0197     parser.add_argument(
0198         "-o",
0199         "--output",
0200         type=Path,
0201         default=Path("geometry_rz.svg"),
0202         help="Output image path (default: geometry_rz.svg)",
0203     )
0204     parser.add_argument(
0205         "--title",
0206         default="Material merge r-z view",
0207         help="Plot title",
0208     )
0209     args = parser.parse_args()
0210
0211     volumes, markers = run(args.geometry, args.output, title=args.title)
0212     print(
0213         f"Drew {len(volumes)} cylinder volumes and {len(markers)} "
0214         f"merge-material markers to {args.output}"
0215     )
0216
0217
0218 if __name__ == "__main__":
0219     main()