def get_previous_ij_list(i0, j0, dirs, lake_fr, glob_lakefr_limit=0.6):
i_list = []
j_list = []
nx, ny = dirs.shape
for di in range(-1, 2):
for dj in range(-1, 2):
if di == 0 and dj == 0:
continue
i = i0 + di
j = j0 + dj
if i >= nx or i < 0:
continue
if j >= ny or j < 0:
continue
if lake_fr[i, j] < glob_lakefr_limit:
continue
i_next, j_next = direction_and_value.to_indices(i, j, dirs[i, j])
if i_next != i0 or j_next != j0:
continue
i_list.append(i)
j_list.append(j)
i_upper, j_upper = get_previous_ij_list(i, j, dirs, lake_fr, glob_lakefr_limit=glob_lakefr_limit)
i_list.extend(i_upper)
j_list.extend(j_upper)
return i_list, j_list
def __init__(self, flow_dirs, nx=None, ny=None,
lons2d=None,
lats2d=None,
accumulation_area_km2=None):
self.cells = []
self.lons2d = lons2d
self.lats2d = lats2d
self.flow_directions = flow_dirs
self.accumulation_area_km2 = accumulation_area_km2
# calculate characteristic distance
if not any([None is arr for arr in [self.lats2d, self.lons2d]]):
v1 = lat_lon.lon_lat_to_cartesian(self.lons2d[0, 0], self.lats2d[0, 0])
v2 = lat_lon.lon_lat_to_cartesian(self.lons2d[1, 1], self.lats2d[1, 1])
dv = np.array(v2) - np.array(v1)
self.characteristic_distance = np.sqrt(np.dot(dv, dv))
x, y, z = lat_lon.lon_lat_to_cartesian(self.lons2d.flatten(), self.lats2d.flatten())
self.kdtree = cKDTree(list(zip(x, y, z)))
if None not in [nx, ny]:
self.nx = nx
self.ny = ny
else:
nx, ny = flow_dirs.shape
self.nx, self.ny = flow_dirs.shape
for i in range(nx):
self.cells.append(list([Cell(i=i, j=j, flow_dir_value=flow_dirs[i, j]) for j in range(ny)]))
self._without_next_mask = np.zeros((nx, ny), dtype=np.int)
self._wo_next_wo_prev_mask = np.zeros((nx, ny), dtype=np.int) # mask of the potential outlets
for i in range(nx):
if i % 100 == 0:
print("Created {}/{}".format(i, nx))
for j in range(ny):
i_next, j_next = direction_and_value.to_indices(i, j, flow_dirs[i][j])
next_cell = None
if 0 <= i_next < nx:
if 0 <= j_next < ny:
next_cell = self.cells[i_next][j_next]
self._without_next_mask[i, j] = int(next_cell is None)
self.cells[i][j].set_next(next_cell)
