pgrid

more progress
parkermac · Sep 7, 2021 · 4630812 · 4630812
1 parent a3f03e7
commit 4630812
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 89 deletions.
diff --git a/lo_tools/lo_tools/zfun.py b/lo_tools/lo_tools/zfun.py
@@ -12,24 +12,19 @@ def interp2(x, y, X, Y, U):
     Interpolate field U(X,Y) to u(x,y).  All grids are required to be plaid and 2D
     """
     if is_plaid(x) and is_plaid(y) and is_plaid(Y) and is_plaid(Y):
-        if False:
-        # slow version
-            # slow version
-            uu = interp_scattered_on_plaid(x, y, X[0,:], Y[:,0], U)
-            u = np.reshape(uu, x.shape)
-        else:
-            # Much Faster
-            xi0, xi1, xf = get_interpolant(x[0,:], X[0,:], extrap_nan=True)
-            yi0, yi1, yf = get_interpolant(y[:,0], Y[:,0], extrap_nan=True)
-            NR, NC = x.shape
-            XF = xf.reshape((1,NC)) * np.ones((NR,1))
-            YF = yf.reshape((NR,1)) * np.ones((1,NC))
-            # bi linear interpolation
-            u00 = U[yi0,:][:,xi0]
-            u10 = U[yi1,:][:,xi0]
-            u01 = U[yi0,:][:,xi1]
-            u11 = U[yi1,:][:,xi1]
-            u = (1-YF)*((1-XF)*u00 + XF*u01) + YF*((1-XF)*u10 + XF*u11)
+        # Much Faster than interp_scatttered_on_plaid()
+        xi0, xi1, xf = get_interpolant(x[0,:], X[0,:], extrap_nan=True)
+        yi0, yi1, yf = get_interpolant(y[:,0], Y[:,0], extrap_nan=True)
+        NR, NC = x.shape
+        XF, YF = np.meshgrid(xf, yf)
+        # XF = xf.reshape((1,NC)) * np.ones((NR,1))
+        # YF = yf.reshape((NR,1)) * np.ones((1,NC))
+        # bi linear interpolation
+        u00 = U[yi0,:][:,xi0]
+        u10 = U[yi1,:][:,xi0]
+        u01 = U[yi0,:][:,xi1]
+        u11 = U[yi1,:][:,xi1]
+        u = (1-YF)*((1-XF)*u00 + XF*u01) + YF*((1-XF)*u10 + XF*u11)
         return u
     else:
         print('grids not plaid')

diff --git a/pgrid/carve_rivers.py b/pgrid/carve_rivers.py
@@ -1,10 +1,5 @@
-# -*- coding: utf-8 -*-
 """
-Created on Wed Jun 15 16:43:32 2016
-
-@author: PM5
-
-This plots river tacks.
+This carves the rivers and generates info for ROMS point sources.
 """
 
 from importlib import reload

diff --git a/pgrid/gfun.py b/pgrid/gfun.py
@@ -27,7 +27,7 @@ def gstart():
     """
     pgdir = Ldir['LOo'] / 'pgrid'
     gdir = pgdir / gridname # where grid.nc will end up
-    ri_dir = Ldir['LOo'] / 'pre' / 'river' / Ldir['gtag'] / 'tracks'
+    ri_dir = Ldir['LOo'] / 'pre' / 'river' / Ldir['gtag']
     Gr ={'gridname': gridname,'pgdir': pgdir, 'gdir': gdir,'ri_dir': ri_dir}
     return Gr
 

diff --git a/pgrid/gfun_plotting.py b/pgrid/gfun_plotting.py
@@ -4,31 +4,6 @@
 import numpy as np
 import pandas as pd
 
