You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

51 lines
1.1 KiB

1 year ago
1 year ago
1 year ago
1 year ago
1 year ago
1 year ago
1 year ago
1 year ago
1 year ago
  1. import xarray as xr
  2. import numpy as np
  3. from datetime import datetime
  4. import xrft
  5. import finufft
  6. def fft(dataArray, **kwargs):
  7. return xrft.fft(dataArray, **kwargs)
  8. def ifft(dataArray, **kwargs):
  9. return xrft.ifft(dataArray, **kwargs)
  10. def fft_nutou(dataArray, modeNum, **kwargs):
  11. data = dataArray.to_numpy()
  12. data = data.astype('complex128')
  13. time = dataArray[dataArray.dims[0]].to_numpy()
  14. if isinstance(time[0], type(np.datetime64(500,'s'))):
  15. time = time.astype(float)
  16. time = time - time[0]
  17. freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1))) * 1e9
  18. else:
  19. time = time.astype(float)
  20. time = time - time[0]
  21. freqUpLim = 1 / np.min(np.abs(time - np.roll(time, 1)))
  22. print(freqUpLim)
  23. time = time / time.max() * 2 * np.pi
  24. # calculate the transform
  25. res = xr.DataArray(
  26. data=finufft.nufft1d1(time, data, modeNum, **kwargs),
  27. dims=['freq'],
  28. coords={
  29. "freq":np.linspace(-freqUpLim/2, freqUpLim/2, modeNum)
  30. }
  31. )
  32. return res
  33. def ifft_nutou(dataArray, modeNum):
  34. pass