Bug fix in visualization.py: visualize dose correctly

Author: gourav3017

portpy/photon/visualization.py

@@ -311,7 +311,11 @@ def plot_2d_slice(my_plan: Plan = None, sol: dict = None, ct: CT = None, structs
         if show_dose:
             dose_1d = sol['inf_matrix'].A @ (sol['optimal_intensity'] * my_plan.get_num_of_fractions())
             dose_3d = sol['inf_matrix'].dose_1d_to_3d(dose_1d=dose_1d)
-            masked = np.ma.masked_where(dose_3d[slice_num, :, :] <= 0, dose_3d[slice_num, :, :])
+            if hasattr(my_plan, 'structures'):
+                body_mask = my_plan.structures.get_structure_mask_3d('BODY')
+                masked = np.ma.masked_where(~body_mask[slice_num, :, :].astype(bool), dose_3d[slice_num, :, :])
+            else:
+                masked = np.ma.masked_where(dose_3d[slice_num, :, :] < 0, dose_3d[slice_num, :, :])
             im = ax.imshow(masked, alpha=0.4, interpolation='none',
                            cmap='rainbow')
 

portpy/photon/vmat_scp/vmat_scp_optimization.py

@@ -919,9 +919,9 @@ def run_sequential_cvx_algo_prediction(self, pred_dose_1d, *args, **kwargs):
                 vmat_params['step_size_b'] = intial_step_size
                 best_obj_value = sol['actual_obj_value']
                 self.arcs.update_best_solution()
+                sol['inner_iteration'] = inner_iteration
                 inner_iteration = inner_iteration + 1
                 sol['accept'] = True
-                sol['inner_iteration'] = inner_iteration
 
             else:
                 if sol['actual_obj_value'] < best_obj_value: