Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/secondlaw-ai/skyrim
π π€ AI weather models united
https://github.com/secondlaw-ai/skyrim
Last synced: 6 days ago
JSON representation
π π€ AI weather models united
- Host: GitHub
- URL: https://github.com/secondlaw-ai/skyrim
- Owner: secondlaw-ai
- License: apache-2.0
- Created: 2024-04-05T15:04:52.000Z (7 months ago)
- Default Branch: master
- Last Pushed: 2024-10-18T01:29:02.000Z (22 days ago)
- Last Synced: 2024-10-21T04:06:38.944Z (18 days ago)
- Language: Jupyter Notebook
- Homepage:
- Size: 28.4 MB
- Stars: 153
- Watchers: 7
- Forks: 7
- Open Issues: 5
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
Awesome Lists containing this project
- Awesome-LWMs - Skyrim
- open-sustainable-technology - Skyrim - Allows you to run any large weather model with a consumer grade GPU. (Atmosphere / Meteorological Observation and Forecast)
README
π₯ Run state-of-the-art large weather models in less than 2 minutes.
πͺοΈ Ensemble and fine-tune (soon) to push the limits on forecasting.
π Simulate extreme weather events!
[![Static Badge](https://img.shields.io/badge/Run-On_Modal-green)](https://github.com/secondlaw-ai/skyrim?tab=readme-ov-file#forecasting-using-modal-recommended)
[![PyPI - Version](https://img.shields.io/pypi/v/skyrim)](https://pypi.org/project/Skyrim/)
[![Twitter Follow](https://img.shields.io/twitter/follow/secondlaw_ai)](https://twitter.com/secondlaw_ai)
![GitHub Repo stars](https://img.shields.io/github/stars/secondlaw-ai/skyrim)
[![GitHub License](https://img.shields.io/github/license/secondlaw-ai/skyrim)](https://github.com/secondlaw-ai/skyrim/blob/master/LICENSE)# Getting Started
Skyrim allows you to run any large weather model with a consumer grade GPU.
Until very recently, weather forecasts were run in 100K+ CPU HPC clusters, solving massive numerical weather models (NWP). Within last 2 years, open-source foundation models trained on weather simulation datasets surpassed the skill level of these numerical models.
Our goal is to make these models accessible by providing a well maintained infrastructure.
## Installation
Clone the repo, set an env (either conda or venv) and then run
```bash
git clone https://github.com/your-repo/skyrim.git
cd skyrim
pip install .
```Depending on your use-case (i.e. AWS storage needs or CDS initial conditions), you may need to fill in a `.env` by `cp .env.example .env`.
## Running Your First Forecast
Skyrim currently supports either running on on [modal](#forecasting-using-modal), on a container βfor instance [vast.ai](#vastai-setup) or [bare metal](#bare-metal)(you will need an NVIDIA GPU with at least 24GB and installation can be long).
Modal is the fastest option, it will run forecasts "serverless" so you don't have to worry about the infrastructure.
### Forecasting using Modal (Recommended):
You will need a [modal](https://modal.com/) key. Run `modal setup` and set it up (<1 min).
Modal comes with $30 free credits and a single forecast costs about 2 cents as of May 2024.
Once you are all good to go, then run:
```bash
modal run skyrim/modal/forecast.py
```This by default uses `pangu` model to forecast for the next 6 hours, starting from yesterday. It gets initial conditions from [NOAA GFS](https://en.wikipedia.org/wiki/Global_Forecast_System) and writes the forecast to a modal volume. You can choose different dates and weather models as shown in [here](#run-forecasts-with-different-models-initial-conditions-dates).
After you have your forecast, you can explore it by running a notebook (without GPU, so cheap) in modal:
```bash
modal run skyrim/modal/forecast.py::run_analysis
```This will output a jupyter notebook link that you can follow and access the forecast. For instance, to read the forecast you can run from the notebook the following:
```
import xarray as xr
forecast = xr.open_dataset('/skyrim/outputs/[forecast_id]/[filename], engine='scipy')
```Once you are done, best is to delete the volume as a daily forecast is about 2GB:
```bash
modal volume rm forecasts /[forecast_id] -r
```If you don't want to use modal volume, and want to aggregate results in a bucket (currently only s3), you just have to run:
```bash
modal run skyrim/modal/forecast.py --output_dir s3://skyrim-dev
```where `skyrim-dev` is the bucket that you want to aggregate the forecasts. By default, `zarr` format is used to store in AWS/GCP so you can read and move only the parts of the forecasts that you need.
See [examples](#examples) section for more.βοΈ
### Forecasting with your own GPUs:
If you are running on your own GPUs, installed either via [bare metal](#bare-metal) or via containers such as [vast.ai](#vast-ai-setup) then you can directly get forecasts as such:
```python
from skyrim.core import Skyrimmodel = Skyrim("pangu")
final_pred, pred_paths = model.predict(
date="20240507", # format: YYYYMMDD, start date of the forecast
time="0000", # format: HHMM, start time of the forecast
lead_time=24 * 7, # in hours, next week
save=True,
)
```To visualise the forecast:
```python
from skyrim.libs.plotting import visualize_rollout
visualize_rollout(output_paths=pred_paths, channels=["u10m", "v10m"], output_dir=".")
```
or you can still use the command line:
```bash
forecast -m graphcast --lead_time 24 --initial_conditions cds --date 20240330`
```See [examples](#examples) section for more.βοΈ
#### vast.ai setup
1. Find a machine you like RTX3090 or above with at least 24GB memory. Make sure you have good bandwith (+500MB/s).
2. Select the instance template from [here](https://cloud.vast.ai/?ref_id=128656&template_id=1883215a8487ec6ea9ad68a7cdb38c5e).
3. Then clone the repo and `pip install . && pip install -r requirements.txt`#### Bare metal
1. You will need a NVIDIA GPU with at least 24GB. We are working on quantization as well so that in the future it would be possible to run simulations with much less compute. Have an environment set with Python == 3.10, Pytorch => 2.2.2 and CUDA +12.x. Or if easier start with the docker image: `nvcr.io/nvidia/pytorch:24.01-py3`.
2. Install conda ([miniconda](https://docs.anaconda.com/free/miniconda/miniconda-install/) for instance). Then run in that environment:```bash
conda create -y -n skyenv python=3.10
conda activate skyenv
conda install eccodes python-eccodes -c conda-forge
pip install . && pip install -r requirements.txt
```## Examples
For each run, you will first pull the initial conditions of your interest (most recent one by default), then the model will run for the desired time step. Initial conditions are pulled from GFS, ECMWF IFS (Operational) or CDS (ERA5 Reanalysis Dataset).
If you are using CDS initial conditions, then you will need a [CDS](https://cds.climate.copernicus.eu/user/login?destination=%2Fcdsapp%23!%2Fdataset%2Freanalysis-era5-single-levels) API key in your `.env` β`cp .env.example` and paste.
All examples can be run using `forecast` or `modal run skyrim/modal/forecast.py`. You just have to make snake case options kebab-case -i.e. `model_name` to `model-name`.
### Example 1: Pick models, initial conditions, lead times
Forecast using `graphcast` model, with ECMWF IFS initial conditions, starting from 2024-04-30T00:00:00 and with a lead time of a week (forecast for the next week, i.e. 168 hours):
```bash
forecast --model_name graphcast --initial_conditions ifs --date 20240403 -output_dir s3://skyrim-dev --lead_time 168
```or in modal:
```bash
modal run skyrim/modal/forecast.py --model-name graphcast --initial-conditions ifs --date 20240403 --output-dir s3://skyrim-dev --lead-time 168
```### Example 2: Store in AWS and then read only what you need
Say you re interested in wind at 37.0344Β° N, 27.4305 E to see if we can kite tomorrow. If we need wind speed, we need to pull wind vectors at about surface level, these are u10m and v10m [components](http://colaweb.gmu.edu/dev/clim301/lectures/wind/wind-uv) of wind. Here is how you go about it:
```bash
modal run skyrim/modal/forecast.py --output-dir s3://[your_bucket]/[optional_path] --lead-time 24
```Then you can read the forecast as below:
```python
import xarray as xr
import pandas as pd
zarr_store_path = "s3://[your_bucket]/[forecast_id]"
forecast = xr.open_dataset(zarr_store_path, engine='zarr') # reads the metadata
df = forecast.sel(lat=37.0344, lon=27.4305, channel=['u10m', 'v10m']).to_pandas()
```Normally each day is about 2GB but using zarr_store you will only fetch what you need.βοΈ
### Example 3: Get predictions in Python
Assuming you have a local gpu set up ready to roll:
```python
from skyrim.core import Skyrimmodel = Skyrim("pangu")
final_pred, pred_paths = model.predict(
date="20240501", # format: YYYYMMDD, start date of the forecast
time="0000", # format: HHMM, start time of the forecast
lead_time=12, # in hours
save=True,
)
akyaka_coords = {"lat": 37.0557, "lon": 28.3242}
wind_speed = final_pred.wind_speed(**akyaka_coords) * 1.94384 # m/s to knots
print(f"Wind speed at Akyaka: {wind_speed:.2f} knots")```
## Supported initial conditions and caveats
1. NOAA GFS
2. ECMWF IFS
3. ERA5 Re-analysis Dataset## Large weather models supported
Currently supported models are:
- [x] [Graphcast](https://arxiv.org/abs/2212.12794)
- [x] [Pangu](https://arxiv.org/abs/2211.02556)
- [x] [Fourcastnet](https://arxiv.org/abs/2202.11214) (v1 & v2)
- [x] [DLWP](https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2020MS002109)
- [x] [(NWP) ECMWF IFS (HRES)](https://www.ecmwf.int/en/forecasts/documentation-and-support/changes-ecmwf-model) -- [notebook](https://github.com/secondlaw-ai/skyrim/blob/master/notebooks/02_ifs_model.ipynb)
- [x] [(NWP) NOAA GFS](https://www.ncei.noaa.gov/products/weather-climate-models/global-forecast) -- [notebook](https://github.com/secondlaw-ai/skyrim/blob/master/notebooks/03_gfs_model.ipynb)
- [x] [(NWP) ECMWF ENS](https://confluence.ecmwf.int/display/FUG/Section+2.1.2.1+ENS+-+Ensemble+Forecasts)
- [ ] [(NWP) ICON](https://www.dwd.de/EN/ourservices/nwp_forecast_data/nwp_forecast_data.html)
- [ ] [Aurora](https://github.com/microsoft/aurora)
- [X] [FuXi](https://www.nature.com/articles/s41612-023-00512-1)
- [X] [FengWu](https://arxiv.org/pdf/2304.02948)
- [ ] [Nano MetNet](https://arxiv.org/abs/2306.06079)### License
For detailed information regarding licensing, please refer to the license details provided on each model's main homepage, which we link to from each of the corresponding components within our repository.
- **Pangu Weather** : [Original](https://github.com/198808xc/Pangu-Weather), [ECMWF](https://github.com/ecmwf-lab/ai-models-panguweather), [NVIDIA](https://github.com/NVIDIA/earth2mip)
- **FourcastNet** : [Original](https://github.com/NVlabs/FourCastNet), [ECMWF](https://github.com/ecmwf-lab/ai-models-fourcastnetv2),[NVIDIA](https://github.com/NVIDIA/earth2mip)
- **Graphcast** : [Original](https://github.com/google-deepmind/graphcast), [ECMWF](https://github.com/ecmwf-lab/ai-models-graphcast), [NVIDIA](https://github.com/NVIDIA/earth2mip)
- **Fuxi**: [Original](https://github.com/tpys/FuXi)
- **FengWu** [Original](https://github.com/OpenEarthLab/FengWu)## Roadmap
- [x] ensemble prediction
- [x] interface to fetch real-time NWP-based predictions, e.g. via ECMWF API.
- [ ] global model performance comparison across various regions and parameters.
- [ ] finetuning api that trains a downstream model on top of features coming from a global/foundation model, that is optimized wrt to a specific criteria and region
- [ ] model quantization and its effect on model efficiency and accuracy.This README will be updated regularly to reflect the progress and integration of new models or features into the library. It serves as a guide for internal development efforts and aids in prioritizing tasks and milestones.
## Development
All in [here](./CONTRIBUTING.md) βοΈ
## Acknowledgements
Skyrim is built on top of NVIDIA's [earth2mip](https://github.com/NVIDIA/earth2mip), [earth2studio](https://github.com/NVIDIA/earth2studio), and ECMWF's [ai-models](https://github.com/ecmwf-lab/ai-models). Definitely check them out!
## Other Useful Resources
- [π Awesome Large Weather Models (LWMs) | AI for Earth (AI4Earth) | AI for Science (AI4Science)](http://github.com/jaychempan/Awesome-LWMs)
- [Climate Data Store](https://cds.climate.copernicus.eu/)
- [Open Climate Fix](https://github.com/openclimatefix)
- [Herbie](https://github.com/blaylockbk/Herbie)