# Path to ROMS grid file

roms_grid = '/short/m68/kaa561/metroms_iceshelf/apps/common/grid/circ30S_quarterdegree.nc'

# Path to ROMS time-averaged file

roms_file = '/short/m68/kaa561/metroms_iceshelf/tmproms/run/intercomparison/2002_2016_avg.nc'

# Path to FESOM mesh directories

fesom_mesh_path_lr = '/short/y99/kaa561/FESOM/mesh/meshA/'

fesom_mesh_path_hr = '/short/y99/kaa561/FESOM/mesh/meshB/'

# Path to FESOM time-averaged ocean files (temp, salt, u, v)

fesom_file_lr_o = '/short/y99/kaa561/FESOM/intercomparison_lowres/output/oce_2002_2016_avg.nc'

fesom_file_hr_o = '/short/y99/kaa561/FESOM/intercomparison_highres/output/oce_2002_2016_avg.nc'

# Path to FESOM time-averaged ice shelf files (wnet)

fesom_file_lr_i = '/short/y99/kaa561/FESOM/intercomparison_lowres/output/wnet_2002_2016_avg.nc'

fesom_file_hr_i = '/short/y99/kaa561/FESOM/intercomparison_highres/output/wnet_2002_2016_avg.nc'

# Name of each region

region_names = ['Filchner-Ronne Ice Shelf', 'Eastern Weddell Region', 'Amery Ice Shelf', 'Australian Sector', 'Ross Sea', 'Amundsen Sea', 'Bellingshausen Sea', 'Larsen Ice Shelves']

num_regions = len(region_names)

# Beginning of filenames for figures

fig_heads = ['filchner_ronne', 'eweddell', 'amery', 'australian', 'ross', 'amundsen', 'bellingshausen', 'larsen']

# Bounds for each region (using polar coordinate transformation as below)

x_min = [-14, -8, 15.25, 12, -9.5, -15.5, -20.25, -22.5]

x_max = [-4.5, 13, 20.5, 25.5, 4, -10.5, -15.5, -14.5]

y_min = [1, 12, 4.75, -20, -13, -11.25, -4.5, 8.3]

y_max = [10, 21, 8, 4, -4.75, -2.25, 7.6, 13]

# Size of each plot in the y direction

ysize = [8, 6, 7, 9, 7, 9, 10, 7]

# Variables to process

var_names = ['vel'] #['bathy', 'draft', 'wct', 'melt', 'temp', 'salt', 'vel']

# Constants

sec_per_year = 365*24*3600

deg2rad = pi/180.0

# Parameters for missing circle in ROMS grid

lon_c = 50

lat_c = -83

radius = 10.1

nbdry = -63+90

# ROMS vertical grid parameters

theta_s = 7.0

theta_b = 2.0

hc = 250

N = 31

# Number of bins in each direction for vector overlay

num_bins = 30

print 'Reading ROMS grid'

# Read the fields we need

id = Dataset(roms_grid, 'r')

roms_lon = id.variables['lon_rho'][:,:]

roms_lat = id.variables['lat_rho'][:,:]

roms_h = id.variables['h'][:,:]

roms_mask = id.variables['mask_rho'][:,:]

roms_zice = id.variables['zice'][:,:]

roms_angle = id.variables['angle'][:,:]

id.close()

# Get land/zice mask

open_ocn = copy(roms_mask)

open_ocn[roms_zice!=0] = 0

land_zice = ma.masked_where(open_ocn==1, open_ocn)

# Convert grid to spherical coordinates

roms_x = -(roms_lat+90)*cos(roms_lon*deg2rad+pi/2)

roms_y = (roms_lat+90)*sin(roms_lon*deg2rad+pi/2)

# Find centre in spherical coordinates

x_c = -(lat_c+90)*cos(lon_c*deg2rad+pi/2)

y_c = (lat_c+90)*sin(lon_c*deg2rad+pi/2)

# Build a regular x-y grid and select the missing circle

x_reg_roms, y_reg_roms = meshgrid(linspace(-nbdry, nbdry, num=1000), linspace(-nbdry, nbdry, num=1000))

land_circle = zeros(shape(x_reg_roms))

