-
Notifications
You must be signed in to change notification settings - Fork 10
/
Copy pathcarve_rivers.py
189 lines (169 loc) · 7.11 KB
/
carve_rivers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
"""
This carves the rivers and generates info for ROMS point sources.
"""
import gfun
Gr =gfun.gstart()
from lo_tools import zfun
from lo_tools import plotting_functions as pfun
import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seawater as sw
import pickle
# load the default choices
dch = pickle.load(open(Gr['gdir'] / 'choices.p', 'rb'))
# get river info
gri_fn = Gr['ri_dir0'] / dch['ctag'] / 'river_info.csv'
ri_df = pd.read_csv(gri_fn, index_col='rname')
# select and increment grid file
in_fn = gfun.select_file(Gr)
out_fn = gfun.increment_filename(in_fn, '_r')
# get the grid data
ds = xr.open_dataset(in_fn)
z = -ds.h.values
mask_rho = ds.mask_rho.values
lon = ds.lon_rho.values
lat = ds.lat_rho.values
X = lon[0,:]
Y = lat[:,0]
def in_domain(x, y, X, Y):
# Utility function to make sure that a point (x, y) is
# in a domain specified by vectors X and Y.
# We actually require the point to be 'pad' in from the edge.
pad = 8 # using 8 should keep us out of the nudging to climatology regions
if x>=X[0+pad] and x<=X[-1-pad] and y>=Y[0+pad] and y<=Y[-1-pad]:
return True
else:
return False
# prepare things for looping over the rivers
ilon_dict = dict()
ilat_dict = dict()
dir_dict = dict()
good_riv = []
roms_info_df = pd.DataFrame()
# loop over rivers
for rn in ri_df.index:
if rn not in dch['excluded_rivers']:
depth = ri_df.loc[rn, 'depth']
try:
fn_tr = Gr['ri_dir0'] / dch['ctag'] / 'tracks' / (rn + '.p')
track_df = pd.read_pickle(fn_tr)
# NOTE: tracks go from ocean to land
x = track_df['lon'].to_numpy()
y = track_df['lat'].to_numpy()
if in_domain(x[0], y[0], X, Y):
# we only consider rivers whose mouth is in the domain
good_riv.append(rn)
print('including ' + rn.title())
II_list = []
JJ_list = []
for ii in range(len(x)-1):
x0 = x[ii]
x1 = x[ii+1]
y0 = y[ii]
y1 = y[ii+1]
if in_domain(x0, y0, X, Y) and in_domain(x1, y1, X, Y):
ix0 = zfun.find_nearest_ind(X,x0)
ix1 = zfun.find_nearest_ind(X,x1)
iy0 = zfun.find_nearest_ind(Y,y0)
iy1 = zfun.find_nearest_ind(Y,y1)
if ix0==ix1 and iy0==iy1:
II = np.array(ix0)
JJ = np.array(iy0)
z[JJ, II] = np.min((z[JJ, II], -depth))
II_list.append(ix0)
JJ_list.append(iy0)
else:
II, JJ = zfun.get_stairstep(ix0, ix1, iy0, iy1)
for pp in range(len(II)):
z[JJ[pp], II[pp]] = np.min((z[JJ[pp], II[pp]], -depth))
II_list += list(II)
JJ_list += list(JJ)
mask_rho[JJ, II] = 1
# form unique entries of lists
JI_ulist = []
for pp in range(len(II_list)):
ji_tup = (JJ_list[pp], II_list[pp])
if ji_tup not in JI_ulist:
JI_ulist.append(ji_tup)
# Then use the last two point to decide on the river direction.
# dx and dy are in: [-1,0,1] and are interpreted as, e.g.
# dx = 1 means that the source is on the E side of the rho cell
# ( column +1 is the step we took to get to the last river channel rho cell).
if len(JI_ulist) >1:
# simple if we have at least two points in JI_ulist
dx = JI_ulist[-1][1] - JI_ulist[-2][1]
dy = JI_ulist[-1][0] - JI_ulist[-2][0]
else:
# more tedious if we have ony one point in JI_ulist
dist, ang = sw.dist([y[0],y[-1]], [x[0], x[-1]])
if -45<ang and ang<=45:
dx=1; dy = 0
elif 45<ang and ang<=135:
dx=0; dy = 1
elif -135<ang and ang<=-45:
dx=0; dy = -1
elif ang<=-135 or ang>135:
dx=-1; dy = 0
else:
print(' Inconsistent angle for %s '.center(60,'*') % (rn))
# Write info for ROMS based on dx and dy.
#
# NOTE: the row_py and col_py that we write to river_info.csv use
# python indexing conventions, meaning that the rho, u, and v grids
# start at [0,0]. Hence when we plot things, as in plot_grid.py, we
# should be able to use [row_py, col_py] directly with the u and v grids.
#
# However, when we later create the source positions for running ROMS,
# e.g. in forcing/riv0/make_forcing_main.py => rivers.nc, we convert these
# values to ROMS indexing conventions, meaning:
# - add 1 to row of N/S sources
# - add 1 to column of E/W sources.
xoff = 0
yoff = 0
if dx==1 and dy==0:
# Source on E side of rho-cell
idir = 0 # 0 = E/W, 1 = N/S (redundant with 'uv')
isign = -1
uv = 'u'
elif dx==-1 and dy==0:
# Source on W side of rho-cell
idir = 0
isign = 1
uv = 'u'
xoff = -1
elif dx==0 and dy==1:
# Source on N side of rho-cell
idir = 1
isign = -1
uv = 'v'
elif dx==0 and dy==-1:
# Source on S side of rho-cell
idir = 1
isign = 1
uv = 'v'
yoff = -1
else:
print(' Inconsistent last points for %s '.center(60,'*') % (rn))
roms_info_df.loc[rn, 'row_py'] = JI_ulist[-1][0] + yoff
roms_info_df.loc[rn, 'col_py'] = JI_ulist[-1][1] + xoff
roms_info_df.loc[rn, 'idir'] = idir
roms_info_df.loc[rn, 'isign'] = isign
roms_info_df.loc[rn, 'uv'] = uv
else:
print(' >> excluding ' + rn.title())
except FileNotFoundError:
pass
# save the updated mask and z
ds.update({'mask_rho': (('eta_rho', 'xi_rho'), mask_rho)})
ds.update({'h': (('eta_rho', 'xi_rho'), -z)})
ds.to_netcdf(out_fn)
ds.close()
# save the roms river info
out_rfn = Gr['gdir'] / 'roms_river_info.csv'
print('\nCreating ' + str(out_rfn))
roms_info_df.index.name = 'rname'
roms_info_df.to_csv(out_rfn)
# here is how you would read the DataFrame back in
# df1 = pd.read_csv(out_rfn, index_col='rname')