diff --git a/.github/workflows/check_formatting.yml b/.github/workflows/check_formatting.yml new file mode 100644 index 00000000..0585c678 --- /dev/null +++ b/.github/workflows/check_formatting.yml @@ -0,0 +1,14 @@ +name: check_formatting +on: [push, pull_request] +jobs: + formatting_job: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: '3.8' + - run: pip install . + - run: pip install black + - run: black tobac --check + diff --git a/setup.py b/setup.py index 5f9cd927..ed14b2ac 100644 --- a/setup.py +++ b/setup.py @@ -1,12 +1,14 @@ from setuptools import setup -setup(name='tobac', - version='1.2', - description='Tracking and object-based analysis of clouds', - url='http://github.com/climate-processes/tobac', - author='Max Heikenfeld', - author_email='max.heikenfeld@physics.ox.ac.uk', - license='GNU', - packages=['tobac'], - install_requires=[], - zip_safe=False) +setup( + name="tobac", + version="1.2", + description="Tracking and object-based analysis of clouds", + url="http://github.com/climate-processes/tobac", + author="Max Heikenfeld", + author_email="max.heikenfeld@physics.ox.ac.uk", + license="GNU", + packages=["tobac"], + install_requires=[], + zip_safe=False, +) diff --git a/tobac/__init__.py b/tobac/__init__.py index f108e7a9..d01e2689 100644 --- a/tobac/__init__.py +++ b/tobac/__init__.py @@ -1,23 +1,63 @@ -#from .tracking import maketrack -from .segmentation import segmentation_3D, segmentation_2D,watershedding_3D,watershedding_2D -from .centerofgravity import calculate_cog,calculate_cog_untracked,calculate_cog_domain -from .plotting import plot_tracks_mask_field,plot_tracks_mask_field_loop,plot_mask_cell_track_follow,plot_mask_cell_track_static,plot_mask_cell_track_static_timeseries -from .plotting import plot_lifetime_histogram,plot_lifetime_histogram_bar,plot_histogram_cellwise,plot_histogram_featurewise -from .plotting import plot_mask_cell_track_3Dstatic,plot_mask_cell_track_2D3Dstatic -from .plotting import plot_mask_cell_individual_static,plot_mask_cell_individual_3Dstatic +# from .tracking import maketrack +from .segmentation import ( + segmentation_3D, + segmentation_2D, + watershedding_3D, + watershedding_2D, +) +from .centerofgravity import ( + calculate_cog, + calculate_cog_untracked, + calculate_cog_domain, +) +from .plotting import ( + plot_tracks_mask_field, + plot_tracks_mask_field_loop, + plot_mask_cell_track_follow, + plot_mask_cell_track_static, + plot_mask_cell_track_static_timeseries, +) +from .plotting import ( + plot_lifetime_histogram, + plot_lifetime_histogram_bar, + plot_histogram_cellwise, + plot_histogram_featurewise, +) +from .plotting import plot_mask_cell_track_3Dstatic, plot_mask_cell_track_2D3Dstatic +from .plotting import ( + plot_mask_cell_individual_static, + plot_mask_cell_individual_3Dstatic, +) from .plotting import animation_mask_field from .plotting import make_map, map_tracks -from .analysis import cell_statistics,cog_cell,lifetime_histogram,histogram_featurewise,histogram_cellwise -from .analysis import calculate_velocity,calculate_distance,calculate_area +from .analysis import ( + cell_statistics, + cog_cell, + lifetime_histogram, + histogram_featurewise, + histogram_cellwise, +) +from .analysis import calculate_velocity, calculate_distance, calculate_area from .analysis import calculate_nearestneighbordistance -from .analysis import velocity_histogram,nearestneighbordistance_histogram,area_histogram +from .analysis import ( + velocity_histogram, + nearestneighbordistance_histogram, + area_histogram, +) from .analysis import calculate_overlap -from .utils import mask_cell,mask_cell_surface,mask_cube_cell,mask_cube_untracked,mask_cube,column_mask_from2D,get_bounding_box -from .utils import mask_features,mask_features_surface,mask_cube_features +from .utils import ( + mask_cell, + mask_cell_surface, + mask_cube_cell, + mask_cube_untracked, + mask_cube, + column_mask_from2D, + get_bounding_box, +) +from .utils import mask_features, mask_features_surface, mask_cube_features -from .utils import add_coordinates,get_spacings +from .utils import add_coordinates, get_spacings from .feature_detection import feature_detection_multithreshold from .tracking import linking_trackpy from .wrapper import maketrack from .wrapper import tracking_wrapper - diff --git a/tobac/analysis.py b/tobac/analysis.py index 1c8bc8e4..11304bdd 100644 --- a/tobac/analysis.py +++ b/tobac/analysis.py @@ -3,162 +3,241 @@ import logging import os -from .utils import mask_cell,mask_cell_surface,mask_cube_cell,get_bounding_box +from .utils import mask_cell, mask_cell_surface, mask_cube_cell, get_bounding_box -def cell_statistics_all(input_cubes,track,mask,aggregators,output_path='./',cell_selection=None,output_name='Profiles',width=10000,z_coord='model_level_number',dimensions=['x','y'],**kwargs): + +def cell_statistics_all( + input_cubes, + track, + mask, + aggregators, + output_path="./", + cell_selection=None, + output_name="Profiles", + width=10000, + z_coord="model_level_number", + dimensions=["x", "y"], + **kwargs +): if cell_selection is None: - cell_selection=np.unique(track['cell']) - for cell in cell_selection : - cell_statistics(input_cubes=input_cubes,track=track, mask=mask, - dimensions=dimensions,aggregators=aggregators,cell=cell, - output_path=output_path,output_name=output_name, - width=width,z_coord=z_coord,**kwargs) - -def cell_statistics(input_cubes,track,mask,aggregators,cell,output_path='./',output_name='Profiles',width=10000,z_coord='model_level_number',dimensions=['x','y'],**kwargs): - from iris.cube import Cube,CubeList + cell_selection = np.unique(track["cell"]) + for cell in cell_selection: + cell_statistics( + input_cubes=input_cubes, + track=track, + mask=mask, + dimensions=dimensions, + aggregators=aggregators, + cell=cell, + output_path=output_path, + output_name=output_name, + width=width, + z_coord=z_coord, + **kwargs + ) + + +def cell_statistics( + input_cubes, + track, + mask, + aggregators, + cell, + output_path="./", + output_name="Profiles", + width=10000, + z_coord="model_level_number", + dimensions=["x", "y"], + **kwargs +): + from iris.cube import Cube, CubeList from iris.coords import AuxCoord - from iris import Constraint,save - + from iris import Constraint, save + # If input is single cube, turn into cubelist if type(input_cubes) is Cube: - input_cubes=CubeList([input_cubes]) - - logging.debug('Start calculating profiles for cell '+str(cell)) - track_i=track[track['cell']==cell] - - cubes_profile={} + input_cubes = CubeList([input_cubes]) + + logging.debug("Start calculating profiles for cell " + str(cell)) + track_i = track[track["cell"] == cell] + + cubes_profile = {} for aggregator in aggregators: - cubes_profile[aggregator.name()]=CubeList() - - for time_i in track_i['time'].values: + cubes_profile[aggregator.name()] = CubeList() + + for time_i in track_i["time"].values: constraint_time = Constraint(time=time_i) - - mask_i=mask.extract(constraint_time) - mask_cell_i=mask_cell(mask_i,cell,track_i,masked=False) - mask_cell_surface_i=mask_cell_surface(mask_i,cell,track_i,masked=False,z_coord=z_coord) - - x_dim=mask_cell_surface_i.coord_dims('projection_x_coordinate')[0] - y_dim=mask_cell_surface_i.coord_dims('projection_y_coordinate')[0] - x_coord=mask_cell_surface_i.coord('projection_x_coordinate') - y_coord=mask_cell_surface_i.coord('projection_y_coordinate') - - if (mask_cell_surface_i.core_data()>0).any(): - box_mask_i=get_bounding_box(mask_cell_surface_i.core_data(),buffer=1) - - box_mask=[[x_coord.points[box_mask_i[x_dim][0]],x_coord.points[box_mask_i[x_dim][1]]], - [y_coord.points[box_mask_i[y_dim][0]],y_coord.points[box_mask_i[y_dim][1]]]] + + mask_i = mask.extract(constraint_time) + mask_cell_i = mask_cell(mask_i, cell, track_i, masked=False) + mask_cell_surface_i = mask_cell_surface( + mask_i, cell, track_i, masked=False, z_coord=z_coord + ) + + x_dim = mask_cell_surface_i.coord_dims("projection_x_coordinate")[0] + y_dim = mask_cell_surface_i.coord_dims("projection_y_coordinate")[0] + x_coord = mask_cell_surface_i.coord("projection_x_coordinate") + y_coord = mask_cell_surface_i.coord("projection_y_coordinate") + + if (mask_cell_surface_i.core_data() > 0).any(): + box_mask_i = get_bounding_box(mask_cell_surface_i.core_data(), buffer=1) + + box_mask = [ + [ + x_coord.points[box_mask_i[x_dim][0]], + x_coord.points[box_mask_i[x_dim][1]], + ], + [ + y_coord.points[box_mask_i[y_dim][0]], + y_coord.points[box_mask_i[y_dim][1]], + ], + ] else: - box_mask=[[np.nan,np.nan],[np.nan,np.nan]] - - x=track_i[track_i['time'].values==time_i]['projection_x_coordinate'].values[0] - y=track_i[track_i['time'].values==time_i]['projection_y_coordinate'].values[0] - - box_slice=[[x-width,x+width],[y-width,y+width]] - - x_min=np.nanmin([box_mask[0][0],box_slice[0][0]]) - x_max=np.nanmax([box_mask[0][1],box_slice[0][1]]) - y_min=np.nanmin([box_mask[1][0],box_slice[1][0]]) - y_max=np.nanmax([box_mask[1][1],box_slice[1][1]]) - - constraint_x=Constraint(projection_x_coordinate=lambda cell: int(x_min) < cell < int(x_max)) - constraint_y=Constraint(projection_y_coordinate=lambda cell: int(y_min) < cell < int(y_max)) - - constraint=constraint_time & constraint_x & constraint_y -# Mask_cell_surface_i=mask_cell_surface(Mask_w_i,cell,masked=False,z_coord='model_level_number') - mask_cell_i=mask_cell_i.extract(constraint) - mask_cell_surface_i=mask_cell_surface_i.extract(constraint) - - input_cubes_i=input_cubes.extract(constraint) + box_mask = [[np.nan, np.nan], [np.nan, np.nan]] + + x = track_i[track_i["time"].values == time_i]["projection_x_coordinate"].values[ + 0 + ] + y = track_i[track_i["time"].values == time_i]["projection_y_coordinate"].values[ + 0 + ] + + box_slice = [[x - width, x + width], [y - width, y + width]] + + x_min = np.nanmin([box_mask[0][0], box_slice[0][0]]) + x_max = np.nanmax([box_mask[0][1], box_slice[0][1]]) + y_min = np.nanmin([box_mask[1][0], box_slice[1][0]]) + y_max = np.nanmax([box_mask[1][1], box_slice[1][1]]) + + constraint_x = Constraint( + projection_x_coordinate=lambda cell: int(x_min) < cell < int(x_max) + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: int(y_min) < cell < int(y_max) + ) + + constraint = constraint_time & constraint_x & constraint_y + # Mask_cell_surface_i=mask_cell_surface(Mask_w_i,cell,masked=False,z_coord='model_level_number') + mask_cell_i = mask_cell_i.extract(constraint) + mask_cell_surface_i = mask_cell_surface_i.extract(constraint) + + input_cubes_i = input_cubes.extract(constraint) for cube in input_cubes_i: - cube_masked=mask_cube_cell(cube,mask_cell_i,cell,track_i) - coords_remove=[] + cube_masked = mask_cube_cell(cube, mask_cell_i, cell, track_i) + coords_remove = [] for coordinate in cube_masked.coords(dim_coords=False): if coordinate.name() not in dimensions: for dim in dimensions: - if set(cube_masked.coord_dims(coordinate)).intersection(set(cube_masked.coord_dims(dim))): + if set(cube_masked.coord_dims(coordinate)).intersection( + set(cube_masked.coord_dims(dim)) + ): coords_remove.append(coordinate.name()) for coordinate in set(coords_remove): - cube_masked.remove_coord(coordinate) - + cube_masked.remove_coord(coordinate) + for aggregator in aggregators: - cube_collapsed=cube_masked.collapsed(dimensions,aggregator,**kwargs) - #remove all collapsed coordinates (x and y dim, scalar now) and keep only time as all these coordinates are useless + cube_collapsed = cube_masked.collapsed(dimensions, aggregator, **kwargs) + # remove all collapsed coordinates (x and y dim, scalar now) and keep only time as all these coordinates are useless for coordinate in cube_collapsed.coords(): if not cube_collapsed.coord_dims(coordinate): - if coordinate.name() is not 'time': + if coordinate.name() is not "time": cube_collapsed.remove_coord(coordinate) logging.debug(str(cube_collapsed)) cubes_profile[aggregator.name()].append(cube_collapsed) + minutes = (track_i["time_cell"] / pd.Timedelta(minutes=1)).values + latitude = track_i["latitude"].values + longitude = track_i["longitude"].values + minutes_coord = AuxCoord(minutes, long_name="cell_time", units="min") + latitude_coord = AuxCoord(latitude, long_name="latitude", units="degrees") + longitude_coord = AuxCoord(longitude, long_name="longitude", units="degrees") - minutes=(track_i['time_cell']/pd.Timedelta(minutes=1)).values - latitude=track_i['latitude'].values - longitude=track_i['longitude'].values - minutes_coord=AuxCoord(minutes,long_name='cell_time',units='min') - latitude_coord=AuxCoord(latitude,long_name='latitude',units='degrees') - longitude_coord=AuxCoord(longitude,long_name='longitude',units='degrees') - for aggregator in aggregators: - cubes_profile[aggregator.name()]=cubes_profile[aggregator.name()].merge() + cubes_profile[aggregator.name()] = cubes_profile[aggregator.name()].merge() for cube in cubes_profile[aggregator.name()]: - cube.add_aux_coord(minutes_coord,data_dims=cube.coord_dims('time')) - cube.add_aux_coord(latitude_coord,data_dims=cube.coord_dims('time')) - cube.add_aux_coord(longitude_coord,data_dims=cube.coord_dims('time')) - os.makedirs(os.path.join(output_path,output_name,aggregator.name()),exist_ok=True) - savefile=os.path.join(output_path,output_name,aggregator.name(),output_name+'_'+ aggregator.name()+'_'+str(int(cell))+'.nc') - save(cubes_profile[aggregator.name()],savefile) - - -def cog_cell(cell,Tracks=None,M_total=None,M_liquid=None, - M_frozen=None, - Mask=None, - savedir=None): - - + cube.add_aux_coord(minutes_coord, data_dims=cube.coord_dims("time")) + cube.add_aux_coord(latitude_coord, data_dims=cube.coord_dims("time")) + cube.add_aux_coord(longitude_coord, data_dims=cube.coord_dims("time")) + os.makedirs( + os.path.join(output_path, output_name, aggregator.name()), exist_ok=True + ) + savefile = os.path.join( + output_path, + output_name, + aggregator.name(), + output_name + "_" + aggregator.name() + "_" + str(int(cell)) + ".nc", + ) + save(cubes_profile[aggregator.name()], savefile) + + +def cog_cell( + cell, + Tracks=None, + M_total=None, + M_liquid=None, + M_frozen=None, + Mask=None, + savedir=None, +): + from iris import Constraint - logging.debug('Start calculating COG for '+str(cell)) - Track=Tracks[Tracks['cell']==cell] - constraint_time=Constraint(time=lambda cell: Track.head(1)['time'].values[0] <= cell <= Track.tail(1)['time'].values[0]) - M_total_i=M_total.extract(constraint_time) - M_liquid_i=M_liquid.extract(constraint_time) - M_frozen_i=M_frozen.extract(constraint_time) - Mask_i=Mask.extract(constraint_time) - - savedir_cell=os.path.join(savedir,'cells',str(int(cell))) - os.makedirs(savedir_cell,exist_ok=True) - savefile_COG_total_i=os.path.join(savedir_cell,'COG_total'+'_'+str(int(cell))+'.h5') - savefile_COG_liquid_i=os.path.join(savedir_cell,'COG_liquid'+'_'+str(int(cell))+'.h5') - savefile_COG_frozen_i=os.path.join(savedir_cell,'COG_frozen'+'_'+str(int(cell))+'.h5') - - Tracks_COG_total_i=calculate_cog(Track,M_total_i,Mask_i) -# Tracks_COG_total_list.append(Tracks_COG_total_i) - logging.debug('COG total loaded for ' +str(cell)) - - Tracks_COG_liquid_i=calculate_cog(Track,M_liquid_i,Mask_i) -# Tracks_COG_liquid_list.append(Tracks_COG_liquid_i) - logging.debug('COG liquid loaded for ' +str(cell)) - Tracks_COG_frozen_i=calculate_cog(Track,M_frozen_i,Mask_i) -# Tracks_COG_frozen_list.append(Tracks_COG_frozen_i) - logging.debug('COG frozen loaded for ' +str(cell)) - - Tracks_COG_total_i.to_hdf(savefile_COG_total_i,'table') - Tracks_COG_liquid_i.to_hdf(savefile_COG_liquid_i,'table') - Tracks_COG_frozen_i.to_hdf(savefile_COG_frozen_i,'table') - logging.debug('individual COG calculated and saved to '+ savedir_cell) - - -def lifetime_histogram(Track,bin_edges=np.arange(0,200,20),density=False,return_values=False): - Track_cell=Track.groupby('cell') - minutes=(Track_cell['time_cell'].max()/pd.Timedelta(minutes=1)).values - hist, bin_edges = np.histogram(minutes, bin_edges,density=density) - bin_centers=bin_edges[:-1]+0.5*np.diff(bin_edges) + + logging.debug("Start calculating COG for " + str(cell)) + Track = Tracks[Tracks["cell"] == cell] + constraint_time = Constraint( + time=lambda cell: Track.head(1)["time"].values[0] + <= cell + <= Track.tail(1)["time"].values[0] + ) + M_total_i = M_total.extract(constraint_time) + M_liquid_i = M_liquid.extract(constraint_time) + M_frozen_i = M_frozen.extract(constraint_time) + Mask_i = Mask.extract(constraint_time) + + savedir_cell = os.path.join(savedir, "cells", str(int(cell))) + os.makedirs(savedir_cell, exist_ok=True) + savefile_COG_total_i = os.path.join( + savedir_cell, "COG_total" + "_" + str(int(cell)) + ".h5" + ) + savefile_COG_liquid_i = os.path.join( + savedir_cell, "COG_liquid" + "_" + str(int(cell)) + ".h5" + ) + savefile_COG_frozen_i = os.path.join( + savedir_cell, "COG_frozen" + "_" + str(int(cell)) + ".h5" + ) + + Tracks_COG_total_i = calculate_cog(Track, M_total_i, Mask_i) + # Tracks_COG_total_list.append(Tracks_COG_total_i) + logging.debug("COG total loaded for " + str(cell)) + + Tracks_COG_liquid_i = calculate_cog(Track, M_liquid_i, Mask_i) + # Tracks_COG_liquid_list.append(Tracks_COG_liquid_i) + logging.debug("COG liquid loaded for " + str(cell)) + Tracks_COG_frozen_i = calculate_cog(Track, M_frozen_i, Mask_i) + # Tracks_COG_frozen_list.append(Tracks_COG_frozen_i) + logging.debug("COG frozen loaded for " + str(cell)) + + Tracks_COG_total_i.to_hdf(savefile_COG_total_i, "table") + Tracks_COG_liquid_i.to_hdf(savefile_COG_liquid_i, "table") + Tracks_COG_frozen_i.to_hdf(savefile_COG_frozen_i, "table") + logging.debug("individual COG calculated and saved to " + savedir_cell) + + +def lifetime_histogram( + Track, bin_edges=np.arange(0, 200, 20), density=False, return_values=False +): + Track_cell = Track.groupby("cell") + minutes = (Track_cell["time_cell"].max() / pd.Timedelta(minutes=1)).values + hist, bin_edges = np.histogram(minutes, bin_edges, density=density) + bin_centers = bin_edges[:-1] + 0.5 * np.diff(bin_edges) if return_values: - return hist,bin_edges,bin_centers,minutes + return hist, bin_edges, bin_centers, minutes else: - return hist,bin_edges,bin_centers - -def haversine(lat1,lon1,lat2,lon2): + return hist, bin_edges, bin_centers + + +def haversine(lat1, lon1, lat2, lon2): """Computes the Haversine distance in kilometres between two points (based on implementation CIS https://github.com/cedadev/cis) :param lat1: first point or points as array, each as array of latitude in degrees :param lon1: first point or points as array, each as array of longitude in degrees @@ -171,88 +250,164 @@ def haversine(lat1,lon1,lat2,lon2): lat2 = np.radians(lat2) lon1 = np.radians(lon1) lon2 = np.radians(lon2) - #print(lat1,lat2,lon1,lon2) - arclen = 2 * np.arcsin(np.sqrt((np.sin((lat2 - lat1) / 2)) ** 2 + np.cos(lat1) * np.cos(lat2) * (np.sin((lon2 - lon1) / 2)) ** 2)) + # print(lat1,lat2,lon1,lon2) + arclen = 2 * np.arcsin( + np.sqrt( + (np.sin((lat2 - lat1) / 2)) ** 2 + + np.cos(lat1) * np.cos(lat2) * (np.sin((lon2 - lon1) / 2)) ** 2 + ) + ) return arclen * RADIUS_EARTH -def calculate_distance(feature_1,feature_2,method_distance=None): + +def calculate_distance(feature_1, feature_2, method_distance=None): """Computes distance between two features based on either lat/lon coordinates or x/y coordinates :param feature_1: first feature or points as array, each as array of latitude, longitude in degrees :param feature_2: second feature or points as array, each as array of latitude, longitude in degrees :return: distance between the two features in metres """ if method_distance is None: - if ('projection_x_coordinate' in feature_1) and ('projection_y_coordinate' in feature_1) and ('projection_x_coordinate' in feature_2) and ('projection_y_coordinate' in feature_2) : - method_distance='xy' - elif ('latitude' in feature_1) and ('longitude' in feature_1) and ('latitude' in feature_2) and ('longitude' in feature_2): - method_distance='latlon' + if ( + ("projection_x_coordinate" in feature_1) + and ("projection_y_coordinate" in feature_1) + and ("projection_x_coordinate" in feature_2) + and ("projection_y_coordinate" in feature_2) + ): + method_distance = "xy" + elif ( + ("latitude" in feature_1) + and ("longitude" in feature_1) + and ("latitude" in feature_2) + and ("longitude" in feature_2) + ): + method_distance = "latlon" else: - raise ValueError('either latitude/longitude or projection_x_coordinate/projection_y_coordinate have to be present to calculate distances') + raise ValueError( + "either latitude/longitude or projection_x_coordinate/projection_y_coordinate have to be present to calculate distances" + ) - if method_distance=='xy': - distance=np.sqrt((feature_1['projection_x_coordinate']-feature_2['projection_x_coordinate'])**2 - +(feature_1['projection_y_coordinate']-feature_2['projection_y_coordinate'])**2) - elif method_distance=='latlon': - distance=1000*haversine(feature_1['latitude'],feature_1['longitude'],feature_2['latitude'],feature_2['longitude']) + if method_distance == "xy": + distance = np.sqrt( + ( + feature_1["projection_x_coordinate"] + - feature_2["projection_x_coordinate"] + ) + ** 2 + + ( + feature_1["projection_y_coordinate"] + - feature_2["projection_y_coordinate"] + ) + ** 2 + ) + elif method_distance == "latlon": + distance = 1000 * haversine( + feature_1["latitude"], + feature_1["longitude"], + feature_2["latitude"], + feature_2["longitude"], + ) else: - raise ValueError('method undefined') + raise ValueError("method undefined") return distance -def calculate_velocity_individual(feature_old,feature_new,method_distance=None): - distance=calculate_distance(feature_old,feature_new,method_distance=method_distance) - diff_time=((feature_new['time']-feature_old['time']).total_seconds()) - velocity=distance/diff_time + +def calculate_velocity_individual(feature_old, feature_new, method_distance=None): + distance = calculate_distance( + feature_old, feature_new, method_distance=method_distance + ) + diff_time = (feature_new["time"] - feature_old["time"]).total_seconds() + velocity = distance / diff_time return velocity -def calculate_velocity(track,method_distance=None): - for cell_i,track_i in track.groupby('cell'): - index=track_i.index.values - for i,index_i in enumerate(index[:-1]): - velocity=calculate_velocity_individual(track_i.loc[index[i]],track_i.loc[index[i+1]],method_distance=method_distance) - track.at[index_i,'v']=velocity + +def calculate_velocity(track, method_distance=None): + for cell_i, track_i in track.groupby("cell"): + index = track_i.index.values + for i, index_i in enumerate(index[:-1]): + velocity = calculate_velocity_individual( + track_i.loc[index[i]], + track_i.loc[index[i + 1]], + method_distance=method_distance, + ) + track.at[index_i, "v"] = velocity return track -def velocity_histogram(track,bin_edges=np.arange(0,30,1),density=False,method_distance=None,return_values=False): - if 'v' not in track.columns: - logging.info('calculate velocities') - track=calculate_velocity(track) - velocities=track['v'].values - hist, bin_edges = np.histogram(velocities[~np.isnan(velocities)], bin_edges,density=density) + +def velocity_histogram( + track, + bin_edges=np.arange(0, 30, 1), + density=False, + method_distance=None, + return_values=False, +): + if "v" not in track.columns: + logging.info("calculate velocities") + track = calculate_velocity(track) + velocities = track["v"].values + hist, bin_edges = np.histogram( + velocities[~np.isnan(velocities)], bin_edges, density=density + ) if return_values: - return hist,bin_edges,velocities + return hist, bin_edges, velocities else: - return hist,bin_edges + return hist, bin_edges -def calculate_nearestneighbordistance(features,method_distance=None): + +def calculate_nearestneighbordistance(features, method_distance=None): from itertools import combinations - features['min_distance']=np.nan - for time_i,features_i in features.groupby('time'): - logging.debug(str(time_i)) - indeces=combinations(features_i.index.values,2) - #Loop over combinations to remove features that are closer together than min_distance and keep larger one (either higher threshold or larger area) - distances=[] - for index_1,index_2 in indeces: - if index_1 is not index_2: - distance=calculate_distance(features_i.loc[index_1],features_i.loc[index_2],method_distance=method_distance) - distances.append(pd.DataFrame({'index_1':index_1,'index_2':index_2,'distance': distance}, index=[0])) - if any([x is not None for x in distances]): - distances=pd.concat(distances, ignore_index=True) - for i in features_i.index: - min_distance=distances.loc[(distances['index_1']==i) | (distances['index_2']==i),'distance'].min() - features.at[i,'min_distance']=min_distance + + features["min_distance"] = np.nan + for time_i, features_i in features.groupby("time"): + logging.debug(str(time_i)) + indeces = combinations(features_i.index.values, 2) + # Loop over combinations to remove features that are closer together than min_distance and keep larger one (either higher threshold or larger area) + distances = [] + for index_1, index_2 in indeces: + if index_1 is not index_2: + distance = calculate_distance( + features_i.loc[index_1], + features_i.loc[index_2], + method_distance=method_distance, + ) + distances.append( + pd.DataFrame( + {"index_1": index_1, "index_2": index_2, "distance": distance}, + index=[0], + ) + ) + if any([x is not None for x in distances]): + distances = pd.concat(distances, ignore_index=True) + for i in features_i.index: + min_distance = distances.loc[ + (distances["index_1"] == i) | (distances["index_2"] == i), + "distance", + ].min() + features.at[i, "min_distance"] = min_distance return features -def nearestneighbordistance_histogram(features,bin_edges=np.arange(0,30000,500),density=False,method_distance=None,return_values=False): - if 'min_distance' not in features.columns: - logging.debug('calculate nearest neighbor distances') - features=calculate_nearestneighbordistance(features,method_distance=method_distance) - distances=features['min_distance'].values - hist, bin_edges = np.histogram(distances[~np.isnan(distances)], bin_edges,density=density) + +def nearestneighbordistance_histogram( + features, + bin_edges=np.arange(0, 30000, 500), + density=False, + method_distance=None, + return_values=False, +): + if "min_distance" not in features.columns: + logging.debug("calculate nearest neighbor distances") + features = calculate_nearestneighbordistance( + features, method_distance=method_distance + ) + distances = features["min_distance"].values + hist, bin_edges = np.histogram( + distances[~np.isnan(distances)], bin_edges, density=density + ) if return_values: - return hist,bin_edges,distances + return hist, bin_edges, distances else: - return hist,bin_edges - + return hist, bin_edges + + # Treatment of 2D lat/lon coordinates to be added: # def calculate_areas_2Dlatlon(latitude_coord,longitude_coord): # lat=latitude_coord.core_data() @@ -260,133 +415,169 @@ def nearestneighbordistance_histogram(features,bin_edges=np.arange(0,30000,500), # area=np.zeros(lat.shape) # dx=np.zeros(lat.shape) # dy=np.zeros(lat.shape) - + # return area -def calculate_area(features,mask,method_area=None): - from tobac.utils import mask_features_surface,mask_features + +def calculate_area(features, mask, method_area=None): + from tobac.utils import mask_features_surface, mask_features from iris import Constraint from iris.analysis.cartography import area_weights - - features['area']=np.nan - - mask_coords=[coord.name() for coord in mask.coords()] + + features["area"] = np.nan + + mask_coords = [coord.name() for coord in mask.coords()] if method_area is None: - if ('projection_x_coordinate' in mask_coords) and ('projection_y_coordinate' in mask_coords): - method_area='xy' - elif ('latitude' in mask_coords) and ('longitude' in mask_coords): - method_area='latlon' + if ("projection_x_coordinate" in mask_coords) and ( + "projection_y_coordinate" in mask_coords + ): + method_area = "xy" + elif ("latitude" in mask_coords) and ("longitude" in mask_coords): + method_area = "latlon" else: - raise ValueError('either latitude/longitude or projection_x_coordinate/projection_y_coordinate have to be present to calculate distances') - logging.debug('calculating area using method '+ method_area) - if method_area=='xy': - if not (mask.coord('projection_x_coordinate').has_bounds() and mask.coord('projection_y_coordinate').has_bounds()): - mask.coord('projection_x_coordinate').guess_bounds() - mask.coord('projection_y_coordinate').guess_bounds() - area=np.outer(np.diff(mask.coord('projection_x_coordinate').bounds,axis=1),np.diff(mask.coord('projection_y_coordinate').bounds,axis=1)) - elif method_area=='latlon': - if (mask.coord('latitude').ndim==1) and (mask.coord('latitude').ndim==1): - if not (mask.coord('latitude').has_bounds() and mask.coord('longitude').has_bounds()): - mask.coord('latitude').guess_bounds() - mask.coord('longitude').guess_bounds() - area=area_weights(mask,normalize=False) - elif mask.coord('latitude').ndim==2 and mask.coord('longitude').ndim==2: - raise ValueError('2D latitude/longitude coordinates not supported yet') + raise ValueError( + "either latitude/longitude or projection_x_coordinate/projection_y_coordinate have to be present to calculate distances" + ) + logging.debug("calculating area using method " + method_area) + if method_area == "xy": + if not ( + mask.coord("projection_x_coordinate").has_bounds() + and mask.coord("projection_y_coordinate").has_bounds() + ): + mask.coord("projection_x_coordinate").guess_bounds() + mask.coord("projection_y_coordinate").guess_bounds() + area = np.outer( + np.diff(mask.coord("projection_x_coordinate").bounds, axis=1), + np.diff(mask.coord("projection_y_coordinate").bounds, axis=1), + ) + elif method_area == "latlon": + if (mask.coord("latitude").ndim == 1) and (mask.coord("latitude").ndim == 1): + if not ( + mask.coord("latitude").has_bounds() + and mask.coord("longitude").has_bounds() + ): + mask.coord("latitude").guess_bounds() + mask.coord("longitude").guess_bounds() + area = area_weights(mask, normalize=False) + elif mask.coord("latitude").ndim == 2 and mask.coord("longitude").ndim == 2: + raise ValueError("2D latitude/longitude coordinates not supported yet") # area=calculate_areas_2Dlatlon(mask.coord('latitude'),mask.coord('longitude')) else: - raise ValueError('latitude/longitude coordinate shape not supported') + raise ValueError("latitude/longitude coordinate shape not supported") else: - raise ValueError('method undefined') - - for time_i,features_i in features.groupby('time'): - logging.debug('timestep:'+ str(time_i)) + raise ValueError("method undefined") + + for time_i, features_i in features.groupby("time"): + logging.debug("timestep:" + str(time_i)) constraint_time = Constraint(time=time_i) - mask_i=mask.extract(constraint_time) + mask_i = mask.extract(constraint_time) for i in features_i.index: - if len(mask_i.shape)==3: - mask_i_surface = mask_features_surface(mask_i, features_i.loc[i,'feature'], z_coord='model_level_number') - elif len(mask_i.shape)==2: - mask_i_surface=mask_features(mask_i,features_i.loc[i,'feature']) - area_feature=np.sum(area*(mask_i_surface.data>0)) - features.at[i,'area']=area_feature + if len(mask_i.shape) == 3: + mask_i_surface = mask_features_surface( + mask_i, features_i.loc[i, "feature"], z_coord="model_level_number" + ) + elif len(mask_i.shape) == 2: + mask_i_surface = mask_features(mask_i, features_i.loc[i, "feature"]) + area_feature = np.sum(area * (mask_i_surface.data > 0)) + features.at[i, "area"] = area_feature return features -def area_histogram(features,mask,bin_edges=np.arange(0,30000,500), - density=False,method_area=None, - return_values=False,representative_area=False): - if 'area' not in features.columns: - logging.info('calculate area') - features=calculate_area(features, mask, method_area) - areas=features['area'].values + +def area_histogram( + features, + mask, + bin_edges=np.arange(0, 30000, 500), + density=False, + method_area=None, + return_values=False, + representative_area=False, +): + if "area" not in features.columns: + logging.info("calculate area") + features = calculate_area(features, mask, method_area) + areas = features["area"].values # restrict to non NaN values: - areas=areas[~np.isnan(areas)] + areas = areas[~np.isnan(areas)] if representative_area: - weights=areas + weights = areas else: - weights=None - hist, bin_edges = np.histogram(areas, bin_edges,density=density,weights=weights) - bin_centers=bin_edges[:-1]+0.5*np.diff(bin_edges) + weights = None + hist, bin_edges = np.histogram(areas, bin_edges, density=density, weights=weights) + bin_centers = bin_edges[:-1] + 0.5 * np.diff(bin_edges) if return_values: - return hist,bin_edges,bin_centers,areas + return hist, bin_edges, bin_centers, areas else: - return hist,bin_edges,bin_centers - -def histogram_cellwise(Track,variable=None,bin_edges=None,quantity='max',density=False): - Track_cell=Track.groupby('cell') - if quantity=='max': - variable_cell=Track_cell[variable].max().values - elif quantity=='min': - variable_cell=Track_cell[variable].min().values - elif quantity=='mean': - variable_cell=Track_cell[variable].mean().values + return hist, bin_edges, bin_centers + + +def histogram_cellwise( + Track, variable=None, bin_edges=None, quantity="max", density=False +): + Track_cell = Track.groupby("cell") + if quantity == "max": + variable_cell = Track_cell[variable].max().values + elif quantity == "min": + variable_cell = Track_cell[variable].min().values + elif quantity == "mean": + variable_cell = Track_cell[variable].mean().values else: - raise ValueError('quantity unknown, must be max, min or mean') - hist, bin_edges = np.histogram(variable_cell, bin_edges,density=density) - bin_centers=bin_edges[:-1]+0.5*np.diff(bin_edges) + raise ValueError("quantity unknown, must be max, min or mean") + hist, bin_edges = np.histogram(variable_cell, bin_edges, density=density) + bin_centers = bin_edges[:-1] + 0.5 * np.diff(bin_edges) + + return hist, bin_edges, bin_centers + - return hist,bin_edges, bin_centers +def histogram_featurewise(Track, variable=None, bin_edges=None, density=False): + hist, bin_edges = np.histogram(Track[variable].values, bin_edges, density=density) + bin_centers = bin_edges[:-1] + 0.5 * np.diff(bin_edges) -def histogram_featurewise(Track,variable=None,bin_edges=None,density=False): - hist, bin_edges = np.histogram(Track[variable].values, bin_edges,density=density) - bin_centers=bin_edges[:-1]+0.5*np.diff(bin_edges) + return hist, bin_edges, bin_centers - return hist,bin_edges, bin_centers -def calculate_overlap(track_1,track_2,min_sum_inv_distance=None,min_mean_inv_distance=None): - cells_1=track_1['cell'].unique() -# n_cells_1_tot=len(cells_1) - cells_2=track_2['cell'].unique() - overlap=pd.DataFrame() - for i_cell_1,cell_1 in enumerate(cells_1): +def calculate_overlap( + track_1, track_2, min_sum_inv_distance=None, min_mean_inv_distance=None +): + cells_1 = track_1["cell"].unique() + # n_cells_1_tot=len(cells_1) + cells_2 = track_2["cell"].unique() + overlap = pd.DataFrame() + for i_cell_1, cell_1 in enumerate(cells_1): for cell_2 in cells_2: - track_1_i=track_1[track_1['cell']==cell_1] - track_2_i=track_2[track_2['cell']==cell_2] - track_1_i=track_1_i[track_1_i['time'].isin(track_2_i['time'])] - track_2_i=track_2_i[track_2_i['time'].isin(track_1_i['time'])] + track_1_i = track_1[track_1["cell"] == cell_1] + track_2_i = track_2[track_2["cell"] == cell_2] + track_1_i = track_1_i[track_1_i["time"].isin(track_2_i["time"])] + track_2_i = track_2_i[track_2_i["time"].isin(track_1_i["time"])] if not track_1_i.empty: - n_overlap=len(track_1_i) - distances=[] + n_overlap = len(track_1_i) + distances = [] for i in range(len(track_1_i)): - distance=calculate_distance(track_1_i.iloc[[i]],track_2_i.iloc[[i]],method_distance='xy') + distance = calculate_distance( + track_1_i.iloc[[i]], track_2_i.iloc[[i]], method_distance="xy" + ) distances.append(distance) -# mean_distance=np.mean(distances) - mean_inv_distance=np.mean(1/(1+np.array(distances)/1000)) -# mean_inv_squaredistance=np.mean(1/(1+(np.array(distances)/1000)**2)) - sum_inv_distance=np.sum(1/(1+np.array(distances)/1000)) -# sum_inv_squaredistance=np.sum(1/(1+(np.array(distances)/1000)**2)) - overlap=overlap.append({'cell_1':cell_1, - 'cell_2':cell_2, - 'n_overlap':n_overlap, -# 'mean_distance':mean_distance, - 'mean_inv_distance':mean_inv_distance, -# 'mean_inv_squaredistance':mean_inv_squaredistance, - 'sum_inv_distance':sum_inv_distance, -# 'sum_inv_squaredistance':sum_inv_squaredistance - },ignore_index=True) + # mean_distance=np.mean(distances) + mean_inv_distance = np.mean(1 / (1 + np.array(distances) / 1000)) + # mean_inv_squaredistance=np.mean(1/(1+(np.array(distances)/1000)**2)) + sum_inv_distance = np.sum(1 / (1 + np.array(distances) / 1000)) + # sum_inv_squaredistance=np.sum(1/(1+(np.array(distances)/1000)**2)) + overlap = overlap.append( + { + "cell_1": cell_1, + "cell_2": cell_2, + "n_overlap": n_overlap, + # 'mean_distance':mean_distance, + "mean_inv_distance": mean_inv_distance, + # 'mean_inv_squaredistance':mean_inv_squaredistance, + "sum_inv_distance": sum_inv_distance, + # 'sum_inv_squaredistance':sum_inv_squaredistance + }, + ignore_index=True, + ) if min_sum_inv_distance: - overlap=overlap[(overlap['sum_inv_distance']>=min_sum_inv_distance)] + overlap = overlap[(overlap["sum_inv_distance"] >= min_sum_inv_distance)] if min_mean_inv_distance: - overlap=overlap[(overlap['mean_inv_distance']>=min_mean_inv_distance)] + overlap = overlap[(overlap["mean_inv_distance"] >= min_mean_inv_distance)] - return overlap \ No newline at end of file + return overlap diff --git a/tobac/centerofgravity.py b/tobac/centerofgravity.py index 544a4ee2..d7d5675d 100644 --- a/tobac/centerofgravity.py +++ b/tobac/centerofgravity.py @@ -1,147 +1,166 @@ import logging -def calculate_cog(tracks,mass,mask): - ''' caluclate centre of gravity and mass forech individual tracked cell in the simulation + +def calculate_cog(tracks, mass, mask): + """caluclate centre of gravity and mass forech individual tracked cell in the simulation Input: tracks: pandas.DataFrame DataFrame containing trajectories of cell centres - mass: iris.cube.Cube + mass: iris.cube.Cube cube of quantity (need coordinates 'time', 'geopotential_height','projection_x_coordinate' and 'projection_y_coordinate') - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int > where belonging to cloud volume, 0 everywhere else ) Output: tracks_out pandas.DataFrame Dataframe containing t,x,y,z positions of centre of gravity and total cloud mass each tracked cells at each timestep - ''' + """ from .utils import mask_cube_cell - from iris import Constraint - - logging.info('start calculating centre of gravity for tracked cells') - - tracks_out=tracks[['time','frame','cell','time_cell']] - - for i_row,row in tracks_out.iterrows(): - cell=row['cell'] - constraint_time=Constraint(time=row['time']) - mass_i=mass.extract(constraint_time) - mask_i=mask.extract(constraint_time) - mass_masked_i=mask_cube_cell(mass_i,mask_i,cell) - x_M,y_M,z_M,mass_M=center_of_gravity(mass_masked_i) - tracks_out.loc[i_row,'x_M']=float(x_M) - tracks_out.loc[i_row,'y_M']=float(y_M) - tracks_out.loc[i_row,'z_M']=float(z_M) - tracks_out.loc[i_row,'mass']=float(mass_M) - - logging.info('Finished calculating centre of gravity for tracked cells') + from iris import Constraint + + logging.info("start calculating centre of gravity for tracked cells") + + tracks_out = tracks[["time", "frame", "cell", "time_cell"]] + + for i_row, row in tracks_out.iterrows(): + cell = row["cell"] + constraint_time = Constraint(time=row["time"]) + mass_i = mass.extract(constraint_time) + mask_i = mask.extract(constraint_time) + mass_masked_i = mask_cube_cell(mass_i, mask_i, cell) + x_M, y_M, z_M, mass_M = center_of_gravity(mass_masked_i) + tracks_out.loc[i_row, "x_M"] = float(x_M) + tracks_out.loc[i_row, "y_M"] = float(y_M) + tracks_out.loc[i_row, "z_M"] = float(z_M) + tracks_out.loc[i_row, "mass"] = float(mass_M) + + logging.info("Finished calculating centre of gravity for tracked cells") return tracks_out - -def calculate_cog_untracked(mass,mask): - ''' caluclate centre of gravity and mass for untracked parts of domain + + +def calculate_cog_untracked(mass, mask): + """caluclate centre of gravity and mass for untracked parts of domain Input: - mass: iris.cube.Cube + mass: iris.cube.Cube cube of quantity (need coordinates 'time', 'geopotential_height','projection_x_coordinate' and 'projection_y_coordinate') - - mask: iris.cube.Cube + + mask: iris.cube.Cube cube containing mask (int > where belonging to cloud volume, 0 everywhere else ) Output: tracks_out pandas.DataFrame Dataframe containing t,x,y,z positions of centre of gravity and total cloud mass for untracked part of dimain - ''' + """ from pandas import DataFrame from .utils import mask_cube_untracked from iris import Constraint - - logging.info('start calculating centre of gravity for untracked parts of the domain') - tracks_out=DataFrame() - time_coord=mass.coord('time') - tracks_out['frame']=range(len(time_coord.points)) - for i_row,row in tracks_out.iterrows(): - time_i=time_coord.units.num2date(time_coord[int(row['frame'])].points[0]) - constraint_time=Constraint(time=time_i) - mass_i=mass.extract(constraint_time) - mask_i=mask.extract(constraint_time) - mass_untracked_i=mask_cube_untracked(mass_i,mask_i) - x_M,y_M,z_M,mass_M=center_of_gravity(mass_untracked_i) - tracks_out.loc[i_row,'time']=time_i - tracks_out.loc[i_row,'x_M']=float(x_M) - tracks_out.loc[i_row,'y_M']=float(y_M) - tracks_out.loc[i_row,'z_M']=float(z_M) - tracks_out.loc[i_row,'mass']=float(mass_M) - - logging.info('Finished calculating centre of gravity for untracked parts of the domain') - + + logging.info( + "start calculating centre of gravity for untracked parts of the domain" + ) + tracks_out = DataFrame() + time_coord = mass.coord("time") + tracks_out["frame"] = range(len(time_coord.points)) + for i_row, row in tracks_out.iterrows(): + time_i = time_coord.units.num2date(time_coord[int(row["frame"])].points[0]) + constraint_time = Constraint(time=time_i) + mass_i = mass.extract(constraint_time) + mask_i = mask.extract(constraint_time) + mass_untracked_i = mask_cube_untracked(mass_i, mask_i) + x_M, y_M, z_M, mass_M = center_of_gravity(mass_untracked_i) + tracks_out.loc[i_row, "time"] = time_i + tracks_out.loc[i_row, "x_M"] = float(x_M) + tracks_out.loc[i_row, "y_M"] = float(y_M) + tracks_out.loc[i_row, "z_M"] = float(z_M) + tracks_out.loc[i_row, "mass"] = float(mass_M) + + logging.info( + "Finished calculating centre of gravity for untracked parts of the domain" + ) + return tracks_out + def calculate_cog_domain(mass): - ''' caluclate centre of gravity and mass for entire domain + """caluclate centre of gravity and mass for entire domain Input: - mass: iris.cube.Cube + mass: iris.cube.Cube cube of quantity (need coordinates 'time', 'geopotential_height','projection_x_coordinate' and 'projection_y_coordinate') Output: tracks_out pandas.DataFrame Dataframe containing t,x,y,z positions of centre of gravity and total cloud mass - ''' + """ from pandas import DataFrame from iris import Constraint - - logging.info('start calculating centre of gravity for entire domain') - - time_coord=mass.coord('time') - - tracks_out=DataFrame() - tracks_out['frame']=range(len(time_coord.points)) - for i_row,row in tracks_out.iterrows(): - time_i=time_coord.units.num2date(time_coord[int(row['frame'])].points[0]) - constraint_time=Constraint(time=time_i) - mass_i=mass.extract(constraint_time) - x_M,y_M,z_M,mass_M=center_of_gravity(mass_i) - tracks_out.loc[i_row,'time']=time_i - tracks_out.loc[i_row,'x_M']=float(x_M) - tracks_out.loc[i_row,'y_M']=float(y_M) - tracks_out.loc[i_row,'z_M']=float(z_M) - tracks_out.loc[i_row,'mass']=float(mass_M) - - logging.info('Finished calculating centre of gravity for entire domain') + + logging.info("start calculating centre of gravity for entire domain") + + time_coord = mass.coord("time") + + tracks_out = DataFrame() + tracks_out["frame"] = range(len(time_coord.points)) + for i_row, row in tracks_out.iterrows(): + time_i = time_coord.units.num2date(time_coord[int(row["frame"])].points[0]) + constraint_time = Constraint(time=time_i) + mass_i = mass.extract(constraint_time) + x_M, y_M, z_M, mass_M = center_of_gravity(mass_i) + tracks_out.loc[i_row, "time"] = time_i + tracks_out.loc[i_row, "x_M"] = float(x_M) + tracks_out.loc[i_row, "y_M"] = float(y_M) + tracks_out.loc[i_row, "z_M"] = float(z_M) + tracks_out.loc[i_row, "mass"] = float(mass_M) + + logging.info("Finished calculating centre of gravity for entire domain") return tracks_out + def center_of_gravity(cube_in): - ''' caluclate centre of gravity and sum of quantity + """caluclate centre of gravity and sum of quantity Input: - cube_in: iris.cube.Cube + cube_in: iris.cube.Cube cube (potentially masked) of quantity (need coordinates 'geopotential_height','projection_x_coordinate' and 'projection_y_coordinate') Output: - x: float - x position of centre of gravity - y: float - y position of centre of gravity - z: float - z position of centre of gravity - variable_sum: float + x: float + x position of centre of gravity + y: float + y position of centre of gravity + z: float + z position of centre of gravity + variable_sum: float sum of quantity of over unmasked part of the cube - ''' + """ from iris.analysis import SUM import numpy as np - cube_sum=cube_in.collapsed(['bottom_top','south_north','west_east'],SUM) - z=cube_in.coord('geopotential_height') - x=cube_in.coord('projection_x_coordinate') - y=cube_in.coord('projection_y_coordinate') - dimensions_collapse=['model_level_number','x','y'] - for coord in cube_in.coords(): - if (coord.ndim>1 and (cube_in.coord_dims(dimensions_collapse[0])[0] in cube_in.coord_dims(coord) or cube_in.coord_dims(dimensions_collapse[1])[0] in cube_in.coord_dims(coord) or cube_in.coord_dims(dimensions_collapse[2])[0] in cube_in.coord_dims(coord))): - cube_in.remove_coord(coord.name()) + + cube_sum = cube_in.collapsed(["bottom_top", "south_north", "west_east"], SUM) + z = cube_in.coord("geopotential_height") + x = cube_in.coord("projection_x_coordinate") + y = cube_in.coord("projection_y_coordinate") + dimensions_collapse = ["model_level_number", "x", "y"] + for coord in cube_in.coords(): + if coord.ndim > 1 and ( + cube_in.coord_dims(dimensions_collapse[0])[0] in cube_in.coord_dims(coord) + or cube_in.coord_dims(dimensions_collapse[1])[0] + in cube_in.coord_dims(coord) + or cube_in.coord_dims(dimensions_collapse[2])[0] + in cube_in.coord_dims(coord) + ): + cube_in.remove_coord(coord.name()) if cube_sum.data > 0: - x=((cube_in*x).collapsed(['model_level_number','x','y'],SUM)/cube_sum).data - y=((cube_in*y).collapsed(['model_level_number','x','y'],SUM)/cube_sum).data - z=((cube_in*z.points).collapsed(['model_level_number','x','y'],SUM)/cube_sum).data + x = ( + (cube_in * x).collapsed(["model_level_number", "x", "y"], SUM) / cube_sum + ).data + y = ( + (cube_in * y).collapsed(["model_level_number", "x", "y"], SUM) / cube_sum + ).data + z = ( + (cube_in * z.points).collapsed(["model_level_number", "x", "y"], SUM) + / cube_sum + ).data else: - x=np.nan - y=np.nan - z=np.nan - variable_sum=cube_sum.data - return(x,y,z,variable_sum) - - - + x = np.nan + y = np.nan + z = np.nan + variable_sum = cube_sum.data + return (x, y, z, variable_sum) diff --git a/tobac/feature_detection.py b/tobac/feature_detection.py index 3e634a96..c24b15fe 100644 --- a/tobac/feature_detection.py +++ b/tobac/feature_detection.py @@ -2,23 +2,32 @@ import numpy as np import pandas as pd -def feature_position(hdim1_indices,hdim2_indeces,region,track_data,threshold_i,position_threshold, target): - ''' + +def feature_position( + hdim1_indices, + hdim2_indeces, + region, + track_data, + threshold_i, + position_threshold, + target, +): + """ function to determine feature position Input: hdim1_indices: list - + hdim2_indeces: list - + region: list list of 2-element tuples track_data: numpy.ndarray 2D numpy array containing the data - + threshold_i: float - + position_threshold: str - + target: str Output: @@ -26,45 +35,48 @@ def feature_position(hdim1_indices,hdim2_indeces,region,track_data,threshold_i,p feature position along 1st horizontal dimension hdim2_index: float feature position along 2nd horizontal dimension - ''' - if position_threshold=='center': + """ + if position_threshold == "center": # get position as geometrical centre of identified region: - hdim1_index=np.mean(hdim1_indices) - hdim2_index=np.mean(hdim2_indeces) - - elif position_threshold=='extreme': - #get position as max/min position inside the identified region: - if target == 'maximum': - index=np.argmax(track_data[region]) - hdim1_index=hdim1_indices[index] - hdim2_index=hdim2_indeces[index] - - if target == 'minimum': - index=np.argmin(track_data[region]) - hdim1_index=hdim1_indices[index] - hdim2_index=hdim2_indeces[index] - - elif position_threshold=='weighted_diff': + hdim1_index = np.mean(hdim1_indices) + hdim2_index = np.mean(hdim2_indeces) + + elif position_threshold == "extreme": + # get position as max/min position inside the identified region: + if target == "maximum": + index = np.argmax(track_data[region]) + hdim1_index = hdim1_indices[index] + hdim2_index = hdim2_indeces[index] + + if target == "minimum": + index = np.argmin(track_data[region]) + hdim1_index = hdim1_indices[index] + hdim2_index = hdim2_indeces[index] + + elif position_threshold == "weighted_diff": # get position as centre of identified region, weighted by difference from the threshold: - weights=abs(track_data[region]-threshold_i) - if sum(weights)==0: - weights=None - hdim1_index=np.average(hdim1_indices,weights=weights) - hdim2_index=np.average(hdim2_indeces,weights=weights) + weights = abs(track_data[region] - threshold_i) + if sum(weights) == 0: + weights = None + hdim1_index = np.average(hdim1_indices, weights=weights) + hdim2_index = np.average(hdim2_indeces, weights=weights) - elif position_threshold=='weighted_abs': + elif position_threshold == "weighted_abs": # get position as centre of identified region, weighted by absolute values if the field: - weights=abs(track_data[region]) - if sum(weights)==0: - weights=None - hdim1_index=np.average(hdim1_indices,weights=weights) - hdim2_index=np.average(hdim2_indeces,weights=weights) + weights = abs(track_data[region]) + if sum(weights) == 0: + weights = None + hdim1_index = np.average(hdim1_indices, weights=weights) + hdim2_index = np.average(hdim2_indeces, weights=weights) else: - raise ValueError('position_threshold must be center,extreme,weighted_diff or weighted_abs') - return hdim1_index,hdim2_index + raise ValueError( + "position_threshold must be center,extreme,weighted_diff or weighted_abs" + ) + return hdim1_index, hdim2_index + -def test_overlap(region_inner,region_outer): - ''' +def test_overlap(region_inner, region_outer): + """ function to test for overlap between two regions (probably scope for further speedup here) Input: region_1: list @@ -75,12 +87,13 @@ def test_overlap(region_inner,region_outer): Output: overlap: bool True if there are any shared points between the two regions - ''' - overlap=frozenset(region_outer).isdisjoint(region_inner) + """ + overlap = frozenset(region_outer).isdisjoint(region_inner) return not overlap -def remove_parents(features_thresholds,regions_i,regions_old): - ''' + +def remove_parents(features_thresholds, regions_i, regions_old): + """ function to remove features whose regions surround newly detected feature regions Input: features_thresholds: pandas.DataFrame @@ -92,30 +105,36 @@ def remove_parents(features_thresholds,regions_i,regions_old): Output: features_thresholds pandas.DataFrame Dataframe containing detected features excluding those that are superseded by newly detected ones - ''' - list_remove=[] - for idx_i,region_i in regions_i.items(): - for idx_old,region_old in regions_old.items(): - if test_overlap(regions_old[idx_old],regions_i[idx_i]): + """ + list_remove = [] + for idx_i, region_i in regions_i.items(): + for idx_old, region_old in regions_old.items(): + if test_overlap(regions_old[idx_old], regions_i[idx_i]): list_remove.append(idx_old) - list_remove=list(set(list_remove)) + list_remove = list(set(list_remove)) # remove parent regions: if features_thresholds is not None: - features_thresholds=features_thresholds[~features_thresholds['idx'].isin(list_remove)] + features_thresholds = features_thresholds[ + ~features_thresholds["idx"].isin(list_remove) + ] return features_thresholds -def feature_detection_threshold(data_i,i_time, - threshold=None, - min_num=0, - target='maximum', - position_threshold='center', - sigma_threshold=0.5, - n_erosion_threshold=0, - n_min_threshold=0, - min_distance=0, - idx_start=0): - ''' + +def feature_detection_threshold( + data_i, + i_time, + threshold=None, + min_num=0, + target="maximum", + position_threshold="center", + sigma_threshold=0.5, + n_erosion_threshold=0, + n_min_threshold=0, + min_distance=0, + idx_start=0, +): + """ function to find features based on individual threshold value: Input: data_i: iris.cube.Cube @@ -139,82 +158,97 @@ def feature_detection_threshold(data_i,i_time, idx_start: int feature id to start with Output: - features_threshold: pandas DataFrame + features_threshold: pandas DataFrame detected features for individual threshold regions: dict dictionary containing the regions above/below threshold used for each feature (feature ids as keys) - ''' + """ from skimage.measure import label from skimage.morphology import binary_erosion # if looking for minima, set values above threshold to 0 and scale by data minimum: - if target == 'maximum': - mask=1*(data_i >= threshold) + if target == "maximum": + mask = 1 * (data_i >= threshold) # if looking for minima, set values above threshold to 0 and scale by data minimum: - elif target == 'minimum': - mask=1*(data_i <= threshold) + elif target == "minimum": + mask = 1 * (data_i <= threshold) # only include values greater than threshold - # erode selected regions by n pixels - if n_erosion_threshold>0: - selem=np.ones((n_erosion_threshold,n_erosion_threshold)) - mask=binary_erosion(mask,selem).astype(np.int64) + # erode selected regions by n pixels + if n_erosion_threshold > 0: + selem = np.ones((n_erosion_threshold, n_erosion_threshold)) + mask = binary_erosion(mask, selem).astype(np.int64) # detect individual regions, label and count the number of pixels included: labels = label(mask, background=0) - values, count = np.unique(labels[:,:].ravel(), return_counts=True) - values_counts=dict(zip(values, count)) + values, count = np.unique(labels[:, :].ravel(), return_counts=True) + values_counts = dict(zip(values, count)) # Filter out regions that have less pixels than n_min_threshold - values_counts={k:v for k, v in values_counts.items() if v>n_min_threshold} - #check if not entire domain filled as one feature - if 0 in values_counts: - #Remove background counts: - values_counts.pop(0) - #create empty list to store individual features for this threshold - list_features_threshold=[] - #create empty dict to store regions for individual features for this threshold - regions=dict() - #create emptry list of features to remove from parent threshold value - #loop over individual regions: - for cur_idx,count in values_counts.items(): - region=labels[:,:] == cur_idx - [hdim1_indices,hdim2_indeces]= np.nonzero(region) - #write region for individual threshold and feature to dict - region_i=list(zip(hdim1_indices,hdim2_indeces)) - regions[cur_idx+idx_start]=region_i + values_counts = {k: v for k, v in values_counts.items() if v > n_min_threshold} + # check if not entire domain filled as one feature + if 0 in values_counts: + # Remove background counts: + values_counts.pop(0) + # create empty list to store individual features for this threshold + list_features_threshold = [] + # create empty dict to store regions for individual features for this threshold + regions = dict() + # create emptry list of features to remove from parent threshold value + # loop over individual regions: + for cur_idx, count in values_counts.items(): + region = labels[:, :] == cur_idx + [hdim1_indices, hdim2_indeces] = np.nonzero(region) + # write region for individual threshold and feature to dict + region_i = list(zip(hdim1_indices, hdim2_indeces)) + regions[cur_idx + idx_start] = region_i # Determine feature position for region by one of the following methods: - hdim1_index,hdim2_index=feature_position(hdim1_indices,hdim2_indeces,region,data_i,threshold,position_threshold,target) - #create individual DataFrame row in tracky format for identified feature - list_features_threshold.append({'frame': int(i_time), - 'idx':cur_idx+idx_start, - 'hdim_1': hdim1_index, - 'hdim_2':hdim2_index, - 'num':count, - 'threshold_value':threshold}) - features_threshold=pd.DataFrame(list_features_threshold) + hdim1_index, hdim2_index = feature_position( + hdim1_indices, + hdim2_indeces, + region, + data_i, + threshold, + position_threshold, + target, + ) + # create individual DataFrame row in tracky format for identified feature + list_features_threshold.append( + { + "frame": int(i_time), + "idx": cur_idx + idx_start, + "hdim_1": hdim1_index, + "hdim_2": hdim2_index, + "num": count, + "threshold_value": threshold, + } + ) + features_threshold = pd.DataFrame(list_features_threshold) else: - features_threshold=pd.DataFrame() - regions=dict() - + features_threshold = pd.DataFrame() + regions = dict() + return features_threshold, regions - -def feature_detection_multithreshold_timestep(data_i,i_time, - threshold=None, - min_num=0, - target='maximum', - position_threshold='center', - sigma_threshold=0.5, - n_erosion_threshold=0, - n_min_threshold=0, - min_distance=0, - feature_number_start=1 - ): - ''' + + +def feature_detection_multithreshold_timestep( + data_i, + i_time, + threshold=None, + min_num=0, + target="maximum", + position_threshold="center", + sigma_threshold=0.5, + n_erosion_threshold=0, + n_min_threshold=0, + min_distance=0, + feature_number_start=1, +): + """ function to find features in each timestep based on iteratively finding regions above/below a set of thresholds Input: data_i: iris.cube.Cube 2D field to perform the feature detection (single timestep) i_time: int - number of the current timestep - + number of the current timestep + threshold: list of floats threshold values used to select target regions to track dxy: float @@ -234,61 +268,74 @@ def feature_detection_multithreshold_timestep(data_i,i_time, feature_number_start: int feature number to start with Output: - features_threshold: pandas DataFrame + features_threshold: pandas DataFrame detected features for individual timestep - ''' + """ from scipy.ndimage.filters import gaussian_filter track_data = data_i.core_data() - track_data=gaussian_filter(track_data, sigma=sigma_threshold) #smooth data slightly to create rounded, continuous field + track_data = gaussian_filter( + track_data, sigma=sigma_threshold + ) # smooth data slightly to create rounded, continuous field # create empty lists to store regions and features for individual timestep - features_thresholds=pd.DataFrame() - for i_threshold,threshold_i in enumerate(threshold): - if (i_threshold>0 and not features_thresholds.empty): - idx_start=features_thresholds['idx'].max()+1 + features_thresholds = pd.DataFrame() + for i_threshold, threshold_i in enumerate(threshold): + if i_threshold > 0 and not features_thresholds.empty: + idx_start = features_thresholds["idx"].max() + 1 else: - idx_start=0 - features_threshold_i,regions_i=feature_detection_threshold(track_data,i_time, - threshold=threshold_i, - sigma_threshold=sigma_threshold, - min_num=min_num, - target=target, - position_threshold=position_threshold, - n_erosion_threshold=n_erosion_threshold, - n_min_threshold=n_min_threshold, - min_distance=min_distance, - idx_start=idx_start - ) + idx_start = 0 + features_threshold_i, regions_i = feature_detection_threshold( + track_data, + i_time, + threshold=threshold_i, + sigma_threshold=sigma_threshold, + min_num=min_num, + target=target, + position_threshold=position_threshold, + n_erosion_threshold=n_erosion_threshold, + n_min_threshold=n_min_threshold, + min_distance=min_distance, + idx_start=idx_start, + ) if any([x is not None for x in features_threshold_i]): - features_thresholds=features_thresholds.append(features_threshold_i) + features_thresholds = features_thresholds.append(features_threshold_i) # For multiple threshold, and features found both in the current and previous step, remove "parent" features from Dataframe - if (i_threshold>0 and not features_thresholds.empty and regions_old): + if i_threshold > 0 and not features_thresholds.empty and regions_old: # for each threshold value: check if newly found features are surrounded by feature based on less restrictive threshold - features_thresholds=remove_parents(features_thresholds,regions_i,regions_old) - regions_old=regions_i + features_thresholds = remove_parents( + features_thresholds, regions_i, regions_old + ) + regions_old = regions_i - logging.debug('Finished feature detection for threshold '+str(i_threshold) + ' : ' + str(threshold_i) ) + logging.debug( + "Finished feature detection for threshold " + + str(i_threshold) + + " : " + + str(threshold_i) + ) return features_thresholds -def feature_detection_multithreshold(field_in, - dxy, - threshold=None, - min_num=0, - target='maximum', - position_threshold='center', - sigma_threshold=0.5, - n_erosion_threshold=0, - n_min_threshold=0, - min_distance=0, - feature_number_start=1 - ): - ''' Function to perform feature detection based on contiguous regions above/below a threshold + +def feature_detection_multithreshold( + field_in, + dxy, + threshold=None, + min_num=0, + target="maximum", + position_threshold="center", + sigma_threshold=0.5, + n_erosion_threshold=0, + n_min_threshold=0, + min_distance=0, + feature_number_start=1, +): + """Function to perform feature detection based on contiguous regions above/below a threshold Input: field_in: iris.cube.Cube 2D field to perform the tracking on (needs to have coordinate 'time' along one of its dimensions) - + thresholds: list of floats threshold values used to select target regions to track dxy: float @@ -306,95 +353,116 @@ def feature_detection_multithreshold(field_in, min_distance: float minimum distance between detected features (m) Output: - features: pandas DataFrame + features: pandas DataFrame detected features - ''' + """ from .utils import add_coordinates - logging.debug('start feature detection based on thresholds') - + logging.debug("start feature detection based on thresholds") + # create empty list to store features for all timesteps - list_features_timesteps=[] + list_features_timesteps = [] # loop over timesteps for feature identification: - data_time=field_in.slices_over('time') - + data_time = field_in.slices_over("time") + # if single threshold is put in as a single value, turn it into a list - if type(threshold) in [int,float]: - threshold=[threshold] - - for i_time,data_i in enumerate(data_time): - time_i=data_i.coord('time').units.num2date(data_i.coord('time').points[0]) - features_thresholds=feature_detection_multithreshold_timestep(data_i,i_time, - threshold=threshold, - sigma_threshold=sigma_threshold, - min_num=min_num, - target=target, - position_threshold=position_threshold, - n_erosion_threshold=n_erosion_threshold, - n_min_threshold=n_min_threshold, - min_distance=min_distance, - feature_number_start=feature_number_start - ) - #check if list of features is not empty, then merge features from different threshold values - #into one DataFrame and append to list for individual timesteps: + if type(threshold) in [int, float]: + threshold = [threshold] + + for i_time, data_i in enumerate(data_time): + time_i = data_i.coord("time").units.num2date(data_i.coord("time").points[0]) + features_thresholds = feature_detection_multithreshold_timestep( + data_i, + i_time, + threshold=threshold, + sigma_threshold=sigma_threshold, + min_num=min_num, + target=target, + position_threshold=position_threshold, + n_erosion_threshold=n_erosion_threshold, + n_min_threshold=n_min_threshold, + min_distance=min_distance, + feature_number_start=feature_number_start, + ) + # check if list of features is not empty, then merge features from different threshold values + # into one DataFrame and append to list for individual timesteps: if not features_thresholds.empty: - #Loop over DataFrame to remove features that are closer than distance_min to each other: - if (min_distance > 0): - features_thresholds=filter_min_distance(features_thresholds,dxy,min_distance) + # Loop over DataFrame to remove features that are closer than distance_min to each other: + if min_distance > 0: + features_thresholds = filter_min_distance( + features_thresholds, dxy, min_distance + ) list_features_timesteps.append(features_thresholds) - - logging.debug('Finished feature detection for ' + time_i.strftime('%Y-%m-%d_%H:%M:%S')) - logging.debug('feature detection: merging DataFrames') + logging.debug( + "Finished feature detection for " + time_i.strftime("%Y-%m-%d_%H:%M:%S") + ) + + logging.debug("feature detection: merging DataFrames") # Check if features are detected and then concatenate features from different timesteps into one pandas DataFrame # If no features are detected raise error if any([not x.empty for x in list_features_timesteps]): - features=pd.concat(list_features_timesteps, ignore_index=True) - features['feature']=features.index+feature_number_start - # features_filtered = features.drop(features[features['num'] < min_num].index) - # features_filtered.drop(columns=['idx','num','threshold_value'],inplace=True) - features=add_coordinates(features,field_in) + features = pd.concat(list_features_timesteps, ignore_index=True) + features["feature"] = features.index + feature_number_start + # features_filtered = features.drop(features[features['num'] < min_num].index) + # features_filtered.drop(columns=['idx','num','threshold_value'],inplace=True) + features = add_coordinates(features, field_in) else: - features=None - logging.info('No features detected') - logging.debug('feature detection completed') + features = None + logging.info("No features detected") + logging.debug("feature detection completed") return features -def filter_min_distance(features,dxy,min_distance): - ''' Function to perform feature detection based on contiguous regions above/below a threshold - Input: - features: pandas DataFrame + +def filter_min_distance(features, dxy, min_distance): + """Function to perform feature detection based on contiguous regions above/below a threshold + Input: + features: pandas DataFrame features dxy: float horzontal grid spacing (m) min_distance: float minimum distance between detected features (m) Output: - features: pandas DataFrame + features: pandas DataFrame features - ''' + """ from itertools import combinations - remove_list_distance=[] - #create list of tuples with all combinations of features at the timestep: - indeces=combinations(features.index.values,2) - #Loop over combinations to remove features that are closer together than min_distance and keep larger one (either higher threshold or larger area) - for index_1,index_2 in indeces: + + remove_list_distance = [] + # create list of tuples with all combinations of features at the timestep: + indeces = combinations(features.index.values, 2) + # Loop over combinations to remove features that are closer together than min_distance and keep larger one (either higher threshold or larger area) + for index_1, index_2 in indeces: if index_1 is not index_2: - features.loc[index_1,'hdim_1'] - distance=dxy*np.sqrt((features.loc[index_1,'hdim_1']-features.loc[index_2,'hdim_1'])**2+(features.loc[index_1,'hdim_2']-features.loc[index_2,'hdim_2'])**2) + features.loc[index_1, "hdim_1"] + distance = dxy * np.sqrt( + (features.loc[index_1, "hdim_1"] - features.loc[index_2, "hdim_1"]) ** 2 + + (features.loc[index_1, "hdim_2"] - features.loc[index_2, "hdim_2"]) + ** 2 + ) if distance <= min_distance: -# logging.debug('distance<= min_distance: ' + str(distance)) - if features.loc[index_1,'threshold_value']>features.loc[index_2,'threshold_value']: + # logging.debug('distance<= min_distance: ' + str(distance)) + if ( + features.loc[index_1, "threshold_value"] + > features.loc[index_2, "threshold_value"] + ): remove_list_distance.append(index_2) - elif features.loc[index_1,'threshold_value']features.loc[index_2,'num']: + elif ( + features.loc[index_1, "threshold_value"] + == features.loc[index_2, "threshold_value"] + ): + if features.loc[index_1, "num"] > features.loc[index_2, "num"]: remove_list_distance.append(index_2) - elif features.loc[index_1,'num'] axis_extent[0]) - & (track['longitude'] < axis_extent[1]) - & (track['latitude'] > axis_extent[2]) - & (track['latitude'] < axis_extent[3])] - - - #Plot tracked features by looping over rows of Dataframe - for i_row,row in track.iterrows(): - feature=row['feature'] - cell=row['cell'] + track = track.loc[ + (track["longitude"] > axis_extent[0]) + & (track["longitude"] < axis_extent[1]) + & (track["latitude"] > axis_extent[2]) + & (track["latitude"] < axis_extent[3]) + ] + + # Plot tracked features by looping over rows of Dataframe + for i_row, row in track.iterrows(): + feature = row["feature"] + cell = row["cell"] if not np.isnan(cell): - color=colors_mask[int(cell%len(colors_mask))] + color = colors_mask[int(cell % len(colors_mask))] if plot_number: - cell_string=' '+str(int(row['cell'])) - axes.text(row['longitude'],row['latitude'],cell_string, - color=color,fontsize=6, clip_on=True) + cell_string = " " + str(int(row["cell"])) + axes.text( + row["longitude"], + row["latitude"], + cell_string, + color=color, + fontsize=6, + clip_on=True, + ) else: - color='grey' + color = "grey" if plot_outline: - mask_i=None + mask_i = None # if mask is 3D, create surface projection, if mask is 2D keep the mask - if mask.ndim==2: - mask_i=mask_features(mask,feature,masked=False) - elif mask.ndim==3: - mask_i=mask_features_surface(mask,feature,masked=False,z_coord='model_level_number') + if mask.ndim == 2: + mask_i = mask_features(mask, feature, masked=False) + elif mask.ndim == 3: + mask_i = mask_features_surface( + mask, feature, masked=False, z_coord="model_level_number" + ) else: - raise ValueError('mask has shape that cannot be understood') + raise ValueError("mask has shape that cannot be understood") # plot countour lines around the edges of the mask - iplt.contour(mask_i,coords=['longitude','latitude'], - levels=[0,feature], - colors=color,linewidths=linewidth_contour, - axes=axes) + iplt.contour( + mask_i, + coords=["longitude", "latitude"], + levels=[0, feature], + colors=color, + linewidths=linewidth_contour, + axes=axes, + ) if plot_marker: - axes.plot(row['longitude'],row['latitude'], - color=color,marker=marker_track,markersize=markersize_track) + axes.plot( + row["longitude"], + row["latitude"], + color=color, + marker=marker_track, + markersize=markersize_track, + ) axes.set_extent(axis_extent) return axes -def animation_mask_field(track,features,field,mask,interval=500,figsize=(10,10),**kwargs): + +def animation_mask_field( + track, features, field, mask, interval=500, figsize=(10, 10), **kwargs +): import cartopy.crs as ccrs import matplotlib.pyplot as plt import matplotlib.animation from iris import Constraint - fig=plt.figure(figsize=figsize) + fig = plt.figure(figsize=figsize) plt.close() def update(time_in): fig.clf() - ax=fig.add_subplot(111,projection=ccrs.PlateCarree()) + ax = fig.add_subplot(111, projection=ccrs.PlateCarree()) constraint_time = Constraint(time=time_in) - field_i=field.extract(constraint_time) - mask_i=mask.extract(constraint_time) - track_i=track[track['time']==time_in] - features_i=features[features['time']==time_in] - #fig1,ax1=plt.subplots(ncols=1, nrows=1,figsize=figsize, subplot_kw={'projection': ccrs.PlateCarree()}) - plot_tobac=plot_tracks_mask_field(track_i,field=field_i,mask=mask_i,features=features_i, - axes=ax, - **kwargs) - ax.set_title('{}'.format(time_in)) - - time=field.coord('time') - datetimes=time.units.num2date(time.points) - animation = matplotlib.animation.FuncAnimation(fig, update,init_func=None, frames=datetimes,interval=interval, blit=False) + field_i = field.extract(constraint_time) + mask_i = mask.extract(constraint_time) + track_i = track[track["time"] == time_in] + features_i = features[features["time"] == time_in] + # fig1,ax1=plt.subplots(ncols=1, nrows=1,figsize=figsize, subplot_kw={'projection': ccrs.PlateCarree()}) + plot_tobac = plot_tracks_mask_field( + track_i, field=field_i, mask=mask_i, features=features_i, axes=ax, **kwargs + ) + ax.set_title("{}".format(time_in)) + + time = field.coord("time") + datetimes = time.units.num2date(time.points) + animation = matplotlib.animation.FuncAnimation( + fig, update, init_func=None, frames=datetimes, interval=interval, blit=False + ) return animation -def plot_mask_cell_track_follow(cell,track, cog, features, mask_total, - field_contour, field_filled, - width=10000, - name= 'test', plotdir='./', - file_format=['png'],figsize=(10/2.54, 10/2.54),dpi=300, - **kwargs): - '''Make plots for all cells centred around cell and with one background field as filling and one background field as contrours + +def plot_mask_cell_track_follow( + cell, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + width=10000, + name="test", + plotdir="./", + file_format=["png"], + figsize=(10 / 2.54, 10 / 2.54), + dpi=300, + **kwargs +): + """Make plots for all cells centred around cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ from iris import Constraint from numpy import unique import os - track_cell=track[track['cell']==cell] - for i_row,row in track_cell.iterrows(): - - constraint_time = Constraint(time=row['time']) - constraint_x = Constraint(projection_x_coordinate = lambda cell: row['projection_x_coordinate']-width < cell < row['projection_x_coordinate']+width) - constraint_y = Constraint(projection_y_coordinate = lambda cell: row['projection_y_coordinate']-width < cell < row['projection_y_coordinate']+width) + track_cell = track[track["cell"] == cell] + for i_row, row in track_cell.iterrows(): + + constraint_time = Constraint(time=row["time"]) + constraint_x = Constraint( + projection_x_coordinate=lambda cell: row["projection_x_coordinate"] - width + < cell + < row["projection_x_coordinate"] + width + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: row["projection_y_coordinate"] - width + < cell + < row["projection_y_coordinate"] + width + ) constraint = constraint_time & constraint_x & constraint_y - mask_total_i=mask_total.extract(constraint) + mask_total_i = mask_total.extract(constraint) if field_contour is None: - field_contour_i=None + field_contour_i = None else: - field_contour_i=field_contour.extract(constraint) + field_contour_i = field_contour.extract(constraint) if field_filled is None: - field_filled_i=None + field_filled_i = None else: - field_filled_i=field_filled.extract(constraint) + field_filled_i = field_filled.extract(constraint) - cells=list(unique(mask_total_i.core_data())) + cells = list(unique(mask_total_i.core_data())) if cell not in cells: cells.append(cell) if 0 in cells: cells.remove(0) - track_i=track[track['cell'].isin(cells)] - track_i=track_i[track_i['time']==row['time']] + track_i = track[track["cell"].isin(cells)] + track_i = track_i[track_i["time"] == row["time"]] if cog is None: - cog_i=None + cog_i = None else: - cog_i=cog[cog['cell'].isin(cells)] - cog_i=cog_i[cog_i['time']==row['time']] + cog_i = cog[cog["cell"].isin(cells)] + cog_i = cog_i[cog_i["time"] == row["time"]] if features is None: - features_i=None + features_i = None else: - features_i=features[features['time']==row['time']] - + features_i = features[features["time"] == row["time"]] fig1, ax1 = plt.subplots(ncols=1, nrows=1, figsize=figsize) fig1.subplots_adjust(left=0.2, bottom=0.15, right=0.85, top=0.80) - - - - datestring_stamp = row['time'].strftime('%Y-%m-%d %H:%M:%S') - celltime_stamp = "%02d:%02d:%02d" % (row['time_cell'].dt.total_seconds() // 3600,(row['time_cell'].dt.total_seconds() % 3600) // 60, row['time_cell'].dt.total_seconds() % 60 ) - title=datestring_stamp + ' , ' + celltime_stamp - datestring_file = row['time'].strftime('%Y-%m-%d_%H%M%S') - - ax1=plot_mask_cell_individual_follow(cell_i=cell,track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - width=width, - axes=ax1,title=title, - **kwargs) + datestring_stamp = row["time"].strftime("%Y-%m-%d %H:%M:%S") + celltime_stamp = "%02d:%02d:%02d" % ( + row["time_cell"].dt.total_seconds() // 3600, + (row["time_cell"].dt.total_seconds() % 3600) // 60, + row["time_cell"].dt.total_seconds() % 60, + ) + title = datestring_stamp + " , " + celltime_stamp + datestring_file = row["time"].strftime("%Y-%m-%d_%H%M%S") + + ax1 = plot_mask_cell_individual_follow( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + width=width, + axes=ax1, + title=title, + **kwargs + ) out_dir = os.path.join(plotdir, name) os.makedirs(out_dir, exist_ok=True) - if 'png' in file_format: - savepath_png = os.path.join(out_dir, name + '_' + datestring_file + '.png') + if "png" in file_format: + savepath_png = os.path.join(out_dir, name + "_" + datestring_file + ".png") fig1.savefig(savepath_png, dpi=dpi) - logging.debug('field_contour field_filled Mask plot saved to ' + savepath_png) - if 'pdf' in file_format: - savepath_pdf = os.path.join(out_dir, name + '_' + datestring_file + '.pdf') + logging.debug( + "field_contour field_filled Mask plot saved to " + savepath_png + ) + if "pdf" in file_format: + savepath_pdf = os.path.join(out_dir, name + "_" + datestring_file + ".pdf") fig1.savefig(savepath_pdf, dpi=dpi) - logging.debug('field_contour field_filled Mask plot saved to ' + savepath_pdf) + logging.debug( + "field_contour field_filled Mask plot saved to " + savepath_pdf + ) plt.close() plt.clf() -def plot_mask_cell_individual_follow(cell_i,track, cog,features, mask_total, - field_contour, field_filled, - axes=None,width=10000, - label_field_contour=None, cmap_field_contour='Blues',norm_field_contour=None, - linewidths_contour=0.8,contour_labels=False, - vmin_field_contour=0,vmax_field_contour=50,levels_field_contour=None,nlevels_field_contour=10, - label_field_filled=None,cmap_field_filled='summer',norm_field_filled=None, - vmin_field_filled=0,vmax_field_filled=100,levels_field_filled=None,nlevels_field_filled=10, - title=None - ): - '''Make individual plot for cell centred around cell and with one background field as filling and one background field as contrours +def plot_mask_cell_individual_follow( + cell_i, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + axes=None, + width=10000, + label_field_contour=None, + cmap_field_contour="Blues", + norm_field_contour=None, + linewidths_contour=0.8, + contour_labels=False, + vmin_field_contour=0, + vmax_field_contour=50, + levels_field_contour=None, + nlevels_field_contour=10, + label_field_filled=None, + cmap_field_filled="summer", + norm_field_filled=None, + vmin_field_filled=0, + vmax_field_filled=100, + levels_field_filled=None, + nlevels_field_filled=10, + title=None, +): + """Make individual plot for cell centred around cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ import numpy as np - from .utils import mask_cell_surface + from .utils import mask_cell_surface from mpl_toolkits.axes_grid1 import make_axes_locatable from matplotlib.colors import Normalize - divider = make_axes_locatable(axes) - x_pos=track[track['cell']==cell_i]['projection_x_coordinate'].item() - y_pos=track[track['cell']==cell_i]['projection_y_coordinate'].item() + x_pos = track[track["cell"] == cell_i]["projection_x_coordinate"].item() + y_pos = track[track["cell"] == cell_i]["projection_y_coordinate"].item() if field_filled is not None: if levels_field_filled is None: - levels_field_filled=np.linspace(vmin_field_filled,vmax_field_filled, nlevels_field_filled) - plot_field_filled = axes.contourf((field_filled.coord('projection_x_coordinate').points-x_pos)/1000, - (field_filled.coord('projection_y_coordinate').points-y_pos)/1000, - field_filled.data, - cmap=cmap_field_filled,norm=norm_field_filled, - levels=levels_field_filled,vmin=vmin_field_filled, vmax=vmax_field_filled) - + levels_field_filled = np.linspace( + vmin_field_filled, vmax_field_filled, nlevels_field_filled + ) + plot_field_filled = axes.contourf( + (field_filled.coord("projection_x_coordinate").points - x_pos) / 1000, + (field_filled.coord("projection_y_coordinate").points - y_pos) / 1000, + field_filled.data, + cmap=cmap_field_filled, + norm=norm_field_filled, + levels=levels_field_filled, + vmin=vmin_field_filled, + vmax=vmax_field_filled, + ) cax_filled = divider.append_axes("right", size="5%", pad=0.1) - norm_filled= Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) - sm_filled= plt.cm.ScalarMappable(norm=norm_filled, cmap = plot_field_filled.cmap) + norm_filled = Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) + sm_filled = plt.cm.ScalarMappable(norm=norm_filled, cmap=plot_field_filled.cmap) sm_filled.set_array([]) - cbar_field_filled = plt.colorbar(sm_filled, orientation='vertical',cax=cax_filled) + cbar_field_filled = plt.colorbar( + sm_filled, orientation="vertical", cax=cax_filled + ) cbar_field_filled.ax.set_ylabel(label_field_filled) cbar_field_filled.set_clim(vmin_field_filled, vmax_field_filled) if field_contour is not None: if levels_field_contour is None: - levels_field_contour=np.linspace(vmin_field_contour, vmax_field_contour, nlevels_field_contour) + levels_field_contour = np.linspace( + vmin_field_contour, vmax_field_contour, nlevels_field_contour + ) if norm_field_contour: - vmin_field_contour=None, - vmax_field_contour=None, - - plot_field_contour = axes.contour((field_contour.coord('projection_x_coordinate').points-x_pos)/1000, - (field_contour.coord('projection_y_coordinate').points-y_pos)/1000, - field_contour.data, - cmap=cmap_field_contour,norm=norm_field_contour, - levels=levels_field_contour,vmin=vmin_field_contour, vmax=vmax_field_contour, - linewidths=linewidths_contour) + vmin_field_contour = (None,) + vmax_field_contour = (None,) + + plot_field_contour = axes.contour( + (field_contour.coord("projection_x_coordinate").points - x_pos) / 1000, + (field_contour.coord("projection_y_coordinate").points - y_pos) / 1000, + field_contour.data, + cmap=cmap_field_contour, + norm=norm_field_contour, + levels=levels_field_contour, + vmin=vmin_field_contour, + vmax=vmax_field_contour, + linewidths=linewidths_contour, + ) if contour_labels: axes.clabel(plot_field_contour, fontsize=10) cax_contour = divider.append_axes("bottom", size="5%", pad=0.1) if norm_field_contour: - vmin_field_contour=None - vmax_field_contour=None - norm_contour=norm_field_contour + vmin_field_contour = None + vmax_field_contour = None + norm_contour = norm_field_contour else: - norm_contour= Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) + norm_contour = Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) - sm_contour= plt.cm.ScalarMappable(norm=norm_contour, cmap = plot_field_contour.cmap) + sm_contour = plt.cm.ScalarMappable( + norm=norm_contour, cmap=plot_field_contour.cmap + ) sm_contour.set_array([]) - cbar_field_contour = plt.colorbar(sm_contour, orientation='horizontal',ticks=levels_field_contour,cax=cax_contour) + cbar_field_contour = plt.colorbar( + sm_contour, + orientation="horizontal", + ticks=levels_field_contour, + cax=cax_contour, + ) cbar_field_contour.ax.set_xlabel(label_field_contour) cbar_field_contour.set_clim(vmin_field_contour, vmax_field_contour) - for i_row, row in track.iterrows(): - cell = int(row['cell']) - if cell==cell_i: - color='darkred' + cell = int(row["cell"]) + if cell == cell_i: + color = "darkred" else: - color='darkorange' - - cell_string=' '+str(int(row['cell'])) - axes.text((row['projection_x_coordinate']-x_pos)/1000, - (row['projection_y_coordinate']-y_pos)/1000, - cell_string,color=color,fontsize=6, clip_on=True) + color = "darkorange" + + cell_string = " " + str(int(row["cell"])) + axes.text( + (row["projection_x_coordinate"] - x_pos) / 1000, + (row["projection_y_coordinate"] - y_pos) / 1000, + cell_string, + color=color, + fontsize=6, + clip_on=True, + ) # Plot marker for tracked cell centre as a cross - axes.plot((row['projection_x_coordinate']-x_pos)/1000, - (row['projection_y_coordinate']-y_pos)/1000, - 'x', color=color,markersize=4) - - - #Create surface projection of mask for the respective cell and plot it in the right color - z_coord = 'model_level_number' - if len(mask_total.shape)==3: - mask_total_i_surface = mask_cell_surface(mask_total, cell, track, masked=False, z_coord=z_coord) - elif len(mask_total.shape)==2: - mask_total_i_surface=mask_total - axes.contour((mask_total_i_surface.coord('projection_x_coordinate').points-x_pos)/1000, - (mask_total_i_surface.coord('projection_y_coordinate').points-y_pos)/1000, - mask_total_i_surface.data, - levels=[0, cell], colors=color, linestyles=':',linewidth=1) + axes.plot( + (row["projection_x_coordinate"] - x_pos) / 1000, + (row["projection_y_coordinate"] - y_pos) / 1000, + "x", + color=color, + markersize=4, + ) + + # Create surface projection of mask for the respective cell and plot it in the right color + z_coord = "model_level_number" + if len(mask_total.shape) == 3: + mask_total_i_surface = mask_cell_surface( + mask_total, cell, track, masked=False, z_coord=z_coord + ) + elif len(mask_total.shape) == 2: + mask_total_i_surface = mask_total + axes.contour( + (mask_total_i_surface.coord("projection_x_coordinate").points - x_pos) + / 1000, + (mask_total_i_surface.coord("projection_y_coordinate").points - y_pos) + / 1000, + mask_total_i_surface.data, + levels=[0, cell], + colors=color, + linestyles=":", + linewidth=1, + ) if cog is not None: for i_row, row in cog.iterrows(): - cell = row['cell'] + cell = row["cell"] - if cell==cell_i: - color='darkred' + if cell == cell_i: + color = "darkred" else: - color='darkorange' + color = "darkorange" # plot marker for centre of gravity as a circle - axes.plot((row['x_M']-x_pos)/1000, (row['y_M']-y_pos)/1000, - 'o', markeredgecolor=color, markerfacecolor='None',markersize=4) + axes.plot( + (row["x_M"] - x_pos) / 1000, + (row["y_M"] - y_pos) / 1000, + "o", + markeredgecolor=color, + markerfacecolor="None", + markersize=4, + ) if features is not None: for i_row, row in features.iterrows(): - color='purple' - axes.plot((row['projection_x_coordinate']-x_pos)/1000, - (row['projection_y_coordinate']-y_pos)/1000, - '+', color=color,markersize=3) - - axes.set_xlabel('x (km)') - axes.set_ylabel('y (km)') - axes.set_xlim([-1*width/1000, width/1000]) - axes.set_ylim([-1*width/1000, width/1000]) - axes.xaxis.set_label_position('top') - axes.xaxis.set_ticks_position('top') - axes.set_title(title,pad=35,fontsize=10,horizontalalignment='left',loc='left') + color = "purple" + axes.plot( + (row["projection_x_coordinate"] - x_pos) / 1000, + (row["projection_y_coordinate"] - y_pos) / 1000, + "+", + color=color, + markersize=3, + ) + + axes.set_xlabel("x (km)") + axes.set_ylabel("y (km)") + axes.set_xlim([-1 * width / 1000, width / 1000]) + axes.set_ylim([-1 * width / 1000, width / 1000]) + axes.xaxis.set_label_position("top") + axes.xaxis.set_ticks_position("top") + axes.set_title(title, pad=35, fontsize=10, horizontalalignment="left", loc="left") return axes -def plot_mask_cell_track_static(cell,track, cog, features, mask_total, - field_contour, field_filled, - width=10000,n_extend=1, - name= 'test', plotdir='./', - file_format=['png'],figsize=(10/2.54, 10/2.54),dpi=300, - **kwargs): - '''Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours + +def plot_mask_cell_track_static( + cell, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + width=10000, + n_extend=1, + name="test", + plotdir="./", + file_format=["png"], + figsize=(10 / 2.54, 10 / 2.54), + dpi=300, + **kwargs +): + """Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ from iris import Constraint from numpy import unique import os - track_cell=track[track['cell']==cell] - x_min=track_cell['projection_x_coordinate'].min()-width - x_max=track_cell['projection_x_coordinate'].max()+width - y_min=track_cell['projection_y_coordinate'].min()-width - y_max=track_cell['projection_y_coordinate'].max()+width - - #set up looping over time based on mask's time coordinate to allow for one timestep before and after the track - time_coord=mask_total.coord('time') - time=time_coord.units.num2date(time_coord.points) - i_start=max(0,np.where(time==track_cell['time'].values[0])[0][0]-n_extend) - i_end=min(len(time)-1,np.where(time==track_cell['time'].values[-1])[0][0]+n_extend+1) - time_cell=time[slice(i_start,i_end)] + + track_cell = track[track["cell"] == cell] + x_min = track_cell["projection_x_coordinate"].min() - width + x_max = track_cell["projection_x_coordinate"].max() + width + y_min = track_cell["projection_y_coordinate"].min() - width + y_max = track_cell["projection_y_coordinate"].max() + width + + # set up looping over time based on mask's time coordinate to allow for one timestep before and after the track + time_coord = mask_total.coord("time") + time = time_coord.units.num2date(time_coord.points) + i_start = max(0, np.where(time == track_cell["time"].values[0])[0][0] - n_extend) + i_end = min( + len(time) - 1, + np.where(time == track_cell["time"].values[-1])[0][0] + n_extend + 1, + ) + time_cell = time[slice(i_start, i_end)] for time_i in time_cell: -# for i_row,row in track_cell.iterrows(): -# time_i=row['time'] -# constraint_time = Constraint(time=row['time']) + # for i_row,row in track_cell.iterrows(): + # time_i=row['time'] + # constraint_time = Constraint(time=row['time']) constraint_time = Constraint(time=time_i) - constraint_x = Constraint(projection_x_coordinate = lambda cell: x_min < cell < x_max) - constraint_y = Constraint(projection_y_coordinate = lambda cell: y_min < cell < y_max) + constraint_x = Constraint( + projection_x_coordinate=lambda cell: x_min < cell < x_max + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: y_min < cell < y_max + ) constraint = constraint_time & constraint_x & constraint_y - mask_total_i=mask_total.extract(constraint) + mask_total_i = mask_total.extract(constraint) if field_contour is None: - field_contour_i=None + field_contour_i = None else: - field_contour_i=field_contour.extract(constraint) + field_contour_i = field_contour.extract(constraint) if field_filled is None: - field_filled_i=None + field_filled_i = None else: - field_filled_i=field_filled.extract(constraint) - - - track_i=track[track['time']==time_i] - - cells_mask=list(unique(mask_total_i.core_data())) - track_cells=track_i.loc[(track_i['projection_x_coordinate'] > x_min) & (track_i['projection_x_coordinate'] < x_max) & (track_i['projection_y_coordinate'] > y_min) & (track_i['projection_y_coordinate'] < y_max)] - cells_track=list(track_cells['cell'].values) - cells=list(set( cells_mask + cells_track )) + field_filled_i = field_filled.extract(constraint) + + track_i = track[track["time"] == time_i] + + cells_mask = list(unique(mask_total_i.core_data())) + track_cells = track_i.loc[ + (track_i["projection_x_coordinate"] > x_min) + & (track_i["projection_x_coordinate"] < x_max) + & (track_i["projection_y_coordinate"] > y_min) + & (track_i["projection_y_coordinate"] < y_max) + ] + cells_track = list(track_cells["cell"].values) + cells = list(set(cells_mask + cells_track)) if cell not in cells: cells.append(cell) if 0 in cells: cells.remove(0) - track_i=track_i[track_i['cell'].isin(cells)] + track_i = track_i[track_i["cell"].isin(cells)] if cog is None: - cog_i=None + cog_i = None else: - cog_i=cog[cog['cell'].isin(cells)] - cog_i=cog_i[cog_i['time']==time_i] + cog_i = cog[cog["cell"].isin(cells)] + cog_i = cog_i[cog_i["time"] == time_i] if features is None: - features_i=None + features_i = None else: - features_i=features[features['time']==time_i] - + features_i = features[features["time"] == time_i] fig1, ax1 = plt.subplots(ncols=1, nrows=1, figsize=figsize) fig1.subplots_adjust(left=0.2, bottom=0.15, right=0.80, top=0.85) - datestring_stamp = time_i.strftime('%Y-%m-%d %H:%M:%S') - if time_i in track_cell['time'].values: - time_cell_i=track_cell[track_cell['time'].values==time_i]['time_cell'] - celltime_stamp = "%02d:%02d:%02d" % (time_cell_i.dt.total_seconds() // 3600, - (time_cell_i.dt.total_seconds() % 3600) // 60, - time_cell_i.dt.total_seconds() % 60 ) + datestring_stamp = time_i.strftime("%Y-%m-%d %H:%M:%S") + if time_i in track_cell["time"].values: + time_cell_i = track_cell[track_cell["time"].values == time_i]["time_cell"] + celltime_stamp = "%02d:%02d:%02d" % ( + time_cell_i.dt.total_seconds() // 3600, + (time_cell_i.dt.total_seconds() % 3600) // 60, + time_cell_i.dt.total_seconds() % 60, + ) else: - celltime_stamp=' - ' - title=datestring_stamp + ' , ' + celltime_stamp - datestring_file = time_i.strftime('%Y-%m-%d_%H%M%S') - - ax1=plot_mask_cell_individual_static(cell_i=cell, - track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - xlim=[x_min/1000,x_max/1000],ylim=[y_min/1000,y_max/1000], - axes=ax1,title=title,**kwargs) + celltime_stamp = " - " + title = datestring_stamp + " , " + celltime_stamp + datestring_file = time_i.strftime("%Y-%m-%d_%H%M%S") + + ax1 = plot_mask_cell_individual_static( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + xlim=[x_min / 1000, x_max / 1000], + ylim=[y_min / 1000, y_max / 1000], + axes=ax1, + title=title, + **kwargs + ) out_dir = os.path.join(plotdir, name) os.makedirs(out_dir, exist_ok=True) - if 'png' in file_format: - savepath_png = os.path.join(out_dir, name + '_' + datestring_file + '.png') + if "png" in file_format: + savepath_png = os.path.join(out_dir, name + "_" + datestring_file + ".png") fig1.savefig(savepath_png, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_png) - if 'pdf' in file_format: - savepath_pdf = os.path.join(out_dir, name + '_' + datestring_file + '.pdf') + logging.debug("Mask static plot saved to " + savepath_png) + if "pdf" in file_format: + savepath_pdf = os.path.join(out_dir, name + "_" + datestring_file + ".pdf") fig1.savefig(savepath_pdf, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_pdf) + logging.debug("Mask static plot saved to " + savepath_pdf) plt.close() plt.clf() -def plot_mask_cell_individual_static(cell_i,track, cog, features, mask_total, - field_contour, field_filled, - axes=None,xlim=None,ylim=None, - label_field_contour=None, cmap_field_contour='Blues',norm_field_contour=None, - linewidths_contour=0.8,contour_labels=False, - vmin_field_contour=0,vmax_field_contour=50,levels_field_contour=None,nlevels_field_contour=10, - label_field_filled=None,cmap_field_filled='summer',norm_field_filled=None, - vmin_field_filled=0,vmax_field_filled=100,levels_field_filled=None,nlevels_field_filled=10, - title=None,feature_number=False - ): - '''Make plots for cell in fixed frame and with one background field as filling and one background field as contrours +def plot_mask_cell_individual_static( + cell_i, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + axes=None, + xlim=None, + ylim=None, + label_field_contour=None, + cmap_field_contour="Blues", + norm_field_contour=None, + linewidths_contour=0.8, + contour_labels=False, + vmin_field_contour=0, + vmax_field_contour=50, + levels_field_contour=None, + nlevels_field_contour=10, + label_field_filled=None, + cmap_field_filled="summer", + norm_field_filled=None, + vmin_field_filled=0, + vmax_field_filled=100, + levels_field_filled=None, + nlevels_field_filled=10, + title=None, + feature_number=False, +): + """Make plots for cell in fixed frame and with one background field as filling and one background field as contrours Input: Output: - ''' + """ import numpy as np - from .utils import mask_features,mask_features_surface + from .utils import mask_features, mask_features_surface from mpl_toolkits.axes_grid1 import make_axes_locatable from matplotlib.colors import Normalize - divider = make_axes_locatable(axes) if field_filled is not None: if levels_field_filled is None: - levels_field_filled=np.linspace(vmin_field_filled,vmax_field_filled, 10) - plot_field_filled = axes.contourf(field_filled.coord('projection_x_coordinate').points/1000, - field_filled.coord('projection_y_coordinate').points/1000, - field_filled.data, - levels=levels_field_filled, norm=norm_field_filled, - cmap=cmap_field_filled, vmin=vmin_field_filled, vmax=vmax_field_filled) - + levels_field_filled = np.linspace(vmin_field_filled, vmax_field_filled, 10) + plot_field_filled = axes.contourf( + field_filled.coord("projection_x_coordinate").points / 1000, + field_filled.coord("projection_y_coordinate").points / 1000, + field_filled.data, + levels=levels_field_filled, + norm=norm_field_filled, + cmap=cmap_field_filled, + vmin=vmin_field_filled, + vmax=vmax_field_filled, + ) cax_filled = divider.append_axes("right", size="5%", pad=0.1) - norm_filled= Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) - sm1= plt.cm.ScalarMappable(norm=norm_filled, cmap = plot_field_filled.cmap) + norm_filled = Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) + sm1 = plt.cm.ScalarMappable(norm=norm_filled, cmap=plot_field_filled.cmap) sm1.set_array([]) - cbar_field_filled = plt.colorbar(sm1, orientation='vertical',cax=cax_filled) + cbar_field_filled = plt.colorbar(sm1, orientation="vertical", cax=cax_filled) cbar_field_filled.ax.set_ylabel(label_field_filled) cbar_field_filled.set_clim(vmin_field_filled, vmax_field_filled) - if field_contour is not None: if levels_field_contour is None: - levels_field_contour=np.linspace(vmin_field_contour, vmax_field_contour, 5) - plot_field_contour = axes.contour(field_contour.coord('projection_x_coordinate').points/1000, - field_contour.coord('projection_y_coordinate').points/1000, - field_contour.data, - cmap=cmap_field_contour,norm=norm_field_contour, - levels=levels_field_contour,vmin=vmin_field_contour, vmax=vmax_field_contour, - linewidths=linewidths_contour) + levels_field_contour = np.linspace( + vmin_field_contour, vmax_field_contour, 5 + ) + plot_field_contour = axes.contour( + field_contour.coord("projection_x_coordinate").points / 1000, + field_contour.coord("projection_y_coordinate").points / 1000, + field_contour.data, + cmap=cmap_field_contour, + norm=norm_field_contour, + levels=levels_field_contour, + vmin=vmin_field_contour, + vmax=vmax_field_contour, + linewidths=linewidths_contour, + ) if contour_labels: axes.clabel(plot_field_contour, fontsize=10) cax_contour = divider.append_axes("bottom", size="5%", pad=0.1) if norm_field_contour: - vmin_field_contour=None - vmax_field_contour=None - norm_contour=norm_field_contour + vmin_field_contour = None + vmax_field_contour = None + norm_contour = norm_field_contour else: - norm_contour= Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) + norm_contour = Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) - sm_contour= plt.cm.ScalarMappable(norm=norm_contour, cmap = plot_field_contour.cmap) + sm_contour = plt.cm.ScalarMappable( + norm=norm_contour, cmap=plot_field_contour.cmap + ) sm_contour.set_array([]) - cbar_field_contour = plt.colorbar(sm_contour, orientation='horizontal',ticks=levels_field_contour,cax=cax_contour) + cbar_field_contour = plt.colorbar( + sm_contour, + orientation="horizontal", + ticks=levels_field_contour, + cax=cax_contour, + ) cbar_field_contour.ax.set_xlabel(label_field_contour) cbar_field_contour.set_clim(vmin_field_contour, vmax_field_contour) for i_row, row in track.iterrows(): - cell = row['cell'] - feature = row['feature'] -# logging.debug("cell: "+ str(row['cell'])) -# logging.debug("feature: "+ str(row['feature'])) + cell = row["cell"] + feature = row["feature"] + # logging.debug("cell: "+ str(row['cell'])) + # logging.debug("feature: "+ str(row['feature'])) - if cell==cell_i: - color='darkred' + if cell == cell_i: + color = "darkred" if feature_number: - cell_string=' '+str(int(cell))+' ('+str(int(feature))+')' + cell_string = " " + str(int(cell)) + " (" + str(int(feature)) + ")" else: - cell_string=' '+str(int(cell)) + cell_string = " " + str(int(cell)) elif np.isnan(cell): - color='gray' + color = "gray" if feature_number: - cell_string=' '+'('+str(int(feature))+')' + cell_string = " " + "(" + str(int(feature)) + ")" else: - cell_string=' ' + cell_string = " " else: - color='darkorange' + color = "darkorange" if feature_number: - cell_string=' '+str(int(cell))+' ('+str(int(feature))+')' + cell_string = " " + str(int(cell)) + " (" + str(int(feature)) + ")" else: - cell_string=' '+str(int(cell)) + cell_string = " " + str(int(cell)) - axes.text(row['projection_x_coordinate']/1000, - row['projection_y_coordinate']/1000, - cell_string,color=color,fontsize=6, clip_on=True) + axes.text( + row["projection_x_coordinate"] / 1000, + row["projection_y_coordinate"] / 1000, + cell_string, + color=color, + fontsize=6, + clip_on=True, + ) # Plot marker for tracked cell centre as a cross - axes.plot(row['projection_x_coordinate']/1000, - row['projection_y_coordinate']/1000, - 'x', color=color,markersize=4) - - - #Create surface projection of mask for the respective cell and plot it in the right color - z_coord = 'model_level_number' - if len(mask_total.shape)==3: - mask_total_i_surface = mask_features_surface(mask_total, feature, masked=False, z_coord=z_coord) - elif len(mask_total.shape)==2: - mask_total_i_surface=mask_features(mask_total, feature, masked=False, z_coord=z_coord) - axes.contour(mask_total_i_surface.coord('projection_x_coordinate').points/1000, - mask_total_i_surface.coord('projection_y_coordinate').points/1000, - mask_total_i_surface.data, - levels=[0, feature], colors=color, linestyles=':',linewidth=1) + axes.plot( + row["projection_x_coordinate"] / 1000, + row["projection_y_coordinate"] / 1000, + "x", + color=color, + markersize=4, + ) + + # Create surface projection of mask for the respective cell and plot it in the right color + z_coord = "model_level_number" + if len(mask_total.shape) == 3: + mask_total_i_surface = mask_features_surface( + mask_total, feature, masked=False, z_coord=z_coord + ) + elif len(mask_total.shape) == 2: + mask_total_i_surface = mask_features( + mask_total, feature, masked=False, z_coord=z_coord + ) + axes.contour( + mask_total_i_surface.coord("projection_x_coordinate").points / 1000, + mask_total_i_surface.coord("projection_y_coordinate").points / 1000, + mask_total_i_surface.data, + levels=[0, feature], + colors=color, + linestyles=":", + linewidth=1, + ) if cog is not None: for i_row, row in cog.iterrows(): - cell = row['cell'] + cell = row["cell"] - if cell==cell_i: - color='darkred' + if cell == cell_i: + color = "darkred" else: - color='darkorange' + color = "darkorange" # plot marker for centre of gravity as a circle - axes.plot(row['x_M']/1000, row['y_M']/1000, - 'o', markeredgecolor=color, markerfacecolor='None',markersize=4) + axes.plot( + row["x_M"] / 1000, + row["y_M"] / 1000, + "o", + markeredgecolor=color, + markerfacecolor="None", + markersize=4, + ) if features is not None: for i_row, row in features.iterrows(): - color='purple' - axes.plot(row['projection_x_coordinate']/1000, - row['projection_y_coordinate']/1000, - '+', color=color,markersize=3) - - axes.set_xlabel('x (km)') - axes.set_ylabel('y (km)') + color = "purple" + axes.plot( + row["projection_x_coordinate"] / 1000, + row["projection_y_coordinate"] / 1000, + "+", + color=color, + markersize=3, + ) + + axes.set_xlabel("x (km)") + axes.set_ylabel("y (km)") axes.set_xlim(xlim) axes.set_ylim(ylim) - axes.xaxis.set_label_position('top') - axes.xaxis.set_ticks_position('top') - axes.set_title(title,pad=35,fontsize=10,horizontalalignment='left',loc='left') + axes.xaxis.set_label_position("top") + axes.xaxis.set_ticks_position("top") + axes.set_title(title, pad=35, fontsize=10, horizontalalignment="left", loc="left") return axes -def plot_mask_cell_track_2D3Dstatic(cell,track, cog, features, mask_total, - field_contour, field_filled, - width=10000,n_extend=1, - name= 'test', plotdir='./', - file_format=['png'],figsize=(10/2.54, 10/2.54),dpi=300, - ele=10,azim=30, - **kwargs): - '''Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours + +def plot_mask_cell_track_2D3Dstatic( + cell, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + width=10000, + n_extend=1, + name="test", + plotdir="./", + file_format=["png"], + figsize=(10 / 2.54, 10 / 2.54), + dpi=300, + ele=10, + azim=30, + **kwargs +): + """Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ from iris import Constraint from numpy import unique import os from mpl_toolkits.mplot3d import Axes3D import matplotlib.gridspec as gridspec - - track_cell=track[track['cell']==cell] - x_min=track_cell['projection_x_coordinate'].min()-width - x_max=track_cell['projection_x_coordinate'].max()+width - y_min=track_cell['projection_y_coordinate'].min()-width - y_max=track_cell['projection_y_coordinate'].max()+width - - #set up looping over time based on mask's time coordinate to allow for one timestep before and after the track - time_coord=mask_total.coord('time') - time=time_coord.units.num2date(time_coord.points) - i_start=max(0,np.where(time==track_cell['time'].values[0])[0][0]-n_extend) - i_end=min(len(time)-1,np.where(time==track_cell['time'].values[-1])[0][0]+n_extend+1) - time_cell=time[slice(i_start,i_end)] + track_cell = track[track["cell"] == cell] + x_min = track_cell["projection_x_coordinate"].min() - width + x_max = track_cell["projection_x_coordinate"].max() + width + y_min = track_cell["projection_y_coordinate"].min() - width + y_max = track_cell["projection_y_coordinate"].max() + width + + # set up looping over time based on mask's time coordinate to allow for one timestep before and after the track + time_coord = mask_total.coord("time") + time = time_coord.units.num2date(time_coord.points) + i_start = max(0, np.where(time == track_cell["time"].values[0])[0][0] - n_extend) + i_end = min( + len(time) - 1, + np.where(time == track_cell["time"].values[-1])[0][0] + n_extend + 1, + ) + time_cell = time[slice(i_start, i_end)] for time_i in time_cell: -# for i_row,row in track_cell.iterrows(): -# time_i=row['time'] -# constraint_time = Constraint(time=row['time']) + # for i_row,row in track_cell.iterrows(): + # time_i=row['time'] + # constraint_time = Constraint(time=row['time']) constraint_time = Constraint(time=time_i) - constraint_x = Constraint(projection_x_coordinate = lambda cell: x_min < cell < x_max) - constraint_y = Constraint(projection_y_coordinate = lambda cell: y_min < cell < y_max) + constraint_x = Constraint( + projection_x_coordinate=lambda cell: x_min < cell < x_max + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: y_min < cell < y_max + ) constraint = constraint_time & constraint_x & constraint_y - mask_total_i=mask_total.extract(constraint) + mask_total_i = mask_total.extract(constraint) if field_contour is None: - field_contour_i=None + field_contour_i = None else: - field_contour_i=field_contour.extract(constraint) + field_contour_i = field_contour.extract(constraint) if field_filled is None: - field_filled_i=None + field_filled_i = None else: - field_filled_i=field_filled.extract(constraint) - - - track_i=track[track['time']==time_i] - - cells_mask=list(unique(mask_total_i.core_data())) - track_cells=track_i.loc[(track_i['projection_x_coordinate'] > x_min) & (track_i['projection_x_coordinate'] < x_max) & (track_i['projection_y_coordinate'] > y_min) & (track_i['projection_y_coordinate'] < y_max)] - cells_track=list(track_cells['cell'].values) - cells=list(set( cells_mask + cells_track )) + field_filled_i = field_filled.extract(constraint) + + track_i = track[track["time"] == time_i] + + cells_mask = list(unique(mask_total_i.core_data())) + track_cells = track_i.loc[ + (track_i["projection_x_coordinate"] > x_min) + & (track_i["projection_x_coordinate"] < x_max) + & (track_i["projection_y_coordinate"] > y_min) + & (track_i["projection_y_coordinate"] < y_max) + ] + cells_track = list(track_cells["cell"].values) + cells = list(set(cells_mask + cells_track)) if cell not in cells: cells.append(cell) if 0 in cells: cells.remove(0) - track_i=track_i[track_i['cell'].isin(cells)] + track_i = track_i[track_i["cell"].isin(cells)] if cog is None: - cog_i=None + cog_i = None else: - cog_i=cog[cog['cell'].isin(cells)] - cog_i=cog_i[cog_i['time']==time_i] + cog_i = cog[cog["cell"].isin(cells)] + cog_i = cog_i[cog_i["time"] == time_i] if features is None: - features_i=None + features_i = None else: - features_i=features[features['time']==time_i] + features_i = features[features["time"] == time_i] - - fig1=plt.figure(figsize=(20 / 2.54, 10 / 2.54)) - fig1.subplots_adjust(left=0.1, bottom=0.15, right=0.9, top=0.9,wspace=0.3, hspace=0.25) + fig1 = plt.figure(figsize=(20 / 2.54, 10 / 2.54)) + fig1.subplots_adjust( + left=0.1, bottom=0.15, right=0.9, top=0.9, wspace=0.3, hspace=0.25 + ) # make two subplots for figure: - gs1 = gridspec.GridSpec(1, 2,width_ratios=[1,1.2]) + gs1 = gridspec.GridSpec(1, 2, width_ratios=[1, 1.2]) fig1.add_subplot(gs1[0]) - fig1.add_subplot(gs1[1], projection='3d') + fig1.add_subplot(gs1[1], projection="3d") ax1 = fig1.get_axes() - - datestring_stamp = time_i.strftime('%Y-%m-%d %H:%M:%S') - if time_i in track_cell['time'].values: - time_cell_i=track_cell[track_cell['time'].values==time_i]['time_cell'] - celltime_stamp = "%02d:%02d:%02d" % (time_cell_i.dt.total_seconds() // 3600, - (time_cell_i.dt.total_seconds() % 3600) // 60, - time_cell_i.dt.total_seconds() % 60 ) + datestring_stamp = time_i.strftime("%Y-%m-%d %H:%M:%S") + if time_i in track_cell["time"].values: + time_cell_i = track_cell[track_cell["time"].values == time_i]["time_cell"] + celltime_stamp = "%02d:%02d:%02d" % ( + time_cell_i.dt.total_seconds() // 3600, + (time_cell_i.dt.total_seconds() % 3600) // 60, + time_cell_i.dt.total_seconds() % 60, + ) else: - celltime_stamp=' - ' - title=datestring_stamp + ' , ' + celltime_stamp - datestring_file = time_i.strftime('%Y-%m-%d_%H%M%S') - - ax1[0]=plot_mask_cell_individual_static(cell_i=cell, - track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - xlim=[x_min/1000,x_max/1000],ylim=[y_min/1000,y_max/1000], - axes=ax1[0],title=title,**kwargs) - - ax1[1]=plot_mask_cell_individual_3Dstatic(cell_i=cell, - track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - xlim=[x_min/1000,x_max/1000],ylim=[y_min/1000,y_max/1000], - axes=ax1[1],title=title, - ele=ele,azim=azim, - **kwargs) + celltime_stamp = " - " + title = datestring_stamp + " , " + celltime_stamp + datestring_file = time_i.strftime("%Y-%m-%d_%H%M%S") + + ax1[0] = plot_mask_cell_individual_static( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + xlim=[x_min / 1000, x_max / 1000], + ylim=[y_min / 1000, y_max / 1000], + axes=ax1[0], + title=title, + **kwargs + ) + + ax1[1] = plot_mask_cell_individual_3Dstatic( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + xlim=[x_min / 1000, x_max / 1000], + ylim=[y_min / 1000, y_max / 1000], + axes=ax1[1], + title=title, + ele=ele, + azim=azim, + **kwargs + ) out_dir = os.path.join(plotdir, name) os.makedirs(out_dir, exist_ok=True) - if 'png' in file_format: - savepath_png = os.path.join(out_dir, name + '_' + datestring_file + '.png') + if "png" in file_format: + savepath_png = os.path.join(out_dir, name + "_" + datestring_file + ".png") fig1.savefig(savepath_png, dpi=dpi) - logging.debug('Mask static 2d/3D plot saved to ' + savepath_png) - if 'pdf' in file_format: - savepath_pdf = os.path.join(out_dir, name + '_' + datestring_file + '.pdf') + logging.debug("Mask static 2d/3D plot saved to " + savepath_png) + if "pdf" in file_format: + savepath_pdf = os.path.join(out_dir, name + "_" + datestring_file + ".pdf") fig1.savefig(savepath_pdf, dpi=dpi) - logging.debug('Mask static 2d/3D plot saved to ' + savepath_pdf) + logging.debug("Mask static 2d/3D plot saved to " + savepath_pdf) plt.close() plt.clf() - -def plot_mask_cell_track_3Dstatic(cell,track, cog, features, mask_total, - field_contour, field_filled, - width=10000,n_extend=1, - name= 'test', plotdir='./', - file_format=['png'],figsize=(10/2.54, 10/2.54),dpi=300, - **kwargs): - '''Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours +def plot_mask_cell_track_3Dstatic( + cell, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + width=10000, + n_extend=1, + name="test", + plotdir="./", + file_format=["png"], + figsize=(10 / 2.54, 10 / 2.54), + dpi=300, + **kwargs +): + """Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ from iris import Constraint from numpy import unique import os from mpl_toolkits.mplot3d import Axes3D - track_cell=track[track['cell']==cell] - x_min=track_cell['projection_x_coordinate'].min()-width - x_max=track_cell['projection_x_coordinate'].max()+width - y_min=track_cell['projection_y_coordinate'].min()-width - y_max=track_cell['projection_y_coordinate'].max()+width - - #set up looping over time based on mask's time coordinate to allow for one timestep before and after the track - time_coord=mask_total.coord('time') - time=time_coord.units.num2date(time_coord.points) - i_start=max(0,np.where(time==track_cell['time'].values[0])[0][0]-n_extend) - i_end=min(len(time)-1,np.where(time==track_cell['time'].values[-1])[0][0]+n_extend+1) - time_cell=time[slice(i_start,i_end)] + track_cell = track[track["cell"] == cell] + x_min = track_cell["projection_x_coordinate"].min() - width + x_max = track_cell["projection_x_coordinate"].max() + width + y_min = track_cell["projection_y_coordinate"].min() - width + y_max = track_cell["projection_y_coordinate"].max() + width + + # set up looping over time based on mask's time coordinate to allow for one timestep before and after the track + time_coord = mask_total.coord("time") + time = time_coord.units.num2date(time_coord.points) + i_start = max(0, np.where(time == track_cell["time"].values[0])[0][0] - n_extend) + i_end = min( + len(time) - 1, + np.where(time == track_cell["time"].values[-1])[0][0] + n_extend + 1, + ) + time_cell = time[slice(i_start, i_end)] for time_i in time_cell: -# for i_row,row in track_cell.iterrows(): -# time_i=row['time'] -# constraint_time = Constraint(time=row['time']) + # for i_row,row in track_cell.iterrows(): + # time_i=row['time'] + # constraint_time = Constraint(time=row['time']) constraint_time = Constraint(time=time_i) - constraint_x = Constraint(projection_x_coordinate = lambda cell: x_min < cell < x_max) - constraint_y = Constraint(projection_y_coordinate = lambda cell: y_min < cell < y_max) + constraint_x = Constraint( + projection_x_coordinate=lambda cell: x_min < cell < x_max + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: y_min < cell < y_max + ) constraint = constraint_time & constraint_x & constraint_y - mask_total_i=mask_total.extract(constraint) + mask_total_i = mask_total.extract(constraint) if field_contour is None: - field_contour_i=None + field_contour_i = None else: - field_contour_i=field_contour.extract(constraint) + field_contour_i = field_contour.extract(constraint) if field_filled is None: - field_filled_i=None + field_filled_i = None else: - field_filled_i=field_filled.extract(constraint) - - - track_i=track[track['time']==time_i] - - cells_mask=list(unique(mask_total_i.core_data())) - track_cells=track_i.loc[(track_i['projection_x_coordinate'] > x_min) & (track_i['projection_x_coordinate'] < x_max) & (track_i['projection_y_coordinate'] > y_min) & (track_i['projection_y_coordinate'] < y_max)] - cells_track=list(track_cells['cell'].values) - cells=list(set( cells_mask + cells_track )) + field_filled_i = field_filled.extract(constraint) + + track_i = track[track["time"] == time_i] + + cells_mask = list(unique(mask_total_i.core_data())) + track_cells = track_i.loc[ + (track_i["projection_x_coordinate"] > x_min) + & (track_i["projection_x_coordinate"] < x_max) + & (track_i["projection_y_coordinate"] > y_min) + & (track_i["projection_y_coordinate"] < y_max) + ] + cells_track = list(track_cells["cell"].values) + cells = list(set(cells_mask + cells_track)) if cell not in cells: cells.append(cell) if 0 in cells: cells.remove(0) - track_i=track_i[track_i['cell'].isin(cells)] + track_i = track_i[track_i["cell"].isin(cells)] if cog is None: - cog_i=None + cog_i = None else: - cog_i=cog[cog['cell'].isin(cells)] - cog_i=cog_i[cog_i['time']==time_i] + cog_i = cog[cog["cell"].isin(cells)] + cog_i = cog_i[cog_i["time"] == time_i] if features is None: - features_i=None + features_i = None else: - features_i=features[features['time']==time_i] - - -# fig1, ax1 = plt.subplots(ncols=1, nrows=1, figsize=figsize) -# fig1.subplots_adjust(left=0.2, bottom=0.15, right=0.80, top=0.85) - fig1, ax1 = plt.subplots(ncols=1, nrows=1, figsize=(10/2.54, 10/2.54), subplot_kw={'projection': '3d'}) - - - datestring_stamp = time_i.strftime('%Y-%m-%d %H:%M:%S') - if time_i in track_cell['time'].values: - time_cell_i=track_cell[track_cell['time'].values==time_i]['time_cell'] - celltime_stamp = "%02d:%02d:%02d" % (time_cell_i.dt.total_seconds() // 3600, - (time_cell_i.dt.total_seconds() % 3600) // 60, - time_cell_i.dt.total_seconds() % 60 ) + features_i = features[features["time"] == time_i] + + # fig1, ax1 = plt.subplots(ncols=1, nrows=1, figsize=figsize) + # fig1.subplots_adjust(left=0.2, bottom=0.15, right=0.80, top=0.85) + fig1, ax1 = plt.subplots( + ncols=1, + nrows=1, + figsize=(10 / 2.54, 10 / 2.54), + subplot_kw={"projection": "3d"}, + ) + + datestring_stamp = time_i.strftime("%Y-%m-%d %H:%M:%S") + if time_i in track_cell["time"].values: + time_cell_i = track_cell[track_cell["time"].values == time_i]["time_cell"] + celltime_stamp = "%02d:%02d:%02d" % ( + time_cell_i.dt.total_seconds() // 3600, + (time_cell_i.dt.total_seconds() % 3600) // 60, + time_cell_i.dt.total_seconds() % 60, + ) else: - celltime_stamp=' - ' - title=datestring_stamp + ' , ' + celltime_stamp - datestring_file = time_i.strftime('%Y-%m-%d_%H%M%S') - - ax1=plot_mask_cell_individual_3Dstatic(cell_i=cell, - track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - xlim=[x_min/1000,x_max/1000],ylim=[y_min/1000,y_max/1000], - axes=ax1,title=title,**kwargs) + celltime_stamp = " - " + title = datestring_stamp + " , " + celltime_stamp + datestring_file = time_i.strftime("%Y-%m-%d_%H%M%S") + + ax1 = plot_mask_cell_individual_3Dstatic( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + xlim=[x_min / 1000, x_max / 1000], + ylim=[y_min / 1000, y_max / 1000], + axes=ax1, + title=title, + **kwargs + ) out_dir = os.path.join(plotdir, name) os.makedirs(out_dir, exist_ok=True) - if 'png' in file_format: - savepath_png = os.path.join(out_dir, name + '_' + datestring_file + '.png') + if "png" in file_format: + savepath_png = os.path.join(out_dir, name + "_" + datestring_file + ".png") fig1.savefig(savepath_png, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_png) - if 'pdf' in file_format: - savepath_pdf = os.path.join(out_dir, name + '_' + datestring_file + '.pdf') + logging.debug("Mask static plot saved to " + savepath_png) + if "pdf" in file_format: + savepath_pdf = os.path.join(out_dir, name + "_" + datestring_file + ".pdf") fig1.savefig(savepath_pdf, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_pdf) + logging.debug("Mask static plot saved to " + savepath_pdf) plt.close() plt.clf() -def plot_mask_cell_individual_3Dstatic(cell_i,track, cog, features, mask_total, - field_contour, field_filled, - axes=None,xlim=None,ylim=None, - label_field_contour=None, cmap_field_contour='Blues',norm_field_contour=None, - linewidths_contour=0.8,contour_labels=False, - vmin_field_contour=0,vmax_field_contour=50,levels_field_contour=None,nlevels_field_contour=10, - label_field_filled=None,cmap_field_filled='summer',norm_field_filled=None, - vmin_field_filled=0,vmax_field_filled=100,levels_field_filled=None,nlevels_field_filled=10, - title=None,feature_number=False, - ele=10.,azim=210. - ): - '''Make plots for cell in fixed frame and with one background field as filling and one background field as contrours +def plot_mask_cell_individual_3Dstatic( + cell_i, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + axes=None, + xlim=None, + ylim=None, + label_field_contour=None, + cmap_field_contour="Blues", + norm_field_contour=None, + linewidths_contour=0.8, + contour_labels=False, + vmin_field_contour=0, + vmax_field_contour=50, + levels_field_contour=None, + nlevels_field_contour=10, + label_field_filled=None, + cmap_field_filled="summer", + norm_field_filled=None, + vmin_field_filled=0, + vmax_field_filled=100, + levels_field_filled=None, + nlevels_field_filled=10, + title=None, + feature_number=False, + ele=10.0, + azim=210.0, +): + """Make plots for cell in fixed frame and with one background field as filling and one background field as contrours Input: Output: - ''' + """ import numpy as np - from .utils import mask_features,mask_features_surface -# from mpl_toolkits.axes_grid1 import make_axes_locatable -# from matplotlib.colors import Normalize + from .utils import mask_features, mask_features_surface + + # from mpl_toolkits.axes_grid1 import make_axes_locatable + # from matplotlib.colors import Normalize from mpl_toolkits.mplot3d import Axes3D axes.view_init(elev=ele, azim=azim) @@ -936,340 +1319,436 @@ def plot_mask_cell_individual_3Dstatic(cell_i,track, cog, features, mask_total, axes.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) axes.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) # make the grid lines transparent - axes.xaxis._axinfo["grid"]['color'] = (1,1,1,0) - axes.yaxis._axinfo["grid"]['color'] = (1,1,1,0) - axes.zaxis._axinfo["grid"]['color'] = (1,1,1,0) + axes.xaxis._axinfo["grid"]["color"] = (1, 1, 1, 0) + axes.yaxis._axinfo["grid"]["color"] = (1, 1, 1, 0) + axes.zaxis._axinfo["grid"]["color"] = (1, 1, 1, 0) if title is not None: - axes.set_title(title,horizontalalignment='left',loc='left') - -# colors_mask = ['pink','darkred', 'orange', 'darkred', 'red', 'darkorange'] - x = mask_total.coord('projection_x_coordinate').points - y = mask_total.coord('projection_y_coordinate').points - z = mask_total.coord('model_level_number').points - -# z = mask_total.coord('geopotential_height').points - zz, yy, xx = np.meshgrid(z, y, x, indexing='ij') -# z_alt = mask_total.coord('geopotential_height').points - - -# divider = make_axes_locatable(axes) - -# if field_filled is not None: -# if levels_field_filled is None: -# levels_field_filled=np.linspace(vmin_field_filled,vmax_field_filled, 10) -# plot_field_filled = axes.contourf(field_filled.coord('projection_x_coordinate').points/1000, -# field_filled.coord('projection_y_coordinate').points/1000, -# field_filled.data, -# levels=levels_field_filled, norm=norm_field_filled, -# cmap=cmap_field_filled, vmin=vmin_field_filled, vmax=vmax_field_filled) - - -# cax_filled = divider.append_axes("right", size="5%", pad=0.1) - -# norm_filled= Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) -# sm1= plt.cm.ScalarMappable(norm=norm_filled, cmap = plot_field_filled.cmap) -# sm1.set_array([]) - -# cbar_field_filled = plt.colorbar(sm1, orientation='vertical',cax=cax_filled) -# cbar_field_filled.ax.set_ylabel(label_field_filled) -# cbar_field_filled.set_clim(vmin_field_filled, vmax_field_filled) - - -# if field_contour is not None: -# if levels_field_contour is None: -# levels_field_contour=np.linspace(vmin_field_contour, vmax_field_contour, 5) -# plot_field_contour = axes.contour(field_contour.coord('projection_x_coordinate').points/1000, -# field_contour.coord('projection_y_coordinate').points/1000, -# field_contour.data, -# cmap=cmap_field_contour,norm=norm_field_contour, -# levels=levels_field_contour,vmin=vmin_field_contour, vmax=vmax_field_contour, -# linewidths=linewidths_contour) - -# if contour_labels: -# axes.clabel(plot_field_contour, fontsize=10) - -# cax_contour = divider.append_axes("bottom", size="5%", pad=0.1) -# if norm_field_contour: -# vmin_field_contour=None -# vmax_field_contour=None -# norm_contour=norm_field_contour -# else: -# norm_contour= Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) -# -# sm_contour= plt.cm.ScalarMappable(norm=norm_contour, cmap = plot_field_contour.cmap) -# sm_contour.set_array([]) -# -# cbar_field_contour = plt.colorbar(sm_contour, orientation='horizontal',ticks=levels_field_contour,cax=cax_contour) -# cbar_field_contour.ax.set_xlabel(label_field_contour) -# cbar_field_contour.set_clim(vmin_field_contour, vmax_field_contour) -# + axes.set_title(title, horizontalalignment="left", loc="left") + + # colors_mask = ['pink','darkred', 'orange', 'darkred', 'red', 'darkorange'] + x = mask_total.coord("projection_x_coordinate").points + y = mask_total.coord("projection_y_coordinate").points + z = mask_total.coord("model_level_number").points + + # z = mask_total.coord('geopotential_height').points + zz, yy, xx = np.meshgrid(z, y, x, indexing="ij") + # z_alt = mask_total.coord('geopotential_height').points + + # divider = make_axes_locatable(axes) + + # if field_filled is not None: + # if levels_field_filled is None: + # levels_field_filled=np.linspace(vmin_field_filled,vmax_field_filled, 10) + # plot_field_filled = axes.contourf(field_filled.coord('projection_x_coordinate').points/1000, + # field_filled.coord('projection_y_coordinate').points/1000, + # field_filled.data, + # levels=levels_field_filled, norm=norm_field_filled, + # cmap=cmap_field_filled, vmin=vmin_field_filled, vmax=vmax_field_filled) + + # cax_filled = divider.append_axes("right", size="5%", pad=0.1) + + # norm_filled= Normalize(vmin=vmin_field_filled, vmax=vmax_field_filled) + # sm1= plt.cm.ScalarMappable(norm=norm_filled, cmap = plot_field_filled.cmap) + # sm1.set_array([]) + + # cbar_field_filled = plt.colorbar(sm1, orientation='vertical',cax=cax_filled) + # cbar_field_filled.ax.set_ylabel(label_field_filled) + # cbar_field_filled.set_clim(vmin_field_filled, vmax_field_filled) + + # if field_contour is not None: + # if levels_field_contour is None: + # levels_field_contour=np.linspace(vmin_field_contour, vmax_field_contour, 5) + # plot_field_contour = axes.contour(field_contour.coord('projection_x_coordinate').points/1000, + # field_contour.coord('projection_y_coordinate').points/1000, + # field_contour.data, + # cmap=cmap_field_contour,norm=norm_field_contour, + # levels=levels_field_contour,vmin=vmin_field_contour, vmax=vmax_field_contour, + # linewidths=linewidths_contour) + + # if contour_labels: + # axes.clabel(plot_field_contour, fontsize=10) + + # cax_contour = divider.append_axes("bottom", size="5%", pad=0.1) + # if norm_field_contour: + # vmin_field_contour=None + # vmax_field_contour=None + # norm_contour=norm_field_contour + # else: + # norm_contour= Normalize(vmin=vmin_field_contour, vmax=vmax_field_contour) + # + # sm_contour= plt.cm.ScalarMappable(norm=norm_contour, cmap = plot_field_contour.cmap) + # sm_contour.set_array([]) + # + # cbar_field_contour = plt.colorbar(sm_contour, orientation='horizontal',ticks=levels_field_contour,cax=cax_contour) + # cbar_field_contour.ax.set_xlabel(label_field_contour) + # cbar_field_contour.set_clim(vmin_field_contour, vmax_field_contour) + # for i_row, row in track.iterrows(): - cell = row['cell'] - feature = row['feature'] -# logging.debug("cell: "+ str(row['cell'])) -# logging.debug("feature: "+ str(row['feature'])) + cell = row["cell"] + feature = row["feature"] + # logging.debug("cell: "+ str(row['cell'])) + # logging.debug("feature: "+ str(row['feature'])) - if cell==cell_i: - color='darkred' + if cell == cell_i: + color = "darkred" if feature_number: - cell_string=' '+str(int(cell))+' ('+str(int(feature))+')' + cell_string = " " + str(int(cell)) + " (" + str(int(feature)) + ")" else: - cell_string=' '+str(int(cell)) + cell_string = " " + str(int(cell)) elif np.isnan(cell): - color='gray' + color = "gray" if feature_number: - cell_string=' '+'('+str(int(feature))+')' + cell_string = " " + "(" + str(int(feature)) + ")" else: - cell_string=' ' + cell_string = " " else: - color='darkorange' + color = "darkorange" if feature_number: - cell_string=' '+str(int(cell))+' ('+str(int(feature))+')' + cell_string = " " + str(int(cell)) + " (" + str(int(feature)) + ")" else: - cell_string=' '+str(int(cell)) - -# axes.text(row['projection_x_coordinate']/1000, -# row['projection_y_coordinate']/1000, -# 0, -# cell_string,color=color,fontsize=6, clip_on=True) - -# # Plot marker for tracked cell centre as a cross -# axes.plot(row['projection_x_coordinate']/1000, -# row['projection_y_coordinate']/1000, -# 0, -# 'x', color=color,markersize=4) - - - #Create surface projection of mask for the respective cell and plot it in the right color -# z_coord = 'model_level_number' -# if len(mask_total.shape)==3: -# mask_total_i_surface = mask_features_surface(mask_total, feature, masked=False, z_coord=z_coord) -# elif len(mask_total.shape)==2: -# mask_total_i_surface=mask_features(mask_total, feature, masked=False, z_coord=z_coord) -# axes.contour(mask_total_i_surface.coord('projection_x_coordinate').points/1000, -# mask_total_i_surface.coord('projection_y_coordinate').points/1000, -# 0, -# mask_total_i_surface.data, -# levels=[0, feature], colors=color, linestyles=':',linewidth=1) + cell_string = " " + str(int(cell)) + + # axes.text(row['projection_x_coordinate']/1000, + # row['projection_y_coordinate']/1000, + # 0, + # cell_string,color=color,fontsize=6, clip_on=True) + + # # Plot marker for tracked cell centre as a cross + # axes.plot(row['projection_x_coordinate']/1000, + # row['projection_y_coordinate']/1000, + # 0, + # 'x', color=color,markersize=4) + + # Create surface projection of mask for the respective cell and plot it in the right color + # z_coord = 'model_level_number' + # if len(mask_total.shape)==3: + # mask_total_i_surface = mask_features_surface(mask_total, feature, masked=False, z_coord=z_coord) + # elif len(mask_total.shape)==2: + # mask_total_i_surface=mask_features(mask_total, feature, masked=False, z_coord=z_coord) + # axes.contour(mask_total_i_surface.coord('projection_x_coordinate').points/1000, + # mask_total_i_surface.coord('projection_y_coordinate').points/1000, + # 0, + # mask_total_i_surface.data, + # levels=[0, feature], colors=color, linestyles=':',linewidth=1) mask_feature = mask_total.data == feature axes.scatter( -# xx[mask_feature]/1000, yy[mask_feature]/1000, zz[mask_feature]/1000, - xx[mask_feature]/1000, yy[mask_feature]/1000, zz[mask_feature], - c=color, marker=',', - s=5,#60000.0 * TWC_i[Mask_particle], - alpha=0.3, cmap='inferno', label=cell_string,rasterized=True) + # xx[mask_feature]/1000, yy[mask_feature]/1000, zz[mask_feature]/1000, + xx[mask_feature] / 1000, + yy[mask_feature] / 1000, + zz[mask_feature], + c=color, + marker=",", + s=5, # 60000.0 * TWC_i[Mask_particle], + alpha=0.3, + cmap="inferno", + label=cell_string, + rasterized=True, + ) axes.set_xlim(xlim) axes.set_ylim(ylim) axes.set_zlim([0, 100]) -# axes.set_zlim([0, 20]) -# axes.set_zticks([0, 5,10,15, 20]) - axes.set_xlabel('x (km)') - axes.set_ylabel('y (km)') + # axes.set_zlim([0, 20]) + # axes.set_zticks([0, 5,10,15, 20]) + axes.set_xlabel("x (km)") + axes.set_ylabel("y (km)") axes.zaxis.set_rotate_label(False) # disable automatic rotation -# axes.set_zlabel('z (km)', rotation=90) - axes.set_zlabel('model level', rotation=90) + # axes.set_zlabel('z (km)', rotation=90) + axes.set_zlabel("model level", rotation=90) return axes - -def plot_mask_cell_track_static_timeseries(cell,track, cog, features, mask_total, - field_contour, field_filled, - track_variable=None,variable=None,variable_ylabel=None,variable_label=[None],variable_legend=False,variable_color=None, - width=10000,n_extend=1, - name= 'test', plotdir='./', - file_format=['png'],figsize=(20/2.54, 10/2.54),dpi=300, - **kwargs): - '''Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours +def plot_mask_cell_track_static_timeseries( + cell, + track, + cog, + features, + mask_total, + field_contour, + field_filled, + track_variable=None, + variable=None, + variable_ylabel=None, + variable_label=[None], + variable_legend=False, + variable_color=None, + width=10000, + n_extend=1, + name="test", + plotdir="./", + file_format=["png"], + figsize=(20 / 2.54, 10 / 2.54), + dpi=300, + **kwargs +): + """Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours Input: Output: - ''' - '''Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours + """ + """Make plots for all cells with fixed frame including entire development of the cell and with one background field as filling and one background field as contrours Input: Output: - ''' + """ from iris import Constraint from numpy import unique import os import pandas as pd - track_cell=track[track['cell']==cell] - x_min=track_cell['projection_x_coordinate'].min()-width - x_max=track_cell['projection_x_coordinate'].max()+width - y_min=track_cell['projection_y_coordinate'].min()-width - y_max=track_cell['projection_y_coordinate'].max()+width - time_min=track_cell['time'].min() -# time_max=track_cell['time'].max() - - track_variable_cell=track_variable[track_variable['cell']==cell] - track_variable_cell['time_cell']=pd.to_timedelta(track_variable_cell['time_cell']) -# track_variable_cell=track_variable_cell[(track_variable_cell['time']>=time_min) & (track_variable_cell['time']<=time_max)] - - #set up looping over time based on mask's time coordinate to allow for one timestep before and after the track - time_coord=mask_total.coord('time') - time=time_coord.units.num2date(time_coord.points) - i_start=max(0,np.where(time==track_cell['time'].values[0])[0][0]-n_extend) - i_end=min(len(time)-1,np.where(time==track_cell['time'].values[-1])[0][0]+n_extend+1) - time_cell=time[slice(i_start,i_end)] + track_cell = track[track["cell"] == cell] + x_min = track_cell["projection_x_coordinate"].min() - width + x_max = track_cell["projection_x_coordinate"].max() + width + y_min = track_cell["projection_y_coordinate"].min() - width + y_max = track_cell["projection_y_coordinate"].max() + width + time_min = track_cell["time"].min() + # time_max=track_cell['time'].max() + + track_variable_cell = track_variable[track_variable["cell"] == cell] + track_variable_cell["time_cell"] = pd.to_timedelta(track_variable_cell["time_cell"]) + # track_variable_cell=track_variable_cell[(track_variable_cell['time']>=time_min) & (track_variable_cell['time']<=time_max)] + + # set up looping over time based on mask's time coordinate to allow for one timestep before and after the track + time_coord = mask_total.coord("time") + time = time_coord.units.num2date(time_coord.points) + i_start = max(0, np.where(time == track_cell["time"].values[0])[0][0] - n_extend) + i_end = min( + len(time) - 1, + np.where(time == track_cell["time"].values[-1])[0][0] + n_extend + 1, + ) + time_cell = time[slice(i_start, i_end)] for time_i in time_cell: constraint_time = Constraint(time=time_i) - constraint_x = Constraint(projection_x_coordinate = lambda cell: x_min < cell < x_max) - constraint_y = Constraint(projection_y_coordinate = lambda cell: y_min < cell < y_max) + constraint_x = Constraint( + projection_x_coordinate=lambda cell: x_min < cell < x_max + ) + constraint_y = Constraint( + projection_y_coordinate=lambda cell: y_min < cell < y_max + ) constraint = constraint_time & constraint_x & constraint_y - mask_total_i=mask_total.extract(constraint) + mask_total_i = mask_total.extract(constraint) if field_contour is None: - field_contour_i=None + field_contour_i = None else: - field_contour_i=field_contour.extract(constraint) + field_contour_i = field_contour.extract(constraint) if field_filled is None: - field_filled_i=None + field_filled_i = None else: - field_filled_i=field_filled.extract(constraint) - - - track_i=track[track['time']==time_i] - cells_mask=list(unique(mask_total_i.core_data())) - track_cells=track_i.loc[(track_i['projection_x_coordinate'] > x_min) & (track_i['projection_x_coordinate'] < x_max) & (track_i['projection_y_coordinate'] > y_min) & (track_i['projection_y_coordinate'] < y_max)] - cells_track=list(track_cells['cell'].values) - cells=list(set( cells_mask + cells_track )) + field_filled_i = field_filled.extract(constraint) + + track_i = track[track["time"] == time_i] + cells_mask = list(unique(mask_total_i.core_data())) + track_cells = track_i.loc[ + (track_i["projection_x_coordinate"] > x_min) + & (track_i["projection_x_coordinate"] < x_max) + & (track_i["projection_y_coordinate"] > y_min) + & (track_i["projection_y_coordinate"] < y_max) + ] + cells_track = list(track_cells["cell"].values) + cells = list(set(cells_mask + cells_track)) if cell not in cells: cells.append(cell) if 0 in cells: cells.remove(0) - track_i=track_i[track_i['cell'].isin(cells)] + track_i = track_i[track_i["cell"].isin(cells)] if cog is None: - cog_i=None + cog_i = None else: - cog_i=cog[cog['cell'].isin(cells)] - cog_i=cog_i[cog_i['time']==time_i] + cog_i = cog[cog["cell"].isin(cells)] + cog_i = cog_i[cog_i["time"] == time_i] if features is None: - features_i=None + features_i = None else: - features_i=features[features['time']==time_i] - + features_i = features[features["time"] == time_i] fig1, ax1 = plt.subplots(ncols=2, nrows=1, figsize=figsize) - fig1.subplots_adjust(left=0.1, bottom=0.15, right=0.90, top=0.85,wspace=0.3) - - datestring_stamp = time_i.strftime('%Y-%m-%d %H:%M:%S') - if time_i in track_cell['time'].values: - time_cell_i=track_cell[track_cell['time'].values==time_i]['time_cell'] - celltime_stamp = "%02d:%02d:%02d" % (time_cell_i.dt.total_seconds() // 3600, - (time_cell_i.dt.total_seconds() % 3600) // 60, - time_cell_i.dt.total_seconds() % 60 ) + fig1.subplots_adjust(left=0.1, bottom=0.15, right=0.90, top=0.85, wspace=0.3) + + datestring_stamp = time_i.strftime("%Y-%m-%d %H:%M:%S") + if time_i in track_cell["time"].values: + time_cell_i = track_cell[track_cell["time"].values == time_i]["time_cell"] + celltime_stamp = "%02d:%02d:%02d" % ( + time_cell_i.dt.total_seconds() // 3600, + (time_cell_i.dt.total_seconds() % 3600) // 60, + time_cell_i.dt.total_seconds() % 60, + ) else: - celltime_stamp=' - ' - title=celltime_stamp + ' , ' + datestring_stamp - datestring_file = time_i.strftime('%Y-%m-%d_%H%M%S') + celltime_stamp = " - " + title = celltime_stamp + " , " + datestring_stamp + datestring_file = time_i.strftime("%Y-%m-%d_%H%M%S") # plot evolving timeseries of variable to second axis: - ax1[0]=plot_mask_cell_individual_static(cell_i=cell, - track=track_i, cog=cog_i,features=features_i, - mask_total=mask_total_i, - field_contour=field_contour_i, field_filled=field_filled_i, - xlim=[x_min/1000,x_max/1000],ylim=[y_min/1000,y_max/1000], - axes=ax1[0],title=title,**kwargs) - - track_variable_past=track_variable_cell[(track_variable_cell['time']>=time_min) & (track_variable_cell['time']<=time_i)] - track_variable_current=track_variable_cell[track_variable_cell['time']==time_i] + ax1[0] = plot_mask_cell_individual_static( + cell_i=cell, + track=track_i, + cog=cog_i, + features=features_i, + mask_total=mask_total_i, + field_contour=field_contour_i, + field_filled=field_filled_i, + xlim=[x_min / 1000, x_max / 1000], + ylim=[y_min / 1000, y_max / 1000], + axes=ax1[0], + title=title, + **kwargs + ) + + track_variable_past = track_variable_cell[ + (track_variable_cell["time"] >= time_min) + & (track_variable_cell["time"] <= time_i) + ] + track_variable_current = track_variable_cell[ + track_variable_cell["time"] == time_i + ] if variable_color is None: - variable_color='navy' + variable_color = "navy" if type(variable) is str: -# logging.debug('variable: '+str(variable)) + # logging.debug('variable: '+str(variable)) if type(variable_color) is str: - variable_color={variable:variable_color} - variable=[variable] - - for i_variable,variable_i in enumerate(variable): - color=variable_color[variable_i] - ax1[1].plot(track_variable_past['time_cell'].dt.total_seconds()/ 60.,track_variable_past[variable_i].values,color=color,linestyle='-',label=variable_label[i_variable]) - ax1[1].plot(track_variable_current['time_cell'].dt.total_seconds()/ 60.,track_variable_current[variable_i].values,color=color,marker='o',markersize=4,fillstyle='full') + variable_color = {variable: variable_color} + variable = [variable] + + for i_variable, variable_i in enumerate(variable): + color = variable_color[variable_i] + ax1[1].plot( + track_variable_past["time_cell"].dt.total_seconds() / 60.0, + track_variable_past[variable_i].values, + color=color, + linestyle="-", + label=variable_label[i_variable], + ) + ax1[1].plot( + track_variable_current["time_cell"].dt.total_seconds() / 60.0, + track_variable_current[variable_i].values, + color=color, + marker="o", + markersize=4, + fillstyle="full", + ) ax1[1].yaxis.tick_right() ax1[1].yaxis.set_label_position("right") - ax1[1].set_xlim([0,2*60]) - ax1[1].set_xticks(np.arange(0,120,15)) - ax1[1].set_ylim([0,max(10,1.25*track_variable_cell[variable].max().max())]) - ax1[1].set_xlabel('cell lifetime (min)') - if variable_ylabel==None: - variable_ylabel=variable + ax1[1].set_xlim([0, 2 * 60]) + ax1[1].set_xticks(np.arange(0, 120, 15)) + ax1[1].set_ylim([0, max(10, 1.25 * track_variable_cell[variable].max().max())]) + ax1[1].set_xlabel("cell lifetime (min)") + if variable_ylabel == None: + variable_ylabel = variable ax1[1].set_ylabel(variable_ylabel) ax1[1].set_title(title) # insert legend, if flag is True if variable_legend: - if (len(variable_label)<5): - ncol=1 + if len(variable_label) < 5: + ncol = 1 else: - ncol=2 - ax1[1].legend(loc='upper right', bbox_to_anchor=(1, 1),ncol=ncol,fontsize=8) + ncol = 2 + ax1[1].legend( + loc="upper right", bbox_to_anchor=(1, 1), ncol=ncol, fontsize=8 + ) out_dir = os.path.join(plotdir, name) os.makedirs(out_dir, exist_ok=True) - if 'png' in file_format: - savepath_png = os.path.join(out_dir, name + '_' + datestring_file + '.png') + if "png" in file_format: + savepath_png = os.path.join(out_dir, name + "_" + datestring_file + ".png") fig1.savefig(savepath_png, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_png) - if 'pdf' in file_format: - savepath_pdf = os.path.join(out_dir, name + '_' + datestring_file + '.pdf') + logging.debug("Mask static plot saved to " + savepath_png) + if "pdf" in file_format: + savepath_pdf = os.path.join(out_dir, name + "_" + datestring_file + ".pdf") fig1.savefig(savepath_pdf, dpi=dpi) - logging.debug('Mask static plot saved to ' + savepath_pdf) + logging.debug("Mask static plot saved to " + savepath_pdf) plt.close() plt.clf() -def map_tracks(track,axis_extent=None,figsize=(10,10),axes=None): - for cell in track['cell'].dropna().unique(): - track_i=track[track['cell']==cell] - axes.plot(track_i['longitude'],track_i['latitude'],'-') + +def map_tracks(track, axis_extent=None, figsize=(10, 10), axes=None): + for cell in track["cell"].dropna().unique(): + track_i = track[track["cell"] == cell] + axes.plot(track_i["longitude"], track_i["latitude"], "-") if axis_extent: axes.set_extent(axis_extent) - axes=make_map(axes) + axes = make_map(axes) return axes + def make_map(axes): import matplotlib.ticker as mticker import cartopy.crs as ccrs from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER - gl = axes.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, - linewidth=2, color='gray', alpha=0.5, linestyle='-') - axes.coastlines('10m') + gl = axes.gridlines( + crs=ccrs.PlateCarree(), + draw_labels=True, + linewidth=2, + color="gray", + alpha=0.5, + linestyle="-", + ) + axes.coastlines("10m") gl.xlabels_top = False gl.ylabels_right = False - gl.xlocator = mticker.MaxNLocator(nbins=5,min_n_ticks=3,steps=None) - gl.ylocator = mticker.MaxNLocator(nbins=5,min_n_ticks=3,steps=None) + gl.xlocator = mticker.MaxNLocator(nbins=5, min_n_ticks=3, steps=None) + gl.ylocator = mticker.MaxNLocator(nbins=5, min_n_ticks=3, steps=None) gl.xformatter = LONGITUDE_FORMATTER gl.yformatter = LATITUDE_FORMATTER - #gl.xlabel_style = {'size': 15, 'color': 'gray'} - #gl.xlabel_style = {'color': 'red', 'weight': 'bold'} + # gl.xlabel_style = {'size': 15, 'color': 'gray'} + # gl.xlabel_style = {'color': 'red', 'weight': 'bold'} return axes -def plot_lifetime_histogram(track,axes=None,bin_edges=np.arange(0,200,20),density=False,**kwargs): - hist, bin_edges,bin_centers = lifetime_histogram(track,bin_edges=bin_edges,density=density) - plot_hist=axes.plot(bin_centers, hist,**kwargs) + +def plot_lifetime_histogram( + track, axes=None, bin_edges=np.arange(0, 200, 20), density=False, **kwargs +): + hist, bin_edges, bin_centers = lifetime_histogram( + track, bin_edges=bin_edges, density=density + ) + plot_hist = axes.plot(bin_centers, hist, **kwargs) return plot_hist -def plot_lifetime_histogram_bar(track,axes=None,bin_edges=np.arange(0,200,20),density=False,width_bar=1,shift=0.5,**kwargs): - hist, bin_edges, bin_centers = lifetime_histogram(track,bin_edges=bin_edges,density=density) - plot_hist=axes.bar(bin_centers+shift,hist,width=width_bar,**kwargs) + +def plot_lifetime_histogram_bar( + track, + axes=None, + bin_edges=np.arange(0, 200, 20), + density=False, + width_bar=1, + shift=0.5, + **kwargs +): + hist, bin_edges, bin_centers = lifetime_histogram( + track, bin_edges=bin_edges, density=density + ) + plot_hist = axes.bar(bin_centers + shift, hist, width=width_bar, **kwargs) return plot_hist -def plot_histogram_cellwise(track,bin_edges,variable,quantity,axes=None,density=False,**kwargs): - hist, bin_edges,bin_centers = histogram_cellwise(track,bin_edges=bin_edges,variable=variable,quantity=quantity,density=density) - plot_hist=axes.plot(bin_centers, hist,**kwargs) + +def plot_histogram_cellwise( + track, bin_edges, variable, quantity, axes=None, density=False, **kwargs +): + hist, bin_edges, bin_centers = histogram_cellwise( + track, + bin_edges=bin_edges, + variable=variable, + quantity=quantity, + density=density, + ) + plot_hist = axes.plot(bin_centers, hist, **kwargs) return plot_hist -def plot_histogram_featurewise(Track,bin_edges,variable,axes=None,density=False,**kwargs): - hist, bin_edges, bin_centers = histogram_featurewise(Track,bin_edges=bin_edges,variable=variable,density=density) - plot_hist=axes.plot(bin_centers, hist,**kwargs) + +def plot_histogram_featurewise( + Track, bin_edges, variable, axes=None, density=False, **kwargs +): + hist, bin_edges, bin_centers = histogram_featurewise( + Track, bin_edges=bin_edges, variable=variable, density=density + ) + plot_hist = axes.plot(bin_centers, hist, **kwargs) return plot_hist diff --git a/tobac/segmentation.py b/tobac/segmentation.py index c09b1c21..2f9d6599 100644 --- a/tobac/segmentation.py +++ b/tobac/segmentation.py @@ -1,22 +1,70 @@ import logging -def segmentation_3D(features,field,dxy,threshold=3e-3,target='maximum',level=None,method='watershed',max_distance=None): - return segmentation(features,field,dxy,threshold=threshold,target=target,level=level,method=method,max_distance=max_distance) -def segmentation_2D(features,field,dxy,threshold=3e-3,target='maximum',level=None,method='watershed',max_distance=None): - return segmentation(features,field,dxy,threshold=threshold,target=target,level=level,method=method,max_distance=max_distance) +def segmentation_3D( + features, + field, + dxy, + threshold=3e-3, + target="maximum", + level=None, + method="watershed", + max_distance=None, +): + return segmentation( + features, + field, + dxy, + threshold=threshold, + target=target, + level=level, + method=method, + max_distance=max_distance, + ) -def segmentation_timestep(field_in,features_in,dxy,threshold=3e-3,target='maximum',level=None,method='watershed',max_distance=None,vertical_coord='auto'): +def segmentation_2D( + features, + field, + dxy, + threshold=3e-3, + target="maximum", + level=None, + method="watershed", + max_distance=None, +): + return segmentation( + features, + field, + dxy, + threshold=threshold, + target=target, + level=level, + method=method, + max_distance=max_distance, + ) + + +def segmentation_timestep( + field_in, + features_in, + dxy, + threshold=3e-3, + target="maximum", + level=None, + method="watershed", + max_distance=None, + vertical_coord="auto", +): """ Function performing watershedding for an individual timestep of the data - + Parameters: - features: pandas.DataFrame + features: pandas.DataFrame features for one specific point in time - field: iris.cube.Cube + field: iris.cube.Cube input field to perform the watershedding on (2D or 3D for one specific point in time) - threshold: float + threshold: float threshold for the watershedding field to be used for the mas target: string switch to determine if algorithm looks strating from maxima or minima in input field (maximum: starting from maxima (default), minimum: starting from minima) @@ -26,7 +74,7 @@ def segmentation_timestep(field_in,features_in,dxy,threshold=3e-3,target='maximu flag determining the algorithm to use (currently watershedding implemented) max_distance: float maximum distance from a marker allowed to be classified as belonging to that cell - + Output: segmentation_out: iris.cube.Cube cloud mask, 0 outside and integer numbers according to track inside the clouds @@ -43,114 +91,136 @@ def segmentation_timestep(field_in,features_in,dxy,threshold=3e-3,target='maximu from copy import deepcopy import numpy as np - # copy feature dataframe for output - features_out=deepcopy(features_in) - # Create cube of the same dimensions and coordinates as input data to store mask: - segmentation_out=1*field_in - segmentation_out.rename('segmentation_mask') - segmentation_out.units=1 - - #Create dask array from input data: - data=field_in.core_data() - - #Set level at which to create "Seed" for each feature in the case of 3D watershedding: + # copy feature dataframe for output + features_out = deepcopy(features_in) + # Create cube of the same dimensions and coordinates as input data to store mask: + segmentation_out = 1 * field_in + segmentation_out.rename("segmentation_mask") + segmentation_out.units = 1 + + # Create dask array from input data: + data = field_in.core_data() + + # Set level at which to create "Seed" for each feature in the case of 3D watershedding: # If none, use all levels (later reduced to the ones fulfilling the theshold conditions) - if level==None: - level=slice(None) + if level == None: + level = slice(None) # transform max_distance in metres to distance in pixels: if max_distance is not None: - max_distance_pixel=np.ceil(max_distance/dxy) + max_distance_pixel = np.ceil(max_distance / dxy) # mask data outside region above/below threshold and invert data if tracking maxima: - if target == 'maximum': - unmasked=data>threshold - data_segmentation=-1*data - elif target == 'minimum': - unmasked=data threshold + data_segmentation = -1 * data + elif target == "minimum": + unmasked = data < threshold + data_segmentation = data else: - raise ValueError('unknown type of target') + raise ValueError("unknown type of target") # set markers at the positions of the features: markers = np.zeros(unmasked.shape).astype(np.int32) - if field_in.ndim==2: #2D watershedding + if field_in.ndim == 2: # 2D watershedding for index, row in features_in.iterrows(): - markers[int(row['hdim_1']), int(row['hdim_2'])]=row['feature'] + markers[int(row["hdim_1"]), int(row["hdim_2"])] = row["feature"] - elif field_in.ndim==3: #3D watershedding - list_coord_names=[coord.name() for coord in field_in.coords()] - #determine vertical axis: - if vertical_coord=='auto': - list_vertical=['z','model_level_number','altitude','geopotential_height'] + elif field_in.ndim == 3: # 3D watershedding + list_coord_names = [coord.name() for coord in field_in.coords()] + # determine vertical axis: + if vertical_coord == "auto": + list_vertical = [ + "z", + "model_level_number", + "altitude", + "geopotential_height", + ] for coord_name in list_vertical: if coord_name in list_coord_names: - vertical_axis=coord_name + vertical_axis = coord_name break elif vertical_coord in list_coord_names: - vertical_axis=vertical_coord + vertical_axis = vertical_coord else: - raise ValueError('Plese specify vertical coordinate') - ndim_vertical=field_in.coord_dims(vertical_axis) - if len(ndim_vertical)>1: - raise ValueError('please specify 1 dimensional vertical coordinate') + raise ValueError("Plese specify vertical coordinate") + ndim_vertical = field_in.coord_dims(vertical_axis) + if len(ndim_vertical) > 1: + raise ValueError("please specify 1 dimensional vertical coordinate") for index, row in features_in.iterrows(): - if ndim_vertical[0]==0: - markers[:,int(row['hdim_1']), int(row['hdim_2'])]=row['feature'] - elif ndim_vertical[0]==1: - markers[int(row['hdim_1']),:, int(row['hdim_2'])]=row['feature'] - elif ndim_vertical[0]==2: - markers[int(row['hdim_1']), int(row['hdim_2']),:]=row['feature'] + if ndim_vertical[0] == 0: + markers[:, int(row["hdim_1"]), int(row["hdim_2"])] = row["feature"] + elif ndim_vertical[0] == 1: + markers[int(row["hdim_1"]), :, int(row["hdim_2"])] = row["feature"] + elif ndim_vertical[0] == 2: + markers[int(row["hdim_1"]), int(row["hdim_2"]), :] = row["feature"] else: - raise ValueError('Segmentations routine only possible with 2 or 3 spatial dimensions') + raise ValueError( + "Segmentations routine only possible with 2 or 3 spatial dimensions" + ) # set markers in cells not fulfilling threshold condition to zero: - markers[~unmasked]=0 - + markers[~unmasked] = 0 + # Turn into np arrays (not necessary for markers) as dask arrays don't yet seem to work for watershedding algorithm - data_segmentation=np.array(data_segmentation) - unmasked=np.array(unmasked) + data_segmentation = np.array(data_segmentation) + unmasked = np.array(unmasked) # perform segmentation: - if method=='watershed': - segmentation_mask = watershed(np.array(data_segmentation),markers.astype(np.int32), mask=unmasked) -# elif method=='random_walker': -# segmentation_mask=random_walker(data_segmentation, markers.astype(np.int32), -# beta=130, mode='bf', tol=0.001, copy=True, multichannel=False, return_full_prob=False, spacing=None) - else: - raise ValueError('unknown method, must be watershed') + if method == "watershed": + segmentation_mask = watershed( + np.array(data_segmentation), markers.astype(np.int32), mask=unmasked + ) + # elif method=='random_walker': + # segmentation_mask=random_walker(data_segmentation, markers.astype(np.int32), + # beta=130, mode='bf', tol=0.001, copy=True, multichannel=False, return_full_prob=False, spacing=None) + else: + raise ValueError("unknown method, must be watershed") # remove everything from the individual masks that is more than max_distance_pixel away from the markers if max_distance is not None: - D=distance_transform_edt((markers==0).astype(int)) - segmentation_mask[np.bitwise_and(segmentation_mask>0, D>max_distance_pixel)]=0 + D = distance_transform_edt((markers == 0).astype(int)) + segmentation_mask[ + np.bitwise_and(segmentation_mask > 0, D > max_distance_pixel) + ] = 0 - #Write resulting mask into cube for output - segmentation_out.data=segmentation_mask + # Write resulting mask into cube for output + segmentation_out.data = segmentation_mask # count number of grid cells asoociated to each tracked cell and write that into DataFrame: values, count = np.unique(segmentation_mask, return_counts=True) - counts=dict(zip(values, count)) - ncells=np.zeros(len(features_out)) - for i,(index,row) in enumerate(features_out.iterrows()): - if row['feature'] in counts.keys(): - ncells=counts[row['feature']] - features_out['ncells']=ncells + counts = dict(zip(values, count)) + ncells = np.zeros(len(features_out)) + for i, (index, row) in enumerate(features_out.iterrows()): + if row["feature"] in counts.keys(): + ncells = counts[row["feature"]] + features_out["ncells"] = ncells + + return segmentation_out, features_out - return segmentation_out,features_out -def segmentation(features,field,dxy,threshold=3e-3,target='maximum',level=None,method='watershed',max_distance=None,vertical_coord='auto'): +def segmentation( + features, + field, + dxy, + threshold=3e-3, + target="maximum", + level=None, + method="watershed", + max_distance=None, + vertical_coord="auto", +): """ Function using watershedding or random walker to determine cloud volumes associated with tracked updrafts - + Parameters: - features: pandas.DataFrame + features: pandas.DataFrame output from trackpy/maketrack - field: iris.cube.Cube - containing the field to perform the watershedding on - threshold: float + field: iris.cube.Cube + containing the field to perform the watershedding on + threshold: float threshold for the watershedding field to be used for the mask - + target: string Switch to determine if algorithm looks strating from maxima or minima in input field (maximum: starting from maxima (default), minimum: starting from minima) @@ -158,50 +228,68 @@ def segmentation(features,field,dxy,threshold=3e-3,target='maximum',level=None,m levels at which to seed the cells for the watershedding algorithm method: str ('method') flag determining the algorithm to use (currently watershedding implemented) - + max_distance: float Maximum distance from a marker allowed to be classified as belonging to that cell - + Output: segmentation_out: iris.cube.Cube Cloud mask, 0 outside and integer numbers according to track inside the cloud """ import pandas as pd from iris.cube import CubeList - - logging.info('Start watershedding 3D') - # check input for right dimensions: - if not (field.ndim==3 or field.ndim==4): - raise ValueError('input to segmentation step must be 3D or 4D including a time dimension') - if 'time' not in [coord.name() for coord in field.coords()]: - raise ValueError("input to segmentation step must include a dimension named 'time'") + logging.info("Start watershedding 3D") + + # check input for right dimensions: + if not (field.ndim == 3 or field.ndim == 4): + raise ValueError( + "input to segmentation step must be 3D or 4D including a time dimension" + ) + if "time" not in [coord.name() for coord in field.coords()]: + raise ValueError( + "input to segmentation step must include a dimension named 'time'" + ) # CubeList and list to store individual segmentation masks and feature DataFrames with information about segmentation - segmentation_out_list=CubeList() - features_out_list=[] - - #loop over individual input timesteps for segmentation: - field_time=field.slices_over('time') - for i,field_i in enumerate(field_time): - time_i=field_i.coord('time').units.num2date(field_i.coord('time').points[0]) - features_i=features.loc[features['time']==time_i] - segmentation_out_i,features_out_i=segmentation_timestep(field_i,features_i,dxy,threshold=threshold,target=target,level=level,method=method,max_distance=max_distance,vertical_coord=vertical_coord) + segmentation_out_list = CubeList() + features_out_list = [] + + # loop over individual input timesteps for segmentation: + field_time = field.slices_over("time") + for i, field_i in enumerate(field_time): + time_i = field_i.coord("time").units.num2date(field_i.coord("time").points[0]) + features_i = features.loc[features["time"] == time_i] + segmentation_out_i, features_out_i = segmentation_timestep( + field_i, + features_i, + dxy, + threshold=threshold, + target=target, + level=level, + method=method, + max_distance=max_distance, + vertical_coord=vertical_coord, + ) segmentation_out_list.append(segmentation_out_i) features_out_list.append(features_out_i) - logging.debug('Finished segmentation for '+time_i.strftime('%Y-%m-%d_%H:%M:%S')) - - #Merge output from individual timesteps: - segmentation_out=segmentation_out_list.merge_cube() - features_out=pd.concat(features_out_list) - - logging.debug('Finished segmentation') - return segmentation_out,features_out - -def watershedding_3D(track,field_in,**kwargs): - kwargs.pop('method',None) - return segmentation_3D(track,field_in,method='watershed',**kwargs) - -def watershedding_2D(track,field_in,**kwargs): - kwargs.pop('method',None) - return segmentation_2D(track,field_in,method='watershed',**kwargs) + logging.debug( + "Finished segmentation for " + time_i.strftime("%Y-%m-%d_%H:%M:%S") + ) + + # Merge output from individual timesteps: + segmentation_out = segmentation_out_list.merge_cube() + features_out = pd.concat(features_out_list) + + logging.debug("Finished segmentation") + return segmentation_out, features_out + + +def watershedding_3D(track, field_in, **kwargs): + kwargs.pop("method", None) + return segmentation_3D(track, field_in, method="watershed", **kwargs) + + +def watershedding_2D(track, field_in, **kwargs): + kwargs.pop("method", None) + return segmentation_2D(track, field_in, method="watershed", **kwargs) diff --git a/tobac/testing.py b/tobac/testing.py index 5299af89..aa4ca751 100644 --- a/tobac/testing.py +++ b/tobac/testing.py @@ -2,63 +2,85 @@ import numpy as np from xarray import DataArray -def make_simple_sample_data_2D(data_type='iris'): + +def make_simple_sample_data_2D(data_type="iris"): """ - function creating a simple dataset to use in tests for tobac. + function creating a simple dataset to use in tests for tobac. The grid has a grid spacing of 1km in both horizontal directions and 100 grid cells in x direction and 500 in y direction. - Time resolution is 1 minute and the total length of the dataset is 100 minutes around a abritraty date (2000-01-01 12:00). + Time resolution is 1 minute and the total length of the dataset is 100 minutes around a abritraty date (2000-01-01 12:00). The longitude and latitude coordinates are added as 2D aux coordinates and arbitrary, but in realisitic range. The data contains a single blob travelling on a linear trajectory through the dataset for part of the time. - + :param data_type: 'iris' or 'xarray' to chose the type of dataset to produce :return: sample dataset as an Iris.Cube.cube or xarray.DataArray """ from iris.cube import Cube - from iris.coords import DimCoord,AuxCoord + from iris.coords import DimCoord, AuxCoord - t_0=datetime.datetime(2000,1,1,12,0,0) - - x=np.arange(0,100e3,1000) - y=np.arange(0,50e3,1000) - t=t_0+np.arange(0,100,1)*datetime.timedelta(minutes=1) - xx,yy=np.meshgrid(x,y) - + t_0 = datetime.datetime(2000, 1, 1, 12, 0, 0) + + x = np.arange(0, 100e3, 1000) + y = np.arange(0, 50e3, 1000) + t = t_0 + np.arange(0, 100, 1) * datetime.timedelta(minutes=1) + xx, yy = np.meshgrid(x, y) - t_temp=np.arange(0,60,1) - track1_t=t_0+t_temp*datetime.timedelta(minutes=1) - x_0_1=10e3 - y_0_1=10e3 - track1_x=x_0_1+30*t_temp*60 - track1_y=y_0_1+14*t_temp*60 - track1_magnitude=10*np.ones(track1_x.shape) + t_temp = np.arange(0, 60, 1) + track1_t = t_0 + t_temp * datetime.timedelta(minutes=1) + x_0_1 = 10e3 + y_0_1 = 10e3 + track1_x = x_0_1 + 30 * t_temp * 60 + track1_y = y_0_1 + 14 * t_temp * 60 + track1_magnitude = 10 * np.ones(track1_x.shape) - data=np.zeros((t.shape[0],y.shape[0],x.shape[0])) - for i_t,t_i in enumerate(t): + data = np.zeros((t.shape[0], y.shape[0], x.shape[0])) + for i_t, t_i in enumerate(t): if np.any(t_i in track1_t): - x_i=track1_x[track1_t==t_i] - y_i=track1_y[track1_t==t_i] - mag_i=track1_magnitude[track1_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) - - t_start=datetime.datetime(1970,1,1,0,0) - t_points=(t-t_start).astype("timedelta64[ms]").astype(int) / 1000 - t_coord=DimCoord(t_points,standard_name='time',var_name='time',units='seconds since 1970-01-01 00:00') - x_coord=DimCoord(x,standard_name='projection_x_coordinate',var_name='x',units='m') - y_coord=DimCoord(y,standard_name='projection_y_coordinate',var_name='y',units='m') - lat_coord=AuxCoord(24+1e-5*xx,standard_name='latitude',var_name='latitude',units='degree') - lon_coord=AuxCoord(150+1e-5*yy,standard_name='longitude',var_name='longitude',units='degree') - sample_data=Cube(data,dim_coords_and_dims=[(t_coord, 0),(y_coord, 1),(x_coord, 2)],aux_coords_and_dims=[(lat_coord, (1,2)),(lon_coord, (1,2))],var_name='w',units='m s-1') - - if data_type=='xarray': - sample_data=DataArray.from_iris(sample_data) - + x_i = track1_x[track1_t == t_i] + y_i = track1_y[track1_t == t_i] + mag_i = track1_magnitude[track1_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) + + t_start = datetime.datetime(1970, 1, 1, 0, 0) + t_points = (t - t_start).astype("timedelta64[ms]").astype(int) / 1000 + t_coord = DimCoord( + t_points, + standard_name="time", + var_name="time", + units="seconds since 1970-01-01 00:00", + ) + x_coord = DimCoord( + x, standard_name="projection_x_coordinate", var_name="x", units="m" + ) + y_coord = DimCoord( + y, standard_name="projection_y_coordinate", var_name="y", units="m" + ) + lat_coord = AuxCoord( + 24 + 1e-5 * xx, standard_name="latitude", var_name="latitude", units="degree" + ) + lon_coord = AuxCoord( + 150 + 1e-5 * yy, standard_name="longitude", var_name="longitude", units="degree" + ) + sample_data = Cube( + data, + dim_coords_and_dims=[(t_coord, 0), (y_coord, 1), (x_coord, 2)], + aux_coords_and_dims=[(lat_coord, (1, 2)), (lon_coord, (1, 2))], + var_name="w", + units="m s-1", + ) + + if data_type == "xarray": + sample_data = DataArray.from_iris(sample_data) + return sample_data -def make_sample_data_2D_3blobs(data_type='iris'): +def make_sample_data_2D_3blobs(data_type="iris"): from iris.cube import Cube - from iris.coords import DimCoord,AuxCoord + from iris.coords import DimCoord, AuxCoord + """ function creating a simple dataset to use in tests for tobac. The grid has a grid spacing of 1km in both horizontal directions and 100 grid cells in x direction and 200 in y direction. @@ -71,158 +93,200 @@ def make_sample_data_2D_3blobs(data_type='iris'): """ - t_0=datetime.datetime(2000,1,1,12,0,0) - - x=np.arange(0,100e3,1000) - y=np.arange(0,200e3,1000) - t=t_0+np.arange(0,100,1)*datetime.timedelta(minutes=1) - xx,yy=np.meshgrid(x,y) - - - t_temp=np.arange(0,60,1) - track1_t=t_0+t_temp*datetime.timedelta(minutes=1) - x_0_1=10e3 - y_0_1=10e3 - track1_x=x_0_1+30*t_temp*60 - track1_y=y_0_1+14*t_temp*60 - track1_magnitude=10*np.ones(track1_x.shape) - - t_temp=np.arange(0,30,1) - track2_t=t_0+(t_temp+40)*datetime.timedelta(minutes=1) - x_0_2=20e3 - y_0_2=10e3 - track2_x=x_0_2+24*(t_temp*60)**2/1000 - track2_y=y_0_2+12*t_temp*60 - track2_magnitude=20*np.ones(track2_x.shape) - - - - t_temp=np.arange(0,20,1) - track3_t=t_0+(t_temp+50)*datetime.timedelta(minutes=1) - x_0_3=70e3 - y_0_3=110e3 - track3_x=x_0_3+20*(t_temp*60)**2/1000 - track3_y=y_0_3+20*t_temp*60 - track3_magnitude=15*np.ones(track3_x.shape) - - - data=np.zeros((t.shape[0],y.shape[0],x.shape[0])) - for i_t,t_i in enumerate(t): + t_0 = datetime.datetime(2000, 1, 1, 12, 0, 0) + + x = np.arange(0, 100e3, 1000) + y = np.arange(0, 200e3, 1000) + t = t_0 + np.arange(0, 100, 1) * datetime.timedelta(minutes=1) + xx, yy = np.meshgrid(x, y) + + t_temp = np.arange(0, 60, 1) + track1_t = t_0 + t_temp * datetime.timedelta(minutes=1) + x_0_1 = 10e3 + y_0_1 = 10e3 + track1_x = x_0_1 + 30 * t_temp * 60 + track1_y = y_0_1 + 14 * t_temp * 60 + track1_magnitude = 10 * np.ones(track1_x.shape) + + t_temp = np.arange(0, 30, 1) + track2_t = t_0 + (t_temp + 40) * datetime.timedelta(minutes=1) + x_0_2 = 20e3 + y_0_2 = 10e3 + track2_x = x_0_2 + 24 * (t_temp * 60) ** 2 / 1000 + track2_y = y_0_2 + 12 * t_temp * 60 + track2_magnitude = 20 * np.ones(track2_x.shape) + + t_temp = np.arange(0, 20, 1) + track3_t = t_0 + (t_temp + 50) * datetime.timedelta(minutes=1) + x_0_3 = 70e3 + y_0_3 = 110e3 + track3_x = x_0_3 + 20 * (t_temp * 60) ** 2 / 1000 + track3_y = y_0_3 + 20 * t_temp * 60 + track3_magnitude = 15 * np.ones(track3_x.shape) + + data = np.zeros((t.shape[0], y.shape[0], x.shape[0])) + for i_t, t_i in enumerate(t): if np.any(t_i in track1_t): - x_i=track1_x[track1_t==t_i] - y_i=track1_y[track1_t==t_i] - mag_i=track1_magnitude[track1_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) + x_i = track1_x[track1_t == t_i] + y_i = track1_y[track1_t == t_i] + mag_i = track1_magnitude[track1_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) if np.any(t_i in track2_t): - x_i=track2_x[track2_t==t_i] - y_i=track2_y[track2_t==t_i] - mag_i=track2_magnitude[track2_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) + x_i = track2_x[track2_t == t_i] + y_i = track2_y[track2_t == t_i] + mag_i = track2_magnitude[track2_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) if np.any(t_i in track3_t): - x_i=track3_x[track3_t==t_i] - y_i=track3_y[track3_t==t_i] - mag_i=track3_magnitude[track3_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.)))\ - - t_start=datetime.datetime(1970,1,1,0,0) - t_points=(t-t_start).astype("timedelta64[ms]").astype(int) / 1000 - t_coord=DimCoord(t_points,standard_name='time',var_name='time',units='seconds since 1970-01-01 00:00') - x_coord=DimCoord(x,standard_name='projection_x_coordinate',var_name='x',units='m') - y_coord=DimCoord(y,standard_name='projection_y_coordinate',var_name='y',units='m') - lat_coord=AuxCoord(24+1e-5*xx,standard_name='latitude',var_name='latitude',units='degree') - lon_coord=AuxCoord(150+1e-5*yy,standard_name='longitude',var_name='longitude',units='degree') - sample_data=Cube(data,dim_coords_and_dims=[(t_coord, 0),(y_coord, 1),(x_coord, 2)],aux_coords_and_dims=[(lat_coord, (1,2)),(lon_coord, (1,2))],var_name='w',units='m s-1') - - if data_type=='xarray': - sample_data=DataArray.from_iris(sample_data) - + x_i = track3_x[track3_t == t_i] + y_i = track3_y[track3_t == t_i] + mag_i = track3_magnitude[track3_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) + t_start = datetime.datetime(1970, 1, 1, 0, 0) + t_points = (t - t_start).astype("timedelta64[ms]").astype(int) / 1000 + t_coord = DimCoord( + t_points, + standard_name="time", + var_name="time", + units="seconds since 1970-01-01 00:00", + ) + x_coord = DimCoord( + x, standard_name="projection_x_coordinate", var_name="x", units="m" + ) + y_coord = DimCoord( + y, standard_name="projection_y_coordinate", var_name="y", units="m" + ) + lat_coord = AuxCoord( + 24 + 1e-5 * xx, standard_name="latitude", var_name="latitude", units="degree" + ) + lon_coord = AuxCoord( + 150 + 1e-5 * yy, standard_name="longitude", var_name="longitude", units="degree" + ) + sample_data = Cube( + data, + dim_coords_and_dims=[(t_coord, 0), (y_coord, 1), (x_coord, 2)], + aux_coords_and_dims=[(lat_coord, (1, 2)), (lon_coord, (1, 2))], + var_name="w", + units="m s-1", + ) + + if data_type == "xarray": + sample_data = DataArray.from_iris(sample_data) + return sample_data -def make_sample_data_2D_3blobs_inv(data_type='iris'): +def make_sample_data_2D_3blobs_inv(data_type="iris"): """ - function creating a version of the dataset created in the function make_sample_cube_2D, but with switched coordinate order for the horizontal coordinates + function creating a version of the dataset created in the function make_sample_cube_2D, but with switched coordinate order for the horizontal coordinates for tests to ensure that this does not affect the results - + :param data_type: 'iris' or 'xarray' to chose the type of dataset to produce :return: sample dataset as an Iris.Cube.cube or xarray.DataArray """ from iris.cube import Cube - from iris.coords import DimCoord,AuxCoord - - t_0=datetime.datetime(2000,1,1,12,0,0) - x=np.arange(0,100e3,1000) - y=np.arange(0,200e3,1000) - t=t_0+np.arange(0,100,1)*datetime.timedelta(minutes=1) - yy,xx=np.meshgrid(y,x) - - - t_temp=np.arange(0,60,1) - track1_t=t_0+t_temp*datetime.timedelta(minutes=1) - x_0_1=10e3 - y_0_1=10e3 - track1_x=x_0_1+30*t_temp*60 - track1_y=y_0_1+14*t_temp*60 - track1_magnitude=10*np.ones(track1_x.shape) - - t_temp=np.arange(0,30,1) - track2_t=t_0+(t_temp+40)*datetime.timedelta(minutes=1) - x_0_2=20e3 - y_0_2=10e3 - track2_x=x_0_2+24*(t_temp*60)**2/1000 - track2_y=y_0_2+12*t_temp*60 - track2_magnitude=20*np.ones(track2_x.shape) - - - - t_temp=np.arange(0,20,1) - track3_t=t_0+(t_temp+50)*datetime.timedelta(minutes=1) - x_0_3=70e3 - y_0_3=110e3 - track3_x=x_0_3+20*(t_temp*60)**2/1000 - track3_y=y_0_3+20*t_temp*60 - track3_magnitude=15*np.ones(track3_x.shape) - - - data=np.zeros((t.shape[0],x.shape[0],y.shape[0])) - for i_t,t_i in enumerate(t): + from iris.coords import DimCoord, AuxCoord + + t_0 = datetime.datetime(2000, 1, 1, 12, 0, 0) + x = np.arange(0, 100e3, 1000) + y = np.arange(0, 200e3, 1000) + t = t_0 + np.arange(0, 100, 1) * datetime.timedelta(minutes=1) + yy, xx = np.meshgrid(y, x) + + t_temp = np.arange(0, 60, 1) + track1_t = t_0 + t_temp * datetime.timedelta(minutes=1) + x_0_1 = 10e3 + y_0_1 = 10e3 + track1_x = x_0_1 + 30 * t_temp * 60 + track1_y = y_0_1 + 14 * t_temp * 60 + track1_magnitude = 10 * np.ones(track1_x.shape) + + t_temp = np.arange(0, 30, 1) + track2_t = t_0 + (t_temp + 40) * datetime.timedelta(minutes=1) + x_0_2 = 20e3 + y_0_2 = 10e3 + track2_x = x_0_2 + 24 * (t_temp * 60) ** 2 / 1000 + track2_y = y_0_2 + 12 * t_temp * 60 + track2_magnitude = 20 * np.ones(track2_x.shape) + + t_temp = np.arange(0, 20, 1) + track3_t = t_0 + (t_temp + 50) * datetime.timedelta(minutes=1) + x_0_3 = 70e3 + y_0_3 = 110e3 + track3_x = x_0_3 + 20 * (t_temp * 60) ** 2 / 1000 + track3_y = y_0_3 + 20 * t_temp * 60 + track3_magnitude = 15 * np.ones(track3_x.shape) + + data = np.zeros((t.shape[0], x.shape[0], y.shape[0])) + for i_t, t_i in enumerate(t): if np.any(t_i in track1_t): - x_i=track1_x[track1_t==t_i] - y_i=track1_y[track1_t==t_i] - mag_i=track1_magnitude[track1_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) + x_i = track1_x[track1_t == t_i] + y_i = track1_y[track1_t == t_i] + mag_i = track1_magnitude[track1_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) if np.any(t_i in track2_t): - x_i=track2_x[track2_t==t_i] - y_i=track2_y[track2_t==t_i] - mag_i=track2_magnitude[track2_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) + x_i = track2_x[track2_t == t_i] + y_i = track2_y[track2_t == t_i] + mag_i = track2_magnitude[track2_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) if np.any(t_i in track3_t): - x_i=track3_x[track3_t==t_i] - y_i=track3_y[track3_t==t_i] - mag_i=track3_magnitude[track3_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))*np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.))) - - t_start=datetime.datetime(1970,1,1,0,0) - t_points=(t-t_start).astype("timedelta64[ms]").astype(int) / 1000 - - t_coord=DimCoord(t_points,standard_name='time',var_name='time',units='seconds since 1970-01-01 00:00') - x_coord=DimCoord(x,standard_name='projection_x_coordinate',var_name='x',units='m') - y_coord=DimCoord(y,standard_name='projection_y_coordinate',var_name='y',units='m') - lat_coord=AuxCoord(24+1e-5*xx,standard_name='latitude',var_name='latitude',units='degree') - lon_coord=AuxCoord(150+1e-5*yy,standard_name='longitude',var_name='longitude',units='degree') + x_i = track3_x[track3_t == t_i] + y_i = track3_y[track3_t == t_i] + mag_i = track3_magnitude[track3_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) + + t_start = datetime.datetime(1970, 1, 1, 0, 0) + t_points = (t - t_start).astype("timedelta64[ms]").astype(int) / 1000 + + t_coord = DimCoord( + t_points, + standard_name="time", + var_name="time", + units="seconds since 1970-01-01 00:00", + ) + x_coord = DimCoord( + x, standard_name="projection_x_coordinate", var_name="x", units="m" + ) + y_coord = DimCoord( + y, standard_name="projection_y_coordinate", var_name="y", units="m" + ) + lat_coord = AuxCoord( + 24 + 1e-5 * xx, standard_name="latitude", var_name="latitude", units="degree" + ) + lon_coord = AuxCoord( + 150 + 1e-5 * yy, standard_name="longitude", var_name="longitude", units="degree" + ) + + sample_data = Cube( + data, + dim_coords_and_dims=[(t_coord, 0), (y_coord, 2), (x_coord, 1)], + aux_coords_and_dims=[(lat_coord, (1, 2)), (lon_coord, (1, 2))], + var_name="w", + units="m s-1", + ) + + if data_type == "xarray": + sample_data = DataArray.from_iris(sample_data) - - sample_data=Cube(data,dim_coords_and_dims=[(t_coord, 0),(y_coord, 2),(x_coord, 1)],aux_coords_and_dims=[(lat_coord, (1,2)),(lon_coord, (1,2))],var_name='w',units='m s-1') - - if data_type=='xarray': - sample_data=DataArray.from_iris(sample_data) - return sample_data -def make_sample_data_3D_3blobs(data_type='iris',invert_xy=False): + +def make_sample_data_3D_3blobs(data_type="iris", invert_xy=False): from iris.cube import Cube - from iris.coords import DimCoord,AuxCoord + from iris.coords import DimCoord, AuxCoord + """ function creating a simple dataset to use in tests for tobac. The grid has a grid spacing of 1km in both horizontal directions and 100 grid cells in x direction and 200 in y direction. @@ -235,95 +299,124 @@ def make_sample_data_3D_3blobs(data_type='iris',invert_xy=False): """ - t_0=datetime.datetime(2000,1,1,12,0,0) - - x=np.arange(0,100e3,1000) - y=np.arange(0,200e3,1000) - z=np.arange(0,20e3,1000) - - t=t_0+np.arange(0,50,2)*datetime.timedelta(minutes=1) - - t_temp=np.arange(0,60,1) - track1_t=t_0+t_temp*datetime.timedelta(minutes=1) - x_0_1=10e3 - y_0_1=10e3 - z_0_1=4e3 - track1_x=x_0_1+30*t_temp*60 - track1_y=y_0_1+14*t_temp*60 - track1_magnitude=10*np.ones(track1_x.shape) - - t_temp=np.arange(0,30,1) - track2_t=t_0+(t_temp+40)*datetime.timedelta(minutes=1) - x_0_2=20e3 - y_0_2=10e3 - z_0_2=6e3 - track2_x=x_0_2+24*(t_temp*60)**2/1000 - track2_y=y_0_2+12*t_temp*60 - track2_magnitude=20*np.ones(track2_x.shape) - - - - t_temp=np.arange(0,20,1) - track3_t=t_0+(t_temp+50)*datetime.timedelta(minutes=1) - x_0_3=70e3 - y_0_3=110e3 - z_0_3=8e3 - track3_x=x_0_3+20*(t_temp*60)**2/1000 - track3_y=y_0_3+20*t_temp*60 - track3_magnitude=15*np.ones(track3_x.shape) - - if invert_xy==False: - zz,yy,xx=np.meshgrid(z,y,x,indexing='ij') - y_dim=2 - x_dim=3 - data=np.zeros((t.shape[0],z.shape[0],y.shape[0],x.shape[0])) + t_0 = datetime.datetime(2000, 1, 1, 12, 0, 0) + + x = np.arange(0, 100e3, 1000) + y = np.arange(0, 200e3, 1000) + z = np.arange(0, 20e3, 1000) + + t = t_0 + np.arange(0, 50, 2) * datetime.timedelta(minutes=1) + + t_temp = np.arange(0, 60, 1) + track1_t = t_0 + t_temp * datetime.timedelta(minutes=1) + x_0_1 = 10e3 + y_0_1 = 10e3 + z_0_1 = 4e3 + track1_x = x_0_1 + 30 * t_temp * 60 + track1_y = y_0_1 + 14 * t_temp * 60 + track1_magnitude = 10 * np.ones(track1_x.shape) + + t_temp = np.arange(0, 30, 1) + track2_t = t_0 + (t_temp + 40) * datetime.timedelta(minutes=1) + x_0_2 = 20e3 + y_0_2 = 10e3 + z_0_2 = 6e3 + track2_x = x_0_2 + 24 * (t_temp * 60) ** 2 / 1000 + track2_y = y_0_2 + 12 * t_temp * 60 + track2_magnitude = 20 * np.ones(track2_x.shape) + + t_temp = np.arange(0, 20, 1) + track3_t = t_0 + (t_temp + 50) * datetime.timedelta(minutes=1) + x_0_3 = 70e3 + y_0_3 = 110e3 + z_0_3 = 8e3 + track3_x = x_0_3 + 20 * (t_temp * 60) ** 2 / 1000 + track3_y = y_0_3 + 20 * t_temp * 60 + track3_magnitude = 15 * np.ones(track3_x.shape) + + if invert_xy == False: + zz, yy, xx = np.meshgrid(z, y, x, indexing="ij") + y_dim = 2 + x_dim = 3 + data = np.zeros((t.shape[0], z.shape[0], y.shape[0], x.shape[0])) else: - zz,xx,yy=np.meshgrid(z,x,y,indexing='ij') - x_dim=2 - y_dim=3 - data=np.zeros((t.shape[0],z.shape[0],x.shape[0],y.shape[0])) + zz, xx, yy = np.meshgrid(z, x, y, indexing="ij") + x_dim = 2 + y_dim = 3 + data = np.zeros((t.shape[0], z.shape[0], x.shape[0], y.shape[0])) - - for i_t,t_i in enumerate(t): + for i_t, t_i in enumerate(t): if np.any(t_i in track1_t): - x_i=track1_x[track1_t==t_i] - y_i=track1_y[track1_t==t_i] - z_i=z_0_1 - mag_i=track1_magnitude[track1_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(zz - z_i, 2.) / (2 * np.power(5e3, 2.))) + x_i = track1_x[track1_t == t_i] + y_i = track1_y[track1_t == t_i] + z_i = z_0_1 + mag_i = track1_magnitude[track1_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) * np.exp( + -np.power(zz - z_i, 2.0) / (2 * np.power(5e3, 2.0)) + ) if np.any(t_i in track2_t): - x_i=track2_x[track2_t==t_i] - y_i=track2_y[track2_t==t_i] - z_i=z_0_2 - mag_i=track2_magnitude[track2_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(zz - z_i, 2.) / (2 * np.power(5e3, 2.))) + x_i = track2_x[track2_t == t_i] + y_i = track2_y[track2_t == t_i] + z_i = z_0_2 + mag_i = track2_magnitude[track2_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) * np.exp( + -np.power(zz - z_i, 2.0) / (2 * np.power(5e3, 2.0)) + ) if np.any(t_i in track3_t): - x_i=track3_x[track3_t==t_i] - y_i=track3_y[track3_t==t_i] - z_i=z_0_3 - mag_i=track3_magnitude[track3_t==t_i] - data[i_t]=data[i_t]+mag_i*np.exp(-np.power(xx - x_i,2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(yy - y_i, 2.) / (2 * np.power(10e3, 2.)))\ - *np.exp(-np.power(zz - z_i, 2.) / (2 * np.power(5e3, 2.))) - - - t_start=datetime.datetime(1970,1,1,0,0) - t_points=(t-t_start).astype("timedelta64[ms]").astype(int) / 1000 - t_coord=DimCoord(t_points,standard_name='time',var_name='time',units='seconds since 1970-01-01 00:00') - z_coord=DimCoord(z,standard_name='geopotential_height',var_name='z',units='m') - y_coord=DimCoord(y,standard_name='projection_y_coordinate',var_name='y',units='m') - x_coord=DimCoord(x,standard_name='projection_x_coordinate',var_name='x',units='m') - lat_coord=AuxCoord(24+1e-5*xx[0],standard_name='latitude',var_name='latitude',units='degree') - lon_coord=AuxCoord(150+1e-5*yy[0],standard_name='longitude',var_name='longitude',units='degree') - sample_data=Cube(data,dim_coords_and_dims=[(t_coord, 0),(z_coord, 1),(y_coord, y_dim),(x_coord, x_dim)],aux_coords_and_dims=[(lat_coord, (2,3)),(lon_coord, (2,3))],var_name='w',units='m s-1') - - if data_type=='xarray': - sample_data=DataArray.from_iris(sample_data) - + x_i = track3_x[track3_t == t_i] + y_i = track3_y[track3_t == t_i] + z_i = z_0_3 + mag_i = track3_magnitude[track3_t == t_i] + data[i_t] = data[i_t] + mag_i * np.exp( + -np.power(xx - x_i, 2.0) / (2 * np.power(10e3, 2.0)) + ) * np.exp(-np.power(yy - y_i, 2.0) / (2 * np.power(10e3, 2.0))) * np.exp( + -np.power(zz - z_i, 2.0) / (2 * np.power(5e3, 2.0)) + ) + + t_start = datetime.datetime(1970, 1, 1, 0, 0) + t_points = (t - t_start).astype("timedelta64[ms]").astype(int) / 1000 + t_coord = DimCoord( + t_points, + standard_name="time", + var_name="time", + units="seconds since 1970-01-01 00:00", + ) + z_coord = DimCoord(z, standard_name="geopotential_height", var_name="z", units="m") + y_coord = DimCoord( + y, standard_name="projection_y_coordinate", var_name="y", units="m" + ) + x_coord = DimCoord( + x, standard_name="projection_x_coordinate", var_name="x", units="m" + ) + lat_coord = AuxCoord( + 24 + 1e-5 * xx[0], standard_name="latitude", var_name="latitude", units="degree" + ) + lon_coord = AuxCoord( + 150 + 1e-5 * yy[0], + standard_name="longitude", + var_name="longitude", + units="degree", + ) + sample_data = Cube( + data, + dim_coords_and_dims=[ + (t_coord, 0), + (z_coord, 1), + (y_coord, y_dim), + (x_coord, x_dim), + ], + aux_coords_and_dims=[(lat_coord, (2, 3)), (lon_coord, (2, 3))], + var_name="w", + units="m s-1", + ) + + if data_type == "xarray": + sample_data = DataArray.from_iris(sample_data) + return sample_data diff --git a/tobac/tests/test_import.py b/tobac/tests/test_import.py index 92e44ec2..5d5b7c68 100644 --- a/tobac/tests/test_import.py +++ b/tobac/tests/test_import.py @@ -1,5 +1,6 @@ import pytest import tobac + def test_dummy_function(): - assert 1==1 + assert 1 == 1 diff --git a/tobac/tests/test_sample_data.py b/tobac/tests/test_sample_data.py index e43f3225..a94e7c8a 100644 --- a/tobac/tests/test_sample_data.py +++ b/tobac/tests/test_sample_data.py @@ -1,63 +1,89 @@ """ Tests for tobac based on simple sample datasets with moving blobs. These tests should be adapted to be more modular in the future. """ -from tobac.testing import make_sample_data_2D_3blobs, make_sample_data_2D_3blobs_inv, make_sample_data_3D_3blobs -from tobac import feature_detection_multithreshold,linking_trackpy,get_spacings,segmentation_2D, segmentation_3D -from iris.analysis import MEAN,MAX,MIN +from tobac.testing import ( + make_sample_data_2D_3blobs, + make_sample_data_2D_3blobs_inv, + make_sample_data_3D_3blobs, +) +from tobac import ( + feature_detection_multithreshold, + linking_trackpy, + get_spacings, + segmentation_2D, + segmentation_3D, +) +from iris.analysis import MEAN, MAX, MIN from pandas.testing import assert_frame_equal from numpy.testing import assert_allclose import pandas as pd + def test_sample_data(): """ Test to make sure that sample datasets in the following tests are set up the right way """ - sample_data=make_sample_data_2D_3blobs() - sample_data_inv=make_sample_data_2D_3blobs_inv() - - assert sample_data.coord('projection_x_coordinate')==sample_data_inv.coord('projection_x_coordinate') - assert sample_data.coord('projection_y_coordinate')==sample_data_inv.coord('projection_y_coordinate') - assert sample_data.coord('time')==sample_data_inv.coord('time') - minimum=sample_data.collapsed(('time','projection_x_coordinate','projection_y_coordinate'),MIN).data - minimum_inv=sample_data_inv.collapsed(('time','projection_x_coordinate','projection_y_coordinate'),MIN).data - assert_allclose(minimum,minimum_inv) - mean=sample_data.collapsed(('time','projection_x_coordinate','projection_y_coordinate'),MEAN).data - mean_inv=sample_data_inv.collapsed(('time','projection_x_coordinate','projection_y_coordinate'),MEAN).data - assert_allclose(mean,mean_inv) + sample_data = make_sample_data_2D_3blobs() + sample_data_inv = make_sample_data_2D_3blobs_inv() + + assert sample_data.coord("projection_x_coordinate") == sample_data_inv.coord( + "projection_x_coordinate" + ) + assert sample_data.coord("projection_y_coordinate") == sample_data_inv.coord( + "projection_y_coordinate" + ) + assert sample_data.coord("time") == sample_data_inv.coord("time") + minimum = sample_data.collapsed( + ("time", "projection_x_coordinate", "projection_y_coordinate"), MIN + ).data + minimum_inv = sample_data_inv.collapsed( + ("time", "projection_x_coordinate", "projection_y_coordinate"), MIN + ).data + assert_allclose(minimum, minimum_inv) + mean = sample_data.collapsed( + ("time", "projection_x_coordinate", "projection_y_coordinate"), MEAN + ).data + mean_inv = sample_data_inv.collapsed( + ("time", "projection_x_coordinate", "projection_y_coordinate"), MEAN + ).data + assert_allclose(mean, mean_inv) + def test_tracking_coord_order(): """ Test a tracking applications to make sure that coordinate order does not lead to different results """ - sample_data=make_sample_data_2D_3blobs() - sample_data_inv=make_sample_data_2D_3blobs_inv() + sample_data = make_sample_data_2D_3blobs() + sample_data_inv = make_sample_data_2D_3blobs_inv() # Keyword arguments for feature detection step: - parameters_features={} - parameters_features['position_threshold']='weighted_diff' - parameters_features['sigma_threshold']=0.5 - parameters_features['min_num']=3 - parameters_features['min_distance']=0 - parameters_features['sigma_threshold']=1 - parameters_features['threshold']=[3,5,10] #m/s - parameters_features['n_erosion_threshold']=0 - parameters_features['n_min_threshold']=3 - - #calculate dxy,dt - dxy,dt=get_spacings(sample_data) - dxy_inv,dt_inv=get_spacings(sample_data_inv) - - #Test that dt and dxy are the same for different order of coordinates - assert_allclose(dxy,dxy_inv) - assert_allclose(dt,dt_inv) - - #Test that dt and dxy are as expected - assert_allclose(dt,60) - assert_allclose(dxy,1000) - - #Find features - Features=feature_detection_multithreshold(sample_data,dxy,**parameters_features) - Features_inv=feature_detection_multithreshold(sample_data_inv,dxy_inv,**parameters_features) - + parameters_features = {} + parameters_features["position_threshold"] = "weighted_diff" + parameters_features["sigma_threshold"] = 0.5 + parameters_features["min_num"] = 3 + parameters_features["min_distance"] = 0 + parameters_features["sigma_threshold"] = 1 + parameters_features["threshold"] = [3, 5, 10] # m/s + parameters_features["n_erosion_threshold"] = 0 + parameters_features["n_min_threshold"] = 3 + + # calculate dxy,dt + dxy, dt = get_spacings(sample_data) + dxy_inv, dt_inv = get_spacings(sample_data_inv) + + # Test that dt and dxy are the same for different order of coordinates + assert_allclose(dxy, dxy_inv) + assert_allclose(dt, dt_inv) + + # Test that dt and dxy are as expected + assert_allclose(dt, 60) + assert_allclose(dxy, 1000) + + # Find features + Features = feature_detection_multithreshold(sample_data, dxy, **parameters_features) + Features_inv = feature_detection_multithreshold( + sample_data_inv, dxy_inv, **parameters_features + ) + # Assert that output of feature detection not empty: assert type(Features) == pd.core.frame.DataFrame assert type(Features_inv) == pd.core.frame.DataFrame @@ -65,93 +91,109 @@ def test_tracking_coord_order(): assert not Features_inv.empty # perform watershedding segmentation - parameters_segmentation={} - parameters_segmentation['target']='maximum' - parameters_segmentation['method']='watershed' - - - segmentation_mask,features_segmentation=segmentation_2D(Features,sample_data,dxy=dxy,**parameters_segmentation) - segmentation_mask_inv,features_segmentation=segmentation_2D(Features_inv,sample_data_inv,dxy=dxy_inv,**parameters_segmentation) - + parameters_segmentation = {} + parameters_segmentation["target"] = "maximum" + parameters_segmentation["method"] = "watershed" + + segmentation_mask, features_segmentation = segmentation_2D( + Features, sample_data, dxy=dxy, **parameters_segmentation + ) + segmentation_mask_inv, features_segmentation = segmentation_2D( + Features_inv, sample_data_inv, dxy=dxy_inv, **parameters_segmentation + ) + # perform trajectory linking - parameters_linking={} - parameters_linking['method_linking']='predict' - parameters_linking['adaptive_stop']=0.2 - parameters_linking['adaptive_step']=0.95 - parameters_linking['extrapolate']=0 - parameters_linking['order']=1 - parameters_linking['subnetwork_size']=100 - parameters_linking['memory']=0 - parameters_linking['time_cell_min']=5*60 - parameters_linking['method_linking']='predict' - parameters_linking['v_max']=100 - parameters_linking['d_min']=2000 - - Track=linking_trackpy(Features,sample_data,dt=dt,dxy=dxy,**parameters_linking) - Track_inv=linking_trackpy(Features_inv,sample_data_inv,dt=dt_inv,dxy=dxy_inv,**parameters_linking) - + parameters_linking = {} + parameters_linking["method_linking"] = "predict" + parameters_linking["adaptive_stop"] = 0.2 + parameters_linking["adaptive_step"] = 0.95 + parameters_linking["extrapolate"] = 0 + parameters_linking["order"] = 1 + parameters_linking["subnetwork_size"] = 100 + parameters_linking["memory"] = 0 + parameters_linking["time_cell_min"] = 5 * 60 + parameters_linking["method_linking"] = "predict" + parameters_linking["v_max"] = 100 + parameters_linking["d_min"] = 2000 + + Track = linking_trackpy(Features, sample_data, dt=dt, dxy=dxy, **parameters_linking) + Track_inv = linking_trackpy( + Features_inv, sample_data_inv, dt=dt_inv, dxy=dxy_inv, **parameters_linking + ) + + def test_tracking_3D(): """ Test a tracking applications to make sure that coordinate order does not lead to different results """ - sample_data=make_sample_data_3D_3blobs() - sample_data_inv=make_sample_data_3D_3blobs(invert_xy=True) + sample_data = make_sample_data_3D_3blobs() + sample_data_inv = make_sample_data_3D_3blobs(invert_xy=True) # Keyword arguments for feature detection step: - parameters_features={} - parameters_features['position_threshold']='weighted_diff' - parameters_features['sigma_threshold']=0.5 - parameters_features['min_num']=3 - parameters_features['min_distance']=0 - parameters_features['sigma_threshold']=1 - parameters_features['threshold']=[3,5,10] #m/s - parameters_features['n_erosion_threshold']=0 - parameters_features['n_min_threshold']=3 - - sample_data_max=sample_data.collapsed('geopotential_height',MAX) - sample_data_max_inv=sample_data.collapsed('geopotential_height',MAX) - - #calculate dxy,dt - dxy,dt=get_spacings(sample_data_max) - dxy_inv,dt_inv=get_spacings(sample_data_max_inv) - - #Test that dt and dxy are the same for different order of coordinates - assert_allclose(dxy,dxy_inv) - assert_allclose(dt,dt_inv) - - #Test that dt and dxy are as expected - assert_allclose(dt,120) - assert_allclose(dxy,1000) - - #Find features - Features=feature_detection_multithreshold(sample_data_max,dxy,**parameters_features) - Features_inv=feature_detection_multithreshold(sample_data_max_inv,dxy_inv,**parameters_features) + parameters_features = {} + parameters_features["position_threshold"] = "weighted_diff" + parameters_features["sigma_threshold"] = 0.5 + parameters_features["min_num"] = 3 + parameters_features["min_distance"] = 0 + parameters_features["sigma_threshold"] = 1 + parameters_features["threshold"] = [3, 5, 10] # m/s + parameters_features["n_erosion_threshold"] = 0 + parameters_features["n_min_threshold"] = 3 + + sample_data_max = sample_data.collapsed("geopotential_height", MAX) + sample_data_max_inv = sample_data.collapsed("geopotential_height", MAX) + + # calculate dxy,dt + dxy, dt = get_spacings(sample_data_max) + dxy_inv, dt_inv = get_spacings(sample_data_max_inv) + + # Test that dt and dxy are the same for different order of coordinates + assert_allclose(dxy, dxy_inv) + assert_allclose(dt, dt_inv) + + # Test that dt and dxy are as expected + assert_allclose(dt, 120) + assert_allclose(dxy, 1000) + + # Find features + Features = feature_detection_multithreshold( + sample_data_max, dxy, **parameters_features + ) + Features_inv = feature_detection_multithreshold( + sample_data_max_inv, dxy_inv, **parameters_features + ) # perform watershedding segmentation - parameters_segmentation={} - parameters_segmentation['target']='maximum' - parameters_segmentation['method']='watershed' + parameters_segmentation = {} + parameters_segmentation["target"] = "maximum" + parameters_segmentation["method"] = "watershed" + + segmentation_mask, features_segmentation = segmentation_3D( + Features, sample_data_max, dxy=dxy, **parameters_segmentation + ) + segmentation_mask_inv, features_segmentation = segmentation_3D( + Features_inv, sample_data_max_inv, dxy=dxy_inv, **parameters_segmentation + ) - segmentation_mask,features_segmentation=segmentation_3D(Features,sample_data_max,dxy=dxy,**parameters_segmentation) - segmentation_mask_inv,features_segmentation=segmentation_3D(Features_inv,sample_data_max_inv,dxy=dxy_inv,**parameters_segmentation) - # perform trajectory linking - parameters_linking={} - parameters_linking['method_linking']='predict' - parameters_linking['adaptive_stop']=0.2 - parameters_linking['adaptive_step']=0.95 - parameters_linking['extrapolate']=0 - parameters_linking['order']=1 - parameters_linking['subnetwork_size']=100 - parameters_linking['memory']=0 - parameters_linking['time_cell_min']=5*60 - parameters_linking['method_linking']='predict' - parameters_linking['v_max']=100 - parameters_linking['d_min']=2000 - - Track=linking_trackpy(Features,sample_data,dt=dt,dxy=dxy,**parameters_linking) - Track_inv=linking_trackpy(Features_inv,sample_data_inv,dt=dt_inv,dxy=dxy_inv,**parameters_linking) + parameters_linking = {} + parameters_linking["method_linking"] = "predict" + parameters_linking["adaptive_stop"] = 0.2 + parameters_linking["adaptive_step"] = 0.95 + parameters_linking["extrapolate"] = 0 + parameters_linking["order"] = 1 + parameters_linking["subnetwork_size"] = 100 + parameters_linking["memory"] = 0 + parameters_linking["time_cell_min"] = 5 * 60 + parameters_linking["method_linking"] = "predict" + parameters_linking["v_max"] = 100 + parameters_linking["d_min"] = 2000 + + Track = linking_trackpy(Features, sample_data, dt=dt, dxy=dxy, **parameters_linking) + Track_inv = linking_trackpy( + Features_inv, sample_data_inv, dt=dt_inv, dxy=dxy_inv, **parameters_linking + ) # Assert that output of feature detection not empty: assert not Track.empty diff --git a/tobac/tracking.py b/tobac/tracking.py index c864e6ef..8302f5f6 100644 --- a/tobac/tracking.py +++ b/tobac/tracking.py @@ -2,20 +2,32 @@ import numpy as np import pandas as pd -def linking_trackpy(features,field_in,dt,dxy, - v_max=None,d_max=None,d_min=None,subnetwork_size=None, - memory=0,stubs=1,time_cell_min=None, - order=1,extrapolate=0, - method_linking='random', - adaptive_step=None,adaptive_stop=None, - cell_number_start=1 - ): + +def linking_trackpy( + features, + field_in, + dt, + dxy, + v_max=None, + d_max=None, + d_min=None, + subnetwork_size=None, + memory=0, + stubs=1, + time_cell_min=None, + order=1, + extrapolate=0, + method_linking="random", + adaptive_step=None, + adaptive_stop=None, + cell_number_start=1, +): """ Function to perform the linking of features in trajectories - + Parameters: - features: pandas.DataFrame - Detected features to be linked + features: pandas.DataFrame + Detected features to be linked v_max: float speed at which features are allowed to move dt: float @@ -25,7 +37,7 @@ def linking_trackpy(features,field_in,dt,dxy, memory int number of output timesteps features allowed to vanish for to be still considered tracked subnetwork_size int - maximim size of subnetwork for linking + maximim size of subnetwork for linking method_detection: str('trackpy' or 'threshold') flag choosing method used for feature detection method_linking: str('predict' or 'random') @@ -34,113 +46,131 @@ def linking_trackpy(features,field_in,dt,dxy, # from trackpy import link_df import trackpy as tp from copy import deepcopy -# from trackpy import filter_stubs -# from .utils import add_coordinates + + # from trackpy import filter_stubs + # from .utils import add_coordinates # calculate search range based on timestep and grid spacing if v_max is not None: - search_range=int(dt*v_max/dxy) - + search_range = int(dt * v_max / dxy) + # calculate search range based on timestep and grid spacing if d_max is not None: - search_range=int(d_max/dxy) - + search_range = int(d_max / dxy) + # calculate search range based on timestep and grid spacing if d_min is not None: - search_range=max(search_range,int(d_min/dxy)) + search_range = max(search_range, int(d_min / dxy)) if time_cell_min: - stubs=np.floor(time_cell_min/dt)+1 - - - logging.debug('stubs: '+ str(stubs)) - - logging.debug('start linking features into trajectories') - - - #If subnetwork size given, set maximum subnet size + stubs = np.floor(time_cell_min / dt) + 1 + + logging.debug("stubs: " + str(stubs)) + + logging.debug("start linking features into trajectories") + + # If subnetwork size given, set maximum subnet size if subnetwork_size is not None: - tp.linking.Linker.MAX_SUB_NET_SIZE=subnetwork_size + tp.linking.Linker.MAX_SUB_NET_SIZE = subnetwork_size # deep copy to preserve features field: - features_linking=deepcopy(features) - - - if method_linking is 'random': -# link features into trajectories: - trajectories_unfiltered = tp.link(features_linking, - search_range=search_range, - memory=memory, - t_column='frame', - pos_columns=['hdim_2','hdim_1'], - adaptive_step=adaptive_step,adaptive_stop=adaptive_stop, - neighbor_strategy='KDTree', link_strategy='auto' - ) - elif method_linking is 'predict': + features_linking = deepcopy(features) + + if method_linking is "random": + # link features into trajectories: + trajectories_unfiltered = tp.link( + features_linking, + search_range=search_range, + memory=memory, + t_column="frame", + pos_columns=["hdim_2", "hdim_1"], + adaptive_step=adaptive_step, + adaptive_stop=adaptive_stop, + neighbor_strategy="KDTree", + link_strategy="auto", + ) + elif method_linking is "predict": pred = tp.predict.NearestVelocityPredict(span=1) - trajectories_unfiltered = pred.link_df(features_linking, search_range=search_range, memory=memory, - pos_columns=['hdim_1','hdim_2'], - t_column='frame', - neighbor_strategy='KDTree', link_strategy='auto', - adaptive_step=adaptive_step,adaptive_stop=adaptive_stop -# copy_features=False, diagnostics=False, -# hash_size=None, box_size=None, verify_integrity=True, -# retain_index=False - ) + trajectories_unfiltered = pred.link_df( + features_linking, + search_range=search_range, + memory=memory, + pos_columns=["hdim_1", "hdim_2"], + t_column="frame", + neighbor_strategy="KDTree", + link_strategy="auto", + adaptive_step=adaptive_step, + adaptive_stop=adaptive_stop + # copy_features=False, diagnostics=False, + # hash_size=None, box_size=None, verify_integrity=True, + # retain_index=False + ) else: - raise ValueError('method_linking unknown') - - + raise ValueError("method_linking unknown") + # Filter trajectories to exclude short trajectories that are likely to be spurious -# trajectories_filtered = filter_stubs(trajectories_unfiltered,threshold=stubs) -# trajectories_filtered=trajectories_filtered.reset_index(drop=True) + # trajectories_filtered = filter_stubs(trajectories_unfiltered,threshold=stubs) + # trajectories_filtered=trajectories_filtered.reset_index(drop=True) # Reset particle numbers from the arbitray numbers at the end of the feature detection and linking to consecutive cell numbers # keep 'particle' for reference to the feature detection step. - trajectories_unfiltered['cell']=None - for i_particle,particle in enumerate(pd.Series.unique(trajectories_unfiltered['particle'])): - cell=int(i_particle+cell_number_start) - trajectories_unfiltered.loc[trajectories_unfiltered['particle']==particle,'cell']=cell - trajectories_unfiltered.drop(columns=['particle'],inplace=True) - - trajectories_bycell=trajectories_unfiltered.groupby('cell') - for cell,trajectories_cell in trajectories_bycell: - logging.debug("cell: "+str(cell)) - logging.debug("feature: "+str(trajectories_cell['feature'].values)) - logging.debug("trajectories_cell.shape[0]: "+ str(trajectories_cell.shape[0])) + trajectories_unfiltered["cell"] = None + for i_particle, particle in enumerate( + pd.Series.unique(trajectories_unfiltered["particle"]) + ): + cell = int(i_particle + cell_number_start) + trajectories_unfiltered.loc[ + trajectories_unfiltered["particle"] == particle, "cell" + ] = cell + trajectories_unfiltered.drop(columns=["particle"], inplace=True) - if trajectories_cell.shape[0] < stubs: - logging.debug("cell" + str(cell)+ " is a stub ("+str(trajectories_cell.shape[0])+ "), setting cell number to Nan..") - trajectories_unfiltered.loc[trajectories_unfiltered['cell']==cell,'cell']=np.nan - - trajectories_filtered=trajectories_unfiltered + trajectories_bycell = trajectories_unfiltered.groupby("cell") + for cell, trajectories_cell in trajectories_bycell: + logging.debug("cell: " + str(cell)) + logging.debug("feature: " + str(trajectories_cell["feature"].values)) + logging.debug("trajectories_cell.shape[0]: " + str(trajectories_cell.shape[0])) + if trajectories_cell.shape[0] < stubs: + logging.debug( + "cell" + + str(cell) + + " is a stub (" + + str(trajectories_cell.shape[0]) + + "), setting cell number to Nan.." + ) + trajectories_unfiltered.loc[ + trajectories_unfiltered["cell"] == cell, "cell" + ] = np.nan - #Interpolate to fill the gaps in the trajectories (left from allowing memory in the linking) - trajectories_filtered_unfilled=deepcopy(trajectories_filtered) + trajectories_filtered = trajectories_unfiltered + # Interpolate to fill the gaps in the trajectories (left from allowing memory in the linking) + trajectories_filtered_unfilled = deepcopy(trajectories_filtered) -# trajectories_filtered_filled=fill_gaps(trajectories_filtered_unfilled,order=order, -# extrapolate=extrapolate,frame_max=field_in.shape[0]-1, -# hdim_1_max=field_in.shape[1],hdim_2_max=field_in.shape[2]) -# add coorinates from input fields to output trajectories (time,dimensions) -# logging.debug('start adding coordinates to trajectories') -# trajectories_filtered_filled=add_coordinates(trajectories_filtered_filled,field_in) -# add time coordinate relative to cell initiation: -# logging.debug('start adding cell time to trajectories') - trajectories_filtered_filled=trajectories_filtered_unfilled - trajectories_final=add_cell_time(trajectories_filtered_filled) + # trajectories_filtered_filled=fill_gaps(trajectories_filtered_unfilled,order=order, + # extrapolate=extrapolate,frame_max=field_in.shape[0]-1, + # hdim_1_max=field_in.shape[1],hdim_2_max=field_in.shape[2]) + # add coorinates from input fields to output trajectories (time,dimensions) + # logging.debug('start adding coordinates to trajectories') + # trajectories_filtered_filled=add_coordinates(trajectories_filtered_filled,field_in) + # add time coordinate relative to cell initiation: + # logging.debug('start adding cell time to trajectories') + trajectories_filtered_filled = trajectories_filtered_unfilled + trajectories_final = add_cell_time(trajectories_filtered_filled) # add coordinate to raw features identified: - logging.debug('start adding coordinates to detected features') - logging.debug('feature linking completed') + logging.debug("start adding coordinates to detected features") + logging.debug("feature linking completed") return trajectories_final -def fill_gaps(t,order=1,extrapolate=0,frame_max=None,hdim_1_max=None,hdim_2_max=None): - ''' add cell time as time since the initiation of each cell + +def fill_gaps( + t, order=1, extrapolate=0, frame_max=None, hdim_1_max=None, hdim_2_max=None +): + """add cell time as time since the initiation of each cell Input: - t: pandas dataframe + t: pandas dataframe trajectories from trackpy order: int Order of polynomial used to extrapolate trajectory into gaps and beyond start and end point @@ -153,71 +183,77 @@ def fill_gaps(t,order=1,extrapolate=0,frame_max=None,hdim_1_max=None,hdim_2_max= hdim_2_max: int size of input data along second horizontal axis Output: - t: pandas dataframe + t: pandas dataframe trajectories from trackpy with with filled gaps and potentially extrapolated - ''' + """ from scipy.interpolate import InterpolatedUnivariateSpline - logging.debug('start filling gaps') - - t_list=[] # empty list to store interpolated DataFrames - + + logging.debug("start filling gaps") + + t_list = [] # empty list to store interpolated DataFrames + # group by cell number and perform process for each cell individually: - t_grouped=t.groupby('cell') - for cell,track in t_grouped: - - # Setup interpolator from existing points (of order given as keyword) - frame_in=track['frame'].values - hdim_1_in=track['hdim_1'].values - hdim_2_in=track['hdim_2'].values + t_grouped = t.groupby("cell") + for cell, track in t_grouped: + + # Setup interpolator from existing points (of order given as keyword) + frame_in = track["frame"].values + hdim_1_in = track["hdim_1"].values + hdim_2_in = track["hdim_2"].values s_x = InterpolatedUnivariateSpline(frame_in, hdim_1_in, k=order) - s_y = InterpolatedUnivariateSpline(frame_in, hdim_2_in, k=order) - + s_y = InterpolatedUnivariateSpline(frame_in, hdim_2_in, k=order) + # Create new index filling in gaps and possibly extrapolating: - index_min=min(frame_in)-extrapolate - index_min=max(index_min,0) - index_max=max(frame_in)+extrapolate - index_max=min(index_max,frame_max) - new_index=range(index_min,index_max+1) # +1 here to include last value - track=track.reindex(new_index) - + index_min = min(frame_in) - extrapolate + index_min = max(index_min, 0) + index_max = max(frame_in) + extrapolate + index_max = min(index_max, frame_max) + new_index = range(index_min, index_max + 1) # +1 here to include last value + track = track.reindex(new_index) + # Interpolate to extended index: - frame_out=new_index - hdim_1_out=s_x(frame_out) - hdim_2_out=s_y(frame_out) - - # Replace fields in data frame with - track['frame']=new_index - track['hdim_1']=hdim_1_out - track['hdim_2']=hdim_2_out - track['cell']=cell - + frame_out = new_index + hdim_1_out = s_x(frame_out) + hdim_2_out = s_y(frame_out) + + # Replace fields in data frame with + track["frame"] = new_index + track["hdim_1"] = hdim_1_out + track["hdim_2"] = hdim_2_out + track["cell"] = cell + # Append DataFrame to list of DataFrames - t_list.append(track) - # Concatenate interpolated trajectories into one DataFrame: - t_out=pd.concat(t_list) + t_list.append(track) + # Concatenate interpolated trajectories into one DataFrame: + t_out = pd.concat(t_list) # Restrict output trajectories to input data in time and space: - t_out=t_out.loc[(t_out['hdim_1']0) & (t_out['hdim_2']>0)] - t_out=t_out.reset_index(drop=True) + t_out = t_out.loc[ + (t_out["hdim_1"] < hdim_1_max) + & (t_out["hdim_2"] < hdim_2_max) + & (t_out["hdim_1"] > 0) + & (t_out["hdim_2"] > 0) + ] + t_out = t_out.reset_index(drop=True) return t_out + def add_cell_time(t): - ''' add cell time as time since the initiation of each cell + """add cell time as time since the initiation of each cell Input: t: pandas DataFrame trajectories with added coordinates Output: - t: pandas dataframe + t: pandas dataframe trajectories with added cell time - ''' - - logging.debug('start adding time relative to cell initiation') - t_grouped=t.groupby('cell') - t['time_cell']=np.nan - for cell,track in t_grouped: - track_0=track.head(n=1) - for i,row in track.iterrows(): - t.loc[i,'time_cell']=row['time']-track_0.loc[track_0.index[0],'time'] + """ + + logging.debug("start adding time relative to cell initiation") + t_grouped = t.groupby("cell") + t["time_cell"] = np.nan + for cell, track in t_grouped: + track_0 = track.head(n=1) + for i, row in track.iterrows(): + t.loc[i, "time_cell"] = row["time"] - track_0.loc[track_0.index[0], "time"] # turn series into pandas timedelta DataSeries - t['time_cell']=pd.to_timedelta(t['time_cell']) + t["time_cell"] = pd.to_timedelta(t["time_cell"]) return t - diff --git a/tobac/utils.py b/tobac/utils.py index ab9be599..11503c82 100644 --- a/tobac/utils.py +++ b/tobac/utils.py @@ -1,231 +1,261 @@ import logging -def column_mask_from2D(mask_2D,cube,z_coord='model_level_number'): - ''' function to turn 2D watershedding mask into a 3D mask of selected columns + +def column_mask_from2D(mask_2D, cube, z_coord="model_level_number"): + """function to turn 2D watershedding mask into a 3D mask of selected columns Input: - cube: iris.cube.Cube + cube: iris.cube.Cube data cube - mask_2D: iris.cube.Cube + mask_2D: iris.cube.Cube 2D cube containing mask (int id for tacked volumes 0 everywhere else) z_coord: str name of the vertical coordinate in the cube Output: - mask_2D: iris.cube.Cube + mask_2D: iris.cube.Cube 3D cube containing columns of 2D mask (int id for tacked volumes 0 everywhere else) - ''' + """ from copy import deepcopy - mask_3D=deepcopy(cube) - mask_3D.rename('segmentation_mask') - dim=mask_3D.coord_dims(z_coord)[0] + + mask_3D = deepcopy(cube) + mask_3D.rename("segmentation_mask") + dim = mask_3D.coord_dims(z_coord)[0] for i in range(len(mask_3D.coord(z_coord).points)): slc = [slice(None)] * len(mask_3D.shape) - slc[dim] = slice(i,i+1) - mask_out=mask_3D[slc] - mask_3D.data[slc]=mask_2D.core_data() + slc[dim] = slice(i, i + 1) + mask_out = mask_3D[slc] + mask_3D.data[slc] = mask_2D.core_data() return mask_3D -def mask_cube_cell(variable_cube,mask,cell,track): - ''' Mask cube for tracked volume of an individual cell +def mask_cube_cell(variable_cube, mask, cell, track): + """Mask cube for tracked volume of an individual cell Input: - variable_cube: iris.cube.Cube + variable_cube: iris.cube.Cube unmasked data cube - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) cell: int interger id of cell to create masked cube for Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube with data for respective cell - ''' + """ from copy import deepcopy - variable_cube_out=deepcopy(variable_cube) - feature_ids=track.loc[track['cell']==cell,'feature'].values - variable_cube_out=mask_cube_features(variable_cube,mask,feature_ids) + + variable_cube_out = deepcopy(variable_cube) + feature_ids = track.loc[track["cell"] == cell, "feature"].values + variable_cube_out = mask_cube_features(variable_cube, mask, feature_ids) return variable_cube_out -def mask_cube_all(variable_cube,mask): - ''' Mask cube for untracked volume + +def mask_cube_all(variable_cube, mask): + """Mask cube for untracked volume Input: - variable_cube: iris.cube.Cube + variable_cube: iris.cube.Cube unmasked data cube - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube for untracked volume - ''' + """ from dask.array import ma from copy import deepcopy - variable_cube_out=deepcopy(variable_cube) - variable_cube_out.data=ma.masked_where(mask.core_data()==0,variable_cube_out.core_data()) + + variable_cube_out = deepcopy(variable_cube) + variable_cube_out.data = ma.masked_where( + mask.core_data() == 0, variable_cube_out.core_data() + ) return variable_cube_out -def mask_cube_untracked(variable_cube,mask): - ''' Mask cube for untracked volume + +def mask_cube_untracked(variable_cube, mask): + """Mask cube for untracked volume Input: - variable_cube: iris.cube.Cube + variable_cube: iris.cube.Cube unmasked data cube - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube for untracked volume - ''' + """ from dask.array import ma from copy import deepcopy - variable_cube_out=deepcopy(variable_cube) - variable_cube_out.data=ma.masked_where(mask.core_data()>0,variable_cube_out.core_data()) + + variable_cube_out = deepcopy(variable_cube) + variable_cube_out.data = ma.masked_where( + mask.core_data() > 0, variable_cube_out.core_data() + ) return variable_cube_out -def mask_cube(cube_in,mask): - ''' Mask cube where mask is larger than zero + +def mask_cube(cube_in, mask): + """Mask cube where mask is larger than zero Input: - cube_in: iris.cube.Cube + cube_in: iris.cube.Cube unmasked data cube - mask: numpy.ndarray or dask.array + mask: numpy.ndarray or dask.array mask to use for masking, >0 where cube is supposed to be masked Output: - cube_out: iris.cube.Cube + cube_out: iris.cube.Cube Masked cube - ''' + """ from dask.array import ma from copy import deepcopy - cube_out=deepcopy(cube_in) - cube_out.data=ma.masked_where(mask!=0,cube_in.core_data()) + + cube_out = deepcopy(cube_in) + cube_out.data = ma.masked_where(mask != 0, cube_in.core_data()) return cube_out -def mask_cell(mask,cell,track,masked=False): - ''' create mask for specific cell + +def mask_cell(mask, cell, track, masked=False): + """create mask for specific cell Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: numpy.ndarray + variable_cube_out: numpy.ndarray Masked cube for untracked volume - ''' - feature_ids=track.loc[track['cell']==cell,'feature'].values - mask_i=mask_features(mask,feature_ids,masked=masked) - return mask_i + """ + feature_ids = track.loc[track["cell"] == cell, "feature"].values + mask_i = mask_features(mask, feature_ids, masked=masked) + return mask_i -def mask_cell_surface(mask,cell,track,masked=False,z_coord='model_level_number'): - '''Create surface projection of mask for individual cell + +def mask_cell_surface(mask, cell, track, masked=False, z_coord="model_level_number"): + """Create surface projection of mask for individual cell Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube for untracked volume - ''' - feature_ids=track.loc[track['cell']==cell,'feature'].values - mask_i_surface=mask_features_surface(mask,feature_ids,masked=masked,z_coord=z_coord) + """ + feature_ids = track.loc[track["cell"] == cell, "feature"].values + mask_i_surface = mask_features_surface( + mask, feature_ids, masked=masked, z_coord=z_coord + ) return mask_i_surface -def mask_cell_columns(mask,cell,track,masked=False,z_coord='model_level_number'): - '''Create mask with entire columns for individual cell + +def mask_cell_columns(mask, cell, track, masked=False, z_coord="model_level_number"): + """Create mask with entire columns for individual cell Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube for untracked volume - ''' - feature_ids=track.loc[track['cell']==cell].loc['feature'] - mask_i=mask_features_columns(mask,feature_ids,masked=masked,z_coord=z_coord) + """ + feature_ids = track.loc[track["cell"] == cell].loc["feature"] + mask_i = mask_features_columns(mask, feature_ids, masked=masked, z_coord=z_coord) return mask_i -def mask_cube_features(variable_cube,mask,feature_ids): - ''' Mask cube for tracked volume of an individual cell + +def mask_cube_features(variable_cube, mask, feature_ids): + """Mask cube for tracked volume of an individual cell Input: - variable_cube: iris.cube.Cube + variable_cube: iris.cube.Cube unmasked data cube - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) cell: int interger id of cell to create masked cube for Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube with data for respective cell - ''' - from dask.array import ma,isin + """ + from dask.array import ma, isin from copy import deepcopy - variable_cube_out=deepcopy(variable_cube) - variable_cube_out.data=ma.masked_where(~isin(mask.core_data(),feature_ids),variable_cube_out.core_data()) + + variable_cube_out = deepcopy(variable_cube) + variable_cube_out.data = ma.masked_where( + ~isin(mask.core_data(), feature_ids), variable_cube_out.core_data() + ) return variable_cube_out -def mask_features(mask,feature_ids,masked=False): - ''' create mask for specific features + +def mask_features(mask, feature_ids, masked=False): + """create mask for specific features Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: numpy.ndarray + variable_cube_out: numpy.ndarray Masked cube for untracked volume - ''' - from dask.array import ma,isin + """ + from dask.array import ma, isin from copy import deepcopy - mask_i=deepcopy(mask) - mask_i_data=mask_i.core_data() - mask_i_data[~isin(mask_i.core_data(),feature_ids)]=0 - if masked: - mask_i.data=ma.masked_equal(mask_i.core_data(),0) - return mask_i - -def mask_features_surface(mask,feature_ids,masked=False,z_coord='model_level_number'): - ''' create surface mask for individual features + + mask_i = deepcopy(mask) + mask_i_data = mask_i.core_data() + mask_i_data[~isin(mask_i.core_data(), feature_ids)] = 0 + if masked: + mask_i.data = ma.masked_equal(mask_i.core_data(), 0) + return mask_i + + +def mask_features_surface( + mask, feature_ids, masked=False, z_coord="model_level_number" +): + """create surface mask for individual features Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: - variable_cube_out: iris.cube.Cube + variable_cube_out: iris.cube.Cube Masked cube for untracked volume - ''' + """ from iris.analysis import MAX - from dask.array import ma,isin + from dask.array import ma, isin from copy import deepcopy - mask_i=deepcopy(mask) -# mask_i.data=[~isin(mask_i.data,feature_ids)]=0 - mask_i_data=mask_i.core_data() - mask_i_data[~isin(mask_i.core_data(),feature_ids)]=0 - mask_i_surface=mask_i.collapsed(z_coord,MAX) + + mask_i = deepcopy(mask) + # mask_i.data=[~isin(mask_i.data,feature_ids)]=0 + mask_i_data = mask_i.core_data() + mask_i_data[~isin(mask_i.core_data(), feature_ids)] = 0 + mask_i_surface = mask_i.collapsed(z_coord, MAX) if masked: - mask_i_surface.data=ma.masked_equal(mask_i_surface.core_data(),0) - return mask_i_surface + mask_i_surface.data = ma.masked_equal(mask_i_surface.core_data(), 0) + return mask_i_surface -def mask_all_surface(mask,masked=False,z_coord='model_level_number'): - ''' create surface mask for individual features + +def mask_all_surface(mask, masked=False, z_coord="model_level_number"): + """create surface mask for individual features Input: - mask: iris.cube.Cube + mask: iris.cube.Cube cube containing mask (int id for tacked volumes 0 everywhere else) Output: mask_i_surface: iris.cube.Cube (2D) Mask with 1 below features and 0 everywhere else - ''' + """ from iris.analysis import MAX - from dask.array import ma,isin + from dask.array import ma, isin from copy import deepcopy - mask_i=deepcopy(mask) - mask_i_surface=mask_i.collapsed(z_coord,MAX) - mask_i_surface_data=mask_i_surface.core_data() - mask_i_surface[mask_i_surface_data>0]=1 + + mask_i = deepcopy(mask) + mask_i_surface = mask_i.collapsed(z_coord, MAX) + mask_i_surface_data = mask_i_surface.core_data() + mask_i_surface[mask_i_surface_data > 0] = 1 if masked: - mask_i_surface.data=ma.masked_equal(mask_i_surface.core_data(),0) - return mask_i_surface + mask_i_surface.data = ma.masked_equal(mask_i_surface.core_data(), 0) + return mask_i_surface # def mask_features_columns(mask,feature_ids,masked=False,z_coord='model_level_number'): -# ''' Mask cube for untracked volume +# ''' Mask cube for untracked volume # Input: -# variable_cube: iris.cube.Cube +# variable_cube: iris.cube.Cube # unmasked data cube -# mask: iris.cube.Cube +# mask: iris.cube.Cube # cube containing mask (int id for tacked volumes 0 everywhere else) # Output: -# variable_cube_out: iris.cube.Cube +# variable_cube_out: iris.cube.Cube # Masked cube for untracked volume # ''' # from iris.analysis import MAX -# import numpy as np +# import numpy as np # from copy import deepcopy # mask_i=deepcopy(mask) # mask_i.data[~np.isin(mask_i.data,feature_ids)]=0 @@ -237,27 +267,25 @@ def mask_all_surface(mask,masked=False,z_coord='model_level_number'): # return mask_i - - -#def constraint_cell(track,mask_cell,width=None,x=None,): +# def constraint_cell(track,mask_cell,width=None,x=None,): # from iris import Constraint # import numpy as np -# +# # time_coord=mask_cell.coord('time') # time_units=time_coord.units -# +# # def time_condition(cell): # return time_units.num2date(track.head(n=1)['time']) <= cell <= time_units.num2date(track.tail(n=1)['time']) # # constraint_time=Constraint(time=time_condition) ## mask_cell_i=mask_cell.extract(constraint_time) # mask_cell_surface_i=mask_cell_surface.extract(constraint_time) -# +# # x_dim=mask_cell_surface_i.coord_dims('projection_x_coordinate')[0] # y_dim=mask_cell_surface_i.coord_dims('projection_y_coordinate')[0] # x_coord=mask_cell_surface_i.coord('projection_x_coordinate') # y_coord=mask_cell_surface_i.coord('projection_y_coordinate') -# +# # if (mask_cell_surface_i.core_data()>0).any(): # box_mask_i=get_bounding_box(mask_cell_surface_i.core_data(),buffer=1) # @@ -275,10 +303,12 @@ def mask_all_surface(mask,masked=False,z_coord='model_level_number'): # # constraint=constraint_time & constraint_x & constraint_y # return constraint - -def add_coordinates(t,variable_cube): + + +def add_coordinates(t, variable_cube): import numpy as np - ''' Function adding coordinates from the tracking cube to the trajectories: time, longitude&latitude, x&y dimensions + + """ Function adding coordinates from the tracking cube to the trajectories: time, longitude&latitude, x&y dimensions Input: t: pandas DataFrame trajectories/features @@ -287,110 +317,131 @@ def add_coordinates(t,variable_cube): Output: t: pandas DataFrame trajectories with added coordinated - ''' + """ from scipy.interpolate import interp2d, interp1d - logging.debug('start adding coordinates from cube') + logging.debug("start adding coordinates from cube") - # pull time as datetime object and timestr from input data and add it to DataFrame: - t['time']=None - t['timestr']=None - - - logging.debug('adding time coordinate') + # pull time as datetime object and timestr from input data and add it to DataFrame: + t["time"] = None + t["timestr"] = None - time_in=variable_cube.coord('time') - time_in_datetime=time_in.units.num2date(time_in.points) - - t["time"]=time_in_datetime[t['frame']] - t["timestr"]=[x.strftime('%Y-%m-%d %H:%M:%S') for x in time_in_datetime[t['frame']]] + logging.debug("adding time coordinate") + + time_in = variable_cube.coord("time") + time_in_datetime = time_in.units.num2date(time_in.points) + + t["time"] = time_in_datetime[t["frame"]] + t["timestr"] = [ + x.strftime("%Y-%m-%d %H:%M:%S") for x in time_in_datetime[t["frame"]] + ] # Get list of all coordinates in input cube except for time (already treated): - coord_names=[coord.name() for coord in variable_cube.coords()] - coord_names.remove('time') - - logging.debug('time coordinate added') - - # chose right dimension for horizontal axis based on time dimension: - ndim_time=variable_cube.coord_dims('time')[0] - if ndim_time==0: - hdim_1=1 - hdim_2=2 - elif ndim_time==1: - hdim_1=0 - hdim_2=2 - elif ndim_time==2: - hdim_1=0 - hdim_2=1 - + coord_names = [coord.name() for coord in variable_cube.coords()] + coord_names.remove("time") + + logging.debug("time coordinate added") + + # chose right dimension for horizontal axis based on time dimension: + ndim_time = variable_cube.coord_dims("time")[0] + if ndim_time == 0: + hdim_1 = 1 + hdim_2 = 2 + elif ndim_time == 1: + hdim_1 = 0 + hdim_2 = 2 + elif ndim_time == 2: + hdim_1 = 0 + hdim_2 = 1 + # create vectors to use to interpolate from pixels to coordinates - dimvec_1=np.arange(variable_cube.shape[hdim_1]) - dimvec_2=np.arange(variable_cube.shape[hdim_2]) + dimvec_1 = np.arange(variable_cube.shape[hdim_1]) + dimvec_2 = np.arange(variable_cube.shape[hdim_2]) # loop over coordinates in input data: for coord in coord_names: - logging.debug('adding coord: '+ coord) + logging.debug("adding coord: " + coord) # interpolate 2D coordinates: - if variable_cube.coord(coord).ndim==1: - - if variable_cube.coord_dims(coord)==(hdim_1,): - f=interp1d(dimvec_1,variable_cube.coord(coord).points,fill_value="extrapolate") - coordinate_points=f(t['hdim_1']) - - if variable_cube.coord_dims(coord)==(hdim_2,): - f=interp1d(dimvec_2,variable_cube.coord(coord).points,fill_value="extrapolate") - coordinate_points=f(t['hdim_2']) + if variable_cube.coord(coord).ndim == 1: + + if variable_cube.coord_dims(coord) == (hdim_1,): + f = interp1d( + dimvec_1, + variable_cube.coord(coord).points, + fill_value="extrapolate", + ) + coordinate_points = f(t["hdim_1"]) + + if variable_cube.coord_dims(coord) == (hdim_2,): + f = interp1d( + dimvec_2, + variable_cube.coord(coord).points, + fill_value="extrapolate", + ) + coordinate_points = f(t["hdim_2"]) # interpolate 2D coordinates: - elif variable_cube.coord(coord).ndim==2: + elif variable_cube.coord(coord).ndim == 2: - if variable_cube.coord_dims(coord)==(hdim_1,hdim_2): - f=interp2d(dimvec_2,dimvec_1,variable_cube.coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_2'],t['hdim_1'])] + if variable_cube.coord_dims(coord) == (hdim_1, hdim_2): + f = interp2d(dimvec_2, dimvec_1, variable_cube.coord(coord).points) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_2"], t["hdim_1"])] - if variable_cube.coord_dims(coord)==(hdim_2,hdim_1): - f=interp2d(dimvec_1,dimvec_2,variable_cube.coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_1'],t['hdim_2'])] + if variable_cube.coord_dims(coord) == (hdim_2, hdim_1): + f = interp2d(dimvec_1, dimvec_2, variable_cube.coord(coord).points) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_1"], t["hdim_2"])] - # interpolate 3D coordinates: + # interpolate 3D coordinates: # mainly workaround for wrf latitude and longitude (to be fixed in future) - - elif variable_cube.coord(coord).ndim==3: - - if variable_cube.coord_dims(coord)==(ndim_time,hdim_1,hdim_2): - f=interp2d(dimvec_2,dimvec_1,variable_cube[0,:,:].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_2'],t['hdim_1'])] - - if variable_cube.coord_dims(coord)==(ndim_time,hdim_2,hdim_1): - f=interp2d(dimvec_1,dimvec_2,variable_cube[0,:,:].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_1'],t['hdim_2'])] - - - if variable_cube.coord_dims(coord)==(hdim_1,ndim_time,hdim_2): - f=interp2d(dimvec_2,dimvec_1,variable_cube[:,0,:].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_2'],t['hdim_1'])] - - if variable_cube.coord_dims(coord)==(hdim_1,hdim_2,ndim_time): - f=interp2d(dimvec_2,dimvec_1,variable_cube[:,:,0].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_2'],t['hdim1'])] - - - if variable_cube.coord_dims(coord)==(hdim_2,ndim_time,hdim_1): - f=interp2d(dimvec_1,dimvec_2,variable_cube[:,0,:].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_1'],t['hdim_2'])] - - if variable_cube.coord_dims(coord)==(hdim_2,hdim_1,ndim_time): - f=interp2d(dimvec_1,dimvec_2,variable_cube[:,:,0].coord(coord).points) - coordinate_points=[f(a,b) for a,b in zip(t['hdim_1'],t['hdim_2'])] + + elif variable_cube.coord(coord).ndim == 3: + + if variable_cube.coord_dims(coord) == (ndim_time, hdim_1, hdim_2): + f = interp2d( + dimvec_2, dimvec_1, variable_cube[0, :, :].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_2"], t["hdim_1"])] + + if variable_cube.coord_dims(coord) == (ndim_time, hdim_2, hdim_1): + f = interp2d( + dimvec_1, dimvec_2, variable_cube[0, :, :].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_1"], t["hdim_2"])] + + if variable_cube.coord_dims(coord) == (hdim_1, ndim_time, hdim_2): + f = interp2d( + dimvec_2, dimvec_1, variable_cube[:, 0, :].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_2"], t["hdim_1"])] + + if variable_cube.coord_dims(coord) == (hdim_1, hdim_2, ndim_time): + f = interp2d( + dimvec_2, dimvec_1, variable_cube[:, :, 0].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_2"], t["hdim1"])] + + if variable_cube.coord_dims(coord) == (hdim_2, ndim_time, hdim_1): + f = interp2d( + dimvec_1, dimvec_2, variable_cube[:, 0, :].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_1"], t["hdim_2"])] + + if variable_cube.coord_dims(coord) == (hdim_2, hdim_1, ndim_time): + f = interp2d( + dimvec_1, dimvec_2, variable_cube[:, :, 0].coord(coord).points + ) + coordinate_points = [f(a, b) for a, b in zip(t["hdim_1"], t["hdim_2"])] # write resulting array or list into DataFrame: - t[coord]=coordinate_points + t[coord] = coordinate_points - logging.debug('added coord: '+ coord) + logging.debug("added coord: " + coord) return t -def get_bounding_box(x,buffer=1): - from numpy import delete,arange,diff,nonzero,array + +def get_bounding_box(x, buffer=1): + from numpy import delete, arange, diff, nonzero, array + """ Calculates the bounding box of a ndarray https://stackoverflow.com/questions/31400769/bounding-box-of-numpy-array """ @@ -398,7 +449,7 @@ def get_bounding_box(x,buffer=1): bbox = [] all_axis = arange(x.ndim) - #loop over dimensions + # loop over dimensions for kdim in all_axis: nk_dim = delete(all_axis, kdim) mask_i = mask.all(axis=tuple(nk_dim)) @@ -406,45 +457,55 @@ def get_bounding_box(x,buffer=1): idx_i = nonzero(dmask_i)[0] # for case where there is no value in idx_i if len(idx_i) == 0: - idx_i=array([0,x.shape[kdim]-1]) + idx_i = array([0, x.shape[kdim] - 1]) # for case where there is only one value in idx_i elif len(idx_i) == 1: - idx_i=array([idx_i,idx_i]) + idx_i = array([idx_i, idx_i]) # make sure there is two values in idx_i elif len(idx_i) > 2: - idx_i=array([idx_i[0],idx_i[-1]]) + idx_i = array([idx_i[0], idx_i[-1]]) # caluclate min and max values for idx_i and append them to list - idx_min=max(0,idx_i[0]+1-buffer) - idx_max=min(x.shape[kdim]-1,idx_i[1]+1+buffer) + idx_min = max(0, idx_i[0] + 1 - buffer) + idx_max = min(x.shape[kdim] - 1, idx_i[1] + 1 + buffer) bbox.append([idx_min, idx_max]) return bbox -def get_spacings(field_in,grid_spacing=None,time_spacing=None): + +def get_spacings(field_in, grid_spacing=None, time_spacing=None): import numpy as np from copy import deepcopy + # set horizontal grid spacing of input data # If cartesian x and y corrdinates are present, use these to determine dxy (vertical grid spacing used to transfer pixel distances to real distances): - coord_names=[coord.name() for coord in field_in.coords()] - - if (('projection_x_coordinate' in coord_names and 'projection_y_coordinate' in coord_names) and (grid_spacing is None)): - x_coord=deepcopy(field_in.coord('projection_x_coordinate')) - x_coord.convert_units('metre') - dx=np.diff(field_in.coord('projection_y_coordinate')[0:2].points)[0] - y_coord=deepcopy(field_in.coord('projection_y_coordinate')) - y_coord.convert_units('metre') - dy=np.diff(field_in.coord('projection_y_coordinate')[0:2].points)[0] - dxy=0.5*(dx+dy) + coord_names = [coord.name() for coord in field_in.coords()] + + if ( + "projection_x_coordinate" in coord_names + and "projection_y_coordinate" in coord_names + ) and (grid_spacing is None): + x_coord = deepcopy(field_in.coord("projection_x_coordinate")) + x_coord.convert_units("metre") + dx = np.diff(field_in.coord("projection_y_coordinate")[0:2].points)[0] + y_coord = deepcopy(field_in.coord("projection_y_coordinate")) + y_coord.convert_units("metre") + dy = np.diff(field_in.coord("projection_y_coordinate")[0:2].points)[0] + dxy = 0.5 * (dx + dy) elif grid_spacing is not None: - dxy=grid_spacing + dxy = grid_spacing else: - ValueError('no information about grid spacing, need either input cube with projection_x_coord and projection_y_coord or keyword argument grid_spacing') - + ValueError( + "no information about grid spacing, need either input cube with projection_x_coord and projection_y_coord or keyword argument grid_spacing" + ) + # set horizontal grid spacing of input data - if (time_spacing is None): + if time_spacing is None: # get time resolution of input data from first to steps of input cube: - time_coord=field_in.coord('time') - dt=(time_coord.units.num2date(time_coord.points[1])-time_coord.units.num2date(time_coord.points[0])).seconds - elif (time_spacing is not None): + time_coord = field_in.coord("time") + dt = ( + time_coord.units.num2date(time_coord.points[1]) + - time_coord.units.num2date(time_coord.points[0]) + ).seconds + elif time_spacing is not None: # use value of time_spacing for dt: - dt=time_spacing - return dxy,dt + dt = time_spacing + return dxy, dt diff --git a/tobac/wrapper.py b/tobac/wrapper.py index fbff488c..0d2e64a8 100644 --- a/tobac/wrapper.py +++ b/tobac/wrapper.py @@ -3,79 +3,92 @@ def tracking_wrapper( - field_in_features, - field_in_segmentation, - time_spacing=None, - grid_spacing=None, - parameters_features=None, - parameters_tracking=None, - parameters_segmentation=None, - ): - + field_in_features, + field_in_segmentation, + time_spacing=None, + grid_spacing=None, + parameters_features=None, + parameters_tracking=None, + parameters_segmentation=None, +): + from .feature_detection import feature_detection_multithreshold from .tracking import linking_trackpy from .segmentation import segmentation_3D, segmentation_2D from .utils import get_spacings - logger = logging.getLogger('trackpy') + logger = logging.getLogger("trackpy") logger.propagate = False logger.setLevel(logging.WARNING) - + ### Prepare Tracking - dxy,dt=get_spacings(field_in_features,grid_spacing=grid_spacing,time_spacing=time_spacing) - + dxy, dt = get_spacings( + field_in_features, grid_spacing=grid_spacing, time_spacing=time_spacing + ) + ### Start Tracking # Feature detection: - - method_detection=parameters_features.pop('method_detection',None) - if method_detection in ["threshold","threshold_multi"]: - features=feature_detection_multithreshold(field_in_features,**parameters_features) - else: - raise ValueError('method_detection unknown, has to be either threshold_multi or threshold') - - method_segmentation=parameters_features.pop('method_segmentation',None) - if method_segmentation == 'watershedding': - if field_in_segmentation.ndim==4: - segmentation_mask,features_segmentation=segmentation_3D(features,field_in_segmentation,**parameters_segmentation) - if field_in_segmentation.ndim==3: - segmentation_mask,features_segmentation=segmentation_2D(features,field_in_segmentation,**parameters_segmentation) - + method_detection = parameters_features.pop("method_detection", None) + if method_detection in ["threshold", "threshold_multi"]: + features = feature_detection_multithreshold( + field_in_features, **parameters_features + ) + else: + raise ValueError( + "method_detection unknown, has to be either threshold_multi or threshold" + ) + + method_segmentation = parameters_features.pop("method_segmentation", None) + + if method_segmentation == "watershedding": + if field_in_segmentation.ndim == 4: + segmentation_mask, features_segmentation = segmentation_3D( + features, field_in_segmentation, **parameters_segmentation + ) + if field_in_segmentation.ndim == 3: + segmentation_mask, features_segmentation = segmentation_2D( + features, field_in_segmentation, **parameters_segmentation + ) # Link the features in the individual frames to trajectories: - method_linking=parameters_features.pop('method_linking',None) + method_linking = parameters_features.pop("method_linking", None) - if method_linking == 'trackpy': - trajectories=linking_trackpy(features,**parameters_tracking) - logging.debug('Finished tracking') + if method_linking == "trackpy": + trajectories = linking_trackpy(features, **parameters_tracking) + logging.debug("Finished tracking") else: - raise ValueError('method_linking unknown, has to be trackpy') - - return features,segmentation_mask,trajectories - - - - -def maketrack(field_in, - grid_spacing=None,time_spacing=None, - target='maximum', - v_max=None,d_max=None, - memory=0,stubs=5, - order=1,extrapolate=0, - method_detection="threshold", - position_threshold='center', - sigma_threshold=0.5, - n_erosion_threshold=0, - threshold=1, min_num=0, - min_distance=0, - method_linking="random", - cell_number_start=1, - subnetwork_size=None, - adaptive_stop=None, - adaptive_step=None, - return_intermediate=False, - ): + raise ValueError("method_linking unknown, has to be trackpy") + + return features, segmentation_mask, trajectories + + +def maketrack( + field_in, + grid_spacing=None, + time_spacing=None, + target="maximum", + v_max=None, + d_max=None, + memory=0, + stubs=5, + order=1, + extrapolate=0, + method_detection="threshold", + position_threshold="center", + sigma_threshold=0.5, + n_erosion_threshold=0, + threshold=1, + min_num=0, + min_distance=0, + method_linking="random", + cell_number_start=1, + subnetwork_size=None, + adaptive_stop=None, + adaptive_step=None, + return_intermediate=False, +): from .feature_detection import feature_detection_multithreshold from .tracking import linking_trackpy @@ -133,74 +146,86 @@ def maketrack(field_in, """ from copy import deepcopy - - logger = logging.getLogger('trackpy') + + logger = logging.getLogger("trackpy") logger.propagate = False logger.setLevel(logging.WARNING) - + ### Prepare Tracking # set horizontal grid spacing of input data # If cartesian x and y corrdinates are present, use these to determine dxy (vertical grid spacing used to transfer pixel distances to real distances): - coord_names=[coord.name() for coord in field_in.coords()] - - if (('projection_x_coordinate' in coord_names and 'projection_y_coordinate' in coord_names) and (grid_spacing is None)): - x_coord=deepcopy(field_in.coord('projection_x_coordinate')) - x_coord.convert_units('metre') - dx=np.diff(field_in.coord('projection_y_coordinate')[0:2].points)[0] - y_coord=deepcopy(field_in.coord('projection_y_coordinate')) - y_coord.convert_units('metre') - dy=np.diff(field_in.coord('projection_y_coordinate')[0:2].points)[0] - dxy=0.5*(dx+dy) + coord_names = [coord.name() for coord in field_in.coords()] + + if ( + "projection_x_coordinate" in coord_names + and "projection_y_coordinate" in coord_names + ) and (grid_spacing is None): + x_coord = deepcopy(field_in.coord("projection_x_coordinate")) + x_coord.convert_units("metre") + dx = np.diff(field_in.coord("projection_y_coordinate")[0:2].points)[0] + y_coord = deepcopy(field_in.coord("projection_y_coordinate")) + y_coord.convert_units("metre") + dy = np.diff(field_in.coord("projection_y_coordinate")[0:2].points)[0] + dxy = 0.5 * (dx + dy) elif grid_spacing is not None: - dxy=grid_spacing + dxy = grid_spacing else: - ValueError('no information about grid spacing, need either input cube with projection_x_coord and projection_y_coord or keyword argument grid_spacing') - + ValueError( + "no information about grid spacing, need either input cube with projection_x_coord and projection_y_coord or keyword argument grid_spacing" + ) + # set horizontal grid spacing of input data - if (time_spacing is None): + if time_spacing is None: # get time resolution of input data from first to steps of input cube: - time_coord=field_in.coord('time') - dt=(time_coord.units.num2date(time_coord.points[1])-time_coord.units.num2date(time_coord.points[0])).seconds - elif (time_spacing is not None): + time_coord = field_in.coord("time") + dt = ( + time_coord.units.num2date(time_coord.points[1]) + - time_coord.units.num2date(time_coord.points[0]) + ).seconds + elif time_spacing is not None: # use value of time_spacing for dt: - dt=time_spacing + dt = time_spacing ### Start Tracking # Feature detection: - if method_detection in ["threshold","threshold_multi"]: - features=feature_detection_multithreshold(field_in=field_in, - threshold=threshold, - dxy=dxy, - target=target, - position_threshold=position_threshold, - sigma_threshold=sigma_threshold, - n_erosion_threshold=n_erosion_threshold) - features_filtered = features.drop(features[features['num'] < min_num].index) + if method_detection in ["threshold", "threshold_multi"]: + features = feature_detection_multithreshold( + field_in=field_in, + threshold=threshold, + dxy=dxy, + target=target, + position_threshold=position_threshold, + sigma_threshold=sigma_threshold, + n_erosion_threshold=n_erosion_threshold, + ) + features_filtered = features.drop(features[features["num"] < min_num].index) else: - raise ValueError('method_detection unknown, has to be either threshold_multi or threshold') - - # Link the features in the individual frames to trajectories: - - trajectories=linking_trackpy(features=features_filtered, - field_in=field_in, - dxy=dxy, - dt=dt, - memory=memory, - subnetwork_size=subnetwork_size, - adaptive_stop=adaptive_stop, - adaptive_step=adaptive_step, - v_max=v_max, - d_max=d_max, - stubs=stubs, - order=order,extrapolate=extrapolate, - method_linking=method_linking, - cell_number_start=1 - ) - - logging.debug('Finished tracking') - - return trajectories,features + raise ValueError( + "method_detection unknown, has to be either threshold_multi or threshold" + ) + # Link the features in the individual frames to trajectories: + trajectories = linking_trackpy( + features=features_filtered, + field_in=field_in, + dxy=dxy, + dt=dt, + memory=memory, + subnetwork_size=subnetwork_size, + adaptive_stop=adaptive_stop, + adaptive_step=adaptive_step, + v_max=v_max, + d_max=d_max, + stubs=stubs, + order=order, + extrapolate=extrapolate, + method_linking=method_linking, + cell_number_start=1, + ) + + logging.debug("Finished tracking") + + return trajectories, features