GSoC 2021: Support of the simulation-based inference with the model fitting toolbox

[akapet00]

This pull request summarizes all the work done during the Google Summer of Code (GSoC) 2021 with International Neuroinformatics Coordinating Facility (INCF) for the project:

Support of the simulation-based inference with the model fitting toolbox

Synopsis

Unlike traditional inverse identification tools that rely on gradient and gradient-free methods, simulation-based inference has been established as the powerful alternative approach that yields twofold improvement over such methods. Firstly, it does not only result in a single set of optimal parameters, rather simulation-based inference acts as if the actual statistical inference is performed and provides an estimate of the full posterior distribution over parameters. Secondly, it exploits prior system knowledge sparsely, using only the most important features to identify mechanistic models which are consistent with the measured data. The aim of the project is to support the simulation-based inference in the brian2modelfitting toolbox by linking it to the sbi [1], PyTorch-based library for simulation-based inference, development of which is coordinated at the Macke lab.

Main results

The entire work done during the GSoC is presented in great detail on the blog on my personal webpage, which is available here. This blog post covers everything: from the work done during the community bonding period, over weekly reports that outline the work done weekly in great detail, to my own experience during the entire period of the project.

Summary of the main results (under each major contribution, there is a short code snippet that describes it):

• development of the Inferencer class that offers an API very similar to that of the Fitter class, but instead of fitting, a neural density estimator is trained under the hood. - PR #49
  This density estimator learns probabilistic mapping between the sampled prior over unknown parameters, and the output data.
  Output data may be either a list of summary features or the time-series of the dynamics of observed state variable(s) in the model. In the case of the latter, the automatic feature extraction is performed where either default or custom embedding network is employed to extract meaningful features from the data traces, thus lowering the overall dimensionality of the output data.

  The Inferencer class supports all currently available estimation methods in sbi:

  • sequential neural posterior estimation (SNPE) [2];
  • sequential neural likelihood estimation (SNLE) [3];
  • sequential neural ratio estimation (SNRE) [4].

  So, the only thing that the user has to do is to instantiate the Inferencer class with required arguments such as the input and output data, and a few optional ones:

  inferencer = Inferencer(dt=dt, model=eqs,
                          input={'I': inp_traces*amp},
                          output={'v': out_traces*mV},
                          features={'v': v_features},
                          method='exponential_euler',
                          threshold='m > 0.5',
                          refractory='m > 0.5',
                          param_init=init_conds)

  The inference process starts when infer method is called:

  posterior = inferencer.infer(n_samples=15_000,
                               n_rounds=1,
                               inference_method='SNPE',
                               density_estimator_model='maf',
                               g_Na=[1*uS, 100*uS],
                               g_K=[0.1*uS, 10*uS])

  Overall idea is to distance the user from the sbi package (and the use of the PyTorch library), and achieve an API very similar to the one to which the brian2modelfitting users are already accustomed through the Fitter class.

• Loading/storing of the training data and the estimated posterior. - PR #52

  inferencer.save_summary_statistics('path_to_file.npz')

  and

  theta, x = inferencer.load_summary_statistics('path_to_file.npz')

• Support for visual analysis of the estimated posterior over free parameters. - PR #58

  # pairwise relationship between parameters
  limits = {'g_Na': [1*uS, 100*uS],
            'g_K': [0.1*uS, 10*uS]}
  labels = {'g_Na': r'$\overline{g}_{Na}$',
            'g_K': r'$\overline{g}_{K}$'}
  fig, ax = inferencer.pairplot(limits=limits,
                                labels=labels,
                                ticks=limits,
                                points=ground_truth_params,
                                points_offdiag={'markersize': 5},
                                points_colors=['C3'],
                                figsize=(6, 6));
  
  # traces/spikes visualization
  inf_traces = inferencer.generate_traces()
  fig, ax = plot_traces(t, inp_traces, out_traces, inf_traces=array(inf_traces/mV))

