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