- implemented the changes in file convert_shp_to_SCRIP_nc.py from branch...

- implemented the changes in file convert_shp_to_SCRIP_nc.py from branch export_weights_to_file and reformatted file

src_python/pre_proc/convert_shp_to_SCRIP_nc.py

@@ -9,18 +9,18 @@ Author      : Robert Schweppe (robert.schweppe@ufz.de)
 Created     : 2019-09-05 11:46
 """
 
+import argparse
+import math
+import warnings
+from itertools import product
+
 # IMPORTS
 import geopandas as gpd
+import numpy as np
 import pandas as pd
+import tqdm
 import xarray as xr
-import numpy as np
-import argparse
-import sys
 from shapely.geometry import Polygon
-import tqdm
-import math
-import warnings
-from itertools import product
 
 # GLOBAL VARIABLES
 MISSING_VALUE = -9999.0
@@ -66,7 +66,7 @@ def parse_args():
 # CLASSES
 class MyGeoDataFrame(gpd.GeoDataFrame):
     def get_SCRIP_vars(self):
-        #self.check_type()
+        # self.check_type()
         print('getting the centroid lon values')
         # centroid_lon = self.check_longitude(self.get_centroid('x'))
         centroid_lon = self.get_centroid('x')
@@ -90,22 +90,23 @@ class MyGeoDataFrame(gpd.GeoDataFrame):
         # first get the number of corners for each polygon
         # subtract 1 because last vertex is double
         print('getting number of vertices for each cell')
-        lengths = [len(item.exterior.coords.xy[0]) - 1 if item.geom_type == 'Polygon' else len(item[0].exterior.coords.xy[0]) - 1 for item in self.geometry]
+        lengths = [len(item.exterior.coords.xy[0]) - 1 if item.geom_type == 'Polygon' else len(
+            item[0].exterior.coords.xy[0]) - 1 for item in self.geometry]
         max_length = max(lengths)
         # init the final arrays and set the default missing value
         corner_lon = np.zeros((len(self.geometry), max_length))
-        corner_lon[corner_lon==0]= MISSING_VALUE
+        corner_lon[corner_lon == 0] = MISSING_VALUE
         corner_lat = np.zeros((len(self.geometry), max_length))
-        corner_lat[corner_lat==0]= MISSING_VALUE
+        corner_lat[corner_lat == 0] = MISSING_VALUE
         # now loop over each polygon and iteratively set the corner values in the target array
         # TODO: sort the nodes so the centroid is always left of the node path
         # https://stackoverflow.com/questions/1165647/how-to-determine-if-a-list-of-polygon-points-are-in-clockwise-order
         for i_item, item in enumerate(self.geometry):
             if item.geom_type == 'Polygon':
-                corner_lon[i_item,:lengths[i_item]], corner_lat[i_item,:lengths[i_item]] = \
+                corner_lon[i_item, :lengths[i_item]], corner_lat[i_item, :lengths[i_item]] = \
                     self._check_order(*item.exterior.coords.xy)
             elif item.geom_type == 'MultiPolygon':
-                corner_lon[i_item,:lengths[i_item]], corner_lat[i_item,:lengths[i_item]] = \
+                corner_lon[i_item, :lengths[i_item]], corner_lat[i_item, :lengths[i_item]] = \
                     self._check_order(*item[0].exterior.coords.xy)
             else:
                 raise
@@ -115,7 +116,7 @@ class MyGeoDataFrame(gpd.GeoDataFrame):
     def check_longitude(self, lon_arg):
         if lon_arg.min() <= 0:
             warnings.warn('Longitude values are ranging from -180 to 180, converting to range 0 to 360')
-            lon_arg= np.where(lon_arg<0,lon_arg+360, lon_arg)
+            lon_arg = np.where(lon_arg < 0, lon_arg + 360, lon_arg)
         return lon_arg
 
     def _check_order(self, lons, lats):
@@ -140,11 +141,11 @@ class MyGeoDataFrame(gpd.GeoDataFrame):
         min_index = np.argmin(lats)
         while True:
             # all neighboring indices
-            indices = [(min_index-1)%len(lats), min_index, (min_index+1)%len(lats)]
+            indices = [(min_index - 1) % len(lats), min_index, (min_index + 1) % len(lats)]
             # calculate determinant as here: https://en.wikipedia.org/wiki/Curve_orientation
             det = \
-                (lons[indices[1]] - lons[indices[0]])*(lats[indices[2]] - lats[indices[0]]) - \
-                (lons[indices[2]] - lons[indices[0]])*(lats[indices[1]] - lats[indices[0]])
+                (lons[indices[1]] - lons[indices[0]]) * (lats[indices[2]] - lats[indices[0]]) - \
+                (lons[indices[2]] - lons[indices[0]]) * (lats[indices[1]] - lats[indices[0]])
             if det == 0:
                 # the three points are colinear, check next vertices
                 min_index = indices[2]
@@ -156,6 +157,7 @@ class MyGeoDataFrame(gpd.GeoDataFrame):
                 # sort ascending
                 return lons[::-1], lats[::-1]
 
+
 def handle_latlon_format(ds):
     coord_var = ['south_north', 'west_east']
     # select all data_vars that have dimension grid_size
@@ -176,13 +178,13 @@ def handle_latlon_format(ds):
             df.columns = [data_var]
         dfs.append(df)
     # we turn the pandas DataFrame in a Geopandas GeoDataFrame and still need geometry
-    gdf = gpd.GeoDataFrame(pd.concat(dfs, axis=1), geometry=[[]]*len(dfs[0]))
+    gdf = gpd.GeoDataFrame(pd.concat(dfs, axis=1), geometry=[[]] * len(dfs[0]))
     # convert units to degree
     # loop over polygons and calculate number of edges and set geometry item
     for i, j in product(ds[coord_var[0]], ds[coord_var[1]]):
         xvals = ds['XLONG_bnds'].isel(**{coord_var[1]: j}).values
         yvals = ds['XLAT_bnds'].isel(**{coord_var[0]: i}).values
-        gdf.loc['{}_{}'.format(float(i),float(j)), 'geometry'] = Polygon([
+        gdf.loc['{}_{}'.format(float(i), float(j)), 'geometry'] = Polygon([
             (xvals[0], yvals[0]),
             (xvals[1], yvals[0]),
             (xvals[1], yvals[1]),
@@ -227,10 +229,10 @@ if __name__ == '__main__':
 
         ds = xr.Dataset(
             data_vars={'grid_corner_lon': (['grid_size', 'grid_corners'], corner_lon),
-             'grid_corner_lat': (['grid_size', 'grid_corners'], corner_lat),
-             'grid_center_lon': (['grid_size'], centroid_lon),
-             'grid_center_lat': (['grid_size'], centroid_lat),
-             },
+                       'grid_corner_lat': (['grid_size', 'grid_corners'], corner_lat),
+                       'grid_center_lon': (['grid_size'], centroid_lon),
+                       'grid_center_lat': (['grid_size'], centroid_lat),
+                       },
             attrs={'title': args.name, 'units': units}
         )
         # add units globally and to each variable
@@ -258,7 +260,7 @@ if __name__ == '__main__':
         EXCLUDE_DIMS = ['grid_center_lon', 'grid_center_lat', 'grid_corner_lon', 'grid_corner_lat', 'grid_imask']
         # all MPI-ICON based variable names and its SCRIP counterparts
         RENAME_VAR_DICT = {'clon': 'grid_center_lon', 'clat': 'grid_center_lat',
-                        'clon_vertices': 'grid_corner_lon', 'clat_vertices': 'grid_corner_lat'}
+                           'clon_vertices': 'grid_corner_lon', 'clat_vertices': 'grid_corner_lat'}
         # all MPI-ICON based dimension names and its SCRIP counterparts
         RENAME_DIM_DICT = {'cell': 'grid_size', 'nv': 'grid_corners'}
         # all MPI-ICON based data variable names
@@ -274,17 +276,18 @@ if __name__ == '__main__':
             # do we have special MPI-ICON names?
             if all(mpi_icon_format):
                 # select and rename
-                ds = ds[list(RENAME_VAR_DICT.keys())+SELECT_VARS].rename({**RENAME_VAR_DICT, **RENAME_DIM_DICT})
+                ds = ds[list(RENAME_VAR_DICT.keys()) + SELECT_VARS].rename({**RENAME_VAR_DICT, **RENAME_DIM_DICT})
                 # set the grid_imask property based on sea_land_mask or to default 1
                 if 'cell_sea_land_mask' in ds:
-                    ds['grid_imask'] = xr.where(ds['cell_sea_land_mask']>1, 1, 0)
+                    ds['grid_imask'] = xr.where(ds['cell_sea_land_mask'] > 1, 1, 0)
                 else:
                     ds['grid_imask'] = (('grid_size',), np.ones_like(ds['grid_center_lon'].values, dtype=int))
             # select all data_vars that have dimension grid_size
-            shp_vars = [data_var for data_var in ds.data_vars if 'grid_size' in ds[data_var].dims and data_var not in EXCLUDE_DIMS]
+            shp_vars = [data_var for data_var in ds.data_vars if
+                        'grid_size' in ds[data_var].dims and data_var not in EXCLUDE_DIMS]
             if not shp_vars:
-                ds['id'] = (('grid_size',), np.array(range(1,len(ds['grid_size'])+1)))
-                shp_vars =['id']
+                ds['id'] = (('grid_size',), np.array(range(1, len(ds['grid_size']) + 1)))
+                shp_vars = ['id']
             # empty container
             dfs = []
             for data_var in shp_vars:
@@ -300,7 +303,10 @@ if __name__ == '__main__':
                     df.columns = [data_var]
                 dfs.append(df)
             # we turn the pandas DataFrame in a Geopandas GeoDataFrame and still need geometry
-            gdf = gpd.GeoDataFrame(pd.concat(dfs, axis=1))
+            gdf = gpd.GeoDataFrame(pd.concat(dfs, axis=1, keys=shp_vars))
+            # merge the MultiIndex into a flat index
+            gdf.columns = ['_'.join((str(item) for item in _)) for _ in gdf.columns]
+
             # convert units to degree
             n_corners = len(ds['grid_corners'])
             for var_name in EXCLUDE_DIMS:
@@ -323,7 +329,7 @@ if __name__ == '__main__':
                         ds['grid_corner_lat'].isel(grid_size=i_polygon)))
 
         # set WGS84 by default
-        gdf.crs = {'init' :'epsg:4326'}
+        gdf.crs = {'init': 'epsg:4326'}
         gdf.to_file(args.output_file)
     else:
         raise Exception('Did not get proper filenames. Script only works for conversion of *.shp->*.nc or vice versa')