1# Copyright 2024 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15"""Fast Fourier Transform operations, the function docs are adapted from Numpy API.""" 16from __future__ import absolute_import 17__all__ = ['fftshift', 'ifftshift', 'fft', 'ifft', 'fft2', 'ifft2', 'fftn', 'ifftn'] 18from mindspore import ops 19 20 21def fftshift(x, axes=None): 22 """ 23 Shift the zero-frequency component to the center of the spectrum. 24 25 Refer to :func:`mindspore.ops.fftshift` for more details. 26 The difference is that `x` corresponds to `input` and `axes` corresponds to `dim`. 27 28 Args: 29 x (Tensor): Input tensor. 30 axes (Union[int, list(int), tuple(int)], optional): Axes over which to shift. 31 Default is ``None`` , which shifts all axes. 32 33 Returns: 34 output (Tensor), the shifted tensor with the same shape and dtype as `x`. 35 36 Supported Platforms: 37 ``Ascend`` ``CPU`` 38 39 Examples: 40 >>> import mindspore.numpy as np 41 >>> from mindspore import dtype as mstype 42 >>> x = np.array([0, 1, 2, 3, 4, -5, -4, -3, -2, -1], dtype=mstype.int32) 43 >>> np.fft.fftshift(x) 44 Tensor(shape=[10], dtype=Int32, value= [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]) 45 """ 46 return ops.fftshift(x, axes) 47 48 49def ifftshift(x, axes=None): 50 """ 51 The inverse of fftshift. 52 53 Refer to :func:`mindspore.ops.ifftshift` for more details. 54 The difference is that `x` corresponds to `input` and `axes` corresponds to `dim`. 55 56 Args: 57 x (Tensor): Input tensor. 58 axes (Union[int, list(int), tuple(int)], optional): Axes over which to shift. 59 Default is ``None`` , which shifts all axes. 60 61 Returns: 62 output (Tensor), the shifted tensor with the same shape and dtype as `x`. 63 64 Supported Platforms: 65 ``Ascend`` ``CPU`` 66 67 Examples: 68 >>> import mindspore.numpy as np 69 >>> from mindspore import dtype as mstype 70 >>> x = np.array([0, 1, 2, 3, 4, -5, -4, -3, -2, -1], dtype=mstype.int32) 71 >>> np.fft.ifftshift(np.fft.fftshift(x)) 72 Tensor(shape=[10], dtype=Int32, value= [ 0, 1, 2, 3, 4, -5, -4, -3, -2, -1]) 73 """ 74 return ops.ifftshift(x, axes) 75 76 77def fft(a, n=None, axis=-1, norm=None): 78 r""" 79 Calculates the one dimensional discrete Fourier transform of `a`. 80 81 Refer to :func:`mindspore.ops.fft` for more details. 82 The difference is that `a` corresponds to `input` and `axis` corresponds to `dim`. 83 84 Args: 85 a (Tensor): The input tensor. 86 Supported dtypes: 87 88 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 89 90 n (int, optional): Length of the transformed `dim` of the result. 91 If given, the size of the `dim` axis will be zero-padded or truncated to `n` before calculating `fft`. 92 Default: ``None`` , which does not need to process `a`. 93 axis (int, optional): Axis over which to compute the `fft`. 94 Default: ``-1`` , which means the last axis of `a` is used. 95 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 96 Three modes are defined as, 97 98 - ``"backward"`` (no normalization). 99 - ``"forward"`` (normalize by :math:`1/n`). 100 - ``"ortho"`` (normalize by :math:`1/\sqrt{n}`). 101 102 Returns: 103 Tensor, The result of `fft()` function. The default is the same shape as `a`. 104 If `n` is given, the size of the `axis` is changed to `n`. 105 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 106 When the `a` is float64 or complex128, the return value type is complex128. 107 108 Supported Platforms: 109 ``Ascend`` ``CPU`` 110 111 Examples: 112 >>> import mindspore.numpy as np 113 >>> input = np.array([ 1.6243454, -0.6117564, -0.5281718, -1.0729686]) 114 >>> np.fft.fft(input) 115 Tensor(shape=[4], dtype=Complex64, value= [-0.588551+0j, 2.15252-0.461212j, 2.7809+0j, 2.15252+0.461212j]) 116 """ 117 return ops.fft(a, n, axis, norm) 118 119 120def ifft(a, n=None, axis=-1, norm=None): 121 r""" 122 Calculates the inverse of `fft()`. 123 124 Refer to :func:`mindspore.ops.ifft` for more details. 125 The difference is that `a` corresponds to `input` and `axis` corresponds to `dim`. 126 127 Args: 128 a (Tensor): The input tensor. 129 Supported dtypes: 130 131 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 132 133 n (int, optional): Length of the transformed `dim` of the result. 134 n (int, optional): Signal length. 135 If given, the input will either be zero-padded or trimmed to this length before computing `ifft`. 136 Default: ``None`` , which does not need to process `a`. 137 axis (int, optional): Axis over which to compute the `ifft`. 138 Default: ``-1`` , which means the last axis of `a` is used. 139 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 140 Three modes are defined as, 141 142 - ``"backward"`` (no normalization). 143 - ``"forward"`` (normalize by :math:`1*n`). 144 - ``"ortho"`` (normalize by :math:`1*\sqrt{n}`). 145 146 Returns: 147 Tensor, The result of `ifft()` function. The default is the same shape as `a`. 148 If `n` is given, the size of the `axis` is changed to `n`. 149 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 150 When the `a` is float64 or complex128, the return value type is complex128. 151 152 Supported Platforms: 153 ``Ascend`` ``CPU`` 154 155 Examples: 156 >>> import mindspore.numpy as np 157 >>> input = np.array([ 1.6243454, -0.6117564, -0.5281718, -1.0729686]) 158 >>> np.fft.ifft(input) 159 Tensor(shape=[4], dtype=Complex64, value= [-0.147138+0j, 0.538129+0.115303j, 0.695225+0j, 0.538129-0.115303j]) 160 """ 161 return ops.ifft(a, n, axis, norm) 162 163 164def rfft(a, n=None, axis=-1, norm=None): 165 r""" 166 Calculates the one dimensional discrete Fourier transform for real input `a`. 167 168 Refer to :func:`mindspore.ops.rfft` for more details. 169 The difference is that `a` corresponds to `input` and `axis` corresponds to `dim`. 170 171 Args: 172 a (Tensor): The input tensor. 173 n (int, optional): Number of points along `axis` in the input to use. 174 If given, the input will either be zero-padded or trimmed to this length before computing `rfft`. 175 Default: ``None``. 176 axis (int, optional): Axis over which to compute the `rfft`. 177 Default: ``-1``, which means the last axis of `a` is used. 178 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"``. 179 Three modes are defined as, 180 181 - ``"backward"`` (no normalization). 182 - ``"forward"`` (normalize by :math:`1/n`). 183 - ``"ortho"`` (normalize by :math:`1/\sqrt{n}`). 184 185 Returns: 186 Tensor, the result of `rfft()` function. 187 188 Supported Platforms: 189 ``Ascend`` ``CPU`` 190 191 Examples: 192 >>> import mindspore 193 >>> from mindspore import Tensor 194 >>> from mindspore import numpy as mnp 195 >>> input = Tensor([1, 2, 3, 4]) 196 >>> y = mnp.fft.rfft(input) 197 >>> print(y) 198 [10.+0.j -2.+2.j -2.+0.j] 199 """ 200 return ops.rfft(a, n, axis, norm) 201 202 203def irfft(a, n=None, axis=-1, norm=None): 204 r""" 205 Calculates the inverse of `rfft()`. 206 207 Refer to :func:`mindspore.ops.irfft` for more details. 208 The difference is that `a` corresponds to `input` and `axis` corresponds to `dim`. 209 210 Args: 211 a (Tensor): The input tensor. 212 n (int, optional): Length of the transformed `dim` of the result. 213 If given, the input will either be zero-padded or trimmed to this length before computing `rfft`. 214 If n is not given, it is taken to be :math:`2*(a.shape[axis]-1)`. 215 Default: ``None``. 216 axis (int, optional): Axis over which to compute the `irfft`. 217 Default: ``-1``, which means the last axis of `a` is used. 218 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"``. 219 Three modes are defined as, 220 221 - ``"backward"`` (no normalization). 222 - ``"forward"`` (normalize by :math:`1/n`). 223 - ``"ortho"`` (normalize by :math:`1/\sqrt{n}`). 224 225 Returns: 226 Tensor, the result of `irfft()` function. 227 228 Supported Platforms: 229 ``Ascend`` ``CPU`` 230 231 Examples: 232 >>> import mindspore 233 >>> from mindspore import Tensor 234 >>> from mindspore import numpy as mnp 235 >>> input = Tensor([1, 2, 3, 4]) 236 >>> y = mnp.fft.irfft(input) 237 >>> print(y) 238 [ 2.5000000e+00 -6.6666669e-01 1.2590267e-15 -1.6666667e-01 239 4.2470195e-16 -6.6666669e-01] 240 """ 241 return ops.irfft(a, n, axis, norm) 242 243 244def fft2(a, s=None, axes=(-2, -1), norm=None): 245 r""" 246 Calculates the two dimensional discrete Fourier transform of `a`. 247 248 Refer to :func:`mindspore.ops.fft2` for more details. 249 The difference is that `a` corresponds to `input` and `axes` corresponds to `dim`. 250 251 Args: 252 a (Tensor): The input tensor. 253 Supported dtypes: 254 255 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 256 257 s (tuple[int], optional): Length of the transformed `axes` of the result. 258 If given, the input will either be zero-padded or trimmed to this length before computing `fft2`. 259 Default: ``None`` , which does not need to process `a`. 260 axes (tuple[int], optional): The dimension along which to take the one dimensional `fft2`. 261 Default: ``(-2, -1)`` , which means transform the last two dimension of `a`. 262 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 263 Three modes are defined as, 264 265 - ``"backward"`` (no normalization). 266 - ``"forward"`` (normalize by :math:`1/n`). 267 - ``"ortho"`` (normalize by :math:`1/\sqrt{n}`). 268 269 Returns: 270 Tensor, The result of `fft2()` function. The default is the same shape as `a`. 271 If `s` is given, the size of the `axes[i]` axis is changed to `s[i]`. 272 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 273 When the `a` is float64 or complex128, the return value type is complex128. 274 275 Supported Platforms: 276 ``Ascend`` ``CPU`` 277 278 Examples: 279 >>> import mindspore.numpy as np 280 >>> a = np.ones((4, 4)) 281 >>> np.fft.fft2(a, s=(4, 4), axes=(0, 1), norm="backward") 282 Tensor(shape=[4, 4], dtype=Complex64, value= 283 [[16+0j, 0+0j, 0+0j, 0+0j], 284 [0+0j, 0+0j, 0+0j, 0+0j], 285 [0+0j, 0+0j, 0+0j, 0+0j], 286 [0+0j, 0+0j, 0+0j, 0+0j]]) 287 """ 288 return ops.fft2(a, s, axes, norm) 289 290 291def ifft2(a, s=None, axes=(-2, -1), norm=None): 292 r""" 293 Calculates the inverse of `fft2()`. 294 295 Refer to :func:`mindspore.ops.ifft2` for more details. 296 The difference is that `a` corresponds to `input` and `axes` corresponds to `dim`. 297 298 Args: 299 a (Tensor): The input tensor. 300 Supported dtypes: 301 302 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 303 304 s (tuple[int], optional): Length of the transformed `axes` of the result. 305 If given, the input will either be zero-padded or trimmed to this length before computing `ifft2`. 306 Default: ``None`` , which does not need to process `a`. 307 axes (tuple[int], optional): The dimension along which to take the one dimensional `ifft2`. 308 Default: ``(-2, -1)`` , which means transform the last two dimension of `a`. 309 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 310 Three modes are defined as, 311 312 - ``"backward"`` (no normalization). 313 - ``"forward"`` (normalize by :math:`1*n`). 314 - ``"ortho"`` (normalize by :math:`1*\sqrt{n}`). 315 316 Returns: 317 Tensor, The result of `ifft2()` function. The default is the same shape as `a`. 318 If `s` is given, the size of the `axes[i]` axis is changed to `s[i]`. 319 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 320 When the `a` is float64 or complex128, the return value type is complex128. 321 322 Supported Platforms: 323 ``Ascend`` ``CPU`` 324 325 Examples: 326 >>> import mindspore.numpy as np 327 >>> a = np.ones((4, 4)) 328 >>> np.fft.ifft2(a, s=(4, 4), axes=(0, 1), norm="backward") 329 Tensor(shape=[4, 4], dtype=Complex64, value= 330 [[1+0j, 0+0j, 0+0j, 0+0j], 331 [0+0j, 0+0j, 0+0j, 0+0j], 332 [0+0j, 0+0j, 0+0j, 0+0j], 333 [0+0j, 0+0j, 0+0j, 0+0j]]) 334 """ 335 return ops.ifft2(a, s, axes, norm) 336 337 338def fftn(a, s=None, axes=None, norm=None): 339 r""" 340 Calculates the N dimensional discrete Fourier transform of `a`. 341 342 Refer to :func:`mindspore.ops.fftn` for more details. 343 The difference is that `a` corresponds to `input` and `axis` corresponds to `dim`. 344 345 Args: 346 a (Tensor): The input tensor. 347 Supported dtypes: 348 349 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 350 351 s (tuple[int], optional): Length of the transformed `axes` of the result. 352 If given, the input will either be zero-padded or trimmed to this length before computing `fftn`. 353 Default: ``None`` , which does not need to process `a`. 354 axes (tuple[int], optional): The dimension along which to take the one dimensional `fftn`. 355 Default: ``None`` , which means transform the all dimension of `a`, 356 or the last `len(s)` dimensions if s is given. 357 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 358 Three modes are defined as, 359 360 - ``"backward"`` (no normalization). 361 - ``"forward"`` (normalize by :math:`1/n`). 362 - ``"ortho"`` (normalize by :math:`1/\sqrt{n}`). 363 364 Returns: 365 Tensor, The result of `fft()` function. The default is the same shape as `a`. 366 If `s` is given, the size of the `axes[i]` axis is changed to `s[i]`. 367 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 368 When the `a` is float64 or complex128, the return value type is complex128. 369 370 Supported Platforms: 371 ``Ascend`` ``CPU`` 372 373 Examples: 374 >>> import mindspore.numpy as np 375 >>> a = np.ones((2, 2, 2)) 376 >>> np.fft.fftn(a, s=(2, 2, 2), axes=(0, 1, 2), norm="backward") 377 Tensor(shape=[2, 2, 2], dtype=Complex64, value= 378 [[[8+0j, 0+0j], 379 [0+0j, 0+0j]], 380 [[0+0j, 0+0j], 381 [0+0j, 0+0j]]]) 382 """ 383 return ops.fftn(a, s, axes, norm) 384 385 386def ifftn(a, s=None, axes=None, norm=None): 387 r""" 388 Calculates the inverse of `fftn()`. 389 390 Refer to :func:`mindspore.ops.ifftn` for more details. 391 The difference is that `a` corresponds to `input` and `axes` corresponds to `dim`. 392 393 Args: 394 a (Tensor): The input tensor. 395 Supported dtypes: 396 397 - Ascend/CPU: int16, int32, int64, float16, float32, float64, complex64, complex128. 398 399 s (tuple[int], optional): Length of the transformed `axes` of the result. 400 If given, the input will either be zero-padded or trimmed to this length before computing `ifftn`. 401 Default: ``None`` , which does not need to process `a`. 402 axes (tuple[int], optional): The dimension along which to take the one dimensional `ifftn`. 403 Default: ``None`` , which means transform the all dimension of `a`, 404 or the last `len(s)` dimensions if s is given. 405 norm (string, optional): Normalization mode. Default: ``None`` that means ``"backward"`` . 406 Three modes are defined as, 407 408 - ``"backward"`` (no normalization). 409 - ``"forward"`` (normalize by :math:`1*n`). 410 - ``"ortho"`` (normalize by :math:`1*\sqrt{n}`). 411 412 Returns: 413 Tensor, The result of `ifftn()` function. The default is the same shape as `a`. 414 If `s` is given, the size of the `axes[i]` axis is changed to `s[i]`. 415 When the `a` is int16, int32, int64, float16, float32, complex64, the return value type is complex64. 416 When the `a` is float64 or complex128, the return value type is complex128. 417 418 Supported Platforms: 419 ``Ascend`` ``CPU`` 420 421 Examples: 422 >>> import mindspore.numpy as np 423 >>> a = np.ones((2, 2, 2)) 424 >>> np.fft.ifftn(a, s=(2, 2, 2), axes=(0, 1, 2), norm="backward") 425 Tensor(shape=[2, 2, 2], dtype=Complex64, value= 426 [[[1+0j, 0+0j], 427 [0+0j, 0+0j]], 428 [[0+0j, 0+0j], 429 [0+0j, 0+0j]]]) 430 """ 431 return ops.ifftn(a, s, axes, norm) 432