implemented simple storm tracking

enstools/feature/identification/_proto_gen/storm_pb2.py

+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tmprt10xltq
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0btmprt10xltq\"\"\n\nProperties\x12\x14\n\x0cmin_pressure\x18\x01 \x02(\x02\")\n\tVoxelData\x12\r\n\x05index\x18\x01 \x03(\r\x12\r\n\x05value\x18\x02 \x01(\x02\"-\n\nVertexData\x12\x10\n\x08position\x18\x01 \x03(\x02\x12\r\n\x05value\x18\x02 \x01(\x02\"+\n\x08\x46\x61\x63\x65\x44\x61ta\x12\x10\n\x08vertices\x18\x01 \x03(\r\x12\r\n\x05value\x18\x02 \x01(\x02\"C\n\x13VoxelRepresentation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12\x1e\n\nvoxel_data\x18\x02 \x03(\x0b\x32\n.VoxelData\"D\n\x12LineRepresentation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12 \n\x0bvertex_data\x18\x02 \x03(\x0b\x32\x0b.VertexData\"f\n\x16\x42oundaryRepresentation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12 \n\x0bvertex_data\x18\x02 \x03(\x0b\x32\x0b.VertexData\x12\x1c\n\tface_data\x18\x03 \x03(\x0b\x32\t.FaceData\"2\n\tGraphNode\x12\x0c\n\x04time\x18\x01 \x02(\t\x12\x17\n\x06object\x18\x02 \x02(\x0b\x32\x07.Object\"K\n\x0fGraphConnection\x12\x1a\n\x06parent\x18\x01 \x02(\x0b\x32\n.GraphNode\x12\x1c\n\x08\x63hildren\x18\x02 \x03(\x0b\x32\n.GraphNode\".\n\x0bObjectGraph\x12\x1f\n\x05\x65\x64ges\x18\x01 \x03(\x0b\x32\x10.GraphConnection\"\xb4\x01\n\x06Object\x12\n\n\x02id\x18\x01 \x02(\x05\x12%\n\x08line_rep\x18\x02 \x03(\x0b\x32\x13.LineRepresentation\x12\'\n\tvoxel_rep\x18\x03 \x01(\x0b\x32\x14.VoxelRepresentation\x12-\n\x0c\x62oundary_rep\x18\x04 \x03(\x0b\x32\x17.BoundaryRepresentation\x12\x1f\n\nproperties\x18\x05 \x01(\x0b\x32\x0b.Properties\"8\n\x08Timestep\x12\x12\n\nvalid_time\x18\x01 \x01(\t\x12\x18\n\x07objects\x18\x02 \x03(\x0b\x32\x07.Object\"\x99\x01\n\x0cTrackableSet\x12\x11\n\tinit_time\x18\x01 \x01(\t\x12\x0e\n\x06member\x18\x02 \x01(\r\x12\r\n\x05level\x18\x03 \x01(\x02\x12\x1c\n\ttimesteps\x18\x04 \x03(\x0b\x32\t.Timestep\x12\x1b\n\x05graph\x18\x05 \x01(\x0b\x32\x0c.ObjectGraph\x12\x1c\n\x06tracks\x18\x06 \x03(\x0b\x32\x0c.ObjectGraph\"_\n\x12\x44\x61tasetDescription\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04\x66ile\x18\x02 \x01(\t\x12\x10\n\x08run_time\x18\x03 \x02(\t\x12\x1b\n\x04sets\x18\x04 \x03(\x0b\x32\r.TrackableSet')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'tmprt10xltq_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _PROPERTIES._serialized_start=15
+  _PROPERTIES._serialized_end=49
+  _VOXELDATA._serialized_start=51
+  _VOXELDATA._serialized_end=92
+  _VERTEXDATA._serialized_start=94
+  _VERTEXDATA._serialized_end=139
+  _FACEDATA._serialized_start=141
+  _FACEDATA._serialized_end=184
+  _VOXELREPRESENTATION._serialized_start=186
+  _VOXELREPRESENTATION._serialized_end=253
+  _LINEREPRESENTATION._serialized_start=255
+  _LINEREPRESENTATION._serialized_end=323
+  _BOUNDARYREPRESENTATION._serialized_start=325
+  _BOUNDARYREPRESENTATION._serialized_end=427
+  _GRAPHNODE._serialized_start=429
+  _GRAPHNODE._serialized_end=479
+  _GRAPHCONNECTION._serialized_start=481
+  _GRAPHCONNECTION._serialized_end=556
+  _OBJECTGRAPH._serialized_start=558
+  _OBJECTGRAPH._serialized_end=604
+  _OBJECT._serialized_start=607
+  _OBJECT._serialized_end=787
+  _TIMESTEP._serialized_start=789
+  _TIMESTEP._serialized_end=845
+  _TRACKABLESET._serialized_start=848
+  _TRACKABLESET._serialized_end=1001
+  _DATASETDESCRIPTION._serialized_start=1003
+  _DATASETDESCRIPTION._serialized_end=1098
+# @@protoc_insertion_point(module_scope)

