#!/usr/bin/env python
# coding: utf-8

# ### Example 2: Nonlinear convection in 2D
# 
# Following the initial convection tutorial with a single state variable $u$, we will now look at non-linear convection (step 6 in the original). This brings one new crucial challenge: computing a pair of coupled equations and thus updating two time-dependent variables $u$ and $v$.
# 
# The full set of coupled equations is now
# 
# \begin{aligned}
# \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} + v \frac{\partial u}{\partial y} = 0 \\
# \\
# \frac{\partial v}{\partial t} + u \frac{\partial v}{\partial x} + v \frac{\partial v}{\partial y} = 0\\
# \end{aligned}
# 
# and rearranging the discretized version gives us an expression for the update of both variables
# 
# \begin{aligned}
# u_{i,j}^{n+1} &= u_{i,j}^n - u_{i,j}^n \frac{\Delta t}{\Delta x} (u_{i,j}^n-u_{i-1,j}^n) - v_{i,j}^n \frac{\Delta t}{\Delta y} (u_{i,j}^n-u_{i,j-1}^n) \\
# \\
# v_{i,j}^{n+1} &= v_{i,j}^n - u_{i,j}^n \frac{\Delta t}{\Delta x} (v_{i,j}^n-v_{i-1,j}^n) - v_{i,j}^n \frac{\Delta t}{\Delta y} (v_{i,j}^n-v_{i,j-1}^n)
# \end{aligned}
# 
# So, for starters we will re-create the original example run in pure NumPy array notation, before demonstrating 
# the Devito version. Let's start again with some utilities and parameters:

# In[1]:


from examples.cfd import plot_field, init_hat
import numpy as np
import sympy
get_ipython().run_line_magic('matplotlib', 'inline')

# Some variable declarations
nx = 101
ny = 101
nt = 80
c = 1.
dx = 2. / (nx - 1)
dy = 2. / (ny - 1)
sigma = .2
dt = sigma * dx


# Let's re-create the initial setup with a 2D "hat function", but this time for two state variables.

# In[2]:


#NBVAL_IGNORE_OUTPUT

# Allocate fields and assign initial conditions
u = np.empty((nx, ny))
v = np.empty((nx, ny))

init_hat(field=u, dx=dx, dy=dy, value=2.)
init_hat(field=v, dx=dx, dy=dy, value=2.)

plot_field(u)


# Now we can create the two stencil expression for our two coupled equations according to the discretized equation above. We again use some simple Dirichlet boundary conditions to keep the values on all sides constant.

# In[3]:


#NBVAL_IGNORE_OUTPUT
for n in range(nt + 1): ##loop across number of time steps
    un = u.copy()
    vn = v.copy()
    u[1:, 1:] = (un[1:, 1:] - 
                 (un[1:, 1:] * c * dt / dy * (un[1:, 1:] - un[1:, :-1])) -
                  vn[1:, 1:] * c * dt / dx * (un[1:, 1:] - un[:-1, 1:]))
    v[1:, 1:] = (vn[1:, 1:] -
                 (un[1:, 1:] * c * dt / dy * (vn[1:, 1:] - vn[1:, :-1])) -
                 vn[1:, 1:] * c * dt / dx * (vn[1:, 1:] - vn[:-1, 1:]))
    
    u[0, :] = 1
    u[-1, :] = 1
    u[:, 0] = 1
    u[:, -1] = 1
    
    v[0, :] = 1
    v[-1, :] = 1
    v[:, 0] = 1
    v[:, -1] = 1
    
plot_field(u)


# Excellent, we again get a wave that resembles the one from the oiginal examples.
# 
# Now we can set up our coupled problem in Devito. Let's start by creating two initial state variables $u$ and $v$, as before, and initialising them with our "hat function.

# In[4]:


#NBVAL_IGNORE_OUTPUT
from devito import Grid, TimeFunction

# First we need two time-dependent data fields, both initialized with the hat function
grid = Grid(shape=(nx, ny), extent=(2., 2.))
u = TimeFunction(name='u', grid=grid)
init_hat(field=u.data[0], dx=dx, dy=dy, value=2.)

v = TimeFunction(name='v', grid=grid)
init_hat(field=v.data[0], dx=dx, dy=dy, value=2.)

plot_field(u.data[0])


