Merge branch 'feature/VHEE_model' into 'develop'

Feature/vhee model

See merge request e040/e0404/pyRadPlan!92

Author: wahln

examples/pencilbeam_vhee.py

@@ -0,0 +1,90 @@
+# %% [markdown]
+"""# Example for VHEE dose calculation using pencilbeam engine."""
+# %% [markdown]
+# This example demonstrates how to use the pyRadPlan library to perform dose calculations using the VHEE model.
+
+# IMPORTANT: Please use the ipopt optimizer, because scipy might not converge or converge to fast.
+# (Will be updated in the future)
+
+# To display this script in a Jupyter Notebook, you need to install jupytext via pip and run the following command.
+# This will create a .ipynb file in the same directory:
+
+# ```bash
+# pip install jupytext
+# jupytext --to notebook path/to/this/file/pencilbeam_proton.py
+
+# %%
+# Import necessary libraries
+import logging
+
+import numpy as np
+
+from pyRadPlan import (
+    IonPlan,
+    generate_stf,
+    calc_dose_influence,
+    fluence_optimization,
+    plot_slice,
+    load_tg119,
+)
+
+from pyRadPlan.optimization.objectives import SquaredDeviation, SquaredOverdosing
+
+logging.basicConfig(level=logging.INFO)
+
+# %%
+# Load TG119 (provided within pyRadPlan)
+ct, cst = load_tg119()
+
+# %% [markdown]
+# In this section, we create a VHEE therapy plan using the FermiEyges machine. </br>
+# We choose "Generic" for FermiEyges or "Focused" for a focused VHEE machine. </br>
+# We set the energy to 200 MeV, the bixel width to 5 mm and select 5 gantry angles with 72 degree spacing.
+# %%
+# Create a plan object - (VHEE | Generic) => FermiEyges
+pln = IonPlan(radiation_mode="VHEE", machine="Generic")  # Alternative: Focused
+pln.prop_opt = {"solver": "ipopt"}  # USE IPOPT!
+
+# Some specific Settings for VHEE
+pln.prop_stf = {
+    "energy": 200,  # set VHEE energy at [100, 150, 200] MeV
+    "bixel_width": 5.0,
+    "gantry_angles": [0, 72, 144, 216, 288],
+    "couch_angles": [0, 0, 0, 0, 0],
+}
+
+# Setting the dose grid resolution
+pln.prop_dose_calc = {"dose_grid": {"resolution": [3.0, 3.0, 3.0]}}
+
+# Generate Steering Geometry ("stf")
+stf = generate_stf(ct, cst, pln)
+
+# Calculate Dose Influence Matrix ("dij")
+dij = calc_dose_influence(ct, cst, stf, pln)
+
+# Optimization
+cst.vois[0].objectives = [SquaredOverdosing(priority=10.0, d_max=1.0)]  # OAR
+cst.vois[1].objectives = [SquaredDeviation(priority=100.0, d_ref=3.0)]  # Target
+cst.vois[2].objectives = [SquaredOverdosing(priority=10.0, d_max=2.0)]  # BODY
+
+# Calculate optimized fluence
+fluence = fluence_optimization(ct, cst, stf, dij, pln)
+
+# Compute the result on the CT grid
+result = dij.compute_result_ct_grid(fluence)
+
+# %% [markdown]
+# Visualize the results
+# %%
+# Choose a slice to visualize
+view_slice = int(np.round(ct.size[2] / 2))
+
+# Visualize
+plot_slice(
+    ct=ct,
+    cst=cst,
+    overlay=result["physical_dose"],
+    view_slice=view_slice,
+    plane="axial",
+    overlay_unit="Gy",
+)

pyRadPlan/data/machines/VHEE_Focused.mat

@@ -183,9 +183,13 @@ def _calc_dose(self, ct: CT, cst: StructureSet, stf: SteeringInformation):
                         # Initialize Ray Geometry
                         curr_ray = self._init_ray(curr_beam, j)
 
+                        # Even if the ray hits nothing, we still emit empty bixel columns
+                        # so that CSC structures and bookkeeping stay consistent.
+                        # The following block might be re-enabled in future to skip empty rays.
+
                         # check if ray hit anything. If so, skip the computation
-                        if all(not arr.size for arr in curr_ray["rad_depths"]):
-                            continue
+                        # if all(not arr.size for arr in curr_ray["rad_depths"]):
+                        #     continue #continue
 
                         # TODO: incorporate scenarios correctly
                         for ct_scen in range(self.mult_scen.num_of_ct_scen):
@@ -458,10 +462,12 @@ def _init_ray(self, beam_info: dict[str], j: int) -> dict[str]:
         ray["ray_index"] = j
         ray["iso_center"] = beam_info["beam"]["iso_center"]
 
