domino.filtering

  1import numpy as np
  2from scipy.ndimage import label
  3import itertools as it
  4import xarray as xr
  5
  6def latlonarea(lat0,lat1,lon0,lon1,r=6731):
  7    theta=np.deg2rad((lat1+lat0)/2)
  8    r2=r**2
  9    dphi=lon1-lon0
 10    dtheta=lat1-lat0
 11    A=r2*np.abs(np.cos(theta))*dtheta*dphi #in km**2
 12    return A
 13
 14def grid_area(da,lat_coord='lat',lon_coord='lon',r=6371):
 15    lat=da[lat_coord].values
 16    lon=da[lon_coord].values
 17    dlon=(lon[1:]-lon[:-1])/2
 18    dlat=(lat[1:]-lat[:-1])/2
 19    dlon=[np.median(dlon),*dlon]
 20    dlat=[np.median(dlat),*dlat]
 21    areas=np.array([[\
 22        latlonarea(la-dla,la+dla,lo-dlo,lo+dlo,r=r)\
 23        for la,dla in zip(lat,dlat)]\
 24        for lo,dlo in zip(lon,dlon)])
 25    area=xr.DataArray(data=areas.T,\
 26        coords={lat_coord:lat,lon_coord:lon})
 27    return area
 28
 29
 30def apply_2d_func_to_da(da,func,*args,dims=None):
 31    da=da.copy(deep=True)
 32    da_dim=None    
 33    if dims is not None:
 34        #if dims is specified, ignore da's without all dims
 35        if not set(dims).issubset(set(da.dims)):
 36            return da
 37        #Put dims at the front
 38        else:
 39            da_dim=list(da.dims)
 40            da=da.transpose(*dims,...)
 41            
 42        #Since funcs are 2D,
 43        #we must add a dummy dim if we only want 1 dim
 44        if len(dims)==1:
 45            da=da.expand_dims('dummy_dim',0).copy(deep=True)
 46            
 47    #We now want to loop over all dims except the first two
 48    #and apply our function
 49    indexers=[np.arange(i) for i in da.shape[2:]]
 50    for i in it.product(*indexers):
 51        da[(...,*i)]=func(da[(...,*i)].values,*args)
 52    try:
 53        da=da.isel(dummy_dim=0)
 54    except Exception as e:
 55        pass
 56    if da_dim is not None:
 57        da=da.transpose(*da_dim)    
 58    return da
 59
 60def get_large_regions(grid,point_thresh):
 61    connected,ncon=label(grid)
 62    connected=connected.astype(float)
 63    ix,n=np.unique(connected,return_counts=True)
 64    for m in ix[n<point_thresh]:
 65        connected[connected==m]=0
 66    return connected>0
 67
 68def get_area_regions(grid,area_thresh,area):
 69    connected,ncon=label(grid)
 70    connected=connected.astype(float)
 71    ix,n=np.unique(connected,return_counts=True)
 72    
 73    connected_da=xr.DataArray(connected,coords=area.coords)
 74    ix_areas=np.array([area.where(connected_da==i).sum().values\
 75              for i in ix])
 76    for m in ix[ix_areas<area_thresh]:
 77        connected[connected==m]=0
 78    return connected>0
 79
 80
 81def da_large_regions(da,n,dims,area_based=False):
 82    
 83    if area_based:
 84        
 85        if len(dims)!=2:
 86            raise(ValueError('Area weighted n currently only supported for two dimensions, assumed to be (lat,lon)'))
 87        lat,lon=dims
 88        #area in units of solid angle
 89        area=grid_area(da,lat_coord=lat,lon_coord=lon,r=1)
 90        
 91        return apply_2d_func_to_da(da,get_area_regions,n,area,dims=dims)
 92    else:
 93        return apply_2d_func_to_da(da,get_large_regions,n,dims=dims)
 94
 95def ds_large_regions(mask_ds,n,dims,area_based=False):
 96        ds=mask_ds.copy()
 97        for var in list(ds.data_vars):
 98            
 99            if (dims is not None) and (not np.all([d in mask_ds[var].dims for d in dims])):
100                pass #ignore variables with missing dims
101            else:
102                ds[var]=da_large_regions(ds[var],n,dims,area_based=area_based)
103        return ds
104
105def convolve_pad(x,N):
106    Y=np.array([np.convolve(y,np.ones(np.min([N,len(y)])),mode='same') for y in x])
107    Y=np.array([np.convolve(y,np.ones(np.min([N,len(y)])),mode='same') for y in Y.T]).T
108    Y[Y>0]=1
109    return Y
110
111def convolve_pad_da(da,n,dims):
112    return apply_2d_func_to_da(da,convolve_pad,n,dims=dims)
113
114def convolve_pad_ds(ds,n,dims):
115    ds=ds.copy()
116    for var in list(ds.data_vars):
117        ds[var]=convolve_pad_da(ds[var],n,dims=dims)
118    return ds
119
120def std_filter(val_ds,std,frac):
121    return np.abs(val_ds)>=frac*std