Skip to content

Inferring Neuronal Properties

Kaiwen Sheng edited this page Nov 22, 2022 · 2 revisions

Setup Model Hypothesis

The model hypothesis is the surrogate of the real neural structure we are targeting. The free parameters of the model are the neural properties to be inferred. Domain Adaptive Neural Inference implements an abstract class for models: BaseGenerator. It takes 5 parameters to initiate an instance: number of sequences, time length of the simulation, time step of the simulation, the number of free parameters to be inferred, the number of features to be extracted from the model responses.

The BaseGenerator has a shared funtion: gen, which takes a batch of free params and stimulus to run the simulation in parallel. It uses pymp to parallel the simuation, and users need to specify the number of threads in this function.

gen function calls gen_single, which is the funtion to simulate the model with an individual set of free parameters and stimulus. gen_single is the user-defined function to describe the dynamics of the model. To ensure it runs correctly in parallel, please avoid using any shared global parameters.

After setting up the gen_single function, users need to define the construct_pilot function which will generate the initial synthesized dataset for training the deep neural network.

For each model inherited from BaseGenerator, it should have param_ranges variable, which is a nd.array with shape (2, param_num). The variable contains the min and max values for each free parameter in the model.

We use a single-compartment Hodgkin-Huxley (HH) model as the model hypothesis, which has been demonstrated to capture a wide range of cortical neuronal dynamics (Pospischil et al., 2008).

The dynamics of the model are as the following:

$$ \begin{align} C_m \frac{dV}{dt} &= \bar g_{\text{Na}}m^3 h(E_\text{Na} - V) + \bar g_\text{K} n^4 (E_\text{K} - V) + \bar g_\text{M} (E_\text{K} - V) + g_L (E_L - V) + I(t) + \sigma \eta(t) \\ \frac{dq}{dt} &= \frac{q_\infty - q}{\tau_q (V)}, q \in {m, h, n, p} \end{align} $$

Register Neural Properties to be Inferred

The neural properties to be inferred in this application are: passive membrane properties ( $C_m$, $g_L$ ), ion channels properties ( $\bar g_\text{Na}$, $\bar g_\text{K}$, $\bar g_\text{M}$, $V_\text{T}$, $\tau_\max$ ), and noise factors ( $\sigma$ ). Firstly, we need to manually set the lower and upper bound of these parameters in generator/HodgkinHuxley.py:

prior_min = np.array([[0.5, 1e-4, 1e-4, 1e-4, 50,   -90,  1e-4, -100]])
prior_max = np.array([[80., 15,   0.6,  0.6,  3000, -40., 0.15, -30]])
self.param_ranges = np.concatenate([prior_min, prior_max], axis=0)[:, :param_num]

Implement Simulation

The simulation is implemented in HodgkinHuxley.gen_single, where it calls gen_single in the same file, which uses numba to accelerate the simulation. The implementation of the HH model is adopted from the tutorial of delfi. HodgkinHuxley.gen_single will return the simulated trace of the model as well as the feature of trace.

# Hodgkin Huxley
def gen_single(self, params, V_0=None):
    stimulus = self.get_stimulus()
    rng = np.random.RandomState()

    rest_len = 1000
    noise = rng.randn(self.time_len + rest_len)
    stimulus = np.concatenate([np.zeros((self.seq_num, rest_len), np.float32), stimulus], axis=1)
    V_0 = self.V_0 if V_0 is None else V_0
    dynamic_data = np.zeros((self.seq_num, self.time_len))
    feature_data = np.zeros((self.feature_num))
    for i in range(self.seq_num):
    dynamic_data[i] = gen_single(stimulus[i], params, self.dt, V_0, noise)[rest_len:].astype(np.float32)
    feature_data = extract_features(dynamic_data)
    return dynamic_data, feature_data

