diff --git a/mlair/data_handler/data_handler_wrf_chem.py b/mlair/data_handler/data_handler_wrf_chem.py index 112f8fba1e37527d83356c4dd6cbab9886c67c49..03346a5c383a308adef71aa9bd24f23a314d62da 100644 --- a/mlair/data_handler/data_handler_wrf_chem.py +++ b/mlair/data_handler/data_handler_wrf_chem.py @@ -5,19 +5,17 @@ import xarray as xr import numpy as np import itertools import matplotlib.pyplot as plt + import dask import inspect -import dask.array as da import os -from mlair.helpers.geofunctions import haversine_dist +from mlair.helpers.geofunctions import haversine_dist, bearing_angle, WindSector from mlair.helpers.helpers import convert2xrda, remove_items -from mlair.helpers import TimeTrackingWrapper - from mlair.data_handler.data_handler_single_station import DataHandlerSingleStation from mlair.data_handler import DefaultDataHandler -from typing import Tuple, Union, List, Dict +from typing import Tuple, Union, Dict import logging import cartopy.crs as ccrs @@ -127,6 +125,13 @@ class BaseWrfChemDataLoader: dist = haversine_dist(lat1=self._data.XLAT, lon1=self._data.XLONG, lat2=lat, lon2=lon) return dist + def get_bearing(self, lat, lon, points_last=True): + bearing = bearing_angle(lat1=lat, lon1=lon, lat2=self._data.XLAT, lon2=self._data.XLONG) + if points_last: + return bearing.transpose(..., 'points') + else: + return bearing + def compute_nearest_icoordinates(self, lat, lon, dim=None): dist = self.get_distances(lat=lat, lon=lon) @@ -158,6 +163,7 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): DEFAULT_ITER_DIM = "points" DEFAULT_WINDOW_DIM = "window" + def __init__(self, coords: Tuple[float_np_xr, float_np_xr], network=DEFAULT_MODEL, target_dim=DEFAULT_TARGET_DIM, target_var=DEFAULT_TARGET_VAR, iter_dim=DEFAULT_ITER_DIM, window_dim=DEFAULT_WINDOW_DIM, @@ -186,9 +192,45 @@ class SingleGridColumnWrfChemDataLoader(BaseWrfChemDataLoader): if external_coords_file is not None: self._apply_external_coordinates() self.rechunk_data(self.rechunk_values) + # self._set_nearest_icoords(dim=[self.logical_x_coord_name, self.logical_y_coord_name]) + # self._set_nearest_coords() + self._set_geoinfos() + print('dummy test') + + def _set_geoinfos(self): + # identify nearest coords self._set_nearest_icoords(dim=[self.logical_x_coord_name, self.logical_y_coord_name]) self._set_nearest_coords() + # get lat and lons + lat = self.get_coordinates()['lat'] + lon = self.get_coordinates()['lon'] + + # calculate distances to points (from points to all grid boxes in lat/lon space) + dist = self.get_distances(lat, lon) + dist.attrs.update(dict(units='km')) + + # calculate bearing angle + bearing = self.get_bearing(lat, lon) + bearing.attrs.update(dict(units='°')) + + # set geoinfo data set + self._geo_infos = xr.Dataset({'dist': dist, 'bearing': bearing}) + + # set sectors + # initialize WindSector() providing helper methods for wind sector calculations + ws = WindSector() + + # expand dataset by new dataarray + self._geo_infos['wind_sectors'] = xr.full_like(bearing, fill_value=np.nan) # xr.DataArray(coords=bearing.coords, dims=bearing.dims) + for i, (k, v) in enumerate(ws.wind_sectore_edges.items()): + self._geo_infos['wind_sectors'] = xr.where(ws.is_in_sector(k, self.geo_infos['bearing']), i, self._geo_infos['wind_sectors']) + # self._geo_infos['wind_sectors'].attrs.update(dict(units='Wind Sector')) + + @property + def geo_infos(self): + return self._geo_infos + def _apply_external_coordinates(self): ds_coords = xr.open_dataset(self.external_coords_file, chunks={'south_north': 36, 'west_east': 40}) data = self._data @@ -348,7 +390,9 @@ class DataHandlerWRF(DefaultDataHandler): if __name__ == '__main__': - def plot_map_proj(data, xlim=None, ylim=None, filename=None, point=None): + def plot_map_proj(data, xlim=None, ylim=None, filename=None, point=None, radius=None, **kwargs): + cmap = kwargs.pop('cmap', plt.cm.Reds) + set_background = kwargs.pop('set_background', False) crs_latlon = ccrs.PlateCarree() if ylim is None: @@ -363,31 +407,76 @@ if __name__ == '__main__': plt.figure() ax = plt.axes(projection=ccrs.PlateCarree()) ax.set_global() - data.squeeze().plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree(), x='XLONG', y='XLAT', - cmap=plt.cm.Reds, ) + pm = data.squeeze().plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree(), x='XLONG', y='XLAT', + cmap=cmap, edgecolors=None, + **kwargs) ax.coastlines(resolution='50m') ##### Idea to draw circle from https://stackoverflow.com/a/58735566 - # circle_points = cartopy.geodesic.Geodesic().circle(lon=lon, lat=lat, radius=radius_in_meters, - # n_samples=n_points, endpoint=False) - # geom = shapely.geometry.Polygon(circle_points) - # ax.add_geometries((geom,), crs=cartopy.crs.PlateCarree(), facecolor='red', edgecolor='none', linewidth=0) #### own try: circle_points = Geodesic().circle(lon=lon, lat=lat, - radius=100000, # radius_in_meters, + radius=radius*1000., # radius_in_meters, ) geom = shapely.geometry.Polygon(circle_points) ax.add_geometries((geom,), crs=ccrs.PlateCarree(), facecolor='none', edgecolor='blue', linewidth=1) # ##### + ax.set_ylim(ylim) ax.set_xlim(xlim) - plt.plot(point[1], point[0], '+k', markersize=7, transform=crs_latlon) + plt.plot(lon, lat, '+k', markersize=7, transform=crs_latlon) + if set_background: + ax.stock_img() plt.tight_layout() plt.savefig(filename) plt.close('all') + + def dummy_plot(data, skipna=None, xlim=None, ylim=None): + plt.figure() + crs_latlon = ccrs.PlateCarree() + ax = plt.axes(projection=ccrs.PlateCarree()) + ax.set_global() + data.sum(dim='points',skipna=skipna).squeeze().plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree(), x='XLONG', y='XLAT', + cmap=plt.cm.Reds, ) + ax.coastlines(resolution='50m') + ax.set_ylim(ylim) + ax.set_xlim(xlim) + plt.tight_layout() + plt.savefig("Example_Multiple_points_test.pdf") + print("plot multiple points") + + + from shapely.geometry import Point, Polygon + import math + + + def sector(center, start_angle, end_angle, radius, steps=200): + def polar_point(origin_point, angle, distance): + return [origin_point.x + math.sin(math.radians(angle)) * distance, + origin_point.y + math.cos(math.radians(angle)) * distance] + + if start_angle > end_angle: + start_angle = start_angle - 360 + else: + pass + step_angle_width = (end_angle - start_angle) / steps + sector_width = (end_angle - start_angle) + segment_vertices = [] + + segment_vertices.append(polar_point(center, 0, 0)) + segment_vertices.append(polar_point(center, start_angle, radius)) + + for z in range(1, steps): + segment_vertices.append((polar_point(center, start_angle + z * step_angle_width, radius))) + segment_vertices.append(polar_point(center, start_angle + sector_width, radius)) + segment_vertices.append(polar_point(center, 0, 0)) + return Polygon(segment_vertices) + + # lat_np = np.array([50.73333, 47.73333]) + # lon_np = np.array([7.1, 8.6]) + lat_np = np.array([50.73333]) lon_np = np.array([7.1]) @@ -428,31 +517,73 @@ if __name__ == '__main__': dist_np = wrf_new.get_distances(lat_np, lon_np) dist_xr = wrf_new.get_distances(lat_xr, lon_xr) dist_xr.attrs.update(dict(units='km')) - dist_xr_set = xr.Dataset({'dist': dist_xr}) - for i, (data, xlim, ylim) in enumerate(((wrf_new._data.T2.isel({'XTIME': 0}), [-42, 66], [23, 80]), - (dist_xr_set.dist, [-42, 66], [23, 80]), - (wrf_new._data.T2.isel({'XTIME': 0}).where( - dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], - [45, 58]), - (dist_xr_set.dist.where( - dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], - [45, 58]), - )): + bearing_np = wrf_new.get_bearing(lat_np, lon_np) + bearing_xr = wrf_new.get_bearing(lat_xr, lon_xr) + bearing_xr.attrs.update(dict(units='°')) + + geo_info = xr.Dataset({'dist': dist_xr, 'bearing': bearing_xr}) + windsect = WindSector() + radius_from_point = 150. # in km + + # dummy_plot(geo_info.bearing.where(dist_xr <= radius_from_point), True, [2, 15], [45, 58]) + + for i, (data, xlim, ylim, kwargs) in enumerate(((wrf_new._data.T2.isel({'XTIME': 0}), [-42, 66], [23, 80], {}), + (wrf_new.geo_infos.dist, [-42, 66], [23, 80], {}), + (wrf_new.geo_infos.bearing, [-42, 66], [23, 80], {}), + (wrf_new.geo_infos.wind_sectors, [-42, 66], [23, 80], { + 'cmap': plt.cm.get_cmap('tab20'), + 'set_background': True, + } + ), + (wrf_new._data.T2.isel({'XTIME': 0}).where( + windsect.is_in_sector('N', wrf_new.geo_infos['bearing'])), + [-42, 66], [23, 80], {}), + (wrf_new._data.T2.isel({'XTIME': 0}).where( + wrf_new.geo_infos.dist.sel({'points': 0}).drop('points') <= radius_from_point), + [2, 15], + [45, 58], {}), + (geo_info.dist.where( + wrf_new.geo_infos.dist.sel({'points': 0}).drop('points') <= radius_from_point), + [2, 15], + [45, 58], {}), + (wrf_new.geo_infos.bearing.where( + dist_xr.sel({'points': 0}).drop('points') <= radius_from_point), + [2, 15], + [45, 58], {}), + (wrf_new.geo_infos.wind_sectors.where( + dist_xr.sel({'points': 0}).drop('points') <= radius_from_point), + [2, 15], + [45, 58], {'cmap': plt.cm.get_cmap('tab20', 16), + 'set_background': True, + }), + (wrf_new.geo_infos.dist.where( + wrf_new.geo_infos.dist.sel({'points': 0}).drop('points') <= radius_from_point).where( + windsect.is_in_sector('N', wrf_new.geo_infos['bearing'])), + [2, 15], + [45, 58], {}), + )): plot_map_proj(data, xlim=xlim, ylim=ylim, - point=[lat_np, lon_np], filename=f'Example_dist{i}.pdf') + # point=[lat_np, lon_np], + point=[wrf_new.get_nearest_coords()['lat'][0], wrf_new.get_nearest_coords()['lon'][0]], + filename=f'Example_dist{i}.pdf', + radius=radius_from_point, + **kwargs, + ) for i, (data, xlim, ylim) in enumerate( ((wrf_new._data.o3.isel({'XTIME': 0, 'bottom_top': 0}), [-42, 66], [23, 80]), - (dist_xr_set.dist, [-42, 66], [23, 80]), + (geo_info.dist, [-42, 66], [23, 80]), (wrf_new._data.o3.isel({'XTIME': 0, 'bottom_top': 0}).where( - dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], [45, 58]), - (dist_xr_set.dist.where(dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], [45, 58]), + dist_xr.sel({'points': 0}).drop('points') <= radius_from_point), [2, 15], [45, 58]), + (geo_info.dist.where(dist_xr.sel({'points': 0}).drop('points') <= radius_from_point), [2, 15], [45, 58]), )): plot_map_proj(data, xlim=xlim, ylim=ylim, - point=[lat_np, lon_np], filename=f'ExampleO3_dist{i}.pdf') + point=[lat_np, lon_np], + filename=f'ExampleO3_dist{i}.pdf', + radius=radius_from_point) ######################### # Larger 4D data use_second_dummy_dataset = False @@ -468,14 +599,14 @@ if __name__ == '__main__': dist_xr = wrf_dh_4d.get_distances(lat_xr, lon_xr) dist_xr.attrs.update(dict(units='km')) - dist_xr_set = xr.Dataset({'dist': dist_xr}) + geo_info = xr.Dataset({'dist': dist_xr}) for i, (data, xlim, ylim) in enumerate(((wrf_dh_4d._data.T2, [-42, 66], [23, 80]), - (dist_xr_set.dist, [-42, 66], [23, 80]), + (geo_info.dist, [-42, 66], [23, 80]), (wrf_dh_4d._data.T2.where( dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], [45, 58]), - (dist_xr_set.dist.where( + (geo_info.dist.where( dist_xr.sel({'points': 0}).drop('points') <= 100), [2, 15], [45, 58]), )):