-
Notifications
You must be signed in to change notification settings - Fork 0
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
added plotting script for baroclinic wave and plot it produces #14
base: main
Are you sure you want to change the base?
Changes from 1 commit
0873051
cd802b3
2f1052d
e06c02e
62b9a9a
72447bc
6568a07
5d2fef8
04565c6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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'] | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is the temperature field actually at 850hPa? I think I'm missing where the interpolation to get that is! It looks like we are taking the surface values |
||
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)} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You could use these domain limits to define the xticks and yticks (see comment below) |
||
colour_schemes = ['jet', 'jet'] | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm actually not personally that keen on these colour schemes (although I know they were used in the original paper!) How would you feel about There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If you want a better colour scheme but still want easy comparison to the original paper, you can add contour lines to your plot at the values of the colour increments in the original paper - this way the contour lines should match the colour boundaries from the original paper. |
||
contours = ['contour a','contour b','contour c ','contour d'] | ||
|
||
# Things that are the same for all subplots | ||
time_idxs = [0, 768, -1] | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. For the plot that we're adding here, I think you want -1, and something separate for the initial conditions |
||
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) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We could add the data format argument here, as for pressure I don't think we want to label the plot to so many significant figures |
||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Somewhere could you restrict the xticks and yticks on the x and y axes so that they are just the limits of the plotted domain? |
||
# 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() | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you fix the typos here? And if you add to this plotting script to include initial fields, you could add a comment about that