Update mutual_information_sensitivity.jl

SciML · Aug 2, 2024 · 8769b16 · 8769b16
1 parent e4d6ffe
commit 8769b16
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/src/mutual_information_sensitivity.jl b/src/mutual_information_sensitivity.jl
@@ -127,10 +127,11 @@ function _total_order_ci(Xi, Y, entropy_Y, discretization_entropy; n_bootstraps
     mi_values = zeros(n_bootstraps)
     est = ValueHistogram(Int(round(sqrt(length(Y)))))
     Y_perm = copy(Y)
+    entropy_Xi = entropy(est, Xi)
     for i in 1:n_bootstraps
         shuffle!(Y_perm)
 
-        conditional_entropy = entropy(est, StateSpaceSet(Xi, Y_perm)) - entropy(est, Xi)
+        conditional_entropy = entropy(est, StateSpaceSet(Xi, Y_perm)) - entropy_Xi
         mi_values[i] = (entropy_Y - conditional_entropy) / (entropy_Y - discretization_entropy)
     end
 
@@ -145,10 +146,12 @@ function _first_order_ci(Xi, Y; n_bootstraps = 100, conf_level = 0.95)
     mi_values = zeros(n_bootstraps)
     est = ValueHistogram(Int(round(sqrt(length(Y)))))
     Y_perm = copy(Y)
+    entropy_Xi = entropy(est, Xi)
+    entropy_Y = entropy(est, Y)
     for i in 1:n_bootstraps
         shuffle!(Y_perm)
-        mutual_information = entropy(est, Xi) + entropy(est, Y_perm) - entropy(est, StateSpaceSet(Xi, Y_perm))
-        mi_values[i] = mutual_information / entropy(est, Y_perm)
+        mutual_information = entropy_Xi + entropy_Y - entropy(est, StateSpaceSet(Xi, Y_perm))
+        mi_values[i] = mutual_information / entropy_Y
     end
 
     α = 1 - conf_level
@@ -162,11 +165,12 @@ function _second_order_ci(Xi, Xj, Y; n_bootstraps = 100, conf_level = 0.95)
     mi_values = zeros(n_bootstraps)
     est = ValueHistogram(Int(round(sqrt(length(Y)))))
     Y_perm = copy(Y)
+    mutual_information = entropy(est, Xi) + entropy(est, Xj) - entropy(est, StateSpaceSet(Xi, Xj))
+    entropy_Y = entropy(est, Y_perm)
     for i in 1:n_bootstraps
         shuffle!(Y_perm)
-        conditional_mutual_information = entropy(est, StateSpaceSet(Xi, Y_perm)) + entropy(est, StateSpaceSet(Xj, Y_perm)) - entropy(est, StateSpaceSet(Xi, Xj, Y_perm)) - entropy(est, Y_perm)
-        mutual_information = entropy(est, Xi) + entropy(est, Xj) - entropy(est, StateSpaceSet(Xi, Xj))
-        mi_values[i] = (conditional_mutual_information - mutual_information) / entropy(est, Y_perm)
+        conditional_mutual_information = entropy(est, StateSpaceSet(Xi, Y_perm)) + entropy(est, StateSpaceSet(Xj, Y_perm)) - entropy(est, StateSpaceSet(Xi, Xj, Y_perm)) - entropy_Y
+        mi_values[i] = (conditional_mutual_information - mutual_information) / entropy_Y
     end
 
     α = 1 - conf_level