<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 2D Channel with Time-Dependent Boundary Conditions
This example demonstrates the simulation of a flow in a 2D channel in a closed rectangular domain using constant and time dependent boundary conditions.
The flow is forced by an initial perturbation in the water elevation field centred in the middle of the channel. The model solves the Shallow Water Equations which are not expanded herein for brevity.
To run this example, it is necessary to install firedrake and Thetis
Link to Tutorials on Thetis Website
We begin with a number of imports. Note that we use from thetis import * which means that we do not need to prepend thetis. all the time. This also automatically includes a from firedrake import * so that we have access to the all the firedrake objects we saw yesterday such as RectangleMesh, Function, FunctionSpace etc. - without having to prepend fd as we did then.
Step1:  We create a rectangular mesh using the builtin Firedrake mesh utility. 
Our domain is 40 km long and 2 km wide, and we will generate 25 elements in the along-channel direction and 2 in the cross-channel direction
Step2:  We can use the built in plot function of firedrake to visualise the mesh.
Step3:  Next we define a bathymetry function in this domain using continuous linear elements, and set the bathymetry to a constant 20 m depth
Step4:  We are now ready to create a 2D solver object. This object can be used to solve the depth-averaged shallow water equations controlled by some options that can be set by the user.
Step5:  First we set the options for the total duration and export intervals. The latter determines how frequently Thetis writes output files.
Step6:  We also select the directory in which the output files will appear, which will be created automatically if it does not yet exist
Step7:  The standard output of Thetis is in the form of a .pvd file with an associated series of .vtu files (and .pvtu when running in parallel). For visualisation we recommend the program paraview. This program is however not available inside this jupyter environment. Therefore we will here use hdf5 output files instead, which can be used to read back in the solution in Firedrake and Thetis (unlike the .pvd+.vtu output which cannot be read back in by Firedrake). This output can also be used as checkpoints, from which Thetis can be restarted. Here we will use the hdf5 output to read it back into Firedrake and plot it.
We select which fields to output ('elev_2d' for elevations and 'uv_2d' for the 2d velocity vector)
Step8:  Next we define the time integrator, and set the time step, which can be chosen freely since Crank-Nicolson is unconditionally stable. Of course it will influence the accuracy of the simulation.
Step9:  Boundary conditions
We will force the model with a prescribed tidal elevation on the right side of the domain ($x=40 $km), and keep a constant elevation on the left ($x=0$).
Boundary conditions are defined for each external boundary using their ID.
In this example we are using a RectangleMesh() which assigns IDs 1, 2, 3, and 4 for the four sides of the rectangle
Step10:  At each boundary we can define the external value of the prognostic variables, i.e. in this case the water elevation and velocity.The value should be either a Firedrake Constant or a Firedrake Function (in case the boundary condition is not uniform in space).
We store the boundary conditions in a dictionary.
Step12:  Above we set the water elevation (using key 'elev') on the left bounday. Alternatively, we could also prescribe the normal velocity (with key 'un'), the full 2D velocity vector ('uv') or the total flux through the boundary ('flux'). For all supported boundary conditions, see module shallowwater_eq.
In order to set time-dependent boundary conditions we first define a python function that evaluates the time dependent variable
Step13:  We then create a Constant object with the initial value, and assign it to the left boundary
Step14:  Boundary conditions are now complete, and we assign them to the solver object
Step16:  Note that if boundary conditions are not assigned for some boundaries (the lateral boundaries 3 and 4 in this case), Thetis assumes impermeable land conditions.
The only missing piece is to add a mechanism that re-evaluates the boundary condition as the simulation progresses. For this purpose we use the optional update_forcings argument of the iterate() method. update_forcings is a python function that updates all time dependent Constants or Functions used to force the model. In this case we only need to update tide_elev_const
Step17:  Finally we pass this callback to the time iterator and run the model
Step18:  While the model is running, Thetis prints some statistics on the command line
Step19:  The following code opens the 9th output. We use os.path.join to combine the different parts (base outputdirectory, hdf5 sub directory, and actual filename) into a filename.
Step20:  Now we can read the solution into the elev and plot it
Step21:  In the same way we can also load and plot the velocity field
Step22:  The following plots the elevations of all exported timesteps
Step25:  Exercises
Step26:  Try to make the following changes
Step27:  For now let's make sure the bathymetry is sufficiently deep, so that the total depth (including the free surface elevation) does not go negative! This avoids us having to deal with wetting and drying.
now let us try a bathymetry, where we do include wetting and drying. you will need the following additional option in your setup
Step28:  To simulate the run up onto a beach, you could try the following bathymetry
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt
from thetis import *
Explanation: 2D Channel with Time-Dependent Boundary Conditions
This example demonstrates the simulation of a flow in a 2D channel in a closed rectangular domain using constant and time dependent boundary conditions.
The flow is forced by an initial perturbation in the water elevation field centred in the middle of the channel. The model solves the Shallow Water Equations which are not expanded herein for brevity.
To run this example, it is necessary to install firedrake and Thetis
Link to Tutorials on Thetis Website
We begin with a number of imports. Note that we use from thetis import * which means that we do not need to prepend thetis. all the time. This also automatically includes a from firedrake import * so that we have access to the all the firedrake objects we saw yesterday such as RectangleMesh, Function, FunctionSpace etc. - without having to prepend fd as we did then.
End of explanation
lx = 40e3
ly = 2e3
nx = 25
ny = 2
mesh2d = RectangleMesh(nx, ny, lx, ly)
Explanation: We create a rectangular mesh using the builtin Firedrake mesh utility. 
Our domain is 40 km long and 2 km wide, and we will generate 25 elements in the along-channel direction and 2 in the cross-channel direction:
End of explanation
fig, ax = plt.subplots(figsize=(12,1))
triplot(mesh2d, axes=ax);
Explanation: We can use the built in plot function of firedrake to visualise the mesh.
End of explanation
P1_2d = FunctionSpace(mesh2d, 'CG', 1)
bathymetry_2d = Function(P1_2d, name='Bathymetry')
depth = 20.0
bathymetry_2d.assign(depth);
Explanation: Next we define a bathymetry function in this domain using continuous linear elements, and set the bathymetry to a constant 20 m depth:
End of explanation
solver_obj = solver2d.FlowSolver2d(mesh2d, bathymetry_2d)
options = solver_obj.options
Explanation: We are now ready to create a 2D solver object. This object can be used to solve the depth-averaged shallow water equations controlled by some options that can be set by the user.
End of explanation
options.simulation_end_time = 12 * 3600
options.simulation_export_time = 1200.0
Explanation: First we set the options for the total duration and export intervals. The latter determines how frequently Thetis writes output files.
End of explanation
options.output_directory = 'outputs_2d_channel'
Explanation: We also select the directory in which the output files will appear, which will be created automatically if it does not yet exist:
End of explanation
options.fields_to_export_hdf5 = ['elev_2d', 'uv_2d']
Explanation: The standard output of Thetis is in the form of a .pvd file with an associated series of .vtu files (and .pvtu when running in parallel). For visualisation we recommend the program paraview. This program is however not available inside this jupyter environment. Therefore we will here use hdf5 output files instead, which can be used to read back in the solution in Firedrake and Thetis (unlike the .pvd+.vtu output which cannot be read back in by Firedrake). This output can also be used as checkpoints, from which Thetis can be restarted. Here we will use the hdf5 output to read it back into Firedrake and plot it.
We select which fields to output ('elev_2d' for elevations and 'uv_2d' for the 2d velocity vector):
End of explanation
options.timestepper_type = 'CrankNicolson'
options.timestep = 200.0
Explanation: Next we define the time integrator, and set the time step, which can be chosen freely since Crank-Nicolson is unconditionally stable. Of course it will influence the accuracy of the simulation.
End of explanation
left_bnd_id = 1
right_bnd_id = 2
Explanation: Boundary conditions
We will force the model with a prescribed tidal elevation on the right side of the domain ($x=40 $km), and keep a constant elevation on the left ($x=0$).
Boundary conditions are defined for each external boundary using their ID.
In this example we are using a RectangleMesh() which assigns IDs 1, 2, 3, and 4 for the four sides of the rectangle:
End of explanation
swe_bnd = {}
swe_bnd[left_bnd_id] = {'elev': Constant(0.0)}
Explanation: At each boundary we can define the external value of the prognostic variables, i.e. in this case the water elevation and velocity.The value should be either a Firedrake Constant or a Firedrake Function (in case the boundary condition is not uniform in space).
We store the boundary conditions in a dictionary.
End of explanation
tide_amp = 0.01  # amplitude of the tide
tide_t = 12 * 3600.  # period of the tide
def tidal_elevation(simulation_time):
    Time-dependent tidal elevation
    elev = tide_amp * sin(2 * pi * simulation_time / tide_t)
    return elev
