pyfvcom諧波分析示例 pyfvcom_harmonic_analysis_example

%matplotlib inline

from __future__ import print_function

import sys

import numpy as np
import matplotlib.pyplot as plt

from mpl_toolkits.basemap import Basemap
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.collections import EllipseCollection
from matplotlib import rcParams
from matplotlib import cm
from netCDF4 import num2date

from PyFVCOM.read_FVCOM_results import ncread
from PyFVCOM.tidal_ellipse import ap2ep
from tappy import TAPPY

# Load model output.
fvcom = 'sample.nc'

varlist = ('lonc', 'latc', 'ua', 'va', 'time')
dims = {'time': ':360'}  # first 15 days at hourly sampling

# Define a plot subset ((xmin, xmax), (ymin, ymax)).
subset = np.array(((-4.2416, -4.0837), (50.2656, 50.3966)))

# Scaling factor for the ellipses. You will need to experiment with this
# value.
scaling = 2000

# Find the model nodes which fall within the subset defined above.
FVCOM = ncread(fvcom, vars=varlist, dims=dims, noisy=False)

# Create a time array for the TAPPy call.
FVCOM['datetimes'] = num2date(FVCOM['time'], 'days since 1858-11-17 00:00:00')
years = [i.year for i in FVCOM['datetimes']]
months = [i.month for i in FVCOM['datetimes']]
days = [i.day for i in FVCOM['datetimes']]
hours = [i.hour for i in FVCOM['datetimes']]
minutes = [i.minute for i in FVCOM['datetimes']]
seconds = [i.second for i in FVCOM['datetimes']]
Times = np.column_stack((years, months, days, hours, minutes, seconds))

# Find the indices of the locations which fall within the subset.
elems = np.where((FVCOM['lonc'] > subset[0].min()) *
                 (FVCOM['lonc'] < subset[0].max()) *
                 (FVCOM['latc'] > subset[1].min()) *
                 (FVCOM['latc'] < subset[1].max()))[0]

# Create dicts for the results.
uharmonics, vharmonics = {}, {}

# Loop over the indices within the region of interest. This will take a 
# while in serial. Implementing this in parallel is left as an exercise
# for the reader!
for i, idx in enumerate(elems):
    
    if i == 0:
        print('{} of {} '.format(i + 1, len(elems)), end='')
    elif (i + 1) % 100 == 0:
        print('\n{} of {} '.format(i + 1, len(elems)), end='')
    elif i + 1 == len(elems):
        print(' done.')
    else:
        print('.', end='')
    sys.stdout.flush()

    # Combine the Times and velocity data.
    u = np.column_stack((Times, FVCOM['ua'][:, idx]))
    v = np.column_stack((Times, FVCOM['va'][:, idx]))

    # Create a dict for the TAPPy results.
    uharm, vharm = {}, {}
    uharm['name'], uharm['speed'], uharm['phase'], uharm['amp'], uharm['infer'] = TAPPY(u)
    vharm['name'], vharm['speed'], vharm['phase'], vharm['amp'], vharm['infer'] = TAPPY(v)

    # Put the results into a meta dict with the key being the current
    # position.
    key = '{}-{}'.format(FVCOM['lonc'][idx], FVCOM['latc'][idx])
    uharmonics[key] = uharm
    vharmonics[key] = vharm

# Create the mapping object ready for the plot.
m = Basemap(llcrnrlon=subset[0].min(),
            llcrnrlat=subset[1].min(),
            urcrnrlon=subset[0].max(),
            urcrnrlat=subset[1].max(),
            rsphere=(6378137.00, 6356752.3142),
            resolution='f',
            area_thresh=0.05,
            projection='merc',
            lat_0=subset[1].mean(),
            lon_0=subset[0].mean(),
            lat_ts=subset[0].mean())

parallels = np.arange(np.floor(subset[1].min()),
                      np.ceil(subset[1].max()), 0.025)
meridians = np.arange(np.floor(subset[0].min()),
                      np.ceil(subset[0].max()), 0.05)

x, y = m(FVCOM['lonc'], FVCOM['latc'])

# Extract all the ellipses for the current subset so we can plot them 
# more easily.

xy = []
widths = []
heights = []
rotations = []
sema = []
for em in elems:
    # Make a key from the original positions. This is the key which we
    # use to extract the harmonic results.
    key = '{}-{}'.format(FVCOM['lonc'][em], FVCOM['latc'][em])
    
    # Find the M2 data position for the current position.
    idx = uharmonics[key]['name'].index('M2')
    
    # Extract the amplitude and phase and calculate the ellipse 
    # properties.
    uZ = uharmonics[key]['amp'][idx]
    vZ = vharmonics[key]['amp'][idx]
    uG = uharmonics[key]['phase'][idx]
    vG = vharmonics[key]['phase'][idx]
    SEMA, ECC, INC, PHA, w = ap2ep(uZ, uG, vZ, vG)

    xy.append((x[em], y[em]))
    # widths and heights are flipped so the rotation is effectively
    # transformed from clockwise from the y-axis to anti-clockwise 
    # relative to the x-axis. This is what EllipseCollection needs
    # for its angles.
    widths.append(scaling * SEMA)
    heights.append(scaling * SEMA * ECC)
    rotations.append(INC)
    # Save the semi-major axis for the colouring the ellipses.
    sema.append(SEMA)
      
# Now plot the M2 results.
rcParams['mathtext.default'] = 'regular'  # sensible font for LaTeX text

fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(111, aspect='equal')

# Make an EllipseCollection.
ellipses = EllipseCollection(widths, heights, rotations, 
                             offsets=xy,
                             units='xy',
                             transOffset=ax.transData)
ellipses.set_array(np.asarray(sema))

# Add coastlines
m.drawmapboundary(zorder=0)
m.drawcoastlines(zorder=1)
m.fillcontinents(zorder=0, color='0.6')
m.drawparallels(parallels, labels=[1, 0, 0, 0], linewidth=0)
m.drawmeridians(meridians, labels=[0, 0, 0, 1], linewidth=0)

# Add ellipses coloured by the semi-major axis magnitude.
ax.add_collection(ellipses)
ellipses.set_linewidth(0)
ellipses.set_cmap(cm.viridis)
ellipses.set_zorder=200
    
# Add a nice colour bar.
div = make_axes_locatable(ax)
cax = div.append_axes("right", size="5%", pad=0.2)
cb = fig.colorbar(ellipses, cax=cax)
cb.set_label("Semi-major axis ($ms^{-1}$)")

ax.set_title('$M_2$ tidal ellipses')