fix(tasks): order readout probabilities by the task's labels

brainscore_core/_model/behavioral.py

@@ -344,7 +344,7 @@ def respond(self, stimuli, channels, regions=()):
                     [str(values['choice'].values[i]) for i in np.asarray(values.values).argmax(1)],
                     [True] * len(stimuli), stimuli, labels, 'readout', raw=None)
                 if self.probability_source == 'readout':
-                    outputs['behavior:probabilities'] = _implied(values, 'readout')
+                    outputs['behavior:probabilities'] = _implied(_in_order(values, labels), 'readout')
         else:
             score_first = path == 'generation' and self.probability_source == 'likelihood'
             remember = owner._memory == 'session'
@@ -408,8 +408,9 @@ def call(row, prompt):
                     outputs['behavior:probabilities'] = _implied_array(
                         _softmax_rows(scores, labels), stimuli, labels, 'likelihood')
                 elif self.probability_source == 'readout':
-                    outputs['behavior:probabilities'] = _implied(self.readout_classifier.predict_proba(
-                        owner._extract_behavioral_features(stimuli)), 'readout')
+                    outputs['behavior:probabilities'] = _implied(_in_order(
+                        self.readout_classifier.predict_proba(owner._extract_behavioral_features(stimuli)),
+                        labels), 'readout')
         if isinstance(response, Choice) and response.reaction_time:
             outputs['behavior:rt'] = _reaction_time(outputs['behavior:probabilities'],
                                                     response.reaction_time, stimuli)
@@ -791,6 +792,12 @@ def _implied_array(probabilities, stimuli, labels, source):
                               attrs={'measure': {'kind': 'implied', 'source': source}})
 
 
+def _in_order(assembly, labels):
+    """Readout columns in the task's label order, as ``behavior:choice`` has them."""
+    present = [label for label in labels if label in set(assembly['choice'].values)]
+    return assembly.sel(choice=present) if len(present) == assembly.sizes['choice'] else assembly
+
+
 def _implied(assembly, source):
     return assembly.assign_attrs(measure={'kind': 'implied', 'source': source})
 

tests/test_task_responses.py

@@ -118,6 +118,16 @@ def test_readout_choice_with_probabilities():
     assert session.collect('behavior:probabilities').values.shape == (4, 2)
 
 
+@pytest.mark.parametrize('labels', [['cat', 'dog'], ['dog', 'cat']])
+def test_readout_probabilities_follow_the_task_label_order(labels):
+    model, _ = _model(readout=True)
+    task = TaskContext(response=Choice(labels, probabilities=True), fitting_stimuli=STIMULI)
+    session = _ask(model, task, ['behavior:choice', 'behavior:probabilities'])
+    choice, probabilities = session.collect('behavior:choice'), session.collect('behavior:probabilities')
+    assert list(probabilities['choice'].values) == list(choice['choice'].values) == labels
+    np.testing.assert_array_equal(probabilities.values.argmax(1), choice.values.argmax(1))
+
+
 @pytest.mark.parametrize('response, channel', [(Text(), 'behavior:text'),
                                                (Choice(['cat', 'dog'], confidence=Scale(1, 5)),
                                                 'behavior:confidence')])