plt.plot(np.arange(0,100000,500),[tidal_elevation(i) for i in np.arange(0,100000,500)])
plt.ylabel('Elev (m)')
plt.xlabel('t (sec)');
Explanation: Above we set the water elevation (using key 'elev') on the left bounday. Alternatively, we could also prescribe the normal velocity (with key 'un'), the full 2D velocity vector ('uv') or the total flux through the boundary ('flux'). For all supported boundary conditions, see module shallowwater_eq.
In order to set time-dependent boundary conditions we first define a python function that evaluates the time dependent variable:
End of explanation
tide_elev_const = Constant(tidal_elevation(0))
swe_bnd[right_bnd_id] = {'elev': tide_elev_const}
Explanation: We then create a Constant object with the initial value, and assign it to the left boundary:
End of explanation
solver_obj.bnd_functions['shallow_water'] = swe_bnd
Explanation: Boundary conditions are now complete, and we assign them to the solver object:
End of explanation
def update_forcings(t_new):
    Callback function that updates all time dependent forcing fields
    uv, elev = solver_obj.fields.solution_2d.split()
    tide_elev_const.assign(tidal_elevation(t_new))
Explanation: Note that if boundary conditions are not assigned for some boundaries (the lateral boundaries 3 and 4 in this case), Thetis assumes impermeable land conditions.
The only missing piece is to add a mechanism that re-evaluates the boundary condition as the simulation progresses. For this purpose we use the optional update_forcings argument of the iterate() method. update_forcings is a python function that updates all time dependent Constants or Functions used to force the model. In this case we only need to update tide_elev_const:
End of explanation
solver_obj.iterate(update_forcings=update_forcings)
Explanation: Finally we pass this callback to the time iterator and run the model
End of explanation
elev = Function(solver_obj.function_spaces.H_2d, name='elev_2d')
uv = Function(solver_obj.function_spaces.U_2d, name='uv_2d')
Explanation: While the model is running, Thetis prints some statistics on the command line:
0     0 T=      0.00 eta norm:     0.0000 u norm:     0.0000  0.00
1     6 T=   1200.00 eta norm:     6.5453 u norm:     4.5832  8.49
2    12 T=   2400.00 eta norm:    17.3241 u norm:    12.1339  0.16
3    18 T=   3600.00 eta norm:    29.9394 u norm:    22.0483  0.16
 ...
