def test_construction(self, sg_data, sg_topology):
filename = sg_data[0]
dataset = sg_data[1]
grid_topology = sg_topology
sg = Grid_S.from_netCDF(filename, dataset, grid_topology=grid_topology)
assert sg.filename == filename
sg2 = Grid_S.from_netCDF(filename)
assert sg2.filename == filename
sg3 = Grid.from_netCDF(filename, dataset, grid_topology=grid_topology)
sg4 = Grid.from_netCDF(filename)
print(sg3.shape)
print(sg4.shape)
assert sg == sg3
assert sg2 == sg4
def test_masked_grid(self, sg_data, sg_topology):
filename = sg_data[0]
dataset = sg_data[1]
grid_topology = sg_topology
sg = Grid_S.from_netCDF(filename, dataset, grid_topology=grid_topology)
assert sg.node_mask is not None
pts = {'on': [1.1,33.6],
'off': [0.9,30.6],
'masked': [1.1,31.6]
}
sg.build_celltree(use_mask=False)
assert sg._cell_tree[1].shape == sg.nodes.reshape(-1,2).shape
on_grid_result = sg.locate_faces(pts['on'])
off_grid_result = sg.locate_faces(pts['off'])
masked_territory_result = sg.locate_faces(pts['masked'])
sg.build_celltree(use_mask=True)
#locate a point that is on the grid in unmasked territory, and make
#sure that with or without the mask returns the same result
assert all(sg.locate_faces(pts['on']) == on_grid_result)
#locate a point off-grid, and make sure the same result is returned
assert all(sg.locate_faces(pts['off']) == off_grid_result)
#masked territory should be different.
assert all(sg.locate_faces(pts['masked']) != masked_territory_result)
#rebuild without the mask, and make sure results match
sg.build_celltree(use_mask=False)
assert all(sg.locate_faces(pts['on']) == on_grid_result)
assert all(sg.locate_faces(pts['off']) == off_grid_result)
assert all(sg.locate_faces(pts['masked']) == masked_territory_result)
def test_regrid_variable_StoS(get_s_depth):
# Time is not present
# Depth is not present
# Grid_S to Grid_S
from gridded.variable import Variable
from gridded.time import Time
from gridded.grids import Grid_S
sd = get_s_depth
grid = sd.grid
data = np.ones((grid.node_lon.shape[0], grid.node_lon.shape[1]))
v1 = Variable(name='v1', grid=grid, data=data, depth=None, time=None)
g2 = Grid_S(
node_lon=(grid.node_lon[0:-1, 0:-1] + grid.node_lon[1:, 1:]) / 2,
node_lat=(grid.node_lat[0:-1, 0:-1] + grid.node_lat[1:, 1:]) / 2)
v2 = utilities.regrid_variable(g2, v1)
# time should be unchanged
assert v2.time is v1.time
# depth should be None
assert v2.depth is v1.depth is None
sz = v1.data.shape[-1]
# data shape should retain the same time/depth dimensions as the original
# except in xy
assert v2.data.shape[-2::] == (sz - 1, sz - 1)
def test_regrid_variable_TDStoS(get_s_depth):
# Time is present
# Depth is present
# Grid_S to Grid_S
from gridded.variable import Variable
from gridded.time import Time
from gridded.grids import Grid_S
sd = get_s_depth
grid = sd.grid
n_levels = sd.num_w_levels
data = np.ones(
(1, n_levels, grid.node_lon.shape[0], grid.node_lon.shape[1]))
for l in range(0, n_levels):
data[0, l] *= l
v1 = Variable(name='v1',
grid=grid,
data=data,
depth=sd,
time=Time.constant_time())
g2 = Grid_S(
node_lon=(grid.node_lon[0:-1, 0:-1] + grid.node_lon[1:, 1:]) / 2,
node_lat=(grid.node_lat[0:-1, 0:-1] + grid.node_lat[1:, 1:]) / 2)
v2 = utilities.regrid_variable(g2, v1)
# time should be unchanged
assert v2.time is v1.time
# number of depth levels should remain unchanged
assert len(v2.depth) == len(v1.depth)
sz = v1.data.shape[-1]
# data shape should retain the same time/depth dimensions as the original
# except in xy
assert v2.data.shape == (v1.data.shape[0], v1.data.shape[1], sz - 1,
sz - 1)
def get_s_depth():
'''
This is setup for a ROMS S-level Depth that is on a square grid with a center
mound. Control vars: sz=xy size, center_el=height of the mound in meters,
d0=general depth in meters, sig=steepness of mound, nz=number of z levels.
'''
sz=40
center_el=10
d0 = 20
sig = 0.75
nz = 11
node_lat, node_lon = np.mgrid[0:sz,0:sz]
b_data = np.empty((sz,sz))
for x in range(0,sz):
for y in range(0,sz):
b_data[x,y] = d0 - center_el*np.exp(-0.1*((x-(sz/2))**2 / 2.*((sig)**2) + (y-(sz/2))**2 / 2.*((sig)**2)))
z_data = np.empty((3,sz,sz))
for t in range(0,3):
for x in range(0,sz):
for y in range(0,sz):
z_data[t,x,y] = (t - 1.)/2.
g = Grid_S(node_lon=node_lon, node_lat=node_lat)
bathy = Variable(name='bathy',
grid = g,
data = b_data)
t_data = np.array([Time.constant_time().data[0] + datetime.timedelta(minutes=10*d) for d in range(0,3)])
zeta = Variable(name='zeta',
time=Time(data=t_data),
grid=g,
data=z_data)
s_w = np.linspace(-1,0,nz)
s_rho = (s_w[0:-1] + s_w[1:]) /2
#equidistant layers, no stretching
Cs_w = np.linspace(-1,0,nz)
Cs_w = 1-1/np.exp(2*Cs_w)
Cs_w /= -Cs_w[0]
Cs_r = (Cs_w[0:-1] + Cs_w[1:]) /2
hc = np.array([0,])
sd = S_Depth(time=zeta.time,
grid=zeta.grid,
bathymetry=bathy,
zeta=zeta,
terms={'s_w':s_w,
's_rho':s_rho,
'Cs_w':Cs_w,
'Cs_r':Cs_r,
'hc':hc})
return sd
def test_masked_grid(self, sg_data, sg_topology):
filename = sg_data[0]
dataset = sg_data[1]
grid_topology = sg_topology
sg = Grid_S.from_netCDF(filename, dataset, grid_topology=grid_topology)
assert sg.node_mask is not None
assert all(sg.node_mask[0, :] == True)
assert all(sg.edge1_mask[-1, :] == True)
sg.build_celltree(grid='node', use_mask=False)
assert sg._cell_trees['node'][1].shape == sg.nodes.reshape(-1, 2).shape
on_grid_result = sg.locate_faces([0.5, 0], 'node')
off_grid_result = sg.locate_faces([0, 0.5], 'node')
masked_territory_result = sg.locate_faces([0.5, 0.9], 'node')
sg.build_celltree(grid='node', use_mask=True)
assert sg._cell_trees['node'][1].shape == (48, 2)
#locate a point that is on the grid in unmasked territory, and make
#sure that with or without the mask returns the same result
assert all(sg.locate_faces([0.5, 0], 'node') == on_grid_result)
#locate a point off-grid, and make sure the same result is returned
assert all(sg.locate_faces([0, 0.5], 'node') == off_grid_result)
#masked territory should be different.
assert all(
sg.locate_faces([0.5, 0.9], 'node') != masked_territory_result)
#rebuild without the mask, and make sure results match
sg.build_celltree(grid='node', use_mask=False)
assert all(sg.locate_faces([0.5, 0], 'node') == on_grid_result)
assert all(sg.locate_faces([0, 0.5], 'node') == off_grid_result)
assert all(
sg.locate_faces([0.5, 0.9], 'node') == masked_territory_result)
sg.build_celltree(grid='node', use_mask=True)
sg.use_masked_boundary = True
sg.build_celltree(grid='node', use_mask=True)
#because masked nodes that are adjacent to at least one unmasked node
#now get unmasked, and this grid has a one-node-thick border,
#all nodes should be unmasked
assert len(np.where(sg._masks['node'][0])[0]) == 0
#behavior should be identical as well
assert all(sg.locate_faces([0.5, 0], 'node') == on_grid_result)
assert all(sg.locate_faces([0, 0.5], 'node') == off_grid_result)
assert all(
sg.locate_faces([0.5, 0.9], 'node') == masked_territory_result)
#rerun with center just to make sure
sg.use_masked_boundary = False
sg.build_celltree(grid='center', use_mask=False)
assert sg._cell_trees['center'][1].shape == sg.centers.reshape(-1,
2).shape
on_grid_result = sg.locate_faces([0.3, 0], 'center')
off_grid_result = sg.locate_faces([-1, 0.5], 'center')
masked_territory_result = sg.locate_faces([0.1, -0.9], 'center')
sg.build_celltree(grid='center', use_mask=True)
#locate a point that is on the grid in unmasked territory, and make
#sure that with or without the mask returns the same result
assert all(sg.locate_faces([0.5, 0], 'center') == on_grid_result)
#locate a point off-grid, and make sure the same result is returned
assert all(sg.locate_faces([-1, 0.5], 'center') == off_grid_result)
#masked territory should be different.
assert all(
sg.locate_faces([0.1, -0.9], 'center') != masked_territory_result)
def gen_vortex_3D(filename=None):
x, y = np.mgrid[-30:30:61j, -30:30:61j]
y = np.ascontiguousarray(y.T)
x = np.ascontiguousarray(x.T)
x_size = 61
y_size = 61
g = Grid_S(node_lon=x, node_lat=y)
g.build_celltree()
lin_nodes = g._cell_tree[1]
lin_faces = np.array([
np.array([([lx, lx + x_size + 1, lx + 1
], [lx, lx + x_size, lx + x_size + 1])
for lx in range(0, x_size - 1, 1)]) + ly * x_size
for ly in range(0, y_size - 1)
])
lin_faces = lin_faces.reshape(-1, 3)
# y += np.sin(x) / 1
# x += np.sin(x) / 5
t0 = datetime(2001, 1, 1, 0, 0)
tarr = [t0 + timedelta(hours=i) for i in range(0, 11)]
angs = -np.arctan2(y, x)
mag = np.sqrt(x**2 + y**2)
vx = np.cos(angs) * mag
vy = np.sin(angs) * mag
vx = vx[np.newaxis, :] * 20
vy = vy[np.newaxis, :] * 20
vw = -0.001
d_scale = [1, 0.5, 0, -0.5, -1]
t_scale = np.linspace(0, 1, 11)
tvx = np.array([vx * t for t in t_scale]).squeeze()
tvy = np.array([vy * t for t in t_scale]).squeeze()
dvx = np.array([vx * s for s in d_scale]).squeeze()
dvy = np.array([vy * s for s in d_scale]).squeeze()
tdvx = np.array([dvx * t for t in t_scale]).squeeze()
tdvy = np.array([dvy * t for t in t_scale]).squeeze()
lin_vx = vx.reshape(-1)
lin_vy = vy.reshape(-1)
lin_tvx = np.array([lin_vx * t for t in t_scale])
lin_tvy = np.array([lin_vy * t for t in t_scale])
lin_dvx = np.array([lin_vx * s for s in d_scale])
lin_dvy = np.array([lin_vy * s for s in d_scale])
lin_tdvx = np.array([lin_dvx * t for t in t_scale])
lin_tdvy = np.array([lin_dvy * t for t in t_scale])
ds = None
if filename is not None:
ds = nc4.Dataset(filename, 'w', diskless=True, persist=True)
ds.createDimension('y', y.shape[0])
ds.createDimension('x', x.shape[1])
ds.createDimension('time', len(tarr))
ds.createDimension('depth', len(d_scale))
ds.createVariable('x', 'f8', dimensions=('x', 'y'))
ds['x'][:] = x
ds.createVariable('y', 'f8', dimensions=('x', 'y'))
ds['y'][:] = y
ds.createVariable('time', 'f8', dimensions=('time'))
ds['time'][:] = nc4.date2num(tarr, 'hours since {0}'.format(t0))
ds['time'].setncattr('units', 'hours since {0}'.format(t0))
ds.createVariable('vx', 'f8', dimensions=('x', 'y'))
ds.createVariable('vy', 'f8', dimensions=('x', 'y'))
ds['vx'][:] = vx
ds['vy'][:] = vy
ds.createVariable('tvx', 'f8', dimensions=('time', 'x', 'y'))
ds.createVariable('tvy', 'f8', dimensions=('time', 'x', 'y'))
ds['tvx'][:] = tvx
ds['tvy'][:] = tvy
ds.createVariable('dvx', 'f8', dimensions=('depth', 'x', 'y'))
ds.createVariable('dvy', 'f8', dimensions=('depth', 'x', 'y'))
ds['dvx'][:] = dvx
ds['dvy'][:] = dvy
ds.createVariable('tdvx', 'f8', dimensions=('time', 'depth', 'x', 'y'))
ds.createVariable('tdvy', 'f8', dimensions=('time', 'depth', 'x', 'y'))
ds['tdvx'][:] = tdvx
ds['tdvy'][:] = tdvy
for v in ds.variables:
if 'v' in v:
ds[v].units = 'm/s'
ds.createDimension('nv', lin_nodes.shape[0])
ds.createDimension('nele', lin_faces.shape[0])
ds.createDimension('two', 2)
ds.createDimension('three', 3)
ds.createVariable('nodes', 'f8', dimensions=('nv', 'two'))
ds.createVariable('faces', 'f8', dimensions=('nele', 'three'))
ds.createVariable('lin_vx', 'f8', dimensions=('nv'))
ds.createVariable('lin_vy', 'f8', dimensions=('nv'))
ds.createVariable('lin_tvx', 'f8', dimensions=('time', 'nv'))
ds.createVariable('lin_tvy', 'f8', dimensions=('time', 'nv'))
ds.createVariable('lin_dvx', 'f8', dimensions=('depth', 'nv'))
ds.createVariable('lin_dvy', 'f8', dimensions=('depth', 'nv'))
ds.createVariable('lin_tdvx', 'f8', dimensions=('time', 'depth', 'nv'))
ds.createVariable('lin_tdvy', 'f8', dimensions=('time', 'depth', 'nv'))
for k, v in {
'nodes': lin_nodes,
'faces': lin_faces,
'lin_vx': lin_vx,
'lin_vy': lin_vy,
'lin_tvx': lin_tvx,
'lin_tvy': lin_tvy,
'lin_dvx': lin_dvx,
'lin_dvy': lin_dvy,
'lin_tdvx': lin_tdvx,
'lin_tdvy': lin_tdvy
}.items():
ds[k][:] = v
if 'lin' in k:
ds[k].units = 'm/s'
Grid._get_grid_type(ds,
grid_topology={
'node_lon': 'x',
'node_lat': 'y'
})
Grid._get_grid_type(ds)
ds.setncattr('grid_type', 'sgrid')
if ds is not None:
# Need to test the dataset...
sgt = {'node_lon': 'x', 'node_lat': 'y'}
sg = Grid.from_netCDF(dataset=ds, grid_topology=sgt, grid_type='sgrid')
ugt = {'nodes': 'nodes', 'faces': 'faces'}
# ug = PyGrid_U(nodes=ds['nodes'][:], faces=ds['faces'][:])
ds.close()
return {
'sgrid': (x, y),
'sgrid_vel': (dvx, dvy),
'sgrid_depth_vel': (tdvx, tdvy),
'ugrid': (lin_nodes, lin_faces),
'ugrid_depth_vel': (lin_tdvx, lin_tdvy)
}