From a238a56834ab45c1c28ed6a07b3ec91e5eb0627e Mon Sep 17 00:00:00 2001
From: Robert Luke <748691+rob-luke@users.noreply.github.com>
Date: Thu, 12 Mar 2020 15:48:21 +1100
Subject: [PATCH] Add interpolation of nirs data. Replaces with nearest

---
 mne/channels/interpolation.py            | 59 ++++++++++++++++++++++++
 mne/channels/tests/test_interpolation.py | 19 ++++++++
 2 files changed, 78 insertions(+)

diff --git a/mne/channels/interpolation.py b/mne/channels/interpolation.py
index 25a185db2c0..9fd1dbe5242 100644
--- a/mne/channels/interpolation.py
+++ b/mne/channels/interpolation.py
@@ -213,3 +213,62 @@ def _interpolate_bads_meg(inst, mode='accurate', origin=(0., 0., 0.04),
     info_to = pick_info(inst.info, picks_bad)
     mapping = _map_meg_channels(info_from, info_to, mode=mode, origin=origin)
     _do_interp_dots(inst, mapping, picks_good, picks_bad)
+
+
+@verbose
+def _interpolate_bads_nirs(inst, method='nearest', verbose=None):
+    """Interpolate bad nirs channels. Simply replaces by closest non bad.
+
+    Parameters
+    ----------
+    inst : mne.io.Raw, mne.Epochs or mne.Evoked
+        The data to interpolate. Must be preloaded.
+    method : str
+        Only the method 'nearest' is currently available. This method replaces
+        each bad channel with the nearest non bad channel.
+    %(verbose)s
+    """
+    from scipy.spatial.distance import pdist, squareform
+    from mne.preprocessing.nirs import _channel_frequencies,\
+        _check_channels_ordered
+
+    freqs = np.unique(_channel_frequencies(inst))
+    # Returns pick of all nirs and ensures channels are correctly ordered
+    picks_nirs = _check_channels_ordered(inst, freqs)
+    nirs_ch_names = [inst.info['ch_names'][p] for p in picks_nirs]
+    bads_nirs = [ch for ch in inst.info['bads'] if ch in nirs_ch_names]
+
+    # select the bad meg channel to be interpolated
+    if len(bads_nirs) == 0:
+        picks_bad = []
+    else:
+        picks_bad = pick_channels(inst.info['ch_names'], bads_nirs,
+                                  exclude=[])
+    bads_mask = [p in picks_bad for p in picks_nirs]
+
+    chs = [inst.info['chs'][i] for i in picks_nirs]
+    locs3d = np.array([ch['loc'][:3] for ch in chs])
+    dist = pdist(locs3d)
+    dist = squareform(dist)
+
+    if method == 'nearest':
+
+        for bad in picks_bad:
+            dists_to_bad = dist[bad]
+            # Ignore distances to self
+            dists_to_bad[dists_to_bad == 0] = np.inf
+            # Ignore distances to other bad channels
+            dists_to_bad[bads_mask] = np.inf
+            # Find closest remaining channels
+            closest_idx = np.where(dists_to_bad == np.amin(dists_to_bad))[0]
+            # Return the same frequency
+            closest_idx = closest_idx[bad % 2]
+            inst._data[bad] = inst._data[closest_idx]
+
+        inst.info['bads'] = []
+
+    else:
+        warn('No interpolation applied. Unknown NIRS interpolation '
+             'method: ' + method)
+
+    return inst
diff --git a/mne/channels/tests/test_interpolation.py b/mne/channels/tests/test_interpolation.py
index eea6b223721..a946289224e 100644
--- a/mne/channels/tests/test_interpolation.py
+++ b/mne/channels/tests/test_interpolation.py
@@ -219,4 +219,23 @@ def test_interpolation_ctf_comp():
     assert raw.info['bads'] == []
 
 
+@testing.requires_testing_data
+def test_interpolation_nirs():
+    """Test interpolating bad nirs channels."""
+    from mne.preprocessing.nirs import optical_density, scalp_coupling_index
+    from mne.datasets.testing import data_path
+    from mne.io import read_raw_nirx
+    from itertools import compress
+    from mne.channels.interpolation import _interpolate_bads_nirs
+
+    fname = op.join(data_path(download=False),
+                    'NIRx', 'nirx_15_2_recording_w_overlap')
+    raw_intensity = read_raw_nirx(fname, preload=False)
+    raw_od = optical_density(raw_intensity)
+    sci = scalp_coupling_index(raw_od)
+    raw_od.info['bads'] = list(compress(raw_od.ch_names, sci < 0.5))
+    _interpolate_bads_nirs(raw_od)
+    assert raw_od.info['bads'] == []
+
+
 run_tests_if_main()