diff --git a/mindscience/sciops/dft/fourier.py b/mindscience/sciops/dft/fourier.py index 81a19d4c60215ff9aaa72472af182f6c94be6fa7..70328735f4b538e047ba3922a47f97ddc3a7ba77 100644 --- a/mindscience/sciops/dft/fourier.py +++ b/mindscience/sciops/dft/fourier.py @@ -23,6 +23,8 @@ from mindspore.ops import operations as P from ...utils.check_func import check_param_no_greater, check_param_value +# pylint: disable=arguments-differ + class MyRoll(nn.Cell): ''' Custom defined roll operator to avoid bug in MindSpore ''' @@ -89,7 +91,7 @@ def convert_shape(shape): elif len(shape) == 1: n, = shape else: - raise TypeError("Only support 1D dct/dst, but got shape {}".format(shape)) + raise TypeError(f"Only support 1D dct/dst, but got shape {shape}") return n @@ -99,7 +101,7 @@ def convert_params(shape, modes, dim): ndim = len(shape) if dim is None: - dim = tuple([n - ndim for n in range(ndim)]) + dim = tuple(n - ndim for n in range(ndim)) else: dim = tuple(np.atleast_1d(dim).astype(int).tolist()) @@ -278,11 +280,11 @@ class RDFTn(_DFTn): Args: shape (tuple): The shape of the dimensions to be transformed, other dimensions need not be included. - dim (tuple): Dimensions to be transformed. Default: None, the leading dimensions will be transformed. + dim (tuple): Dimensions to be transformed. Default: None, the trailing dimensions will be transformed. norm (str): Normalization mode, should be one of 'forward', 'backward', 'ortho'. Default: 'backward', same as torch.fft.rfftn - modes (tuple, int, None): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. + modes (tuple, int, None): The length of the output transform axis. + The `modes` must be no greater than half of the dimension of input 'x'. compute_dtype (mindspore.dtype): The type of input tensor. Default: mindspore.float32. Inputs: @@ -301,7 +303,7 @@ class RDFTn(_DFTn): >>> from mindspore import ops >>> from mindflow.core import RDFTn >>> ar = ops.rand((2, 32, 512)) - >>> dft_cell = RDFTn(x.shape[-2:]) + >>> dft_cell = RDFTn(ar.shape[-2:]) >>> br, bi = dft_cell(ar) >>> print(br.shape) (2, 32, 257) @@ -320,11 +322,11 @@ class IRDFTn(_DFTn): Args: shape (tuple): The shape of the dimensions to be transformed, other dimensions need not be included. - dim (tuple): Dimensions to be transformed. Default: None, the leading dimensions will be transformed. + dim (tuple): Dimensions to be transformed. Default: None, the trailing dimensions will be transformed. norm (str): Normalization mode, should be one of 'forward', 'backward', 'ortho'. Default: 'backward', same as torch.fft.irfftn - modes (tuple, int, None): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. + modes (tuple, int, None): The length of the output transform axis. + The `modes` must be no greater than half of the dimension of input 'x'. compute_dtype (mindspore.dtype): The type of input tensor. Default: mindspore.float32. Inputs: @@ -342,10 +344,11 @@ class IRDFTn(_DFTn): Examples: >>> from mindspore import ops >>> from mindflow.core import IRDFTn + >>> full_shape = (2, 32, 512) >>> ar = ops.rand((2, 32, 257)) >>> ai = ops.rand((2, 32, 257)) - >>> dft_cell = IRDFTn(x.shape[-2:]) - >>> br = dft_cell(ar) + >>> dft_cell = IRDFTn(full_shape[-2:]) + >>> br = dft_cell(ar, ai) >>> print(br.shape) (2, 32, 512) """ @@ -372,11 +375,11 @@ class DFTn(_DFTn): Args: shape (tuple): The shape of the dimensions to be transformed, other dimensions need not be included. - dim (tuple): Dimensions to be transformed. Default: None, the leading dimensions will be transformed. + dim (tuple): Dimensions to be transformed. Default: None, the trailing dimensions will be transformed. norm (str): Normalization mode, should be one of 'forward', 'backward', 'ortho'. Default: 'backward', same as torch.fft.irfftn - modes (tuple, int, None): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. + modes (tuple, int, None): The length of the output transform axis. + The `modes` must be no greater than half of the dimension of input 'x'. compute_dtype (mindspore.dtype): The type of input tensor. Default: mindspore.float32. Inputs: @@ -395,7 +398,7 @@ class DFTn(_DFTn): >>> from mindflow.cell import DFTn >>> ar = ops.rand((2, 32, 512)) >>> ai = ops.rand((2, 32, 512)) - >>> dft_cell = DFTn(x.shape[-2:]) + >>> dft_cell = DFTn(ar.shape[-2:]) >>> br, bi = dft_cell(ar, ai) >>> print(br.shape) (2, 32, 512) @@ -423,11 +426,11 @@ class IDFTn(DFTn): Args: shape (tuple): The shape of the dimensions to be transformed, other dimensions need not be included. - dim (tuple): Dimensions to be transformed. Default: None, the leading dimensions will be transformed. + dim (tuple): Dimensions to be transformed. Default: None, the trailing dimensions will be transformed. norm (str): Normalization mode, should be one of 'forward', 'backward', 'ortho'. Default: 'backward', same as torch.fft.irfftn - modes (tuple, int, None): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. + modes (tuple, int, None): The length of the output transform axis. + The `modes` must be no greater than half of the dimension of input 'x'. compute_dtype (mindspore.dtype): The type of input tensor. Default: mindspore.float32. Inputs: @@ -446,7 +449,7 @@ class IDFTn(DFTn): >>> from mindflow.cell import DFTn >>> ar = ops.rand((2, 32, 512)) >>> ai = ops.rand((2, 32, 512)) - >>> dft_cell = DFTn(x.shape[-2:]) + >>> dft_cell = DFTn(ar.shape[-2:]) >>> br, bi = dft_cell(ar, ai) >>> print(br.shape) (2, 32, 512) @@ -486,7 +489,7 @@ class DCT(nn.Cell): >>> from mindspore import ops >>> from mindflow.cell import DCT >>> a = ops.rand((2, 32, 512)) - >>> dft_cell = DCT(x.shape[-1:]) + >>> dft_cell = DCT(a.shape[-1:]) >>> b = dft_cell(a) >>> print(b.shape) (2, 32, 512) @@ -538,7 +541,7 @@ class IDCT(nn.Cell): >>> from mindspore import ops >>> from mindflow.cell import IDCT >>> a = ops.rand((2, 32, 512)) - >>> dft_cell = IDCT(x.shape[-1:]) + >>> dft_cell = IDCT(a.shape[-1:]) >>> b = dft_cell(a) >>> print(b.shape) (2, 32, 512) @@ -602,7 +605,7 @@ class DST(nn.Cell): >>> from mindspore import ops >>> from mindflow.cell import DST >>> a = ops.rand((2, 32, 512)) - >>> dft_cell = DST(x.shape[-1:]) + >>> dft_cell = DST(a.shape[-1:]) >>> b = dft_cell(a) >>> print(b.shape) (2, 32, 512) @@ -646,7 +649,7 @@ class IDST(nn.Cell): >>> from mindspore import ops >>> from mindflow.cell import IDST >>> a = ops.rand((2, 32, 512)) - >>> dft_cell = IDST(x.shape[-1:]) + >>> dft_cell = IDST(a.shape[-1:]) >>> b = dft_cell(a) >>> print(b.shape) (2, 32, 512) diff --git a/mindscience/solvers/cbs.py b/mindscience/solvers/cbs.py index 1ad9108f54b6ae8ded4a67485e17e62479d2d546..91aace3de8bb97753b41f8ab11aa1555fab5dc01 100644 --- a/mindscience/solvers/cbs.py +++ b/mindscience/solvers/cbs.py @@ -290,7 +290,7 @@ class CBS(nn.Cell): return num_real, num_imag def den(x): - return sum([(alpha * x) ** i / float(factorial(i)) for i in range(rampup + 1)]) * factorial(rampup) + return sum((alpha * x) ** i / float(factorial(i)) for i in range(rampup + 1)) * factorial(rampup) def transform_fun(x): num_real, num_imag = num(x) @@ -310,7 +310,7 @@ class CBS(nn.Cell): # Reference: https://github.com/ucl-bug/jwave/blob/9d114ac6acade3c866e78a9984be47833119a9f9/ # jwave/acoustics/time_harmonic.py#L153 - diff *= (diff > 0).astype(ms.float32) / 4. + diff *= (diff > 0).astype(ms.float32) dist = ops.norm(diff, dim=0) k_k0_real, k_k0_imag = transform_fun(dist) diff --git a/tests/solver/test_cbs.py b/tests/solver/test_cbs.py index ca7e745ff6463639e44edf7efebe409c3bc04063..bfb130d2e40b5172bfb432ca1e6c91e30b5ef188 100644 --- a/tests/solver/test_cbs.py +++ b/tests/solver/test_cbs.py @@ -68,8 +68,8 @@ def test_solve_2d(mode): n_steps = len(errs) step_time = time_spent / n_steps - assert n_steps <= 180 - assert isclose((ur - ui).std(), 0.029885478, rel_tol=1e-3, abs_tol=1e-3) + assert n_steps <= 260 + assert isclose((ur - ui).std(), 0.029885478, rel_tol=1e-3, abs_tol=1e-3), (ur - ui).std() assert time_spent <= 60 assert step_time <= 0.5 @@ -106,8 +106,8 @@ def test_grad_2d(mode): loss, grad = grd_func(c_star, f_star) time_spent = toc() - tic - assert isclose(loss, 8.4663, rel_tol=1e-3, abs_tol=1e-3) - assert isclose(grad.std(), 0.01240166, rel_tol=1e-3, abs_tol=1e-3) + assert isclose(loss, 6.23106, rel_tol=1e-3, abs_tol=1e-3), loss + assert isclose(grad.std(), 0.0087822, rel_tol=1e-3, abs_tol=1e-3), grad.std() assert time_spent <= 30 @@ -157,7 +157,7 @@ def test_solve_3d(mode): step_time = time_spent / n_steps assert n_steps <= 180 - assert isclose((ur - ui).std(), 0.00265927, rel_tol=1e-3, abs_tol=1e-3) + assert isclose((ur - ui).std(), 0.00265927, rel_tol=1e-3, abs_tol=1e-3), (ur - ui).std() assert time_spent <= 60 assert step_time <= 0.5 @@ -194,6 +194,6 @@ def test_grad_3d(mode): loss, grad = grd_func(c_star, f_star) time_spent = toc() - tic - assert isclose(loss, 7.84422, rel_tol=1e-3, abs_tol=1e-3) - assert isclose(grad.std(), 0.00207235, rel_tol=1e-3, abs_tol=1e-3) + assert isclose(loss, 4.43072, rel_tol=1e-3, abs_tol=1e-3), loss + assert isclose(grad.std(), 0.00207235, rel_tol=1e-3, abs_tol=1e-3), grad.std() assert time_spent <= 30