refactor/machines

Co-authored-by: Niklas Wahl <[email protected]>

Author: FlorianPigged

pyRadPlan/dose/engines/_base_pencilbeam_particle.py

@@ -13,7 +13,11 @@
 
 from pyRadPlan.ct import CT
 from pyRadPlan.stf import SteeringInformation
-from pyRadPlan.machines import IonAccelerator, IonPencilBeamKernel, LateralCutOff
+from pyRadPlan.machines.particles import (
+    ParticleAccelerator,
+    ParticlePencilBeamKernel,
+    LateralCutOff,
+)
 from pyRadPlan.cst import StructureSet
 from ._base_pencilbeam import PencilBeamEngineAbstract
 
@@ -183,7 +187,7 @@ def _init_bixel(self, curr_ray, k):
         bixel["energy_ix"] = energy_ix
 
         # Get the kernel for the current energy
-        tmp_machine = cast(IonAccelerator, self._machine)
+        tmp_machine = cast(ParticleAccelerator, self._machine)
         bixel["kernel"] = tmp_machine.get_kernel_by_index(energy_ix)
 
         bixel["range_shifter"] = curr_ray["beamlets"][k]["range_shifter"]
@@ -218,7 +222,7 @@ def _init_bixel(self, curr_ray, k):
         return bixel
 
     def _interpolate_kernels_in_depth(self, bixel):
-        kernel = cast(IonPencilBeamKernel, bixel["kernel"])
+        kernel = cast(ParticlePencilBeamKernel, bixel["kernel"])
         depths = kernel.depths
 
         # Add potential offset
@@ -299,7 +303,7 @@ def _get_ray_geometry_from_beam(self, ray: dict[str], beam_info: dict[str]):
         ]
 
     def _get_bixel_indices_on_ray(self, curr_bixel, curr_ray):
-        kernel = cast(IonPencilBeamKernel, curr_bixel["kernel"])
+        kernel = cast(ParticlePencilBeamKernel, curr_bixel["kernel"])
 
         # Create offset vector to account for additional offsets modeled in the base data
         # and a potential range shifter
@@ -437,7 +441,7 @@ def _init_beam(
         beam_info = super()._init_beam(dij, ct, cst, stf, i)
 
         # Sanity Check
-        assert isinstance(self._machine, IonAccelerator)
+        assert isinstance(self._machine, ParticleAccelerator)
 
         # Assuming currBeam is part of beam_info
         curr_beam = beam_info["beam"]
@@ -557,7 +561,7 @@ def _allocate_let_container(self, dij):
 
     def _calc_lateral_particle_cut_off(self, cut_off_level, stf_element):
         # Sanity Checks
-        assert isinstance(self._machine, IonAccelerator)
+        assert isinstance(self._machine, ParticleAccelerator)
 
         if len(stf_element) > 1 and not isinstance(stf_element, dict):
             raise ValueError(
@@ -798,7 +802,7 @@ def _calc_lateral_particle_cut_off(self, cut_off_level, stf_element):
     def _init_ray(self, beam_info: dict[str], j: int) -> dict[str]:
         ray = super()._init_ray(beam_info, j)
 
-        self._machine = cast(IonAccelerator, self._machine)
+        self._machine = cast(ParticleAccelerator, self._machine)
 
         # Calculate initial sigma for all bixels on the current ray
         # TODO: here [ray] since calc_sigma_ini takes multiple rays (why?)
@@ -853,7 +857,7 @@ def _fill_dij(
 
     def _get_lateral_distance_from_dose_cutoff_on_ray(self, ray: dict):
         # Find index of maximum used energy (round to keV for numerical reasons)
-        self._machine = cast(IonAccelerator, self._machine)
+        self._machine = cast(ParticleAccelerator, self._machine)
 
         max_energy = max([beamlet["energy"] for beamlet in ray["beamlets"]])
         kernel = self._machine.get_kernel_by_energy(max_energy)

pyRadPlan/dose/engines/_hongpb.py

@@ -1,7 +1,7 @@
 from typing import cast
 import numpy as np
 from ._base_pencilbeam_particle import ParticlePencilBeamEngineAbstract
-from pyRadPlan.machines import IonPencilBeamKernel
+from pyRadPlan.machines import ParticlePencilBeamKernel
 
 
 class ParticleHongPencilBeamEngine(ParticlePencilBeamEngineAbstract):
@@ -14,7 +14,7 @@ class ParticleHongPencilBeamEngine(ParticlePencilBeamEngineAbstract):
     def _calc_particle_bixel(self, bixel):
         kernels = self._interpolate_kernels_in_depth(bixel)
 
-        pb_kernel = cast(IonPencilBeamKernel, bixel["kernel"])
+        pb_kernel = cast(ParticlePencilBeamKernel, bixel["kernel"])
 
         # Lateral Component
         if self.lateral_model == "single":

pyRadPlan/machines/__init__.py

@@ -1,21 +1,36 @@
-"""Treatment machines and beam models."""
+"""Machine definitions for external beam radiotherapy."""
 
-from ._base import Machine, ExternalBeamMachine
-from ._photons import PhotonLINAC, PhotonSVDKernel
-from ._ions import IonAccelerator, IonPencilBeamKernel, LateralCutOff
+from .base import (
+    Machine,
+    ExternalBeamMachine,
+    InternalBeamMachine,
+    register_machine,
+)
+from .photons import PhotonLINAC, PhotonSVDKernel
+from .particles import (
+    ParticlePencilBeamKernel,
+    LateralCutOff,
+    IonAccelerator,
+)
 from ._validate import validate_machine
 from ._load import load_from_name, load_machine_from_mat, load_machine
 
+register_machine(PhotonLINAC)
+register_machine(IonAccelerator)
+# register_machine(VHEEAccelerator)
+
 __all__ = [
     "Machine",
     "ExternalBeamMachine",
+    "InternalBeamMachine",
     "PhotonLINAC",
     "PhotonSVDKernel",
-    "load_from_name",
-    "validate_machine",
-    "IonAccelerator",
-    "IonPencilBeamKernel",
+    "ParticlePencilBeamKernel",
     "LateralCutOff",
+    "IonAccelerator",
+    "validate_machine",
+    "load_from_name",
     "load_machine_from_mat",
     "load_machine",
+    "register_machine",
 ]