Examples/tests/test_geometry.py

0001 import pytest
0002 import acts
0003 import functools
0004 from acts.examples import GenericDetector
0005 from acts.examples.odd import getOpenDataDetector
0006 import json
0007
0008 from helpers import dd4hepEnabled
0009
0010
0011 @pytest.mark.parametrize(
0012     "detectorFactory,aligned,nobj",
0013     [
0014         (functools.partial(GenericDetector, gen3=False), True, 450),
0015         pytest.param(
0016             functools.partial(GenericDetector, gen3=True),
0017             True,
0018             2,  # Gen3 geometry visualiztion produces a single file + materials
0019         ),
0020         pytest.param(
0021             getOpenDataDetector,
0022             True,
0023             540,
0024             marks=[
0025                 pytest.mark.skipif(not dd4hepEnabled, reason="DD4hep not set up"),
0026                 pytest.mark.slow,
0027                 pytest.mark.odd,
0028             ],
0029         ),
0030     ],
0031     ids=[
0032         "generic",
0033         "generic-gen3",
0034         "odd",
0035     ],
0036 )
0037 @pytest.mark.slow
0038 def test_geometry_example(detectorFactory, aligned, nobj, tmp_path):
0039     detector = detectorFactory()
0040     trackingGeometry = detector.trackingGeometry()
0041     decorators = detector.contextDecorators()
0042
0043     from geometry import runGeometry
0044
0045     json_dir = tmp_path / "json"
0046     csv_dir = tmp_path / "csv"
0047     obj_dir = tmp_path / "obj"
0048
0049     for d in (json_dir, csv_dir, obj_dir):
0050         d.mkdir()
0051
0052     events = 5
0053
0054     kwargs = dict(
0055         trackingGeometry=trackingGeometry,
0056         decorators=decorators,
0057         events=events,
0058         outputDir=tmp_path,
0059     )
0060
0061     runGeometry(outputJson=True, **kwargs)
0062     runGeometry(outputJson=False, **kwargs)
0063
0064     assert len(list(obj_dir.iterdir())) == nobj
0065
0066     assert len(list(csv_dir.iterdir())) == 3 * events
0067
0068     detector_files = [csv_dir / f"event{i:>09}-detectors.csv" for i in range(events)]
0069     for detector_file in detector_files:
0070         assert detector_file.exists()
0071         assert detector_file.stat().st_size > 200
0072
0073     contents = [f.read_text() for f in detector_files]
0074     ref = contents[0]
0075     for c in contents[1:]:
0076         if aligned:
0077             assert c == ref, "Detector writeout is expected to be identical"
0078         else:
0079             assert c != ref, "Detector writeout is expected to be different"
0080
0081     if aligned:
0082         for f in [json_dir / f"event{i:>09}-detector.json" for i in range(events)]:
0083             assert detector_file.exists()
0084             with f.open() as fh:
0085                 data = json.load(fh)
0086                 assert data
0087         material_file = tmp_path / "geometry-map.json"
0088         assert material_file.exists()
0089         assert material_file.stat().st_size > 200
0090
0091
0092 class CountingVisitor(acts.TrackingGeometryMutableVisitor):
0093     def __init__(self):
0094         super().__init__()
0095         self.num_surfaces = 0
0096         self.num_layers = 0
0097         self.num_volumes = 0
0098         self.num_portals = 0
0099         self.num_boundary_surfaces = 0
0100
0101     def visitSurface(self, surface: acts.Surface):
0102         self.num_surfaces += 1
0103
0104     def visitLayer(self, layer: acts.Layer):
0105         self.num_layers += 1
0106
0107     def visitVolume(self, volume: acts.Volume):
0108         self.num_volumes += 1
0109
0110     def visitPortal(self, portal: acts.Portal):
0111         self.num_portals += 1
0112
0113     def visitBoundarySurface(self, boundary: acts.BoundarySurfaceT_TrackingVolume):
0114         self.num_boundary_surfaces += 1
0115
0116
0117 def test_geometry_visitor(trk_geo):
0118     visitor = CountingVisitor()
0119     trk_geo.apply(visitor)
0120
0121     assert visitor.num_surfaces == 19078
0122     assert visitor.num_layers == 111
0123     assert visitor.num_volumes == 18
0124     assert visitor.num_portals == 0
0125     assert visitor.num_boundary_surfaces == 67
0126
0127
0128 @pytest.mark.skipif(not dd4hepEnabled, reason="DD4hep not set up")
0129 @pytest.mark.odd
0130 def test_odd_gen1():
0131     with getOpenDataDetector(gen3=False) as detector:
0132         trackingGeometry = detector.trackingGeometry()
0133
0134         visitor = CountingVisitor()
0135         trackingGeometry.apply(visitor)
0136
0137         assert visitor.num_surfaces == 19264
0138         assert visitor.num_layers == 142
0139         assert visitor.num_volumes == 32
0140         assert visitor.num_portals == 0  # Gen1: will have no portals
0141         assert visitor.num_boundary_surfaces == 126  # Gen1: will have boundary surfaces
0142
0143
0144 @pytest.mark.skipif(not dd4hepEnabled, reason="DD4hep not set up")
0145 @pytest.mark.odd
0146 def test_odd_gen3():
0147     with getOpenDataDetector(gen3=True) as detector:
0148         trackingGeometry = detector.trackingGeometry()
0149
0150         visitor = CountingVisitor()
0151         trackingGeometry.apply(visitor)
0152
0153         assert visitor.num_surfaces == 9
0154         assert visitor.num_layers == 0  # Gen3: no layers
0155         assert visitor.num_volumes == 2
0156         assert visitor.num_portals == 9  # Gen3: will have portals
0157         assert visitor.num_boundary_surfaces == 0  # Gen3: no boundary surfaces