Updated notebooks and created utils

gourav3017 · gourav3017 · commit 32e5626b5338 · 2024-10-01T12:24:33.000-04:00
diff --git a/compress_rtp/compress_rtp_optimization.py b/compress_rtp/compress_rtp_optimization.py
@@ -9,15 +9,6 @@
 import cvxpy as cp
 import numpy as np
 from copy import deepcopy
-try:
-    from sklearn.utils.extmath import randomized_svd
-except ImportError:
-    pass
-import scipy
-try:
-    import pywt
-except ImportError:
-    pass
 
 
 class CompressRTPOptimization(Optimization):
@@ -165,82 +156,4 @@ def create_cvxpy_problem_compressed(self, S=None, H=None, W=None):
                                         <= limit / num_fractions]
 
         constraints += [Wx == (W @ x)]
-        print('Constraints done')
-
-    def get_sparse_plus_low_rank(self, A=None, thresold_perc=1, rank=5):
-        """
-            :param A: dose influence matrix
-            :param thresold_perc: thresold percentage. Default to 1% of max(A)
-            :type rank: rank of L = A-S.
-            :returns: S, H, W using randomized svd
-        """
-        if A is None:
-            A = deepcopy(self.inf_matrix.A)
-        tol = np.max(A) * thresold_perc * 0.01
-        # S = S*0
-        S = np.where(A > tol, A, 0)
-        if rank == 0:
-            H = np.zeros((A.shape[0], 1))
-            W = np.zeros((1, A.shape[1]))
-        else:
-            print('Running svd..')
-            [U, svd_S, V] = randomized_svd(A - S, n_components=rank + 1, random_state=0)
-            print('svd done!')
-            H = U[:, :rank]
-            W = np.diag(svd_S[:rank]) @ V[:rank, :]
-        S = scipy.sparse.csr_matrix(S)
-        return S, H, W
-
-    def get_low_dim_basis(self, inf_matrix: InfluenceMatrix = None, compression: str = 'wavelet'):
-        """
-        :param inf_matrix: an object of class InfluenceMatrix for the specified plan
-        :param compression: the compression method
-        :type compression: str
-        :return: a list that contains the dimension reduction basis in the format of array(float)
-        """
-        if inf_matrix is None:
-            inf_matrix = self.inf_matrix
-        low_dim_basis = {}
-        num_of_beams = len(inf_matrix.beamlets_dict)
-        num_of_beamlets = inf_matrix.beamlets_dict[num_of_beams - 1]['end_beamlet_idx'] + 1
-        beam_id = [inf_matrix.beamlets_dict[i]['beam_id'] for i in range(num_of_beams)]
-        beamlets = inf_matrix.get_bev_2d_grid(beam_id=beam_id)
-        index_position = list()
-        for ind in range(num_of_beams):
-            low_dim_basis[beam_id[ind]] = []
-            for i in range(inf_matrix.beamlets_dict[ind]['start_beamlet_idx'],
-                           inf_matrix.beamlets_dict[ind]['end_beamlet_idx'] + 1):
-                index_position.append((np.where(beamlets[ind] == i)[0][0], np.where(beamlets[ind] == i)[1][0]))
-        if compression == 'wavelet':
-            max_dim_0 = np.max([beamlets[ind].shape[0] for ind in range(num_of_beams)])
-            max_dim_1 = np.max([beamlets[ind].shape[1] for ind in range(num_of_beams)])
-            beamlet_2d_grid = np.zeros((int(np.ceil(max_dim_0 / 2)), int(np.ceil(max_dim_1 / 2))))
-            for row in range(beamlet_2d_grid.shape[0]):
-                for col in range(beamlet_2d_grid.shape[1]):
-                    beamlet_2d_grid[row][col] = 1
-                    approximation_coeffs = pywt.idwt2((beamlet_2d_grid, (None, None, None)), 'sym4',
-                                                      mode='periodization')
-                    horizontal_coeffs = pywt.idwt2((None, (beamlet_2d_grid, None, None)), 'sym4', mode='periodization')
-                    for b in range(num_of_beams):
-                        if ((2 * row + 1 < beamlets[b].shape[0] and 2 * col + 1 < beamlets[b].shape[1] and
-                             beamlets[b][2 * row + 1][2 * col + 1] != -1) or
-                                (2 * row + 1 < beamlets[b].shape[0] and 2 * col < beamlets[b].shape[1] and
-                                 beamlets[b][2 * row + 1][2 * col] != -1) or
-                                (2 * row < beamlets[b].shape[0] and 2 * col + 1 < beamlets[b].shape[1] and
-                                 beamlets[b][2 * row][2 * col + 1] != -1) or
-                                (2 * row < beamlets[b].shape[0] and 2 * col < beamlets[b].shape[1] and
-                                 beamlets[b][2 * row][2 * col] != -1)):
-                            approximation = np.zeros(num_of_beamlets)
-                            horizontal = np.zeros(num_of_beamlets)
-                            for ind in range(inf_matrix.beamlets_dict[b]['start_beamlet_idx'],
-                                             inf_matrix.beamlets_dict[b]['end_beamlet_idx'] + 1):
-                                approximation[ind] = approximation_coeffs[index_position[ind]]
-                                horizontal[ind] = horizontal_coeffs[index_position[ind]]
-                            low_dim_basis[beam_id[b]].append(np.transpose(np.stack([approximation, horizontal])))
-                    beamlet_2d_grid[row][col] = 0
-        for b in beam_id:
-            low_dim_basis[b] = np.concatenate(low_dim_basis[b], axis=1)
-            u, s, vh = scipy.sparse.linalg.svds(low_dim_basis[b], k=min(low_dim_basis[b].shape[0], low_dim_basis[b].shape[1]) - 1)
-            ind = np.where(s > 0.0001)
-            low_dim_basis[b] = u[:, ind[0]]
-        return np.concatenate([low_dim_basis[b] for b in beam_id], axis=1)
+        print('Constraints done')
diff --git a/compress_rtp/utils/__init__.py b/compress_rtp/utils/__init__.py
diff --git a/compress_rtp/utils/get_low_dim_basis.py b/compress_rtp/utils/get_low_dim_basis.py
@@ -0,0 +1,64 @@
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from portpy.photon.influence_matrix import InfluenceMatrix
+import numpy as np
+import scipy
+try:
+    import pywt
+except ImportError:
+    pass
+
+
+def get_low_dim_basis(inf_matrix: InfluenceMatrix, compression: str = 'wavelet'):
+    """
+    :param inf_matrix: an object of class InfluenceMatrix for the specified plan
+    :param compression: the compression method
+    :type compression: str
+    :return: a list that contains the dimension reduction basis in the format of array(float)
+    """
+
+    low_dim_basis = {}
+    num_of_beams = len(inf_matrix.beamlets_dict)
+    num_of_beamlets = inf_matrix.beamlets_dict[num_of_beams - 1]['end_beamlet_idx'] + 1
+    beam_id = [inf_matrix.beamlets_dict[i]['beam_id'] for i in range(num_of_beams)]
+    beamlets = inf_matrix.get_bev_2d_grid(beam_id=beam_id)
+    index_position = list()
+    for ind in range(num_of_beams):
+        low_dim_basis[beam_id[ind]] = []
+        for i in range(inf_matrix.beamlets_dict[ind]['start_beamlet_idx'],
+                       inf_matrix.beamlets_dict[ind]['end_beamlet_idx'] + 1):
+            index_position.append((np.where(beamlets[ind] == i)[0][0], np.where(beamlets[ind] == i)[1][0]))
+    if compression == 'wavelet':
+        max_dim_0 = np.max([beamlets[ind].shape[0] for ind in range(num_of_beams)])
+        max_dim_1 = np.max([beamlets[ind].shape[1] for ind in range(num_of_beams)])
+        beamlet_2d_grid = np.zeros((int(np.ceil(max_dim_0 / 2)), int(np.ceil(max_dim_1 / 2))))
+        for row in range(beamlet_2d_grid.shape[0]):
+            for col in range(beamlet_2d_grid.shape[1]):
+                beamlet_2d_grid[row][col] = 1
+                approximation_coeffs = pywt.idwt2((beamlet_2d_grid, (None, None, None)), 'sym4',
+                                                  mode='periodization')
+                horizontal_coeffs = pywt.idwt2((None, (beamlet_2d_grid, None, None)), 'sym4', mode='periodization')
+                for b in range(num_of_beams):
+                    if ((2 * row + 1 < beamlets[b].shape[0] and 2 * col + 1 < beamlets[b].shape[1] and
+                         beamlets[b][2 * row + 1][2 * col + 1] != -1) or
+                            (2 * row + 1 < beamlets[b].shape[0] and 2 * col < beamlets[b].shape[1] and
+                             beamlets[b][2 * row + 1][2 * col] != -1) or
+                            (2 * row < beamlets[b].shape[0] and 2 * col + 1 < beamlets[b].shape[1] and
+                             beamlets[b][2 * row][2 * col + 1] != -1) or
+                            (2 * row < beamlets[b].shape[0] and 2 * col < beamlets[b].shape[1] and
+                             beamlets[b][2 * row][2 * col] != -1)):
+                        approximation = np.zeros(num_of_beamlets)
+                        horizontal = np.zeros(num_of_beamlets)
+                        for ind in range(inf_matrix.beamlets_dict[b]['start_beamlet_idx'],
+                                         inf_matrix.beamlets_dict[b]['end_beamlet_idx'] + 1):
+                            approximation[ind] = approximation_coeffs[index_position[ind]]
+                            horizontal[ind] = horizontal_coeffs[index_position[ind]]
+                        low_dim_basis[beam_id[b]].append(np.transpose(np.stack([approximation, horizontal])))
+                beamlet_2d_grid[row][col] = 0
+    for b in beam_id:
+        low_dim_basis[b] = np.concatenate(low_dim_basis[b], axis=1)
+        u, s, vh = scipy.sparse.linalg.svds(low_dim_basis[b], k=min(low_dim_basis[b].shape[0], low_dim_basis[b].shape[1]) - 1)
+        ind = np.where(s > 0.0001)
+        low_dim_basis[b] = u[:, ind[0]]
+    return np.concatenate([low_dim_basis[b] for b in beam_id], axis=1)
diff --git a/compress_rtp/utils/get_sparse_plus_low_rank.py b/compress_rtp/utils/get_sparse_plus_low_rank.py
@@ -0,0 +1,29 @@
+import numpy as np
+import scipy
+
+try:
+    from sklearn.utils.extmath import randomized_svd
+except ImportError:
+    pass
+
+
+def get_sparse_plus_low_rank(A: np.ndarray, thresold_perc: float = 1, rank: int = 5):
+    """
+        :param A: dose influence matrix
+        :param thresold_perc: thresold percentage. Default to 1% of max(A)
+        :type rank: rank of L = A-S.
+        :returns: S, H, W using randomized svd
+    """
+    tol = np.max(A) * thresold_perc * 0.01
+    S = np.where(A > tol, A, 0)
+    if rank == 0:
+        S = scipy.sparse.csr_matrix(S)
+        return S
+    else:
+        print('Running svd..')
+        [U, svd_S, V] = randomized_svd(A - S, n_components=rank + 1, random_state=0)
+        print('svd done!')
+        H = U[:, :rank]
+        W = np.diag(svd_S[:rank]) @ V[:rank, :]
+        S = scipy.sparse.csr_matrix(S)
+        return S, H, W
diff --git a/examples/fluence_wavelets.ipynb b/examples/fluence_wavelets.ipynb
@@ -26,7 +26,7 @@
    "outputs": [],
    "source": [
     "import portpy.photon as pp\n",
-    "from compress_rtp.compress_rtp_optimization import CompressRTPOptimization\n",
+    "from compress_rtp.utils.get_low_dim_basis import get_low_dim_basis\n",
     "import cvxpy as cp\n",
     "import matplotlib.pyplot as plt"
    ]