land_circle = ma.masked_where(sqrt((x_reg_roms-x_c)**2 + (y_reg_roms-y_c)**2) > radius, land_circle)

print 'Building FESOM low-res mesh'

elements_lr = fesom_grid(fesom_mesh_path_lr, circumpolar=True)

# Make patches for all elements, ice shelf elements, and open ocean elements

patches_lr = []

patches_shelf_lr = []

patches_ocn_lr = []

for elm in elements_lr:

coord = transpose(vstack((elm.x, elm.y)))

patches_lr.append(Polygon(coord, True, linewidth=0.))

if elm.cavity:

patches_shelf_lr.append(Polygon(coord, True, linewidth=0.))

else:

patches_ocn_lr.append(Polygon(coord, True, linewidth=0.))

print 'Building FESOM high-res mesh'

elements_hr = fesom_grid(fesom_mesh_path_hr, circumpolar=True)

patches_hr = []

patches_shelf_hr = []

patches_ocn_hr = []

for elm in elements_hr:

coord = transpose(vstack((elm.x, elm.y)))

patches_hr.append(Polygon(coord, True, linewidth=0.))

if elm.cavity:

patches_shelf_hr.append(Polygon(coord, True, linewidth=0.))

else:

patches_ocn_hr.append(Polygon(coord, True, linewidth=0.))

for var in var_names:

print 'Processing variable ' + var

print 'Reading ROMS fields'

if var == 'draft':

# Swap sign on existing zice field; nothing more to read

roms_data = -1*roms_zice

elif var == 'bathy':

# Point to h field and mask out land mask; nothing more to read

roms_data = ma.masked_where(roms_mask==0, roms_h)

elif var == 'wct':

# Add h (positive) and zice (negative); nothing more to read

roms_data = roms_h + roms_zice

else:

id = Dataset(roms_file, 'r')

if var == 'melt':

# Convert from m/s to m/y

roms_data = id.variables['m'][0,:,:]*sec_per_year

elif var in ['temp', 'salt']:

# Bottom layer

roms_data = id.variables[var][0,0,:,:]

elif var == 'vel':

# Read full 3D u and v

u_3d_tmp = id.variables['u'][0,:,:,:]

v_3d_tmp = id.variables['v'][0,:,:,:]

# Get integrands on 3D grid; we only care about dz

dx, dy, dz, z = cartesian_grid_3d(roms_lon, roms_lat, roms_h, roms_zice, theta_s, theta_b, hc, N)

# Unrotate each vertical level

u_3d = ma.empty(shape(dz))

v_3d = ma.empty(shape(dz))

num_lat_u = size(u_3d_tmp,1)

num_lon_u = size(u_3d_tmp,2)

num_lat_v = size(v_3d_tmp,1)

num_lon_v = size(v_3d_tmp,2)

for k in range(N):

# Extend into land mask before interpolation to rho-grid so

# the land mask doesn't change in the final plot

for j in range(1,num_lat_u-1):

for i in range(1,num_lon_u-1):

# Check for masked points

if u_3d_tmp[k,j,i] is ma.masked:

# Look at 4 neighbours

neighbours = ma.array([u_3d_tmp[k,j-1,i], u_3d_tmp[k,j,i-1], u_3d_tmp[k,j+1,i], u_3d_tmp[k,j,i+1]])

# Find how many of them are unmasked

num_unmasked = MaskedArray.count(neighbours)

if num_unmasked > 0:

# There is at least one unmasked neighbour;

# set u_3d_tmp to their average

u_3d_tmp[k,j,i] = sum(neighbours)/num_unmasked

# Repeat for v

for j in range(1,num_lat_v-1):

for i in range(1,num_lon_v-1):

if v_3d_tmp[k,j,i] is ma.masked:

neighbours = ma.array([v_3d_tmp[k,j-1,i], v_3d_tmp[k,j,i-1], v_3d_tmp[k,j+1,i], v_3d_tmp[k,j,i+1]])

num_unmasked = MaskedArray.count(neighbours)

if num_unmasked > 0:

v_3d_tmp[k,j,i] = sum(neighbours)/num_unmasked

# Interpolate to rho grid and rotate

u_k, v_k = rotate_vector_roms(u_3d_tmp[k,:,:], v_3d_tmp[k,:,:], roms_angle)

