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

# ### Example 5: Laplace equation
# 
# In this tutorial we will look at constructing the steady-state heat example using the Laplace equation. In contrast to the previous tutorials this example is entirely driven by the prescribed Dirichlet and Neumann boundary conditions, instead of an initial condition. We will also demonstrate how to use Devito to solve a steady-state problem without time derivatives and how to switch buffers explicitly without having to re-compile the kernel.
# 
# First, we again define our governing equation:
# $$\frac{\partial ^2 p}{\partial x^2} + \frac{\partial ^2 p}{\partial y^2} = 0$$
# 
# We are again discretizing second-order derivatives using a central difference scheme to construct a diffusion problem (see tutorial 3). This time we have no time-dependent term in our equation though, since there is no term $p_{i,j}^{n+1}$. This means that we are simply updating our field variable $p$ over and over again, until we have reached an equilibrium state. In a discretized form, after rearranging to update the central point $p_{i,j}^n$ we have
# $$p_{i,j}^n = \frac{\Delta x^2(p_{i+1,j}^n+p_{i-1,j}^n)+\Delta y^2(p_{i,j+1}^n + p_{i,j-1}^n)}{2(\Delta x^2 + \Delta y^2)}$$
# 
# And, as always, we first re-create the original implementation to see what we are aiming for. Here we initialize the field $p$ to $0$ and apply the following boundary conditions:
# 
# $p=0$ at $x=0$
#  
# $p=y$ at $x=2$
#  
# $\frac{\partial p}{\partial y}=0$ at $y=0, \ 1$
# 
# 

# In[1]:


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

# Some variable declarations
nx = 31
ny = 31
c = 1
dx = 2. / (nx - 1)
dy = 1. / (ny - 1)


# The first of these boundary conditions we have seen previously and is set by enforcing a constant value.
# The second will be implemented by creating a line of $p=y$ from $y=0$ to the edge of the domain along this boundary. The last of these boundary conditions are implemented using a numerical trick whereby we copy the second to last value in the array into the last value at each step. This numerically imposes the zero derivative boundary condition.

# In[2]:


def laplace2d(p, bc_y, dx, dy, l1norm_target):
    l1norm = 1
    pn = np.empty_like(p)

    while l1norm > l1norm_target:
        pn = p.copy()
        p[1:-1, 1:-1] = ((dx**2 * (pn[2:, 1:-1] + pn[0:-2, 1:-1]) +
                         dy**2 * (pn[1:-1, 2:] + pn[1:-1, 0:-2])) /
                        (2 * (dx**2 + dy**2)))
            
        p[0, :] = 0  # p = 0 @ x = 0
        p[-1, :] = bc_y  # p = y @ x = 2
        p[:, 0] = p[:, 1]  # dp/dy = 0 @ y = 0
        p[:, -1] = p[:, -2]  # dp/dy = 0 @ y = 1
        l1norm = (np.sum(np.abs(p[:]) - np.abs(pn[:])) /
                  np.sum(np.abs(pn[:])))
     
    return p


# In[3]:


#NBVAL_IGNORE_OUTPUT

# Our initial condition is 0 everywhere, except at the boundary
p = np.zeros((ny, nx))

# Boundary conditions
bc_right = np.linspace(0, 1, ny)
p[0, :] = 0  # p = 0 @ x = 0
p[-1, :] = bc_right  # p = y @ x = 2
p[:, 0] = p[:, 1]  # dp/dy = 0 @ y = 0
p[:, -1] = p[:, -2]  # dp/dy = 0 @ y = 1

plot_field(p, ymax=1.0, view=(30, 225))


# In[4]:


#NBVAL_IGNORE_OUTPUT

p = laplace2d(p, bc_right, dx, dy, 1e-4)
plot_field(p, ymax=1.0, view=(30, 225))


# OK, nice. Now, to re-create this example in Devito we need to look a little bit further under the hood. There are two things that make this different to the examples we covered so far:
# * We have no time dependence in the `p` field, but we still need to advance the state of p in between buffers. So, instead of using `TimeFunction` objects that provide multiple data buffers for timestepping schemes, we will use `Function` objects that have no time dimension and only allocate a single buffer according to the space dimensions. However, since we are still implementing a pseudo-timestepping loop, we will need two objects, say `p` and `pn`, to act as alternating buffers.
# * If we're using two different symbols to denote our buffers, any operator we create will only perform a single timestep. This is desired though, since we need to check a convergence criterion outside of the main stencil update to determine when we stop iterating. As a result we will need to call the operator repeatedly after instantiating it outside the convergence loop.
# 
# So, how do we make sure our operator doesn't accidentally overwrite values in the same buffer? Well, we can again let SymPy reorganise our Laplace equation based on `pn` to generate the stencil, but when we create the update expression, we set the LHS to our second buffer variable `p`.

# In[5]:


from devito import Grid, Function, Eq, solve

# Create two explicit buffers for pseudo-timestepping
grid = Grid(shape=(nx, ny), extent=(1., 2.))
p = Function(name='p', grid=grid, space_order=2)
pn = Function(name='pn', grid=grid, space_order=2)

