Skip to content

Commit 156af53

Browse files
vtavanavtavana
committed
update _scipy_fft.py
1 parent d4929a3 commit 156af53

File tree

5 files changed

+198
-287
lines changed

5 files changed

+198
-287
lines changed

‎mkl_fft/_fft_utils.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#!/usr/bin/env python
2+
# Copyright (c) 2025, Intel Corporation
3+
#
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted provided that the following conditions are met:
6+
#
7+
# * Redistributions of source code must retain the above copyright notice,
8+
# this list of conditions and the following disclaimer.
9+
# * Redistributions in binary form must reproduce the above copyright
10+
# notice, this list of conditions and the following disclaimer in the
11+
# documentation and/or other materials provided with the distribution.
12+
# * Neither the name of Intel Corporation nor the names of its contributors
13+
# may be used to endorse or promote products derived from this software
14+
# without specific prior written permission.
15+
#
16+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
17+
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
19+
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
20+
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
21+
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
22+
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
23+
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
24+
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
import numpy as np
28+
29+
__all__ = ["_check_norm", "_compute_fwd_scale"]
30+
31+
32+
def _check_norm(norm):
33+
if norm not in (None, "ortho", "forward", "backward"):
34+
raise ValueError(
35+
f"Invalid norm value {norm} should be None, 'ortho', 'forward', "
36+
"or 'backward'."
37+
)
38+
39+
40+
def _compute_fwd_scale(norm, n, shape):
41+
_check_norm(norm)
42+
if norm in (None, "backward"):
43+
return 1.0
44+
45+
ss = n if n is not None else shape
46+
nn = np.prod(ss)
47+
fsc = 1 / nn if nn != 0 else 1
48+
if norm == "forward":
49+
return fsc
50+
else: # norm == "ortho"
51+
return np.sqrt(fsc)

‎mkl_fft/_float_utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,5 +91,5 @@ def __supported_array_or_not_implemented(x):
9191
if hasattr(np, "complex256"):
9292
black_list.append(np.complex256)
9393
if __x.dtype in black_list:
94-
return NotImplemented
94+
raise NotImplementedError
9595
return __x

‎mkl_fft/_numpy_fft.py

Lines changed: 56 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -74,32 +74,53 @@
7474
import warnings
7575

7676
import numpy as np
77-
from numpy import array, conjugate, prod, sqrt, take
7877

79-
from . import _float_utils
8078
from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module
79+
from ._fft_utils import _check_norm, _compute_fwd_scale
80+
from ._float_utils import __downcast_float128_array
8181

8282

