import xarray as xr
import matplotlib.pyplot as plt

#
PV=xr.open_dataset('/projects/NS9869K/noresm/cases/BLOM_channel/postproc/PV_diffusivity.nc',chunks={'experiments':1})
K=xr.open_dataset('/projects/NS9869K/noresm/cases/BLOM_channel/postproc/Buoyancy_diffusivity.nc',chunks={'experiments':1})
# Define buoyancy diffusivity from the flux gradient relation
# but limit the selection to consider only cases when there is
# a 'finite' gradient
kT=-(K.VTflux.where(K.dTdy>1E-8)/K.dTdy.where(K.dTdy>1E-8)).mean('time')
# PV diffusivity is already calculated based on PV flux and gradient,
# but just pick the same time steps here.
kPV=-(PV.K_PV.integrate('freq_x').where(K.dTdy>1E-8)).mean('time')
# Require PV diffusivity to be positive before taking the time mean
kPV2=-(PV.K_PV.integrate('freq_x').where(-PV.K_PV.integrate('freq_x')>0)).mean('time')
# Restrict to both of the above conditions
kPV3=-(PV.K_PV.integrate('freq_x').where(K.dTdy>1E-8).where(-PV.K_PV.integrate('freq_x')>0)).mean('time')
#
plt.figure()
plt.semilogy(K.y/1E3,kT.rolling(y=3,center=True).mean().T,'C0')
plt.semilogy(K.y/1E3,kPV.rolling(y=3,center=True).mean().T,'C1')
plt.semilogy(K.y/1E3,kPV2.rolling(y=3,center=True).mean().T,'C2')
plt.semilogy(K.y/1E3,kPV3.rolling(y=3,center=True).mean().T,'C3')