From a32b6e1ef8cb9a6488d1ec4b20c701b23f55e3bb Mon Sep 17 00:00:00 2001 From: Surbhi Goel Date: Mon, 14 Oct 2024 11:43:17 +0100 Subject: [PATCH 1/3] Added output generation code and data files --- newCAM_emulation/NN_pred.py | 59 ++++++++++++++++++++++++++++++++----- newCAM_emulation/train.py | 12 ++++++-- 2 files changed, 60 insertions(+), 11 deletions(-) diff --git a/newCAM_emulation/NN_pred.py b/newCAM_emulation/NN_pred.py index 3824e44..f0d6544 100644 --- a/newCAM_emulation/NN_pred.py +++ b/newCAM_emulation/NN_pred.py @@ -8,6 +8,7 @@ import torch.nn.functional as nnF import torchvision from loaddata import data_loader, newnorm +from savedata import save_netcdf_file from torch import nn from torch.utils.data import DataLoader from torchvision import datasets, transforms @@ -24,6 +25,8 @@ Initialize Hyperparameters """ ilev = 93 +ilev_94 = 94 + dim_NN = 8 * ilev + 4 dim_NNout = 2 * ilev @@ -33,8 +36,8 @@ ## load mean and std for normalization -fm = np.load("Demodata/mean_demo.npz") -fs = np.load("Demodata/std_demo.npz") +fm = np.load("Demodata/mean_demo_sub.npz") +fs = np.load("Demodata/std_demo_sub.npz") Um = fm["U"] Vm = fm["V"] @@ -73,14 +76,16 @@ optimizer = torch.optim.Adam(GWnet.parameters(), lr=learning_rate) -s_list = list(range(5, 6)) +s_list = list(range(6, 7)) + +data_vars = [] for iter in s_list: if iter > 0: GWnet.load_state_dict(torch.load("./conv_torch.pth")) GWnet.eval() print("data loader iteration", iter) - filename = "./Demodata/Demo_timestep_" + str(iter).zfill(3) + ".nc" + filename = "Demodata/newCAM_demo_sub_" + str(iter).zfill(1) + ".nc" F = nc.Dataset(filename) PS = np.asarray(F["PS"][0, :]) @@ -116,10 +121,10 @@ NM = np.asarray(F["NMBV"][0, :, :]) NM = newnorm(NM, NMm, NMs) - UTGWSPEC = np.asarray(F["BUTGWSPEC"][0, :, :]) + UTGWSPEC = np.asarray(F["UTGWSPEC"][0, :, :]) UTGWSPEC = newnorm(UTGWSPEC, UTGWSPECm, UTGWSPECs) - VTGWSPEC = np.asarray(F["BVTGWSPEC"][0, :, :]) + VTGWSPEC = np.asarray(F["VTGWSPEC"][0, :, :]) VTGWSPEC = newnorm(VTGWSPEC, VTGWSPECm, VTGWSPECs) print("shape of PS", np.shape(PS)) @@ -134,13 +139,32 @@ print("shape of UTGWSPEC", np.shape(UTGWSPEC)) print("shape of VTGWSPEC", np.shape(VTGWSPEC)) + output_filename = f"Demodata/normalized_data_{iter}.nc" + save_netcdf_file( + output_filename, + lat, + lon, + ilev, + ilev_94, + PS, + Z3, + U, + V, + T, + DSE, + RHOI, + NETDT, + NM, + 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)) - data = Model.myDataset(X=x_test, Y=y_test) test_loader = DataLoader(data, batch_size=len(data), shuffle=False) print(test_loader) @@ -155,4 +179,23 @@ print("shape of truth ", np.shape(truth)) print("shape of prediction", np.shape(predict)) - np.save("./pred_data_" + str(iter) + ".npy", predict) + # np.save("./pred_data_" + str(iter) + ".npy", predict) + output_filename = f"Demodata/predicted_data_{iter}.nc" + save_netcdf_file( + output_filename, + lat, + lon, + ilev, + ilev_94, + PS, + Z3, + U, + V, + T, + DSE, + RHOI, + NETDT, + NM, + UTGWSPEC, + VTGWSPEC, + ) diff --git a/newCAM_emulation/train.py b/newCAM_emulation/train.py index 75a521f..10b67f3 100644 --- a/newCAM_emulation/train.py +++ b/newCAM_emulation/train.py @@ -4,6 +4,7 @@ import netCDF4 as nc import numpy as np import torch +import xarray as xr from loaddata import data_loader, newnorm from torch import nn from torch.backends import mps @@ -55,8 +56,8 @@ def early_stop(self, validation_loss): ## load mean and std for normalization -fm = np.load("../Demodata/mean_demo_sub.npz") -fs = np.load("../Demodata/std_demo_sub.npz") +fm = np.load("Demodata/mean_demo_sub.npz") +fs = np.load("Demodata/std_demo_sub.npz") Um = fm["U"] Vm = fm["V"] @@ -88,6 +89,8 @@ def early_stop(self, validation_loss): ilev = 93 +ilev_94 = 94 + dim_NN = int(8 * ilev + 4) dim_NNout = int(2 * ilev) @@ -103,11 +106,14 @@ def early_stop(self, validation_loss): s_list = list(range(1, 6)) +data_vars = [] + for iter in s_list: if iter > 1: model.load_state_dict(torch.load("conv_torch.pth")) + print("data loader iteration", iter) - filename = "../Demodata/newCAM_demo_sub_" + str(iter).zfill(1) + ".nc" + filename = "Demodata/newCAM_demo_sub_" + str(iter).zfill(1) + ".nc" print("working on: ", filename) F = nc.Dataset(filename) From ffc61550d935a754469a959f785c30f8ab4a5521 Mon Sep 17 00:00:00 2001 From: Surbhi Goel Date: Thu, 24 Oct 2024 14:47:23 -0600 Subject: [PATCH 2/3] Generated new normalised input data using single column mean/std, updated NN_pred_new --- newCAM_emulation/NN_pred.py | 17 ++++--- newCAM_emulation/loaddata.py | 89 +++++++++++++++++++++++++++--------- 2 files changed, 78 insertions(+), 28 deletions(-) diff --git a/newCAM_emulation/NN_pred.py b/newCAM_emulation/NN_pred.py index f0d6544..ce281c3 100644 --- a/newCAM_emulation/NN_pred.py +++ b/newCAM_emulation/NN_pred.py @@ -36,14 +36,17 @@ ## load mean and std for normalization -fm = np.load("Demodata/mean_demo_sub.npz") -fs = np.load("Demodata/std_demo_sub.npz") +# fm = np.load("Demodata/mean_demo_sub.npz") +# fs = np.load("Demodata/std_demo_sub.npz") + +fm = nc.Dataset('/glade/derecho/scratch/yqsun/archive/meanstd/mid-top-CAM_mean_SCM.nc') +fs = nc.Dataset('/glade/derecho/scratch/yqsun/archive/meanstd/mid-top-CAM_std_SCM.nc') Um = fm["U"] Vm = fm["V"] Tm = fm["T"] DSEm = fm["DSE"] -NMm = fm["NM"] +NMm = fm["NMBV"] NETDTm = fm["NETDT"] Z3m = fm["Z3"] RHOIm = fm["RHOI"] @@ -57,7 +60,7 @@ Vs = fs["V"] Ts = fs["T"] DSEs = fs["DSE"] -NMs = fs["NM"] +NMs = fs["NMBV"] NETDTs = fs["NETDT"] Z3s = fs["Z3"] RHOIs = fs["RHOI"] @@ -76,13 +79,15 @@ optimizer = torch.optim.Adam(GWnet.parameters(), lr=learning_rate) -s_list = list(range(6, 7)) +s_list = list(range(1, 6)) +model_path = "/glade/derecho/scratch/jatkinson/archive/CAM_GW_zero-tend/atm/hist/CAM_GW_zero-tend.cam.h1.1979-01-01-00000.nc" data_vars = [] for iter in s_list: if iter > 0: - GWnet.load_state_dict(torch.load("./conv_torch.pth")) + # GWnet.load_state_dict(torch.load("./conv_torch.pth")) + GWnet = torch.jit.load(model_path) GWnet.eval() print("data loader iteration", iter) filename = "Demodata/newCAM_demo_sub_" + str(iter).zfill(1) + ".nc" diff --git a/newCAM_emulation/loaddata.py b/newCAM_emulation/loaddata.py index baccddf..b3c8198 100644 --- a/newCAM_emulation/loaddata.py +++ b/newCAM_emulation/loaddata.py @@ -22,17 +22,36 @@ def newnorm(var, varm, varstd): ------- numpy.ndarray: Normalized variable(s). """ + + print(f'Var shape - {var.shape}') + print(f'Varm shape - {varm.shape}') + print(f'Varstd shape - {varstd.shape}') + + var = var[:,0] + varm = np.array(varm) + varstd = np.array(varstd) + dim = varm.size - if dim > 1: - vara = var - varm[:, :] - varstdmax = varstd + + # if dim > 1: + # vara = var - varm[:, :] + # varstdmax = varstd.copy() + # varstdmax[varstd == 0.0] = 1.0 + # tmp = vara / varstdmax[:, :] + if dim == 1: + # var = var[:,0] + vara = var - varm[:] + varstdmax = varstd.copy() varstdmax[varstd == 0.0] = 1.0 - tmp = vara / varstdmax[:, :] + tmp = vara / varstdmax[:] else: tmp = (var - varm) / varstd return tmp + + + def data_loader(U, V, T, DSE, NM, NETDT, Z3, RHOI, PS, lat, lon, UTGWSPEC, VTGWSPEC): """ Load and preprocess input data for neural network training. @@ -56,26 +75,52 @@ def data_loader(U, V, T, DSE, NM, NETDT, Z3, RHOI, PS, lat, lon, UTGWSPEC, VTGWS ------- tuple: A tuple containing the input data and target data arrays. """ - Ncol = U.shape[1] + # print(f'U-shape in Dataloader{U.shape}') + # Ncol = U.shape[1] # Nlon = U.shape[2] # Ncol = Nlat*Nlon + # print (ilev) + # print(Ncol) + + x_train = np.zeros([dim_NN]) + y_train = np.zeros([dim_NNout]) + + # x_train[0:ilev, :] = U.reshape(ilev, Ncol) + # x_train[ilev : 2 * ilev, :] = V.reshape(ilev, Ncol) + # x_train[2 * ilev : 3 * ilev, :] = T.reshape(ilev, Ncol) + # x_train[3 * ilev : 4 * ilev, :] = DSE.reshape(ilev, Ncol) + # x_train[4 * ilev : 5 * ilev, :] = NM.reshape(ilev, Ncol) + # x_train[5 * ilev : 6 * ilev, :] = NETDT.reshape(ilev, Ncol) + # x_train[6 * ilev : 7 * ilev, :] = Z3.reshape(ilev, Ncol) + # x_train[7 * ilev : 8 * ilev + 1, :] = RHOI.reshape(ilev + 1, Ncol) + # x_train[8 * ilev + 1 : 8 * ilev + 2, :] = PS.reshape(1, Ncol) + # x_train[8 * ilev + 2 : 8 * ilev + 3, :] = lat.reshape(1, Ncol) + # x_train[8 * ilev + 3 : ilev * ilev + 4, :] = lon.reshape(1, Ncol) + + # y_train[0:ilev, :] = UTGWSPEC.reshape(ilev, Ncol) + # y_train[ilev : 2 * ilev, :] = VTGWSPEC.reshape(ilev, Ncol) + + # Populate the input data (x_train) + x_train[0:ilev] = U # U: shape (ilev,) + x_train[ilev:2 * ilev] = V + x_train[2 * ilev:3 * ilev] = T + x_train[3 * ilev:4 * ilev] = DSE + x_train[4 * ilev:5 * ilev] = NM + x_train[5 * ilev:6 * ilev] = NETDT + x_train[6 * ilev:7 * ilev] = Z3 + x_train[7 * ilev:8 * ilev] = RHOI[:ilev] # First 93 levels of RHOI + + # Handle the extra 94th level of RHOI + x_train[8 * ilev] = RHOI[ilev] # The 94th level of RHOI (extra level) + + # Add scalar values for PS, latitude, and longitude + x_train[8 * ilev + 1] = PS # Surface pressure (scalar) + x_train[8 * ilev + 2] = lat # Latitude (scalar) + x_train[8 * ilev + 3] = lon # Longitude (scalar) + + # Populate the target data (y_train) + y_train[0:ilev] = UTGWSPEC + y_train[ilev:2 * ilev] = VTGWSPEC - x_train = np.zeros([dim_NN, Ncol]) - y_train = np.zeros([dim_NNout, Ncol]) - - x_train[0:ilev, :] = U.reshape(ilev, Ncol) - x_train[ilev : 2 * ilev, :] = V.reshape(ilev, Ncol) - x_train[2 * ilev : 3 * ilev, :] = T.reshape(ilev, Ncol) - x_train[3 * ilev : 4 * ilev, :] = DSE.reshape(ilev, Ncol) - x_train[4 * ilev : 5 * ilev, :] = NM.reshape(ilev, Ncol) - x_train[5 * ilev : 6 * ilev, :] = NETDT.reshape(ilev, Ncol) - x_train[6 * ilev : 7 * ilev, :] = Z3.reshape(ilev, Ncol) - x_train[7 * ilev : 8 * ilev + 1, :] = RHOI.reshape(ilev + 1, Ncol) - x_train[8 * ilev + 1 : 8 * ilev + 2, :] = PS.reshape(1, Ncol) - x_train[8 * ilev + 2 : 8 * ilev + 3, :] = lat.reshape(1, Ncol) - x_train[8 * ilev + 3 : ilev * ilev + 4, :] = lon.reshape(1, Ncol) - - y_train[0:ilev, :] = UTGWSPEC.reshape(ilev, Ncol) - y_train[ilev : 2 * ilev, :] = VTGWSPEC.reshape(ilev, Ncol) return x_train, y_train From 4c76150bec8e0883ae51851c54bedd31bfcd187e Mon Sep 17 00:00:00 2001 From: Surbhi Goel Date: Thu, 24 Oct 2024 14:50:29 -0600 Subject: [PATCH 3/3] Updated NN_pred_new to read single column data --- newCAM_emulation/NN_pred_new.py | 122 ++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 newCAM_emulation/NN_pred_new.py diff --git a/newCAM_emulation/NN_pred_new.py b/newCAM_emulation/NN_pred_new.py new file mode 100644 index 0000000..158de8d --- /dev/null +++ b/newCAM_emulation/NN_pred_new.py @@ -0,0 +1,122 @@ +"""Prediction module for the neural network.""" + +import matplotlib.pyplot as plt +import Model +import netCDF4 as nc +import numpy as np +import torch +import torch.nn.functional as nnF +import torchvision +from loaddata import data_loader, newnorm +# from savedata import save_netcdf_file +from torch import nn +from torch.utils.data import DataLoader +from torchvision import datasets, transforms +from torchvision.utils import save_image + +""" +Determine if any GPUs are available +""" +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +print(device) + + +""" +Initialize Hyperparameters +""" +ilev = 93 +ilev_94 = 94 + +dim_NN = 8 * ilev + 4 +dim_NNout = 2 * ilev + +batch_size = 8 +learning_rate = 1e-4 +num_epochs = 1 + + + +""" +Initialize the network and the Adam optimizer +""" +GWnet = Model.FullyConnected() + +optimizer = torch.optim.Adam(GWnet.parameters(), lr=learning_rate) + +model_path = "/glade/derecho/scratch/yqsun/archive/meanstd/saved_convNN_SCM.pt" + +data_vars = [] + +# GWnet.load_state_dict(torch.load("./conv_torch.pth")) +GWnet = torch.jit.load(model_path) +GWnet.eval() + +# Load the .npz file +npz_file = np.load('training_data.npz') + +# List all variables stored in the .npz file +print("Variables in the .npz file:") +print(npz_file.files) + +# Access each variable +U = npz_file['U'] +V = npz_file['V'] +T = npz_file['T'] +Z3 = npz_file['Z3'] +RHOI = npz_file['RHOI'] +DSE = npz_file['DSE'] +NMBV = npz_file['NMBV'] +NETDT = npz_file['NETDT'] +PS = npz_file['PS'] +lat = npz_file['lat'] +lon = npz_file['lon'] +UTGWSPEC = npz_file['UTGWSPEC'] +VTGWSPEC = npz_file['VTGWSPEC'] + +# Print shapes or inspect the loaded variables +print(f"U shape: {U.shape}") +print(f"U shape: {U.size}") +print(f"V shape: {V.shape}") +print(f"T shape: {T.shape}") +print(f"Z3 shape: {Z3.shape}") +print(f"Lat shape: {lat.shape}") +print(f"Lon shape: {lon.shape}") +print(f"PS shape: {PS.shape}") +print(f"RHOI shape: {RHOI.shape}") + + +time_dim, ilev_dim, column_dim = U.shape +print(f"Data shape (time, levels, columns): {U.shape}") + +#Loop over timesteps and columns +for t in range(time_dim): # Iterate over time + for col in range(column_dim): # Iterate over columns (latitude/longitude) + # print(f"Processing time {t}, column {col}") + + # Extract all levels for the current time and column + U_sc = U[t, :, col] # Shape: (ilev,) + V_sc = V[t, :, col] # Shape: (ilev,) + T_sc = T[t, :, col] # Shape: (ilev,) + Z3_sc = Z3[t, :, col] # Shape: (ilev,) + RHOI_sc = RHOI[t, :, col] # Shape: (ilev_94,) for RHOI + DSE_sc = DSE[t, :, col] # Shape: (ilev,) + NMBV_sc = NMBV[t, :, col] # Shape: (ilev,) + NETDT_sc = NETDT[t, :, col] # Shape: (ilev,) + PS_sc = PS[t, col] # Shape: (1,) for pressure surface + + lat_sc = lat[col] # Latitude for the column + lon_sc = lon[col] # Longitude for the column + UTGWSPEC_sc = UTGWSPEC[t, :, col] # Shape: (ilev,) + VTGWSPEC_sc = VTGWSPEC[t, :, col] # Shape: (ilev,) + + ##Call the data_loader for the current slices + x_test, y_test = data_loader(U_sc, V_sc, T_sc, DSE_sc, NMBV_sc, NETDT_sc, Z3_sc, RHOI_sc, PS_sc, lat_sc, lon_sc, UTGWSPEC_sc, VTGWSPEC_sc) + # print(f'xtest: {x_test.shape}') + # print(f'ytest: {y_test.shape}') + + data = Model.myDataset(X=x_test, Y=y_test) + # test_loader = DataLoader(data, batch_size=len(data), shuffle=False) + # print(test_loader) + + + \ No newline at end of file