Merge pull request #810 from effigies/fix/X_patterns

ENH: Handle wildcards in model X

Author: effigies

bids/modeling/statsmodels.py

@@ -4,6 +4,8 @@
 from collections import namedtuple, OrderedDict, Counter, defaultdict
 import itertools
 from functools import reduce
+import re
+import fnmatch
 
 import numpy as np
 import pandas as pd
@@ -520,6 +522,15 @@ def get_collections(self, **filters):
         return [c for c in self._collections if matches_entities(c, filters)]
 
 
+def expand_wildcards(selectors, pool):
+    out = list(selectors)
+    for spec in selectors:
+        if re.search(r'[\*\?\[\]]', spec):
+            idx = out.index(spec)
+            out[idx:idx + 1] = fnmatch.filter(pool, spec)
+    return out
+
+
 class BIDSStatsModelsNodeOutput:
     """Represents a single node in a BIDSStatsModelsGraph.
 
@@ -597,6 +608,8 @@ def merge_dfs(a, b):
             else:
                 var_names.append(int_name)
 
+        var_names = expand_wildcards(var_names, df.columns)
+
         # Verify all X names are actually present
         missing = list(set(var_names) - set(df.columns))
         if missing:
@@ -642,7 +655,9 @@ def _collections_to_dfs(self, collections):
 
             # Take the intersection of variables and Model.X (var_names), ignoring missing
             # variables (usually contrasts)
-            coll.variables = {v: coll.variables[v] for v in var_names if v in coll.variables}
+            coll.variables = {v: coll.variables[v]
+                              for v in expand_wildcards(var_names, coll.variables)
+                              if v in coll.variables}
             if not coll.variables:
                 continue
 

bids/modeling/tests/test_statsmodels.py

@@ -10,7 +10,7 @@
 import pytest
 
 from bids.modeling import BIDSStatsModelsGraph
-from bids.modeling.statsmodels import ContrastInfo
+from bids.modeling.statsmodels import ContrastInfo, expand_wildcards
 from bids.layout import BIDSLayout
 from bids.tests import get_test_data_path
 from bids.variables import BIDSVariableCollection
@@ -152,3 +152,20 @@ def test_entire_graph_smoketest(graph):
     assert model_spec.X.shape == (2, 1)
     assert model_spec.Z is None
     assert not set(model_spec.terms.keys()) - {"RT", "gain", "RT:gain"}
+
+
+def test_expand_wildcards():
+    # No wildcards == no modification
+    assert expand_wildcards(["a", "b"], ["a", "c"]) == ["a", "b"]
+    # No matches == removal
+    assert expand_wildcards(["a", "b*"], ["a", "c"]) == ["a"]
+    # Matches expand in-place
+    assert expand_wildcards(["a*", "b"], ["a", "c"]) == ["a", "b"]
+    assert expand_wildcards(["a*", "b"], ["a0", "c", "a1", "a2"]) == ["a0", "a1", "a2", "b"]
+    # Some examples
+    assert expand_wildcards(
+        ["trial_type.*"], ["trial_type.A", "trial_type.B"]
+    ) == ["trial_type.A", "trial_type.B"]
+    assert expand_wildcards(
+        ["non_steady_state*"], ["non_steady_state00", "non_steady_state01"]
+    ) == ["non_steady_state00", "non_steady_state01"]