83-
def _compute_fwd_scale(norm, n, shape):
84-
_check_norm(norm)
85-
if norm in (None, "backward"):
86-
return 1.0
87-
88-
ss = n if n is not None else shape
89-
nn = prod(ss)
90-
fsc = 1 / nn if nn != 0 else 1
91-
if norm == "forward":
92-
return fsc
93-
else: # norm == "ortho"
94-
return sqrt(fsc)
95-
96-
97-
def _check_norm(norm):
98-
if norm not in (None, "ortho", "forward", "backward"):
99-
raise ValueError(
100-
f"Invalid norm value {norm} should be None, 'ortho', 'forward', "
101-
"or 'backward'."
83+
# copied from: https://github.com/numpy/numpy/blob/main/numpy/fft/_pocketfft.py
84+
def _cook_nd_args(a, s=None, axes=None, invreal=False):
85+
if s is None:
86+
shapeless = True
87+
if axes is None:
88+
s = list(a.shape)
89+
else:
90+
s = np.take(a.shape, axes)
91+
else:
92+
shapeless = False
93+
s = list(s)
94+
if axes is None:
95+
if not shapeless and np.__version__ >= "2.0":
96+
msg = (
97+
"`axes` should not be `None` if `s` is not `None` "
98+
"(Deprecated in NumPy 2.0). In a future version of NumPy, "
99+
"this will raise an error and `s[i]` will correspond to "
100+
"the size along the transformed axis specified by "
101+
"`axes[i]`. To retain current behaviour, pass a sequence "
102+
"[0, ..., k-1] to `axes` for an array of dimension k."
103+
)
104+
warnings.warn(msg, DeprecationWarning, stacklevel=3)
105+
axes = list(range(-len(s), 0))
106+
if len(s) != len(axes):
107+
raise ValueError("Shape and axes have different lengths.")
108+
if invreal and shapeless:
109+
s[-1] = (a.shape[axes[-1]] - 1) * 2
110+
if None in s and np.__version__ >= "2.0":
111+
msg = (
112+
"Passing an array containing `None` values to `s` is "
113+
"deprecated in NumPy 2.0 and will raise an error in "
114+
"a future version of NumPy. To use the default behaviour "
115+
"of the corresponding 1-D transform, pass the value matching "
116+
"the default for its `n` parameter. To use the default "
117+
"behaviour for every axis, the `s` argument can be omitted."
102118
)
119+
warnings.warn(msg, DeprecationWarning, stacklevel=3)
120+
# use the whole input array along axis `i` if `s[i] == -1 or None`
121+
s = [a.shape[_a] if _s in [-1, None] else _s for _s, _a in zip(s, axes)]
122+
123+
return s, axes
103124

104125

105126
def _swap_direction(norm):
@@ -218,7 +239,7 @@ def fft(a, n=None, axis=-1, norm=None):
218239
219240
"""
220241

221-
x = _float_utils.__downcast_float128_array(a)
242+
x = __downcast_float128_array(a)
222243
fsc = _compute_fwd_scale(norm, n, x.shape[axis])
223244

224245
return trycall(mkl_fft.fft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
@@ -311,7 +332,7 @@ def ifft(a, n=None, axis=-1, norm=None):
311332
312333
"""
313334

314-
x = _float_utils.__downcast_float128_array(a)
335+
x = __downcast_float128_array(a)
315336
fsc = _compute_fwd_scale(norm, n, x.shape[axis])
316337

317338
return trycall(mkl_fft.ifft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
@@ -402,7 +423,7 @@ def rfft(a, n=None, axis=-1, norm=None):
402423
403424
"""
404425

405-
x = _float_utils.__downcast_float128_array(a)
426+
x = __downcast_float128_array(a)
406427
fsc = _compute_fwd_scale(norm, n, x.shape[axis])
407428

408429
return trycall(mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc})
@@ -495,7 +516,7 @@ def irfft(a, n=None, axis=-1, norm=None):
495516
496517
"""
497518

498-
x = _float_utils.__downcast_float128_array(a)
519+
x = __downcast_float128_array(a)
499520
fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))
500521

