From 07e904fa3efff4d9bff08077c80f5c0890042d55 Mon Sep 17 00:00:00 2001
From: Taylor Salo <salot@pennmedicine.upenn.edu>
Date: Tue, 3 Sep 2024 12:15:29 -0400
Subject: [PATCH] Write out denoised data to requested output spaces (#44)

---
 src/fmripost_aroma/cli/parser.py     | 33 +++++++------
 src/fmripost_aroma/workflows/base.py | 70 +++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 18 deletions(-)

diff --git a/src/fmripost_aroma/cli/parser.py b/src/fmripost_aroma/cli/parser.py
index 1603698..a14f046 100644
--- a/src/fmripost_aroma/cli/parser.py
+++ b/src/fmripost_aroma/cli/parser.py
@@ -37,12 +37,11 @@ def _build_parser(**kwargs):
     from functools import partial
     from pathlib import Path
 
+    from niworkflows.utils.spaces import OutputReferencesAction
     from packaging.version import Version
 
     from fmripost_aroma.cli.version import check_latest, is_flagged
 
-    # from niworkflows.utils.spaces import OutputReferencesAction
-
     class ToDict(Action):
         def __call__(self, parser, namespace, values, option_string=None):
             d = {}
@@ -326,21 +325,21 @@ def _bids_filter(value, parser):
         ),
     )
     # Disable output spaces until warping works
-    # g_conf.add_argument(
-    #     '--output-spaces',
-    #     nargs='*',
-    #     action=OutputReferencesAction,
-    #     help="""\
-    # Standard and non-standard spaces to resample denoised functional images to. \
-    # Standard spaces may be specified by the form \
-    # ``<SPACE>[:cohort-<label>][:res-<resolution>][...]``, where ``<SPACE>`` is \
-    # a keyword designating a spatial reference, and may be followed by optional, \
-    # colon-separated parameters. \
-    # Non-standard spaces imply specific orientations and sampling grids. \
-    # For further details, please check out \
-    # https://fmriprep.readthedocs.io/en/%s/spaces.html"""
-    #    % (currentv.base_version if is_release else 'latest'),
-    # )
+    g_conf.add_argument(
+        '--output-spaces',
+        nargs='*',
+        action=OutputReferencesAction,
+        help="""\
+Standard and non-standard spaces to resample denoised functional images to. \
+Standard spaces may be specified by the form \
+``<SPACE>[:cohort-<label>][:res-<resolution>][...]``, where ``<SPACE>`` is \
+a keyword designating a spatial reference, and may be followed by optional, \
+colon-separated parameters. \
+Non-standard spaces imply specific orientations and sampling grids. \
+For further details, please check out \
+https://fmriprep.readthedocs.io/en/%s/spaces.html"""
+        % (currentv.base_version if is_release else 'latest'),
+    )
     g_conf.add_argument(
         '--dummy-scans',
         required=False,
diff --git a/src/fmripost_aroma/workflows/base.py b/src/fmripost_aroma/workflows/base.py
index 3cf831b..b0a5540 100644
--- a/src/fmripost_aroma/workflows/base.py
+++ b/src/fmripost_aroma/workflows/base.py
@@ -36,6 +36,7 @@
 
 import yaml
 from nipype.pipeline import engine as pe
+from niworkflows.interfaces.utility import KeySelect
 from packaging.version import Version
 
 from fmripost_aroma import config
@@ -488,8 +489,20 @@ def init_single_run_wf(bold_file):
     func_fit_reports_wf.inputs.inputnode.anat_dseg = functional_cache['anat_dseg']
     workflow.connect([(mni6_buffer, func_fit_reports_wf, [('bold', 'inputnode.bold_mni6')])])
 
-    if config.workflow.denoise_method:
+    if config.workflow.denoise_method and spaces.get_spaces():
         # Now denoise the output-space BOLD data using ICA-AROMA
+        from smriprep.workflows.outputs import init_template_iterator_wf
+
+        templates = spaces.get_spaces()
+        template_iterator_wf = init_template_iterator_wf(
+            spaces=spaces,
+            sloppy=config.execution.sloppy,
+        )
+        template_iterator_wf.inputs.inputnode.anat2std_xfm = functional_cache[
+            'anat2outputspaces_xfm'
+        ]
+        template_iterator_wf.inputs.inputnode.template = templates
+
         denoise_wf = init_denoise_wf(bold_file=bold_file, metadata=bold_metadata)
         denoise_wf.inputs.inputnode.skip_vols = skip_vols
         denoise_wf.inputs.inputnode.space = 'MNI152NLin6Asym'
@@ -505,8 +518,63 @@ def init_single_run_wf(bold_file):
                 ('outputnode.mixing', 'inputnode.mixing'),
                 ('outputnode.aroma_features', 'inputnode.classifications'),
             ]),
+            (template_iterator_wf, denoise_wf, [
+                ('outputnode.space', 'inputnode.space'),
+                ('outputnode.cohort', 'inputnode.cohort'),
+                ('outputnode.res', 'inputnode.res'),
+            ]),
         ])  # fmt:skip
 
+        if functional_cache['bold_outputspaces']:
+            # No transforms necessary
+            std_buffer = pe.Node(
+                KeySelect(
+                    fields=['bold', 'bold_mask'],
+                    keys=[str(space) for space in spaces.references],
+                ),
+                name='std_buffer',
+            )
+            std_buffer.inputs.bold = functional_cache['bold_outputspaces']
+            std_buffer.inputs.bold_mask = functional_cache['bold_mask_outputspaces']
+            workflow.connect([
+                (template_iterator_wf, std_buffer, [('outputnode.space', 'key')]),
+                (std_buffer, denoise_wf, [
+                    ('bold', 'inputnode.bold_file'),
+                    ('bold_mask', 'inputnode.bold_mask'),
+                ]),
+            ])  # fmt:skip
+        else:
+            # Warp native BOLD to requested output spaces
+            xfms = [
+                functional_cache['hmc'],
+                functional_cache['boldref2fmap'],
+                functional_cache['bold2anat'],
+            ]
+            all_xfms = pe.Node(niu.Merge(2), name='all_xfms')
+            all_xfms.inputs.in1 = xfms
+            workflow.connect([
+                (template_iterator_wf, all_xfms, [('outputnode.anat2std_xfm', 'in2')]),
+            ])  # fmt:skip
+
+            resample_std_wf = init_resample_volumetric_wf(
+                bold_file=bold_file,
+                functional_cache=functional_cache,
+                run_stc=False,
+                name=_get_wf_name(bold_file, 'resample_std'),
+            )
+            workflow.connect([
+                (template_iterator_wf, resample_std_wf, [
+                    ('outputnode.space', 'inputnode.space'),
+                    ('outputnode.res', 'inputnode.res'),
+                    ('outputnode.cohort', 'inputnode.cohort'),
+                ]),
+                (all_xfms, resample_std_wf, [('out', 'inputnode.transforms')]),
+                (resample_std_wf, denoise_wf, [
+                    ('outputnode.bold_std', 'inputnode.bold'),
+                    ('outputnode.bold_mask_std', 'inputnode.bold_mask'),
+                ]),
+            ])  # fmt:skip
+
     # Fill-in datasinks seen so far
     for node in workflow.list_node_names():
         if node.split('.')[-1].startswith('ds_'):