Skip to content
Snippets Groups Projects

Add a bit of documentation

Open Kilian Lieret requested to merge add-documentation-b-bf-reweighting into master
1 file
+ 17
0
Compare changes
  • Side-by-side
  • Inline
@@ -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
Loading