geant4-python-application 0.1.0__cp312-cp312-macosx_11_0_arm64.whl

geant4_python_application/cad.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+"""CAD -> GDML helpers via tessellated solids (no external CAD kernel).
+
+Provides:
+- mesh_to_gdml(vertices, faces): build GDML with a <tessellated> solid.
+- stl_ascii_to_gdml(stl_text): parse ASCII STL -> GDML (triangles only).
+Both produce a self-contained GDML string with world box + placed volume,
+suitable for Application(gdml=...).
+"""
+
+
+def mesh_to_gdml(
+    vertices: list[tuple[float, float, float]],
+    faces: list[tuple[int, int, int]],
+    *,
+    name: str = "cad_solid",
+    material: str = "G4_WATER",
+    world_size_mm: float = 1000.0,
+    unit: str = "mm",
+) -> str:
+    if len(vertices) == 0 or len(faces) == 0:
+        msg = "vertices and faces must be non-empty"
+        raise ValueError(msg)
+    pos_lines = []
+    for i, (x, y, z) in enumerate(vertices):
+        pos_lines.append(
+            f'  <position name="v{i}" unit="{unit}" x="{x}" y="{y}" z="{z}" />'
+        )
+    tri_lines = []
+    for a, b, c in faces:
+        if min(a, b, c) < 0 or max(a, b, c) >= len(vertices):
+            msg = f"face index out of range: {(a, b, c)}"
+            raise ValueError(msg)
+        tri_lines.append(
+            f'   <triangular vertex1="v{a}" vertex2="v{b}" vertex3="v{c}" />'
+        )
+    positions = "\n".join(pos_lines)
+    triangles = "\n".join(tri_lines)
+    return f"""<?xml version="1.0" encoding="UTF-8" ?>
+<gdml xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="http://service-spi.web.cern.ch/service-spi/app/releases/GDML/schema/gdml.xsd">
+ <define>
+{positions}
+  <position name="cadpos" unit="{unit}" x="0" y="0" z="0" />
+ </define>
+ <materials>
+  <element Z="1" formula="H" name="H"><atom value="1.0"/></element>
+ </materials>
+ <solids>
+  <box lunit="{unit}" name="world" x="{world_size_mm}" y="{world_size_mm}" z="{world_size_mm}" />
+  <tessellated aunit="degree" lunit="{unit}" name="{name}">
+{triangles}
+  </tessellated>
+ </solids>
+ <structure>
+  <volume name="CadVolume">
+   <materialref ref="{material}"/>
+   <solidref ref="{name}"/>
+  </volume>
+  <volume name="World">
+   <materialref ref="G4_AIR"/>
+   <solidref ref="world"/>
+   <physvol name="cad">
+    <volumeref ref="CadVolume"/>
+    <positionref ref="cadpos"/>
+   </physvol>
+  </volume>
+ </structure>
+ <setup name="Default" version="1.0"><world ref="World"/></setup>
+</gdml>
+"""
+
+
+def cube_mesh(size_mm: float = 100.0) -> tuple[list, list]:
+    """Return (vertices, faces) for an axis-aligned cube centred at origin."""
+    h = size_mm / 2.0
+    v = [
+        (-h, -h, -h),
+        (h, -h, -h),
+        (h, h, -h),
+        (-h, h, -h),
+        (-h, -h, h),
+        (h, -h, h),
+        (h, h, h),
+        (-h, h, h),
+    ]
+    f = [
+        (0, 1, 2),
+        (0, 2, 3),  # -z
+        (4, 6, 5),
+        (4, 7, 6),  # +z
+        (0, 4, 5),
+        (0, 5, 1),  # -y
+        (3, 2, 6),
+        (3, 6, 7),  # +y
+        (0, 3, 7),
+        (0, 7, 4),  # -x
+        (1, 5, 6),
+        (1, 6, 2),  # +x
+    ]
+    return v, f
+
+
+def stl_ascii_to_gdml(stl_text: str, **kwargs) -> str:
+    """Parse minimal ASCII STL (facet normal/vertex lines) into GDML."""
+    verts: list[tuple[float, float, float]] = []
+    index: dict[tuple[float, float, float], int] = {}
+    faces: list[tuple[int, int, int]] = []
+    cur: list[int] = []
+    for line in stl_text.splitlines():
+        s = line.strip().split()
+        if len(s) >= 4 and s[0] == "vertex":
+            p = (float(s[1]), float(s[2]), float(s[3]))
+            if p not in index:
+                index[p] = len(verts)
+                verts.append(p)
+            cur.append(index[p])
+            if len(cur) == 3:
+                faces.append((cur[0], cur[1], cur[2]))
+                cur = []
+    if not faces:
+        msg = "no facets parsed from STL text"
+        raise ValueError(msg)
+    return mesh_to_gdml(verts, faces, **kwargs)

