sbi-dev / sbi

Simulation-based inference toolkit
https://sbi-dev.github.io/sbi/
Apache License 2.0
591 stars 152 forks source link

Hodgkin-Huxley model parameter estimation for observed experimental traces. #294

Closed ybernaerts closed 4 years ago

ybernaerts commented 4 years ago

Hey guys,

first of all, it's a pleasure to experiment with the sbi package. Great work on having state of the art simulator based inference algorithms in one package.

Since I have gotten only unsatisfactory results (with multiple packages) for quite some time now, I wanted to get your opinion on strategies to efficiently 1.-- create 1 compartment models and 2.-- run the inference for actually observed experimental traces. Feel free to comment on every step that I employ to infer parameters for mouse motor cortical neuron responses I see to current clamp (Patch-seq) experimental protocols.

I check the neuron cell type: Pvalb (fast spiking) e.g. and 2 channels (Na+ / K+) should in principle be suffient (I also use 'sufficient 1 compartment models from Pospischil et al.); Sst and perhaps I need an adapting K+ current (g_M in your tutorial); a bursting Vip or Sst neuron and I'd add one or two Ca++ channels (also explained in Pospischil et al.). First of all, the models in Pospischil et al. were done for ferret cortical neurons, not for mice, but I nevertheless expect these 1 compartment models to be rather sufficient to describe also mouse cortical neurons.

Then I'd check my experimentally observed membrane time constant and input resistance and derive a membrane area from it. Concretely: if I assume a specific membrane resistance c_m of 1 microF/cm*2 , I'd adapt the specific membrane resistance r_m (kOhm cm*2) so that c_m r_m gives me the experimentally observed tau (ms). With that derived r_m, the observed input resistance R_i should equal r_m * area, so I adapt my 1 compartment area so that it is satisfied. ` def syn_current(duration=800, dt=0.04, t_on = 100, curr_level = 1e-4, exp_input_res=270, exp_tau=10, seed=None):

exp_input_res in MOhms!

# exp_tau in ms!
duration = duration
t_off = duration - t_on
t = np.arange(0, duration+dt, dt)

# external current
r_m = exp_tau*1e3            # specific membrane resistance (Ohm*cm**2) (so that r_m*c_m= exp_tau, 
                             # c_m assumed to be 1 microF/cm**2)

A_soma = r_m/(exp_input_res*1e6)  # cm2
#A_soma=np.pi*((70.*1e-4)**2)
I = np.zeros_like(t)
I[int(np.round(t_on/dt)):int(np.round(t_off/dt))] = curr_level/A_soma # muA/cm2
return I, t_on, t_off, dt, t, A_soma`

What about the temperature of the experiment? Optimally one should adapt the rate constant equations with the Q10 factor so that the rate of change of open probabilities of ion channels is temperature scaled. Models in Pospischil et al. are optimised for T1 = 36 degree Celcius, my experiments are either done at T2 = 25 or 34 degree Celcius. I adapt, therefore, dx/dt = scale * ... with the scale being Q10**(T2-T1). Here x represents m, h, p, ... I hope I added this scale at the right place in your exponential Euler method (see code below).

`def HHsimulator(V0, params, dt, t, I, seed=None): """Simulates the Hodgkin-Huxley model for a specified time duration and current

    Parameters
    ----------
    V0 : float
        Voltage at first time step
    params : np.array, 1d of length dim_param
        Parameter vector
    dt : float
        Timestep
    t : array
        Numpy array with the time steps
    I : array
        Numpy array with the input current
    seed : int
    """

gbar_Na = params[0, 0]  # mS/cm2
gbar_Na.astype(float)
gbar_K = params[0, 1]  # mS/cm2
gbar_K.astype(float)
g_leak = params[0,2] # mS/cm2
g_leak.astype(float)
gbar_M = params[0,3] # mS/cm2
gbar_M.astype(float)
tau_max = params[0,3] # ms
tau_max.astype(float)
#tau_max=800
#gbar_M=0.0 # mS/cm2

# fixed parameters
#g_leak = 0.1  # mS/cm2
#gbar_M = 0.07  # mS/cm2
#tau_max = 6e2  # ms
Vt = -60.0  # mV
nois_fact = 0.1  # uA/cm2
E_leak = np.mean(voltage_obs[0:2500, curr_index])  # mV
C = 1  # uF/cm2
E_Na = 53  # mV
E_K = -107  # mV
Q10=2.5

T_1 = 36 # °C, from paper MartinPospischil et al.
T_2 = 34 # °C, experiment was actually done at 34 °C
t_adj_factor = Q10**((T_2-T_1)/10) # temperature coeff., https://en.wikipedia.org/wiki/Q10_(temperature_coefficient)
tstep = float(dt)

if seed is not None:
    rng = np.random.RandomState(seed=seed)
else:
    rng = np.random.RandomState()

####################################
# 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.0
    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.0
    return 0.128 * np.exp(-v1 / 18.0)

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

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

def beta_n(x):
    v1 = x - Vt - 10.0
    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.0
    return 1.0 / (1.0 + np.exp(-0.1 * v1))

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

####################################
# simulation from initial point
V = np.zeros_like(t)  # voltage
n = np.zeros_like(t)
m = np.zeros_like(t)
h = np.zeros_like(t)
p = np.zeros_like(t)

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, t.shape[0]):
    tau_V_inv = (
        (m[i - 1] ** 3) * gbar_Na * h[i - 1]
        + (n[i - 1] ** 4) * gbar_K
        + g_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
        + g_leak * E_leak
        + gbar_M * p[i - 1] * E_K
        + I[i - 1]
        + nois_fact * rng.randn() / (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*t_adj_factor / tau_n(V[i])))
    m[i] = m_inf(V[i]) + (m[i - 1] - m_inf(V[i])) * np.exp((-tstep*t_adj_factor / tau_m(V[i])))
    h[i] = h_inf(V[i]) + (h[i - 1] - h_inf(V[i])) * np.exp((-tstep*t_adj_factor / tau_h(V[i])))
    p[i] = p_inf(V[i]) + (p[i - 1] - p_inf(V[i])) * np.exp((-tstep*t_adj_factor / tau_p(V[i])))

return np.array(V).reshape(-1, 1)`

And then inference right? Well, no matter what I seem to opt for (BDNs, MAFs, SNLE or SNPE) I seem to either, run into posteriors that have no prior support (i.e. the summary statistics I observe lead to parameter values outside of the prior support), or I do get some posterior that gives me traces rather 'not so close' to what I experimentally observe, in terms of action potential (AP) shape, or even firing rate.

Here's an example of an Sst neuron responding to a 100 pA positive current injection. Unbenannt What I seem to struggle with is that I don't have simulations arriving at the same 'height' meaning that the derived membrane area might be wrong but even when I try to scale the simulation input current (which is the same as if one wants to scale the area), or even try to infer the scale, I get no sensible results. What about the Nernst potential for E_Na+ and E_K+? Perhaps they are dissimilar to ferret conditions. I tried to use different values for them too, however, or infer them, but to no avail.

So I'm reaching a bit of a dead end. Now, a HH-like model is highly nonlinear and a true challenge for these inference algorithms. I think it is fair to say that we've not completely solved this open problem. In Fig. 4 of "Training deep neural density estimators to 2 identify mechanistic models of neural dynamics", you've fitted Allen Institute data and, although, in the right direction, AP shapes and adapting firing patterns can be very different between high posterior samples and experimentally observed traces. I was also surprised to see that you've used simpler BDNs for experimentally observed traces and more flexible MAFs for simulated data.

Any thoughts, comments, ideas are highly appreciated!

ybernaerts commented 4 years ago

I barely believe it but at the time of writing I finally got a posterior sample which came close for the first time to an experimentally observed trace: Unbenannt_1 But, most posterior samples still did have no APs at all (I used the BDN framework, inspired by your paper, 1 round and trained on 10000 samples). Perhaps I'm slowly moving in better directions. Anyway, thoughts and ideas still highly appreciated.

ppjgoncalves commented 4 years ago

Hi Yves,

Thanks for the words of encouragement and the questions. I will here only reply to the sbi-related issues.

So I'm reaching a bit of a dead end. Now, a HH-like model is highly nonlinear and a true challenge for these inference algorithms. I think it is fair to say that we've not completely solved this open problem. In Fig. 4 of "Training deep neural density estimators to 2 identify mechanistic models of neural dynamics", you've fitted Allen Institute data and, although, in the right direction, AP shapes and adapting firing patterns can be very different between high posterior samples and experimentally observed traces.

In https://www.biorxiv.org/content/10.1101/838383v3, you are right that the high-probability (under the posterior) samples are not always in very close agreement with the experimental data from Allen Institute. In principle, this could be an indication that (1) the posterior is not well inferred or/and (2) the model is not flexible enough to fit the experimental data. At this point, we cannot tell with full certainty what is the main cause. However, we have run another algorithm (SMC-ABC) on the same cells (and with much higher number of simulations than with SNPE), and the posteriors were quite similar to the ones obtained with SNPE, suggesting that the limitations on this problem are more on the modelling side than on the inference side.

I was also surprised to see that you've used simpler BDNs for experimentally observed traces and more flexible MAFs for simulated data.

There is another important distinction between simulated data and experimental data: in the simulated data case, we ran one round SNPE, whereas we ran 2 rounds SNPE for the experimental data. Indeed, for the experimental data, we empirically found that multi-round SNPE improved substantially the quality of the posteriors. However, currently, multi-round SNPE with MAFs has some "leakage" issues (see issue https://github.com/mackelab/sbi/issues/287). Therefore, for our problem, we used multi-round SNPE with MDNs. Note that MDNs are planned to be implemented in sbi in the next few weeks.

But, most posterior samples still did have no APs at all (I used the BDN framework, inspired by your paper, 1 round and trained on 10000 samples). Perhaps I'm slowly moving in better directions.

Do you really mean 10k samples? If so, I would suggest you to increase the number of simulations. In our case, for each experimental trace, we ran 2 rounds SNPE, each round with 125000 simulations.

ybernaerts commented 4 years ago

Hi Pedro,

thanks for the fast reply and encouraging suggested adaptations!

suggesting that the limitations on this problem are more on the modelling side than on the inference side.

This might indeed suggest that, for some neurons, we need an additional compartment modelling e.g. the axon initial segment, or simply current diffustion to dendrites and axons.

However, currently, multi-round SNPE with MAFs has some "leakage" issues (see issue #287).

Does using MAFS with SNLE have the same issue?

we ran 2 rounds SNPE, each round with 125000 simulations.

I will try 2 rounds of SNPE with much more samples and MDNs.

Note that MDNs are planned to be implemented in sbi in the next few weeks.

To me it seems implemented already? Are you sure it's not in there yet?

ppjgoncalves commented 4 years ago

Hi Yves,

This might indeed suggest that, for some neurons, we need an additional compartment modelling e.g. the axon initial segment, or simply current diffustion to dendrites and axons.

Yes, absolutely. An additional compartment would perhaps be a good idea. But even before that, do increase the number of simulations for training SNPE. This might lead to substantially better results.

Does using MAFS with SNLE have the same issue?

It will not have the same issue, but it might have other issues. We have never tried SNLE nor SNRE in this problem.

To me it seems implemented already? Are you sure it's not in there yet?

Sorry, I was not very clear. You are right that MDNs are implemented. Let me be more specific. At the moment, SNPE is implemented with atomic proposals (see https://arxiv.org/abs/1905.07488), which leads to the "leakage" issues I mentioned (independently of using MDNs or MAFs). In the next few weeks, the non-atomic version of SNPE (https://arxiv.org/abs/1905.07488) will be implemented: this one does not have the serious "leakage" issues but only works for MDNs (not MAFs).

ybernaerts commented 4 years ago

Sorry, I was not very clear. You are right that MDNs are implemented. Let me be more specific. At the moment, SNPE is implemented with atomic proposals (see https://arxiv.org/abs/1905.07488), which leads to the "leakage" issues I mentioned (independently of using MDNs or MAFs). In the next few weeks, the non-atomic version of SNPE (https://arxiv.org/abs/1905.07488) will be implemented: this one does not have the serious "leakage" issues but only works for MDNs (not MAFs).

I see what you mean. 2 rounds with 100000 simulations each and these 'leakage' issues seem to occur.

But even before that, do increase the number of simulations for training SNPE. This might lead to substantially better results.

I'll keep it to one round of SNPE MDN 100000 simulations for now then :). Happy to hear you guys are working on the non-atomic version. Could it be that you posted the same reference twice? I read that reference, but perhaps I'm less aware of a reference referring to non-atomic SNPE.

Cheers!

ppjgoncalves commented 4 years ago

I'll keep it to one round of SNPE MDN 100000 simulations for now then :). Happy to hear you guys are working on the non-atomic version. Could it be that you posted the same reference twice? I read that reference, but perhaps I'm less aware of a reference referring to non-atomic SNPE.

I posted the same reference twice, and it was on purpose. ;) The paper has both types of proposals (see beginning of section 3, and sections 3.1 and 3.2 of the paper).

ybernaerts commented 4 years ago

I see. A closer look is required then. Thanks again Pedro!

I hope to post a good posterior sample here soon, to partially resolve this issue.

michaeldeistler commented 4 years ago

Hi Yves, cool that you're working with sbi! :)

I've started implementing the non-atomic version of APT. You can track progress here. If everything goes well, it will be done end of next week.

All the best and see you soon in Tübingen Michael

ybernaerts commented 4 years ago

Hey Michael,

Sounds great! I'll track you closely then ;)

See you in Tübingen indeed!

jan-matthis commented 4 years ago

https://github.com/mackelab/sbi/pull/301 got merged, I'll close this issue for now. Feel free to reopen with further questions