Make d columns configurable

Just a bit of cleanup/flexibility/catching another warning (that it raised when I used it with dask).

rdstar/mc_reweighting/BFReweighting/d_bf_reweighting.py

@@ -4,7 +4,6 @@ Tool for the Reweighting of Branching Fractions of D -> X
 Patrick Ecker 2020
 """
 
-import copy
 import logging
 import warnings
 import itertools
@@ -27,10 +26,8 @@ __all__ = [
 
 
 class DBFReweighter:
-    _required_cols = [
-        "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomdstIndex__bc__bc__bc",
-        "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomcMother__bomdstIndex__bc__bc__bc__bc",
-    ]  # type: List[str]
+    b_mdst_col = "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomdstIndex__bc__bc__bc"
+    b_mdst_tau_col = "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomcMother__bomdstIndex__bc__bc__bc__bc"
 
     _d_depth_conditions = {
         1: (
@@ -75,9 +72,9 @@ class DBFReweighter:
         self.signal_side_only = signal_side_only
         assert all(col in self.df.columns for col in DBFReweighter.get_required_columns())
 
-    @staticmethod
-    def get_required_columns() -> List[str]:
-        req_cols = copy.copy(DBFReweighter._required_cols)
+    @classmethod
+    def get_required_columns(cls) -> List[str]:
+        req_cols = [cls.b_mdst_col, cls.b_mdst_tau_col]
         req_cols.extend(
             [
                 c
@@ -110,10 +107,8 @@ class DBFReweighter:
         return self._d_depth_conditions[depth][b_mdst - 1]
 
     def _set_d_depth(self) -> None:
-        b_mdst_col = "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomdstIndex__bc__bc__bc"
-        b_mdst_tau_col = "daughter__bo1__cm__spdaughter__bo1__cm__spmcMother__bomcMother__bomdstIndex__bc__bc__bc__bc"
-        required_cols = [b_mdst_col, b_mdst_tau_col]
-        assert all(col in self._required_cols for col in required_cols), required_cols
+        required_cols = [self.b_mdst_col, self.b_mdst_tau_col]
+        assert all(col in self.get_required_columns() for col in required_cols), required_cols
 
         # define temporary columns
         temp_b1_d_depth_col = "b1_d_depth"
@@ -162,12 +157,12 @@ class DBFReweighter:
                 assert not pd.isnull(self.df[col]).any(), col
         else:
             # Now check for the Signal B Meson
-            cond = self.df[b_mdst_col].isin([1, 2])
-            cond_tau = self.df[b_mdst_tau_col].isin([1, 2])
+            cond = self.df[self.b_mdst_col].isin([1, 2])
+            cond_tau = self.df[self.b_mdst_tau_col].isin([1, 2])
 
             self.df.loc[:, self._temp_sig_b_mdst_index_col] = np.nan
-            self.df.loc[cond, self._temp_sig_b_mdst_index_col] = self.df[b_mdst_col]
-            self.df.loc[cond_tau, self._temp_sig_b_mdst_index_col] = self.df[b_mdst_tau_col]
+            self.df.loc[cond, self._temp_sig_b_mdst_index_col] = self.df[self.b_mdst_col]
+            self.df.loc[cond_tau, self._temp_sig_b_mdst_index_col] = self.df[self.b_mdst_tau_col]
             self.df.loc[(~cond) & (~cond_tau), self._temp_sig_b_mdst_index_col] = 0
             assert not pd.isnull(self.df[self._temp_sig_b_mdst_index_col]).any(), self._temp_sig_b_mdst_index_col
 
@@ -236,8 +231,11 @@ class DBFReweighter:
                 # Actual command:
                 self.df["temp"] = list(zip(*[self.df["temp" + var] for var in decay_vars]))
                 for w in warning:
-                    assert issubclass(w.category, SettingWithCopyWarning)
-                    assert "A value is trying to be set on a copy" in str(w.message)
+                    is_setting_with_copy_warning = issubclass(w.category, SettingWithCopyWarning)
+                    is_setting_with_copy_warning &= "A value is trying to be set on a copy" in str(w.message)
+                    is_deprecation_warning = issubclass(w.category, DeprecationWarning)
+                    if not is_setting_with_copy_warning | is_deprecation_warning:
+                        raise ValueError(f"Caught warning of category {w.category} and message {w.message}")
 
             self.df.loc[conds, decay_col] = self.df["temp"]
             for var in decay_vars: