Skip to content

Instantly share code, notes, and snippets.

@andersy005

andersy005/moc_1deg_fortranless.ipynb

Forked from emaroon/moc_1deg_fortranless.ipynb
Created August 23, 2019 22:14
Show Gist options
  • Save andersy005/01f09e77408a3ccc5e16b72e70e993f9 to your computer and use it in GitHub Desktop.
Save andersy005/01f09e77408a3ccc5e16b72e70e993f9 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
moc_1deg_fortranless.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"#load modules\n",
"import xarray as xr \n",
"import glob \n",
"import numpy as np \n",
"import os \t \n",
"import time as timer\n",
"import pop_tools\n",
"from dask_jobqueue import PBSCluster,SLURMCluster\n",
"from dask.distributed import Client, LocalCluster\n",
"import dask.array as da"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Set Options\n",
"zcoord=False # True-->compute MOC(z), False-->compute MOC(sigma2)\n",
"debug=False # If zcoord=True, compute MOCdiff\n",
" # If zcoord=False, compute vertical sums\n",
"computew=False # Only applies for zcoord=True. True--> w will be computed from div(u,v)\n"
]
},
{
"cell_type": "code",
"execution_count": 62,
"metadata": {},
"outputs": [],
"source": [
"# Define input/output streams\n",
"#in_dir='/glade/u/home/yeager/analysis/python/POP_MOC/tests/data/'\n",
"#out_dir='/glade/scratch/yeager/POP_MOC/'\n",
"out_dir='/glade/scratch/emaroon/moctest/'\n",
"#in_file = [in_dir+'cesm_pop.h_g17.nc']\n",
"\n",
"in_dir='/glade/p/cgd/oce/projects/JRA55/IAF/g20a10.GIAF_JRA.gx1v7.C03/ocn/hist/'\n",
"#in_file=sorted(glob.glob(in_dir+'*.pop.h.000[1]-0[1,2].nc'))\n",
"in_file=sorted(glob.glob(in_dir+'*.pop.h.000[1]-0[1,2,3,4].nc'))\n",
"\n",
"moc_template_file = ['/glade/u/home/yeager/analysis/python/POP_MOC/moc_template.nc']\n",
"if zcoord:\n",
" out_file = out_dir+'cesm_pop.h_g17.MOCz.python.nc'\n",
" if computew:\n",
" out_file = out_dir+'cesm_pop.h_g17.MOCz.python.computew.nc'\n",
" if debug:\n",
" moc_template_file = in_file\n",
"else:\n",
" out_file = out_dir+'cesm_pop.h_g17.MOCsig2.python_first4.nc'\n",
" out_file_db = out_dir+'cesm_pop.h_g17.MOCsig2.python.debug.nc'"
]
},
{
"cell_type": "code",
"execution_count": 63,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"moc_offline already compiled\n"
]
}
],
"source": [
"# Import offline MOC routines written in fortran (compile with f2py if necessary)\n",
"f90mocroutines='./MOCoffline.POP_1deg.f90'\n",
"if not os.path.isfile('moc_offline_1deg.cpython-37m-x86_64-linux-gnu.so'): \n",
" print('MOCoffline compiling')\n",
" os.system('f2py -c '+f90mocroutines+' -m moc_offline_1deg')\n",
"else: print('moc_offline already compiled')\n",
"import moc_offline_1deg\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"#specifications for dask cluster\n",
"projectCode='NCGD0011'\n",
"dav=True\n",
"chey=False\n",
"local = False\n",
"\n",
"if dav:\n",
" ncpu=2\n",
" memory='200GB'\n",
" numNodes=1\n",
"elif chey:\n",
" ncpu=36\n",
" numNodes=1\n",
" memory='120GB' \n",
"elif local: \n",
" ncpu=1\n",
" \n",
" \n",
"if dav:\n",
" cluster = SLURMCluster(project=projectCode, cores=ncpu, memory=memory,\\\n",
" walltime='01:30:00', processes=ncpu,)\n",
"elif chey:\n",
" cluster = PBSCluster(project=projectCode, cores=ncpu, memory=memory,\\\n",
" walltime='01:30:00', processes=ncpu, queue='regular')\n",
"elif local: cluster = LocalCluster(n_workers=ncpu,threads_per_worker=1,diagnostics_port=8787)\n",
"\n",
"client = Client(cluster)\n",
"\n",
"if dav or chey:\n",
" cluster.scale(numNodes*ncpu) "
]
},
{
"cell_type": "code",
"execution_count": 64,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Worker tcp://10.12.205.27:38729 restart in Job 3198906. This can be due to memory issue.\n"
]
},
{
"data": {
"text/html": [
"<table style=\"border: 2px solid white;\">\n",
"<tr>\n",
"<td style=\"vertical-align: top; border: 0px solid white\">\n",
"<h3>Client</h3>\n",
"<ul>\n",
" <li><b>Scheduler: </b>tcp://10.12.205.11:38025\n",
" <li><b>Dashboard: </b><a href='http://10.12.205.11:8787/status' target='_blank'>http://10.12.205.11:8787/status</a>\n",
"</ul>\n",
"</td>\n",
"<td style=\"vertical-align: top; border: 0px solid white\">\n",
"<h3>Cluster</h3>\n",
"<ul>\n",
" <li><b>Workers: </b>2</li>\n",
" <li><b>Cores: </b>2</li>\n",
" <li><b>Memory: </b>200.00 GB</li>\n",
"</ul>\n",
"</td>\n",
"</tr>\n",
"</table>"
],
"text/plain": [
"<Client: scheduler='tcp://10.12.205.11:38025' processes=0 cores=0>"
]
},
"execution_count": 64,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"client.restart()\n",
"#client"
]
},
{
"cell_type": "code",
"execution_count": 65,
"metadata": {},
"outputs": [],
"source": [
"# Define MOC template\n",
"tchunk = 1\n",
"ds=xr.concat([xr.open_dataset(ff,chunks={'time':tchunk}) for ff in moc_template_file],dim='time')\n",
"ds = ds.chunk({'time':tchunk})\n",
"#ds = xr.open_dataset(moc_template_file)\n",
"#moc = ds['MOC'].isel(moc_comp=0)\n",
"moc = ds['MOC']\n",
"time = ds['time'] \n",
"lat_aux_grid = ds['lat_aux_grid'] \n",
"lat_aux_grid.encoding['_FillValue']=None # because xarray is weird\n",
"transport_regions = ds['transport_regions'].isel(time=0) \n",
"moc_components = ds['moc_components'].isel(time=0) \n",
"ncomp = moc_components.shape[0]\n",
"ntr = transport_regions.shape[0]\n",
"###FIX MONDAY\n",
"ntr = 2\n",
"\n",
"nyaux = lat_aux_grid.shape[0]"
]
},
{
"cell_type": "code",
"execution_count": 66,
"metadata": {},
"outputs": [],
"source": [
"# Open a POP history file requiring MOC\n",
"ds=xr.concat([xr.open_dataset(ff,chunks={'time':tchunk}) for ff in in_file],dim='time')\n",
"ds = ds.chunk({'time':tchunk})\n",
"#ds = xr.open_dataset(in_file)\n",
"pd = ds['PD']\n",
"pd=pd.drop(['ULAT','ULONG']) # this is a python bug that we are correcting\n",
"temp = ds['TEMP']\n",
"temp=temp.drop(['ULAT','ULONG'])\n",
"salt = ds['SALT']\n",
"salt=salt.drop(['ULAT','ULONG'])\n",
"u_e = ds['UVEL']/100\n",
"u_e=u_e.drop(['TLAT','TLONG'])\n",
"u_e.attrs['units']='m/s'\n",
"#u_i = ds['UISOP']/100\n",
"#u_i=u_i.drop(['ULAT','TLONG'])\n",
"#u_i.attrs['units']='m/s'\n",
"#u_s = ds['USUBM']/100\n",
"#u_s=u_s.drop(['ULAT','TLONG'])\n",
"#u_s.attrs['units']='m/s'\n",
"v_e = ds['VVEL']/100\n",
"v_e=v_e.drop(['TLAT','TLONG'])\n",
"v_e.attrs['units']='m/s'\n",
"#v_i = ds['VISOP']/100\n",
"#v_i=v_i.drop(['TLAT','ULONG'])\n",
"#v_i.attrs['units']='m/s'\n",
"#v_s = ds['VSUBM']/100\n",
"#v_s=v_s.drop(['TLAT','ULONG'])\n",
"#v_s.attrs['units']='m/s'\n",
"w_e = ds['WVEL']/100\n",
"w_e=w_e.drop(['ULAT','ULONG'])\n",
"w_e.attrs['units']='m/s'\n",
"#w_i = ds['WISOP']/100\n",
"#w_i=w_i.drop(['ULAT','ULONG'])\n",
"#w_i.attrs['units']='m/s'\n",
"#w_s = ds['WSUBM']/100\n",
"#w_s=w_s.drop(['ULAT','ULONG'])\n",
"#w_s.attrs['units']='m/s'\n",
"\n",
"ulat = ds['ULAT']\n",
"ulon = ds['ULONG']\n",
"tlat = ds['TLAT']\n",
"tlon = ds['TLONG']\n",
"kmt = ds['KMT'].isel(time=0)\n",
"kmu = ds['KMU'].isel(time=0)\n",
"dxu = ds['DXU'].isel(time=0)/100\n",
"dxu.attrs['units']='m'\n",
"dyu = ds['DYU'].isel(time=0)/100\n",
"dyu.attrs['units']='m'\n",
"hte = ds['HTE']/100\n",
"hte.attrs['units']='m'\n",
"htn = ds['HTN']/100\n",
"htn.attrs['units']='m'\n",
"rmask = ds['REGION_MASK'].isel(time=0)\n",
"tarea = ds['TAREA'].isel(time=0)/100/100\n",
"tarea.attrs['units']='m^2'\n",
"uarea = ds['UAREA'].isel(time=0)/100/100\n",
"uarea.attrs['units']='m^2'\n",
"time = ds['time']\n",
"time.encoding['_FillValue']=None # because xarray is weird\n",
"z_t = ds['z_t']/100\n",
"z_t.attrs['units']='m'\n",
"z_w = ds['z_w']/100\n",
"z_w.attrs['units']='m'\n",
"dz = ds['dz'].isel(time=0)/100\n",
"dz.attrs['units']='m'\n",
"dzw = ds['dzw'].isel(time=0)/100\n",
"dzw.attrs['units']='m'\n",
"hu = ds['HU'].isel(time=0)/100\n",
"hu.attrs['units']='m'\n",
"ht = ds['HT'].isel(time=0)/100\n",
"ht.attrs['units']='m'\n",
"dims = np.shape(temp)\n",
"nt = dims[0]\n",
"nz = dims[1]\n",
"ny = dims[2]\n",
"nx = dims[3]\n",
"km = int(np.max(kmt).values)\n",
"mval=pd.encoding['_FillValue']"
]
},
{
"cell_type": "code",
"execution_count": 67,
"metadata": {},
"outputs": [],
"source": [
"# Create a k-index array for masking purposes\n",
"kji = np.indices((nz,ny,nx))\n",
"kindices = kji[0,:,:,:] + 1\n",
"\n",
"coords3={k:v for k,v in temp.coords.items() if k!='time'}\n",
"kindices = xr.DataArray(kji[0,:,:,:],dims=temp.dims[1:],coords=coords3)\n",
"kindices =kindices.chunk({'nlat':ny})"
]
},
{
"cell_type": "code",
"execution_count": 68,
"metadata": {},
"outputs": [],
"source": [
"# Zero out Eulerian velocity missing values\n",
"u_e=u_e.fillna(0)\n",
"v_e=v_e.fillna(0)\n",
"w_e=w_e.fillna(0)\n"
]
},
{
"cell_type": "code",
"execution_count": 69,
"metadata": {},
"outputs": [],
"source": [
"# grid-oriented volume fluxes \n",
"ueflux_z = u_e*dyu*dz # m^3/s\n",
"veflux_z = v_e*dxu*dz # m^3/s\n",
"weflux_z = w_e*tarea # m^3/s\n",
"#uiflux_z = u_i*hte*dz # m^3/s\n",
"#viflux_z = v_i*htn*dz # m^3/s\n",
"#wiflux_z = w_i*tarea # m^3/s\n",
"#usflux_z = u_s*hte*dz # m^3/s\n",
"#vsflux_z = v_s*htn*dz # m^3/s\n",
"#wsflux_z = w_s*tarea # m^3/s"
]
},
{
"cell_type": "code",
"execution_count": 70,
"metadata": {},
"outputs": [],
"source": [
"# Define top/bottom depths of POP T-grid\n",
"z_bot=z_w.values\n",
"z_bot=z_w.values+dz.values\n",
"z_top=z_w.values"
]
},
{
"cell_type": "code",
"execution_count": 94,
"metadata": {},
"outputs": [],
"source": [
"#function to wrap fortran fluxconv and sgsfluxconv\n",
"#def fluxconv(stacked_sigma2_l,ue_l,ve_l,targ_ztop1_l,targ_zbot1_l,targ_ztop2_l,targ_zbot2_l,dxu_l,dyu_l,dz_l,tarea_l,kmt_l,ztop_l,zbot_l,vdim,sgs_l):\n",
"def fluxconv(stacked_sigma2_l,stacked_sigma2_l2,ue_l,ve_l,dxu_l,dyu_l,dz_l,tarea_l,kmt_l,ztop_l,zbot_l,vdim,sgs_l):\n",
" tmpkmt=np.transpose(kmt_l,axes=[1,0])\n",
" tmpu=np.transpose(ue_l,axes=[3,2,1,0])\n",
" tmpv=np.transpose(ve_l,axes=[3,2,1,0])\n",
" \n",
" targ_ztop1_l=np.transpose(stacked_sigma2_l[:,:,:,:,0],axes=[3,2,1,0])\n",
" targ_zbot1_l=np.transpose(stacked_sigma2_l[:,:,:,:,1],axes=[3,2,1,0])\n",
"\n",
" if sgs_l:\n",
" targ_ztop2_l=np.transpose(stacked_sigma2_l2[:,:,:,:,0],axes=[3,2,1,0])\n",
" targ_zbot2_l=np.transpose(stacked_sigma2_l2[:,:,:,:,1],axes=[3,2,1,0])\n",
" utmp,vtmp,wtmp = moc_offline_1deg.sgsfluxconv(tmpkmt,ztop_l,zbot_l,dz_l,targ_ztop1_l,targ_ztop2_l,targ_zbot1_l,\\\n",
" targ_zbot2_l,tmpu,tmpv,mval,[tchunk,nz,ny,nx,vdim])\n",
" else: \n",
" ##PUT PYTHON TRANSLATION IN HERE\n",
" \n",
" utmp,vtmp,wtmp = moc_offline_1deg.fluxconv(tmpkmt,ztop_l,zbot_l,dz_l,targ_ztop1_l,targ_zbot1_l,tmpu,tmpv,mval,\\\n",
" [tchunk,nz,ny,nx,vdim])\n",
" \n",
" stacked=np.stack([np.transpose(utmp,axes=[3,2,1,0]),np.transpose(vtmp,axes=[3,2,1,0]),np.transpose(wtmp,axes=[3,2,1,0])],axis=4)\n",
"\n",
" return stacked"
]
},
{
"cell_type": "code",
"execution_count": 95,
"metadata": {},
"outputs": [],
"source": [
"#function for removing density inversions in a map_blocks friendly way\n",
"def correct_inversions(sigma2):\n",
" sh=np.shape(sigma2)\n",
" km_l=sh[1]\n",
" \n",
" for iz in range(km_l-1,0,-1):\n",
" tmp1 = sigma2[:,iz,:,:]\n",
" tmp2 = sigma2[:,iz-1,:,:]\n",
" tmp3 = (~np.isnan(tmp1)) & (~np.isnan(tmp2)) & np.greater(tmp2,tmp1) \n",
" tmp2[tmp3] = tmp1[tmp3]-1.e-5\n",
" sigma2[:,iz-1,:,:] = tmp2\n",
" \n",
" return sigma2"
]
},
{
"cell_type": "code",
"execution_count": 96,
"metadata": {},
"outputs": [],
"source": [
"#function for wrapping fortran sig2z for map_blocks call\n",
"def sig2z(sigma2,z_top_l,z_bot_l,sig2_top_l,sig2_bot_l,kmu_l):\n",
" sigma2[np.isnan(sigma2)]=mval \n",
" tmpsig=np.transpose(sigma2,axes=[3,2,1,0])\n",
" tmpkmu=np.transpose(kmu_l,axes=[1,0])\n",
" \n",
" ##PUT PYTHON TRANSLATION IN HERE\n",
" targ_ztop,targ_zbot = moc_offline_1deg.sig2z(tmpsig,tmpkmu,z_top_l,z_bot_l,sig2_top_l,sig2_bot_l,\\\n",
" mval,[tchunk,nz,ny,nx,nsig])\n",
"\n",
" targ_ztop_t=np.transpose(targ_ztop,axes=[3,2,1,0])\n",
" targ_zbot_t=np.transpose(targ_zbot,axes=[3,2,1,0])\n",
" stacked=np.stack([targ_ztop_t,targ_zbot_t],axis=4)\n",
"\n",
" return stacked"
]
},
{
"cell_type": "code",
"execution_count": 97,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/glade/u/home/emaroon/miniconda3/lib/python3.7/site-packages/dask/array/blockwise.py:204: UserWarning: The da.atop function has moved to da.blockwise\n",
" warnings.warn(\"The da.atop function has moved to da.blockwise\")\n"
]
}
],
"source": [
"if zcoord:\n",
" # Define target vertical coordinates for MOC computation\n",
" # zcoord: use POP T-grid vertical coordinates\n",
" mocz=xr.DataArray(np.append(z_top[0],z_bot),dims=['moc_z'],attrs={'long_name':'depth from surface to top of layer','units':'m','positive':'down'})\n",
" mocz.encoding['_FillValue']=None\n",
" mocnz=nz+1\n",
"\n",
"# Define vertical volume flux \n",
" if computew:\n",
" targnz=nz\n",
" #targ_ztop = np.zeros((nx,ny,targnz,nt)) + z_top[None,None,:,None]\n",
" #targ_zbot = np.zeros((nx,ny,targnz,nt)) + z_bot[None,None,:,None]\n",
" targ_ztop = np.zeros((nt,targnz,ny,nx)) + z_top[None,:,None,None]\n",
" targ_zbot = np.zeros((nt,targnz,ny,nx)) + z_bot[None,:,None,None]\n",
" targ_ztop_tu = targ_ztop\n",
" targ_zbot_tu = targ_zbot\n",
" targ_ztop_ut = targ_ztop\n",
" targ_zbot_ut = targ_zbot\n",
" \n",
" uvwstack = da.map_blocks(fluxconv,ueflux_z.data,veflux_z.data,\\\n",
" dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
" targ_ztop1_l=targ_ztop,targ_ztop2_l=None,targ_zbot1_l=targ_zbot,targ_zbot2_l=None,\\\n",
" ztop_l=z_top,zbot_l=z_bot,vdim=targnz,sgs_l=False,\\\n",
" dtype=np.float32,chunks=(tchunk,nz,ny,nx,3),new_axis=4)\n",
" \n",
" utmp=uvwstack[:,:,:,:,0]\n",
" vtmp=uvwstack[:,:,:,:,1]\n",
" wtmp=uvwstack[:,:,:,:,2]\n",
" \n",
" ueflux=xr.DataArray(utmp,dims=['time','z_t','nlat','nlon'], \\\n",
" coords={'time':time,'z_t':z_t,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
" name='uedydz',attrs={'units':'m^3/s'})\n",
" veflux=xr.DataArray(vtmp,dims=['time','z_t','nlat','nlon'], \\\n",
" coords={'time':time,'z_t':z_t,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
" name='vedxdz',attrs={'units':'m^3/s'})\n",
" weflux=xr.DataArray(wtmp,dims=['time','z_top','nlat','nlon'], \\\n",
" coords={'time':time,'z_top':z_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
" name='wedxdy',attrs={'units':'m^3/s'}) \n",
" \n",
"# uvwstack = da.map_blocks(fluxconv,uiflux_z.data,viflux_z.data,\\\n",
"# dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
"# targ_ztop1_l=targ_ztop_ut,targ_ztop2_l=targ_ztop_tu,targ_zbot1_l=targ_zbot_ut,targ_zbot2_l=targ_zbot_tu,\\\n",
"# ztop_l=z_top,zbot_l=z_bot,vdim=targnz,sgs_l=True,\\\n",
"# dtype=np.float32,chunks=(tchunk,nz,ny,nx,3),new_axis=4)\n",
" \n",
"# utmp=uvwstack[:,:,:,:,0]\n",
"# vtmp=uvwstack[:,:,:,:,1]\n",
"# wtmp=uvwstack[:,:,:,:,2]\n",
"\n",
"# uiflux=xr.DataArray(utmp,dims=['time','z_t','nlat','nlon'], \\\n",
"# coords={'time':time,'z_t':z_t,'TLAT':(('nlat','nlon'),tlat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='uidydz',attrs={'units':'m^3/s'})\n",
"# viflux=xr.DataArray(vtmp,dims=['time','z_t','nlat','nlon'], \\\n",
"# coords={'time':time,'z_t':z_t,'ULAT':(('nlat','nlon'),ulat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='vidxdz',attrs={'units':'m^3/s'})\n",
"# wiflux=xr.DataArray(wtmp,dims=['time','z_top','nlat','nlon'], \\\n",
"# coords={'time':time,'z_top':z_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='widxdy',attrs={'units':'m^3/s'})\n",
"\n",
" \n",
"# uvwstack = da.map_blocks(fluxconv,usflux_z.data,vsflux_z.data,\\\n",
"# dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
"# targ_ztop1_l=targ_ztop_ut,targ_ztop2_l=targ_ztop_tu,targ_zbot1_l=targ_zbot_ut,targ_zbot2_l=targ_zbot_tu,\\\n",
"# ztop_l=z_top,zbot_l=z_bot,vdim=targnz,sgs_l=True,\\\n",
"# dtype=np.float32,chunks=(tchunk,nz,ny,nx,3),new_axis=4)\n",
" \n",
"# utmp=uvwstack[:,:,:,:,0]\n",
"# vtmp=uvwstack[:,:,:,:,1]\n",
"# wtmp=uvwstack[:,:,:,:,2]\n",
"\n",
"# usflux=xr.DataArray(utmp,dims=['time','z_t','nlat','nlon'], \\\n",
"# coords={'time':time,'z_t':z_t,'TLAT':(('nlat','nlon'),tlat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='uidydz',attrs={'units':'m^3/s'})\n",
"# vsflux=xr.DataArray(vtmp,dims=['time','z_t','nlat','nlon'], \\\n",
"# coords={'time':time,'z_t':z_t,'ULAT':(('nlat','nlon'),ulat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='vidxdz',attrs={'units':'m^3/s'})\n",
"# wsflux=xr.DataArray(wtmp,dims=['time','z_top','nlat','nlon'], \\\n",
"# coords={'time':time,'z_top':z_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='widxdy',attrs={'units':'m^3/s'})\n",
" \n",
" else:\n",
" ueflux=ueflux_z\n",
" veflux=veflux_z\n",
" weflux=weflux_z\n",
" #uiflux=uiflux_z\n",
" #viflux=viflux_z\n",
" #wiflux=wiflux_z\n",
" #usflux=usflux_z\n",
" #vsflux=vsflux_z\n",
" #wsflux=wsflux_z\n",
"else:\n",
" # Define target vertical coordinates for MOC computation\n",
" # not zcoord: define a set of well-chosen sigma2 levels \n",
" tmp1 = np.linspace(28,34.8,35)\n",
" tmp2 = np.linspace(35,35.9,10)\n",
" tmp3 = np.linspace(36,38.0,41)\n",
" sig2 = np.concatenate((tmp1,tmp2,tmp3))\n",
" nsig = len(sig2)\n",
" sig2_top = sig2.copy()\n",
" tmp = 0.5*(sig2[1:]-sig2[0:-1])\n",
" sig2_top[0] = sig2[0]-tmp[0]\n",
" sig2_top[1:] = sig2[0:-1]+tmp\n",
" sig2_bot = sig2.copy()\n",
" sig2_bot[-1] = sig2[-1]+tmp[-1]\n",
" sig2_bot[0:-1] = sig2[0:-1]+tmp\n",
" mocz=xr.DataArray(np.append(sig2_top[0],sig2_bot),dims=['sigma'],attrs={'long_name':'Sigma2 at top of layer','units':'kg/m^3'})\n",
" mocz.encoding['_FillValue']=None\n",
" mocnz=nsig+1\n",
"\n",
" # Compute POP sigma2 field (NOTE: this is only necessary if zcoord=False)\n",
" # using pop_tools:\n",
" depth=xr.DataArray(np.ones(np.shape(salt))*2000.,dims=salt.dims,coords=salt.coords)\n",
" sigma2=pop_tools.eos(salt=salt,temp=temp,depth=depth)\n",
" sigma2 = sigma2-1000.\n",
" \n",
" # convert to DataArray & regrid from T-grid to: U-grid and T/U-combo grids\n",
" sigma2 = xr.DataArray(sigma2.data,name='Sigma2',dims=pd.dims,coords=pd.coords)\n",
" sigma2.attrs=pd.attrs\n",
" sigma2.attrs['long_name']='Sigma referenced to 2000m'\n",
" sigma2.attrs['units']='kg/m^3'\n",
" tmp=sigma2\n",
" tmpe=tmp.roll(nlon=-1,roll_coords=False) # wraparound shift to west, without changing coords\n",
" tmpn=tmp.shift(nlat=-1) # shift to south, without changing coords\n",
" tmpne=tmpn.roll(nlon=-1,roll_coords=False) # wraparound shift to west, without changing coords\n",
" tmpall=xr.concat([tmp,tmpe,tmpn,tmpne],dim='dummy')\n",
" sigma2=tmpall.mean('dummy')\n",
" sigma2.attrs=tmp.attrs # sigma2 now on ULONG,ULAT\n",
" del tmp,tmpe,tmpn,tmpne,tmpall\n",
" \n",
" # apply U-grid mask\n",
" #mask=kindices>kmu.values[None,:,:]\n",
" #sigma2.values[mask[None,:,:,:]]=np.nan\n",
" mask=(kindices>kmu)\n",
" sigma2=sigma2.where(~mask)\n",
"\n",
" sigma2_UT = xr.DataArray(sigma2.data,name='Sigma2',dims=pd.dims,coords=u_i.coords)\n",
" sigma2_UT.attrs=pd.attrs\n",
" sigma2_UT.attrs['long_name']='Sigma referenced to 2000m'\n",
" sigma2_UT.attrs['units']='kg/m^3'\n",
" tmp=sigma2_UT\n",
" tmps=tmp.shift(nlat=1) # shift to north, without changing coords\n",
" tmpall=xr.concat([tmp,tmps],dim='dummy')\n",
" sigma2_UT=tmpall.mean('dummy')\n",
" sigma2_UT.attrs=tmp.attrs # sigma2_UT now on ULONG,TLAT\n",
" del tmp,tmps,tmpall \n",
" \n",
" sigma2_TU = xr.DataArray(sigma2.data,name='Sigma2',dims=pd.dims,coords=v_i.coords)\n",
" sigma2_TU.attrs=pd.attrs\n",
" sigma2_TU.attrs['long_name']='Sigma referenced to 2000m'\n",
" sigma2_TU.attrs['units']='kg/m^3'\n",
" tmp=sigma2_TU\n",
" tmpw=tmp.roll(nlon=1,roll_coords=False) # shift to east, without changing coords\n",
" tmpall=xr.concat([tmp,tmpw],dim='dummy')\n",
" sigma2_TU=tmpall.mean('dummy')\n",
" sigma2_TU.attrs=tmp.attrs # sigma2_TU now on TLONG,ULAT\n",
" del tmp,tmpw,tmpall\n",
" \n",
" \n",
" # remove density inversions \n",
" sigma2=sigma2.data.map_blocks(correct_inversions,dtype=np.float32)\n",
" sigma2_UT=sigma2_UT.data.map_blocks(correct_inversions,dtype=np.float32)\n",
" sigma2_TU=sigma2_TU.data.map_blocks(correct_inversions,dtype=np.float32)\n",
" \n",
" # debug test: read in sigma2 from NCL code:\n",
" #tmpds = xr.open_dataset('/glade/scratch/yeager/g.DPLE.GECOIAF.T62_g16.009.chey/g.DPLE.GECOIAF.T62_g16.009.chey.pop.h.0301-01.MOCsig2.debug.nc')\n",
" #sigma2.values=tmpds.pdu.values\n",
" # NOTE: this test shows that sigma2 differences are to blame for python/NCL discrepancies in MOC(sig2)!\n",
"\n",
" # Find sigma2 layer interface depths on U-grid\n",
" stacked_sigma2_array=sigma2.map_blocks(sig2z, z_top_l=z_t.values, z_bot_l=z_bot, sig2_top_l=sig2_top, \\\n",
" sig2_bot_l=sig2_bot,kmu_l=kmu.values, dtype=np.float32,\\\n",
" chunks=(tchunk,nsig,ny,nx,2),new_axis=4)\n",
"\n",
" \n",
" stacked_sigma2_ut_array=sigma2_UT.map_blocks(sig2z, z_top_l=z_t.values, z_bot_l=z_bot, sig2_top_l=sig2_top, \\\n",
" sig2_bot_l=sig2_bot,kmu_l=kmu.values, dtype=np.float32,\\\n",
" chunks=(tchunk,nsig,ny,nx,2),new_axis=4)\n",
" \n",
"# targ_ztop_ut=stacked_sigma2_ut_array[:,:,:,:,0]\n",
" # targ_zbot_ut=stacked_sigma2_ut_array[:,:,:,:,1]\n",
" \n",
" \n",
" stacked_sigma2_tu_array=sigma2_TU.map_blocks(sig2z, z_top_l=z_t.values, z_bot_l=z_bot, sig2_top_l=sig2_top, \\\n",
" sig2_bot_l=sig2_bot,kmu_l=kmu.values, dtype=np.float32,\\\n",
" chunks=(tchunk,nsig,ny,nx,2),new_axis=4)\n",
"\n",
" # targ_ztop_tu=stacked_sigma2_tu_array[:,:,:,:,0]\n",
" # targ_zbot_tu=stacked_sigma2_tu_array[:,:,:,:,1]\n",
"\n",
"\n",
" # Calculate horizontal & vertical volume fluxes:\n",
" \n",
" # uvwstack = da.map_blocks(fluxconv,ueflux_z.data,veflux_z.data,targ_ztop1_l=targ_ztop,targ_zbot1_l=targ_zbot,\\\n",
" # dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
" # targ_ztop2_l=None,targ_zbot2_l=None,\\\n",
" # ztop_l=z_top,zbot_l=z_bot,vdim=nsig,sgs_l=False,\\\n",
" # dtype=np.float32,chunks=(tchunk,nsig,ny,nx,3),new_axis=4)\n",
"\n",
" uvwstack = da.map_blocks(fluxconv,stacked_sigma2_array,stacked_sigma2_array,ueflux_z.data,veflux_z.data,\\\n",
" dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
" ztop_l=z_top,zbot_l=z_bot,vdim=nsig,sgs_l=False,\\\n",
" dtype=np.float32,chunks=(tchunk,nsig,ny,nx,3))\n",
" \n",
" utmp=uvwstack[:,:,:,:,0]\n",
" vtmp=uvwstack[:,:,:,:,1]\n",
" wtmp=uvwstack[:,:,:,:,2]\n",
"\n",
"\n",
" ueflux_sig=xr.DataArray(utmp,dims=['time','sigma','nlat','nlon'], \\\n",
" coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
" name='uedydz_sig',attrs={'units':'m^3/s'})\n",
" veflux_sig=xr.DataArray(vtmp,dims=['time','sigma','nlat','nlon'], \\\n",
" coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
" name='vedxdz_sig',attrs={'units':'m^3/s'})\n",
" weflux_sig=xr.DataArray(wtmp,dims=['time','sigma_top','nlat','nlon'], \\\n",
" coords={'time':time,'sigma_top':sig2_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
" name='wedxdy_sig',attrs={'units':'m^3/s'})\n",
"\n",
"\n",
"# uvwstack = da.map_blocks(fluxconv,stacked_sigma2_ut_array,stacked_sigma2_tu_array,uiflux_z.data,viflux_z.data,\\\n",
"# dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
"# ztop_l=z_top,zbot_l=z_bot,vdim=nsig,sgs_l=True,\\\n",
"# dtype=np.float32,chunks=(tchunk,nsig,ny,nx,3))\n",
"\n",
" # uvwstack.persist() \n",
"\n",
"# utmp=uvwstack[:,:,:,:,0]\n",
"# vtmp=uvwstack[:,:,:,:,1]\n",
"# wtmp=uvwstack[:,:,:,:,2]\n",
"\n",
"\n",
"# uiflux_sig=xr.DataArray(utmp,dims=['time','sigma','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='uidydz_sig',attrs={'units':'m^3/s'})\n",
"# viflux_sig=xr.DataArray(vtmp,dims=['time','sigma','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='vidxdz_sig',attrs={'units':'m^3/s'})\n",
"# wiflux_sig=xr.DataArray(wtmp,dims=['time','sigma_top','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma_top':sig2_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='widxdy_sig',attrs={'units':'m^3/s'})\n",
" \n",
"# uvwstack = da.map_blocks(fluxconv,stacked_sigma2_ut_array,stacked_sigma2_tu_array,usflux_z.data,vsflux_z.data,\\\n",
"# dxu_l=dxu.values,dyu_l=dyu.values,tarea_l=tarea.values,kmt_l=kmt.data,dz_l=dz.values,\\\n",
"# ztop_l=z_top,zbot_l=z_bot,vdim=nsig,sgs_l=True,\\\n",
"# dtype=np.float32,chunks=(tchunk,nsig,ny,nx,3))\n",
"\n",
"\n",
" # uvwstack.persist() \n",
"\n",
"# utmp=uvwstack[:,:,:,:,0]\n",
"# vtmp=uvwstack[:,:,:,:,1]\n",
"# wtmp=uvwstack[:,:,:,:,2]\n",
" \n",
"# usflux_sig=xr.DataArray(utmp,dims=['time','sigma','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='usdydz_sig',attrs={'units':'m^3/s'})\n",
"# vsflux_sig=xr.DataArray(vtmp,dims=['time','sigma','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma':sig2,'ULAT':(('nlat','nlon'),ulat),'ULONG':(('nlat','nlon'),ulon)}, \\\n",
"# name='vsdxdz_sig',attrs={'units':'m^3/s'})\n",
"# wsflux_sig=xr.DataArray(wtmp,dims=['time','sigma_top','nlat','nlon'], \\\n",
"# coords={'time':time,'sigma_top':sig2_top,'TLAT':(('nlat','nlon'),tlat),'TLONG':(('nlat','nlon'),tlon)}, \\\n",
"# name='wsdxdy_sig',attrs={'units':'m^3/s'})\n",
"\n",
" ueflux=ueflux_sig\n",
" veflux=veflux_sig\n",
" weflux=weflux_sig\n",
"# uiflux=uiflux_sig\n",
"# viflux=viflux_sig\n",
"# wiflux=wiflux_sig\n",
"# usflux=usflux_sig\n",
"# vsflux=vsflux_sig\n",
"# wsflux=wsflux_sig\n",
" \n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 98,
"metadata": {},
"outputs": [],
"source": [
"#wrapping function for fortran zonal sum and AMOC to work with map_blocks\n",
"def wzonalsum(wflux_l,amoc_s_l,rmlak_l,tlat_l,lat_aux_grid_l):\n",
" \n",
"\n",
" \n",
" tmptlat = np.transpose(tlat.values,axes=[1,0])\n",
" tmpw = np.transpose(np.where(~np.isnan(wflux_l),wflux_l,mval),axes=[3,2,1,0]) \n",
" ##PUT PYTHON TRANSLATION IN HERE\n",
" tmpmoc = moc_offline_1deg.wzonalsum(tmptlat,lat_aux_grid_l,rmlak_l,tmpw,mval,[nyaux,nx,ny,nz,tchunk,ntr])\n",
" \n",
" # b. integrate in meridional direction\n",
" tmpmoc[tmpmoc==mval]=np.nan\n",
" mocsum=np.nancumsum(np.transpose(tmpmoc,axes=[3,2,1,0]),axis=3)\n",
" mocsum = mocsum*1.0e-6 \n",
" moc_l = np.transpose(mocsum,axes=[0,2,3,1])\n",
" \n",
" moc_l[:,:,:,1] = moc_l[:,:,:,1] + amoc_s_l[:,:,None]\n",
" moc_l = np.single(np.append(moc_l,np.zeros((tchunk,1,nyaux,ntr)),axis=1))\n",
" \n",
" return moc_l\n",
" \n"
]
},
{
"cell_type": "code",
"execution_count": 99,
"metadata": {},
"outputs": [],
"source": [
"# a. integrate w_sigma in zonal direction\n",
"rmlak = np.zeros((nx,ny,2),dtype=np.int)\n",
"tmprmask = np.transpose(rmask.values,axes=[1,0])\n",
"rmlak[:,:,0] = np.where(tmprmask>0,1,0)\n",
"rmlak[:,:,1] = np.where((tmprmask>=6) & (tmprmask<=11),1,0)\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 100,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"lat_aux_atl_start= 84\n"
]
}
],
"source": [
"lat_aux_atl_start = ny\n",
"for n in range(1,ny):\n",
" section = (rmlak[:,n-1,1] == 1)\n",
" if (section.any()): \n",
" lat_aux_atl_start = n-2\n",
" break\n",
"print(\"lat_aux_atl_start= \",lat_aux_atl_start)"
]
},
{
"cell_type": "code",
"execution_count": 101,
"metadata": {},
"outputs": [],
"source": [
"# Compute MOC and AMOC\n",
"# a. zonal integral of Atlantic points\n",
"veflux = veflux.where(veflux<1e30)\n",
"#vsflux = vsflux.where(vsflux<1e30)\n",
"#viflux = viflux.where(viflux<1e30)\n",
"atlmask=xr.DataArray(np.where(rmask==6,1,0),dims=['nlat','nlon'])\n",
"atlmask=atlmask.roll(nlat=-1,roll_coords=False)\n",
"veflux_xint=veflux.where(atlmask==1).sum(dim='nlon')\n",
"#viflux_xint=viflux.where(atlmask==1).sum(dim='nlon')\n",
"#vsflux_xint=vsflux.where(atlmask==1).sum(dim='nlon')\n",
"if zcoord: sumdim='z_t'\n",
"else: sumdim='sigma'\n",
"amoc_s_e=-veflux_xint.sel(nlat=lat_aux_atl_start).sortby(sumdim,ascending=False).cumsum(dim=sumdim)\n",
"#amoc_s_i=-viflux_xint.sel(nlat=lat_aux_atl_start).sortby(sumdim,ascending=False).cumsum(dim=sumdim)\n",
"#amoc_s_s=-vsflux_xint.sel(nlat=lat_aux_atl_start).sortby(sumdim,ascending=False).cumsum(dim=sumdim)\n",
"\n",
"amoc_s_e = amoc_s_e.sortby(sumdim,ascending=True)*1.e-6\n",
"#amoc_s_i = amoc_s_i.sortby(sumdim,ascending=True)*1.e-6\n",
"#amoc_s_s = amoc_s_s.sortby(sumdim,ascending=True)*1.e-6\n"
]
},
{
"cell_type": "code",
"execution_count": 102,
"metadata": {},
"outputs": [],
"source": [
"# b. integrate wflux in zonal direction (zonal sum of wflux already in m^3/s)\n",
"\n",
"tmpmoc_e=da.map_blocks(wzonalsum,weflux.data,amoc_s_e.data,rmlak_l=rmlak,tlat_l=tlat.values,lat_aux_grid_l=lat_aux_grid.values,\\\n",
" dtype=np.float32,chunks=(tchunk,mocnz,nyaux,ntr),drop_axis=[1,3],new_axis=[1,3])\n",
"\n",
"#tmpmoc_i=da.map_blocks(wzonalsum,wiflux.data,amoc_s_i.data,rmlak_l=rmlak,tlat_l=tlat.values,lat_aux_grid_l=lat_aux_grid.values,\\\n",
"# dtype=np.float32,chunks=(tchunk,mocnz,nyaux,ntr),drop_axis=[1,3],new_axis=[1,3])\n",
"\n",
"#tmpmoc_s=da.map_blocks(wzonalsum,wsflux.data,amoc_s_s.data,rmlak_l=rmlak,tlat_l=tlat.values,lat_aux_grid_l=lat_aux_grid.values,\\\n",
"# dtype=np.float32,chunks=(tchunk,mocnz,nyaux,ntr),drop_axis=[1,3],new_axis=[1,3])\n"
]
},
{
"cell_type": "code",
"execution_count": 103,
"metadata": {},
"outputs": [],
"source": [
"# c. integrate in meridional direction\n",
"if zcoord:\n",
" outdims=['time','transport_reg','moc_comp','moc_z','lat_aux_grid']\n",
" outcoords={'time':time,'transport_regions':transport_regions,'moc_comp':moc_components,'moc_z':mocz,'lat_aux_grid':lat_aux_grid}\n",
"else:\n",
" outdims=['time','transport_reg','moc_comp','sigma','lat_aux_grid'] \n",
" outcoords={'time':time,'transport_regions':transport_regions,'moc_comp':moc_components,'sigma':mocz,'lat_aux_grid':lat_aux_grid}\n",
" \n",
"#stacked=da.stack([tmpmoc_e,tmpmoc_i,tmpmoc_s],axis=4)\n",
"#transposed=da.transpose(stacked,axes=[0,3,4,1,2])\n",
"transposed=da.transpose(tmpmoc_e,axes=[0,3,1,2])\n",
"\n",
"mocout=xr.DataArray(transposed, dims=outdims,coords=outcoords,name='MOC')"
]
},
{
"cell_type": "code",
"execution_count": 104,
"metadata": {},
"outputs": [],
"source": [
"mocout = mocout.where(mocout<mval)\n",
"mocout.attrs={'units':'Sverdrups','long_name':'Meridional Overturning Circulation'}\n",
"mocout.encoding['_FillValue']=mval"
]
},
{
"cell_type": "code",
"execution_count": 105,
"metadata": {},
"outputs": [],
"source": [
"#8. Write output to netcdf\n",
"out_ds=mocout.to_dataset(name='MOC')\n",
"if zcoord and debug:\n",
" mocdiff=mocout.copy()\n",
" moc['moc_z']=mocout['moc_z']\n",
" mocdiff = mocout - moc\n",
" out_ds['MOCdiff']=mocdiff\n",
"out_ds.to_netcdf(out_file)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def sig2z_fort(nt,nz,ny,nx,nsig,popsig,kmu,z_t,zbot,slevst,slevsb,zoutt,zoutb,mval):\n",
"\n",
" zoutt=np.ones((nx,ny,nsig,nt))*mval\n",
" zoutb=np.ones((nx,ny,nsig,nt))*mval\n",
"\n",
" for iy in range(ny):\n",
" for ix in range(nx):\n",
" if popsig[ix,iy,0,0]!=mval:\n",
" for it in range(nt):\n",
" popsig1d=popsig[ix,iy,:,it]\n",
" kmax=kmu[ix,iy]\n",
" hu=zbot[kmax]\n",
" sigmin=popsig1d[0]\n",
" sigmax=popsig1d[kmax]\n",
"\n",
" \n",
" for iss in range(nsig):\n",
" if (slevst[iss]>=sigmin) and slevst[iss]<=sigmax:\n",
" for iz in range(1,kmax):\n",
" if (slevst[iss]>popsig1d[iz-1]) and (slevst[iss]<popsig1d[iz]):\n",
" dsig=slevst[iss]-popsig1d[iz-1]\n",
" dzdzsig=(z_t[iz]-z_t[iz-1])/(popsig1d[iz]-popsig1d[iz-1])\n",
" zoutt[ix,iy,iss,it]=z_t[iz-1]+dzdsig*dsig\n",
" if slevst[iss]==sigmax:\n",
" zoutt[ix,iy,iss,it]=hu\n",
" \n",
" \n",
" if (slevsb[iss]>=sigmin) and (slevsb[iss])<=sigmax:\n",
" for iz in range(1,kmax):\n",
" if (slevsb[iss]>popsig1d[iz-1]) and (slevsb[iss]<popsig1d[iz]):\n",
" dsig=slevsb[iss]-popsig1d[iz-1]\n",
" dzdsig=(z_1[iz]-z_t[iz-1])/(popsig1d[iz]-popsig1d[iz-1])\n",
" zoutb[ix,iy,iss,it] = z_t[iz-1]+dzdsig*dsig\n",
" if (slevsb[iss]==sigmax):\n",
" zoutb[ix,iy,iss,it]=hu\n",
" \n",
"\n",
" if (zoutb[ix,iy,iss,it]!=mval) and (zoutt[ix,iy,iss,itz]==mval):\n",
" zoutt[ix,iy,iss,it]=0.0\n",
" if (zoutt[ix,iy,iss,it]!=mval) and (zoutb[ix,iy,iss,it]==mval):\n",
" zoutb[ix,iy,iss,it]=hu\n",
" \n",
" \n",
" if (slevst[iss]<=sigmin) and (slevsb[iss]>=sigmax):\n",
" zoutt[ix,iy,iss,it]=0.0\n",
" zoutb[ix,iy,iss,it]=hu\n",
" \n",
" if (slevst[0]>sigmax):\n",
" zoutt[ix,iy,0,it]=0.0\n",
" zoutt[ix,iy,1:nsig,it]=mval\n",
" zoutb[ix,iy,0,it]=hu\n",
" zoutb[ix,iy,1:nsig,it]=mval\n",
" if (slevsb[nsig]<sigmin):\n",
" zoutt[ix,iy,nsig,it]=0.0\n",
" zoutt[ix,iy,0:(nsig-1),it]=mval\n",
" zoutb[ix,iy,nsig,it]=hu\n",
" zoutb[ix,iy,0:(nsig-1),it]=mval\n",
"\n",
"\n",
" if (zoutb[ix,iy,0,it]!=mval):\n",
" zoutt[ix,iy,0,it]=0.0\n",
" if (zoutt[ix,iy,nsig,it]!=mval):\n",
" zoutb[ix,iy,nsig,it]=hu\n",
" \n",
" anyzeros=np.sum(zoutt[ix,iy,:,it]==0.0)\n",
" anyhus=np.sum(zoutb[ix,iy,:,it]==hu)\n",
" if (anyzeros==0) or (anyhus==0):\n",
" print('sig2z Error at i,j,hu',ix,iy,hu)\n",
" if (ix==280) and (iy==331):\n",
" print('zoutt=',zoutt[ix,iy,:,it])\n",
" print('zoutb=',zoutb[ix,iy,:,it])\n",
" else: \n",
" print('land at i,j =',ix,iy)\n",
" \n",
" \n",
" return zoutt, zoutb"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"def fluxconv_fort(nt,nz,ny,nx,targnz,kmt,ztop,zbot,dz,targztop,targzbot,workuNE,workvNE):\n",
" \n",
" ixx = np.zeros(4,dtype='int')\n",
" iyy = np.zeros(4,dtype='int')\n",
" tmpu=np.zeros((2,targnz))\n",
" tmpv=np.zeros((2,targnz))\n",
" wgt=np.zeros(targnz)\n",
" \n",
" for it in range(nt):\n",
" for iy in range(1,ny):\n",
" for ix in range(nx):\n",
" if kmt[ix,iy]!=0:\n",
" if kmt[ix,iy]>0:\n",
" ixx[0]=ix\n",
" ixx[3]=ix\n",
" if ix!=1: \n",
" ixx[1]=ix-1\n",
" ixx[2]=ix-1\n",
" else:\n",
" ixx[1]=nx\n",
" ixx[2]=nx\n",
" iyy[0]=iy\n",
" iyy[1]=iy\n",
" iyy[2]=iy-1\n",
" iyy[3]=iy-1\n",
"\n",
" for ij in range(4):\n",
" tmpu[ij,:]=mval\n",
" tmpv[ij,:]=mval\n",
" for iz in range(nz):\n",
" zb=zbot[iz]\n",
" zt=zbot[iz]\n",
" if (iz<=kmt[ix,iy]):\n",
" udydz=workuNE[ixx[ij],iyy[ij],iz,it]\n",
" vdxdz=workvNE[ixx[ij],iyy[ij],iz,it]\n",
" else:\n",
" udydz=0.\n",
" vdxdz=0.\n",
" for iz2 in range(targnz):\n",
" zbsig=targzbot[ixx[ij],iyy[ij],iz2,it]\n",
" ztsig=targztop[ixx[ij],iyy[ij],iz2,it]\n",
" wgt[iz2]=0.0\n",
" if zbsig<1.e15:\n",
" dzsig=zbsig-ztsig\n",
"\n",
" \n",
" if (zb<=zbsig) and (zt>=ztsig):\n",
" wgt[iz2]=1.0\n",
" if (zb>=ztsig) and (zt<ztsig):\n",
" wgt[iz2]=(zb-ztsig)/dz[iz]\n",
" if (zt<=zbsig) and (zb>zbsig):\n",
" wgt[iz2]=(zbsig-zt)/dz[iz]\n",
" if (zt<ztsig) and (zb>zbsig):\n",
" wgt[iz2]=dzsig/dz[iz]\n",
" \n",
"\n",
" sumwgt=np.sum(wgt)\n",
" if (sumwgt!=0.0): \n",
" for iz2 in range(targnz):\n",
" if wgt[iz2]!=0:\n",
" wgt[iz2]=wgt[iz2]/sumwgt\n",
" if tmpu[ij,iz2]==mval:\n",
" tmpu[ij,iz2]=wgt[iz2]*udydz\n",
" else:\n",
" tmpu[ij,iz2]=tmpu[ij,iz2]+wgt[iz2]*udydz\n",
" if tmpv[ij,iz2]==mval:\n",
" tmpv[ij,iz2]=wgt[iz2]*vdxdz\n",
" else:\n",
" tmpv[ij,iz2]=tmpv[ij,iz2]+wgt[iz2]*vdxdz\n",
"\n",
" for iz2 in range(targnz):\n",
" uout[ix,iy,iz2,it]=tmpu[0,iz2]\n",
" vout[ix,iy,iz2,it]=tmpv[0,iz2]\n",
" if (np.sum(tmpv[:,iz2])==mval*4) and (np.sum(tmpu[:,iz2])==mval*4):\n",
" wout[ix,iy,iz2,it]=0.0\n",
" else:\n",
" tmpu_iz=tmpu[:,iz2]\n",
" tmpu_iz[tmpu_iz==mval]=0.0\n",
" tmpv_iz=tmpv[:,iz2]\n",
" tmpv_iz[tmpv_iz==mval]=0.0\n",
" fueE=0.5*(tmpu[0,iz2]+tmpu[3,iz2])\n",
" fueW=0.5*(tmpu[1,iz2]+tmpu[2,iz2])\n",
" fvnN=0.5*(tmpv[0,iz2]+tmpv[1,iz2])\n",
" fvnS=0.5*(tmpv[2,iz2]+tmpv[3,iz2])\n",
" wout[ix,iy,iz2,it] = -(fueE-fueW+fvnN-fvnS)\n",
" \n",
" for iz2 in range(targnz-1):\n",
" wout[ix,iy,targnz-iz2,it]=wout[ix,iy,targnz-iz2,it]+wout[ix,iy,targnz-iz2+1,it] \n",
" \n",
" return uout, vout, wout"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def wzonalsum_fort(nyaux,nx,ny,nz,nt,ntr,tlat,lat_aux_grid, rmlak,wflux,wmsg):\n",
"\n",
" \n",
" for ir in range(ntr):\n",
" for iyy in range(1,nyaux):\n",
" for iy in range(ny):\n",
" for ix in range(nx):\n",
" wflux_tmp=wflux[ix,iy,:,:]\n",
" wflux_tmp=wflux_tmp.where(tlat[ix,iy]>=lataux_grid[iyy-1] and tlat[ix,iy]<lat_aux_grid[iyy] and rmlak[ix,iy,ir]==1)\n",
" wfluxzonsum[:,:,ir,:]=np.sum(wflux_tmp,axis=2)\n",
" \n",
" \n",
" return wfluxzonalsum"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 106,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 65,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 79,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.3"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
@andersy005
Copy link
Author

  • zcoord=True and compute_w=True ----> wzonal and fluxconv
  • zcoord=True and compute_w=False ---> wzonal
  • else: wzonal, fluxconv, sig2z

@andersy005
Copy link
Author 

#in_dir='/glade/u/home/yeager/analysis/python/POP_MOC/tests/data/'
#in_file = [in_dir+'cesm_pop.h_g17.nc']

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment