Rework README: crazyflow layout, badges, animated splat header

Author: amacati

README.md

@@ -1,64 +1,123 @@
+<div align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="assets/logo_emerges_dark.gif">
+    <img src="assets/logo_emerges_light.gif" alt="splax" width="560">
+  </picture>
+</div>
+
+--------------------------------------------------------------------------------
+
+<div align="center">
+
 # splax
 
-High-performance gaussian splatting in JAX, implemented in [NVIDIA Warp](https://github.com/NVIDIA/warp) —
-no system CUDA toolchain needed. Differentiable, vmap-batch-native (single
-kernel launch over the whole batch), faster than the CUDA reference it
-started from on every benchmarked config.
+**Differentiable 3D gaussian splatting for JAX, with rasterization kernels written in [NVIDIA Warp](https://github.com/NVIDIA/warp).**
 
-- `splax.inference.render` — grad-free rendering (residuals dropped, tight
-  intersection, batchable via `jax.vmap` up to 1000+ parallel views).
-- `splax.training.render` (= `splax.render`) — differentiable w.r.t. means,
-  scales, quats, colors, opacities.
-- `splax.io.write_ply` / `scripts/train_lego.py --out-ply` — fit splats to
-  posed images and export standard 3DGS `.ply`.
-- `jaxsplat/` — the vendored CUDA reference implementation (parity + grad
-  baseline), kept intact.
+[![Python](https://img.shields.io/badge/python-3.12-blue.svg)](https://www.python.org)
+[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
+[![Ruff](https://github.com/amacati/splax/actions/workflows/ruff.yml/badge.svg)](https://github.com/amacati/splax/actions/workflows/ruff.yml)
+[![Docs](https://github.com/amacati/splax/actions/workflows/docs.yml/badge.svg)](https://amacati.github.io/splax)
 
-## Setup
+</div>
 
-Requires an NVIDIA GPU.
+splax renders and trains 3D gaussian splats inside JAX. Batched rendering and training run through `jax.vmap`, `jax.grad`, and `jax.jit`, with no system CUDA toolchain needed.
 
-```sh
-uv sync                    # jax + warp + deps into .venv
-# optional: build the CUDA reference backend (hermetic toolchain, no sudo)
-scripts/setup_cuda.sh
-CUDACXX=$PWD/.cuda/bin/nvcc CUDAToolkit_ROOT=$PWD/.cuda \
-    uv pip install --no-build-isolation -e ./jaxsplat
+## Examples
+
+Render a `.ply` scene from one camera.
+
+```python
+import jax.numpy as jnp
+import splax
+
+means, scales, quats, colors, opacities = splax.load_ply("scene.ply")
+img = splax.inference.render(
+    means, scales, quats, colors, opacities,
+    viewmat=viewmat, background=jnp.ones(3),
+    img_shape=(H, W), f=(fx, fy), c=(W // 2, H // 2),
+    glob_scale=1.0, clip_thresh=0.01,
+)  # (H, W, 3)
+```
+
+Batch over a stack of camera poses with `jax.vmap`. One batched kernel launch, not a Python loop.
+
+```python
+import jax
+
+frames = jax.vmap(lambda vm: splax.inference.render(
+    means, scales, quats, colors, opacities,
+    viewmat=vm, background=jnp.ones(3), img_shape=(H, W),
+    f=(fx, fy), c=(W // 2, H // 2), glob_scale=1.0, clip_thresh=0.01,
+))(viewmats)  # (B, H, W, 3)
+```
+
+Take gradients through the differentiable renderer with `jax.grad`. `splax.render` is `splax.training.render` and differentiates w.r.t. means, scales, quats, colors, opacities.
+
+```python
+import jax
+
+def loss(means, scales, quats, colors, opacities):
+    img = splax.render(
+        means, scales, quats, colors, opacities,
+        viewmat=viewmat, background=jnp.ones(3), img_shape=(H, W),
+        f=(fx, fy), c=(W // 2, H // 2), glob_scale=1.0, clip_thresh=0.01,
+    )
+    return jnp.mean((img - target) ** 2)
+
+grads = jax.grad(loss, argnums=(0, 1, 2, 3, 4))(means, scales, quats, colors, opacities)
 ```
 
+## Documentation
+
+Full documentation lives at [amacati.github.io/splax](https://amacati.github.io/splax): installation, a quickstart, a user guide for rendering, training, batching, and IO, and the API reference.
+
+## Why
+
+Gaussian splatting lives mostly in PyTorch and hand-written CUDA. splax puts it inside JAX so splat rendering composes with `jax.vmap`, `jax.grad`, and `jax.jit` and drops into research pipelines that already run on JAX, without leaving the ecosystem for the render step.
+
+## Architecture
+
+The renderer is not pure JAX because the core of splatting does not map to XLA primitives. Rasterization is tile-binned with a data-dependent sort of gaussian-tile intersections, per-pixel early termination once transmittance saturates, and a memory-frugal backward that recomputes the blend instead of storing per-pixel state. splax implements the projection, rasterization, and their backward passes as Warp kernels and wires them into JAX through FFI custom calls under `jax.custom_vjp`. The kernels are batch-native: `jax.vmap` maps to a single batched launch (camera id folded into the sort key) rather than a sequential per-sample loop.
+
+## Relation to jaxsplat
+
+splax started from [jaxsplat](https://github.com/yklcs/jaxsplat) as the reference and parity baseline. The rasterizer was rewritten from CUDA to Warp to drop the system toolchain, then extended with the feature and performance work below.
+
 ## Improvements
 
-Ported from gsplat and the papers behind it (credit per item; details in `reports/`):
+Ported from [gsplat](https://github.com/nerfstudio-project/gsplat) and the papers behind it, which inspired most of the performance work (credit per item):
+
+- Native multi-camera batched rendering, one launch with the camera id folded into the sort key (gsplat)
+- Opacity-aware tight tile intersection (StopThePop, Speedy-Splat, gsplat #927)
+- Packed 32-bit sort keys with quantized depth
+- Persistent sort and bin scratch across frames (gsplat caching allocator design)
+- Cooperative shared-memory tile blending with block-vote early exit (3DGS, gsplat)
+- Fixed-budget MCMC training with static shapes, relocation plus covariance noise (gsplat MCMCStrategy, Kheradmand et al. 2024)
+- Opacity and scale regularizers (gsplat mcmc preset)
+- Progressive resolution fine-tuning (coarse-to-fine, 3DGS)
+- Per-parameter Adam learning-rate schedules (gsplat, 3DGS)
+- L1 plus D-SSIM photometric loss (3DGS, gsplat, via dm-pix)
+- Camera pose gradients via `render(..., diff_wrt=("viewmat",))` (gsplat projection backward)
+- Batch-native backward passes, `jax.vmap(jax.grad(render))` runs as one batched launch (gsplat)
+- Batched training steps with sqrt-batch learning-rate scaling (gsplat `batch_size` and `steps_scaler`)
+- Anti-aliased opacity compensation (Mip-Splatting, gsplat), depth regularization from COLMAP points (gsplat `depth_loss`), per-image exposure correction (gsplat appearance optimization), all opt-in
 
-- Native multi-camera batched rendering, single launch with camera id in the sort key — up to 5× per-frame at batch 8, ~60k renders/s at 128² batch 1024 (gsplat)
-- Opacity-aware tight tile intersection, −40–81% fewer intersections and −15–76% frame time on real trained scenes (lego 313k, drone 150k) (StopThePop §B.2; Speedy-Splat SNUGBOX/AccuTile, gsplat #927)
-- Persistent signature-keyed sort/bin scratch, kills per-frame realloc (gsplat caching-allocator design)
-- Cooperative shared-memory tile blending with block-vote early exit (3DGS; gsplat)
-- Verified negatives, kept documented: 64-thread/4-px CTA blend (1.5–3.7× slower) and warp-level backward reduction (1.9× slower on dense drone; wins random 1.05–1.80×) — both occupancy/barrier-limited under warp-lang, re-confirmed on clean lego/drone scenes (gsplat designs; reports/phase5_o4.md, phase6b.md, phase6g_lego_reverification.md)
-- Fixed-budget MCMC training in JAX (static shapes): dead-gaussian relocation + per-step covariance noise, +1.9 dB on the lego fit (gsplat MCMCStrategy; Kheradmand et al. 2024)
-- Opacity/scale regularizers from the MCMC paper — opacity_reg is load-bearing for a fixed budget with no pruning, +8.9 dB on the lego fit (gsplat mcmc preset)
-- Progressive 400²→800² resolution fine-tune, +5–6 dB by resolving fine detail at the eval resolution (coarse-to-fine, 3DGS)
-- Per-parameter Adam LR schedules, means exponentially decayed ~100× over training (gsplat/3DGS create_splats_with_optimizers)
-- L1 + 0.2·D-SSIM photometric loss mix, +0.7 dB on the lego fit (3DGS; gsplat, via dm-pix)
-- Camera-pose (viewmat) gradients via `render(..., diff_wrt=("viewmat",))` — dedicated projection-backward kernels sharing all vjp math; recovers a 5°/0.2 pose perturbation to <1e-3°/5e-7 on lego. The 12-float v_viewmat reduction uses a block `tile_sum`, 22–113× faster than plain atomics at 50k–742k gaussians (gsplat ProjectionEWA3DGSFused bwd, reports/phase6e.md)
-- Batch-native backward passes — `jax.vmap(jax.grad(render))` runs one batched backward launch (no per-sample loop) for every `diff_wrt`, up to 8.4× faster than the loop and ~2.1× for simultaneous multi-pose recovery (reports/phase8a_batched_grads.md)
-- Batched training steps + √batch LR scaling — B views/step amortize per-step fixed cost, −12% wall at PSNR parity on the mid-N (150k) drone fit; scale-dependent, so B=1 stays default (gsplat `batch_size`/`steps_scaler`, reports/phase8l_batched_training.md)
+## Installation
 
-## Benchmarks & reports
+Requires an NVIDIA GPU and a CUDA-enabled JAX (`jax[cuda12]`, pulled in as a dependency).
 
 ```sh
-.venv/bin/python scripts/benchmark.py <label> --backend {warp,cuda} [--mem]   # real scene: lego.ply by default
-.venv/bin/python scripts/benchmark_train.py <label>       # training steps
-.venv/bin/python scripts/plot_results.py <label> [<label2>]
-.venv/bin/python scripts/render_lego.py                   # correctness gate vs ground truth
+uv pip install "git+https://github.com/amacati/splax"
 ```
 
-Per-phase reports with figures live in `reports/` (phase1–4: setup, JAX FFI
-migration, vmap, Warp port; phase5_*: one report per ported gsplat
-optimization, including verified negatives; phase6*: training support).
+Developer setup with [pixi](https://pixi.sh/), which installs splax editable with the dev tooling:
+
+```sh
+git clone https://github.com/amacati/splax.git
+cd splax
+pixi shell
+```
 
 ## License
 
-MIT (see [LICENSE](LICENSE)).
-gsplat-derived portions under Apache-2.0 ([licenses/Apache-2.0.txt](licenses/Apache-2.0.txt)).
-Started from [jaxsplat](https://github.com/yklcs/jaxsplat), rewritten in Warp to drop the CUDA toolchain, with optimizations ported from [gsplat](https://github.com/nerfstudio-project/gsplat) and the papers behind it (3DGS, StopThePop, Speedy-Splat).
+MIT (see [LICENSE](LICENSE)). gsplat-derived portions are under Apache-2.0 ([licenses/Apache-2.0.txt](licenses/Apache-2.0.txt)).