Implementation

Full simulation file

HRB.py

Step-by-step

As in Your first simulation, we create a new file HRB.py in Simulations/.

We then implement the HRB equations:

Non-linear term

This corresponds to \(\partial_t u_i = \mathbb{P}\left[-u_j\partial_ju_i\right]_i, \partial_t\theta_i = -u_i\partial_i\theta + u_z\)

def equation_nonlinear(_t: float, grid: Grid, _simu_params: dict) -> dict[str, np.ndarray]:
    M = grid.maths
    ux, uy, uz, theta = grid.field("ux", "uy", "uz", "theta")

    # Convolutions
    uxdxux, uydyux, uzdzux, uxdxuy, uydyuy, uzdzuy, uxdxuz, uydyuz, uzdzuz, uxdxtheta, uydytheta, uzdztheta = M.convolve_batch(
        (
            (ux, M.dx * ux),
            (uy, M.dy * ux),
            (uz, M.dz * ux),
            (ux, M.dx * uy),
            (uy, M.dy * uy),
            (uz, M.dz * uy),
            (ux, M.dx * uz),
            (uy, M.dy * uz),
            (uz, M.dz * uz),
            (ux, M.dx * theta),
            (uy, M.dy * theta),
            (uz, M.dz * theta),
        )
    )

    # w/o pressure
    dux_dt = -uxdxux - uydyux - uzdzux
    duy_dt = -uxdxuy - uydyuy - uzdzuy
    duz_dt = -uxdxuz - uydyuz - uzdzuz + theta
    dtheta_dt = -uxdxtheta - uydytheta - uzdztheta + uz

    # Add pressure
    dux_dt, duy_dt, duz_dt = grid.maths.P_projector([dux_dt, duy_dt, duz_dt])

    return {"ux": dux_dt, "uy": duy_dt, "uz": duz_dt, "theta": dtheta_dt}

Linear term

We then add the linear viscous/friction part \(\partial_t u_i = \mathbb{P}\left[\frac{Pr}{Ra}\nabla^2u_i-fu_i\delta_{k\approx k_{min}}\right]_i, \partial_t\theta_i = \frac{\nabla^2\theta}{\sqrt{RaPr}}-f\delta_{k\approx k_{min}}\):

def equation_linear(_t: float, grid: Grid, simu_params: dict) -> dict[str, np.ndarray]:
    M = grid.maths
    Pr, Ra = simu_params["Pr"], simu_params["Ra"]

    f = simu_params["f"] * np.ones_like(grid.ks_modulus)
    f[grid.ks_modulus > grid.k_min * grid.l ** 3] = 0

    visc = np.sqrt(Pr / Ra) * M.laplacian - f
    visc_theta = 1 / np.sqrt(Ra * Pr) * M.laplacian - f

    return {"ux": visc, "uy": visc, "uz": visc, "theta": visc_theta}

Initial conditions

We initilize fields with random values at large scales:

def initial_conditions(fields: dict[str, np.ndarray], grid: Grid, _simu_params: dict) -> dict[str, np.ndarray]:
    grid = grid.to_new_size_empty(N_points)
    ks = grid.ks_modulus

    wx = np.zeros_like(ks, dtype="complex")
    wy = np.zeros_like(ks, dtype="complex")
    wz = np.zeros_like(ks, dtype="complex")
    theta = randcomplex_like(ks) * np.ones_like(ks, dtype="complex") / 100
    theta[ks > grid.k_min * 4] = 0
    ux, uy, uz = grid.maths.rot3D_inv([wx, wy, wz])

    fields["theta"] = theta
    fields["ux"] = ux
    fields["uy"] = uy
    fields["uz"] = uz

    return fields

Grid size update

We update the grid size based on the fraction of energy contained in the outermost layers:

def update_gridsize(grid: Grid) -> int | None:
    """update the grid size based on the fraction of energy contained in the outermost layers"""
    E = grid.physics.energy()
    ux, uy, uz, theta = grid.field("ux", "uy", "uz", "theta")
    mask = grid.ks_modulus > grid.k_min * grid.l ** (grid.N_points - 1)
    # grid
    comp = np.max(np.abs(ux[mask]) + np.abs(uy[mask]) + np.abs(uz[mask]) + np.abs(theta[mask]))
    if comp / np.sqrt(E) > 1e-100:
        return grid.N_points + 1
    if comp / np.sqrt(E) < 1e-170 and grid.N_points > 5:
        return grid.N_points - 1

Core simulation code

We then call the solver with a given set of parameters:

Pr = 1
Ra = 1e12
print(f"Chosen parameters Pr={Pr}, Ra={Ra}")
fields = ["ux", "uy", "uz", "theta"]
D = 3
l_params = {"plastic": False, "a": 1, "b": 2}
n_threads_convolution = 6
simu_params = {"Pr": Pr, "Ra": Ra, "f": 1}
N_points = 13
rtol = 1e-1

save_path = "results/saveRB"
loadfromsave = False

solver = Solver(
    fields_names=fields,
    equation_nl=equation_nonlinear,
    D=D,
    equation_l=equation_linear,
    l_params=l_params,
    simu_params=simu_params,
    n_threads=n_threads_convolution,
)
solver.solve(
    save_path=save_path,
    initial_conditions=initial_conditions,
    save_one_in=100,
    update_gridsize_cb=update_gridsize,
    solver_params={"rtol": rtol, "adapt_cutoff": 1e-5},
)