def test_at(self):
curr_file = os.path.join(test_dir, 'staggered_sine_channel.nc')
u = Variable.from_netCDF(filename=curr_file, varname='u_rho')
v = Variable.from_netCDF(filename=curr_file, varname='v_rho')
points = np.array(([0, 0, 0], [np.pi, 1, 0], [2 * np.pi, 0, 0]))
time = datetime.datetime.now()
res = u.at(points, time)
assert all(res == [1, 1, 1])
print(np.cos(points[:, 0] / 2) / 2)
assert all(np.isclose(v.at(points, time),
np.cos(points[:, 0] / 2) / 2))
#reorganize the points so they are arrays of components
#results of .at should follow the same format
points = np.array([
[0, np.pi / 2, np.pi, 1.5 * np.pi, 2 * np.pi], #x
[0, 0.5, 1, 0.5, 0], #y
[0, 0, 0, 0, 0]
]) #z
assert all(u.at(points, time) == [1, 1, 1, 1, 1])
print(np.cos(points[0] / 2) / 2)
assert all(np.isclose(v.at(points, time), np.cos(points[0] / 2) / 2))
def test_write_with_depths(two_triangles, ncds):
"""Tests writing a netcdf file with depth data."""
fname, ncds = ncds
grid = two_triangles
grid.mesh_name = 'mesh1'
gds = Dataset(grid=grid)
# Create a Variable object for the depths:
depths = Variable(name='depth', location='node', data=[1.0, 2.0, 3.0, 4.0])
depths.attributes['units'] = 'm'
depths.attributes['standard_name'] = 'sea_floor_depth_below_geoid'
depths.attributes['positive'] = 'down'
gds.variables['depth'] = depths
gds.save(ncds, format='netcdf4')
ds = netCDF4.Dataset(fname)
assert nc_has_variable(ds, 'mesh1')
assert nc_has_variable(ds, 'depth')
assert nc_var_has_attr_vals(
ds, 'depth', {
'coordinates': 'mesh1_node_lon mesh1_node_lat',
'location': 'node',
'mesh': 'mesh1'
})
ds.close()
def test_construction(self):
curr_file = os.path.join(test_dir, 'staggered_sine_channel.nc')
u = Variable.from_netCDF(filename=curr_file, varname='u_rho')
v = Variable.from_netCDF(filename=curr_file, varname='v_rho')
gvp = VectorVariable(name='velocity',
units='m/s',
time=u.time,
variables=[u, v])
assert gvp.name == 'velocity'
assert gvp.units == 'm/s'
assert gvp.varnames[0] == 'u_rho'
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 _load_variables(self, ds):
"""
load up the variables in the nc file
"""
variables = {}
for k in ds.variables.keys():
# find which netcdf variables are used to define the grid
is_not_grid_attr = all([
k not in str(v).split()
for v in self.grid.grid_topology.values()
])
if is_not_grid_attr:
ncvar = ds[k]
# find the location of the variable
try:
location = ncvar.location
except AttributeError:
# that didn't work, need to try to infer it
location = self.grid.infer_location(ncvar)
if location is not None:
try:
ln = ds[k].long_name
except AttributeError: # no long_name attribute
ln = ds[k].name # use the name attribute
# fixme: Variable.from_netCDF should really be able to figure out the location itself
# maybe we need multiple Variable subclasses for different grid types?
variables[k] = Variable.from_netCDF(
dataset=ds,
name=ln,
varname=k,
grid=self.grid,
location=location,
)
return variables
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 test_write_with_bound_data(two_triangles, ncds):
"""
Tests writing a netcdf file with data on the boundaries
suitable for boundary conditions, for example fluxes.
"""
fname, ncds = ncds
grid = two_triangles
grid.mesh_name = 'mesh3'
gds = Dataset(grid=grid)
# Add the boundary definitions:
grid.boundaries = [(0, 1), (0, 2), (1, 3), (2, 3)]
# Create a Variable object for boundary conditions:
bnds = Variable('bnd_cond', location='boundary', data=[0, 1, 0, 0])
bnds.attributes['long_name'] = 'model boundary conditions'
bnds.attributes['flag_values'] = '0 1'
bnds.attributes['flag_meanings'] = 'no_flow_boundary open_boundary'
gds.variables['bnd_cond'] = bnds
gds.save(ncds)
ds = netCDF4.Dataset(fname)
assert nc_has_variable(ds, 'mesh3')
assert nc_has_variable(ds, 'bnd_cond')
assert nc_var_has_attr_vals(
ds, 'mesh3', {'boundary_node_connectivity': 'mesh3_boundary_nodes'})
assert nc_var_has_attr_vals(
ds, 'bnd_cond', {
'location': 'boundary',
'flag_values': '0 1',
'flag_meanings': 'no_flow_boundary open_boundary',
'mesh': 'mesh3',
})
# There should be no coordinates attribute or variable for the
# boundaries as there is no boundaries_coordinates defined.
assert not nc_has_variable(ds, 'mesh_boundary_lon')
assert not nc_has_variable(ds, 'mesh_boundary_lat')
assert not nc_var_has_attr(ds, 'bnd_cond', 'coordinates')
ds.close()
def test_write_with_edge_data(two_triangles, ncds):
"""Tests writing a netcdf file with data on the edges (fluxes, maybe?)."""
fname, ncds = ncds
two_triangles.mesh_name = 'mesh2'
gds = Dataset(grid=two_triangles)
# Create a Variable object for fluxes:
flux = Variable('flux', location='edge', data=[
0.0,
0.0,
4.1,
0.0,
5.1,
])
flux.attributes['units'] = 'm^3/s'
flux.attributes['long_name'] = 'volume flux between cells'
flux.attributes['standard_name'] = 'ocean_volume_transport_across_line'
gds.variables['flux'] = flux
# Add coordinates for edges.
gds.grid.build_edge_coordinates()
gds.save(ncds)
ds = netCDF4.Dataset(fname)
assert nc_has_variable(ds, 'mesh2')
assert nc_has_variable(ds, 'flux')
assert nc_var_has_attr_vals(
ds, 'flux', {
'coordinates': 'mesh2_edge_lon mesh2_edge_lat',
'location': 'edge',
'units': 'm^3/s',
'mesh': 'mesh2'
})
assert np.array_equal(ds.variables['mesh2_edge_lon'],
gds.grid.edge_coordinates[:, 0])
assert np.array_equal(ds.variables['mesh2_edge_lat'],
gds.grid.edge_coordinates[:, 1])
ds.close()
def test_construction(self, sg_data):
fn = sg_data[0]
sinusoid = sg_data[1]
data = sinusoid['u'][:]
grid = Grid.from_netCDF(dataset=sinusoid)
time = None
u = Variable(name='u',
units='m/s',
data=data,
grid=grid,
time=time,
data_file='staggered_sine_channel.nc',
grid_file='staggered_sine_channel.nc')
curr_file = os.path.join(test_dir, 'staggered_sine_channel.nc')
k = Variable.from_netCDF(filename=curr_file, varname='u', name='u')
assert k.name == u.name
assert k.units == 'm/s'
# fixme: this was failing
# assert k.time == u.time
assert k.data[0, 0] == u.data[0, 0]
def test_write_with_velocities(two_triangles, ncds):
"""Tests writing a netcdf file with velocities on the faces."""
fname, ncds = ncds
two_triangles.mesh_name = 'mesh2'
gds = Dataset(grid=two_triangles)
# Create a Variable object for u velocity:
u_vel = Variable('u', location='face', data=[1.0, 2.0])
u_vel.attributes['units'] = 'm/s'
u_vel.attributes['standard_name'] = 'eastward_sea_water_velocity'
gds.variables['u'] = u_vel
# Create a Variable object for v velocity:
v_vel = Variable('v', location='face', data=[3.2, 4.3])
v_vel.attributes['units'] = 'm/s'
v_vel.attributes['standard_name'] = 'northward_sea_water_velocity'
gds.variables['v'] = v_vel
# Add coordinates for face data.
gds.grid.build_face_coordinates()
gds.save(ncds)
ds = netCDF4.Dataset(fname)
assert nc_has_variable(ds, 'mesh2')
assert nc_has_variable(ds, 'u')
assert nc_has_variable(ds, 'v')
assert nc_var_has_attr_vals(
ds, 'u', {
'coordinates': 'mesh2_face_lon mesh2_face_lat',
'location': 'face',
'mesh': 'mesh2',
})
ds.close()
def _load_variables(self, ds):
"""
load up the variables in the nc file
"""
variables = {}
for k in ds.variables.keys():
is_not_grid_attr = all([
k not in str(v).split()
for v in self.grid.grid_topology.values()
])
if is_not_grid_attr and self.grid.infer_location(
ds[k]) is not None:
try:
ln = ds[k].long_name
except: # fixme: what Exception are we expecting???
ln = ds[k].name
variables[k] = Variable.from_netCDF(
dataset=ds,
name=ln,
varname=k,
grid=self.grid,
)
return variables
def _load_variables(self, ds):
# fixme: need a way to do this for non-compliant files
variables = {}
for k in ds.variables.keys():
is_not_grid_attr = all([
k not in str(v).split()
for v in self.grid.grid_topology.values()
])
if is_not_grid_attr and self.grid.infer_location(
ds[k]) is not None:
try:
ln = ds[k].long_name
except:
ln = ds[k].name
variables[k] = Variable.from_netCDF(
dataset=ds,
name=ln,
varname=k,
grid=self.grid,
)
return variables
def regrid_variable(grid, o_var, location='node'):
from gridded.variable import Variable
from gridded.grids import Grid_S, Grid_U
from gridded.depth import S_Depth, Depth
"""
Takes a Variable or VectorVariable and interpolates the data onto grid.
You may pass a location ('nodes', 'faces', 'edge1', 'edge2) and the
variable will be interpolated there if possible
If no location is passed, the variable will be interpolated to the
nodes of this grid. If the Variable's grid and this grid are the same, this
function will return the Variable unchanged.
If this grid covers area that the source grid does not, all values
in this area will be masked. If regridding from cell centers to the nodes,
The values of any border point not within will be equal to the value at the
center of the border cell.
NOTE: This function will load the ENTIRE data space of the source Variable.
Make sure that this is reasonable. If it is not, consider pre-slicing
the time and depth dimensions of the source Variable to what you need.
"""
dest_points = None
if location == 'node':
dest_points = grid.nodes
if location == 'face' or location == 'center':
if grid.face_coordinates is None and isinstance(grid, Grid_U):
grid.build_face_coordinates()
dest_points = grid.face_coordinates
else:
dest_points = grid.centers
dest_points = dest_points.reshape(-1, 2)
if 'edge' in location:
raise NotImplementedError(
"Cannot regrid variable to edges at this time")
dest_indices = o_var.grid.locate_faces(dest_points, 'node')
if np.all(dest_indices == -1):
raise ValueError("Grid {0} has no destination points overlapping\
the grid of the source variable {1}".format(grid, o_var))
n_depth = None
if o_var.depth is not None:
if isinstance(o_var.depth, S_Depth):
n_depth = _regrid_s_depth(grid, o_var.depth)
elif isinstance(o_var.depth, Depth):
n_depth = o_var.depth
else:
raise NotImplementedError("Can only regrid sigma depths for now")
xy_shp = grid._get_grid_vars(location)[0].shape
d_shp = len(n_depth) if n_depth is not None else None
t_shp = len(o_var.time) if o_var.time is not None else None
n_shape = xy_shp
if d_shp:
n_shape = (d_shp, ) + n_shape
if t_shp:
n_shape = (t_shp, ) + n_shape
n_data = np.empty(n_shape)
location_shp = grid.node_lon.shape
pts = np.zeros((dest_points.shape[0], 3))
pts[:, 0:2] = dest_points
if o_var.time is not None:
for t_idx, t in enumerate(o_var.time.data):
if n_depth is not None and isinstance(n_depth, S_Depth):
for lev_idx, lev_data in enumerate(o_var.depth.get_section(t)):
lev = Variable(name='level{0}'.format(lev_idx),
data=lev_data,
grid=o_var.grid)
zs = lev.at(pts, t)
pts[:, 2] = zs
n_data[t_idx, lev_idx] = o_var.at(pts,
t).reshape(location_shp)
else:
n_data[t_idx] = o_var.at(pts, t).reshape(location_shp)
else:
n_data = o_var.at(pts, None).reshape(location_shp)
n_var = Variable(name='regridded {0}'.format(o_var.name),
grid=grid,
time=o_var.time,
depth=n_depth,
data=n_data,
units=o_var.units)
return n_var
def test_write_everything(twenty_one_triangles, ncds):
"""An example with all features enabled, and a less trivial grid."""
fname, ncds = ncds
grid = twenty_one_triangles
grid.mesh_name = 'mesh'
gds = Dataset(grid=grid)
grid.build_edges()
grid.build_face_face_connectivity()
grid.build_edge_coordinates()
grid.build_face_coordinates()
grid.build_boundary_coordinates()
# Depth on the nodes.
depths = Variable('depth', location='node', data=np.linspace(1, 10, 20))
depths.attributes['units'] = 'm'
depths.attributes['standard_name'] = 'sea_floor_depth_below_geoid'
depths.attributes['positive'] = 'down'
gds.variables['depth'] = depths
# Create a Variable object for u velocity:
u_vel = Variable('u', location='face', data=np.sin(np.linspace(3, 12, 21)))
u_vel.attributes['units'] = 'm/s'
u_vel.attributes['standard_name'] = 'eastward_sea_water_velocity'
gds.variables['u'] = u_vel
# Create a Variable object for v velocity:
v_vel = Variable('v',
location='face',
data=np.sin(np.linspace(12, 15, 21)))
v_vel.attributes['units'] = 'm/s'
v_vel.attributes['standard_name'] = 'northward_sea_water_velocity'
gds.variables['v'] = v_vel
# Fluxes on the edges:
flux = Variable('flux', location='edge', data=np.linspace(1000, 2000, 41))
flux.attributes['units'] = 'm^3/s'
flux.attributes['long_name'] = 'volume flux between cells'
flux.attributes['standard_name'] = 'ocean_volume_transport_across_line'
gds.variables['flux'] = flux
# Boundary conditions:
bounds = np.zeros((19, ), dtype=np.uint8)
bounds[7] = 1
bnds = Variable('bnd_cond', location='boundary', data=bounds)
bnds.attributes['long_name'] = 'model boundary conditions'
bnds.attributes['flag_values'] = '0 1'
bnds.attributes['flag_meanings'] = 'no_flow_boundary open_boundary'
gds.variables['bnd_cond'] = bnds
gds.save(ncds)
ds = netCDF4.Dataset(fname)
# Now the tests:
assert nc_has_variable(ds, 'mesh')
assert nc_has_variable(ds, 'depth')
assert nc_var_has_attr_vals(ds, 'depth', {
'coordinates': 'mesh_node_lon mesh_node_lat',
'location': 'node'
})
assert nc_has_variable(ds, 'u')
assert nc_has_variable(ds, 'v')
assert nc_var_has_attr_vals(
ds, 'u', {
'coordinates': 'mesh_face_lon mesh_face_lat',
'location': 'face',
'mesh': 'mesh'
})
assert nc_var_has_attr_vals(
ds, 'v', {
'coordinates': 'mesh_face_lon mesh_face_lat',
'location': 'face',
'mesh': 'mesh'
})
assert nc_has_variable(ds, 'flux')
assert nc_var_has_attr_vals(
ds, 'flux', {
'coordinates': 'mesh_edge_lon mesh_edge_lat',
'location': 'edge',
'units': 'm^3/s',
'mesh': 'mesh'
})
assert nc_has_variable(ds, 'mesh')
assert nc_has_variable(ds, 'bnd_cond')
assert nc_var_has_attr_vals(
ds, 'mesh', {'boundary_node_connectivity': 'mesh_boundary_nodes'})
assert nc_var_has_attr_vals(
ds, 'bnd_cond', {
'location': 'boundary',
'flag_values': '0 1',
'flag_meanings': 'no_flow_boundary open_boundary',
'mesh': 'mesh'
})
ds.close()
# And make sure pyugrid can reload it!
gds = Dataset(fname)
grid = gds.grid
# And that some things are the same.
# NOTE: more testing might be good here.
# maybe some grid comparison functions?
assert grid.mesh_name == 'mesh'
assert len(grid.nodes) == 20
depth = gds.variables['depth']
assert depth.attributes['units'] == 'm'
u = gds.variables['u']
assert u.attributes['units'] == 'm/s'
