def water_balance(avg_hs, avg_rp, avg_dis, area, cell, ofstorage, label_discharge="dis_0"):
"""All output units are in mm"""
P = avg_hs["pre_0"]
AET = avg_hs["aet_0"]
R = avg_hs["rch_0"]
OF = avg_hs["run_0"]
Qgw = avg_hs["gdh_0"]*1000.0
TL = avg_hs["tls_0"]*1000/np.power(cell,2)
Peff = avg_hs["pre_0"] - avg_hs["inf_0"]
Qgw = avg_hs["gdh_0"]*1000.
Q = avg_dis[label_discharge]*1000/area
factor = avg_hs["tls_0"]/avg_rp["tls_0"]
AER = avg_rp["aet_0"]*factor
Rfch = avg_rp["fch_0"]*factor
Rdch = avg_hs["rch_0"] - Rfch
ofstorage = ofstorage*factor*1000/np.power(cell,2)
Qbc = avg_hs["chb_0"]*1000.
EGW = avg_hs["egw_0"]
MBuz = avg_hs["inf_0"] - avg_hs["aet_0"] - Rdch
MB = (avg_hs["pre_0"] - avg_hs["aet_0"] - avg_hs["egw_0"]
- Q - AER - ofstorage - avg_hs["twsc_0"] - Qbc)
# store fluxes in a dictionary
fluxes = {"P": [P, 100.0],
"AET": [AET, AET*100./P], "AER": [AER, AER*100./P], "EGW": [EGW, EGW*100./P],
"R": [R, R*100./P], "Rfch": [Rfch, Rfch*100./P], "Rdch": [Rdch, Rdch*100./P],
"Q": [Q, Q*100./P], "Qgw": [Qgw, Qgw*100./P],
"Peff": [Peff, Peff*100./P],
"OF": [OF, OF*100./P], "TL": [TL, TL*100./P],
"TSWS":[ofstorage, ofstorage*100./P],
"TWSA":[avg_hs["twsc_0"],avg_hs["twsc_0"]*100./P],
"Qbc": [Qbc, Qbc*100./P],
"MB": [MB, MB*100./P],
"MBuz": [MBuz, MBuz*100./P]}
return pd.DataFrame.from_dict(fluxes)
def calculate_water_balance_from_csv(fhillslope, friparian, fdischarge, area, cell, label_discharge="dis_0"):
avg_hs = pd.read_csv(fhillslope).sum()
avg_rp = pd.read_csv(friparian).sum()
avg_dis = pd.read_csv(fdischarge).sum()
avg_rp_mean = pd.read_csv(friparian).mean()
ofstorage = avg_rp_mean["ssz_0"]
return water_balance(avg_hs, avg_rp, avg_dis, area, cell, ofstorage, label_discharge)