# Using the two `TimeFunction` objects we can again derive our discretized equation, rearrange for the forward stencil point in time and define our variable update expression - only we have to do everything twice now! We again use forward differences for time via `u.dt` and backward differences in space via `u.dxl` and `u.dyl` to match the original tutorial.

# In[5]:


from devito import Eq, solve

eq_u = Eq(u.dt + u*u.dxl + v*u.dyl)
eq_v = Eq(v.dt + u*v.dxl + v*v.dyl)

# We can use the same SymPy trick to generate two
# stencil expressions, one for each field update.
stencil_u = solve(eq_u, u.forward)
stencil_v = solve(eq_v, v.forward)
update_u = Eq(u.forward, stencil_u, subdomain=grid.interior)
update_v = Eq(v.forward, stencil_v, subdomain=grid.interior)

print("U update:\n%s\n" % update_u)
print("V update:\n%s\n" % update_v)


# We then set Dirichlet boundary conditions at all sides of the domain to $1$.

# In[6]:


x, y = grid.dimensions
t = grid.stepping_dim
bc_u = [Eq(u[t+1, 0, y], 1.)]  # left
bc_u += [Eq(u[t+1, nx-1, y], 1.)]  # right
bc_u += [Eq(u[t+1, x, ny-1], 1.)]  # top
bc_u += [Eq(u[t+1, x, 0], 1.)]  # bottom
bc_v = [Eq(v[t+1, 0, y], 1.)]  # left
bc_v += [Eq(v[t+1, nx-1, y], 1.)]  # right
bc_v += [Eq(v[t+1, x, ny-1], 1.)]  # top
bc_v += [Eq(v[t+1, x, 0], 1.)]  # bottom


# And finally we can put it all together to build an operator and solve our coupled problem.

# In[7]:


#NBVAL_IGNORE_OUTPUT
from devito import Operator

# Reset our data field and ICs
init_hat(field=u.data[0], dx=dx, dy=dy, value=2.)
init_hat(field=v.data[0], dx=dx, dy=dy, value=2.)

op = Operator([update_u, update_v] + bc_u + bc_v)
op(time=nt, dt=dt)

plot_field(u.data[0])


# Excellent, we have now a scalar implementation of a convection problem, but this can be written as a single vectorial equation:
# 
# $\frac{d U}{dt} + \nabla(U)U = 0$
# 
# Let's now use devito vectorial utilities and implement the vectorial equation

# In[8]:


from devito import VectorTimeFunction, grad

U = VectorTimeFunction(name='U', grid=grid)
init_hat(field=U[0].data[0], dx=dx, dy=dy, value=2.)
init_hat(field=U[1].data[0], dx=dx, dy=dy, value=2.)

plot_field(U[1].data[0])

eq_u = Eq(U.dt + grad(U)*U)


# We now have a vectorial equation. Unlike in the previous case, we do not need to play with left/right derivatives
# as the automated staggering of the vectorial function takes care of this.

# In[9]:


eq_u


# Then we set the nboundary conditions

# In[10]:


x, y = grid.dimensions
t = grid.stepping_dim
bc_u = [Eq(U[0][t+1, 0, y], 1.)]  # left
bc_u += [Eq(U[0][t+1, nx-1, y], 1.)]  # right
bc_u += [Eq(U[0][t+1, x, ny-1], 1.)]  # top
bc_u += [Eq(U[0][t+1, x, 0], 1.)]  # bottom
bc_v = [Eq(U[1][t+1, 0, y], 1.)]  # left
bc_v += [Eq(U[1][t+1, nx-1, y], 1.)]  # right
bc_v += [Eq(U[1][t+1, x, ny-1], 1.)]  # top
bc_v += [Eq(U[1][t+1, x, 0], 1.)]  # bottom


# In[11]:


# We can use the same SymPy trick to generate two
# stencil expressions, one for each field update.
stencil_U = solve(eq_u, U.forward)
update_U = Eq(U.forward, stencil_U, subdomain=grid.interior)


# And we have the updated (stencil) as a vectorial equation once again

# In[12]:


update_U


# We finally run the operator

# In[13]:


#NBVAL_IGNORE_OUTPUT
op = Operator([update_U] + bc_u + bc_v)
op(time=nt, dt=dt)

# The result is indeed the expected one.
plot_field(U[0].data[0])


# In[14]:


from devito import norm
assert np.isclose(norm(u), norm(U[0]), rtol=1e-5, atol=0)