Skip to content

Commit

Permalink
Merge branch 'develop' into feature/oops_app
Browse files Browse the repository at this point in the history
  • Loading branch information
guillaumevernieres committed Jul 17, 2023
2 parents c77a979 + a184e68 commit 7e9289e
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 96 deletions.
97 changes: 1 addition & 96 deletions scripts/exgdas_global_marine_analysis_vrfy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
102 changes: 102 additions & 0 deletions ush/soca_vrfy.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 7e9289e

Please sign in to comment.