Merge branch 'aew_development' of...

Merge branch 'aew_development' of https://gitlab.physik.uni-muenchen.de/Christoph.Fischer/enstools-feature into aew_development

enstools/feature/identification/aew_vortices/configuration.py

@@ -36,8 +36,8 @@ base_threshold = {'cv': 3e-5, 'cva': 1e-5, 'rv': 3e-5}
 # time of interest, if None all
 # june-oct is AEW season
 
-start_date = '2021-08-23T00:00'  # '2022-08-01T00:00' # None # '2008-08-01T00:00' #  # '2008-08-01T00:00'
-end_date = '2021-09-03T00:00'  # '2022-08-15T00:00'  # None # '2008-08-15T00:00' # None # '2008-08-03T00:00'
+start_date = '2021-08-27T00:00'  # '2022-08-01T00:00' # None # '2008-08-01T00:00' #  # '2008-08-01T00:00'
+end_date = '2021-08-28T10:00'  # '2022-08-15T00:00'  # None # '2008-08-15T00:00' # None # '2008-08-03T00:00'
 
 radius = 500000  # 500 km
 prominence_radius = 400000  # 300 km

enstools/feature/identification/aew_vortices/identification.py

+import IMTreatment.vortex_detection.vortex_detection
 import skimage.morphology
 
 from enstools.feature.identification import IdentificationStrategy
@@ -22,6 +23,12 @@ from skimage.draw import line
 from IMTreatment.vortex_criterions import get_lambda2, get_gamma
 from IMTreatment.core import VectorField, ScalarField, SpatialScalarFields, SpatialVectorFields
 
+def round_off_rating(number):
+    """Round a number to the closest half integer.
+    """
+    return round(number * 2) / 2
+
+
 class AEWVortexIdentification(IdentificationStrategy):
 
     def __init__(self, **kwargs):
@@ -93,12 +100,14 @@ class AEWVortexIdentification(IdentificationStrategy):
 
         # BPF (optional) TODO
         if self.bpf:
+            print(dataset)
             dataset = bpf(dataset, u_name, v_name, min_days=2, max_days=6)
             self.config.u_dim = 'u_bpf'
             self.config.v_dim = 'v_bpf'
             u_name = 'u_bpf'
             v_name = 'v_bpf'
 
+        print(dataset)
         self.json_desc = interpolate_wavetroughs(self.json_desc, african_easterly_waves_pb2, self.pb_dataset)
 
         # create field to use
@@ -106,34 +115,45 @@ class AEWVortexIdentification(IdentificationStrategy):
         dataset = dataset.transpose(..., lat_str, lon_str)
 
         data_field = dataset[self.field]
-        dataset['vortices'] = xr.zeros_like(dataset[self.field], dtype=int)
+        # dataset['vortices'] = xr.zeros_like(dataset[self.field], dtype=int)
+        dataset['saddle'] = xr.zeros_like(dataset[self.field], dtype=int)
+        dataset['foci_c'] = xr.zeros_like(dataset[self.field], dtype=int)
 
-        """
         l2 = xr.zeros_like(dataset[u_name], dtype=float)
+        # TODO compute CPs in parallel in identify()
         for time in dataset.time.values:
-            print(time)
-            dt = dataset.sel(level=850, time=time)
-
-            sx = ScalarField()
-            sx.import_from_arrays(dt.latitude.values.tolist(), dt.longitude.values.tolist(), dt[u_name].data)
-            sy = ScalarField()
-            sy.import_from_arrays(dt.latitude.values.tolist(), dt.longitude.values.tolist(), dt[v_name].data)
-
-            v = VectorField()
-            v.import_from_sfs(sx, sy)
-
-            cur_l2 = get_gamma(v)
-            print(cur_l2)
-            l2.loc[dict(level=850, time=time)] = cur_l2.values
-        print(l2)
-        dataset['lambda2'] = l2
-        dataset['lambda2'].attrs['long_name'] = 'lambda2'
-        dataset['lambda2'].attrs['standard_name'] = 'l2'
-        dataset['lambda2'].attrs['units'] = 'dunno'
-        """
+            for lev in dataset.level.values:
+                print(time)
+                print(lev)
+                dt = dataset.sel(level=lev, time=time)
+                cp = dt.saddle
+                cp_c = dt.foci_c
+
+                sx = ScalarField()
+                sx.import_from_arrays(dt.latitude.values.tolist(), dt.longitude.values.tolist(), dt[u_name].data)
+                sy = ScalarField()
+                sy.import_from_arrays(dt.latitude.values.tolist(), dt.longitude.values.tolist(), dt[v_name].data)
+
+                v = VectorField()
+                v.import_from_sfs(sx, sy)
+
+                cpp = IMTreatment.vortex_detection.vortex_detection.get_critical_points(v)
+                foci = cpp.sadd[0].xy
+                foci_c = cpp.foc_c[0].xy
+
+                for i in range(foci.shape[0]):
+                    cp.loc[dict(longitude=round_off_rating(foci[i][1]), latitude=round_off_rating(foci[i][0]))] = 1
+
+                for i in range(foci_c.shape[0]):
+                    cp_c.loc[dict(longitude=round_off_rating(foci_c[i][1]), latitude=round_off_rating(foci_c[i][0]))] = 1
+
+                dt['saddle'] = cp
+                dt['foci_c'] = cp_c
+                dataset.loc[dict(level=lev, time=time)] = dt
 
         dataset = create_wt_troughs_and_areas(dataset, self.json_desc, self.circles, False, d=self.config.radius)
 
+        """
         # search local max.
         neighborhood = generate_binary_structure(2, 2)
         # TODO depending on num dimensions. here 4d, with latlon back.
@@ -145,6 +165,7 @@ class AEWVortexIdentification(IdentificationStrategy):
         # max_lon = find_peaks(data_field.values, np.greater, axis=-1) #  scipy.signal.find_peaks
 
         dataset['vortices'].data = max_mask
+        """
 
         # out: set data to 1, fakes to 2
         dataset['wt_rea'] = xr.where(dataset.wt_rea > 0, 1, dataset.wt_rea)
@@ -156,48 +177,24 @@ class AEWVortexIdentification(IdentificationStrategy):
         thr = self.config.base_threshold[self.field]
         if self.bpf:
             thr = thr / 2 # TODO?
+
+        """
         dataset['vortices'] = xr.where(dataset[self.field] > thr, dataset.vortices, 0)
         # FILTER by area -> in vicinity of WT
         dataset['vortices'] = xr.where(dataset['infl_rea'] != 0, dataset.vortices, 0)
         dataset['vortices'] = dataset.vortices.transpose(..., lat_str, lon_str)
-
-
-        # self.config.prominence_radius
-
-
-
-        # filter vortices somehow, before? prob need after
-        # add them to WT descs.
-
-        # filter max points by WT mask
-        # interpolate_wave_areas() # for gaps in tracks. -> also add json info? maybe create WT area ourself.
-        #     for each child in track node: if dt > 6h: add inbetween "fake" wt to json
-        # put lows in.
-        # maybe depending on whether they form a track?
-
-
-        #if self.config.cv_name not in dataset.data_vars:
-        #    print("Curvature Vorticity not found, trying to compute it out of U and V...")
-        #    dataset = compute_cv(dataset, u_name, v_name, self.config.cv_name)
-
-        # make dataset to 2.5 (or same as cv_clim)
-        # dataset = dataset.interp({lat_str: cv_clim.coords[lat_str], lon_str: cv_clim.coords[lon_str]})
-
-        # make sure that lat and lon are last two dimensions
-        # dataset = dataset.transpose(..., lat_str, lon_str)
-
-
-        # dataset['wavetroughs'] = wt
+        """
 
         return dataset
 
     def identify(self, data_chunk: xr.Dataset, **kwargs):
 
         objs = []
-        v = data_chunk.vortices
+        # v = data_chunk.vortices
+        cp = data_chunk.saddle
 
         ds = data_chunk.stack(z=["latitude", "longitude"])
-        ds = ds.where(ds.vortices > 0, drop=True)
+        ds = ds.where(ds.saddle > 0, drop=True)
 
         # TODO maybe create WT+area also in here, so we have direct reference to IDs
         # CLUSTER
@@ -205,7 +202,7 @@ class AEWVortexIdentification(IdentificationStrategy):
         #                    -> or discard if any peak within 2deg is higher? or is this the same?
         # For each vortex...
         for z in range(len(ds.z.data)):
-            if ds.vortices[z].item() == 0:
+            if ds.saddle[z].item() == 0:
                 print("skip?")
                 continue
             lon_vortex = ds.longitude[z].item()
@@ -222,9 +219,9 @@ class AEWVortexIdentification(IdentificationStrategy):
 
                 distance_c = circle.sel(longitude=lon_vc, latitude=lat_vc).item()
                 # if distance in promenance radius and other point has higher vorticity -> delete this point.
-                if distance_c < self.config.prominence_radius and data_c > data_vertex: # TODO what happens if two with exactly the same cv?
+                if distance_c < self.config.prominence_radius and data_c > data_vertex:
                     # vortex in vicinity is higher. discard this.
-                    v.loc[dict(longitude=lon_vortex, latitude=lat_vortex)] = 0
+                    cp.loc[dict(longitude=lon_vortex, latitude=lat_vortex)] = 0
                     break
 
         # data_chunk.vortices.values = skimage.morphology.binary_dilation(data_chunk.vortices.values)
@@ -235,25 +232,6 @@ class AEWVortexIdentification(IdentificationStrategy):
 
 
 
-        """
-        id_ = 1
-        for path in []:
-
-            # get new object, set id
-            # o = self.get_new_object()
-            # o.id = id_
-            # populate it
-            # populate_object(o.properties, path, self.config)
-
-            # add to objects if keep
-            # if not self.keep_wavetrough(o.properties):
-            #     continue
-
-            # objs.append(o)
-            # id_ += 1
-            break
-        """
-
         return data_chunk, objs
 
     def postprocess(self, dataset: xr.Dataset, data_desc, **kwargs):

enstools/feature/identification/aew_vortices/plotting.py

@@ -131,7 +131,9 @@ def plot_kw_style(dataset, dataset_desc, config, lev=700):
 
         # plot vortices
         # ds_t.sel(level=700).vortices.plot.contourf(levels=[-0.5,0.5,99], colors=('#00000000', 'blue'), subplot_kws={'transform_first': True}, add_colorbar=False)
-        ds_t.sel(level=lv).vortices.plot.contourf(levels=[-0.5,0.5,99], colors=('#00000000', 'orange'), subplot_kws={'transform_first': True}, add_colorbar=False)
+        ds_t.sel(level=lv).saddle.plot.contourf(levels=[-0.5,0.5,99], colors=('#00000000', 'green'), subplot_kws={'transform_first': True}, add_colorbar=False)
+        ds_t.sel(level=lv).foci_c.plot.contourf(levels=[-0.5, 0.5, 99], colors=('#00000000', 'grey'),
+                                            subplot_kws={'transform_first': True}, add_colorbar=False)
 
         # ds_t.prec_rate_rea.plot.contourf(levels=levels_rain, extend='max', subplot_kws={'transform_first': True},
         #                  cmap=rain_cm, norm=norm, add_colorbar=False)