optmapper 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- optmapper/__init__.py +7 -0
- optmapper/_version.py +24 -0
- optmapper/cli.py +150 -0
- optmapper/config.py +342 -0
- optmapper/geometry.py +56 -0
- optmapper/log_utils.py +29 -0
- optmapper/macro.py +136 -0
- optmapper/schedulers.py +108 -0
- optmapper/sites/local.yaml +3 -0
- optmapper/sites/nersc.yaml +23 -0
- optmapper/sites/slurm.yaml +12 -0
- optmapper/templates/optmap.mac +30 -0
- optmapper/work.py +240 -0
- optmapper-0.1.0.dist-info/METADATA +58 -0
- optmapper-0.1.0.dist-info/RECORD +19 -0
- optmapper-0.1.0.dist-info/WHEEL +5 -0
- optmapper-0.1.0.dist-info/entry_points.txt +3 -0
- optmapper-0.1.0.dist-info/licenses/LICENSE +674 -0
- optmapper-0.1.0.dist-info/top_level.txt +1 -0
optmapper/macro.py
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
1
|
+
"""Composition of the remage macro from a template and the production configuration."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import importlib
|
|
6
|
+
import inspect
|
|
7
|
+
import string
|
|
8
|
+
import tempfile
|
|
9
|
+
from collections.abc import Callable, Iterator
|
|
10
|
+
from contextlib import contextmanager
|
|
11
|
+
from importlib import resources
|
|
12
|
+
from pathlib import Path
|
|
13
|
+
|
|
14
|
+
import pint
|
|
15
|
+
|
|
16
|
+
from .config import Config, ConfigError, EmissionConfig, OptmapConfig
|
|
17
|
+
|
|
18
|
+
ureg = pint.get_application_registry()
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _quantity(value: str | float, default_unit: str, what: str) -> pint.Quantity:
|
|
22
|
+
try:
|
|
23
|
+
q = ureg.Quantity(value)
|
|
24
|
+
except (pint.PintError, ValueError) as e:
|
|
25
|
+
msg = f"cannot interpret {what} = {value!r} as a quantity: {e}"
|
|
26
|
+
raise ConfigError(msg) from e
|
|
27
|
+
return ureg.Quantity(q.m, default_unit) if q.dimensionless else q
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def confinement_commands(optmap: OptmapConfig) -> list[str]:
|
|
31
|
+
"""Primary vertex confinement to the optical map volume(s) and bounds."""
|
|
32
|
+
volumes = [optmap.volume] if isinstance(optmap.volume, str) else optmap.volume
|
|
33
|
+
cmds = [
|
|
34
|
+
"/RMG/Generator/Confine Volume",
|
|
35
|
+
"/RMG/Generator/Confinement/ForceContainmentCheck",
|
|
36
|
+
]
|
|
37
|
+
|
|
38
|
+
if optmap.shape != "none":
|
|
39
|
+
cmds += ["/RMG/Generator/Confinement/SamplingMode IntersectPhysicalWithGeometrical"]
|
|
40
|
+
cmds += [f"/RMG/Generator/Confinement/Physical/AddVolume {v}" for v in volumes]
|
|
41
|
+
if optmap.shape == "none":
|
|
42
|
+
return cmds
|
|
43
|
+
|
|
44
|
+
geo = "/RMG/Generator/Confinement/Geometrical"
|
|
45
|
+
center = [(lo + hi) / 2 for lo, hi in optmap.range_in_m]
|
|
46
|
+
size = [hi - lo for lo, hi in optmap.range_in_m]
|
|
47
|
+
cmds += [f"{geo}/AddSolid {optmap.shape.capitalize()}"]
|
|
48
|
+
cmds += [f"{geo}/CenterPosition{ax} {c:g} m" for ax, c in zip("XYZ", center, strict=True)]
|
|
49
|
+
|
|
50
|
+
if optmap.shape == "cylinder":
|
|
51
|
+
cmds += [
|
|
52
|
+
f"{geo}/Cylinder/InnerRadius {optmap.inner_radius_in_m:g} m",
|
|
53
|
+
f"{geo}/Cylinder/OuterRadius {size[0] / 2:g} m",
|
|
54
|
+
f"{geo}/Cylinder/Height {size[2]:g} m",
|
|
55
|
+
]
|
|
56
|
+
else:
|
|
57
|
+
cmds += [f"{geo}/Box/{ax}Length {s:g} m" for ax, s in zip("XYZ", size, strict=True)]
|
|
58
|
+
return cmds
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def load_spectrum_function(spec: str) -> Callable[[str, bool], None]:
|
|
62
|
+
"""Import a spectrum function from its dotted path, e.g. ``pygeomoptics.lar.g4gps_...``."""
|
|
63
|
+
module_name, _, func_name = spec.rpartition(".")
|
|
64
|
+
try:
|
|
65
|
+
func = getattr(importlib.import_module(module_name), func_name)
|
|
66
|
+
except (ImportError, AttributeError, ValueError) as e:
|
|
67
|
+
msg = f"cannot import emission spectrum function '{spec}': {e}"
|
|
68
|
+
raise ConfigError(msg) from e
|
|
69
|
+
if not callable(func) or len(inspect.signature(func).parameters) != 2:
|
|
70
|
+
msg = f"'{spec}' must be a function f(filename, output_macro), like pygeomoptics' g4gps_*"
|
|
71
|
+
raise ConfigError(msg)
|
|
72
|
+
return func
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
@contextmanager
|
|
76
|
+
def _optics_plugin(plugin: str | None) -> Iterator[None]:
|
|
77
|
+
"""Apply a pygeom-optics plugin, restoring the original material properties afterwards."""
|
|
78
|
+
if plugin is None:
|
|
79
|
+
yield
|
|
80
|
+
return
|
|
81
|
+
from pygeomoptics import store # noqa: PLC0415
|
|
82
|
+
|
|
83
|
+
store.load_user_material_code(plugin)
|
|
84
|
+
try:
|
|
85
|
+
yield
|
|
86
|
+
finally:
|
|
87
|
+
store.reset_all_to_original()
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def energy_commands(emission: EmissionConfig, optics_plugin: str | None = None) -> list[str]:
|
|
91
|
+
"""GPS energy distribution of the optical photons.
|
|
92
|
+
|
|
93
|
+
The pygeom-optics `optics_plugin`, if any, is applied while generating the spectrum, so
|
|
94
|
+
that it sees the same material properties as the geometry.
|
|
95
|
+
"""
|
|
96
|
+
if emission.gaussian is not None:
|
|
97
|
+
try:
|
|
98
|
+
with ureg.context("sp"):
|
|
99
|
+
mean = _quantity(emission.gaussian["mean"], "eV", "gaussian mean").to("eV")
|
|
100
|
+
sigma = _quantity(emission.gaussian["sigma"], "eV", "gaussian sigma").to("eV")
|
|
101
|
+
except pint.DimensionalityError as e:
|
|
102
|
+
msg = f"emission.gaussian: {e}"
|
|
103
|
+
raise ConfigError(msg) from e
|
|
104
|
+
return [
|
|
105
|
+
"/gps/ene/type Gauss",
|
|
106
|
+
f"/gps/ene/mono {mean.m:.6g} eV",
|
|
107
|
+
f"/gps/ene/sigma {sigma.m:.6g} eV",
|
|
108
|
+
]
|
|
109
|
+
|
|
110
|
+
assert emission.spectrum is not None
|
|
111
|
+
func = load_spectrum_function(emission.spectrum)
|
|
112
|
+
with _optics_plugin(optics_plugin), tempfile.TemporaryDirectory() as tmpdir:
|
|
113
|
+
fn = Path(tmpdir) / "spectrum.mac"
|
|
114
|
+
func(str(fn), True)
|
|
115
|
+
return fn.read_text().strip().splitlines()
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def render_macro(cfg: Config, n_events: int) -> str:
|
|
119
|
+
"""Return the full remage macro for this production."""
|
|
120
|
+
if cfg.advanced.template is not None:
|
|
121
|
+
template = Path(cfg.advanced.template).read_text()
|
|
122
|
+
else:
|
|
123
|
+
template = (resources.files("optmapper") / "templates" / "optmap.mac").read_text()
|
|
124
|
+
|
|
125
|
+
try:
|
|
126
|
+
return string.Template(template).substitute(
|
|
127
|
+
name=cfg.name,
|
|
128
|
+
pre_init_commands="\n".join(cfg.advanced.pre_init_commands),
|
|
129
|
+
confinement_commands="\n".join(confinement_commands(cfg.optmap)),
|
|
130
|
+
energy_commands="\n".join(energy_commands(cfg.emission, cfg.geometry.optics_plugin)),
|
|
131
|
+
commands="\n".join(cfg.advanced.commands),
|
|
132
|
+
n_events=n_events,
|
|
133
|
+
)
|
|
134
|
+
except (KeyError, ValueError) as e:
|
|
135
|
+
msg = f"invalid placeholder in macro template: {e}"
|
|
136
|
+
raise ConfigError(msg) from e
|
optmapper/schedulers.py
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""Job submission backends."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import logging
|
|
6
|
+
import shlex
|
|
7
|
+
import subprocess
|
|
8
|
+
from abc import ABC, abstractmethod
|
|
9
|
+
from collections.abc import Mapping, Sequence
|
|
10
|
+
from dataclasses import dataclass, field
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
|
|
13
|
+
log = logging.getLogger(__name__)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class Job:
|
|
18
|
+
name: str
|
|
19
|
+
command: list[str]
|
|
20
|
+
log_file: Path
|
|
21
|
+
"""Log file of the job, ``%j`` is replaced by the job ID."""
|
|
22
|
+
options: Mapping = field(default_factory=dict)
|
|
23
|
+
"""Backend specific options, e.g. ``sbatch`` options for Slurm."""
|
|
24
|
+
dependencies: Sequence[str] = ()
|
|
25
|
+
"""IDs of jobs that must have terminated (successfully or not) before this one starts."""
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class Scheduler(ABC):
|
|
29
|
+
def __init__(self, dry_run: bool = False):
|
|
30
|
+
self.dry_run = dry_run
|
|
31
|
+
|
|
32
|
+
@abstractmethod
|
|
33
|
+
def submit(self, job: Job) -> str:
|
|
34
|
+
"""Submit `job` and return its ID."""
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class SlurmScheduler(Scheduler):
|
|
38
|
+
"""Submit jobs with ``sbatch --wrap``, no job script files are written."""
|
|
39
|
+
|
|
40
|
+
@staticmethod
|
|
41
|
+
def sbatch_command(job: Job) -> list[str]:
|
|
42
|
+
cmd = [
|
|
43
|
+
"sbatch",
|
|
44
|
+
"--parsable",
|
|
45
|
+
f"--job-name={job.name}",
|
|
46
|
+
f"--output={job.log_file}",
|
|
47
|
+
f"--error={job.log_file}",
|
|
48
|
+
]
|
|
49
|
+
for k, v in job.options.items():
|
|
50
|
+
if v is True:
|
|
51
|
+
cmd.append(f"--{k}")
|
|
52
|
+
elif v is not None and v is not False:
|
|
53
|
+
cmd.append(f"--{k}={v}")
|
|
54
|
+
if job.dependencies:
|
|
55
|
+
cmd.append(f"--dependency=afterany:{':'.join(job.dependencies)}")
|
|
56
|
+
cmd += ["--wrap", shlex.join(job.command)]
|
|
57
|
+
return cmd
|
|
58
|
+
|
|
59
|
+
def submit(self, job: Job) -> str:
|
|
60
|
+
cmd = self.sbatch_command(job)
|
|
61
|
+
if self.dry_run:
|
|
62
|
+
log.info("would run: %s", shlex.join(cmd))
|
|
63
|
+
return f"<{job.name}>"
|
|
64
|
+
|
|
65
|
+
log.debug("running: %s", shlex.join(cmd))
|
|
66
|
+
try:
|
|
67
|
+
res = subprocess.run(cmd, check=True, capture_output=True, text=True)
|
|
68
|
+
except FileNotFoundError as e:
|
|
69
|
+
msg = "sbatch not found, is this a Slurm cluster?"
|
|
70
|
+
raise RuntimeError(msg) from e
|
|
71
|
+
except subprocess.CalledProcessError as e:
|
|
72
|
+
msg = f"sbatch failed for job {job.name}: {e.stderr.strip()}"
|
|
73
|
+
raise RuntimeError(msg) from e
|
|
74
|
+
# --parsable prints "jobid[;cluster]"
|
|
75
|
+
return res.stdout.strip().split(";")[0]
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class LocalScheduler(Scheduler):
|
|
79
|
+
"""Run jobs one after the other in the foreground, on the current machine.
|
|
80
|
+
|
|
81
|
+
Jobs are executed at submission time, so dependencies are always satisfied.
|
|
82
|
+
Like for ``afterany`` in Slurm, a job runs even if its dependencies failed.
|
|
83
|
+
"""
|
|
84
|
+
|
|
85
|
+
def __init__(self, dry_run: bool = False):
|
|
86
|
+
super().__init__(dry_run)
|
|
87
|
+
self.failed: list[str] = []
|
|
88
|
+
self._count = 0
|
|
89
|
+
|
|
90
|
+
def submit(self, job: Job) -> str:
|
|
91
|
+
self._count += 1
|
|
92
|
+
job_id = str(self._count)
|
|
93
|
+
log_file = Path(str(job.log_file).replace("%j", job_id))
|
|
94
|
+
if self.dry_run:
|
|
95
|
+
log.info("would run: %s > %s", shlex.join(job.command), log_file)
|
|
96
|
+
return job_id
|
|
97
|
+
|
|
98
|
+
log.info("running job %s, log in %s", job.name, log_file)
|
|
99
|
+
with log_file.open("w") as f:
|
|
100
|
+
ret = subprocess.run(job.command, stdout=f, stderr=subprocess.STDOUT, check=False)
|
|
101
|
+
if ret.returncode != 0:
|
|
102
|
+
log.error("job %s failed with exit code %d", job.name, ret.returncode)
|
|
103
|
+
self.failed.append(job.name)
|
|
104
|
+
return job_id
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def get_scheduler(name: str, dry_run: bool = False) -> Scheduler:
|
|
108
|
+
return {"slurm": SlurmScheduler, "local": LocalScheduler}[name](dry_run=dry_run)
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
# NERSC Perlmutter CPU nodes, LEGEND allocation.
|
|
2
|
+
scheduler: slurm
|
|
3
|
+
scratch_dir: $PSCRATCH/optmapper
|
|
4
|
+
slurm:
|
|
5
|
+
node:
|
|
6
|
+
account: m2676
|
|
7
|
+
qos: regular
|
|
8
|
+
constraint: cpu
|
|
9
|
+
licenses: cfs,scratch
|
|
10
|
+
nodes: 1
|
|
11
|
+
ntasks: 1
|
|
12
|
+
ntasks-per-node: 1
|
|
13
|
+
cpus-per-task: 256
|
|
14
|
+
time: "04:00:00"
|
|
15
|
+
# merging maps needs a few cores only: use the shared QOS instead of a full node
|
|
16
|
+
merge:
|
|
17
|
+
account: m2676
|
|
18
|
+
qos: shared
|
|
19
|
+
constraint: cpu
|
|
20
|
+
licenses: cfs,scratch
|
|
21
|
+
ntasks: 1
|
|
22
|
+
cpus-per-task: 16
|
|
23
|
+
time: "01:00:00"
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
# generic Slurm cluster. Start from this preset when porting to a new facility,
|
|
2
|
+
# options under slurm.node and slurm.merge are passed verbatim to sbatch.
|
|
3
|
+
scheduler: slurm
|
|
4
|
+
slurm:
|
|
5
|
+
node:
|
|
6
|
+
nodes: 1
|
|
7
|
+
ntasks: 1
|
|
8
|
+
exclusive: true
|
|
9
|
+
merge:
|
|
10
|
+
nodes: 1
|
|
11
|
+
ntasks: 1
|
|
12
|
+
cpus-per-task: 8
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
# remage macro for optical map production, generated by optmapper
|
|
2
|
+
# production: ${name}
|
|
3
|
+
|
|
4
|
+
/RMG/Processes/OpticalPhysics true
|
|
5
|
+
/RMG/Processes/OpticalPhysicsMaxOneWLSPhoton true
|
|
6
|
+
|
|
7
|
+
/RMG/Geometry/RegisterDetectorsFromGDML Optical
|
|
8
|
+
|
|
9
|
+
/RMG/Manager/PrintProgressModulo 100000
|
|
10
|
+
/RMG/Geometry/GDMLDisableOverlapCheck true
|
|
11
|
+
|
|
12
|
+
${pre_init_commands}
|
|
13
|
+
|
|
14
|
+
/run/initialize
|
|
15
|
+
|
|
16
|
+
/RMG/Output/NtuplePerDetector false
|
|
17
|
+
/RMG/Output/Vertex/StorePrimaryParticleInformation false
|
|
18
|
+
/RMG/Output/Vertex/StoreSinglePrecisionEnergy
|
|
19
|
+
/RMG/Output/Vertex/StoreSinglePrecisionPosition
|
|
20
|
+
|
|
21
|
+
${confinement_commands}
|
|
22
|
+
|
|
23
|
+
/RMG/Generator/Select GPS
|
|
24
|
+
/gps/particle opticalphoton
|
|
25
|
+
${energy_commands}
|
|
26
|
+
/gps/ang/type iso
|
|
27
|
+
|
|
28
|
+
${commands}
|
|
29
|
+
|
|
30
|
+
/run/beamOn ${n_events}
|
optmapper/work.py
ADDED
|
@@ -0,0 +1,240 @@
|
|
|
1
|
+
"""Worker process, runs on the compute nodes."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import logging
|
|
7
|
+
import os
|
|
8
|
+
import re
|
|
9
|
+
import resource
|
|
10
|
+
import shlex
|
|
11
|
+
import shutil
|
|
12
|
+
import subprocess
|
|
13
|
+
import time
|
|
14
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
15
|
+
from contextlib import contextmanager
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
|
|
18
|
+
import dbetto
|
|
19
|
+
|
|
20
|
+
from .config import Config
|
|
21
|
+
from .log_utils import setup_log
|
|
22
|
+
|
|
23
|
+
log = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
@contextmanager
|
|
27
|
+
def _step(what: str):
|
|
28
|
+
log.info("%s...", what)
|
|
29
|
+
start = time.monotonic()
|
|
30
|
+
yield
|
|
31
|
+
log.info("%s: done in %.0f s", what, time.monotonic() - start)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _natural_key(path: Path) -> list:
|
|
35
|
+
return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", path.name)]
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _available_cpus() -> int:
|
|
39
|
+
if "SLURM_CPUS_ON_NODE" in os.environ:
|
|
40
|
+
return int(os.environ["SLURM_CPUS_ON_NODE"])
|
|
41
|
+
return len(os.sched_getaffinity(0))
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def merge_maps(cfg: Config, inputs: list[Path], output: Path, workdir: Path) -> None:
|
|
45
|
+
"""Merge optical maps into `output`, temporary files are written to `workdir`."""
|
|
46
|
+
from reboost.optmap.create import merge_optical_maps # noqa: PLC0415
|
|
47
|
+
|
|
48
|
+
output.parent.mkdir(parents=True, exist_ok=True)
|
|
49
|
+
tmp = workdir / f".{output.name}.part"
|
|
50
|
+
tmp.unlink(missing_ok=True)
|
|
51
|
+
if len(inputs) == 1:
|
|
52
|
+
shutil.copyfile(inputs[0], tmp)
|
|
53
|
+
else:
|
|
54
|
+
merge_optical_maps(
|
|
55
|
+
[str(p) for p in inputs],
|
|
56
|
+
str(tmp),
|
|
57
|
+
cfg.optmap.reboost_settings(),
|
|
58
|
+
n_procs=cfg.processing.merge_procs,
|
|
59
|
+
)
|
|
60
|
+
shutil.move(tmp, output)
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def _reboost_create_args(cfg: Config, workdir: Path) -> list[str]:
|
|
64
|
+
"""Options of ``reboost-optmap create``, its input files are written to `workdir`."""
|
|
65
|
+
settings = workdir / "optmap-settings.json"
|
|
66
|
+
dbetto.utils.write_dict(cfg.optmap.reboost_settings(), str(settings))
|
|
67
|
+
args = [
|
|
68
|
+
*("--bufsize", str(cfg.processing.bufsize), "create"),
|
|
69
|
+
*("--settings", str(settings), "--geom", str(cfg.gdml_file)),
|
|
70
|
+
*("--n-procs", str(cfg.processing.procs_per_map)),
|
|
71
|
+
]
|
|
72
|
+
if cfg.optmap.detectors is not None:
|
|
73
|
+
detectors = workdir / "optmap-detectors.json"
|
|
74
|
+
dbetto.utils.write_dict(cfg.optmap.detectors, str(detectors))
|
|
75
|
+
args += ["--detectors", str(detectors)]
|
|
76
|
+
return args
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def _create_map(create_args: list[str], inputs: list[Path], output: Path) -> None:
|
|
80
|
+
"""Create a map with ``reboost-optmap``, in a separate process as it spawns its own pool."""
|
|
81
|
+
cmd = ["reboost-optmap", *create_args, "--", *map(str, inputs), str(output)]
|
|
82
|
+
res = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
|
83
|
+
if res.returncode != 0:
|
|
84
|
+
msg = f"map creation failed for {output.name}:\n{res.stdout}{res.stderr}"
|
|
85
|
+
raise RuntimeError(msg)
|
|
86
|
+
log.info("created %s from %d stp files", output.name, len(inputs))
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def run_node(cfg: Config, index: int) -> None:
|
|
90
|
+
"""Simulate photons on this node and build the node optical map."""
|
|
91
|
+
from .macro import render_macro # noqa: PLC0415
|
|
92
|
+
|
|
93
|
+
job_id = os.environ.get("SLURM_JOB_ID", str(os.getpid()))
|
|
94
|
+
scratch_base = os.path.expandvars(
|
|
95
|
+
cfg.execution.scratch_dir or str(Path(cfg.output_dir) / "scratch")
|
|
96
|
+
)
|
|
97
|
+
scratch = Path(scratch_base) / cfg.name / f"node{index:04d}-{job_id}"
|
|
98
|
+
# a requeued Slurm job keeps its ID, do not pick up files from the previous attempt
|
|
99
|
+
shutil.rmtree(scratch, ignore_errors=True)
|
|
100
|
+
scratch.mkdir(parents=True)
|
|
101
|
+
os.environ.setdefault("NUMBA_CACHE_DIR", str(scratch / ".numba-cache"))
|
|
102
|
+
|
|
103
|
+
nprocs = cfg.statistics.processes_per_node or _available_cpus()
|
|
104
|
+
n_primaries = cfg.statistics.primaries_per_process
|
|
105
|
+
log.info(
|
|
106
|
+
"node %d (job %s): %d run(s) x %d remage processes x %.3g primaries, scratch in %s",
|
|
107
|
+
index,
|
|
108
|
+
job_id,
|
|
109
|
+
cfg.statistics.runs_per_node,
|
|
110
|
+
nprocs,
|
|
111
|
+
n_primaries,
|
|
112
|
+
scratch,
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
macro = scratch / f"{cfg.name}.mac"
|
|
116
|
+
macro.write_text(render_macro(cfg, n_primaries))
|
|
117
|
+
if index == 0:
|
|
118
|
+
shutil.copyfile(macro, Path(cfg.output_dir) / macro.name)
|
|
119
|
+
|
|
120
|
+
proc = cfg.processing
|
|
121
|
+
map_dir = scratch / "maps"
|
|
122
|
+
map_dir.mkdir()
|
|
123
|
+
create_args = _reboost_create_args(cfg, scratch)
|
|
124
|
+
|
|
125
|
+
for run in range(cfg.statistics.runs_per_node):
|
|
126
|
+
stp_dir = scratch / "stp" / f"run{run:03d}"
|
|
127
|
+
stp_dir.mkdir(parents=True)
|
|
128
|
+
|
|
129
|
+
cmd = [
|
|
130
|
+
"remage",
|
|
131
|
+
*cfg.advanced.remage_args,
|
|
132
|
+
"--gdml-files",
|
|
133
|
+
str(cfg.gdml_file),
|
|
134
|
+
"--output-file",
|
|
135
|
+
str(stp_dir / f"{cfg.name}-node{index:04d}-run{run:03d}-tier_stp.lh5"),
|
|
136
|
+
"--procs",
|
|
137
|
+
str(nprocs),
|
|
138
|
+
"--flat-output",
|
|
139
|
+
"--",
|
|
140
|
+
str(macro),
|
|
141
|
+
]
|
|
142
|
+
with _step(f"run {run}: remage simulation"):
|
|
143
|
+
log.debug("running: %s", shlex.join(cmd))
|
|
144
|
+
subprocess.run(cmd, check=True)
|
|
145
|
+
|
|
146
|
+
stp_files = sorted(stp_dir.glob("*.lh5"), key=_natural_key)
|
|
147
|
+
if not stp_files:
|
|
148
|
+
msg = f"remage did not produce any output file in {stp_dir}"
|
|
149
|
+
raise RuntimeError(msg)
|
|
150
|
+
|
|
151
|
+
chunks = [
|
|
152
|
+
stp_files[i : i + proc.stp_files_per_map]
|
|
153
|
+
for i in range(0, len(stp_files), proc.stp_files_per_map)
|
|
154
|
+
]
|
|
155
|
+
with (
|
|
156
|
+
_step(f"run {run}: creating {len(chunks)} map(s) from {len(stp_files)} stp files"),
|
|
157
|
+
ThreadPoolExecutor(proc.parallel_maps) as pool,
|
|
158
|
+
):
|
|
159
|
+
futures = [
|
|
160
|
+
pool.submit(_create_map, create_args, chunk, map_dir / f"run{run:03d}-{k:04d}.lh5")
|
|
161
|
+
for k, chunk in enumerate(chunks)
|
|
162
|
+
]
|
|
163
|
+
for f in futures:
|
|
164
|
+
f.result()
|
|
165
|
+
|
|
166
|
+
if proc.keep_stp:
|
|
167
|
+
dest = Path(cfg.output_dir) / "stp" / f"node{index:04d}"
|
|
168
|
+
with _step(f"run {run}: copying stp files to {dest}"):
|
|
169
|
+
dest.mkdir(parents=True, exist_ok=True)
|
|
170
|
+
for stp_file in stp_files:
|
|
171
|
+
shutil.copy2(stp_file, dest)
|
|
172
|
+
shutil.rmtree(stp_dir)
|
|
173
|
+
|
|
174
|
+
output = cfg.node_map(index)
|
|
175
|
+
with _step(f"merging intermediate maps into {output}"):
|
|
176
|
+
merge_maps(cfg, sorted(map_dir.glob("*.lh5"), key=_natural_key), output, scratch)
|
|
177
|
+
|
|
178
|
+
shutil.rmtree(scratch)
|
|
179
|
+
log.info(
|
|
180
|
+
"node %d done, max RSS of child processes: %.1f GiB",
|
|
181
|
+
index,
|
|
182
|
+
resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss / 1024**2,
|
|
183
|
+
)
|
|
184
|
+
|
|
185
|
+
|
|
186
|
+
def run_merge(cfg: Config) -> None:
|
|
187
|
+
"""Merge the node maps into the final map, tolerating failed nodes."""
|
|
188
|
+
nodes = [cfg.node_map(i) for i in range(cfg.statistics.nodes)]
|
|
189
|
+
found = [p for p in nodes if p.is_file()]
|
|
190
|
+
missing = [p for p in nodes if not p.is_file()]
|
|
191
|
+
if not found:
|
|
192
|
+
msg = f"all node jobs failed, no node map to merge: see the node job logs in {cfg.log_dir}"
|
|
193
|
+
raise RuntimeError(msg)
|
|
194
|
+
if missing:
|
|
195
|
+
log.warning(
|
|
196
|
+
"%d of %d node maps are missing (failed jobs?), see the logs in %s: %s",
|
|
197
|
+
len(missing),
|
|
198
|
+
len(nodes),
|
|
199
|
+
cfg.log_dir,
|
|
200
|
+
", ".join(p.name for p in missing),
|
|
201
|
+
)
|
|
202
|
+
|
|
203
|
+
with _step(f"merging {len(found)} node maps into {cfg.final_map}"):
|
|
204
|
+
merge_maps(cfg, found, cfg.final_map, cfg.final_map.parent)
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
def work_cli(argv: list[str] | None = None) -> None:
|
|
208
|
+
parser = argparse.ArgumentParser(
|
|
209
|
+
prog="optmapper-work",
|
|
210
|
+
description="%(prog)s: optical map production worker, usually started by optmapper",
|
|
211
|
+
)
|
|
212
|
+
parser.add_argument("--verbose", "-v", action="count", default=0, help="increase verbosity")
|
|
213
|
+
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
214
|
+
|
|
215
|
+
def _add_config(p):
|
|
216
|
+
p.add_argument(
|
|
217
|
+
"--config", required=True, help="resolved production configuration file (JSON)"
|
|
218
|
+
)
|
|
219
|
+
|
|
220
|
+
node_parser = subparsers.add_parser("node", help="run the full production on one node")
|
|
221
|
+
_add_config(node_parser)
|
|
222
|
+
node_parser.add_argument("--index", type=int, required=True, help="node index")
|
|
223
|
+
|
|
224
|
+
merge_parser = subparsers.add_parser("merge", help="merge the node maps into the final map")
|
|
225
|
+
_add_config(merge_parser)
|
|
226
|
+
|
|
227
|
+
args = parser.parse_args(argv)
|
|
228
|
+
setup_log((None, logging.INFO, logging.DEBUG)[min(args.verbose, 2)])
|
|
229
|
+
|
|
230
|
+
cfg = Config.load(args.config)
|
|
231
|
+
os.environ.update(cfg.execution.env)
|
|
232
|
+
|
|
233
|
+
if args.command == "node":
|
|
234
|
+
run_node(cfg, args.index)
|
|
235
|
+
elif args.command == "merge":
|
|
236
|
+
run_merge(cfg)
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
if __name__ == "__main__":
|
|
240
|
+
work_cli()
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: optmapper
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Production of LEGEND optical maps on batch systems
|
|
5
|
+
Author-email: Luigi Pertoldi <gipert@pm.me>
|
|
6
|
+
Maintainer: The LEGEND Collaboration
|
|
7
|
+
License-Expression: GPL-3.0
|
|
8
|
+
Project-URL: Homepage, https://github.com/legend-exp/optmapper
|
|
9
|
+
Project-URL: Bug Tracker, https://github.com/legend-exp/optmapper/issues
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: Operating System :: POSIX
|
|
12
|
+
Classifier: Operating System :: Unix
|
|
13
|
+
Classifier: Programming Language :: Python
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering
|
|
17
|
+
Requires-Python: >=3.11
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
License-File: LICENSE
|
|
20
|
+
Requires-Dist: colorlog
|
|
21
|
+
Requires-Dist: dbetto
|
|
22
|
+
Requires-Dist: legend-pygeom-optics>=0.15.0
|
|
23
|
+
Requires-Dist: pint
|
|
24
|
+
Requires-Dist: reboost>=1.4
|
|
25
|
+
Provides-Extra: all
|
|
26
|
+
Requires-Dist: optmapper[docs,test]; extra == "all"
|
|
27
|
+
Provides-Extra: docs
|
|
28
|
+
Requires-Dist: furo; extra == "docs"
|
|
29
|
+
Requires-Dist: myst-parser; extra == "docs"
|
|
30
|
+
Requires-Dist: sphinx; extra == "docs"
|
|
31
|
+
Requires-Dist: sphinx-copybutton; extra == "docs"
|
|
32
|
+
Provides-Extra: test
|
|
33
|
+
Requires-Dist: pre-commit; extra == "test"
|
|
34
|
+
Requires-Dist: pytest>=6.0; extra == "test"
|
|
35
|
+
Requires-Dist: pytest-cov; extra == "test"
|
|
36
|
+
Dynamic: license-file
|
|
37
|
+
|
|
38
|
+
# optmapper
|
|
39
|
+
|
|
40
|
+

|
|
41
|
+
[](https://github.com/legend-exp/optmapper/actions)
|
|
42
|
+
[](https://github.com/pre-commit/pre-commit)
|
|
43
|
+
[](https://app.codecov.io/gh/legend-exp/optmapper)
|
|
44
|
+

|
|
45
|
+

|
|
46
|
+

|
|
47
|
+
[](https://optmapper.readthedocs.io)
|
|
48
|
+
|
|
49
|
+
_optmapper_ produces LEGEND optical maps on batch systems. An optical map gives,
|
|
50
|
+
for each position in a volume (e.g. the liquid argon), the probability that an
|
|
51
|
+
optical photon emitted there is detected by each optical detector. Building one
|
|
52
|
+
requires simulating billions of photons with
|
|
53
|
+
[remage](https://remage.readthedocs.io/en/stable/) and histogramming the result
|
|
54
|
+
with [reboost](https://reboost.readthedocs.io/en/stable/). _optmapper_ takes a
|
|
55
|
+
single configuration file and runs the whole chain on one or many compute nodes.
|
|
56
|
+
|
|
57
|
+
For more information see our dedicated
|
|
58
|
+
[documentation](https://optmapper.readthedocs.io/en/stable/)!
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
optmapper/__init__.py,sha256=a14TEcUQN9VjmFAsSR9oU6SgZbolz34DNqxA96m1924,167
|
|
2
|
+
optmapper/_version.py,sha256=n_5vdJsPNu7wZ57LGuRL585uvll-hiuvZUBWzdG0RQU,520
|
|
3
|
+
optmapper/cli.py,sha256=jWpd74k-BhIU1PWRQBCo4-gZzQgCrg3mA0qLFXIp_ow,5147
|
|
4
|
+
optmapper/config.py,sha256=xCbslmxxZLqO_6rJmqGwym-IRDIPVA8O-rdwQSOjJ44,13208
|
|
5
|
+
optmapper/geometry.py,sha256=7Izc1WedlnR6q2VfkZ1fJJF6gY4MWhgD24w5nPMXOTk,1833
|
|
6
|
+
optmapper/log_utils.py,sha256=Cd5vFu3STrBzvQmkgpICLpnjteGO00BD--2UDITqQzE,768
|
|
7
|
+
optmapper/macro.py,sha256=ROLRiVLk3dKDlKyB8kSyi9tE8dQlwyNzf3IhMe210g0,5257
|
|
8
|
+
optmapper/schedulers.py,sha256=2iDD54A_gYYQj-YFwIX7KEVDBtC7YeTPHIViZgqQemk,3643
|
|
9
|
+
optmapper/work.py,sha256=PMGI6mXobjNmV__EGveA_FE7AiuZdW9sNEJHLsJtKMY,8417
|
|
10
|
+
optmapper/sites/local.yaml,sha256=-E6i1vzSOJYwOWlrZNvMyqALvJ56RlVS1xQLPg459r4,160
|
|
11
|
+
optmapper/sites/nersc.yaml,sha256=-HBeTnACwZJcc-LLnktVEnzqIuj9ih0vFV9bk69j0UU,520
|
|
12
|
+
optmapper/sites/slurm.yaml,sha256=6G2SR52GhvP3ZwhFvtAU1SSi1ZrYQtEok5QOV1_jHzI,290
|
|
13
|
+
optmapper/templates/optmap.mac,sha256=hFYKvjtnZgUpqgk-6hsFsWdhYZL3c7k3Lo2QUEdLYjU,697
|
|
14
|
+
optmapper-0.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
|
|
15
|
+
optmapper-0.1.0.dist-info/METADATA,sha256=j1Uab-Bnv5fy185ynKUranJ5NSEo2RJhX9vK_qM1vow,2915
|
|
16
|
+
optmapper-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
|
|
17
|
+
optmapper-0.1.0.dist-info/entry_points.txt,sha256=DEM7sd9pnWOrNAhDHSFxeoBC9OfKhZtHnscRvcy4oeg,99
|
|
18
|
+
optmapper-0.1.0.dist-info/top_level.txt,sha256=paE_szQAS0XPUNrxr6QJCcIb7JKN96ve-dX-Uls3Ers,10
|
|
19
|
+
optmapper-0.1.0.dist-info/RECORD,,
|