diff --git a/doc/changes/dev/14245.newfeature.rst b/doc/changes/dev/14245.newfeature.rst new file mode 100644 index 00000000000..f19d0dba6a1 --- /dev/null +++ b/doc/changes/dev/14245.newfeature.rst @@ -0,0 +1 @@ +Speed up channel selection and projector construction by replacing linear channel-name lookups with dictionaries, by `Bruno Aristimunha`_. diff --git a/mne/_fiff/pick.py b/mne/_fiff/pick.py index 0526b99ca5a..78fc05d8c28 100644 --- a/mne/_fiff/pick.py +++ b/mne/_fiff/pick.py @@ -319,13 +319,17 @@ def pick_channels(ch_names, include, exclude=(), ordered=True, *, verbose=None): include = list(ch_names) if not isinstance(exclude, list): exclude = list(exclude) + # ch_names is unique (checked above), so a lookup table is safe here; the + # list scans this replaces made the loop quadratic in the channel count + name_to_idx = {name: ii for ii, name in enumerate(ch_names)} + exclude_set = set(exclude) sel, missing = list(), list() for name in include: - if name in ch_names: - if name not in exclude: - sel.append(ch_names.index(name)) - else: + idx = name_to_idx.get(name) + if idx is None: missing.append(name) + elif name not in exclude_set: + sel.append(idx) if len(missing) and ordered: raise ValueError( f"Missing channels from ch_names required by include:\n{missing}" @@ -1408,13 +1412,20 @@ def _picks_str_to_idx( # second: match all to channel names # + # setdefault keeps the first occurrence, so this matches list.index() + # exactly even for duplicate names (which are rejected further down, so + # the difference is not reachable today -- it just keeps the swap honest) + name_to_idx = {} + for ii, name in enumerate(info["ch_names"]): + name_to_idx.setdefault(name, ii) bad_names = [] picks_name = list() for pick in picks: - try: - picks_name.append(info["ch_names"].index(pick)) - except ValueError: + idx = name_to_idx.get(pick) + if idx is None: bad_names.append(pick) + else: + picks_name.append(idx) # # third: match all to types diff --git a/mne/_fiff/proj.py b/mne/_fiff/proj.py index 16418e2d601..bba3d6797f2 100644 --- a/mne/_fiff/proj.py +++ b/mne/_fiff/proj.py @@ -889,11 +889,15 @@ def _make_projector(projs, ch_names, bads=(), include_active=True, inplace=False # the projection vectors omitting bad channels sel = [] vecsel = [] - p_set = set(p["data"]["col_names"]) # faster membership access + # map name -> position once; .index() here made this loop quadratic + # in the channel count (~16x slower at 306 channels, ~47x at 1000) + p_idx = {name: i for i, name in enumerate(p["data"]["col_names"])} for c, name in enumerate(ch_names): - if name not in bads and name in p_set: - sel.append(c) - vecsel.append(p["data"]["col_names"].index(name)) + if name not in bads: + vi = p_idx.get(name) + if vi is not None: + sel.append(c) + vecsel.append(vi) # If there is something to pick, pickit nrow = p["data"]["nrow"] diff --git a/mne/_fiff/tests/test_pick.py b/mne/_fiff/tests/test_pick.py index c958d072e4d..ff4ab67f34a 100644 --- a/mne/_fiff/tests/test_pick.py +++ b/mne/_fiff/tests/test_pick.py @@ -762,3 +762,28 @@ def test_get_channel_types_equiv(meg, eeg, ordered): types = np.array(raw.get_channel_types(picks=picks)) types_iter = np.array([channel_type(raw.info, idx) for idx in picks]) assert_array_equal(types, types_iter) + + +def test_pick_channels_matches_by_name(): + """Test picking maps names to positions regardless of the order given.""" + ch_names = ["a", "b", "c", "d"] + # include is out of order, repeats a name, and names one that is excluded + sel = pick_channels(ch_names, ["d", "b", "b", "c"], exclude=["c"], ordered=False) + assert_array_equal(sel, [1, 3]) + # with ordered=True the caller's order is kept, duplicates and all + sel = pick_channels(ch_names, ["d", "b", "b"], ordered=True) + assert_array_equal(sel, [3, 1, 1]) + # a name that is not present is an error, not a silent skip + with pytest.raises(ValueError, match="Missing channels"): + pick_channels(ch_names, ["a", "nope"], ordered=True) + + +def test_picks_to_idx_duplicate_names(): + """Test a repeated channel name resolves to its first position.""" + with pytest.warns(RuntimeWarning, match="not unique"): + info = create_info(["a", "b", "a"], 100.0, "eeg") + # "b" is unambiguous; picking it must not be shifted by the duplicate "a" + assert_array_equal(_picks_to_idx(info, ["b"]), [1]) + # an ambiguous name is rejected rather than silently resolved + with pytest.raises(ValueError, match="could not be interpreted"): + _picks_to_idx(info, ["a"]) diff --git a/mne/tests/test_proj.py b/mne/tests/test_proj.py index 316258c6eb0..7e879e710a8 100644 --- a/mne/tests/test_proj.py +++ b/mne/tests/test_proj.py @@ -832,3 +832,18 @@ def test_compute_proj_explained_variance(): if proj["desc"].split("-")[0] == type_ ] assert n_vector_ <= sum(explained_var) + + +def test_make_projector_matches_by_name(): + """Test projector columns are matched by channel name, not by position.""" + ch_names = ["EEG 001", "EEG 002", "EEG 003"] + # col_names deliberately permuted, so a positional match gives the wrong + # vector: the direction is [2, 3, 1] in ch_names order, not [1, 2, 3] + proj = _make_test_proj(["EEG 003", "EEG 001", "EEG 002"], [1.0, 2.0, 3.0], "perm") + P, nproj, _ = make_projector([proj], ch_names) + assert nproj == 1 + direction = np.array([2.0, 3.0, 1.0]) + assert_allclose(P @ direction, 0.0, atol=1e-12) + # and a channel the projector does not mention is left untouched + P, _, _ = make_projector([proj], ch_names + ["EEG 004"]) + assert_allclose(P[3], [0, 0, 0, 1], atol=1e-12)