Added trajectory plot (Figure 5).

trajectory_plot.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+from knn_regressors import active_knn_regressor, oracle_knn_regressor
+from SIR_simulation import Sampler
+
+plt.rc('font', size=22)  # Increase matplotlib default font size
+np.random.seed(0)  # Fix random seed for reproducibility
+
+n = 2**8
+
+sim_len = 20
+sampler = Sampler(sim_len=sim_len)
+
+Xs_ground_truth, Ys_ground_truth, Zs_ground_truth = sampler.generate_joint_samples(1)
+Ys_ground_truth = Ys_ground_truth.squeeze()
+
+# Sample passive dataset
+Xs, Ys, Zs = sampler.generate_joint_samples(n)
+_, Xs_active, Ys_active, Zs_active = active_knn_regressor(sampler, Xs_ground_truth, n, method='cv')
+
+plt.figure(figsize=(12, 12))
+plt.subplot(3, 2, 1)
+plt.title('Passive')
+plt.plot(np.tile(range(1, sim_len+1), (n, 1)).transpose(), Xs[:, 1*sim_len:2*sim_len].transpose(), c='b', alpha=0.2, lw=2, label='Passive')
+plt.plot(range(1, sim_len+1), Xs_ground_truth[:, 1*sim_len:2*sim_len].squeeze(), c='k', lw=4, label='Ground Truth')
+plt.ylabel('# Infected ($I$)')
+plt.subplot(3, 2, 2)
+plt.title('Active')
+plt.plot(np.tile(range(1, sim_len+1), (n, 1)).transpose(), Xs_active[:, 1*sim_len:2*sim_len].transpose(), c='g', alpha=0.2, lw=2, label='Active')
+plt.plot(range(1, sim_len+1), Xs_ground_truth[:, 1*sim_len:2*sim_len].squeeze(), c='k', lw=4, label='Ground Truth')
+
+plt.subplot(3, 2, 3)
+plt.plot(np.tile(range(1, sim_len+1), (n, 1)).transpose(), Xs[:, 3*sim_len:4*sim_len].transpose(), c='b', alpha=0.2, lw=2, label='Passive')
+plt.plot(range(1, sim_len+1), Xs_ground_truth[:, 3*sim_len:4*sim_len].squeeze(), c='k', lw=4, label='Ground Truth')
+plt.ylabel('# Deceased ($D$)')
+plt.xlabel('Time ($t$)')
+plt.subplot(3, 2, 4)
+plt.plot(np.tile(range(1, sim_len+1), (n, 1)).transpose(), Xs_active[:, 3*sim_len:4*sim_len].transpose(), c='g', alpha=0.2, lw=2, label='Active')
+plt.plot(range(1, sim_len+1), Xs_ground_truth[:, 3*sim_len:4*sim_len].squeeze(), c='k', lw=4, label='Ground Truth')
+plt.xlabel('Time ($t$)')
+
+labels = ['$\\beta$', '$\gamma$', '$\delta$']
+def set_axis_style(ax):
+  ax.xaxis.set_tick_params(direction='out')
+  ax.xaxis.set_ticks_position('bottom')
+  ax.set_xticks(np.arange(1, len(labels) + 1), labels=labels)
+  ax.set_xlim(0.25, len(labels) + 0.75)
+
+plt.subplot(3, 2, 5)
+plt.violinplot(dataset=Zs[:, :3])#, c='b', alpha=0.2, lw=2)
+plt.plot([1, 2, 3], Zs_ground_truth.transpose()[:3], c='k', lw=3)
+plt.ylim((0, 1))
+plt.ylabel('Value')
+plt.xlabel('Model Parameter')
+set_axis_style(plt.gca())
+
+plt.subplot(3, 2, 6)
+plt.violinplot(dataset=Zs_active[:, :3])#, c='g', alpha=0.2, lw=2)
+plt.plot([1, 2, 3], Zs_ground_truth.transpose()[:3], c='k', lw=3)
+plt.ylim((0, 1))
+plt.xlabel('Model Parameter')
+set_axis_style(plt.gca())
+
+plt.show()