Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added output generation code and data files #29

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Added output generation code and data files #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a32b6e1
Added output generation code and data files
Oct 14, 2024
ffc6155
Generated new normalised input data using single column mean/std, upd…
Oct 24, 2024
4c76150
Updated NN_pred_new to read single column data
Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 59 additions & 11 deletions newCAM_emulation/NN_pred.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -24,6 +25,8 @@
Initialize Hyperparameters
"""
ilev = 93
ilev_94 = 94

dim_NN = 8 * ilev + 4
dim_NNout = 2 * ilev

Expand All @@ -33,14 +36,17 @@


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

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"]
Expand All @@ -54,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"]
Expand All @@ -73,14 +79,18 @@
optimizer = torch.optim.Adam(GWnet.parameters(), lr=learning_rate)


s_list = list(range(5, 6))
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/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, :])
Expand Down Expand Up @@ -116,10 +126,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))
Expand All @@ -134,13 +144,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)
Expand All @@ -155,4 +184,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,
)
122 changes: 122 additions & 0 deletions newCAM_emulation/NN_pred_new.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)



89 changes: 67 additions & 22 deletions newCAM_emulation/loaddata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Loading
Loading