<commit_before>from gusto import *
from firedrake import CubedSphereMesh, ExtrudedMesh, Expression, \
    VectorFunctionSpace, FunctionSpace, Function, SpatialCoordinate, \
    as_vector, DumbCheckpoint, interpolate, CellVolume, sqrt
from firedrake import exp, acos, cos, sin, parameters
from firedrake.petsc import PETSc
from argparse import ArgumentParser
import numpy as np
import sys


parameters["pyop2_options"]["lazy_evaluation"] = False


PETSc.Log.begin()
parser = ArgumentParser(description=("""
DCMIP Test 3-1 Non-orographic gravity waves on a small planet.
"""), add_help=False)

parser.add_argument("--hybridization",
                    action="store_true",
                    help="Use a hybridized compressible solver.")

parser.add_argument("--test",
                    action="store_true",
                    help="Enable a quick test run.")

parser.add_argument("--profile",
                    action="store_true",
                    help="Turn on profiling.")

parser.add_argument("--dumpfreq",
                    default=100,
                    type=int,
                    action="store",
                    help="Dump frequency of output files.")

parser.add_argument("--cfl",
                    default=1.,
                    type=float,
                    action="store",
                    help="CFL number to run at (determines dt).")

parser.add_argument("--tmax",
                    default=100.,
                    type=float,
                    action="store",
                    help="Max time (s). Max test time was set to 3600s.")

parser.add_argument("--refinements",
                    default=4,
                    type=int,
                    action="store",
                    help="Resolution scaling parameter.")

parser.add_argument("--layers",
                    default=16,
                    type=int,
                    action="store",
                    help="Number of vertical layers.")

parser.add_argument("--debug",
                    action="store_true",
                    help="Turn on KSP monitors")

parser.add_argument("--rtol",
                    default=1.0e-6,
                    type=float,
                    help="Rtolerance for the linear solve.")

parser.add_argument("--help",
                    action="store_true",
                    help="Show help.")


args, _ = parser.parse_known_args()

if args.help:
    help = parser.format_help()
    PETSc.Sys.Print("%s\n" % help)
    sys.exit(1)

tmax = args.tmax                # Maximum time (s)
nlayers = args.layers           # Number of vertical layers
refinements = args.refinements  # Number of horiz. cells = 20*(4^refinements)

hybrid = bool(args.hybridization)
PETSc.Sys.Print("""
Problem parameters:\n
Test case DCMIP 3-1: Non-orographic gravity waves on a small planet.\n
Hybridized compressible solver: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Profiling: %s,\n
Max time: %s,\n
Dump frequency: %s.\n
""" % (hybrid, refinements, nlayers,
       bool(args.profile), args.tmax, args.dumpfreq))

# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6               # Radius of the Earth (m)
X = 125.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 constant 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                     # Latitudinal centerpoint of Theta' (equator)
deltaTheta = 1.0                # Maximum amplitude of Theta' (K)
L_z = 20000.0                   # Vertical wave length of the Theta' perturb.

# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
                    refinement_level=refinements,
                    degree=2)

# Horizontal Courant (advective) number
cell_vs = interpolate(CellVolume(m),
                      FunctionSpace(m, "DG", 0))
a_min = cell_vs.dat.data.min()
a_max = cell_vs.dat.data.max()
dx_min = sqrt(a_min)
dx_max = sqrt(a_max)
dx_avg = (dx_min + dx_max)/2.0
u_max = u_0

dt = args.cfl * (dx_avg / 343.0)

if args.profile:
    tmax = 5*dt

if args.test:
    tmax = dt

PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" % (dt,
                                                                      tmax))

# Build volume mesh
z_top = 1.0e4            # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers,
                    layer_height=z_top/nlayers,
                    extrusion_type="radial")

# Space for initialising velocity (using this ensures things are in layers)
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)

# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
z_expr = Expression("sqrt(x[0]*x[0] + x[1]*x[1] + x[2]*x[2]) - a", a=a)
z = Function(W_CG1).interpolate(z_expr)<commit_msg>Take integer value of dt<commit_after>from gusto import *
from firedrake import CubedSphereMesh, ExtrudedMesh, Expression, \
    VectorFunctionSpace, FunctionSpace, Function, SpatialCoordinate, \
    as_vector, DumbCheckpoint, interpolate, CellVolume, sqrt
from firedrake import exp, acos, cos, sin, parameters
from firedrake.petsc import PETSc
from argparse import ArgumentParser
import numpy as np
import sys


parameters["pyop2_options"]["lazy_evaluation"] = False


PETSc.Log.begin()
parser = ArgumentParser(description=("""
DCMIP Test 3-1 Non-orographic gravity waves on a small planet.
"""), add_help=False)

parser.add_argument("--hybridization",
                    action="store_true",
                    help="Use a hybridized compressible solver.")

parser.add_argument("--test",
                    action="store_true",
                    help="Enable a quick test run.")

parser.add_argument("--profile",
                    action="store_true",
                    help="Turn on profiling.")

parser.add_argument("--dumpfreq",
                    default=100,
                    type=int,
                    action="store",
                    help="Dump frequency of output files.")

parser.add_argument("--cfl",
                    default=1.,
                    type=float,
                    action="store",
                    help="CFL number to run at (determines dt).")

parser.add_argument("--tmax",
                    default=100.,
                    type=float,
                    action="store",
                    help="Max time (s). Max test time was set to 3600s.")

parser.add_argument("--refinements",
                    default=4,
                    type=int,
                    action="store",
                    help="Resolution scaling parameter.")

parser.add_argument("--layers",
                    default=16,
                    type=int,
                    action="store",
                    help="Number of vertical layers.")

parser.add_argument("--debug",
                    action="store_true",
                    help="Turn on KSP monitors")

parser.add_argument("--rtol",
                    default=1.0e-6,
                    type=float,
                    help="Rtolerance for the linear solve.")

parser.add_argument("--help",
                    action="store_true",
                    help="Show help.")


args, _ = parser.parse_known_args()

if args.help:
    help = parser.format_help()
    PETSc.Sys.Print("%s\n" % help)
    sys.exit(1)

tmax = args.tmax                # Maximum time (s)
nlayers = args.layers           # Number of vertical layers
refinements = args.refinements  # Number of horiz. cells = 20*(4^refinements)

hybrid = bool(args.hybridization)
PETSc.Sys.Print("""
Problem parameters:\n
Test case DCMIP 3-1: Non-orographic gravity waves on a small planet.\n
Hybridized compressible solver: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Profiling: %s,\n
Max time: %s,\n
Dump frequency: %s.\n
""" % (hybrid, refinements, nlayers,
       bool(args.profile), args.tmax, args.dumpfreq))

# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6               # Radius of the Earth (m)
X = 125.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 constant 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                     # Latitudinal centerpoint of Theta' (equator)
deltaTheta = 1.0                # Maximum amplitude of Theta' (K)
L_z = 20000.0                   # Vertical wave length of the Theta' perturb.

# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
                    refinement_level=refinements,
                    degree=2)

# Horizontal Courant (advective) number
cell_vs = interpolate(CellVolume(m),
                      FunctionSpace(m, "DG", 0))
a_min = cell_vs.dat.data.min()
a_max = cell_vs.dat.data.max()
dx_min = sqrt(a_min)
dx_max = sqrt(a_max)
dx_avg = (dx_min + dx_max)/2.0
u_max = u_0

# Take integer value
dt = int(args.cfl * (dx_avg / 343.0))

if args.profile:
    tmax = 5*dt

if args.test:
    tmax = dt

PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" % (dt,
                                                                      tmax))

# Build volume mesh
z_top = 1.0e4            # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers,
                    layer_height=z_top/nlayers,
                    extrusion_type="radial")

# Space for initialising velocity (using this ensures things are in layers)
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)

# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
z_expr = Expression("sqrt(x[0]*x[0] + x[1]*x[1] + x[2]*x[2]) - a", a=a)
z = Function(W_CG1).interpolate(z_expr)
{'fraction_non_alphanumeric': 0.10938768148615745, 'fraction_numerical': 0.0183966092524123, 'mean_word_length': 1.973987664253151, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2265318', 'n_tokens_mistral': 3496, 'n_tokens_neox': 3204, 'n_words': 888}