u_3d[k,:,:] = u_k

v_3d[k,:,:] = v_k

# Vertically average u and v

u_rho = sum(u_3d*dz, axis=0)/sum(dz, axis=0)

v_rho = sum(v_3d*dz, axis=0)/sum(dz, axis=0)

# Get speed

roms_data = sqrt(u_rho**2 + v_rho**2)

# Mask out land

u_rho = ma.masked_where(roms_mask==0, u_rho)

v_rho = ma.masked_where(roms_mask==0, v_rho)

roms_data = ma.masked_where(roms_mask==0, roms_data)

id.close()

if var in ['draft', 'melt', 'wct']:

# Mask out open ocean

roms_data = ma.masked_where(roms_zice==0, roms_data)

print 'Reading FESOM low-res fields'

if var not in ['draft', 'bathy', 'wct']:

if var == 'melt':

id = Dataset(fesom_file_lr_i, 'r')

# Convert from m/s to m/y

node_data_lr = id.variables['wnet'][0,:]*sec_per_year

elif var in ['temp', 'salt']:

id = Dataset(fesom_file_lr_o, 'r')

# Read full 3D field for now

node_data_lr = id.variables[var][0,:]

elif var == 'vel':

id = Dataset(fesom_file_lr_o, 'r')

# The overlaid vectors are based on nodes not elements, so many

# of the fesom_grid data structures fail to apply and we need to

# read some of the FESOM grid files again.

# Read the cavity flag for each 2D surface node

fesom_cavity_lr = []

f = open(fesom_mesh_path_lr + 'cavity_flag_nod2d.out', 'r')

for line in f:

tmp = int(line)

if tmp == 1:

fesom_cavity_lr.append(True)

elif tmp == 0:

fesom_cavity_lr.append(False)

else:

print 'Problem'

#return

f.close()

# Save the number of 2D nodes

fesom_n2d_lr = len(fesom_cavity_lr)

# Read rotated lat and lon for each node, also depth

f = open(fesom_mesh_path_lr + 'nod3d.out', 'r')

f.readline()

rlon_lr = []

rlat_lr = []

node_depth_lr = []

for line in f:

tmp = line.split()

lon_tmp = float(tmp[1])

lat_tmp = float(tmp[2])

node_depth_tmp = -1*float(tmp[3])

if lon_tmp < -180:

lon_tmp += 360

elif lon_tmp > 180:

lon_tmp -= 360

rlon_lr.append(lon_tmp)

rlat_lr.append(lat_tmp)

node_depth_lr.append(node_depth_tmp)

f.close()

# For lat and lon, only care about the 2D nodes (the first

# fesom_n2d indices)

rlon_lr = array(rlon_lr[0:fesom_n2d_lr])

rlat_lr = array(rlat_lr[0:fesom_n2d_lr])

node_depth_lr = array(node_depth_lr)

# Unrotate longitude

fesom_lon_lr, fesom_lat_lr = unrotate_grid(rlon_lr, rlat_lr)

# Calculate polar coordinates of each node

fesom_x_lr = -(fesom_lat_lr+90)*cos(fesom_lon_lr*deg2rad+pi/2)

fesom_y_lr = (fesom_lat_lr+90)*sin(fesom_lon_lr*deg2rad+pi/2)

# Read lists of which nodes are directly below which

f = open(fesom_mesh_path_lr + 'aux3d.out', 'r')

max_num_layers_lr = int(f.readline())

node_columns_lr = zeros([fesom_n2d_lr, max_num_layers_lr])

for n in range(fesom_n2d_lr):

for k in range(max_num_layers_lr):

node_columns_lr[n,k] = int(f.readline())

node_columns_lr = node_columns_lr.astype(int)

f.close()

# Now we can actually read the data

# Read full 3D field for both u and v

node_ur_3d_lr = id.variables['u'][0,:]

node_vr_3d_lr = id.variables['v'][0,:]

# Vertically average

node_ur_lr = zeros(fesom_n2d_lr)

node_vr_lr = zeros(fesom_n2d_lr)

for n in range(fesom_n2d_lr):

# Integrate udz, vdz, and dz over this water column