-def get_plon_plat(using_old_grid, ds):
-    if using_old_grid==True:
-        # because older grids did not have lon,lat_psi_ex we create this
-        # as an extension of lon,lat_psi
-        plon0 = ds.variables['lon_psi'][:]
-        plat0 = ds.variables['lat_psi'][:]
-        dx = plon0[0,1] - plon0[0,0]
-        dy = plat0[1,0] - plat0[0,0]
-        ny0, nx0 = plon0.shape
-        plon = np.nan * np.ones((ny0+2, nx0+2))
-        plat = np.nan * np.ones((ny0+2, nx0+2))
-        plon[1:-1, 1:-1] = plon0
-        plat[1:-1, 1:-1] = plat0
-        plon[:,0] = plon0[0,0] - dx
-        plon[:,-1] = plon0[0,-1] + dx
-        plon[0,:] = plon[1,:]
-        plon[-1,:] = plon[-2,:]
-        plat[0,:] = plat0[0,0] - dy
-        plat[-1,:] = plat0[-1,0] + dy
-        plat[:,0] = plat[:,1]
-        plat[:,-1] = plat[:,-2]
-    elif using_old_grid==False:
-        plon = ds.variables['lon_psi_ex'][:]
-        plat = ds.variables['lat_psi_ex'][:]
-    return (plon, plat)
 
 def get_grids(ds):
     lon_dict = dict()
@@ -42,43 +17,43 @@ def get_grids(ds):
     return (lon_dict, lat_dict, mask_dict)
 
 def add_river_tracks(Gr, ds, ax):
-    # add river tracks and endpoints
-    in_rfn = Gr['gdir'] / 'river_info.csv'
-    try:
-        df = pd.read_csv(in_rfn, index_col='rname')
-    except FileNotFoundError:
-        return
-    uv_dict = df['uv']
-    row_dict_py = df['row_py']
-    col_dict_py = df['col_py']
-    isign_dict = df['isign']
-    idir_dict = df['idir']
-    # get grids
-    lon_dict, lat_dict, mask_dict = get_grids(ds)
-    lonu = lon_dict['u']
-    latu = lat_dict['u']
-    lonv = lon_dict['v']
-    latv = lat_dict['v']
+    # add river tracks
+    gri_fn = Gr['ri_dir'] / 'river_info.csv'
+    df = pd.read_csv(gri_fn, index_col='rname')
+    # uv_dict = df['uv']
+    # row_dict_py = df['row_py']
+    # col_dict_py = df['col_py']
+    # isign_dict = df['isign']
+    # idir_dict = df['idir']
+    # # get grids
+    # lon_dict, lat_dict, mask_dict = get_grids(ds)
+    # lonu = lon_dict['u']
+    # latu = lat_dict['u']
+    # lonv = lon_dict['v']
+    # latv = lat_dict['v']
     # plot river tracks
     for rn in df.index:
-        fn_tr = Gr['ri_dir'] + 'tracks/' + rn + '.csv'
-        df_tr = pd.read_csv(fn_tr, index_col='ind')
-        x = df_tr['lon'].values
-        y = df_tr['lat'].values
+        fn_tr = Gr['ri_dir'] / 'tracks' / (rn + '.p')
+        try:
+            track_df = pd.read_pickle(fn_tr)
+        except FileNotFoundError:
+            return
+        x = track_df['lon'].to_numpy()
+        y = track_df['lat'].to_numpy()
         ax.plot(x, y, '-r', linewidth=2)
         ax.plot(x[-1], y[-1], '*r')    
