# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # This notebook demos a shallow water stencil

# We first write our shallow water stencil:

# +
from dusk.script import *


@stencil
def shallow_water(hC: Field[Cell], hC_t: Field[Cell],
                     vC: Field[Cell], vC_t: Field[Cell],
                     uC: Field[Cell], uC_t: Field[Cell],
                     hC_x: Field[Cell], hC_y: Field[Cell], uvC_div: Field[Cell],
                     hE: Field[Edge], vE: Field[Edge], uE: Field[Edge],
                     nx: Field[Edge], ny: Field[Edge], L: Field[Edge], alpha: Field[Edge],
                     boundary_edges: Field[Edge], boundary_cells: Field[Cell],
                     A: Field[Cell], edge_orientation: Field[Cell > Edge],
                     Grav: Field[Cell]):

    with levels_downward:
        # todo: implement shallow water solver
        pass



# -

# Then we can use dusk's Python API to convert the stencils to SIR. This API can also invoke dawn to compile SIR to C++ which we will write to `shallow_water_cxx-naive.cpp`:

from dusk.transpile import callable_to_pyast, pyast_to_sir, sir_to_json
with open("shallow_water.sir", "w+") as f:
    sir = pyast_to_sir(callable_to_pyast(shallow_water))
    f.write(sir_to_json(sir))
# !dawn-opt shallow_water.sir | dawn-codegen -b naive-ico -o shallow_water_cxx-naive.cpp

# The generated C++ code also requires a driver which is already setup for this demo. With the driver code we can generate an executable `runner`:

# !make

# First, we put the runner into test mode to ensure that the computed kernel is correct:

# !./runner test

# If the tester reported that your dusk stencil works correctly, you can now run the complete stencil and visualize. It takes quite a while to run. Its finished after 400 time steps.

# !./runner run

# +
# %%capture
import prep_animation as pa
import matplotlib.animation as animation
from IPython.display import HTML

animate = animation.FuncAnimation(pa.fig, pa.update_plot, len(pa.out_files), fargs=(pa.zarray, pa.plot), interval = 50)
# -

HTML(animate.to_html5_video())

StarcoderdataJupyterScriptsDedupFiltered