udz_col = 0

vdz_col = 0

dz_col = 0

for k in range(max_num_layers_lr-1):

if node_columns_lr[n,k+1] == -999:

# Reached the bottom

break

# Trapezoidal rule

top_id = node_columns_lr[n,k]

bot_id = node_columns_lr[n,k+1]

dz_tmp = node_depth_lr[bot_id-1] - node_depth_lr[top_id-1]

udz_col += 0.5*(node_ur_3d_lr[top_id-1]+node_ur_3d_lr[bot_id-1])*dz_tmp

vdz_col += 0.5*(node_vr_3d_lr[top_id-1]+node_vr_3d_lr[bot_id-1])*dz_tmp

dz_col += dz_tmp

# Convert from integrals to averages

node_ur_lr[n] = udz_col/dz_col

node_vr_lr[n] = vdz_col/dz_col

# Unrotate

node_u_lr, node_v_lr = unrotate_vector(rlon_lr, rlat_lr, node_ur_lr, node_vr_lr)

# Calculate speed

node_data_lr = sqrt(node_u_lr**2 + node_v_lr**2)

id.close()

# Calculate given field at each element

fesom_data_lr = []

for elm in elements_lr:

# For each element, append the mean value for the 3 component Nodes

# Restrict to ice shelf cavities for draft, melt, wct

if elm.cavity or var not in ['draft', 'melt', 'wct']:

if var == 'draft':

# Ice shelf draft is depth of surface layer

fesom_data_lr.append(mean([elm.nodes[0].depth, elm.nodes[1].depth, elm.nodes[2].depth]))

elif var == 'bathy':

# Bathymetry is depth of bottom layer

fesom_data_lr.append(mean([elm.nodes[0].find_bottom().depth, elm.nodes[1].find_bottom().depth, elm.nodes[2].find_bottom().depth]))

elif var == 'wct':

# Water column thickness is depth of bottom layer minus

# depth of surface layer

fesom_data_lr.append(mean([elm.nodes[0].find_bottom().depth - elm.nodes[0].depth, elm.nodes[1].find_bottom().depth - elm.nodes[1].depth, elm.nodes[2].find_bottom().depth - elm.nodes[2].depth]))

elif var in ['melt', 'vel']:

# Surface nodes

fesom_data_lr.append(mean([node_data_lr[elm.nodes[0].id], node_data_lr[elm.nodes[1].id], node_data_lr[elm.nodes[2].id]]))

elif var in ['temp', 'salt']:

# Bottom nodes

fesom_data_lr.append(mean([node_data_lr[elm.nodes[0].find_bottom().id], node_data_lr[elm.nodes[1].find_bottom().id], node_data_lr[elm.nodes[2].find_bottom().id]]))

print 'Reading FESOM high-res fields'

# As before

if var not in ['draft', 'bathy', 'wct']:

if var == 'melt':

id = Dataset(fesom_file_hr_i, 'r')

node_data_hr = id.variables['wnet'][0,:]*sec_per_year

elif var in ['temp', 'salt']:

id = Dataset(fesom_file_hr_o, 'r')

node_data_hr = id.variables[var][0,:]

elif var == 'vel':

id = Dataset(fesom_file_hr_o, 'r')

fesom_cavity_hr = []

f = open(fesom_mesh_path_hr + 'cavity_flag_nod2d.out', 'r')

for line in f:

tmp = int(line)

if tmp == 1:

fesom_cavity_hr.append(True)

elif tmp == 0:

fesom_cavity_hr.append(False)

else:

print 'Problem'

#return

f.close()

fesom_n2d_hr = len(fesom_cavity_hr)

f = open(fesom_mesh_path_hr + 'nod3d.out', 'r')

f.readline()

rlon_hr = []

rlat_hr = []

node_depth_hr = []

for line in f:

tmp = line.split()

lon_tmp = float(tmp[1])

lat_tmp = float(tmp[2])

node_depth_tmp = -1*float(tmp[3])

if lon_tmp < -180:

lon_tmp += 360

elif lon_tmp > 180:

lon_tmp -= 360

rlon_hr.append(lon_tmp)

rlat_hr.append(lat_tmp)

node_depth_hr.append(node_depth_tmp)