The first column is the export index, the second one the number of executed
 time steps, followed by the simulation time. eta norm and u norm are
 the L2 norms of the elevation and depth averaged velocity fields, respectively.
 The last column stands for the (approximate) wall-clock time between exports.
The simulation terminates once the end time is reached.
Visualising the output
Information on how to visualise results can be found here outputs and visualization page.
Here we will use the hdf5 output (also used for checkpointing) to read back in the solution at different timesteps and plot it using Firedrake. First we need an elevation function, and a velocity function from the same function space that we used in the simulation:
End of explanation
idx = 9 # which hdf5 do we want to open
filename =  os.path.join(options.output_directory, 'hdf5','Elevation2d_%05d' % idx)
dc = DumbCheckpoint(filename, mode=FILE_READ)
Explanation: The following code opens the 9th output. We use os.path.join to combine the different parts (base outputdirectory, hdf5 sub directory, and actual filename) into a filename.
End of explanation
dc.load(elev)
fig, ax = plt.subplots(figsize=(12,2))
tricontourf(elev, axes=ax, cmap=matplotlib.cm.coolwarm, levels=50)
plt.axis('equal');
Explanation: Now we can read the solution into the elev and plot it:
End of explanation
idx = 9 # which hdf5 do we want to open
filename = os.path.join(options.output_directory, 'hdf5','Velocity2d_%05d' % idx)
dc = DumbCheckpoint(filename, mode=FILE_READ)
dc.load(uv)
fig, ax = plt.subplots(figsize=(12,2))
quiver(uv, axes=ax, cmap=matplotlib.cm.coolwarm)
plt.axis('equal');
Explanation: In the same way we can also load and plot the velocity field:
End of explanation
last_idx = solver_obj.i_export
matplotlib.rcParams['figure.max_open_warning'] = last_idx+1  # avoid warning
for idx in range(last_idx+1):
    filename =  os.path.join(options.output_directory, 'hdf5','Elevation2d_%05d' % idx)
    dc = DumbCheckpoint(filename, mode=FILE_READ)
    dc.load(elev)
    dc.close()
    fig, ax = plt.subplots(figsize=(16,1))
    tricontourf(elev, axes=ax, cmap=matplotlib.cm.coolwarm, levels=50)
    # Firedrake sets an automatic colorbar range which therefore changes per timestep
    # instead, we want to fix it:
    cbar = fig.axes[0].collections[-1].set_clim(-tide_amp, tide_amp)
    plt.axis('equal')
