Don't export excluded annotations if not in config

Author: adelavega

autonima.egg-info/SOURCES.txt

@@ -60,6 +60,7 @@ tests/test_multi_annotation.py
 tests/test_objective_in_prompt.py
 tests/test_parallel_screening.py
 tests/test_pipeline_retrieval.py
+tests/test_pipeline_stage_cutoff.py
 tests/test_pubmed.py
 tests/test_retrieval.py
 tests/test_screening.py

autonima/coordinates/nimads_models.py

@@ -340,16 +340,32 @@ def create_annotations_from_results(
     """
     if not annotation_results:
         return []
-    
+
+    # Keep annotations aligned to the exported NiMADS studyset.
+    # This avoids emitting notes for analyses that were filtered out
+    # (e.g., excluded studies when export_excluded_studies is false).
+    studyset_analysis_ids = {
+        analysis.id
+        for study in studyset.studies
+        for analysis in study.analyses
+    }
+
     # Group annotation results by annotation name and analysis_id
     annotations_by_name = {}
-    all_analysis_ids = set()
+    annotation_names = set()
 
     for result in annotation_results:
+        annotation_names.add(result.annotation_name)
+        if result.analysis_id not in studyset_analysis_ids:
+            continue
         if result.annotation_name not in annotations_by_name:
             annotations_by_name[result.annotation_name] = {}
         annotations_by_name[result.annotation_name][result.analysis_id] = result.include
-        all_analysis_ids.add(result.analysis_id)
+    
+    # Ensure all annotation columns are represented in note_keys, even if
+    # some analyses were filtered out.
+    for annotation_name in annotation_names:
+        annotations_by_name.setdefault(annotation_name, {})
 
     # Create note_keys with all annotation names as boolean columns
     note_keys = {}
@@ -366,8 +382,8 @@ def create_annotations_from_results(
         studyset_id=studyset_id
     )
 
-    # Create notes for each analysis with all annotation decisions
-    for analysis_id in sorted(all_analysis_ids):
+    # Create notes for each exported analysis with all annotation decisions.
+    for analysis_id in sorted(studyset_analysis_ids):
         note_data = {}
 
         # For each annotation type, check if this analysis was included
@@ -511,4 +527,4 @@ def sanitize_annotation_dict(annotation_dict: dict) -> dict:
             if 'analysis' in note and note['analysis']:
                 # Sanitize the analysis ID (which may contain quotes from the name)
                 note['analysis'] = sanitize_analysis_name(note['analysis'])
-    return annotation_dict
+    return annotation_dict

tests/test_nimads_annotations.py

@@ -0,0 +1,69 @@
+from autonima.annotation.schema import AnnotationDecision
+from autonima.coordinates.nimads_models import (
+    convert_to_nimads_studyset,
+    create_annotations_from_results,
+)
+from autonima.coordinates.schema import Analysis, CoordinatePoint
+from autonima.models.types import Study, StudyStatus
+
+
+def _make_study(pmid: str, status: StudyStatus) -> Study:
+    return Study(
+        pmid=pmid,
+        title=f"Study {pmid}",
+        abstract="Abstract",
+        authors=["Doe"],
+        journal="Journal",
+        publication_date="2020",
+        status=status,
+        analyses=[
+            Analysis(
+                name="contrast",
+                description="desc",
+                table_id="table_1",
+                points=[
+                    CoordinatePoint(
+                        coordinates=[1.0, 2.0, 3.0],
+                        space="MNI",
+                        values=[],
+                    )
+                ],
+            )
+        ],
+    )
+
+
+def test_create_annotations_from_results_filters_out_non_studyset_analyses():
+    included = _make_study("100", StudyStatus.INCLUDED_FULLTEXT)
+    excluded = _make_study("200", StudyStatus.EXCLUDED_FULLTEXT)
+
+    # Simulate export_excluded_studies: false by building the studyset from
+    # only included studies.
+    studyset = convert_to_nimads_studyset("studyset_1", [included])
+
+    annotation_results = [
+        AnnotationDecision(
+            annotation_name="risky_dm",
+            analysis_id="100_analysis_0",
+            study_id="100",
+            include=True,
+            reasoning="included study analysis",
+            model_used="test-model",
+        ),
+        AnnotationDecision(
+            annotation_name="risky_dm",
+            analysis_id="200_analysis_0",
+            study_id="200",
+            include=True,
+            reasoning="excluded study analysis",
+            model_used="test-model",
+        ),
+    ]
+
+    annotation = create_annotations_from_results(
+        "studyset_1", studyset, annotation_results
+    )
+
+    exported_analysis_ids = {note.analysis_id for note in annotation.notes}
+    assert exported_analysis_ids == {"100_analysis_0"}
+    assert annotation.notes[0].note["risky_dm"] is True