diff --git a/Squeegee/README b/Squeegee/README new file mode 100644 index 0000000..d15675e --- /dev/null +++ b/Squeegee/README @@ -0,0 +1,10 @@ +#centos +source /modules/centos7/conda/Feb2021/etc/profile.d/conda.sh +conda activate production +python3 -m pip install cf_xarray==0.7.4 --user +python3 -m pip install cartopy==0.21 --user + +#redhat (Helene sitt env) +source /modules/rhel8/conda/install/etc/profile.d/conda.sh +conda activate /lustre/storeB/project/nwp/H2O/wp4/conda/xesmf_envn +python -c "import esmpy" diff --git a/Squeegee/minimal_squeegee_runoff_drainage.py b/Squeegee/minimal_squeegee_runoff_drainage.py new file mode 100644 index 0000000..a9bd3a3 --- /dev/null +++ b/Squeegee/minimal_squeegee_runoff_drainage.py @@ -0,0 +1,355 @@ +import sys, os, glob +import os +import numpy as np +import xarray as xr +import matplotlib +import matplotlib.pyplot as plt +#matplotlib.use('Agg') +import pandas as pd +from shapely.geometry import Polygon, MultiPolygon, shape, Point +import geopandas as gp +import xesmf as xe +import pyproj +import requests +import cartopy.crs as ccrs +from cartopy.io.img_tiles import Stamen +import cartopy.feature as cfeature +tiler = Stamen('terrain-background') + + +def convert_3D_2D(geometry): + ''' + Takes a GeoSeries of 3D Multi/Polygons (has_z) and returns a list of 2D Multi/Polygons + ''' + new_geo = [] + for p in geometry: + if p.has_z: + if p.geom_type == 'Polygon': + lines = [xy[:2] for xy in list(p.exterior.coords)] + new_p = Polygon(lines) + new_geo.append(new_p) + elif p.geom_type == 'MultiPolygon': + new_multi_p = [] + for ap in p: + lines = [xy[:2] for xy in list(ap.exterior.coords)] + new_p = Polygon(lines) + new_multi_p.append(new_p) + new_geo.append(MultiPolygon(new_multi_p)) + return new_geo + + +def getDomainData(expdirs): + + # Assuming same forcing + #FORCING=xr.open_dataset(expdirs+'forcing/2023050100/FORCING.nc', cache=False,use_cftime=False) + FORCING=xr.open_dataset('/lustre/storeB/users/josteinbl/sfx_data/LDAS_NOR/archive/2022/12/12/06/raw.nc', cache=False,use_cftime=False) + print(FORCING.latitude.min(), FORCING.latitude.max()) + print(FORCING.longitude.min(), FORCING.longitude.max()) + + fPREP=xr.open_dataset('/lustre/storeB/users/josteinbl/sfx_data/LDAS_NOR/'+'archive/2023/04/01/00/SURFOUT.nc', cache=False,use_cftime=False) + fPGD=xr.open_dataset('/lustre/storeB/users/josteinbl/sfx_data/LDAS_NOR/'+'climate/PGD.nc', cache=False,use_cftime=False) + + # Add coords etc so that surfex files may be processed + R_pysurfex=6371000 # and met nordic + R_arome =6371229 + Rpy=6.371229e+06 + proj_string = "+proj=lcc +lat_0=63 +lon_0=15 +lat_1=63 +lat_2=63 +units=m +no_defs +R=" + str(Rpy) + myP=pyproj.Proj(proj_string) + + lcc = ccrs.LambertConformal(#globe=globe, # no datumshift in GCMs + central_longitude=15, central_latitude=63, + standard_parallels=(63,63))#, + geodetic=ccrs.Geodetic() #default WGS84 + + return myP, proj_string, fPREP, fPGD, FORCING, lcc + + +def readStationList(stationfile): + + #read a list of stations + stationst=stationfile + stations=pd.read_csv(stationst, header=0,index_col=0,sep=None, + dtype={'GauID': 'str'}) + stations=stations.drop_duplicates(subset=('GauID'), keep='last') + + # Nr. GauID Lat Long Ara_km2 Group + daily = 1440 + station='2.25.0'# stations.GauID[2] + + stations['GauID']=stations.GauID+'.0' + + return stations + + +def readShapeFiles(shapefiles, stations): + + # or Outlets to sea defined for Europe CCM2 +# ccm2_basins = gp.read_file(shapefiles,layer='SEAOUTLETS') + #ccm2sub = ccm2_basins.cx[FORCING.LON.min().values:FORCING.LON.max().values,FORCING.LAT.min().values:FORCING.LAT.max().values] + + # or cathements where NVE have measuring stations + #regs = gp.read_file('/lustre/storeB/project/nwp/H2O/wp4/RRdata/shapefiles/utm33shp/NVEData/Hydrologi/Hydrologi_TotalNedborfeltMalestasjon.shp') + regs2 = gp.read_file(shapefiles) + + print("prior") + print(regs2.columns) + regs2=regs2.loc[regs2['stID'].isin(stations.GauID)] + + print("posterior") + #inspect the shape file + print(regs2.crs) + print(regs2.columns) + + #remove extra dim in nve shape-files + regs2.geometry = convert_3D_2D(regs2.geometry) + + plot_polygons=False + #make a plot of the polygons: + if plot_polygons==True: + + fig = plt.figure(figsize=(4,8)) + ax1 = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree(), frameon=False) + + ax1.coastlines(resolution='10m') + toplot= regs2.cx[FORCING.longitude.min().values:FORCING.longitude.max().values,FORCING.latitude.min().values:FORCING.latitude.max().values] + + toplot.plot(column='stID',kind='geo',edgecolor="g",linewidth=0.2, + ax=ax1,legend=True,cmap='Reds',transform=ccrs.Geodetic()) + #,legend_kwds={"orientation": "horizontal", "pad": 0.01, "label": "VassOmrNr"}) + #ccm2sub.plot(column='AREA_KM2',kind='geo',edgecolor="g",linewidth=0.2,ax=ax1,legend=True,cmap='Reds',legend_kwds={"orientation": "horizontal", "pad": 0.01}) + ax1.axis('off') + plt.title("") + plt.savefig("polygons.png") + plt.show() + + return regs2 + +def readModelData(expdirs, expdirfs, fPREP, fPGD, ncstart, expnam, variables): + + numruns=len(expnam) + + #----------------------------------------- + # read data in this way because a large amount of data + def preprocess(ds): + return ds[variables] + + dictds={} + + for i in np.arange(0,numruns): + dictds[expnam[i]] = xr.open_mfdataset(expdirs+expdirfs[i]+ncstart[i]+ + 'sel2023??.nc', + cache=False, + preprocess=preprocess, + concat_dim="time", + combine="nested", + chunks={'time': 1}) + print(expnam[i] + ' given to: ' +expdirs+expdirfs[i]+ncstart[i]+ + 'sel2023??.nc') + + print(dictds) + + return dictds + +def generateWeights(dictds, fPREP, fPGD, FORCING, regs2, weightfile): + + # xesmf needs dataset to generate weights + TSr=dictds[expnam[-1]]['RUNOFFC_ISBA'].isel(time=0).to_dataset() + print('TSr') + print(TSr) + + xfalseaR, yfalseaR = myP(fPREP.LONORI.data, fPREP.LATORI.data,inverse=False) + #Surfex makes all coords positive by setting the origin to the lower left corner + + TSr=TSr.assign_coords(x= fPREP.XX[0,:].data+xfalseaR) + TSr=TSr.assign_coords(y= fPREP.YY[:,0].data+yfalseaR) + TSr = TSr.rename({'xx': 'x', 'yy': 'y'}) + + dx, dy = fPGD.DX.data[0,0], fPGD.DY.data[0,0] + Xcorners=np.arange(TSr['x'].data[0]-dx/2., TSr['x'].data[-1]+3*dx/2., dx) + Ycorners=np.arange(TSr['y'].data[0]-dy/2., TSr['y'].data[-1]+3*dy/2., dy) + #Lon2, Lat2 = myP(fPREP.XX.data+xfalseaR,fPREP.YY.data+yfalseaR,inverse=True) + Lon2b, Lat2b = myP(*np.meshgrid(Xcorners,Ycorners),inverse=True) + + TSr.coords['xb'] = (Xcorners) + TSr.coords['yb'] = (Ycorners) + TSr.coords['lat_b'] = (('yb','xb'),Lat2b) + TSr.coords['lon_b'] = (('yb','xb'),Lon2b) + TSr.set_coords(['lat_b','lon_b']) + + print("Forcing lat data") + print(np.shape(FORCING.latitude.data)) + #print(TSr.coords['lat_b']) + + TSr.coords['lat'] = (('y','x'),FORCING.latitude.data) + TSr.coords['lon'] = (('y','x'),FORCING.longitude.data) + TSr.lon.attrs=FORCING.longitude.attrs + TSr.lat.attrs=FORCING.longitude.attrs + TSr.set_coords(['lat','lon']) + + TSr['mask']=TSr.RUNOFFC_ISBA.isnull() + TSr['mask'].values=np.where(~TSr.RUNOFFC_ISBA.isnull(),1,0) + + TSr.attrs['pyproj_srs']=proj_string + + #make an outer domain polygon to crop the input shape file + + dompoly=Polygon(zip([TSr['lon_b'][0,0].data,TSr['lon_b'][0,-1].data,TSr['lon_b'][-1,-1].data, + TSr['lon_b'][-1,0].data],[TSr['lat_b'][0,0].data,TSr['lat_b'][0,-1].data, + TSr['lat_b'][-1,-1].data,TSr['lat_b'][-1,0].data])) + + gdom = gp.GeoSeries([dompoly]) + + # assing defined polygon to a new dataframe + pol_gpd= gp.GeoDataFrame() + pol_gpd['geometry'] = None + pol_gpd.loc[0,'geometry'] = dompoly + pol_gpd.crs=regs2.crs #hope okay + + # crop shape file to domain, remove basins > 2 km2 + result = gp.sjoin(regs2, pol_gpd, how='inner')#, op='within') + resultn=result[result.areal_km2 > 2] + resultn=resultn[resultn.stID!='2.11.0'] #only missing data + resultn.to_csv('basins_used.csv') + + # This generates the weights, takes time the first time only + # ... + savg = xe.SpatialAverager(TSr, resultn.geometry, geom_dim_name="stID", + filename=weightfile,reuse_weights=True) + + plotweights=False + + # this takes a lot of memory for all these basins, not recommended + if plotweights: + w = xr.DataArray( + savg.weights.toarray().reshape(resultn.geometry.size, *TSr.lat.shape), + dims=("stID", *TSr.lat.dims), + coords=dict(stID=out.stID, **TSr.lon.coords), + ) + + plt.subplots_adjust(top=0.9) + facets = w.plot(col="stID", col_wrap=6, aspect=2, vmin=0, vmax=0.05) + facets.cbar.set_label("Averaging weights") + plt.savefig('weights.png') + + return TSr, savg, resultn + +#store the dfs in dict + +# here the sparse matrix with grid cell weights for each ploygon is generated +# see xesmf doc: +# https://pangeo-xesmf.readthedocs.io/en/latest/notebooks/Spatial_Averaging.html +# the matrix can be stored as a nc-file and re-used, so that the program runs +# very fast once it is generated. +# The weights are specific to the domain, projection, resolution, and to +# options like +#&nam_pgd_arrange_cover +# lwater_to_nature = .true. +# ltown_to_rock = .true. +# which alters where RUNOFFC etc is defined + + +def run(dictds, expname, TSr, savg, resultn): + + numruns=len(expnam) + cachrunoff={} + + for i in np.arange(0,numruns): + + tmp1=dictds[expnam[i]]['RUNOFFC_ISBA'] +dictds[expnam[i]]['DRAINC_ISBA'] + + #some experiments have accumlated runoff (i.e. do not use option + # LRESETCUMUL = .true. in NAM_WRITE_DIAG_SURFn + # see https://www.umr-cnrm.fr/surfex/spip.php?article406 ) + # so need to deaccumulate + #if expnam[i] in ['LDAS_FC']: + # tmp1=tmp1.diff(dim='time') + tmp1=tmp1.to_dataset(name=expnam[i]) + tmp1['mask']=TSr.RUNOFFC_ISBA.isnull() + tmp1['mask'].values=np.where(~TSr.RUNOFFC_ISBA.isnull(),1,0) + out=savg(tmp1[expnam[i]]) + #out = out.assign_coords(roms_id=xr.DataArray(resultn.index.values, dims=("roms_id",))) + out = out.assign_coords(stID=xr.DataArray(resultn["stID"], dims=("stID",))) + print("out") + print(out) + cachrunoff[expnam[i]]=out.to_pandas() + rainfRun= cachrunoff[expnam[i]]=out.to_pandas() + print("rainfRun") + print(rainfRun.keys()) + + rainfRun[:].plot() + #rainfRun.plot() + plt.title(expnam[i], fontsize=12) + plt.show() + cachrunoff[expnam[i]].to_csv(expnam[i]+'.csv') + + +if __name__ == "__main__": + + expdirs='/lustre/storeB/users/josteinbl/sfx_data/' + #forcingfile = '/lustre/storeB/project/nwp/H2O/wp4/FORCING/527_450_2020/FORCING_527_450_202010.nc' + stationfile = '/lustre/storeB/project/nwp/H2O/wp4/RRdata/stations_Huang20.txt' +# shapefiles = '/lustre/storeB/project/nwp/H2O/wp4/RRdata/shapefiles/CCM2/ccm21/WGS84_W2008.gdb' + shapefiles = '/lustre/storeB/project/nwp/H2O/wp4/RRdata/shapefiles/latlonHyd/latlonhydorder/NVEData/Hydrologi/Hydrologi_TotalNedborfeltMalestasjon.shp' + #shapefiles = '/lustre/storeB/project/nwp/H2O/wp4/RRdata/shapefiles/latlonREG/Nedborfelt/Nedborfelt_Vassdragsomr.shp' + + weightfile = '/lustre/storeB/users/josteinbl/sfx_data/LDAS_NOR/spatial_avg_4catchments_527nature2lim.nc' + + # folder where runs are, this could have been a long list of different exps. + #expdirfs=['LDAS_NOR/', 'LDAS_NOR_eps_05/'] + expdirfs=['LDAS_FC/'] + # start of nc-file-output from runs + #openoflsel201903.nc + + # Analysis mode: + #ncstart= ['openofl', 'openofl'] + ncstart= ['openofl'] + # choose a name for legend in plots + #expnam=['LDAS_NOR', 'LDAS_NOR_eps_05'] + expnam=['LDAS_FC'] + + variables = ['RUNOFFC_ISBA', 'DRAINC_ISBA'] + + myP, proj_string, fPREP, fPGD, FORCING, lcc = getDomainData(expdirs) + stations = readStationList(stationfile) + + regs2 = readShapeFiles(shapefiles, stations) + + dictds = readModelData(expdirs, expdirfs, fPREP, fPGD, ncstart, expnam, variables) + + TSr, savg, resultn = generateWeights(dictds, fPREP, fPGD, FORCING, regs2, weightfile) + + run(dictds, expnam, TSr, savg, resultn) + + plot_catchments = False + + if plot_catchments == True: + + #for sid in regs2.stID: + #for sid in stations["GauID"]: + + fig = plt.figure(figsize=(12,7)) + #gs = fig.add_gridspec(1, 1) + #ax0=fig.add_subplot(gs[0,0], projection=ccrs.PlateCarree()) + #ax0.set_extent([FORCING.longitude.data.min()-15, FORCING.longitude.data.max()+15, FORCING.latitude.data.min()-15, FORCING.latitude.data.max()+15], crs=ccrs.PlateCarree()) + #ax0.add_feature(cfeature.LAND) + #ax0.add_feature(cfeature.COASTLINE) + #ax0.add_image(tiler, 7) + #gl = ax0.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, + # linewidth=2, color='gray', alpha=0.5, linestyle='--') + #ax0.plot(FORCING.longitude[[0,0,-1,-1,0],[0,-1,-1,0,0]].values.diagonal(),FORCING.latitude[[0,0,-1,-1,0],[0,-1,-1,0,0]].values.diagonal(),color='blue', linewidth=2,transform=ccrs.Geodetic()) + #gl.top_labels = gl.right_labels = False + ax0b = fig.add_subplot(projection=lcc) + ax0b.set_extent([TSr.x.min(), TSr.x.max(), TSr.y.min(), TSr.y.max()], crs=lcc) + gl = ax0b.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, + linewidth=0.8, color='gray', alpha=0.5, linestyle='--') + gl.top_labels = gl.right_labels = False + ax0b.add_image(tiler, 7) + #regs2[regs2.stID==sid].boundary.plot(ax=ax0b,transform=ccrs.Geodetic(),color='r') + regs2.boundary.plot(ax=ax0b,transform=ccrs.Geodetic(),color='r') + #plt.title(sid+' '+str(resultn.areal_km2[resultn.stID==sid])+' km2') + #regs2.boundary.plot(ax=ax0b,transform=ccrs.Geodetic(),color='r') + plt.tight_layout() + #plt.savefig("/lustre/storeB/users/josteinbl/TOPD/Figures/"+sid+'_snowDA_test.png') + #plt.savefig('/lustre/storeB/users/josteinbl/TOPD/Figures/All_catchments_snowDA_test.png') + #plt.close() + #plt.show()