TST: add test for MPI and Network refactoring

Testing with MPI is a little tricky because separate processes need
to be spawned. Also hyperthreading (Intel procs) will give a false
number of cpus with multiprocessing.cpu_count. We need to use psutil
to get the number of physical cores that MPI can actually use.

.travis.yml

@@ -24,7 +24,7 @@ install:
     - source activate testenv
     - conda install --yes scipy numpy matplotlib
     - pip install flake8 pytest pytest-cov
-    - pip install mpi4py mne
+    - pip install mpi4py mne psutil
     - pip install coverage coveralls
     - |
       git clone https://github.com/neuronsimulator/nrn

hnn_core/tests/test_compare_hnn.py

@@ -1,14 +1,23 @@
-import os.path as op
+from hnn_core import simulate_dipole, Network
 
-from numpy import loadtxt
-from numpy.testing import assert_array_equal
 
-from mne.utils import _fetch_file
-import hnn_core
-from hnn_core import simulate_dipole, Params, Network
+def run_simulation(params, dpl_master):
+    net = Network(params)
+    dpls = simulate_dipole(net)
+
+    return dpls
 
 
 def test_hnn_core():
+    import os.path as op
+
+    from numpy import loadtxt
+    from numpy.testing import assert_array_equal
+
+    from mne.utils import _fetch_file
+    import hnn_core
+    from hnn_core import Params
+
     """Test to check if MNE neuron does not break."""
     # small snippet of data on data branch for now. To be deleted
     # later. Data branch should have only commit so it does not
@@ -21,15 +30,39 @@ def test_hnn_core():
 
     hnn_core_root = op.join(op.dirname(hnn_core.__file__), '..')
 
+    # default params
     params_fname = op.join(hnn_core_root, 'param', 'default.json')
     params = Params(params_fname)
 
-    net = Network(params)
-    dpl = simulate_dipole(net)[0]
+    # run the simulation
+    dpl = run_simulation(params, dpl_master)[0]
 
+    # write the dipole to a file and compare
     fname = './dpl2.txt'
     dpl.write(fname)
 
     dpl_pr = loadtxt(fname)
     assert_array_equal(dpl_pr[:, 2], dpl_master[:, 2])  # L2
     assert_array_equal(dpl_pr[:, 3], dpl_master[:, 3])  # L5
+
+
+if __name__ == '__main__':
+    # started as an MPI child from test_mpi_simulation.py
+
+    from mpi4py import MPI
+    comm = MPI.Comm.Get_parent()
+
+    # receive params and dpl_master
+    (params, dpl_master) = comm.bcast(comm.Get_rank(), root=0)
+
+    # run the simulation
+    dpl = run_simulation(params, dpl_master)[0]
+
+    if comm.Get_rank() == 0:
+        # send results back to parent
+        comm.send(dpl, dest=0)
+
+    comm.Barrier()
+    comm.Disconnect()
+    MPI.Finalize()
+    exit(0)  # stop the child

hnn_core/tests/test_mpi_simulation.py

@@ -0,0 +1,61 @@
+import os.path as op
+
+from numpy import loadtxt
+from numpy.testing import assert_array_equal
+
+from mne.utils import _fetch_file
+import hnn_core
+from hnn_core import Params
+
+
+def test_simulate_dipole_mpi():
+    """Test to check if simulate dipole can be called with MPI."""
+    from mpi4py import MPI
+    from psutil import cpu_count
+
+    # need to get physical CPU cores using psutil and leave one CPU
+    # for the current script (bug in OpenMPI < 4.0) or MPI will fail
+    # to spawn nrniv workers
+    n_core = cpu_count(logical=False) - 1
+    if n_core < 1:
+        n_core = 1
+
+    # spawn NEURON sim
+    args = ['-python', '-mpi', '-nobanner', 'python',
+            op.join(op.dirname(hnn_core.__file__), 'tests',
+                    'test_compare_hnn.py')]
+
+    mpiinfo = MPI.Info().Create()
+    mpiinfo.Set('env', 'OMPI_MCA_btl=self,vader,tcp')
+    mpiinfo.Set('env', 'OMPI_MCA_btl_base_warn_component_unused=0')
+    mpiinfo.Set('env', 'OMPI_MCA_rmaps_base_oversubscribe=1')
+    child = MPI.COMM_SELF.Spawn('nrniv', args=args, maxprocs=n_core,
+                                info=mpiinfo)
+
+    # get dipole data and params
+    data_url = ('https://raw.githubusercontent.com/hnnsolver/'
+                'hnn-core/test_data/dpl.txt')
+    if not op.exists('dpl.txt'):
+        _fetch_file(data_url, 'dpl.txt')
+    dpl_master = loadtxt('dpl.txt')
+
+    hnn_core_root = op.join(op.dirname(hnn_core.__file__), '..')
+
+    params_fname = op.join(hnn_core_root, 'param', 'default.json')
+    params = Params(params_fname)
+
+    # send params and dipole data to spawned nrniv procs to start sim
+    simdata = (params, dpl_master)
+    child.bcast(simdata, root=MPI.ROOT)
+
+    # wait to recevie results from child rank 0
+    dpl = child.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG)
+    child.Barrier()
+    child.Disconnect()
+
+    fname = './dpl2.txt'
+    dpl.write(fname)
+
+    dpl_pr = loadtxt(fname)
+    assert_array_equal(dpl_pr[:, 2], dpl_master[:, 2])  # L2
+    assert_array_equal(dpl_pr[:, 3], dpl_master[:, 3])  # L5

hnn_core/tests/test_network.py

@@ -13,7 +13,7 @@ def test_network():
     params_fname = op.join(hnn_core_root, 'param', 'default.json')
     params = Params(params_fname)
     net = Network(deepcopy(params))
-    net.build_in_neuron()
+    net.build_in_neuron()  # needed to populate net.cells
     for p in params:
         assert params[p] == net.params[p]
     assert len(params) == len(net.params)