FIX: allow reading experimental dipole

doc/whats_new.rst

@@ -62,6 +62,8 @@ Bug
 
 - Fix overlapping non-cell-specific drive gid assignment over different ranks in `~hnn_core.MPIBackend`, by `Ryan Thorpe`_ and `Mainak Jas`_ in `#399 <https://github.com/jonescompneurolab/hnn-core/pull/399>`_
 
+- Allow :func:`~hnn_core.read_dipoles` to read dipole from a file with only two columns (`times` and `data`), by `Mainak Jas`_ in `#421 <https://github.com/jonescompneurolab/hnn-core/pull/421>`_
+
 API
 ~~~
 - New API for defining cell-cell connections. Custom connections can be added with :func:`~hnn_core.Network.add_connection`, by `Nick Tolley`_ in `#276 <https://github.com/jonescompneurolab/hnn-core/pull/276>`_

examples/howto/optimize_evoked.py

@@ -19,8 +19,8 @@
 # Let us import hnn_core
 
 import hnn_core
-from hnn_core.dipole import simulate_dipole
-from hnn_core import read_params, Dipole, MPIBackend, jones_2009_model
+from hnn_core.dipole import read_dipole, simulate_dipole
+from hnn_core import read_params, MPIBackend, jones_2009_model
 
 hnn_core_root = op.join(op.dirname(hnn_core.__file__))
 
@@ -38,9 +38,7 @@
 data_url = ('https://raw.githubusercontent.com/jonescompneurolab/hnn/master/'
             'data/MEG_detection_data/S1_SupraT.txt')
 urlretrieve(data_url, 'S1_SupraT.txt')
-extdata = np.loadtxt('S1_SupraT.txt')
-exp_dpl = Dipole(extdata[:, 0], np.c_[extdata[:, 1],
-                 extdata[:, 1], extdata[:, 1]])
+exp_dpl = read_dipole('S1_SupraT.txt')
 
 ###############################################################################
 # Read the base parameters from a file

hnn_core/dipole.py

@@ -108,7 +108,7 @@ def read_dipole(fname):
         The instance of Dipole class
     """
     dpl_data = np.loadtxt(fname, dtype=float)
-    dpl = Dipole(dpl_data[:, 0], dpl_data[:, 1:4])
+    dpl = Dipole(dpl_data[:, 0], dpl_data[:, 1:])
     return dpl
 
 
@@ -137,15 +137,14 @@ def average_dipoles(dpls):
                              " trials. Cannot reaverage" %
                              (dpl_idx, dpl.nave))
 
-    agg_avg = np.mean(np.array([dpl.data['agg'] for dpl in dpls]), axis=0)
-    L2_avg = np.mean(np.array([dpl.data['L2'] for dpl in dpls]), axis=0)
-    L5_avg = np.mean(np.array([dpl.data['L5'] for dpl in dpls]), axis=0)
-
-    avg_dpl_data = np.c_[agg_avg,
-                         L2_avg,
-                         L5_avg]
-
-    avg_dpl = Dipole(dpls[0].times, avg_dpl_data)
+    avg_data = list()
+    layers = dpl.data.keys()
+    for layer in layers:
+        avg_data.append(
+            np.mean(np.array([dpl.data[layer] for dpl in dpls]), axis=0)
+        )
+    avg_data = np.c_[avg_data].T
+    avg_dpl = Dipole(dpls[0].times, avg_data)
 
     # set nave to the number of trials averaged in this dipole
     avg_dpl.nave = len(dpls)
@@ -232,9 +231,10 @@ class Dipole(object):
     ----------
     times : array (n_times,)
         The time vector (in ms)
-    data : array (n_times x 3)
-        The data. The first column represents 'agg',
-        the second 'L2' and the last one 'L5'
+    data : array, shape (n_times x n_layers)
+        The data. The first column represents 'agg' (the total diple),
+        the second 'L2' layer and the last one 'L5' layer. For experimental
+        data, it can contain only one column.
     nave : int
         Number of trials that were averaged to produce this Dipole. Defaults
         to 1
@@ -256,7 +256,15 @@ class Dipole(object):
 
     def __init__(self, times, data, nave=1):  # noqa: D102
         self.times = np.array(times)
-        self.data = {'agg': data[:, 0], 'L2': data[:, 1], 'L5': data[:, 2]}
+
+        if data.ndim == 1:
+            data = data[:, None]
+
+        if data.shape[1] == 3:
+            self.data = {'agg': data[:, 0], 'L2': data[:, 1], 'L5': data[:, 2]}
+        elif data.shape[1] == 1:
+            self.data = {'agg': data[:, 0]}
+
         self.nave = nave
         self.sfreq = 1000. / (times[1] - times[0])  # NB assumes len > 1
         self.scale_applied = 1  # for visualisation
@@ -569,7 +577,11 @@ def write(self, fname):
             warnings.warn("Saving Dipole to file that is an average of %d"
                           " trials" % self.nave)
 
-        X = np.r_[[self.times, self.data['agg'], self.data['L2'],
-                   self.data['L5']]].T
-        np.savetxt(fname, X, fmt=['%3.3f', '%5.4f', '%5.4f', '%5.4f'],
-                   delimiter='\t')
+        X = [self.times]
+        fmt = ['%3.3f']
+        for data in self.data.values():
+            X.append(data)
+            fmt.append('%5.4f')
+        X = np.r_[X].T
+
+        np.savetxt(fname, X, fmt=fmt, delimiter='\t')

hnn_core/tests/test_dipole.py

@@ -40,6 +40,8 @@ def test_dipole(tmpdir, run_hnn_core_fixture):
 
     dipole.plot(show=False)
     plot_dipole([dipole, dipole], show=False)
+
+    # Test IO
     dipole.write(dpl_out_fname)
     dipole_read = read_dipole(dpl_out_fname)
     assert_allclose(dipole_read.times, dipole.times, rtol=0, atol=0.00051)
@@ -74,6 +76,15 @@ def test_dipole(tmpdir, run_hnn_core_fixture):
     for dpl_key in dpl_avg.data.keys():
         assert_allclose(dpl_1[0].data[dpl_key] / 2., dpl_avg.data[dpl_key])
 
+    # Test experimental dipole
+    dipole_exp = Dipole(times, data[:, 1])
+    dipole_exp.write(dpl_out_fname)
+    dipole_exp_read = read_dipole(dpl_out_fname)
+    assert_allclose(dipole_exp.data['agg'], dipole_exp_read.data['agg'],
+                    rtol=1e-2)
+    dipole_exp_avg = average_dipoles([dipole_exp, dipole_exp])
+    assert_allclose(dipole_exp.data['agg'], dipole_exp_avg.data['agg'])
+
     # XXX all below to be deprecated in 0.3
     dpls_raw, net = run_hnn_core_fixture(backend='joblib', n_jobs=1,
                                          reduced=True, record_isoma=True,