def make_nino34_timeseries(file, mean_ssts, n34_lat_start_idx, n34_lat_end_idx, n34_lon_start_idx, n34_lon_end_idx):
# Make full timeseries in nino34 region, subtract the mean sst for that
# month.
# FIXME: need to check that file contains monthly means.
with nc.Dataset(file) as f:
periods = time_dim_to_pandas_periods(f.variables["time"])
temp_var = f.variables["temp"]
ssts = []
for t in range(temp_var.shape[0]):
sst = np.mean(temp_var[t, 0, n34_lat_start_idx:n34_lat_end_idx, n34_lon_start_idx:n34_lon_end_idx])
ssts.append(mean_ssts[periods[t].month - 1] - sst)
ts = pd.Series(ssts, periods)
print("^", end="")
sys.stdout.flush()
return ts
def make_amoc_idx_maps(file):
"""
Make amoc index maps for every time point in file, also returns the max of
the amoc stream function for all time points.
"""
with nc.Dataset(file) as f:
lons = f.variables['geolon_t']
lats = f.variables['geolat_t']
time_var = f.variables['time']
temp_var = f.variables['temp']
nh_mask = get_nh_mask(temp_var[0, 0, :, :].mask)
ty_trans = f.variables['ty_trans']
dzt = f.variables['dzt']
amoc_idx_map = np.ma.zeros(temp_var[:,0,:,:].shape)
amoc_psi_max = []
for t in range(time_var.shape[0]):
amoc, atlantic_mask, _, _, _ = \
calc_atlantic_moc(ty_trans[t, :, :, :],
dzt[t, :, :, :], lons, lats)
amoc_psi_max.append(np.max(amoc))
# Get surface temp spatial mean in the NH
nh_sst_mean = np.mean(np.ma.masked_array(temp_var[t, 0, :, :],
mask=nh_mask))
# Map of the AMOC index, see (1)
amoc_idx_map[t, :, :] = temp_var[t, 0, :, :] - nh_sst_mean
amoc_idx_map[t, :, :].mask = atlantic_mask[0, :, :]
periods = time_dim_to_pandas_periods(time_var)
amoc_psi_max_ts = pd.Series(amoc_psi_max, periods)
time_check = time_var[:]
print('+', end='')
sys.stdout.flush()
# Pass back time variable just to check that everything is put back in the
# correct order.
return amoc_idx_map, amoc_psi_max_ts, time_check, nh_sst_mean, nh_mask
def main():
parser = argparse.ArgumentParser()
parser.add_argument('input_files', nargs='+',
help='The MOM ocean_scalar.nc input data files.')
parser.add_argument('--field', default='temp_global_ave',
help='The fields included in this plot.')
parser.add_argument('--output_dir', default='./',
help='Directory where plots will be written.')
args = parser.parse_args()
title = None
ylabel = None
ts = pd.Series()
# Go through input files one at a time building the timeseries as we go.
for file in args.input_files:
with nc.Dataset(file) as f:
time_var = f.variables['time']
data_var = f.variables[args.field]
title = data_var.long_name
ylabel = data_var.units
# Calculate the times/dates, these will be our indices.
periods = time_dim_to_pandas_periods(f.variables['time'])
data = f.variables[args.field][:]
assert(data.shape[1] == 1)
data = data.flatten()
new_ts = pd.Series(data, periods)
ts = ts.append(new_ts)
ts = ts.sort_index()
plot = ts.plot()
plt.xlabel('Time (years)')
plt.ylabel(ylabel)
plt.title(title)
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(9,4.5)
plt.savefig(os.path.join(args.output_dir, '{}.png'.format(args.field)))
def make_timeseries(args):
"""
Make a timeseries of func(var_name) for all time points in input_file. Only the
points where mask > 0 are used.
"""
input_file, var_name, mask, mask_var = args
func = np.mean
with nc.Dataset(input_file) as f:
periods = time_dim_to_pandas_periods(f.variables['time'])
vals = []
for t in range(len(periods)):
var = f.variables[var_name][t, :]
vals.append(func(var[np.where(mask != 0)]))
ts = pd.Series(vals, periods)
print('%', end='')
sys.stdout.flush()
return ts
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)
