64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
from scipy._lib._array_api import array_namespace
|
|
import numpy as np
|
|
from . import _pocketfft
|
|
|
|
__all__ = ['dct', 'idct', 'dst', 'idst', 'dctn', 'idctn', 'dstn', 'idstn']
|
|
|
|
|
|
def _execute(pocketfft_func, x, type, s, axes, norm,
|
|
overwrite_x, workers, orthogonalize):
|
|
xp = array_namespace(x)
|
|
x = np.asarray(x)
|
|
y = pocketfft_func(x, type, s, axes, norm,
|
|
overwrite_x=overwrite_x, workers=workers,
|
|
orthogonalize=orthogonalize)
|
|
return xp.asarray(y)
|
|
|
|
|
|
def dctn(x, type=2, s=None, axes=None, norm=None,
|
|
overwrite_x=False, workers=None, *, orthogonalize=None):
|
|
return _execute(_pocketfft.dctn, x, type, s, axes, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def idctn(x, type=2, s=None, axes=None, norm=None,
|
|
overwrite_x=False, workers=None, *, orthogonalize=None):
|
|
return _execute(_pocketfft.idctn, x, type, s, axes, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def dstn(x, type=2, s=None, axes=None, norm=None,
|
|
overwrite_x=False, workers=None, orthogonalize=None):
|
|
return _execute(_pocketfft.dstn, x, type, s, axes, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def idstn(x, type=2, s=None, axes=None, norm=None,
|
|
overwrite_x=False, workers=None, *, orthogonalize=None):
|
|
return _execute(_pocketfft.idstn, x, type, s, axes, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def dct(x, type=2, n=None, axis=-1, norm=None,
|
|
overwrite_x=False, workers=None, orthogonalize=None):
|
|
return _execute(_pocketfft.dct, x, type, n, axis, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def idct(x, type=2, n=None, axis=-1, norm=None,
|
|
overwrite_x=False, workers=None, orthogonalize=None):
|
|
return _execute(_pocketfft.idct, x, type, n, axis, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def dst(x, type=2, n=None, axis=-1, norm=None,
|
|
overwrite_x=False, workers=None, orthogonalize=None):
|
|
return _execute(_pocketfft.dst, x, type, n, axis, norm,
|
|
overwrite_x, workers, orthogonalize)
|
|
|
|
|
|
def idst(x, type=2, n=None, axis=-1, norm=None,
|
|
overwrite_x=False, workers=None, orthogonalize=None):
|
|
return _execute(_pocketfft.idst, x, type, n, axis, norm,
|
|
overwrite_x, workers, orthogonalize)
|