<jupyter_start><jupyter_text>Movies for vertical Flux<jupyter_code>#KRM

import numpy as np

import matplotlib.pyplot as plt

import matplotlib.colors as mcolors

import matplotlib as mpl
#%matplotlib inline
from math import *

import scipy.io

import scipy as spy

from netCDF4 import Dataset

import pylab as pl

import os 

import sys

import seaborn as sns
lib_path = os.path.abspath('../../Building_canyon/BuildCanyon/PythonModulesMITgcm') # Add absolute path to my python scripts
sys.path.append(lib_path)

import ReadOutTools_MITgcm as rout 

# Functions
def vertTracBox(Mask, W,zlim=30,ylim=230):
    mask_expand = np.expand_dims(Mask[zlim,ylim:,:],0)
    mask_expand = mask_expand + np.zeros(W[:,zlim,ylim:,:].shape)

    vert = W[:,zlim,ylim:,:]
    VertTrans= np.ma.masked_array(vert,mask = mask_expand)
    return VertTrans

# Functions
def vertTracBoxDom(Mask, W,zlim=30):
    mask_expand = np.expand_dims(Mask[zlim,:,:],0)
    mask_expand = mask_expand + np.zeros(W[:,zlim,:,:].shape)

    vert = W[:,zlim,:,:]
    VertTrans= np.ma.masked_array(vert,mask = mask_expand)
    return VertTrans


NoCCanyonGrid='/ocean/kramosmu/MITgcm/TracerExperiments/3DDIFF/run06/gridGlob.nc'
NoCCanyonGridOut = Dataset(NoCCanyonGrid)

NoCCanyonState='/ocean/kramosmu/MITgcm/TracerExperiments/3DDIFF/run06/stateGlob.nc'
NoCCanyonStateOut = Dataset(NoCCanyonState)

FluxTR01NoC = '/ocean/kramosmu/MITgcm/TracerExperiments/3DDIFF/run06/FluxTR01Glob.nc'
NoCFluxOut = Dataset(FluxTR01NoC)

CanyonGrid='/ocean/kramosmu/MITgcm/TracerExperiments/CNTDIFF/run03/gridGlob.nc' # this has a canyon but calling it no canyon to use old code
CanyonGridOut = Dataset(CanyonGrid)

CanyonState='/ocean/kramosmu/MITgcm/TracerExperiments/CNTDIFF/run03/stateGlob.nc'
CanyonStateOut = Dataset(CanyonState)

FluxTR01 = '/ocean/kramosmu/MITgcm/TracerExperiments/CNTDIFF/run03/FluxTR01Glob.nc'
FluxOut = Dataset(FluxTR01)


#for dimobj in CanyonStateOut.variables.values():
#    print dimobj
# General input

nx = 360
ny = 360
nz = 90
nt = 19 # t dimension size 

z = CanyonStateOut.variables['Z']
Time = CanyonStateOut.variables['T']

xc = rout.getField(CanyonGrid, 'XC') # x coords tracer cells
yc = rout.getField(CanyonGrid, 'YC') # y coords tracer cells

bathy = rout.getField(CanyonGrid, 'Depth')

hFacC = rout.getField(CanyonGrid, 'HFacC')
MaskC = rout.getMask(CanyonGrid, 'HFacC')
MaskCNoC = rout.getMask(NoCCanyonGrid, 'HFacC')



WTRAC = rout.getField(FluxTR01,'ADVrTr01') # 
WTRACNoC = rout.getField(FluxTR01NoC,'ADVrTr01') # 

#Vertbase = vertTracAlong(MaskC,VTRAC,zlim=30,ylim=230)*1000.0
#VertbaseNoC = vertTracAlong(MaskCNoC,VTRACNoC,zlim=30,ylim=230)*1000.0
#Anom = (Vertbase-VertbaseNoC)

Vertbase = vertTracBoxDom(MaskC, WTRAC,zlim=30)*1000.0
VertbaseNoC = vertTracBoxDom(MaskCNoC, WTRACNoC,zlim=30)*1000.0
Anom = (Vertbase-VertbaseNoC)

minT = Vertbase.min()
maxT = Vertbase.max()

minTNoC = VertbaseNoC.min()
maxTNoC = VertbaseNoC.max()

minTAnom = Anom.min()
maxTAnom = Anom.max()

print(minT, maxT)
print(minTNoC, maxTNoC)
print(minTAnom,maxTAnom)

import matplotlib.animation as animation
sns.set()
sns.set_style('white')
sns.set_context("talk")


#divmap = sns.diverging_palette(255, 100, l=60, n=7, center="dark", as_cmap=True)
def Plot1(t,ax1):
    ax1.clear()
    csU = np.linspace(minT,-minT,num=31)
    Base = ax1.contourf(xc, yc,Vertbase[t,:,:],csU,cmap='RdYlGn')
    
    if t == 1:    
        cbar=plt.colorbar(Base,ax=ax1,ticks=[np.arange(minT,-minT,100000)])
        cbar.set_label(' mol/s')
        
    #CS = ax1.contour(yc[100:-1,200],z[:58],Uplot[:58,100:]/Umax,csU2,colors='k',linewidths=[0.75] )
    ax1.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))
    ax1.set_xlabel('Along-shore distance [km]')
    ax1.set_xticks([10000,20000,30000,40000,50000,60000,70000,80000,90000,100000,110000,120000])
    ax1.set_xticklabels(['10','20','30','40', '50', '60', '70', '80','90','100','110','120'])
    ax1.set_ylabel('Cross-shore distance [km]')
    ax1.set_yticks([10000,20000,30000,40000,50000,60000,70000,80000,90000])
    ax1.set_yticklabels(['10','20','30','40', '50', '60', '70', '80','90'])
    ax1.set_title('Base case - %.1f  m Vertical flux at day %0.1f' %(t/2.0+0.5, z[30]))