f.close()

rlon_hr = array(rlon_hr[0:fesom_n2d_hr])

rlat_hr = array(rlat_hr[0:fesom_n2d_hr])

node_depth_hr = array(node_depth_hr)

fesom_lon_hr, fesom_lat_hr = unrotate_grid(rlon_hr, rlat_hr)

fesom_x_hr = -(fesom_lat_hr+90)*cos(fesom_lon_hr*deg2rad+pi/2)

fesom_y_hr = (fesom_lat_hr+90)*sin(fesom_lon_hr*deg2rad+pi/2)

f = open(fesom_mesh_path_hr + 'aux3d.out', 'r')

max_num_layers_hr = int(f.readline())

node_columns_hr = zeros([fesom_n2d_hr, max_num_layers_hr])

for n in range(fesom_n2d_hr):

for k in range(max_num_layers_hr):

node_columns_hr[n,k] = int(f.readline())

node_columns_hr = node_columns_hr.astype(int)

f.close()

node_ur_3d_hr = id.variables['u'][0,:]

node_vr_3d_hr = id.variables['v'][0,:]

node_ur_hr = zeros(fesom_n2d_hr)

node_vr_hr = zeros(fesom_n2d_hr)

for n in range(fesom_n2d_hr):

udz_col = 0

vdz_col = 0

dz_col = 0

for k in range(max_num_layers_hr-1):

if node_columns_hr[n,k+1] == -999:

break

top_id = node_columns_hr[n,k]

bot_id = node_columns_hr[n,k+1]

dz_tmp = node_depth_hr[bot_id-1] - node_depth_hr[top_id-1]

udz_col += 0.5*(node_ur_3d_hr[top_id-1]+node_ur_3d_hr[bot_id-1])*dz_tmp

vdz_col += 0.5*(node_vr_3d_hr[top_id-1]+node_vr_3d_hr[bot_id-1])*dz_tmp

dz_col += dz_tmp

node_ur_hr[n] = udz_col/dz_col

node_vr_hr[n] = vdz_col/dz_col

node_u_hr, node_v_hr = unrotate_vector(rlon_hr, rlat_hr, node_ur_hr, node_vr_hr)

node_data_hr = sqrt(node_u_hr**2 + node_v_hr**2)

id.close()

fesom_data_hr = []

for elm in elements_hr:

if elm.cavity or var not in ['draft', 'melt', 'wct']:

if var == 'draft':

fesom_data_hr.append(mean([elm.nodes[0].depth, elm.nodes[1].depth, elm.nodes[2].depth]))

elif var == 'bathy':

fesom_data_hr.append(mean([elm.nodes[0].find_bottom().depth, elm.nodes[1].find_bottom().depth, elm.nodes[2].find_bottom().depth]))

elif var == 'wct':

fesom_data_hr.append(mean([elm.nodes[0].find_bottom().depth - elm.nodes[0].depth, elm.nodes[1].find_bottom().depth - elm.nodes[1].depth, elm.nodes[2].find_bottom().depth - elm.nodes[2].depth]))

elif var in ['melt', 'vel']:

fesom_data_hr.append(mean([node_data_hr[elm.nodes[0].id], node_data_hr[elm.nodes[1].id], node_data_hr[elm.nodes[2].id]]))

elif var in ['temp', 'salt']:

fesom_data_hr.append(mean([node_data_hr[elm.nodes[0].find_bottom().id], node_data_hr[elm.nodes[1].find_bottom().id], node_data_hr[elm.nodes[2].find_bottom().id]]))

# Loop over regions

for index in range(num_regions):

print 'Processing ' + region_names[index]

# Set up a grey square for FESOM to fill the background with land

x_reg_fesom, y_reg_fesom = meshgrid(linspace(x_min[index], x_max[index], num=100), linspace(y_min[index], y_max[index], num=100))

land_square = zeros(shape(x_reg_fesom))

# Find bounds on variable in this region, for both ROMS and FESOM

# Start with ROMS

loc = (roms_x >= x_min[index])*(roms_x <= x_max[index])*(roms_y >= y_min[index])*(roms_y <= y_max[index])

var_min = amin(roms_data[loc])

var_max = amax(roms_data[loc])

