Statistics of cortical representational drift can enable robust readout.

1. # %% [markdown]
2. # # Comparing sudden and gradual models of drift in a simple toy model
3. #
4. # This notebook steps through generating the figures from _Statistics of cortical representational drift can enable robust readout_ (Micou & O'Leary).
5. # This notebook is primarily concerned with reproducibility--the paper provides a more comprehensive explanation of the model.
6. #
7. # The model itself is implemented within the files in this repository, which we import below to plot the key figures.
9. # %%
10. # Import the high-level plotting for this notebook (which internally call out to the model)
11. from drift_model.model_overview import make_ratemaps, make_correlation_vs_parameter_difference, make_example_pdfs, make_correlation_vs_parameter_scatter, plot_noise_free_correlation_difference, plot_residual_distribution
12. from drift_model.drift_equivalence import make_drift_equivalence, get_params_for_drift_rate, plot_drift_equivalence, plot_multiple_drift_equivalences, plot_approach_to_chance
13. from drift_model.adaptive_decoding import simulate_over_time, plot_decoder_error_against_time, simulate_over_population_sizes_and_drift_rates, plot_decoder_error_vs_parameters, simulate_over_noise_levels, plot_error_against_noise_level
14. from drift_model.scaled_correlation_changes import plot_drift_in_scaled_comparison, plot_gmm_fits, fit_gmms_over_drift_rates, plot_mu2_values
16. # Import additional utilities for the notebook
17. import numpy as np
18. import os
19. %matplotlib inline
21. # Make a directory to save .svg versions of the figures
22. OUTPUT_FOLDER = 'out'
23. if not os.path.exists(OUTPUT_FOLDER):
24. os.mkdir(OUTPUT_FOLDER)
26. # %% [markdown]
27. # ## Encoding of a stimulus in the population
28. #
29. # Some neuron $i$ in the population is exposed to a stimulus $s$ and produces a response $x_i$ that depends on its tuning parameter $θ_i$ and some random perturbation parameterised by the noise level $σ$. The stimulus has both a magnitude and an orientation.
30. # The parameter $θ_i$ describes the orientation of maximum response for each neuron.
31. # Taking $M$ samples of the stimulus, we can visualise the how each neuron responds to the stimulus by plotting a ratemap (its mean response at each combination of stimulus magnitude and orientation).
32. # This plot is used in Figure 2.b in the paper.
34. # %%
35. # Reproduce the heatmaps for Fig. 2b
36. example_neuron_tuning_parameters = np.array([0, np.pi/4, -np.pi*2/3])
37. noise_sigma = 1.0
38. num_samples_M = 10000
40. make_ratemaps(example_neuron_tuning_parameters, noise_sigma, num_samples_M, save_fig=f'{OUTPUT_FOLDER}/fig_2b')
42. # %% [markdown]
43. # ## Pairwise correlation of neurons
44. # If we take a pair of neurons and measure the correlation of their activities over the course of $M$ samples, we find that neurons with more similar values of $θ$ are more highly correlated with each other. However, due to the observation noise and the finite number of samples, the correlation between a pair of neurons at some angular separation $\Deltaθ$ between their preferred firing orientations is a probability distribution. We visualise this below.
45. # This plot is used in Figure 2.c in the paper.
47. # %%
48. # Reproduce the pairwise correlations between neurons at some separation Δθ for Fig. 2c
49. noise_sigma = 1.0
50. num_samples_M = 10000
52. num_points = 101
53. samples_per_point = 20000
55. make_correlation_vs_parameter_difference(num_points, samples_per_point, noise_sigma, num_samples_M, save_fig=f'{OUTPUT_FOLDER}/fig_2c')
57. # %% [markdown]
58. # ## Sudden and gradual drift
59. #
60. # We consider two models of drift. Gradual drift (parameterised by $\rho$) makes alterations to all neuron tuning parameters in the population drawn from a Gaussian distribution, while sudden drift (parameterised by $\alpha$) selects a subset of the population and retunes them drawing from a uniform distribution.
61. # We show probability distributions that describe the distribution of relative change of tuning parameter from one day to the next for both of these types of drift. This plot is used in Figure 3.b in the paper.
62. #
63. # Note that using a big value of $\rho$ will cause the tuning parameter to 'wrap around' the unit circle. You can play with visualising that here by increasing the value of $\rho$. For the drift rate of 0.1, which is what we show in most of Figure 3, the value of $\rho$ is just below 0.75, and this wrapping around is entirely negligible.
64. # The largest value of $\rho$ we use anywhere in the paper is 1.1 (for the very fast drift rate of 0.2 in Figure 3.g), and at this point the wrapping around is just about noticeable, so pushing the adaptive decoder to even faster rates of drift risks unfairly disadvantaging the gradual drift model.
66. # %%
67. # Reproduce the PDFs of each drift type from Fig. 3b
68. alpha = 0.25
69. rho = 0.75
70. make_example_pdfs(alpha, rho, save_fig=f'{OUTPUT_FOLDER}/fig_3b')
72. # %% [markdown]
73. # ## Drift rate equivalence
74. #
75. # We define the drift rate of a population to be the number of days of drift it takes to scramble a population to within an $\epsilon$ of a completely randomised population. We use this to construct an equivalence between the gradual and sudden drift types. This plot is used in Figure 3.c of the paper.
77. # %%
78. # Reproduce the equivalence between drift types from Fig. 3c
80. # Generate the equivalence
81. num_equivalence_samples = 500
82. num_simulations_per_sample = 1000000
83. equivalence_mapping = make_drift_equivalence(num_samples=num_equivalence_samples,
84. epsilon=0.05,
85. num_simulations=num_simulations_per_sample)
87. # Show an example of extracting a single alpha and rho for a given drift rate
88. target_drift_rate = 0.1
89. actual_drift_rate, alpha, rho = get_params_for_drift_rate(target_drift_rate, equivalence_mapping)
90. print(f"Drift rate={actual_drift_rate}, α={alpha}, ρ={rho}")
92. # Plot the entire equivalence
93. plot_drift_equivalence(equivalence_mapping, save_fig=f'{OUTPUT_FOLDER}/fig_3c')
95. # %% [markdown]
96. # ## Adaptive decoding works best for sudden drift
97. #
98. # We use an adaptive decoder to attempt to track the changing tuning parameters of neurons in the population as they undergo drift.
99. # We find that this adaptive decoding strategy works better with the sudden drift model than for the gradual drift model at equivalent rates of drift. This code reproduces Figure 3.f from the paper.
100. #
101. # Note: the adaptive decoder relies on minimising a cost function. This estimates the parameter on some day $T$ based on some prior belief of the estimate on day $T-1$. If the adaptive decoder diverges sufficiently from the ground truth, as in the case of the gradual drift simulations below, then the estimates on $T-1$ can become poor and occasionally result in the solver not neatly converging. This isn't a problem if it takes place after the estimate is already poor, but would be a methodological issue if it took place early on (after just a day or two of drift). The code below will surface any solver warnings that occur before the 5th day of drift (and silence them afterwards to keep this notebook clean, but in general this warning threshold can be increased to the whole 50 days and still generate very few warning messages over all 100 simulations).
103. # %%
104. # Reproduce the adaptive decoding over time in Fig. 3f
105. num_simulations = 100
106. days_to_simulate = 51
107. num_neurons = 100
108. noise_sigma = 1.0
109. num_samples_M = 10000
110. drift_rate = 0.1
111. _, alpha, rho = get_params_for_drift_rate(drift_rate, equivalence_mapping)
113. # Convergence warning threshold: this will report any solver convergence failures that occur before the configured day, and ignore any after
114. # (see note above: as the estimates from adapting to gradual drift diverge from the truth solving eventually becomes ill-condiditioned)
115. solver_warning_cutoff_day = 5
117. sim_results = simulate_over_time(alpha,
118. rho,
119. num_neurons,
120. days_to_simulate,
121. num_samples_M,
122. noise_sigma,
123. num_simulations,
124. solver_warning_cutoff_day)
125. plot_decoder_error_against_time(sim_results, save_fig=f'{OUTPUT_FOLDER}/fig_3f')
127. # %% [markdown]
128. # ## Sudden drift retains its advantage across drift rates and population sizes
129. # We can explore the parameter space of population sizes and of drift rates to confirm that sudden drift maintains an advantage over gradual drift and that our observations are not sensitive to these choices of parameter.
130. # This plot is used in Figure 3.g of the paper.
132. # %%
133. # Reproduce the adaptive decoding over a range of drift rates and population sizes from Fig. 3g
135. simulations_per_datapoint = 100
136. target_drift_rates = np.array([0.05, 0.1, 0.2])
137. population_sizes = np.array([25, 50, 75, 100, 125, 150, 175, 200])
138. noise_sigma = 1.0
139. time_window_M = 10000
140. days_to_simulate = 26 # drift is applied between days (i.e. day 0 has no drift), so this is 25 days of drift
142. drift_rates, alphas, rhos = get_params_for_drift_rate(target_drift_rates, equivalence_mapping)
144. results = simulate_over_population_sizes_and_drift_rates(alphas,
145. rhos,
146. population_sizes,
147. days_to_simulate,
148. time_window_M,
149. noise_sigma,
150. simulations_per_datapoint)
152. plot_decoder_error_vs_parameters(results, drift_rates, population_sizes, save_fig=f'{OUTPUT_FOLDER}/fig_3g')
154. # %% [markdown]
155. # ## Adaptive decoding works best for sudden drift over a range of noise levels
156. #
157. # The adaptive decoder is more susceptible to confusing daily random variations in activity due to noise with an underlying change in tuning.
158. # We consider the accuracy of the parameter estimates after the very first day of drift to demonstrate the relative sensitivity of the gradual drift scenario to the amount of observational noise.
159. # This plot is used in Figure 3.e of the paper.
161. # %%
162. # Reproduce the decoder error after the initial day of drift as a function of noise level sigma, as in Fig. 3e
164. noise_levels = np.array([0.2, 0.4, 0.8, 1.6, 2.4, 4.8, 9.6, 19.2, 38.4, 76.8])
165. time_window_M = 1000000
167. drift_rate = 0.1
168. _, alpha, rho = get_params_for_drift_rate(drift_rate, equivalence_mapping)
169. simulations_per_datapoint = 100
170. population_size_N = 100
172. noise_results = simulate_over_noise_levels(alpha, rho, population_size_N, time_window_M, noise_levels, simulations_per_datapoint)
174. plot_error_against_noise_level(noise_results, time_window_M, noise_levels, save_fig=f'{OUTPUT_FOLDER}/fig_3e')
176. # %% [markdown]
177. # ## Visualisation of changes in the pairwise correlation between neurons
178. #
179. # One way we can visualise the differences between sudden and gradual drift is to consider how the correlation between pairs of neurons in the populations change. We consider this for our simulated model before moving to considering it for in vivo data in the paper.
180. # This plot is used in Figure 3.d of the paper.
182. # %%
183. # Reproduce the scaled correlation change ΔC from Fig. 3d
185. days_to_plot = [1, 5, 100]
186. alpha, rho = 0.0577, 0.3336 # drift rate of 0.02
188. # Use a lower correlation threshold but a relatively un-noisy sim to reduce computation time
189. corr_threshold = 0.05
190. time_window_M = 1000
191. noise_sigma = 0.5
193. plot_drift_in_scaled_comparison(alpha, rho, days_to_plot, corr_threshold, time_window_M, noise_sigma, save_fig=f'{OUTPUT_FOLDER}/fig_3d')
195. # %% [markdown]
196. # ## Mixture of Gaussian fits to distributions of correlation change
197. # During the early stages of drift, the distributions of scaled correlation changes can be well-approximated by a mixture of Gaussians.
198. # The Gaussian with the larger mean in the mixture, $\mu_2$, has different behaviour with respect to time for each type of drift.
199. # This plot is used in Figures 4.a and 4.b in the paper.
201. # %%
202. # Reproduce the example 2-component GMM fits from Fig. 4a and 4b
204. drift_rate = 0.025
205. _, alpha, rho = get_params_for_drift_rate(drift_rate, equivalence_mapping)
207. num_days_to_plot = 6
208. corr_threshold = 0.05
210. plot_gmm_fits(alpha, rho, num_days_to_plot, corr_threshold, save_fig=f'{OUTPUT_FOLDER}/fig_4ab')
212. # %% [markdown]
213. # ## Approximate behaviour of $\mu_2$ as a function of drift
214. #
215. # We can approximate the behaviour of $\mu_2$ as an exponential approach to a final value for gradual drift, and as an entirely flat line independent of elapsed days for sudden drift. This holds across a variety of drift rates.
216. #
217. # This plot is used in Figures 4.e and 4.f in the paper.
219. # %%
220. # Reproduce the approximations of mu2's location as a function of time and drift rate from Fig. 4e and 4f
222. Tcs = np.array([5, 10, 20, 40, 50, 60, 80])
223. drift_rates = 1 / Tcs
224. corr_threshold = 0.05
225. days_to_simulate = 8
226. noise_sigma = 0.5
227. time_window_M = 1000
229. drift_rates, mu2_values_sudden, mu2_values_gradual = fit_gmms_over_drift_rates(drift_rates, equivalence_mapping, noise_sigma, time_window_M, corr_threshold, days_to_simulate)
230. plot_mu2_values(drift_rates, mu2_values_sudden, mu2_values_gradual, save_fig=f'{OUTPUT_FOLDER}/fig_4ef')
232. # %% [markdown]
233. # ## Supplementary figures
234. # ### Accuracy of the approximation of the expected correlation between neurons (and choice of M and σ)
235. #
236. # We rely on an analytical form for the expected correlation between pairs of neurons. It allows us to compute the Jacobian quickly, and therefore makes the optimisation problems above tractable. This analytical form is very, very close to the values we get from simulation (if they were not, the solver would fail due to a disagreement between gradient and cost function), and it gets better for large $M$.
237. # The plots below show the distribution of this correlation for a larger variety of $M$ and $σ$, as well as how close to normal the distribution around the expected value is.
238. #
239. # These plots are used in Figure S1 of the paper.
241. # %%
242. # Reproduce Fig. S1a: show distributions at a variety of M and σ
243. num_points = 201
244. samples_per_point = 100
245. sigmas = [0.5, 1.0, 2.0]
246. num_samples_M = [1000, 10000, 100000]
248. make_correlation_vs_parameter_scatter(num_points, samples_per_point, sigmas, num_samples_M, save_fig=f'{OUTPUT_FOLDER}/fig_s1a')
250. # %%
251. # Reproduce Fig. S1b: show how the expected correlation gets more accurate as a function of M
253. num_samples_M = [100, 1000, 10000, 100000]
254. num_points = 201
255. samples_per_point = 10000
256. plot_noise_free_correlation_difference(num_points, samples_per_point, num_samples_M, save_fig=f'{OUTPUT_FOLDER}/fig_s2b')
258. # %%
259. # Reproduce Fig. S1c: a visualisation that the distribution is normally distributed about the expected value
261. noise_sigma = 1.0
262. num_points = 201
263. samples_per_point = 10000
264. time_window_M = 15 # this works even for low M (~10-100), so can reduce computation time
266. plot_residual_distribution(time_window_M, noise_sigma, num_points, samples_per_point, save_fig=f'{OUTPUT_FOLDER}/fig_s2c')
268. # %% [markdown]
269. # ### Insensitivity of drift equivalence to choice of 𝜖
270. # We use a drift rate equivalence based on setting $\epsilon=0.05$. Below, we regenerate the equivalence for $\epsilon=0.01$ and $\epsilon=0.1$ to show that this equivalence is not particularly sensitive to this choice of value.
271. # We also show the approach of a population's state to chance as a function of days of drift for a variety of drift rates.
272. #
273. # These plots are used in Figure S2.a-b.
275. # %%
276. # Reproduce Fig. S2b, showing that various choices of epsilon produce comparable equivalence mappings
278. num_equivalence_samples = 100
279. num_simulations_per_sample = 1000000
281. epsilons = [0.1, 0.05, 0.001]
282. equivalences = [make_drift_equivalence(lower_rho=0.2, upper_rho=2.25, num_samples=num_equivalence_samples, epsilon=e, num_simulations=num_simulations_per_sample) for e in epsilons]
283. plot_multiple_drift_equivalences(equivalences, epsilons, save_fig=f'{OUTPUT_FOLDER}/fig_s2b')
285. # %%
286. # Reproduce Fig. S2a, showing how a population gradually approaches a completely scrambled state
288. num_days = 41
289. Tcs = np.array([5, 10, 20]) # corresponding to drift rates of 0.2, 0.1, 0.05 respectively
290. num_sims = 10000000
292. plot_approach_to_chance(equivalence_mapping, Tcs, num_days, num_sims, save_fig=f'{OUTPUT_FOLDER}/fig_s2a')
294. # %% [markdown]
295. # ### Alternative definition of drift based on disruption to representation on first day of drift
296. #
297. # The definition of drift rate equivalence used above relies on how many days it takes to scramble a population to within an $\epsilon$ of chance. We show below that are results also hold if we used an alternative definition: a choice of parameters such that, after one day, the mean error of parameter estimate for a fixed decoder should be equal between drift types.
298. #
299. # These plots are used in Figure S2.c-d.
301. # %%
302. # Reproduce the alternative definition of drift equivalence from Fig. S2c
304. alpha, rho = 0.38, 0.74242424 # Manually enforce equal drift after first day
306. noise_levels = np.array([0.2, 0.4, 0.8, 1.6, 2.4, 4.8, 9.6, 19.2, 38.4, 76.8])
307. time_window_M = 1000000
308. simulations_per_datapoint = 100
309. population_size_N = 100
311. noise_results = simulate_over_noise_levels(alpha, rho, population_size_N, time_window_M, noise_levels, simulations_per_datapoint)
313. plot_error_against_noise_level(noise_results, time_window_M, noise_levels, save_fig=f'{OUTPUT_FOLDER}/fig_s2c')
315. # %%
316. # Reproduce the adaptive decoding over time for the alternative drift equivalence in Fig. S2d
318. num_simulations = 100
319. days_to_simulate = 51
320. num_neurons = 100
321. noise_sigma = 1.0
322. num_samples_M = 10000
323. alpha, rho = 0.38, 0.74242424 # Manually enforce equal drift after first day
324. solver_warning_cutoff_day = 5
326. sim_results = simulate_over_time(alpha,
327. rho,
328. num_neurons,
329. days_to_simulate,
330. num_samples_M,
331. noise_sigma,
332. num_simulations,
333. solver_warning_cutoff_day)
334. plot_decoder_error_against_time(sim_results, save_fig=f'{OUTPUT_FOLDER}/fig_s2d')

drift_demo.ipynb at commit 92c5386, no license · at the source