[python-package] Add tests for plotting ylim, xlim, max_num_features, and pre-allocated ax

Author: DeborahOlaboye

tests/python_package_test/test_plotting.py

@@ -5,7 +5,12 @@
 from sklearn.model_selection import train_test_split
 
 import lightgbm as lgb
-from lightgbm.compat import GRAPHVIZ_INSTALLED, MATPLOTLIB_INSTALLED, PANDAS_INSTALLED, pd_DataFrame
+from lightgbm.compat import (
+    GRAPHVIZ_INSTALLED,
+    MATPLOTLIB_INSTALLED,
+    PANDAS_INSTALLED,
+    pd_DataFrame,
+)
 
 if MATPLOTLIB_INSTALLED:
     import matplotlib
@@ -19,17 +24,23 @@
 
 @pytest.fixture(scope="module")
 def breast_cancer_split():
-    return train_test_split(*load_breast_cancer(return_X_y=True), test_size=0.1, random_state=1)
+    return train_test_split(
+        *load_breast_cancer(return_X_y=True), test_size=0.1, random_state=1
+    )
 
 
 def _categorical_data(category_values_lower_bound, category_values_upper_bound):
     X, y = load_breast_cancer(return_X_y=True)
     X_df = pd.DataFrame()
     rnd = np.random.RandomState(0)
-    n_cat_values = rnd.randint(category_values_lower_bound, category_values_upper_bound, size=X.shape[1])
+    n_cat_values = rnd.randint(
+        category_values_lower_bound, category_values_upper_bound, size=X.shape[1]
+    )
     for i in range(X.shape[1]):
         bins = np.linspace(0, 1, num=n_cat_values[i] + 1)
-        X_df[f"cat_col_{i}"] = pd.qcut(X[:, i], q=bins, labels=range(n_cat_values[i])).as_unordered()
+        X_df[f"cat_col_{i}"] = pd.qcut(
+            X[:, i], q=bins, labels=range(n_cat_values[i])
+        ).as_unordered()
     return X_df, y
 
 
@@ -68,7 +79,9 @@ def test_plot_importance(params, breast_cancer_split, train_data):
     for patch in ax1.patches:
         assert patch.get_facecolor() == (1.0, 0, 0, 1.0)  # red
 
-    ax2 = lgb.plot_importance(gbm0, color=["r", "y", "g", "b"], title=None, xlabel=None, ylabel=None)
+    ax2 = lgb.plot_importance(
+        gbm0, color=["r", "y", "g", "b"], title=None, xlabel=None, ylabel=None
+    )
     assert isinstance(ax2, matplotlib.axes.Axes)
     assert ax2.get_title() == ""
     assert ax2.get_xlabel() == ""
@@ -80,7 +93,10 @@ def test_plot_importance(params, breast_cancer_split, train_data):
     assert ax2.patches[3].get_facecolor() == (0, 0, 1.0, 1.0)  # b
 
     ax3 = lgb.plot_importance(
-        gbm0, title="t @importance_type@", xlabel="x @importance_type@", ylabel="y @importance_type@"
+        gbm0,
+        title="t @importance_type@",
+        xlabel="x @importance_type@",
+        ylabel="y @importance_type@",
     )
     assert isinstance(ax3, matplotlib.axes.Axes)
     assert ax3.get_title() == "t @importance_type@"
@@ -97,20 +113,53 @@ def test_plot_importance(params, breast_cancer_split, train_data):
     assert len(ax4.patches) <= 30
 
     with pytest.raises(TypeError, match="xlim must be a tuple of 2 elements."):
-        lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, xlim="not a tuple")
+        lgb.plot_importance(
+            gbm0, title=None, xlabel=None, ylabel=None, xlim="not a tuple"
+        )
 
