dist plots for demo

mlair/data_handler/data_handler_wrf_chem.py

@@ -160,21 +160,25 @@ def kdtree_fast(latvar, lonvar, lat0, lon0):
 
 if __name__ == '__main__':
 
-    def plot_map_proj(data, xlim=None, ylim=None, filename=None):
+    def plot_map_proj(data, xlim=None, ylim=None, filename=None, point=None):
+        crs_latlon = ccrs.PlateCarree()
         if ylim is None:
             ylim = [-90, 90]
         if xlim is None:
             xlim = [-180, 180]
         if filename is None:
             filename = 'test_fig.pdf'
-        plt.figure(figsize=(14, 6))
+        # plt.figure(figsize=(14, 6))
+        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, )
-        ax.coastlines()
+        ax.coastlines(resolution='50m')
         ax.set_ylim(ylim)
         ax.set_xlim(xlim)
+        plt.plot(point[1], point[0], 'bo', markersize=7, transform=crs_latlon)
+        plt.tight_layout()
         plt.savefig(filename)
         plt.close('all')
 
@@ -189,18 +193,29 @@ if __name__ == '__main__':
     wrf_dh.rechunk_data({"XTIME": 1, "y": 36, "x": 40})
     T2 = wrf_dh._data.T2
 
-    icoords = dask.compute(wrf_dh.get_nearest_coordinates(50.733, 7.10))[0]
-    # BN_T2_test = 279.95132
-    # BN_T2 = wrf_dh._data.isel(icoords).T2
-
-    lat_np = np.array([50.73333, 55.73333])
-    lon_np = np.array([7.1, 7.1])
-    lat_xr = xr.DataArray(lat_np, dims=["points"], coords={'points': [0, 1]})
-    lon_xr = xr.DataArray(lon_np, dims=["points"], coords={'points': [0, 1]})
+    lat_np = np.array([50.73333])
+    lon_np = np.array([7.1])
+    icoords = dask.compute(wrf_dh.get_nearest_coordinates(lat_np, lon_np))[0]
+    lat_xr = xr.DataArray(lat_np, dims=["points"], coords={'points': range(len(lat_np))})
+    lon_xr = xr.DataArray(lon_np, dims=["points"], coords={'points': range(len(lon_np))})
 
     dist_np = wrf_dh.get_distances(lat_np, lon_np)
     dist_xr = wrf_dh.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_dh._data.T2, [-42, 66], [23, 80]),
+                                            (dist_xr_set.dist, [-42, 66], [23, 80]),
+                                            (wrf_dh._data.T2.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]),
+                                            )):
+        plot_map_proj(data, xlim=xlim,
+                      ylim=ylim,
+                      point=[lat_np, lon_np], filename=f'test_dist{i}.pdf')
+
 
+    plot_map_proj(wrf_dh._data.T2.where(dist_xr.sel({'points': 0}).drop('points') <= 100), xlim=[-0, 15], ylim=[40, 58],
+                  point=[lat_np, lon_np], filename='test_dist.pdf')
     wrf_dh_18 = WrfChemDataHandler(data_path='/home/felix/Data/WRF-Chem/test_data_aura/',
                                    common_file_starter='wrfout_d01_2018-08-01',
                                    window_history_size=12,