def Dist_HBV2(lakecell,q_lake,DEM,flow_acc,flow_acc_plan, sp_prec, sp_et, sp_temp, sp_pars, p2, init_st=None,
ll_temp=None, q_0=None):
'''
Make spatially distributed HBV in the SM and UZ
interacting cells
'''
n_steps = sp_prec.shape[2] + 1 # no of time steps =length of time series +1
# intiialise vector of nans to fill states
dummy_states = np.empty([n_steps, 5]) # [sp,sm,uz,lz,wc]
dummy_states[:] = np.nan
# Get the mask
mask, no_val = GISpy.get_mask(DEM)
x_ext, y_ext = mask.shape # shape of the fpl raster (rows, columns)-------------- rows are x and columns are y
# y_ext, x_ext = mask.shape # shape of the fpl raster (rows, columns)------------ should change rows are y and columns are x
# Get deltas of pixel
geo_trans = DEM.GetGeoTransform() # get the coordinates of the top left corner and cell size [x,dx,y,dy]
dx = np.abs(geo_trans[1])/1000.0 # dx in Km
dy = np.abs(geo_trans[-1])/1000.0 # dy in Km
px_area = dx*dy # area of the cell
# Enumerate the total number of pixels in the catchment
tot_elem = np.sum(np.sum([[1 for elem in mask_i if elem != no_val] for mask_i in mask])) # get row by row and search [mask_i for mask_i in mask]
# total pixel area
px_tot_area = tot_elem*px_area # total area of pixels
# Get number of non-value data
st = [] # Spatially distributed states
q_lz = []
q_uz = []
#------------------------------------------------------------------------------
for x in range(x_ext): # no of rows
st_i = []
q_lzi = []
q_uzi = []
# q_out_i = []
# run all cells in one row ----------------------------------------------------
for y in range(y_ext): # no of columns
if mask [x, y] != no_val: # only for cells in the domain
# Calculate the states per cell
# TODO optimise for multiprocessing these loops
_, _st, _uzg, _lzg = HBV.simulate_new_model(avg_prec = sp_prec[x, y,:],
temp = sp_temp[x, y,:],
et = sp_et[x, y,:],
par = sp_pars[x, y, :],
p2 = p2,
init_st = init_st,
ll_temp = None,
q_0 = q_0,
extra_out = True)
# # append column after column in the same row -----------------
st_i.append(np.array(_st))
#calculate upper zone Q = K1*(LZ_int_1)
q_lz_temp=np.array(sp_pars[x, y, 6])*_lzg
q_lzi.append(q_lz_temp)
# calculate lower zone Q = k*(UZ_int_3)**(1+alpha)
q_uz_temp = np.array(sp_pars[x, y, 5])*(np.power(_uzg, (1.0 + sp_pars[x, y, 7])))
q_uzi.append(q_uz_temp)
# print("total = "+str(fff)+"/"+str(tot_elem)+" cell, row= "+str(x+1)+" column= "+str(y+1) )
else: # if the cell is novalue-------------------------------------
# Fill the empty cells with a nan vector
st_i.append(dummy_states) # fill all states(5 states) for all time steps = nan
q_lzi.append(dummy_states[:,0]) # q lower zone =nan for all time steps = nan
q_uzi.append(dummy_states[:,0]) # q upper zone =nan for all time steps = nan
# store row by row-------- ----------------------------------------------------
# st.append(st_i) # state variables
st.append(st_i) # state variables
q_lz.append(np.array(q_lzi)) # lower zone discharge mm/timestep
q_uz.append(np.array(q_uzi)) # upper zone routed discharge mm/timestep
#------------------------------------------------------------------------------
# convert to arrays
st = np.array(st)
q_lz = np.array(q_lz)
q_uz = np.array(q_uz)
#%% convert quz from mm/time step to m3/sec
area_coef=p2[1]/px_tot_area
q_uz=q_uz*px_area*area_coef/(p2[0]*3.6)
no_cells=list(set([flow_acc_plan[i,j] for i in range(x_ext) for j in range(y_ext) if not np.isnan(flow_acc_plan[i,j])]))
# no_cells=list(set([int(flow_acc_plan[i,j]) for i in range(x_ext) for j in range(y_ext) if flow_acc_plan[i,j] != no_val]))
no_cells.sort()
#%% routing lake discharge with DS cell k & x and adding to cell Q
q_lake=Routing.muskingum(q_lake,q_lake[0],sp_pars[lakecell[0],lakecell[1],10],sp_pars[lakecell[0],lakecell[1],11],p2[0])
q_lake=np.append(q_lake,q_lake[-1])
# both lake & Quz are in m3/s
#new
q_uz[lakecell[0],lakecell[1],:]=q_uz[lakecell[0],lakecell[1],:]+q_lake
#%% cells at the divider
q_uz_routed=np.zeros_like(q_uz)*np.nan
# for all cell with 0 flow acc put the q_uz
for x in range(x_ext): # no of rows
for y in range(y_ext): # no of columns
if mask [x, y] != no_val and flow_acc_plan[x, y]==0:
q_uz_routed[x,y,:]=q_uz[x,y,:]
#%% new
for j in range(1,len(no_cells)): #2):#
for x in range(x_ext): # no of rows
for y in range(y_ext): # no of columns
# check from total flow accumulation
if mask [x, y] != no_val and flow_acc_plan[x, y]==no_cells[j]:
# print(no_cells[j])
q_r=np.zeros(n_steps)
for i in range(len(flow_acc[str(x)+","+str(y)])): # no_cells[j]
# bring the indexes of the us cell
x_ind=flow_acc[str(x)+","+str(y)][i][0]
y_ind=flow_acc[str(x)+","+str(y)][i][1]
# sum the Q of the US cells (already routed for its cell)
# route first with there own k & xthen sum
q_r=q_r+Routing.muskingum(q_uz_routed[x_ind,y_ind,:],q_uz_routed[x_ind,y_ind,0],sp_pars[x_ind,y_ind,10],sp_pars[x_ind,y_ind,11],p2[0])
# q=q_r
# add the routed upstream flows to the current Quz in the cell
q_uz_routed[x,y,:]=q_uz[x,y,:]+q_r
#%% check if the max flow _acc is at the outlet
# if tot_elem != np.nanmax(flow_acc_plan):
# raise ("flow accumulation plan is not correct")
# outlet is the cell that has the max flow_acc
outlet=np.where(flow_acc_plan==np.nanmax(flow_acc_plan)) #np.nanmax(flow_acc_plan)
outletx=outlet[0][0]
outlety=outlet[1][0]
#%%
q_lz = np.array([np.nanmean(q_lz[:,:,i]) for i in range(n_steps)]) # average of all cells (not routed mm/timestep)
# convert Qlz to m3/sec
q_lz = q_lz* p2[1]/ (p2[0]*3.6) # generation
q_out = q_lz + q_uz_routed[outletx,outlety,:]
return q_out, st, q_uz_routed, q_lz, q_uz
