diff --git a/compressible/dry_baroclinic_wave_temperature_pressure_plot.png b/compressible/dry_baroclinic_wave_temperature_pressure_plot.png new file mode 100644 index 0000000..c3546c0 Binary files /dev/null and b/compressible/dry_baroclinic_wave_temperature_pressure_plot.png differ diff --git a/compressible/plot_dry_baroclinic_sphere.py b/compressible/plot_dry_baroclinic_sphere.py new file mode 100644 index 0000000..1876935 --- /dev/null +++ b/compressible/plot_dry_baroclinic_sphere.py @@ -0,0 +1,121 @@ +""" +Plot for a 2x1 grpah of surface tempeorature and surface pressure fields. +""" +import numpy as np +import matplotlib.pyplot as plt +from netCDF4 import Dataset +from tomplot import (tomplot_contours, tomplot_cmap, + plot_contoured_field, add_colorbar_ax, + regrid_vertical_slice, tomplot_field_title, + extract_gusto_vertical_slice, + reshape_gusto_data, extract_gusto_field, + extract_gusto_coords, area_restriction) + + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these should not be relative to this file +results_dir = '/data/home/dw603/firedrake-08_05_24/src/gusto/PhD/results/baroclinic_wave_order_1_1_with_vort' # This needs to point to a data file +results_file_name = f'{results_dir}/field_output.nc' +plot_dir = f'{results_dir}/plots' +plot_stem = f'{plot_dir}/output_testing' + +# ---------------------------------------------------------------------------- # +# Things that should be altered based on the plot +# ---------------------------------------------------------------------------- # +# Specify lists for variables that are different between subplots +field_names = ['Temperature', 'Pressure_Vt'] +titles = ['Temperature (850 HPa)', 'Surface Pressure'] +slice_along_values = ['z', 'z'] + +field_labels = [r'$T \ / $K', r'$P \ / $Pa'] + +remove_contour_vals = [None, None] +domain_limit = {'X' : (0, 180), 'Y' : (0, 90)} +colour_schemes = ['jet', 'jet'] +contours = ['contour a','contour b','contour c ','contour d'] + +# Things that are the same for all subplots +time_idxs = [0, 768, -1] +contour_method = 'tricontour' +slice_at_vals = [0, 0 , 0, 0] +# 1D grids for vertical regridding +coords_lon_1d = np.linspace(-180, 180, 50) +coords_lat_1d = np.linspace(-90, 90, 50) +# Dictionary to hold plotting grids -- keys are "slice_along" values +plotting_grids = {'lat': coords_lon_1d, 'lon': coords_lat_1d} +# Level for horizontal slices +levels = [0, 0] +# ---------------------------------------------------------------------------- # +# Things that are likely the same for all plots +# ---------------------------------------------------------------------------- # + +data_file = Dataset(results_file_name, 'r') +for time_idx in time_idxs: + fig, axarray = plt.subplots(2, 1, figsize=(16, 8), sharey='row') + + # Loop through subplots + for i, (ax, field_name, field_label, colour_scheme, slice_along, remove_contour, slice_at, title, level, contour) in \ + enumerate(zip(axarray.flatten(), field_names, field_labels, + colour_schemes, slice_along_values, remove_contour_vals, slice_at_vals, titles, levels, contours)): + # ------------------------------------------------------------------------ # + # Data extraction + # ------------------------------------------------------------------------ # + if slice_along == 'z': + # Extraction + field_full = extract_gusto_field(data_file, field_name, time_idx) + coords_X_full, coords_Y_full, coords_Z_full = \ + extract_gusto_coords(data_file, field_name) + + # Reshape + field_full, coords_X_full, coords_Y_full, _ = \ + reshape_gusto_data(field_full, coords_X_full, + coords_Y_full, coords_Z_full) + + # Domain restriction + field_data, coords_hori, coords_Z = \ + area_restriction(field_full[:, level], coords_X_full[:, level], + coords_Y_full[:, level], domain_limit) + else: + orig_field_data, orig_coords_X, orig_coords_Y, orig_coords_Z = \ + extract_gusto_vertical_slice(data_file, field_name, time_idx, + slice_along=slice_along, slice_at=slice_at) + + # Slices need regridding as points don't cleanly live along lon or lat = 0.0 + field_data, coords_hori, coords_Z = regrid_vertical_slice(plotting_grids[slice_along], + slice_along, slice_at, + orig_coords_X, orig_coords_Y, + orig_coords_Z, orig_field_data) + time = data_file['time'][time_idx] + time_in_days = time / (24*60*60) + # ------------------------------------------------------------------------ # + # Plot data + # ------------------------------------------------------------------------ # + contours = tomplot_contours(field_data) + if field_name == 'Temperature': + contours = np.arange(220, 320, 10) + + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=remove_contour) + cf, _ = plot_contoured_field(ax, coords_hori, coords_Z, field_data, + contour_method, contours, cmap=cmap, + line_contours=lines) + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + # Don't add ylabels unless left-most subplots + ylabel = True if i % 3 == 0 else None + ylabelpad = -30 if i > 2 else -10 + + tomplot_field_title(ax, title, minmax=True, field_data=field_data) + + # These subplots tend to be quite clustered together, so move them apart a bit + fig.subplots_adjust(wspace=0.3, hspace=0.3) + fig.suptitle(f'Baroclinic Wave at {time_in_days} days') + # ---------------------------------------------------------------------------- # + # Save figure + # ---------------------------------------------------------------------------- # + plot_name = f'{plot_stem}_{time_in_days}_days.png' + print(f'Saving figure to {plot_name}') + fig.savefig(plot_name, bbox_inches='tight') + plt.close() + +