diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml new file mode 100644 index 0000000..fe312a1 --- /dev/null +++ b/.github/workflows/lint.yaml @@ -0,0 +1,30 @@ +name: Python Linting CI + +on: + push: + branches: [ $default-branch ] + pull_request: + types: + - synchronize + - opened + - reopened + +jobs: + linting: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: pip install .[lint] + # annotate each step with `if: always` to run all regardless + - name: Check code formatting with ruff + if: always() + run: ruff format --diff newCAM_emulation/ + - name: Lint with ruff using pyproject.toml configuration + if: always() + run: ruff check newCAM_emulation/ \ No newline at end of file diff --git a/README.md b/README.md index ee234ad..e1e06bb 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,8 @@ -# Overview -The repository contains the code for a machine learning model that emulates the climatic process of gravity wave drag (GWD, both zonal and meridional). The model is a part of parameterization scheme where smaller and highly dynamical climatic processes are emulated using neural networks. +# newCAM-Emulation + +This is a DNN written with PyTorch to Emulate the gravity wave drag (GWD, both zonal and meridional) in the CAM model. +The repository contains the code for a machine learning model that emulates the climatic process of gravity wave drag (GWD, both zonal and meridional). +The model is a part of parameterization scheme where smaller and highly dynamical climatic processes are emulated using neural networks. Gravity waves, also called buyoncy waves are formed due to displacement of air in the atmosphere instigated by differnt physical mechanisms, such as moist convection, orographic lifting, shear unstability etc. These waves can propagate both vertically and horizontally through the lift and drag mechanism respectively. This ML model focuses on the drag component of gravity waves. diff --git a/newCAM_emulation/Model.py b/newCAM_emulation/Model.py index d259d93..ccdd690 100644 --- a/newCAM_emulation/Model.py +++ b/newCAM_emulation/Model.py @@ -1,26 +1,36 @@ +"""Neural Network model for the CAM-EM.""" + import netCDF4 as nc import numpy as np import scipy.stats as st -import xarray as xr - import torch +import xarray as xr from torch import nn -import torch.nn.utils.prune as prune -from torch.utils.data import DataLoader -from torch.utils.data import Dataset +from torch.nn.utils import prune +from torch.utils.data import DataLoader, Dataset + # Required for feeding the data iinto NN. class myDataset(Dataset): - def __init__(self, X, Y): + """ + Dataset class for loading features and labels. + Args: + X (numpy.ndarray): Input features. + Y (numpy.ndarray): Corresponding labels. + """ + + def __init__(self, X, Y): + """Create an instance of myDataset class.""" self.features = torch.tensor(X, dtype=torch.float64) self.labels = torch.tensor(Y, dtype=torch.float64) def __len__(self): + """Return the number of samples in the dataset.""" return len(self.features.T) def __getitem__(self, idx): - + """Return a sample from the dataset.""" feature = self.features[:, idx] label = self.labels[:, idx] @@ -29,12 +39,23 @@ def __getitem__(self, idx): # The NN model. class FullyConnected(nn.Module): + """ + Fully connected neural network model. + + The model consists of multiple fully connected layers with SiLU activation function. + + Attributes + ---------- + linear_stack (torch.nn.Sequential): Sequential container for layers. + """ + def __init__(self): + """Create an instance of FullyConnected NN model.""" super(FullyConnected, self).__init__() - ilev=93 + ilev = 93 self.linear_stack = nn.Sequential( - nn.Linear(8*ilev+4, 500, dtype=torch.float64), + nn.Linear(8 * ilev + 4, 500, dtype=torch.float64), nn.SiLU(), nn.Linear(500, 500, dtype=torch.float64), nn.SiLU(), @@ -58,16 +79,38 @@ def __init__(self): nn.SiLU(), nn.Linear(500, 500, dtype=torch.float64), nn.SiLU(), - nn.Linear(500, 2*ilev, dtype=torch.float64), + nn.Linear(500, 2 * ilev, dtype=torch.float64), ) def forward(self, X): + """ + Forward pass through the network. + Args: + X (torch.Tensor): Input tensor. + + Returns + ------- + torch.Tensor: Output tensor. + """ return self.linear_stack(X) # training loop def train_loop(dataloader, model, loss_fn, optimizer): + """ + Training loop. + + Args: + dataloader (DataLoader): DataLoader for training data. + model (nn.Module): Neural network model. + loss_fn (torch.nn.Module): Loss function. + optimizer (torch.optim.Optimizer): Optimizer. + + Returns + ------- + float: Average training loss. + """ size = len(dataloader.dataset) avg_loss = 0 for batch, (X, Y) in enumerate(dataloader): @@ -90,6 +133,18 @@ def train_loop(dataloader, model, loss_fn, optimizer): # validating loop def val_loop(dataloader, model, loss_fn): + """ + Validation loop. + + Args: + dataloader (DataLoader): DataLoader for validation data. + model (nn.Module): Neural network model. + loss_fn (torch.nn.Module): Loss function. + + Returns + ------- + float: Average validation loss. + """ avg_loss = 0 with torch.no_grad(): for batch, (X, Y) in enumerate(dataloader): @@ -101,6 +156,3 @@ def val_loop(dataloader, model, loss_fn): avg_loss /= len(dataloader) return avg_loss - - - diff --git a/newCAM_emulation/NN_pred.py b/newCAM_emulation/NN_pred.py index d5fc280..8b20304 100644 --- a/newCAM_emulation/NN_pred.py +++ b/newCAM_emulation/NN_pred.py @@ -1,23 +1,17 @@ +"""Prediction module for the neural network.""" -""" -The following is an import of PyTorch libraries. -""" +import matplotlib.pyplot as plt +import Model +import netCDF4 as nc +import numpy as np import torch -import torch.nn as nn import torch.nn.functional as nnF -from torch.utils.data import DataLoader import torchvision +from loaddata import data_loader, newnorm +from torch import nn +from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchvision.utils import save_image -import matplotlib.pyplot as plt -import numpy as np -import random -import netCDF4 as nc -import Model -from loaddata import newnorm, data_loader - - - """ Determine if any GPUs are available @@ -132,7 +126,7 @@ VTGWSPEC = np.asarray(F['BVTGWSPEC'][0,:,:]) VTGWSPEC = newnorm(VTGWSPEC, VTGWSPECm, VTGWSPECs) - + print('shape of PS',np.shape(PS)) print('shape of Z3',np.shape(Z3)) @@ -146,8 +140,9 @@ print('shape of UTGWSPEC',np.shape(UTGWSPEC)) print('shape of VTGWSPEC',np.shape(VTGWSPEC)) - x_test,y_test = data_loader (U,V,T, DSE, NM, NETDT, Z3, RHOI, PS,lat,lon,UTGWSPEC, VTGWSPEC) - + x_test,y_test = data_loader (U,V,T, DSE, NM, NETDT, Z3, + RHOI, PS,lat,lon,UTGWSPEC, VTGWSPEC) + print('shape of x_test', np.shape(x_test)) print('shape of y_test', np.shape(y_test)) @@ -166,10 +161,10 @@ print(np.corrcoef(truth.flatten(), predict.flatten())[0, 1]) print('shape of truth ',np.shape(truth)) print('shape of prediction',np.shape(predict)) - + np.save('./pred_data_' + str(iter) + '.npy', predict) - - + + diff --git a/newCAM_emulation/loaddata.py b/newCAM_emulation/loaddata.py index c774b96..859bf26 100644 --- a/newCAM_emulation/loaddata.py +++ b/newCAM_emulation/loaddata.py @@ -1,3 +1,5 @@ +"""Implementing data loader for training neural network.""" + import numpy as np ilev = 93 @@ -5,6 +7,17 @@ dim_NNout =int(2*ilev) def newnorm(var, varm, varstd): + """Normalizes the input variable(s) using mean and standard deviation. + + Args: + var (numpy.ndarray): Input variable(s) to be normalized. + varm (numpy.ndarray): Mean of the variable(s). + varstd (numpy.ndarray): Standard deviation of the variable(s). + + Returns + ------- + numpy.ndarray: Normalized variable(s). + """ dim=varm.size if dim > 1 : vara = var - varm[:, :] @@ -17,11 +30,32 @@ def newnorm(var, varm, varstd): def data_loader (U,V,T, DSE, NM, NETDT, Z3, RHOI, PS, lat, lon, UTGWSPEC, VTGWSPEC): + """ + Loads and preprocesses input data for neural network training. + Args: + U (numpy.ndarray): Zonal wind component. + V (numpy.ndarray): Meridional wind component. + T (numpy.ndarray): Temperature. + DSE (numpy.ndarray): Dry static energy. + NM (numpy.ndarray): Northward mass flux. + NETDT (numpy.ndarray): Net downward total radiation flux. + Z3 (numpy.ndarray): Geopotential height. + RHOI (numpy.ndarray): Air density. + PS (numpy.ndarray): Surface pressure. + lat (numpy.ndarray): Latitude. + lon (numpy.ndarray): Longitude. + UTGWSPEC (numpy.ndarray): Target zonal wind spectral component. + VTGWSPEC (numpy.ndarray): Target meridional wind spectral component. + + Returns + ------- + tuple: A tuple containing the input data and target data arrays. + """ Ncol = U.shape[1] #Nlon = U.shape[2] #Ncol = Nlat*Nlon - + x_train = np.zeros([dim_NN,Ncol]) y_train = np.zeros([dim_NNout,Ncol]) diff --git a/newCAM_emulation/train.py b/newCAM_emulation/train.py index 7379560..198097e 100644 --- a/newCAM_emulation/train.py +++ b/newCAM_emulation/train.py @@ -1,28 +1,45 @@ -import matplotlib -import matplotlib.pyplot as plt +"""Training script for the neural network.""" + +import Model import netCDF4 as nc import numpy as np -import scipy.stats as st -import xarray as xr - import torch +from loaddata import data_loader, newnorm from torch import nn -import torch.nn.utils.prune as prune +from torch.backends import mps +from torch.cuda import is_available from torch.utils.data import DataLoader -from torch.utils.data import Dataset -import Model -from loaddata import newnorm, data_loader +if is_available(): + DEVICE = "cuda" +elif mps.is_available(): + DEVICE = "mps" +else: + DEVICE = "cpu" +print(f"Using device: {DEVICE}") class EarlyStopper: + """Class for implementing early stopping during training.""" + def __init__(self, patience=1, min_delta=0): + """Create an instance of EarlyStopper class.""" self.patience = patience self.min_delta = min_delta self.counter = 0 self.min_validation_loss = np.inf def early_stop(self, validation_loss): + """ + Check if early stopping condition is met. + + Args: + validation_loss (float): Loss value on the validation set. + + Returns + ------- + bool: True if early stopping condition is met, False otherwise. + """ if validation_loss < self.min_validation_loss: self.min_validation_loss = validation_loss self.counter = 0 @@ -98,52 +115,58 @@ def early_stop(self, validation_loss): F = nc.Dataset(filename) PS = np.asarray(F['PS'][0,:]) PS = newnorm(PS, PSm, PSs) - + Z3 = np.asarray(F['Z3'][0,:,:]) Z3 = newnorm(Z3, Z3m, Z3s) - + U = np.asarray(F['U'][0,:,:]) U = newnorm(U, Um, Us) - + V = np.asarray(F['V'][0,:,:]) V = newnorm(V, Vm, Vs) - + T = np.asarray(F['T'][0,:,:]) T = newnorm(T, Tm, Ts) - + lat = F['lat'] lat = newnorm(lat, np.mean(lat), np.std(lat)) - + lon = F['lon'] lon = newnorm(lon, np.mean(lon), np.std(lon)) - + DSE = np.asarray(F['DSE'][0,:,:]) DSE = newnorm(DSE, DSEm, DSEs) - + RHOI = np.asarray(F['RHOI'][0,:,:]) RHOI = newnorm(RHOI, RHOIm, RHOIs) - + NETDT = np.asarray(F['NETDT'][0,:,:]) NETDT = newnorm(NETDT, NETDTm, NETDTs) - + NM = np.asarray(F['NMBV'][0,:,:]) NM = newnorm(NM, NMm, NMs) - + UTGWSPEC = np.asarray(F['UTGWSPEC'][0,:,:]) UTGWSPEC = newnorm(UTGWSPEC, UTGWSPECm, UTGWSPECs) - + VTGWSPEC = np.asarray(F['VTGWSPEC'][0,:,:]) VTGWSPEC = newnorm(VTGWSPEC, VTGWSPECm, VTGWSPECs) - - x_train,y_train = data_loader(U,V,T, DSE, NM, NETDT, Z3, RHOI, PS,lat,lon,UTGWSPEC, VTGWSPEC) + + x_train,y_train = data_loader(U,V,T, DSE, NM, NETDT, Z3, + RHOI, PS,lat,lon,UTGWSPEC, VTGWSPEC) data = Model.myDataset(X=x_train, Y=y_train) batch_size = 128 - split_data = torch.utils.data.random_split(data, [0.75, 0.25], generator=torch.Generator().manual_seed(42)) - train_dataloader = DataLoader(split_data[0], batch_size=batch_size, shuffle=True) - val_dataloader = DataLoader(split_data[1], batch_size=len(split_data[1]), shuffle=True) + split_data = torch.utils.data.random_split(data, [0.75, 0.25], + generator=torch.Generator().manual_seed(42)) + train_dataloader = DataLoader(split_data[0], + batch_size=batch_size, + shuffle=True) + val_dataloader = DataLoader(split_data[1], + batch_size=len(split_data[1]), + shuffle=True) # training early_stopper = EarlyStopper(patience=5, min_delta=0) # Note the hyper parameters. @@ -153,11 +176,11 @@ def early_stop(self, validation_loss): print(val_losses[-1]) print('counter=' + str(early_stopper.counter)) train_loss = Model.train_loop(train_dataloader, model, nn.MSELoss(), optimizer) - + train_losses.append(train_loss) val_loss = Model.val_loop(val_dataloader, model, nn.MSELoss()) val_losses.append(val_loss) if early_stopper.early_stop(val_loss): print("BREAK!") break - + diff --git a/pyproject.toml b/pyproject.toml index 572ca9a..93d8d9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,12 +37,7 @@ dependencies = [ [project.optional-dependencies] lint = [ - "black>=24.1.0", - "pylint", - # "mypy>=1.0.0", - # "pytest>=7.2.0", - # "pytest-mock", - "pydocstyle", + "ruff>=0.3.2", ] [project.urls] @@ -59,3 +54,19 @@ where = ["."] # list of folders that contain the packages (["."] by default) include = ["newCAM_emulation*"] # package names should match these glob patterns (["*"] by default) exclude = ["Demodata/*"] # exclude packages matching these glob patterns (empty by default) namespaces = false # to disable scanning PEP 420 namespaces (true by default) + +[tool.ruff] +# Run linting and formatting on notebooks +extend-include = ["*.ipynb"] + +[tool.ruff.lint] +# Enable: D: `pydocstyle`, PL: `pylint`, I: `isort`, W: `pycodestyle whitespace` +# NPY: `numpy`, +select = ["D", "PL", "I", "E", "W", "NPY" ] + +# Enable D417 (Missing argument description) on top of the NumPy convention. +extend-select = ["D417"] + +[tool.ruff.lint.pydocstyle] +# Use NumPy convention for checking docstrings +convention = "numpy" \ No newline at end of file