from numba import njit
@njit
def gen_single(stimulus, params, dt, V_0, noise):
    param_num = params.shape[0]
    assert param_num in [3, 5, 8]
    """ 3 Params Model """
    gbar_Na = params[0]  # mS/cm2
    gbar_K = params[1]  # mS/cm2
    gbar_leak = params[2]  # mS/cm2

    if param_num >= 5:
        """ 5 Params Model """
        gbar_M = params[3]  # mS/cm2
        tau_max = params[4]  # ms
    else:
        gbar_M = 0.07  # mS/cm2
        tau_max = 6e2  # ms

    if param_num >= 7:
        """ 7 Params Model """
        Vt = params[5]  # mV
        nois_fact = params[6]  # uA/cm2
        E_leak = params[7]  # mV
    else:
        Vt = -60.  # mV
        nois_fact = 0.1  # uA/cm2
        E_leak = -70.  # mV

    # fixed parameters
    C = 1.  # uF/cm2
    E_Na = 53  # mV
    E_K = -107  # mV
    V0 = V_0  # mV

    tstep = dt

    ####################################
    # kinetics
    def efun(z):
        if np.abs(z) < 1e-4:
            return 1 - z / 2
        else:
            return z / (np.exp(z) - 1)

    def alpha_m(x):
        v1 = x - Vt - 13.
        return 0.32 * efun(-0.25 * v1) / 0.25

    def beta_m(x):
        v1 = x - Vt - 40
        return 0.28 * efun(0.2 * v1) / 0.2

    def alpha_h(x):
        v1 = x - Vt - 17.
        return 0.128 * np.exp(-v1 / 18.)

    def beta_h(x):
        v1 = x - Vt - 40.
        return 4.0 / (1 + np.exp(-0.2 * v1))

    def alpha_n(x):
        v1 = x - Vt - 15.
        return 0.032 * efun(-0.2 * v1) / 0.2

    def beta_n(x):
        v1 = x - Vt - 10.
        return 0.5 * np.exp(-v1 / 40)

    # steady-states and time constants
    def tau_n(x):
        return 1 / (alpha_n(x) + beta_n(x))

    def n_inf(x):
        return alpha_n(x) / (alpha_n(x) + beta_n(x))

    def tau_m(x):
        return 1 / (alpha_m(x) + beta_m(x))

    def m_inf(x):
        return alpha_m(x) / (alpha_m(x) + beta_m(x))

    def tau_h(x):
        return 1 / (alpha_h(x) + beta_h(x))

    def h_inf(x):
        return alpha_h(x) / (alpha_h(x) + beta_h(x))

    # slow non-inactivating K+
    def p_inf(x):
        v1 = x + 35.
        return 1.0 / (1. + np.exp(-0.1 * v1))

    def tau_p(x):
        v1 = x + 35.
        return tau_max / (3.3 * np.exp(0.05 * v1) + np.exp(-0.05 * v1))

    time_len = stimulus.shape[-1]
    ####################################
    # simulation from initial point
    V = np.zeros(time_len)  # voltage
    n = np.zeros(time_len)
    m = np.zeros(time_len)
    h = np.zeros(time_len)
    p = np.zeros(time_len)

    V[0] = float(V0)
    n[0] = n_inf(V[0])
    m[0] = m_inf(V[0])
    h[0] = h_inf(V[0])
    p[0] = p_inf(V[0])

    for i in range(1, time_len):
        tau_V_inv = ((m[i - 1]**3) * gbar_Na * h[i - 1] +
                     (n[i - 1]**4) * gbar_K + gbar_leak + gbar_M * p[i - 1]) / C
        V_inf = (
            (m[i - 1]**3) * gbar_Na * h[i - 1] * E_Na +
            (n[i - 1]**4) * gbar_K * E_K + gbar_leak * E_leak +
            gbar_M * p[i - 1] * E_K + stimulus[i - 1] + nois_fact * noise[i] /
            (tstep**0.5)) / (tau_V_inv * C)
        V[i] = V_inf + (V[i - 1] - V_inf) * np.exp(-tstep * tau_V_inv)
        n[i] = n_inf(
            V[i]) + (n[i - 1] - n_inf(V[i])) * np.exp(-tstep / tau_n(V[i]))
        m[i] = m_inf(
            V[i]) + (m[i - 1] - m_inf(V[i])) * np.exp(-tstep / tau_m(V[i]))
        h[i] = h_inf(
            V[i]) + (h[i - 1] - h_inf(V[i])) * np.exp(-tstep / tau_h(V[i]))
        p[i] = p_inf(
            V[i]) + (p[i - 1] - p_inf(V[i])) * np.exp(-tstep / tau_p(V[i]))

    return V

Register Stimulus Protocols

The we setup the stimulus protocol for the model to be a step current with the amplitude of 0.13 nA for 1s. The membrane area of the neuron was set to be identical for all neurons, which is 8.315e-5 cm2. The timestep used here is 0.2 ms.

amp = 0.13        
stimulus = np.zeros((seq_num, time_len))
stimulus[0, 600:5600] = amp * 1e-3 / (0.0153 / 184)

Define Biophysical Features for Similarity Calculation

The biophysical features calculated from the neural responses are used to calculate how similar the simulated neuron responses are to the experimental data. The features are task-dependent. In this application, there are 13 biophysical features. The detailed description of the features can be found in (Sheng et al., 2022) and the generator/HodgkinHuxley.py: extract_features(dynamic_data).

The similarity is calculated based on the z-scored mean absolute error (MAE) between features of experimental data and synthesized data.

$$ \sum_i{\frac{\left|fea^{\text{syn}}_i - fea^{\text{exp}}_i\right|}{\sigma_i}} $$

where $\sigma_i$ is the minimal tolerance of the difference of the feature, see Register Experimental Data for details.

Construct Synthesized Dataset

To generate the pilot synthesized dataset of the HH model, we implement the construct_pilot function to save the data in a h5 file format. The h5 data should contain the simulated model responses, stored by dynamic_data, the corresponding free parameters, stored by param_data, and the features for similarity calculation, stored by feature_data. We apply a data-washing procedure to get rid of models without spikes or with spontaneous firing before stimulus onset.

