def tracer_sphere(tmpdir, degree):
radius = 1
mesh = IcosahedralSphereMesh(radius=radius, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
# Parameters chosen so that dt != 1
# Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0)
# to demonstrate that transport is working correctly
dt = pi / 3. * 0.02
output = OutputParameters(dirname=str(tmpdir), dumpfreq=15)
state = State(mesh, dt=dt, output=output)
umax = 1.0
uexpr = as_vector([-umax * x[1] / radius, umax * x[0] / radius, 0.0])
tmax = pi / 2
f_init = exp(-x[2]**2 - x[0]**2)
f_end = exp(-x[2]**2 - x[1]**2)
tol = 0.05
return TracerSetup(state, tmax, f_init, f_end, "BDM", degree, uexpr, umax,
radius, tol)
def setup_DGadvection(element, vector=False):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 1.0
dt = pi / 3 * 0.01
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ["u", "D"]
timestepping = TimesteppingParameters(dt=dt)
state = ShallowWaterState(
mesh, vertical_degree=None, horizontal_degree=2, family="BDM", timestepping=timestepping, fieldlist=fieldlist
)
# interpolate initial conditions
u0 = Function(state.V[0], name="velocity")
x = SpatialCoordinate(mesh)
uexpr = as_vector([-x[1], x[0], 0.0])
u0.project(uexpr)
if vector:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = VectorFunctionSpace(mesh, "DG", 1)
Space = VectorFunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression(("exp(-pow(x[2],2) - pow(x[1],2))", "0.0", "0.0"))
f_end = Function(Space)
f_end_expr = Expression(("exp(-pow(x[2],2) - pow(x[0],2))", "0", "0"))
else:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = FunctionSpace(mesh, "DG", 1)
Space = FunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression("exp(-pow(x[2],2) - pow(x[1],2))")
f_end = Function(Space)
f_end_expr = Expression("exp(-pow(x[2],2) - pow(x[0],2))")
f.interpolate(fexpr)
f_end.interpolate(f_end_expr)
return state, BrokenSpace, u0, f, f_end
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24. * 60. * 60.
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements,
degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
dt = 3600.
output = OutputParameters(dirname=dirname + "/sw_linear_w2",
steady_state_error_fields=['u', 'D'],
dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
state = State(mesh, dt=dt, output=output, parameters=parameters)
# Coriolis
Omega = parameters.Omega
fexpr = 2 * Omega * x[2] / R
eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr)
# interpolate initial conditions
# Initial/current conditions
u0 = state.fields("u")
D0 = state.fields("D")
u_max = 2 * pi * R / (12 * day
) # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
g = parameters.g
Dexpr = H - ((R * Omega * u_max) * (x[2] * x[2] / (R * R))) / g
u0.project(uexpr)
D0.interpolate(Dexpr)
transport_schemes = [ForwardEuler(state, "D")]
# build time stepper
stepper = CrankNicolson(state, eqns, transport_schemes)
return stepper, 2 * day
def state(tmpdir, geometry):
"""
returns an instance of the State class, having set up either spherical
geometry or 2D vertical slice geometry
"""
output = OutputParameters(dirname=str(tmpdir), dumplist=["f"], dumpfreq=15)
if geometry == "sphere":
mesh = IcosahedralSphereMesh(radius=1, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
family = "BDM"
vertical_degree = None
fieldlist = ["u", "D"]
dt = pi / 3. * 0.01
uexpr = as_vector([-x[1], x[0], 0.0])
if geometry == "slice":
m = PeriodicIntervalMesh(15, 1.)
mesh = ExtrudedMesh(m, layers=15, layer_height=1. / 15.)
family = "CG"
vertical_degree = 1
fieldlist = ["u", "rho", "theta"]
dt = 0.01
x = SpatialCoordinate(mesh)
uexpr = as_vector([1.0, 0.0])
timestepping = TimesteppingParameters(dt=dt)
state = State(mesh,
vertical_degree=vertical_degree,
horizontal_degree=1,
family=family,
timestepping=timestepping,
output=output,
fieldlist=fieldlist)
u0 = state.fields("u")
u0.project(uexpr)
return state
ref_dt = {3: 900., 4: 450., 5: 225., 6: 112.5}
tmax = 50*day
# setup shallow water parameters
R = 6371220.
H = 5960.
# setup input that doesn't change with ref level or dt
fieldlist = ['u', 'D']
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
for ref_level, dt in ref_dt.items():
dirname = "sw_W5_ref%s_dt%s" % (ref_level, dt)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=ref_level, degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname, dumplist_latlon=['D'], dumpfreq=100)
diagnostic_fields = [Sum('D', 'topography')]
state = State(mesh, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
def run_profliler(hybridization,
model_degree,
model_family,
mesh_degree,
cfl,
refinements,
layers,
debug,
rtol,
flexsolver=True,
stronger_smoother=False,
suppress_data_output=False):
nlayers = layers # Number of vertical layers
refinements = refinements # Number of horiz. cells = 20*(4^refinements)
hybrid = bool(hybridization)
# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6 # Radius of the Earth (m)
X = 1.0 # Reduced-size Earth reduction factor
a = a_ref / X # Scaled radius of planet (m)
g = parameters.g # Acceleration due to gravity (m/s^2)
N = parameters.N # Brunt-Vaisala frequency (1/s)
p_0 = parameters.p_0 # Reference pressure (Pa, not hPa)
c_p = parameters.cp # SHC of dry air at const. pressure (J/kg/K)
R_d = parameters.R_d # Gas constant for dry air (J/kg/K)
kappa = parameters.kappa # R_d/c_p
T_eq = 300.0 # Isothermal atmospheric temperature (K)
p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator
u_0 = 20.0 # Maximum amplitude of the zonal wind (m/s)
d = 5000.0 # Width parameter for Theta'
lamda_c = 2.0 * np.pi / 3.0 # Longitudinal centerpoint of Theta'
phi_c = 0.0 # Lat. centerpoint of Theta' (equator)
deltaTheta = 1.0 # Maximum amplitude of Theta' (K)
L_z = 20000.0 # Vert. wave length of the Theta' perturb.
gamma = (1 - kappa) / kappa
cs = sqrt(c_p * T_eq / gamma) # Speed of sound in an air parcel
if model_family == "RTCF":
# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
elif model_family == "RT" or model_family == "BDFM":
m = IcosahedralSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
else:
raise ValueError("Unknown family: %s" % model_family)
cell_vs = interpolate(CellVolume(m), FunctionSpace(m, "DG", 0))
a_max = fmax(cell_vs)
dx_max = sqrt(a_max)
u_max = u_0
PETSc.Sys.Print("\nDetermining Dt from specified horizontal CFL: %s" % cfl)
dt = int(cfl * (dx_max / cs))
# Height position of the model top (m)
z_top = 1.0e4
deltaz = z_top / nlayers
vertical_cfl = dt * (cs / deltaz)
PETSc.Sys.Print("""
Problem parameters:\n
Profiling linear solver for the compressible Euler equations.\n
Speed of sound in compressible atmosphere: %s,\n
Hybridized compressible solver: %s,\n
Model degree: %s,\n
Model discretization: %s,\n
Mesh degree: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Dx (max, m): %s,\n
Dz (m): %s,\n
Dt (s): %s,\n
horizontal CFL: %s,\n
vertical CFL: %s.
""" % (cs, hybrid, model_degree, model_family, mesh_degree, refinements,
nlayers, dx_max, deltaz, dt, cfl, vertical_cfl))
# Build volume mesh
mesh = ExtrudedMesh(m,
layers=nlayers,
layer_height=deltaz,
extrusion_type="radial")
x = SpatialCoordinate(mesh)
# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
W_Q1 = FunctionSpace(mesh, "CG", 1)
z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a
z = Function(W_Q1).interpolate(z_expr)
lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]))
lat = Function(W_Q1).interpolate(lat_expr)
lon = Function(W_Q1).interpolate(atan_2(x[1], x[0]))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=1, maxi=1)
dirname = 'meanflow_ref'
if hybrid:
dirname += '_hybridization'
# No output
output = OutputParameters(dumpfreq=3600,
dirname=dirname,
perturbation_fields=['theta', 'rho'],
dump_vtus=False,
dump_diagnostics=False,
checkpoint=False,
log_level='INFO')
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
vertical_degree=model_degree,
horizontal_degree=model_degree,
family=model_family,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Initial conditions
u0 = state.fields.u
theta0 = state.fields.theta
rho0 = state.fields.rho
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
x = SpatialCoordinate(mesh)
# Random velocity field
CG2 = VectorFunctionSpace(mesh, "CG", 2)
urand = Function(CG2)
urand.dat.data[:] += np.random.randn(*urand.dat.data.shape)
u0.project(urand)
# Surface temperature
G = g**2 / (N**2 * c_p)
Ts_expr = G + (T_eq - G) * exp(-(u_max * N**2 / (4 * g * g)) * u_max *
(cos(2.0 * lat) - 1.0))
Ts = Function(W_Q1).interpolate(Ts_expr)
# Surface pressure
ps_expr = p_eq * exp((u_max / (4.0 * G * R_d)) * u_max *
(cos(2.0 * lat) - 1.0)) * (Ts / T_eq)**(1.0 / kappa)
ps = Function(W_Q1).interpolate(ps_expr)
# Background pressure
p_expr = ps * (1 + G / Ts * (exp(-N**2 * z / g) - 1))**(1.0 / kappa)
p = Function(W_Q1).interpolate(p_expr)
# Background temperature
Tb_expr = G * (1 - exp(N**2 * z / g)) + Ts * exp(N**2 * z / g)
Tb = Function(W_Q1).interpolate(Tb_expr)
# Background potential temperature
thetab_expr = Tb * (p_0 / p)**kappa
thetab = Function(W_Q1).interpolate(thetab_expr)
theta_b = Function(theta0.function_space()).interpolate(thetab)
rho_b = Function(rho0.function_space())
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a * acos(sin_tmp + cos_tmp * cos(lon - lamda_c))
s = (d**2) / (d**2 + r**2)
theta_pert = deltaTheta * s * sin(2 * np.pi * z / L_z)
theta0.interpolate(theta_pert)
# Compute the balanced density
PETSc.Sys.Print("Computing balanced density field...\n")
# Use vert. hybridization preconditioner for initialization
pi_params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
'vert_hybridization': {
'ksp_type': 'gmres',
'pc_type': 'gamg',
'pc_gamg_sym_graph': True,
'ksp_rtol': 1e-12,
'ksp_atol': 1e-12,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_max_it': 3,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
if debug:
pi_params['vert_hybridization']['ksp_monitor_true_residual'] = None
compressible_hydrostatic_balance(state,
theta_b,
rho_b,
top=False,
pi_boundary=(p / p_0)**kappa,
solve_for_rho=False,
params=pi_params)
# Random potential temperature perturbation
theta0.assign(0.0)
theta0.dat.data[:] += np.random.randn(len(theta0.dat.data))
# Random density field
rho0.assign(0.0)
rho0.dat.data[:] += np.random.randn(len(rho0.dat.data))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt, equation_form="advective")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn, subcycles=2)))
advected_fields.append(
("theta", SSPRK3(state, theta0, thetaeqn, subcycles=2)))
# Set up linear solver
if hybrid:
outer_solver_type = "Hybrid_SCPC"
PETSc.Sys.Print("""
Setting up hybridized solver on the traces.""")
if flexsolver:
inner_solver_type = "fgmres_gamg_gmres_smoother"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {
'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
else:
inner_solver_type = "fgmres_ml_richardson"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'ml',
'pc_mg_cycles': 1,
'pc_ml_maxNlevels': 25,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_richardson_scale': 0.8,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
if stronger_smoother:
inner_parameters['mg_levels']['ksp_max_it'] = 5
inner_solver_type += "_stronger"
if debug:
PETSc.Sys.Print("""Debugging on.""")
inner_parameters['ksp_monitor_true_residual'] = None
PETSc.Sys.Print("Inner solver: %s" % inner_solver_type)
# Use Firedrake static condensation interface
solver_parameters = {
'mat_type': 'matfree',
'pmat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
'condensed_field': inner_parameters
}
linear_solver = HybridizedCompressibleSolver(
state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
else:
outer_solver_type = "gmres_SchurPC"
PETSc.Sys.Print("""
Setting up GCR fieldsplit solver with Schur complement PC.""")
solver_parameters = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'fgmres',
'ksp_max_it': 100,
'ksp_rtol': rtol,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
},
'fieldsplit_1': {
'ksp_type': 'preonly',
'ksp_max_it': 30,
'ksp_monitor_true_residual': None,
'pc_type': 'hypre',
'pc_hypre_type': 'boomeramg',
'pc_hypre_boomeramg_max_iter': 1,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_coarsen_type': 'Falgout',
'pc_hypre_boomeramg_smooth_type': 'Euclid',
'pc_hypre_boomeramg_eu_bj': 1,
'pc_hypre_boomeramg_interptype': 'classical',
'pc_hypre_boomeramg_P_max': 0,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_strong_threshold': 0.25,
'pc_hypre_boomeramg_max_levels': 25,
'pc_hypre_boomeramg_no_CF': False
}
}
inner_solver_type = "hypre"
if debug:
solver_parameters['ksp_monitor_true_residual'] = None
linear_solver = CompressibleSolver(state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
param_info = ParameterInfo(dt=dt,
deltax=dx_max,
deltaz=deltaz,
horizontal_courant=cfl,
vertical_courant=vertical_cfl,
family=model_family,
model_degree=model_degree,
mesh_degree=mesh_degree,
solver_type=outer_solver_type,
inner_solver_type=inner_solver_type)
# Build profiler
profiler = Profiler(parameterinfo=param_info,
state=state,
advected_fields=advected_fields,
linear_solver=linear_solver,
forcing=compressible_forcing,
suppress_data_output=suppress_data_output)
PETSc.Sys.Print("Starting profiler...\n")
profiler.run(t=0, tmax=dt)
def setup_sw(dirname, euler_poincare):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=1500.)
output = OutputParameters(dirname=dirname + "/sw",
dumplist_latlon=['D', 'D_error'],
steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [
RelativeVorticity(),
AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity', 'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity', 'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity', 'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')
]
state = State(mesh,
vertical_degree=None,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
Omega = parameters.Omega
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max * u_max / 2.0) * (x[2] * x[2] /
(R * R))) / g
# Coriolis
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", Function(V))
f.interpolate(fexpr) # Coriolis frequency (1/s)
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
if euler_poincare:
ueqn = EulerPoincare(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=True)
else:
ueqn = VectorInvariant(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=False)
Deqn = AdvectionEquation(state,
D0.function_space(),
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("D", SSPRK3(state, D0, Deqn)))
linear_solver = ShallowWaterSolver(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2 * u_max / R) * x[2] / R
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr + f) / D0)
return stepper, 0.25 * day
def setup_sw(dirname, dt, u_transport_option):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity',
'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity',
'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity',
'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Omega = parameters.Omega
fexpr = 2*Omega*x[2]/R
eqns = ShallowWaterEquations(state, family="BDM", degree=1,
fexpr=fexpr,
u_transport_option=u_transport_option)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2*u_max/R)*x[2]/R
f = state.fields("coriolis")
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr+f)/D0)
return state, eqns
from gusto import *
from firedrake import IcosahedralSphereMesh, Expression, SpatialCoordinate, \
Constant, as_vector
from math import pi
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24.*60.*60.
u_0 = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements, degree=3)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname='sw_linear_w2', steady_state_dump_err={'u':True, 'D':True})
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = ShallowWaterState(mesh, vertical_degree=None, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
fname = 'Galewsky_ec_flux'
write_latlon = True
nc_diag = {'Energy': 0., 'Enstrophy': 0., 'Mass': 0.}
h5_safe, nc_safe = True, True
if COMM_WORLD.Get_rank() == 0:
print("Simulation for Galewsky test case, model u-ad_flux", "\n",
"| Discr details: Poisson implicit, no EP, TM=0.5", "\n", "| dt", dt,
"| tmax", tmax, " | refinement level", ref_level, "\n", "| maxk",
maxk, "| dfr field, nc_h5:", field_dumpfreq, nc_h5_dumpfreq, "\n",
"| pickup", pickup, "| Diagnostics:", tuple(nc_diag.keys()))
print("Starting Initial condition, and function setup at", ctime())
R, Omega = 6371220., 7.292e-5
mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2)
mesh.init_cell_orientations(SpatialCoordinate(mesh))
x = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
def latlon_coords(mesh):
"""Compute latitude-longitude coordinates given Cartesian ones"""
x0, y0, z0 = SpatialCoordinate(mesh)
unsafe = z0 / sqrt(x0 * x0 + y0 * y0 + z0 * z0)
safe = Min(Max(unsafe, -1.0), 1.0) # avoid silly roundoff errors
theta = asin(safe) # latitude
lamda = atan_2(y0, x0) # longitude
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24.*60.*60.
u_0 = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements, degree=3)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_dump_err={'u':True,'D':True}, dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = ShallowWaterState(mesh, vertical_degree=None, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
g = parameters.g
Omega = parameters.Omega
# Coriolis expression
R = Constant(R)
Omega = Constant(parameters.Omega)
x = SpatialCoordinate(mesh)
fexpr = 2*Omega*x[2]/R
V = FunctionSpace(mesh, "CG", 1)
state.f = Function(V).interpolate(fexpr) # Coriolis frequency (1/s)
u_max = Constant(u_0)
# interpolate initial conditions
# Initial/current conditions
u0, D0 = Function(state.V[0]), Function(state.V[1])
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = Constant(parameters.g)
Dexpr = - ((R * Omega * u_max)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([u0, D0])
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["D"] = NoAdvection(state)
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state, linear=True)
# build time stepper
stepper = Timestepper(state, advection_dict, linear_solver,
sw_forcing)
return stepper, 2*day
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24. * 60. * 60.
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements,
degree=3)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname=dirname + "/sw_linear_w2",
steady_state_error_fields=['u', 'D'],
dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Coriolis
Omega = parameters.Omega
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", Function(V))
f.interpolate(fexpr) # Coriolis frequency (1/s)
# interpolate initial conditions
# Initial/current conditions
u0 = state.fields("u")
D0 = state.fields("D")
u_max = 2 * pi * R / (12 * day
) # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
g = parameters.g
Dexpr = -((R * Omega * u_max) * (x[2] * x[2] / (R * R))) / g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
Deqn = LinearAdvection(state,
D0.function_space(),
state.parameters.H,
ibp=IntegrateByParts.ONCE,
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("D", ForwardEuler(state, D0, Deqn)))
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state, linear=True)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
return stepper, 2 * day
