Initial OW patient ranking pipeline

align_system/algorithms/alignment_adm_component.py

@@ -7,6 +7,7 @@
 
 log = logging.getLogger(__name__)
 med_urg_str = "medical"
+attr_str = "attribute"
 
 
 class AlignmentADMComponent(ADMComponent):
@@ -57,7 +58,7 @@ def run(self,
 
 class MedicalOnlyAlignmentADMComponent(ADMComponent):
     def run_returns(self):
-        return ('chosen_choice', 'best_sample_idx')
+        return ('chosen_choice', 'best_sample_idx', 'medical_urgency_info')
 
     def run(
         self,
@@ -110,7 +111,11 @@ def _get_best_sample_idx(average_urgency, samples):
 
             return best_idx
 
-        return (selected_choice, _get_best_sample_idx(max_urg, attribute_prediction_scores[selected_choice]))
+        return (
+            selected_choice,
+            _get_best_sample_idx(max_urg, attribute_prediction_scores[selected_choice]),
+            med_urg,
+        )
 
 
 def _handle_single_value(predictions):
@@ -496,6 +501,52 @@ def _composite_probs(self, p_matrix):
         scores = np.sum(log_odds, axis=1)
         return self._stable_softmax(scores)
 
+    def run_returns(self):
+        return ('chosen_choice', 'best_sample_idx', 'p_choices', 'alignment_info')
+
+    def _compute_p_choose_a(
+        self, kdma, intercept, medical_weight, attr_weight, opt_a, opt_b,
+    ):
+        # Provided by ADEPT 2025-12-12
+        # MF updated 2026-01-21
+        scaling = {
+                "affiliation": {
+                    med_urg_str: [0.589, 0.330],
+                    attr_str: [0.703, 0.365],
+                },
+                "merit": {
+                    med_urg_str: [0.576, 0.339],
+                    attr_str: [0.671, 0.381],
+                },
+                "personal_safety": {
+                    med_urg_str: [0.228, 0.287],
+                    attr_str: [0.777, 0.309],
+                },
+                "search": {
+                    med_urg_str: [0.263, 0.365],
+                    attr_str: [0.286, 0.325],
+                },
+            }
+        if kdma not in scaling:
+            raise RuntimeError(f"No z-scaling values provided for {kdma}")
+        scaling = scaling[kdma]
+
+        def _apply_z_scaling(key, raw_value):
+            return (raw_value - scaling[key][0]) / scaling[key][1]
+
+        # Apply z-scaling
+        a_med = _apply_z_scaling(med_urg_str, opt_a[med_urg_str])
+        a_attr = _apply_z_scaling(attr_str, opt_a[kdma])
+        b_med = _apply_z_scaling(med_urg_str, opt_b[med_urg_str])
+        b_attr = _apply_z_scaling(attr_str, opt_b[kdma])
+
+        medical_delta = a_med - b_med
+        attr_score = a_attr - b_attr
+
+        # Compute p_choose_a
+        y_ij = intercept + medical_weight*medical_delta + attr_weight*attr_score
+        return math.exp(y_ij) / (1 + math.exp(y_ij))
+
     def run(
         self,
         attribute_prediction_scores,
@@ -516,12 +567,14 @@ def run(
 
         # Only one option, decision always has to be the same
         if len(choices) == 1:
+            p_choices = np.ones((1,)).tolist()
             return (
                 predictions[0]["choice"],
                 0,  # TODO: best sample index
+                p_choices,
                 {
                     "source": type(self).__name__,
-                    "p_choices": np.ones((1,)),
+                    "p_choices": p_choices,
                 },
             )
 
@@ -557,10 +610,9 @@ def run(
                         flip_order = True
                         raw_medical_delta *= -1
                     primary, secondary = (opt_a, opt_b) if not flip_order else (opt_b, opt_a)
-                    raw_attr_score = secondary[kdma] if kdma == "search" else primary[kdma]
 
                     p_choose_primary = self._compute_p_choose_a(
-                        kdma, intercept, medical_weight, attr_weight, raw_medical_delta, raw_attr_score)
+                        kdma, intercept, medical_weight, attr_weight, primary, secondary)
 
                     p_matrix[choice_idx_a][choice_idx_b] = p_choose_primary if not flip_order else 1 - p_choose_primary
                     p_matrix[choice_idx_b][choice_idx_a] = 1 - p_choose_primary if not flip_order else p_choose_primary
@@ -570,11 +622,14 @@ def run(
             # TODO: Figure out what it means to be the best prediction for this alignment function
             best_sample_idx = 0
 
+            max_idx = np.argmax(p_choices)
+            p_choices = p_choices.tolist()
+
             alignment_info = {
                 "source": type(self).__name__,
-                "p_choices": p_choices.tolist(),
+                "p_choices": p_choices,
             }
 
-            max_idx = np.argmax(p_choices)
 
-            return (predictions[max_idx]["choice"], best_sample_idx, alignment_info)
+
+            return (predictions[max_idx]["choice"], best_sample_idx, p_choices, alignment_info)