linting

Author: jmafoster1

causal_testing/generation/enum_gen.py

@@ -1,16 +1,18 @@
 """This module contains the class EnumGen, which allows us to easily create generating uniform distributions from enums."""
 
-from scipy.stats import rv_discrete
 from enum import Enum
+from scipy.stats import rv_discrete
 import numpy as np
 
 
 class EnumGen(rv_discrete):
-    """This class allows us to easily create generating uniform distributions from enums. This is helpful for generating concrete test inputs from abstract test cases."""
+    """This class allows us to easily create generating uniform distributions
+    from enums. This is helpful for generating concrete test inputs from
+    abstract test cases."""
 
     def __init__(self, datatype: Enum):
-        self.dt = dict(enumerate(datatype, 1))
-        self.inverse_dt = {v: k for k, v in self.dt.items()}
+        self.datatype = dict(enumerate(datatype, 1))
+        self.inverse_dt = {v: k for k, v in self.datatype.items()}
 
     def ppf(self, q):
         """Percent point function (inverse of `cdf`) at q of the given RV.
@@ -23,7 +25,7 @@ def ppf(self, q):
         k : array_like
             Quantile corresponding to the lower tail probability, q.
         """
-        return np.vectorize(self.dt.get)(np.ceil(len(self.dt) * q))
+        return np.vectorize(self.datatype.get)(np.ceil(len(self.datatype) * q))
 
     def cdf(self, q):
         """
@@ -37,4 +39,4 @@ def cdf(self, q):
         cdf : ndarray
             Cumulative distribution function evaluated at `x`
         """
-        return np.vectorize(self.inverse_dt.get)(q) / len(self.dt)
+        return np.vectorize(self.inverse_dt.get)(q) / len(self.datatype)

causal_testing/testing/causal_test_outcome.py

@@ -26,7 +26,7 @@ class SomeEffect(CausalTestOutcome):
     """An extension of TestOutcome representing that the expected causal effect should not be zero."""
 
     def apply(self, res: CausalTestResult) -> bool:
-        if res.test_value.type == "ate" or res.test_value.type == "coefficient":
+        if res.test_value.type in {"ate", "coefficient"}:
             return (0 < res.ci_low() < res.ci_high()) or (res.ci_low() < res.ci_high() < 0)
         if res.test_value.type == "risk_ratio":
             return (1 < res.ci_low() < res.ci_high()) or (res.ci_low() < res.ci_high() < 1)
@@ -38,7 +38,7 @@ class NoEffect(CausalTestOutcome):
 
     def apply(self, res: CausalTestResult) -> bool:
         print("RESULT", res)
-        if res.test_value.type == "ate" or res.test_value.type == "coefficient":
+        if res.test_value.type in {"ate", "coefficient"}:
             return (res.ci_low() < 0 < res.ci_high()) or (abs(res.test_value.value) < 1e-10)
         if res.test_value.type == "risk_ratio":
             return (res.ci_low() < 1 < res.ci_high()) or np.isclose(res.test_value.value, 1.0, atol=1e-10)
@@ -72,7 +72,7 @@ def apply(self, res: CausalTestResult) -> bool:
             return False
         if res.test_value.type in {"ate", "coefficient"}:
             return res.test_value.value > 0
-        elif res.test_value.type == "risk_ratio":
+        if res.test_value.type == "risk_ratio":
             return res.test_value.value > 1
         raise ValueError(f"Test Value type {res.test_value.type} is not valid for this TestOutcome")
 

causal_testing/testing/estimators.py

@@ -529,7 +529,8 @@ def add_modelling_assumptions(self):
 
     def estimate_coefficient(self, df):
         """
-        Estimate the linear regression coefficient of the treatment on the outcome.
+        Estimate the linear regression coefficient of the treatment on the
+        outcome.
         """
         # Estimate the total effect of instrument I on outcome Y = abI + c1
         ab = sm.OLS(df[self.outcome], df[[self.instrument]]).fit().params[self.instrument]
@@ -541,6 +542,10 @@ def estimate_coefficient(self, df):
         return ab / a
 
     def estimate_unit_ate(self, bootstrap_size=100):
+        """
+        Estimate the unit ate (i.e. coefficient) of the treatment on the
+        outcome.
+        """
         bootstraps = sorted(
             [self.estimate_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)]
         )