-        if uv_dict[rn] == 'u' and isign_dict[rn] == 1:
-            ax.plot(lonu[row_dict_py[rn], col_dict_py[rn]],
-                    latu[row_dict_py[rn], col_dict_py[rn]], '>r')
-        elif uv_dict[rn] == 'u' and isign_dict[rn] == -1:
-            ax.plot(lonu[row_dict_py[rn], col_dict_py[rn]],
-                    latu[row_dict_py[rn], col_dict_py[rn]], '<r')
-        elif uv_dict[rn] == 'v' and isign_dict[rn] == 1:
-            ax.plot(lonv[row_dict_py[rn], col_dict_py[rn]],
-                    latv[row_dict_py[rn], col_dict_py[rn]], '^b')
-        elif uv_dict[rn] == 'v' and isign_dict[rn] == -1:
-            ax.plot(lonv[row_dict_py[rn], col_dict_py[rn]],
-                    latv[row_dict_py[rn], col_dict_py[rn]], 'vb')
+        # if uv_dict[rn] == 'u' and isign_dict[rn] == 1:
+        #     ax.plot(lonu[row_dict_py[rn], col_dict_py[rn]],
+        #             latu[row_dict_py[rn], col_dict_py[rn]], '>r')
+        # elif uv_dict[rn] == 'u' and isign_dict[rn] == -1:
+        #     ax.plot(lonu[row_dict_py[rn], col_dict_py[rn]],
+        #             latu[row_dict_py[rn], col_dict_py[rn]], '<r')
+        # elif uv_dict[rn] == 'v' and isign_dict[rn] == 1:
+        #     ax.plot(lonv[row_dict_py[rn], col_dict_py[rn]],
+        #             latv[row_dict_py[rn], col_dict_py[rn]], '^b')
+        # elif uv_dict[rn] == 'v' and isign_dict[rn] == -1:
+        #     ax.plot(lonv[row_dict_py[rn], col_dict_py[rn]],
+        #             latv[row_dict_py[rn], col_dict_py[rn]], 'vb')
 
 def edit_mask_river_tracks(Gr, NR, ax):
     # add river tracks and endpoints for edit_mask.py

diff --git a/pgrid/gfun_utility.py b/pgrid/gfun_utility.py
@@ -68,6 +68,9 @@ def load_bathy_nc(t_fn):
     tlon_vec = ds['lon'].values
     tlat_vec = ds['lat'].values
     tz = ds['z'].values
+    # There is a bug in xarray with these files: it does
+    # not set masked regions to nan.  So we do it by hand.
+    tz[tz>1e6] = np.nan
     ds.close()
     return tlon_vec, tlat_vec, tz
 

diff --git a/pgrid/plot_grid.py b/pgrid/plot_grid.py
@@ -19,6 +19,7 @@
     Gr = gfun.gstart()
 from lo_tools import plotting_functions as pfun
 import gfun_plotting as gfp
+reload(gfp)
 
 import matplotlib.pyplot as plt
 import xarray as xr
@@ -76,15 +77,15 @@
 ax1 = fig.add_subplot(1,NC,1)
 cmap1 = plt.get_cmap(name='gist_earth') # terrain, viridis
 cs = ax1.pcolormesh(plon, plat, zm,
-                   vmin=-300, vmax=10, cmap = cmap1)
+                   vmin=-300, vmax=100, cmap = cmap1)
 fig.colorbar(cs, ax=ax1, extend='both')
 pfun.add_coast(ax1)
 pfun.dar(ax1)
 ax1.axis(ax_lims)
 ax1.set_title(Gr['gridname'] + '/' + fn)
 ax1.text(.95, .05, str(mask_rho.shape), horizontalalignment='right',
          transform=ax1.transAxes)                   
-#gfp.add_river_tracks(Gr, ds, ax1)
+gfp.add_river_tracks(Gr, ds, ax1)
 
 if flag_show_sections:
     color_list = ['orange', 'gold', 'greenyellow', 'lightgreen',

diff --git a/pgrid/start_grid.py b/pgrid/start_grid.py
@@ -69,10 +69,7 @@
     m = np.ones_like(lon)
     for t_fn in dch['t_list']:
         print('\nOPENING BATHY FILE: ' + t_fn.name)
-        if t_fn.name == 'topo15.nc':
-            tlon_vec, tlat_vec, tz = gfu.load_bathy2(t_fn, Lon_vec, Lat_vec)
-        else:
-            tlon_vec, tlat_vec, tz = gfu.load_bathy_nc(t_fn)
+        tlon_vec, tlat_vec, tz = gfu.load_bathy_nc(t_fn)
         tlon, tlat = np.meshgrid(tlon_vec, tlat_vec)
         z_part = zfun.interp2(lon, lat, tlon, tlat, tz)
         # put good values of z_part in z