models.py

"""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(
 "NN_weights_YOG_convection.nc"
 )
 net.save("nn_state.pt")
 print("Model saved to nn_state.pt")