Auto-format with Black

ppi_py/baselines.py

@@ -309,6 +309,7 @@ def classical_poisson_ci(X, Y, alpha=0.1, alternative="two-sided"):
         pointest, np.sqrt(np.diag(cov_mat) / n), alpha, alternative
     )
 
+
 """
     BOOTSTRAP CI
 

ppi_py/ppi.py

@@ -1069,6 +1069,7 @@ def ppi_logistic_ci(
         alternative=alternative,
     )
 
+
 def ppi_poisson_pointestimate(
     X,
     Y,
@@ -1102,8 +1103,12 @@ def ppi_poisson_pointestimate(
     d = X.shape[1]
     N = Yhat_unlabeled.shape[0]
     w = np.ones(n) if w is None else w / w.sum() * n
-    w_unlabeled = np.ones(N) if w_unlabeled is None else w_unlabeled / w_unlabeled.sum() * N
-
+    w_unlabeled = (
+        np.ones(N)
+        if w_unlabeled is None
+        else w_unlabeled / w_unlabeled.sum() * N
+    )
+
     if optimizer_options is None:
         optimizer_options = {"ftol": 1e-15}
     if "ftol" not in optimizer_options.keys():
@@ -1131,14 +1136,15 @@ def poisson_loss(_theta):
             / N
             * np.sum(
                 w_unlabeled
-                * (np.exp(X_unlabeled @ _theta) - Yhat_unlabeled * (X_unlabeled @ _theta))
+                * (
+                    np.exp(X_unlabeled @ _theta)
+                    - Yhat_unlabeled * (X_unlabeled @ _theta)
+                )
             )
             - lam_curr
             / n
             * np.sum(w * (np.exp(X @ _theta) - Yhat * (X @ _theta)))
-            + 1
-            / n
-            * np.sum(w * (np.exp(X @ _theta) - Y * (X @ _theta)))
+            + 1 / n * np.sum(w * (np.exp(X @ _theta) - Y * (X @ _theta)))
         )
 
     def poisson_grad(_theta):
@@ -1147,14 +1153,8 @@ def poisson_grad(_theta):
             / N
             * X_unlabeled.T
             @ (w_unlabeled * (np.exp(X_unlabeled @ _theta) - Yhat_unlabeled))
-            - lam_curr
-            / n
-            * X.T
-            @ (w * (np.exp(X @ _theta) - Yhat))
-            + 1
-            / n
-            * X.T
-            @ (w * (np.exp(X @ _theta) - Y))
+            - lam_curr / n * X.T @ (w * (np.exp(X @ _theta) - Yhat))
+            + 1 / n * X.T @ (w * (np.exp(X @ _theta) - Y))
         )
 
     ppi_pointest = minimize(
@@ -1240,7 +1240,11 @@ def _poisson_get_stats(
     d = X.shape[1]
     N = Yhat_unlabeled.shape[0]
     w = np.ones(n) if w is None else w / w.sum() * n
-    w_unlabeled = np.ones(N) if w_unlabeled is None else w_unlabeled / w_unlabeled.sum() * N
+    w_unlabeled = (
+        np.ones(N)
+        if w_unlabeled is None
+        else w_unlabeled / w_unlabeled.sum() * N
+    )
 
     mu = np.exp(X @ pointest)
     mu_til = np.exp(X_unlabeled @ pointest)
@@ -1316,7 +1320,11 @@ def ppi_poisson_ci(
     d = X.shape[1]
     N = Yhat_unlabeled.shape[0]
     w = np.ones(n) if w is None else w / w.sum() * n
-    w_unlabeled = np.ones(N) if w_unlabeled is None else w_unlabeled / w_unlabeled.sum() * N
+    w_unlabeled = (
+        np.ones(N)
+        if w_unlabeled is None
+        else w_unlabeled / w_unlabeled.sum() * N
+    )
     use_unlabeled = lam != 0
 
     ppi_pointest = ppi_poisson_pointestimate(
@@ -1377,6 +1385,7 @@ def ppi_poisson_ci(
         alternative=alternative,
     )
 
+
 """
     PPBOOT
 

tests/test_poisson.py

@@ -7,6 +7,7 @@
     PPI tests for Poisson regression
 """
 
+
 def test_ppi_poisson_pointestimate_debias():
     # Make a synthetic regression problem
     n = 100
@@ -34,6 +35,7 @@ def test_ppi_poisson_pointestimate_debias():
         beta_prediction - beta
     )  # Makes it less biased
 
+
 def test_ppi_poisson_pointestimate_recovers():
     # Make a synthetic regression problem
     n = 10000
@@ -58,6 +60,7 @@ def test_ppi_poisson_pointestimate_recovers():
     # Check that the point estimate is close to the true beta
     assert np.linalg.norm(beta_ppi_pointestimate - beta) < 0.2
 
+
 def ppi_poisson_ci_subtest(i, alphas, n=1000, N=10000, d=1, epsilon=0.02):
     includeds = np.zeros(len(alphas))
     # Make a synthetic regression problem
@@ -86,6 +89,7 @@ def ppi_poisson_ci_subtest(i, alphas, n=1000, N=10000, d=1, epsilon=0.02):
         )
     return includeds
 
+
 def test_ppi_poisson_ci_parallel():
     n = 1000
     N = 10000
@@ -111,10 +115,12 @@ def test_ppi_poisson_ci_parallel():
     failed = np.any((total_includeds / num_trials) < (1 - alphas - epsilon))
     assert not failed
 
+
 """
     Baseline tests
 """
 
+
 def classical_poisson_ci_subtest(i, alphas, n, d, epsilon):
     includeds = np.zeros(len(alphas))
     # Make a synthetic regression problem
@@ -130,6 +136,7 @@ def classical_poisson_ci_subtest(i, alphas, n, d, epsilon):
         )
     return includeds
 
+
 def test_classical_poisson_ci_parallel():
     n = 1000
     d = 2
@@ -152,4 +159,4 @@ def test_classical_poisson_ci_parallel():
 
     print("Classical: ", (total_includeds / num_trials))
     failed = np.any((total_includeds / num_trials) < (1 - alphas - epsilon))
-    assert not failed
+    assert not failed