Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Keith Battocchi <kbattocchi@gmail.com>

Author: kbattocchi

econml/dr/_drlearner.py

@@ -219,14 +219,14 @@ def _compute_trim_mask(self, T_pred):
         trim_mask : ndarray of shape (n_samples,), dtype bool
             True for samples to keep, False for samples to trim.
         """
-        # Check for multiple treatments with trimming and bail early
-        if self._trimming_threshold is not None:
-            if T_pred.shape[1] > 2:
-                warn("Sample trimming is currently only supported for binary treatment "
-                     "(single treatment plus control).", UserWarning)
-                return np.ones(T_pred.shape[0], dtype=bool)
+        # If no trimming threshold specified, keep all samples
+        if self._trimming_threshold is None:
+            return np.ones(T_pred.shape[0], dtype=bool)
 
-        elif self._trimming_threshold is None:
+        # Check for multiple treatments with trimming and bail early
+        if T_pred.shape[1] > 2:
+            warn("Sample trimming is currently only supported for binary treatment "
+                 "(single treatment plus control).", UserWarning)
             return np.ones(T_pred.shape[0], dtype=bool)
 
         if self._trimming_threshold == 'auto':

econml/tests/test_drlearner.py

@@ -1296,7 +1296,7 @@ def test_auto_trimming_uniform_propensity_full_range(self):
         Y = T + np.random.randn(n) * 0.5
 
         # Use a flexible model that can learn the propensity accurately
-        # Gradient boosting can fit the non-linear relationship well
+# The propensity is linear in the feature, so just use LinearRegression as the propensity model even though it's a regressor rather than a classifier.
         est = DRLearner(
             model_propensity=LinearRegression(),
             model_regression=LinearRegression(),
@@ -1331,7 +1331,7 @@ def test_auto_trimming_uniform_propensity_narrow_range(self):
         # Generate outcome
         Y = T + np.random.randn(n) * 0.5
 
-        # Use a flexible model that can learn the propensity accurately
+# The propensity is linear in the feature, so just use LinearRegression as the propensity model even though it's a regressor rather than a classifier.
         est = DRLearner(
             model_propensity=LinearRegression(),
             model_regression=LinearRegression(),