Vasilis/docs

econml/inference/_bootstrap.py

@@ -33,12 +33,15 @@ class BootstrapEstimator:
         This object must support a `fit` method which takes numpy arrays with consistent first dimensions
         as arguments.
 
-    n_bootstrap_samples : int
+    n_bootstrap_samples : int, default: 100
         How many draws to perform.
 
     n_jobs: int, default: None
         The maximum number of concurrently running jobs, as in joblib.Parallel.
 
+    verbose: int, default: 0
+        Verbosity level
+
     compute_means : bool, default: True
         Whether to pass calls through to the underlying collection and return the mean.  Setting this
         to ``False`` can avoid ambiguities if the wrapped object itself has method names with an `_interval` suffix.
@@ -50,10 +53,16 @@ class BootstrapEstimator:
         assuming the replicates are normally distributed.
     """
 
-    def __init__(self, wrapped, n_bootstrap_samples=1000, n_jobs=None, compute_means=True, bootstrap_type='pivot'):
+    def __init__(self, wrapped,
+                 n_bootstrap_samples=100,
+                 n_jobs=None,
+                 verbose=0,
+                 compute_means=True,
+                 bootstrap_type='pivot'):
         self._instances = [clone(wrapped, safe=False) for _ in range(n_bootstrap_samples)]
         self._n_bootstrap_samples = n_bootstrap_samples
         self._n_jobs = n_jobs
+        self._verbose = verbose
         self._compute_means = compute_means
         self._bootstrap_type = bootstrap_type
         self._wrapped = wrapped
@@ -109,7 +118,7 @@ def convertArg(arg, inds):
             else:  # arg was a scalar, so we shouldn't have converted it
                 return arg
 
-        self._instances = Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=3)(
+        self._instances = Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=self._verbose)(
             delayed(fit)(obj,
                          *[convertArg(arg, inds) for arg in args],
                          **{arg: convertArg(named_args[arg], inds) for arg in named_args})
@@ -130,7 +139,7 @@ def __getattr__(self, name):
 
         def proxy(make_call, name, summary):
             def summarize_with(f):
-                results = np.array(Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=3)(
+                results = np.array(Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=self._verbose)(
                     (f, (obj, name), {}) for obj in self._instances)), f(self._wrapped, name)
                 return summary(*results)
             if make_call:

econml/inference/_inference.py

@@ -68,6 +68,9 @@ class BootstrapInference(Inference):
     n_jobs: int, optional (default -1)
         The maximum number of concurrently running jobs, as in joblib.Parallel.
 
+    verbose: int, default: 0
+        Verbosity level
+
     bootstrap_type: 'percentile', 'pivot', or 'normal', default 'pivot'
         Bootstrap method used to compute results.
         'percentile' will result in using the empiracal CDF of the replicated computations of the statistics.
@@ -76,14 +79,15 @@ class BootstrapInference(Inference):
         'normal' will instead compute a pivot interval assuming the replicates are normally distributed.
     """
 
-    def __init__(self, n_bootstrap_samples=100, n_jobs=-1, bootstrap_type='pivot'):
+    def __init__(self, n_bootstrap_samples=100, n_jobs=-1, bootstrap_type='pivot', verbose=0):
         self._n_bootstrap_samples = n_bootstrap_samples
         self._n_jobs = n_jobs
         self._bootstrap_type = bootstrap_type
+        self._verbose = verbose
 
     def fit(self, estimator, *args, **kwargs):
         est = BootstrapEstimator(estimator, self._n_bootstrap_samples, self._n_jobs, compute_means=False,
-                                 bootstrap_type=self._bootstrap_type)
+                                 bootstrap_type=self._bootstrap_type, verbose=self._verbose)
         est.fit(*args, **kwargs)
         self._est = est
         self._d_t = estimator._d_t

econml/tests/test_bootstrap.py

@@ -274,7 +274,7 @@ def test_stratify(self):
         Y = [1, 2, 3, 4, 5, 6]
         X = np.array([1, 1, 2, 2, 1, 2]).reshape(-1, 1)
         est = LinearDML(model_y=LinearRegression(), model_t=LogisticRegression(), discrete_treatment=True)
-        inference = BootstrapInference(n_bootstrap_samples=5)
+        inference = BootstrapInference(n_bootstrap_samples=5, n_jobs=-1, verbose=0)
         est.fit(Y, T, inference=inference)
         est.const_marginal_effect_interval()
 
@@ -292,7 +292,7 @@ def test_stratify_orthoiv(self):
         X = np.array([1, 1, 2, 2, 1, 2, 1, 2]).reshape(-1, 1)
         est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(),
                                       flexible_model_effect=LinearRegression(), cv=2)
-        inference = BootstrapInference(n_bootstrap_samples=20)
+        inference = BootstrapInference(n_bootstrap_samples=20, n_jobs=-1, verbose=3)
         est.fit(Y, T, Z=Z, X=X, inference=inference)
         est.const_marginal_effect_interval(X)
 
@@ -303,7 +303,7 @@ def test_all_kinds(self):
         est = LinearDML(cv=2)
         for kind in ['percentile', 'pivot', 'normal']:
             with self.subTest(kind=kind):
-                inference = BootstrapInference(n_bootstrap_samples=5, bootstrap_type=kind)
+                inference = BootstrapInference(n_bootstrap_samples=5, n_jobs=-1, verbose=0, bootstrap_type=kind)
                 est.fit(Y, T, inference=inference)
                 i = est.const_marginal_effect_interval()
                 inf = est.const_marginal_effect_inference()