• Multi-objective optimization support. - PR #59

  inferencer = Inferencer(dt=dt, model=eqs,
                          input={'I': inp_traces*amp},
                          output={'v': out_traces*mV, 'm': Na_activation_dynamics},  # 2 state variables are considered...
                          features={'v': v_features, 'm': m_features},  # ... which means that 2 list of summary features should also be considered
                          method='exponential_euler',
                          threshold='m > 0.5',
                          refractory='m > 0.5',
                          param_init=init_conds)

• Inference by observing spikes either instead or together with the dynamics of different state variables in the model. - PR #60

  inferencer = Inferencer(dt=dt, model=eqs,
                          input={'I': inp_traces*amp},
                          output={'v': out_traces*mV, 'spikes': spike_times_list},
                          features={'v': v_features, 'spikes': s_features},
                          method='exponential_euler',
                          threshold='m > 0.5',
                          refractory='m > 0.5',
                          param_init=init_conds)

• GPU support for the inference process. - PR #61

  posterior = inferencer.infer(..., sbi_device='gpu', ...)

  This feature should be used carefully as discussed in the PR and relevant issues mentioned in the PR.

• Updated documentation with examples and the comprehensive tutorial, currently available here.

• Added examples:

  • IF_sbi.py - inference of two unknown parameters of an integrate-and-fire neuron model by observing only spike trains
  • hh_sbi_simple_interface.py - inference on real data traces by using a Hodgkin-Huxley model with three unknown parameters with simple interface by using infer method in the Inferencer class
  • hh_sbi_flexible_interface.py - multi-round inference on real data traces by using the Hodgkin-Huxley model with three unknown parameters with flexible interface by using infer_step method in the Inferencer class manually
  • hh_sbi_synthetic_traces.py - automatic feature extraction and inference on synthetic data traces by using the Hodgkin-Huxley model

• Updated tests.

What's left to do

Even though all the features defined at the very beginning of the project, both through the project proposal, available here, and internally between Marcel and me, are done, the test suite for the Inferencer class is far from comprehensive and updates and refinement are needed.

Additional work

Additional development, testing and various examples that are not officially included in the PR are located in the gsoc_dump_repo on my personal GitHub account, available here.

Progress tracker

During the project (period of 10 weeks), the entire work is accompanied by weekly reports that outline in detail what has been done during the current week, what is planned for the next week and whether there are any obstacles. The reports are available here.

For a more concise review, it is sufficient to refer to Projects -> sbi integration (GSoC 2021) in the brian2modelfitting GitHub repository, available here. Under Done section, there are listed all issues and pull requests that resolve those issues, described earlier in the section Main results. This PR will be listed under the In progress section 🙂

Thanks

Many thanks to @mstimberg for his immense patience and time for my (often silly) questions 🙇‍♂️

References

[1] Alvaro Tejero-Cantero, Jan Boelts, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pedro J Goncalves, David S Greenberg, Jakob H Macke. sbi: A toolkit for simulation-based inference, 2020, Journal of Open Source Software 5:2505, doi:10.21105/joss.02505.

[2] David S. Greenberg, Marcel Nonnenmacher, Jakob H. Macke. Automatic posterior transformation for likelihood-free inference, 2019, 36th International Conference on Machine Learning, PMLR 97:2404-2414.

[3] George Papamakarios, David C. Sterratt, Iain Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows, 2019, 22nd International Conference on Artificial Intelligence and Statistics, PMLR 89:837-848.

[4] Conor Durkan, Iain Murray, George Papamakarios. On contrastive learning for likelihood-free inference, 2020, 37th International Conference on Machine Learning, PMLR 119:2771-2781.

[mstimberg]

Thanks for summarizing everything so nicely, this looks great! It'll take me some time to review everything in detail before doing a final merge and the release, but I am already very happy with the results of this GSoC 😊

[mstimberg]

Apologies that this review is taking ages... There are still a few things that I'd like to work on (mostly the tests and trying it out on real-life examples), but I can do this in the master branch as well, so I'll take the important symbolic step of merging this pull request now 😄
Thanks again for all your work on this @antelk