# Create Laplace equation base on `pn`
eqn = Eq(pn.laplace, subdomain=grid.interior)
# Let SymPy solve for the central stencil point
stencil = solve(eqn, pn)
# Now we let our stencil populate our second buffer `p`
eq_stencil = Eq(p, stencil)

# In the resulting stencil `pn` is exclusively used on the RHS
# and `p` on the LHS is the grid the kernel will update
print("Update stencil:\n%s\n" % eq_stencil)


# Now we can add our boundary conditions. We have already seen how to prescribe constant Dirichlet BCs by simply setting values using the low-level notation. This time we will go a little further by setting a prescribed profile, which we create first as a custom 1D symbol and supply with the BC values. For this we need to create a `Function` object that has a different shape than our general `grid`, so instead of the grid we provide an explicit pair of dimension symbols and the according shape for the data.

# In[6]:


x, y = grid.dimensions
bc_right = Function(name='bc_right', shape=(ny, ), dimensions=(y, ))
bc_right.data[:] = np.linspace(0, 1, ny)


# Now we can create a set of expressions for the BCs again, where we want prescribed values on the right and left of our grid. For the Neumann BCs along the top and bottom boundaries we simply copy the second row from the outside into the outermost row, just as the original tutorial did. Using these expressions and our stencil update we can now create an operator.

# In[7]:


from devito import Operator

# Create boundary condition expressions
bc = [Eq(p[0, y], 0.)]  # p = 0 @ x = 0
bc += [Eq(p[nx-1, y], bc_right[y])]  # p = y @ x = 2
bc += [Eq(p[x, 0], p[x, 1])]  # dp/dy = 0 @ y = 0
bc += [Eq(p[x, ny-1], p[x, ny-2])]  # dp/dy = 0 @ y = 1

# Now we can build the operator that we need
op = Operator(expressions=[eq_stencil] + bc)


# We can now use this single-step operator repeatedly in a Python loop, where we can arbitrarily execute other code in between invocations. This allows us to update our L1 norm and check for convergence. Using our pre-compiled operator now comes down to a single function call that supplies the relevant data symbols. One thing to note is that we now do exactly the same thing as the original NumPy loop, in that we deep-copy the data between each iteration of the loop, which we will look at after this.

# In[8]:


get_ipython().run_cell_magic('time', '', "#NBVAL_IGNORE_OUTPUT\n\n# Silence the runtime performance logging\nfrom devito import configuration\nconfiguration['log-level'] = 'ERROR'\n\n# Initialise the two buffer fields\np.data[:] = 0.\np.data[-1, :] = np.linspace(0, 1, ny)\npn.data[:] = 0.\npn.data[-1, :] = np.linspace(0, 1, ny)\n\n# Visualize the initial condition\nplot_field(p.data, ymax=1.0, view=(30, 225))\n\n# Run the convergence loop with deep data copies\nl1norm_target = 1.e-4\nl1norm = 1\nwhile l1norm > l1norm_target:\n    # This call implies a deep data copy\n    pn.data[:] = p.data[:]\n    op(p=p, pn=pn)\n    \n    l1norm = (np.sum(np.abs(p.data[:]) - np.abs(pn.data[:])) /\n              np.sum(np.abs(pn.data[:])))\n\n# Visualize the converged steady-state\nplot_field(p.data, ymax=1.0, view=(30, 225))\n")


# One crucial detail about the code above is that the deep data copy between iterations will really hurt performance if we were to run this on a large grid. However, we have already seen how we can match data symbols to symbolic names when calling the pre-compiled operator, which we can now use to actually switch the roles of `pn` and `p` between iterations, eg. `op(p=pn, pn=p)`. Thus, we can implement a simple buffer-switching scheme by simply testing for odd and even time-steps, without ever having to shuffle data around. You can see below that this is leads to a reduction in runtime.

# In[9]:


get_ipython().run_cell_magic('time', '', '#NBVAL_IGNORE_OUTPUT\n\n# Initialise the two buffer fields\np.data[:] = 0.\np.data[-1, :] = np.linspace(0, 1, ny)\npn.data[:] = 0.\npn.data[-1, :] = np.linspace(0, 1, ny)\n\n# Visualize the initial condition\nplot_field(p.data, ymax=1.0, view=(30, 225))\n\n# Run the convergence loop by explicitly flipping buffers\nl1norm_target = 1.e-4\nl1norm = 1\ncounter = 0\nwhile l1norm > l1norm_target:\n    # Determine buffer order\n    if counter % 2 == 0:\n        _p = p\n        _pn = pn\n    else:\n        _p = pn\n        _pn = p\n\n    # Apply operator\n    op(p=_p, pn=_pn)\n    \n    # Compute L1 norm\n    l1norm = (np.sum(np.abs(_p.data[:]) - np.abs(_pn.data[:])) /\n              np.sum(np.abs(_pn.data[:])))\n    counter += 1\n     \nplot_field(p.data, ymax=1.0, view=(30, 225))\n')