501522
return trycall(
@@ -581,9 +602,9 @@ def hfft(a, n=None, axis=-1, norm=None):
581602
"""
582603

583604
norm = _swap_direction(norm)
584-
x = _float_utils.__downcast_float128_array(a)
585-
x = array(x, copy=True, dtype=complex)
586-
conjugate(x, out=x)
605+
x = __downcast_float128_array(a)
606+
x = np.array(x, copy=True, dtype=complex)
607+
np.conjugate(x, out=x)
587608
fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))
588609

589610
return trycall(
@@ -651,61 +672,18 @@ def ihfft(a, n=None, axis=-1, norm=None):
651672

652673
# The copy may be required for multithreading.
653674
norm = _swap_direction(norm)
654-
x = _float_utils.__downcast_float128_array(a)
655-
x = array(x, copy=True, dtype=float)
675+
x = __downcast_float128_array(a)
676+
x = np.array(x, copy=True, dtype=float)
656677
fsc = _compute_fwd_scale(norm, n, x.shape[axis])
657678

658679
output = trycall(
659680
mkl_fft.rfft, (x,), {"n": n, "axis": axis, "fwd_scale": fsc}
660681
)
661682

662-
conjugate(output, out=output)
683+
np.conjugate(output, out=output)
663684
return output
664685

665686

666-
# copied from: https://github.com/numpy/numpy/blob/main/numpy/fft/_pocketfft.py
667-
def _cook_nd_args(a, s=None, axes=None, invreal=False):
668-
if s is None:
669-
shapeless = True
670-
if axes is None:
671-
s = list(a.shape)
672-
else:
673-
s = take(a.shape, axes)
674-
else:
675-
shapeless = False
676-
s = list(s)
677-
if axes is None:
678-
if not shapeless and np.__version__ >= "2.0":
679-
msg = (
680-
"`axes` should not be `None` if `s` is not `None` "
681-
"(Deprecated in NumPy 2.0). In a future version of NumPy, "
682-
"this will raise an error and `s[i]` will correspond to "
683-
"the size along the transformed axis specified by "
684-
"`axes[i]`. To retain current behaviour, pass a sequence "
685-
"[0, ..., k-1] to `axes` for an array of dimension k."
686-
)
687-
warnings.warn(msg, DeprecationWarning, stacklevel=3)
688-
axes = list(range(-len(s), 0))
689-
if len(s) != len(axes):
690-
raise ValueError("Shape and axes have different lengths.")
691-
if invreal and shapeless:
692-
s[-1] = (a.shape[axes[-1]] - 1) * 2
693-
if None in s and np.__version__ >= "2.0":
694-
msg = (
695-
"Passing an array containing `None` values to `s` is "
696-
"deprecated in NumPy 2.0 and will raise an error in "
697-
"a future version of NumPy. To use the default behaviour "
698-
"of the corresponding 1-D transform, pass the value matching "
699-
"the default for its `n` parameter. To use the default "
700-
"behaviour for every axis, the `s` argument can be omitted."
701-
)
702-
warnings.warn(msg, DeprecationWarning, stacklevel=3)
703-
# use the whole input array along axis `i` if `s[i] == -1 or None`
704-
s = [a.shape[_a] if _s in [-1, None] else _s for _s, _a in zip(s, axes)]
705-
706-
return s, axes
707-
708-
709687
def fftn(a, s=None, axes=None, norm=None):
710688
"""
711689
Compute the N-dimensional discrete Fourier Transform.
@@ -806,7 +784,7 @@ def fftn(a, s=None, axes=None, norm=None):
806784
807785
"""
808786

809-
x = _float_utils.__downcast_float128_array(a)
787+
x = __downcast_float128_array(a)
810788
s, axes = _cook_nd_args(x, s, axes)
811789
fsc = _compute_fwd_scale(norm, s, x.shape)
812790

@@ -913,7 +891,7 @@ def ifftn(a, s=None, axes=None, norm=None):
913891
914892
"""
915893

916-
x = _float_utils.__downcast_float128_array(a)
894+
x = __downcast_float128_array(a)
917895
s, axes = _cook_nd_args(x, s, axes)
918896
fsc = _compute_fwd_scale(norm, s, x.shape)
919897

@@ -1201,7 +1179,7 @@ def rfftn(a, s=None, axes=None, norm=None):
12011179
12021180
"""
12031181

1204-
x = _float_utils.__downcast_float128_array(a)
1182+
x = __downcast_float128_array(a)
12051183
s, axes = _cook_nd_args(x, s, axes)
12061184
fsc = _compute_fwd_scale(norm, s, x.shape)
12071185

@@ -1345,7 +1323,7 @@ def irfftn(a, s=None, axes=None, norm=None):
13451323
13461324
"""
13471325

1348-
x = _float_utils.__downcast_float128_array(a)
1326+
x = __downcast_float128_array(a)
13491327
s, axes = _cook_nd_args(x, s, axes, invreal=True)
13501328
fsc = _compute_fwd_scale(norm, s, x.shape)
13511329

0 commit comments

Comments
 (0)