OceanParcels recipes¶

import numpy as np
import xarray as xr
import os
from matplotlib import pyplot as plt, animation
from datetime import datetime, timedelta
from dateutil.parser import parse
from IPython.display import HTML
from salishsea_tools import nc_tools, places

from parcels import FieldSet, Field, VectorField, ParticleSet, JITParticle, ErrorCode, AdvectionRK4, AdvectionRK4_3D

%matplotlib inline

plt.rcParams['font.size'] = 12

def fieldset_from_nemo(daterange, coords, flat=True):
    """Generate a fieldset from a hourly SalishSeaCast forcing fields
    over daterange.
    """

    # Generate sequential list of forcing file prefixes
    prefixes = [
        nc_tools.get_hindcast_prefix(daterange[0] + timedelta(days=d))
        for d in range(np.diff(daterange)[0].days + 1)
    ]

    # Predefine fieldset argument dictionaries
    filenames, variables, dimensions = {}, {}, {}

    # Define dict fields for each variable
    for var, name in zip(['U', 'V', 'W'], ['vozocrtx', 'vomecrty', 'vovecrtz']):
        
        # Exclude vertical velocity if 2D
        if flat:
            if var == 'W': break

        # Dict of filenames containing the coordinate and forcing variables
        datafiles = [prefix + f'_grid_{var}.nc' for prefix in prefixes]
        filenames[var] = {'lon': coords, 'lat': coords, 'data': datafiles}

        # NEMO variable name
        variables[var] = name

        # Dict of NEMO coordinate names (f-points)
        dimensions[var] = {'lon': 'glamf', 'lat': 'gphif', 'time': 'time_counter'}
        
        # Add depth fields if 3D (f-points are on W grid)
        if not flat:
            filenames[var]['depth'] = prefixes[0] + '_grid_W.nc'
            dimensions[var]['depth'] = 'depthw'

    # Load NEMO forcing into fieldset
    field_set = FieldSet.from_nemo(filenames, variables, dimensions, allow_time_extrapolation=True)
    
    return field_set


def append_auxiliary_forcing(fieldset, model, daterange, coords_path):
    """Append hourly GEM or WW3 forcing fields to a fieldset over daterange
    """

    # Generate sequential list of forcing files
    datafiles = [
        getattr(nc_tools, f'get_{model}_path')(daterange[0] + timedelta(days=d))
        for d in range(np.diff(daterange)[0].days + 1)
    ]
    
    # Assign variable suffix and dimensions definitions according to model
    if model == 'GEM':
        suffix = '_wind'
        dimensions = {'lon': 'nav_lon', 'lat': 'nav_lat', 'time': 'time_counter'}
    elif model == 'WW3':
        suffix = 'uss'
        dimensions = {'lon': 'longitude', 'lat': 'latitude', 'time': 'time'}
    else:
        raise ValueError(f'Unknown surface forcing model: {model}')
    
    # Define coordinates file path and create if necessary
    #  - only method I've had success with for correcting lons
    coords = os.path.join(coords_path, f'{model}_grid.nc')
    if not os.path.exists(coords):
        with xr.open_dataset(datafiles[0]) as ds:
            data_vars = [dimensions['time']] + list(ds.data_vars)
            if model == 'GEM':
                for key in ['lon', 'lat']: data_vars.remove(dimensions[key])
            ds = ds.drop_vars(data_vars)
            ds.update({dimensions['lon']: ds[dimensions['lon']] - 360})
            ds.to_netcdf(coords)
    
    # Filenames dict
    filenames = {'lon': coords, 'lat': coords, 'data': datafiles}
    
    # Load u velocity
    u = Field.from_netcdf(
        filenames, f'u{suffix}', dimensions, fieldtype='U', allow_time_extrapolation=True,
    )
    
    # Load v velocity with u grid
    v = Field.from_netcdf(
        filenames, f'v{suffix}', dimensions, fieldtype='V', allow_time_extrapolation=True,
        grid=u.grid, dataFiles=u.dataFiles,
    )

    # Add velocity fields to fieldset
    fieldset.add_field(u)
    fieldset.add_field(v)

    # Make new vector field from u and v for use in kernels
    u, v = [getattr(fieldset, f'{var}{suffix}') for var in ('u', 'v')]
    fieldset.add_vector_field(VectorField(f"UV_{model}", u, v))

def WindDrift(particle, fieldset, time):
    (uw, vw) = fieldset.UV_GEM[time, particle.depth, particle.lat, particle.lon]
    particle.lon += uw * 0.03 * particle.dt
    particle.lat += vw * 0.03 * particle.dt


