From 8769b16e28b5bbb9c9c59d803264025de6f87e00 Mon Sep 17 00:00:00 2001
From: max-de-rooij <54850292+max-de-rooij@users.noreply.github.com>
Date: Fri, 2 Aug 2024 13:33:53 +0200
Subject: [PATCH] Update mutual_information_sensitivity.jl

---
 src/mutual_information_sensitivity.jl | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/mutual_information_sensitivity.jl b/src/mutual_information_sensitivity.jl
index 512fba5d..46c8236e 100644
--- 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