geant4_python_application/detector.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from geant4_python_application import Application
+from geant4_python_application.application import Message
+
+
+class Detector:
+    def __init__(self, application: Application):
+        self._application = application
+
+    def check_overlaps(self) -> Detector:
+        return self._application._send_and_recv(
+            Message("detector", "check_overlaps", (), {})
+        )
+
+    @property
+    def gdml(self) -> str:
+        return self._application._send_and_recv(Message("detector", "get_gdml", (), {}))
+
+    @property
+    def logical_volumes(self):
+        return self._application._send_and_recv(
+            Message("detector", "get_logical_volumes", (), {})
+        )
+
+    @property
+    def physical_volumes(self):
+        return self._application._send_and_recv(
+            Message("detector", "get_physical_volumes", (), {})
+        )
+
+    @property
+    def materials(self):
+        return self._application._send_and_recv(
+            Message("detector", "get_materials", (), {})
+        )
+
+    @property
+    def sensitive_volumes(self):
+        return self._application._send_and_recv(
+            Message("detector", "get_sensitive_volumes", (), {})
+        )
+
+    @sensitive_volumes.setter
+    def sensitive_volumes(self, volumes):
+        self._application._send_and_recv(
+            Message("detector", "set_sensitive_volumes", (volumes,), {})
+        )
+
+    @property
+    def magnetic_field(self) -> tuple[float, float, float]:
+        """Uniform global magnetic field vector in tesla."""
+        return tuple(
+            self._application._send_and_recv(
+                Message("detector", "get_magnetic_field", (), {})
+            )
+        )
+
+    @magnetic_field.setter
+    def magnetic_field(self, field_tesla: tuple[float, float, float]):
+        self._application._send_and_recv(
+            Message("detector", "set_magnetic_field", (field_tesla,), {})
+        )
+
+    @property
+    def electric_field(self) -> tuple[float, float, float]:
+        """Uniform global electric field vector in kV/cm."""
+        return tuple(
+            self._application._send_and_recv(
+                Message("detector", "get_electric_field", (), {})
+            )
+        )
+
+    @electric_field.setter
+    def electric_field(self, field_kv_per_cm: tuple[float, float, float]):
+        self._application._send_and_recv(
+            Message("detector", "set_electric_field", (field_kv_per_cm,), {})
+        )
+
+    def material_from_volume(self, volume: str) -> str:
+        return self._application._send_and_recv(
+            Message("detector", "get_material_from_volume", (volume,), {})
+        )
+
+    def logical_volume_from_physical(self, physical: str) -> str:
+        return self._application._send_and_recv(
+            Message(
+                "detector", "get_logical_volume_from_physical_volume", (physical,), {}
+            )
+        )
+
+    def physical_volumes_from_logical(self, logical: str) -> set[str]:
+        return self._application._send_and_recv(
+            Message(
+                "detector", "get_physical_volumes_from_logical_volume", (logical,), {}
+            )
+        )

