Skip to content

Add Hermitian FFT to scipy interface #161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 23, 2025
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,13 @@
:code:`mkl_fft` changelog
=========================

1.3.14 (04/10/2025)
[dev] (MM/DD/YY)
==================

scipy interface :code:`mkl_fft.interfaces.scipy_fft` now includes Hermitian FFT functions:
:code:`hfft`, :code:`ihfft`, :code:`hfftn`, :code:`ihfftn`, :code:`hfft2`, and :code:`ihfft2`

1.3.14 (04/11/2025)
===================

resolves gh-152 by adding an explicit :code:`mkl-service` dependency to :code:`mkl-fft` when building the wheel
Expand Down
14 changes: 13 additions & 1 deletion mkl_fft/_fft_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@

import numpy as np

__all__ = ["_check_norm", "_compute_fwd_scale"]
__all__ = ["_check_norm", "_compute_fwd_scale", "_swap_direction"]


def _check_norm(norm):
Expand All @@ -49,3 +49,15 @@ def _compute_fwd_scale(norm, n, shape):
return fsc
else: # norm == "ortho"
return np.sqrt(fsc)


def _swap_direction(norm):
_check_norm(norm)
_swap_direction_map = {
"backward": "forward",
None: "forward",
"ortho": "ortho",
"forward": "backward",
}

return _swap_direction_map[norm]
18 changes: 2 additions & 16 deletions mkl_fft/_numpy_fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@
import numpy as np

from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module
from ._fft_utils import _check_norm, _compute_fwd_scale
from ._fft_utils import _compute_fwd_scale, _swap_direction
from ._float_utils import _downcast_float128_array


Expand Down Expand Up @@ -124,18 +124,6 @@ def _cook_nd_args(a, s=None, axes=None, invreal=False):
return s, axes


def _swap_direction(norm):
_check_norm(norm)
_swap_direction_map = {
"backward": "forward",
None: "forward",
"ortho": "ortho",
"forward": "backward",
}

return _swap_direction_map[norm]


def trycall(func, args, kwrds):
try:
res = func(*args, **kwrds)
Expand Down Expand Up @@ -604,7 +592,7 @@ def hfft(a, n=None, axis=-1, norm=None):

norm = _swap_direction(norm)
x = _downcast_float128_array(a)
x = np.array(x, copy=True, dtype=complex)
x = np.array(x, copy=True)
np.conjugate(x, out=x)
fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))

Expand Down Expand Up @@ -671,10 +659,8 @@ def ihfft(a, n=None, axis=-1, norm=None):

"""

# The copy may be required for multithreading.
norm = _swap_direction(norm)
x = _downcast_float128_array(a)
x = np.array(x, copy=True, dtype=float)
fsc = _compute_fwd_scale(norm, n, x.shape[axis])

output = trycall(
Expand Down
106 changes: 93 additions & 13 deletions mkl_fft/_scipy_fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
import numpy as np

from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module
from ._fft_utils import _compute_fwd_scale
from ._fft_utils import _compute_fwd_scale, _swap_direction
from ._float_utils import _supported_array_or_not_implemented

__doc__ = """
Expand Down Expand Up @@ -125,6 +125,12 @@ def set_workers(n_workers):
"irfft2",
"rfftn",
"irfftn",
"hfft",
"ihfft",
"hfft2",
"ihfft2",
"hfftn",
"ihfftn",
"get_workers",
"set_workers",
"DftiBackend",
Expand Down Expand Up @@ -231,7 +237,7 @@ def _validate_input(a):


