Improve horseshoe prior demo: pedagogical visualizations, bug fix, Python 3.12 + CUDA 12 environment

README.rst

@@ -88,6 +88,42 @@ For development:
     $ pip install -e ".[dev]"
 
 
+Environment Setup (Python 3.12 + CUDA 12)
+------------------------------------------
+
+**Python 3.12 (CPU only)**
+
+.. code-block:: console
+
+    $ pip install bayesian_feature_selection
+
+This installs JAX ≥ 0.7.0 and NumPyro ≥ 0.15.0 automatically.
+
+**Python 3.12 + CUDA 12 (GPU)**
+
+Install the package, then upgrade JAX with the CUDA 12 backend:
+
+.. code-block:: console
+
+    $ pip install bayesian_feature_selection
+    $ pip install "jax[cuda12]"
+
+Verify the setup:
+
+.. code-block:: python
+
+    import jax
+    print(jax.devices())  # should show CudaDevice(id=0) when GPU is available
+
+**Important version notes**
+
+* NumPyro ≥ 0.15.0 requires JAX ≥ 0.7.0.
+* JAX ≥ 0.10.0 removed internal symbols that NumPyro ≤ 0.20.1 depends on;
+  use JAX 0.7.x – 0.9.x until NumPyro releases a compatible update.
+* The ``gpu`` extra (``pip install bayesian_feature_selection[gpu]``) installs
+  ``jax[cuda12]`` and is the recommended way to enable GPU support.
+
+
 Credits
 -------
 

pyproject.toml

@@ -20,14 +20,15 @@ classifiers = [
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
 ]
 license = {text = "MIT license"}
 dependencies = [
     "typer>=0.9.0",
     "rich>=13.0.0",
-    "numpyro>=0.13.0",
-    "jax>=0.4.0",
-    "jaxlib>=0.4.0",
+    "numpyro>=0.15.0",
+    "jax>=0.7.0",
+    "jaxlib>=0.7.0",
     "arviz>=0.16.0",
     "pandas>=2.0.0",
     "numpy>=1.24.0",
@@ -47,7 +48,7 @@ dev = [
     "ipython",
 ]
 gpu = [
-    "jax[cuda12_pip]",
+    "jax[cuda12]",
 ]
 
 [project.scripts]

src/bayesian_feature_selection/visualization.py

@@ -12,7 +12,8 @@ def plot_feature_importance(
     importance_df: pd.DataFrame,
     output_dir: Path,
     feature_names: Optional[List[str]] = None,
-    top_n: int = 20
+    top_n: int = 20,
+    method: str = "beta",
 ):
     """
     Plot feature importance with credible intervals.
@@ -27,6 +28,9 @@ def plot_feature_importance(
         Feature names
     top_n : int
         Number of top features to plot
+    method : str
+        Feature selection method used ('beta', 'lambda', or 'both').
+        Determines which inclusion probability column to display.
     """
     # Select top features
     top_features = importance_df.head(top_n).copy()
@@ -81,7 +85,14 @@ def plot_feature_importance(
     fig, ax = plt.subplots(figsize=(10, 8))
 
     colors = top_features["selected"].map({True: "green", False: "red"})
-    ax.barh(y_pos, top_features["inclusion_prob"], color=colors, alpha=0.6)
+    # Select the correct inclusion-probability column based on the method used
+    if method == "lambda":
+        inc_col = "lambda_inclusion_prob"
+    elif method == "both":
+        inc_col = "combined_inclusion_prob"
+    else:  # default: "beta"
+        inc_col = "beta_inclusion_prob"
+    ax.barh(y_pos, top_features[inc_col], color=colors, alpha=0.6)
 
     ax.set_yticks(y_pos)
     if feature_names is not None: