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