def main(inputfile, outputfile, **kwargs):
# output figure type and name
fig_name = os.path.splitext(outputfile)[0]
fig_type = os.path.splitext(outputfile)[1][1:]
# create the plot workstation
wks = ngl.open_wks(fig_type, fig_name)
# Increase maximum buffer memory
increase_workspace_memory(value=10000000000)
# Read image data
ds = xr.open_dataset(inputfile)
band1 = ds['Band1'].values
band2 = ds['Band2'].values
band3 = ds['Band3'].values
# Plot image
# mpCenterLonF=150 for Pacific; 340 for Atlantic; 200 Pacific v2
# mpCenterLatF=30 for Pacific; 45 for Atlantic; 40 Pacific v2
pl = plot_image_gdal(wks, band1, band2, band3, mpCenterLonF=340, mpCenterLatF=45)
# Finalize plot
ngl.frame(wks)
ngl.end()
def main(test_files,
cntl_files,
vname,
fig_name,
test_name='Model',
cntl_name='Obs',
**kwargs):
# Load datasets
datasets = [open_dataset(f) for f in (test_files, cntl_files)]
data_arrays = [get_data(ds, vname) for ds in datasets]
coords = [get_coords(ds) for ds in datasets]
weights = [get_area_weights(ds) for ds in datasets]
# Compute area averages
means = [
area_average(da, wgt, dims=[d for d in da.dims if d != 'time'])
for (da, wgt) in zip(data_arrays, weights)
]
# Make line plots of area averages
times = [da['time'] for da in data_arrays]
labels = (test_name, cntl_name)
plot_format = fig_name.split('.')[-1]
wks = ngl.open_wks(plot_format, fig_name)
def main(datafile, varname, gridfile=None):
# Read data
ds_data = xarray.open_dataset(datafile)
data = ds_data[varname]
if gridfile is not None:
ds_grid = xarray.open_dataset(gridfile).rename({'grid_size': 'ncol'})
lon = ds_grid['grid_corner_lon']
lat = ds_grid['grid_corner_lat']
else:
lon = ds_data['lon']
lat = ds_data['lat']
# Make sure we don't have time or level dimensions
if 'time' in data.dims:
data = data.isel(time=0).squeeze()
if 'lev' in data.dims:
data = data.isel(lev=-1).squeeze()
if 'lon' in ds_grid and 'lat' in ds_grid:
x = ds_grid['lon']
y = ds_grid['lat']
elif 'grid_corner_lon' in ds_grid and 'grid_corner_lat' in ds_grid:
x = ds_grid['grid_corner_lon'].rename({'grid_size': 'ncol'})
y = ds_grid['grid_corner_lat'].rename({'grid_size': 'ncol'})
# Setup the canvas
plot_format='png'
plot_name='./' + varname
wks = ngl.open_wks(plot_format, plot_name)
# Make plot
plot = plot_unstructured(
wks, x.values, y.values, data.values,
mpGeophysicalLineColor='white',
lbOrientation='horizontal',
lbTitleString='%s (%s)'%(data.long_name, data.units),
cnFillMode='RasterFill',
cnLineLabelsOn=False, cnLinesOn=False,
)
# Close things
ngl.end()
def plot_frame(lon, lat, data, frame_name,
lat_min=None, lat_max=None, lon_min=None, lon_max=None,
**kwargs):
# Open figure
wks = ngl.open_wks('png', os.path.splitext(frame_name)[0])
# Plot data
pl = plot_map(
wks, lon, lat, data,
#tiMainString=f'{str(data.time.values)}',
lbOrientation='horizontal',
#lbTitleString='%s (%s)'%(data.long_name, data.units),
cnFillMode='RasterFill',
**kwargs
)
# Save figure
ngl.draw(pl)
ngl.destroy(wks)
return frame_name
def main(varname,
plotname,
*datafiles,
gridfile=None,
time_index=None,
vmin=None,
vmax=None,
**kwargs):
# Read data
ds_data = xarray.open_mfdataset(
sorted(datafiles),
chunks={'time': 1},
drop_variables=('P3_input_dim', 'P3_output_dim'),
)
data = get_data(ds_data, varname)
if gridfile is not None:
ds_grid = xarray.open_dataset(gridfile).rename({'grid_size': 'ncol'})
if 'lon' in ds_grid and 'lat' in ds_grid:
x = ds_grid['lon']
y = ds_grid['lat']
elif 'grid_corner_lon' in ds_grid and 'grid_corner_lat' in ds_grid:
x = ds_grid['grid_corner_lon']
y = ds_grid['grid_corner_lat']
else:
raise RuntimeError('No valid coordinates in grid file.')
else:
x = ds_data['lon']
y = ds_data['lat']
# Make sure we don't have time or level dimensions
if 'time' in data.dims:
if time_index is None:
data = data.mean(dim='time', keep_attrs=True).squeeze()
else:
data = data.isel(time=int(time_index))
if 'lev' in data.dims:
data = data.isel(lev=-1).squeeze()
if 'time' in x.dims: x = x.isel(time=0)
if 'time' in y.dims: y = y.isel(time=0)
# Setup the canvas
wks = ngl.open_wks(
os.path.splitext(plotname)[1][1:],
os.path.splitext(plotname)[0])
# Get contour levels; the explicit type casting deals with problems calling
# this standalone code using subprocess.run() with string arguments, where
# all kwargs are going to be interpreted as strings
if vmin is None: vmin = data.min().values
if vmax is None: vmax = data.max().values
if float(vmin) < 0 and float(vmax) > 0:
*__, clevels = nice_cntr_levels(float(vmin),
float(vmax),
returnLevels=True,
max_steps=13,
aboutZero=True)
else:
*__, clevels = nice_cntr_levels(float(vmin),
float(vmax),
returnLevels=True,
max_steps=13)
kwargs['cnLevels'] = clevels #get_contour_levels(data)
kwargs['cnLevelSelectionMode'] = 'ExplicitLevels'
# Make plot
if 'lbTitleString' not in kwargs.keys():
kwargs['lbTitleString'] = f'{data.long_name} ({data.units})'
plot = plot_map(wks,
x.values,
y.values,
data.values,
mpGeophysicalLineColor='white',
lbOrientation='horizontal',
cnFillMode='RasterFill',
cnLineLabelsOn=False,
cnLinesOn=False,
**kwargs)
ngl.destroy(wks)
def main(test_files,
cntl_files,
varname,
plotname,
test_name="Model",
cntl_name="Obs",
time_offsets=[None, None],
map_file=None,
test_grid=None,
cntl_grid=None,
t1=None,
t2=None,
verbose=False,
**kwargs):
# Read data
if verbose:
print('Open datasets...')
sys.stdout.flush()
datasets = [
open_dataset(files, time_offset=time_offset, chunks={'time': 1})
for files, time_offset in zip((test_files, cntl_files), time_offsets)
]
grids = [
xarray.open_mfdataset(f) if f is not None else None
for f in (test_grid, cntl_grid)
]
# Subset data
if verbose:
print('Subset consistent time periods...')
sys.stdout.flush()
if t1 is None: t1 = max([ds.time[0].values for ds in datasets])
if t2 is None: t2 = min([ds.time[-1].values for ds in datasets])
if verbose:
print('Comparing period {} to {}'.format(str(t1), str(t2)))
sys.stdout.flush()
datasets_subset = [ds.sel(time=slice(str(t1), str(t2))) for ds in datasets]
# Compute time average
if verbose:
print('Compute time averages...')
sys.stdout.flush()
datasets_ta = [
ds.mean(dim='time', keep_attrs=True) for ds in datasets_subset
]
# Select data
if verbose:
print('Select data...')
sys.stdout.flush()
data_arrays = [mask_all_zero(get_data(ds, varname)) for ds in datasets_ta]
#data_arrays = [get_data(ds, varname) for ds in datasets]
coords = [
get_coords(ds, ds_grid) for ds, ds_grid in zip(datasets_ta, grids)
]
#coords = [get_coords(ds) for ds in datasets]
weights = [get_area_weights(ds) for ds in datasets_ta]
# Try explicitly evaluating dask graph to avoid dask getting confused later
# and trying to load entire dataset into memory. By this point we should
# have reduced the data arrays to a manageable size, so loading this into
# memory should not be a problem. TODO: this feels like a hack I should not
# have to do, is something else going on here?
# NOTE: doing this here seems to speed things up quite a bit too.
#print('Execute dask computations (data_arrays)...'); sys.stdout.flush()
#data_arrays = [d.compute() for d in data_arrays]
#print('Execute dask computations (coords)...'); sys.stdout.flush()
#coords = [[c.compute() for c in coord] for coord in coords]
#print('Execute dask computations (weights)...'); sys.stdout.flush()
#weights = [w.compute() for w in weights]
for i in range(len(data_arrays)):
data_arrays[i].load()
coords[i][0].load()
coords[i][1].load()
weights[i].load()
# Compute differences
if verbose:
print('Compute differences...')
sys.stdout.flush()
data_arrays.append(
compute_differences(data_arrays[0], data_arrays[1], *coords[0],
*coords[1], map_file))
coords.append(coords[1])
weights.append(weights[1])
# Compute global statistics for plot labels
if verbose:
print('Compute statistics...')
sys.stdout.flush()
means = [area_average(d, w) for (d, w) in zip(data_arrays, weights)]
# Get common contour levels
if verbose:
print('Find good clevels for plotting...')
sys.stdout.flush()
clevels = get_contour_levels([d for d in data_arrays[:-1]])
dlevels = get_contour_levels([
data_arrays[-1],
], aboutZero=True)
levels_list = [clevels, clevels, dlevels]
# Make plots
if verbose:
print('Make plots...')
sys.stdout.flush()
plot_format = plotname.split('.')[-1]
wks = ngl.open_wks(plot_format, plotname)
diff_name = f'{test_name} minus {cntl_name}'
labels = [
f'{label} ({m.values:.1f})'
for (label, m) in zip((test_name, cntl_name, diff_name), means)
]
plots = [
plot_panel(wks, d, *c, tiMainString=label, cnLevels=levels, **kwargs)
for (d, c, label,
levels) in zip(data_arrays, coords, labels, levels_list)
]
# Panel plots
if verbose:
print('Closing plot workspace and writing figures...')
sys.stdout.flush()
ngl.panel(wks, plots, [1, len(plots)])
ngl.destroy(wks)
# Finally, trim whitespace from our figures
if verbose:
print('Trimming whitespace from figure...')
sys.stdout.flush()
subprocess.call(f'convert -trim {plotname} {plotname}'.split(' '))