Use fANOVA with Spearman fallback for parameter sorting, add sample image

Author: yasumorishima

package/visualization/plot_beeswarm/README.md

@@ -30,11 +30,11 @@ plot_beeswarm = mod.plot_beeswarm
 
 
 def objective(trial: optuna.trial.Trial) -> float:
-    x = trial.suggest_float("x", -5.0, 5.0)
-    y = trial.suggest_float("y", -5.0, 5.0)
-    z = trial.suggest_float("z", -5.0, 5.0)
-    w = trial.suggest_float("w", -5.0, 5.0)
-    return x**2 + 0.5 * y**2 + 0.1 * z**2 + 0.01 * w**2
+    x = trial.suggest_float("x", 0.0, 10.0)
+    y = trial.suggest_float("y", 0.0, 10.0)
+    z = trial.suggest_float("z", 0.0, 10.0)
+    w = trial.suggest_float("w", 0.0, 10.0)
+    return 1.0 * x + 0.5 * y + 0.1 * z + 0.01 * w
 
 
 study = optuna.create_study()

package/visualization/plot_beeswarm/__init__.py

@@ -131,30 +131,36 @@ def _compute_density_jitter(
 
 
 def _sort_params_by_importance(
+    study: optuna.Study,
     param_names: list[str],
     param_values: dict[str, np.ndarray],
     objectives: np.ndarray,
 ) -> list[str]:
-    """Sort parameters by absolute Spearman-like correlation with objective.
+    """Sort parameters by importance (most important at top of plot).
 
-    Parameters with stronger monotonic relationships with the objective
-    are placed at the top (higher index = plotted higher), matching the
-    convention of SHAP beeswarm plots.
+    Tries Optuna's built-in fANOVA-based importance first (handles non-monotonic
+    relationships). Falls back to Spearman rank correlation if sklearn is not
+    installed.
     """
+    # Try Optuna's fANOVA (requires sklearn).
+    try:
+        importances = optuna.importance.get_param_importances(study)
+        return sorted(param_names, key=lambda n: importances.get(n, 0.0))
+    except (ImportError, RuntimeError):
+        pass
+
+    # Fallback: absolute Spearman rank correlation.
     correlations: dict[str, float] = {}
     for name in param_names:
         vals = param_values[name]
-        # Use rank correlation (Spearman) without scipy.
         valid = np.isfinite(vals) & np.isfinite(objectives)
         if valid.sum() < 3:
             correlations[name] = 0.0
             continue
         v = vals[valid]
         o = objectives[valid]
-        # Rank the arrays.
         v_rank = np.argsort(np.argsort(v)).astype(float)
         o_rank = np.argsort(np.argsort(o)).astype(float)
-        # Pearson correlation of ranks = Spearman correlation.
         v_rank -= v_rank.mean()
         o_rank -= o_rank.mean()
         denom = np.sqrt((v_rank**2).sum() * (o_rank**2).sum())
@@ -163,7 +169,6 @@ def _sort_params_by_importance(
         else:
             correlations[name] = abs(float(np.dot(v_rank, o_rank) / denom))
 
-    # Sort ascending so that the most important param is at the top of the plot.
     return sorted(param_names, key=lambda n: correlations.get(n, 0.0))
 
 
@@ -215,7 +220,7 @@ def plot_beeswarm(
         raise ValueError("No valid parameters found in completed trials.")
 
     # Sort parameters by importance (least important at bottom).
-    sorted_params = _sort_params_by_importance(param_names, param_values, objectives)
+    sorted_params = _sort_params_by_importance(study, param_names, param_values, objectives)
 
     # Resolve colormap.
     cmap: Colormap = cm.get_cmap(color_map)

package/visualization/plot_beeswarm/example.py

@@ -21,13 +21,13 @@
 
 
 def objective(trial: optuna.trial.Trial) -> float:
-    """Multi-parameter objective with clear parameter-objective relationships."""
-    x = trial.suggest_float("x", -5.0, 5.0)
-    y = trial.suggest_float("y", -5.0, 5.0)
-    z = trial.suggest_float("z", -5.0, 5.0)
-    w = trial.suggest_float("w", -5.0, 5.0)
+    """Multi-parameter objective with clear monotonic relationships."""
+    x = trial.suggest_float("x", 0.0, 10.0)
+    y = trial.suggest_float("y", 0.0, 10.0)
+    z = trial.suggest_float("z", 0.0, 10.0)
+    w = trial.suggest_float("w", 0.0, 10.0)
     # x has the strongest influence, w the weakest.
-    return x**2 + 0.5 * y**2 + 0.1 * z**2 + 0.01 * w**2
+    return 1.0 * x + 0.5 * y + 0.1 * z + 0.01 * w
 
 
 if __name__ == "__main__":