# Modify with FESOM

# Low-res

i = 0

for elm in elements_lr:

if elm.cavity or var not in ['draft', 'melt', 'wct']:

if any(elm.x >= x_min[index]) and any(elm.x <= x_max[index]) and any(elm.y >= y_min[index]) and any(elm.y <= y_max[index]):

if fesom_data_lr[i] < var_min:

var_min = fesom_data_lr[i]

if fesom_data_lr[i] > var_max:

var_max = fesom_data_lr[i]

i += 1

# High-res

i = 0

for elm in elements_hr:

if elm.cavity or var not in ['draft', 'melt', 'wct']:

if any(elm.x >= x_min[index]) and any(elm.x <= x_max[index]) and any(elm.y >= y_min[index]) and any(elm.y <= y_max[index]):

if fesom_data_hr[i] < var_min:

var_min = fesom_data_hr[i]

if fesom_data_hr[i] > var_max:

var_max = fesom_data_hr[i]

i += 1

if var == 'melt':

# Special colour map

if var_min < 0:

# There is refreezing here; include blue for elements < 0

cmap_vals = array([var_min, 0, 0.25*var_max, 0.5*var_max, 0.75*var_max, var_max])

cmap_colors = [(0.26, 0.45, 0.86), (1, 1, 1), (1, 0.9, 0.4), (0.99, 0.59, 0.18), (0.5, 0.0, 0.08), (0.96, 0.17, 0.89)]

cmap_vals_norm = (cmap_vals - var_min)/(var_max - var_min)

cmap_list = []

for i in range(size(cmap_vals)):

cmap_list.append((cmap_vals_norm[i], cmap_colors[i]))

mf_cmap = LinearSegmentedColormap.from_list('melt_freeze', cmap_list)

else:

# No refreezing

cmap_vals = array([0, 0.25*var_max, 0.5*var_max, 0.75*var_max, var_max])

cmap_colors = [(1, 1, 1), (1, 0.9, 0.4), (0.99, 0.59, 0.18), (0.5, 0.0, 0.08), (0.96, 0.17, 0.89)]

cmap_vals_norm = cmap_vals/var_max

cmap_list = []

for i in range(size(cmap_vals)):

cmap_list.append((cmap_vals_norm[i], cmap_colors[i]))

mf_cmap = LinearSegmentedColormap.from_list('melt_freeze', cmap_list)

colour_map = mf_cmap

elif var == 'vel':

colour_map = 'cool'

else:

colour_map = 'jet'

if var == 'vel':

# Make vectors for overlay

# Set up bins (edges)

x_bins = linspace(x_min[index], x_max[index], num=num_bins+1)

y_bins = linspace(y_min[index], y_max[index], num=num_bins+1)

# Calculate centres of bins (for plotting)

x_centres = 0.5*(x_bins[:-1] + x_bins[1:])

y_centres = 0.5*(y_bins[:-1] + y_bins[1:])

# ROMS

# First set up arrays to integrate velocity in each bin

# Simple averaging of all the points inside each bin

roms_u = zeros([size(y_centres), size(x_centres)])

roms_v = zeros([size(y_centres), size(x_centres)])

roms_num_pts = zeros([size(y_centres), size(x_centres)])

# First convert to polar coordinates, rotate to account for

# longitude in circumpolar projection, and convert back to vector

# components

theta_roms = arctan2(v_rho, u_rho)

theta_circ_roms = theta_roms - roms_lon*deg2rad

u_circ_roms = roms_data*cos(theta_circ_roms) # roms_data is speed

v_circ_roms = roms_data*sin(theta_circ_roms)

# Loop over all points (can't find a better way to do this)

for j in range(size(roms_data,0)):

for i in range(size(roms_data,1)):

# Make sure data isn't masked (i.e. land)

if u_circ_roms[j,i] is not ma.masked:

# Check if we're in the region of interest

if roms_x[j,i] > x_min[index] and roms_x[j,i] < x_max[index] and roms_y[j,i] > y_min[index] and roms_y[j,i] < y_max[index]:

# Figure out which bins this falls into

x_index = nonzero(x_bins > roms_x[j,i])[0][0]-1

