FIX: Select reference resolution based on grayordinates

nibabies/tests/test_config.py

@@ -104,6 +104,23 @@ def test_config_spaces():
     ] == ['MNIInfant_cohort-1_res-native', 'MNI152NLin6Asym_res-2', 'MNIInfant_cohort-1_res-2']
     _reset_config()
 
+    config.execution.output_spaces = None
+    config.workflow.cifti_output = '170k'
+    spaces = _load_spaces(1)
+
+    assert [str(s) for s in spaces.get_standard(full_spec=True)] == [
+        'MNIInfant:cohort-1:res-native',  # Default output space
+        'MNI152NLin6Asym:res-1',
+        'MNIInfant:cohort-1:res-1',
+    ]
+
+    assert [
+        format_reference((s.fullname, s.spec))
+        for s in spaces.references
+        if s.standard and s.dim == 3
+    ] == ['MNIInfant_cohort-1_res-native', 'MNI152NLin6Asym_res-1', 'MNIInfant_cohort-1_res-1']
+    _reset_config()
+
 
 @pytest.mark.parametrize(
     ('master_seed', 'ants_seed', 'numpy_seed'), [(1, 17612, 8272), (100, 19094, 60232)]

nibabies/workflows/base.py

@@ -515,10 +515,12 @@ def init_single_subject_wf(
         ])  # fmt:skip
 
         if cifti_output and 'MNIInfant' in [ref.space for ref in spaces.references]:
+            mniinfant_res = 2 if config.workflow.cifti_output == '91k' else 1
+
             select_MNIInfant_xfm = pe.Node(
                 KeySelect(
                     fields=['anat2std_xfm', 'std2anat_xfm'],
-                    key=get_MNIInfant_key(spaces),
+                    key=get_MNIInfant_key(spaces, mniinfant_res),
                 ),
                 name='select_MNIInfant_xfm',
                 run_without_submitting=True,
@@ -840,7 +842,7 @@ def init_workflow_spaces(execution_spaces: SpatialReferences, age_months: int):
         spaces.add(Reference('MNI152NLin6Asym', {'res': vol_res}))
         # Ensure a non-native version of MNIInfant is added as a target
         cohort = cohort_by_months('MNIInfant', age_months)
-        spaces.add(Reference('MNIInfant', {'cohort': cohort, 'res': 2}))
+        spaces.add(Reference('MNIInfant', {'cohort': cohort, 'res': vol_res}))
 
     return spaces
 
@@ -950,15 +952,10 @@ def get_estimator(layout, fname):
     return field_source
 
 
-def get_MNIInfant_key(spaces: SpatialReferences) -> str:
+def get_MNIInfant_key(spaces: SpatialReferences, res: str | int) -> str:
     """Parse spaces and return matching MNIInfant space, including cohort."""
-    key = None
-    for space in spaces.references:
-        # str formats as <reference.name>:<reference.spec>
-        if 'MNIInfant' in str(space) and 'res-2' in str(space):
-            key = str(space)
-            break
-
-    if key is None:
-        raise KeyError(f'MNIInfant (resolution 2x2x2) not found in SpatialReferences: {spaces}')
-    return key
+    for ref in spaces.references:
+        if ref.space == 'MNIInfant' and f'res-{res}' in str(ref):
+            return ref.fullname
+
+    raise KeyError(f'MNIInfant (resolution {res}) not found in SpatialReferences: {spaces}')

nibabies/workflows/bold/base.py

@@ -558,7 +558,8 @@ def init_bold_wf(
             ]),
         ])  # fmt:skip
 
-    if config.workflow.cifti_output:
+    cifti_output = config.workflow.cifti_output
+    if cifti_output:
         from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
 
         from nibabies.workflows.bold.alignment import (
@@ -581,7 +582,7 @@ def init_bold_wf(
         )
 
         bold_fsLR_resampling_wf = init_bold_fsLR_resampling_wf(
-            grayord_density=config.workflow.cifti_output,
+            grayord_density=cifti_output,
             omp_nthreads=omp_nthreads,
             mem_gb=mem_gb['resampled'],
         )
@@ -615,7 +616,7 @@ def init_bold_wf(
         subcortical_mni_alignment_wf = init_subcortical_mni_alignment_wf()
 
         bold_grayords_wf = init_bold_grayords_wf(
-            grayord_density=config.workflow.cifti_output,
+            grayord_density=cifti_output,
             repetition_time=all_metadata[0]['RepetitionTime'],
         )
 
@@ -624,7 +625,7 @@ def init_bold_wf(
                 base_directory=output_dir,
                 dismiss_entities=DEFAULT_DISMISS_ENTITIES,
                 space='fsLR',
-                density=config.workflow.cifti_output,
+                density=cifti_output,
                 suffix='bold',
                 compress=False,
                 TaskName=all_metadata[0].get('TaskName'),
@@ -635,7 +636,8 @@ def init_bold_wf(
         )
         ds_bold_cifti.inputs.source_file = bold_file
 
-        inputnode.inputs.mniinfant_mask = get_MNIInfant_mask(spaces)
+        mniinfant_res = 2 if cifti_output == '91k' else 1
+        inputnode.inputs.mniinfant_mask = get_MNIInfant_mask(spaces, mniinfant_res)
 
         workflow.connect([
             # Resample BOLD to MNI152NLin6Asym, may duplicate bold_std_wf above
@@ -747,11 +749,11 @@ def init_bold_wf(
     ])  # fmt:skip
 
     # MG: Carpetplot workflow only work with CIFTI
-    if config.workflow.cifti_output:
+    if cifti_output:
         carpetplot_wf = init_carpetplot_wf(
             mem_gb=mem_gb['resampled'],
             metadata=all_metadata[0],
-            cifti_output=config.workflow.cifti_output,
+            cifti_output=cifti_output,
             name='carpetplot_wf',
         )
 
@@ -847,24 +849,20 @@ def _read_json(in_file):
     return loads(Path(in_file).read_text())
 
 
-def get_MNIInfant_mask(spaces: 'SpatialReferences') -> str:
+def get_MNIInfant_mask(spaces: 'SpatialReferences', res: str | int) -> str:
     """Parse spaces and return matching MNIInfant space, including cohort."""
     import templateflow.api as tf
 
-    mask = None
     for ref in spaces.references:
-        # str formats as <reference.name>:<reference.spec>
-        if ref.space == 'MNIInfant' and ref.spec.get('res', '') != 'native':
-            mask = str(
+        if ref.space == 'MNIInfant' and f'res-{res}' in str(ref):
+            return str(
                 tf.get(
                     'MNIInfant',
                     cohort=ref.spec['cohort'],
-                    resolution=1,
+                    resolution=res,
                     desc='brain',
                     suffix='mask',
                 )
             )
 
-    if mask is None:
-        raise FileNotFoundError('MNIInfant brain mask not found.')
-    return mask
+    raise FileNotFoundError(f'MNIInfant mask (resolution {res}) not found.')