@@ -124,7 +124,7 @@
     "my_plan = pp.Plan(ct=ct, structs=structs, beams=beams, inf_matrix=inf_matrix, clinical_criteria=clinical_criteria)\n",
     "\n",
     "# create cvxpy problem using the clinical criteria and optimization parameters\n",
-    "opt = CompressRTPOptimization(my_plan, opt_params=opt_params)\n",
+    "opt = pp.Optimization(my_plan, opt_params=opt_params)\n",
     "opt.create_cvxpy_problem()\n",
     "sol_no_quad_no_wav = opt.solve(solver='MOSEK', verbose=False)"
    ]
@@ -145,7 +145,7 @@
    "outputs": [],
    "source": [
     "# creating the wavelet incomplete basis representing a low dimensional subspace for dimension reduction\n",
-    "wavelet_basis = opt.get_low_dim_basis(inf_matrix=inf_matrix, 'wavelet')\n",
+    "wavelet_basis = get_low_dim_basis(inf_matrix=inf_matrix, compression='wavelet')\n",
     "# Smoothness Constraint\n",
     "y = cp.Variable(wavelet_basis.shape[1])\n",
     "opt.constraints += [wavelet_basis @ y == opt.vars['x']]\n",
diff --git a/examples/fluence_wavelets.py b/examples/fluence_wavelets.py
@@ -5,7 +5,7 @@
 """
 
 import portpy.photon as pp
-from compress_rtp.compress_rtp_optimization import CompressRTPOptimization
+from compress_rtp.utils.get_low_dim_basis import get_low_dim_basis
 import cvxpy as cp
 import matplotlib.pyplot as plt
 
@@ -64,14 +64,14 @@ def ex_wavelet():
     my_plan = pp.Plan(ct=ct, structs=structs, beams=beams, inf_matrix=inf_matrix, clinical_criteria=clinical_criteria)
 
     # create cvxpy problem using the clinical criteria and optimization parameters
-    opt = CompressRTPOptimization(my_plan, opt_params=opt_params)
+    opt = pp.Optimization(my_plan, opt_params=opt_params)
     opt.create_cvxpy_problem()
     sol_no_quad_no_wav = opt.solve(solver='MOSEK', verbose=False)
 
     # - With wavelet constraint
 
     # creating the wavelet incomplete basis representing a low dimensional subspace for dimension reduction
-    wavelet_basis = opt.get_low_dim_basis()
+    wavelet_basis = get_low_dim_basis(inf_matrix=inf_matrix, compression='wavelet')
     # Smoothness Constraint
     y = cp.Variable(wavelet_basis.shape[1])
     opt.constraints += [wavelet_basis @ y == opt.vars['x']]
@@ -88,7 +88,7 @@ def ex_wavelet():
             opt_params['objective_functions'][i]['weight'] = 10
 
     # create cvxpy problem using the clinical criteria and optimization parameters
-    opt = CompressRTPOptimization(my_plan, opt_params=opt_params)
+    opt = pp.Optimization(my_plan, opt_params=opt_params)
     opt.create_cvxpy_problem()
 
     sol_quad_no_wav = opt.solve(solver='MOSEK', verbose=False)
@@ -154,7 +154,7 @@ def ex_wavelet():
     
     '''
     # visualize plan metrics based upon clinical criteria
-    pp.Evaluation.plan_metrics(my_plan,
+    pp.Evaluation.display_clinical_criteria(my_plan,
                               sol=[sol_no_quad_no_wav, sol_no_quad_with_wav, sol_quad_no_wav, sol_quad_with_wav],
                               sol_names=['no_quad_no_wav', 'no_quad_with_wav', 'quad_no_wav', 'quad_with_wav'])
 
diff --git a/examples/matrix_spare_plus_low_rank.py b/examples/matrix_spare_plus_low_rank.py
@@ -6,10 +6,9 @@
 import os
 import portpy.photon as pp
 from compress_rtp.compress_rtp_optimization import CompressRTPOptimization
+from compress_rtp.utils.get_sparse_plus_low_rank import get_sparse_plus_low_rank
 import matplotlib.pyplot as plt
 from copy import deepcopy
-import numpy as np
-import scipy
 
 
 def sparse_plus_low_rank():
@@ -25,7 +24,7 @@ def sparse_plus_low_rank():
     data = pp.DataExplorer(data_dir=data_dir)
 
     # pick a patient from the existing patient list to get detailed info (e.g., beam angles, structures).
-    data.patient_id = 'Lung_Patient_3'
+    data.patient_id = 'Lung_Patient_2'
 
     ct = pp.CT(data)
     structs = pp.Structures(data)
@@ -71,18 +70,16 @@ def sparse_plus_low_rank():
     # run optimization with naive thresold of 1% of max(A) and no low rank
     # create cvxpy problem using the clinical criteria and optimization parameters
     A = deepcopy(inf_matrix.A)
-    tol = np.max(A) * 1 * 0.01
-    S = np.where(A > tol, A, 0)
-    S = scipy.sparse.csr_matrix(S)
+    S = get_sparse_plus_low_rank(A=A, thresold_perc=1, rank=0)
     inf_matrix.A = S
     opt = pp.Optimization(my_plan, inf_matrix=inf_matrix, opt_params=opt_params)
     opt.create_cvxpy_problem()
     sol_sparse = opt.solve(solver='MOSEK', verbose=True)
 
     # run optimization with thresold of 1% and rank 5
     # create cvxpy problem using the clinical criteria and optimization parameters
+    S, H, W = get_sparse_plus_low_rank(A=A, thresold_perc=1, rank=5)
     opt = CompressRTPOptimization(my_plan, opt_params=opt_params)
-    S, H, W = opt.get_sparse_plus_low_rank(A=A, thresold_perc=1, rank=5)
     opt.create_cvxpy_problem_compressed(S=S, H=H, W=W)
 
     # run imrt fluence map optimization using cvxpy and one of the supported solvers and save the optimal solution in sol
@@ -95,22 +92,19 @@ def sparse_plus_low_rank():
 
     """
 
-    fig, ax = plt.subplots(figsize=(12, 8))
+    fig, ax = plt.subplots(1, 2, figsize=(20, 8))
     struct_names = ['PTV', 'ESOPHAGUS', 'HEART', 'CORD', 'LUNGS_NOT_GTV']
     dose_1d_sparse = (S @ sol_sparse['optimal_intensity']) * my_plan.get_num_of_fractions()
-    dose_1d_full = (A @ sol_sparse['optimal_intensity']) * my_plan.get_num_of_fractions()
-    ax = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_sparse, struct_names=struct_names, style='dotted', ax=ax, norm_flag=True)
-    ax = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_full, struct_names=struct_names, style='solid', ax=ax, norm_flag=True)
-    ax.set_title("sparse_vs_full")
-    plt.show(block=False)
+    dose_1d_full_sparse = (A @ sol_sparse['optimal_intensity']) * my_plan.get_num_of_fractions()
+    ax0 = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_sparse, struct_names=struct_names, style='dotted', ax=ax[0], norm_flag=True)
+    ax0 = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_full_sparse, struct_names=struct_names, style='solid', ax=ax0, norm_flag=True)
+    ax0.set_title("sparse_vs_full")
 