geant4_python_application/distributed.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+"""MPI / distributed helpers with graceful fallback (no hard mpi4py dep)."""
+
+import os
+
+
+def get_rank_size() -> tuple[int, int]:
+    """Return (rank, size) from mpi4py if available, else env vars, else (0,1)."""
+    try:
+        from mpi4py import MPI  # type: ignore
+
+        comm = MPI.COMM_WORLD
+        return int(comm.Get_rank()), int(comm.Get_size())
+    except Exception:
+        pass
+    for rank_var, size_var in [
+        ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
+        ("PMI_RANK", "PMI_SIZE"),
+        ("MPI_RANK", "MPI_SIZE"),
+    ]:
+        if rank_var in os.environ and size_var in os.environ:
+            try:
+                return int(os.environ[rank_var]), int(os.environ[size_var])
+            except ValueError:
+                continue
+    return 0, 1
+
+
+def split_count(n_total: int, rank: int, size: int) -> tuple[int, int]:
+    """Return (start, count) for rank under even split of n_total events."""
+    if n_total < 0 or rank < 0 or size <= 0 or rank >= size:
+        msg = f"invalid split args: n={n_total} rank={rank} size={size}"
+        raise ValueError(msg)
+    base, rem = divmod(n_total, size)
+    start = rank * base + min(rank, rem)
+    count = base + (1 if rank < rem else 0)
+    return start, count
+
+
+def local_count(n_total: int) -> tuple[int, int, int]:
+    """Return (rank, size, count) for this process."""
+    rank, size = get_rank_size()
+    _, count = split_count(n_total, rank, size)
+    return rank, size, count

geant4_python_application/events.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import awkward as ak
+import numpy as np
+import vector
+
+
+class EventRecord(ak.Record):
+    def energy_in_volume(self, volume: str) -> float:
+        energy = self.track.step.energy[self.track.step.volume == volume]
+        # collapse to a scalar
+        return ak.sum(energy)
+
+
+class EventArray(ak.Array):
+    def energy_in_volume(self, volume: str) -> ak.Array:
+        energy = self.track.step.energy[self.track.step.volume == volume]
+        # collapse to an array of scalars (one per event)
+        return ak.sum(ak.flatten(energy, axis=-1), axis=1)
+
+    def hits(self, volume: str | None = None) -> ak.Array:
+        hits = ak.flatten(self.track.step, axis=-1)
+        hits = hits[hits.energy > 0]
+        if volume is not None:
+            hits = hits[hits.volume == volume]
+        # Keep available hit fields; event-field selection may omit time or position.
+        fields = [
+            field for field in ("position", "time", "energy") if field in hits.fields
+        ]
+        hits = hits[fields]
+        return ak.Array(hits, with_name="Geant4Hits")
+
+
+class HitsRecord(ak.Record):
+    # This is really expensive!
+    def electrons(self, work_function: float = 20.0 / 1e3) -> ak.Array:
+        n_electrons = np.round(self.energy / work_function)
+        return ak.flatten(
+            ak.Array(
+                [
+                    [
+                        {
+                            "position": self.position[i],
+                            "time": self.time[i],
+                        }
+                        for _ in range(int(n_electrons[i]))
+                    ]
+                    for i in range(len(self.energy))
+                ],
+                with_name="Geant4Electrons",
+            )
+        )
+
+
+class ElectronsRecord(ak.Array):
+    def drift(
+        self,
+        *,
+        drift_velocity: float,
+        longitudinal_diffusion: float,
+        transversal_diffusion: float,
+        plane_point: tuple[float, float, float],
+        plane_normal: tuple[float, float, float],
+    ) -> ElectronsRecord:
+        # TODO: not working yet
+        plane_normal = vector.obj(
+            x=plane_normal[0], y=plane_normal[1], z=plane_normal[2]
+        ).unit()
+        assert np.isclose(plane_normal.mag, 1.0)
+        plane_point = vector.obj(x=plane_point[0], y=plane_point[1], z=plane_point[2])
+        distance_to_plane = (self.position - plane_point) * plane_normal
+        self.time += distance_to_plane / drift_velocity + np.random.normal(
+            0, longitudinal_diffusion * distance_to_plane, len(self.time)
+        )
+        self.position += -distance_to_plane * plane_normal
+        self.position += np.random.normal(
+            0, transversal_diffusion * distance_to_plane, len(self.position)
+        )
+        return self
+
+
+ak.behavior["Geant4Event"] = EventRecord
+ak.behavior["*", "Geant4Event"] = EventArray
+
+ak.behavior["Geant4Hits"] = HitsRecord
+
+ak.behavior["Geant4Electrons"] = ElectronsRecord

geant4_python_application/files/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

geant4_python_application/files/datasets.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+import concurrent.futures
+import hashlib
+import os
+import shutil
+import subprocess
+import tarfile
+import tempfile
+import time
+from collections import namedtuple
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+import geant4_python_application
+
+url = "https://cern.ch/geant4-data/datasets"
+
+
+# It is discouraged to use the python package directory to store data
+# data_dir = os.path.join(os.path.dirname(__file__), "geant4/data")
+# another idea is to use 'platformdirs' to store data in a platform-specific location
+def application_data_directory() -> str:
+    return os.path.join(
+        geant4_python_application.application_directory(),
+        geant4_python_application.__name__,
+        "geant4",
+        geant4_python_application.geant4_version,
+        "data",
+    )
+
+
+def geant4_data_directory() -> str:
+    try:
+        # capture stdout, save into a str
+        geant4_prefix = subprocess.run(
+            ["geant4-config", "--prefix"], capture_output=True, check=False
+        )
+        geant4_data_path = os.path.join(
+            geant4_prefix.stdout.decode().strip(), "share", "Geant4", "data"
+        )
+        if not os.path.exists(geant4_data_path):
+            return ""
+
+        return geant4_data_path
+
+    except FileNotFoundError:
+        return ""
+
+
+# the datasets versions should be updated with each Geant4 version (remember to update the checksum too!)
+# https://geant4.web.cern.ch/download/11.2.2.html#datasets
+Dataset = namedtuple("Dataset", ["name", "version", "filename", "env", "md5sum"])
+
+datasets = (
+    Dataset(
+        name="G4NDL",
+        version="4.7.1",
+        filename="G4NDL",
+        env="G4NEUTRONHPDATA",
+        md5sum="54f0ed3995856f02433d42ec96d70bc6",
+    ),
+    Dataset(
+        name="G4EMLOW",
+        version="8.5",
+        filename="G4EMLOW",
+        env="G4LEDATA",
+        md5sum="146d0625d8d39f294056e1618271bc46",
+    ),
+    Dataset(
+        name="PhotonEvaporation",
+        version="5.7",
+        filename="G4PhotonEvaporation",
+        env="G4LEVELGAMMADATA",
+        md5sum="81ff27deb23af4aa225423e6b3a06b39",
+    ),
+    Dataset(
+        name="RadioactiveDecay",
+        version="5.6",
+        filename="G4RadioactiveDecay",
+        env="G4RADIOACTIVEDATA",
+        md5sum="acc1dbeb87b6b708b2874ced729a3a8f",
+    ),
+    Dataset(
+        name="G4PARTICLEXS",
+        version="4.0",
+        filename="G4PARTICLEXS",
+        env="G4PARTICLEXSDATA",
+        md5sum="d82a4d171d50f55864e28b6cd6f433c0",
+    ),
+    Dataset(
+        name="G4PII",
+        version="1.3",
+        filename="G4PII",
+        env="G4PIIDATA",
+        md5sum="05f2471dbcdf1a2b17cbff84e8e83b37",
+    ),
+    Dataset(
+        name="RealSurface",
+        version="2.2",
+        filename="G4RealSurface",
+        env="G4REALSURFACEDATA",
+        md5sum="ea8f1cfa8d8aafd64b71fb30b3e8a6d9",
+    ),
+    Dataset(
+        name="G4SAIDDATA",
+        version="2.0",
+        filename="G4SAIDDATA",
+        env="G4SAIDXSDATA",
+        md5sum="d5d4e9541120c274aeed038c621d39da",
+    ),
+    Dataset(
+        name="G4ABLA",
+        version="3.3",
+        filename="G4ABLA",
+        env="G4ABLADATA",
+        md5sum="b25d093339e1e4532e31038653580ca6",
+    ),
+    Dataset(
+        name="G4INCL",
+        version="1.2",
+        filename="G4INCL",
+        env="G4INCLDATA",
+        md5sum="0a76df936839bb557dae7254117eb58e",
+    ),
+    Dataset(
+        name="G4ENSDFSTATE",
+        version="2.3",
+        filename="G4ENSDFSTATE",
+        env="G4ENSDFSTATEDATA",
+        md5sum="6f18fce8f217e7aaeaa3711be9b2c7bf",
+    ),
+)
+
+
+def _dataset_url(dataset: Dataset) -> str:
+    return f"{url}/{dataset.filename}.{dataset.version}.tar.gz"
+
+
+def _get_dataset_download_size(dataset: Dataset) -> int:
+    r = requests.head(_dataset_url(dataset))
+    r.raise_for_status()
+    return int(r.headers.get("content-length", 0))
+
+
+def _get_total_download_size(datasets_to_download: list[Dataset] = datasets) -> int:
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=len(datasets_to_download)
+    ) as executor:
+        futures = [
+            executor.submit(_get_dataset_download_size, dataset)
+            for dataset in datasets_to_download
+        ]
+        return sum(f.result() for f in concurrent.futures.as_completed(futures))
+
+
+chunk_size = 1024 * 1024
+max_attempts = 5
+
+
+def _download_extract_dataset(dataset: Dataset, pbar: tqdm):
+    filename = dataset.filename
+    urlpath = _dataset_url(dataset)
+
+    with tempfile.TemporaryFile() as f:
+        downloaded = 0
+        for attempt in range(max_attempts):
+            # resume where the previous attempt died instead of restarting a multi-GB download
+            headers = {"Range": f"bytes={downloaded}-"} if downloaded else {}
+            try:
+                r = requests.get(urlpath, stream=True, headers=headers, timeout=60)
+                r.raise_for_status()
+                if downloaded and r.status_code != 206:
+                    # server ignored the range request, so start over
+                    f.seek(0)
+                    f.truncate()
+                    pbar.update(-downloaded)
+                    downloaded = 0
+                for chunk in r.iter_content(chunk_size=chunk_size):
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    pbar.update(len(chunk))
+                break
+            except requests.RequestException:
+                if attempt == max_attempts - 1:
+                    raise
+                time.sleep(2**attempt)
+
+        f.seek(0)
+        md5 = hashlib.md5()
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            md5.update(chunk)
+        if md5.hexdigest() != dataset.md5sum:
+            print(
+                f"MD5 checksum mismatch for {filename}: got {md5.hexdigest()} but expected {dataset.md5sum}"
+            )
+            raise RuntimeError(f"MD5 checksum mismatch for {filename}")
+
+        f.seek(0)
+        with tarfile.open(fileobj=f, mode="r:gz") as tar:
+            tar.extractall(application_data_directory())
+
+
+def missing_datasets(directory: str) -> list[Dataset]:
+    datasets_to_download = []
+    for dataset in datasets:
+        path = os.path.join(directory, dataset.name + dataset.version)
+        if not os.path.exists(path):
+            datasets_to_download.append(dataset)
+    return datasets_to_download
+
+
+def check_datasets(throw: bool = False) -> bool:
+    datasets_to_download = missing_datasets(directory=application_data_directory())
+    if datasets_to_download:
+        # if not found in the application data directory, try to find them in the Geant4 installation data directory
+        if geant4_data_directory():
+            datasets_to_download = missing_datasets(directory=geant4_data_directory())
+
+    if datasets_to_download:
+        if throw:
+            raise RuntimeError(
+                f"Missing Geant4 datasets: {', '.join([f'{dataset.name}@v{dataset.version}' for dataset in datasets_to_download])}"
+            )
+        return False
+    return True
+
+
+def install_datasets(show_progress: bool = True):
+    # first try to see if the datasets are installed in the application directory
+    datasets_to_download = missing_datasets(directory=application_data_directory())
+    if not datasets_to_download:
+        # datasets are installed in application directory
+        os.environ["GEANT4_DATA_DIR"] = application_data_directory()
+        return
+
+    if geant4_data_directory():
+        datasets_to_download = missing_datasets(directory=geant4_data_directory())
+        if not datasets_to_download:
+            # datasets are installed in Geant4 data directory, no need to install them in the application data directory
+            os.environ["GEANT4_DATA_DIR"] = geant4_data_directory()
+            return
+
+    # download and extract the datasets to the application data directory
+    os.environ["GEANT4_DATA_DIR"] = application_data_directory()
+
+    os.makedirs(application_data_directory(), exist_ok=True)
+    if show_progress:
+        print(
+            f"""
+Geant4 datasets will be installed to "{application_data_directory()}".
+This may take a while but only needs to be done once.
+You can override the default location by calling `application_directory(path)` or `application_directory(temp=True)` to use a temporary directory.
+The following Geant4 datasets will be installed: {", ".join([f"{dataset.name}@v{dataset.version}" for dataset in datasets_to_download])}"""
+        )
+
+    with tqdm(
+        total=_get_total_download_size(datasets_to_download),
+        desc="Downloading Geant4 datasets",
+        disable=not show_progress,
+        unit="B",
+        unit_scale=True,
+    ) as pbar:
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=len(datasets_to_download)
+        ) as executor:
+            futures = [
+                executor.submit(_download_extract_dataset, dataset, pbar)
+                for dataset in datasets_to_download
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                # surface download/extraction errors instead of reporting a partial install as success
+                future.result()
+
+    if show_progress:
+        total_size_gb = sum(
+            fp.stat().st_size for fp in Path(application_data_directory()).rglob("*")
+        ) / (1024**3)
+        print(f"Geant4 datasets size on disk after extraction: {total_size_gb:.2f}GB")
+
+
+def uninstall_datasets():
+    dir_to_remove = os.path.dirname(application_data_directory())
+    package_dir = os.path.dirname(__file__)
+
+    if not os.path.relpath(package_dir, dir_to_remove).startswith(".."):
+        raise RuntimeError(
+            f"Refusing to remove {dir_to_remove} because it is not a subdirectory of {package_dir}"
+        )
+    shutil.rmtree(dir_to_remove, ignore_errors=True)
+
+
+def reinstall_datasets():
+    uninstall_datasets()
+    install_datasets()

geant4_python_application/files/directories.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import os
+import pathlib
+import tempfile
+
+import platformdirs
+
+import geant4_python_application
+
+_app_name = geant4_python_application.__name__
+# TODO: get from pyproject.toml
+_app_author = "lobis"
+
+_dirs = platformdirs.AppDirs(_app_name, _app_author)
+
+# Shared by every environment of the user, so the ~2 GB of datasets download
+# only once. GEANT4_PYTHON_APPLICATION_DIR overrides it, e.g. on clusters.
+_application_directory = os.environ.get(
+    "GEANT4_PYTHON_APPLICATION_DIR", _dirs.user_data_dir
+)
+
+
+def application_directory(path: str | None = None, *, temp: bool = False) -> str:
+    global _application_directory
+
+    if temp and path:
+        raise ValueError("Cannot set both temp and path options")
+    if temp:
+        _application_directory = tempfile.gettempdir()
+        return _application_directory
+    if path:
+        # override the default application directory
+        # make sure path exists, otherwise create it. Throw if it cannot be created
+        pathlib.Path(path).mkdir(parents=True, exist_ok=True)
+        if not os.path.isdir(path):
+            raise ValueError(f"Cannot create application directory: {path}")
+        _application_directory = path
+
+    return _application_directory