enstools/feature/identification/storm/__init__.py

+from .identification import StormIdentification

enstools/feature/identification/storm/identification.py

+from ..identification import IdentificationStrategy
+
+import xarray as xr
+from random import randrange
+from scipy import ndimage as ndi
+import numpy as np
+from enstools.feature.identification.storm.util import *
+
+
+class StormIdentification(IdentificationStrategy):
+
+    def __init__(self, start_pressure_thr=965, step_pressure_thr=2, size_thr=10000, amplitude_thr=2, **kwargs):
+        # Constructor. Called from example_template.py, parameters can be passed and set here.
+
+        self.start_pressure_thr = start_pressure_thr
+        self.step_pressure_thr = step_pressure_thr
+        self.size_thr = size_thr
+        self.amplitude_thr = amplitude_thr
+
+        pass
+
+    def precompute(self, dataset: xr.Dataset, **kwargs):
+        plt.switch_backend('agg')  # this is thread safe matplotlib but cant display.
+
+        dataset['storm_areas'] = xr.zeros_like(dataset['msl'], dtype=int)
+        return dataset
+
+    def identify(self, dataset: xr.Dataset, **kwargs):
+
+        msl_data = dataset.msl
+        considered_mask = xr.zeros_like(msl_data, dtype=bool)
+
+        pressure_thrs = np.arange(self.start_pressure_thr, 1000, self.step_pressure_thr)
+
+        storm_id = 1
+        obj_list = []
+        for pressure_thr in pressure_thrs:
+
+            lp_area = get_msl_areas_by_hPa_thr(msl_data, pressure_thr)
+
+            storm_areas = split_areas(lp_area)
+
+            for storm_area in storm_areas:
+
+                if np.any(np.logical_and(considered_mask.data, storm_area.data)):
+                    continue
+
+                storm_size = get_storm_size_km2(storm_area)
+                if storm_size < self.size_thr:
+                    continue
+
+                if not has_local_minimum(storm_area, msl_data):
+                    continue
+
+                storm_amplitude = get_storm_amplitude_hPa(storm_area, msl_data)
+                if storm_amplitude < self.amplitude_thr:
+                    continue
+
+                # keep storm
+                dataset['storm_areas'].data[storm_area] = storm_id
+                considered_mask.data[storm_area] = True
+
+                # get an instance of a new object, can pass an ID or set in manually afterwards
+                obj = self.get_new_object()
+                # set some ID to it
+                obj.id = storm_id
+
+                # get properties of object and populate them (like defined in template.proto)
+                properties = obj.properties
+
+                properties.min_pressure = get_min_pressure_hPa(storm_area, msl_data)
+
+                obj_list.append(obj)
+                storm_id += 1
+
+        # return the dataset (can be changed here), and the list of objects
+        return dataset, obj_list
+
+    def postprocess(self, dataset: xr.Dataset, obj_desc, **kwargs):
+
+        # obj_desc = self.make_ids_unique(obj_desc)
+        return dataset, obj_desc

enstools/feature/identification/storm/run_identify.py

+# Usage Example
+
+from enstools.feature.pipeline import FeaturePipeline
+from enstools.feature.identification.storm import StormIdentification
+from enstools.feature.tracking.overlap_tracking import OverlapTracking
+from enstools.feature.identification._proto_gen import storm_pb2
+from os.path import expanduser
+from enstools.feature.identification.storm.util import plot_timestep
+
+# set the pb_reference to the compiled pb2.py file (see proto_gen directory)
+pipeline = FeaturePipeline(storm_pb2, processing_mode='2d')
+
+# change this to an identification strategy that actually does something: existing one or implement your own
+i_strat = StormIdentification(some_parameter='foo') # set the Identification strategy
+t_strat = OverlapTracking(field_name='storm_areas') # set the tracking strategy
+
+pipeline.set_identification_strategy(i_strat)
+pipeline.set_tracking_strategy(t_strat) # or None as argument if no tracking
+
+path = expanduser("~") + '/phd/data/framework_example_ds/zeynep.nc'
+pipeline.set_data_path(path)
+
+# execute pipeline
+pipeline.execute()
+
+od = pipeline.get_object_desc()
+
+reanalysis_set = od.sets[0]
+
+# subset = get_subset_by_description(pipeline.get_data(), reanalysis_set, '2d')
+
+for ts in reanalysis_set.timesteps:
+    ds_t = pipeline.get_data().sel(time=ts.valid_time)
+
+    plot_timestep(ds_t, str(ts.valid_time)[:13] + '.png')
+
+
+out_json_path = path[:-3] + '_desc.json'
+out_ds_path = path[:-3] + '_desc.nc'
+out_graph_path = path[:-3] + '_graph.json'
+
+pipeline.save_result(description_type='json', description_path=out_json_path, dataset_path=out_ds_path) # dataset_path=...
+
+

