def open_data_with_lts_and_path():
ds = open_data('training').sel(time=slice(120,140))
p = open_data('pressure')
# make necessary computations
ds['p'] = p
ds['lat'] = ngaqua_y_to_lat(ds.y)
# compute LTS and Mid Trop moisture
ds['lts'] = lower_tropospheric_stability(ds.TABS, p, ds.SST, ds.Ps)
ds['path'] = midtropospheric_moisture(ds.QV, p, bottom=850, top=600)
return ds
def get_data():
variables = ['PW']
# open NN run
run = runs['debias']
nn = run.data_2d.rename({'NPNN': 'net_precip'})
# open NN run
run = runs['unstable']
unstable = run.data_2d.rename({'NPNN': 'net_precip'})
time = float(unstable.time[-1])
print(time)
# open microphysics
run = runs['micro']
micro = run.data_2d
micro['net_precip'] = micro.Prec - lhf_to_evap(micro.LHF)
# open NGAqua
ng = open_data('ngaqua_2d')
ng['net_precip'] = ng.Prec - lhf_to_evap(ng.LHF)
# make sure the x and y value agree
ng = ng.assign(x=nn.x, y=nn.y)
runs_at_time = {
'NG-Aqua': ng[variables].interp(time=time),
'NN-Lower': nn[variables].interp(time=time),
'Base': micro[variables].interp(time=time),
f'NN-All': unstable[variables].interp(time=time)
}
return xr.concat(list(runs_at_time.values()),
dim=list(runs_at_time.keys()))
def get_ng_and_semiprog():
# open model
model = common.get_model('NN-Lower')
# get data
ds = open_data('training').sel(time=slice(100,115))
def integrate_moist(src):
return (src * ds.layer_mass).sum('z')/1000
q2 = compute_apparent_source(ds.QT, 86400 * ds.FQT)
ng = xr.Dataset({
netprec_name: -integrate_moist(q2),
pw_name: integrate_moist(ds.QT)
})
# semiprognostic
predicted_srcs = model.predict(ds)
ds = xr.Dataset({
netprec_name: -integrate_moist(predicted_srcs['QT']),
pw_name: integrate_moist(ds.QT)
})
return ng, ds
def plot(data):
p = open_data('pressure')
fig, axs = plt.subplots(1, 4, sharex=True, sharey=True,
constrained_layout=True,
figsize=(common.textwidth, 2))
axs.shape = (1, -1)
abcd = 'abcd'
m = common.get_vmax(data)
kw = dict(levels=common.diverging_levels(25, 5), cmap='RdBu_r')
axs[0,0].invert_yaxis()
for k in range(4):
v = data.isel(step=k).squeeze()
# import pdb; pdb.set_trace()
im = axs[0,k].contourf(
v.time, p, v.T, **kw)
axs[0,k].set_title(f"{abcd[k]}) {get_title(k)}", loc='left')
# v.plot(col='step', x='time')
plt.colorbar(im, ax=axs, orientation='horizontal',
shrink=.3, aspect=2)
# axs[0,0].yaxis.set_major_locator(plt.MaxNLocator(4))
common.label_outer_axes(axs, "day", "p (mb)")
def get_data():
variables = ['PW', 'net_precip']
times = slice(100, 120)
# open NN run
run = runs['debias']
nn = run.data_2d.rename({'NPNN': 'net_precip'})
# open microphysics
run = runs['micro']
micro = run.data_2d
micro['net_precip'] = micro.Prec - lhf_to_evap(micro.LHF)
# open NGAqua
ng = open_data('ngaqua_2d')
ng['net_precip'] = ng.Prec - lhf_to_evap(ng.LHF)
# make sure the x and y value agree
ng = ng.assign(x=nn.x, y=nn.y)
plotme = xr.concat(
[ng[variables].interp(time=nn.time), nn[variables], micro[variables]],
dim=['NG-Aqua', 'NN-Lower', 'Base'])
return plotme.sel(time=times).mean('x')
def get_data():
ds = open_data('training')
time = ds.time[0] + 2
# compute expected_precip
model = torch.load('../../nn/NNLower/5.pkl')
data_at_time = ds.sel(time=time).load()
neural_net_srcs = model.call_with_xr(data_at_time)
semi_prognostic = -integrate_q2(neural_net_srcs['QT'], ds.layer_mass)
# net precip from model
net_precip = runs['debias'].data_2d.NPNN
# net precip micro
data_2d = runs['micro'].data_2d
micro = data_2d.Prec - lhf_to_evap(data_2d.LHF)
evap = lhf_to_evap(ds.LHF)
net_precip_truth = ds.Prec - evap
return xr.Dataset({
'NG-Aqua': net_precip_truth.interp(time=time),
'SemiProg': semi_prognostic.interp(time=time),
'DEBIAS': net_precip.interp(time=time),
'MICRO': micro.interp(time=time),
}).compute()
def get_data(start_time=100, end_time=120):
ng = open_data('ngaqua_2d')
debias = runs['debias'].data_2d
time_slice = slice(start_time, end_time)
#
ng, semiprog = get_ng_and_semiprog()
# select the data
debias = debias.sel(time=time_slice)
debias = xr.Dataset({
pw_name: debias.PW,
netprec_name: debias.NPNN
})
# merge the data into dataframe
ng = ng.to_dataframe().assign(run='NG-Aqua')
debias = debias.to_dataframe().assign(run='NN-Lower')
semiprog = semiprog.to_dataframe().assign(run='NN-Lower-semi')
df = pd.concat([ng, debias, semiprog])
return df.reset_index().set_index(['run', 'time', 'x', 'y'])
def get_base_state_from_training_data():
ds = open_data('training')
mean = ds.isel(y=32, time=slice(0, 10)).mean(['x', 'time'])
base_state = {}
for key in ['SLI', 'QT', 'SOLIN', 'SST']:
base_state[key] = xarray2torch(mean[key])
return base_state
def get_error(model):
from src.data import open_data
data = open_data("training")
data = data.isel(time=slice(0, 100)).compute()
srcs = model.predict(data)
q2 = compute_apparent_source(data.QT, data.FQT * 86400)
q1 = compute_apparent_source(data.SLI, data.FSLI * 86400)
return xr.Dataset({
'QT': q2 - srcs.QT,
'SLI': q1 - srcs.SLI
}).dropna('time')
def get_ng():
ds = open_data('training').sel(time=slice(100, 115), y=slice(4.5e6, 5.5e6))
def integrate_moist(src):
return (src * ds.layer_mass).sum('z') / 1000
q2 = compute_apparent_source(ds.QT, 86400 * ds.FQT)
ng = xr.Dataset({
netprec_name: -integrate_moist(q2),
pw_name: integrate_moist(ds.QT)
})
return ng.to_dataframe()
def get_precipitation():
net_precip_nn = runs['debias'].data_2d.NPNN
ds = open_data('training')
evap = lhf_to_evap(ds.LHF)
net_precip_truth = ds.Prec - evap
dsm = runs['micro'].data_2d
net_precip_control = dsm.Prec - lhf_to_evap(dsm.LHF)
time = net_precip_nn.time
return (net_precip_truth.interp(time=time), net_precip_nn,
net_precip_control)
def get_data():
ds = open_data('training')
run = runs['debias']
nn = run.data_3d
common_variables = list(set(nn.data_vars) & set(ds.data_vars))
plotme = xr.concat([avg(nn[common_variables]),
avg(ds[common_variables])],
dim=['NN', 'NG-Aqua'])
length_domain = 160e3 * len(ds.x)
plotme['stream_function'] = (plotme.V *
ds.layer_mass[0]).cumsum('z') * length_domain
plotme = plotme.assign(p=plotme.p[0]).swap_dims({'z': 'p'})
return plotme
def get_data(model="../../nn/277/1.pkl", **kwargs):
# open model and training data
model = torch.load(model)
ds = open_data('training')
# select x=0, y=32
index = {'x': 0, 'y': 32}
index = {key: slice(val, val + 1) for key, val in index.items()}
location = ds.isel(**index)
# run the single column models
merged = predict_for_each_time(model, location, **kwargs)
# compute Q2 for comparison
location_subset = location.sel(time=merged.time)
q2 = compute_apparent_source(location_subset.QT,
location_subset.FQT * 86400)
plotme = xr.concat(
[q2.assign_coords(step='Truth'), merged.FQTNN], dim='step')
return plotme
def get_data():
# open model
model = torch.load(args.model)
# get data
ds = open_data('training').sel(time=slice(100, 115), y=slice(4.5e6, 5.5e6))
usrf = np.sqrt(ds.U.isel(z=0)**2 + ds.V.isel(z=0)**2)
ds['surface_wind_speed'] = usrf
def integrate_moist(src):
return (src * ds.layer_mass).sum('z') / 1000
# semiprognostic
predicted_srcs = model.predict(ds)
ds = xr.Dataset({
netprec_name: -integrate_moist(predicted_srcs['QT']),
pw_name: integrate_moist(ds.QT)
})
return ds.to_dataframe()
def get_data():
model = common.get_model('NN-Lower')
ds = open_data('training')\
.chunk({'time': 1})\
.pipe(assign_apparent_sources)
outputs = map_dataset(ds, model.call_with_xr, 'time')
for key in outputs:
ds['F' + key + 'NN'] = outputs[key]
output = xr.Dataset()
output['net_moist_nn'] = integrate_q2(ds['FQTNN'], ds.layer_mass)
output['net_heat_nn'] = integrate_q1(ds['FSLINN'], ds.layer_mass)
output['net_moist'] = integrate_q2(ds['Q2'], ds.layer_mass)
output['net_heat'] = integrate_q1(ds['Q1'], ds.layer_mass)
output['Q1'] = ds['Q1']
output['Q2'] = ds['Q2']
output['Q1nn'] = ds['FSLINN']
output['Q2nn'] = ds['FQTNN']
return output
def get_data():
ds = open_data('training')
variables = ['QT', 'QP', 'SLI', 'U', 'V']
return ds[variables].apply(compute_acf)
def get_wave_from_training_data(src=None):
ds = open_data('training')
mean = ds.isel(y=32, time=slice(0, 10)).mean(['x', 'time'])
return wave_from_xarray(mean, src)
def get_plot_manager(model):
from src.data import open_data
dataset = open_data('training')
return TrainingPlotManager(ex, model, dataset)
def get_time_series_data():
ds = open_data('training')
index = {'x': 0, 'y': 32, 'z': 20}
location = ds.isel(**index)
return location.to_dataframe()
def to_pressure_dim(ds):
from src.data import open_data
return ds.assign_coords(p=open_data('pressure')).swap_dims({'z': 'p'})
from jacobian import *
from src.data import open_data
import common
import pandas as pd
# bootstrap sample size
n = 20
hatch_threshold = 10
# compute jacobians
training = open_data('training')
training['region'] = common.get_regions(training.y)
tropics = training.isel(y=slice(30,34)).load()
tropics['time_of_day'] = tropics.time % 1
p = tropics.p[0].values
model = common.get_model('NN-All')
samples = list(bootstrap_samples(tropics, n))
jacobians = [get_jacobian(model, sample) for sample in samples]
# make plot
fig, axs = plt.subplots(
4, 5, figsize=(common.textwidth, common.textwidth-2), sharex=True, sharey=True)
plt.rcParams['hatch.color'] = '0.5'
axs[0,0].invert_yaxis()
axs[0,0].invert_xaxis()
norm = SymLogNorm(1, 2, vmin=-1e5, vmax=1e5)
for ax, jac in zip(axs.flat, jacobians):
qt_qt = jac['QT']['QT'].detach().numpy()