<jupyter_start><jupyter_text>Max temperature<jupyter_code>erai_deg1_arr = erai_deg1[['MAXT']].to_array('chan').transpose('Time','lat','lon','chan')
ras_deg14_arr = ras_deg14[['MAXT']].to_array('chan').transpose('Time','lat','lon','chan')
lr_train, lr_test, hr_train, hr_test = get_train_test_splits(erai_deg1_arr, ras_deg14_arr, southeast_us, scale=4)
data_fold = maxt.preprocess_fold_maxt(((lr_train, hr_train),(lr_test, hr_test)))
train_lo, train_hi = data_fold.train
test_lo, test_hi = data_fold.test
N_train, N_test = train_lo.Time.size, test_lo.Time.size
(wt, ht), (wt_hi, ht_hi) = train_lo.shape[1:3], train_hi.shape[1:3]
monthly_means_lo, monthly_means_hi = data_fold.monthly_means
test_ds = data_fold.test_dataset(batch_size=100, buffer_size=200,
                                        map_fn_lo=upsample(wt_hi, ht_hi, tf.image.ResizeMethod.NEAREST_NEIGHBOR),
                                        mode='test')
model_joint = build_jflvm((None,wt_hi,ht_hi,1), scale=4, layers=3, depth=8,
                          min_filters=32, max_filters=256,
                          dnet_layers=3, dnet_filters=64)
