<commit_before>from dcore import *
from firedrake import Expression, \
    VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh
from firedrake import exp, sin, ds_b
import numpy as np


def setup_sk(dirname):
    nlayers = 10  # horizontal layers
    columns = 30  # number of columns
    L = 1.e5
    m = PeriodicIntervalMesh(columns, L)

    # build volume mesh
    H = 1.0e4  # Height position of the model top
    mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)

    # Space for initialising velocity
    W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
    W_CG1 = FunctionSpace(mesh, "CG", 1)

    # vertical coordinate and normal
    z = Function(W_CG1).interpolate(Expression("x[1]"))
    k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))

    fieldlist = ['u', 'rho', 'theta']
    timestepping = TimesteppingParameters(dt=15.0)
    output = OutputParameters(dirname=dirname+"/sk_linear", dumplist=['u'], dumpfreq=150)
    diagnostics = Diagnostics(*fieldlist)
    parameters = CompressibleParameters()

    state = CompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
                              family="CG",
                              z=z, k=k,
                              timestepping=timestepping,
                              output=output,
                              parameters=parameters,
                              diagnostics=diagnostics,
                              fieldlist=fieldlist)

    # Initial conditions
    u0, theta0, rho0 = Function(state.V[0]), Function(state.V[2]), Function(state.V[1])

    # Thermodynamic constants required for setting initial conditions
    # and reference profiles
    g = parameters.g
    N = parameters.N
    p_0 = parameters.p_0
    c_p = parameters.cp
    R_d = parameters.R_d
    kappa = parameters.kappa

    # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
    Tsurf = 300.
    thetab = Tsurf*exp(N**2*z/g)

    theta_b = Function(state.V[2]).interpolate(thetab)
    rho_b = Function(state.V[1])

    # Calculate hydrostatic Pi
    W = MixedFunctionSpace((state.Vv,state.V[1]))
    v, pi = TrialFunctions(W)
    dv, dpi = TestFunctions(W)

    n = FacetNormal(mesh)

    alhs = (
        (c_p*inner(v,dv) - c_p*div(dv*theta_b)*pi)*dx
        + dpi*div(theta_b*v)*dx
    )

    arhs = (
        - g*inner(dv,k)*dx
        - c_p*inner(dv,n)*theta_b*ds_b  # bottom surface value pi = 1.
    )
    bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), "top")]

    w = Function(W)
    PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)

    params = {'pc_type': 'fieldsplit',
              'pc_fieldsplit_type': 'schur',
              'ksp_type': 'gmres',
              'ksp_monitor_true_residual': True,
              'ksp_max_it': 100,
              'ksp_gmres_restart': 50,
              'pc_fieldsplit_schur_fact_type': 'FULL',
              'pc_fieldsplit_schur_precondition': 'selfp',
              'fieldsplit_0_ksp_type': 'richardson',
              'fieldsplit_0_ksp_max_it': 5,
              'fieldsplit_0_pc_type': 'bjacobi',
              'fieldsplit_0_sub_pc_type': 'ilu',
              'fieldsplit_1_ksp_type': 'richardson',
              'fieldsplit_1_ksp_max_it': 5,
              "fieldsplit_1_ksp_monitor_true_residual": True,
              'fieldsplit_1_pc_type': 'bjacobi',
              'fieldsplit_1_sub_pc_type': 'ilu'}

    PiSolver = LinearVariationalSolver(PiProblem,
                                       solver_parameters=params)

    PiSolver.solve()
    v, Pi = w.split()

    rho_b.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta_b)

    W_DG1 = FunctionSpace(mesh, "DG", 1)
    x = Function(W_DG1).interpolate(Expression("x[0]"))
    a = 5.0e3
    deltaTheta = 1.0e-2
    theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
    theta0.interpolate(theta_b + theta_pert)
    rho0.assign(rho_b)

    state.initialise([u0, rho0, theta0])
    state.set_reference_profiles(rho_b, theta_b)
    state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}

    # Set up advection schemes
    advection_list = []
    velocity_advection = NoAdvection(state)
    advection_list.append((velocity_advection, 0))
    rho_advection = LinearAdvection_V3(state, state.V[1], rho_b)
    advection_list.append((rho_advection, 1))
    theta_advection = LinearAdvection_Vt(state, state.V[2], theta_b)
    advection_list.append((theta_advection, 2))

    # Set up linear solver
    schur_params = {'pc_type': 'fieldsplit',
                    'pc_fieldsplit_type': 'schur',
                    'ksp_type': 'gmres',
                    'ksp_monitor_true_residual': True,
                    'ksp_max_it': 100,
                    'ksp_gmres_restart': 50,
                    'pc_fieldsplit_schur_fact_type': 'FULL',
                    'pc_fieldsplit_schur_precondition': 'selfp',
                    'fieldsplit_0_ksp_type': 'richardson',
                    'fieldsplit_0_ksp_max_it': 5,
                    'fieldsplit_0_pc_type': 'bjacobi',
                    'fieldsplit_0_sub_pc_type': 'ilu',
                    'fieldsplit_1_ksp_type': 'richardson',
                    'fieldsplit_1_ksp_max_it': 5,
                    "fieldsplit_1_ksp_monitor_true_residual": True,
                    'fieldsplit_1_pc_type': 'gamg',
                    'fieldsplit_1_pc_gamg_sym_graph': True,
                    'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
                    'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
                    'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
                    'fieldsplit_1_mg_levels_ksp_max_it': 5,
                    'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
                    'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}

    linear_solver = CompressibleSolver(state, params=schur_params)

    # Set up forcing
    compressible_forcing = CompressibleForcing(state)

    # build time stepper
    stepper = Timestepper(state, advection_list, linear_solver,
                          compressible_forcing)

    return stepper


def run_sk_linear(dirname):

    stepper = setup_sk(dirname)
    stepper.run(t=0, tmax=dt)


def test_sk(tmpdir):

    dirname = str(tmpdir)
    run_sk_linear(dirname)
<commit_msg>fixing test<commit_after>from dcore import *
from firedrake import Expression, \
    VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh
from firedrake import exp, sin, ds_b
import numpy as np


def setup_sk(dirname):
    nlayers = 10  # horizontal layers
    columns = 30  # number of columns
    L = 1.e5
    m = PeriodicIntervalMesh(columns, L)
    dt = 15.0

    # build volume mesh
    H = 1.0e4  # Height position of the model top
    mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)

    # Space for initialising velocity
    W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
    W_CG1 = FunctionSpace(mesh, "CG", 1)

    # vertical coordinate and normal
    z = Function(W_CG1).interpolate(Expression("x[1]"))
    k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))

    fieldlist = ['u', 'rho', 'theta']
    timestepping = TimesteppingParameters(dt=dt)
    output = OutputParameters(dirname=dirname+"/sk_linear", dumplist=['u'], dumpfreq=150)
    diagnostics = Diagnostics(*fieldlist)
    parameters = CompressibleParameters()

    state = CompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
                              family="CG",
                              z=z, k=k,
                              timestepping=timestepping,
                              output=output,
                              parameters=parameters,
                              diagnostics=diagnostics,
                              fieldlist=fieldlist)

    # Initial conditions
    u0, theta0, rho0 = Function(state.V[0]), Function(state.V[2]), Function(state.V[1])

    # Thermodynamic constants required for setting initial conditions
    # and reference profiles
    g = parameters.g
    N = parameters.N
    p_0 = parameters.p_0
    c_p = parameters.cp
    R_d = parameters.R_d
    kappa = parameters.kappa

    # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
    Tsurf = 300.
    thetab = Tsurf*exp(N**2*z/g)

    theta_b = Function(state.V[2]).interpolate(thetab)
    rho_b = Function(state.V[1])

    # Calculate hydrostatic Pi
    W = MixedFunctionSpace((state.Vv,state.V[1]))
    v, pi = TrialFunctions(W)
    dv, dpi = TestFunctions(W)

    n = FacetNormal(mesh)

    alhs = (
        (c_p*inner(v,dv) - c_p*div(dv*theta_b)*pi)*dx
        + dpi*div(theta_b*v)*dx
    )

    arhs = (
        - g*inner(dv,k)*dx
        - c_p*inner(dv,n)*theta_b*ds_b  # bottom surface value pi = 1.
    )
    bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), "top")]

    w = Function(W)
    PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)

    params = {'pc_type': 'fieldsplit',
              'pc_fieldsplit_type': 'schur',
              'ksp_type': 'gmres',
              'ksp_monitor_true_residual': True,
              'ksp_max_it': 100,
              'ksp_gmres_restart': 50,
              'pc_fieldsplit_schur_fact_type': 'FULL',
              'pc_fieldsplit_schur_precondition': 'selfp',
              'fieldsplit_0_ksp_type': 'richardson',
              'fieldsplit_0_ksp_max_it': 5,
              'fieldsplit_0_pc_type': 'bjacobi',
              'fieldsplit_0_sub_pc_type': 'ilu',
              'fieldsplit_1_ksp_type': 'richardson',
              'fieldsplit_1_ksp_max_it': 5,
              "fieldsplit_1_ksp_monitor_true_residual": True,
              'fieldsplit_1_pc_type': 'bjacobi',
              'fieldsplit_1_sub_pc_type': 'ilu'}

    PiSolver = LinearVariationalSolver(PiProblem,
                                       solver_parameters=params)

    PiSolver.solve()
    v, Pi = w.split()

    rho_b.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta_b)

    W_DG1 = FunctionSpace(mesh, "DG", 1)
    x = Function(W_DG1).interpolate(Expression("x[0]"))
    a = 5.0e3
    deltaTheta = 1.0e-2
    theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
    theta0.interpolate(theta_b + theta_pert)
    rho0.assign(rho_b)

    state.initialise([u0, rho0, theta0])
    state.set_reference_profiles(rho_b, theta_b)
    state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}

    # Set up advection schemes
    advection_list = []
    velocity_advection = NoAdvection(state)
    advection_list.append((velocity_advection, 0))
    rho_advection = LinearAdvection_V3(state, state.V[1], rho_b)
    advection_list.append((rho_advection, 1))
    theta_advection = LinearAdvection_Vt(state, state.V[2], theta_b)
    advection_list.append((theta_advection, 2))

    # Set up linear solver
    schur_params = {'pc_type': 'fieldsplit',
                    'pc_fieldsplit_type': 'schur',
                    'ksp_type': 'gmres',
                    'ksp_monitor_true_residual': True,
                    'ksp_max_it': 100,
                    'ksp_gmres_restart': 50,
                    'pc_fieldsplit_schur_fact_type': 'FULL',
                    'pc_fieldsplit_schur_precondition': 'selfp',
                    'fieldsplit_0_ksp_type': 'richardson',
                    'fieldsplit_0_ksp_max_it': 5,
                    'fieldsplit_0_pc_type': 'bjacobi',
                    'fieldsplit_0_sub_pc_type': 'ilu',
                    'fieldsplit_1_ksp_type': 'richardson',
                    'fieldsplit_1_ksp_max_it': 5,
                    "fieldsplit_1_ksp_monitor_true_residual": True,
                    'fieldsplit_1_pc_type': 'gamg',
                    'fieldsplit_1_pc_gamg_sym_graph': True,
                    'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
                    'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
                    'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
                    'fieldsplit_1_mg_levels_ksp_max_it': 5,
                    'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
                    'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}

    linear_solver = CompressibleSolver(state, params=schur_params)

    # Set up forcing
    compressible_forcing = CompressibleForcing(state)

    # build time stepper
    stepper = Timestepper(state, advection_list, linear_solver,
                          compressible_forcing)

    return stepper, dt


def run_sk_linear(dirname):

    stepper, tmax = setup_sk(dirname)
    stepper.run(t=0, tmax=tmax)


def test_sk(tmpdir):

    dirname = str(tmpdir)
    run_sk_linear(dirname)

{'fraction_non_alphanumeric': 0.11319708491761724, 'fraction_numerical': 0.018377693282636248, 'mean_word_length': 2.418629840238289, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5289969', 'n_tokens_mistral': 4545, 'n_tokens_neox': 4229, 'n_words': 699}