-    fig, ax = plt.subplots(figsize=(12, 8))
-    struct_names = ['PTV', 'ESOPHAGUS', 'HEART', 'CORD', 'LUNGS_NOT_GTV']
     dose_1d_slr = (S @ sol_slr['optimal_intensity'] + H @ (W @ sol_slr['optimal_intensity'])) * my_plan.get_num_of_fractions()
-    dose_1d_full = (A @ sol_slr['optimal_intensity']) * my_plan.get_num_of_fractions()
-    ax = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_slr, struct_names=struct_names, style='dashed', ax=ax, norm_flag=True)
-    ax = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_full, struct_names=struct_names, style='solid', ax=ax, norm_flag=True)
-    ax.set_title("slr_vs_full")
+    dose_1d_full_slr = (A @ sol_slr['optimal_intensity']) * my_plan.get_num_of_fractions()
+    ax1 = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_slr, struct_names=struct_names, style='dashed', ax=ax[1], norm_flag=True)
+    ax1 = pp.Visualization.plot_dvh(my_plan, dose_1d=dose_1d_full_slr, struct_names=struct_names, style='solid', ax=ax1, norm_flag=True)
+    ax1.set_title("slr_vs_full")
     plt.show(block=False)
 
     """ 
@@ -122,7 +116,7 @@ def sparse_plus_low_rank():
     """
 
     # visualize plan metrics based upon clinical criteria
-    pp.Evaluation.display_clinical_criteria(my_plan, dose_1d=[dose_1d_sparse, dose_1d_slr], in_browser=True)
+    pp.Evaluation.display_clinical_criteria(my_plan, dose_1d=[dose_1d_full_sparse, dose_1d_full_slr], sol_names=['Without compression', 'With compression'])
 
     """ 
     5) saving and loading the plan for future use (utils) 
diff --git a/examples/matrix_sparse_plus_low_rank.ipynb b/examples/matrix_sparse_plus_low_rank.ipynb
diff --git a/examples/sparse_plus_low_rank.ipynb b/examples/sparse_plus_low_rank.ipynb
