import matplotlib.pyplot as plt
import numpy as np
from numpy import pi
%matplotlib inline
from pyqg import sqg_model

"""
Explanation: Surface Quasi-Geostrophic (SQG) Model
Here will will use pyqg to reproduce the results of the paper: <br /> 
I. M. Held, R. T. Pierrehumbert, S. T. Garner and K. L. Swanson (1985). Surface quasi-geostrophic dynamics. Journal of Fluid Mechanics, 282, pp 1-20 [doi:: http://dx.doi.org/10.1017/S0022112095000012)
End of explanation
"""


# create the model object
year = 1.
m = sqg_model.SQGModel(L=2.*pi,nx=512, tmax = 26.005,
        beta = 0., Nb = 1., H = 1., f_0 = 1, dt = 0.005,
                     taveint=1, twrite=400, ntd=4)
# in this example we used ntd=4, four threads
# if your machine has more (or fewer) cores available, you could try changing it

"""
Explanation: Surface quasi-geostrophy (SQG) is a relatively simple model that describes surface intensified flows due to buoyancy. One of it's advantages is that it only has two spatial dimensions but describes a three-dimensional solution.
If we define $b$ to be the buoyancy, then the evolution equation for buoyancy at each the top and bottom surface is
$$
\partial_t b + J(\psi, b) = 0.
$$
The invertibility relation between the streamfunction, $\psi$, and the buoyancy, $b$, is hydrostatic balance
$$
b = \partial_z \psi. 
$$
Using the fact that the Potential Vorticity is exactly zero in the interior of the domain and that the domain is semi-infinite, yields that the inversion in Fourier space is,
$$
\hat b = \frac{\kappa N}{f_0} \hat \psi.
$$
Held et al. studied several different cases, the first of which was the nonlinear evolution of an elliptical vortex. There are several other cases that they studied and people are welcome to adapt the code to study those as well.  But here we focus on this first example for pedagogical reasons.
End of explanation
"""


# Choose ICs from Held et al. (1995)
# case i) Elliptical vortex
x = np.linspace(m.dx/2,2*np.pi,m.nx) - np.pi
y = np.linspace(m.dy/2,2*np.pi,m.ny) - np.pi
x,y = np.meshgrid(x,y)

qi = -np.exp(-(x**2 + (4.0*y)**2)/(m.L/6.0)**2)

# initialize the model with that initial condition
m.set_q(qi[np.newaxis,:,:])

# Plot the ICs
plt.rcParams['image.cmap'] = 'RdBu'
plt.clf()
p1 = plt.imshow(m.q.squeeze() + m.beta * m.y)
plt.title('b(x,y,t=0)')
plt.colorbar()
plt.clim([-1, 0])
plt.xticks([])
plt.yticks([])
plt.show()

"""
Explanation: Initial condition
The initial condition is an elliptical vortex, 
$$
b = 0.01 \exp( - (x^2 + (4y)^2)/(L/y)^2
$$
where $L$ is the length scale of the vortex in the $x$ direction.  The amplitude is 0.01, which sets the strength and speed of the vortex.  The aspect ratio in this example is $4$ and gives rise to an instability.  If you reduce this ratio sufficiently you will find that it is stable.  Why don't you try it and see for yourself?
End of explanation
"""


for snapshot in m.run_with_snapshots(tsnapstart=0., tsnapint=400*m.dt):
    plt.clf()
    p1 = plt.imshow(m.q.squeeze() + m.beta * m.y)
    #plt.clim([-30., 30.])
    plt.title('t='+str(m.t))
    plt.colorbar()
    plt.clim([-1, 0])
    plt.xticks([])
    plt.yticks([])
    plt.show()

"""
Explanation: Runing the model
Here we demonstrate how to use the run_with_snapshots feature to periodically stop the model and perform some action (in this case, visualization).
End of explanation
"""


