Add epidemic threshold (#35)

* add epidemic threshold lines

* fix label

collect_clustered_network.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from joblib import Parallel, delayed
+from numpy.linalg import eigh
 
 from lcs import *
 
@@ -58,6 +59,9 @@ def get_metrics(f, dir, c_dict, b_dict, s_dict, r_dict):
     m = dict()
     m["rho"] = density(A)
     m["rho-samples"] = density(samples.mean(axis=0))
+    m["mean-degree"] = degrees(A).mean()
+    m["mean-squared-degree"] = (degrees(A) ** 2).mean()
+    m["pf-eigenvalue"] = eigh(A)[0][-1]
     m["ps"] = posterior_similarity(samples, A)
     m["sps"] = samplewise_posterior_similarity(samples, A)
     m["fs"] = f_score(samples, A)
@@ -100,6 +104,9 @@ def get_metrics(f, dir, c_dict, b_dict, s_dict, r_dict):
 data["size"] = list(s_dict)
 data["rho"] = np.zeros((n_c, n_b, n_s, n_r))
 data["rho-samples"] = np.zeros((n_c, n_b, n_s, n_r))
+data["mean-degree"] = np.zeros((n_c, n_b, n_s, n_r))
+data["mean-squared-degree"] = np.zeros((n_c, n_b, n_s, n_r))
+data["pf-eigenvalue"] = np.zeros((n_c, n_b, n_s, n_r))
 data["ps"] = np.zeros((n_c, n_b, n_s, n_r))
 data["sps"] = np.zeros((n_c, n_b, n_s, n_r))
 data["fs"] = np.zeros((n_c, n_b, n_s, n_r))

collect_cm.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from joblib import Parallel, delayed
+from numpy.linalg import eigh
 
 from lcs import *
 
@@ -58,6 +59,9 @@ def get_metrics(f, dir, c_dict, b_dict, a_dict, r_dict):
     m = dict()
     m["rho"] = density(A)
     m["rho-samples"] = density(samples.mean(axis=0))
+    m["mean-degree"] = degrees(A).mean()
+    m["mean-squared-degree"] = (degrees(A) ** 2).mean()
+    m["pf-eigenvalue"] = eigh(A)[0][-1]
     m["ps"] = posterior_similarity(samples, A)
     m["sps"] = samplewise_posterior_similarity(samples, A)
     m["fs"] = f_score(samples, A)
@@ -95,6 +99,9 @@ def get_metrics(f, dir, c_dict, b_dict, a_dict, r_dict):
 data["alpha"] = list(a_dict)
 data["rho"] = np.zeros((n_c, n_b, n_a, n_r))
 data["rho-samples"] = np.zeros((n_c, n_b, n_a, n_r))
+data["mean-degree"] = np.zeros((n_c, n_b, n_a, n_r))
+data["mean-squared-degree"] = np.zeros((n_c, n_b, n_a, n_r))
+data["pf-eigenvalue"] = np.zeros((n_c, n_b, n_a, n_r))
 data["ps"] = np.zeros((n_c, n_b, n_a, n_r))
 data["sps"] = np.zeros((n_c, n_b, n_a, n_r))
 data["fs"] = np.zeros((n_c, n_b, n_a, n_r))

collect_erdos-renyi.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from joblib import Parallel, delayed
+from numpy.linalg import eigh
 
 from lcs import *
 
@@ -58,6 +59,9 @@ def get_metrics(f, dir, c_dict, b_dict, p_dict, r_dict):
     m = dict()
     m["rho"] = density(A)
     m["rho-samples"] = density(samples.mean(axis=0))
+    m["mean-degree"] = degrees(A).mean()
+    m["mean-squared-degree"] = (degrees(A) ** 2).mean()
+    m["pf-eigenvalue"] = eigh(A)[0][-1]
     m["ps"] = posterior_similarity(samples, A)
     m["sps"] = samplewise_posterior_similarity(samples, A)
     m["fs"] = f_score(samples, A)
@@ -95,6 +99,9 @@ def get_metrics(f, dir, c_dict, b_dict, p_dict, r_dict):
 data["p"] = list(p_dict)
 data["rho"] = np.zeros((n_c, n_b, n_p, n_r))
 data["rho-samples"] = np.zeros((n_c, n_b, n_p, n_r))
+data["mean-degree"] = np.zeros((n_c, n_b, n_p, n_r))
+data["mean-squared-degree"] = np.zeros((n_c, n_b, n_p, n_r))
+data["pf-eigenvalue"] = np.zeros((n_c, n_b, n_p, n_r))
 data["ps"] = np.zeros((n_c, n_b, n_p, n_r))
 data["sps"] = np.zeros((n_c, n_b, n_p, n_r))
 data["fs"] = np.zeros((n_c, n_b, n_p, n_r))

collect_sbm.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from joblib import Parallel, delayed
+from numpy.linalg import eigh
 
 from lcs import *
 
@@ -58,6 +59,9 @@ def get_metrics(f, dir, c_dict, b_dict, e_dict, r_dict):
     m = dict()
     m["rho"] = density(A)
     m["rho-samples"] = density(samples.mean(axis=0))
+    m["mean-degree"] = degrees(A).mean()
+    m["mean-squared-degree"] = (degrees(A) ** 2).mean()
+    m["pf-eigenvalue"] = eigh(A)[0][-1]
     m["ps"] = posterior_similarity(samples, A)
     m["sps"] = samplewise_posterior_similarity(samples, A)
     m["fs"] = f_score(samples, A)
@@ -95,6 +99,9 @@ def get_metrics(f, dir, c_dict, b_dict, e_dict, r_dict):
 data["epsilon"] = list(e_dict)
 data["rho"] = np.zeros((n_c, n_b, n_e, n_r))
 data["rho-samples"] = np.zeros((n_c, n_b, n_e, n_r))
+data["mean-degree"] = np.zeros((n_c, n_b, n_e, n_r))
+data["mean-squared-degree"] = np.zeros((n_c, n_b, n_e, n_r))
+data["pf-eigenvalue"] = np.zeros((n_c, n_b, n_e, n_r))
 data["ps"] = np.zeros((n_c, n_b, n_e, n_r))
 data["sps"] = np.zeros((n_c, n_b, n_e, n_r))
 data["fs"] = np.zeros((n_c, n_b, n_e, n_r))

collect_watts-strogatz.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from joblib import Parallel, delayed
+from numpy.linalg import eigh
 
 from lcs import *
 
@@ -58,6 +59,9 @@ def get_metrics(f, dir, c_dict, b_dict, p_dict, r_dict):
     m = dict()
     m["rho"] = density(A)
     m["rho-samples"] = density(samples.mean(axis=0))
+    m["mean-degree"] = degrees(A).mean()
+    m["mean-squared-degree"] = (degrees(A) ** 2).mean()
+    m["pf-eigenvalue"] = eigh(A)[0][-1]
     m["ps"] = posterior_similarity(samples, A)
     m["sps"] = samplewise_posterior_similarity(samples, A)
     m["fs"] = f_score(samples, A)
@@ -95,6 +99,9 @@ def get_metrics(f, dir, c_dict, b_dict, p_dict, r_dict):
 data["p"] = list(p_dict)
 data["rho"] = np.zeros((n_c, n_b, n_p, n_r))
 data["rho-samples"] = np.zeros((n_c, n_b, n_p, n_r))
+data["mean-degree"] = np.zeros((n_c, n_b, n_p, n_r))
+data["mean-squared-degree"] = np.zeros((n_c, n_b, n_p, n_r))
+data["pf-eigenvalue"] = np.zeros((n_c, n_b, n_p, n_r))
 data["ps"] = np.zeros((n_c, n_b, n_p, n_r))
 data["sps"] = np.zeros((n_c, n_b, n_p, n_r))
 data["fs"] = np.zeros((n_c, n_b, n_p, n_r))

plot_fig1.py

@@ -212,7 +212,7 @@
     va="top",
 )
 
