<jupyter_start><jupyter_text>Sanity checks for the dataloaderS.Rasp developed the dataloader, the pytorch tool to wrap all the data laoding, preparation and batch generation... **Special focus should be given on th patches and times**<jupyter_code>%load_ext autoreload
%autoreload 2
import xarray as xr
import numpy as np
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import cartopy.feature as cfeature
states_provinces = cfeature.NaturalEarthFeature(
        category='cultural',
        name='admin_1_states_provinces_lines',
        scale='50m',
        facecolor='none')
import hvplot.xarray
import cartopy.crs as crs
# add src directory
import sys
sys.path.append('..')
from src.dataloader import *<jupyter_output><empty_output><jupyter_text>General functionality of data loader<jupyter_code>ds = TiggeMRMSDataset(
    tigge_dir='/datadrive/tigge/32km/',
    tigge_vars=['total_precipitation'],
    mrms_dir='/datadrive/mrms/4km/RadarOnly_QPE_06H/',
    rq_fn='/datadrive/mrms/4km/RadarQuality.nc',
#     const_fn='/datadrive/tigge/32km/constants.nc',
#     const_vars=['orog', 'lsm'],
    val_days=7,
    split='train',
    tp_log=1
) 
print(ds)
idx= 1200 # which patch index to use
X, y = ds[idx] # X: tigge, y : radar

plt.pcolor(X[0]) 
plt.colorbar()
plt.figure()
plt.pcolor(y[0])
plt.colorbar()<jupyter_output><empty_output><jupyter_text>Check TIGGE data<jupyter_code>tigge_tp = xr.open_mfdataset('/datadrive/tigge/32km/total_precipitation/*.nc').tp.diff('lead_time')
tigge_tp
tigge_tp.<jupyter_output><empty_output><jupyter_text>Compare Time steps: the time-idx value corresponds to the time given by the overlap_times array. Note that the leadtime needs to be subtracted. Given the different precip scalings, the two images are similar!  --> Time seems to be alright between manually loaded and ds.tigge data<jupyter_code>idx= 3140 # which patch index to use
from  matplotlib.patches import Rectangle
tigge_sel = tigge_tp.sel(lead_time=np.timedelta64(12, 'h'))
t_idx, lat_idx, lon_idx = ds.idxs[idx] # take indices from selected patch idx

fig = plt.figure(figsize=[25,10])
tigge_sel  = tigge_sel.sel(init_time=ds.overlap_times[t_idx] - np.timedelta64(12, 'h'))
tigge_sel.plot(cmap='viridis', levels=np.arange(0,20))

patch_lat = tigge_sel.lat.max().values+0.32
patch_lon = tigge_sel.lon.min().values

