File indexing completed on 2026-08-16 08:19:18
0001 import os
0002 from pathlib import Path
0003 from typing import Optional
0004
0005 import acts
0006 import acts.examples
0007
0008
0009
0010 COLLIDERML_DATASET_ID = "CERN/ColliderML-Release-1"
0011
0012
0013 def getColliderMLDirectory() -> Path:
0014 """Return the local ColliderML cache directory.
0015
0016 Uses `$COLLIDERML_DATA_DIR` if set, otherwise falls back to
0017 `~/.cache/colliderml`. This matches the canonical `colliderml` Python
0018 library (`colliderml.core.hf_download.default_data_dir`), so a cache
0019 populated with `colliderml download ...` is found without any further
0020 configuration.
0021
0022 Raises `RuntimeError` if the resulting directory does not exist.
0023 """
0024 env = os.environ.get("COLLIDERML_DATA_DIR")
0025 if env:
0026 path = Path(env).expanduser().resolve()
0027 else:
0028 path = (Path.home() / ".cache" / "colliderml").resolve()
0029
0030 if not path.is_dir():
0031 raise RuntimeError(
0032 f"ColliderML data directory not found at {path}. "
0033 f"Set $COLLIDERML_DATA_DIR to point at an existing ColliderML "
0034 f"cache, or populate the default location with "
0035 f"'colliderml download ...'."
0036 )
0037 return path
0038
0039
0040 def _sanitizedDatasetId(datasetId: str) -> str:
0041 return datasetId.replace("/", "__")
0042
0043
0044 def getColliderMLObjectDirectory(
0045 object: str,
0046 channel: str = "ttbar",
0047 pileup: str = "pu0",
0048 dataDir: Optional[Path] = None,
0049 datasetId: str = COLLIDERML_DATASET_ID,
0050 ) -> Path:
0051 """Resolve the local directory holding one ColliderML "object" (e.g.
0052 `particles`, `tracker_hits`, `tracks`) for a given channel/pileup.
0053
0054 Parameters
0055 ----------
0056 object : str
0057 Object name, e.g. "particles", "tracker_hits", "tracks".
0058 channel : str
0059 Physics channel/process, e.g. "ttbar".
0060 pileup : str
0061 Pileup token, e.g. "pu0" or "pu200".
0062 dataDir : Path, None
0063 Base cache directory. If not given, uses `getColliderMLDirectory()`.
0064 datasetId : str
0065 HuggingFace dataset id the config was downloaded from.
0066
0067 A directory downloaded via the `colliderml` library's `download_config`
0068 lands at `<dataDir>/<datasetId sanitized>/<config>/data/<config>`. Some
0069 hand-prepared samples (e.g. the ACTS CI sample) omit the dataset-id
0070 level, so both layouts are tried.
0071 """
0072 if dataDir is None:
0073 dataDir = getColliderMLDirectory()
0074 else:
0075 dataDir = Path(dataDir)
0076
0077 config = f"{channel}_{pileup}_{object}"
0078 candidates = [
0079 dataDir / _sanitizedDatasetId(datasetId) / config / "data" / config,
0080 dataDir / config / "data" / config,
0081 ]
0082 for candidate in candidates:
0083 if candidate.is_dir():
0084 return candidate
0085
0086 searched = "\n".join(f" {c}" for c in candidates)
0087 raise RuntimeError(
0088 f"ColliderML config '{config}' not found. Searched:\n{searched}\n"
0089 f"Download it with:\n"
0090 f" colliderml download --channels {channel} --pileup {pileup} --objects {object}"
0091 )
0092
0093
0094 def addColliderML(
0095 s: acts.examples.Sequencer,
0096 trackingGeometry: acts.TrackingGeometry,
0097 channel: str = "ttbar",
0098 pileup: str = "pu0",
0099 dataDir: Optional[Path] = None,
0100 particlesDir: Optional[Path] = None,
0101 hitsDir: Optional[Path] = None,
0102 tracksDir: Optional[Path] = None,
0103 readTracks: bool = True,
0104 geoIdMapPath: Optional[Path] = None,
0105 geoIdMapSourcePrefix: str = "gen1",
0106 geoIdMapTargetPrefix: str = "gen3",
0107 outputParticles: str = "particles",
0108 outputSimHits: str = "simhits",
0109 outputMeasurements: str = "measurements",
0110 outputMeasurementSubset: str = "measurement_subset",
0111 outputMeasSimHitsMap: str = "measurement_simhits_map",
0112 outputMeasParticlesMap: str = "measurement_particles_map",
0113 outputParticleMeasurementsMap: str = "particle_measurements_map",
0114 outputTracks: str = "colliderml_tracks",
0115 logLevel: Optional[acts.logging.Level] = None,
0116 ) -> acts.examples.Sequencer:
0117 """Read a ColliderML sample and convert it to ACTS EDM collections.
0118
0119 Parameters
0120 ----------
0121 s : Sequencer
0122 the sequencer to add the ColliderML reading steps to
0123 trackingGeometry : TrackingGeometry
0124 tracking geometry to build measurements against
0125 channel, pileup : str
0126 selects the sample, e.g. channel="ttbar", pileup="pu0". Only used to
0127 resolve `particlesDir`/`hitsDir`/`tracksDir` via
0128 `getColliderMLObjectDirectory` when those are not given explicitly.
0129 dataDir : Path, None
0130 base ColliderML cache directory, forwarded to
0131 `getColliderMLObjectDirectory`. None triggers `getColliderMLDirectory()`.
0132 particlesDir, hitsDir, tracksDir : Path, None
0133 raw overrides for the resolved directories, bypassing channel/pileup/
0134 dataDir resolution entirely.
0135 readTracks : bool
0136 if True, also read and convert the dataset's own published tracks
0137 (output under `outputTracks`, default "colliderml_tracks"). Not all
0138 samples ship these.
0139 geoIdMapPath : Path, None
0140 CSV mapping detector geometry ids to `trackingGeometry` geometry ids,
0141 e.g. produced by `generate_geoid_map.py`. None uses `trackingGeometry`
0142 directly.
0143 """
0144 from acts.examples.arrow import (
0145 ColliderMLRelease1InputConverter,
0146 ParquetReader,
0147 )
0148
0149 customLogLevel = acts.examples.defaultLogging(s, logLevel)
0150
0151 if particlesDir is None:
0152 particlesDir = getColliderMLObjectDirectory(
0153 "particles", channel=channel, pileup=pileup, dataDir=dataDir
0154 )
0155 if hitsDir is None:
0156 hitsDir = getColliderMLObjectDirectory(
0157 "tracker_hits", channel=channel, pileup=pileup, dataDir=dataDir
0158 )
0159 if readTracks and tracksDir is None:
0160 tracksDir = getColliderMLObjectDirectory(
0161 "tracks", channel=channel, pileup=pileup, dataDir=dataDir
0162 )
0163
0164 readerCollections = {
0165 "cml_particles": str(particlesDir),
0166 "cml_hits": str(hitsDir),
0167 }
0168 readerSchemas = {
0169 "cml_particles": ColliderMLRelease1InputConverter.particleSchema(),
0170 "cml_hits": ColliderMLRelease1InputConverter.hitSchema(),
0171 }
0172 if readTracks:
0173 readerCollections["cml_tracks"] = str(tracksDir)
0174 readerSchemas["cml_tracks"] = ColliderMLRelease1InputConverter.tracksSchema()
0175
0176 s.addReader(
0177 ParquetReader(
0178 level=customLogLevel(),
0179 collections=readerCollections,
0180 expectedSchemas=readerSchemas,
0181 )
0182 )
0183
0184 converterKWArgs = dict(
0185 inputParticlesTable="cml_particles",
0186 inputHitsTable="cml_hits",
0187 outputParticles=outputParticles,
0188 outputSimHits=outputSimHits,
0189 outputMeasurements=outputMeasurements,
0190 outputMeasurementSubset=outputMeasurementSubset,
0191 outputMeasSimHitsMap=outputMeasSimHitsMap,
0192 outputMeasParticlesMap=outputMeasParticlesMap,
0193 outputParticleMeasurementsMap=outputParticleMeasurementsMap,
0194 trackingGeometry=trackingGeometry,
0195 )
0196 if geoIdMapPath is not None:
0197 converterKWArgs["geoIdMapPath"] = geoIdMapPath
0198 converterKWArgs["geoIdMapSourcePrefix"] = geoIdMapSourcePrefix
0199 converterKWArgs["geoIdMapTargetPrefix"] = geoIdMapTargetPrefix
0200 if readTracks:
0201 converterKWArgs["inputTracksTable"] = "cml_tracks"
0202 converterKWArgs["outputTracks"] = outputTracks
0203
0204 s.addAlgorithm(
0205 ColliderMLRelease1InputConverter(level=customLogLevel(), **converterKWArgs)
0206 )
0207
0208 return s