-        if "num_of_bixels_per_ray" not in beam_info["beam"]:
-            ray["num_of_bixels"] = 1
-        else:
+        if "num_of_bixels_per_ray" in beam_info["beam"]:
             ray["num_of_bixels"] = beam_info["beam"]["num_of_bixels_per_ray"][j]
+        else:
+            # Fallback: use the actual number of beamlets on this ray
+            # This is needed for machines like "Focused" that don't precompute counts.
+            ray["num_of_bixels"] = len(ray.get("beamlets", [])) or 0
 
         ray["source_point_bev"] = beam_info["beam"]["source_point_bev"]
         ray["sad"] = beam_info["beam"]["sad"]
@@ -713,59 +719,45 @@ def _fill_dij(
         dict
             The updated dose influence matrix.
         """
-        # Only fill if we actually had bixel (indices) to compute
-        if bixel and "ix" in bixel and bixel["ix"].any():
-            sub_scen_idx = [
-                np.unravel_index(scen_idx, self.mult_scen.scen_mask.shape)[i]
-                for i in range(self.mult_scen.scen_mask.ndim)
-            ]
-            sub_scen_idx = tuple(sub_scen_idx)
-
-            for q_name in self._computed_quantities:
-                if self._calc_dose_direct:
-                    # We accumulate the current bixel to the container
+        # Determine if we have entries to store for this bixel
+        ix_size = 0
+        if bixel and "ix" in bixel:
+            ix_val = bixel["ix"]
+            try:
+                ix_size = ix_val.size
+            except AttributeError:
+                ix_size = len(ix_val)
+
+        sub_scen_idx = tuple(np.unravel_index(scen_idx, self.mult_scen.scen_mask.shape))
+
+        for q_name in self._computed_quantities:
+            if self._calc_dose_direct:
+                if ix_size > 0:
                     dij[q_name][sub_scen_idx][bixel["ix"], curr_beam_idx] += (
                         bixel["weight"] * bixel[q_name]
                     )
-                else:
-                    # first we fill the index pointer array
-                    dij[q_name][sub_scen_idx]["indptr"][bixel_counter + 1] = (
-                        dij[q_name][sub_scen_idx]["indptr"][bixel_counter] + bixel["ix"].size
-                    )
-
-                    # If we haven't preallocated enough, we need to expand the arrays
-                    if (
-                        dij[q_name][sub_scen_idx]["data"].size
-                        < dij[q_name][sub_scen_idx]["nnz"] + bixel["ix"].size
-                    ):
+            else:
+                data_dict = dij[q_name][sub_scen_idx]
+                # Advance column pointer even when ix_size == 0 to keep CSC valid
+                start = data_dict["indptr"][bixel_counter]
+                end = start + ix_size
+                data_dict["indptr"][bixel_counter + 1] = end
+
+                if ix_size > 0:
+                    need_nnz = data_dict["nnz"] + ix_size
+                    if data_dict["data"].size < need_nnz:
                         logger.debug("Resizing data and indices arrays for %s...", q_name)
-
-                        # We estimate the required size by the remaining
-                        # number of beamlets / columns in the sparse matrix
-                        # 1 additional beamlet is added
-                        shape_resize = (self._num_of_columns_dij - bixel_counter) * (
-                            bixel["ix"].size + 1
-                        )
-                        shape_resize += dij[q_name][sub_scen_idx]["data"].size
-                        dij[q_name][sub_scen_idx]["data"].resize((shape_resize,), refcheck=False)
-                        dij[q_name][sub_scen_idx]["indices"].resize(
-                            (shape_resize,), refcheck=False
+                        grow = max(
+                            ix_size, (self._num_of_columns_dij - bixel_counter) * (ix_size + 1)
                         )
+                        new_size = data_dict["data"].size + grow
+                        data_dict["data"].resize((new_size,), refcheck=False)
+                        data_dict["indices"].resize((new_size,), refcheck=False)
 
-                    # Fill the corresponding values and indices using indptr
-                    dij[q_name][sub_scen_idx]["data"][
-                        dij[q_name][sub_scen_idx]["indptr"][bixel_counter] : dij[q_name][
-                            sub_scen_idx
-                        ]["indptr"][bixel_counter + 1]
-                    ] = bixel[q_name]
-                    dij[q_name][sub_scen_idx]["indices"][
-                        dij[q_name][sub_scen_idx]["indptr"][bixel_counter] : dij[q_name][
-                            sub_scen_idx
-                        ]["indptr"][bixel_counter + 1]
-                    ] = bixel["ix"]
-
-                    # Store how many nnzs we actually have in the matrix
-                    dij[q_name][sub_scen_idx]["nnz"] += bixel["ix"].size
+                    # Fill values and indices into the allocated slice
+                    data_dict["data"][start:end] = bixel[q_name]
+                    data_dict["indices"][start:end] = bixel["ix"]
+                    data_dict["nnz"] += ix_size
 
         # Bookkeeping of bixel numbers
         # remember beam and bixel number
@@ -886,6 +878,8 @@ def calc_geo_dists(
         # Normalize the vector
         b = b / np.linalg.norm(b)
 
+        # Cross product
+        cross_ab = np.cross(a, b)
         # Define rotation matrix
         if np.all(a == b):
             rot_coords_temp = rot_coords_bev[rad_depth_ix, :]
@@ -897,10 +891,10 @@ def ssc(v: np.ndarray) -> np.ndarray:
 
             derived_rot_mat = (
                 np.eye(3)
-                + ssc(np.cross(a, b))
-                + np.dot(ssc(np.cross(a, b)), ssc(np.cross(a, b)))
+                + ssc(cross_ab)
+                + np.dot(ssc(cross_ab), ssc(cross_ab))
                 * (1 - np.dot(a, b))
-                / (np.linalg.norm(np.cross(a, b)) ** 2)
+                / (np.linalg.norm(cross_ab) ** 2)
             )
             rot_coords_temp = np.dot(rot_coords_bev[rad_depth_ix, :], derived_rot_mat)
 

pyRadPlan/data/machines/VHEE_Generic.mat

@@ -52,7 +52,7 @@ class ParticlePencilBeamEngineAbstract(PencilBeamEngineAbstract):
     calc_let: bool
     calc_bio_dose: bool
     air_offset_correction: bool
-    lateral_model: Literal["auto", "single", "double", "multi", "fastest"]
+    lateral_model: Literal["auto", "single", "double", "multi", "fastest", "singleXY"]
     cut_off_method: Literal["integral", "relative"]
 
     def __init__(self, pln):
@@ -81,6 +81,7 @@ def _choose_lateral_model(self):
             "single": self._machine.has_single_gaussian_kernel,
             "double": self._machine.has_double_gaussian_kernel,
             "multi": self._machine.has_multi_gaussian_kernel,
+            "singleXY": self._machine.has_focused_gaussian_kernel,
         }
 
         # First we validate the users choice
@@ -204,15 +205,16 @@ def _init_bixel(self, curr_ray, k):
         # indices to be used) after cutoff. Storing these allows us to
         # use indexing for performance and avoid too many copies
         self._get_bixel_indices_on_ray(bixel, curr_ray)
-        if "ix" not in bixel or bixel["ix"].size == 0:
-            return bixel
 
         # Get quantities 1:1 from ray. Here we trust Python's memory
         # management to not copy the arrays until they are modified.
         # This allows us to efficiently access them by indexing in the
         # bixel computation
         bixel["radial_dist_sq"] = curr_ray["radial_dist_sq"][bixel["sub_ix"]]
         bixel["rad_depths"] = curr_ray["rad_depths"][bixel["sub_ix"]]
+        # Propagate lateral distances if available (needed for singleXY focused model)
+        if "lat_dists" in curr_ray:
+            bixel["lat_dists"] = curr_ray["lat_dists"][bixel["sub_ix"]]
         # TODO:
         # if self.calc_bio_dose:
         #     bixel["v_tissue_index"] = curr_ray["v_tissue_index"][bixel["sub_ix"]]
@@ -238,6 +240,9 @@ def _interpolate_kernels_in_depth(self, bixel):
         # Lateral Kernel Model
         if self.lateral_model == "single":
             used_kernels["sigma"] = kernel.sigma
+        elif self.lateral_model == "singleXY":
+            used_kernels["sigma_x"] = kernel.sigma_x
+            used_kernels["sigma_y"] = kernel.sigma_y
         elif self.lateral_model == "double":
             used_kernels["sigma_1"] = kernel.sigma_1
             used_kernels["sigma_2"] = kernel.sigma_2
@@ -276,7 +281,7 @@ def _get_ray_geometry_from_beam(self, ray: dict[str], beam_info: dict[str]):
         lateral_ray_cutoff = self._get_lateral_distance_from_dose_cutoff_on_ray(ray)
 
         # Ray tracing for beam i and ray j
-        ix, radial_dist_sq, _, _ = self.calc_geo_dists(
+        ix, radial_dist_sq, lat_dists, _ = self.calc_geo_dists(
             beam_info["bev_coords"],
             ray["source_point_bev"],
             ray["target_point_bev"],
@@ -292,6 +297,8 @@ def _get_ray_geometry_from_beam(self, ray: dict[str], beam_info: dict[str]):
         ray["radial_dist_sq"] = [
             radial_dist_sq[beam_ix[ix]] for beam_ix in beam_info["valid_coords"]
         ]
+        if lat_dists is not None:
+            ray["lat_dists"] = [lat_dists[beam_ix[ix]] for beam_ix in beam_info["valid_coords"]]
 
         ray["valid_coords_all"] = np.any(np.vstack(ray["valid_coords"]), axis=1)
 
@@ -590,6 +597,39 @@ def _calc_lateral_particle_cut_off(self, cut_off_level, stf_element):
         dr = np.diff(v_x)
         radial_dist_sq = r_mid**2
 
+        # Handle different lateral models for integration
+        if self.lateral_model == "singleXY":
+            v_theta = np.linspace(0, 2 * np.pi, 361)
+            v_theta = v_theta[:-1]  # remove last to avoid duplication at 2*pi
+            # angular bin size
+            d_theta_angle = (v_theta[1] - v_theta[0]) * np.ones_like(v_theta)
+
+            # Store original radial mid vector length before expansion
+            n_radial_bins = r_mid.size  # equals numel(vR)-1 in MATLAB
+            n_angles = v_theta.size
+
+            r_mid_mesh, theta_mesh = np.meshgrid(r_mid, v_theta)
+
+            # Flatten r_mid_mesh
+            r_mid = r_mid_mesh.ravel(order="F")
+            radial_dist_sq = r_mid**2
+
+            # Expand dr /repeat
+            dr = np.repeat(dr, n_angles)
+
+            # Flatten theta
+            theta = theta_mesh.ravel(order="F")
+
+            # Expand d_theta: replicate angular step sequence for each radial bin
+            d_theta = np.tile(d_theta_angle, n_radial_bins)
+
+            # Lateral distances per sample point (vector of radii times unit vectors)
+            sqrt_r2 = np.sqrt(radial_dist_sq)
+            lat_dists = sqrt_r2[:, np.newaxis] * np.column_stack((np.cos(theta), np.sin(theta)))
+        else:
+            lat_dists = np.tile(np.sqrt(radial_dist_sq / 2)[:, np.newaxis], (1, 2))
+            d_theta = 2 * np.pi
+
         # Number of depth points for which a lateral cutoff is determined
         num_depth_val = 35
 
@@ -753,6 +793,7 @@ def _calc_lateral_particle_cut_off(self, cut_off_level, stf_element):
                     "energy_ix": energy_ix,
                     "kernel": base_kernel,
                     "radial_dist_sq": radial_dist_sq,
+                    "lat_dists": lat_dists,
                     "sigma_ini_sq": largest_sigma_sq4unique_energies[
                         cnt - 1
                     ],  # TODO: check if this result is correct (rounded but may be right)
@@ -770,9 +811,8 @@ def _calc_lateral_particle_cut_off(self, cut_off_level, stf_element):
                 self._calc_particle_bixel(bixel)
                 dose_r = bixel["physical_dose"]
 
-                # Do an integration check that the radial dose integrates to the tabulated
-                # integrated depth dose
-                cum_area = np.cumsum(2 * np.pi * r_mid * dose_r * dr)
+                cum_area = np.cumsum(r_mid * dose_r * dr * d_theta)
+
                 relative_tolerance = 0.5  # in [%]
 
                 if abs((cum_area[-1] / idd[j]) - 1) * 100 > relative_tolerance:

pyRadPlan/dose/engines/_base_pencilbeam.py

@@ -8,7 +8,7 @@ class ParticleHongPencilBeamEngine(ParticlePencilBeamEngineAbstract):
     # constants
     short_name = "HongPB"
     name = "Hong Particle Pencil-Beam"
-    possible_radiation_modes = ["protons", "helium", "carbon"]
+    possible_radiation_modes = ["protons", "helium", "carbon", "VHEE"]
 
     # private methods
     def _calc_particle_bixel(self, bixel):
@@ -46,6 +46,21 @@ def _calc_particle_bixel(self, bixel):
                 ),
                 axis=1,
             )
+        elif self.lateral_model == "singleXY":
+            # Extract squared distances in the two lateral axes
+            x_sq = bixel["lat_dists"][:, 0] ** 2
+            y_sq = bixel["lat_dists"][:, 1] ** 2
+
+            sigma_sq_x = kernels["sigma_x"] ** 2 + bixel["sigma_ini_sq"]
+            sigma_sq_y = kernels["sigma_y"] ** 2 + bixel["sigma_ini_sq"]
+            sigma_x = np.sqrt(sigma_sq_x)
+            sigma_y = np.sqrt(sigma_sq_y)
+
+            # Anisotropic 2D Gaussian
+            lateral = np.exp(-(x_sq / (2 * sigma_sq_x) + y_sq / (2 * sigma_sq_y))) / (
+                2 * np.pi * sigma_x * sigma_y
+            )
+
         else:
             raise ValueError("Invalid Lateral Model")