y_index = nonzero(y_bins > roms_y[j,i])[0][0]-1

# Integrate

roms_u[y_index, x_index] += u_circ_roms[j,i]

roms_v[y_index, x_index] += v_circ_roms[j,i]

roms_num_pts[y_index, x_index] += 1

# Convert from sums to averages

# First mask out points with no data

roms_u = ma.masked_where(roms_num_pts==0, roms_u)

roms_v = ma.masked_where(roms_num_pts==0, roms_v)

# Divide everything else by the number of points

flag = roms_num_pts > 0

roms_u[flag] = roms_u[flag]/roms_num_pts[flag]

roms_v[flag] = roms_v[flag]/roms_num_pts[flag]

# FESOM low-res

fesom_u_lr = zeros([size(y_centres), size(x_centres)])

fesom_v_lr = zeros([size(y_centres), size(x_centres)])

fesom_num_pts_lr = zeros([size(y_centres), size(x_centres)])

theta_fesom_lr = arctan2(node_v_lr, node_u_lr)

theta_circ_fesom_lr = theta_fesom_lr - fesom_lon_lr*deg2rad

u_circ_fesom_lr = node_data_lr*cos(theta_circ_fesom_lr) # node_data is speed

v_circ_fesom_lr = node_data_lr*sin(theta_circ_fesom_lr)

# Loop over 2D nodes to fill in the velocity bins as before

for n in range(fesom_n2d_lr):

if fesom_x_lr[n] > x_min[index] and fesom_x_lr[n] < x_max[index] and fesom_y_lr[n] > y_min[index] and fesom_y_lr[n] < y_max[index]:

x_index = nonzero(x_bins > fesom_x_lr[n])[0][0]-1

y_index = nonzero(y_bins > fesom_y_lr[n])[0][0]-1

fesom_u_lr[y_index, x_index] += u_circ_fesom_lr[n]

fesom_v_lr[y_index, x_index] += v_circ_fesom_lr[n]

fesom_num_pts_lr[y_index, x_index] += 1

fesom_u_lr = ma.masked_where(fesom_num_pts_lr==0, fesom_u_lr)

fesom_v_lr = ma.masked_where(fesom_num_pts_lr==0, fesom_v_lr)

flag = fesom_num_pts_lr > 0

fesom_u_lr[flag] = fesom_u_lr[flag]/fesom_num_pts_lr[flag]

fesom_v_lr[flag] = fesom_v_lr[flag]/fesom_num_pts_lr[flag]

# FESOM high-res

fesom_u_hr = zeros([size(y_centres), size(x_centres)])

fesom_v_hr = zeros([size(y_centres), size(x_centres)])

fesom_num_pts_hr = zeros([size(y_centres), size(x_centres)])

theta_fesom_hr = arctan2(node_v_hr, node_u_hr)

theta_circ_fesom_hr = theta_fesom_hr - fesom_lon_hr*deg2rad

u_circ_fesom_hr = node_data_hr*cos(theta_circ_fesom_hr) # node_data is speed

v_circ_fesom_hr = node_data_hr*sin(theta_circ_fesom_hr)

# Loop over 2D nodes to fill in the velocity bins as before

for n in range(fesom_n2d_hr):

if fesom_x_hr[n] > x_min[index] and fesom_x_hr[n] < x_max[index] and fesom_y_hr[n] > y_min[index] and fesom_y_hr[n] < y_max[index]:

x_index = nonzero(x_bins > fesom_x_hr[n])[0][0]-1

y_index = nonzero(y_bins > fesom_y_hr[n])[0][0]-1

fesom_u_hr[y_index, x_index] += u_circ_fesom_hr[n]

fesom_v_hr[y_index, x_index] += v_circ_fesom_hr[n]

fesom_num_pts_hr[y_index, x_index] += 1

fesom_u_hr = ma.masked_where(fesom_num_pts_hr==0, fesom_u_hr)

fesom_v_hr = ma.masked_where(fesom_num_pts_hr==0, fesom_v_hr)

flag = fesom_num_pts_hr > 0

fesom_u_hr[flag] = fesom_u_hr[flag]/fesom_num_pts_hr[flag]

fesom_v_hr[flag] = fesom_v_hr[flag]/fesom_num_pts_hr[flag]