-    gbm2 = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1, importance_type="gain")
+    # test ylim parameter
+    ax5 = lgb.plot_importance(gbm0, title=None, xlabel=None, ylabel=None, ylim=(-1, 30))
+    assert isinstance(ax5, matplotlib.axes.Axes)
+    assert ax5.get_ylim() == (-1, 30)
+
+    with pytest.raises(TypeError, match="ylim must be a tuple of 2 elements."):
+        lgb.plot_importance(gbm0, ylim="not a tuple")
+
+    # test max_num_features parameter
+    ax6 = lgb.plot_importance(gbm0, max_num_features=5)
+    assert isinstance(ax6, matplotlib.axes.Axes)
+    assert len(ax6.patches) == 5
+
+    # test providing pre-allocated ax with figsize
+    fig, ax_prealloc = matplotlib.pyplot.subplots(1, 1, figsize=(12, 8))
+    ax7 = lgb.plot_importance(gbm0, ax=ax_prealloc, figsize=(6, 4))
+    assert ax7 is ax_prealloc
+    # when ax is provided, figsize should be ignored, so figure size remains (12, 8)
+    assert ax7.get_figure().get_figwidth() == 12
+    assert ax7.get_figure().get_figheight() == 8
+
+    gbm2 = lgb.LGBMClassifier(
+        n_estimators=10, num_leaves=3, verbose=-1, importance_type="gain"
+    )
     gbm2.fit(X_train, y_train)
 
     def get_bounds_of_first_patch(axes):
         return axes.patches[0].get_extents().bounds
 
     first_bar1 = get_bounds_of_first_patch(lgb.plot_importance(gbm1))
-    first_bar2 = get_bounds_of_first_patch(lgb.plot_importance(gbm1, importance_type="split"))
-    first_bar3 = get_bounds_of_first_patch(lgb.plot_importance(gbm1, importance_type="gain"))
+    first_bar2 = get_bounds_of_first_patch(
+        lgb.plot_importance(gbm1, importance_type="split")
+    )
+    first_bar3 = get_bounds_of_first_patch(
+        lgb.plot_importance(gbm1, importance_type="gain")
+    )
     first_bar4 = get_bounds_of_first_patch(lgb.plot_importance(gbm2))
-    first_bar5 = get_bounds_of_first_patch(lgb.plot_importance(gbm2, importance_type="split"))
-    first_bar6 = get_bounds_of_first_patch(lgb.plot_importance(gbm2, importance_type="gain"))
+    first_bar5 = get_bounds_of_first_patch(
+        lgb.plot_importance(gbm2, importance_type="split")
+    )
+    first_bar6 = get_bounds_of_first_patch(
+        lgb.plot_importance(gbm2, importance_type="gain")
+    )
 
     assert first_bar1 == first_bar2
     assert first_bar1 == first_bar5
@@ -153,7 +202,13 @@ def test_plot_split_value_histogram(params, breast_cancer_split, train_data):
         assert patch.get_facecolor() == (1.0, 0, 0, 1.0)  # red
 
     ax2 = lgb.plot_split_value_histogram(
-        gbm0, 27, bins=10, color=["r", "y", "g", "b"], title=None, xlabel=None, ylabel=None
+        gbm0,
+        27,
+        bins=10,
+        color=["r", "y", "g", "b"],
+        title=None,
+        xlabel=None,
+        ylabel=None,
     )
     assert isinstance(ax2, matplotlib.axes.Axes)
     assert ax2.get_title() == ""
@@ -165,14 +220,22 @@ def test_plot_split_value_histogram(params, breast_cancer_split, train_data):
     assert ax2.patches[2].get_facecolor() == (0, 0.5, 0, 1.0)  # g
     assert ax2.patches[3].get_facecolor() == (0, 0, 1.0, 1.0)  # b
 
+    # test xlim and ylim parameters
+    ax3 = lgb.plot_split_value_histogram(gbm0, 27, xlim=(0, 100), ylim=(0, 50))
+    assert isinstance(ax3, matplotlib.axes.Axes)
+    assert ax3.get_xlim() == (0, 100)
+    assert ax3.get_ylim() == (0, 50)
+
     with pytest.raises(
-        ValueError, match="Cannot plot split value histogram, because feature 0 was not used in splitting"
+        ValueError,
+        match="Cannot plot split value histogram, because feature 0 was not used in splitting",
     ):
         lgb.plot_split_value_histogram(gbm0, 0)  # was not used in splitting
 
 
 @pytest.mark.skipif(
-    not MATPLOTLIB_INSTALLED or not GRAPHVIZ_INSTALLED, reason="matplotlib or graphviz is not installed"
+    not MATPLOTLIB_INSTALLED or not GRAPHVIZ_INSTALLED,
+    reason="matplotlib or graphviz is not installed",
 )
 def test_plot_tree(breast_cancer_split):
     X_train, _, y_train, _ = breast_cancer_split
