|
24 | 24 | # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
25 | 25 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
26 | 26 |
|
27 | | -from ._fft_utils import _c2c_fftnd_impl, _c2r_fftnd_impl, _r2c_fftnd_impl |
| 27 | +from ._fft_utils import ( |
| 28 | + _c2c_fftnd_impl, |
| 29 | + _c2r_fftnd_impl, |
| 30 | + _compute_fwd_scale, |
| 31 | + _r2c_fftnd_impl, |
| 32 | +) |
28 | 33 |
|
29 | 34 | # pylint: disable=no-name-in-module |
30 | 35 | from ._pydfti import _c2c_fft1d_impl, _c2r_fft1d_impl, _r2c_fft1d_impl |
|
45 | 50 | ] |
46 | 51 |
|
47 | 52 |
|
48 | | -def fft(x, n=None, axis=-1, fwd_scale=1.0, out=None): |
49 | | - return _c2c_fft1d_impl( |
50 | | - x, n=n, axis=axis, out=out, direction=+1, fsc=fwd_scale |
51 | | - ) |
| 53 | +def fft(x, n=None, axis=-1, norm=None, out=None): |
| 54 | + fsc = _compute_fwd_scale(norm, n, x.shape[axis]) |
| 55 | + return _c2c_fft1d_impl(x, n=n, axis=axis, out=out, direction=+1, fsc=fsc) |
52 | 56 |
|
53 | 57 |
|
54 | | -def ifft(x, n=None, axis=-1, fwd_scale=1.0, out=None): |
55 | | - return _c2c_fft1d_impl( |
56 | | - x, n=n, axis=axis, out=out, direction=-1, fsc=fwd_scale |
57 | | - ) |
| 58 | +def ifft(x, n=None, axis=-1, norm=None, out=None): |
| 59 | + fsc = _compute_fwd_scale(norm, n, x.shape[axis]) |
| 60 | + return _c2c_fft1d_impl(x, n=n, axis=axis, out=out, direction=-1, fsc=fsc) |
58 | 61 |
|
59 | 62 |
|
60 | | -def fft2(x, s=None, axes=(-2, -1), fwd_scale=1.0, out=None): |
61 | | - return fftn(x, s=s, axes=axes, out=out, fwd_scale=fwd_scale) |
| 63 | +def fft2(x, s=None, axes=(-2, -1), norm=None, out=None): |
| 64 | + return fftn(x, s=s, axes=axes, norm=norm, out=out) |
62 | 65 |
|
63 | 66 |
|
64 | | -def ifft2(x, s=None, axes=(-2, -1), fwd_scale=1.0, out=None): |
65 | | - return ifftn(x, s=s, axes=axes, out=out, fwd_scale=fwd_scale) |
| 67 | +def ifft2(x, s=None, axes=(-2, -1), norm=None, out=None): |
| 68 | + return ifftn(x, s=s, axes=axes, norm=norm, out=out) |
66 | 69 |
|
67 | 70 |
|
68 | | -def fftn(x, s=None, axes=None, fwd_scale=1.0, out=None): |
69 | | - return _c2c_fftnd_impl( |
70 | | - x, s=s, axes=axes, out=out, direction=+1, fsc=fwd_scale |
71 | | - ) |
| 71 | +def fftn(x, s=None, axes=None, norm=None, out=None): |
| 72 | + fsc = _compute_fwd_scale(norm, s, x.shape) |
| 73 | + return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=+1, fsc=fsc) |
72 | 74 |
|
73 | 75 |
|
74 | | -def ifftn(x, s=None, axes=None, fwd_scale=1.0, out=None): |
75 | | - return _c2c_fftnd_impl( |
76 | | - x, s=s, axes=axes, out=out, direction=-1, fsc=fwd_scale |
77 | | - ) |
| 76 | +def ifftn(x, s=None, axes=None, norm=None, out=None): |
| 77 | + fsc = _compute_fwd_scale(norm, s, x.shape) |
| 78 | + return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=-1, fsc=fsc) |
78 | 79 |
|
79 | 80 |
|
80 | | -def rfft(x, n=None, axis=-1, fwd_scale=1.0, out=None): |
81 | | - return _r2c_fft1d_impl(x, n=n, axis=axis, out=out, fsc=fwd_scale) |
| 81 | +def rfft(x, n=None, axis=-1, norm=None, out=None): |
| 82 | + fsc = _compute_fwd_scale(norm, n, x.shape[axis]) |
| 83 | + return _r2c_fft1d_impl(x, n=n, axis=axis, out=out, fsc=fsc) |
82 | 84 |
|
83 | 85 |
|
84 | | -def irfft(x, n=None, axis=-1, fwd_scale=1.0, out=None): |
85 | | - return _c2r_fft1d_impl(x, n=n, axis=axis, out=out, fsc=fwd_scale) |
| 86 | +def irfft(x, n=None, axis=-1, norm=None, out=None): |
| 87 | + fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1)) |
| 88 | + return _c2r_fft1d_impl(x, n=n, axis=axis, out=out, fsc=fsc) |
86 | 89 |
|
87 | 90 |
|
88 | | -def rfft2(x, s=None, axes=(-2, -1), fwd_scale=1.0, out=None): |
89 | | - return rfftn(x, s=s, axes=axes, out=out, fwd_scale=fwd_scale) |
| 91 | +def rfft2(x, s=None, axes=(-2, -1), norm=None, out=None): |
| 92 | + return rfftn(x, s=s, axes=axes, norm=norm, out=out) |
90 | 93 |
|
91 | 94 |
|
92 | | -def irfft2(x, s=None, axes=(-2, -1), fwd_scale=1.0, out=None): |
93 | | - return irfftn(x, s=s, axes=axes, out=out, fwd_scale=fwd_scale) |
| 95 | +def irfft2(x, s=None, axes=(-2, -1), norm=None, out=None): |
| 96 | + return irfftn(x, s=s, axes=axes, norm=norm, out=out) |
94 | 97 |
|
95 | 98 |
|
96 | | -def rfftn(x, s=None, axes=None, fwd_scale=1.0, out=None): |
97 | | - return _r2c_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fwd_scale) |
| 99 | +def rfftn(x, s=None, axes=None, norm=None, out=None): |
| 100 | + fsc = _compute_fwd_scale(norm, s, x.shape) |
| 101 | + return _r2c_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc) |
98 | 102 |
|
99 | 103 |
|
100 | | -def irfftn(x, s=None, axes=None, fwd_scale=1.0, out=None): |
101 | | - return _c2r_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fwd_scale) |
| 104 | +def irfftn(x, s=None, axes=None, norm=None, out=None): |
| 105 | + fsc = _compute_fwd_scale(norm, s, x.shape) |
| 106 | + return _c2r_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc) |
0 commit comments