add pylint fixes

Author: juanitorduz

pymc_bart/pgbart.py

@@ -27,7 +27,13 @@
 
 from pymc_bart.bart import BARTRV
 from pymc_bart.split_rules import ContinuousSplitRule
-from pymc_bart.tree import Node, Tree, get_depth, get_idx_left_child, get_idx_right_child
+from pymc_bart.tree import (
+    Node,
+    Tree,
+    get_depth,
+    get_idx_left_child,
+    get_idx_right_child,
+)
 
 
 class ParticleTree:
@@ -110,7 +116,7 @@ class PGBART(ArrayStepShared):
     generates_stats = True
     stats_dtypes = [{"variable_inclusion": object, "tune": bool}]
 
-    def __init__(
+    def __init__(  # noqa: PLR0915
         self,
         vars=None,  # pylint: disable=redefined-builtin
         num_particles: int = 10,
@@ -544,11 +550,10 @@ def draw_leaf_value(
 
     if y_mu_pred.size == 1:
         mu_mean = np.full(shape, y_mu_pred.item() / m) + norm
+    elif y_mu_pred.size < 3 or response == "constant":
+        mu_mean = fast_mean(y_mu_pred) / m + norm
     else:
-        if y_mu_pred.size < 3 or response == "constant":
-            mu_mean = fast_mean(y_mu_pred) / m + norm
-        else:
-            mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m, norm=norm)
+        mu_mean, linear_params = fast_linear_fit(x=x_mu, y=y_mu_pred, m=m, norm=norm)
 
     return mu_mean, linear_params
 

pymc_bart/utils.py

@@ -279,15 +279,14 @@ def identity(x):
             if var in var_discrete:
                 axes[count].plot(new_x, p_di.mean(0), "o", color=color_mean)
                 axes[count].plot(new_x, p_di.T, ".", color=color, alpha=alpha)
+            elif smooth:
+                x_data, y_data = _smooth_mean(new_x, p_di, "ice", smooth_kwargs)
+                axes[count].plot(x_data, y_data.mean(1), color=color_mean)
+                axes[count].plot(x_data, y_data, color=color, alpha=alpha)
             else:
-                if smooth:
-                    x_data, y_data = _smooth_mean(new_x, p_di, "ice", smooth_kwargs)
-                    axes[count].plot(x_data, y_data.mean(1), color=color_mean)
-                    axes[count].plot(x_data, y_data, color=color, alpha=alpha)
-                else:
-                    idx = np.argsort(new_x)
-                    axes[count].plot(new_x[idx], p_di.mean(0)[idx], color=color_mean)
-                    axes[count].plot(new_x[idx], p_di.T[idx], color=color, alpha=alpha)
+                idx = np.argsort(new_x)
+                axes[count].plot(new_x[idx], p_di.mean(0)[idx], color=color_mean)
+                axes[count].plot(new_x[idx], p_di.T[idx], color=color, alpha=alpha)
             axes[count].set_xlabel(x_labels[i_var])
 
             count += 1
@@ -515,13 +514,12 @@ def _get_axes(
             for i in range(n_plots, len(axes)):
                 fig.delaxes(axes[i])
             axes = axes[:n_plots]
+    elif isinstance(ax, np.ndarray):
+        axes = ax
+        fig = ax[0].get_figure()
     else:
-        if isinstance(ax, np.ndarray):
-            axes = ax
-            fig = ax[0].get_figure()
-        else:
-            axes = [ax]
-            fig = ax.get_figure()  # type: ignore
+        axes = [ax]
+        fig = ax.get_figure()  # type: ignore
 
     return fig, axes, shape
 
@@ -694,7 +692,7 @@ def _smooth_mean(
     return x_data, y_data
 
 
-def plot_variable_importance(
+def plot_variable_importance(  # noqa: PLR0915
     idata: az.InferenceData,
     bartrv: Variable,
     X: npt.NDArray[np.float_],

pyproject.toml

@@ -1,17 +1,21 @@
 [tool.pytest.ini_options]
 minversion = "6.0"
-xfail_strict=true
-addopts = [
-    "-vv",
-    "--color=yes",
-]
+xfail_strict = true
+addopts = ["-vv", "--color=yes"]
 
 [tool.ruff]
 line-length = 100
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "PL", "UP", "W"]
 ignore-init-module-imports = true
+ignore = [
+  "PLR2004", # Checks for the use of unnamed numerical constants ("magic") values in comparisons.
+]
+
+[tool.ruff.lint.pylint]
+max-args = 19
+max-branches = 15
 
 [tool.ruff.extend-per-file-ignores]
 "docs/conf.py" = ["E501", "F541"]