Merge pull request #7 from nipype/develop

updated auto-generated code with function tasks

Author: tclose

pydra/tasks/mriqc/interfaces/bids/iqm_file_sink.py

@@ -73,7 +73,7 @@ def function(
         fields = list(set(fields) - set(self_dict["inputs"].copyable_trait_names()))
         self_dict["_input_names"] = fields
         undefined_traits = {
-            key: _add_field(key, add_trait=add_trait, _outputs=_outputs)
+            key: _add_field(key, _outputs=_outputs, add_trait=add_trait)
             for key in fields
         }
         self_dict["inputs"].trait_set(trait_change_notify=False, **undefined_traits)
@@ -82,7 +82,7 @@ def function(
             self_dict["_always_run"] = True
         self_dict = {}
         out_file = _gen_outfile(
-            in_file=in_file, out_dir=out_dir, dismiss_entities=dismiss_entities
+            dismiss_entities=dismiss_entities, out_dir=out_dir, in_file=in_file
         )
 
         if root is not attrs.NOTHING:
@@ -151,14 +151,14 @@ def function(
         return out_file
 
 
-def _add_field(name, value=attrs.NOTHING, add_trait=None, _outputs=None):
+def _add_field(name, value=attrs.NOTHING, _outputs=None, add_trait=None):
     self_dict = {}
     self_dict["inputs"].add_trait(name, traits.Any)
     _outputs[name] = value
     return value
 
 
-def _gen_outfile(in_file=None, out_dir=None, dismiss_entities=None):
+def _gen_outfile(dismiss_entities=None, out_dir=None, in_file=None):
     out_dir = Path()
     if out_dir is not attrs.NOTHING:
         out_dir = Path(out_dir)

pydra/tasks/mriqc/workflows/anatomical/base.py

@@ -51,17 +51,7 @@ def anat_qc_workflow(
     wf_inputs_metadata=None,
     wf_species="human",
     wf_template_id="MNI152NLin2009cAsym",
-) -> [
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-]:
+) -> tuple[ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any]:
     """
     One-subject-one-session-one-run pipeline to extract the NR-IQMs from
     anatomical images
@@ -151,10 +141,10 @@ def anat_qc_workflow(
     norm = workflow.add(
         spatial_normalization(
             wf_species=wf_species,
+            nipype_omp_nthreads=nipype_omp_nthreads,
             wf_template_id=wf_template_id,
             exec_ants_float=exec_ants_float,
             exec_debug=exec_debug,
-            nipype_omp_nthreads=nipype_omp_nthreads,
             modality=modality,
             name="norm",
         )
@@ -196,9 +186,9 @@ def anat_qc_workflow(
     # Reports
     anat_report_wf = workflow.add(
         init_anat_report_wf(
-            wf_species=wf_species,
             exec_verbose_reports=exec_verbose_reports,
             exec_work_dir=exec_work_dir,
+            wf_species=wf_species,
             in_ras=to_ras.out_file,
             headmask=hmsk.out_file,
             airmask=amw.air_mask,
@@ -235,12 +225,12 @@ def anat_qc_workflow(
         # fmt: on
     outputs_["norm_report"] = norm.report
     outputs_["iqmswf_noise_report"] = iqmswf.noise_report
+    outputs_["anat_report_wf_airmask_report"] = anat_report_wf.airmask_report
+    outputs_["anat_report_wf_bg_report"] = anat_report_wf.bg_report
     outputs_["anat_report_wf_bmask_report"] = anat_report_wf.bmask_report
     outputs_["anat_report_wf_artmask_report"] = anat_report_wf.artmask_report
-    outputs_["anat_report_wf_headmask_report"] = anat_report_wf.headmask_report
-    outputs_["anat_report_wf_bg_report"] = anat_report_wf.bg_report
-    outputs_["anat_report_wf_airmask_report"] = anat_report_wf.airmask_report
     outputs_["anat_report_wf_zoom_report"] = anat_report_wf.zoom_report
+    outputs_["anat_report_wf_headmask_report"] = anat_report_wf.headmask_report
     outputs_["anat_report_wf_segm_report"] = anat_report_wf.segm_report
 
     return tuple(outputs_)
@@ -252,7 +242,7 @@ def airmsk_wf(
     in_file: ty.Any = attrs.NOTHING,
     ind2std_xfm: ty.Any = attrs.NOTHING,
     name="AirMaskWorkflow",
-) -> ["ty.Any", "ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any, ty.Any]:
     """
     Calculate air, artifacts and "hat" masks to evaluate noise in the background.
 
@@ -314,7 +304,7 @@ def headmsk_wf(
     name="HeadMaskWorkflow",
     omp_nthreads=1,
     wf_species="human",
-) -> ["ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any]:
     """
     Computes a head mask as in [Mortamet2009]_.
 
@@ -337,56 +327,32 @@ def _select_wm(inlist):
         return [f for f in inlist if "WM" in f][0]
 
     enhance = workflow.add(
-        FunctionTask(
-            func=_enhance,
-            input_spec=SpecInfo(
-                name="FunctionIn",
-                bases=(BaseSpec,),
-                fields=[("in_file", ty.Any), ("wm_tpm", ty.Any)],
-            ),
-            output_spec=SpecInfo(
-                name="FunctionOut", bases=(BaseSpec,), fields=[("out_file", ty.Any)]
-            ),
-            in_file=in_file,
-            wm_tpm=in_tpms,
-        ),
+        python.define(
+            _enhance,
+            inputs={"in_file": ty.Any, "wm_tpm": ty.Any},
+            outputs={"out_file": ty.Any},
+        )(in_file=in_file, wm_tpm=in_tpms),
         name="enhance",
     )
     gradient = workflow.add(
-        FunctionTask(
-            func=image_gradient,
-            input_spec=SpecInfo(
-                name="FunctionIn",
-                bases=(BaseSpec,),
-                fields=[("in_file", ty.Any), ("brainmask", ty.Any), ("sigma", ty.Any)],
-            ),
-            output_spec=SpecInfo(
-                name="FunctionOut", bases=(BaseSpec,), fields=[("out_file", ty.Any)]
-            ),
-            brainmask=brainmask,
-            in_file=enhance.out_file,
-        ),
+        python.define(
+            image_gradient,
+            inputs={"in_file": ty.Any, "brainmask": ty.Any, "sigma": ty.Any},
+            outputs={"out_file": ty.Any},
+        )(brainmask=brainmask, in_file=enhance.out_file),
         name="gradient",
     )
     thresh = workflow.add(
-        FunctionTask(
-            func=gradient_threshold,
-            input_spec=SpecInfo(
-                name="FunctionIn",
-                bases=(BaseSpec,),
-                fields=[
-                    ("in_file", ty.Any),
-                    ("brainmask", ty.Any),
-                    ("aniso", ty.Any),
-                    ("thresh", ty.Any),
-                ],
-            ),
-            output_spec=SpecInfo(
-                name="FunctionOut", bases=(BaseSpec,), fields=[("out_file", ty.Any)]
-            ),
-            brainmask=brainmask,
-            in_file=gradient.out_file,
-        ),
+        python.define(
+            gradient_threshold,
+            inputs={
+                "in_file": ty.Any,
+                "brainmask": ty.Any,
+                "aniso": ty.Any,
+                "thresh": ty.Any,
+            },
+            outputs={"out_file": ty.Any},
+        )(brainmask=brainmask, in_file=gradient.out_file),
         name="thresh",
     )
     if wf_species != "human":
@@ -412,7 +378,7 @@ def init_brain_tissue_segmentation(
     name="brain_tissue_segmentation",
     nipype_omp_nthreads=12,
     std_tpms: ty.Any = attrs.NOTHING,
-) -> ["ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any]:
     """
     Setup a workflow for brain tissue segmentation.
 
@@ -454,15 +420,12 @@ def _format_tpm_names(in_files, fname_string=None):
         return file_format, out_files
 
     format_tpm_names = workflow.add(
-        FunctionTask(
+        python.define(
+            _format_tpm_names,
+            inputs={"in_files": ty.Any},
+            outputs={"file_format": ty.Any},
+        )(
             execution={"keep_inputs": True, "remove_unnecessary_outputs": False},
-            func=_format_tpm_names,
-            input_spec=SpecInfo(
-                name="FunctionIn", bases=(BaseSpec,), fields=[("in_files", ty.Any)]
-            ),
-            output_spec=SpecInfo(
-                name="FunctionOut", bases=(BaseSpec,), fields=[("file_format", ty.Any)]
-            ),
             in_files=std_tpms,
         ),
         name="format_tpm_names",
@@ -510,7 +473,7 @@ def spatial_normalization(
     nipype_omp_nthreads=12,
     wf_species="human",
     wf_template_id="MNI152NLin2009cAsym",
-) -> ["ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any]:
     """Create a simplified workflow to perform fast spatial normalization."""
     from pydra.tasks.niworkflows.interfaces.reportlets.registration import (
         SpatialNormalizationRPT as RobustMNINormalization,
@@ -595,7 +558,7 @@ def compute_iqms(
     segmentation: ty.Any = attrs.NOTHING,
     std_tpms: ty.Any = attrs.NOTHING,
     wf_species="human",
-) -> ["ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any]:
     """
     Setup the workflow that actually computes the IQMs.
 

pydra/tasks/mriqc/workflows/anatomical/output.py

@@ -30,7 +30,7 @@ def init_anat_report_wf(
     name: str = "anat_report_wf",
     segmentation: ty.Any = attrs.NOTHING,
     wf_species="human",
-) -> ["ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any]:
     """
     Generate the components of the individual report.
 

pydra/tasks/mriqc/workflows/diffusion/base.py

@@ -44,17 +44,8 @@ def dmri_qc_workflow(
     wf_inputs_metadata=None,
     wf_species="human",
     wf_template_id="MNI152NLin2009cAsym",
-) -> [
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
-    "ty.Any",
+) -> tuple[
+    ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any
 ]:
     """
     Initialize the dMRI-QC workflow.
@@ -209,12 +200,12 @@ def dmri_qc_workflow(
     # EPI to MNI registration
     spatial_norm = workflow.add(
         epi_mni_align(
-            nipype_omp_nthreads=nipype_omp_nthreads,
             wf_species=wf_species,
-            wf_template_id=wf_template_id,
+            nipype_omp_nthreads=nipype_omp_nthreads,
             nipype_nprocs=nipype_nprocs,
-            exec_debug=exec_debug,
+            wf_template_id=wf_template_id,
             exec_ants_float=exec_ants_float,
+            exec_debug=exec_debug,
             epi_mean=dwi_ref.out_file,
             epi_mask=dmri_bmsk.out_mask,
             name="spatial_norm",
@@ -249,11 +240,11 @@ def dmri_qc_workflow(
     # Generate outputs
     dwi_report_wf = workflow.add(
         init_dwi_report_wf(
-            wf_species=wf_species,
-            wf_biggest_file_gb=wf_biggest_file_gb,
             exec_verbose_reports=exec_verbose_reports,
-            wf_fd_thres=wf_fd_thres,
             exec_work_dir=exec_work_dir,
+            wf_species=wf_species,
+            wf_fd_thres=wf_fd_thres,
+            wf_biggest_file_gb=wf_biggest_file_gb,
             in_bdict=shells.b_dict,
             brain_mask=dmri_bmsk.out_mask,
             in_avgmap=averages.out_file,
@@ -313,14 +304,14 @@ def averages_out_file_to_iqms_wf_in_b0_callable(in_: ty.Any) -> ty.Any:
     # fmt: on
     outputs_["iqms_wf_out_file"] = iqms_wf.out_file
     outputs_["iqms_wf_noise_floor"] = iqms_wf.noise_floor
-    outputs_["dwi_report_wf_noise_report"] = dwi_report_wf.noise_report
+    outputs_["dwi_report_wf_heatmap_report"] = dwi_report_wf.heatmap_report
+    outputs_["dwi_report_wf_spikes_report"] = dwi_report_wf.spikes_report
+    outputs_["dwi_report_wf_fa_report"] = dwi_report_wf.fa_report
+    outputs_["dwi_report_wf_carpet_report"] = dwi_report_wf.carpet_report
     outputs_["dwi_report_wf_md_report"] = dwi_report_wf.md_report
-    outputs_["dwi_report_wf_bmask_report"] = dwi_report_wf.bmask_report
+    outputs_["dwi_report_wf_noise_report"] = dwi_report_wf.noise_report
     outputs_["dwi_report_wf_snr_report"] = dwi_report_wf.snr_report
-    outputs_["dwi_report_wf_carpet_report"] = dwi_report_wf.carpet_report
-    outputs_["dwi_report_wf_fa_report"] = dwi_report_wf.fa_report
-    outputs_["dwi_report_wf_spikes_report"] = dwi_report_wf.spikes_report
-    outputs_["dwi_report_wf_heatmap_report"] = dwi_report_wf.heatmap_report
+    outputs_["dwi_report_wf_bmask_report"] = dwi_report_wf.bmask_report
 
     return tuple(outputs_)
 
@@ -332,7 +323,7 @@ def hmc_workflow(
     name="dMRI_HMC",
     reference: ty.Any = attrs.NOTHING,
     wf_fd_radius=50,
-) -> ["ty.Any", "ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any, ty.Any]:
     """
     Create a :abbr:`HMC (head motion correction)` workflow for dMRI.
 
@@ -406,7 +397,7 @@ def epi_mni_align(
     nipype_omp_nthreads=12,
     wf_species="human",
     wf_template_id="MNI152NLin2009cAsym",
-) -> ["ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any]:
     """
     Estimate the transform that maps the EPI space into MNI152NLin2009cAsym.
 
@@ -487,9 +478,7 @@ def epi_mni_align(
                 suffix="mask",
             )[0]
         )
-        bspline_grid = workflow.add(
-            FunctionTask(func=_bspline_grid), name="bspline_grid"
-        )
+        bspline_grid = workflow.add(python.define(_bspline_grid)(), name="bspline_grid")
         # fmt: off
         bspline_grid.inputs.in_file = epi_mean
         n4itk.inputs.args = bspline_grid.out
@@ -557,7 +546,7 @@ def compute_iqms(
     qspace_neighbors: ty.Any = attrs.NOTHING,
     spikes_mask: ty.Any = attrs.NOTHING,
     wm_mask: ty.Any = attrs.NOTHING,
-) -> ["ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any]:
     """
     Initialize the workflow that actually computes the IQMs.
 
@@ -582,7 +571,7 @@ def compute_iqms(
     # from mriqc.workflows.utils import _tofloat, get_fwhmx
 
     estimate_sigma = workflow.add(
-        FunctionTask(func=_estimate_sigma, in_file=in_noise, mask=brain_mask),
+        python.define(_estimate_sigma)(in_file=in_noise, mask=brain_mask),
         name="estimate_sigma",
     )
     measures = workflow.add(

pydra/tasks/mriqc/workflows/diffusion/output.py

@@ -1,7 +1,7 @@
 import attrs
 import logging
 from pathlib import Path
-from pydra.compose import workflow
+from pydra.compose import python, workflow
 from pydra.tasks.mriqc.workflows.diffusion.output import init_dwi_report_wf
 from pydra.tasks.nireports.interfaces.dmri import DWIHeatmap
 from pydra.tasks.nireports.interfaces.reporting.base import (
@@ -41,7 +41,7 @@ def init_dwi_report_wf(
     wf_biggest_file_gb=1,
     wf_fd_thres=0.2,
     wf_species="human",
-) -> ["ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any", "ty.Any"]:
+) -> tuple[ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any, ty.Any]:
     """
     Write out individual reportlets.
 
@@ -118,7 +118,7 @@ def _gen_entity(inlist):
     outputs_['md_report'] = mosaic_md.out_file
     # fmt: on
     get_wm = workflow.add(
-        FunctionTask(func=_get_wm, in_file=in_parcellation), name="get_wm"
+        python.define(_get_wm)(in_file=in_parcellation), name="get_wm"
     )
     plot_heatmap = workflow.add(
         DWIHeatmap(