formatting

Author: SvenKlaassen

doubleml/did/did_multi.py

@@ -1040,24 +1040,25 @@ def plot_effects(
         """
         if self.framework is None:
             raise ValueError("Apply fit() before plot_effects().")
-        
+
         if result_type not in ["effect", "rv", "est_bounds", "ci_bounds"]:
             raise ValueError("result_type must be either 'effect', 'rv', 'est_bounds' or 'ci_bounds'.")
-        
+
         if result_type != "effect" and self._framework.sensitivity_params is None:
-            raise ValueError(f"result_type='{result_type}' requires sensitivity analysis. "
-                           "Please call sensitivity_analysis() first.")
-        
+            raise ValueError(
+                f"result_type='{result_type}' requires sensitivity analysis. " "Please call sensitivity_analysis() first."
+            )
+
         df = self._create_ci_dataframe(level=level, joint=joint)
 
         # Set default y_label and title based on result_type
         label_configs = {
             "effect": {"y_label": "Effect", "title": "Estimated ATTs by Group"},
             "rv": {"y_label": "Robustness Value", "title": "Robustness Values by Group"},
             "est_bounds": {"y_label": "Estimate Bounds", "title": "Estimate Bounds by Group"},
-            "ci_bounds": {"y_label": "Confidence Interval Bounds", "title": "Confidence Interval Bounds by Group"}
+            "ci_bounds": {"y_label": "Confidence Interval Bounds", "title": "Confidence Interval Bounds by Group"},
         }
-        
+
         config = label_configs[result_type]
         y_label = y_label if y_label is not None else config["y_label"]
         title = title if title is not None else config["title"]
@@ -1198,10 +1199,20 @@ def _plot_single_group(self, ax, period_df, period, result_type, colors, is_date
         plot_configs = {
             "effect": {"plot_col": "Estimate", "err_col_upper": "CI Upper", "err_col_lower": "CI Lower", "s_val": 30},
             "rv": {"plot_col": "RV", "plot_col_2": "RVa", "s_val": 50},
-            "est_bounds": {"plot_col": "Estimate", "err_col_upper": "Estimate Upper Bound", "err_col_lower": "Estimate Lower Bound", "s_val": 30},
-            "ci_bounds": {"plot_col": "Estimate", "err_col_upper": "CI Upper Bound", "err_col_lower": "CI Lower Bound", "s_val": 30}
+            "est_bounds": {
+                "plot_col": "Estimate",
+                "err_col_upper": "Estimate Upper Bound",
+                "err_col_lower": "Estimate Lower Bound",
+                "s_val": 30,
+            },
+            "ci_bounds": {
+                "plot_col": "Estimate",
+                "err_col_upper": "CI Upper Bound",
+                "err_col_lower": "CI Lower Bound",
+                "s_val": 30,
+            },
         }
-        
+
         config = plot_configs[result_type]
         plot_col = config["plot_col"]
         plot_col_2 = config.get("plot_col_2")
@@ -1240,13 +1251,17 @@ def _plot_single_group(self, ax, period_df, period, result_type, colors, is_date
                         markeredgewidth=1,
                         linewidth=1,
                     )
-                    
+
                 elif result_type == "rv":
                     ax.scatter(
-                        category_data["jittered_x"], category_data[plot_col_2], color=colors[category_name], alpha=0.8, s=s_val,
-                        marker="s"
+                        category_data["jittered_x"],
+                        category_data[plot_col_2],
+                        color=colors[category_name],
+                        alpha=0.8,
+                        s=s_val,
+                        marker="s",
                     )
-        
+
         # Format axes
         if is_datetime:
             period_str = np.datetime64(period, self._dml_data.datetime_unit)
@@ -1519,4 +1534,4 @@ def _create_ci_dataframe(self, level=0.95, joint=True):
             df["CI Upper Bound"] = self.framework.sensitivity_params["ci"]["upper"]
             df["Estimate Lower Bound"] = self.framework.sensitivity_params["theta"]["lower"]
             df["Estimate Upper Bound"] = self.framework.sensitivity_params["theta"]["upper"]
-        return df
+        return df

doubleml/did/tests/test_did_multi_plot.py

@@ -195,102 +195,98 @@ def test_plot_effects_result_types(doubleml_did_fixture):
     fig_effect, axes_effect = dml_obj.plot_effects(result_type="effect")
     assert isinstance(fig_effect, plt.Figure)
     assert isinstance(axes_effect, list)
-    
+
     # Check that the default y-label is set correctly
     assert axes_effect[0].get_ylabel() == "Effect"
-    
+
     plt.close("all")
 
 
 @pytest.mark.ci
 def test_plot_effects_result_type_rv(doubleml_did_fixture):
     """Test plot_effects with result_type='rv' (requires sensitivity analysis)."""
     dml_obj = doubleml_did_fixture["model"]
-    
+
     # Perform sensitivity analysis first
     dml_obj.sensitivity_analysis(cf_y=0.03, cf_d=0.03)
-    
+
     # Test result_type='rv'
     fig_rv, axes_rv = dml_obj.plot_effects(result_type="rv")
     assert isinstance(fig_rv, plt.Figure)
     assert isinstance(axes_rv, list)
-    
+
     # Check that the y-label is set correctly
     assert axes_rv[0].get_ylabel() == "Robustness Value"
-    
+
     plt.close("all")
 
 
 @pytest.mark.ci
 def test_plot_effects_result_type_est_bounds(doubleml_did_fixture):
     """Test plot_effects with result_type='est_bounds' (requires sensitivity analysis)."""
     dml_obj = doubleml_did_fixture["model"]
-    
+
     # Perform sensitivity analysis first
     dml_obj.sensitivity_analysis(cf_y=0.03, cf_d=0.03)
-    
+
     # Test result_type='est_bounds'
     fig_est, axes_est = dml_obj.plot_effects(result_type="est_bounds")
     assert isinstance(fig_est, plt.Figure)
     assert isinstance(axes_est, list)
-    
+
     # Check that the y-label is set correctly
     assert axes_est[0].get_ylabel() == "Estimate Bounds"
-    
+
     plt.close("all")
 
 
 @pytest.mark.ci
 def test_plot_effects_result_type_ci_bounds(doubleml_did_fixture):
     """Test plot_effects with result_type='ci_bounds' (requires sensitivity analysis)."""
     dml_obj = doubleml_did_fixture["model"]
-    
+
     # Perform sensitivity analysis first
     dml_obj.sensitivity_analysis(cf_y=0.03, cf_d=0.03)
-    
+
     # Test result_type='ci_bounds'
     fig_ci, axes_ci = dml_obj.plot_effects(result_type="ci_bounds")
     assert isinstance(fig_ci, plt.Figure)
     assert isinstance(axes_ci, list)
-    
+
     # Check that the y-label is set correctly
     assert axes_ci[0].get_ylabel() == "Confidence Interval Bounds"
-    
+
     plt.close("all")
 
 
 @pytest.mark.ci
 def test_plot_effects_result_type_invalid(doubleml_did_fixture):
     """Test plot_effects with invalid result_type."""
     dml_obj = doubleml_did_fixture["model"]
-    
+
     # Test with invalid result_type
     with pytest.raises(ValueError, match="result_type must be either"):
         dml_obj.plot_effects(result_type="invalid_type")
-    
+
     plt.close("all")
 
 
 @pytest.mark.ci
 def test_plot_effects_result_type_with_custom_labels(doubleml_did_fixture):
     """Test plot_effects with result_type and custom labels."""
     dml_obj = doubleml_did_fixture["model"]
-    
+
     # Perform sensitivity analysis first
     dml_obj.sensitivity_analysis(cf_y=0.03, cf_d=0.03)
-    
+
     # Test result_type with custom labels
     custom_title = "Custom Sensitivity Plot"
     custom_ylabel = "Custom Bounds Label"
-    
-    fig, axes = dml_obj.plot_effects(
-        result_type="est_bounds",
-        title=custom_title,
-        y_label=custom_ylabel
-    )
-    
+
+    fig, axes = dml_obj.plot_effects(result_type="est_bounds", title=custom_title, y_label=custom_ylabel)
+
     assert isinstance(fig, plt.Figure)
     assert fig._suptitle.get_text() == custom_title
     assert axes[0].get_ylabel() == custom_ylabel
-    
+
     plt.close("all")