From 051addffefa67d46e3050535c543b444018b4b73 Mon Sep 17 00:00:00 2001 From: lizhihao Date: Tue, 12 Sep 2023 10:08:54 +0800 Subject: [PATCH] [SCIAI] add license for all py files and add release --- SciAI/NOTICE | 2 +- SciAI/RELEASE.md | 17 + SciAI/RELEASE_CN.md | 17 + SciAI/sciai/__init__.py | 40 +- SciAI/sciai/architecture/__init__.py | 14 + SciAI/sciai/architecture/activation.py | 14 + SciAI/sciai/architecture/basic_block.py | 14 + .../architecture/neural_operators/__init__.py | 32 +- .../architecture/neural_operators/dft.py | 1314 +++++++++-------- .../architecture/neural_operators/fno1d.py | 262 ++-- .../architecture/neural_operators/fno2d.py | 268 ++-- .../architecture/neural_operators/fno3d.py | 274 ++-- .../architecture/neural_operators/kno1d.py | 212 +-- .../architecture/neural_operators/kno2d.py | 14 + .../architecture/neural_operators/m2k.py | 14 + .../architecture/neural_operators/pdenet.py | 458 +++--- .../architecture/transformer/__init__.py | 14 + SciAI/sciai/architecture/transformer/layer.py | 14 + SciAI/sciai/architecture/transformer/vit.py | 14 + SciAI/sciai/common/__init__.py | 14 + SciAI/sciai/common/dataset.py | 14 + SciAI/sciai/common/initializer.py | 14 + SciAI/sciai/common/optimizer.py | 14 + SciAI/sciai/common/train_cell.py | 14 + SciAI/sciai/context/__init__.py | 14 + SciAI/sciai/context/context.py | 14 + SciAI/sciai/model/phygeonet/config.yaml | 2 +- SciAI/sciai/model/phygeonet/src/plot.py | 3 + SciAI/sciai/operators/__init__.py | 14 + SciAI/sciai/operators/derivatives.py | 14 + SciAI/sciai/operators/jacobian_weights.py | 14 + SciAI/sciai/utils/__init__.py | 14 + SciAI/sciai/utils/check_utils.py | 14 + SciAI/sciai/utils/file_utils.py | 14 + SciAI/sciai/utils/log_utils.py | 14 + SciAI/sciai/utils/math_utils.py | 14 + SciAI/sciai/utils/ms_utils.py | 14 + SciAI/sciai/utils/plot_utils.py | 14 + SciAI/sciai/utils/python_utils.py | 14 + SciAI/sciai/utils/register_utils.py | 122 +- SciAI/sciai/utils/time_utils.py | 14 + SciAI/sciai/version.py | 20 +- SciAI/setup.py | 4 +- SciAI/tutorial/example_grad_net.py | 194 +-- SciAI/tutorial/example_load_ckpt.py | 94 +- SciAI/tutorial/example_net.py | 170 +-- .../st/sciai/architecture/test_activation.py | 2 +- .../st/sciai/architecture/test_basic_block.py | 2 +- tests/st/sciai/common/test_dataset.py | 2 +- tests/st/sciai/common/test_initializer.py | 2 +- tests/st/sciai/common/test_optimizer.py | 2 +- tests/st/sciai/common/test_train_cell.py | 2 +- tests/st/sciai/context/test_context.py | 2 +- tests/st/sciai/operators/test_derivatives.py | 2 +- tests/st/sciai/operators/test_jacobian.py | 2 +- tests/st/sciai/test_utils/basic_nets.py | 2 +- tests/st/sciai/test_utils/func_utils.py | 2 +- tests/st/sciai/test_utils/test_base.py | 2 +- tests/st/sciai/utils/test_check_utils.py | 2 +- tests/st/sciai/utils/test_log_utils.py | 2 +- tests/st/sciai/utils/test_ms_utils.py | 2 +- tests/st/sciai/utils/test_plot_utils.py | 2 +- tests/st/sciai/utils/test_python_utils.py | 2 +- tests/st/sciai/utils/test_time_utils.py | 2 +- 64 files changed, 2237 insertions(+), 1682 deletions(-) create mode 100644 SciAI/RELEASE.md create mode 100644 SciAI/RELEASE_CN.md diff --git a/SciAI/NOTICE b/SciAI/NOTICE index 79b7aa2c8..a853eac93 100644 --- a/SciAI/NOTICE +++ b/SciAI/NOTICE @@ -1,3 +1,3 @@ MindSpore SciAI -Copyright 2019-2022 Huawei Technologies Co., Ltd +Copyright 2019-2023 Huawei Technologies Co., Ltd diff --git a/SciAI/RELEASE.md b/SciAI/RELEASE.md new file mode 100644 index 000000000..5281dc2c7 --- /dev/null +++ b/SciAI/RELEASE.md @@ -0,0 +1,17 @@ +# MindSpore SciAI 0.1.0 Release Notes + +Initial release of MindSpore SciAI. + +[查看中文](./RELEASE_CN.md) + +### Major Feature + +- [STABLE] [`High-frequency models library`](./sciai/model): Built-in 60+ high-frequency models which ranks No.1 in the world in terms of coverage. High-frequency models cover from physics-informed (such as PINNs, DeepRitz and PFNN, etc.) and neural operators (such as FNO, DeepONet). +- [STABLE] [`High-level API`](./README.md#Quick-Start): Developers and users with high-level APIs (for example AutoModel) allow an immediate deployment. +- [STABLE] [`Basic API`](./sciai): Provides basic APIs (such as Fourier neural operators, adaptive activation functions and high-order differentiation, etc.), which are convenient for users to build, train and inference neural network models of AI4SCI. + +### Contributors + +Thanks goes to these wonderful people: + +yufan, wangzidong, yangkang, lujiale, lizhihao, huyintong, zhangrenyuan, mazhiming. \ No newline at end of file diff --git a/SciAI/RELEASE_CN.md b/SciAI/RELEASE_CN.md new file mode 100644 index 000000000..6a17d9cc4 --- /dev/null +++ b/SciAI/RELEASE_CN.md @@ -0,0 +1,17 @@ +# MindSpore SciAI 0.1.0 Release Notes + +Initial release of MindSpore SciAI. + +[View English](./RELEASE.md) + +### 主要特性 + +- [STABLE] [`高频模型库`](./sciai/model): 内置了60+高频模型，覆盖度全球第一。高频模型如物理感知（如PINNs、DeepRitz以及PFNN）和神经算子（如FNO、DeepONet）等主流模型。 +- [STABLE] [`高阶API`](./README_CN.md#快速使用): 开发者和用户通过调用高阶API（如AutoModel），实现开箱即用。 +- [STABLE] [`基础API`](./sciai): 提供了基础API（如傅里叶神经算子、自适应激活函数、高阶微分等），便于用户搭建、训练和推理AI4SCI神经网络模型。 + +### 贡献者 + +感谢以下开发者做出的贡献： + +yufan, wangzidong, yangkang, lujiale, lizhihao, huyintong, zhangrenyuan, mazhiming. \ No newline at end of file diff --git a/SciAI/sciai/__init__.py b/SciAI/sciai/__init__.py index 2e0f00fbe..ea3a68c60 100644 --- a/SciAI/sciai/__init__.py +++ b/SciAI/sciai/__init__.py @@ -1,13 +1,27 @@ -"""sciai modules""" -from sciai import operators, architecture, common, context, utils -from sciai.version import __version__ -from sciai.model import AutoModel - -__all__ = [] -__all__.extend(__version__) -__all__.extend(architecture.__all__) -__all__.extend(common.__all__) -__all__.extend(context.__all__) -__all__.extend(operators.__all__) -__all__.extend(utils.__all__) -__all__.extend(["AutoModel"]) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""sciai modules""" +from sciai import operators, architecture, common, context, utils +from sciai.version import __version__ +from sciai.model import AutoModel + +__all__ = [] +__all__.extend(__version__) +__all__.extend(architecture.__all__) +__all__.extend(common.__all__) +__all__.extend(context.__all__) +__all__.extend(operators.__all__) +__all__.extend(utils.__all__) +__all__.extend(["AutoModel"]) diff --git a/SciAI/sciai/architecture/__init__.py b/SciAI/sciai/architecture/__init__.py index 2363fad31..fcb00610b 100644 --- a/SciAI/sciai/architecture/__init__.py +++ b/SciAI/sciai/architecture/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """sciai architecture""" from .activation import Swish, SReLU, get_activation, AdaptActivation from .basic_block import MLP, MLPAAF, MLPShortcut, MSE, SSE, FirstOutputCell, NoArgNet, Normalize diff --git a/SciAI/sciai/architecture/activation.py b/SciAI/sciai/architecture/activation.py index cca99269c..1325b59b2 100644 --- a/SciAI/sciai/architecture/activation.py +++ b/SciAI/sciai/architecture/activation.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """activation""" import copy from numbers import Number diff --git a/SciAI/sciai/architecture/basic_block.py b/SciAI/sciai/architecture/basic_block.py index f386df77a..daba34f40 100644 --- a/SciAI/sciai/architecture/basic_block.py +++ b/SciAI/sciai/architecture/basic_block.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """basic block""" from numbers import Number from types import FunctionType diff --git a/SciAI/sciai/architecture/neural_operators/__init__.py b/SciAI/sciai/architecture/neural_operators/__init__.py index 02a7d6db5..a8d4526bb 100644 --- a/SciAI/sciai/architecture/neural_operators/__init__.py +++ b/SciAI/sciai/architecture/neural_operators/__init__.py @@ -1,9 +1,23 @@ -"""neural operators""" -from .fno1d import FNO1D -from .fno2d import FNO2D -from .fno3d import FNO3D -from .kno1d import KNO1D -from .kno2d import KNO2D -from .pdenet import PDENet - -__all__ = ["FNO1D", "FNO2D", "FNO3D", "KNO1D", "KNO2D", "PDENet"] +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""neural operators""" +from .fno1d import FNO1D +from .fno2d import FNO2D +from .fno3d import FNO3D +from .kno1d import KNO1D +from .kno2d import KNO2D +from .pdenet import PDENet + +__all__ = ["FNO1D", "FNO2D", "FNO3D", "KNO1D", "KNO2D", "PDENet"] diff --git a/SciAI/sciai/architecture/neural_operators/dft.py b/SciAI/sciai/architecture/neural_operators/dft.py index 7b250ff68..b4b814f26 100644 --- a/SciAI/sciai/architecture/neural_operators/dft.py +++ b/SciAI/sciai/architecture/neural_operators/dft.py @@ -1,650 +1,664 @@ -"""dft""" -import mindspore as ms -import mindspore.numpy as mnp -from mindspore import nn, ops, Tensor, Parameter -import numpy as np -from scipy.linalg import dft - -from sciai.utils.check_utils import _check_param_no_greater, _check_value_in, _check_type, _check_param_even - - -class DFT1d(nn.Cell): - """One dimensional Discrete Fourier Transformation""" - - def __init__(self, n, modes, last_index, idx=0, inv=False, dtype=ms.float32): - super().__init__() - self.n = n - self.dft_mat = dft(n, scale="sqrtn") - self.modes = modes - self.last_index = last_index - self.inv = inv - self.idx = idx - - self.dft_mode_mat_upper = self.dft_mat[:, :modes] - self.a_re_upper = Tensor(self.dft_mode_mat_upper.real, dtype=dtype) - self.a_im_upper = Tensor(self.dft_mode_mat_upper.imag, dtype=dtype) - - self.dft_mode_mat_lower = self.dft_mat[:, -modes:] - self.a_re_lower = Tensor(self.dft_mode_mat_lower.real, dtype=dtype) - self.a_im_lower = Tensor(self.dft_mode_mat_lower.imag, dtype=dtype) - self.concat = ops.Concat(axis=-1) - - if self.inv: - self.a_re_upper = self.a_re_upper.T - self.a_im_upper = -self.a_im_upper.T - if last_index: - if modes == n // 2 + 1: - self.dft_mat_res = self.dft_mat[:, -modes + 2:] - else: - self.dft_mat_res = self.dft_mat[:, -modes + 1:] - - mat = Tensor(np.zeros(n), dtype=dtype).reshape(n, 1) - self.a_re_res = mnp.flip(Tensor(self.dft_mat_res.real, dtype=dtype), axis=-1) - self.a_im_res = mnp.flip(Tensor(self.dft_mat_res.imag, dtype=dtype), axis=-1) - if modes == n // 2 + 1: - self.a_re_res = self.concat((mat, self.a_re_res, mat)) - self.a_im_res = self.concat((mat, self.a_im_res, mat)) - else: - self.a_re_res = self.concat((mat, self.a_re_res)) - self.a_im_res = self.concat((mat, self.a_im_res)) - - self.a_re_res = self.a_re_res.T - self.a_im_res = -self.a_im_res.T - else: - self.a_re_res = self.a_re_lower.T - self.a_im_res = -self.a_im_lower.T - - if (self.n - 2 * self.modes) > 0: - self.mat = ops.zeros((self.n - 2 * self.modes), dtype) - - def swap_axes(self, x_re, x_im): - return x_re.swapaxes(-1, self.idx), x_im.swapaxes(-1, self.idx) - - def complex_matmul(self, x_re, x_im, a_re, a_im): - y_re = ops.matmul(x_re, a_re) - ops.matmul(x_im, a_im) - y_im = ops.matmul(x_im, a_re) + ops.matmul(x_re, a_im) - return y_re, y_im - - def construct(self, x): - """construct""" - x_re, x_im = x - - if not self.inv: - x_re, x_im = self.swap_axes(x_re, x_im) - y_re, y_im = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_upper, a_im=self.a_im_upper) - - if not self.last_index: - y_re2, y_im2 = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_lower, a_im=self.a_im_lower) - - if self.n == self.modes * 2: - y_re = self.concat((y_re, y_re2)) - y_im = self.concat((y_im, y_im2)) - else: - dims = x_re.shape[:-1] - length = len(dims) - mat = self.mat - for i in range(length - 1, -1, -1): - mat = mat.expand_dims(0).repeat(dims[i], 0) - y_re = self.concat((y_re, mat, y_re2)) - y_im = self.concat((y_im, mat, y_im2)) - - y_re, y_im = self.swap_axes(y_re, y_im) - return y_re, y_im - - x_re, x_im = self.swap_axes(x_re, x_im) - y_re, y_im = self.complex_matmul(x_re=x_re[..., :self.modes], x_im=x_im[..., :self.modes], - a_re=self.a_re_upper, a_im=self.a_im_upper) - y_re, y_im = self.swap_axes(y_re, y_im) - - if self.last_index: - y_re_res, y_im_res = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_res, a_im=-self.a_im_res) - else: - y_re_res, y_im_res = self.complex_matmul(x_re=x_re[..., -self.modes:], x_im=x_im[..., -self.modes:], - a_re=self.a_re_res, a_im=self.a_im_res) - - y_re_res, y_im_res = self.swap_axes(y_re_res, y_im_res) - return y_re + y_re_res, y_im + y_im_res - - -class DFTn(nn.Cell): - """N dimensional Discrete Fourier Transformation""" - - def __init__(self, shape, modes, dim=None, inv=False, dtype=ms.float32): - super().__init__() - if dim is None: - dim = range(len(shape)) - self.dft1_seq = nn.SequentialCell() - last_index = [False for _ in range(len(shape))] - last_index[-1] = True - for dim_id, idx in enumerate(dim): - dft1d = DFT1d(n=shape[dim_id], modes=modes[dim_id], last_index=last_index[dim_id], idx=idx, inv=inv, - dtype=dtype) - self.dft1_seq.append(dft1d) - - def construct(self, x): - """construct""" - return self.dft1_seq(x) - - -def _dftn(shape, modes, dim=None, dtype=ms.float32): - return DFTn(shape=shape, modes=modes, dim=dim, inv=False, dtype=dtype) - - -def _idftn(shape, modes, dim=None, dtype=ms.float32): - return DFTn(shape=shape, modes=modes, dim=dim, inv=True, dtype=dtype) - - -def dft3(shape, modes, dim=(-3, -2, -1), dtype=ms.float32): - r""" - Calculate three-dimensional discrete Fourier transform. Corresponding to the rfftn operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (ms.dtype): The type of input tensor. Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 3-D tuple of Tensor. It's a complex, - including x real and imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 3. - - Examples: - >>> import numpy as np - >>> import mindspore as ms - >>> from mindspore import Tensor, ops - >>> from sciai.architecture.neural_operators.dft import dft3 - >>> array = np.ones((6, 6, 6)) * np.arange(1, 7) - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = x_re - >>> dft3_cell = dft3(shape=array.shape, modes=(2, 2, 2), dtype=ms.float32) - >>> ret, _ = dft3_cell((x_re, x_im)) - >>> print(ret) - [[[ 5.1439293e+01 -2.0076393e+01] - [ 7.9796671e-08 -1.9494735e-08] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 9.0537789e-08 1.0553553e-07] - [ 3.3567730e-07 1.0368046e-07]] - - [[ 4.7683722e-07 -3.1770034e-07] - [ 6.5267522e-15 -2.7775875e-15] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [-2.1755840e-15 -1.5215135e-15] - [ 3.6259736e-15 -4.0336615e-15]] - - [[ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00]] - - [[ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00]] - - [[ 1.1920930e-07 -5.1619136e-08] - [-3.6259733e-16 -1.0747753e-15] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 3.6259733e-16 -1.8129867e-16] - [ 3.6259733e-16 -1.4373726e-15]] - - [[ 5.9604650e-07 -2.5809570e-07] - [ 8.7023360e-15 -1.9812689e-15] - [ 0.0000000e+00 0.0000000e+00] - [ 0.0000000e+00 0.0000000e+00] - [ 2.9007787e-15 7.2519467e-16] - [ 8.7023360e-15 -1.7869532e-15]]] - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=tuple) - _check_value_in(len(shape), "shape length", 3) - _check_value_in(len(modes), "modes length", 3) - _check_param_even(shape, "shape") - _check_param_no_greater(modes[0], "mode1", shape[0] // 2) - _check_param_no_greater(modes[1], "mode2", shape[1] // 2) - _check_param_no_greater(modes[2], "mode3", shape[2] // 2 + 1) - return _dftn(shape, modes, dim=dim, dtype=dtype) - - -def idft3(shape, modes, dim=(-3, -2, -1), dtype=ms.float32): - r""" - Calculate three-dimensional discrete Fourier transform. Corresponding to the irfftn operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (ms.dtype): The type of input tensor. Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 3-D tuple of Tensor. It's a complex, including x real and - imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 3. - - Examples: - >>> import numpy as np - >>> from mindspore import Tensor, ops - >>> import mindspore.common.dtype as ms - >>> from sciai.architecture.neural_operators.dft import idft3 - >>> array = np.ones((2, 2, 2)) * np.arange(1, 3) - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = ops.zeros_like(x_re) - >>> idft3_cell = idft3(shape=(6, 6, 6), modes=(2, 2, 2), dtype=ms.float32) - >>> ret, _ = idft3_cell((x_re, x_im)) - >>> print(ret) - [[[ 5.44331074e+00 3.26598644e+00 -1.08866215e+00 -3.26598644e+00 -1.08866215e+00 3.26598644e+00] - [ 2.04124165e+00 2.04124165e+00 4.08248246e-01 -1.22474492e+00 -1.22474492e+00 4.08248365e-01] - [-6.80413842e-01 -1.22474492e+00 -6.80413783e-01 4.08248305e-01 9.52579379e-01 4.08248246e-01] - [ 0.00000000e+00 -2.30921616e-16 -2.30921616e-16 6.53092730e-32 2.30921616e-16 2.30921616e-16] - [-6.80413842e-01 4.08248246e-01 9.52579379e-01 4.08248305e-01 -6.80413783e-01 -1.22474492e+00] - [ 2.04124165e+00 4.08248365e-01 -1.22474492e+00 -1.22474492e+00 4.08248246e-01 2.04124165e+00]] - ...... - [[ 2.04124165e+00 4.08248544e-01 -1.22474492e+00 -1.22474504e+00 4.08248186e-01 2.04124165e+00] - [ 1.02062082e+00 6.12372518e-01 -2.04124182e-01 -6.12372518e-01 -2.04124182e-01 6.12372518e-01] - [-5.10310411e-01 -5.10310411e-01 -1.02062061e-01 3.06186229e-01 3.06186229e-01 -1.02062091e-01] - [-7.21630050e-17 -1.29893429e-16 -7.21630183e-17 4.32978030e-17 1.01028220e-16 4.32978163e-17] - [-6.08337416e-08 4.08248246e-01 4.08248305e-01 3.65002428e-08 -4.08248246e-01 -4.08248305e-01] - [ 5.10310471e-01 -3.06186140e-01 -7.14434564e-01 -3.06186318e-01 5.10310352e-01 9.18558717e-01]]] - - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=tuple) - _check_value_in(len(shape), "shape length", 3) - _check_value_in(len(modes), "modes length", 3) - _check_param_even(shape, "shape") - _check_param_no_greater(modes[0], "mode1", shape[0] // 2) - _check_param_no_greater(modes[1], "mode2", shape[1] // 2) - _check_param_no_greater(modes[2], "mode3", shape[2] // 2 + 1) - return _idftn(shape, modes, dim=dim, dtype=dtype) - - -def dft2(shape, modes, dim=(-2, -1), dtype=ms.float32): - """ - Calculate two-dimensional discrete Fourier transform. Corresponding to the rfft2 operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, - including x real and imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 2. - - Examples: - >>> import numpy as np - >>> from mindspore import Tensor, ops - >>> import mindspore.common.dtype as ms - >>> from sciai.architecture.neural_operators.dft import dft2 - >>> array = np.ones((5, 5)) * np.arange(1, 6) - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = x_re - >>> dft2_cell = dft2(shape=array.shape, modes=(2, 2), dtype=ms.float32) - >>> ret, _ = dft2_cell((x_re, x_im)) - >>> print(ret) - [[ 1.5000000e+01 -5.9409552e+00] - [-2.4656805e-07 7.6130398e-08] - [ 0.0000000e+00 0.0000000e+00] - [-1.9992007e-07 7.3572544e-08] - [-2.4656805e-07 7.6130398e-08]] - - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=tuple) - _check_value_in(len(shape), "shape length", 2) - _check_value_in(len(modes), "modes length", 2) - _check_param_even(shape, "shape") - _check_param_no_greater(modes[0], "mode1", shape[0] // 2) - _check_param_no_greater(modes[1], "mode2", shape[1] // 2 + 1) - return _dftn(shape, modes, dim=dim, dtype=dtype) - - -def idft2(shape, modes, dim=(-2, -1), dtype=ms.float32): - """ - Calculate two-dimensional discrete Fourier transform. Corresponding to the irfft2 operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, - including x real and imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 2. - - Examples: - >>> import numpy as np - >>> from mindspore import Tensor, ops - >>> import mindspore.common.dtype as ms - >>> from sciai.architecture.neural_operators.dft import idft2 - >>> array = np.ones((2, 2)) * np.arange(1, 3) - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = ops.zeros_like(x_re) - >>> idft2_cell = idft2(shape=(5, 5), modes=(2, 2), dtype=ms.float32) - >>> ret, _ = idft2_cell((x_re, x_im)) - >>> print(ret) - [[ 3.9999998 1.7888544 -1.7888546 -1.7888546 1.7888544 ] - [ 0.80901694 0.80901694 -0.08541022 -0.6381966 -0.08541021] - [-0.30901706 -0.8618034 -0.30901694 0.5854102 0.5854101 ] - [-0.30901706 0.5854101 0.5854102 -0.30901694 -0.8618034 ] - [ 0.80901694 -0.08541021 -0.6381966 -0.08541022 0.80901694]] - - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=tuple) - _check_value_in(len(shape), "shape length", 2) - _check_value_in(len(modes), "modes length", 2) - _check_param_even(shape, "shape") - _check_param_no_greater(modes[0], "mode1", shape[0] // 2) - _check_param_no_greater(modes[1], "mode2", shape[1] // 2 + 1) - return _idftn(shape, modes, dim=dim, dtype=dtype) - - -def dft1(shape, modes, dim=(-1,), dtype=ms.float32): - """ - Calculate one-dimensional discrete Fourier transform. Corresponding to the rfft operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (int): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (:class:`ms.dtype`): The type of input tensor. - Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, - including x real and imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 1. - - Examples: - >>> from mindspore import Tensor, ops - >>> import mindspore.common.dtype as ms - >>> from sciai.architecture.neural_operators.dft import dft1 - >>> array = [_ for _ in range(5)] - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = ops.zeros_like(x_re) - >>> dft1_cell = dft1(shape=(len(x_re),), modes=2, dtype=ms.float32) - >>> ret, _ = dft1_cell((x_re, x_im)) - >>> print(ret) - [ 4.4721355 -1.1180341] - - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=int) - _check_value_in(len(shape), "shape length", 1) - _check_param_even(shape, "shape") - _check_param_no_greater(modes, "mode1", shape[0] // 2 + 1) - modes = (modes,) - return _dftn(shape, modes, dim=dim, dtype=dtype) - - -def idft1(shape, modes, dim=(-1,), dtype=ms.float32): - """ - Calculate one-dimensional discrete Fourier transform. Corresponding to the irfft operator in torch. - - Args: - shape (tuple): Dimension of the input 'x'. - modes (int): The length of the output transform axis. The `modes` must be no greater than half of the - dimension of input 'x'. - dim (tuple): Dimensions to be transformed. - dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. - - Inputs: - - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, - including x real and imaginary. Tensor of shape :math:`(*, *)`. - - Returns: - Complex tensor with the same shape of input x. - - Raises: - TypeError: If `shape` is not a tuple. - ValueError: If the length of `shape` is no equal to 1. - - Examples: - >>> from mindspore import Tensor, ops - >>> import mindspore.common.dtype as ms - >>> from sciai.architecture.neural_operators.dft import idft1 - >>> array = [_ for _ in range(2)] - >>> x_re = Tensor(array, dtype=ms.float32) - >>> x_im = x_re - >>> idft1_cell = idft1(shape=(len(x_re),), modes=2, dtype=ms.float32) - >>> ret, _ = idft1_cell((x_re, x_im)) - >>> print(ret) - [ 0.8944272 -0.5742576 -1.2493379 -0.19787574 1.127044 ] - - """ - _check_type(shape, "shape", target_type=tuple) - _check_type(modes, "modes", target_type=int) - _check_value_in(len(shape), "shape length", 1) - _check_param_even(shape, "shape") - _check_param_no_greater(modes, "mode1", shape[0] // 2 + 1) - modes = (modes,) - return _idftn(shape, modes, dim=dim, dtype=dtype) - - -class SpectralConv1dDft(nn.Cell): - """1D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" - - def __init__(self, in_channels, out_channels, modes1, resolution, dtype=ms.float32): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.modes1 = modes1 - self.resolution = resolution - self.dtype = dtype - self.cast = ops.Cast() - - self.scale = (1. / (in_channels * out_channels)) - self.w_re, self.w_im = [_base_param(self.scale, in_channels, out_channels, dtype, modes1) for _ in range(2)] - self.dft1_cell = dft1(shape=(self.resolution,), modes=modes1, dtype=dtype) - self.idft1_cell = idft1(shape=(self.resolution,), modes=modes1, dtype=dtype) - - @staticmethod - def mul1d(inputs, weights): - weights = weights.expand_dims(0) - inputs = inputs.expand_dims(2) - out = inputs * weights - return out.sum(1) - - def construct(self, x: Tensor): - """construct""" - x_re = x - x_im = ops.zeros_like(x_re) - x_ft_re, x_ft_im = self.dft1_cell((x_re, x_im)) - w_re = self.cast(self.w_re, self.dtype) - w_im = self.cast(self.w_im, self.dtype) - out_ft_re = self.mul1d(x_ft_re[:, :, :self.modes1], w_re) - self.mul1d(x_ft_im[:, :, :self.modes1], w_im) - out_ft_im = self.mul1d(x_ft_re[:, :, :self.modes1], w_im) + self.mul1d(x_ft_im[:, :, :self.modes1], w_re) - - x, _ = self.idft1_cell((out_ft_re, out_ft_im)) - return x - - -class SpectralConv2dDft(nn.Cell): - """2D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" - - def __init__(self, in_channels, out_channels, modes1, modes2, column_resolution, raw_resolution, dtype=ms.float16): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.modes1 = modes1 - self.modes2 = modes2 - self.column_resolution = column_resolution - self.raw_resolution = raw_resolution - self.dtype = dtype - - self.scale = (1. / (in_channels * out_channels)) - - self.w_re1, self.w_im1, self.w_re2, self.w_im2 \ - = [_base_param(self.scale, in_channels, out_channels, dtype, modes1, modes2) for _ in range(4)] - - self.dft2_cell = dft2(shape=(column_resolution, raw_resolution), modes=(modes1, modes2), dtype=dtype) - self.idft2_cell = idft2(shape=(column_resolution, raw_resolution), modes=(modes1, modes2), dtype=dtype) - self.mat = ops.zeros((1, out_channels, column_resolution - 2 * modes1, modes2), dtype) - self.concat = ops.Concat(-2) - - @staticmethod - def mul2d(inputs, weights): - weight = weights.expand_dims(0) - data = inputs.expand_dims(2) - out = weight * data - return out.sum(1) - - def construct(self, x: Tensor): - """construct""" - x_re = x - x_im = ops.zeros_like(x_re) - x_ft_re, x_ft_im = self.dft2_cell((x_re, x_im)) - - out_ft_re1 = self.mul2d(x_ft_re[:, :, :self.modes1, :self.modes2], self.w_re1) \ - - self.mul2d(x_ft_im[:, :, :self.modes1, :self.modes2], self.w_im1) - out_ft_im1 = self.mul2d(x_ft_re[:, :, :self.modes1, :self.modes2], self.w_im1) \ - + self.mul2d(x_ft_im[:, :, :self.modes1, :self.modes2], self.w_re1) - - out_ft_re2 = self.mul2d(x_ft_re[:, :, -self.modes1:, :self.modes2], self.w_re2) \ - - self.mul2d(x_ft_im[:, :, -self.modes1:, :self.modes2], self.w_im2) - out_ft_im2 = self.mul2d(x_ft_re[:, :, -self.modes1:, :self.modes2], self.w_im2) \ - + self.mul2d(x_ft_im[:, :, -self.modes1:, :self.modes2], self.w_re2) - - batch_size = x.shape[0] - mat = self.mat.repeat(batch_size, 0) - out_re = self.concat((out_ft_re1, mat, out_ft_re2)) - out_im = self.concat((out_ft_im1, mat, out_ft_im2)) - - x, _ = self.idft2_cell((out_re, out_im)) - return x - - -class SpectralConv3d(nn.Cell): - """3D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" - - def __init__(self, in_channels, out_channels, modes1, modes2, modes3, - column_resolution, row_resolution, bar_resolution, dtype=ms.float32): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - # Number of Fourier modes to multiply, at most floor(N/2) + 1 - self.modes1 = modes1 - self.modes2 = modes2 - self.modes3 = modes3 - self.column_resolution = column_resolution - self.row_resolution = row_resolution - self.bar_resolution = bar_resolution - self.dtype = dtype - - self.scale = (1 / (in_channels * out_channels)) - - self.w_re1, self.w_im1, self.w_re2, self.w_im2, self.w_re3, self.w_im3, self.w_re4, self.w_im4 \ - = [_base_param(self.scale, in_channels, out_channels, dtype, modes1, modes2, modes3) for _ in range(8)] - - self.dft3_cell = dft3(shape=(column_resolution, row_resolution, bar_resolution), - modes=(modes1, modes2, modes3), - dtype=dtype) - self.idft3_cell = idft3(shape=(column_resolution, row_resolution, bar_resolution), - modes=(modes1, modes2, modes3), - dtype=dtype) - self.mat_x = ops.zeros((1, out_channels, column_resolution - 2 * modes1, modes2, modes3), dtype) - self.mat_y = ops.zeros((1, out_channels, column_resolution, row_resolution - 2 * modes2, modes3), dtype) - self.concat = ops.Concat(-2) - - def _base_parameter(self): - return Parameter(Tensor( - self.scale * np.random.rand(self.in_channels, self.out_channels, self.modes1, self.modes2, self.modes3), - dtype=self.dtype), requires_grad=True) - - # Complex multiplication - def mul3d(self, inputs, weights): - weight = weights.expand_dims(0) - data = inputs.expand_dims(2) - out = weight * data - return out.sum(1) - - def construct(self, x: Tensor): - """construct""" - x_re = x - x_im = ops.zeros_like(x_re) - x_ft_re, x_ft_im = self.dft3_cell((x_re, x_im)) - - out_ft_re1 = self.mul3d(x_ft_re[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_re1) \ - - self.mul3d(x_ft_im[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_im1) - out_ft_im1 = self.mul3d(x_ft_re[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_im1) \ - + self.mul3d(x_ft_im[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_re1) - - out_ft_re2 = self.mul3d(x_ft_re[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_re2) \ - - self.mul3d(x_ft_im[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_im2) - out_ft_im2 = self.mul3d(x_ft_re[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_im2) \ - + self.mul3d(x_ft_im[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_re2) - - out_ft_re3 = self.mul3d(x_ft_re[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_re3) \ - - self.mul3d(x_ft_im[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_im3) - out_ft_im3 = self.mul3d(x_ft_re[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_im3) \ - + self.mul3d(x_ft_im[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_re3) - - out_ft_re4 = self.mul3d(x_ft_re[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_re4) \ - - self.mul3d(x_ft_im[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_im4) - out_ft_im4 = self.mul3d(x_ft_re[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_im4) \ - + self.mul3d(x_ft_im[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_re4) - - batch_size = x.shape[0] - mat_x = self.mat_x.repeat(batch_size, 0) - mat_y = self.mat_y.repeat(batch_size, 0) - - out_re1 = ops.concat((out_ft_re1, mat_x, out_ft_re2), -3) - out_im1 = ops.concat((out_ft_im1, mat_x, out_ft_im2), -3) - - out_re2 = ops.concat((out_ft_re3, mat_x, out_ft_re4), -3) - out_im2 = ops.concat((out_ft_im3, mat_x, out_ft_im4), -3) - out_re = ops.concat((out_re1, mat_y, out_re2), -2) - out_im = ops.concat((out_im1, mat_y, out_im2), -2) - x, _ = self.idft3_cell((out_re, out_im)) - return x - - -def _base_param(scale, in_channels, out_channels, dtype, *modes): - return Parameter(Tensor(scale * np.random.rand(in_channels, out_channels, *modes), dtype=dtype), requires_grad=True) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""dft""" +import mindspore as ms +import mindspore.numpy as mnp +from mindspore import nn, ops, Tensor, Parameter +import numpy as np +from scipy.linalg import dft + +from sciai.utils.check_utils import _check_param_no_greater, _check_value_in, _check_type, _check_param_even + + +class DFT1d(nn.Cell): + """One dimensional Discrete Fourier Transformation""" + + def __init__(self, n, modes, last_index, idx=0, inv=False, dtype=ms.float32): + super().__init__() + self.n = n + self.dft_mat = dft(n, scale="sqrtn") + self.modes = modes + self.last_index = last_index + self.inv = inv + self.idx = idx + + self.dft_mode_mat_upper = self.dft_mat[:, :modes] + self.a_re_upper = Tensor(self.dft_mode_mat_upper.real, dtype=dtype) + self.a_im_upper = Tensor(self.dft_mode_mat_upper.imag, dtype=dtype) + + self.dft_mode_mat_lower = self.dft_mat[:, -modes:] + self.a_re_lower = Tensor(self.dft_mode_mat_lower.real, dtype=dtype) + self.a_im_lower = Tensor(self.dft_mode_mat_lower.imag, dtype=dtype) + self.concat = ops.Concat(axis=-1) + + if self.inv: + self.a_re_upper = self.a_re_upper.T + self.a_im_upper = -self.a_im_upper.T + if last_index: + if modes == n // 2 + 1: + self.dft_mat_res = self.dft_mat[:, -modes + 2:] + else: + self.dft_mat_res = self.dft_mat[:, -modes + 1:] + + mat = Tensor(np.zeros(n), dtype=dtype).reshape(n, 1) + self.a_re_res = mnp.flip(Tensor(self.dft_mat_res.real, dtype=dtype), axis=-1) + self.a_im_res = mnp.flip(Tensor(self.dft_mat_res.imag, dtype=dtype), axis=-1) + if modes == n // 2 + 1: + self.a_re_res = self.concat((mat, self.a_re_res, mat)) + self.a_im_res = self.concat((mat, self.a_im_res, mat)) + else: + self.a_re_res = self.concat((mat, self.a_re_res)) + self.a_im_res = self.concat((mat, self.a_im_res)) + + self.a_re_res = self.a_re_res.T + self.a_im_res = -self.a_im_res.T + else: + self.a_re_res = self.a_re_lower.T + self.a_im_res = -self.a_im_lower.T + + if (self.n - 2 * self.modes) > 0: + self.mat = ops.zeros((self.n - 2 * self.modes), dtype) + + def swap_axes(self, x_re, x_im): + return x_re.swapaxes(-1, self.idx), x_im.swapaxes(-1, self.idx) + + def complex_matmul(self, x_re, x_im, a_re, a_im): + y_re = ops.matmul(x_re, a_re) - ops.matmul(x_im, a_im) + y_im = ops.matmul(x_im, a_re) + ops.matmul(x_re, a_im) + return y_re, y_im + + def construct(self, x): + """construct""" + x_re, x_im = x + + if not self.inv: + x_re, x_im = self.swap_axes(x_re, x_im) + y_re, y_im = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_upper, a_im=self.a_im_upper) + + if not self.last_index: + y_re2, y_im2 = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_lower, a_im=self.a_im_lower) + + if self.n == self.modes * 2: + y_re = self.concat((y_re, y_re2)) + y_im = self.concat((y_im, y_im2)) + else: + dims = x_re.shape[:-1] + length = len(dims) + mat = self.mat + for i in range(length - 1, -1, -1): + mat = mat.expand_dims(0).repeat(dims[i], 0) + y_re = self.concat((y_re, mat, y_re2)) + y_im = self.concat((y_im, mat, y_im2)) + + y_re, y_im = self.swap_axes(y_re, y_im) + return y_re, y_im + + x_re, x_im = self.swap_axes(x_re, x_im) + y_re, y_im = self.complex_matmul(x_re=x_re[..., :self.modes], x_im=x_im[..., :self.modes], + a_re=self.a_re_upper, a_im=self.a_im_upper) + y_re, y_im = self.swap_axes(y_re, y_im) + + if self.last_index: + y_re_res, y_im_res = self.complex_matmul(x_re=x_re, x_im=x_im, a_re=self.a_re_res, a_im=-self.a_im_res) + else: + y_re_res, y_im_res = self.complex_matmul(x_re=x_re[..., -self.modes:], x_im=x_im[..., -self.modes:], + a_re=self.a_re_res, a_im=self.a_im_res) + + y_re_res, y_im_res = self.swap_axes(y_re_res, y_im_res) + return y_re + y_re_res, y_im + y_im_res + + +class DFTn(nn.Cell): + """N dimensional Discrete Fourier Transformation""" + + def __init__(self, shape, modes, dim=None, inv=False, dtype=ms.float32): + super().__init__() + if dim is None: + dim = range(len(shape)) + self.dft1_seq = nn.SequentialCell() + last_index = [False for _ in range(len(shape))] + last_index[-1] = True + for dim_id, idx in enumerate(dim): + dft1d = DFT1d(n=shape[dim_id], modes=modes[dim_id], last_index=last_index[dim_id], idx=idx, inv=inv, + dtype=dtype) + self.dft1_seq.append(dft1d) + + def construct(self, x): + """construct""" + return self.dft1_seq(x) + + +def _dftn(shape, modes, dim=None, dtype=ms.float32): + return DFTn(shape=shape, modes=modes, dim=dim, inv=False, dtype=dtype) + + +def _idftn(shape, modes, dim=None, dtype=ms.float32): + return DFTn(shape=shape, modes=modes, dim=dim, inv=True, dtype=dtype) + + +def dft3(shape, modes, dim=(-3, -2, -1), dtype=ms.float32): + r""" + Calculate three-dimensional discrete Fourier transform. Corresponding to the rfftn operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (ms.dtype): The type of input tensor. Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 3-D tuple of Tensor. It's a complex, + including x real and imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 3. + + Examples: + >>> import numpy as np + >>> import mindspore as ms + >>> from mindspore import Tensor, ops + >>> from sciai.architecture.neural_operators.dft import dft3 + >>> array = np.ones((6, 6, 6)) * np.arange(1, 7) + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = x_re + >>> dft3_cell = dft3(shape=array.shape, modes=(2, 2, 2), dtype=ms.float32) + >>> ret, _ = dft3_cell((x_re, x_im)) + >>> print(ret) + [[[ 5.1439293e+01 -2.0076393e+01] + [ 7.9796671e-08 -1.9494735e-08] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 9.0537789e-08 1.0553553e-07] + [ 3.3567730e-07 1.0368046e-07]] + + [[ 4.7683722e-07 -3.1770034e-07] + [ 6.5267522e-15 -2.7775875e-15] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [-2.1755840e-15 -1.5215135e-15] + [ 3.6259736e-15 -4.0336615e-15]] + + [[ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00]] + + [[ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00]] + + [[ 1.1920930e-07 -5.1619136e-08] + [-3.6259733e-16 -1.0747753e-15] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 3.6259733e-16 -1.8129867e-16] + [ 3.6259733e-16 -1.4373726e-15]] + + [[ 5.9604650e-07 -2.5809570e-07] + [ 8.7023360e-15 -1.9812689e-15] + [ 0.0000000e+00 0.0000000e+00] + [ 0.0000000e+00 0.0000000e+00] + [ 2.9007787e-15 7.2519467e-16] + [ 8.7023360e-15 -1.7869532e-15]]] + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=tuple) + _check_value_in(len(shape), "shape length", 3) + _check_value_in(len(modes), "modes length", 3) + _check_param_even(shape, "shape") + _check_param_no_greater(modes[0], "mode1", shape[0] // 2) + _check_param_no_greater(modes[1], "mode2", shape[1] // 2) + _check_param_no_greater(modes[2], "mode3", shape[2] // 2 + 1) + return _dftn(shape, modes, dim=dim, dtype=dtype) + + +def idft3(shape, modes, dim=(-3, -2, -1), dtype=ms.float32): + r""" + Calculate three-dimensional discrete Fourier transform. Corresponding to the irfftn operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (ms.dtype): The type of input tensor. Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 3-D tuple of Tensor. It's a complex, including x real and + imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 3. + + Examples: + >>> import numpy as np + >>> from mindspore import Tensor, ops + >>> import mindspore.common.dtype as ms + >>> from sciai.architecture.neural_operators.dft import idft3 + >>> array = np.ones((2, 2, 2)) * np.arange(1, 3) + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = ops.zeros_like(x_re) + >>> idft3_cell = idft3(shape=(6, 6, 6), modes=(2, 2, 2), dtype=ms.float32) + >>> ret, _ = idft3_cell((x_re, x_im)) + >>> print(ret) + [[[ 5.44331074e+00 3.26598644e+00 -1.08866215e+00 -3.26598644e+00 -1.08866215e+00 3.26598644e+00] + [ 2.04124165e+00 2.04124165e+00 4.08248246e-01 -1.22474492e+00 -1.22474492e+00 4.08248365e-01] + [-6.80413842e-01 -1.22474492e+00 -6.80413783e-01 4.08248305e-01 9.52579379e-01 4.08248246e-01] + [ 0.00000000e+00 -2.30921616e-16 -2.30921616e-16 6.53092730e-32 2.30921616e-16 2.30921616e-16] + [-6.80413842e-01 4.08248246e-01 9.52579379e-01 4.08248305e-01 -6.80413783e-01 -1.22474492e+00] + [ 2.04124165e+00 4.08248365e-01 -1.22474492e+00 -1.22474492e+00 4.08248246e-01 2.04124165e+00]] + ...... + [[ 2.04124165e+00 4.08248544e-01 -1.22474492e+00 -1.22474504e+00 4.08248186e-01 2.04124165e+00] + [ 1.02062082e+00 6.12372518e-01 -2.04124182e-01 -6.12372518e-01 -2.04124182e-01 6.12372518e-01] + [-5.10310411e-01 -5.10310411e-01 -1.02062061e-01 3.06186229e-01 3.06186229e-01 -1.02062091e-01] + [-7.21630050e-17 -1.29893429e-16 -7.21630183e-17 4.32978030e-17 1.01028220e-16 4.32978163e-17] + [-6.08337416e-08 4.08248246e-01 4.08248305e-01 3.65002428e-08 -4.08248246e-01 -4.08248305e-01] + [ 5.10310471e-01 -3.06186140e-01 -7.14434564e-01 -3.06186318e-01 5.10310352e-01 9.18558717e-01]]] + + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=tuple) + _check_value_in(len(shape), "shape length", 3) + _check_value_in(len(modes), "modes length", 3) + _check_param_even(shape, "shape") + _check_param_no_greater(modes[0], "mode1", shape[0] // 2) + _check_param_no_greater(modes[1], "mode2", shape[1] // 2) + _check_param_no_greater(modes[2], "mode3", shape[2] // 2 + 1) + return _idftn(shape, modes, dim=dim, dtype=dtype) + + +def dft2(shape, modes, dim=(-2, -1), dtype=ms.float32): + """ + Calculate two-dimensional discrete Fourier transform. Corresponding to the rfft2 operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, + including x real and imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 2. + + Examples: + >>> import numpy as np + >>> from mindspore import Tensor, ops + >>> import mindspore.common.dtype as ms + >>> from sciai.architecture.neural_operators.dft import dft2 + >>> array = np.ones((5, 5)) * np.arange(1, 6) + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = x_re + >>> dft2_cell = dft2(shape=array.shape, modes=(2, 2), dtype=ms.float32) + >>> ret, _ = dft2_cell((x_re, x_im)) + >>> print(ret) + [[ 1.5000000e+01 -5.9409552e+00] + [-2.4656805e-07 7.6130398e-08] + [ 0.0000000e+00 0.0000000e+00] + [-1.9992007e-07 7.3572544e-08] + [-2.4656805e-07 7.6130398e-08]] + + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=tuple) + _check_value_in(len(shape), "shape length", 2) + _check_value_in(len(modes), "modes length", 2) + _check_param_even(shape, "shape") + _check_param_no_greater(modes[0], "mode1", shape[0] // 2) + _check_param_no_greater(modes[1], "mode2", shape[1] // 2 + 1) + return _dftn(shape, modes, dim=dim, dtype=dtype) + + +def idft2(shape, modes, dim=(-2, -1), dtype=ms.float32): + """ + Calculate two-dimensional discrete Fourier transform. Corresponding to the irfft2 operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (tuple): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, + including x real and imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 2. + + Examples: + >>> import numpy as np + >>> from mindspore import Tensor, ops + >>> import mindspore.common.dtype as ms + >>> from sciai.architecture.neural_operators.dft import idft2 + >>> array = np.ones((2, 2)) * np.arange(1, 3) + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = ops.zeros_like(x_re) + >>> idft2_cell = idft2(shape=(5, 5), modes=(2, 2), dtype=ms.float32) + >>> ret, _ = idft2_cell((x_re, x_im)) + >>> print(ret) + [[ 3.9999998 1.7888544 -1.7888546 -1.7888546 1.7888544 ] + [ 0.80901694 0.80901694 -0.08541022 -0.6381966 -0.08541021] + [-0.30901706 -0.8618034 -0.30901694 0.5854102 0.5854101 ] + [-0.30901706 0.5854101 0.5854102 -0.30901694 -0.8618034 ] + [ 0.80901694 -0.08541021 -0.6381966 -0.08541022 0.80901694]] + + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=tuple) + _check_value_in(len(shape), "shape length", 2) + _check_value_in(len(modes), "modes length", 2) + _check_param_even(shape, "shape") + _check_param_no_greater(modes[0], "mode1", shape[0] // 2) + _check_param_no_greater(modes[1], "mode2", shape[1] // 2 + 1) + return _idftn(shape, modes, dim=dim, dtype=dtype) + + +def dft1(shape, modes, dim=(-1,), dtype=ms.float32): + """ + Calculate one-dimensional discrete Fourier transform. Corresponding to the rfft operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (int): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (:class:`ms.dtype`): The type of input tensor. + Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, + including x real and imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 1. + + Examples: + >>> from mindspore import Tensor, ops + >>> import mindspore.common.dtype as ms + >>> from sciai.architecture.neural_operators.dft import dft1 + >>> array = [_ for _ in range(5)] + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = ops.zeros_like(x_re) + >>> dft1_cell = dft1(shape=(len(x_re),), modes=2, dtype=ms.float32) + >>> ret, _ = dft1_cell((x_re, x_im)) + >>> print(ret) + [ 4.4721355 -1.1180341] + + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=int) + _check_value_in(len(shape), "shape length", 1) + _check_param_even(shape, "shape") + _check_param_no_greater(modes, "mode1", shape[0] // 2 + 1) + modes = (modes,) + return _dftn(shape, modes, dim=dim, dtype=dtype) + + +def idft1(shape, modes, dim=(-1,), dtype=ms.float32): + """ + Calculate one-dimensional discrete Fourier transform. Corresponding to the irfft operator in torch. + + Args: + shape (tuple): Dimension of the input 'x'. + modes (int): The length of the output transform axis. The `modes` must be no greater than half of the + dimension of input 'x'. + dim (tuple): Dimensions to be transformed. + dtype (:class:`ms.dtype`): The type of input tensor. Default: ms.float32. + + Inputs: + - **x** (Tensor, Tensor): The input data. It's 2-D tuple of Tensor. It's a complex, + including x real and imaginary. Tensor of shape :math:`(*, *)`. + + Returns: + Complex tensor with the same shape of input x. + + Raises: + TypeError: If `shape` is not a tuple. + ValueError: If the length of `shape` is no equal to 1. + + Examples: + >>> from mindspore import Tensor, ops + >>> import mindspore.common.dtype as ms + >>> from sciai.architecture.neural_operators.dft import idft1 + >>> array = [_ for _ in range(2)] + >>> x_re = Tensor(array, dtype=ms.float32) + >>> x_im = x_re + >>> idft1_cell = idft1(shape=(len(x_re),), modes=2, dtype=ms.float32) + >>> ret, _ = idft1_cell((x_re, x_im)) + >>> print(ret) + [ 0.8944272 -0.5742576 -1.2493379 -0.19787574 1.127044 ] + + """ + _check_type(shape, "shape", target_type=tuple) + _check_type(modes, "modes", target_type=int) + _check_value_in(len(shape), "shape length", 1) + _check_param_even(shape, "shape") + _check_param_no_greater(modes, "mode1", shape[0] // 2 + 1) + modes = (modes,) + return _idftn(shape, modes, dim=dim, dtype=dtype) + + +class SpectralConv1dDft(nn.Cell): + """1D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" + + def __init__(self, in_channels, out_channels, modes1, resolution, dtype=ms.float32): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.modes1 = modes1 + self.resolution = resolution + self.dtype = dtype + self.cast = ops.Cast() + + self.scale = (1. / (in_channels * out_channels)) + self.w_re, self.w_im = [_base_param(self.scale, in_channels, out_channels, dtype, modes1) for _ in range(2)] + self.dft1_cell = dft1(shape=(self.resolution,), modes=modes1, dtype=dtype) + self.idft1_cell = idft1(shape=(self.resolution,), modes=modes1, dtype=dtype) + + @staticmethod + def mul1d(inputs, weights): + weights = weights.expand_dims(0) + inputs = inputs.expand_dims(2) + out = inputs * weights + return out.sum(1) + + def construct(self, x: Tensor): + """construct""" + x_re = x + x_im = ops.zeros_like(x_re) + x_ft_re, x_ft_im = self.dft1_cell((x_re, x_im)) + w_re = self.cast(self.w_re, self.dtype) + w_im = self.cast(self.w_im, self.dtype) + out_ft_re = self.mul1d(x_ft_re[:, :, :self.modes1], w_re) - self.mul1d(x_ft_im[:, :, :self.modes1], w_im) + out_ft_im = self.mul1d(x_ft_re[:, :, :self.modes1], w_im) + self.mul1d(x_ft_im[:, :, :self.modes1], w_re) + + x, _ = self.idft1_cell((out_ft_re, out_ft_im)) + return x + + +class SpectralConv2dDft(nn.Cell): + """2D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" + + def __init__(self, in_channels, out_channels, modes1, modes2, column_resolution, raw_resolution, dtype=ms.float16): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.modes1 = modes1 + self.modes2 = modes2 + self.column_resolution = column_resolution + self.raw_resolution = raw_resolution + self.dtype = dtype + + self.scale = (1. / (in_channels * out_channels)) + + self.w_re1, self.w_im1, self.w_re2, self.w_im2 \ + = [_base_param(self.scale, in_channels, out_channels, dtype, modes1, modes2) for _ in range(4)] + + self.dft2_cell = dft2(shape=(column_resolution, raw_resolution), modes=(modes1, modes2), dtype=dtype) + self.idft2_cell = idft2(shape=(column_resolution, raw_resolution), modes=(modes1, modes2), dtype=dtype) + self.mat = ops.zeros((1, out_channels, column_resolution - 2 * modes1, modes2), dtype) + self.concat = ops.Concat(-2) + + @staticmethod + def mul2d(inputs, weights): + weight = weights.expand_dims(0) + data = inputs.expand_dims(2) + out = weight * data + return out.sum(1) + + def construct(self, x: Tensor): + """construct""" + x_re = x + x_im = ops.zeros_like(x_re) + x_ft_re, x_ft_im = self.dft2_cell((x_re, x_im)) + + out_ft_re1 = self.mul2d(x_ft_re[:, :, :self.modes1, :self.modes2], self.w_re1) \ + - self.mul2d(x_ft_im[:, :, :self.modes1, :self.modes2], self.w_im1) + out_ft_im1 = self.mul2d(x_ft_re[:, :, :self.modes1, :self.modes2], self.w_im1) \ + + self.mul2d(x_ft_im[:, :, :self.modes1, :self.modes2], self.w_re1) + + out_ft_re2 = self.mul2d(x_ft_re[:, :, -self.modes1:, :self.modes2], self.w_re2) \ + - self.mul2d(x_ft_im[:, :, -self.modes1:, :self.modes2], self.w_im2) + out_ft_im2 = self.mul2d(x_ft_re[:, :, -self.modes1:, :self.modes2], self.w_im2) \ + + self.mul2d(x_ft_im[:, :, -self.modes1:, :self.modes2], self.w_re2) + + batch_size = x.shape[0] + mat = self.mat.repeat(batch_size, 0) + out_re = self.concat((out_ft_re1, mat, out_ft_re2)) + out_im = self.concat((out_ft_im1, mat, out_ft_im2)) + + x, _ = self.idft2_cell((out_re, out_im)) + return x + + +class SpectralConv3d(nn.Cell): + """3D Fourier layer. It does DFT, linear transform, and Inverse DFT.""" + + def __init__(self, in_channels, out_channels, modes1, modes2, modes3, + column_resolution, row_resolution, bar_resolution, dtype=ms.float32): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + # Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = modes1 + self.modes2 = modes2 + self.modes3 = modes3 + self.column_resolution = column_resolution + self.row_resolution = row_resolution + self.bar_resolution = bar_resolution + self.dtype = dtype + + self.scale = (1 / (in_channels * out_channels)) + + self.w_re1, self.w_im1, self.w_re2, self.w_im2, self.w_re3, self.w_im3, self.w_re4, self.w_im4 \ + = [_base_param(self.scale, in_channels, out_channels, dtype, modes1, modes2, modes3) for _ in range(8)] + + self.dft3_cell = dft3(shape=(column_resolution, row_resolution, bar_resolution), + modes=(modes1, modes2, modes3), + dtype=dtype) + self.idft3_cell = idft3(shape=(column_resolution, row_resolution, bar_resolution), + modes=(modes1, modes2, modes3), + dtype=dtype) + self.mat_x = ops.zeros((1, out_channels, column_resolution - 2 * modes1, modes2, modes3), dtype) + self.mat_y = ops.zeros((1, out_channels, column_resolution, row_resolution - 2 * modes2, modes3), dtype) + self.concat = ops.Concat(-2) + + def _base_parameter(self): + return Parameter(Tensor( + self.scale * np.random.rand(self.in_channels, self.out_channels, self.modes1, self.modes2, self.modes3), + dtype=self.dtype), requires_grad=True) + + # Complex multiplication + def mul3d(self, inputs, weights): + weight = weights.expand_dims(0) + data = inputs.expand_dims(2) + out = weight * data + return out.sum(1) + + def construct(self, x: Tensor): + """construct""" + x_re = x + x_im = ops.zeros_like(x_re) + x_ft_re, x_ft_im = self.dft3_cell((x_re, x_im)) + + out_ft_re1 = self.mul3d(x_ft_re[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_re1) \ + - self.mul3d(x_ft_im[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_im1) + out_ft_im1 = self.mul3d(x_ft_re[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_im1) \ + + self.mul3d(x_ft_im[:, :, :self.modes1, :self.modes2, :self.modes3], self.w_re1) + + out_ft_re2 = self.mul3d(x_ft_re[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_re2) \ + - self.mul3d(x_ft_im[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_im2) + out_ft_im2 = self.mul3d(x_ft_re[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_im2) \ + + self.mul3d(x_ft_im[:, :, -self.modes1:, :self.modes2, :self.modes3], self.w_re2) + + out_ft_re3 = self.mul3d(x_ft_re[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_re3) \ + - self.mul3d(x_ft_im[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_im3) + out_ft_im3 = self.mul3d(x_ft_re[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_im3) \ + + self.mul3d(x_ft_im[:, :, :self.modes1, -self.modes2:, :self.modes3], self.w_re3) + + out_ft_re4 = self.mul3d(x_ft_re[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_re4) \ + - self.mul3d(x_ft_im[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_im4) + out_ft_im4 = self.mul3d(x_ft_re[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_im4) \ + + self.mul3d(x_ft_im[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.w_re4) + + batch_size = x.shape[0] + mat_x = self.mat_x.repeat(batch_size, 0) + mat_y = self.mat_y.repeat(batch_size, 0) + + out_re1 = ops.concat((out_ft_re1, mat_x, out_ft_re2), -3) + out_im1 = ops.concat((out_ft_im1, mat_x, out_ft_im2), -3) + + out_re2 = ops.concat((out_ft_re3, mat_x, out_ft_re4), -3) + out_im2 = ops.concat((out_ft_im3, mat_x, out_ft_im4), -3) + out_re = ops.concat((out_re1, mat_y, out_re2), -2) + out_im = ops.concat((out_im1, mat_y, out_im2), -2) + x, _ = self.idft3_cell((out_re, out_im)) + return x + + +def _base_param(scale, in_channels, out_channels, dtype, *modes): + return Parameter(Tensor(scale * np.random.rand(in_channels, out_channels, *modes), dtype=dtype), requires_grad=True) diff --git a/SciAI/sciai/architecture/neural_operators/fno1d.py b/SciAI/sciai/architecture/neural_operators/fno1d.py index b2cf64080..714f0456e 100644 --- a/SciAI/sciai/architecture/neural_operators/fno1d.py +++ b/SciAI/sciai/architecture/neural_operators/fno1d.py @@ -1,124 +1,138 @@ -"""FNO1d""" -import mindspore as ms -from mindspore import ops, nn, Tensor - -from sciai.architecture.neural_operators.dft import SpectralConv1dDft -from sciai.utils.check_utils import _check_type -from sciai.utils.math_utils import _get_grid_1d - - -class FNOBlock(nn.Cell): - def __init__(self, in_channels, out_channels, modes1, resolution=1024, gelu=True, dtype=ms.float32): - super().__init__() - self.conv = SpectralConv1dDft(in_channels, out_channels, modes1, resolution, dtype=dtype) - self.w = nn.Conv1d(in_channels, out_channels, 1).to_float(dtype) - self.act = ops.GeLU() if gelu else ops.Identity() - - def construct(self, x): - return self.act(self.conv(x) + self.w(x)) + x - - -class FNO1D(nn.Cell): - r""" - The 1-dimensional Fourier Neural Operator (FNO1D) contains a lifting layer, - multiple Fourier layers and a decoder layer. - The details can be found in `Fourier neural operator for - parametric partial differential equations `_. - - Args: - in_channels (int): The number of channels in the input space. - out_channels (int): The number of channels in the output space. - resolution (int): The spatial resolution of the input. - modes (int): The number of low-frequency components to keep. - channels (int): The number of channels after dimension lifting of the input. Default: 20. - depths (int): The number of FNO layers. Default: 4. - mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. - dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. - `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. - Default: `ms.float32`. - - Inputs: - - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. - - Outputs: - Tensor, the output of FNO network. - - - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, out\_channels)`. - - Raises: - TypeError: If `in_channels` is not an int. - TypeError: If `out_channels` is not an int. - TypeError: If `resolution` is not an int. - TypeError: If `modes` is not an int. - ValueError: If `modes` is less than 1. - - Supported Platforms: - `Ascend` `GPU` - - Examples: - >>> import numpy as np - >>> from mindspore.common.initializer import initializer, Normal - >>> from sciai.architecture.neural_operators import FNO1D - >>> B, W, C = 32, 1024, 1 - >>> x = initializer(Normal(), [B, W, C]) - >>> net = FNO1D(in_channels=1, out_channels=1, resolution=64, modes=12) - >>> output = net(x) - >>> print(output.shape) - (32, 1024, 1) - - """ - - def __init__(self, - in_channels, - out_channels, - resolution, - modes, - channels=20, - depths=4, - mlp_ratio=4, - dtype=ms.float32): - super().__init__() - _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) - _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) - _check_type(resolution, "resolution", target_type=int, exclude_type=bool) - _check_type(modes, "modes", target_type=int, exclude_type=bool) - if modes < 1: - raise ValueError( - "modes must at least 1, but got mode: {}".format(modes)) - self.modes1 = modes - self.channels = channels - self.fc_channel = mlp_ratio * channels - self.fc0 = nn.Dense(in_channels + 1, self.channels).to_float(dtype) - self.layers = depths - - self.fno_seq = nn.SequentialCell() - for _ in range(self.layers - 1): - self.fno_seq.append( - FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) - self.fno_seq.append( - FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) - - self.fc1 = nn.Dense(self.channels, self.fc_channel).to_float(dtype) - self.fc2 = nn.Dense(self.fc_channel, out_channels).to_float(dtype) - - self.grid = Tensor(_get_grid_1d(resolution), dtype=ms.float32) - self.concat = ops.Concat(axis=-1) - self.transpose = ops.Transpose() - self.act = ops.GeLU() - - def construct(self, x: Tensor): - """construct""" - batch_size = x.shape[0] - - grid = self.grid.repeat(batch_size, axis=0) - x = self.concat((x, grid)) - x = self.fc0(x) - x = self.transpose(x, (0, 2, 1)) - - x = self.fno_seq(x) - - x = self.transpose(x, (0, 2, 1)) - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - return x +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""FNO1d""" +import mindspore as ms +from mindspore import ops, nn, Tensor + +from sciai.architecture.neural_operators.dft import SpectralConv1dDft +from sciai.utils.check_utils import _check_type +from sciai.utils.math_utils import _get_grid_1d + + +class FNOBlock(nn.Cell): + def __init__(self, in_channels, out_channels, modes1, resolution=1024, gelu=True, dtype=ms.float32): + super().__init__() + self.conv = SpectralConv1dDft(in_channels, out_channels, modes1, resolution, dtype=dtype) + self.w = nn.Conv1d(in_channels, out_channels, 1).to_float(dtype) + self.act = ops.GeLU() if gelu else ops.Identity() + + def construct(self, x): + return self.act(self.conv(x) + self.w(x)) + x + + +class FNO1D(nn.Cell): + r""" + The 1-dimensional Fourier Neural Operator (FNO1D) contains a lifting layer, + multiple Fourier layers and a decoder layer. + The details can be found in `Fourier neural operator for + parametric partial differential equations `_. + + Args: + in_channels (int): The number of channels in the input space. + out_channels (int): The number of channels in the output space. + resolution (int): The spatial resolution of the input. + modes (int): The number of low-frequency components to keep. + channels (int): The number of channels after dimension lifting of the input. Default: 20. + depths (int): The number of FNO layers. Default: 4. + mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. + dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. + `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. + Default: `ms.float32`. + + Inputs: + - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. + + Outputs: + Tensor, the output of FNO network. + + - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, out\_channels)`. + + Raises: + TypeError: If `in_channels` is not an int. + TypeError: If `out_channels` is not an int. + TypeError: If `resolution` is not an int. + TypeError: If `modes` is not an int. + ValueError: If `modes` is less than 1. + + Supported Platforms: + `Ascend` `GPU` + + Examples: + >>> import numpy as np + >>> from mindspore.common.initializer import initializer, Normal + >>> from sciai.architecture.neural_operators import FNO1D + >>> B, W, C = 32, 1024, 1 + >>> x = initializer(Normal(), [B, W, C]) + >>> net = FNO1D(in_channels=1, out_channels=1, resolution=64, modes=12) + >>> output = net(x) + >>> print(output.shape) + (32, 1024, 1) + + """ + + def __init__(self, + in_channels, + out_channels, + resolution, + modes, + channels=20, + depths=4, + mlp_ratio=4, + dtype=ms.float32): + super().__init__() + _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) + _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) + _check_type(resolution, "resolution", target_type=int, exclude_type=bool) + _check_type(modes, "modes", target_type=int, exclude_type=bool) + if modes < 1: + raise ValueError( + "modes must at least 1, but got mode: {}".format(modes)) + self.modes1 = modes + self.channels = channels + self.fc_channel = mlp_ratio * channels + self.fc0 = nn.Dense(in_channels + 1, self.channels).to_float(dtype) + self.layers = depths + + self.fno_seq = nn.SequentialCell() + for _ in range(self.layers - 1): + self.fno_seq.append( + FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) + self.fno_seq.append( + FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) + + self.fc1 = nn.Dense(self.channels, self.fc_channel).to_float(dtype) + self.fc2 = nn.Dense(self.fc_channel, out_channels).to_float(dtype) + + self.grid = Tensor(_get_grid_1d(resolution), dtype=ms.float32) + self.concat = ops.Concat(axis=-1) + self.transpose = ops.Transpose() + self.act = ops.GeLU() + + def construct(self, x: Tensor): + """construct""" + batch_size = x.shape[0] + + grid = self.grid.repeat(batch_size, axis=0) + x = self.concat((x, grid)) + x = self.fc0(x) + x = self.transpose(x, (0, 2, 1)) + + x = self.fno_seq(x) + + x = self.transpose(x, (0, 2, 1)) + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x diff --git a/SciAI/sciai/architecture/neural_operators/fno2d.py b/SciAI/sciai/architecture/neural_operators/fno2d.py index c6d065f38..bfc7951f8 100644 --- a/SciAI/sciai/architecture/neural_operators/fno2d.py +++ b/SciAI/sciai/architecture/neural_operators/fno2d.py @@ -1,127 +1,141 @@ -"""FNO2d""" -import mindspore as ms -from mindspore import ops, nn, Tensor -from mindspore.ops import operations as P - -from sciai.architecture.neural_operators.dft import SpectralConv2dDft -from sciai.utils.check_utils import _check_type -from sciai.utils.math_utils import _get_grid_2d - - -class FNOBlock(nn.Cell): - """FNO layer""" - - def __init__(self, in_channels, out_channels, modes1, resolution=211, gelu=True, dtype=ms.float16): - super().__init__() - self.conv = SpectralConv2dDft(in_channels, out_channels, modes1, modes1, resolution, resolution, dtype=dtype) - self.w = nn.Conv2d(in_channels, out_channels, 1, weight_init='HeUniform').to_float(dtype) - self.act = ops.GeLU() if gelu else ops.Identity() - - def construct(self, x): - return self.act(self.conv(x) + self.w(x)) - - -class FNO2D(nn.Cell): - r""" - The 2-dimensional Fourier Neural Operator (FNO2D) contains a lifting layer, - multiple Fourier layers and a decoder layer. - The details can be found in `Fourier neural operator for parametric - partial differential equations `_. - - Args: - in_channels (int): The number of channels in the input space. - out_channels (int): The number of channels in the output space. - resolution (int): The spatial resolution of the input. - modes (int): The number of low-frequency components to keep. - channels (int): The number of channels after dimension lifting of the input. Default: 20. - depths (int): The number of FNO layers. Default: 4. - mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. - dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. - `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. - Default: `ms.float32`. - - Inputs: - - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, resolution, in\_channels)`. - - Outputs: - Tensor, the output of this FNO network. - - - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, resolution, out\_channels)`. - - Raises: - TypeError: If `in_channels` is not an int. - TypeError: If `out_channels` is not an int. - TypeError: If `resolution` is neither an int nor a tuple of int. - TypeError: If `modes` is not an int. - ValueError: If `modes` is less than 1. - - Supported Platforms: - `Ascend` `GPU` - - Examples: - >>> import numpy as np - >>> from mindspore.common.initializer import initializer, Normal - >>> from sciai.architecture.neural_operators import FNO2D - >>> B, H, W, C = 32, 64, 64, 1 - >>> x = initializer(Normal(), [B, H, W, C]) - >>> net = FNO2D(in_channels=1, out_channels=1, resolution=64, modes=12) - >>> output = net(x) - >>> print(output.shape) - (32, 64, 64, 1) - """ - - def __init__(self, - in_channels, - out_channels, - resolution, - modes, - channels=20, - depths=4, - mlp_ratio=4, - dtype=ms.float32): - super().__init__() - _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) - _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) - _check_type(resolution, "resolution", target_type=(int, tuple), exclude_type=bool) - _check_type(modes, "modes", target_type=int, exclude_type=bool) - if modes < 1: - raise ValueError( - "modes must at least 1, but got mode: {}".format(modes)) - - self.modes1 = modes - self.channels = channels - self.fc_channel = mlp_ratio * channels - self.fc0 = nn.Dense(in_channels + 2, self.channels, has_bias=False).to_float(dtype) - self.layers = depths - - self.fno_seq = nn.SequentialCell() - for _ in range(self.layers - 1): - self.fno_seq.append( - FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) - self.fno_seq.append( - FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) - - self.fc1 = nn.Dense(self.channels, self.fc_channel, has_bias=False).to_float(dtype) - self.fc2 = nn.Dense(self.fc_channel, out_channels, has_bias=False).to_float(dtype) - - self.grid = Tensor(_get_grid_2d(resolution), dtype=ms.float32) - self.concat = ops.Concat(axis=-1) - self.transpose = P.Transpose() - self.act = ops.GeLU() - - def construct(self, x: Tensor): - """construct""" - batch_size = x.shape[0] - - grid = self.grid.repeat(batch_size, axis=0) - x = self.concat((x, grid)) - x = self.fc0(x) - x = self.transpose(x, (0, 3, 1, 2)) - - x = self.fno_seq(x) - - x = self.transpose(x, (0, 2, 3, 1)) - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - return x +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""FNO2d""" +import mindspore as ms +from mindspore import ops, nn, Tensor +from mindspore.ops import operations as P + +from sciai.architecture.neural_operators.dft import SpectralConv2dDft +from sciai.utils.check_utils import _check_type +from sciai.utils.math_utils import _get_grid_2d + + +class FNOBlock(nn.Cell): + """FNO layer""" + + def __init__(self, in_channels, out_channels, modes1, resolution=211, gelu=True, dtype=ms.float16): + super().__init__() + self.conv = SpectralConv2dDft(in_channels, out_channels, modes1, modes1, resolution, resolution, dtype=dtype) + self.w = nn.Conv2d(in_channels, out_channels, 1, weight_init='HeUniform').to_float(dtype) + self.act = ops.GeLU() if gelu else ops.Identity() + + def construct(self, x): + return self.act(self.conv(x) + self.w(x)) + + +class FNO2D(nn.Cell): + r""" + The 2-dimensional Fourier Neural Operator (FNO2D) contains a lifting layer, + multiple Fourier layers and a decoder layer. + The details can be found in `Fourier neural operator for parametric + partial differential equations `_. + + Args: + in_channels (int): The number of channels in the input space. + out_channels (int): The number of channels in the output space. + resolution (int): The spatial resolution of the input. + modes (int): The number of low-frequency components to keep. + channels (int): The number of channels after dimension lifting of the input. Default: 20. + depths (int): The number of FNO layers. Default: 4. + mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. + dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. + `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. + Default: `ms.float32`. + + Inputs: + - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, resolution, in\_channels)`. + + Outputs: + Tensor, the output of this FNO network. + + - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, resolution, out\_channels)`. + + Raises: + TypeError: If `in_channels` is not an int. + TypeError: If `out_channels` is not an int. + TypeError: If `resolution` is neither an int nor a tuple of int. + TypeError: If `modes` is not an int. + ValueError: If `modes` is less than 1. + + Supported Platforms: + `Ascend` `GPU` + + Examples: + >>> import numpy as np + >>> from mindspore.common.initializer import initializer, Normal + >>> from sciai.architecture.neural_operators import FNO2D + >>> B, H, W, C = 32, 64, 64, 1 + >>> x = initializer(Normal(), [B, H, W, C]) + >>> net = FNO2D(in_channels=1, out_channels=1, resolution=64, modes=12) + >>> output = net(x) + >>> print(output.shape) + (32, 64, 64, 1) + """ + + def __init__(self, + in_channels, + out_channels, + resolution, + modes, + channels=20, + depths=4, + mlp_ratio=4, + dtype=ms.float32): + super().__init__() + _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) + _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) + _check_type(resolution, "resolution", target_type=(int, tuple), exclude_type=bool) + _check_type(modes, "modes", target_type=int, exclude_type=bool) + if modes < 1: + raise ValueError( + "modes must at least 1, but got mode: {}".format(modes)) + + self.modes1 = modes + self.channels = channels + self.fc_channel = mlp_ratio * channels + self.fc0 = nn.Dense(in_channels + 2, self.channels, has_bias=False).to_float(dtype) + self.layers = depths + + self.fno_seq = nn.SequentialCell() + for _ in range(self.layers - 1): + self.fno_seq.append( + FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) + self.fno_seq.append( + FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) + + self.fc1 = nn.Dense(self.channels, self.fc_channel, has_bias=False).to_float(dtype) + self.fc2 = nn.Dense(self.fc_channel, out_channels, has_bias=False).to_float(dtype) + + self.grid = Tensor(_get_grid_2d(resolution), dtype=ms.float32) + self.concat = ops.Concat(axis=-1) + self.transpose = P.Transpose() + self.act = ops.GeLU() + + def construct(self, x: Tensor): + """construct""" + batch_size = x.shape[0] + + grid = self.grid.repeat(batch_size, axis=0) + x = self.concat((x, grid)) + x = self.fc0(x) + x = self.transpose(x, (0, 3, 1, 2)) + + x = self.fno_seq(x) + + x = self.transpose(x, (0, 2, 3, 1)) + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x diff --git a/SciAI/sciai/architecture/neural_operators/fno3d.py b/SciAI/sciai/architecture/neural_operators/fno3d.py index 10f80bc2d..1d0c45fb5 100644 --- a/SciAI/sciai/architecture/neural_operators/fno3d.py +++ b/SciAI/sciai/architecture/neural_operators/fno3d.py @@ -1,130 +1,144 @@ -"""FNO3D""" -import mindspore as ms -from mindspore import ops, nn, Tensor - -from sciai.architecture.neural_operators.dft import SpectralConv3d -from sciai.utils.check_utils import _check_type -from sciai.utils.math_utils import _get_grid_3d, _to_3tuple - - -class FNOBlock(nn.Cell): - """FNO layer""" - - def __init__(self, in_channels, out_channels, modes1, resolution=1024, gelu=True, dtype=ms.float32): - super().__init__() - resolution = _to_3tuple(resolution) - self.conv = SpectralConv3d(in_channels, out_channels, modes1, modes1, - modes1, resolution[0], resolution[1], resolution[2], dtype=dtype) - self.w = nn.Conv3d(in_channels, out_channels, 1).to_float(dtype) - self.act = ops.GeLU() if gelu else ops.Identity() - - def construct(self, x): - """residual output""" - return self.act(self.conv(x) + self.w(x)) + x - - -class FNO3D(nn.Cell): - r""" - The 3-dimensional Fourier Neural Operator (FNO3D) contains a lifting layer, - multiple Fourier layers and a decoder layer. - The details can be found in `Fourier neural operator for - parametric partial differential equations `_. - - Args: - in_channels (int): The number of channels in the input space. - out_channels (int): The number of channels in the output space. - resolution (Union[Number, tuple[Number]]): The spatial resolution of the input. - modes (int): The number of low-frequency components to keep. - channels (int): The number of channels after dimension lifting of the input. Default: 20. - depths (int): The number of FNO layers. Default: 4. - mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. - dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. - `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. - Default: `ms.float32`. - - Inputs: - - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, resolution, resolution, in\_channels)`. - - Outputs: - Tensor, the output of this FNO network. - - - **output** (Tensor) -Tensor of shape - - :math:`(batch\_size, resolution, resolution, resolution, out\_channels)`. - - Raises: - TypeError: If `in_channels` is not an int. - TypeError: If `out_channels` is not an int. - TypeError: If `resolution` is neither an int nor a tuple of int. - TypeError: If `modes` is not an int. - ValueError: If `modes` is less than 1. - - Supported Platforms: - `Ascend` `GPU` - - Examples: - >>> import numpy as np - >>> from mindspore.common.initializer import initializer, Normal - >>> from sciai.architecture.neural_operators import FNO3D - >>> B, H, W, L, C = 2, 64, 64, 64, 1 - >>> x = initializer(Normal(), [B, C, H, W, L]) - >>> net = FNO3D(in_channels=1, out_channels=1, resolution=64, modes=12) - >>> output = net(x) - >>> print(output.shape) - (2, 64, 64, 64, 1) - """ - - def __init__(self, - in_channels, - out_channels, - resolution, - modes, - channels=20, - depths=4, - mlp_ratio=4, - dtype=ms.float32): - super().__init__() - _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) - _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) - _check_type(resolution, "resolution", target_type=(int, tuple), exclude_type=bool) - _check_type(modes, "modes", target_type=int, exclude_type=bool) - if modes < 1: - raise ValueError( - "modes must at least 1, but got mode: {}".format(modes)) - self.modes1 = modes - self.channels = channels - self.fc_channel = mlp_ratio * channels - self.fc0 = nn.Dense(in_channels + 3, self.channels, has_bias=False).to_float(dtype) - self.layers = depths - - self.fno_seq = nn.SequentialCell() - for _ in range(self.layers - 1): - self.fno_seq.append( - FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) - self.fno_seq.append( - FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) - - self.fc1 = nn.Dense(self.channels, self.fc_channel, has_bias=False).to_float(dtype) - self.fc2 = nn.Dense(self.fc_channel, out_channels, has_bias=False).to_float(dtype) - - self.grid = Tensor(_get_grid_3d(resolution), dtype=dtype) - self.concat = ops.Concat(axis=-1) - self.transpose = ops.Transpose() - self.act = ops.GeLU() - - def construct(self, x: Tensor): - """construct""" - batch_size = x.shape[0] - - grid = self.grid.repeat(batch_size, axis=0) - x = self.concat((x, grid)) - x = self.fc0(x) - x = self.transpose(x, (0, 4, 1, 2, 3)) - - x = self.fno_seq(x) - - x = self.transpose(x, (0, 2, 3, 4, 1)) - x = self.fc1(x) - x = self.act(x) - x = self.fc2(x) - return x +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""FNO3D""" +import mindspore as ms +from mindspore import ops, nn, Tensor + +from sciai.architecture.neural_operators.dft import SpectralConv3d +from sciai.utils.check_utils import _check_type +from sciai.utils.math_utils import _get_grid_3d, _to_3tuple + + +class FNOBlock(nn.Cell): + """FNO layer""" + + def __init__(self, in_channels, out_channels, modes1, resolution=1024, gelu=True, dtype=ms.float32): + super().__init__() + resolution = _to_3tuple(resolution) + self.conv = SpectralConv3d(in_channels, out_channels, modes1, modes1, + modes1, resolution[0], resolution[1], resolution[2], dtype=dtype) + self.w = nn.Conv3d(in_channels, out_channels, 1).to_float(dtype) + self.act = ops.GeLU() if gelu else ops.Identity() + + def construct(self, x): + """residual output""" + return self.act(self.conv(x) + self.w(x)) + x + + +class FNO3D(nn.Cell): + r""" + The 3-dimensional Fourier Neural Operator (FNO3D) contains a lifting layer, + multiple Fourier layers and a decoder layer. + The details can be found in `Fourier neural operator for + parametric partial differential equations `_. + + Args: + in_channels (int): The number of channels in the input space. + out_channels (int): The number of channels in the output space. + resolution (Union[Number, tuple[Number]]): The spatial resolution of the input. + modes (int): The number of low-frequency components to keep. + channels (int): The number of channels after dimension lifting of the input. Default: 20. + depths (int): The number of FNO layers. Default: 4. + mlp_ratio (int): The number of channels lifting ratio of the decoder layer. Default: 4. + dtype (dtype.Number): The computation type of dense. It should be `ms.float16` or `ms.float32`. + `ms.float32` is recommended for the GPU backend, and `ms.float16` is recommended for the Ascend backend. + Default: `ms.float32`. + + Inputs: + - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, resolution, resolution, in\_channels)`. + + Outputs: + Tensor, the output of this FNO network. + + - **output** (Tensor) -Tensor of shape + + :math:`(batch\_size, resolution, resolution, resolution, out\_channels)`. + + Raises: + TypeError: If `in_channels` is not an int. + TypeError: If `out_channels` is not an int. + TypeError: If `resolution` is neither an int nor a tuple of int. + TypeError: If `modes` is not an int. + ValueError: If `modes` is less than 1. + + Supported Platforms: + `Ascend` `GPU` + + Examples: + >>> import numpy as np + >>> from mindspore.common.initializer import initializer, Normal + >>> from sciai.architecture.neural_operators import FNO3D + >>> B, H, W, L, C = 2, 64, 64, 64, 1 + >>> x = initializer(Normal(), [B, C, H, W, L]) + >>> net = FNO3D(in_channels=1, out_channels=1, resolution=64, modes=12) + >>> output = net(x) + >>> print(output.shape) + (2, 64, 64, 64, 1) + """ + + def __init__(self, + in_channels, + out_channels, + resolution, + modes, + channels=20, + depths=4, + mlp_ratio=4, + dtype=ms.float32): + super().__init__() + _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) + _check_type(out_channels, "out_channels", target_type=int, exclude_type=bool) + _check_type(resolution, "resolution", target_type=(int, tuple), exclude_type=bool) + _check_type(modes, "modes", target_type=int, exclude_type=bool) + if modes < 1: + raise ValueError( + "modes must at least 1, but got mode: {}".format(modes)) + self.modes1 = modes + self.channels = channels + self.fc_channel = mlp_ratio * channels + self.fc0 = nn.Dense(in_channels + 3, self.channels, has_bias=False).to_float(dtype) + self.layers = depths + + self.fno_seq = nn.SequentialCell() + for _ in range(self.layers - 1): + self.fno_seq.append( + FNOBlock(self.channels, self.channels, modes1=self.modes1, resolution=resolution, dtype=dtype)) + self.fno_seq.append( + FNOBlock(self.channels, self.channels, self.modes1, resolution=resolution, gelu=False, dtype=dtype)) + + self.fc1 = nn.Dense(self.channels, self.fc_channel, has_bias=False).to_float(dtype) + self.fc2 = nn.Dense(self.fc_channel, out_channels, has_bias=False).to_float(dtype) + + self.grid = Tensor(_get_grid_3d(resolution), dtype=dtype) + self.concat = ops.Concat(axis=-1) + self.transpose = ops.Transpose() + self.act = ops.GeLU() + + def construct(self, x: Tensor): + """construct""" + batch_size = x.shape[0] + + grid = self.grid.repeat(batch_size, axis=0) + x = self.concat((x, grid)) + x = self.fc0(x) + x = self.transpose(x, (0, 4, 1, 2, 3)) + + x = self.fno_seq(x) + + x = self.transpose(x, (0, 2, 3, 4, 1)) + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x diff --git a/SciAI/sciai/architecture/neural_operators/kno1d.py b/SciAI/sciai/architecture/neural_operators/kno1d.py index b38197ffb..67fe8e4d9 100644 --- a/SciAI/sciai/architecture/neural_operators/kno1d.py +++ b/SciAI/sciai/architecture/neural_operators/kno1d.py @@ -1,99 +1,113 @@ -"""KNO1D""" -import mindspore.common.dtype as mstype -from mindspore import ops, nn, Tensor - -from sciai.utils.check_utils import _check_type -from .dft import SpectralConv1dDft - - -class KNO1D(nn.Cell): - r""" - The 1-dimensional Koopman Neural Operator (KNO1D) contains a encoder layer and a decoder layer, - multiple Koopman layers. - The details can be found in `KoopmanLab: machine learning for solving complex physics equations - `_. - - Args: - in_channels (int): The number of channels in the input space. Default: ``1``. - channels (int): The number of channels after dimension lifting of the input. Default: ``32``. - modes (int): The number of low-frequency components to keep. Default: ``16``. - resolution (int): The spatial resolution of the input. Default: ``1024``. - depths (int): The number of KNO layers. Default: ``4``. - compute_dtype (dtype.Number): The computation type of dense. Default: ``mstype.float16``. - Should be ``mstype.float32`` or ``mstype.float16``. mstype.float32 is recommended for - the GPU backend, mstype.float16 is recommended for the Ascend backend. - - Inputs: - - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. - - Outputs: - Tensor, the output of this KNO network. - - - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. - - Raises: - TypeError: If `in_channels` is not an int. - TypeError: If `channels` is not an int. - TypeError: If `modes` is not an int. - TypeError: If `depths` is not an int. - TypeError: If `resolution` is not an int. - - Supported Platforms: - ``Ascend`` ``GPU`` - - Examples: - >>> import numpy as np - >>> from sciai.architecture.neural_operators import KNO1D - >>> input_ = Tensor(np.ones([32, 1024, 1]), mstype.float32) - >>> net = KNO1D() - >>> x, x_reconstruct = net(input_) - >>> print(x.shape, x_reconstruct.shape) - (32, 1024, 1) (32, 1024, 1) - """ - - def __init__(self, - in_channels=1, - channels=32, - modes=16, - depths=4, - resolution=1024, - compute_dtype=mstype.float32): - super().__init__() - _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) - _check_type(channels, "channels", target_type=int, exclude_type=bool) - _check_type(modes, "modes", target_type=int, exclude_type=bool) - _check_type(depths, "depths", target_type=int, exclude_type=bool) - _check_type(resolution, "resolution", target_type=int, exclude_type=bool) - self.in_channels = in_channels - self.channels = channels - self.modes = modes - self.depths = depths - self.resolution = resolution - self.enc = nn.Dense(in_channels, channels, has_bias=True) - self.dec = nn.Dense(channels, in_channels, has_bias=True) - self.koopman_layer = SpectralConv1dDft(channels, channels, modes, resolution, dtype=compute_dtype) - self.w0 = nn.Conv1d(channels, channels, 1, has_bias=True) - - def construct(self, x: Tensor): - """KNO1D forward function. - - Args: - x (Tensor): Input Tensor. - """ - # reconstruct - x_reconstruct = self.enc(x) - x_reconstruct = ops.tanh(x_reconstruct) - x_reconstruct = self.dec(x_reconstruct) - - # predict - x = self.enc(x) - x = ops.tanh(x) - x = x.transpose(0, 2, 1) - x_w = x - for _ in range(self.depths): - x1 = self.koopman_layer(x) - x = ops.tanh(x + x1) - x = ops.tanh(self.w0(x_w) + x) - x = x.transpose(0, 2, 1) - x = self.dec(x) - return x, x_reconstruct +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""KNO1D""" +import mindspore.common.dtype as mstype +from mindspore import ops, nn, Tensor + +from sciai.utils.check_utils import _check_type +from .dft import SpectralConv1dDft + + +class KNO1D(nn.Cell): + r""" + The 1-dimensional Koopman Neural Operator (KNO1D) contains a encoder layer and a decoder layer, + multiple Koopman layers. + The details can be found in `KoopmanLab: machine learning for solving complex physics equations + `_. + + Args: + in_channels (int): The number of channels in the input space. Default: ``1``. + channels (int): The number of channels after dimension lifting of the input. Default: ``32``. + modes (int): The number of low-frequency components to keep. Default: ``16``. + resolution (int): The spatial resolution of the input. Default: ``1024``. + depths (int): The number of KNO layers. Default: ``4``. + compute_dtype (dtype.Number): The computation type of dense. Default: ``mstype.float16``. + Should be ``mstype.float32`` or ``mstype.float16``. mstype.float32 is recommended for + the GPU backend, mstype.float16 is recommended for the Ascend backend. + + Inputs: + - **x** (Tensor) - Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. + + Outputs: + Tensor, the output of this KNO network. + + - **output** (Tensor) -Tensor of shape :math:`(batch\_size, resolution, in\_channels)`. + + Raises: + TypeError: If `in_channels` is not an int. + TypeError: If `channels` is not an int. + TypeError: If `modes` is not an int. + TypeError: If `depths` is not an int. + TypeError: If `resolution` is not an int. + + Supported Platforms: + ``Ascend`` ``GPU`` + + Examples: + >>> import numpy as np + >>> from sciai.architecture.neural_operators import KNO1D + >>> input_ = Tensor(np.ones([32, 1024, 1]), mstype.float32) + >>> net = KNO1D() + >>> x, x_reconstruct = net(input_) + >>> print(x.shape, x_reconstruct.shape) + (32, 1024, 1) (32, 1024, 1) + """ + + def __init__(self, + in_channels=1, + channels=32, + modes=16, + depths=4, + resolution=1024, + compute_dtype=mstype.float32): + super().__init__() + _check_type(in_channels, "in_channels", target_type=int, exclude_type=bool) + _check_type(channels, "channels", target_type=int, exclude_type=bool) + _check_type(modes, "modes", target_type=int, exclude_type=bool) + _check_type(depths, "depths", target_type=int, exclude_type=bool) + _check_type(resolution, "resolution", target_type=int, exclude_type=bool) + self.in_channels = in_channels + self.channels = channels + self.modes = modes + self.depths = depths + self.resolution = resolution + self.enc = nn.Dense(in_channels, channels, has_bias=True) + self.dec = nn.Dense(channels, in_channels, has_bias=True) + self.koopman_layer = SpectralConv1dDft(channels, channels, modes, resolution, dtype=compute_dtype) + self.w0 = nn.Conv1d(channels, channels, 1, has_bias=True) + + def construct(self, x: Tensor): + """KNO1D forward function. + + Args: + x (Tensor): Input Tensor. + """ + # reconstruct + x_reconstruct = self.enc(x) + x_reconstruct = ops.tanh(x_reconstruct) + x_reconstruct = self.dec(x_reconstruct) + + # predict + x = self.enc(x) + x = ops.tanh(x) + x = x.transpose(0, 2, 1) + x_w = x + for _ in range(self.depths): + x1 = self.koopman_layer(x) + x = ops.tanh(x + x1) + x = ops.tanh(self.w0(x_w) + x) + x = x.transpose(0, 2, 1) + x = self.dec(x) + return x, x_reconstruct diff --git a/SciAI/sciai/architecture/neural_operators/kno2d.py b/SciAI/sciai/architecture/neural_operators/kno2d.py index a6e87790b..eeeebe226 100644 --- a/SciAI/sciai/architecture/neural_operators/kno2d.py +++ b/SciAI/sciai/architecture/neural_operators/kno2d.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """KNO2D""" import mindspore.common.dtype as mstype from mindspore import ops, nn, Tensor diff --git a/SciAI/sciai/architecture/neural_operators/m2k.py b/SciAI/sciai/architecture/neural_operators/m2k.py index 9827166b5..e158e8762 100644 --- a/SciAI/sciai/architecture/neural_operators/m2k.py +++ b/SciAI/sciai/architecture/neural_operators/m2k.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """Moment(sum rules) and Kernel(convolution kernel) converter""" import mindspore as ms from mindspore import nn, ops, Tensor diff --git a/SciAI/sciai/architecture/neural_operators/pdenet.py b/SciAI/sciai/architecture/neural_operators/pdenet.py index 2e6d4ba19..7e475070d 100644 --- a/SciAI/sciai/architecture/neural_operators/pdenet.py +++ b/SciAI/sciai/architecture/neural_operators/pdenet.py @@ -1,222 +1,236 @@ -"""pde net model""" -import mindspore as ms -import mindspore.numpy as ms_np -from mindspore import nn, ops, Parameter - -from sciai.architecture.neural_operators.m2k import _M2K -from sciai.utils.check_utils import _check_type - - -def _count_num_filter(max_order): - count = 0 - for i in range(max_order + 1): - for j in range(max_order + 1): - if i + j <= max_order: - count += 1 - return count - - -class PDENet(nn.Cell): - r""" - The PDE-Net model. - - PDE-Net is a feed-forward deep network to fulfill two objectives at the same time: to accurately predict dynamics of - complex systems and to uncover the underlying hidden PDE models. The basic idea is to learn differential operators - by learning convolution kernels (filters), and apply neural networks or other machine learning methods to - approximate the unknown nonlinear responses. A special feature of the proposed PDE-Net is that all filters are - properly constrained, which enables us to easily identify the governing PDE models while still maintaining the - expressive and predictive power of the network. These constrains are carefully designed by fully exploiting the - relation between the orders of differential operators and the orders of sum rules of filters (an important concept - originated from wavelet theory). - - For more details, please refers to the paper `PDE-Net: Learning PDEs from Data - `_. - - Args: - height (int): The height number of the input and output tensor of the PDE-Net. - width (int): The width number of the input and output tensor of the PDE-Net. - channels (int): The channel number of the input and output tensor of the PDE-Net. - kernel_size (int): Specifies the height and width of the 2D convolution kernel. - max_order (int): The max order of the PDE models. - dx (float): The spatial resolution of x dimension. Default: ``0.01``. - dy (float): The spatial resolution of y dimension. Default: ``0.01``. - dt (float): The time step of the PDE-Net. Default: ``0.01``. - periodic (bool): Specifies whether periodic pad is used with convolution kernels. Default: ``True``. - enable_moment (bool): Specifies whether the convolution kernels are constrained by moments. Default: ``True``. - if_fronzen (bool): Specifies whether the moment is frozen. Default: ``False``. - - Inputs: - - **input** (Tensor) - Tensor of shape :math:`(batch\_size, channels, height, width)`. - - Outputs: - Tensor, has the same shape as `input` with data type of float32. - - Raises: - TypeError: If `height`, `width`, `channels`, `kernel_size` or `max_order` is not an int. - TypeError: If `periodic`, `enable_moment`, `if_fronzen` is not a bool. - - Supported Platforms: - ``Ascend`` ``GPU`` - - Examples: - >>> import numpy as np - >>> from mindspore import Tensor - >>> import mindspore.common.dtype as mstype - >>> from sciai.architecture.neural_operators import PDENet - >>> input = Tensor(np.random.rand(1, 2, 16, 16), mstype.float32) - >>> net = PDENet(16, 16, 2, 5, 3, 2) - >>> output = net(input) - >>> print(output.shape) - (1, 2, 16, 16) - """ - - def __init__(self, - height, - width, - channels, - kernel_size, - max_order, - dx=0.01, - dy=0.01, - dt=0.01, - periodic=True, - enable_moment=True, - if_fronzen=False): - """Initialize PDE-Net.""" - super().__init__() - _check_type(height, "height", target_type=int) - _check_type(width, "width", target_type=int) - _check_type(channels, "channels", target_type=int) - _check_type(kernel_size, "kernel_size", target_type=int) - _check_type(max_order, "max_order", target_type=int) - _check_type(periodic, "periodic", target_type=bool) - _check_type(enable_moment, "enable_moment", target_type=bool) - _check_type(if_fronzen, "if_fronzen", target_type=bool) - self.in_c = channels - self.out_c = channels - self.periodic = periodic - self.kernel_size = kernel_size - self.max_order = max_order - self.delta_t = dt - self.h = height - self.w = width - self.dtype = ms.float32 - self.num_filter = _count_num_filter(max_order) - self.dx = dx - self.dy = dy - self.enable_moment = enable_moment - self.if_fronzen = if_fronzen - self.padding = int((self.kernel_size - 1) / 2) - if self.enable_moment: - self._init_moment() - else: - self.id_conv = nn.Conv2d(self.in_c, self.out_c, kernel_size=self.kernel_size, pad_mode='valid') - self.fd_conv = nn.Conv2d(self.in_c, self.num_filter - 1, kernel_size=self.kernel_size, pad_mode='valid') - self.m2k = _M2K((self.kernel_size, self.kernel_size)) - self.idx2ij = {} - if self.periodic: - self.pad = [self.padding, self.padding, self.padding, self.padding] - self.padding = 0 - self.coe_param = Parameter(ops.UniformReal(seed=2)((self.num_filter - 1, self.h, self.w))) - - def construct(self, x): - return self._one_step_forward(x) - - @property - def coe(self): - return self.coe_param - - def _one_step_forward(self, x): - """one step forward""" - if self.periodic: - x = self._periodicpad(x) - - cast = ops.Cast() - x = cast(x, self.dtype) - - if self.enable_moment: - if self.if_fronzen: - cur_moment = self.raw_moment - else: - cur_moment = self.moment * self.mask + self.raw_moment - kernel = [] - for idx in range(cur_moment.shape[0]): - kernel.append(self.m2k(cur_moment[idx])) - kernel = ops.Stack()(kernel).astype(self.dtype) - kernel = self.scale * kernel - id_kernel = kernel[0].reshape((1, 1, self.kernel_size, self.kernel_size)) - fd_kernel = kernel[1:].reshape((self.num_filter - 1, 1, self.kernel_size, self.kernel_size)) - id_conv2d = ops.Conv2D(out_channel=id_kernel.shape[0], kernel_size=self.kernel_size, pad=self.padding) - fd_conv2d = ops.Conv2D(out_channel=fd_kernel.shape[0], kernel_size=self.kernel_size, pad=self.padding) - id_out = id_conv2d(x, id_kernel) - fd_out = fd_conv2d(x, fd_kernel) - else: - id_out = self.id_conv(x) - fd_out = self.fd_conv(x) - - f = 0 - for idx in range(fd_out.shape[1]): - if idx == 0: - f = self.coe[idx] * fd_out[:, 0:1, :, :] - else: - f = f + self.coe[idx] * fd_out[:, idx:(idx + 1), :, :] - - out = id_out + f * self.delta_t - return out - - def _init_moment(self): - """init moment""" - raw_moment = ms_np.zeros((self.num_filter, self.kernel_size, self.kernel_size)) - mask = ms_np.ones((self.num_filter, self.kernel_size, self.kernel_size)) - scale = ms_np.ones((self.num_filter,)) - - self.idx2ij = {} - idx = 0 - for o1 in range(self.max_order + 1): - for o2 in range(o1 + 1): - i = o1 - o2 - j = o2 - self.idx2ij[str(idx)] = (i, j,) - raw_moment[idx, j, i] = 1 - scale[idx] = 1.0 / (self.dx ** i * self.dy ** j) - for p in range(i + j + 1): - for q in range(i + j + 1): - if p + q <= (i + j): - mask[idx, p, q] = 0 - idx += 1 - - scale = scale.reshape([self.num_filter, 1, 1]) - self.raw_moment = raw_moment - self.mask = mask - self.scale = scale - self.moment = Parameter(raw_moment) - - def _periodicpad(self, x): - """periodic pad""" - cast = ops.Cast() - x = cast(x, self.dtype) - x_dim = len(x.shape) - inputs = ops.Transpose()(x, tuple(range(x_dim - 1, -1, -1))) - i = 0 - periodic_pad = self.pad - for _ in periodic_pad: - if i + 2 >= len(periodic_pad): - break - pad_value = periodic_pad[i] - pad_next_value = periodic_pad[i + 1] - permute = list(range(x_dim)) - permute[i] = 0 - permute[0] = i - permute_tuple = tuple(permute) - inputs = ops.Transpose()(inputs, permute_tuple) - inputlist = [inputs] - if pad_value > 0: - inputlist = [inputs[-pad_value:, :, :, :], inputs] - if pad_next_value > 0: - inputlist = inputlist + [inputs[0:pad_next_value, :, :, :]] - if pad_value + pad_next_value > 0: - inputs = ops.Concat()(inputlist) - inputs = ops.Transpose()(inputs, permute_tuple) - i += 1 - x = ops.Transpose()(inputs, tuple(range(x_dim - 1, -1, -1))) - return x +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""pde net model""" +import mindspore as ms +import mindspore.numpy as ms_np +from mindspore import nn, ops, Parameter + +from sciai.architecture.neural_operators.m2k import _M2K +from sciai.utils.check_utils import _check_type + + +def _count_num_filter(max_order): + count = 0 + for i in range(max_order + 1): + for j in range(max_order + 1): + if i + j <= max_order: + count += 1 + return count + + +class PDENet(nn.Cell): + r""" + The PDE-Net model. + + PDE-Net is a feed-forward deep network to fulfill two objectives at the same time: to accurately predict dynamics of + complex systems and to uncover the underlying hidden PDE models. The basic idea is to learn differential operators + by learning convolution kernels (filters), and apply neural networks or other machine learning methods to + approximate the unknown nonlinear responses. A special feature of the proposed PDE-Net is that all filters are + properly constrained, which enables us to easily identify the governing PDE models while still maintaining the + expressive and predictive power of the network. These constrains are carefully designed by fully exploiting the + relation between the orders of differential operators and the orders of sum rules of filters (an important concept + originated from wavelet theory). + + For more details, please refers to the paper `PDE-Net: Learning PDEs from Data + `_. + + Args: + height (int): The height number of the input and output tensor of the PDE-Net. + width (int): The width number of the input and output tensor of the PDE-Net. + channels (int): The channel number of the input and output tensor of the PDE-Net. + kernel_size (int): Specifies the height and width of the 2D convolution kernel. + max_order (int): The max order of the PDE models. + dx (float): The spatial resolution of x dimension. Default: ``0.01``. + dy (float): The spatial resolution of y dimension. Default: ``0.01``. + dt (float): The time step of the PDE-Net. Default: ``0.01``. + periodic (bool): Specifies whether periodic pad is used with convolution kernels. Default: ``True``. + enable_moment (bool): Specifies whether the convolution kernels are constrained by moments. Default: ``True``. + if_fronzen (bool): Specifies whether the moment is frozen. Default: ``False``. + + Inputs: + - **input** (Tensor) - Tensor of shape :math:`(batch\_size, channels, height, width)`. + + Outputs: + Tensor, has the same shape as `input` with data type of float32. + + Raises: + TypeError: If `height`, `width`, `channels`, `kernel_size` or `max_order` is not an int. + TypeError: If `periodic`, `enable_moment`, `if_fronzen` is not a bool. + + Supported Platforms: + ``Ascend`` ``GPU`` + + Examples: + >>> import numpy as np + >>> from mindspore import Tensor + >>> import mindspore.common.dtype as mstype + >>> from sciai.architecture.neural_operators import PDENet + >>> input = Tensor(np.random.rand(1, 2, 16, 16), mstype.float32) + >>> net = PDENet(16, 16, 2, 5, 3, 2) + >>> output = net(input) + >>> print(output.shape) + (1, 2, 16, 16) + """ + + def __init__(self, + height, + width, + channels, + kernel_size, + max_order, + dx=0.01, + dy=0.01, + dt=0.01, + periodic=True, + enable_moment=True, + if_fronzen=False): + """Initialize PDE-Net.""" + super().__init__() + _check_type(height, "height", target_type=int) + _check_type(width, "width", target_type=int) + _check_type(channels, "channels", target_type=int) + _check_type(kernel_size, "kernel_size", target_type=int) + _check_type(max_order, "max_order", target_type=int) + _check_type(periodic, "periodic", target_type=bool) + _check_type(enable_moment, "enable_moment", target_type=bool) + _check_type(if_fronzen, "if_fronzen", target_type=bool) + self.in_c = channels + self.out_c = channels + self.periodic = periodic + self.kernel_size = kernel_size + self.max_order = max_order + self.delta_t = dt + self.h = height + self.w = width + self.dtype = ms.float32 + self.num_filter = _count_num_filter(max_order) + self.dx = dx + self.dy = dy + self.enable_moment = enable_moment + self.if_fronzen = if_fronzen + self.padding = int((self.kernel_size - 1) / 2) + if self.enable_moment: + self._init_moment() + else: + self.id_conv = nn.Conv2d(self.in_c, self.out_c, kernel_size=self.kernel_size, pad_mode='valid') + self.fd_conv = nn.Conv2d(self.in_c, self.num_filter - 1, kernel_size=self.kernel_size, pad_mode='valid') + self.m2k = _M2K((self.kernel_size, self.kernel_size)) + self.idx2ij = {} + if self.periodic: + self.pad = [self.padding, self.padding, self.padding, self.padding] + self.padding = 0 + self.coe_param = Parameter(ops.UniformReal(seed=2)((self.num_filter - 1, self.h, self.w))) + + def construct(self, x): + return self._one_step_forward(x) + + @property + def coe(self): + return self.coe_param + + def _one_step_forward(self, x): + """one step forward""" + if self.periodic: + x = self._periodicpad(x) + + cast = ops.Cast() + x = cast(x, self.dtype) + + if self.enable_moment: + if self.if_fronzen: + cur_moment = self.raw_moment + else: + cur_moment = self.moment * self.mask + self.raw_moment + kernel = [] + for idx in range(cur_moment.shape[0]): + kernel.append(self.m2k(cur_moment[idx])) + kernel = ops.Stack()(kernel).astype(self.dtype) + kernel = self.scale * kernel + id_kernel = kernel[0].reshape((1, 1, self.kernel_size, self.kernel_size)) + fd_kernel = kernel[1:].reshape((self.num_filter - 1, 1, self.kernel_size, self.kernel_size)) + id_conv2d = ops.Conv2D(out_channel=id_kernel.shape[0], kernel_size=self.kernel_size, pad=self.padding) + fd_conv2d = ops.Conv2D(out_channel=fd_kernel.shape[0], kernel_size=self.kernel_size, pad=self.padding) + id_out = id_conv2d(x, id_kernel) + fd_out = fd_conv2d(x, fd_kernel) + else: + id_out = self.id_conv(x) + fd_out = self.fd_conv(x) + + f = 0 + for idx in range(fd_out.shape[1]): + if idx == 0: + f = self.coe[idx] * fd_out[:, 0:1, :, :] + else: + f = f + self.coe[idx] * fd_out[:, idx:(idx + 1), :, :] + + out = id_out + f * self.delta_t + return out + + def _init_moment(self): + """init moment""" + raw_moment = ms_np.zeros((self.num_filter, self.kernel_size, self.kernel_size)) + mask = ms_np.ones((self.num_filter, self.kernel_size, self.kernel_size)) + scale = ms_np.ones((self.num_filter,)) + + self.idx2ij = {} + idx = 0 + for o1 in range(self.max_order + 1): + for o2 in range(o1 + 1): + i = o1 - o2 + j = o2 + self.idx2ij[str(idx)] = (i, j,) + raw_moment[idx, j, i] = 1 + scale[idx] = 1.0 / (self.dx ** i * self.dy ** j) + for p in range(i + j + 1): + for q in range(i + j + 1): + if p + q <= (i + j): + mask[idx, p, q] = 0 + idx += 1 + + scale = scale.reshape([self.num_filter, 1, 1]) + self.raw_moment = raw_moment + self.mask = mask + self.scale = scale + self.moment = Parameter(raw_moment) + + def _periodicpad(self, x): + """periodic pad""" + cast = ops.Cast() + x = cast(x, self.dtype) + x_dim = len(x.shape) + inputs = ops.Transpose()(x, tuple(range(x_dim - 1, -1, -1))) + i = 0 + periodic_pad = self.pad + for _ in periodic_pad: + if i + 2 >= len(periodic_pad): + break + pad_value = periodic_pad[i] + pad_next_value = periodic_pad[i + 1] + permute = list(range(x_dim)) + permute[i] = 0 + permute[0] = i + permute_tuple = tuple(permute) + inputs = ops.Transpose()(inputs, permute_tuple) + inputlist = [inputs] + if pad_value > 0: + inputlist = [inputs[-pad_value:, :, :, :], inputs] + if pad_next_value > 0: + inputlist = inputlist + [inputs[0:pad_next_value, :, :, :]] + if pad_value + pad_next_value > 0: + inputs = ops.Concat()(inputlist) + inputs = ops.Transpose()(inputs, permute_tuple) + i += 1 + x = ops.Transpose()(inputs, tuple(range(x_dim - 1, -1, -1))) + return x diff --git a/SciAI/sciai/architecture/transformer/__init__.py b/SciAI/sciai/architecture/transformer/__init__.py index 353c1b243..e14dc4229 100644 --- a/SciAI/sciai/architecture/transformer/__init__.py +++ b/SciAI/sciai/architecture/transformer/__init__.py @@ -1,2 +1,16 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """vit""" from .vit import ViT diff --git a/SciAI/sciai/architecture/transformer/layer.py b/SciAI/sciai/architecture/transformer/layer.py index 251ca041d..c7e6f6e9d 100644 --- a/SciAI/sciai/architecture/transformer/layer.py +++ b/SciAI/sciai/architecture/transformer/layer.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """transformer layers""" import mindspore as ms from mindspore import Parameter, Tensor, ops, nn diff --git a/SciAI/sciai/architecture/transformer/vit.py b/SciAI/sciai/architecture/transformer/vit.py index 4cf840ad4..b1d0fa3f9 100644 --- a/SciAI/sciai/architecture/transformer/vit.py +++ b/SciAI/sciai/architecture/transformer/vit.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """The ViT model""" import mindspore as ms import mindspore.ops.operations as P diff --git a/SciAI/sciai/common/__init__.py b/SciAI/sciai/common/__init__.py index d6a8eb8d5..8a7e5ec83 100644 --- a/SciAI/sciai/common/__init__.py +++ b/SciAI/sciai/common/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """sciai common""" from .dataset import DatasetGenerator, Sampler from .initializer import LeCunNormal, LeCunUniform, StandardUniform, XavierTruncNormal diff --git a/SciAI/sciai/common/dataset.py b/SciAI/sciai/common/dataset.py index 3ac80b892..b175b9ffe 100644 --- a/SciAI/sciai/common/dataset.py +++ b/SciAI/sciai/common/dataset.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """dataset""" import mindspore as ms import numpy as np diff --git a/SciAI/sciai/common/initializer.py b/SciAI/sciai/common/initializer.py index d151874a8..9c16dc779 100644 --- a/SciAI/sciai/common/initializer.py +++ b/SciAI/sciai/common/initializer.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """initializer""" import numpy as np import scipy.stats diff --git a/SciAI/sciai/common/optimizer.py b/SciAI/sciai/common/optimizer.py index 4f8e48fcd..74dd8ee0c 100644 --- a/SciAI/sciai/common/optimizer.py +++ b/SciAI/sciai/common/optimizer.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """optimizer""" import time from functools import reduce diff --git a/SciAI/sciai/common/train_cell.py b/SciAI/sciai/common/train_cell.py index f5019145f..96ec2b2af 100644 --- a/SciAI/sciai/common/train_cell.py +++ b/SciAI/sciai/common/train_cell.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """train_cell""" import os from numbers import Number diff --git a/SciAI/sciai/context/__init__.py b/SciAI/sciai/context/__init__.py index 8d6a497db..9ad238cc8 100644 --- a/SciAI/sciai/context/__init__.py +++ b/SciAI/sciai/context/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """sciai context""" from .context import init_project, set_context_auto diff --git a/SciAI/sciai/context/context.py b/SciAI/sciai/context/context.py index e5a75b001..67271032f 100644 --- a/SciAI/sciai/context/context.py +++ b/SciAI/sciai/context/context.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """context""" import mindspore as ms diff --git a/SciAI/sciai/model/phygeonet/config.yaml b/SciAI/sciai/model/phygeonet/config.yaml index 57f7d4fe4..fc823d4e4 100644 --- a/SciAI/sciai/model/phygeonet/config.yaml +++ b/SciAI/sciai/model/phygeonet/config.yaml @@ -7,7 +7,7 @@ save_fig: true save_ckpt_path: ./checkpoints load_ckpt_path: ./checkpoints/model_final.ckpt log_path: ./logs -figures_path: /figures +figures_path: ./figures load_data_path: ./data/case0 save_data_path: ./data/case0 lr: 0.001 diff --git a/SciAI/sciai/model/phygeonet/src/plot.py b/SciAI/sciai/model/phygeonet/src/plot.py index daa3775e5..09ca295d4 100644 --- a/SciAI/sciai/model/phygeonet/src/plot.py +++ b/SciAI/sciai/model/phygeonet/src/plot.py @@ -1,4 +1,6 @@ """plot functions for phygeonet""" +import os + import matplotlib.pyplot as plt import numpy as np import tikzplotlib @@ -26,6 +28,7 @@ def plot_train(args, ev_hist, m_res_hist, time_spent): tikzplotlib.save(f'{args.figures_path}/error.tikz') ev_hist = np.asarray(ev_hist) m_res_hist = np.asarray(m_res_hist) + os.makedirs(args.save_data_path, exist_ok=True) np.savetxt(f'{args.save_data_path}/ev_hist.txt', ev_hist) np.savetxt(f'{args.save_data_path}/m_res_hist.txt', m_res_hist) np.savetxt(f'{args.save_data_path}/time_spent.txt', np.zeros([2, 2]) + time_spent) diff --git a/SciAI/sciai/operators/__init__.py b/SciAI/sciai/operators/__init__.py index c89d470ee..014140f42 100644 --- a/SciAI/sciai/operators/__init__.py +++ b/SciAI/sciai/operators/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """sciai operators""" from .derivatives import grad from .jacobian_weights import JacobianWeights diff --git a/SciAI/sciai/operators/derivatives.py b/SciAI/sciai/operators/derivatives.py index 8aecf1541..fdcd6bcf2 100644 --- a/SciAI/sciai/operators/derivatives.py +++ b/SciAI/sciai/operators/derivatives.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== # -*- coding: utf-8 -*- """derivatives""" from mindspore import nn diff --git a/SciAI/sciai/operators/jacobian_weights.py b/SciAI/sciai/operators/jacobian_weights.py index 768ba2be2..98f96f3c6 100644 --- a/SciAI/sciai/operators/jacobian_weights.py +++ b/SciAI/sciai/operators/jacobian_weights.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """jacobian weights""" from types import FunctionType, MethodType diff --git a/SciAI/sciai/utils/__init__.py b/SciAI/sciai/utils/__init__.py index b71a84fe1..48634a98f 100644 --- a/SciAI/sciai/utils/__init__.py +++ b/SciAI/sciai/utils/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """sciai utils""" from .check_utils import to_tuple from .file_utils import make_sciai_dirs diff --git a/SciAI/sciai/utils/check_utils.py b/SciAI/sciai/utils/check_utils.py index 457bd377b..0e5b2390f 100644 --- a/SciAI/sciai/utils/check_utils.py +++ b/SciAI/sciai/utils/check_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """check utils""" from __future__ import absolute_import diff --git a/SciAI/sciai/utils/file_utils.py b/SciAI/sciai/utils/file_utils.py index e469fad24..cdce485d8 100644 --- a/SciAI/sciai/utils/file_utils.py +++ b/SciAI/sciai/utils/file_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """file utils""" import os diff --git a/SciAI/sciai/utils/log_utils.py b/SciAI/sciai/utils/log_utils.py index 8f10e938d..0c86d755b 100644 --- a/SciAI/sciai/utils/log_utils.py +++ b/SciAI/sciai/utils/log_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """log utils""" import logging import os diff --git a/SciAI/sciai/utils/math_utils.py b/SciAI/sciai/utils/math_utils.py index 6a4cde018..86f5c39b7 100644 --- a/SciAI/sciai/utils/math_utils.py +++ b/SciAI/sciai/utils/math_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """math operators""" import mindspore.numpy as mnp import numpy as np diff --git a/SciAI/sciai/utils/ms_utils.py b/SciAI/sciai/utils/ms_utils.py index facb87f70..6850b632b 100644 --- a/SciAI/sciai/utils/ms_utils.py +++ b/SciAI/sciai/utils/ms_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """ms utils""" import numbers import random diff --git a/SciAI/sciai/utils/plot_utils.py b/SciAI/sciai/utils/plot_utils.py index 65ffcfdcd..e6fc39543 100644 --- a/SciAI/sciai/utils/plot_utils.py +++ b/SciAI/sciai/utils/plot_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """plot utils""" import json import os diff --git a/SciAI/sciai/utils/python_utils.py b/SciAI/sciai/utils/python_utils.py index 86c16d3a0..2d46669cc 100644 --- a/SciAI/sciai/utils/python_utils.py +++ b/SciAI/sciai/utils/python_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """python utils""" import argparse import glob diff --git a/SciAI/sciai/utils/register_utils.py b/SciAI/sciai/utils/register_utils.py index ce6666d7f..356cae7d7 100644 --- a/SciAI/sciai/utils/register_utils.py +++ b/SciAI/sciai/utils/register_utils.py @@ -1,54 +1,68 @@ -"""Register for network model functions""" -from sciai.utils import print_log - - -class FunctionType: - """function types for each network model""" - def __init__(self): - pass - - TRAINER = 'trainer' - VAL = 'validation' - TRAIN_PREPARE = 'train_prepare' - EVAL_PREPARE = 'eval_prepare' - - -class Register: - """Register class""" - def __init__(self): - pass - - dict_functions = {FunctionType.TRAINER: {}, FunctionType.VAL: {}, FunctionType.TRAIN_PREPARE: {}} - - @classmethod - def register(cls, target, function_type=FunctionType.TRAINER): - """wrapper fot registration""" - def register_item(func_type, key, value): - cls.dict_functions[func_type][key] = value - return value - - if function_type not in cls.dict_functions.keys(): - raise ValueError(f"unsupported function type: {function_type}") - - if callable(target): - return register_item(function_type, target.__name__, target) - if isinstance(target, str): - return lambda x: register_item(function_type, target, x) - raise ValueError("target should be function or string") - - @classmethod - def get_functions(cls, model_name: str, function_type=FunctionType.TRAINER): - """return the function according to model name and function type""" - if function_type not in cls.dict_functions.keys(): - print_log("The available function types:") - for key in cls.dict_functions: - print_log(key) - raise ValueError("Invalid function type, check the available model names above") - - if model_name not in cls.dict_functions.get(function_type).keys(): - print_log("The available networks:") - for key in cls.dict_functions.get(function_type): - print_log(key) - raise ValueError("Invalid model name, check the available model names above") - - return cls.dict_functions.get(function_type).get(model_name) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Register for network model functions""" +from sciai.utils import print_log + + +class FunctionType: + """function types for each network model""" + def __init__(self): + pass + + TRAINER = 'trainer' + VAL = 'validation' + TRAIN_PREPARE = 'train_prepare' + EVAL_PREPARE = 'eval_prepare' + + +class Register: + """Register class""" + def __init__(self): + pass + + dict_functions = {FunctionType.TRAINER: {}, FunctionType.VAL: {}, FunctionType.TRAIN_PREPARE: {}} + + @classmethod + def register(cls, target, function_type=FunctionType.TRAINER): + """wrapper fot registration""" + def register_item(func_type, key, value): + cls.dict_functions[func_type][key] = value + return value + + if function_type not in cls.dict_functions.keys(): + raise ValueError(f"unsupported function type: {function_type}") + + if callable(target): + return register_item(function_type, target.__name__, target) + if isinstance(target, str): + return lambda x: register_item(function_type, target, x) + raise ValueError("target should be function or string") + + @classmethod + def get_functions(cls, model_name: str, function_type=FunctionType.TRAINER): + """return the function according to model name and function type""" + if function_type not in cls.dict_functions.keys(): + print_log("The available function types:") + for key in cls.dict_functions: + print_log(key) + raise ValueError("Invalid function type, check the available model names above") + + if model_name not in cls.dict_functions.get(function_type).keys(): + print_log("The available networks:") + for key in cls.dict_functions.get(function_type): + print_log(key) + raise ValueError("Invalid model name, check the available model names above") + + return cls.dict_functions.get(function_type).get(model_name) diff --git a/SciAI/sciai/utils/time_utils.py b/SciAI/sciai/utils/time_utils.py index aa4377a4f..dea87b9a7 100644 --- a/SciAI/sciai/utils/time_utils.py +++ b/SciAI/sciai/utils/time_utils.py @@ -1,3 +1,17 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== """time utils""" import time from datetime import datetime, timezone diff --git a/SciAI/sciai/version.py b/SciAI/sciai/version.py index 9ada633eb..3951e4bb2 100644 --- a/SciAI/sciai/version.py +++ b/SciAI/sciai/version.py @@ -1,3 +1,17 @@ -"""SciAI version""" - -__version__ = '0.1.0' +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""SciAI version""" + +__version__ = '0.1.0' diff --git a/SciAI/setup.py b/SciAI/setup.py index a1840d650..e29647c52 100644 --- a/SciAI/setup.py +++ b/SciAI/setup.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================ +# ============================================================================== """Setup.""" import os diff --git a/SciAI/tutorial/example_grad_net.py b/SciAI/tutorial/example_grad_net.py index ce66b7065..a1dcdf7b5 100644 --- a/SciAI/tutorial/example_grad_net.py +++ b/SciAI/tutorial/example_grad_net.py @@ -1,97 +1,97 @@ -# Copyright 2021 Huawei Technologies Co., Ltd -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -""" An example neural network - -This is an example of preparing and training a simple neural networks with 2 hidden layers. -We aim to give you a quick insight of how to use SciAI to train your NNs to predict function, -which is defined by differential equation, and make it easier for you to understand the models provided in SciAI. -""" - -import numpy as np - -from mindspore import nn, ops -from sciai.architecture import MSE, MLP -from sciai.common import TrainCellWithCallBack -from sciai.common.initializer import XavierTruncNormal -from sciai.context import init_project -from sciai.operators import grad -from sciai.utils import to_tensor, print_log, log_config - -init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. -log_config("./logs") # Configure the logging path - - -class ExampleLoss(nn.Cell): - """ Loss definition class""" - def __init__(self, network): - """ - Everything besides input data should be initialized in '__init__', such as networks, gradient operators etc. - Args: - network: the neural networks to be trained. - """ - super().__init__() - self.network = network - self.dy_dx = grad(net=self.network, output_index=0, input_index=0) - self.mse = MSE() - - def construct(self, x, x_bc, y_bc_true): - """ - In this method, we define how the loss is calculated. - No ground truth is given, what we know is PDEs and its boundary conditions - Args: - x: input of the neural networks - x_bc: boundary positions - y_bc_true: true value at boundary positions - - Returns: - Mean-Squared of PDE residual error and BC error - """ - y = self.network(x) - dy_dx = self.dy_dx(x) - domain_res = dy_dx - 2 * ops.div(y, x) + ops.mul(ops.pow(x, 2), ops.pow(y, 2)) # PDE residual error - - y_bc = self.network(x_bc) - bc_res = y_bc_true - y_bc - return self.mse(domain_res) + 10 * self.mse(bc_res) - - -def func(x): - """ - The function to be learned to, which is the explicit solution of following Bernoulli equation (PDE): - y' - 2y / x + x^2 * y^2 = 0 - with boundary condition: y(1) = 1, y(0) = 0 - """ - return x ** 2 / (0.2 * x ** 5 + 0.8) - - -layers = [1, 5, 5, 1] -example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") -example_loss = ExampleLoss(example_net) -example_optimizer = nn.Adam(example_loss.trainable_params(), learning_rate=1e-3) -example_trainer = TrainCellWithCallBack(example_loss, example_optimizer, loss_interval=100, time_interval=100) - -x_train = np.random.rand(100, 1) -x_bc_train = np.array([[0], [1]]) -y_bc_train = np.array([[0], [1]]) -x_train, x_bc_train, y_bc_train = to_tensor((x_train, x_bc_train, y_bc_train)) # In MindSpore, we should feed data with `ms.Tensor` type. - -for _ in range(10000): - example_trainer(x_train, x_bc_train, y_bc_train) - -x_val = np.random.rand(5, 1) -y_true = func(x_val) -y_pred = example_net(to_tensor(x_val)).asnumpy() -print_log("y_true: ", y_true) -print_log("y_pred: ", y_pred) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" An example neural network + +This is an example of preparing and training a simple neural networks with 2 hidden layers. +We aim to give you a quick insight of how to use SciAI to train your NNs to predict function, +which is defined by differential equation, and make it easier for you to understand the models provided in SciAI. +""" + +import numpy as np + +from mindspore import nn, ops +from sciai.architecture import MSE, MLP +from sciai.common import TrainCellWithCallBack +from sciai.common.initializer import XavierTruncNormal +from sciai.context import init_project +from sciai.operators import grad +from sciai.utils import to_tensor, print_log, log_config + +init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. +log_config("./logs") # Configure the logging path + + +class ExampleLoss(nn.Cell): + """ Loss definition class""" + def __init__(self, network): + """ + Everything besides input data should be initialized in '__init__', such as networks, gradient operators etc. + Args: + network: the neural networks to be trained. + """ + super().__init__() + self.network = network + self.dy_dx = grad(net=self.network, output_index=0, input_index=0) + self.mse = MSE() + + def construct(self, x, x_bc, y_bc_true): + """ + In this method, we define how the loss is calculated. + No ground truth is given, what we know is PDEs and its boundary conditions + Args: + x: input of the neural networks + x_bc: boundary positions + y_bc_true: true value at boundary positions + + Returns: + Mean-Squared of PDE residual error and BC error + """ + y = self.network(x) + dy_dx = self.dy_dx(x) + domain_res = dy_dx - 2 * ops.div(y, x) + ops.mul(ops.pow(x, 2), ops.pow(y, 2)) # PDE residual error + + y_bc = self.network(x_bc) + bc_res = y_bc_true - y_bc + return self.mse(domain_res) + 10 * self.mse(bc_res) + + +def func(x): + """ + The function to be learned to, which is the explicit solution of following Bernoulli equation (PDE): + y' - 2y / x + x^2 * y^2 = 0 + with boundary condition: y(1) = 1, y(0) = 0 + """ + return x ** 2 / (0.2 * x ** 5 + 0.8) + + +layers = [1, 5, 5, 1] +example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") +example_loss = ExampleLoss(example_net) +example_optimizer = nn.Adam(example_loss.trainable_params(), learning_rate=1e-3) +example_trainer = TrainCellWithCallBack(example_loss, example_optimizer, loss_interval=100, time_interval=100) + +x_train = np.random.rand(100, 1) +x_bc_train = np.array([[0], [1]]) +y_bc_train = np.array([[0], [1]]) +x_train, x_bc_train, y_bc_train = to_tensor((x_train, x_bc_train, y_bc_train)) # In MindSpore, we should feed data with `ms.Tensor` type. + +for _ in range(10000): + example_trainer(x_train, x_bc_train, y_bc_train) + +x_val = np.random.rand(5, 1) +y_true = func(x_val) +y_pred = example_net(to_tensor(x_val)).asnumpy() +print_log("y_true: ", y_true) +print_log("y_pred: ", y_pred) diff --git a/SciAI/tutorial/example_load_ckpt.py b/SciAI/tutorial/example_load_ckpt.py index a16f05e79..3e395dbf2 100644 --- a/SciAI/tutorial/example_load_ckpt.py +++ b/SciAI/tutorial/example_load_ckpt.py @@ -1,47 +1,47 @@ -# Copyright 2021 Huawei Technologies Co., Ltd -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -""" An example neural network - -This is an example of preparing and loading checkpoints for a simple neural networks with 2 hidden layers. -We aim to give you a quick insight of how to load pre-trained checkpoints into your NNs, -and make it easier for you to understand the models provided in SciAI. -""" - -import numpy as np - -import mindspore as ms -from sciai.architecture import MLP -from sciai.common.initializer import XavierTruncNormal -from sciai.context import init_project -from sciai.utils import to_tensor, print_log, log_config - -init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. -log_config("./logs") # Configure the logging path - - -def func(x): - """The function to be learned to""" - return x[:, 0:1] ** 2 + np.sin(x[:, 1:2]) - - -layers = [2, 5, 5, 1] -example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") -ms.load_checkpoint("./checkpoints/example_net_10000.ckpt", example_net) - -x_val = np.random.rand(5, 2) -y_true = func(x_val) -y_pred = example_net(to_tensor(x_val)).asnumpy() -print_log("y_true: ", y_true) -print_log("y_pred: ", y_pred) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" An example neural network + +This is an example of preparing and loading checkpoints for a simple neural networks with 2 hidden layers. +We aim to give you a quick insight of how to load pre-trained checkpoints into your NNs, +and make it easier for you to understand the models provided in SciAI. +""" + +import numpy as np + +import mindspore as ms +from sciai.architecture import MLP +from sciai.common.initializer import XavierTruncNormal +from sciai.context import init_project +from sciai.utils import to_tensor, print_log, log_config + +init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. +log_config("./logs") # Configure the logging path + + +def func(x): + """The function to be learned to""" + return x[:, 0:1] ** 2 + np.sin(x[:, 1:2]) + + +layers = [2, 5, 5, 1] +example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") +ms.load_checkpoint("./checkpoints/example_net_10000.ckpt", example_net) + +x_val = np.random.rand(5, 2) +y_true = func(x_val) +y_pred = example_net(to_tensor(x_val)).asnumpy() +print_log("y_true: ", y_true) +print_log("y_pred: ", y_pred) diff --git a/SciAI/tutorial/example_net.py b/SciAI/tutorial/example_net.py index cc0bf5bfd..1e26a59cf 100644 --- a/SciAI/tutorial/example_net.py +++ b/SciAI/tutorial/example_net.py @@ -1,85 +1,85 @@ -# Copyright 2021 Huawei Technologies Co., Ltd -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -""" An example neural network - -This is an example of preparing and training a simple neural networks with 2 hidden layers. -We aim to give you a quick insight of how to use SciAI to train your NNs, -and make it easier for you to understand the models provided in SciAI. -""" - -import numpy as np - -from mindspore import nn -from sciai.architecture import MSE, MLP -from sciai.common import TrainCellWithCallBack -from sciai.common.initializer import XavierTruncNormal -from sciai.utils import to_tensor, print_log, log_config -from sciai.context.context import init_project - -init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. -log_config("./logs") # Configure the logging path - - -class ExampleLoss(nn.Cell): - """ Loss definition class""" - def __init__(self, network): - """ - Everything besides input data should be initialized in '__init__', such as networks, gradient operators etc. - Args: - network: the neural networks to be trained. - """ - super().__init__() - self.network = network - self.mse = MSE() - - def construct(self, x, y_real): - """ - In this method, we define how the loss is calculated. - Args: - x: input of the neural networks - y_real: ground truth - - Returns: - Mean-Squared-Error loss - """ - y_predict = self.network(x) - return self.mse(y_predict - y_real) - - -def func(x): - """The function to be learned to""" - return x[:, 0:1] ** 2 + np.sin(x[:, 1:2]) - - -layers = [2, 5, 5, 1] -example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") -example_loss = ExampleLoss(example_net) -example_optimizer = nn.Adam(example_loss.trainable_params(), learning_rate=1e-3) -example_trainer = TrainCellWithCallBack(example_loss, example_optimizer, - loss_interval=100, time_interval=100, - ckpt_interval=10000) - -x_train = np.random.rand(1000, 2) -y_true = func(x_train) -x_train, y_true = to_tensor((x_train, y_true)) # In MindSpore, we should feed data with `ms.Tensor` type. - -for _ in range(10001): - example_trainer(x_train, y_true) - -x_val = np.random.rand(5, 2) -y_true = func(x_val) -y_pred = example_net(to_tensor(x_val)).asnumpy() -print_log("y_true: ", y_true) -print_log("y_pred: ", y_pred) +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" An example neural network + +This is an example of preparing and training a simple neural networks with 2 hidden layers. +We aim to give you a quick insight of how to use SciAI to train your NNs, +and make it easier for you to understand the models provided in SciAI. +""" + +import numpy as np + +from mindspore import nn +from sciai.architecture import MSE, MLP +from sciai.common import TrainCellWithCallBack +from sciai.common.initializer import XavierTruncNormal +from sciai.utils import to_tensor, print_log, log_config +from sciai.context.context import init_project + +init_project() # Get the correct platform automatically and set to GRAPH_MODE by default. +log_config("./logs") # Configure the logging path + + +class ExampleLoss(nn.Cell): + """ Loss definition class""" + def __init__(self, network): + """ + Everything besides input data should be initialized in '__init__', such as networks, gradient operators etc. + Args: + network: the neural networks to be trained. + """ + super().__init__() + self.network = network + self.mse = MSE() + + def construct(self, x, y_real): + """ + In this method, we define how the loss is calculated. + Args: + x: input of the neural networks + y_real: ground truth + + Returns: + Mean-Squared-Error loss + """ + y_predict = self.network(x) + return self.mse(y_predict - y_real) + + +def func(x): + """The function to be learned to""" + return x[:, 0:1] ** 2 + np.sin(x[:, 1:2]) + + +layers = [2, 5, 5, 1] +example_net = MLP(layers=layers, weight_init=XavierTruncNormal(), bias_init='zeros', activation="tanh") +example_loss = ExampleLoss(example_net) +example_optimizer = nn.Adam(example_loss.trainable_params(), learning_rate=1e-3) +example_trainer = TrainCellWithCallBack(example_loss, example_optimizer, + loss_interval=100, time_interval=100, + ckpt_interval=10000) + +x_train = np.random.rand(1000, 2) +y_true = func(x_train) +x_train, y_true = to_tensor((x_train, y_true)) # In MindSpore, we should feed data with `ms.Tensor` type. + +for _ in range(10001): + example_trainer(x_train, y_true) + +x_val = np.random.rand(5, 2) +y_true = func(x_val) +y_pred = example_net(to_tensor(x_val)).asnumpy() +print_log("y_true: ", y_true) +print_log("y_pred: ", y_pred) diff --git a/tests/st/sciai/architecture/test_activation.py b/tests/st/sciai/architecture/test_activation.py index 9f655d90b..bef2ad2bd 100644 --- a/tests/st/sciai/architecture/test_activation.py +++ b/tests/st/sciai/architecture/test_activation.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/architecture/test_basic_block.py b/tests/st/sciai/architecture/test_basic_block.py index 055fdab02..7b9079d86 100644 --- a/tests/st/sciai/architecture/test_basic_block.py +++ b/tests/st/sciai/architecture/test_basic_block.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/common/test_dataset.py b/tests/st/sciai/common/test_dataset.py index 5800c8ed2..4057c3884 100644 --- a/tests/st/sciai/common/test_dataset.py +++ b/tests/st/sciai/common/test_dataset.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/common/test_initializer.py b/tests/st/sciai/common/test_initializer.py index e0a79cb2a..450087c29 100644 --- a/tests/st/sciai/common/test_initializer.py +++ b/tests/st/sciai/common/test_initializer.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/common/test_optimizer.py b/tests/st/sciai/common/test_optimizer.py index 2cdc550a2..8170e3f2f 100644 --- a/tests/st/sciai/common/test_optimizer.py +++ b/tests/st/sciai/common/test_optimizer.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/common/test_train_cell.py b/tests/st/sciai/common/test_train_cell.py index 62abe1394..0adc98838 100644 --- a/tests/st/sciai/common/test_train_cell.py +++ b/tests/st/sciai/common/test_train_cell.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/context/test_context.py b/tests/st/sciai/context/test_context.py index 25a9f9754..37f529451 100644 --- a/tests/st/sciai/context/test_context.py +++ b/tests/st/sciai/context/test_context.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/operators/test_derivatives.py b/tests/st/sciai/operators/test_derivatives.py index c93d38fcc..6d8a8c63e 100644 --- a/tests/st/sciai/operators/test_derivatives.py +++ b/tests/st/sciai/operators/test_derivatives.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/operators/test_jacobian.py b/tests/st/sciai/operators/test_jacobian.py index 29640340c..12d8f9f53 100644 --- a/tests/st/sciai/operators/test_jacobian.py +++ b/tests/st/sciai/operators/test_jacobian.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/test_utils/basic_nets.py b/tests/st/sciai/test_utils/basic_nets.py index 2bf42f394..335b8676b 100644 --- a/tests/st/sciai/test_utils/basic_nets.py +++ b/tests/st/sciai/test_utils/basic_nets.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/test_utils/func_utils.py b/tests/st/sciai/test_utils/func_utils.py index 8f8e052a8..9485f539f 100644 --- a/tests/st/sciai/test_utils/func_utils.py +++ b/tests/st/sciai/test_utils/func_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/test_utils/test_base.py b/tests/st/sciai/test_utils/test_base.py index c0d3fe490..c0941f0df 100644 --- a/tests/st/sciai/test_utils/test_base.py +++ b/tests/st/sciai/test_utils/test_base.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_check_utils.py b/tests/st/sciai/utils/test_check_utils.py index b588aaefd..26f0f0630 100644 --- a/tests/st/sciai/utils/test_check_utils.py +++ b/tests/st/sciai/utils/test_check_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_log_utils.py b/tests/st/sciai/utils/test_log_utils.py index 7a4aea977..65748add3 100644 --- a/tests/st/sciai/utils/test_log_utils.py +++ b/tests/st/sciai/utils/test_log_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_ms_utils.py b/tests/st/sciai/utils/test_ms_utils.py index 3db0aa729..0da4eb06a 100644 --- a/tests/st/sciai/utils/test_ms_utils.py +++ b/tests/st/sciai/utils/test_ms_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_plot_utils.py b/tests/st/sciai/utils/test_plot_utils.py index f8258c911..93d9d80db 100644 --- a/tests/st/sciai/utils/test_plot_utils.py +++ b/tests/st/sciai/utils/test_plot_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_python_utils.py b/tests/st/sciai/utils/test_python_utils.py index bdadb4325..d35da1468 100644 --- a/tests/st/sciai/utils/test_python_utils.py +++ b/tests/st/sciai/utils/test_python_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/st/sciai/utils/test_time_utils.py b/tests/st/sciai/utils/test_time_utils.py index 56bc11718..d820c95cf 100644 --- a/tests/st/sciai/utils/test_time_utils.py +++ b/tests/st/sciai/utils/test_time_utils.py @@ -1,4 +1,4 @@ -# Copyright 2021 Huawei Technologies Co., Ltd +# Copyright 2023 Huawei Technologies Co., Ltd # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. -- Gitee