-
Notifications
You must be signed in to change notification settings - Fork 13
/
filtering.py
80 lines (60 loc) · 1.99 KB
/
filtering.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import xarray as xr
import numpy as np
try:
from astropy.convolution import convolve_fft, Gaussian1DKernel, Gaussian2DKernel
astropy = True
except ImportError:
astropy = None
from xarrayutils.utilities import detect_dtype
def filter_1D(data, std, dim="time", dtype=None):
if astropy is None:
raise RuntimeError(
"Module `astropy` not found. Please install optional dependency with `conda install -c conda-forge astropy"
)
if dtype is None:
dtype = detect_dtype(data)
kernel = Gaussian1DKernel(std)
def smooth_raw(data):
raw_data = getattr(data, "values", data)
result = convolve_fft(raw_data, kernel, boundary="wrap")
result[np.isnan(raw_data)] = np.nan
return result
def temporal_smoother(data):
dims = [dim]
return xr.apply_ufunc(
smooth_raw,
data,
vectorize=True,
dask="parallelized",
input_core_dims=[dims],
output_core_dims=[dims],
output_dtypes=[dtype],
)
return temporal_smoother(data)
def filter_2D(data, std, dim, dtype=None):
if astropy is None:
raise RuntimeError(
"Module `astropy` not found. Please install optional dependency with `conda install -c conda-forge astropy"
)
if dtype is None:
dtype = detect_dtype(data)
kernel = Gaussian2DKernel(std)
def smooth_raw(data):
raw_data = getattr(data, "values", data)
result = convolve_fft(raw_data, kernel, boundary="wrap")
result[np.isnan(raw_data)] = np.nan
return result
def smoother(data):
dims = dim
# this is different from the 1d case
return xr.apply_ufunc(
smooth_raw,
data,
vectorize=True,
dask="parallelized",
input_core_dims=[dims],
output_core_dims=[dims],
output_dtypes=[dtype],
)
return smoother(data)
# TODO spatial filter