def test_truncation(da_urho,
da_vrho,
trunc_vals=[10, 20, 30, 40, 50, 100, 200]):
for idx, trunc in enumerate(trunc_vals):
print(f'Done {int(100*idx/da_urho.time.values.shape[0])} %')
da_urho_ = da_urho.isel(time=slice(0, 8))
da_vrho_ = da_vrho.isel(time=slice(0, 8))
w = VectorWind(da_urho_, da_vrho_)
da_urho_ = w.truncate(da_urho_, truncation=trunc)
da_vrho_ = w.truncate(da_vrho_, truncation=trunc)
lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4)
ftle = lcs(u=da_urho_,
v=da_vrho_,
isglobal=True,
s=4e6,
interp_to_common_grid=True,
traj_interp_order=3)
ftle = .5 * np.log(ftle)
toplot = ftle
fig, ax = plt.subplots(1,
1,
subplot_kw={'projection': ccrs.Robinson()})
toplot.plot(ax=ax,
transform=ccrs.PlateCarree(),
robust=True,
cbar_kwargs={'shrink': .6},
vmin=.5,
vmax=2)
ax.coastlines()
ax.set_title(f'N96 truncated at t{trunc}')
plt.savefig(f'upscale/figs/trunc_test/{trunc:03d}.png',
dpi=600,
boundary='trim',
bbox_inches='tight')
plt.close()
def calc_ftle(da_urho_, da_vrho_, truncation=20):
print('********************************** \n'
'Truncating wind and computing FTLE\n'
'**********************************')
w = VectorWind(da_urho_, da_vrho_)
da_urho_ = w.truncate(da_urho_, truncation=20)
da_vrho_ = w.truncate(da_vrho_, truncation=20)
lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4)
ftle = lcs(u=da_urho_,
v=da_vrho_,
isglobal=True,
s=4e6,
interp_to_common_grid=True,
traj_interp_order=3)
ftle = .5 * np.log(ftle)
return ftle
def __call__(self,
ds: xr.Dataset = None,
u: xr.DataArray = None,
v: xr.DataArray = None,
verbose=True,
s=None,
resample=None,
s_is_error=False,
isglobal=False,
return_traj=False,
interp_to_common_grid=True,
traj_interp_order=3,
truncation=20) -> xr.DataArray:
"""
:param ds: xarray.Dataset, optional
xarray dataset containing u and v as variables. Mutually exclusive with parameters u and v.
:param u: xarray.Dataarray, optional
xarray datarray containing u-wind component. Mutually exclusive with parameter ds.
:param v: xarray.Dataarray, optional
xarray datarray containing u-wind component. Mutually exclusive with parameter ds.
:param verbose: bool, optional
Whether to print intermediate values
:return: xarray.Dataarray with the Finite-Time Lyapunov vector
>>> subtimes_len = 1
>>> timestep = -6*3600 # 6 hours in seconds
>>> lcs = LCS(timestep=timestep, timedim='time', subtimes_len=subtimes_len)
>>> ds = sampleData()
>>> ftle = lcs(ds, verbose=False)
"""
global verboseprint
print('!' * 100)
verboseprint = print if verbose else lambda *a, **k: None
timestep = self.timestep
timedim = self.timedim
self.verbose = verbose
if isinstance(ds, xr.Dataset):
u = ds.u.copy()
v = ds.v.copy()
elif isinstance(ds, str):
ds = xr.open_dataset(ds)
u = ds.u.copy()
v = ds.v.copy()
if isinstance(resample, str):
u = u.resample({self.timedim: resample}).interpolate('linear')
v = v.resample({self.timedim: resample}).interpolate('linear')
timestep = np.sign(timestep) * (
u[self.timedim].values[1] - u[self.timedim].values[0]
).astype('timedelta64[s]').astype('float')
u_dims = u.dims
v_dims = v.dims
assert set(u_dims) == set(v_dims), "u and v dims are different"
assert set(u_dims) == {
'latitude', 'longitude', timedim
}, 'array dims should be latitude and longitude only'
# Ascending order is required
# TODO: Add checks for continuity
u = u.sortby('latitude')
u = u.sortby('longitude')
v = v.sortby('latitude')
v = v.sortby('longitude')
if isglobal:
if interp_to_common_grid:
lats = np.linspace(-89.75, 89.75, 180 * 2)
lons = np.linspace(-180, 179.5, 360 * 2 + 1)
u_reindex = u.reindex(latitude=lats,
longitude=lons,
method='nearest')
v_reindex = v.reindex(latitude=lats,
longitude=lons,
method='nearest')
u_interp = u.interp(latitude=lats,
longitude=lons,
method='linear')
v_interp = v.interp(latitude=lats,
longitude=lons,
method='linear')
u = u_interp.where(~xr.ufuncs.isnan(u_interp), u_reindex)
v = v_interp.where(~xr.ufuncs.isnan(v_interp), v_reindex)
if truncation is not None:
w = VectorWind(u, v)
u = w.truncate(u, truncation=truncation)
v = w.truncate(v, truncation=truncation)
cyclic_xboundary = True
self.subdomain = None
else:
cyclic_xboundary = False
if s is None:
s = int(10 * u.isel({
timedim: 0
}).size * u.isel({
timedim: 0
}).std())
print(f'using s = ' + str(s / 1e6) + '1e6')
verboseprint("*---- Parcel propagation ----*")
x_departure, y_departure = parcel_propagation(
u,
v,
timestep,
propdim=self.timedim,
SETTLS_order=self.SETTLS_order,
verbose=verbose,
cyclic_xboundary=cyclic_xboundary,
return_traj=return_traj,
interp_order=traj_interp_order,
copy=True)
if return_traj:
x_trajs = x_departure.copy()
y_trajs = y_departure.copy()
x_departure = x_departure.isel({timedim: -1})
y_departure = y_departure.isel({timedim: -1})
verboseprint("*---- Computing deformation tensor ----*")
def_tensor = flowmap_gradient(x_departure,
y_departure,
sigma=self.gauss_sigma)
if isinstance(self.subdomain, dict):
def_tensor = latlonsel(def_tensor, **self.subdomain)
def_tensor = def_tensor.stack({'points': ['latitude', 'longitude']})
def_tensor = def_tensor.dropna(dim='points')
# eigenvalues = xr.apply_ufunc(lambda x: compute_eigenvalues(x), def_tensor.groupby('points'),
# dask='parallelized',
# output_dtypes=[float]
# )
verboseprint("*---- Computing eigenvalues ----*")
vals = def_tensor.transpose(..., 'points').values
vals = vals.reshape([3, 3, def_tensor.shape[-1]])
def_tensor_norm = norm(vals, axis=(0, 1), ord=2)
def_tensor_norm = def_tensor.isel(
derivatives=0).drop('derivatives').copy(data=def_tensor_norm)
verboseprint("*---- Done eigenvalues ----*")
def_tensor_norm = def_tensor_norm.unstack('points')
timestamp = u[self.timedim].values[-1] if np.sign(
timestep) == 1 else u[self.timedim].values[0]
def_tensor_norm['time'] = timestamp
eigenvalues = def_tensor_norm.expand_dims(self.timedim)
if self.return_dpts and return_traj:
return eigenvalues, x_departure, y_departure, x_trajs, y_trajs
elif self.return_dpts:
return eigenvalues, x_departure, y_departure
elif return_traj:
return eigenvalues, x_trajs, y_trajs
else:
return eigenvalues