-measure = "auprc"
+measure = "auroc"
 
 with open("Data/zkc_infer_vs_tmax.json") as file:
     data = json.load(file)
@@ -241,7 +241,7 @@
     hdi_b[i] = b
 
 ax3.fill_between(tmax, hdi_a, hdi_b, alpha=0.3, color="C1", edgecolor="none")
-ax3.set_ylabel("AUPRC")
+ax3.set_ylabel(measure.upper())
 ax3.set_xlabel(r"$t_{\mathregular{max}}$")
 ax3.set_xlim([10, 10**4])
 ax3.set_xticks(
@@ -258,7 +258,7 @@
 
 ax3.legend(
     loc="lower right",
-    bbox_to_anchor=(0.87, -0.05, 0.2, 0.2),
+    bbox_to_anchor=(0.87, -0.07, 0.2, 0.2),
     markerfirst=False,
     frameon=False,
     handlelength=0.8,

plot_fig2.py

@@ -1,5 +1,6 @@
 import json
 
+import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 import xgi
@@ -8,10 +9,10 @@
 import fig_settings as fs
 from lcs import *
 
+gamma = 0.1
+measure = "auroc"
 axis_limits = [0, 1]
 
-measure = "auprc"
-
 fs.set_fonts()
 fs.set_colors()
 cmap = fs.cmap
@@ -86,17 +87,18 @@ def visualize_networks(i, ax):
         data = json.load(file)
     var = np.array(data[keys[i]], dtype=float)
     b = np.array(data["beta"], dtype=float)
+    l = np.array(data["pf-eigenvalue"], dtype=float)[0].mean(axis=0).mean(axis=1)
 
     performance = np.array(data[measure], dtype=float)
 
-    # plot the difference in auprc
+    # plot the difference in auroc
     mean_difference = performance[1].mean(axis=2).T - performance[0].mean(axis=2).T
     ax = fig.add_subplot(gs[1, i])
     im1 = ax.imshow(
         to_imshow_orientation(mean_difference),
         extent=(min(var), max(var), min(b), max(b)),
-        vmin=-0.5,
-        vmax=0.5,
+        vmin=-0.3,
+        vmax=0.3,
         aspect="auto",
         cmap=cmap,
     )
@@ -105,6 +107,10 @@ def visualize_networks(i, ax):
     ax.set_xticks(xticks[i], xticklabels[i])
     ax.set_yticks([0, 0.5, 1], [0, 0.5, 1])
 
+    beta_c = gamma / l
+    for r0 in range(1, 10, 2):
+        ax.plot(var, r0 * beta_c, "r-", linewidth=0.5)
+
     if i == 0:
         ax.set_ylabel(r"$\beta$")
     else:
@@ -122,8 +128,8 @@ def visualize_networks(i, ax):
     im2 = ax.imshow(
         to_imshow_orientation(mean_difference),
         extent=(min(var), max(var), min(b), max(b)),
-        vmin=-0.5,
-        vmax=0.5,
+        vmin=-0.3,
+        vmax=0.3,
         aspect="auto",
         cmap=cmap,
     )
@@ -142,12 +148,12 @@ def visualize_networks(i, ax):
 cbar_ax1 = fig.add_axes([0.85, 0.4, 0.015, 0.25])
 cbar = fig.colorbar(im1, cax=cbar_ax1)
 cbar.set_label(rf"$\Delta$ {measure.upper()}", rotation=270, labelpad=10)
-cbar_ax1.set_yticks([-0.5, 0, 0.5])
+cbar_ax1.set_yticks([-0.3, 0, 0.3])
 
 cbar_ax2 = fig.add_axes([0.86, 0.1, 0.015, 0.25])
 cbar = fig.colorbar(im2, cax=cbar_ax2)
 cbar.set_label(r"$\Delta\, \phi_{\rho}$", rotation=270, labelpad=10)
-cbar_ax2.set_yticks([-0.5, 0, 0.5])
+cbar_ax2.set_yticks([-0.3, 0, 0.3])
 
 for i, m in enumerate(models):
     ax = fig.add_subplot(gs[0, i])

plot_fig4.py

@@ -8,7 +8,7 @@
 import fig_settings as fs
 from lcs import *
 
-measure = "auprc"
+measure = "auroc"
 axis_limits = [0, 1]
 
 axislabel_fontsize = 20

plot_fig5.py

@@ -8,7 +8,8 @@
 import fig_settings as fs
 from lcs import *
 
-measure = "auprc"
+gamma = 0.1
+measure = "auroc"
 
 axis_limits = [0, 1]
 
@@ -90,17 +91,18 @@ def visualize_networks(i, ax):
         data = json.load(file)
     var = np.array(data[keys[i]], dtype=float)
     b = np.array(data["beta"], dtype=float)
+    l = np.array(data["pf-eigenvalue"], dtype=float)[0].mean(axis=0).mean(axis=1)
 
     performance = np.array(data[measure], dtype=float)
 
-    # plot the difference in auprc
+    # plot the difference in auroc
     mean_difference = performance[2].mean(axis=2).T - performance[0].mean(axis=2).T
     ax = fig.add_subplot(gs[1, i])
     im1 = ax.imshow(
         to_imshow_orientation(mean_difference),
         extent=(min(var), max(var), min(b), max(b)),
-        vmin=-0.5,
-        vmax=0.5,
+        vmin=-0.3,
+        vmax=0.3,
         aspect="auto",
         cmap=cmap,
     )
@@ -109,6 +111,10 @@ def visualize_networks(i, ax):
     ax.set_xticks(xticks[i], xticklabels[i])
     ax.set_yticks([0, 0.5, 1], [0, 0.5, 1])
 
+    beta_c = gamma / l
+    for r0 in range(1, 10, 2):
+        ax.plot(var, r0 * beta_c, "r-", linewidth=0.5)
+
     if i == 0:
         ax.set_ylabel(r"$\beta$")
     else:
@@ -126,8 +132,8 @@ def visualize_networks(i, ax):
     im2 = ax.imshow(
         to_imshow_orientation(mean_difference),
         extent=(min(var), max(var), min(b), max(b)),
-        vmin=-0.5,
-        vmax=0.5,
+        vmin=-0.3,
+        vmax=0.3,
         aspect="auto",
         cmap=cmap,
     )
@@ -146,12 +152,12 @@ def visualize_networks(i, ax):
 cbar_ax1 = fig.add_axes([0.85, 0.4, 0.015, 0.25])
 cbar = fig.colorbar(im1, cax=cbar_ax1)
 cbar.set_label(rf"$\Delta$ {measure.upper()}", rotation=270, labelpad=10)
-cbar_ax1.set_yticks([-0.5, 0, 0.5])
+cbar_ax1.set_yticks([-0.3, 0, 0.3])
 
 cbar_ax2 = fig.add_axes([0.85, 0.1, 0.015, 0.25])
 cbar = fig.colorbar(im2, cax=cbar_ax2)
 cbar.set_label(r"$\Delta\, \phi_{\rho}$", rotation=270, labelpad=10)
-cbar_ax2.set_yticks([-0.5, 0, 0.5])
+cbar_ax2.set_yticks([-0.3, 0, 0.3])
 
 for i, m in enumerate(models):
     ax = fig.add_subplot(gs[0, i])