diff --git a/README.md b/README.md index 0ceeb50e..5a962bbf 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,13 @@ -# NSF NCAR MILES Community Runnable Earth Digital Intelligence Twin (CREDIT) +# NSF NCAR MILES Community Research Earth Digital Intelligence Twin (CREDIT) ## About -CREDIT is a package to train and run neural networks -that can emulate full NWP models by predicting -the next state of the atmosphere given the current state. +CREDIT is a research platform to train and run neural networks that can emulate full NWP models by predicting +the next state of the atmosphere given the current state. The platform is still under very active development. +If you are interested in using or contributing to CREDIT, please reach out to David John Gagne (dgagne@ucar.edu). + ## NSF-NCAR Derecho Installation -Currently, the framework for running miles-credit in parallel is centered around NSF-NCAR's Derecho HPC. Derecho requires building several miles-credit dependent packages locally, including PyTorch, to enable correct MPI configuration. To begin, create a clone of the pre-built miles-credit environment, which contains compatiable versions of torch, torch-vision, numpy, and others. +Currently, the framework for running miles-credit in parallel is centered around NSF NCAR's Derecho HPC. Derecho requires building several miles-credit dependent packages locally, including PyTorch, to enable correct MPI configuration. To begin, create a clone of the pre-built miles-credit environment, which contains compatiable versions of torch, torch-vision, numpy, and others. ```bash module purge @@ -72,14 +73,14 @@ python applications/train.py -c config/vit.yml Or use a fancier [variation](https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/rvt.py) ```bash -python applications/train.py -c config/rvt.yml +python applications/train.py -c config/wxformer_1dg_test.yml ``` ## Launch with PBS on Casper or Derecho Adjust the PBS settings in a configuration file for either casper or derecho. Then, submit the job via ```bash -python applications/train.py -c config/vit.yml -l 1 +python applications/train.py -c config/wxformer_1dg_test.yml -l 1 ``` The launch script may be found in the save location that you set in the configation file. The automatic launch script generation will take care of MPI calls and other complexities if you are using more than 1 GPU. @@ -88,5 +89,12 @@ The launch script may be found in the save location that you set in the configat The predict field in the config file allows one to speficy start and end dates to roll-out a trained model. To generate a forecast, ```bash -python applications/rollout_to_netcdf.py -c config/vit.yml +python applications/rollout_to_netcdf.py -c config/wxformer_1dg_test.yml ``` + +# Support +This software is based upon work supported by the NSF National Center for Atmospheric Research, a major facility sponsored by the +U.S. National Science Foundation under Cooperative Agreement No. 1852977 and managed by the University Corporation for Atmospheric Research. Any opinions, findings and conclusions or recommendations +expressed in this material do not necessarily reflect the views of NSF. Additional support for development was provided by +The NSF AI Institute for Research on Trustworthy AI for Weather, Climate, and Coastal Oceanography (AI2ES) with grant +number RISE-2019758. diff --git a/credit/VERSION b/credit/VERSION index 57c6e875..a73a8518 100644 --- a/credit/VERSION +++ b/credit/VERSION @@ -1 +1 @@ -2023.1.0 +2024.1.0 diff --git a/credit/models/crossformer.py b/credit/models/crossformer.py index 27993730..eb941bb6 100644 --- a/credit/models/crossformer.py +++ b/credit/models/crossformer.py @@ -345,31 +345,62 @@ def forward(self, x): class CrossFormer(BaseModel): def __init__( self, - image_height=640, - patch_height=1, - image_width=1280, - patch_width=1, - frames=2, - channels=4, - surface_channels=7, - input_only_channels=3, - output_only_channels=0, - levels=15, - dim=(64, 128, 256, 512), - depth=(2, 2, 8, 2), - dim_head=32, - global_window_size=(5, 5, 2, 1), - local_window_size=10, - cross_embed_kernel_sizes=((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), - cross_embed_strides=(4, 2, 2, 2), - attn_dropout=0.0, - ff_dropout=0.0, - use_spectral_norm=True, - interp=True, - padding_conf=None, - post_conf=None, + image_height: int = 640, + patch_height: int = 1, + image_width: int = 1280, + patch_width: int = 1, + frames: int = 2, + channels: int = 4, + surface_channels: int = 7, + input_only_channels: int = 3, + output_only_channels: int = 0, + levels: int = 15, + dim: tuple = (64, 128, 256, 512), + depth: tuple = (2, 2, 8, 2), + dim_head: int = 32, + global_window_size: tuple = (5, 5, 2, 1), + local_window_size: int = 10, + cross_embed_kernel_sizes: tuple = ((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), + cross_embed_strides: tuple = (4, 2, 2, 2), + attn_dropout: float = 0.0, + ff_dropout: float = 0.0, + use_spectral_norm: bool = True, + interp: bool = True, + padding_conf: dict = None, + post_conf: dict = None, **kwargs, ): + """ + CrossFormer is the base architecture for the WXFormer model. It uses convolutions and long and short distance + attention layers in the encoder layer and then uses strided transpose convolution blocks for the decoder + layer. + + Args: + image_height (int): number of grid cells in the south-north direction. + patch_height (int): number of grid cells within each patch in the south-north direction. + image_width (int): number of grid cells in the west-east direction. + patch_width (int): number of grid cells within each patch in the west-east direction. + frames (int): number of time steps being used as input + channels (int): number of 3D variables. Default is 4 for our ERA5 configuration (U, V, T, and Q) + surface_channels (int): number of surface (single-level) variables. + input_only_channels (int): number of variables only used as input to the ML model (e.g., forcing variables) + output_only_channels (int):number of variables that are only output by the model (e.g., diagnostic variables). + levels (int): number of vertical levels for each 3D variable (should be the same across frames) + dim (tuple): output dimensions of hidden state of each conv/transformer block in the encoder + depth (tuple): number of attention blocks per encoder layer + dim_head (int): dimension of each attention head. + global_window_size (tuple): number of grid cells between cells in long range attention + local_window_size (tuple): number of grid cells between cells in short range attention + cross_embed_kernel_sizes (tuple): width of the cross embed kernels in each layer + cross_embed_strides (tuple): stride of convolutions in each block + attn_dropout (float): dropout rate for attention layout + ff_dropout (float): dropout rate for feedforward layers. + use_spectral_norm (bool): whether to use spectral normalization + interp (bool): whether to use interpolation + padding_conf (dict): padding configuration + post_conf (dict): configuration for postblock processing + **kwargs: + """ super().__init__() dim = tuple(dim) diff --git a/docs/source/_static/credit_logo.graffle b/docs/source/_static/credit_logo.graffle new file mode 100644 index 00000000..abe44efb Binary files /dev/null and b/docs/source/_static/credit_logo.graffle differ diff --git a/docs/source/_static/credit_logo.png b/docs/source/_static/credit_logo.png new file mode 100644 index 00000000..308cd421 Binary files /dev/null and b/docs/source/_static/credit_logo.png differ diff --git a/docs/source/contrib.rst b/docs/source/contrib.rst new file mode 100644 index 00000000..75c99d7c --- /dev/null +++ b/docs/source/contrib.rst @@ -0,0 +1,446 @@ +.. _contributing: + +=========================== +Contributing to miles-credit +=========================== + +************ +Introduction +************ + +Thank you for considering making a contribution to ``miles-credit``! There are +many ways to contribute to this project, including reporting bugs, requesting +additional or new functionality, improving our documentation, or contributing +your own code and we appreciate all of them. + +If you have any questions, please feel free to reach out to us on `GitHub +Issues `__. You can also +reach us by email at milescore@ucar.edu. + +************** +Where to start +************** + +Look through our open issues and see if there is anything you would like to take +on! We recommend working with core developers to implement new functionality. We +can help you get started and make sure your code is consistent with the rest of +the project. + +Also check out any beginner-friendly issues we have tagged with `good first +issue `__. + +We do not officially "assign" issues to contributors, but if you are interested +in working on an issue, please comment on the issue to let us know you are +working on it. This will help us avoid duplicate work. + +The code for ``miles-credit`` is hosted on GitHub. If you do not have one, you +will need to create a `free GitHub account `__. +The `GitHub Quickstart Guide +`__ is a great place to get +started with git and GitHub. + +************** +Reporting bugs +************** + +Something not working as expected? We would love to hear about it! Please report +any bugs you find by opening an issue on GitHub. + +When reporting a bug, please include as much information as possible. This will +help us reproduce the bug and fix it efficiently. For more information on how to +write a good bug report, see this stackoverflow post on `how to make a good bug +report `__. + +*********************** +Requesting new features +*********************** + +Have an idea for a new feature? Please let us know in the Issues section. + + +*********************** +Improving Documentation +*********************** + +We are always looking for ways to improve our documentation. If you find +something that is unclear or confusing, please let us know by opening an issue. +To contribute to our documentation yourself, see the `Documentation`_ section of +this guide. + +***************************** +Development workflow overview +***************************** + +This is a brief overview of the development workflow we use for ``miles-credit``. +A more detailed description of each step is provided in following sections. + +**Get set up to develop on your local machine** + +#. `Fork and clone the repository`_. +#. `Create a development environment`_. +#. `Create a branch for your changes`_. +#. `Install pre-commit hooks`_. + +**Make your changes** + +#. `Understanding the codebase`_. +#. `Write and run tests`_. +#. :ref:`Generate ` and :ref:`check ` the documentation. + +**Contribute your code** + +#. `Push your changes to your fork`_. +#. `Open a pull request`_. +#. `Address feedback`_. +#. Wait for your pull request to be merged. +#. `Delete your branch`_. + +******************************************* +Get set up to develop on your local machine +******************************************* + +Fork and clone the repository +----------------------------- + +Get started by forking the NCAR/miles-credit repository on GitHub. To do this, +find the "Fork" button near the top of the page and click it. This will create a +copy of the project under your personal github account. + +Next, clone your forked copy to your local machine. + +.. code-block:: bash + + git clone https://github.com/your-user-name/miles-credit.git + + +Enter the project folder and set the upstream remote to the NCAR/miles-credit +repository. This will allow you to keep your fork up to date with the main +repository. + +.. code-block:: bash + + cd miles-credit git remote add upstream https://github.com/NCAR/miles-credit.git + +For more information, see the `GitHub quickstart section on forking a repository +`__. + +Create a development environment +-------------------------------- + +To run and test any changes you make in ``miles-credit``, you will need to create +a development environment. We recommend installing and using `conda +`__ +and/or `mamba +`__. + +Use the following commands to create a new conda environment to develop +``miles-credit`` in. + +.. code-block:: bash + + # Create a new conda environment + conda create -c conda-forge -n credit python=3.11 + + # Use the environment file to populate the environment with the required dependencies + conda env update -f environment_gpu.yml + + # Activate your new environment + conda activate credit + + # Install your local copy of miles-credit in interactive mode + pip install -e . + +To test your new install, open a python session and try importing +``credit.``. You can also try printing the version number, which should be +unique to the latest commit on your fork. + +.. code-block:: python + + >>> import credit + >>> credit.__version__ + '2024.1.0' + + + + +Create a branch for your changes +-------------------------------- + +We highly recommend creating a new branch on your fork for each new feature or +bug that you work on. + +To create and check out a new branch, use the following command: + +.. code-block:: bash + + git checkout -b + +Track upstream changes for `git pull` and `git push` + +.. code-block:: bash + + git branch --set-upstream-to=origin/ + +You can see a list of all branches in your local repository by running: + +.. code-block:: bash + + git branch + +For more information on branching, check out this `learn git branching +`__ interactive tool. + +Install pre-commit hooks +------------------------ + +``miles-credit`` uses pre-commit hooks to ensure a standardized base-level code +formatting and style. + +The ``pre-commit`` package is installed by default when using the +``build_envs/environment.yml`` file. To set up the pre-commit hooks, run the +following command from the root of the repository: + +.. code-block:: bash + + pre-commit install + +Now, whenever you commit changes, the pre-commit hooks will run and may make +small modifications to your code. If the pre-commit hooks make any changes, you +will need to re-add the files and commit them again in order to successfully make +the commit. + +To manually run the pre-commit hooks, use the following command: + +.. code-block:: bash + + pre-commit run --all-files + +If `codespell` flags on a jargon specific term, you can add the term to the `ignore-words-list` in`pyproject.toml`. Any words added to the `ignore-words-list` should be lower-case. + +You can skip the pre-commit hooks by adding the ``--no-verify`` flag to your +commit command like this: + +.. code-block:: bash + + git commit -m "your commit message" --no-verify + +For more information on pre-commit hooks, see the `pre-commit documentation `__. + + +***************** +Make your changes +***************** + +After you're all set up to develop ``miles-credit``, you can start making your +changes. This section describes where, how, and what to change to add your +contributions to the ``miles-credit`` codebase. + + +Understanding the codebase +-------------------------- + +The ``miles-credit`` top-level directory is organized as follows: + +.. code-block:: bash + + miles-credit + ├── applications + ├── config + ├── docs + ├── credit + │ └── models + │ └── datasets + │ └── metadata + │ └── trainers + └── tests + + +* The ``applications`` directory contains scripts for training and running models +developed in CREDIT. + +* The ``docs`` directory contains the ``sphinx`` documentation for + ``miles-credit``. + +* The ``credit`` directory, contains the library code for the + package. This is the place to add new functionality. The code + is organized into modules, each of which is contained in its own file. It is + recommended that you add new functionality to an existing file, though it may + be appropriate to make a new file. + +* The ``tests`` directory contains the unit tests for ``miles-credit``. Most + modules in ``credit`` has a corresponding test module in the ``tests`` + directory. + + + +Write and run tests +------------------- + +``miles-credit`` uses `pytest `__ for unit tests, so we +encourage you to write new tests using ``pytest`` as well. + +To run the tests locally, use the following command from the root of the +repository: + +.. code-block:: bash + + pytest + +To run a specific test, use the following command: + +.. code-block:: bash + + pytest tests/test_mod.py::test_func + +These tests will also run automatically when you open a pull request using +GitHub Actions and the ``.github/workflows/ci.yml`` file. + +See the `pytest documentation `__ for more information. + + +************* +Documentation +************* + +``miles-credit`` uses `sphinx `__ and +`ReadTheDocs `__ to build and host the +documentation. + + +Docstrings +---------- + +The most common situation in which you will need to add to the documentation is +through docstrings. + +``miles-credit`` uses `Google +`__ style docstrings. + +To include your docstring documentation in the API reference, you will need to +add it to either the ``docs/internal_api/index.rst`` or +``docs/user_api/index.rst`` file, depending on whether the function is intended +for internal or external use. + +Editing other documentation files +--------------------------------- + +We welcome changes and improvements to all parts of our documentation (including +this guide)! You can find these files in the ``docs`` directory. + +These files are mainly written in `reStructuredText +`__, +but additional file types such as ``.md`` and ``.ipynb`` are also used. + +Important documentation files to know about include: + +* ``docs/index.rst``: This file is the main page of the documentation. Files + added to ``toctree`` blocks in this file will be included in the documentation + as top-level subpages. + +* ``docs/contrib.rst``: This file is the source for this guide! + +* ``docs/conf.py``: This file contains the configuration for building the documentation. + +* ``docs/examples/*.ipynb``, ``docs/examples.rst``, and ``docs/gallery.yml``: + These files are used to generate the jupyter notebook examples in the + documentation. Notebooks in the ``docs/examples/`` directory are added to the + documentation by adding them to the ``toctree`` in ``docs/examples.rst`` and + linked to their cover picture by adding them to the ``docs/gallery.yml`` + file. + +See the `sphinx documentation `__ for +more information about writing sphinx documentation. + +.. _generate-docs: + +Generate the documentation locally +---------------------------------- + +To generate the documentation locally, follow the steps below. + +#. Create and activate the ``gc-docs`` conda environment using the ``build_envs/docs.yml`` file. +#. Enter the ``docs`` directory. +#. Run ``make html`` or to build the documentation. +#. Open ``docs/_build/html/index.html`` in your browser to view the documentation. + +.. _check-docs: + +Check the documentation +----------------------- + +As well as checking local documentation generation, you should also check the +preview documentation generated as part of a PR. To do this, scroll down to the +"checks" section of the PR and click on the "Details" link next to the +"docs/readthedocs.org:miles-credit" check. This will take you to the +corresponding build on ReadTheDocs, where you can view the documentation built +from your PR and see any warnings or errors on your build. + +******************** +Contribute your code +******************** + +Once you have prepared your changes and are ready for them to be reviewed by the +GeoCAT team, you can open a pull request. This section describes how to open a +pull request and what to expect after you open it. + +Push your changes to your fork +------------------------------ + +Once you have made your changes locally, you will need to push them to your +branch on your fork on GitHub. To do this, use the following command: + +.. code-block:: bash + + git push + +From here, you can request that your changes be merged into the main repository in the form of a pull request. + +Open a pull request +------------------- + +GitHub has extensive `pull request guides and documentation +`__ that we recommend. This section +describes the basics for our workflow. + +From your branch on your fork, open the "Pull requests" tab and click the "New +pull request" button. Make sure the "base repository" is "NCAR/miles-credit" and +the "base" branch is set to "main", with the "head repository" and "compare" +branch set to your fork and prepared branch, respectively. + +From this page, you can see a view of the changes you have made in your branch. + +We recommend adding a short, descriptive title to your pull request. The body of +the pull request should autofill with our pull request template, which has it's +own set of directions. Please fill out the relevant sections of the template, +including adding a more detailed description of your changes. + +Once you have filled out the template, click the "Create pull request" button. +This will open your pull request on the ``miles-credit`` repository. + +If you want to open a pull request but are not ready for it to be reviewed, you +can open the pull request as a draft. This is also a good way to get feedback on +your work that might not be ready to contribute yet. + +Address feedback +---------------- + +After you open your pull request, the CREDIT team will review it and +may provide feedback like asking for changes or suggesting improvements. You can +address this feedback by making changes to your branch and pushing them to your +fork. The pull request will automatically update with your changes. + +The CREDIT team appreciates your contributions and the interactive process of +reviewing pull requests, and will do our best to review your pull request in a +timely manner. It is totally normal to have to make several rounds of changes to +your pull request before it is ready to be merged, especially if you are new to +the project. + +Once your pull request is approved by a core maintainer and passes the relevant +checks, it will be merged into the main repository! + + +Delete your branch +------------------ + +We recommend deleting your branch after your pull request is merged. This will +help keep your fork clean and organized, but is not required. \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst index b1798aa7..9e294a23 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -6,7 +6,7 @@ MILES-CREDIT documentation! ======================================== -The Community Research Earth Digital Intelligent Twin (CREDIT) project is a +The NSF NCAR Community Research Earth Digital Intelligent Twin (CREDIT) project is a research platform developed by the NSF NCAR Machine Integration and Learning for Earth Systems, or `MILES `_ group. CREDIT allows users to train, run, and evaluate AI numerical weather prediction models. @@ -17,7 +17,7 @@ allows users to train, run, and evaluate AI numerical weather prediction models. :caption: Contents: Getting Started - + Contributing Indices and tables diff --git a/requirements.txt b/requirements.txt index caea1c43..69d4a027 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -numpy +numpy<2 pandas matplotlib cartopy