Merge pull request #18 from sijiaww/feature/colgen-dose-mimic

Add prediction-based dose mimicking APIs to VMAT column generation

Author: gourav3017

portpy/photon/vmat_scp/vmat_optimization_col_gen.py

@@ -44,7 +44,7 @@ class VmatOptimizationColGen(Optimization):
     def __init__(self, my_plan: Plan, arcs=None, inf_matrix: InfluenceMatrix = None,
                  clinical_criteria: ClinicalCriteria = None,
                  opt_params: dict = None, vars: dict = None, sol=None,
-                 is_corr: bool = False, delta: np.ndarray = None):
+                 is_corr: bool = False, delta: np.ndarray = None, predicted_dose = None):
 
         # combine steps objective functions into one
         if 'col_gen' in opt_params:
@@ -79,6 +79,7 @@ def __init__(self, my_plan: Plan, arcs=None, inf_matrix: InfluenceMatrix = None,
         self.constraints_ind_map = {}
         if delta is None:
             self.delta = np.zeros(self.inf_matrix.A.shape[0])
+        self.predicted_dose = predicted_dose
 
     def calculate_scores_grad_new(self, A, A_indices, sol):
         """
@@ -107,6 +108,8 @@ def calculate_scores_grad_new(self, A, A_indices, sol):
 
         # Initialize gradients
         num_beamlets = A.shape[1]
+        ptv_dm_grad = np.zeros(num_beamlets)
+        oar_dm_grad = np.zeros(num_beamlets)
         overdose_grad = np.zeros(num_beamlets)
         underdose_grad = np.zeros(num_beamlets)
         quadratic_grad = np.zeros(num_beamlets)
@@ -118,8 +121,33 @@ def calculate_scores_grad_new(self, A, A_indices, sol):
         inf_matrix = self.my_plan.inf_matrix
         num_fractions = self.my_plan.get_num_of_fractions()
 
+        if self.predicted_dose is not None:
+            pred_dose_per_frac = self.predicted_dose / num_fractions
+
         for obj in obj_funcs:
-            if obj['type'] == 'quadratic-overdose':
+            if obj['type'] == "PTV-dose-mimicking":
+                struct = obj['structure_name']
+                if struct in self.my_plan.structures.get_structures():
+                    vox = inf_matrix.get_opt_voxels_idx(struct)
+                    if len(vox) == 0:
+                        continue
+                    weight = float(obj['weight'])
+                    influence = A[vox, :] 
+                    ptv_dm_grad += (2 * weight / len(vox)) * (influence.T @ (current_dose[vox] - pred_dose_per_frac[vox]))
+            elif obj['type'] == "OAR-dose-mimicking":
+                target_voxels = self.inf_matrix.get_opt_voxels_idx('PTV') 
+                all_vox = self.my_plan.inf_matrix.get_opt_voxels_idx('BODY') 
+                oar_voxels = np.setdiff1d(all_vox, target_voxels)
+                if oar_voxels.size == 0:
+                    continue
+                weight = obj['weight']
+                influence = A[oar_voxels, :]                        
+                d_now = current_dose[oar_voxels]                   
+                d_hat = pred_dose_per_frac[oar_voxels]              
+                over = d_now - d_hat
+                over[over < 0.0] = 0.0                       
+                oar_dm_grad += (2.0 * weight / len(oar_voxels)) * (influence.T @ over)
+            elif obj['type'] == 'quadratic-overdose':
                 struct = obj['structure_name']
                 if struct not in structures.get_structures():
                     continue
@@ -567,10 +595,34 @@ def solve_rmp(self, *args, **kwargs):
         obj = self.obj
         constraints = self.constraints
 
+        if self.predicted_dose is not None:
+            pred_dose_per_fraction = self.predicted_dose / num_fractions
+
 
         # obj += [(1 / B.shape[0]) * cp.sum_squares(cp.multiply(np.sqrt(self.weights), B @ y - p))]
         for i in range(len(obj_funcs)):
-            if obj_funcs[i]['type'] == 'quadratic-overdose':
+            if obj_funcs[i]['type'] == "PTV-dose-mimicking":
+                struct = obj_funcs[i]['structure_name']
+                if struct in self.my_plan.structures.get_structures():
+                    vox = st.get_opt_voxels_idx(struct)
+                    if len(vox) == 0:
+                        continue
+                    weight = float(obj_funcs[i]['weight'])
+                    obj += [
+                        (weight / len(vox)) * cp.sum_squares(B[vox, :] @ y - pred_dose_per_fraction[vox])
+                    ]
+            elif obj_funcs[i]['type'] == "OAR-dose-mimicking":
+                target_voxels = self.inf_matrix.get_opt_voxels_idx('PTV')
+                all_vox = self.my_plan.inf_matrix.get_opt_voxels_idx('BODY')
+                oar_voxels = np.setdiff1d(all_vox, target_voxels)
+                weight = float(obj_funcs[i]['weight'])
+                dO = B[oar_voxels, :] @ y
+                s_var = cp.Variable(len(oar_voxels), nonneg=True)
+                constraints += [ dO <= pred_dose_per_fraction[oar_voxels] + s_var ]
+                obj += [
+                    (weight / len(oar_voxels)) * cp.sum_squares(s_var)
+                ]
+            elif obj_funcs[i]['type'] == 'quadratic-overdose':
                 if obj_funcs[i]['structure_name'] in my_plan.structures.get_structures():
                     struct = obj_funcs[i]['structure_name']
                     if len(st.get_opt_voxels_idx(struct)) == 0:  # check if there are any opt voxels for the structure

portpy/photon/vmat_scp/vmat_scp_optimization.py

@@ -1220,7 +1220,8 @@ def run_sequential_cvx_algo_prediction(self, pred_dose_1d, *args, **kwargs):
         inner_iteration = int(0)
         best_obj_value = 0
         vmat_params = self.vmat_params
-        self.arcs.get_initial_leaf_pos(initial_leaf_pos=vmat_params['initial_leaf_pos'])
+        if not vmat_params['initial_leaf_pos'].lower() == 'cg':
+            self.arcs.get_initial_leaf_pos(initial_leaf_pos=vmat_params['initial_leaf_pos'])
         self.create_cvx_params()
         sol_convergence = []
         while True: