#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Filters
"""
import logging
import xarray as xr
import numpy as np
from xoa.filter import erode_mask
from .cf import get_x_dim_name, get_y_dim_name
LOGGER = logging.getLogger(__name__)
[docs]def erode_coast(da, kernel=3, niter=1):
"""
Erode the coast using a local average
Parameters
----------
da : xarray.DataArray
Data array to erode.
kernel : int, dict, tuple, optional
Size of kernel in X and Y directions.
Returns
-------
xarray.DataArray.
"""
# Masked data?
mask = da.isnull().data
if not mask.any():
return da
valid = ~mask
# Local import to always import scipy
from scipy.signal import convolve
# kernel
if not isinstance(kernel, dict):
x_dim = get_x_dim_name(da)
y_dim = get_y_dim_name(da)
if isinstance(kernel, int):
kernel = (kernel, kernel)
kernel = {x_dim: kernel[0], y_dim: kernel[1]}
kk = ()
for dim in da.dims:
kk += kernel.get(dim, 1),
kernel = np.ones(kk, dtype='l')
# Convolve
vnum = da.fillna(0.)
for i in range(niter):
vnum = convolve(vnum, kernel, mode='same')
vdenom = convolve((~mask).astype('l'), kernel, mode='same')
mask = vdenom == 0
vdenom[mask] = 1
vnum /= vdenom
del vdenom
# Output
return da.where(valid, np.where(mask, np.nan, vnum))
[docs]def erode_mask_vec(ds, kernel, param=None, until=5):
"""
Erode the mask of a vectorized parameter
Parameters
----------
da : xarray.DataArray
Data array to erode.
kernel : int, dict, tuple, optional
Size of kernel in X and Y directions.
param : str, optional
Type of vectorized parameter
Returns
-------
xarray.DataArray.
"""
LOGGER.info(f"Erode mask on {param} vectorized variables")
if param == 'wind':
tau = (ds.ewss**2 + ds.nsss**2)**0.5
si10 = ds.si10**2
si10 = si10.where(si10 > 1e-6, 1e-6)
rhocd = tau / si10
rhocd1 = rhocd.where(ds.mask == 1)
rhocd2 = erode_mask(rhocd1, kernel=kernel, until=until)
maskd = rhocd2.where(xr.ufuncs.isfinite(rhocd2), np.nan)
maskd = maskd.where(xr.ufuncs.isnan(maskd), 1)
w = rhocd2 / rhocd
w = w.where(w < 1, 1)
for var in ['ewss', 'nsss']:
ds[var] = ds[var] * w * maskd
return ds