#!/usr/bin/env python2
# -*- coding: utf-8 -*-
import numpy as np
import scipy.spatial as sp
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
def nemo_bdy_order(fname):
"""
Determine the ordering and breaks in BDY files to aid plotting.
This function takes the i/j coordinates from BDY files and orders them sequentially
making it easier to visualise sections along the open boundary. Breaks in the open
boundary are also determined (i.e. where the distance between two points > 2**0.5)
Args:
fname (str) : filename of BDY file
Returns:
bdy_ind (dict): re-ordered indices
bdy_dst (dict): distance (in model coords) between points
bdy_brk (dict): location of the break points in the open boundary
"""
# open file pointer and extract data
rootgrp = Dataset(fname, "r", format="NETCDF4")
nbi = np.squeeze(rootgrp.variables['nbidta'][:, :]) - 1 # subtract 1 for python indexing
nbj = np.squeeze(rootgrp.variables['nbjdta'][:, :]) - 1
nbr = np.squeeze(rootgrp.variables['nbrdta'][:, :]) - 1
lon = np.squeeze(rootgrp.variables['nav_lon'][:, :])
lat = np.squeeze(rootgrp.variables['nav_lat'][:, :])
rootgrp.close()
rw = np.amax(nbr) + 1
bdy_ind = {}
bdy_brk = {}
bdy_dst = {}
nbdy = []
for r in range(rw):
nbdy.append(np.sum(nbr == r))
# TODO: deal with domain that spans wrap
# start with outer rim and work in
for r in range(rw):
# set initial constants
ind = nbr == r
nbi_tmp = nbi[ind]
nbj_tmp = nbj[ind]
count = 1
id_order = np.zeros((nbdy[r], 1), dtype=int) - 1
id_order[0,] = 0
flag = False
mark = 0
source_tree = sp.cKDTree(zip(nbi_tmp, nbj_tmp), balanced_tree=False, compact_nodes=False)
# order bdy entries
while count < nbdy[r]:
nn_dist, nn_id = source_tree.query(zip(nbi_tmp[id_order[count - 1]], nbj_tmp[id_order[count - 1]]),
k=3, distance_upper_bound=2.9)
if np.sum(id_order == nn_id[0, 1]) == 1: # is the nearest point already in the list?
if np.sum(id_order == nn_id[0, 2]) == 1: # is the 2nd nearest point already in the list?
if flag == False: # then we've found an end point and we can start the search in earnest!
flag = True
id_order[mark] = id_order[count - 1] # reset values
id_order[mark + 1:] = -1 # reset values
count = mark + 1 # reset counter
else:
i = 0 # should this be zero?
t = count
while count == t:
if np.sum(id_order == i) == 1:
i += 1
else:
id_order[count] = i
flag = False
mark = count
count += 1
elif nn_id[0, 2] == nbdy[r]:
i = 0
t = count
while count == t:
if np.sum(id_order == i) == 1:
i += 1
else:
id_order[count] = i
flag = False
mark = count
count += 1
else:
id_order[count] = nn_id[0, 2]
count += 1
else:
id_order[count] = nn_id[0, 1]
count += 1
bdy_ind[r] = id_order
bdy_dst[r] = np.sqrt((np.diff(np.hstack((nbi_tmp[id_order], nbj_tmp[id_order])), axis=0) ** 2).sum(axis=1))
bdy_brk[r], = np.where(bdy_dst[r] > 2 ** 0.5)
bdy_brk[r] += 1
bdy_brk[r] = np.insert(bdy_brk[r], 0, 0) # insert start point
bdy_brk[r] = np.insert(bdy_brk[r], len(bdy_brk[r]), len(id_order)) # insert end point
bdy_dst[r] = np.insert(np.cumsum(bdy_dst[r]), 0, 0)
if r == 2: # change to a valid rw number to get a visual output (outer most rw is zero)
f, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
plt.scatter(nbi_tmp[bdy_ind[r][:]], nbj_tmp[bdy_ind[r][:]], s=1, marker='.')
for t in np.arange(0, len(bdy_ind[r]), 20):
plt.text(nbi_tmp[bdy_ind[r][t]], nbj_tmp[bdy_ind[r][t]], t)
return bdy_ind, bdy_dst, bdy_brk
def plot_bdy(fname, bdy_ind, bdy_dst, bdy_brk, varnam, t, rw):
"""
Determine the ordering and breaks in BDY files to aid plotting.
This function takes the i/j coordinates from BDY files and orders them sequentially
making it easier to visualise sections along the open boundary. Breaks in the open
boundary are also determined (i.e. where the distance between two points > 2**0.5)
Args:
fname (str) : filename of BDY file
Returns:
bdy_ind (dict): re-ordered indices
bdy_dst (dict): distance (in model coords) between points
bdy_brk (dict): location of the break points in the open boundary
"""
# need to write in a check that either i or j are single values
rootgrp = Dataset(fname, "r", format="NETCDF4")
var = np.squeeze(rootgrp.variables[varnam][t, :])
nbr = np.squeeze(rootgrp.variables['nbrdta'][:, :]) - 1
var = var[:, nbr == rw]
# let us use gdept as the central depth irrespective of whether t, u or v
try:
gdep = np.squeeze(rootgrp.variables['deptht'][:, :, :])
except KeyError:
try:
gdep = np.squeeze(rootgrp.variables['gdept'][:, :, :])
except KeyError:
try:
gdep = np.squeeze(rootgrp.variables['depthu'][:, :, :])
except KeyError:
try:
gdep = np.squeeze(rootgrp.variables['depthv'][:, :, :])
except KeyError:
print 'depth variable not found'
rootgrp.close()
jpk = len(gdep[:, 0])
nsc = len(bdy_brk[rw][:]) - 1
dep = {}
dta = {}
# divide data up into sections and re-order
for n in range(nsc):
dta[n] = np.squeeze(var[:, bdy_ind[rw][bdy_brk[rw][n]:bdy_brk[rw][n + 1]]])
dep[n] = np.squeeze(gdep[:, bdy_ind[rw][bdy_brk[rw][n]:bdy_brk[rw][n + 1]]])
# loop over number of sections and plot
f, ax = plt.subplots(nrows=1, ncols=1, figsize=(11, 4))
ax.plot(dta[0][10, :])
ax.set_title('BDY points along section: 1, depth level: 10')
f, ax = plt.subplots(nrows=nsc, ncols=1, figsize=(14, 10 * nsc))
for n in range(nsc):
print('NSC '+str(n))
plt.sca(ax[n])
# from gdep create some pseudo w points
gdept = dep[n][:, :]
coords = np.arange(0, len(gdept[0, :]))
gdepw = np.zeros((len(gdept[:, 0]) + 1, len(gdept[0, :])))
for z in range(jpk):
gdepw[z + 1, :] = gdept[z, :] + (gdept[z, :] - gdepw[z, :])
gdepvw = np.zeros((len(gdept[:, 0]) + 1, len(gdept[0, :]) + 1))
# TODO: put in an adjustment for zps levels
gdepvw[:, 1:-1] = (gdepw[:, :-1] + gdepw[:, 1:]) / 2
gdepvw[:, 0] = gdepvw[:, 1]
gdepvw[:, -1] = gdepvw[:, -2]
# create a pseudo bathymetry from the depth data
bathy = np.zeros_like(coords)
mbath = np.sum(dta[n].mask == 0, axis=0)
for i in range(len(coords)):
bathy[i] = gdepw[mbath[i], i]
bathy_patch = Polygon(np.vstack((np.hstack((coords[0], coords, coords[-1])),
np.hstack((np.amax(bathy[:]), bathy, np.amax(bathy[:]))))).T,
closed=True,
facecolor=(0.8, 0.8, 0), alpha=0, edgecolor=None)
# Add patch to axes
ax[n].add_patch(bathy_patch)
ax[n].set_title('BDY points along section: ' + str(n))
patches = []
colors = []
for i in range(len(coords)):
for k in np.arange(jpk - 2, -1, -1):
if dta[n][k, i] > -10:
x = [coords[i] - 0.5, coords[i], coords[i] + 0.5,
coords[i] + 0.5, coords[i], coords[i] - 0.5, coords[i] - 0.5]
y = [gdepvw[k + 1, i], gdepw[k + 1, i], gdepvw[k + 1, i + 1],
gdepvw[k, i + 1], gdepw[k, i], gdepvw[k, i], gdepvw[k + 1, i]]
polygon = Polygon(np.vstack((x, y)).T, True)
patches.append(polygon)
colors = np.append(colors, dta[n][k, i])
# for i in range(len(coords)):
# #print(i)
# for k in np.arange(jpk - 2, -1, -1):
# x = [coords[i] - 0.5, coords[i], coords[i] + 0.5,
# coords[i] + 0.5, coords[i], coords[i] - 0.5, coords[i] - 0.5]
# y = [gdepvw[k + 1, i], gdepw[k + 1, i], gdepvw[k + 1, i + 1],
# gdepvw[k, i + 1], gdepw[k, i], gdepvw[k, i], gdepvw[k + 1, i]]
# plt.plot(x, y, 'k-', linewidth=0.1)
# plt.plot(coords[i], gdept[k, i], 'k.', markersize=1)
plt.plot(coords, bathy, '-', color=(0.4, 0, 0))
p = PatchCollection(patches, alpha=0.8, edgecolor='none')
p.set_array(np.array(colors))
ax[n].add_collection(p)
f.colorbar(p, ax=ax[n])
ax[n].set_ylim((0, np.max(bathy)))
ax[n].invert_yaxis()
return f
fname = '/Users/thopri/Projects/PyNEMO/outputs/NNA_R12_bdyT_y1979m11.nc'
ind, dst, brk = nemo_bdy_order(fname)
f = plot_bdy(fname, ind, dst, brk, 'votemper', 0, 0)
plt.show()