ENH: Restore correlation plots

poldracklab · Dec 19, 2018 · 760830d · 760830d
1 parent f165116
commit 760830d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/fitlins/interfaces/visualizations.py b/fitlins/interfaces/visualizations.py
@@ -61,17 +61,24 @@ def _visualize(self, data, out_name):
 
 
 class DesignCorrelationPlotInputSpec(VisualizationInputSpec):
-    explanatory_variables = traits.Range(
-        low=0, desc='Number of explanatory variables')
+    contrast_info = traits.List(traits.Dict)
 
 
 class DesignCorrelationPlot(Visualization):
     input_spec = DesignCorrelationPlotInputSpec
 
     def _visualize(self, data, out_name):
+        contrast_matrix = pd.DataFrame({c['name']: c['weights'][0]
+                                        for c in self.inputs.contrast_info})
+        all_cols = list(data.columns)
+        evs = set(contrast_matrix.index)
+        if set(contrast_matrix.index) != all_cols[:len(evs)]:
+            ev_cols = [col for col in all_cols if col in evs]
+            confound_cols = [col for col in all_cols if col not in evs]
+            data = data[ev_cols + confound_cols]
         plot_and_save(out_name, plot_corr_matrix,
                       data.drop(columns='constant').corr(),
-                      self.inputs.explanatory_variables)
+                      len(evs))
 
 
 class ContrastMatrixPlotInputSpec(VisualizationInputSpec):

diff --git a/fitlins/workflows/base.py b/fitlins/workflows/base.py
@@ -92,6 +92,11 @@ def init_fitlins_wf(bids_dir, derivatives, out_dir, space, exclude_pattern=None,
         iterfield='data',
         name='plot_design')
 
+    plot_corr = pe.MapNode(
+        DesignCorrelationPlot(image_type='svg'),
+        iterfield=['data', 'contrast_info'],
+        name='plot_corr')
+
     plot_l1_contrast_matrix = pe.MapNode(
         ContrastMatrixPlot(image_type='svg'),
         iterfield=['data', 'contrast_info'],
@@ -104,6 +109,13 @@ def init_fitlins_wf(bids_dir, derivatives, out_dir, space, exclude_pattern=None,
         run_without_submitting=True,
         name='ds_design')
 
+    ds_corr = pe.MapNode(
+        BIDSDataSink(base_directory=out_dir, fixed_entities={'suffix': 'corr'},
+                     path_patterns=image_pattern),
+        iterfield=['entities', 'in_file'],
+        run_without_submitting=True,
+        name='ds_corr')
+
     ds_l1_contrasts = pe.MapNode(
         BIDSDataSink(base_directory=out_dir, fixed_entities={'suffix': 'contrasts'},
                      path_patterns=image_pattern),
@@ -185,9 +197,13 @@ def init_fitlins_wf(bids_dir, derivatives, out_dir, space, exclude_pattern=None,
                 (select_entities, ds_design, [('out', 'entities')]),
                 (plot_design, ds_design, [('figure', 'in_file')]),
                 (select_contrasts, plot_l1_contrast_matrix,  [('out', 'contrast_info')]),
+                (select_contrasts, plot_corr,  [('out', 'contrast_info')]),
                 (model, plot_l1_contrast_matrix,  [('design_matrix', 'data')]),
+                (model, plot_corr,  [('design_matrix', 'data')]),
                 (select_entities, ds_l1_contrasts, [('out', 'entities')]),
+                (select_entities, ds_corr, [('out', 'entities')]),
                 (plot_l1_contrast_matrix, ds_l1_contrasts,  [('figure', 'in_file')]),
+                (plot_corr, ds_corr,  [('figure', 'in_file')]),
             ])
 
         #  Set up higher levels