ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.modules.ecp

    # Copyright 2017. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import os
    import h5py
    import math
    import pandas as pd
    from neuron import h
    import numpy as np

    from bmtk.simulator.bionet.modules.sim_module import SimulatorMod
    from bmtk.utils.sonata.utils import add_hdf5_magic, add_hdf5_version

    pc = h.ParallelContext()
    MPI_RANK = int(pc.id())
    N_HOSTS = int(pc.nhost())

    class EcpMod(SimulatorMod):
        def __init__(self, tmp_dir, file_name, electrode_positions, contributions_dir=None, cells=None, variable_name='v',
                     electrode_channels=None, cell_bounds=None, minimum_distance=None):
            self._ecp_output = file_name if os.path.isabs(file_name) else os.path.join(tmp_dir, file_name)
            self._positions_file = electrode_positions
            self._tmp_outputdir = tmp_dir

            if contributions_dir is not None:
                self._save_individ_cells = True
                self._contributions_dir = contributions_dir if os.path.isabs(contributions_dir) else os.path.join(tmp_dir, contributions_dir)
            else:
                self._save_individ_cells = False
                self._contributions_dir = None

            if cell_bounds is None:
                self._cells = cells
            else:
                self._cells = list(np.arange(cell_bounds[0],cell_bounds[1]+1))
            if minimum_distance and minimum_distance != 'auto':
                try:
                    minimum_distance = max(float(minimum_distance), 0)
                except Exception as e:
                    minimum_distance = 0
            self.minimum_distance = minimum_distance
            self._rel = None
            self._fih1 = None
            self._rel_nsites = 0
            self._block_size = 0
            self._saved_gids = {}
            self._nsteps = 0

            self._tstep = 0  # accumlative time step
            # self._rel_time = 0  #
            self._block_step = 0  # time step within the given block of time
            self._tstep_start_block = 0
            self._data_block = None
            self._cell_var_files = {}

            self._tmp_ecp_file = self._get_tmp_fname(MPI_RANK)
            self._tmp_ecp_handle = None
            # self._tmp_ecp_dataset = None

            self._local_gids = []

        def _get_tmp_fname(self, rank):
            return os.path.join(self._tmp_outputdir, 'tmp_{}_{}'.format(MPI_RANK, os.path.basename(self._ecp_output)))

        def _create_ecp_file(self, sim):
            dt = sim.dt
            tstop = sim.tstop
            self._nsteps = int(round(tstop/dt))

            # create file to temporary store ecp data on each rank
            self._tmp_ecp_handle = h5py.File(self._tmp_ecp_file, 'a')
            self._tmp_ecp_handle.create_dataset('/ecp/data', (self._nsteps, self._rel_nsites), maxshape=(None, self._rel_nsites),
                                                dtype=np.float64,
                                                chunks=True)

            # only the primary node will need to save the final ecp
            if MPI_RANK == 0:
                with h5py.File(self._ecp_output, 'w') as f5:
                    add_hdf5_magic(f5)
                    add_hdf5_version(f5)
                    f5.create_dataset('/ecp/data', (self._nsteps, self._rel_nsites), maxshape=(None, self._rel_nsites),
                                      dtype=np.float64,
                                      chunks=True)
                    f5['/ecp/data'].attrs['units'] = 'mV'
                    #f5.attrs['dt'] = dt
                    #f5.attrs['tstart'] = 0.0
                    #f5.attrs['tstop'] = tstop

                    f5.create_dataset('/ecp/time', (3,), data=(0.0, sim.tstop, sim.dt))
                    f5['/ecp/time'].attrs['units'] = 'ms'

                    # Save channels. Current we record from all channels, may want to be more selective in the future.
                    f5.create_dataset('/ecp/channel_id', data=np.arange(self._rel.nsites))

            pc.barrier()

        def _create_cell_file(self, gid):
            if not self._save_individ_cells:
                return

            file_name = os.path.join(self._contributions_dir, '{}.h5'.format(int(gid)))
            file_h5 = h5py.File(file_name, 'a')
            self._cell_var_files[gid] = file_h5
            file_h5.create_dataset('/ecp/data', (self._nsteps, self._rel_nsites), maxshape=(None, self._rel_nsites), chunks=True)
            # self._cell_var_files[gid] = file_h5['ecp']

        def _calculate_ecp(self, sim):
            self._rel = RecXElectrode(self._positions_file, self.minimum_distance)
            for gid in self._local_gids:
                cell = sim.net.get_cell_gid(gid)
                #cell = sim.net.get_local_cell(gid)
                # cell = sim.net.cells[gid]
                self._rel.calc_transfer_resistance(gid, cell.seg_coords)

            self._rel_nsites = self._rel.nsites
            sim.h.cvode.use_fast_imem(1)  # make i_membrane_ a range variable

            def set_pointers():
                for gid, cell in sim.net.get_local_cells().items():
                #for gid, cell in sim.net.local_cells.items():
                    # for gid, cell in sim.net.cells.items():
                    cell.set_im_ptr()
            self._fih1 = sim.h.FInitializeHandler(0, set_pointers)

        def _save_block(self, interval):
            """Add """
            itstart, itend = interval
            self._tmp_ecp_handle['/ecp/data'][itstart:itend, :] += self._data_block[0:(itend - itstart), :]
            self._tmp_ecp_handle.flush()
            self._data_block[:] = 0.0

        def _save_ecp(self, sim):
            """Save ECP from each rank to disk into a single file"""
            block_size = sim.nsteps_block
            nblocks, remain = divmod(self._nsteps, block_size)
            ivals = [i*block_size for i in range(nblocks+1)]
            if remain != 0:
                ivals.append(self._nsteps)

            for rank in range(N_HOSTS):  # iterate over the ranks
                if rank == MPI_RANK:  # wait until finished with a particular rank
                    with h5py.File(self._ecp_output, 'a') as ecp_f5:
                        for i in range(len(ivals)-1):
                            ecp_f5['/ecp/data'][ivals[i]:ivals[i+1], :] += self._tmp_ecp_handle['/ecp/data'][ivals[i]:ivals[i+1], :]

                pc.barrier()

        def _save_cell_vars(self, interval):
            itstart, itend = interval

            for gid, data in self._saved_gids.items():
                h5_file = self._cell_var_files[gid]
                h5_file['/ecp/data'][itstart:itend, :] = data[0:(itend-itstart), :]
                h5_file.flush()
                data[:] = 0.0

        def _delete_tmp_files(self):
            if os.path.exists(self._tmp_ecp_file):
                self._tmp_ecp_handle.close()
                os.remove(self._tmp_ecp_file)

        def initialize(self, sim):
            if self._contributions_dir and (not os.path.exists(self._contributions_dir)) and MPI_RANK == 0:
                os.makedirs(self._contributions_dir)
            pc.barrier()

            self._block_size = sim.nsteps_block

            # Get list of gids being recorded
            selected_gids = set(sim.net.get_node_set(self._cells).gids())
            self._local_gids = list(set(sim.biophysical_gids) & selected_gids)

            self._calculate_ecp(sim)
            self._create_ecp_file(sim)

            # ecp data
            self._data_block = np.zeros((self._block_size, self._rel_nsites))

            # create list of all cells whose ecp values will be saved separetly
            self._saved_gids = {gid: np.empty((self._block_size, self._rel_nsites))
                                for gid in self._local_gids}
            for gid in self._saved_gids.keys():
                self._create_cell_file(gid)

            pc.barrier()

        def step(self, sim, tstep):
            for gid in self._local_gids:  # compute ecp only from the biophysical cells
                cell = sim.net.get_cell_gid(gid)
                im = cell.get_im()
                tr = self._rel.get_transfer_resistance(gid)
                ecp = np.dot(tr, im)

                if gid in self._saved_gids.keys():
                    # save individual contribution
                    self._saved_gids[gid][self._block_step, :] = ecp

                # add to total ecp contribution
                self._data_block[self._block_step, :] += ecp

            self._block_step += 1

        def block(self, sim, block_interval):
            self._save_block(block_interval)
            if self._save_individ_cells:
                self._save_cell_vars(block_interval)

            self._block_step = 0
            self._tstep_start_block = self._tstep

        def finalize(self, sim):
            if self._block_step > 0:
                # just in case the simulation doesn't end on a block step
                self.block(sim, (sim.n_steps - self._block_step, sim.n_steps))

            self._save_ecp(sim)
            self._delete_tmp_files()
            pc.barrier()

    class RecXElectrode(object):
        """Extracellular electrode

        """

        def __init__(self, positions, minimum_distance):
            """Create an array"""
            # self.conf = conf
            electrode_file = positions  # self.conf["recXelectrode"]["positions"]

            # convert coordinates to ndarray, The first index is xyz and the second is the channel number
            el_df = pd.read_csv(electrode_file, sep=' ')
            self.pos = el_df[['x_pos', 'y_pos', 'z_pos']].T.values
            #self.pos = el_df.as_matrix(columns=['x_pos', 'y_pos', 'z_pos']).T
            self.nsites = self.pos.shape[1]
            # self.conf['run']['nsites'] = self.nsites  # add to the config
            self.transfer_resistances = {}  # V_e = transfer_resistance*Im
            self.minimum_distance = minimum_distance

        def drift(self):
            # will include function to model electrode drift
            pass

        def get_transfer_resistance(self, gid):
            return self.transfer_resistances[gid]

        def calc_transfer_resistance(self, gid, seg_coords):
            """Precompute mapping from segment to electrode locations"""
            sigma = 0.3  # mS/mm

            r05 = (seg_coords.p0 + seg_coords.p1) / 2
            dl = seg_coords.p1 - seg_coords.p0
            if self.minimum_distance == 'auto':
                rm = seg_coords.d / 2  # set minimum distance to the radius of each segment
            else:
                rm = self.minimum_distance  # set a common minimum distance

            nseg = r05.shape[1]

            tr = np.zeros((self.nsites, nseg))

            for j in range(self.nsites):  # calculate mapping for each site on the electrode
                rel = np.expand_dims(self.pos[:, j], axis=1)  # coordinates of a j-th site on the electrode
                rel_05 = rel - r05  # distance between electrode and segment centers

                # compute dot product column-wise, the resulting array has as many columns as original
                r2 = np.einsum('ij,ij->j', rel_05, rel_05)

                # compute dot product column-wise, the resulting array has as many columns as original
                rlldl = np.einsum('ij,ij->j', rel_05, dl)
                dlmag = np.linalg.norm(dl, axis=0)  # length of each segment
                rll = abs(rlldl / dlmag)  # component of r parallel to the segment axis it must be always positive
                rT2 = r2 - rll ** 2  # square of perpendicular component
                up = rll + dlmag / 2
                low = rll - dlmag / 2
                if self.minimum_distance:
                    # if electrode projection on segment axis is within distance rm to the segment
                    # check perpendicular distance to the segment axis and set lower limit to rm
                    np.fmax(rT2, rm * rm, out=rT2, where=low < rm)
                num = up + np.sqrt(up ** 2 + rT2)
                den = low + np.sqrt(low ** 2 + rT2)
                tr[j, :] = np.log(num / den) / dlmag  # units of (um) use with im_ (total seg current)

            tr *= 1 / (4 * math.pi * sigma)
            self.transfer_resistances[gid] = tr