# Create  Rectangle patches
delta = ds.patch_size/100 # patch size in deg
for i in range(0,12):
    for j in range(0,9):
        xy = np.array((patch_lon+i*(delta+0.32),patch_lat-(j+1)*(delta+0.32)))
        rect = Rectangle(xy ,delta, delta , linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(i,j), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')
from  matplotlib.patches import Rectangle
fig = plt.figure(figsize=[25,10])
ds.tigge.isel(valid_time=t_idx,variable=0).plot()

patch_lat = tigge_sel.lat.max().values +0.32
patch_lon = tigge_sel.lon.min().values

# Create  Rectangle patches
delta = ds.patch_size/100 # patch size in deg
for i in range(0,12):
    for j in range(0,9):
        xy = np.array((patch_lon+i*(delta+0.32),patch_lat-(j+1)*(delta+0.32)))
        rect = Rectangle(xy ,delta, +delta , linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(i,j), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')<jupyter_output><empty_output><jupyter_text>Now compare ds.tigge with ds[patch_idx]<jupyter_code>idx= 3590 # which patch index to use
test = ds.tigge.isel(
    valid_time=t_idx,
    lat=slice(0 * ds.patch_tigge, (0+1) * ds.patch_tigge + ds.pad_tigge*2),
    lon=slice(0 * ds.patch_tigge, (0+1) * ds.patch_tigge + ds.pad_tigge*2)
)
test.isel(lat=0).plot()
test = ds.tigge.isel(
    valid_time=t_idx,
    lat=slice(1 * ds.patch_tigge, (1+1) * ds.patch_tigge + ds.pad_tigge*2),
    lon=slice(0 * ds.patch_tigge, (0+1) * ds.patch_tigge + ds.pad_tigge*2)
)

test.isel(lat=-1).plot()
from  matplotlib.patches import Rectangle
fig = plt.figure(figsize=[25,10])
t_idx, lat_idx, lon_idx = ds.idxs[idx] # take indices from selected patch idx

ds.tigge.isel(valid_time=t_idx,variable=0).plot()

patch_lat = tigge_sel.lat.max().values+0.32
patch_lon = tigge_sel.lon.min().values

# Create  Rectangle patches
delta = ds.patch_size/100 # patch size in deg
for lat_idx,lon_idx in ds.idxs[:,1:3]:
        test = ds.tigge.isel(
            valid_time=t_idx,
            lat=slice(lat_idx * ds.patch_tigge, (lat_idx+1) * ds.patch_tigge + ds.pad_tigge*2),
            lon=slice(lon_idx * ds.patch_tigge, (lon_idx+1) * ds.patch_tigge + ds.pad_tigge*2)
        )

        #print(j)
        #print(test.lat.max('lat').values)
        xy = np.array([test.lon.min('lon'),test.lat.min('lat').values])
        rect = Rectangle(xy ,delta, delta , linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(lon_idx,lat_idx), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')
        plt.ylim([20,51])


from  matplotlib.patches import Rectangle
fig = plt.figure(figsize=[25,10])
t_idx, lat_idx, lon_idx = ds.idxs[idx] # take indices from selected patch idx

ds.tigge.isel(valid_time=t_idx,variable=0).plot()

patch_lat = tigge_sel.lat.max().values+0.32
patch_lon = tigge_sel.lon.min().values

pad_tigge = 2
# Create  Rectangle patches
delta = ds.patch_size/100 +pad_tigge*2*0.32 # patch size in deg
for lat_idx,lon_idx in ds.idxs[:,1:3]:
        test = ds.tigge.isel(
            valid_time=t_idx,
            lat=slice(lat_idx * ds.patch_tigge, (lat_idx+1) * ds.patch_tigge + pad_tigge*2),
            lon=slice(lon_idx * ds.patch_tigge, (lon_idx+1) * ds.patch_tigge + pad_tigge*2)
        )

        #print(j)
        #print(test.lat.max('lat').values)
        xy = np.array([test.lon.min('lon'),test.lat.min('lat').values])
        rect = Rectangle(xy ,delta, delta, linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(lon_idx,lat_idx), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')
        plt.ylim([20,51])


from  matplotlib.patches import Rectangle
fig = plt.figure(figsize=[25,10])
t_idx, lat_idx, lon_idx = ds.idxs[idx] # take indices from selected patch idx

ds.tigge.isel(valid_time=t_idx,variable=0).plot()

patch_lat = tigge_sel.lat.max().values+0.32
patch_lon = tigge_sel.lon.min().values

# Create  Rectangle patches
delta = ds.patch_size/100 # patch size in deg
for i in range(0,10):
    for j in range(0,5):
        test = ds.tigge.isel(
            valid_time=t_idx,
            lat=slice(j * ds.patch_tigge, (j+1) * ds.patch_tigge + ds.pad_tigge*2),
            lon=slice(i * ds.patch_tigge, (i+1) * ds.patch_tigge + ds.pad_tigge*2)
        )

        #print(j)
        #print(test.lat.max('lat').values)
        xy = np.array([test.lon.min('lon'),test.lat.min('lat').values])
        print(xy)
        rect = Rectangle(xy ,delta, delta , linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(i,j), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')


np.array([test.lon.min('lon'),test.lat.min('lat').values])[0]
test = ds.tigge.isel(
            valid_time=t_idx,
            lat=slice(j * ds.patch_tigge, (j+1) * ds.patch_tigge + ds.pad_tigge*2),
            lon=slice(i * ds.patch_tigge, (i+1) * ds.patch_tigge + ds.pad_tigge*2)
        )
test.plot()
ds.        xy = np.array((test.lon.min('lon').values), test.lat.min('lat').values)
test = ds.tigge.isel(valid_time=t_idx,variable=0).sel(lon=slice(xy[0],xy[0]+delta), lat = slice(xy[1]+delta, xy[1]))
test.plot(levels=np.arange(0,1,.05))
X,_ = ds[idx]
plt.pcolor(X[0])

X = ds.tigge.isel(
            valid_time=t_idx,
            lat=slice(lat_idx * self.patch_tigge, (lat_idx+1) * self.patch_tigge + self.pad_tigge*2),
            lon=slice(lon_idx * self.patch_tigge, (lon_idx+1) * self.patch_tigge + self.pad_tigge*2)
        ).values
tigge_tp<jupyter_output><empty_output><jupyter_text>Check MRMS Radar data<jupyter_code>idx=3590
t_idx, lat_idx,lon_idx = ds.idxs[idx]
X,y = ds[idx]
plt.pcolor(y[0])
ds.mrms.isel(time=t_idx).plot()
fig = plt.figure(figsize=[25,10])
ds.mrms.isel(time=t_idx).plot()

# Create  Rectangle patches
delta = ds.patch_size/100 # patch size in deg
for lat_idx,lon_idx in ds.idxs[:,1:3]:
        test = ds.mrms.isel(
            time=t_idx,
            lat=slice(lat_idx * ds.patch_mrms, (lat_idx+1) * ds.patch_mrms + ds.pad_mrms*2),
            lon=slice(lon_idx * ds.patch_mrms, (lon_idx+1) * ds.patch_mrms + ds.pad_mrms*2)
        )

        #print(j)
        #print(test.lat.max('lat').values)
        xy = np.array([test.lon.min('lon'),test.lat.min('lat').values])
        rect = Rectangle(xy ,delta, delta , linewidth=1,edgecolor='r',facecolor='none')
        fig.gca().add_patch(rect)
        fig.gca().annotate('(%i,%i)'%(lon_idx,lat_idx), xy+delta/2, color='w', weight='bold', 
                fontsize=16, ha='center', va='center')
        plt.ylim([20,51])


lon_idx
lat_idx<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16395035078251485, 'fraction_numerical': 0.02547220723151646, 'mean_word_length': 3.8335941575378194, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '430614', 'license': 'MIT', 'n_tokens_mistral': 3811, 'n_tokens_neox': 3597, 'n_words': 612, 'path': 'notebooks/stephan_notebooks/MH_notebooks/01_Check_dataloader.ipynb', 'repo_name': 'raspstephan/nwp-downscale', 'size': '9265'}