def Plot2(t,ax2):
    ax2.clear()
    csU = np.linspace(minTNoC,-minTNoC,num=31)
    Base = ax2.contourf(xc, yc,VertbaseNoC[t,:,:],csU,cmap='RdYlGn')
    
    if t == 1:    
        cbar=plt.colorbar(Base,ax=ax2,ticks=[np.arange(minTNoC,-minTNoC,100000)])
        cbar.set_label('mol/s')
        
    #CS = ax1.contour(yc[100:-1,200],z[:58],Uplot[:58,100:]/Umax,csU2,colors='k',linewidths=[0.75] )
    ax2.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))
    ax2.set_xlabel('Along-shore distance [km]')
    ax2.set_xticks([10000,20000,30000,40000,50000,60000,70000,80000,90000,100000,110000,120000])
    ax2.set_xticklabels(['10','20','30','40', '50', '60', '70', '80','90','100','110','120'])
    ax2.set_ylabel('Cross-shore distance [km]')
    ax2.set_yticks([10000,20000,30000,40000,50000,60000,70000,80000,90000])
    ax2.set_yticklabels(['10','20','30','40', '50', '60', '70', '80','90'])
    ax2.set_title('Enhanced $K_v$ ($10^{-3}m^2s^{-1}$) in within canyon')

def Plot3(t,ax3):
    ax3.clear()
    csU = np.linspace(-maxTAnom,maxTAnom,num=31)
    Base = ax3.contourf(xc, yc,Anom[t,:,:],csU,cmap='RdYlBu')
    
    if t == 1:    
        cbar=plt.colorbar(Base,ax=ax3,ticks=[np.arange(minTAnom,-minTAnom,10000)])
        cbar.set_label('mol/s')
        
    #CS = ax3.contour(yc[100:-1,200],z[:58],Uplot[:58,100:]/Umax,csU2,colors='k',linewidths=[0.75] )
    ax3.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))
    ax3.set_xlabel('Along-shore distance [km]')
    ax3.set_xticks([10000,20000,30000,40000,50000,60000,70000,80000,90000,100000,110000,120000])
    ax3.set_xticklabels(['10','20','30','40', '50', '60', '70', '80','90','100','110','120'])
    
    ax3.set_ylabel('Cross-shore distance [km]')
    ax3.set_yticks([10000,20000,30000,40000,50000,60000,70000,80000,90000])
    ax3.set_yticklabels(['10','20','30','40', '50', '60', '70', '80','90'])
    
    ax3.set_title('Difference')

def Plot4(t,ax4):
    ax4.clear()
    csU = np.linspace(minTNoC,-minTNoC,num=31)
    Base = ax4.contourf(xc[100:-100,100:-100], yc[100:-100,100:-100],VertbaseNoC[t,100:-100,100:-100],csU,cmap='RdYlGn')
    line = ax4.plot(xc[230,100:-100], yc[230,100:-100],'k',linewidth=2)
    if t == 1:    
        cbar=plt.colorbar(Base,ax=ax4,ticks=[np.arange(minTNoC,-minTNoC,75000)])
        cbar.set_label('mol/s')
        
    #CS = ax3.contour(yc[100:-1,200],z[:58],Uplot[:58,100:]/Umax,csU2,colors='k',linewidths=[0.75] )
    ax4.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))
    ax4.set_xlabel('Along-shore distance [km]')
    ax4.set_xticks([45000,55000,65000,75000])
    ax4.set_xticklabels(['45', '55', '65', '75'])
    
    ax4.set_ylabel('Cross-shore distance [km]')
    ax4.set_yticks([40000,50000,60000])
    ax4.set_yticklabels(['40', '50', '60'])
    
    ax4.set_title('Zoom enhanced $K_v$')



## Animation
#N=5
xslice=180
yslice=235
zslice= 29 # shelf break index
zslice2= 23 

#Empty figures
fig,((ax1,ax2),(ax3,ax4))  = plt.subplots(2, 2,figsize=(30,30))

#Initial image
def init():
    Plot1(0,ax1)
    Plot2(0,ax2)
    Plot3(0,ax3)
    Plot4(0,ax4)
    
    plt.tight_layout()
    #return[ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9]
    

def animate(tt):
    Plot1(tt,ax1)
    Plot2(tt,ax2)
    Plot3(tt,ax3)
    Plot4(tt,ax4)
    plt.tight_layout()
#The animation function (max frames=47)   


anim = animation.FuncAnimation(fig, animate, init_func=init,frames=18, interval = 200,blit=False, repeat=False)

##A line that makes it all work
mywriter = animation.FFMpegWriter()

##Save in current folder
anim.save('FluxVert_Z30_Tr01_Base_and_3DDiffRun06.mp4',writer=mywriter,fps=0.2, dpi = 300, bitrate = 1000000, codec = "libx264")

plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1828556936342886, 'fraction_numerical': 0.12718742074562517, 'mean_word_length': 4.993161094224924, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '487864', 'license': 'Apache-2.0', 'n_tokens_mistral': 4068, 'n_tokens_neox': 3449, 'n_words': 456, 'path': 'MoviesNotebooks/MoviesVerticalFlux.ipynb', 'repo_name': 'UBC-MOAD/outputanalysisnotebooks', 'size': '7886'}