def get_bathymetry(bathymetry_file, mesh2d, minimum_depth=5.0, project=False):
"""Interpolates/projects bathymetry from a raster to P1 field."""
retrieve_bath_file(bathymetry_file)
d = Dataset(bathymetry_file)
x = d['x'][:]
y = d['y'][:]
bath = -d['bathymetry'][:]
if np.ma.isMaskedArray(bath):
bath = bath.filled(minimum_depth)
bath[~np.isfinite(bath)] = minimum_depth
interpolator = scipy.interpolate.RegularGridInterpolator((x, y), bath.T)
P1_2d = get_functionspace(mesh2d, 'CG', 1)
bathymetry2d = Function(P1_2d, name='bathymetry')
if project:
# interpolate on a high order mesh
P3_2d = get_functionspace(mesh2d, 'CG', 3)
P3_2d_v = get_functionspace(mesh2d, 'CG', 3, vector=True)
bathymetry2d_ho = Function(P3_2d, name='bathymetry')
coords_ho = Function(P3_2d_v).interpolate(SpatialCoordinate(mesh2d))
interpolate_onto(interpolator, bathymetry2d_ho, coords_ho, minimum_depth)
# project on P1
bathymetry2d.project(bathymetry2d_ho)
shallow_ix = bathymetry2d.dat.data < minimum_depth
bathymetry2d.dat.data[shallow_ix] = minimum_depth
else:
interpolate_onto(interpolator, bathymetry2d, mesh2d.coordinates, minimum_depth)
return bathymetry2d
def test():
mesh2d = Mesh('mesh_cre-plume_03_normal.msh')
p1 = get_functionspace(mesh2d, 'CG', 1)
p1v = get_functionspace(mesh2d, 'CG', 1, vector=True)
elev_field = Function(p1, name='elevation')
uv_field = Function(p1v, name='transport')
sim_tz = timezone.FixedTimeZone(-8, 'PST')
init_date = datetime.datetime(2006, 5, 15, tzinfo=sim_tz)
# tide_cls = FES2004TidalBoundaryForcing
tide_cls = TPXOTidalBoundaryForcing
tbnd = tide_cls(elev_field,
init_date,
to_latlon,
COORDSYS,
uv_field=uv_field,
data_dir='forcings',
constituents=['M2', 'K1'],
boundary_ids=[2, 5, 7])
tbnd.set_tidal_field(0.0)
elev_outfn = 'tmp/tidal_elev.pvd'
uv_outfn = 'tmp/tidal_uv.pvd'
print('Saving to {:} {:}'.format(elev_outfn, uv_outfn))
elev_out = File(elev_outfn)
uv_out = File(uv_outfn)
for t in np.linspace(0, 12 * 3600., 49):
tbnd.set_tidal_field(t)
if elev_field.function_space().mesh().comm.rank == 0:
print('t={:7.1f} elev: {:7.1f} uv: {:7.1f}'.format(
t, norm(elev_field), norm(uv_field)))
elev_out.write(elev_field)
uv_out.write(uv_field)
def u(mesh, spaces):
(name, order), (vname, vorder) = spaces
return utility.get_functionspace(mesh,
name,
order,
vname,
vorder,
vector=True)
def compute_hcc_metric(nelem, nlayers, slope, deform='uniform'):
mesh2d = UnitSquareMesh(nelem, nelem)
mesh = ExtrudedMesh(mesh2d, nlayers, 1.0 / nlayers)
mesh.coordinates.dat.data[:, 2] += -1.0
xyz = SpatialCoordinate(mesh)
if deform == 'uniform':
mesh.coordinates.dat.data[:, 2] = (
mesh.coordinates.dat.data[:, 2] +
slope * mesh.coordinates.dat.data[:, 0])
else:
mesh.coordinates.dat.data[:, 2] = (
mesh.coordinates.dat.data[:, 2] + slope *
mesh.coordinates.dat.data[:, 0] * mesh.coordinates.dat.data[:, 2])
P1DG = utility.get_functionspace(mesh, 'DG', 1)
f_hcc = Function(P1DG, name='hcc_metric_3d')
# emulate solver object
solver_obj = utility.AttrDict()
solver_obj.mesh = mesh
solver_obj.comm = comm
solver_obj.fields = utility.AttrDict()
solver_obj.function_spaces = utility.AttrDict()
solver_obj.fields.hcc_metric_3d = f_hcc
solver_obj.fields.z_coord_3d = Function(P1DG, name='z_coord_3d')
solver_obj.fields.z_coord_3d.interpolate(xyz[2])
solver_obj.function_spaces.P1DG = P1DG
utility.Mesh3DConsistencyCalculator(solver_obj).solve()
hcc_min = f_hcc.dat.data.min()
hcc_max = f_hcc.dat.data.max()
return hcc_min, hcc_max
def test_interpolator():
"""
Test NCOM 3d interpolator.
.. note::
The following NCOM output files must be present:
./forcings/ncom/model_h.nc
./forcings/ncom/model_lat.nc
./forcings/ncom/model_ang.nc
./forcings/ncom/model_lon.nc
./forcings/ncom/model_zm.nc
./forcings/ncom/2006/s3d/s3d.glb8_2f_2006050100.nc
./forcings/ncom/2006/s3d/s3d.glb8_2f_2006050200.nc
./forcings/ncom/2006/t3d/t3d.glb8_2f_2006050100.nc
./forcings/ncom/2006/t3d/t3d.glb8_2f_2006050200.nc
./forcings/ncom/2006/u3d/u3d.glb8_2f_2006050100.nc
./forcings/ncom/2006/u3d/u3d.glb8_2f_2006050200.nc
./forcings/ncom/2006/v3d/v3d.glb8_2f_2006050100.nc
./forcings/ncom/2006/v3d/v3d.glb8_2f_2006050200.nc
./forcings/ncom/2006/ssh/ssh.glb8_2f_2006050100.nc
./forcings/ncom/2006/ssh/ssh.glb8_2f_2006050200.nc
"""
# load and extrude mesh
from bathymetry import get_bathymetry, smooth_bathymetry, smooth_bathymetry_at_bnd
nlayers, surf_elem_height, max_z_stretch = (9, 5.0, 4.0)
mesh2d = Mesh('mesh_cre-plume_03_normal.msh')
# interpolate bathymetry and smooth it
bathymetry_2d = get_bathymetry('bathymetry_utm_large.nc', mesh2d, project=False)
bathymetry_2d = smooth_bathymetry(
bathymetry_2d, delta_sigma=1.0, bg_diff=0,
alpha=1e2, exponent=2.5,
minimum_depth=3.5, niter=30)
bathymetry_2d = smooth_bathymetry_at_bnd(bathymetry_2d, [2, 7])
# 3d mesh vertical stretch factor
z_stretch_fact_2d = Function(bathymetry_2d.function_space(), name='z_stretch')
# 1.0 (sigma mesh) in shallow areas, 4.0 in deep ocean
z_stretch_fact_2d.project(-ln(surf_elem_height/bathymetry_2d)/ln(nlayers))
z_stretch_fact_2d.dat.data[z_stretch_fact_2d.dat.data < 1.0] = 1.0
z_stretch_fact_2d.dat.data[z_stretch_fact_2d.dat.data > max_z_stretch] = max_z_stretch
extrude_options = {
'z_stretch_fact': z_stretch_fact_2d,
}
mesh = extrude_mesh_sigma(mesh2d, nlayers, bathymetry_2d,
**extrude_options)
p1_2d = get_functionspace(mesh2d, 'CG', 1)
p1 = get_functionspace(mesh, 'CG', 1)
# make functions
salt = Function(p1, name='salinity')
temp = Function(p1, name='temperature')
uvel = Function(p1, name='u-velocity')
vvel = Function(p1, name='v-velocity')
elev = Function(p1_2d, name='elevation')
sim_tz = timezone.FixedTimeZone(-8, 'PST')
init_date = datetime.datetime(2006, 5, 1, tzinfo=sim_tz)
interp = NCOMInterpolator(
p1_2d, p1, [salt, temp, uvel, vvel, elev],
['Salinity', 'Temperature', 'U_Velocity', 'V_Velocity', 'Surface_Elevation'],
['s3d', 't3d', 'u3d', 'v3d', 'ssh'],
to_latlon, 'forcings/ncom',
'{year:04d}/{fieldstr:}/{fieldstr:}.glb8_2f_{year:04d}{month:02d}{day:02d}00.nc',
init_date, COORDSYS, verbose=True
)
interp.set_fields(0.0)
salt_fn = 'tmp/salt.pvd'
temp_fn = 'tmp/temp.pvd'
uvel_fn = 'tmp/uvel.pvd'
vvel_fn = 'tmp/vvel.pvd'
elev_fn = 'tmp/elev.pvd'
print('Saving output to {:} {:} {:} {:} {:}'.format(salt_fn, temp_fn, uvel_fn, vvel_fn, elev_fn))
out_salt = File(salt_fn)
out_temp = File(temp_fn)
out_uvel = File(uvel_fn)
out_vvel = File(vvel_fn)
out_elev = File(elev_fn)
out_salt.write(salt)
out_temp.write(temp)
out_uvel.write(uvel)
out_vvel.write(vvel)
out_elev.write(elev)
dt = 3*3600.
for i in range(8):
print('Time step {:}'.format(i))
interp.set_fields(i*dt)
out_salt.write(salt)
out_temp.write(temp)
out_uvel.write(uvel)
out_vvel.write(vvel)
out_elev.write(elev)
def test_implicit_friction(do_export=False, do_assert=True):
# set mesh resolution
scale = 1000.0
reso = 2.5 * scale
layers = 50
depth = 15.0
# generate unit mesh and transform its coords
x_max = 5.0 * scale
x_min = -5.0 * scale
lx = (x_max - x_min)
n_x = int(lx / reso)
mesh2d = RectangleMesh(n_x, n_x, lx, lx, reorder=True)
# move mesh, center to (0,0)
mesh2d.coordinates.dat.data[:, 0] -= lx / 2
mesh2d.coordinates.dat.data[:, 1] -= lx / 2
mesh = ExtrudedMesh(mesh2d, layers=50, layer_height=-depth / layers)
if do_export:
out_file = File('implicit_bf_sol.pvd')
# ----- define function spaces
deg = 1
p1dg = get_functionspace(mesh, 'DG', 1)
p1dgv = get_functionspace(mesh, 'DG', 1, vector=True)
u_h_elt = FiniteElement('RT', triangle, deg + 1, variant='equispaced')
u_v_elt = FiniteElement('DG', interval, deg, variant='equispaced')
u_elt = HDiv(TensorProductElement(u_h_elt, u_v_elt))
# for vertical velocity component
w_h_elt = FiniteElement('DG', triangle, deg, variant='equispaced')
w_v_elt = FiniteElement('CG', interval, deg + 1, variant='equispaced')
w_elt = HDiv(TensorProductElement(w_h_elt, w_v_elt))
# in deformed mesh horiz. velocity must actually live in U + W
uw_elt = EnrichedElement(u_elt, w_elt)
# final spaces
v = FunctionSpace(mesh, uw_elt) # uv
solution = Function(v, name='velocity')
solution_new = Function(v, name='new velocity')
solution_p1_dg = Function(p1dgv, name='velocity p1dg')
viscosity_v = Function(p1dg, name='viscosity')
elev_slope = -1.0e-5
source = Constant((-9.81 * elev_slope, 0, 0))
z0 = 1.5e-3
kappa = 0.4
drag = (kappa / np.log((depth / layers) / z0))**2
bottom_drag = Constant(drag)
u_bf = 0.035 # NOTE tuned to produce ~correct viscosity profile
x, y, z = SpatialCoordinate(mesh)
viscosity_v.project(kappa * u_bf * -z * (depth + z + z0) / (depth + z0))
print('Cd {:}'.format(drag))
print('u_bf {:}'.format(u_bf))
print('nu {:}'.format(viscosity_v.dat.data.min(),
viscosity_v.dat.data.max()))
# --- solve mom eq
test = TestFunction(v)
normal = FacetNormal(mesh)
def rhs(solution, sol_old):
# source term (external pressure gradient
f = inner(source, test) * dx
# vertical diffusion (integrated by parts)
f += -viscosity_v * inner(Dx(solution, 2), Dx(test, 2)) * dx
# interface term
diff_flux = viscosity_v * Dx(solution, 2)
f += (dot(avg(diff_flux), test('+')) * normal[2]('+') +
dot(avg(diff_flux), test('-')) * normal[2]('-')) * dS_h
# symmetric interior penalty stabilization
l = Constant(depth / layers)
nb_neigh = 2
o = 1
d = 3
sigma = Constant((o + 1) * (o + d) / d * nb_neigh / 2) / l
gamma = sigma * avg(viscosity_v)
f += gamma * dot(
jump(solution),
test('+') * normal[2]('+') + test('-') * normal[2]('-')) * dS_h
# boundary term
uv_bot_old = sol_old + Dx(sol_old, 2) * l * 0.5
uv_bot = solution + Dx(solution, 2) * l * 0.5 # solver fails
uv_mag = sqrt(uv_bot_old[0]**2 + uv_bot_old[1]**2) + Constant(1e-12)
bnd_flux = bottom_drag * uv_mag * uv_bot
ds_bottom = ds_t
f += dot(bnd_flux, test) * normal[2] * ds_bottom
return f
# ----- define solver
sp = {}
dt = 3600.0
time_steps = 13
dt_const = Constant(dt)
# Backward Euler
f = (inner(solution_new, test) * dx - inner(solution, test) * dx -
dt_const * rhs(solution_new, solution))
prob = NonlinearVariationalProblem(f, solution_new)
solver = LinearVariationalSolver(prob, solver_parameters=sp)
if do_export:
out_file.write(solution)
# ----- solve
t = 0
for it in range(1, time_steps + 1):
t = it * dt
t0 = time_mod.clock()
solver.solve()
solution.assign(solution_new)
t1 = time_mod.clock()
if do_export:
out_file.write(solution)
print('{:4d} T={:9.1f} s cpu={:.2f} s'.format(it, t, t1 - t0))
if do_assert:
target_u_min = 0.4
target_u_max = 1.0
target_u_tol = 5e-2
target_zero = 1e-6
solution_p1_dg.project(solution)
uvw = solution_p1_dg.dat.data
w_max = np.max(np.abs(uvw[:, 2]))
v_max = np.max(np.abs(uvw[:, 1]))
print('w {:}'.format(w_max))
print('v {:}'.format(v_max))
assert w_max < target_zero, 'z velocity component too large'
assert v_max < target_zero, 'y velocity component too large'
u_min = uvw[:, 0].min()
u_max = uvw[:, 0].max()
print('u {:} {:}'.format(u_min, u_max))
assert np.abs(
u_min - target_u_min) < target_u_tol, 'minimum u velocity is wrong'
assert np.abs(
u_max - target_u_max) < target_u_tol, 'maximum u velocity is wrong'
print('*** PASSED ***')
def test_implicit_diffusion(do_export=False, do_assert=True):
# set mesh resolution
scale = 1000.0
reso = 2.5 * scale
layers = 25 # 50
depth = 50
# generate unit mesh and transform its coords
x_max = 5.0 * scale
x_min = -5.0 * scale
lx = (x_max - x_min)
ly = 0.7 * lx
n_x = int(lx / reso)
area = lx * ly
mesh2d = RectangleMesh(n_x, n_x, lx, lx, reorder=True)
# move mesh, center at (0,0)
mesh2d.coordinates.dat.data[:, 0] -= lx / 2
mesh2d.coordinates.dat.data[:, 1] -= lx / 2
dz = depth / layers
v_elem_size = Constant(dz)
mesh = ExtrudedMesh(mesh2d, layers=layers, layer_height=-depth / layers)
if do_export:
sol_file = File('sol.pvd')
ana_file = File('ana_sol.pvd')
# define function spaces
fam = 'DG'
deg = 1
fs = get_functionspace(mesh, fam, deg)
solution = Function(fs, name='tracer')
solution_new = Function(fs, name='new tracer')
ana_sol = Function(fs, name='analytical tracer')
nu_v = 1.0
diffusivity_v = Constant(nu_v)
test = TestFunction(fs)
normal = FacetNormal(mesh)
# setup simulation time
t_init = 5.0
t = t_init
t_const = Constant(t)
t_end = 100.0
dt = dz * dz / nu_v / 10
n_iter = np.ceil(t_end / dt)
dt = t_end / n_iter
print('dt {:}'.format(dt))
dt_const = Constant(dt)
# analytical solution
u_max = 1.0
u_min = -1.0
z0 = -depth / 2.0
x, y, z = SpatialCoordinate(mesh)
ana_sol_expr = 0.5 * (u_max + u_min) - 0.5 * (u_max - u_min) * erf(
(z - z0) / sqrt(4 * diffusivity_v * t_const))
# initial condition
solution.project(ana_sol_expr)
ana_sol.project(ana_sol_expr)
def rhs(solution):
# vertical diffusion operator integrated by parts
# f = -diffusivity_v*inner(Dx(solution, 2), Dx(test, 2)) * dx
# interface term
# diffFlux = diffusivity_v*Dx(solution, 2)
# f += (dot(avg(diffFlux), test('+'))*normal[2]('+') +
# dot(avg(diffFlux), test('-'))*normal[2]('-')) * dS_h
# symmetric interior penalty stabilization
# L = Constant(depth/layers)
# nbNeigh = 2
# d = 3
# sigma = Constant((deg + 1)*(deg + d)/d * nbNeigh / 2) / L
# gamma = sigma*avg(diffusivity_v)
# jump_test = test('+')*normal[2]('+') + test('-')*normal[2]('-')
# f += gamma * dot(jump(solution), jump_test) * dS_h
def grad_v(a):
return as_vector((0, 0, Dx(a, 2)))
n = as_vector((0, 0, normal[2]))
grad_test = grad_v(test)
diff_flux = diffusivity_v * grad_v(solution)
diff_flux_jump = diffusivity_v * jump(solution, n)
f = -inner(grad_test, diff_flux) * dx
# Interior penalty as in horizontal case
degree_h, degree_v = fs.ufl_element().degree()
dim = 3.0
sigma = (degree_v + 1.0) * (degree_v + dim) / dim / v_elem_size
# sigma = 1.0/v_elem_size
alpha = avg(sigma)
ds_interior = (dS_h)
f += -alpha * inner(jump(test, n), diff_flux_jump) * ds_interior
f += +inner(avg(grad_test), diff_flux_jump) * ds_interior
f += +inner(jump(test, n), avg(diff_flux)) * ds_interior
# symmetric boundary terms
f += inner(test, dot(diff_flux, n)) * (ds_t + ds_b)
return f
# define solver
sp = {}
sp['ksp_atol'] = 1e-20
sp['ksp_rtol'] = 1e-20
sp['snes_rtol'] = 1e-20
sp['snes_atol'] = 1e-20
# sp['ksp_monitor'] = True
# sp['ksp_monitor_true_residual'] = True
# sp['snes_converged_reason'] = True
# sp['ksp_converged_reason'] = True
# Cr-Ni
# f = (inner(solution_new, test)*dx - inner(solution, test)*dx -
# dt_const*RHS(0.5*solution + 0.5*solution_new))
# prob = NonlinearVariationalProblem(f, solution_new)
# solver = LinearVariationalSolver(prob, solver_parameters=sp)
# Backward Euler
# f = (inner(solution_new, test)*dx - inner(solution, test)*dx -
# dt_const*RHS(solution_new))
# prob = NonlinearVariationalProblem(f, solution_new)
# solver = LinearVariationalSolver(prob, solver_parameters=sp)
# From DIRK(2,3,2) IMEX scheme in Ascher et al. (1997)
# This method has the Butcher tableau
#
# gamma | gamma 0
# 1 | 1-gamma gamma
# -------------------------
# | 0.5 0.5
# with
# gamma = (2 + sqrt(2))/2
#
solution_1 = Function(fs, name='tracer K1')
solution_2 = Function(fs, name='tracer K2')
gamma = Constant((2.0 + np.sqrt(2.0)) / 2.0)
f1 = (inner(solution_1, test) * dx - inner(solution, test) * dx -
gamma * dt_const * rhs(solution_1))
prob1 = NonlinearVariationalProblem(f1, solution_1)
solver1 = LinearVariationalSolver(prob1, solver_parameters=sp)
f2 = (inner(solution_2, test) * dx - inner(solution, test) * dx -
(1.0 - gamma) * dt_const * rhs(solution_1) -
(gamma) * dt_const * rhs(solution_2))
prob2 = NonlinearVariationalProblem(f2, solution_2)
solver2 = LinearVariationalSolver(prob2, solver_parameters=sp)
f = (inner(solution_new, test) * dx - inner(solution, test) * dx -
dt_const * (0.5 * rhs(solution_1) + 0.5 * rhs(solution_2)))
prob = NonlinearVariationalProblem(f, solution_new)
solver = LinearVariationalSolver(prob, solver_parameters=sp)
if do_export:
sol_file.write(solution)
ana_file.write(ana_sol)
print('sol {:} {:}'.format(solution.dat.data.min(),
solution.dat.data.max()))
print('ana {:} {:}'.format(ana_sol.dat.data.min(), ana_sol.dat.data.max()))
# time loop
while t < t_end + t_init:
# solve
solver1.solve()
solver2.solve()
solver.solve()
solution.assign(solution_new)
t += dt
# update analytical solution
t_const.assign(t)
ana_sol.project(ana_sol_expr)
if do_export:
sol_file.write(solution)
ana_file.write(ana_sol)
print('sol {:} {:}'.format(solution.dat.data.min(),
solution.dat.data.max()))
print('ana {:} {:}'.format(ana_sol.dat.data.min(), ana_sol.dat.data.max()))
l2_err = errornorm(ana_sol, solution) / area
print('L2 error: {:}'.format(l2_err))
if do_assert:
l2_threshold = 1e-4
assert l2_err < l2_threshold, 'L2 error exceeds threshold'
def p1(mesh, spaces):
(name, order), (vname, vorder) = spaces
return utility.get_functionspace(mesh, name, order, vname, vorder)
def p1_2d(mesh2d, spaces):
(name, order), (vname, vorder) = spaces
return utility.get_functionspace(mesh2d, name, order)
def mesh(mesh2d):
fs = utility.get_functionspace(mesh2d, 'CG', 1)
bathymetry_2d = Function(fs).assign(1.0)
n_layers = 10
return utility.extrude_mesh_sigma(mesh2d, n_layers, bathymetry_2d)
def compute(refinement=1, order=1, do_export=False):
print('--- soving refinement {:}'.format(refinement))
n = 5 * refinement
mesh = UnitSquareMesh(n, n)
family = 'DG'
p0dg = get_functionspace(mesh, family, order - 1)
p1dg = get_functionspace(mesh, family, order)
p1dg_v = get_functionspace(mesh, family, order, vector=True)
p1dg_ho = get_functionspace(mesh, family, order + 2)
p1dg_v_ho = get_functionspace(mesh, family, order + 2, vector=True)
lx = 1.0
x, y = SpatialCoordinate(mesh)
uv_expr = as_vector((sin(0.2 * pi * (3.0 * x + 1.0 * y) / lx),
0.2 * sin(0.2 * pi * (1.0 * x + 3.0 * y) / lx)))
div_expr = 0.12 * pi * cos(0.2 * pi *
(1.0 * x + 3.0 * y) / lx) / lx + 0.6 * pi * cos(
0.2 * pi * (3.0 * x + 1.0 * y) / lx) / lx
div_uv = Function(p1dg, name='div')
div_uv.project(div_expr)
div_uv_source = Function(p0dg, name='div')
div_uv_source.project(div_expr)
uv = Function(p1dg_v, name='uv')
uv.project(uv_expr)
div_ana = Function(p1dg_ho, name='div_ho')
div_ana.project(div_expr)
uv_ana = Function(p1dg_v_ho, name='uv+ho')
uv_ana.project(uv_expr)
if do_export:
print('analytical div {:} {:}'.format(div_uv.dat.data.min(),
div_uv.dat.data.max()))
# export analytical solutions
out_uv = File('uv.pvd')
out_div = File('div.pvd')
out_uv.write(uv)
out_div.write(div_uv)
div_uv.assign(0)
test = TestFunction(p1dg)
tri = TrialFunction(p1dg)
normal = FacetNormal(mesh)
# solve div_uv = div(uv)
a = inner(test, tri) * dx
# div(uv) point-wise
# l = inner(div(uv), test)*dx
# div(uv) integrated by parts
l = -inner(grad(test), uv) * dx + dot(avg(uv), jump(
test, normal)) * dS + test * dot(uv, normal) * ds
# analytical source
# l = inner(test, div_uv_source)*dx
solve(a == l, div_uv)
if do_export:
print('numerical div {:} {:}'.format(div_uv.dat.data.min(),
div_uv.dat.data.max()))
# export numerical solutions
out_uv.write(uv)
out_div.write(div_uv)
l2err_uv = errornorm(uv_ana, uv)
l2err_div = errornorm(div_ana, div_uv)
print('L2 norm uv {:}'.format(l2err_uv))
print('L2 norm div'.format(l2err_div))
return l2err_uv, l2err_div
def test():
"""
Tests atmospheric model data interpolation.
.. note::
The following files must be present
forcings/atm/wrf/wrf_air.2015_05_16.nc
forcings/atm/wrf/wrf_air.2015_05_17.nc
forcings/atm/nam/nam_air.local.2006_05_01.nc
forcings/atm/nam/nam_air.local.2006_05_02.nc
"""
mesh2d = Mesh('mesh_cre-plume_03_normal.msh')
comm = mesh2d.comm
p1 = get_functionspace(mesh2d, 'CG', 1)
p1v = get_functionspace(mesh2d, 'CG', 1, vector=True)
windstress_2d = Function(p1v, name='wind stress')
atmpressure_2d = Function(p1, name='atm pressure')
sim_tz = timezone.FixedTimeZone(-8, 'PST')
# WRF
# init_date = datetime.datetime(2015, 5, 16, tzinfo=sim_tz)
# pattern = 'forcings/atm/wrf/wrf_air.2015_*_*.nc'
# atm_time_step = 3600. # for verification only
# test_atm_file = 'forcings/atm/wrf/wrf_air.2015_05_16.nc'
# NAM
init_date = datetime.datetime(2006, 5, 1, tzinfo=sim_tz)
pattern = 'forcings/atm/nam/nam_air.local.2006_*_*.nc'
atm_time_step = 3*3600.
test_atm_file = 'forcings/atm/nam/nam_air.local.2006_05_01.nc'
atm_interp = ATMInterpolator(p1, windstress_2d, atmpressure_2d,
to_latlon,
pattern, init_date, COORDSYS, verbose=True)
# create a naive interpolation for first file
xy = SpatialCoordinate(p1.mesh())
fsx = Function(p1).interpolate(xy[0]).dat.data_with_halos
fsy = Function(p1).interpolate(xy[1]).dat.data_with_halos
mesh_lonlat = []
for node in range(len(fsx)):
lat, lon = to_latlon(fsx[node], fsy[node])
mesh_lonlat.append((lon, lat))
mesh_lonlat = np.array(mesh_lonlat)
ncfile = netCDF4.Dataset(test_atm_file)
itime = 6
grid_lat = ncfile['lat'][:].ravel()
grid_lon = ncfile['lon'][:].ravel()
grid_lonlat = np.array((grid_lon, grid_lat)).T
grid_pres = ncfile['prmsl'][itime, :, :].ravel()
pres = scipy.interpolate.griddata(grid_lonlat, grid_pres, mesh_lonlat, method='linear')
grid_uwind = ncfile['uwind'][itime, :, :].ravel()
uwind = scipy.interpolate.griddata(grid_lonlat, grid_uwind, mesh_lonlat, method='linear')
grid_vwind = ncfile['vwind'][itime, :, :].ravel()
vwind = scipy.interpolate.griddata(grid_lonlat, grid_vwind, mesh_lonlat, method='linear')
vrot = coordsys.VectorCoordSysRotation(coordsys.LL_WGS84, COORDSYS, mesh_lonlat[:, 0], mesh_lonlat[:, 1])
uwind, vwind = vrot(uwind, vwind)
u_stress, v_stress = compute_wind_stress(uwind, vwind)
# compare
atm_interp.set_fields(itime*atm_time_step - 8*3600.) # NOTE timezone offset
assert np.allclose(pres, atmpressure_2d.dat.data_with_halos)
assert np.allclose(u_stress, windstress_2d.dat.data_with_halos[:, 0])
# write fields to disk for visualization
pres_fn = 'tmp/AtmPressure2d.pvd'
wind_fn = 'tmp/WindStress2d.pvd'
print('Saving output to {:} {:}'.format(pres_fn, wind_fn))
out_pres = File(pres_fn)
out_wind = File(wind_fn)
hours = 24*1.5
granule = 4
simtime = np.arange(granule*hours)*3600./granule
i = 0
for t in simtime:
atm_interp.set_fields(t)
norm_atm = norm(atmpressure_2d)
norm_wind = norm(windstress_2d)
if comm.rank == 0:
print('{:} {:} {:} {:}'.format(i, t, norm_atm, norm_wind))
out_pres.write(atmpressure_2d)
out_wind.write(windstress_2d)
i += 1