docs: prepare v0.3.0 release with documentation site and uv migration

- Add documentation site with GitHub Pages deployment
  - Core concepts guide (batch-once-then-mask pattern)
  - API reference pages for models, acquisition, data, metrics
  - Updated getting started guide with correct examples

- Migrate CI/CD to uv package manager
  - ci.yml: use astral-sh/setup-uv with caching
  - release.yml: use uv for builds
  - New docs.yml workflow for GitHub Pages deployment

- Update README
  - Emphasize Flax NNX as the modern API
  - Simplify with links to documentation
  - uv-first installation instructions

- Update CLAUDE.md with new modules and APIs
- Update CHANGELOG.md for v0.3.0

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

Author: HFooladi

.github/workflows/ci.yml

@@ -16,19 +16,19 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
-        python-version: ${{ matrix.python-version }}
+        enable-cache: true
+
+    - name: Set up Python ${{ matrix.python-version }}
+      run: uv python install ${{ matrix.python-version }}
 
     - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install .[dev]
+      run: uv pip install -e .[dev] --system
 
     - name: Run tests
-      run: |
-        pytest tests/
+      run: pytest tests/
 
     - name: Check code formatting and linting
       run: |
@@ -41,18 +41,21 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
-    - name: Set up Python
-      uses: actions/setup-python@v5
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
-        python-version: "3.11"
+        enable-cache: true
+
+    - name: Set up Python
+      run: uv python install 3.11
 
     - name: Build package
       run: |
-        pip install build
+        uv pip install build --system
         python -m build
 
     - name: Store build artifacts
       uses: actions/upload-artifact@v4
       with:
         name: dist
-        path: dist/
+        path: dist/

.github/workflows/docs.yml

@@ -0,0 +1,39 @@
+name: Deploy Docs
+
+on:
+  push:
+    branches: [main]
+  release:
+    types: [published]
+  workflow_dispatch:
+
+permissions:
+  contents: write
+  pages: write
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+
+    - name: Set up Python
+      run: uv python install 3.11
+
+    - name: Install dependencies
+      run: uv pip install -e .[docs] --system
+
+    - name: Configure git
+      run: |
+        git config user.name "github-actions[bot]"
+        git config user.email "github-actions[bot]@users.noreply.github.com"
+
+    - name: Deploy docs
+      run: mkdocs gh-deploy --force

.github/workflows/release.yml

@@ -10,20 +10,21 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
-    - name: Set up Python
-      uses: actions/setup-python@v5
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
-        python-version: "3.11"
+        enable-cache: true
 
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build twine
+    - name: Set up Python
+      run: uv python install 3.11
+
+    - name: Install build tools
+      run: uv pip install build twine --system
 
     - name: Build and publish
       env:
         TWINE_USERNAME: __token__
         TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
       run: |
         python -m build
