Run formatting

workflows/_benchmark_utils.py

@@ -8,7 +8,7 @@ def generate_dataset(n_points: int, seed: int) -> None:
 
     X = rng.normal(size=(n_points, 3))
     A = rng.binomial(1, p=[0.1, 0.4, 0.1, 0.5], size=(n_points, 4))
-    index = np.argsort(["".join(map(str,  a)) for a in A])
+    index = np.argsort(["".join(map(str, a)) for a in A])
     A = A[index, :]
 
     states = np.hstack((A, X))
@@ -22,14 +22,18 @@ def generate_dataset(n_points: int, seed: int) -> None:
     effect_size = 4.0
 
     for i in range(n_covariates):
-        coefs[i*n_genes_per_covariate:(i+1)*n_genes_per_covariate, i] = effect_size
+        coefs[i * n_genes_per_covariate : (i + 1) * n_genes_per_covariate, i] = (
+            effect_size
+        )
 
     if n_additional_genes > 0:
-        coefs[-n_additional_genes:, :] = effect_size * rng.binomial(1, 0.5, size=(n_additional_genes, n_covariates))
+        coefs[-n_additional_genes:, :] = effect_size * rng.binomial(
+            1, 0.5, size=(n_additional_genes, n_covariates)
+        )
 
     offset = -5
     logits = offset + np.einsum("nf,gf->ng", states, coefs)
-    ps = 1/(1 + np.exp(-logits))
+    ps = 1 / (1 + np.exp(-logits))
     Y = rng.binomial(1, ps)
 
     return {
@@ -44,34 +48,34 @@ def generate_dataset(n_points: int, seed: int) -> None:
 def calculate_probabilities(samples):
     counts = defaultdict(int)
     total_samples = len(samples)
-    
+
     for s in samples:
         counts[s] += 1
-    
+
     probabilities = {k: v / total_samples for k, v in counts.items()}
     return probabilities
 
 
 def mutual_information(X_samples, Y_samples):
     assert len(X_samples) == len(Y_samples), "Mismatched sample sizes"
-    
+
     K = len(X_samples)
-    
+
     # Joint probabilities P(X, Y)
     joint_samples = [(tuple(x), tuple(y)) for x, y in zip(X_samples, Y_samples)]
     joint_probabilities = calculate_probabilities(joint_samples)
-    
+
     # Marginal probabilities P(X) and P(Y)
     X_probabilities = calculate_probabilities([tuple(x) for x in X_samples])
     Y_probabilities = calculate_probabilities([tuple(y) for y in Y_samples])
-    
+
     MI = 0
-    
+
     for (x, y), p_xy in joint_probabilities.items():
         p_x = X_probabilities.get(x, 0)
         p_y = Y_probabilities.get(y, 0)
-        
+
         if p_x > 0 and p_y > 0:
             MI += p_xy * np.log2(p_xy / (p_x * p_y))
-    
-    return MI
+
+    return MI