def process_contour(self, ds: xr.Dataset, y_ind, x_ind):
""" Redefine contour so that each point on the contour defined by
y_ind and x_ind is seperated from its neighbours by a single index change
in y or x, but not both.
example: convert y_ind = [10,11], x_ind = [1,2] to y_ind = [10,10], x_ind = [1,2]
or y_ind = [10,11], x_ind = [1,1]
Parameters
----------
nemo : xarray.Dataset
xarray Dataset from supplied nemo object
y_ind : numpy.ndarray
1d array of y indices defining the contour on the model grid
x_ind : numpy.ndarray
1d array of x indices defining the contour on the model grid
Returns
-------
y_ind : numpy.ndarray
processed y indices of the contour on the model grid
x_ind : numpy.ndarray
processed x indices of the contour on the model grid
"""
try:
y_ind = np.asarray(y_ind)
x_ind = np.asarray(x_ind)
# When replacing diagonal segments in the contour, pick the path that is
# closest to the contour isobath depth
option1 = uf.fabs(ds.bathymetry[xr.DataArray(y_ind+1), xr.DataArray(x_ind)] - self.depth)
option0 = uf.fabs(ds.bathymetry[xr.DataArray(y_ind), xr.DataArray(x_ind+1)] - self.depth)
add_new_y_point = xr.where(option1 <= option0, 1, 0 )
spacing = np.abs( np.diff(y_ind) ) + np.abs( np.diff(x_ind) )
if spacing.max() > 2:
raise ValueError("The contour is not continuous. The contour must be defined on "
"adjacent grid points.")
spacing[spacing!=2]=0
doublespacing = np.nonzero( spacing )[0]
for ispacing in doublespacing[::-1]:
if add_new_y_point[ispacing]:
y_ind = np.insert( y_ind, ispacing+1, y_ind[ispacing+1] )
x_ind = np.insert( x_ind, ispacing+1, x_ind[ispacing] )
else:
y_ind = np.insert( y_ind, ispacing+1, y_ind[ispacing] )
x_ind = np.insert( x_ind, ispacing+1, x_ind[ispacing+1] )
# Remove any repeated points caused by the rounding of the indices
nonrepeated_idx = np.nonzero( np.abs( np.diff(y_ind) ) + np.abs( np.diff(x_ind) ) )
y_ind = np.concatenate( (y_ind[nonrepeated_idx], [y_ind[-1]]) )
x_ind = np.concatenate( (x_ind[nonrepeated_idx], [x_ind[-1]]) )
return (y_ind, x_ind)
except ValueError:
print(traceback.format_exc())
def cfl_check(dataset, model_params):
sagp.get_grid_sizes(dataset, model_params)
dataset_2 = xruf.fabs(dataset['ucomp']) * model_params[
'delta_t'] * model_params['res'] / model_params['planet_radius']
dataset['cfl'] = (('time'), dataset_2.max(dim=('pfull', 'lat', 'lon')))
def _get_earth_mask(lat):
"""Identify earth/space pixels
Returns:
Mask (1=earth, 0=space)
"""
logger.debug('Computing earth mask')
return xu.fabs(lat) <= 90
def normalized_copy(data):
"""
Return a copy of data, with the absolute taken and normalized to 0-1.
The maximum across all regions and timesteps is used to normalize.
"""
ds = data.copy(deep=True) # Work off a copy
data_vars_in_t = [v for v in time_clustering._get_datavars(data)
if 't' in data[v].dims]
for var in data_vars_in_t:
for y in ds.coords['y'].values:
# Get max across all regions to normalize against
norm_max = fabs(ds[var].loc[{'y': y}]).max()
for x in ds.coords['x'].values:
df = ds[var].loc[{'x': x, 'y': y}]
ds[var].loc[{'x': x, 'y': y}] = fabs(df) / norm_max
return ds
def getEndOfSeason(self, array, sMmaxgreen):
# compute the minimum in the season start
seasonEnd_Range=array.sel(t=slice(self.eStart,self.eEnd))
seasonEnd_MinGreennessIdx=seasonEnd_Range.argmin('t')
seasonEnd_DateAtMin,seasonEnd_MinGreenness=seasonEnd_Range.t[seasonEnd_MinGreennessIdx].dt.dayofyear,seasonEnd_Range.isel(t=seasonEnd_MinGreennessIdx)
# Calculate the greenness value corresponding to the start of the season
seasonEnd_Greenness = seasonEnd_MinGreenness + ((sMmaxgreen - seasonEnd_MinGreenness) * (self.tEos / 100.0))
# Get the closest date to this greenness
#for i in range(len(seasonEnd_Range[:])): seasonEnd_Range[i]=seasonEnd_Range[i]-seasonEnd_Greenness
seasonEnd_Range=fabs(seasonEnd_Range-seasonEnd_Greenness)
seasonEnd_Idx=seasonEnd_Range.where(seasonEnd_Range.t.dt.dayofyear<=seasonEnd_DateAtMin).argmin('t',skipna=True)
seasonEnd_Date=seasonEnd_Range.t[seasonEnd_Idx].dt.dayofyear
return seasonEnd_Date
def _double_t_test(self, t0, dof):
"""
Double tailed student t test.
**Arguments:**
*array*
t0 field `xarray.DataArray`;
dof `int`;
**Return:**
*array*
p: p values;
"""
from scipy.stats import t
from xarray.ufuncs import fabs
pvalue = 2 * t.sf(fabs(t0), dof) * xr.ones_like(t0)
return pvalue
def plot_3d_isosurface(X, Y, Z, Q, level, is_ordered=False, fig=None,
**kwargs):
"""Plot a the surface of a given level value of Q
Assumes Q approximately decreases/increases monotonically with depth
Inputs
------
X, Y, Z: 1D or 3D arrays
The grid associated with Q
Q: 3D array
The quantity from which to find the isosurface
level: float
The value of the isosurface
is_ordered: bool
Set to true if Q(X, Y, Z) not Q(Z, Y, X), which is MITgcm default
fig: figure instance
If no figure instance is passed in, a new one is created
kwargs: dict
Values to pass to ?
Output
------
fig: figure instance
surf: contour surface artist
"""
fig = plt.figure() if fig is None else fig
ax = fig.gca(projection='3d')
Z_inds = np.argmin(fabs(Q - level), axis=0)
Z_surface = Z.data[[Z_inds]].squeeze()
print(Z_surface.shape)
if X.ndim == 1 and Y.ndim == 1:
X, Y = np.meshgrid(X, Y)
# Quick hack to avoid edge effects
inds = np.s_[5:-5, 5:-5]
ax.plot_surface(X[inds], Y[inds], Z_surface[inds])
return Z_inds
def plot_lid_driven_cavity_frame(Re, n):
ds = xr.open_dataset(f"lid_driven_cavity_Re{Re}.nc", decode_times=False)
Ny = ds.yC.size
Nz = ds.zC.size
fig, axes = plt.subplots(nrows=3, ncols=2, figsize=(16, 16), dpi=200)
plt.subplots_adjust(hspace=0.25)
fig.suptitle(f"Lid-driven cavity, Re = {Re}, t = {ds.time[n].values:.2f}",
fontsize=16)
ax_v_line, ax_v_mesh = axes[0, 0], axes[0, 1]
ax_w_line, ax_w_mesh = axes[1, 0], axes[1, 1]
ax_ζ_line, ax_ζ_mesh = axes[2, 0], axes[2, 1]
v_line = ds.v.isel(time=n, yF=Ny // 2)
ax_v_line.plot(y_Ghia,
v_Ghia[Re],
label="Ghia et al. (1982)",
color="tab:blue",
linestyle="",
marker="o",
fillstyle="none")
ax_v_line.plot(ds.zC,
v_line.values.flatten(),
label="Oceananigans.jl",
color="tab:blue")
ax_v_line.legend(loc="lower left",
bbox_to_anchor=(0, 1.01, 1, 0.2),
ncol=2,
frameon=False)
ax_v_line.set_xlabel("z")
ax_v_line.set_ylabel("v")
ax_v_line.set_xlim([0, 1])
w_line = ds.w.isel(time=n, zF=Nz // 2)
ax_w_line.plot(z_Ghia,
w_Ghia[Re],
label="Ghia et al. (1982)",
color="tab:orange",
linestyle="",
marker="o",
fillstyle="none")
ax_w_line.plot(ds.yC,
w_line.values.flatten(),
label="Oceananigans.jl",
color="tab:orange")
ax_w_line.legend(loc="lower left",
bbox_to_anchor=(0, 1.01, 1, 0.2),
ncol=2,
frameon=False)
ax_w_line.set_xlabel("y")
ax_w_line.set_ylabel("w")
ax_w_line.set_xlim([0, 1])
v = ds.v.isel(time=n).squeeze()
img_v = v.plot.pcolormesh(ax=ax_v_mesh,
vmin=-1,
vmax=1,
cmap=cmocean.cm.balance,
add_colorbar=False)
fig.colorbar(img_v, ax=ax_v_mesh, extend="both")
ax_v_mesh.axvline(x=0.5, color="tab:blue", alpha=0.5)
ax_v_mesh.set_title("v-velocity")
ax_v_mesh.set_xlabel("y")
ax_v_mesh.set_ylabel("z")
ax_v_mesh.set_aspect("equal")
w = ds.w.isel(time=n).squeeze()
img_w = w.plot.pcolormesh(ax=ax_w_mesh,
vmin=-1,
vmax=1,
cmap=cmocean.cm.balance,
add_colorbar=False)
fig.colorbar(img_w, ax=ax_w_mesh, extend="both")
ax_w_mesh.axhline(y=0.5, color="tab:orange", alpha=0.5)
ax_w_mesh.set_title("w-velocity")
ax_w_mesh.set_xlabel("y")
ax_w_mesh.set_ylabel("z")
ax_w_mesh.set_aspect("equal")
ζ_line = fabs(ds.ζ.isel(time=n, zF=Nz))
ax_ζ_line.plot(y_ζ_Ghia,
ζ_Ghia[Re],
label="Ghia et al. (1982)",
color="tab:purple",
linestyle="",
marker="o",
fillstyle="none")
ax_ζ_line.plot(ds.yF,
ζ_line.values.flatten(),
label="Oceananigans.jl",
color="tab:purple")
ax_ζ_line.legend(loc="lower left",
bbox_to_anchor=(0, 1.01, 1, 0.2),
ncol=2,
frameon=False)
ax_ζ_line.set_xlabel("y")
ax_ζ_line.set_ylabel("vorticity $|\zeta$|")
ax_ζ_line.set_xlim([0, 1])
ax_ζ_line.set_ylim(bottom=0)
ζ = ds.ζ.isel(time=n).squeeze()
img_ζ = ζ.plot.pcolormesh(ax=ax_ζ_mesh,
cmap=cmocean.cm.curl,
extend="both",
add_colorbar=False,
norm=colors.SymLogNorm(base=10,
linthresh=1,
vmin=-1e2,
vmax=1e2))
fig.colorbar(img_ζ, ax=ax_ζ_mesh, extend="both")
ax_ζ_mesh.axhline(y=1, color="tab:purple", alpha=1.0)
ax_ζ_mesh.set_title("vorticity")
ax_ζ_mesh.set_xlabel("y")
ax_ζ_mesh.set_ylabel("z")
ax_ζ_mesh.set_aspect("equal")
print(f"Saving lid-driven cavity Re={Re} frame {n}/{ds.time.size-1}...")
plt.savefig(f"lid_driven_cavity_Re{Re}_{n:05d}.png")
plt.close("all")
