PyMPDATA_examples.burgers_equation.burgers_equation

Solution for the Burgers equation solution with MPDATA

 1"""
 2Solution for the Burgers equation solution with MPDATA
 3"""
 4
 5import numpy as np
 6
 7from PyMPDATA import Options, ScalarField, Solver, Stepper, VectorField
 8from PyMPDATA.boundary_conditions import Constant
 9
10OPTIONS = Options(nonoscillatory=False, infinite_gauge=True)
11
12
13def initialize_simulation(nt, nx, t_max):
14    """
15    Initializes simulation variables and returns them.
16    """
17    dt = t_max / nt
18    courants_x, dx = np.linspace(-1, 1, nx + 1, endpoint=True, retstep=True)
19    x = courants_x[:-1] + dx / 2
20    u0 = -np.sin(np.pi * x)
21
22    stepper = Stepper(options=OPTIONS, n_dims=1)
23    advectee = ScalarField(
24        data=u0, halo=OPTIONS.n_halo, boundary_conditions=(Constant(0), Constant(0))
25    )
26    advector = VectorField(
27        data=(np.full(courants_x.shape, 0.0),),
28        halo=OPTIONS.n_halo,
29        boundary_conditions=(Constant(0), Constant(0)),
30    )
31    solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
32    return dt, dx, x, advectee, advector, solver
33
34
35def update_advector_n(vel, dt, dx, slice_idx):
36    """
37    Computes and returns the updated advector_n.
38    """
39    indices = np.arange(slice_idx.start, slice_idx.stop)
40    return 0.5 * ((vel[indices] - vel[indices - 1]) / 2 + vel[:-1]) * dt / dx
41
42
43def run_numerical_simulation(*, nt, nx, t_max):
44    """
45    Runs the numerical simulation and returns (states, x, dt, dx).
46    """
47    dt, dx, x, advectee, advector, solver = initialize_simulation(nt, nx, t_max)
48    states = []
49    vel = advectee.get()
50    advector_n_1 = 0.5 * (vel[:-1] + np.diff(vel) / 2) * dt / dx
51    assert np.all(advector_n_1 <= 1)
52    i = slice(1, len(vel))
53
54    for _ in range(nt):
55        vel = advectee.get()
56        advector_n = update_advector_n(vel, dt, dx, i)
57        advector.get_component(0)[1:-1] = 0.5 * (3 * advector_n - advector_n_1)
58        assert np.all(advector.get_component(0) <= 1)
59
60        solver.advance(n_steps=1)
61        advector_n_1 = advector_n.copy()
62        states.append(solver.advectee.get().copy())
63
64    return np.array(states), x, dt, dx
OPTIONS = <PyMPDATA.options.Options object>
def initialize_simulation(nt, nx, t_max): 
14def initialize_simulation(nt, nx, t_max):
15    """
16    Initializes simulation variables and returns them.
17    """
18    dt = t_max / nt
19    courants_x, dx = np.linspace(-1, 1, nx + 1, endpoint=True, retstep=True)
20    x = courants_x[:-1] + dx / 2
21    u0 = -np.sin(np.pi * x)
22
23    stepper = Stepper(options=OPTIONS, n_dims=1)
24    advectee = ScalarField(
25        data=u0, halo=OPTIONS.n_halo, boundary_conditions=(Constant(0), Constant(0))
26    )
27    advector = VectorField(
28        data=(np.full(courants_x.shape, 0.0),),
29        halo=OPTIONS.n_halo,
30        boundary_conditions=(Constant(0), Constant(0)),
31    )
32    solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
33    return dt, dx, x, advectee, advector, solver

Initializes simulation variables and returns them.

def update_advector_n(vel, dt, dx, slice_idx): 
36def update_advector_n(vel, dt, dx, slice_idx):
37    """
38    Computes and returns the updated advector_n.
39    """
40    indices = np.arange(slice_idx.start, slice_idx.stop)
41    return 0.5 * ((vel[indices] - vel[indices - 1]) / 2 + vel[:-1]) * dt / dx

Computes and returns the updated advector_n.

def run_numerical_simulation(*, nt, nx, t_max): 
44def run_numerical_simulation(*, nt, nx, t_max):
45    """
46    Runs the numerical simulation and returns (states, x, dt, dx).
47    """
48    dt, dx, x, advectee, advector, solver = initialize_simulation(nt, nx, t_max)
49    states = []
50    vel = advectee.get()
51    advector_n_1 = 0.5 * (vel[:-1] + np.diff(vel) / 2) * dt / dx
52    assert np.all(advector_n_1 <= 1)
53    i = slice(1, len(vel))
54
55    for _ in range(nt):
56        vel = advectee.get()
57        advector_n = update_advector_n(vel, dt, dx, i)
58        advector.get_component(0)[1:-1] = 0.5 * (3 * advector_n - advector_n_1)
59        assert np.all(advector.get_component(0) <= 1)
60
61        solver.advance(n_steps=1)
62        advector_n_1 = advector_n.copy()
63        states.append(solver.advectee.get().copy())
64
65    return np.array(states), x, dt, dx

Runs the numerical simulation and returns (states, x, dt, dx).