Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Change recommended default of 0.1 static vars #37

Merged
merged 1 commit into from
Sep 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions aurora/model/aurora.py
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,4 @@ def configure_activation_checkpointing(self):
patch_size=10,
encoder_depths=(6, 8, 8),
decoder_depths=(8, 8, 6),
# One particular static variable requires a different normalisation.
surf_stats={"z": (-3.270407e03, 6.540335e04)},
)
5 changes: 4 additions & 1 deletion docs/example_hres_0.1.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,10 @@
"source": [
"from aurora import AuroraHighRes, rollout\n",
"\n",
"model = AuroraHighRes()\n",
"model = AuroraHighRes(\n",
" # Use manually downloaded and regridded static variables.\n",
" surf_stats={\"z\": (-3.270407e03, 6.540335e04)},\n",
")\n",
"model.load_checkpoint(\"wbruinsma/aurora\", \"aurora-0.1-finetuned.ckpt\")\n",
"\n",
"model.eval()\n",
Expand Down
28 changes: 20 additions & 8 deletions docs/models.md
Original file line number Diff line number Diff line change
Expand Up @@ -157,23 +157,35 @@ For optimal performance, the model requires the following variables and pressure

### Static Variables

Aurora 0.1° Fine-Tuned requires

Due to differences between implementations of regridding methods, it is recommended to use
[the exact static variables which we used during training](https://huggingface.co/microsoft/aurora/blob/main/aurora-0.1-static.pickle).

It is also possible to use the
[static variables from IFS HRES analysis](https://rda.ucar.edu/datasets/ds113.1/) regridded
to 0.1° resolution.
Because of differences between implementations of regridding methods, the resulting static
variables might not be exactly equal to the ones we used during training.
For this reason we also uploaded
[the exact static variables which we used during training](https://huggingface.co/microsoft/aurora/blob/main/aurora-0.1-static.pickle).
To use these, you must remove an exception to the normalisation by instantiating
the model in the following way:
However, these static variables will not be exactly equal to the ones we used, which might impact
performance.
If you download the static variables yourself, you must adjust the normalisation statistics.
You can do that in the following way:

```python
from aurora import AuroraHighRes

model = AuroraHighRes(surf_stats=None) # Use static variables from HF repo.
model = AuroraHighRes(
# Use manually downloaded and regridded static variables.
surf_stats={"z": (-3.270407e03, 6.540335e04)},
)

model.load_checkpoint("microsoft/aurora", "aurora-0.1-finetuned.ckpt")
```

The specific values above should work reasonably.
<!-- Jupyter book complains that the below link doesn't work, but it does. -->
See [the API](api.rst#aurora.Aurora.__init__) for a description of `surf_vars`.
Generally, the first value in the tuple should be `min(static_z)`
and the second value `max(static_z) - min(static_z)`.

### Notes

[Like for Aurora 0.25° Fine-Tuned](lora-or-no-lora),
Expand Down
Loading