From 3310ccd923e1b8eb76823cfb34744534e3e340a7 Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 4 Nov 2026 13:42:21 -0500 Subject: [PATCH] fix(tasks): order readout probabilities by the task's labels behavior:choice followed Choice.labels but behavior:probabilities followed the readout classifier's own label order, so the two assemblies' .values disagreed column for column. --- brainscore_core/_model/behavioral.py | 13 ++++++++++--- tests/test_task_responses.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/brainscore_core/_model/behavioral.py b/brainscore_core/_model/behavioral.py index 7ad05e1a..841f08b9 100644 --- a/brainscore_core/_model/behavioral.py +++ b/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}) diff --git a/tests/test_task_responses.py b/tests/test_task_responses.py index 6db13331..372a19af 100644 --- a/tests/test_task_responses.py +++ b/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')])