diff --git a/case_studies/compressible_euler/moist_skamarock_klemp.py b/case_studies/compressible_euler/moist_skamarock_klemp.py new file mode 100644 index 0000000..122be96 --- /dev/null +++ b/case_studies/compressible_euler/moist_skamarock_klemp.py @@ -0,0 +1,301 @@ +""" +A moist version of the non-hydrostatic vertical slice gravity wave test case of +Skamarock and Klemp, as described in Bendall et al, 2020: +``A compatible finite‐element discretisation for the moist compressible Euler +equations'', QJRMS. + +The gravity wave propagates through a saturated and cloudy atmosphere, so that +evaporation and condensation impact the development of the wave. +""" +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter + +from petsc4py import PETSc +PETSc.Sys.popErrorHandler() +from firedrake import ( + as_vector, SpatialCoordinate, PeriodicIntervalMesh, ExtrudedMesh, exp, sin, + Function, pi, errornorm, TestFunction, dx, BrokenElement, FunctionSpace, + NonlinearVariationalProblem, NonlinearVariationalSolver +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + Perturbation, thermodynamics, CompressibleParameters, + CompressibleEulerEquations, RungeKuttaFormulation, CompressibleSolver, + WaterVapour, CloudWater, Theta_e, Recoverer, saturated_hydrostatic_balance, + EmbeddedDGOptions, ForwardEuler, SaturationAdjustment, SubcyclingOptions +) + +moist_skamarock_klemp_defaults = { + 'ncolumns': 150, + 'nlayers': 10, + 'dt': 60.0, + 'tmax': 3600., + 'dumpfreq': 25, + 'dirname': 'moist_skamarock_klemp' +} + + +def moist_skamarock_klemp( + ncolumns=moist_skamarock_klemp_defaults['ncolumns'], + nlayers=moist_skamarock_klemp_defaults['nlayers'], + dt=moist_skamarock_klemp_defaults['dt'], + tmax=moist_skamarock_klemp_defaults['tmax'], + dumpfreq=moist_skamarock_klemp_defaults['dumpfreq'], + dirname=moist_skamarock_klemp_defaults['dirname'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 3.0e5 # Width of domain (m) + domain_height = 1.0e4 # Height of domain (m) + Tsurf = 300. # Temperature at surface (K) + wind_initial = 20. # Initial wind in x direction (m/s) + pert_width = 5.0e3 # Width parameter of perturbation (m) + deltaTheta = 1.0e-2 # Magnitude of theta perturbation (K) + N = 0.01 # Brunt-Vaisala frequency (1/s) + total_water = 0.02 # total moisture mixing ratio, in kg/kg + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + alpha = 0.5 + u_eqn_type = 'vector_invariant_form' + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + parameters = CompressibleParameters() + tracers = [WaterVapour(), CloudWater()] + eqns = CompressibleEulerEquations( + domain, parameters, active_tracers=tracers, + u_transport_option=u_eqn_type + ) + + # I/O + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True + ) + + diagnostic_fields = [Theta_e(eqns), Perturbation('Theta_e')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + Vt_brok = FunctionSpace( + mesh, BrokenElement(domain.spaces("theta").ufl_element()) + ) + embedded_dg = EmbeddedDGOptions(embedding_space=Vt_brok) + subcycling_opts = SubcyclingOptions(subcycle_by_courant=0.25) + transported_fields = [ + SSPRK3( + domain, "u", subcycling_options=subcycling_opts, + rk_formulation=RungeKuttaFormulation.predictor + ), + SSPRK3( + domain, "rho", subcycling_options=subcycling_opts, + rk_formulation=RungeKuttaFormulation.linear + ), + SSPRK3( + domain, "theta", + subcycling_options=subcycling_opts, options=embedded_dg + ), + SSPRK3( + domain, "water_vapour", + subcycling_options=subcycling_opts, options=embedded_dg + ), + SSPRK3( + domain, "cloud_water", + subcycling_options=subcycling_opts, options=embedded_dg + ) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho", advective_then_flux=True), + DGUpwind(eqns, "theta"), + DGUpwind(eqns, "water_vapour"), + DGUpwind(eqns, "cloud_water") + ] + + # Linear solver + tau_values = {'rho': 1.0, 'theta': 1.0} + linear_solver = CompressibleSolver(eqns, alpha=alpha, tau_values=tau_values) + + # Physics schemes + physics_schemes = [ + (SaturationAdjustment(eqns), ForwardEuler(domain)) + ] + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, physics_schemes=physics_schemes, + alpha=alpha, reference_update_freq=1, accelerator=True + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") + + # spaces + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Thermodynamic constants required for setting initial conditions + # and reference profiles + g = parameters.g + + x, z = SpatialCoordinate(mesh) + quadrature_degree = (4, 4) + dxp = dx(degree=(quadrature_degree)) + + # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) + theta_e = Function(Vt).interpolate(Tsurf*exp(N**2*z/g)) + water_t = Function(Vt).assign(total_water) + u0.project(as_vector([wind_initial, 0.0])) + + # Calculate hydrostatic fields + saturated_hydrostatic_balance(eqns, stepper.fields, theta_e, water_t) + + # Add perturbation to theta_e ---------------------------------------------- + # Need to store background state for theta_e + theta_e_b = stepper.fields("Theta_e_bar", space=Vt) + theta_e_b.assign(theta_e) + theta_e_pert = ( + deltaTheta * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + theta_e.interpolate(theta_e + theta_e_pert) + + # Find perturbed rho, water_v and theta_vd --------------------------------- + # expressions for finding theta0 and water_v0 from theta_e and water_t + rho_averaged = Function(Vt) + rho_recoverer = Recoverer(rho0, rho_averaged) + + # make expressions to evaluate residual + exner_expr = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) + p_expr = thermodynamics.p(eqns.parameters, exner_expr) + T_expr = thermodynamics.T(eqns.parameters, theta0, exner_expr, water_v0) + theta_e_expr = thermodynamics.theta_e(eqns.parameters, T_expr, p_expr, water_v0, water_t) + msat_expr = thermodynamics.r_sat(eqns.parameters, T_expr, p_expr) + + # Create temporary functions for iterating towards correct initial conditions + water_v_h = Function(Vt) + theta_h = Function(Vt) + theta_e_test = Function(Vt) + rho_h = Function(Vr) + zeta = TestFunction(Vr) + + # Set existing ref variables (which define the pressure which is unchanged) + rho_b = Function(Vr).assign(rho0) + theta_b = Function(Vt).assign(theta0) + + rho_form = zeta * rho_h * theta0 * dxp - zeta * rho_b * theta_b * dxp + rho_prob = NonlinearVariationalProblem(rho_form, rho_h) + rho_solver = NonlinearVariationalSolver(rho_prob) + + max_outer_solve_count = 40 + max_theta_solve_count = 15 + max_inner_solve_count = 5 + delta = 0.8 + + # We have to simultaneously solve for the initial rho, theta_vd and water_v + # This is done by nested fixed point iterations + for _ in range(max_outer_solve_count): + + rho_solver.solve() + rho0.assign(rho0 * (1 - delta) + delta * rho_h) + rho_recoverer.project() + + theta_e_test.interpolate(theta_e_expr) + if errornorm(theta_e_test, theta_e) < 1e-6: + break + + for _ in range(max_theta_solve_count): + theta_h.interpolate(theta_e / theta_e_expr * theta0) + theta0.assign(theta0 * (1 - delta) + delta * theta_h) + + # break when close enough + if errornorm(theta_e_test, theta_e) < 1e-6: + break + for _ in range(max_inner_solve_count): + water_v_h.interpolate(msat_expr) + water_v0.assign(water_v0 * (1 - delta) + delta * water_v_h) + + # break when close enough + theta_e_test.interpolate(theta_e_expr) + if errornorm(theta_e_test, theta_e) < 1e-6: + break + + water_c0.assign(water_t - water_v0) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncolumns', + help="The number of columns in the vertical slice mesh.", + type=int, + default=moist_skamarock_klemp_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=moist_skamarock_klemp_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=moist_skamarock_klemp_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=moist_skamarock_klemp_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=moist_skamarock_klemp_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=moist_skamarock_klemp_defaults['dirname'] + ) + args, unknown = parser.parse_known_args() + + moist_skamarock_klemp(**vars(args)) diff --git a/case_studies/compressible_euler/mountain_hydrostatic.py b/case_studies/compressible_euler/mountain_hydrostatic.py new file mode 100644 index 0000000..9b0d14c --- /dev/null +++ b/case_studies/compressible_euler/mountain_hydrostatic.py @@ -0,0 +1,322 @@ +""" +The hydrostatic 1 metre high mountain test case from Melvin et al, 2010: +``An inherently mass-conserving iterative semi-implicit semi-Lagrangian +discretization of the non-hydrostatic vertical-slice equations.'', QJRMS. + +This test describes a wave over a 1m high mountain in an isothermal atmosphere. +The domain is larger than the "non-hydrostatic mountain" case, so the solutions +between the hydrostatic and non-hydrostatic equations should be similar. + +The setup used here uses the order 1 finite elements. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + as_vector, VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh, sqrt, + SpatialCoordinate, exp, Function, Mesh, Constant +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, ZComponent, Perturbation, Temperature, Exner, + SpongeLayerParameters, CompressibleParameters, CompressibleSolver, logger, + compressible_hydrostatic_balance, MinKernel, MaxKernel, + HydrostaticCompressibleEulerEquations, CompressibleEulerEquations, + hydrostatic_parameters +) + +mountain_hydrostatic_defaults = { + 'ncolumns': 100, + 'nlayers': 50, + 'dt': 12.5, + 'tmax': 15000., + 'dumpfreq': 600, + 'dirname': 'mountain_hydrostatic', + 'hydrostatic': False +} + + +def mountain_hydrostatic( + ncolumns=mountain_hydrostatic_defaults['ncolumns'], + nlayers=mountain_hydrostatic_defaults['nlayers'], + dt=mountain_hydrostatic_defaults['dt'], + tmax=mountain_hydrostatic_defaults['tmax'], + dumpfreq=mountain_hydrostatic_defaults['dumpfreq'], + dirname=mountain_hydrostatic_defaults['dirname'], + hydrostatic=mountain_hydrostatic_defaults['hydrostatic'] +): + + # ------------------------------------------------------------------------ # + # Parameters for test case + # ------------------------------------------------------------------------ # + + domain_width = 240000. # width of domain in x direction, in m + domain_height = 50000. # height of model top, in m + a = 10000. # scale width of mountain, in m + hm = 1. # height of mountain, in m + Tsurf = 250. # temperature of surface, in K + initial_wind = 20.0 # initial horizontal wind, in m/s + sponge_depth = 20000.0 # depth of sponge layer, in m + g = 9.80665 # acceleration due to gravity, in m/s^2 + cp = 1004. # specific heat capacity at constant pressure + mu_dt = 0.3 # parameter for strength of sponge layer, no units + exner_surf = 1.0 # maximum value of Exner pressure at surface + max_inner_iters = 10 # maximum number of hydrostatic balance iterations + tolerance = 1e-7 # tolerance for hydrostatic balance iteration + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + u_eqn_type = 'vector_invariant_form' + alpha = 0.55 # Necessary to absorb grid scale waves + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + # Make normal extruded mesh which will be distorted to describe the mountain + base_mesh = PeriodicIntervalMesh(ncolumns, domain_width) + ext_mesh = ExtrudedMesh( + base_mesh, layers=nlayers, layer_height=domain_height/nlayers + ) + Vc = VectorFunctionSpace(ext_mesh, "DG", 2) + + # Describe the mountain + xc = domain_width/2. + x, z = SpatialCoordinate(ext_mesh) + zs = hm * a**2 / ((x - xc)**2 + a**2) + xexpr = as_vector( + [x, z + ((domain_height - z) / domain_height) * zs] + ) + + # Make new mesh + new_coords = Function(Vc).interpolate(xexpr) + mesh = Mesh(new_coords) + mesh._base_mesh = base_mesh # Force new mesh to inherit original base mesh + domain = Domain(mesh, dt, "CG", element_order) + + # Equation + parameters = CompressibleParameters(g=g, cp=cp) + sponge = SpongeLayerParameters( + H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt + ) + if hydrostatic: + eqns = HydrostaticCompressibleEulerEquations( + domain, parameters, sponge_options=sponge, + u_transport_option=u_eqn_type + ) + else: + eqns = CompressibleEulerEquations( + domain, parameters, sponge_options=sponge, + u_transport_option=u_eqn_type + ) + + # I/O + # Adjust default directory name + if hydrostatic and dirname == mountain_hydrostatic_defaults['dirname']: + dirname = f'hyd_switch_{dirname}' + + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=True, dump_nc=True + ) + diagnostic_fields = [ + ZComponent('u', space=domain.spaces('theta')), + Perturbation('theta'), Exner(parameters), Temperature(eqns) + ] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + theta_opts = SUPGOptions() + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp) + ] + + # Linear solver + if hydrostatic: + linear_solver = CompressibleSolver( + eqns, alpha, solver_parameters=hydrostatic_parameters, + overwrite_solver_parameters=True + ) + else: + linear_solver = CompressibleSolver(eqns, alpha) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver, alpha=alpha + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Isothermal initial conditions, means we don't know initial theta + x, z = SpatialCoordinate(mesh) + + # Calculate hydrostatic exner + exner = Function(Vr) + rho_b = Function(Vr) + theta_b = Function(Vt) + + # Set up kernels to evaluate global minima and maxima of fields + min_kernel = MinKernel() + max_kernel = MaxKernel() + + # First guess for theta -- should be pretty good as comes from hydrostatic + # balance + # dp/dz = - rho*g + # = - p*g/(Rd*T) + # p = p0 exp[-g*z/(Rd*T)] + # theta = T*(p0/p)**-(Rd/cp) + # theta = T*exp[g*z/(cp*T)] + g = parameters.g + cp = parameters.cp + N = g / sqrt(cp*Tsurf) + theta_b.interpolate(Tsurf*exp(N**2*z/g)) + + # First solve hydrostatic balance that gives Exner = 1 at bottom boundary + # This gives us a guess for the top boundary condition + bottom_boundary = Constant(exner_surf, domain=mesh) + logger.info(f'Solving hydrostatic with bottom Exner of {exner_surf}') + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary + ) + + # Solve hydrostatic balance again, but now use minimum value from first + # solve as the *top* boundary condition for Exner + top_value = min_kernel.apply(exner) + top_boundary = Constant(top_value, domain=mesh) + logger.info(f'Solving hydrostatic with top Exner of {top_value}') + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + ) + + max_bottom_value = max_kernel.apply(exner) + + # Now we iterate, adjusting the top boundary condition, until this gives + # a maximum value of 1.0 at the surface + lower_top_guess = 0.9*top_value + upper_top_guess = 1.2*top_value + for i in range(max_inner_iters): + # If max bottom Exner value is equal to desired value, stop iteration + if abs(max_bottom_value - exner_surf) < tolerance: + break + + # Make new guess by average of previous guesses + top_guess = 0.5*(lower_top_guess + upper_top_guess) + top_boundary.assign(top_guess) + + logger.info( + f'Solving hydrostatic balance iteration {i}, with top ' + + f'Exner value of {top_guess}' + ) + + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary + ) + + max_bottom_value = max_kernel.apply(exner) + + # Adjust guesses based on new value + if max_bottom_value < exner_surf: + lower_top_guess = top_guess + else: + upper_top_guess = top_guess + + logger.info(f'Final max bottom Exner value of {max_bottom_value}') + + # Perform a final solve to obtain hydrostatically balanced rho + compressible_hydrostatic_balance( + eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary, + solve_for_rho=True + ) + + theta0.assign(theta_b) + rho0.assign(rho_b) + u0.project(as_vector([initial_wind, 0.0]), bcs=eqns.bcs['u']) + + stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncolumns', + help="The number of columns in the vertical slice mesh.", + type=int, + default=mountain_hydrostatic_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=mountain_hydrostatic_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=mountain_hydrostatic_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=mountain_hydrostatic_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=mountain_hydrostatic_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=mountain_hydrostatic_defaults['dirname'] + ) + parser.add_argument( + '--hydrostatic', + help=( + "Whether to use the hydrostatic switch to emulate the " + + "hydrostatic equations. Otherwise use the full non-hydrostatic" + + "equations." + ), + action="store_true", + default=mountain_hydrostatic_defaults['hydrostatic'] + ) + args, unknown = parser.parse_known_args() + + mountain_hydrostatic(**vars(args)) diff --git a/case_studies/compressible_euler/mountain_nonhydrostatic.py b/case_studies/compressible_euler/mountain_nonhydrostatic.py index f73b61c..a99bc45 100644 --- a/case_studies/compressible_euler/mountain_nonhydrostatic.py +++ b/case_studies/compressible_euler/mountain_nonhydrostatic.py @@ -3,7 +3,9 @@ ``An inherently mass-conserving iterative semi-implicit semi-Lagrangian discretization of the non-hydrostatic vertical-slice equations.'', QJRMS. -This test describes a wave over a mountain in a non-hydrostatic atmosphere. +This test describes a wave over a 1m high mountain. The domain is smaller than +that in the "non-hydrostatic mountain" case, so the solutions between the +hydrostatic and non-hydrostatic equations should be different. The setup used here uses the order 1 finite elements. """ @@ -15,19 +17,21 @@ ) from gusto import ( Domain, CompressibleParameters, CompressibleSolver, logger, - CompressibleEulerEquations, OutputParameters, IO, SSPRK3, - DGUpwind, SemiImplicitQuasiNewton, compressible_hydrostatic_balance, - SpongeLayerParameters, Exner, ZComponent, Perturbation, - SUPGOptions, TrapeziumRule, MaxKernel, MinKernel + CompressibleEulerEquations, HydrostaticCompressibleEulerEquations, + OutputParameters, IO, SSPRK3, DGUpwind, SemiImplicitQuasiNewton, + compressible_hydrostatic_balance, SpongeLayerParameters, Exner, ZComponent, + Perturbation, SUPGOptions, TrapeziumRule, MaxKernel, MinKernel, + hydrostatic_parameters ) mountain_nonhydrostatic_defaults = { - 'ncolumns': 180, - 'nlayers': 70, + 'ncolumns': 90, + 'nlayers': 35, 'dt': 5.0, 'tmax': 9000., 'dumpfreq': 450, - 'dirname': 'mountain_nonhydrostatic' + 'dirname': 'mountain_nonhydrostatic', + 'hydrostatic': False } @@ -37,7 +41,8 @@ def mountain_nonhydrostatic( dt=mountain_nonhydrostatic_defaults['dt'], tmax=mountain_nonhydrostatic_defaults['tmax'], dumpfreq=mountain_nonhydrostatic_defaults['dumpfreq'], - dirname=mountain_nonhydrostatic_defaults['dirname'] + dirname=mountain_nonhydrostatic_defaults['dirname'], + hydrostatic=mountain_nonhydrostatic_defaults['hydrostatic'] ): # ------------------------------------------------------------------------ # @@ -54,7 +59,7 @@ def mountain_nonhydrostatic( sponge_depth = 10000.0 # depth of sponge layer, in m g = 9.80665 # acceleration due to gravity, in m/s^2 cp = 1004. # specific heat capacity at constant pressure - sponge_mu = 0.15 # parameter for strength of sponge layer, in J/kg/K + mu_dt = 0.15 # parameter for strength of sponge layer, no units exner_surf = 1.0 # maximum value of Exner pressure at surface max_iterations = 10 # maximum number of hydrostatic balance iterations tolerance = 1e-7 # tolerance for hydrostatic balance iteration @@ -62,8 +67,10 @@ def mountain_nonhydrostatic( # ------------------------------------------------------------------------ # # Our settings for this set up # ------------------------------------------------------------------------ # + element_order = 1 u_eqn_type = 'vector_invariant_form' + alpha = 0.5 # ------------------------------------------------------------------------ # # Set up model objects @@ -94,19 +101,29 @@ def mountain_nonhydrostatic( # Equation parameters = CompressibleParameters(g=g, cp=cp) sponge = SpongeLayerParameters( - H=domain_height, z_level=domain_height-sponge_depth, mubar=sponge_mu/dt - ) - eqns = CompressibleEulerEquations( - domain, parameters, sponge_options=sponge, u_transport_option=u_eqn_type + H=domain_height, z_level=domain_height-sponge_depth, mubar=mu_dt/dt ) + if hydrostatic: + eqns = HydrostaticCompressibleEulerEquations( + domain, parameters, sponge_options=sponge, + u_transport_option=u_eqn_type + ) + else: + eqns = CompressibleEulerEquations( + domain, parameters, sponge_options=sponge, + u_transport_option=u_eqn_type + ) # I/O + # Adjust default directory name + if hydrostatic and dirname == mountain_nonhydrostatic_defaults['dirname']: + dirname = f'hyd_switch_{dirname}' + output = OutputParameters( dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True ) diagnostic_fields = [ - Exner(parameters), ZComponent('u'), Perturbation('theta'), - Perturbation('rho') + Exner(parameters), ZComponent('u'), Perturbation('theta') ] io = IO(domain, output, diagnostic_fields=diagnostic_fields) @@ -124,12 +141,18 @@ def mountain_nonhydrostatic( ] # Linear solver - linear_solver = CompressibleSolver(eqns) + if hydrostatic: + linear_solver = CompressibleSolver( + eqns, alpha, solver_parameters=hydrostatic_parameters, + overwrite_solver_parameters=True + ) + else: + linear_solver = CompressibleSolver(eqns, alpha) # Time stepper stepper = SemiImplicitQuasiNewton( eqns, io, transported_fields, transport_methods, - linear_solver=linear_solver + linear_solver=linear_solver, alpha=alpha ) # ------------------------------------------------------------------------ # @@ -277,6 +300,16 @@ def mountain_nonhydrostatic( type=str, default=mountain_nonhydrostatic_defaults['dirname'] ) + parser.add_argument( + '--hydrostatic', + help=( + "Whether to use the hydrostatic switch to emulate the " + + "hydrostatic equations. Otherwise use the full non-hydrostatic" + + "equations." + ), + action="store_true", + default=mountain_nonhydrostatic_defaults['hydrostatic'] + ) args, unknown = parser.parse_known_args() mountain_nonhydrostatic(**vars(args)) diff --git a/case_studies/compressible_euler/skamarock_klemp_hydrostatic.py b/case_studies/compressible_euler/skamarock_klemp_hydrostatic.py new file mode 100644 index 0000000..9a8706e --- /dev/null +++ b/case_studies/compressible_euler/skamarock_klemp_hydrostatic.py @@ -0,0 +1,236 @@ +""" +This example uses the compressible Euler equations to solve the hydrostatic +vertical slice gravity wave test case of Skamarock and Klemp, 1994: +``Efficiency and Accuracy of the Klemp-Wilhelmson Time-Splitting Technique'', +MWR. + +The domain is larger than the "non-hydrostatic" gravity wave test, so that the +hydrostatic and non-hydrostatic solutions are similar. The domain is one layer +thick in the y-direction, and the Coriolis force is included, which is balanced +by an additional prescribed pressure gradient. + +Potential temperature is transported using SUPG, and the degree 1 elements are +used. +""" + +from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter +from firedrake import ( + as_vector, SpatialCoordinate, PeriodicRectangleMesh, ExtrudedMesh, exp, sin, + Function, pi +) +from gusto import ( + Domain, IO, OutputParameters, SemiImplicitQuasiNewton, SSPRK3, DGUpwind, + TrapeziumRule, SUPGOptions, Perturbation, CompressibleParameters, + CompressibleEulerEquations, HydrostaticCompressibleEulerEquations, + CompressibleSolver, compressible_hydrostatic_balance, hydrostatic_parameters +) + +skamarock_klemp_hydrostatic_defaults = { + 'ncolumns': 80, + 'nlayers': 10, + 'dt': 250.0, + 'tmax': 60000., + 'dumpfreq': 120, + 'dirname': 'skamarock_klemp_hydrostatic', + 'hydrostatic': False +} + + +def skamarock_klemp_hydrostatic( + ncolumns=skamarock_klemp_hydrostatic_defaults['ncolumns'], + nlayers=skamarock_klemp_hydrostatic_defaults['nlayers'], + dt=skamarock_klemp_hydrostatic_defaults['dt'], + tmax=skamarock_klemp_hydrostatic_defaults['tmax'], + dumpfreq=skamarock_klemp_hydrostatic_defaults['dumpfreq'], + dirname=skamarock_klemp_hydrostatic_defaults['dirname'], + hydrostatic=skamarock_klemp_hydrostatic_defaults['hydrostatic'] +): + + # ------------------------------------------------------------------------ # + # Test case parameters + # ------------------------------------------------------------------------ # + + domain_width = 6.0e6 # Width of domain in x direction (m) + domain_length = 1.0e4 # Length of domain in y direction (m) + domain_height = 1.0e4 # Height of domain (m) + Tsurf = 300. # Temperature at surface (K) + wind_initial = 20. # Initial wind in x direction (m/s) + pert_width = 1.0e5 # Width parameter of perturbation (m) + deltaTheta = 1.0e-2 # Magnitude of theta perturbation (K) + N = 0.01 # Brunt-Vaisala frequency (1/s) + Omega = 0.5e-4 # Planetary rotation rate (1/s) + pressure_gradient_y = -1.0e-4*20 # Prescribed force in y direction (m/s^2) + + # ------------------------------------------------------------------------ # + # Our settings for this set up + # ------------------------------------------------------------------------ # + + element_order = 1 + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain -- 3D volume mesh + base_mesh = PeriodicRectangleMesh( + ncolumns, 1, domain_width, domain_length, quadrilateral=True + ) + mesh = ExtrudedMesh(base_mesh, nlayers, layer_height=domain_height/nlayers) + domain = Domain(mesh, dt, "RTCF", element_order) + + # Equation + parameters = CompressibleParameters(Omega=Omega) + balanced_pg = as_vector((0., pressure_gradient_y, 0.)) + if hydrostatic: + eqns = HydrostaticCompressibleEulerEquations( + domain, parameters, extra_terms=[("u", balanced_pg)] + ) + else: + eqns = CompressibleEulerEquations( + domain, parameters, extra_terms=[("u", balanced_pg)] + ) + + # I/O + # Adjust default directory name + if hydrostatic and dirname == skamarock_klemp_hydrostatic_defaults['dirname']: + dirname = f'hyd_switch_{dirname}' + + output = OutputParameters( + dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True, + dumplist=['u', 'theta'], + ) + diagnostic_fields = [Perturbation('theta')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + theta_opts = SUPGOptions() + transported_fields = [ + TrapeziumRule(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts) + ] + transport_methods = [ + DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp) + ] + + # Linear solver + if hydrostatic: + linear_solver = CompressibleSolver( + eqns, solver_parameters=hydrostatic_parameters, + overwrite_solver_parameters=True + ) + else: + linear_solver = CompressibleSolver(eqns) + + # Time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transported_fields, transport_methods, + linear_solver=linear_solver + ) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") + + # Thermodynamic constants required for setting initial conditions + # and reference profiles + g = parameters.g + + x, _, z = SpatialCoordinate(mesh) + + # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) + thetab = Tsurf*exp(N**2*z/g) + + theta_b = Function(Vt).interpolate(thetab) + rho_b = Function(Vr) + + theta_pert = ( + deltaTheta * sin(pi*z/domain_height) + / (1 + (x - domain_width/2)**2 / pert_width**2) + ) + theta0.interpolate(theta_b + theta_pert) + + compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) + + rho0.assign(rho_b) + u0.project(as_vector([wind_initial, 0.0, 0.0])) + + stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) + + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(t=0, tmax=tmax) + +# ---------------------------------------------------------------------------- # +# MAIN +# ---------------------------------------------------------------------------- # + + +if __name__ == "__main__": + + parser = ArgumentParser( + description=__doc__, + formatter_class=ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + '--ncolumns', + help="The number of columns in the vertical slice mesh.", + type=int, + default=skamarock_klemp_hydrostatic_defaults['ncolumns'] + ) + parser.add_argument( + '--nlayers', + help="The number of layers for the mesh.", + type=int, + default=skamarock_klemp_hydrostatic_defaults['nlayers'] + ) + parser.add_argument( + '--dt', + help="The time step in seconds.", + type=float, + default=skamarock_klemp_hydrostatic_defaults['dt'] + ) + parser.add_argument( + "--tmax", + help="The end time for the simulation in seconds.", + type=float, + default=skamarock_klemp_hydrostatic_defaults['tmax'] + ) + parser.add_argument( + '--dumpfreq', + help="The frequency at which to dump field output.", + type=int, + default=skamarock_klemp_hydrostatic_defaults['dumpfreq'] + ) + parser.add_argument( + '--dirname', + help="The name of the directory to write to.", + type=str, + default=skamarock_klemp_hydrostatic_defaults['dirname'] + ) + parser.add_argument( + '--hydrostatic', + help=( + "Whether to use the hydrostatic switch to emulate the " + + "hydrostatic equations. Otherwise use the full non-hydrostatic" + + "equations." + ), + action="store_true", + default=skamarock_klemp_hydrostatic_defaults['hydrostatic'] + ) + args, unknown = parser.parse_known_args() + + skamarock_klemp_hydrostatic(**vars(args)) diff --git a/case_studies/compressible_euler/test_compressible_euler.py b/case_studies/compressible_euler/test_compressible_euler.py index 2f54384..47a2318 100644 --- a/case_studies/compressible_euler/test_compressible_euler.py +++ b/case_studies/compressible_euler/test_compressible_euler.py @@ -1,8 +1,12 @@ from dry_baroclinic_sphere import dry_baroclinic_sphere from moist_baroclinic_channel import moist_baroclinic_channel from moist_bryan_fritsch import moist_bryan_fritsch +from moist_skamarock_klemp import moist_skamarock_klemp +from mountain_hydrostatic import mountain_hydrostatic from mountain_nonhydrostatic import mountain_nonhydrostatic from robert_bubble import robert_bubble +from skamarock_klemp_hydrostatic import skamarock_klemp_hydrostatic +import pytest def test_dry_baroclinic_sphere(): @@ -39,6 +43,43 @@ def test_moist_bryan_fritsch(): ) +def test_moist_skamarock_klemp(): + moist_skamarock_klemp( + ncolumns=30, + nlayers=5, + dt=6.0, + tmax=60.0, + dumpfreq=10, + dirname='pytest_moist_skamarock_klemp' + ) + + +def test_mountain_hydrostatic(): + mountain_hydrostatic( + ncolumns=20, + nlayers=10, + dt=5.0, + tmax=50.0, + dumpfreq=10, + dirname='pytest_mountain_hydrostatic', + hydrostatic=False + ) + + +# Hydrostatic switch not currently working +@pytest.mark.xfail +def test_hyd_switch_mountain_hydrostatic(): + mountain_hydrostatic( + ncolumns=20, + nlayers=10, + dt=5.0, + tmax=50.0, + dumpfreq=10, + dirname='pytest_hyd_switch_mountain_hydrostatic', + hydrostatic=True + ) + + def test_mountain_nonhydrostatic(): mountain_nonhydrostatic( ncolumns=20, @@ -46,7 +87,22 @@ def test_mountain_nonhydrostatic(): dt=5.0, tmax=10.0, dumpfreq=2, - dirname='pytest_mountain_nonhydrostatic' + dirname='pytest_mountain_nonhydrostatic', + hydrostatic=False + ) + + +# Hydrostatic switch not currently working +@pytest.mark.xfail +def test_hyd_switch_mountain_nonhydrostatic(): + mountain_nonhydrostatic( + ncolumns=20, + nlayers=10, + dt=5.0, + tmax=10.0, + dumpfreq=2, + dirname='pytest_hyd_switch_mountain_nonhydrostatic', + hydrostatic=True ) @@ -59,3 +115,29 @@ def test_robert_bubble(): dumpfreq=2, dirname='pytest_robert_bubble' ) + + +def test_skamarock_klemp_hydrostatic(): + skamarock_klemp_hydrostatic( + ncolumns=20, + nlayers=4, + dt=6.0, + tmax=60.0, + dumpfreq=10, + dirname='pytest_skamarock_klemp_hydrostatic', + hydrostatic=False + ) + + +# Hydrostatic equations not currently working +@pytest.mark.xfail +def test_hyd_switch_skamarock_klemp_hydrostatic(): + skamarock_klemp_hydrostatic( + ncolumns=20, + nlayers=4, + dt=6.0, + tmax=60.0, + dumpfreq=10, + dirname='pytest_hyd_switch_skamarock_klemp_hydrostatic', + hydrostatic=True + ) diff --git a/figures/compressible_euler/moist_skamarock_klemp_final.png b/figures/compressible_euler/moist_skamarock_klemp_final.png new file mode 100644 index 0000000..a2e40f3 Binary files /dev/null and b/figures/compressible_euler/moist_skamarock_klemp_final.png differ diff --git a/figures/compressible_euler/moist_skamarock_klemp_initial.png b/figures/compressible_euler/moist_skamarock_klemp_initial.png new file mode 100644 index 0000000..7efb6ab Binary files /dev/null and b/figures/compressible_euler/moist_skamarock_klemp_initial.png differ diff --git a/figures/compressible_euler/mountain_hydrostatic_final.png b/figures/compressible_euler/mountain_hydrostatic_final.png new file mode 100644 index 0000000..1cdc263 Binary files /dev/null and b/figures/compressible_euler/mountain_hydrostatic_final.png differ diff --git a/figures/compressible_euler/mountain_hydrostatic_initial.png b/figures/compressible_euler/mountain_hydrostatic_initial.png new file mode 100644 index 0000000..ed6574e Binary files /dev/null and b/figures/compressible_euler/mountain_hydrostatic_initial.png differ diff --git a/figures/compressible_euler/mountain_nonhydrostatic_final.png b/figures/compressible_euler/mountain_nonhydrostatic_final.png index bc8f912..7734b76 100644 Binary files a/figures/compressible_euler/mountain_nonhydrostatic_final.png and b/figures/compressible_euler/mountain_nonhydrostatic_final.png differ diff --git a/figures/compressible_euler/mountain_nonhydrostatic_initial.png b/figures/compressible_euler/mountain_nonhydrostatic_initial.png index 9e2f94c..60773a6 100644 Binary files a/figures/compressible_euler/mountain_nonhydrostatic_initial.png and b/figures/compressible_euler/mountain_nonhydrostatic_initial.png differ diff --git a/figures/compressible_euler/skamarock_klemp_hydrostatic_final.png b/figures/compressible_euler/skamarock_klemp_hydrostatic_final.png new file mode 100644 index 0000000..4fd5908 Binary files /dev/null and b/figures/compressible_euler/skamarock_klemp_hydrostatic_final.png differ diff --git a/figures/compressible_euler/skamarock_klemp_hydrostatic_initial.png b/figures/compressible_euler/skamarock_klemp_hydrostatic_initial.png new file mode 100644 index 0000000..53dc937 Binary files /dev/null and b/figures/compressible_euler/skamarock_klemp_hydrostatic_initial.png differ diff --git a/plotting/compressible_euler/plot_moist_skamarock_klemp.py b/plotting/compressible_euler/plot_moist_skamarock_klemp.py new file mode 100644 index 0000000..af7fde3 --- /dev/null +++ b/plotting/compressible_euler/plot_moist_skamarock_klemp.py @@ -0,0 +1,193 @@ +""" +Plots the moist Skamarock-Klemp gravity wave in a vertical slice. + +This plots the initial conditions @ t = 0 s, with +(a) theta_e perturbation, (b) theta +and the final state @ t = 3000 s, with +(a) theta_e perturbation, +(b) a 1D slice through the wave +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +import pandas as pd +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, tomplot_field_title, extract_gusto_coords, + extract_gusto_field, reshape_gusto_data, add_colorbar_fig +) + +test = 'moist_skamarock_klemp' + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file +results_file_name = f'{abspath(dirname(__file__))}/../../../results/{test}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/compressible_euler/{test}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['Theta_e_perturbation', 'Theta_e'] +init_colour_schemes = ['YlOrRd', 'Purples'] +init_field_labels = [r'$\Delta\theta$ (K)', r'$\theta$ (K)'] +init_contours = [np.linspace(0.0, 0.01, 11), np.linspace(300, 335, 8)] +init_contours_to_remove = [None, None] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_name = 'Theta_e_perturbation' +final_colour_scheme = 'RdBu_r' +final_field_label = r'$\Delta\theta$ (K)' +final_contours = np.linspace(-3.0e-3, 3.0e-3, 13) +final_contour_to_remove = 0.0 + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +ylims = [0, 10.0] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +xlims = [0, 300.0] + +fig, axarray = plt.subplots(1, 2, figsize=(12, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, init_contours_to_remove)): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.2) +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +x_offset = -3000.0*20/1000.0 +xlims = [-x_offset, 300.0-x_offset] + +fig, axarray = plt.subplots(2, 1, figsize=(8, 8), sharex='all') +time_idx = -1 + +# Data extraction ---------------------------------------------------------- +field_data = extract_gusto_field(data_file, final_field_name, time_idx=time_idx) +coords_X, coords_Y = extract_gusto_coords(data_file, final_field_name) +time = data_file['time'][time_idx] + +# Wave has wrapped around periodic boundary, so shift the coordinates +coords_X = np.where(coords_X < xlims[0], coords_X + 300.0, coords_X) + +# Sort data given the change in coordinates +data_dict = { + 'X': coords_X, + 'Y': coords_Y, + 'field': field_data +} +data_frame = pd.DataFrame(data_dict) +data_frame.sort_values(by=['X', 'Y'], inplace=True) +coords_X = data_frame['X'].values[:] +coords_Y = data_frame['Y'].values[:] +field_data = data_frame['field'].values[:] + +# Plot 2D data ----------------------------------------------------------------- +ax = axarray[0] + +cmap, lines = tomplot_cmap( + final_contours, final_colour_scheme, remove_contour=final_contour_to_remove +) +cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, final_contours, + cmap=cmap, line_contours=lines +) + +add_colorbar_fig( + fig, cf, final_field_label, ax_idxs=[0], location='right', cbar_labelpad=-40 +) +tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data +) + +ax.set_ylabel(r'$z$ (km)', labelpad=-20) +ax.set_ylim(ylims) +ax.set_yticks(ylims) +ax.set_yticklabels(ylims) + +# Plot 1D data ----------------------------------------------------------------- +ax = axarray[1] + +field_data, coords_X, coords_Y = reshape_gusto_data(field_data, coords_X, coords_Y) + +# Determine midpoint index +mid_idx = np.floor_divide(np.shape(field_data)[1], 2) +slice_height = coords_Y[0, mid_idx] + +ax.plot(coords_X[:, mid_idx], field_data[:, mid_idx], color='black') + +tomplot_field_title( + ax, r'$z$' + f' = {slice_height} km' +) + +theta_lims = [np.min(final_contours), np.max(final_contours)] + +ax.set_ylabel(final_field_label, labelpad=-20) +ax.set_ylim(theta_lims) +ax.set_yticks(theta_lims) +ax.set_yticklabels(theta_lims) + +ax.set_xlabel(r'$x$ (km)', labelpad=-10) +ax.set_xlim(xlims) +ax.set_xticks(xlims) +ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(hspace=0.2) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/compressible_euler/plot_mountain_hydrostatic.py b/plotting/compressible_euler/plot_mountain_hydrostatic.py new file mode 100644 index 0000000..00156fe --- /dev/null +++ b/plotting/compressible_euler/plot_mountain_hydrostatic.py @@ -0,0 +1,202 @@ +""" +Plots the hydrostatic mountain test case. + +This plots: +(a) w @ t = 15000 s, (b) theta @ t = 15000 s +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +from netCDF4 import Dataset +import numpy as np +import pandas as pd +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, tomplot_field_title, tomplot_contours, + extract_gusto_coords, extract_gusto_field, reshape_gusto_data +) + +test = 'mountain_hydrostatic' + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file +results_file_name = f'{abspath(dirname(__file__))}/../../../results/{test}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/compressible_euler/{test}' + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['u_z', 'theta_perturbation', 'u_z', 'theta_perturbation'] +final_colour_schemes = ['PiYG', 'RdBu_r', 'PiYG', 'RdBu_r'] +final_field_labels = [ + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)', + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)' +] +final_contours = [ + np.linspace(-0.012, 0.012, 13), np.linspace(-0.14, 0.14, 15), + np.linspace(-0.004, 0.004, 21), np.linspace(-0.02, 0.02, 21) +] + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +initial_field_names = ['Exner', 'theta', 'Temperature'] +initial_colour_schemes = ['PuBu', 'Reds', 'RdYlBu_r'] +initial_field_labels = [r'$\Pi$', r'$\theta$ (K)', r'$T$ (K)'] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'contour' # Need to use this method to show mountains! +xlims = [0., 240.] +ylims = [0., 50.] + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # + +fig, axarray = plt.subplots(1, 3, figsize=(18, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, colour_scheme, field_label) in \ + enumerate(zip( + axarray.flatten(), initial_field_names, initial_colour_schemes, + initial_field_labels + )): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] + + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + + if field_name == 'Temperature': + contours = np.linspace(250.0, 266.0, 9) + else: + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme) + + # Plot data ---------------------------------------------------------------- + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom') + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.suptitle(f't = {time:.1f} s') +fig.subplots_adjust(wspace=0.25) +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +xlims_zoom = [80., 160.] +ylims_zoom = [0., 12.] + +fig, axarray = plt.subplots(2, 2, figsize=(18, 12), sharex='row', sharey='row') +time_idx = -1 + +for i, (ax, field_name, colour_scheme, field_label, contours) in \ + enumerate(zip( + axarray.flatten(), final_field_names, final_colour_schemes, + final_field_labels, final_contours + )): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X, coords_Y = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] + + # Filter data for panels that are zoomed in mountain region + if i in [2, 3]: + data_dict = { + 'X': coords_X, + 'Y': coords_Y, + 'field': field_data + } + data_frame = pd.DataFrame(data_dict) + + data_frame = data_frame[ + (data_frame['X'] >= xlims_zoom[0]) + & (data_frame['X'] <= xlims_zoom[1]) + & (data_frame['Y'] >= ylims_zoom[0]) + & (data_frame['Y'] <= ylims_zoom[1]) + ] + field_data = data_frame['field'].values[:] + coords_X = data_frame['X'].values[:] + coords_Y = data_frame['Y'].values[:] + + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=0.0) + + # Plot data ---------------------------------------------------------------- + cf, _ = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom') + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, minmax_format='.3f' + ) + + # Labels ------------------------------------------------------------------- + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + + if i in [0, 1]: + ax.set_xlim(xlims) + ax.set_ylim(ylims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + else: + ax.set_xlim(xlims_zoom) + ax.set_ylim(ylims_zoom) + ax.set_xticks(xlims_zoom) + ax.set_xticklabels(xlims_zoom) + + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + elif i == 2: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_yticks(ylims_zoom) + ax.set_yticklabels(ylims_zoom) + + +# Save figure ------------------------------------------------------------------ +fig.suptitle(f't = {time:.1f} s') +fig.subplots_adjust(wspace=0.25) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/compressible_euler/plot_mountain_nonhydrostatic.py b/plotting/compressible_euler/plot_mountain_nonhydrostatic.py index 2f79ddd..4da4276 100644 --- a/plotting/compressible_euler/plot_mountain_nonhydrostatic.py +++ b/plotting/compressible_euler/plot_mountain_nonhydrostatic.py @@ -7,10 +7,12 @@ from os.path import abspath, dirname import matplotlib.pyplot as plt from netCDF4 import Dataset +import numpy as np +import pandas as pd from tomplot import ( set_tomplot_style, tomplot_cmap, plot_contoured_field, add_colorbar_ax, tomplot_field_title, tomplot_contours, - extract_gusto_coords, extract_gusto_field, + extract_gusto_coords, extract_gusto_field, reshape_gusto_data ) test = 'mountain_nonhydrostatic' @@ -20,15 +22,22 @@ # ---------------------------------------------------------------------------- # # When copying this example these paths need editing, which will usually involve # removing the abspath part to set directory paths relative to this file -results_file_name = f'{abspath(dirname(__file__))}/../../results/{test}/field_output.nc' +results_file_name = f'{abspath(dirname(__file__))}/../../../results/{test}/field_output.nc' plot_stem = f'{abspath(dirname(__file__))}/../../figures/compressible_euler/{test}' # ---------------------------------------------------------------------------- # # Final plot details # ---------------------------------------------------------------------------- # -final_field_names = ['u_z', 'theta_perturbation'] -final_colour_schemes = ['PiYG', 'RdBu_r'] -final_field_labels = [r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)'] +final_field_names = ['u_z', 'theta_perturbation', 'u_z', 'theta_perturbation'] +final_colour_schemes = ['PiYG', 'RdBu_r', 'PiYG', 'RdBu_r'] +final_field_labels = [ + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)', + r'$w$ (m s$^{-1}$)', r'$\Delta\theta$ (K)' +] +final_contours = [ + np.linspace(-0.04, 0.04, 21), np.linspace(-0.025, 0.025, 21), + np.linspace(-0.04, 0.04, 21), np.linspace(-0.025, 0.025, 21) +] # ---------------------------------------------------------------------------- # # Initial plot details @@ -40,7 +49,7 @@ # ---------------------------------------------------------------------------- # # General options # ---------------------------------------------------------------------------- # -contour_method = 'tricontour' +contour_method = 'contour' # Need to use this method to show mountains! xlims = [0., 144.] ylims = [0., 35.] @@ -51,6 +60,7 @@ # ---------------------------------------------------------------------------- # # INITIAL PLOTTING # ---------------------------------------------------------------------------- # + fig, axarray = plt.subplots(1, 2, figsize=(18, 6), sharex='all', sharey='all') time_idx = 0 @@ -65,6 +75,9 @@ coords_X, coords_Y = extract_gusto_coords(data_file, field_name) time = data_file['time'][time_idx] + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + contours = tomplot_contours(field_data) cmap, lines = tomplot_cmap(contours, colour_scheme) @@ -100,13 +113,16 @@ # ---------------------------------------------------------------------------- # # FINAL PLOTTING # ---------------------------------------------------------------------------- # -fig, axarray = plt.subplots(1, 2, figsize=(18, 6), sharex='all', sharey='all') +xlims_zoom = [60., 95.] +ylims_zoom = [0., 12.] + +fig, axarray = plt.subplots(2, 2, figsize=(18, 12), sharex='row', sharey='row') time_idx = -1 -for i, (ax, field_name, colour_scheme, field_label) in \ +for i, (ax, field_name, colour_scheme, field_label, contours) in \ enumerate(zip( axarray.flatten(), final_field_names, final_colour_schemes, - final_field_labels + final_field_labels, final_contours )): # Data extraction ---------------------------------------------------------- @@ -114,7 +130,28 @@ coords_X, coords_Y = extract_gusto_coords(data_file, field_name) time = data_file['time'][time_idx] - contours = tomplot_contours(field_data) + # Filter data for panels that are zoomed in mountain region + if i in [2, 3]: + data_dict = { + 'X': coords_X, + 'Y': coords_Y, + 'field': field_data + } + data_frame = pd.DataFrame(data_dict) + + data_frame = data_frame[ + (data_frame['X'] >= xlims_zoom[0]) + & (data_frame['X'] <= xlims_zoom[1]) + & (data_frame['Y'] >= ylims_zoom[0]) + & (data_frame['Y'] <= ylims_zoom[1]) + ] + field_data = data_frame['field'].values[:] + coords_X = data_frame['X'].values[:] + coords_Y = data_frame['Y'].values[:] + + field_data, coords_X, coords_Y = \ + reshape_gusto_data(field_data, coords_X, coords_Y) + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=0.0) # Plot data ---------------------------------------------------------------- @@ -129,16 +166,29 @@ ) # Labels ------------------------------------------------------------------- + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + + if i in [0, 1]: + ax.set_xlim(xlims) + ax.set_ylim(ylims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + else: + ax.set_xlim(xlims_zoom) + ax.set_ylim(ylims_zoom) + ax.set_xticks(xlims_zoom) + ax.set_xticklabels(xlims_zoom) + if i == 0: ax.set_ylabel(r'$z$ (km)', labelpad=-20) - ax.set_ylim(ylims) ax.set_yticks(ylims) ax.set_yticklabels(ylims) - ax.set_xlabel(r'$x$ (km)', labelpad=-10) - ax.set_xlim(xlims) - ax.set_xticks(xlims) - ax.set_xticklabels(xlims) + elif i == 2: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_yticks(ylims_zoom) + ax.set_yticklabels(ylims_zoom) + # Save figure ------------------------------------------------------------------ fig.suptitle(f't = {time:.1f} s') diff --git a/plotting/compressible_euler/plot_skamarock_klemp_hydrostatic.py b/plotting/compressible_euler/plot_skamarock_klemp_hydrostatic.py new file mode 100644 index 0000000..10add8d --- /dev/null +++ b/plotting/compressible_euler/plot_skamarock_klemp_hydrostatic.py @@ -0,0 +1,200 @@ +""" +Plots the hydrostatic Skamarock-Klemp gravity wave in a vertical slice. + +This plots the initial conditions @ t = 0 s, with +(a) theta perturbation, (b) theta +and the final state @ t = 60,000 s, with +(a) theta perturbation, +(b) a 1D slice through the wave +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +import pandas as pd +from tomplot import ( + set_tomplot_style, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, tomplot_field_title, extract_gusto_vertical_slice, + add_colorbar_fig, reshape_gusto_data +) + +test = 'skamarock_klemp_hydrostatic' + +# ---------------------------------------------------------------------------- # +# Directory for results and plots +# ---------------------------------------------------------------------------- # +# When copying this example these paths need editing, which will usually involve +# removing the abspath part to set directory paths relative to this file +results_file_name = f'{abspath(dirname(__file__))}/../../../results/{test}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/compressible_euler/{test}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['theta_perturbation', 'theta'] +init_colour_schemes = ['YlOrRd', 'Purples'] +init_field_labels = [r'$\Delta\theta$ (K)', r'$\theta$ (K)'] +init_contours = [np.linspace(0.0, 0.01, 11), np.linspace(300, 335, 8)] +init_contours_to_remove = [None, None] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_name = 'theta_perturbation' +final_colour_scheme = 'RdBu_r' +final_field_label = r'$\Delta\theta$ (K)' +final_contours = np.linspace(-3.5e-3, 3.5e-3, 15) +final_contour_to_remove = 0.0 + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +ylims = [0, 10.0] +slice_along = 'y' +slice_at = 2.113249 + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +xlims = [0, 6000.0] +fig, axarray = plt.subplots(1, 2, figsize=(12, 6), sharex='all', sharey='all') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, to_remove) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, init_contours_to_remove)): + + # Data extraction ---------------------------------------------------------- + field_data, coords_X, _, coords_Z = \ + extract_gusto_vertical_slice( + data_file, field_name, time_idx=time_idx, + slice_along=slice_along, slice_at=slice_at) + + time = data_file['time'][time_idx] + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Z, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10 + ) + tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(ylims) + ax.set_yticks(ylims) + ax.set_yticklabels(ylims) + + ax.set_xlabel(r'$x$ (km)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.2) +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +x_offset = -60000.0*20/1000.0 +xlims = [-x_offset, 6000.0-x_offset] + +fig, axarray = plt.subplots(2, 1, figsize=(8, 8), sharex='all') +time_idx = -1 + +# Data extraction ---------------------------------------------------------- +field_data, coords_X, _, coords_Z = \ + extract_gusto_vertical_slice( + data_file, final_field_name, time_idx=time_idx, + slice_along=slice_along, slice_at=slice_at + ) +time = data_file['time'][time_idx] + +# Wave has wrapped around periodic boundary, so shift the coordinates +coords_X = np.where(coords_X < xlims[0], coords_X + 6000.0, coords_X) + +# Sort data given the change in coordinates +data_dict = { + 'X': coords_X.flatten(), + 'Z': coords_Z.flatten(), + 'field': field_data.flatten() +} +data_frame = pd.DataFrame(data_dict) +data_frame.sort_values(by=['X', 'Z'], inplace=True) +coords_X = data_frame['X'].values[:] +coords_Z = data_frame['Z'].values[:] +field_data = data_frame['field'].values[:] + +# Plot 2D data ----------------------------------------------------------------- +ax = axarray[0] + +cmap, lines = tomplot_cmap( + final_contours, final_colour_scheme, remove_contour=final_contour_to_remove +) +cf, lines = plot_contoured_field( + ax, coords_X, coords_Z, field_data, contour_method, final_contours, + cmap=cmap, line_contours=lines +) + +add_colorbar_fig( + fig, cf, final_field_label, ax_idxs=[0], location='right', cbar_labelpad=-40 +) +tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data +) + +ax.set_ylabel(r'$z$ (km)', labelpad=-20) +ax.set_ylim(ylims) +ax.set_yticks(ylims) +ax.set_yticklabels(ylims) + +# Plot 1D data ----------------------------------------------------------------- +ax = axarray[1] + +field_data, coords_X, coords_Y = reshape_gusto_data(field_data, coords_X, coords_Z) + +# Determine midpoint index +mid_idx = np.floor_divide(np.shape(field_data)[1], 2) +slice_height = coords_Y[0, mid_idx] + +ax.plot(coords_X[:, mid_idx], field_data[:, mid_idx], color='black') + +tomplot_field_title( + ax, r'$z$' + f' = {slice_height:.1f} km' +) + +theta_lims = [np.min(final_contours), np.max(final_contours)] + +ax.set_ylabel(final_field_label, labelpad=-20) +ax.set_ylim(theta_lims) +ax.set_yticks(theta_lims) +ax.set_yticklabels(theta_lims) + +ax.set_xlabel(r'$x$ (km)', labelpad=-10) +ax.set_xlim(xlims) +ax.set_xticks(xlims) +ax.set_xticklabels(xlims) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(hspace=0.2) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close()