def StokesDrift(particle, fieldset, time):
    (us, vs) = fieldset.UV_WW3[time, particle.depth, particle.lat, particle.lon]
    particle.lon += us * particle.dt
    particle.lat += vs * particle.dt


def DeleteParticle(particle, fieldset, time):
    print(f'Particle {particle.id} lost !! [{particle.lon}, {particle.lat}, {particle.depth}, {particle.time}]')
    particle.delete()

# Paths and filenames
paths = {
    'coords': '/data/bmoorema/MEOPAR/grid/coordinates_seagrid_SalishSea201702.nc',
    'mask': '/data/bmoorema/MEOPAR/grid/mesh_mask201702.nc',
    'results': '/data/bmoorema/results/parcels/sandbox',
}

# Load coords and mask files and extract grid variables
coords, mask = [xr.open_dataset(paths[key], decode_times=False) for key in ('coords', 'mask')]
gridlon, gridlat = [coords[key][0, ...].values for key in ('glamt', 'gphit')]
tmask = mask.tmask[0, 0, ...].values

# Define release parameters
location = 'Sandheads'
n = 100   # number of particles
r = 50   # radius of particle cloud [m]

# Start time, duration and timestep
start = datetime(2019, 1, 1, 12, 30, 0)
duration = timedelta(days=3)
dt = timedelta(seconds=90)

# Create Gaussian distribution around release point
mean, cov = [0, 0], [[r**2, 0], [0, r**2]]
x_offset, y_offset = np.random.multivariate_normal(mean, cov, n).T
lon, lat = places.PLACES[location]['lon lat']
lons = lon + x_offset / 111000 / np.cos(np.deg2rad(lat))
lats = lat + y_offset / 111000

# Forcing daterange (I add 1-day buffers)
daterange = [start - timedelta(days=1), start + duration + timedelta(days=1)]

# Output file prefix
strings = [location] + [t.strftime('%Y%m%dT%H%M%S') for t in (start, start + duration)]
prefix = os.path.join(paths['results'], '_'.join(strings))

%%capture

# Make initial figure
fig, axs = plt.subplots(1, 2, figsize=(15, 8), gridspec_kw={'wspace': 0.1})
l1, = axs[0].plot([], [], 'ko')
l2, = axs[1].plot([], [], 'ko')
t = axs[0].text(0.02, 0.02, '', transform=axs[0].transAxes)
data = {}
for ax, config in zip(axs, ['NEMO_2D', 'NEMOGEMWW3_2D']):
    data[config] = xr.open_dataset(prefix + f'_{config}.nc')
    ax.contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray')
    ax.contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')
    ax.set_xlim([-124.3, -123])
    ax.set_ylim([48.7, 49.6])
    ax.set_title(config)
    ax.set_aspect(1/np.sin(np.deg2rad(49)))
axs[1].yaxis.set_ticklabels('')

# Init function
def init():
    t.set_text('')
    for l in [l1, l2]:
        l.set_data([], [])
    return l1, l2, t,

# Animate function
def animate(hour):
    tstamp = data['NEMO_2D'].time[0, hour].values.astype('datetime64[s]').astype(datetime)
    t.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))
    for l, config in zip([l1, l2], ['NEMO_2D', 'NEMOGEMWW3_2D']):
        l.set_data(data[config].lon[:, hour], data[config].lat[:, hour])
    return l1, l2, t,

# Build animation
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)

# Load SalishSeaCast results into fieldset
fieldset = fieldset_from_nemo(daterange, paths['coords'], flat=False)

%%capture

# Make initial figure
fig, ax = plt.subplots(figsize=(8, 8))
cax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
l = ax.scatter([], [], s=50, c=[], vmin=0, vmax=10, edgecolor='k')
t = ax.text(0.02, 0.02, '', transform=ax.transAxes)
data = xr.open_dataset(prefix + '_NEMO_3D.nc')
ax.contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray')
ax.contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124.3, -123])
ax.set_ylim([48.7, 49.6])
ax.set_title('NEMO_3D')
ax.set_aspect(1/np.sin(np.deg2rad(49)))
fig.colorbar(l, cax=cax, label='Depth [m]')

# Init function
def init():
    t.set_text('')
    l.set_offsets(np.empty((0, 2)))
    l.set_array(np.empty(0))
    return l, t,

# Animate function
def animate(hour):
    tstamp = data.time[0, hour].values.astype('datetime64[s]').astype(datetime)
    t.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))
    l.set_offsets(np.vstack([data.lon[:, hour], data.lat[:, hour]]).T)
    l.set_array(data.z[:, hour])
    return l, t,

# Build animation
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)