From a184e688b7aeca9631a138d39b89adeb34387008 Mon Sep 17 00:00:00 2001 From: AndrewEichmann-NOAA <58948505+AndrewEichmann-NOAA@users.noreply.github.com> Date: Mon, 17 Jul 2023 12:41:18 -0400 Subject: [PATCH] initial changes for vrfy refactoring (#532) * initial changes for refactoring * python coding norms * removed paleolithic header * python coding norms --- scripts/exgdas_global_marine_analysis_vrfy.py | 97 +---------------- ush/soca_vrfy.py | 102 ++++++++++++++++++ 2 files changed, 103 insertions(+), 96 deletions(-) create mode 100755 ush/soca_vrfy.py diff --git a/scripts/exgdas_global_marine_analysis_vrfy.py b/scripts/exgdas_global_marine_analysis_vrfy.py index 04f56afd5..ed68fb100 100755 --- a/scripts/exgdas_global_marine_analysis_vrfy.py +++ b/scripts/exgdas_global_marine_analysis_vrfy.py @@ -19,108 +19,13 @@ import os import numpy as np -import matplotlib.pyplot as plt -import xarray as xr -import cartopy -import cartopy.crs as ccrs import gen_eva_obs_yaml import marine_eva_post import diag_statistics +from soca_vrfy import plot_config, plot_horizontal_slice, plot_zonal_slice import subprocess from datetime import datetime, timedelta -projs = {'North': ccrs.NorthPolarStereo(), - 'South': ccrs.SouthPolarStereo(), - 'Global': ccrs.Mollweide(central_longitude=-150)} - - -def plot_config(grid_file=[], data_file=[], - variable=[], levels=[], bounds=[], colormap=[], comout=[], lats=[]): - """ - Prepares the configuration for the plotting functions below - """ - config = {} - config['grid file'] = grid_file - config['fields file'] = data_file - config['variable'] = variable - config['levels'] = levels - config['bounds'] = bounds - config['colormap'] = colormap - config['lats'] = lats - config['comout'] = comout - config['max depth'] = 5000.0 - config['proj'] = 'Global' - return config - - -def plot_horizontal_slice(config): - """ - pcolormesh of a horizontal slice of an ocean field - """ - grid = xr.open_dataset(config['grid file']) - data = xr.open_dataset(config['fields file']) - - dirname = os.path.join(config['comout'], config['variable']) - os.makedirs(dirname, exist_ok=True) - - if config['variable'] in ['Temp', 'Salt', 'u', 'v']: - level = config['levels'][0] - slice_data = np.squeeze(data[config['variable']])[level, :, :] - label_colorbar = config['variable']+' Level '+str(level) - figname = os.path.join(dirname, config['variable']+'_Level_'+str(level)) - else: - slice_data = np.squeeze(data[config['variable']]) - label_colorbar = config['variable'] - figname = os.path.join(dirname, config['variable']+'_'+config['proj']) - - bounds = config['bounds'] - - fig, ax = plt.subplots(figsize=(8, 5), subplot_kw={'projection': projs[config['proj']]}) - plt.pcolormesh(np.squeeze(grid.lon), - np.squeeze(grid.lat), - slice_data, - vmin=bounds[0], vmax=bounds[1], - transform=ccrs.PlateCarree(), - cmap=config['colormap']) - - plt.colorbar(label=label_colorbar, shrink=0.5, orientation='horizontal') - ax.coastlines() # TODO: make this work on hpc - ax.gridlines(draw_labels=True) - if config['proj'] == 'South': - ax.set_extent([-180, 180, -90, -50], ccrs.PlateCarree()) - if config['proj'] == 'North': - ax.set_extent([-180, 180, 50, 90], ccrs.PlateCarree()) - # ax.add_feature(cartopy.feature.LAND) # TODO: make this work on hpc - plt.savefig(figname, bbox_inches='tight', dpi=600) - - -def plot_zonal_slice(config): - """ - pcolormesh of a zonal slice of an ocean field - """ - lat = float(config['lats'][0]) - grid = xr.open_dataset(config['grid file']) - data = xr.open_dataset(config['fields file']) - lat_index = np.argmin(np.array(np.abs(np.squeeze(grid.lat)[:, 0]-lat))) - slice_data = np.squeeze(np.array(data[config['variable']]))[:, lat_index, :] - depth = np.squeeze(np.array(grid['h']))[:, lat_index, :] - depth[np.where(np.abs(depth) > 10000.0)] = 0.0 - depth = np.cumsum(depth, axis=0) - bounds = config['bounds'] - x = np.tile(np.squeeze(grid.lon[:, lat_index]), (np.shape(depth)[0], 1)) - fig, ax = plt.subplots(figsize=(8, 5)) - plt.pcolormesh(x, -depth, slice_data, - vmin=bounds[0], vmax=bounds[1], - cmap=config['colormap']) - plt.colorbar(label=config['variable']+' Lat '+str(lat), shrink=0.5, orientation='horizontal') - ax.set_ylim(-config['max depth'], 0) - dirname = os.path.join(config['comout'], config['variable']) - os.makedirs(dirname, exist_ok=True) - figname = os.path.join(dirname, config['variable'] + - 'zonal_lat_'+str(int(lat)) + '_' + str(int(config['max depth'])) + 'm') - plt.savefig(figname, bbox_inches='tight', dpi=600) - - comout = os.getenv('COM_OCEAN_ANALYSIS') com_ice_history = os.getenv('COM_ICE_HISTORY_PREV') com_ocean_history = os.getenv('COM_OCEAN_HISTORY_PREV') diff --git a/ush/soca_vrfy.py b/ush/soca_vrfy.py new file mode 100755 index 000000000..f06e4cd4f --- /dev/null +++ b/ush/soca_vrfy.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 + +# make plots for marine analysis + +import matplotlib.pyplot as plt +import xarray as xr +import cartopy +import cartopy.crs as ccrs +import numpy as np +import os + + +projs = {'North': ccrs.NorthPolarStereo(), + 'South': ccrs.SouthPolarStereo(), + 'Global': ccrs.Mollweide(central_longitude=-150)} + + +def plot_config(grid_file=[], data_file=[], + variable=[], levels=[], bounds=[], colormap=[], comout=[], lats=[]): + """ + Prepares the configuration for the plotting functions below + """ + config = {} + config['grid file'] = grid_file + config['fields file'] = data_file + config['variable'] = variable + config['levels'] = levels + config['bounds'] = bounds + config['colormap'] = colormap + config['lats'] = lats + config['comout'] = comout + config['max depth'] = 5000.0 + config['proj'] = 'Global' + return config + + +def plot_horizontal_slice(config): + """ + pcolormesh of a horizontal slice of an ocean field + """ + grid = xr.open_dataset(config['grid file']) + data = xr.open_dataset(config['fields file']) + + dirname = os.path.join(config['comout'], config['variable']) + os.makedirs(dirname, exist_ok=True) + + if config['variable'] in ['Temp', 'Salt', 'u', 'v']: + level = config['levels'][0] + slice_data = np.squeeze(data[config['variable']])[level, :, :] + label_colorbar = config['variable']+' Level '+str(level) + figname = os.path.join(dirname, config['variable']+'_Level_'+str(level)) + else: + slice_data = np.squeeze(data[config['variable']]) + label_colorbar = config['variable'] + figname = os.path.join(dirname, config['variable']+'_'+config['proj']) + + bounds = config['bounds'] + + fig, ax = plt.subplots(figsize=(8, 5), subplot_kw={'projection': projs[config['proj']]}) + plt.pcolormesh(np.squeeze(grid.lon), + np.squeeze(grid.lat), + slice_data, + vmin=bounds[0], vmax=bounds[1], + transform=ccrs.PlateCarree(), + cmap=config['colormap']) + + plt.colorbar(label=label_colorbar, shrink=0.5, orientation='horizontal') + ax.coastlines() # TODO: make this work on hpc + ax.gridlines(draw_labels=True) + if config['proj'] == 'South': + ax.set_extent([-180, 180, -90, -50], ccrs.PlateCarree()) + if config['proj'] == 'North': + ax.set_extent([-180, 180, 50, 90], ccrs.PlateCarree()) + # ax.add_feature(cartopy.feature.LAND) # TODO: make this work on hpc + plt.savefig(figname, bbox_inches='tight', dpi=600) + + +def plot_zonal_slice(config): + """ + pcolormesh of a zonal slice of an ocean field + """ + lat = float(config['lats'][0]) + grid = xr.open_dataset(config['grid file']) + data = xr.open_dataset(config['fields file']) + lat_index = np.argmin(np.array(np.abs(np.squeeze(grid.lat)[:, 0]-lat))) + slice_data = np.squeeze(np.array(data[config['variable']]))[:, lat_index, :] + depth = np.squeeze(np.array(grid['h']))[:, lat_index, :] + depth[np.where(np.abs(depth) > 10000.0)] = 0.0 + depth = np.cumsum(depth, axis=0) + bounds = config['bounds'] + x = np.tile(np.squeeze(grid.lon[:, lat_index]), (np.shape(depth)[0], 1)) + fig, ax = plt.subplots(figsize=(8, 5)) + plt.pcolormesh(x, -depth, slice_data, + vmin=bounds[0], vmax=bounds[1], + cmap=config['colormap']) + plt.colorbar(label=config['variable']+' Lat '+str(lat), shrink=0.5, orientation='horizontal') + ax.set_ylim(-config['max depth'], 0) + dirname = os.path.join(config['comout'], config['variable']) + os.makedirs(dirname, exist_ok=True) + figname = os.path.join(dirname, config['variable'] + + 'zonal_lat_'+str(int(lat)) + '_' + str(int(config['max depth'])) + 'm') + plt.savefig(figname, bbox_inches='tight', dpi=600)