# Plot

fig = figure(figsize=(20, ysize[index]))

fig.patch.set_facecolor('white')

# MetROMS

ax = fig.add_subplot(1,3,1, aspect='equal')

# First shade land and zice in grey

contourf(roms_x, roms_y, land_zice, 1, colors=(('0.6', '0.6', '0.6')))

# Fill in the missing circle

contourf(x_reg_roms, y_reg_roms, land_circle, 1, colors=(('0.6', '0.6', '0.6')))

# Now shade the data

pcolor(roms_x, roms_y, roms_data, vmin=var_min, vmax=var_max, cmap=colour_map)

if var == 'vel':

# Overlay vectors

quiver(x_centres, y_centres, roms_u, roms_v, scale=1.5, headwidth=6, headlength=7, color='black')

xlim([x_min[index], x_max[index]])

ylim([y_min[index], y_max[index]])

axis('off')

title('MetROMS', fontsize=24)

# FESOM low-res

ax = fig.add_subplot(1,3,2, aspect='equal')

# Start with land background

contourf(x_reg_fesom, y_reg_fesom, land_square, 1, colors=(('0.6', '0.6', '0.6')))

# Add elements

if var in ['draft', 'melt', 'wct']:

# Ice shelf elements only

img = PatchCollection(patches_shelf_lr, cmap=colour_map)

else:

img = PatchCollection(patches_lr, cmap=colour_map)

img.set_array(array(fesom_data_lr))

img.set_edgecolor('face')

img.set_clim(vmin=var_min, vmax=var_max)

ax.add_collection(img)

if var in ['draft', 'melt', 'wct']:

# Mask out the open ocean in white

overlay = PatchCollection(patches_ocn_lr, facecolor=(1,1,1))

overlay.set_edgecolor('face')

ax.add_collection(overlay)

if var == 'vel':

# Overlay vectors

quiver(x_centres, y_centres, fesom_u_lr, fesom_v_lr, scale=1.5, headwidth=6, headlength=7, color='black')

xlim([x_min[index], x_max[index]])

ylim([y_min[index], y_max[index]])

axis('off')

title('FESOM (low-res)', fontsize=24)

# FESOM high-res

ax = fig.add_subplot(1,3,3, aspect='equal')

contourf(x_reg_fesom, y_reg_fesom, land_square, 1, colors=(('0.6', '0.6', '0.6')))

if var in ['draft', 'melt', 'wct']:

# Ice shelf elements only

img = PatchCollection(patches_shelf_hr, cmap=colour_map)

else:

img = PatchCollection(patches_hr, cmap=colour_map)

img.set_array(array(fesom_data_hr))

img.set_edgecolor('face')

img.set_clim(vmin=var_min, vmax=var_max)

ax.add_collection(img)

if var in ['draft', 'melt', 'wct']:

overlay = PatchCollection(patches_ocn_hr, facecolor=(1,1,1))

overlay.set_edgecolor('face')

ax.add_collection(overlay)

if var == 'vel':

# Overlay vectors

quiver(x_centres, y_centres, fesom_u_hr, fesom_v_hr, scale=1.5, headwidth=6, headlength=7, color='black')

xlim([x_min[index], x_max[index]])

ylim([y_min[index], y_max[index]])

axis('off')

title('FESOM (high-res)', fontsize=24)

# Colourbar on the right

cbaxes = fig.add_axes([0.92, 0.2, 0.01, 0.6])

cbar = colorbar(img, cax=cbaxes)

cbar.ax.tick_params(labelsize=20)

# Main title

if var == 'draft':

title_string = ' draft (m)'

elif var == 'bathy':

title_string = ' bathymetry (m)'

elif var == 'wct':

title_string = ' water column thickness (m)'

elif var == 'melt':

title_string = ' melt rate (m/y)'

elif var == 'temp':

title_string = r' bottom water temperature ($^{\circ}$C)'

elif var == 'salt':

title_string = ' bottom water salinity (psu)'

elif var == 'vel':

title_string = ' vertically averaged ocean velocity (m/s)'

suptitle(region_names[index] + title_string, fontsize=30)

subplots_adjust(wspace=0.05)

#fig.show()