#!/usr/bin/python3
"""
Make a stream function for a vector field.

We are going to cheat a bit for simplicity - make the stream function first
then derive vector field from it.
"""

import numpy as np
from matplotlib import pyplot as plt

# make a 1001 x 1001 grid
x = np.linspace(0,10,1001)
y = np.linspace(0,10,1001)[:,None]

dx = x[1]-x[0]
dy = y[1,0] - y[0,0]

extent = (x[0] - 0.5*dx, x[-1] + 0.5*dx, y[0,0] - 0.5*dy, y[-1,0] + 0.5*dy)

def vort(x0,y0,r):
	return -np.log(np.sqrt((x - x0)**2 + (y - y0)**2 + r**2))

# construct some stream function with
# overall background flow in x direction
streamfunc = 0*x + 0.15*y
# put an irrotational vortex (singular) with negative vorticity
streamfunc -= 0.50 * vort(3.5,4.5,0)
# make a spread-out vortex with positive vorticity
streamfunc += 0.40 * vort(7,6,0.5)

# Create a hole in the domain by cutting off the singular vortex
streamfunc[(streamfunc < (+0.0)) & (y>2)] = np.nan
streammin = np.nanmin(streamfunc)
streammax = np.nanmax(streamfunc)

#calculate velocities; just forward difference for laziness.
ux = np.diff(streamfunc,axis=0) / dy
uy = np.diff(streamfunc,axis=1) / -dx


Nlevels = 9
contour_levels = (np.arange(0.5,(Nlevels+1))/(Nlevels+1)) * (streammax - streammin) + streammin


fig = plt.figure(figsize=(2,4), dpi=300)

# bottom axes: 3D plot 
ax = fig.add_axes((0,0,1,0.5),projection = '3d', computed_zorder=False)
ax.view_init(30, -60, 0)
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
ax.plot_surface(x, y, streamfunc, cmap='RdYlGn',
                       linewidth=0, antialiased=False,
                       )
ax.contour(x,y[:,0],streamfunc, levels=contour_levels, colors='k', linestyles='--', linewidths=0.5, zorder=1000)

ax = plt.axes((0,0.5,0.997,0.5))
ax.set_xticks([])
ax.set_yticks([])
# draw the hole in the domain as gray
plt.imshow(np.isnan(streamfunc), origin='lower', extent=extent, cmap='gray_r', vmin=0,vmax=4)

# plot every 67th velocity arrow
slicer = slice(33,None,67)
plt.quiver(x[slicer], y[slicer,:], ux[slicer,slicer], uy[slicer,slicer], scale=8., width=0.015, color='red')
plt.contour(x,y[:,0],streamfunc, levels=contour_levels, colors='k', linestyles='--', linewidths=0.75)
fig.savefig('stream function.png')


fig = plt.figure(figsize=(2,2), dpi=300)
ax = plt.axes((0,0,0.997,1))
plt.imshow(streamfunc, extent=extent, origin='lower')
plt.contour(x,y[:,0],streamfunc, levels=contour_levels, colors='k', linestyles='--', linewidths=0.75)
fig.savefig('stream function imshow.png')