33
33
import numpy as np
34
34
35
35
from . import _pydfti as mkl_fft # pylint: disable=no-name-in-module
36
- from ._fft_utils import _compute_fwd_scale
36
+ from ._fft_utils import _compute_fwd_scale , _swap_direction
37
37
from ._float_utils import _supported_array_or_not_implemented
38
38
39
39
__doc__ = """
@@ -125,6 +125,12 @@ def set_workers(n_workers):
125
125
"irfft2" ,
126
126
"rfftn" ,
127
127
"irfftn" ,
128
+ "hfft" ,
129
+ "ihfft" ,
130
+ "hfft2" ,
131
+ "ihfft2" ,
132
+ "hfftn" ,
133
+ "ihfftn" ,
128
134
"get_workers" ,
129
135
"set_workers" ,
130
136
"DftiBackend" ,
@@ -231,7 +237,7 @@ def _validate_input(a):
231
237
232
238
233
239
def fft (
234
- a , n = None , axis = - 1 , norm = None , overwrite_x = False , workers = None , plan = None
240
+ a , n = None , axis = - 1 , norm = None , overwrite_x = False , workers = None , * , plan = None
235
241
):
236
242
_check_plan (plan )
237
243
x = _validate_input (a )
@@ -244,7 +250,7 @@ def fft(
244
250
245
251
246
252
def ifft (
247
- a , n = None , axis = - 1 , norm = None , overwrite_x = False , workers = None , plan = None
253
+ a , n = None , axis = - 1 , norm = None , overwrite_x = False , workers = None , * , plan = None
248
254
):
249
255
_check_plan (plan )
250
256
x = _validate_input (a )
@@ -263,6 +269,7 @@ def fft2(
263
269
norm = None ,
264
270
overwrite_x = False ,
265
271
workers = None ,
272
+ * ,
266
273
plan = None ,
267
274
):
268
275
@@ -284,6 +291,7 @@ def ifft2(
284
291
norm = None ,
285
292
overwrite_x = False ,
286
293
workers = None ,
294
+ * ,
287
295
plan = None ,
288
296
):
289
297
@@ -299,7 +307,14 @@ def ifft2(
299
307
300
308
301
309
def fftn (
302
- a , s = None , axes = None , norm = None , overwrite_x = False , workers = None , plan = None
310
+ a ,
311
+ s = None ,
312
+ axes = None ,
313
+ norm = None ,
314
+ overwrite_x = False ,
315
+ workers = None ,
316
+ * ,
317
+ plan = None ,
303
318
):
304
319
_check_plan (plan )
305
320
x = _validate_input (a )
@@ -312,7 +327,14 @@ def fftn(
312
327
313
328
314
329
def ifftn (
315
- a , s = None , axes = None , norm = None , overwrite_x = False , workers = None , plan = None
330
+ a ,
331
+ s = None ,
332
+ axes = None ,
333
+ norm = None ,
334
+ overwrite_x = False ,
335
+ workers = None ,
336
+ * ,
337
+ plan = None ,
316
338
):
317
339
_check_plan (plan )
318
340
x = _validate_input (a )
@@ -324,7 +346,7 @@ def ifftn(
324
346
)
325
347
326
348
327
- def rfft (a , n = None , axis = - 1 , norm = None , workers = None , plan = None ):
349
+ def rfft (a , n = None , axis = - 1 , norm = None , workers = None , * , plan = None ):
328
350
_check_plan (plan )
329
351
x = _validate_input (a )
330
352
fsc = _compute_fwd_scale (norm , n , x .shape [axis ])
@@ -333,7 +355,7 @@ def rfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
333
355
return mkl_fft .rfft (x , n = n , axis = axis , fwd_scale = fsc )
334
356
335
357
336
- def irfft (a , n = None , axis = - 1 , norm = None , workers = None , plan = None ):
358
+ def irfft (a , n = None , axis = - 1 , norm = None , workers = None , * , plan = None ):
337
359
_check_plan (plan )
338
360
x = _validate_input (a )
339
361
fsc = _compute_fwd_scale (norm , n , 2 * (x .shape [axis ] - 1 ))
@@ -342,17 +364,15 @@ def irfft(a, n=None, axis=-1, norm=None, workers=None, plan=None):
342
364
return mkl_fft .irfft (x , n = n , axis = axis , fwd_scale = fsc )
343
365
344
366
345
- def rfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , plan = None ):
346
-
367
+ def rfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , * , plan = None ):
347
368
return rfftn (a , s = s , axes = axes , norm = norm , workers = workers , plan = plan )
348
369
349
370
350
- def irfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , plan = None ):
351
-
371
+ def irfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , * , plan = None ):
352
372
return irfftn (a , s = s , axes = axes , norm = norm , workers = workers , plan = plan )
353
373
354
374
355
- def rfftn (a , s = None , axes = None , norm = None , workers = None , plan = None ):
375
+ def rfftn (a , s = None , axes = None , norm = None , workers = None , * , plan = None ):
356
376
_check_plan (plan )
357
377
x = _validate_input (a )
358
378
s , axes = _cook_nd_args (x , s , axes )
@@ -362,11 +382,71 @@ def rfftn(a, s=None, axes=None, norm=None, workers=None, plan=None):
362
382
return mkl_fft .rfftn (x , s , axes , fwd_scale = fsc )
363
383
364
384
365
- def irfftn (a , s = None , axes = None , norm = None , workers = None , plan = None ):
385
+ def irfftn (a , s = None , axes = None , norm = None , workers = None , * , plan = None ):
386
+ _check_plan (plan )
387
+ x = _validate_input (a )
388
+ s , axes = _cook_nd_args (x , s , axes , invreal = True )
389
+ fsc = _compute_fwd_scale (norm , s , x .shape )
390
+
391
+ with Workers (workers ):
392
+ return mkl_fft .irfftn (x , s , axes , fwd_scale = fsc )
393
+
394
+
395
+ def hfft (a , n = None , axis = - 1 , norm = None , workers = None , * , plan = None ):
396
+ _check_plan (plan )
397
+ x = _validate_input (a )
398
+ norm = _swap_direction (norm )
399
+ x = np .array (x , copy = True )
400
+ np .conjugate (x , out = x )
401
+ fsc = _compute_fwd_scale (norm , n , 2 * (x .shape [axis ] - 1 ))
402
+
403
+ with Workers (workers ):
404
+ return mkl_fft .irfft (x , n = n , axis = axis , fwd_scale = fsc )
405
+
406
+
407
+ def ihfft (a , n = None , axis = - 1 , norm = None , workers = None , * , plan = None ):
408
+ _check_plan (plan )
409
+ x = _validate_input (a )
410
+ norm = _swap_direction (norm )
411
+ fsc = _compute_fwd_scale (norm , n , x .shape [axis ])
412
+
413
+ with Workers (workers ):
414
+ result = mkl_fft .rfft (x , n = n , axis = axis , fwd_scale = fsc )
415
+
416
+ np .conjugate (result , out = result )
417
+ return result
418
+
419
+
420
+ def hfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , * , plan = None ):
421
+ return hfftn (a , s = s , axes = axes , norm = norm , workers = workers , plan = plan )
422
+
423
+
424
+ def ihfft2 (a , s = None , axes = (- 2 , - 1 ), norm = None , workers = None , * , plan = None ):
425
+ return ihfftn (a , s = s , axes = axes , norm = norm , workers = workers , plan = plan )
426
+
427
+
428
+ def hfftn (a , s = None , axes = None , norm = None , workers = None , * , plan = None ):
366
429
_check_plan (plan )
367
430
x = _validate_input (a )
431
+ norm = _swap_direction (norm )
432
+ x = np .array (x , copy = True )
433
+ np .conjugate (x , out = x )
368
434
s , axes = _cook_nd_args (x , s , axes , invreal = True )
369
435
fsc = _compute_fwd_scale (norm , s , x .shape )
370
436
371
437
with Workers (workers ):
372
438
return mkl_fft .irfftn (x , s , axes , fwd_scale = fsc )
439
+
440
+
441
+ def ihfftn (a , s = None , axes = None , norm = None , workers = None , * , plan = None ):
442
+ _check_plan (plan )
443
+ x = _validate_input (a )
444
+ norm = _swap_direction (norm )
445
+ s , axes = _cook_nd_args (x , s , axes )
446
+ fsc = _compute_fwd_scale (norm , s , x .shape )
447
+
448
+ with Workers (workers ):
449
+ result = mkl_fft .rfftn (x , s , axes , fwd_scale = fsc )
450
+
451
+ np .conjugate (result , out = result )
452
+ return result
0 commit comments