enstools/feature/identification/storm/storm.proto

+syntax = "proto2";
+
+message Properties {
+  required float min_pressure = 1;
+
+}
+

enstools/feature/identification/storm/util.py

+from matplotlib import pyplot as plt
+from cartopy import crs as ccrs
+import cartopy.feature as cfeature
+import numpy as np
+from matplotlib import colors as mcolors
+from scipy import ndimage as ndi
+import xarray as xr
+from pyproj import Geod
+from shapely.geometry import Polygon
+from shapely.ops import orient
+
+binary_grey = mcolors.ListedColormap(['white', 'grey'])
+
+def plot_timestep(ds, file_name):
+
+    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(11, 4), subplot_kw=dict(projection=ccrs.PlateCarree()))
+
+    levels = np.linspace(97000, 104000, 14)
+
+    ds.storm_areas.plot.contourf(levels=[0.5, 1000], cmap=binary_grey, add_colorbar=False) # , transform=ccrs.NorthPolarStereo())
+    ds.msl.plot.contour(levels=levels, vmin=0, extend='max', cmap='Blues')
+
+    ax.coastlines()
+    ax.add_feature(cfeature.BORDERS.with_scale('50m'))
+    # ax.set_extent(extent, crs=ccrs.PlateCarree())
+    # yt1 = ax.set_yticks(y_ticks, crs=ccrs.PlateCarree())
+    # xt1 = ax.set_xticks(x_ticks, crs=ccrs.PlateCarree())
+
+
+
+    print("Save to " + file_name)
+    plt.savefig(file_name, format='png')
+
+    plt.figure().clear()
+    plt.close()
+    plt.cla()
+    plt.clf()
+
+
+def get_min_pressure_hPa(storm_area, msl_data):
+    wh = msl_data.data[storm_area.data]
+    return np.amin(wh) / 100.0
+
+def get_msl_areas_by_hPa_thr(dataset, hPa):
+    return dataset < hPa * 100.0
+
+def split_areas(areas):
+    # split dataarray areas into list of cohesive regions (DAs)
+    lp_area_lab = xr.zeros_like(areas, dtype=int)
+
+    lp_area_lab.data, num_features = ndi.label(areas)
+    a_list = []
+    for i in range(1, num_features + 1):
+        a_list.append(lp_area_lab == i)
+
+    return a_list
+
+
+def get_storm_size_km2(storm_da):
+
+    # create contour
+    res = storm_da.plot.contour(levels=[0.5, 1000])
+    paths = res.collections[0].get_paths()
+
+    if len(paths) > 1:
+        print("Multi path?") # TODO
+    path = paths[0]
+
+    #for v_idx in range(len(path.vertices) - 1):
+    #            start_lonlat = path.vertices[v_idx][0], path.vertices[v_idx][1]
+    #            end_lonlat = path.vertices[v_idx + 1][0], path.vertices[v_idx + 1][1]
+
+    node_list = path.vertices
+
+    polygon = Polygon(node_list)
+    geod = Geod(ellps="WGS84")
+    poly_area, poly_perimeter = geod.geometry_area_perimeter(orient(polygon))
+
+    return poly_area / 1000000.0 # in km^2
+
+def has_local_minimum(storm_area, msl_data):
+
+    all_msl = msl_data.data[storm_area.data]
+    min_all_data = np.amin(all_msl)
+
+    inner_area = ndi.binary_erosion(storm_area.data)
+    inner_msl = msl_data.data[inner_area]
+
+    if inner_msl.size == 0:
+        return False # too small
+
+    min_inner_data = np.amin(inner_msl)
+    return min_all_data == min_inner_data # true if min not on border
+
+
+def get_storm_amplitude_hPa(storm_area, msl_data):
+
+    all_msl = msl_data.data[storm_area.data]
+    min_p = np.amin(all_msl)
+    avg_p = np.mean(all_msl)
+    ampl_p = avg_p - min_p
+
+    return ampl_p / 100.0
+