{"id":19199812,"url":"https://github.com/rose-stl-lab/spherical-dyffusion","last_synced_at":"2025-09-14T03:15:50.709Z","repository":{"id":260411015,"uuid":"881219185","full_name":"Rose-STL-Lab/spherical-dyffusion","owner":"Rose-STL-Lab","description":"[NeurIPS 24] Probablistic Emulation of a Global Climate Model with Spherical DYffusion","archived":false,"fork":false,"pushed_at":"2025-03-07T21:31:57.000Z","size":345,"stargazers_count":26,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-04-10T11:38:41.759Z","etag":null,"topics":["ai-for-science","climate-modeling","diffusion","dyffusion","large-ensembles","neurips","neurips-2024","pytorch","pytorch-lightning"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2406.14798","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Rose-STL-Lab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-10-31T06:03:43.000Z","updated_at":"2025-03-09T07:23:40.000Z","dependencies_parsed_at":"2024-10-31T07:17:28.700Z","dependency_job_id":"9bf8f0d5-199f-472b-80a4-9891f7e7a92d","html_url":"https://github.com/Rose-STL-Lab/spherical-dyffusion","commit_stats":null,"previous_names":["rose-stl-lab/spherical-dyffusion"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/Rose-STL-Lab/spherical-dyffusion","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Rose-STL-Lab%2Fspherical-dyffusion","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Rose-STL-Lab%2Fspherical-dyffusion/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Rose-STL-Lab%2Fspherical-dyffusion/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Rose-STL-Lab%2Fspherical-dyffusion/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Rose-STL-Lab","download_url":"https://codeload.github.com/Rose-STL-Lab/spherical-dyffusion/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Rose-STL-Lab%2Fspherical-dyffusion/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":275055832,"owners_count":25397741,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-14T02:00:10.474Z","response_time":75,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["ai-for-science","climate-modeling","diffusion","dyffusion","large-ensembles","neurips","neurips-2024","pytorch","pytorch-lightning"],"created_at":"2024-11-09T12:29:04.759Z","updated_at":"2025-09-14T03:15:50.659Z","avatar_url":"https://github.com/Rose-STL-Lab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Probabilistic Emulation of a Global Climate Model with Spherical DYffusion (NeurIPS 2024, Spotlight)\n\n\u003ca href=\"https://pytorch.org/get-started/locally/\"\u003e\u003cimg alt=\"Python\" src=\"https://img.shields.io/badge/-Python 3.9+-blue?style=for-the-badge\u0026logo=python\u0026logoColor=white\"\u003e\u003c/a\u003e\n\u003ca href=\"https://pytorch.org/get-started/locally/\"\u003e\u003cimg alt=\"PyTorch\" src=\"https://img.shields.io/badge/-PyTorch -ee4c2c?style=for-the-badge\u0026logo=pytorch\u0026logoColor=white\"\u003e\u003c/a\u003e\n\u003ca href=\"https://pytorchlightning.ai/\"\u003e\u003cimg alt=\"Lightning\" src=\"https://img.shields.io/badge/-Lightning-792ee5?style=for-the-badge\u0026logo=pytorchlightning\u0026logoColor=white\"\u003e\u003c/a\u003e\n\u003ca href=\"https://hydra.cc/\"\u003e\u003cimg alt=\"Config: hydra\" src=\"https://img.shields.io/badge/config-hydra-89b8cd?style=for-the-badge\u0026labelColor=gray\"\u003e\u003c/a\u003e\n\u003ca href=\"https://github.com/Rose-STL-Lab/dyffusion/blob/main/LICENSE\"\u003e\u003cimg alt=\"License\" src=\"https://img.shields.io/github/license/huggingface/datasets.svg?color=blue\"\u003e\u003c/a\u003e\n\n\u003ch3\u003e ✨Official implementation of our \u003ca href=\"https://openreview.net/forum?id=Ib2iHIJRTh\"\u003eSpherical DYffusion\u003c/a\u003e paper✨ \u003c/h3\u003e\n\n[//]: # ([![Watch the video]\u0026#40;https://img.youtube.com/vi/Hac_xGsJ1qY/hqdefault.jpg\u0026#41;]\u0026#40;https://youtu.be/Hac_xGsJ1qY\u0026#41;)\n\n## | Environment Setup\n\nWe recommend installing in a virtual environment from PyPi or Conda. Then, run:\n\n    python3 -m pip install .[dev]\n    python3 -m pip install --no-deps nvidia-modulus@git+https://github.com/ai2cm/modulus.git@94f62e1ce2083640829ec12d80b00619c40a47f8\n\nAlternatively, use the provided [environment/install_dependencies.sh](environment/install_dependencies.sh) script.\n\nNote that for some compute setups you may want to install pytorch first for proper GPU support.\nFor more details about installing [PyTorch](https://pytorch.org/get-started/locally/), please refer to their official documentation.\n\n\n## | Dataset\n\nThe final\ntraining and validation data can be downloaded from Google Cloud Storage following the instructions\nof the ACE paper at https://zenodo.org/records/10791087. The data are licensed under Creative\nCommons Attribution 4.0 International.\n\nThe dataset statistics that you'll need to run the experiments are saved in the [data_statistics](data_statistics) directory.\n\n## | Checkpoints\n\nModel weights are available at [https://huggingface.co/salv47/spherical-dyffusion](https://huggingface.co/salv47/spherical-dyffusion/tree/main).\n\n## | Running experiments\n\n### Inference\n\nFirstly, download the validation data as instructed in the [Dataset](#dataset) section.\n\nSecondly, use the `run_inference.py` script with a corresponding configuration file.\nThe configurations files used for our paper can be found in the [src/configs/inference](src/configs/inference) directory.\nThat is, you can run inference with the following command:\n\n    python run_inference.py \u003cpath-to-inference-config\u003e.yaml\n\nThe available inference configurations are:\n- [ckpts_from_huggingface_debug.yaml](src/configs/inference/ckpts_from_huggingface_debug.yaml): Short inference meant for debugging with checkpoints downloaded from Hugging Face.\n- [ckpts_from_huggingface_10years.yaml](src/configs/inference/ckpts_from_huggingface_10years.yaml): 10-year-long inference with checkpoints downloaded from Hugging Face.\n\nTo use these configs, **you need to correctly specify the `dataset.data_path` parameter in the configuration file to point to the downloaded validation data.**\n\nYou may also need to specify the `overrides.datamodule_config.data_dir_stats` to point to the [data statistics](data_statistics) directory.\n### Training\n\nWe use [Hydra](https://hydra.cc/) for configuration management and [PyTorch Lightning](https://www.pytorchlightning.ai/) for training.\nWe recommend familiarizing yourself with these tools before running training experiments.\n\nThe basic configuration files used for our paper can be found in the [src/configs](src/configs) directory.\nTo use them, please first specify the correct data path's in the [src/configs/datamodule/fv3gfs_prescriptive_only.yaml](src/configs/datamodule/fv3gfs_prescriptive_only.yaml) config file.\nSpecifically, set ``datamodule.data_dir=\u003cpath-to-dataset\u003e`` and ``datamodule.data_dir_stats=\u003cpath-to-data_statistics\u003e`` appropriately.\n\nTo run training, you can use the `run.py` script with a corresponding experiment configuration file.\nFor example, to train Spherical DYffusion with the default configuration, first train the interpolator with:\n\n    python run.py experiment=fv3gfs_interpolation\n\nAfter that, train the forecaster with:\n\n    python run.py experiment=fv3gfs_dyffusion diffusion.interpolator_run_id=\u003cinterpolator_run_id\u003e\n\nwhere `\u003cinterpolator_run_id\u003e` is the Weights \u0026 Biases run ID of the interpolator experiment.\nTo train the basic UNet-based DYffusion, you would additionally append `model=unet` to both commands.\n\n\n### Tips \u0026 Tricks\n\n\u003cdetails\u003e\n    \u003csummary\u003eMemory Considerations and OOM Errors\u003c/summary\u003e\n\nTo control memory usage and avoid OOM errors, you can adjust the training batch size and evaluation batch size:\n\n**For training**, you can adjust the `datamodule.batch_size_per_gpu` parameter. \nNote that this will automatically adjust `trainer.accumulate_grad_batches` to keep the effective batch size (set by `datamodule.batch_size`) constant (so it need to be divisible by `datamodule.batch_size_per_gpu`).\n\n**For evaluation** or OOMs during validation, you can adjust the `datamodule.eval_batch_size` parameter. \nNote that the effective validation-time batch size is `datamodule.eval_batch_size * module.num_predictions`. Be mindful of that when choosing `eval_batch_size`. You can control how many ensemble members to run in memory \nat once with `module.num_predictions_in_memory`.\n\nBesides those main knobs, you may turn on mixed precision training with `trainer.precision=16` to reduce memory usage and\nmay also adjust the `datamodule.num_workers` parameter to control the number of data loading processes.\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eWandb Integration\u003c/summary\u003e\n\nWe use [Weights \u0026 Biases](https://wandb.ai/) for logging and checkpointing.\nPlease set your wandb username/entity with one of the following options:\n- Edit the [src/configs/local/default.yaml](src/configs/local/default.yaml) file (recommended, local for you only).\n- Edit the [src/configs/logger/wandb.yaml](src/configs/logger/wandb.yaml) file.\n- as a command line argument (e.g. `python run.py logger.wandb.entity=my_username`).\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eCheckpointing\u003c/summary\u003e\n\nBy default, checkpoints are saved locally in the `\u003cwork_dir\u003e/checkpoints` directory in the root of the repository, \nwhich you can control with the `work_dir=\u003cpath\u003e` argument.\n\nWhen using the wandb logger (default), checkpoints may be saved to wandb (`logger.wandb.save_to_wandb`) or S3 storage (`logger.wandb.save_to_s3_bucket`).\nSet these to `False` to disable saving them to wandb or S3.\nIf disabling both (only save checkpoints locally), make sure to set `logger.wandb.save_best_ckpt=False logger.wandb.save_last_ckpt=False`.\nYou can set these preferences in your [local config](src/configs/local/default.yaml) file\n(see [src/configs/local/example_local_config.yaml](src/configs/local/example_local_config.yaml) for an example).\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eDebugging\u003c/summary\u003e\n\nFor minimal data and model size, you can use the following:\n\n    python run.py ++model.debug_mode=True ++datamodule.debug_mode=True\n\nNote that the model and datamodule need to support to appropriately handle the debug mode.\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eCode Quality\u003c/summary\u003e\n\nCode quality is automatically checked when pushing to the repository.\nHowever, it is recommended that you also run the checks locally with `make quality`.\n\nTo automatically fix some issues (as much as possible), run:\n\n    make style\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003ehydra.errors.InstantiationException\u003c/summary\u003e\n\nThe ``hydra.errors.InstantiationException`` itself is not very informative, \nso you need to look at the preceding exception(s) (i.e. scroll up) to see what went wrong.\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eLocal Configurations\u003c/summary\u003e\n\nYou can use a local config file that, defines the local data dir, working dir etc., by putting a ``default.yaml`` config \nin the [src/configs/local/](src/configs/local) subdirectory. Hydra searches for \u0026 uses by default the file configs/local/default.yaml, if it exists.\nYou may take inspiration from the [example_local_config.yaml](src/configs/local/example_local_config.yaml) file.\n\u003c/details\u003e\n\n## | Citation\n\n    @inproceedings{cachay2024spherical,\n        title={Probablistic Emulation of a Global Climate Model with Spherical {DY}ffusion},\n        author={Salva R{\\\"u}hling Cachay and Brian Henn and Oliver Watt-Meyer and Christopher S. Bretherton and Rose Yu},\n        booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},\n        year={2024},\n        url={https://openreview.net/forum?id=Ib2iHIJRTh}\n    }\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frose-stl-lab%2Fspherical-dyffusion","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Frose-stl-lab%2Fspherical-dyffusion","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frose-stl-lab%2Fspherical-dyffusion/lists"}