-        twine upload dist/*
+        twine upload dist/*

CHANGELOG.md

@@ -17,6 +17,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ---
 
+## [0.3.0] - 2025-01-26
+
+### Added
+- **Documentation Site** deployed to GitHub Pages
+  - Core concepts guide explaining batch-once-then-mask pattern
+  - Full API reference with mkdocstrings
+  - Installation and quick start guides
+
+- **Calibration Metrics** (`molax/metrics/`)
+  - `expected_calibration_error` for measuring uncertainty quality
+  - `compute_calibration_curve` for calibration analysis
+  - `negative_log_likelihood` metric
+  - `calibration_report` for comprehensive analysis
+  - Visualization tools: `plot_calibration_curve`, `plot_reliability_diagram`
+
+### Changed
+- **Migrated to uv** for package management
+  - CI/CD workflows now use `astral-sh/setup-uv@v5`
+  - Installation instructions updated to uv-first
+  - Faster dependency resolution and caching
+
+- Simplified README with links to documentation site
+
+---
+
 ## [0.2.0] - 2025-01-26
 
 ### Added
@@ -53,6 +78,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ---
 
-[Unreleased]: https://github.com/HFooladi/molax/compare/v0.2.0...HEAD
+[Unreleased]: https://github.com/HFooladi/molax/compare/v0.3.0...HEAD
+[0.3.0]: https://github.com/HFooladi/molax/compare/v0.2.0...v0.3.0
 [0.2.0]: https://github.com/HFooladi/molax/compare/v0.1.0...v0.2.0
 [0.1.0]: https://github.com/HFooladi/molax/releases/tag/v0.1.0

CLAUDE.md

@@ -9,8 +9,8 @@ Molax is a high-performance JAX framework for molecular active learning. It uses
 ## Quick Commands
 
 ```bash
-# Setup
-source env.sh  # or: pip install -e .[dev]
+# Setup (using uv)
+source env.sh
 
 # Test
 pytest tests/
@@ -19,9 +19,15 @@ pytest tests/
 ruff format .
 ruff check .
 
+# Build docs locally
+uv pip install -e .[docs]
+mkdocs serve
+
 # Run examples
 python examples/simple_active_learning.py
 python examples/active_learning_benchmark.py
+python examples/ensemble_demo.py
+python examples/evidential_demo.py
 ```
 
 ## Architecture
@@ -59,17 +65,22 @@ SMILES string
 jraph.GraphsTuple (single molecule)
     ↓ jraph.batch()
 jraph.GraphsTuple (batched - all molecules as one graph)
-    ↓ UncertaintyGCN
+    ↓ UncertaintyGCN / DeepEnsemble / EvidentialGCN
 (mean, variance) predictions
 ```
 
 ### Key Files
 
 | File | Purpose |
 |------|---------|
-| `molax/models/gcn.py` | `GCNConfig`, `UncertaintyGCN`, `train_step`, `eval_step` |
+| `molax/models/gcn.py` | `GCNConfig`, `UncertaintyGCN`, `MolecularGCN`, `train_step`, `eval_step` |
+| `molax/models/ensemble.py` | `EnsembleConfig`, `DeepEnsemble` for ensemble uncertainty |
+| `molax/models/evidential.py` | `EvidentialConfig`, `EvidentialGCN` for evidential uncertainty |
 | `molax/utils/data.py` | `MolecularDataset`, `smiles_to_jraph`, `batch_graphs` |
-| `molax/acquisition/uncertainty.py` | `uncertainty_sampling`, `diversity_sampling`, `combined_acquisition` |
+| `molax/acquisition/uncertainty.py` | `uncertainty_sampling`, `ensemble_uncertainty_sampling`, `evidential_uncertainty_sampling` |
+| `molax/acquisition/diversity.py` | `diversity_sampling` |
+| `molax/metrics/calibration.py` | `expected_calibration_error`, `calibration_report` |
+| `molax/metrics/visualization.py` | `plot_calibration_curve`, `plot_reliability_diagram` |
 
 ### Model API
 
@@ -88,6 +99,55 @@ model = UncertaintyGCN(config, rngs=nnx.Rngs(0))
 mean, variance = model(batched_graphs, training=True)
 ```
 
+### Ensemble API
+
+```python
+from molax.models.ensemble import EnsembleConfig, DeepEnsemble
+
+config = EnsembleConfig(
+    node_features=6,
+    hidden_features=[64, 64],
+    out_features=1,
+    n_members=5,
+)
+ensemble = DeepEnsemble(config, rngs=nnx.Rngs(0))
+
+# Returns separate epistemic and aleatoric uncertainty
+mean, epistemic_var, aleatoric_var = ensemble(batched_graphs, training=False)
+```
+
+### Evidential API
+
+```python
+from molax.models.evidential import EvidentialConfig, EvidentialGCN
+
+config = EvidentialConfig(
+    node_features=6,
+    hidden_features=[64, 64],
+    out_features=1,
+)
+model = EvidentialGCN(config, rngs=nnx.Rngs(0))
+
+# Single forward pass for both uncertainties
+mean, aleatoric_var, epistemic_var = model(batched_graphs, training=False)
+```
+
+### Calibration Metrics
+
+```python
+from molax.metrics import expected_calibration_error, calibration_report
+from molax.metrics.visualization import plot_calibration_curve
+
+# Compute ECE
+ece = expected_calibration_error(predictions, variances, targets)
+
+# Generate full report
+report = calibration_report(predictions, variances, targets)
+
+# Visualize
+fig = plot_calibration_curve(predictions, variances, targets)
+```
+
 ### Optimizer Pattern (Flax 0.11+)
 
 ```python
@@ -104,7 +164,10 @@ optimizer.update(model, grads)
 ```bash
 pytest tests/ -v                    # All tests
 pytest tests/test_gcn.py -v         # Model tests
+pytest tests/test_ensemble.py -v    # Ensemble tests
+pytest tests/test_evidential.py -v  # Evidential tests
 pytest tests/test_acquisition.py -v # Acquisition tests
+pytest tests/test_calibration.py -v # Calibration tests
 ```
 
 ## Dependencies