Upload models

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020, 2023 Janni Yuval and Institute of Computing for Climate Science
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,59 @@
+# Convection Parameterization in CAM
+
+Note that this repository and code is still work in progress and undergoing significant development.
+Once a useable release is produced it will be tagged.
+
+## Description
+This repository contains code as part of an effort to deploy machine learning (ML) models of geophysical parameterisations into the [Community Earth System Model (CESM)](https://www.cesm.ucar.edu/).
+This work is part of the [M<sup>2</sup>LInES](https://m2lines.github.io/) project aiming to improve performance of climate models using ML models for subgrid parameterizations.
+
+A Neural Net providing a subgrid parameterization of atmospheric convection in a [single column model](https://www.arm.gov/publications/proceedings/conf04/extended_abs/randall_da.pdf) has been developed and successfully deployed as part of an atmospheric simulation.
+The work is described in a [GRL paper](https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2020GL091363) with [accompanying code available](https://github.com/yaniyuval/Neural_nework_parameterization/tree/v.1.0.3). The repository contains the neural net and its implementation into a simple system for atmospheric modelling, [SAM](http://rossby.msrc.sunysb.edu/~marat/SAM.html).
+
+The aims of this repository are to:
+1. develop a standalone fortran module based on this neural net that can be used elsewhere,
+2. deploy the module in another atmospheric model, and
+3. evaluate its performance.
+
+We may also perform an investigation into interfacing the pytorch implementation of the Neural Net using the [pytorch-fortran bridging code](https://github.com/Cambridge-ICCS/fortran-pytorch-lib) developed at the [Institute of Computing for Climate Science](https://cambridge-iccs.github.io/).
+
+The model will first be deployed into the [Single Column Atmospheric Model (SCAM)](https://www.cesm.ucar.edu/models/simple/scam) - a single column version of the CESM.
+We plan to evaluate performance using SCAM in the gateIII configuration for tropical convection in a similar manner described by the [SCAM6 pulication in JAMES](https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2018MS001578).
+This will compare model performance to data from an intense observation period (IOP) described in an [AMS publication](https://journals.ametsoc.org/view/journals/atsc/36/1/1520-0469_1979_036_0053_saposs_2_0_co_2.xml).
+
+Long term developments of this project will seek to re-deploy more complex ML parameterizations into mode complex atmospheric models such as the [Community Atmospheric Model (CAM)](https://www.cesm.ucar.edu/models/cam) part of the CESM.
+
+
+## Repository structure
+
+```
+├── NN_module
+│   └── ...
+└── torch_nets
+    └── ...
+```
+
+
+
+### Contents
+
+### `NN_module/`
+This folder contains the fortran neural net extracted from the [code referenced above](https://github.com/yaniyuval/Neural_nework_parameterization/tree/v.1.0.3), along with any dependencies, that may be compiled as a standalone fortran module.
+
+Currently there is code that can be built on CSD3 using the included shell script.
+
+This now needs cleaning up, testing, and a proper makefile creating (see open issues #9 and #10).
+
+### ``torch_nets/``
+The directory contains the PyTorch versions of the neural networks we are interested in.
+
+
+## Contributing
+
+This repository is currently private as it is new and work in progress.
+Open tickets can be viewed at ['Issues'](https://github.com/m2lines/convection-parameterization-in-CAM/issues).
+
+To contribute find a relevant issue or open a new one and assign yourself to work on it.
+Then create a branch in which to add your contribution and open a pull request.
+Once ready assign a reviewer and request a code review.
+Merging should _only_ be performed once a code review has taken place.

models.py

@@ -0,0 +1,196 @@
+"""Neural network architectures."""
+
+from typing import Optional
+
+import netCDF4 as nc  # type: ignore
+import torch
+from torch import nn, Tensor
+
+
+class ANN(nn.Sequential):
+    """Model used in the paper.
+
+    Paper: https://doi.org/10.1029/2020GL091363
+
+
+    Parameters
+    ----------
+    n_in : int
+        Number of input features.
+    n_out : int
+        Number of output features.
+    n_layers : int
+        Number of layers.
+    neurons : int
+        The number of neurons in the hidden layers.
+    dropout : float
+        The dropout probability to apply in the hidden layers.
+    device : str
+        The device to put the model on.
+    features_mean : ndarray
+        The mean of the input features.
+    features_std : ndarray
+        The standard deviation of the input features.
+    outputs_mean : ndarray
+        The mean of the output features.
+    outputs_std : ndarray
+        The standard deviation of the output features.
+    output_groups : ndarray
+        The number of output features in each group of the ouput.
+
+    Notes
+    -----
+    If you are doing inference, always remember to put the model in eval model,
+    by using ``model.eval()``, so the dropout layers are turned off.
+
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments,too-many-locals
+        self,
+        n_in: int = 61,
+        n_out: int = 148,
+        n_layers: int = 5,
+        neurons: int = 128,
+        dropout: float = 0.0,
+        device: str = "cpu",
+        features_mean: Optional[Tensor] = None,
+        features_std: Optional[Tensor] = None,
+        outputs_mean: Optional[Tensor] = None,
+        outputs_std: Optional[Tensor] = None,
+        output_groups: Optional[list] = None,
+    ):
+        """Initialize the ANN model."""
+        dims = [n_in] + [neurons] * (n_layers - 1) + [n_out]
+        layers = []
+
+        for i in range(n_layers):
+            layers.append(nn.Linear(dims[i], dims[i + 1]))
+            if i < n_layers - 1:
+                layers.append(nn.ReLU())  # type: ignore
+                layers.append(nn.Dropout(dropout))  # type: ignore
+
+        super().__init__(*layers)
+
+        fmean = fstd = omean = ostd = None
+
+        if features_mean is not None:
+            assert features_std is not None
+            assert len(features_mean) == len(features_std)
+            fmean = torch.tensor(features_mean)
+            fstd = torch.tensor(features_std)
+
+        if outputs_mean is not None:
+            assert outputs_std is not None
+            assert len(outputs_mean) == len(outputs_std)
+            if output_groups is None:
+                omean = torch.tensor(outputs_mean)
+                ostd = torch.tensor(outputs_std)
+            else:
+                assert len(output_groups) == len(outputs_mean)
+                omean = torch.tensor(
+                    [x for x, g in zip(outputs_mean, output_groups) for _ in range(g)]
+                )
+                ostd = torch.tensor(
+                    [x for x, g in zip(outputs_std, output_groups) for _ in range(g)]
+                )
+
+        self.register_buffer("features_mean", fmean)
+        self.register_buffer("features_std", fstd)
+        self.register_buffer("outputs_mean", omean)
+        self.register_buffer("outputs_std", ostd)
+
+        self.to(torch.device(device))
+
+    def forward(self, input: Tensor):  # pylint: disable=redefined-builtin
+        """Pass the input through the model.
+
+        Override the forward method of nn.Sequential to add normalization
+        to the input and denormalization to the output.
+
+        Parameters
+        ----------
+        input : Tensor
+            A mini-batch of inputs.
+
+        Returns
+        -------
+        Tensor
+            The model output.
+
+        """
+        if self.features_mean is not None:
+            input = (input - self.features_mean) / self.features_std
+
+        # pass the input through the layers using nn.Sequential.forward
+        output = super().forward(input)
+
+        if self.outputs_mean is not None:
+            output = output * self.outputs_std + self.outputs_mean
+
+        return output
+
+    def load(self, path: str) -> "ANN":
+        """Load the model from a checkpoint.
+
+        Parameters
+        ----------
+        path : str
+            The path to the checkpoint.
+
+        """
+        state = torch.load(path)
+        for key in ["features_mean", "features_std", "outputs_mean", "outputs_std"]:
+            if key in state and getattr(self, key) is None:
+                setattr(self, key, state[key])
+        self.load_state_dict(state)
+        return self
+
+    def save(self, path: str):
+        """Save the model to a checkpoint.
+
+        Parameters
+        ----------
+        path : str
+            The path to save the checkpoint to.
+
+        """
+        torch.save(self.state_dict(), path)
+
+
+def load_from_netcdf_params(nc_file: str, dtype: str = "float32") -> ANN:
+    """Load the model with weights and biases from the netcdf file.
+
+    Parameters
+    ----------
+    nc_file : str
+        The netcdf file containing the parameters.
+    dtype : str
+        The data type to cast the parameters to.
+
+    """
+    data_set = nc.Dataset(nc_file)  # pylint: disable=no-member
+
+    model = ANN(
+        features_mean=data_set["fscale_mean"][:].astype(dtype),
+        features_std=data_set["fscale_stnd"][:].astype(dtype),
+        outputs_mean=data_set["oscale_mean"][:].astype(dtype),
+        outputs_std=data_set["oscale_stnd"][:].astype(dtype),
+        output_groups=[30, 29, 29, 30, 30],
+    )
+
+    for i, layer in enumerate(l for l in model.modules() if isinstance(l, nn.Linear)):
+        layer.weight.data = torch.tensor(data_set[f"w{i+1}"][:].astype(dtype))
+        layer.bias.data = torch.tensor(data_set[f"b{i+1}"][:].astype(dtype))
+
+    return model
+
+
+if __name__ == "__main__":
+    # Load the model from the netcdf file and save it to a checkpoint.
+    net = load_from_netcdf_params(
+        "qobsTTFFFFFTF30FFTFTF30TTFTFTFFF80FFTFTTF2699FFFF_X01_no_qp_no_adv_"
+        "surf_F_Tin_qin_disteq_O_Trad_rest_Tadv_qadv_qout_qsed_RESCALED_7epochs"
+        "_no_drop_REAL_NN_layers5in61out148_BN_F_te70.nc"
+    )
+    net.save("nn_state.pt")
+    print("Model saved to nn_state.pt")

nn_state.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:552b224668a40820e9afd0e4f83053dbc9f4ee7c75814ad32ed2805041afb1e1
+size 312574

requirements.txt

@@ -0,0 +1,7 @@
+torch
+black
+pytest
+pydocstyle
+pylint
+mypy
+netcdf4

test_python_net.py

@@ -0,0 +1,40 @@
+"""A smoke test for the ANN model.
+
+This test checks that the model can be loaded from a weights file in both pt format and
+netcdf format and that they produce the expected output when given an input of all ones.
+This ensures that it is equivalent to the Fortran NN model.
+"""
+
+import os
+from pathlib import Path
+import torch
+import numpy as np
+from models import ANN, load_from_netcdf_params
+
+
+os.chdir(Path(__file__).parent)
+
+expected = np.loadtxt("nn_ones.txt").astype(np.float32)
+# nn_ones.txt is the output of the Fortran NN model given an input of all ones.
+
+model1 = ANN().load("nn_state.pt")  # load from the pytorch weights
+model2 = load_from_netcdf_params(
+    "qobsTTFFFFFTF30FFTFTF30TTFTFTFFF80FFTFTTF2699FFFF_X01_no_qp_no_adv_"
+    "surf_F_Tin_qin_disteq_O_Trad_rest_Tadv_qadv_qout_qsed_RESCALED_7epochs"
+    "_no_drop_REAL_NN_layers5in61out148_BN_F_te70.nc"
+)  # load from the NetCDF weights of the pretrained Fortran NN model
+# file created at https://github.com/yaniyuval/Neural_nework_parameterization/blob/f81f5f695297888f0bd1e0e61524590b4566bf03/NN_training/src/ml_train_nn.py#L417 # pylint: disable=line-too-long
+# (which the naming scheme integrating information about the training setup, see e.g., https://github.com/yaniyuval/Neural_nework_parameterization/blob/f81f5f695297888f0bd1e0e61524590b4566bf03/NN_training/src/ml_train_nn.py#L263-L265) # pylint: disable=line-too-long
+# This Neural Net can be found at https://github.com/yaniyuval/Neural_nework_parameterization/tree/f81f5f695297888f0bd1e0e61524590b4566bf03/NNs # pylint: disable=line-too-long
+
+
+x = torch.ones(61)
+
+actual1 = model1.forward(x).detach().numpy()
+actual2 = model2.forward(x).detach().numpy()
+
+assert np.all(actual1 == actual2)
+assert np.allclose(expected, actual1, atol=3e-8, rtol=2e-6)
+# Values of atol and rtol are chosen to be the lowest that still pass the test.
+
+print("Smoke tests passed")