Merge pull request #45 from biolab/apply_domian

[FIX] owcohorts: transform domain before computing
biolab · Apr 2, 2022 · 6b1d39f · 6b1d39f
2 parents 15ce516 + 954f9e3
commit 6b1d39f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/orangecontrib/survival_analysis/modeling/cox.py b/orangecontrib/survival_analysis/modeling/cox.py
@@ -11,6 +11,7 @@
     MISSING_SURVIVAL_DATA,
 )
 
+
 def to_data_frame(table: Table) -> pd.DataFrame:
     columns = table.domain.attributes + table.domain.class_vars
     df = pd.DataFrame({col.name: table.get_column_view(col)[0] for col in columns})

diff --git a/orangecontrib/survival_analysis/widgets/owchorts.py b/orangecontrib/survival_analysis/widgets/owchorts.py
@@ -42,11 +42,15 @@ def cutoff_by_log_rank_optimization(durations, events, callback, values):
     return cutoff
 
 
-def cox_risk_score(cox_model: CoxRegressionModel, data: Table):
+def cox_risk_score(cox_model: CoxRegressionModel, domain: Domain, data: Table):
+    data = data.transform(domain)
     return cox_model.predict(data.X)
 
 
-def stratify(stratify_on: ContinuousVariable, splitting_criteria: int, callback, data: Table):
+def stratify(
+    stratify_on: ContinuousVariable, splitting_criteria: int, domain: Domain, callback, data: Table
+):
+    data = data.transform(domain)
     stratify_on = stratify_on.compute_value(data)
 
     def check_unique_values(split_by, values):
@@ -117,12 +121,16 @@ def __init__(self, *args, **kwargs):
         )
 
         box = gui.vBox(self.controlArea, 'Splitting Criteria', margin=0)
-        self.radio_buttons = gui.radioButtons(box, self, 'splitting_criteria', callback=self.commit.deferred)
+        self.radio_buttons = gui.radioButtons(
+            box, self, 'splitting_criteria', callback=self.commit.deferred
+        )
 
         for _, opt in self.splitting_criteria_options:
             gui.appendRadioButton(self.radio_buttons, opt)
 
-        self.commit_button = gui.auto_commit(self.controlArea, self, 'auto_commit', '&Commit', box=False)
+        self.commit_button = gui.auto_commit(
+            self.controlArea, self, 'auto_commit', '&Commit', box=False
+        )
 
     @Inputs.learner
     def set_learner(self, learner: CoxRegressionLearner):
@@ -156,11 +164,15 @@ def callback():
         if self.stratify_on == StratifyOn.CoxRiskScore:
             cox_model = self.learner(data)
             _, risk_score_label = self.stratify_on_options[self.stratify_on]
-            risk_score_var = ContinuousVariable(risk_score_label, compute_value=partial(cox_risk_score, cox_model))
+            risk_score_var = ContinuousVariable(
+                risk_score_label, compute_value=partial(cox_risk_score, cox_model, data.domain)
+            )
             risk_group_var = DiscreteVariable(
                 'Cohorts',
                 values=['Low risk', 'High risk'],
-                compute_value=partial(stratify, risk_score_var, self.splitting_criteria, callback),
+                compute_value=partial(
+                    stratify, risk_score_var, self.splitting_criteria, data.domain, callback
+                ),
             )
 
             cohort_vars = (