55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
|
import xarray as xr
|
||
|
|
import pandas as pd
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
import xrft
|
||
|
|
import finufft
|
||
|
|
|
||
|
|
|
||
|
|
class FFTAnalyser():
|
||
|
|
|
||
|
|
def __init__(self) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def fft(self, dataArray, **kwargs):
|
||
|
|
return xrft.fft(dataArray, **kwargs)
|
||
|
|
|
||
|
|
def ifft(self, dataArray, **kwargs):
|
||
|
|
return xrft.ifft(dataArray, **kwargs)
|
||
|
|
|
||
|
|
def fft_nutou(self, dataArray, modeNum):
|
||
|
|
|
||
|
|
data = dataArray.to_numpy()
|
||
|
|
data = data.astype('complex128')
|
||
|
|
|
||
|
|
time = dataArray[dataArray.dims[0]].to_numpy()
|
||
|
|
|
||
|
|
if isinstance(time[0], type(np.datetime64(500,'s'))):
|
||
|
|
time = time.astype(float)
|
||
|
|
time = time - time[0]
|
||
|
|
freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1))) * 1e9
|
||
|
|
else:
|
||
|
|
time = time.astype(float)
|
||
|
|
time = time - time[0]
|
||
|
|
freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1)))
|
||
|
|
|
||
|
|
print(freqUpLim)
|
||
|
|
|
||
|
|
time = time / time.max() * 2 * np.pi
|
||
|
|
|
||
|
|
# calculate the transform
|
||
|
|
res = xr.DataArray(
|
||
|
|
data=finufft.nufft1d1(time, data, modeNum),
|
||
|
|
dims=['freq'],
|
||
|
|
coords={
|
||
|
|
"freq":np.linspace(-freqUpLim/2, freqUpLim/2, modeNum)
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
return res
|
||
|
|
|
||
|
|
def ifft_nutou(self, dataArray, modeNum):
|
||
|
|
pass
|
||
|
|
|