def main():
parser = argparse.ArgumentParser()
parser.add_argument('--inputs_A', nargs='+', help='The first input.')
parser.add_argument('--inputs_B', nargs='+', help='The second input.')
parser.add_argument('--var', default='temp', help='The variable to plot.')
parser.add_argument('--model', default='ocean',
help='Whether this is an ocean or atmosphere plot.')
parser.add_argument('--output_file', default='spatial_difference.png',
help='File Name of the plot.')
parser.add_argument('--title', default=None,
help='Title of the plot.')
args = parser.parse_args()
with nc.Dataset(args.inputs_A[0]) as f:
if args.model == 'ocean':
lats = f.variables['geolat_t'][:]
lons = f.variables['geolon_t'][:]
else:
lons, lats = np.meshgrid(f.variables['lon'][:], f.variables['lat'][:])
dims = len(f.variables[args.var].shape)
z_dim = None
if dims == 4:
if args.model == 'ocean':
z_dim='st_ocean'
else:
z_dim='z16_p_level'
_, tmp_a = run_ncra(args.inputs_A, [args.var], z_dim=z_dim)
_, tmp_b = run_ncra(args.inputs_B, [args.var], z_dim=z_dim)
fa = nc.Dataset(tmp_a)
fb = nc.Dataset(tmp_b)
diff = fa.variables[args.var][:] - fb.variables[args.var][:]
fa.close()
fb.close()
os.remove(tmp_a)
os.remove(tmp_b)
if len(diff.shape) == 3:
diff = diff[0, :, :]
elif len(diff.shape) == 4:
diff = diff[0, 0, :, :]
if not args.title:
args.title = '{}: A - B'.format(args.var)
do_global_plot(diff, args.title, lats, lons, args.output_file)
return 0
def visit_data_file(args):
"""
Visit a data file and collect/calculate the following:
1. AMOC maximum timeseries
2. AMOC mean timeseries
3. AMOC index timeseries (see (1) for definition)
4. AMOC psi timeseries as a numpy array, shape (t, depth, lat)
5. Surface plot of difference between SST and NH spatial mean SST,
shape (t, lat, lon)
"""
file, grid_def_file, do_depth_correlation_plot, \
do_surface_correlation_plot = args
# Try with annual data.
use_annual = False
if use_annual:
_, tmp_file = run_ncra(file, ['geolon_t', 'geolat_t', 'time', 'temp', 'ty_trans'])
f = nc.Dataset(tmp_file)
else:
f = nc.Dataset(file)
lons = f.variables['geolon_t']
lats = f.variables['geolat_t']
time_var = f.variables['time']
temp_var = f.variables['temp']
ty_trans = f.variables['ty_trans']
depths = np.cumsum(f.variables['st_ocean'][:])
t_dim = len(f.dimensions['time'])
z_dim = len(f.dimensions['st_ocean'])
x_dim = len(f.dimensions['xt_ocean'])
y_dim = len(f.dimensions['yt_ocean'])
gf = nc.Dataset(grid_def_file)
areas = gf.variables['area_t'][:]
gf.close()
nh_mask = get_nh_mask(temp_var[0, 0, :, :].mask, lats)
atlantic_mask = get_atlantic_mask(ty_trans[0, :, :].mask, lons, lats)
lat_start, lat_end, lon_start, \
lon_end = get_indices_for_amoc_idx_region(lons, lats)
amoc_max_ts = []
amoc_mean_ts = []
amoc_idx_ts = []
amoc_psi_ts = np.ma.zeros((t_dim, z_dim, y_dim))
sst_nh_diff_ts = np.ma.zeros((t_dim, y_dim, x_dim))
for t in range(time_var.shape[0]):
# Get surface temp spatial mean in the NH
nh_sst_mean = np.ma.average(np.ma.masked_array(temp_var[t, 0, :, :],
mask=nh_mask), weights=areas)
# Get AMOC max and mean.
amoc_psi = calc_atlantic_moc(ty_trans[t, :, :, :], lons, lats)
amoc_idx_sst = np.average(temp_var[t, 0, lat_start:lat_end, lon_start:lon_end],
weights=areas[lat_start:lat_end, lon_start:lon_end])
# Calculate the AMOC index
amoc_idx_ts.append(amoc_idx_sst - nh_sst_mean)
amoc_max_ts.append(max_within_region(amoc_psi, 500.0, 35.0, depths, lats))
amoc_mean_ts.append(np.mean(amoc_psi))
# Get AMOC psi timeseries
if do_depth_correlation_plot:
amoc_psi_ts[t, :, :] = amoc_psi[:, :]
# Get the surface difference between temp and NH SST mean
if do_surface_correlation_plot:
sst_nh_diff_ts[t, :, :] = \
np.ma.masked_array(temp_var[t, 0, :, :] - nh_sst_mean,
mask=atlantic_mask[0, :, :])
# Add time dim to pandas timeseries
periods = time_dim_to_pandas_periods(time_var)
if use_annual:
periods = [periods[len(periods) / 2]]
amoc_idx_ts = pd.Series(amoc_idx_ts, periods)
amoc_max_ts = pd.Series(amoc_max_ts, periods)
amoc_mean_ts = pd.Series(amoc_mean_ts, periods)
f.close()
if use_annual:
os.remove(tmp_file)
ret = [None, None, None, None, None]
ret[0] = amoc_idx_ts
ret[1] = amoc_max_ts
ret[2] = amoc_mean_ts
if do_depth_correlation_plot:
ret[3] = amoc_psi_ts
if do_surface_correlation_plot:
ret[4] = sst_nh_diff_ts
print('^', end='')
return tuple(ret)
