ALLEN INSTITUTE

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

    import os
    import numpy as np
    import pandas as pd
    import json
    from six import string_types

    from bmtk.simulator.bionet.io_tools import io

    class BaseWaveform(object):
        """Abstraction of waveform class to ensure calculate method is implemented"""

        def calculate(self, simulation_time):
            raise NotImplementedError("Implement specific waveform calculation")

    class BaseWaveformType(object):
        """Specific waveform type"""
        def __init__(self, waveform_config):
            self.amp = float(waveform_config["amp"]) # units? mA?
            self.delay = float(waveform_config["del"]) # ms
            self.duration = float(waveform_config["dur"]) # ms

        def is_active(self, simulation_time):
            stop_time = self.delay + self.duration
            return self.delay < simulation_time < stop_time

    class WaveformTypeDC(BaseWaveformType, BaseWaveform):
        """DC (step) waveform"""
        def __init__(self, waveform_config):
            super(WaveformTypeDC, self).__init__(waveform_config)

        def calculate(self, t): # TODO better name
            if self.is_active(t):
                return self.amp
            else:
                return 0

    class WaveformTypeSin(BaseWaveformType, BaseWaveform):
        """Sinusoidal waveform"""
        def __init__(self, waveform_config):
            super(WaveformTypeSin, self).__init__(waveform_config)
            self.freq = float(waveform_config["freq"])  # Hz
            self.phase_offset = float(waveform_config.get("phase", np.pi))  # radians, optional
            self.amp_offset = float(waveform_config.get("offset", 0))  # units? mA? optional

        def calculate(self, t):  # TODO better name
            if self.is_active(t):
                f = self.freq / 1000. # Hz to mHz
                a = self.amp
                return a * np.sin(2 * np.pi * f * t + self.phase_offset) + self.amp_offset
            else:
                return 0

    class WaveformCustom(BaseWaveform):
        """Custom waveform defined by csv file"""
        def __init__(self, waveform_file):
            self.definition = pd.read_csv(waveform_file, sep=' ')

        def calculate(self, t):
            return np.interp(t, self.definition["time"], self.definition["amplitude"])

    class ComplexWaveform(BaseWaveform):
        """Superposition of simple waveforms"""
        def __init__(self, el_collection):
            self.electrodes = el_collection

        def calculate(self, t):
            val = 0
            for el in self.electrodes:
                val += el.calculate(t)

            return val

    # mapping from 'shape' code to subclass, always lowercase
    shape_classes = {
        'dc': WaveformTypeDC,
        'sin': WaveformTypeSin,
    }

    def stimx_waveform_factory(waveform):
        """
        Factory to create correct waveform class based on conf.
        Supports json config in conf as well as string pointer to a file.
        :rtype: BaseWaveformType
        """
        if isinstance(waveform, string_types):
            # if waveform_conf is str or unicode assume to be name of file in stim_dir
            # waveform_conf = str(waveform_conf)   # make consistent
            file_ext = os.path.splitext(waveform)
            if file_ext[-1] == '.csv':
                return WaveformCustom(waveform)

            elif file_ext[-1] == '.json':
                with open(waveform, 'r') as f:
                    waveform = json.load(f)
            else:
                io.log_warning('Unknown filetype for waveform')

        shape_key = waveform["shape"].lower()

        if shape_key not in shape_classes:
            io.log_warning("Waveform shape not known")  # throw error?

        Constructor = shape_classes[shape_key]
        return Constructor(waveform)

    def iclamp_waveform_factory(conf):
        """
        Factory to create correct waveform class based on conf.
        Supports json config in conf as well as string pointer to a file.
        :rtype: BaseWaveformType
        """
        iclamp_waveform_conf = conf["iclamp"]

        shape_key = iclamp_waveform_conf["shape"].lower()

        if shape_key not in shape_classes:
            io.log_warning('iclamp waveform shape not known')  # throw error?

        Constructor = shape_classes[shape_key]
        return Constructor(iclamp_waveform_conf)