Source code for main

# -*- coding: utf-8 -*-
r"""
Run a parallel ring network of ball-and-stick cells and write the spike times
to a file.

To execute with MPI across 2 processors::
       
    $ mpiexec -n 2 python main.py
    
To modify any simulation parameters, specify individual variables encased in
a string to a sim_var dictionary variable::

    $ mpiexec -n 2 python main.py "sim_var['N']=6" \
       "sim_var['spike_out_file']='N6.spk'"
       
To visualize a few cells, playing back from a recorded spike file, specify
a list of cells in the ``playback`` variable and run from an interactive
python console::

    $ python -i main.py "sim_var['playback']=[0,3]"
    
.. Note::
    
    This code requires ``numpy`` and ``mpi4py`` and assumes NEURON is
    also installed in the Python environment.
    
.. seealso::
    
    Hines M.L. and Carnevale N.T, 
    Translating network models to parallel hardware in NEURON,
    Journal of Neuroscience Methods 169 (2008) 425-455).

AUTHORS:

- THOMAS MCTAVISH (2010-11-04): initial version.

"""
import sys, os

def process_args(in_argv):
    """Process additional arguments on command line. Any Python statement can
    be executed, but arguments need to be formatted as strings and each one
    in the form ``"sim_var['<var>']=x"`` will run
    the simulation with a modification of parameters."""
    sim_var = {'N' : 5, # Number of cells
            'stim_w' : 0.004, # Stimulus weight
            'stim_spike_num' : 1, # Number of spikes in the stimulator
            'syn_w' : 0.01, # Synaptic weights
            'syn_delay' : 5, # Synaptic delay
            'spike_out_file' : 'out.spk', # Spike output file
            'playback' : [], # List of cells to playback in non-parallel mode
            }

    if in_argv is not None:
        global_dict = globals()
        local_dict = locals()
        
        for arg in in_argv:
            try:
                exec(arg, global_dict, local_dict)
            except:
                print "WARNING: Do not know how to process statement \
                        \'{0}\'".format(item)
                pass
        if 'sim_var' in local_dict:
            # Replace default values with those in the local_dict
            sim_var.update(local_dict['sim_var'])

    return sim_var

def plot_vecs(vecs, ids):
    """Plot time series data.
    
    :param vecs: Dict containing vectors. Assumes there is at least a 't' key.
    :param ids: Cell ids to plot.
    """
    from matplotlib import pyplot
    
    fig = pyplot.figure()
    ax = fig.add_axes([.1,.1,.8,.8])
    plts = []
    legtext = []
    for idx in ids:
        plts.extend(ax.plot(vecs['t'], vecs[idx]))
        legtext.append("cell %d" %(idx))
    ax.legend(plts, legtext)
    pyplot.show()
        
def run(argv = None):
    """Run a ring network simulation. Additional arguments on command line can
    be any Python statement to execute, but strings in the form 
    ``"sim_var['<var>']=x"`` will run the simulation with a modification of 
    those parameters."""
    sim_var = process_args(argv)
    
    if len(sim_var['playback']) == 0:
        try:
            from mpi4py import MPI # Must come before importing NEURON
        except ImportError as (strerror):
            print "*****\nWARNING: {0}:\n  Running in serial mode.\n*****" \
            .format(strerror)
            pass
                

    from neuron import h as nrn
    import ring

    nrn.load_file("stdrun.hoc") # For when we run simulations

    pc = nrn.ParallelContext()

    net = ring.Ring(sim_var['N'], \
            sim_var['stim_w'], \
            sim_var['stim_spike_num'], \
            sim_var['syn_w'], \
            sim_var['syn_delay'], \
            sim_var['playback'])

    pc.set_maxstep(10)
    nrn.stdinit()
    nrn.dt = 0.025 # Fixed dt
    nrn.tstop = 100

    if len(sim_var['playback']) > 0:
        import numpy
        from itertools import izip
        rec = {}
        for (cell, gid) in izip(net.cells, net.playlist):
            vec = nrn.Vector()
            vec.record(cell.soma(0.5)._ref_v)
            rec[gid] = vec

        vec = nrn.Vector()
        vec.record(nrn._ref_t)
        rec['t'] = vec
    
        npspikevec, npgidvec = numpy.loadtxt('sorted_' + 
                                             sim_var['spike_out_file'], 
                                             unpack=True)
        
        spikevec = nrn.Vector()
        gidvec = nrn.Vector()
        
        patternstim = nrn.PatternStim()
        patternstim.play(spikevec.from_python(npspikevec), 
                         gidvec.from_python(npgidvec))
        nrn.run()
        plot_vecs(rec, sim_var['playback'])
    else:
        pc.psolve(100)
        net.write_spikes(sim_var['spike_out_file'])
        pc.runworker()    
        pc.done()
    
        # After we are done, re-sort the file by spike times.
        exec_cmd = 'sort -k 1n,1n -k 2n,2n ' + sim_var['spike_out_file'] + \
                ' > ' + 'sorted_' + sim_var['spike_out_file']
        os.system(exec_cmd)
        nrn.quit()
        
if __name__ == '__main__':
    """
    Runs when called from the command line. Optional arguments may specify 
    different param files and Python commands.
    """
    idx=1
    for item in sys.argv:
        if item.startswith('main'): # Ignore everything leading up to ourselves
            break
        idx += 1
    run(sys.argv[idx:])