def construct_pilot(self, num, store_file, NUM_THREAD=1):
    cnt = 0
    param_arr = np.zeros((num, self.param_num), dtype=np.float32)
    dynamic_arr = np.zeros((num, self.seq_num, self.time_len), dtype=np.float32)
    feature_arr = np.zeros((num, self.feature_num), dtype=np.float32)

    while cnt < num:
        param = np.random.uniform(self.param_ranges[0], self.param_ranges[1], size=(num, self.param_num)).astype(np.float32)

        dynamic, feature = self.gen(num-cnt, param, NUM_THREAD)
        good_idx = np.arange(dynamic.shape[0]).astype(np.int32)

        # for data wash
        if True:
            spike_idx = np.argwhere(feature[..., 0] < 1)[:, 0]
            spon_idx = np.argwhere(feature[..., -1] > 0)[:, 0]
            wash_idx = np.unique(np.concatenate([spike_idx, spon_idx]))
            good_idx = np.delete(np.arange(dynamic.shape[0]), wash_idx).astype(np.int32)

            _end = min(num, cnt + good_idx.shape[0])
            param_arr[cnt:_end] = param[good_idx[:_end - cnt]]
            dynamic_arr[cnt:_end] = dynamic[good_idx[:_end - cnt]]
            feature_arr[cnt:_end] = feature[good_idx[:_end - cnt]]
            cnt = _end

    param_arr -= self.param_ranges[0]
    param_arr /= self.param_ranges[1] - self.param_ranges[0]

    with h5.File(store_file, 'w') as file:
        file.create_dataset('dynamic_data', data=dynamic_arr)
        file.create_dataset('param_data', data=param_arr)
        file.create_dataset('feature_data', data=feature_arr)

Finally run the following command in generator directory to generate the synthesized dataset. Please change the number of threads based on your own machine.

python run.py data -t 7000 -d 0.2 -num 1000000 -seq 1 -param_num 8 -feature_num 13 -threads 100 -store_file /your/data/path/HH_8_param.h5 -neural_model HH

Register Experimental Data

Another half of the inputs to Domain Adaptive Neural Inference is the experimental dataset. The dataset is defined by setting up construct_exp in BaseGenerator, which will return 4 variables: experimental data, the features of the experimental data, the zscore factors to calculate feature similarity. For inference on neuronal or microcircuit properties in (Sheng et al., 2022), the features and factors are used to select similar samples.

The experimental data we used in (Sheng et al., 2022) are obtained from Allen Cell Types Database. We chose experimental data with long squared stimulus protocol at 130 pA and eliminated those without spikes under this stimulus protocol. So we got 1052 data in total. The data are stored in a h5 file, and we should implement the generator/HodgkinHuxley.py: construct_exp function to get access to the experimental data.

def construct_exp(self):
    with h5.File('/you/data/path/AllenData_130pA.h5', 'r') as f:
        exp_dynamic_data = f.get('dynamic_data')[:]
        if len(exp_dynamic_data.shape) == 2:
        	exp_dynamic_data = exp_dynamic_data[:, None]
        exp_feature_data = f.get('feature_data')[:]
        exp_feature_data[exp_feature_data == -1e8] = 0
        feature_ratio = np.array([0.5, 2, 2, 2, 5e-2, 1e-4, 5e-3, 1e-3, 1e-2, 5e-4, 1e-3, 2, 100])
	return exp_dynamic_data, exp_feature_data, feature_ratio

The feature_ratio here denotes the minimal tolerance of each feature (i.e., $\sigma$).

Run!

So far, we've already prepared everythingto infer the neural properties. Now, we need to register the neural simulator we've constructed in main.py.

def construct_data(model_name, seq_num, time_len, dt, param_num, feature_num):
    if model_name == 'HH':
        from generator.HodgkinHuxley import HodgkinHuxley
        model = HodgkinHuxley(seq_num, time_len, dt, param_num, feature_num, V_0=-65)
    else:
        raise NotImplementedError('Only accept model: HH, but got {}'.format(model_name))
        
    exp_dynamic_data, exp_feature_data, feature_ratio = model.construct_exp()
        
    return model, exp_dynamic_data, exp_feature_data, feature_ratio

Then run the command line:

CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --node_rank=0 --nnodes=1 --master_port=12345 --master_addr=localhost main.py -t 7000 -b 50 -seq 1 -e 100 -dir /you/save/path -lr 1e-4 -data /you/data/path/HH_8_param.h5 -size 1000000 -syn_model HH -out 8 -fdim 13 -dynamic_key dynamic_data -target_key param_data -feature_key feature_data -subdata 0 -cnn_type resnet -time_type lstm -cnn_config 50 -st 1 -st_thre 3 -st_bank_size 50 -da 1 -da_type DANN -da_n_layer 1 -da_layer_offset 0 -da_factor 0.0001

You can view the intermediate results via tensorboard or in your /you/save/path.

References

Pospischil, M., Toledo-Rodriguez, M., Monier, C., Piwkowska, Z., Bal, T., Frégnac, Y., Markram, H. and Destexhe, A., 2008. Minimal Hodgkin–Huxley type models for different classes of cortical and thalamic neurons. Biological cybernetics, 99(4), pp.427-441.