fix: update uncertainty_gcn_demo.py to use current jraph-based API

The demo was using an outdated API (UncertaintyGCNConfig, adjacency matrices).
Updated to use current API (GCNConfig, jraph graphs, smiles_to_jraph).

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: HFooladi

examples/uncertainty_gcn_demo.py

@@ -2,115 +2,123 @@
 Example demonstrating how to use UncertaintyGCN for molecular property prediction.
 
 This script shows how to:
-1. Create a simple molecule graph
+1. Load molecules and convert to jraph graphs
 2. Initialize and use an UncertaintyGCN model
 3. Interpret uncertainty in predictions
 """
 
-import jax
+import flax.nnx as nnx
 import jax.numpy as jnp
+import jraph
 import matplotlib.pyplot as plt
-import numpy as np
-from flax import nnx
-
-from molax.models.gcn import UncertaintyGCN, UncertaintyGCNConfig
-
-# Set random seed for reproducibility
-key = jax.random.PRNGKey(42)
-rngs = nnx.Rngs(0, params=1, dropout=2)
-
-# Create a simple molecule graph (6 atoms with 4 features per atom)
-num_atoms = 6
-in_features = 4
-
-# 1. Create random atom features
-key, subkey = jax.random.split(key)
-atom_features = jax.random.normal(subkey, (num_atoms, in_features))
-
-# 2. Create a ring-shaped molecular graph
-adjacency_matrix = np.zeros((num_atoms, num_atoms))
-for i in range(num_atoms):
-    # Connect each atom to its neighbors (creating a ring)
-    adjacency_matrix[i, (i + 1) % num_atoms] = 1.0
-    adjacency_matrix[i, (i - 1) % num_atoms] = 1.0
-adjacency_matrix = jnp.array(adjacency_matrix)
-
-print("atom_features.shape", atom_features.shape)
-print("adjacency_matrix.shape", adjacency_matrix.shape)
-
-# 3. Create and initialize the UncertaintyGCN model
-config = UncertaintyGCNConfig(
-    in_features=in_features,  # Number of input features per atom
-    hidden_features=[32, 16, 8],  # GCN layer sizes
-    out_features=1,  # Single property prediction
-    dropout_rate=0.1,  # Dropout for regularization
-    n_heads=2,
-    rngs=rngs,
-)
-
-model = UncertaintyGCN(config)
-
-
-# 4. Make a prediction with uncertainty
-mean, variance = model(atom_features, adjacency_matrix)
 
-print(f"Predicted property value: {mean[0]:.4f}")
-print(f"Prediction uncertainty (variance): {variance[0]:.4f}")
-print(
-    f"95% confidence interval: ({mean[0] - 1.96 * jnp.sqrt(variance[0]):.4f}, "
-    f"{mean[0] + 1.96 * jnp.sqrt(variance[0]):.4f})"
+from molax.models.gcn import GCNConfig, UncertaintyGCN
+from molax.utils.data import smiles_to_jraph
+
+# Create some sample molecules with varying complexity
+molecules = [
+    ("C", "methane"),
+    ("CC", "ethane"),
+    ("CCC", "propane"),
+    ("CCCC", "butane"),
+    ("c1ccccc1", "benzene"),
+    ("CCO", "ethanol"),
+    ("CC(=O)O", "acetic acid"),
+    ("c1ccc(O)cc1", "phenol"),
+]
+
+print("=" * 60)
+print("UncertaintyGCN Demo: Molecular Property Prediction")
+print("=" * 60)
+
+# Convert SMILES to jraph graphs
+graphs = [smiles_to_jraph(smi) for smi, _ in molecules]
+batched_graphs = jraph.batch(graphs)
+
+print(f"\nLoaded {len(molecules)} molecules")
+print(f"Node features: {graphs[0].nodes.shape[1]}")
+
+# Create and initialize the UncertaintyGCN model
+config = GCNConfig(
+    node_features=graphs[0].nodes.shape[1],
+    hidden_features=[32, 16],
+    out_features=1,
+    dropout_rate=0.1,
 )
-
-# 6. Demonstrate uncertainty behavior with modified inputs
-test_points = 50
-scaling_factors = jnp.linspace(0.1, 10.0, test_points)
-means = []
-uncertainties = []
-
-for scale in scaling_factors:
-    # Scale the input features to create increasingly out-of-distribution examples
-    scaled_features = atom_features * scale
-    mean, var = model(scaled_features, adjacency_matrix)
-    means.append(mean[0])
-    uncertainties.append(
-        jnp.sqrt(var[0])
-    )  # Use standard deviation for easier interpretation
-
-# Convert to arrays
-means = jnp.array(means)
-uncertainties = jnp.array(uncertainties)
-
-# 7. Visualize how uncertainty changes with input distribution shift
-plt.figure(figsize=(10, 6))
-plt.plot(scaling_factors, means, "b-", label="Prediction")
-plt.fill_between(
-    scaling_factors,
-    means - 1.96 * uncertainties,
-    means + 1.96 * uncertainties,
-    alpha=0.3,
-    color="b",
-    label="95% Confidence Interval",
-)
-plt.xlabel("Input Scaling Factor")
-plt.ylabel("Predicted Property")
-plt.title("Prediction with Uncertainty for Different Input Scales")
-plt.legend()
-plt.grid(True)
-plt.savefig("examples/uncertainty_demo.png")
+model = UncertaintyGCN(config, rngs=nnx.Rngs(42))
+
+print("\nModel configuration:")
+print(f"  Hidden layers: {config.hidden_features}")
+print(f"  Dropout rate: {config.dropout_rate}")
+
+# Make predictions with uncertainty
+print("\n" + "-" * 60)
+print("Predictions with Uncertainty")
+print("-" * 60)
+
+mean, variance = model(batched_graphs, training=False)
+mean = mean.squeeze(-1)
+variance = variance.squeeze(-1)
+
+print(f"{'Molecule':<15} {'Mean':>10} {'Std Dev':>10} {'95% CI'}")
+print("-" * 60)
+
+for i, (_, name) in enumerate(molecules):
+    m = float(mean[i])
+    std = float(jnp.sqrt(variance[i]))
+    ci_low = m - 1.96 * std
+    ci_high = m + 1.96 * std
+    print(f"{name:<15} {m:>10.4f} {std:>10.4f} [{ci_low:.2f}, {ci_high:.2f}]")
+
+# Demonstrate MC Dropout uncertainty
+print("\n" + "-" * 60)
+print("MC Dropout Uncertainty (10 samples)")
+print("-" * 60)
+
+mc_predictions = []
+for _ in range(10):
+    pred, _ = model(batched_graphs, training=True)  # Dropout active
+    mc_predictions.append(pred.squeeze(-1))
+
+mc_predictions = jnp.stack(mc_predictions)
+mc_mean = jnp.mean(mc_predictions, axis=0)
+mc_std = jnp.std(mc_predictions, axis=0)
+
+print(f"{'Molecule':<15} {'MC Mean':>10} {'MC Std':>10}")
+print("-" * 40)
+
+for i, (_, name) in enumerate(molecules):
+    print(f"{name:<15} {float(mc_mean[i]):>10.4f} {float(mc_std[i]):>10.4f}")
+
+# Visualize predictions
+print("\n" + "-" * 60)
+print("Creating visualization...")
+
+fig, ax = plt.subplots(figsize=(10, 6))
+
+x = range(len(molecules))
+names = [name for _, name in molecules]
+means = [float(mean[i]) for i in range(len(molecules))]
+stds = [float(jnp.sqrt(variance[i])) for i in range(len(molecules))]
+
+ax.bar(x, means, yerr=[1.96 * s for s in stds], capsize=5, alpha=0.7)
+ax.set_xticks(x)
+ax.set_xticklabels(names, rotation=45, ha="right")
+ax.set_ylabel("Predicted Value")
+ax.set_title("UncertaintyGCN Predictions with 95% Confidence Intervals")
+ax.grid(axis="y", alpha=0.3)
+
+plt.tight_layout()
+plt.savefig("examples/uncertainty_demo.png", dpi=150)
 plt.close()
 
-print("\nGenerating out-of-distribution examples:")
-for scale in [0.1, 1.0, 10.0]:
-    scaled_features = atom_features * scale
-    mean, var = model(scaled_features, adjacency_matrix)
-    std = jnp.sqrt(var[0])
-    print(
-        f"Scale {scale:.1f}: {mean[0]:.4f} ± {std:.4f}  "
-        f"(95% CI: {mean[0] - 1.96 * std:.4f} to {mean[0] + 1.96 * std:.4f})"
-    )
-
-print("\nDemo completed. Saved visualization to 'uncertainty_demo.png'")
-print(
-    "This demonstrates how uncertainty increases as inputs become more "
-    "out-of-distribution."
-)
+print("Saved visualization to 'examples/uncertainty_demo.png'")
+
+print("\n" + "=" * 60)
+print("Demo completed successfully!")
+print("=" * 60)
+print("\nKey takeaways:")
+print("- UncertaintyGCN outputs both mean prediction and variance")
+print("- Variance head predicts aleatoric (data) uncertainty")
+print("- MC Dropout provides epistemic (model) uncertainty")
+print("- 95% CI = mean ± 1.96 * std")