# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] Collapsed="false"
# # Creating an unbiased train/valid split
#
# 1. How biased is a simple <2016/2016 split? Answer: There doesn't seem to be a large bias. This means I can simply use this split to test for overfitting.

# + Collapsed="false"
# %load_ext autoreload
# %autoreload 2

# + Collapsed="false" jupyter={"outputs_hidden": true}
from src.train_nn import *
import matplotlib.pyplot as plt

# + Collapsed="false" jupyter={"outputs_hidden": true}
limit_mem()

# + Collapsed="false"
datadir = '/data/weather-benchmark/5.625deg/'

# + Collapsed="false"
z = xr.open_mfdataset(f'{datadir}geopotential_500/*.nc', combine='by_coords')
t = xr.open_mfdataset(f'{datadir}temperature_850/*.nc', combine='by_coords')
ds = xr.merge([z, t])

# + [markdown] Collapsed="false"
# ## How biased is a simple <2016/2016 split?

# + Collapsed="false"
ds_train = ds.sel(time=slice('1979', '2015'))
ds_valid = ds.sel(time=slice('2016', '2016'))
ds_test = ds.sel(time=slice('2017', '2018'))

# + Collapsed="false"
dic = {'z': None, 't': None}

# + Collapsed="false"
batch_size = 32
lead_time = 3*24

# + Collapsed="false"
dummy = DataGenerator(ds_train.isel(time=slice(0, None, 5*12)), dic, lead_time, batch_size=batch_size, shuffle=False)
dg_train = DataGenerator(ds_train, dic, lead_time, batch_size=batch_size, mean=dummy.mean,
                         std=dummy.std, shuffle=False)
dg_valid = DataGenerator(ds_valid, dic, lead_time, batch_size=batch_size, mean=dg_train.mean,
                         std=dg_train.std, shuffle=False)
dg_test =  DataGenerator(ds_test, dic, lead_time, batch_size=batch_size, mean=dg_train.mean,
                         std=dg_train.std, shuffle=False)
dg_fake =  DataGenerator(ds_test.sel(time='2018'), dic, lead_time, batch_size=batch_size, mean=dg_train.mean,
                         std=dg_train.std)

# + Collapsed="false"
model = build_cnn([64, 64, 64, 64, 2], [5, 5, 5, 5, 5], input_shape=(32, 64, 2), activation='elu')
model.compile(keras.optimizers.Adam(), 'mse')

# + Collapsed="false"
model.fit_generator(dg_fake, epochs=1)

# + Collapsed="false"
model.evaluate_generator(dg_train, verbose=1)

# + Collapsed="false"
model.evaluate_generator(dg_valid, verbose=1)

# + Collapsed="false"
model.evaluate_generator(dg_test, verbose=1)

# + Collapsed="false"
# Create predictions from an untrained model
preds_train = create_predictions(model, dg_train)

# + Collapsed="false"
preds_valid = create_predictions(model, dg_valid)
preds_test = create_predictions(model, dg_test)

# + Collapsed="false"
rmse_train = compute_weighted_rmse(preds_train.isel(time=slice(0, None, 36)), ds).load(); rmse_train

# + Collapsed="false"
rmse_valid = compute_weighted_rmse(preds_valid, ds).load(); rmse_valid

# + Collapsed="false"
rmse_test = compute_weighted_rmse(preds_test, ds).load(); rmse_test

# + Collapsed="false"
len(dg_train), len(dg_valid)

# + Collapsed="false"
mses_train = []
for i in range(len(dg_valid))[::-1]:
    x, y = dg_train[-i]
    p = model.predict_on_batch(x)
    mses_train.append(((y - p)**2).mean())

# + Collapsed="false"
mses_valid = []
for i in range(len(dg_valid)):
    x, y = dg_valid[i]
    p = model.predict_on_batch(x)
    mses_valid.append(((y - p)**2).mean())

# + Collapsed="false"
plt.plot(mses_train)
plt.plot(mses_valid)

# + Collapsed="false"


StarcoderdataJupyterScriptsDedupFiltered