###############################################################################
#
# Plot from VDS
#
# Author: Aiden Huffman
# Date: March 1st, 2023
# 
# Description: Plot results from the virtual data set
#
###############################################################################

import numpy as np
import matplotlib.pyplot as plt
import h5py

from scipy.interpolate import RegularGridInterpolator

file = h5py.File('simulation/VDS.h5')

# Get grid/collocation/interpolation points
x = file['x']
z = file['z']

X, Z = np.meshgrid(x, z)

# Build Uniform Grid - Can't make streamline plots without
# uniform grids and Chebyshev nodes are not uniform, we need
# to interpolate
xu = np.linspace(np.amin(x), np.amax(x), len(x))
zu = np.linspace(np.amin(z), np.amax(z), len(z))

Xu, Zu = np.meshgrid(xu, zu)

# Get our simulation data [time, x, z]
T = file['T']
ux = file['u_x']
uz = file['u_z']

plt.figure()
plt.pcolormesh(X, Z, T[-1,:,:].T-Z, cmap='RdYlBu_r')
plt.colorbar()
plt.title("Temperature Profile Relative to Steady-State")
plt.xlabel("$x$")
plt.ylabel("$z$")
plt.savefig("rb_temp_plot.png")

interp_u_x = RegularGridInterpolator((x,z), ux[-1,:,:])
interp_u_z = RegularGridInterpolator((x,z), uz[-1,:,:])

plt.figure()
plt.pcolormesh(X, Z, T[-1,:,:].T-Z, cmap='RdYlBu_r')
plt.colorbar()
plt.streamplot(Xu, Zu, interp_u_x((Xu, Zu)), interp_u_z((Xu, Zu)), color='k')
plt.title("Temperature Profile Relative to Steady-State")
plt.xlabel("$x$")
plt.ylabel("$z$")
plt.savefig("rb_temp_w_streamline.png")