@@ -194,7 +257,9 @@ def test_create_tree_digraph(tmp_path, breast_cancer_split):
     X_train, _, y_train, _ = breast_cancer_split
 
     constraints = [-1, 1] * int(X_train.shape[1] / 2)
-    gbm = lgb.LGBMClassifier(n_estimators=10, num_leaves=3, verbose=-1, monotone_constraints=constraints)
+    gbm = lgb.LGBMClassifier(
+        n_estimators=10, num_leaves=3, verbose=-1, monotone_constraints=constraints
+    )
     gbm.fit(X_train, y_train)
 
     with pytest.raises(IndexError, match="tree_index is out of range."):
@@ -389,7 +454,9 @@ def test_example_case_in_tree_digraph():
         while "decision_type" in node:  # iterate through the splits
             split_index = node["split_index"]
 
-            node_in_graph = [n for n in gbody if f"split{split_index}" in n and "->" not in n]
+            node_in_graph = [
+                n for n in gbody if f"split{split_index}" in n and "->" not in n
+            ]
             assert len(node_in_graph) == 1
             seen_indices.add(gbody.index(node_in_graph[0]))
 
@@ -420,14 +487,22 @@ def test_example_case_in_tree_digraph():
         assert "color=blue" in leaf_in_graph[0]
         assert len(edge_to_leaf) == 1
         assert "color=blue" in edge_to_leaf[0]
-        seen_indices.update([gbody.index(leaf_in_graph[0]), gbody.index(edge_to_leaf[0])])
+        seen_indices.update(
+            [gbody.index(leaf_in_graph[0]), gbody.index(edge_to_leaf[0])]
+        )
 
         # check that the rest of the elements have black color
-        remaining_elements = [e for i, e in enumerate(graph.body) if i not in seen_indices and "graph" not in e]
+        remaining_elements = [
+            e
+            for i, e in enumerate(graph.body)
+            if i not in seen_indices and "graph" not in e
+        ]
         assert all("color=black" in e for e in remaining_elements)
 
         # check that we got to the expected leaf
-        expected_leaf = bst.predict(example_case, start_iteration=i, num_iteration=1, pred_leaf=True)[0]
+        expected_leaf = bst.predict(
+            example_case, start_iteration=i, num_iteration=1, pred_leaf=True
+        )[0]
         assert leaf_index == expected_leaf
     assert makes_categorical_splits
 
@@ -464,7 +539,9 @@ def test_plot_metrics(params, breast_cancer_split, train_data):
         num_boost_round=10,
         callbacks=[lgb.record_evaluation(evals_result0)],
     )
-    with pytest.warns(UserWarning, match="More than one metric available, picking one to plot."):
+    with pytest.warns(
+        UserWarning, match="More than one metric available, picking one to plot."
+    ):
         ax0 = lgb.plot_metric(evals_result0)
     assert isinstance(ax0, matplotlib.axes.Axes)
     assert ax0.get_title() == "Metric during training"
@@ -521,7 +598,12 @@ def test_plot_metrics(params, breast_cancer_split, train_data):
         assert not grid_line.get_visible()
 
     evals_result1 = {}
-    lgb.train(params, train_data, num_boost_round=10, callbacks=[lgb.record_evaluation(evals_result1)])
+    lgb.train(
+        params,
+        train_data,
+        num_boost_round=10,
+        callbacks=[lgb.record_evaluation(evals_result1)],
+    )
     with pytest.raises(ValueError, match="eval results cannot be empty."):
         lgb.plot_metric(evals_result1)
 
@@ -535,3 +617,11 @@ def test_plot_metrics(params, breast_cancer_split, train_data):
     legend_items = ax4.get_legend().get_texts()
     assert len(legend_items) == 1
     assert legend_items[0].get_text() == "valid_0"
+
+    # test xlim and ylim parameters
+    ax5 = lgb.plot_metric(
+        evals_result0, metric="binary_logloss", xlim=(0, 15), ylim=(0, 1)
+    )
+    assert isinstance(ax5, matplotlib.axes.Axes)
+    assert ax5.get_xlim() == (0, 15)
+    assert ax5.get_ylim() == (0, 1)