bdy_var_plot.py

#!/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_y2017m11.nc'
print(fname)
ind, dst, brk = nemo_bdy_order(fname)
f = plot_bdy(fname, ind, dst, brk, 'thetao', 0, 0)

plt.show()