Expand the ensemble instructions

README.md

@@ -27,7 +27,7 @@ simply several draws — which here is just the batch dimension.
 | [WeatherNext2](https://huggingface.co/kashif/weathernext2) | 0.25° (721×1440) | 40,962 | 183.8M | yes |
 | [WeatherNextCyclones](https://huggingface.co/kashif/weathernext-cyclones) | 0.25° (721×1440) | 40,962 | 183.8M | no |
 
-These weights correspond to `WeatherNext2_<2025_model1`, trained on data through 2024 and fine-tuned for
+These weights correspond to `WeatherNext2_<2025_model{1..4}`, trained on data through 2024 and fine-tuned for
 initialization from operational ECMWF HRES analysis. All four independently trained members (`model1`–`model4`) are included; see [Ensembles](#ensembles).
 
 ## Usage
@@ -36,7 +36,7 @@ initialization from operational ECMWF HRES analysis. All four independently trai
 pip install transformers torch scipy
 ```
 
-The model works in a normalized space; [`WeatherNext2Processor`] owns everything physical — the per-variable
+The model works in a normalized space; `WeatherNext2Processor` owns everything physical — the per-variable
 normalization statistics, the calendar-derived forcings, and the residual connection back to an atmospheric state.
 
 ```python
@@ -78,33 +78,41 @@ for step in range(20):  # 5 days
 
 ## Ensembles
 
-There are two independent ensembles here, and useful forecasts combine both.
+There are two independent ensembles here, and the operational product combines both.
 
-**1. Noise ensemble (within one checkpoint).** This is the FGN mechanism: each member is one draw of the
-32-dimensional noise vector through the same weights. Members are the batch axis, and stay independent through an
-autoregressive rollout.
+### 1. Noise ensemble (within one checkpoint)
+
+This is the FGN mechanism: each member is one draw of the 32-dimensional noise vector through the *same* weights.
+Members ride on the batch axis and stay independent through an autoregressive rollout.
 
 ```python
 members = 8
 inputs = processor(state, seconds_since_epoch=valid_time)
-inputs = {key: value.repeat(members, *([1] * (value.ndim - 1))) for key, value in inputs.items()}
+batched = {key: value.repeat(members, *([1] * (value.ndim - 1))) for key, value in inputs.items()}
 with torch.no_grad():
-    outputs = model(**inputs, generator=torch.Generator().manual_seed(0))
+    outputs = model(**batched, generator=torch.Generator().manual_seed(0))
+# outputs.prediction is (members, channels, lat, lon)
 ```
 
-At 0.25° one member needs roughly 50 GB, so batching all members at once usually will not fit on a single device.
-Loop over draws instead (or shard them across devices); the result is identical.
+At 0.25° a single member needs roughly 50 GB, so batching all of them at once will usually not fit on one device.
+Looping over draws gives identical results with a constant memory footprint:
 
 ```python
+single = processor(state, seconds_since_epoch=valid_time)   # batch of 1
 predictions = []
 for member in range(members):
     noise = torch.randn(1, model.config.noise_channels, generator=torch.Generator().manual_seed(member))
     with torch.no_grad():
-        predictions.append(model(**inputs, noise=noise).prediction)
+        predictions.append(model(**single, noise=noise).prediction)
 ```
 
-**2. Multi-model ensemble (across checkpoints).** The released product is four independently trained networks. Member 1
-is at the repository root; all four are also available as subfolders, so you can loop uniformly.
+Seeding per member (rather than drawing from one stream) means the first N members are reproducible regardless of how
+many you end up running — the same property the original implementation gets from `jax.random.fold_in`.
+
+### 2. Multi-model ensemble (across checkpoints)
+
+The released product is four independently trained networks. Member 1 is at the repository root; all four are also
+available as subfolders, so you can loop uniformly.
 
 ```python
 REPO = "kashif/weathernext2"
@@ -113,15 +121,44 @@ def load_member(member: int, revision: str = "main"):
     return WeatherNext2ForWeatherForecasting.from_pretrained(
         REPO, subfolder=f"model{member}", revision=revision
     ).eval()
-
-models = [load_member(i) for i in range(1, 5)]   # or a subset, they are ~700 MB each
 ```
 
-`subfolder` works the same way for [`WeatherNext2Processor`] and `AutoConfig`, since each subfolder holds its own
-`config.json` and `preprocessor_config.json`. The processors are identical across members, so loading one is enough.
+`subfolder` composes with `revision`, and works the same way for `WeatherNext2Processor` and `AutoConfig` — each
+subfolder carries its own `config.json` and `preprocessor_config.json`. The processors are identical across members,
+so loading one is enough.
+
+### Putting them together
+
+The full ensemble is `num_models × num_noise_draws` trajectories. Loading one member at a time keeps peak memory at
+roughly one model:
+
+```python
+import numpy as np
+import torch
+
+processor = WeatherNext2Processor.from_pretrained(REPO)
+inputs = processor(state, seconds_since_epoch=valid_time)
+
+forecasts = []
+for member in range(1, 5):
+    model = load_member(member)
+    for draw in range(4):
+        noise = torch.randn(
+            1, model.config.noise_channels,
+            generator=torch.Generator().manual_seed(1000 * member + draw),
+        )
+        with torch.no_grad():
+            prediction = model(**inputs, noise=noise).prediction
+        forecasts.append(processor.postprocess(prediction, state)["2m_temperature"])
+    del model  # free before loading the next member
+
+stack = np.concatenate(forecasts, axis=0)   # (16, lat, lon)
+ensemble_mean = stack.mean(axis=0)
+ensemble_spread = stack.std(axis=0)
+```
 
-Combining the two gives the full ensemble: `num_models x num_noise_draws` trajectories, from which you take the
-ensemble mean, spread, or whatever quantiles the downstream product needs.
+For multi-step forecasts each trajectory carries its own state, so keep one `state` per member and advance them
+separately (or keep members on the batch axis, which `advance_state` handles for you).
 
 ## Model details