Add a bit of documentation

rdstar/mc_reweighting/BFReweighting/b_bf_reweighting.py

@@ -167,6 +167,11 @@ class BMesonBFReweighter:
         data_type_column_name: str,
         rescale_others: bool,
     ) -> None:
+        """A class to reweight B branching fractions.
+
+        :param data_type_column_name: Column that contains the data type (e.g. "Charged")
+        :param rescale_others: Apply correction factor for charged/mixed B production
+        """
         assert all([dt in signal_mc_bfs.keys() for dt in [v for vs in self.double_counting_cases.values() for v in vs]])
 
         flavour_keys = tuple(["b_charged", "b_mixed"])  # type: Tuple[str, ...]
@@ -193,6 +198,7 @@ class BMesonBFReweighter:
         return tuple([*self._required_cols, self.data_type_col])
 
     def add_bf_weights(self, df: pd.DataFrame) -> None:
+        """Adds BF weight to column defined in class attribute."""
         assert self.b_bf_weight_col_name not in df.columns, self.b_bf_weight_col_name
         assert all(c in df.columns for c in self.required_cols), [c for c in self.required_cols if c not in df.columns]
 
@@ -220,6 +226,10 @@ class BMesonBFReweighter:
         self._remove_special_gen_ids(df=df)
 
     def add_bf_weights_for_datatype(self, df: pd.DataFrame, data_type: str) -> None:
+        """If the whole dataframe belongs to this `data_type`, this will be the
+        same as add_bf_weights.
+        We still check that the data type column is set and agrees with `data_type`.
+        """
         assert self.b_bf_weight_col_name not in df.columns, self.b_bf_weight_col_name
         assert all(c in df.columns for c in self.required_cols), [c for c in self.required_cols if c not in df.columns]
 
@@ -289,6 +299,8 @@ class BMesonBFReweighter:
 
     @staticmethod
     def get_b_flavour_from_data_type(data_type: str) -> str:
+        """Helper function to infer b flavor from data type.
+        E.g. Charged -> b_charged. Raises RuntimeError if impossible."""
         flavour = get_b_flavour_from_data_type(data_type=data_type)  # type: str
         if flavour == "Charged":
             return "b_charged"
@@ -299,6 +311,7 @@ class BMesonBFReweighter:
 
     @staticmethod
     def add_neutral_weights(df: pd.DataFrame) -> None:
+        """Assigns unit weights to the weight column."""
         assert BMesonBFReweighter.b_bf_weight_col_name not in df.columns, BMesonBFReweighter.b_bf_weight_col_name
         assert all(col in df.columns for col in BMesonBFReweighter._required_cols)
 
@@ -306,6 +319,7 @@ class BMesonBFReweighter:
 
     @staticmethod
     def _get_neutral_weights(df: pd.DataFrame) -> np.ndarray:
+        """Assigns unit weights to the weight column, e. for non B events."""
         assert BMesonBFReweighter.b1_gen_id_col_name in df.columns, BMesonBFReweighter.b1_gen_id_col_name
         neutral_weights = np.ones_like(df[BMesonBFReweighter.b1_gen_id_col_name].values)  # type: np.ndarray
         assert len(neutral_weights) == len(np.unique(df.index)), (len(neutral_weights), len(np.unique(df.index)))
@@ -344,6 +358,9 @@ class BMesonBFReweighter:
         df.loc[~(_b2_do_not_reweighted | _b2_reweighted), self._own_b2_gen_id_col] = self.own_others_id
 
     def _remove_special_gen_ids(self, df: pd.DataFrame) -> None:
+        """Remove additional columns that were set during computation but
+        aren't required afterwards.
+        """
         assert self._own_b1_gen_id_col in df.columns, self._own_b1_gen_id_col
         assert self._own_b2_gen_id_col in df.columns, self._own_b2_gen_id_col