constants_class.py

import os

import matplotlib.pyplot as plt
import pandas as pd


class ModelParameters:
    def __init__(self, io_name: str = "baseline"):
        """Define all the constant parameters for this project"""

        # If data is outside of open-kbp-opt, change `parent_data_directory` to where data is stored.
        self.primary_directory = f"{os.getcwd()}/open-kbp-opt-data"  # directory where everything is stored

        self._set_directories(io_name)
        self._set_plan_criteria()
        self._set_plot_fonts()

        # Placeholders for directories
        self.plan_dose_from_pred_dir = None
        self.plan_fluence_from_pred_dir = None
        self.plan_gap_from_pred_dir = None
        self.plan_weights_from_pred_dir = None

        # Set beamlet/fluence shape parameters
        self.fluence_shape = (64, 64)
        self.max_beamlet_row = 32
        self.max_fluence_column = 32
        self.num_beam_angles = 9
        self.beam_view_shape = 64
        self.patient_number = None
        self.opt_name = None
        self.prediction_name = None

        # Set patient parameters
        self.patient_shape = (128, 128, 128)
        self.rois = {
            "oars": ["Brainstem", "SpinalCord", "RightParotid", "LeftParotid", "Esophagus", "Larynx", "Mandible"],
            "targets": ["PTV56", "PTV63", "PTV70"],
        }

        self.all_rois = self._get_all_rois()
        self.num_rois = len(self.all_rois)
        self.plan_criteria_dict = pd.Series(
            {
                ("D_0.1_cc", "Brainstem"): 50,
                ("D_0.1_cc", "SpinalCord"): 45,
                ("D_0.1_cc", "Mandible"): 73.5,
                ("mean", "RightParotid"): 26,
                ("mean", "LeftParotid"): 26,
                ("mean", "Esophagus"): 45,
                ("mean", "Larynx"): 45,
                ("D_99", "PTV56"): -53.2,  # negative sign used to communicate that higher values better than lower
                ("D_99", "PTV63"): -59.85,  # negative sign used to communicate that higher values better than lower
                ("D_99", "PTV70"): -66.5,  # negative sign used to communicate that higher values better than lower
            }
        )  # Set patient view shapes
        self.roi_series = pd.Series(dict((i, k) for k, v in self.rois.items() for i in v))

        self.dose_shape = (*self.patient_shape, 1)
        self.ct_shape = (*self.patient_shape, 1)
        self.roi_masks_shape = (*self.patient_shape, self.num_rois)

        # Set maximum limits
        self.max_dose = 82
        self.max_beam = 15

        # Set statistics
        self.p_thresh = 0.05

        # Dictionary for plots axis
        self.optimization_short_hands_dict = {
            "Prediction": "Prediction",
            "absolute_mean": "MeanAbs",
            "absolute_max": "MaxAbs",
            "relative_mean": "MeanRel",
            "relative_max": "MaxRel",
        }

        self.dvh_metric_axis_dict = {
            "D_0.1_cc": "D$_\mathrm{0.1cc}$",
            "mean": "D$_\mathrm{mean}$",
            "D_99": "D$_\mathrm{99}$",
            "D_95": "D$_\mathrm{95}$",
            "D_1": "D$_\mathrm{1}$",
        }

        self.structure_printing = {
            "SpinalCord": "Spinal cord",
            "RightParotid": "Right parotid",
            "LeftParotid": "Left parotid",
            "oars": "OARs",
            "targets": "Targets",
            "AllCriteria": "All ROIs",
            "opt": "Optimization",
        }

        self.rois_plotting_order = {
            "oars": ["Brainstem", "Spinal cord", "Right parotid", "Left parotid", "Esophagus", "Larynx", "Mandible"],
            "targets": ["PTV56", "PTV63", "PTV70"],
        }

    def _get_all_rois(self):
        return sum(map(list, self.rois.values()), [])

    def _set_directories(self, io_name: str = "baseline"):
        # Define directory where given data is stored
        self.reference_data_dir = f"{self.primary_directory}/reference-plans"
        self.prediction_dir = f"{self.primary_directory}/paper-predictions"
        # path where any data generated by this code (e.g., predictions, models) are stored
        self.plans_dir = f"{self.primary_directory}/paper-plans"
        self.optimized_plans_dir = f"{self.plans_dir}/{io_name}"
        self.plan_dose_dir = f"{self.optimized_plans_dir}/plan-dose"
        self.plan_fluence_dir = f"{self.optimized_plans_dir}/plan-fluence"
        self.plan_gap_dir = f"{self.optimized_plans_dir}/plan-gap"
        self.plan_weights_dir = f"{self.optimized_plans_dir}/plan-weights"
        # Directories for summarizing results
        self.results_dir = f"{self.primary_directory}/results"
        self.results_data_dir = f"{self.primary_directory}/results-data"

        # Initialize attribute names for prediction specific directories
        self.plan_dose_from_pred_dir = None
        self.plan_fluence_from_pred_dir = None
        self.plan_gap_from_pred_dir = None

    def _set_plan_criteria(self):
        self.dvh_columns_better = [
            "('D_99', 'PTV56')",
            "('D_99', 'PTV63')",
            "('D_99', 'PTV70')",
            "('D_95', 'PTV56')",
            "('D_95', 'PTV63')",
            "('D_95', 'PTV70')",
        ]

    def _set_plot_fonts(self):
        self.ticks_font_size = 10
        self.legend_font_size = 10
        self.label_font_size = 10
        self.line_width = 6.1759606299

        plt.rcParams.update(
            {
                "axes.labelsize": self.label_font_size,
                "xtick.labelsize": self.ticks_font_size,
                "ytick.labelsize": self.ticks_font_size,
                "legend.fontsize": self.legend_font_size,
                "axes.titlesize": self.label_font_size,
                "font.family": "serif",
                "font.serif": ["Times New Roman"],
                "mathtext.fontset": "dejavuserif",
            }
        )

    def reset_plot_fonts(self):
        self._set_plot_fonts()

    def make_optimization_directories(self, prediction_name: str) -> None:
        """
        Makes the directories to store all of the data generated by the optimization process

        Args:
            prediction_name: name of the prediction set
        """
        # Define paths for set of directories
        self.get_optimization_directories(prediction_name)

        # Make directories
        os.makedirs(self.plan_dose_from_pred_dir, exist_ok=True)
        os.makedirs(self.plan_fluence_from_pred_dir, exist_ok=True)
        os.makedirs(self.plan_gap_from_pred_dir, exist_ok=True)
        os.makedirs(self.plan_weights_from_pred_dir, exist_ok=True)

    def get_optimization_directories(self, prediction_name: str, opt_name: str or None = None) -> None:
        """
        Makes the paths for all data, which should be made into directories
        Args:
            prediction_name: name of the prediction method
            opt_name: name of optimization method
        """
        if opt_name is not None:
            self._set_directories(opt_name)
        self.opt_name = opt_name
        self.prediction_name = prediction_name
        self.plan_dose_from_pred_dir = f"{self.plan_dose_dir}/{prediction_name}"
        self.plan_fluence_from_pred_dir = f"{self.plan_fluence_dir}/{prediction_name}"
        self.plan_gap_from_pred_dir = f"{self.plan_gap_dir}/{prediction_name}"
        self.plan_weights_from_pred_dir = f"{self.plan_weights_dir}/{prediction_name}"

    def set_patient(self, patient_number: str) -> None:
        """Set the current patient number/id"""
        self.patient_number = patient_number

    def check_patient(self) -> bool:
        """Check if a plan has already been generated for the requested prediction, prediction and patient"""
        return os.path.exists(f"{self.plan_fluence_from_pred_dir}/{self.patient_number}.csv")

    def set_directories_for_new_io(self, prediction_name: str, opt_name: str) -> None:
        """
        Set up the directory that will contain plans for the set of <prediction_name> dose predictions and the
        optimization model named <opt_name>

        Args:
            prediction_name: name of the prediction model that generated predictions
            opt_name:  name of the optimization model that generates plans
        """
        self._set_directories(opt_name)
        self.make_optimization_directories(prediction_name)