fix mixins of custom tasks (I hope)

uhh-cms · Nov 25, 2024 · 5defd46 · 5defd46
1 parent 09a4ab4
commit 5defd46
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 31 deletions.
diff --git a/hbw/tasks/ml.py b/hbw/tasks/ml.py
@@ -22,7 +22,6 @@
     CalibratorsMixin,
     ProducersMixin,
     MLModelTrainingMixin,
-    MLModelMixin,
     MLModelsMixin,
     MLModelDataMixin,
     SelectorStepsMixin,
@@ -403,13 +402,9 @@ def workflow_run(self):
 
 
 class MLEvaluationSingleFold(
-    # NOTE: this should probably be a MLModelTrainingMixin, but I'll postpone this until the MultiConfigTask
-    # is implemented
+    # NOTE: mixins might need fixing, needs to be checked
     HBWTask,
-    MLModelMixin,
-    ProducersMixin,
-    SelectorMixin,
-    CalibratorsMixin,
+    MLModelTrainingMixin,
     law.LocalWorkflow,
     RemoteWorkflow,
 ):
@@ -541,13 +536,9 @@ def run(self):
 
 
 class PlotMLResultsSingleFold(
-    # NOTE: this should probably be a MLModelTrainingMixin, but I'll postpone this until the MultiConfigTask
-    # is implemented
+    # NOTE: mixins might need fixing, needs to be checked
     HBWTask,
-    MLModelMixin,
-    ProducersMixin,
-    SelectorMixin,
-    CalibratorsMixin,
+    MLModelTrainingMixin,
     law.LocalWorkflow,
     RemoteWorkflow,
 ):

diff --git a/hbw/tasks/optimization.py b/hbw/tasks/optimization.py
@@ -14,10 +14,11 @@
 
 from columnflow.tasks.framework.base import Requirements
 from columnflow.tasks.framework.mixins import (
-    SelectorMixin,
-    CalibratorsMixin,
-    ProducersMixin,
-    MLModelMixin,
+    # SelectorMixin,
+    # CalibratorsMixin,
+    # ProducersMixin,
+    # MLModelMixin,
+    MLModelTrainingMixin,
 )
 # from columnflow.tasks.framework.remote import RemoteWorkflow
 from columnflow.util import DotDict
@@ -97,12 +98,11 @@ def run(self):
 
 
 class Optimizer(
+    # NOTE: mixins might need fixing, needs to be tested
     HBWTask,
-    MLModelMixin,
-    ProducersMixin,
-    SelectorMixin,
-    CalibratorsMixin,
+    MLModelTrainingMixin,
     law.LocalWorkflow,
+    # RemoteWorkflow,
 ):
     """
     Workflow that runs optimization. Needs to be run from within the sandbox
@@ -191,12 +191,11 @@ def run(self):
 
 
 class Objective(
+    # NOTE: mixins might need fixing, needs to be tested
     HBWTask,
-    MLModelMixin,
-    ProducersMixin,
-    SelectorMixin,
-    CalibratorsMixin,
+    MLModelTrainingMixin,
     law.LocalWorkflow,
+    # RemoteWorkflow,
 ):
     """
     Objective to optimize.
@@ -281,11 +280,9 @@ def run(self):
 
 
 class DummyObjective(
+    # NOTE: mixins might need fixing, needs to be tested
     HBWTask,
-    MLModelMixin,
-    ProducersMixin,
-    SelectorMixin,
-    CalibratorsMixin,
+    MLModelTrainingMixin,
     law.LocalWorkflow,
     # RemoteWorkflow,
 ):

diff --git a/hbw/tasks/plotting.py b/hbw/tasks/plotting.py
@@ -10,13 +10,14 @@
 import law
 import order as od
 
-from columnflow.tasks.framework.base import Requirements, ShiftTask
+from columnflow.tasks.framework.base import Requirements, MultiConfigTask
 from columnflow.tasks.framework.mixins import (
     CalibratorsMixin, SelectorStepsMixin, ProducersMixin, MLModelsMixin,
     CategoriesMixin,
 )
 from columnflow.tasks.framework.plotting import (
     PlotBase, PlotBase1D, ProcessPlotSettingMixin, VariablePlotSettingMixin,
+    PlotShiftMixin,
 )
 from columnflow.tasks.framework.decorators import view_output_plots
 from columnflow.tasks.framework.remote import RemoteWorkflow
@@ -108,7 +109,7 @@ def plot_multi_weight_producer(
 
 class PlotVariablesMultiWeightProducer(
     HBWTask,
-    ShiftTask,
+    PlotShiftMixin,
     VariablePlotSettingMixin,
     ProcessPlotSettingMixin,
     PlotBase1D,
@@ -117,6 +118,7 @@ class PlotVariablesMultiWeightProducer(
     ProducersMixin,
     SelectorStepsMixin,
     CalibratorsMixin,
+    MultiConfigTask,
     law.LocalWorkflow,
     RemoteWorkflow,
 ):

diff --git a/hbw/tasks/postfit_plots.py b/hbw/tasks/postfit_plots.py
@@ -12,6 +12,7 @@
 import luigi
 import order as od
 
+from columnflow.tasks.framework.base import ConfigTask
 from columnflow.tasks.framework.mixins import (
     InferenceModelMixin, MLModelsMixin, ProducersMixin, SelectorStepsMixin,
     CalibratorsMixin,
@@ -107,6 +108,7 @@ def plot_postfit_shapes(
 
 
 class PlotPostfitShapes(
+    # NOTE: mixins might be wrong and could (should?) be extended to MultiConfigTask
     HBWTask,
     PlotBase1D,
     # to correctly setup our InferenceModel, we need all these mixins, but hopefully, all these
@@ -116,6 +118,7 @@ class PlotPostfitShapes(
     ProducersMixin,
     SelectorStepsMixin,
     CalibratorsMixin,
+    ConfigTask,
 ):
     """
     Task that creates Postfit shape plots based on a fit_diagnostics file.