model_joint.load('data/saved_models/glow-jflvm/maxt-seus/ckpt', 8)
model_vars = model_joint.trainable_variables
param_count = sum([tf.math.reduce_prod(var.shape) for var in model_vars]).numpy()
print(f'Loaded JFLVM with {param_count} trainable parameters')
x, y = next(test_ds.__iter__())
y_pred = model_joint.predict_y(x)
(x_, p_x), (y_, p_y) = model_joint.sample(n=1000, return_log_prob=True)
(z_x, p_zx), (z_y, p_zy) = model_joint.encode_x(x, return_log_prob=True), model_joint.encode_y(y, return_log_prob=True)
plt.figure(figsize=(9,5))
sns.jointplot(x=p_x, y=p_y, kind='hex').set_axis_labels('$\\log p(x)$', '$\\log p(y)$')
p_xs = []
p_ys = []
for x, y in tqdm(test_ds, total=hr_test.Time.size//100+1):
    y_, _, p_x = model_joint.predict_y(x, return_log_prob=True)
    x_, _, p_y = model_joint.predict_x(y, return_log_prob=True)
    p_xs.append(p_x)
    p_ys.append(p_y)
p_xs = tf.concat(p_xs, axis=0)
p_ys = tf.concat(p_ys, axis=0)
plt.figure(figsize=(9,5))
sns.jointplot(x=p_xs, y=p_ys, kind='hex').set_axis_labels('$\\log p(x)$', '$\\log p(y)$')
ax = plt.gca()
x_samples, y_samples, (p_sx, p_sy, p_x, p_y) = sample_prediction(x[:1], model_joint, sigma=1.0, n=1000)
print(p_x, p_y)
sns.jointplot(x=p_sx, y=p_sy, kind='hex').set_axis_labels('$\\log p(x)$', '$\\log p(y)$')
z_x, log_prob_x = model_joint.encode_x(x, return_log_prob=True)
z_y, log_prob_y = model_joint.encode_y(y, return_log_prob=True)
z_yp, log_prob_yp = model_joint.encode_y(y_samples, return_log_prob=True)
plt.figure(figsize=(3*8,6))
plt.subplot(1,3,1)
plt.imshow(tf.squeeze(tf.image.resize(x[0], (wt,ht), method='nearest')).numpy(), origin='lower')
plt.subplot(1,3,2)
plt.imshow(tf.squeeze(y[0]).numpy(), origin='lower')
plt.subplot(1,3,3)
plt.imshow(tf.squeeze(y_pred[0]).numpy(), origin='lower')
zs = slerp(z_y[:1], z_y[9:10], steps=5)
x_interp, p_xi = model_joint.decode_x(zs, return_log_prob=True)
y_interp, p_yi = model_joint.decode_y(zs, return_log_prob=True)
x_interp = tf.image.resize(x_interp, (wt, ht), method='nearest')
fig, axs, plot_fn = image_map_factory(2, 5, figsize=(5,3))
plot_fn(axs[0,0], y_interp[0].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap='viridis')
plot_fn(axs[0,1], y_interp[1].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap='viridis')
plot_fn(axs[0,2], y_interp[2].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap='viridis')
plot_fn(axs[0,3], y_interp[3].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap='viridis')
plot_fn(axs[0,4], y_interp[4].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap='viridis')
plot_fn(axs[1,0], x_interp[0].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap='viridis')
plot_fn(axs[1,1], x_interp[1].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap='viridis')
plot_fn(axs[1,2], x_interp[2].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap='viridis')
plot_fn(axs[1,3], x_interp[3].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap='viridis')
cs = plot_fn(axs[1,4], x_interp[4].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap='viridis')
fig.colorbar(cs, ax=axs.ravel().tolist(), orientation='vertical', shrink=0.6, pad=0.01).set_label('daily max temperature anomalies (K)')
plt.show()<jupyter_output><empty_output><jupyter_text>Precipitation<jupyter_code>erai_deg1_arr = prcp.preprocess_dataset(erai_deg1[['PRCP']], erai_deg1['PRCP'].shape)
erai_deg1_arr = erai_deg1_arr.to_array('chan').transpose('Time','lat','lon','chan')
ras_deg14_arr = prcp.preprocess_dataset(ras_deg14[['PRCP']], ras_deg14['PRCP'].shape)
ras_deg14_arr = ras_deg14_arr.to_array('chan').transpose('Time','lat','lon','chan')
lr_train, lr_test, hr_train, hr_test = get_train_test_splits(erai_deg1_arr, ras_deg14_arr, southeast_us, scale=4)
data_fold = prcp.preprocess_fold_prcp(((lr_train, hr_train),(lr_test, hr_test)))
train_lo_prcp, train_hi_prcp = data_fold.train
test_lo_prcp, test_hi_prcp = data_fold.test
N_train, N_test = train_lo_prcp.Time.size, test_lo_prcp.Time.size
(wt, ht), (wt_hi, ht_hi) = train_lo_prcp.shape[1:3], train_hi_prcp.shape[1:3]
test_ds_prcp = data_fold.test_dataset(batch_size=1, buffer_size=1000,
                                      map_fn_lo=upsample(wt_hi, ht_hi, tf.image.ResizeMethod.NEAREST_NEIGHBOR),
                                      mode='test')
model_joint_prcp = build_jflvm((None,wt_hi,ht_hi,1), scale=4, layers=3, depth=4,
                               min_filters=32, max_filters=256,
                               dnet_layers=3, dnet_filters=64)
model_joint_prcp.load('data/saved_models/glow-jflvm/prcp-seus/ckpt', 10)
model_vars = model_joint_prcp.trainable_variables
param_count = sum([tf.math.reduce_prod(var.shape) for var in model_vars]).numpy()
print(f'Loaded JFLVM with {param_count} trainable parameters')
for var in model_vars:
    print(f'{var.name}: {var.shape} {tf.math.reduce_mean(var):.3f} +/- {tf.math.reduce_std(var):.2f}')
t = 229 #np.random.choice(range(hr_test.Time.size))
print(t, hr_test.Time[t].values)
x, y = next(test_ds_prcp.skip(t).__iter__())
y_pred = model_joint_prcp.predict_y(x[:1])
y_samples,_ = model_joint_prcp.sample_predict_y(x[:1], temperature=0.7, n=10)
y_samples = y_samples**3
y_samples = tf.where(y_samples > 5.0, y_samples, 0.0)
x = tf.where(x > 1.0, x**3, 0.0)
y = tf.where(y > 1.0, y**3, 0.0)
y_pred = tf.where(y_pred > 1.0, y_pred**3, 0.0)
x = tf.image.resize(x, (wt,ht), method='nearest')
fig, axs, plot_fn = image_map_factory(2, 4, figsize=(6,3))
pmap = prcp_cmap()
axs[0,0].set_visible(False)
axs[0,-1].set_visible(False)
plot_fn(axs[0,1], x[0].numpy().squeeze(), lr_test.lat, lr_test.lon, title='ERA-I', cmap=pmap)
plot_fn(axs[0,2], y[0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='WRF-4', cmap=pmap)
plot_fn(axs[1,0], y_pred[0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Predicted', cmap=pmap)
plot_fn(axs[1,1], y_samples[0,0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 1', cmap=pmap)
plot_fn(axs[1,2], y_samples[0,1].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 2', cmap=pmap)
cs = plot_fn(axs[1,3], y_samples[0,2].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 3', cmap=pmap)
fig.colorbar(cs, ax=axs.ravel().tolist(), orientation='vertical', shrink=0.75, pad=0.01).set_label('daily precipitation (mm)')
t = 208
print(hr_test.Time[t].values)
x, y = next(test_ds_prcp.skip(t).unbatch().batch(10).__iter__())
z_x, log_prob_x = model_joint_prcp.encode_x(x, return_log_prob=True)
z_y, log_prob_y = model_joint_prcp.encode_y(y, return_log_prob=True)
zs = slerp(z_y[:1], z_y[9:10], steps=5)
x_interp, p_xi = model_joint_prcp.decode_x(zs, return_log_prob=True)
y_interp, p_yi = model_joint_prcp.decode_y(zs, return_log_prob=True)
x_interp = tf.where(x_interp > 1.0, x_interp**3, 0.0)
y_interp = tf.where(y_interp > 1.0, y_interp**3, 0.0)
x_interp = tf.image.resize(x_interp, (wt,ht), method='nearest')
fig, axs, plot_fn = image_map_factory(2, 5, figsize=(5,3))
plot_fn(axs[0,0], y_interp[0].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap=pmap)
plot_fn(axs[0,1], y_interp[1].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap=pmap)
plot_fn(axs[0,2], y_interp[2].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap=pmap)
plot_fn(axs[0,3], y_interp[3].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap=pmap)
plot_fn(axs[0,4], y_interp[4].numpy().squeeze(), hr_test.lat, hr_test.lon, cmap=pmap)
plot_fn(axs[1,0], x_interp[0].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap=pmap)
plot_fn(axs[1,1], x_interp[1].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap=pmap)
plot_fn(axs[1,2], x_interp[2].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap=pmap)
plot_fn(axs[1,3], x_interp[3].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap=pmap)
cs = plot_fn(axs[1,4], x_interp[4].numpy().squeeze(), lr_test.lat, lr_test.lon, cmap=pmap)
fig.colorbar(cs, ax=axs.ravel().tolist(), orientation='vertical', shrink=0.6, pad=0.01).set_label('daily precipitation (mm)')
plt.show()<jupyter_output><empty_output><jupyter_text>Anomaly likelihood analysis<jupyter_code>upper_quartile = test_hi.chunk({'Time': -1}).quantile(0.75, dim='Time')
lower_quartile = test_hi.chunk({'Time': -1}).quantile(0.25, dim='Time')
out_iqr_counts = xr.ufuncs.logical_or(test_hi > upper_quartile, test_hi < lower_quartile).sum(dim=['lat','lon']).compute()
inds = out_iqr_counts.squeeze().argsort(axis=0)
lr_sorted = test_lo[inds.values]
lr_vals = tf.image.resize(lr_sorted.values.astype(np.float32), (wt_hi,ht_hi), method='nearest')
zp, p_x_maxt = model_joint.encode_x(lr_vals, return_log_prob=True)
_, p_y_maxt = model_joint.decode_y(zp, return_log_prob=True)
prcp_data = xr.where(test_hi_prcp > 1.0, test_hi_prcp**3, 0.0)
upper_quartile = prcp_data.chunk({'Time': -1}).quantile(0.75, dim='Time')
lower_quartile = prcp_data.chunk({'Time': -1}).quantile(0.25, dim='Time')
out_iqr_counts_prcp = xr.ufuncs.logical_or(prcp_data > upper_quartile, prcp_data < lower_quartile).sum(dim=['lat','lon']).compute()
inds_prcp = out_iqr_counts_prcp.squeeze().argsort(axis=0)
lr_sorted = test_lo_prcp[inds_prcp.values]
lr_vals = tf.image.resize(lr_sorted.values.astype(np.float32), (wt_hi,ht_hi), method='nearest')
zp, p_x_prcp = model_joint.encode_x(lr_vals, return_log_prob=True)
_, p_y_prcp = model_joint.decode_y(zp, return_log_prob=True)
anomaly_counts_maxt = out_iqr_counts[inds.values].squeeze()
anomaly_counts_prcp = out_iqr_counts_prcp[inds_prcp.values].squeeze()
fig, axs = plt.subplots(1,2,figsize=(2*6,4))
fig.add_subplot(111, frameon=False)
# hide tick and tick label of the big axis
plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)
plt.xlabel("Spatial anomaly count")
sns.scatterplot(x=anomaly_counts_maxt.values, y=p_y_maxt, ax=axs[0], alpha=0.75)
axs[0].set(ylabel='Log likelihood', title='Max temperature')
axs[1].set(title='Precipitation')
sns.scatterplot(x=anomaly_counts_prcp.values, y=p_y_prcp, ax=axs[1], alpha=0.75)
plt.tight_layout()<jupyter_output><empty_output><jupyter_text>Preciptation WRF-8 ($\frac{1}{8}^{\circ}$)<jupyter_code>erai_deg1_pnw = pacific_nw(erai_deg1[['PRCP']]) 
ras_deg18_pnw = pacific_nw(ras_deg18[['PRCP']], scale_factor=8) 
erai_deg1_pnw = prcp.preprocess_dataset(erai_deg1_pnw, erai_deg1_pnw['PRCP'].shape)
erai_deg1_arr = erai_deg1_pnw.to_array('chan').transpose('Time','lat','lon','chan')
ras_deg18_arr = prcp.preprocess_dataset(ras_deg18_pnw, ras_deg18_pnw['PRCP'].shape)
ras_deg18_arr = ras_deg18_arr.to_array('chan').transpose('Time','lat','lon','chan')
lr_train = erai_deg1_arr.isel(Time=slice(0,erai_deg1_arr.Time.size-2*365+146))
lr_test = erai_deg1_arr.isel(Time=slice(erai_deg1_arr.Time.size-2*365+146, erai_deg1_arr.Time.size+1))
hr_train = ras_deg18_arr.isel(Time=slice(0,ras_deg18_arr.Time.size-2*365+146))
hr_test = ras_deg18_arr.isel(Time=slice(ras_deg18_arr.Time.size-2*365+146, ras_deg18_arr.Time.size+1))
data_fold = prcp.preprocess_fold_prcp(((lr_train, hr_train),(lr_test, hr_test)))
train_lo_prcp, train_hi_prcp = data_fold.train
test_lo_prcp, test_hi_prcp = data_fold.test
N_train, N_test = train_lo_prcp.Time.size, test_lo_prcp.Time.size
(wt, ht), (wt_hi, ht_hi) = train_lo_prcp.shape[1:3], train_hi_prcp.shape[1:3]
test_ds_prcp = data_fold.test_dataset(batch_size=1, buffer_size=100,
                                      map_fn_lo=upsample(wt_hi, ht_hi, tf.image.ResizeMethod.NEAREST_NEIGHBOR),
                                      mode='test')
model_joint_prcp = build_jflvm((None,wt_hi,ht_hi,1), scale=8, layers=4, depth=4,
                               min_filters=32, max_filters=256,
                               dnet_layers=3, dnet_filters=64)
model_joint_prcp.load('data/saved_models/glow-jflvm/prcp-pnw/8x/ckpt', 5)
t = 573 #np.random.choice(range(147, hr_test.Time.size))
print(t, hr_test.Time[t].values)
x, y = next(test_ds_prcp.skip(t).__iter__())
y_pred = model_joint_prcp.predict_y(x[:1])
y_samples,_ = model_joint_prcp.sample_predict_y(x[:1], temperature=0.7, n=10)
y_samples = y_samples**3
y_samples = tf.where(y_samples > 2.0, y_samples, 0.0)
x = tf.where(x > 1.0, x**3, 0.0)
y = tf.where(y > 1.0, y**3, 0.0)
y_pred = y_pred**3
y_pred = tf.where(y_pred > 2.0, y_pred, 0.0)
x = tf.image.resize(x, (wt,ht), method='nearest')
fig, axs, plot_fn = image_map_factory(2, 3, figsize=(6,3))
pmap = prcp_cmap()
plot_fn(axs[0,0], x[0].numpy().squeeze(), lr_test.lat, lr_test.lon, title='ERA-I', cmap=pmap)
plot_fn(axs[0,1], y[0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='WRF-8', cmap=pmap)
plot_fn(axs[0,2], y_pred[0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Predicted', cmap=pmap)
plot_fn(axs[1,0], y_samples[0,0].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 1', cmap=pmap)
plot_fn(axs[1,1], y_samples[0,1].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 2', cmap=pmap)
cs = plot_fn(axs[1,2], y_samples[0,2].numpy().squeeze(), hr_test.lat, hr_test.lon, title='Sample 3', cmap=pmap)
fig.colorbar(cs, ax=axs.ravel().tolist(), orientation='vertical', shrink=0.75, pad=0.01).set_label('daily precipitation (mm)')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18481824661439772, 'fraction_numerical': 0.03492516037063435, 'mean_word_length': 6.584324324324324, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '52642', 'license': 'MIT', 'n_tokens_mistral': 6450, 'n_tokens_neox': 6235, 'n_words': 796, 'path': 'qualitative-analysis.ipynb', 'repo_name': 'bgroenks96/generative-downscaling', 'size': '14030'}