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")