def fft(
a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None
a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, plan=None
):
_check_plan(plan)
x = _validate_input(a)
Expand All @@ -244,7 +250,7 @@ def fft(


def ifft(
a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, plan=None
a, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, plan=None
):
_check_plan(plan)
x = _validate_input(a)
Expand All @@ -263,6 +269,7 @@ def fft2(
norm=None,
overwrite_x=False,
workers=None,
*,
plan=None,
):

Expand All @@ -284,6 +291,7 @@ def ifft2(
norm=None,
overwrite_x=False,
workers=None,
*,
plan=None,
):

Expand All @@ -299,7 +307,14 @@ def ifft2(


def fftn(
a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None
a,
s=None,
axes=None,
norm=None,
overwrite_x=False,
workers=None,
*,
plan=None,
):
_check_plan(plan)
x = _validate_input(a)
Expand All @@ -312,7 +327,14 @@ def fftn(


def ifftn(
a, s=None, axes=None, norm=None, overwrite_x=False, workers=None, plan=None
a,
s=None,
axes=None,
norm=None,
overwrite_x=False,
workers=None,
*,
plan=None,
):
_check_plan(plan)
x = _validate_input(a)
Expand All @@ -324,7 +346,7 @@ def ifftn(
)


def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
def rfft(a, n=None, axis=-1, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
fsc = _compute_fwd_scale(norm, n, x.shape[axis])
Expand All @@ -333,7 +355,7 @@ def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
return mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc)


def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
def irfft(a, n=None, axis=-1, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))
Expand All @@ -342,17 +364,15 @@ def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
return mkl_fft.irfft(x, n=n, axis=axis, fwd_scale=fsc)


def rfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None):

def rfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, *, plan=None):
return rfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan)


def irfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, plan=None):

def irfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, *, plan=None):
return irfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan)


def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None):
def rfftn(a, s=None, axes=None, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
s, axes = _cook_nd_args(x, s, axes)
Expand All @@ -362,11 +382,71 @@ def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None):
return mkl_fft.rfftn(x, s, axes, fwd_scale=fsc)


def irfftn(a, s=None, axes=None, norm=None, workers=None, plan=None):
def irfftn(a, s=None, axes=None, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
s, axes = _cook_nd_args(x, s, axes, invreal=True)
fsc = _compute_fwd_scale(norm, s, x.shape)

with Workers(workers):
return mkl_fft.irfftn(x, s, axes, fwd_scale=fsc)


def hfft(a, n=None, axis=-1, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
norm = _swap_direction(norm)
x = np.array(x, copy=True)
np.conjugate(x, out=x)
fsc = _compute_fwd_scale(norm, n, 2 * (x.shape[axis] - 1))

with Workers(workers):
return mkl_fft.irfft(x, n=n, axis=axis, fwd_scale=fsc)


def ihfft(a, n=None, axis=-1, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
norm = _swap_direction(norm)
fsc = _compute_fwd_scale(norm, n, x.shape[axis])

with Workers(workers):
result = mkl_fft.rfft(x, n=n, axis=axis, fwd_scale=fsc)

np.conjugate(result, out=result)
return result


def hfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, *, plan=None):
return hfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan)


def ihfft2(a, s=None, axes=(-2, -1), norm=None, workers=None, *, plan=None):
return ihfftn(a, s=s, axes=axes, norm=norm, workers=workers, plan=plan)


def hfftn(a, s=None, axes=None, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
norm = _swap_direction(norm)
x = np.array(x, copy=True)
np.conjugate(x, out=x)
s, axes = _cook_nd_args(x, s, axes, invreal=True)
fsc = _compute_fwd_scale(norm, s, x.shape)

with Workers(workers):
return mkl_fft.irfftn(x, s, axes, fwd_scale=fsc)


def ihfftn(a, s=None, axes=None, norm=None, workers=None, *, plan=None):
_check_plan(plan)
x = _validate_input(a)
norm = _swap_direction(norm)
s, axes = _cook_nd_args(x, s, axes)
fsc = _compute_fwd_scale(norm, s, x.shape)

with Workers(workers):
result = mkl_fft.rfftn(x, s, axes, fwd_scale=fsc)

np.conjugate(result, out=result)
return result
Loading
Loading