Merge pull request #224 from poldracklab/enh/model_ouputs

ENH: Save model level output (r-squared)
poldracklab · Apr 6, 2020 · abe803d · abe803d
2 parents 2d2fc55 + 29a2a0d
commit abe803d
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 3 deletions.
diff --git a/fitlins/interfaces/abstract.py b/fitlins/interfaces/abstract.py
@@ -42,6 +42,8 @@ class EstimatorOutputSpec(TraitedSpec):
     zscore_maps = traits.List(File)
     pvalue_maps = traits.List(File)
     contrast_metadata = traits.List(traits.Dict)
+    model_maps = traits.List(File)
+    model_metadata = traits.List(traits.Dict)
 
 
 class FirstLevelEstimatorInterface(BaseInterface):

diff --git a/fitlins/interfaces/nistats.py b/fitlins/interfaces/nistats.py
@@ -122,18 +122,33 @@ def _run_interface(self, runtime):
         smoothing_fwhm = self.inputs.smoothing_fwhm
         if not isdefined(smoothing_fwhm):
             smoothing_fwhm = None
+
         flm = level1.FirstLevelModel(
+            minimize_memory=False,
             mask_img=mask_file, smoothing_fwhm=smoothing_fwhm)
         flm.fit(img, design_matrices=mat)
 
+        out_ents = self.inputs.contrast_info[0]['entities']
+        fname_fmt = os.path.join(runtime.cwd, '{}_{}.nii.gz').format
+
+        # Save model level images
+        model_maps = []
+        model_metadata = []
+        for output in ['r_square']:
+            model_metadata.append(
+                {'stat': output,
+                 **out_ents}
+                )
+            fname = fname_fmt('model', output)
+            getattr(flm, output)[0].to_filename(fname)
+            model_maps.append(fname)
+
         effect_maps = []
         variance_maps = []
         stat_maps = []
         zscore_maps = []
         pvalue_maps = []
         contrast_metadata = []
-        out_ents = self.inputs.contrast_info[0]['entities']
-        fname_fmt = os.path.join(runtime.cwd, '{}_{}.nii.gz').format
         for name, weights, contrast_type in prepare_contrasts(
               self.inputs.contrast_info, mat.columns):
             contrast_metadata.append(
@@ -160,6 +175,8 @@ def _run_interface(self, runtime):
         self._results['zscore_maps'] = zscore_maps
         self._results['pvalue_maps'] = pvalue_maps
         self._results['contrast_metadata'] = contrast_metadata
+        self._results['model_maps'] = model_maps
+        self._results['model_metadata'] = model_metadata
 
         return runtime
 

diff --git a/fitlins/workflows/base.py b/fitlins/workflows/base.py
@@ -115,7 +115,10 @@ def _deindex(tsv):
         '[sub-{subject}_][ses-{session}_]task-{task}[_acq-{acquisition}]' \
         '[_rec-{reconstruction}][_run-{run}][_echo-{echo}][_space-{space}]_' \
         'contrast-{contrast}_stat-{stat<effect|variance|z|p|t|F|FEMA>}_statmap.nii.gz'
-
+    model_map_pattern = '[sub-{subject}/][ses-{session}/]' \
+        '[sub-{subject}_][ses-{session}_]task-{task}[_acq-{acquisition}]' \
+        '[_rec-{reconstruction}][_run-{run}][_echo-{echo}][_space-{space}]_' \
+        'stat-{stat<rSquare>}_statmap.nii.gz'
     # Set up general interfaces
     #
     # HTML snippets to be included directly in report, not
@@ -264,6 +267,18 @@ def _deindex(tsv):
             name='ds_{}_contrast_plots'.format(level))
 
         if ix == 0:
+            ds_model_maps = pe.Node(
+                BIDSDataSink(base_directory=out_dir,
+                             path_patterns=model_map_pattern),
+                run_without_submitting=True,
+                name='ds_{}_model_maps'.format(level))
+
+            collate_mm = pe.Node(
+                MergeAll(['model_maps', 'model_metadata'],
+                         check_lengths=(not drop_missing)),
+                name='collate_mm_{}'.format(level),
+                run_without_submitting=True)
+
             wf.connect([
                 (loader, select_entities, [('entities', 'inlist')]),
                 (select_entities, getter,  [('out', 'entities')]),
@@ -277,6 +292,10 @@ def _deindex(tsv):
                 (select_entities, ds_corr, [('out', 'entities')]),
                 (plot_l1_contrast_matrix, ds_l1_contrasts,  [('figure', 'in_file')]),
                 (plot_corr, ds_corr,  [('figure', 'in_file')]),
+                (model, collate_mm, [('model_maps', 'model_maps'),
+                                     ('model_metadata', 'model_metadata')]),
+                (collate_mm, ds_model_maps, [('model_maps', 'in_file'),
+                                             ('model_metadata', 'entities')]),
             ])
 
         #  Set up higher levels