# Some imports we will need below
import numpy as np
from devito import *
import matplotlib.pyplot as plt 
%matplotlib inline

nx, ny = 100, 100
grid = Grid(shape=(nx, ny))

u = TimeFunction(name='u', grid=grid, space_order=2, save=200)
c = Constant(name='c')

eqn = Eq(u.dt, c * u.laplace)

step = Eq(u.forward, solve(eqn, u.forward))

op = Operator([step])

xx, yy = np.meshgrid(np.linspace(0., 1., nx, dtype=np.float32),
                     np.linspace(0., 1., ny, dtype=np.float32))
r = (xx - .5)**2. + (yy - .5)**2.
# Inserting the ring
u.data[0, np.logical_and(.05 <= r, r <= .1)] = 1.

stats = op.apply(dt=5e-05, c=0.5)

plt.rcParams['figure.figsize'] = (20, 20)
for i in range(1, 6):
    plt.subplot(1, 6, i)
    plt.imshow(u.data[(i-1)*40])
plt.show()

# Instead of `platform=nvidiaX`, you may run your Python code with
# the environment variable `DEVITO_PLATFORM=nvidiaX`

# We also need the `gpu-fit` option to tell Devito that `u` will definitely
# fit in the GPU memory. This is necessary every time a TimeFunction with
# `save != None` is used. Otherwise, Devito could generate code such that
# `u` gets streamed between the CPU and the GPU, but for this advanced 
# feature you will need `devitopro`.
op = Operator([step], platform='nvidiaX', opt=('advanced', {'gpu-fit': u}))

# Uncomment and run only if Devito was installed with GPU support.
# stats = op.apply(dt=5e-05, c=0.5)

print(op)