Explanation: The following plots the elevations of all exported timesteps:
End of explanation
# setup the mesh:
lx = 40e3
ly = 2e3
nx = 25
ny = 2
mesh2d = RectangleMesh(nx, ny, lx, ly)
# setup the bathymetry:
P1_2d = FunctionSpace(mesh2d, 'CG', 1)
bathymetry_2d = Function(P1_2d, name='Bathymetry')
depth = 20.0
bathymetry_2d.assign(depth);
# setup the solver object and set some options
solver_obj = solver2d.FlowSolver2d(mesh2d, bathymetry_2d)
options = solver_obj.options
options.simulation_end_time = 12 * 3600
options.simulation_export_time = 1200.0
options.output_directory = 'outputs_2d_channel'
options.fields_to_export_hdf5 = ['elev_2d', 'uv_2d']
options.timestepper_type = 'CrankNicolson'
options.timestep = 200.0
# setup boundary conditions:
left_bnd_id = 1
right_bnd_id = 2
swe_bnd = {}
swe_bnd[left_bnd_id] = {'elev': Constant(0.0)}
tide_amp = 0.01  # amplitude of the tide
tide_t = 12 * 3600.  # period of the tide
def tidal_elevation(simulation_time):
    Time-dependent tidal elevation
    elev = tide_amp * sin(2 * pi * simulation_time / tide_t)
    return elev
tide_elev_const = Constant(tidal_elevation(0))
swe_bnd[right_bnd_id] = {'elev': tide_elev_const}
solver_obj.bnd_functions['shallow_water'] = swe_bnd
def update_forcings(t_new):
    Callback function that updates all time dependent forcing fields
    uv, elev = solver_obj.fields.solution_2d.split()
    tide_elev_const.assign(tidal_elevation(t_new))
solver_obj.iterate(update_forcings=update_forcings)
Explanation: Exercises:
In the exercises we will try to make a number of changes to the above setup. For ease of use we have copied all the steps above in a single cell, so that you can easily repeat them and make your own adaptations:
End of explanation
x, y = SpatialCoordinate(mesh2d)
bathymetry_2d.interpolate(depth * (1-0.5*max_value(cos((x-lx/2)/lx*pi*2), 0)));
Explanation: Try to make the following changes:
apply a prescribed velocity boundary condition at the open boundary instead of a prescribed elevation
change the bathymetry to be spatially varying, for instance use:
End of explanation
options.use_wetting_and_drying = True
Explanation: For now let's make sure the bathymetry is sufficiently deep, so that the total depth (including the free surface elevation) does not go negative! This avoids us having to deal with wetting and drying.
now let us try a bathymetry, where we do include wetting and drying. you will need the following additional option in your setup:
End of explanation
x, y = SpatialCoordinate(mesh2d)
bathymetry_2d.interpolate(depth * x/lx);
Explanation: To simulate the run up onto a beach, you could try the following bathymetry:
End of explanation
<END_TASK>