Merge pull request #54 from vyron-arvanitis/dev-emanuel

Feedbackround I

Author: angelos-asln

.gitignore

@@ -75,4 +75,4 @@ target/
 .LSOverride
 
 # auto-generated files
-bde/_version.py
+bde/_version.py

.idea/bde.iml

@@ -1,4 +1,4 @@
-bde: Bayesian Deep Ensembles for scikit-learn and JAX
+bde: Bayesian Deep Ensembles for scikit-learn
 ====================================================
 
 ![tests](https://github.com/scikit-learn-contrib/bde/actions/workflows/python-app.yml/badge.svg)
@@ -13,6 +13,14 @@ both scikit-learn and JAX. It exposes estimators that plug into scikit-learn
 pipelines while leveraging JAX for accelerator-backed training, sampling, and
 uncertainty estimation.
 
+In particular, **bde** implements **Microcanonical Langevin Ensembles (MILE)** as
+introduced in [*Microcanonical Langevin Ensembles: Advancing the Sampling of Bayesian Neural Networks* (ICLR 2025)](https://arxiv.org/abs/2502.06335).
+A conceptual overview of MILE is shown below (in general implementation details of this package are not exactly matching this diagram):
+
+<div style="width: 60%; margin: auto;">
+    <img src="doc/_static/img/flowchart.png" alt="MILE Overview" style="width: 100%;">
+</div>
+
 Installation
 ------------
 
@@ -42,10 +50,11 @@ Example Usage
 
 Minimal runnable scripts live in `examples/`, and the snippets below highlight the
 most common regression and classification workflows. When running outside those
-scripts, remember to set the XLA device count so JAX allocates enough host devices:
-NOTE MENTNION THAT HTHE EXPORT NEEDS TO BE DONE BEFORE JAX
+scripts, remember to set the XLA device count so JAX allocates enough host devices (
+this needs to be done before importing JAX):
+
 ```
-export XLA_FLAGS="--xla_force_host_platform_device_count=8" 
+export XLA_FLAGS="--xla_force_host_platform_device_count=8"
 ```
 
 Adjust the value to match the number of CPU (or GPU) devices you plan to use.
@@ -67,16 +76,13 @@ from bde.loss import GaussianNLL
 
 
 data = fetch_openml(name="airfoil_self_noise", as_frame=True)
-X = data.data.values  # shape (1503, 5)
-y = data.target.values.reshape(-1, 1)  # shape (1503, 1)
 
+X = data.data.values
+y = data.target.values.reshape(-1, 1)
 X_train, X_test, y_train, y_test = train_test_split(
-   X,
-   y,
-   test_size=0.2,
-   random_state=0,
+    X, y, test_size=0.2, random_state=42
 )
-# Normalize data
+
 Xmu, Xstd = jnp.mean(X_train, 0), jnp.std(X_train, 0) + 1e-8
 Ymu, Ystd = jnp.mean(y_train, 0), jnp.std(y_train, 0) + 1e-8
 
@@ -87,26 +93,29 @@ yte = (y_test - Ymu) / Ystd
 
 regressor = BdeRegressor(
     hidden_layers=[16, 16],
-    n_members=20,
+    n_members=8,
     seed=0,
     loss=GaussianNLL(),
     epochs=200,
     lr=1e-3,
-    warmup_steps=500,
-    n_samples=100,
-    n_thinning=1,
+    warmup_steps=5000, # 50k in the original paper
+    n_samples=2000, # 10k in the original paper
+    n_thinning=2,
     patience=10,
 )
 
 regressor.fit(x=Xtr, y=ytr)
 
-mean, std = regressor.predict(jnp.array(X_test), mean_and_std=True)
-mu, intervals = regressor.predict(Xte, credible_intervals=[0.9, 0.95])
-raw = regressor.predict(Xte, raw=True)
-print("RSME: ", root_mean_squared_error(y_true=yte, y_pred=mean))
-score = regressor.score(Xtr, ytr)
-print(f"the sklearn score is {score}")
+means, sigmas = regressor.predict(Xte, mean_and_std=True)
+
+print("RSME: ", root_mean_squared_error(y_true=yte, y_pred=means))
 
+mean, intervals = regressor.predict(Xte, credible_intervals=[0.1, 0.9])
+
+lower = intervals[0]
+upper = intervals[1]
+coverage = jnp.mean((yte.ravel() >= lower) & (yte.ravel() <= upper))
+print(f"Coverage of the 80% credible interval: {coverage * 100:.2f}%")
 
 ```
 
@@ -125,32 +134,32 @@ from bde.loss import CategoricalCrossEntropy
 
 iris = load_iris()
 X = iris.data.astype("float32")
-y = iris.target.astype("int32").ravel()
+y = iris.target.astype("int32").ravel()  # 0, 1, 2
 X_train, X_test, y_train, y_test = train_test_split(
-    X, y, test_size=0.2, random_state=42)
+    X, y, test_size=0.2, random_state=42
+)
+
 classifier = BdeClassifier(
-    n_members=2,
+    n_members=4,
     hidden_layers=[16, 16],
     seed=0,
     loss=CategoricalCrossEntropy(),
     activation="relu",
-    epochs=4,
+    epochs=100,
     lr=1e-3,
-    warmup_steps=50,
-    n_samples=2,
+    warmup_steps=400, # very few steps required for this simple dataset
+    n_samples=100,
     n_thinning=1,
-    patience=2
-    )
+    patience=10,
+)
+
 classifier.fit(x=X_train, y=y_train)
+
 preds = classifier.predict(X_test)
 probs = classifier.predict_proba(X_test)
-score = classifier.score(X_train, y_train)
-raw = classifier.predict(X_test, raw=True)
-print("Predicted class probabilities:\n", probs)
-print("Predicted class labels:\n", preds)
-print("True labels:\n", y_test)
-print(f"the sklearn score is {score}")
-print(f"The shape of the raw predictions are {raw.shape}")
+
+accuracy = jnp.mean(preds == y_test)
+print(f"Test accuracy: {accuracy * 100:.2f}%")
 ```
 
 Workflow
@@ -210,7 +219,3 @@ flowchart TD
     Cache --> EvalCall --> MakePred --> Predictor --> Outputs
     Posterior --> Predictor
 ```
-
-
-Mathematical Background
------------------------

.idea/dictionaries/project.xml

@@ -588,6 +588,8 @@ def predict(
         x: ArrayLike,
         mean_and_std: bool = False,
         credible_intervals: list[float] | None = None,
+        # Docstring necessary to explain this parameter which
+        # actually lists quantiles not the intervals
         raw: bool = False,
     ):
         out = self._evaluate(

.idea/inspectionProfiles/profiles_settings.xml

@@ -118,6 +118,13 @@ def _predict_regression(
         if mean_and_std:
             out["std"] = std_total
         if credible_intervals:
+            # this diregards the sigmas and uses only the mus which is an ok strategy
+            # (the cheapest one, could be a default)
+            # if you want to incorporate the sigmas you could sample for example n
+            # (n=1 for more than 1000 samples and n=10 for less than 1000 samples)
+            # predictions from each mu,sigma (gaussian) pair and then compute the
+            # quantiles over all these sampled predictions - likely a better
+            # strategy but more expensive.
             qs = jnp.quantile(mu, q=jnp.array(credible_intervals), axis=(0, 1))
             out["credible_intervals"] = qs
         if raw: