Introduction

Simulation of complex physical systems described by partial differential equations (PDEs) has important applications1 in multiple fields, such as chemical engineering2, aerospace engineering3 and climate prediction4. Nowadays, the shortage of fresh water is a growing problem, which is further exacerbated by population growth and economic development. Desalination is an effective way to increase local fresh water supplies5. Membrane-based separation technologies, such as, reverse osmosis (RO) have been widely applied for seawater and brackish water desalination6. In recent years, ultrapermeable membranes (UPMs) based on carbon-nanotubes7, graphene8, aquaporins9, conjugated-polymer10, fluorous oligoamide nanochannel11 and enhanced polyamide12, have attracted widespread attention due to their significant potential on enhancing water production efficacy. However, there are still some difficulties in the large-scale production of UPMs. And although UPMs have many advantages, the intensified concentration polarization (CP) and membrane fouling make them impractical to operate at a high water flux, due to the performance limitations of current membrane module on fluid mechanics and transport13. In addition, UPMs have a limiting flux, which is influenced by the mass transfer coefficient14. Therefore, the advantage of UPMs may not be as significant as originally envisioned in practice15 and UPMs may make less impact on process efficiency than the redesign of membrane module in industrial applications16. Because the redesign of membrane module could improve mass transfer coefficient and reduce pressure drop, it is crucial for maximizing the returns of UPM systems17. Optimal design of membrane module involves solving numerous three-dimensional (3D) multi-physics models in feed spacer-filled channels, which is computationally expensive using traditional numerical methods. Therefore, to address the complicated nonlinear channel flow and mass transport in RO, the development of an efficient computational approach is of great significance.

The models used to describe the complicated 3D multi-physics problems are mainly divided into two types: phenomenological models and mechanistic models18. The separation performance of the membranes is represented by measurable parameters in phenomenological models, which regard the membranes as “black boxes”. While these models can effectively characterize performance, they do not offer a mechanistic explanation. In contrast, mechanistic models facilitate a more fundamental understanding of membranes transport by connecting the separation performance observed in the membranes to their physical and chemical properties as well as those of the solute materials. However, it is difficult to build an exact mechanistic model for a complex process. Traditional computational methods, such as computational fluid dynamics (CFD), describe the multi-physics problems using many nonlinear PDEs and then apply numerical algorithms to solve the problems in discrete meshes. For these type of problems, numerical methods, including finite difference method (FDM)19, finite volume method (FVM)20 and finite element method (FEM)21, are generally computationally expensive22. In comparison, particle methods, such as Lattice Boltzmann method (LBM)23 and smoothed particle hydrodynamics (SPH)24, are advantageous in solving the multi-physics problems with complex structures. However, their computational cost increases with the complexity of the computational domains25. Advances in computing power and the advent of supercomputers have accelerated the calculation. Nevertheless, the precise 3D simulation of RO using these traditional methods is still a daunting task, due to the complex 3D structures and nonlinear items in PDEs.

Deep learning methods have been proven to have universal nonlinear fitting ability and can approximate the governing equations of complex systems26. There are three main types of deep learning methods used to solve PDEs27, finite-dimensional operators, neural-FEM and neural operators. Finite-dimensional operators use deep convolutional neural networks (CNNs) to approximate the solution operator in finite-dimensional Euclidean space25,28,29, and are widely used in fluid flow problems. However, these methods are mesh-dependent and are limited by the size of discretization and geometry of the training data27. Neural-FEM methods, such as physics-informed neural network30, deep Galerkin method26 and deep Ritz method31, directly parameterize the solution function as a neural network. Because the networks are constructed based on mathematical knowledge, they are interpretable. However, Neural-FEM methods are designed to model one specific instance of PDE, and they are computationally expensive to train new neural networks when the parameters of PDE change. Neural operators learn PDE operators, such as DeepONet32 and Fourier neural operator (FNO)27, which are mesh-independent and only need to be trained once even when the parameters of the PDE change. However, the lack of interpretability and difficulty in assessing model generalization limit their application to practical engineering problems.

FNO constructs mappings between function spaces by parameterizing an integral kernel in Fourier space and shows a strong ability to solve parametric PDEs. The zero-shot super-resolution of FNO allows a higher resolution prediction when trained with lower resolution data. This feature is important for physical problems which need direct numerical simulation at a small scale1. However, FNO can only handle regular and fixed size data, while the computational results of complex physical problems are usually unstructured and the sizes of computational domains are unfixed. In order to tackle such problems, fully convolutional neural network (FCNN) replaces fully connected layers with convolutional layers33, while spatial pyramid pooling (SPP) could generate a fixed-length representation regardless of the input size of data34, but SPP is not suitable for engineering problems where the size of the computational domains needs to be maintained. Recurrent neural network35 and transformer model36 could tackle sequences of different lengths, and are applied to solve 2D and 3D problems37,38. However, due to the strong spatial correlation between data points in this problem, converting 3D data into 1D sequence degrades prediction performance, even with positional embedding and attention modules. Moreover, acquiring large datasets to meet training requirements for engineering problems remains challenging. Therefore, in this work, the idea of FCNN33 is borrowed to modify the framework of FNO. The accuracy of the modified FNO-based method is evaluated by the high-fidelity CFD simulation data computed by large-scale high-throughput computation using supercomputing39.

Results

3D multi-physical simulations with modified FNO-based method

The 3D multi-physics problems considered in this work are solved on the complex structures via a new deep learning method, which is shown in Fig. 1. High-fidelity CFD data are obtained by solving forward problems on the “Tianhe-2” supercomputer using a high-throughput strategy. In our previous study40 on brackish water RO desalination, a hybrid model was established in coupling a 3D CFD model with a one-dimensional (1D) system-level model. Computing by the hybrid model, the calculated transmembrane pressure had less than 5% deviation and the total flow rate was almost identical with industrial data41. The industrial data were obtained from an industrial RO train in Chino Desalter I located in Chino, California. The CFD models we use in this work apply the same modeling techniques and governing equations as those used in our previous work40, the only difference between them is we modify initial and boundary conditions to adapt different processes. These 3D CFD models are discretized by unstructured meshes, which cannot be handled by FNO directly. Although the body-fitted curvilinear interpolation method can render more details42, it is challenging for complex geometries. To this end, linear interpolation is applied to convert the unstructured data to regularly organized data. To reflect the data variation in a small range, logarithmic intervals are used for interpolation in the concentration problem, and two separate networks are used for fluid flow and mass transfer. In addition, in order to deal with different sizes of input data, the fully connected layers in the original framework of FNO27 are replaced by convolution layers whose convolution kernel is 1 × 1 × 1. As shown in Fig. 1, the framework of the modified FNO-based method includes both convolution layers and Fourier layers: the convolution layers map the input spaces to high-dimensional spaces and reduce the dimensionality to that required for the output function spaces, while the Fourier layers extract information.

Fig. 1: Fast computational strategy based on high-fidelity CFD data.
figure 1

High throughput computing strategy is applied to solve 960 CFD models using the “Tianhe-2” supercomputer. The obtained high-fidelity CFD data are interpolated to obtain structured data for training the modified FNO-based model.

Full size image

The methodology involves training deep neural networks to take inlet velocity, geometry, initial concentration and coordinates of each model as input, and then producing the distribution of velocity, pressure and concentration as output. 960 CFD models with \(1.5\times {10}^{6}-7.6\times {10}^{6}\) unstructured data points are interpolated into \(2.3\times {10}^{4}-5.9\times {10}^{5}\) organized points to train and evaluate the modified FNO-based method, where 582 models are training data, 194 models are validation data and 184 models are testing data. The deep networks are trained for 500 epochs, and each network takes ~5.5 h on the computer configured with Intel(R) Xeon(R) CPU E5-2640 v4 and NVIDIA GeForce RTX 2080 Ti GPU.

The CPU time required for each CFD model simulation ranges from 0.6 to 12.9 h using 4 computational nodes on the “Tianhe-2” supercomputer. And the average time of each model predicted by the modified FNO-based method is 0.04 s using a single computer. Therefore, the modified FNO-based method is approximately \({10}^{5}-{10}^{6}\) times faster without losing accuracy. In addition, our strategy proposed in this work requires 7 times less training time than the multilayer perceptron (MLP)43 which needs 40 h training time using Intel Xeon Gold 6130 CPU and NVIDIA Tesla V100-SXM2 GPU, whereas offering better precision for the same type of problems.

Evaluation of the predictions with the modified FNO-based method

Structure similarity index measure (SSIM)44, root mean squared error (RMSE) and coefficient of determination (\({R}^{2}\)) are used to evaluate the difference between the predicted results of the modified FNO-based method and computational results of CFD. RMSE and \({R}^{2}\) focus on the point-to-point difference, while SSIM is a perceptual model that pays more attention to the visual difference of human eyes.

In order to eliminate the impact of the order of magnitude on the evaluation, the data for comparison are normalized to [0, 1] according to the maximum and minimum values of the CFD results. The evaluation of SSIM is the mean value of three cross sections at \(\frac{1}{4}{L}_{x},\frac{2}{4}{L}_{x},\frac{3}{4}{L}_{x}\), respectively. The comparisons between CFD simulation and FNO prediction are summarized in Supplementary Tables S1, S2 and Fig. 2. Compared with CFD, the average accuracy of the modified FNO-based method evaluated by SSIM, RMSE and \({R}^{2}\) is greater than 0.990, 0.970 and 0.961 in velocity, greater than 0.994, 0.967 and 0.977 in pressure, and greater than 0.989, 0.995 and 0.995 in concentration, respectively. The accuracy of MLP43, evaluated by RMSE, is 0.935, 0.983 and 0.951 for velocity, pressure and concentration, respectively. Thus, the modified FNO-based method offers better precision. According to Fig. 2, 99% of models have errors between 0.95 and 1 in SSIM. For RMSE, 83% and 17% of models fall in the ranges of 0.95–1 and 0.9–0.95 respectively. And for R2, 86% and 11% of models fall in the ranges of 0.95–1 and 0.9–0.95 respectively. However, there are a few models that deviate slightly. The main reason is that these models are constructed with unrealistic design parameters for the feed spacer, whose calculated velocity, pressure and concentration fields have greatly deviation than the others. Another reason is that pressure is trained with velocity in this work. However, their values differ by three orders of magnitude.

Fig. 2: Comparison between results of the modified FNO-based method and CFD data.
figure 2

a Structure similarity index measure (SSIM); b Root mean squared error (RMSE); c Coefficient of determination (\({R}^{2}\)). Data points represent different geometries. And red stars represent the average values.

Full size image

The three components of velocity calculated by CFD and predicted by the modified FNO-based method are shown in Supplementary videos 1–3 in the supplementary materials. And the results of 3 geometries of typical feed spacers are shown in Figs. 3–5 as examples. According to the results shown in the figures, the modified FNO-based method can well predict the three components of velocity, and the existence of high velocity zones. Some differences with the CFD results are still noted, which may be due to the loss of details caused by the selection of Fourier modes.

Fig. 3: Velocity in the x direction in 3 geometries of typical feed spacers predicted by the modified FNO-based method and calculated by CFD.
figure 3

a–c are the results calculated by CFD, while d–f are the results predicted by the modified FNO-based method of the corresponding geometries.

Full size image
Fig. 4: Velocity in the y direction in 3 geometries of typical feed spacers predicted by the modified FNO-based method and calculated by CFD.
figure 4

a–c are the results calculated by CFD, while d–f are the results predicted by the modified FNO-based method of the corresponding geometries.

Full size image
Fig. 5: Velocity in the z direction in 3 geometries of typical feed spacers predicted by the modified FNO-based method and calculated by CFD.
figure 5

a–c are the results calculated by CFD, while d–f are the results predicted by the modified FNO-based method of the corresponding geometries.

Full size image

The pressure calculated by CFD and predicted by the modified FNO-based method are shown in Supplementary video 4 in the supplementary materials. And the results of 3 geometries of typical feed spacers are shown in Fig. 6 as examples. The local pressure distribution along the transversal line of the membrane surface at \(z=0.5H,y=0.25{L}_{y},0\le x\le {L}_{x}\) is shown in Supplementary Fig. S1. According to the results shown in Fig. 6 and Supplementary Fig. S1, the modified FNO-based method accurately predicts the pressure distribution.

Fig. 6: Pressure in 3 geometries of typical feed spacers predicted by the modified FNO-based method and calculated by CFD.
figure 6

a–c are the results calculated by CFD, while d–f are the results predicted by the modified FNO-based method of the corresponding geometries.

Full size image

The concentration calculated by CFD and predicted by the modified FNO-based method are shown in supplementary video 5 in the supplementary materials. And the results of 3 geometries of typical feed spacers are shown in Fig. 7 as examples. According to the results shown in the figure, the modified FNO-based method can predict the CP phenomenon, but there are some differences with the shape of the boundary layer calculated by the CFD model. This may be caused by the loss of data when interpolating the unstructured mesh into the structured mesh, or the limited capability of the data-driven training method to learn abrupt changes at the boundary. The CP factor45, \({\rm{CPF}}={{\rm{e}}}^{{J}_{{\rm{w}}}/k}\), is influenced by the permeation flux (\({J}_{{\rm{w}}}\)) and the mass transfer coefficient (\(k\)). Therefore, if \({J}_{{\rm{w}}}\) is enhanced, and \(k\) is maintained, the CP phenomenon will worsen. Improving the cross velocity (\(\bar{u}\)) could increase \(k\). The previous study45 indicated that doubling \(k\) using a commercial membrane module would lead to a 9-fold friction loss that must be overcome by the pump. Module design could enhance \(k\) with a moderate friction loss penalty, and thereby enable to improve \({J}_{{\rm{w}}}\). However, optimal design of membrane module requires to iteratively solve many high-fidelity multi-physics models, which is computationally intensive and time-consuming using CFD simulations. Herein, the computational time of the modified FNO-based method is approximately \({10}^{5}-{10}^{6}\) times faster than the CFD approach while both methods have an equivalent accuracy. It enables fast computation for complex geometries and therefore greatly reduces the computational time of optimal design of membrane modules.

Fig. 7: Concentration in 3 geometries of typical feed spacers predicted by the modified FNO-based method and calculated by CFD.
figure 7

a–c are the results calculated by CFD, while d–f are the results predicted by the modified FNO-based method of the corresponding geometries.

Full size image

Zero-shot super-resolution

The research of Li et al.27 shows that FNO is mesh-independent, and it can be trained on a low resolution dataset and evaluated on a higher resolution dataset without seeing any higher resolution data. This feature is called zero-shot super-resolution. In this work, we study on whether the modification changes this feature of FNO.

We double the amount of data in all three directions, i.e., we change the amount of data in the new dataset to \((2{N}_{{\rm{x}}})\times (2{N}_{{\rm{y}}})\times (2{N}_{{\rm{z}}})\), where \({N}_{{\rm{x}}},{N}_{{\rm{y}}},{N}_{{\rm{z}}}\) denote the amount of data of the original training dataset in the x, y and z directions respectively.

The comparison between the higher resolution dataset and the CFD data are shown in Supplementary Tables S3 and S4. Compared with Supplementary Tables S1 and S2, the accuracy of the modified FNO-based method in the higher resolution dataset is slightly lower than that in the original dataset for most models, whereas differentiable errors are noticed in some of the cases. This may be because the higher resolution data was not seen during the training. Another reason is that the truncated modes in the Fourier layers, which are utilized for extracting information, are selected by the smallest size in the original dataset. However, the difference between the sizes of different models is quite large, and the increase in the amount of data make the difference even larger. The small number of the truncated modes will result in the Fourier layers not being able to extract enough information for prediction46. The research of Qin et al. based on traditional FNO also shows the negative impact on the prediction when training on resolutions lower than testing data47.

The three components of velocity of the higher resolution dataset predicted by the modified FNO-based method are shown in Supplementary Fig. S2. By comparing with the computational results of CFD and the predicted results of original dataset, we could conclude that although the modified FNO-based method is trained on the lower resolution dataset, it can still predict the distribution of velocity well and the existence of the high velocity zones in the higher resolution dataset. However, as with the predicted result of the original dataset, there are still some numerical differences.

The pressure and concentration of the higher resolution dataset predicted by the modified FNO-based method are shown in Supplementary Fig. S3. Compared with the computational results of CFD and the predicted results of the original dataset, the modified FNO-based method can predict the distribution of pressure and the CP phenomenon in the higher resolution dataset well, but there are still some differences with the shape of the boundary layer calculated by CFD. This may be due to the loss of information in the original resolution dataset used in the training process, which has a negative effect on the prediction accuracy of the higher resolution dataset.

In summary, the modified FNO-based method proposed in this work is mesh-independent and has the feature of zero-shot super-resolution. In this work, we simplify the computational domain to contain only a few feed spacer cells. But there are more than 1 million feed spacer cells in a typical membrane module45. Because of the complex structures of feed spacers, the finer meshes are required to obtain high-fidelity results. Therefore, it is computationally expensive for the simplified computational domain and impractical for the whole systems using CFD approach. And a recalculation is required in order to obtain finer results. In contrast, zero-shot super-resolution allows the modified FNO-based method obtain finer results in very short time without retraining. Although full-size industrial membrane module simulations have not been validated, this feature offers opportunities for future research in this domain.

Discussion

Solving the multi-physics forward and inverse problems is very crucial for the optimal design of RO desalination processes. However, it is computationally expensive using established numerical methods. Aiming at developing a universal and efficient computational method for 3D multi-physics problems with different boundary conditions and sizes of computational domains, the modified FNO-based method is proposed. The solution quality and efficiency of the newly developed method have also been validated by tests with the high-fidelity CFD simulation data. The following conclusions are drawn:

  1. (1)

    The modified FNO-based method can tackle problems with different boundary conditions and computational domain sizes. Moreover, it can achieve a second-level calculation which is \(5\times {10}^{4}-{10}^{6}\) times increase in computational speed compared to FEM, while maintaining the solution quality.

  2. (2)

    The modified FNO-based method can precisely predict the distribution of velocity, pressure and concentration of most models. Its average prediction accuracy, evaluated by RMSE, and compared with CFD results, stands at 97%, 98.5% and 98.4% for the three components of velocity, and 96.7% and 99.5% for pressure and concentration, respectively.

  3. (3)

    The modification proposed in this work does not change the feature of mesh-independence and zero-shot super-resolution of FNO. The modified FNO-based method can be trained on lower resolution dataset and evaluated on higher resolution dataset.

Because of the fast computation and the high prediction accuracy, the modified FNO-based method is expected to replace traditional CFD methods as an intelligent solver for the optimal design of membrane modules. The modified FNO-based method can quickly design optimal schemes for improving energy efficiency with respect to various feed conditions. And although the modified FNO-based method is purposed to solve UPMs problems in this work, its capability to resolve nonlinear channel flow and mass transfer model enables direct extension to other membrane systems, which have the similar fundamental problems. And we will compare the modified FNO-based method with other deep learning methods in the future. However, generalization of this method is limited by the data-driven method. In addition, the use of linear interpolation to change unstructured data to structured data may lead to the loss of information, which will make the prediction less accurate in problems which focus on details, such as the prediction of concentration. Future improvements to this method could be:

  1. (1)

    Adding physical information constraints to the modified FNO-based method to improve the generalization ability of the method, and incorporating flow field predictions into concentration network to further improve the accuracy of concentration prediction.

  2. (2)

    Adding layer normalization to the network and adjusting weights in the loss function for fluid flow problem to further improve the accuracy of prediction.

  3. (3)

    Exploring improved methods to interpolate unstructured data into structured data, such as radial basis functions interpolation48. Alternatively, investigating point cloud input methods by leveraging techniques from point cloud transformer (PCT)49 and voxel set transformer50, where point clouds are processed by coordinate-based input embedding and neighbor embedding, or voxelized to extract voxel features.

Methods

3D multi-physics modeling

In this work, the multi-physics problem of RO includes nonlinear channel flow and mass transport. Navier–Stokes equations and diffusion-convection equation are applied to describe the fluid flow

$$\left\{\begin{array}{ll} \rho ({\bf{u}}\cdot \nabla ){\bf{u}}=\nabla \cdot \left[-p{\bf{I}}+\mu (\nabla {\bf{u}}+\left(\nabla {\bf{u}}\right)^{{\rm{T}}})\right],\rho \nabla \cdot {\bf{u}}=0,\quad {\rm{in}}\,\Omega ,\\ {\bf{u}}=0,\,{\rm{on}}\,{\Gamma }_{{\rm{U}}},{\Gamma }_{{\rm{D}}},{\Gamma }_{{\rm{S}}},\\ {\bf{u}}=-{U}_{0}{\bf{n}},{U}_{0}=1.5{u}_{0}\left(1-\left(\frac{z}{2{D}_{{\rm{tot}}}}\right)^{2}\right),{\rm{on}}\,{\Gamma }_{{\rm{I}}},\\ \left[-p{\bf{I}}+\mu \left(\nabla {\bf{u}}+{(\nabla {\bf{u}})}^{{\rm{T}}}\right)\right]{\bf{n}}=-{\hat{p}}_{0}{\bf{n}},{\hat{p}}_{0}\le 0,{\rm{on}}\,{\Gamma }_{{\rm{O}}},\\ {{\bf{u}}}_{{\rm{src}}}={{\bf{u}}}_{{\rm{dst}}},{p}_{{\rm{src}}}={p}_{{\rm{dst}}},{\rm{on}}\,{\Gamma }_{{\rm{L}}},{\Gamma }_{{\rm{R}}},\end{array}\right.$$
(1)

and mass transfer characteristics.

$$\nabla \cdot ({D}_{{\rm{s}}}\nabla c)-{\bf{u}}\cdot \nabla c=0,{\rm{in}}\,\Omega ,$$
(2)

where ρ (kg m−3) and μ (Pa s) are the density and viscosity of the fluid. u (m s−1) is the velocity vector. p (Pa) denotes hydraulic pressure. The design parameter vectors \({\boldsymbol{\gamma }}\) consist of geometric parameters \({{\boldsymbol{\beta }}}_{1}=[{L}_{{\rm{tot}}},{D}_{1,{\rm{A}}},{D}_{2,{\rm{A}}},{D}_{1,{\rm{B}}},{D}_{2,{\rm{B}}},{D}_{{\rm{tot}}},\alpha ]\) and average inlet velocity magnitude \({\bar{u}}_{{\rm{ave}},{\rm{in}}}={U}_{0}\) in this work, while \({L}_{{i},{\rm{A}}}={L}_{{i},{\rm{B}}}(i=1,2,3,4,5)\) are assumed to be constant and \({L}_{{\rm{tot}},{\rm{A}}}={L}_{{\rm{tot}},{\rm{B}}}={L}_{{\rm{tot}}}\). \({D}_{{\rm{s}}}\) (m2 s−1) and c (mol m−3) denote diffusivity and molar concentration. \({\varGamma }_{\cdot }\) represents the boundary, and the specific boundaries are shown in Supplementary Fig. S4.

Supercomputing-accelerated 3D CFD simulations

In our previous work39, a Latin hypercube sampling-like method was used to obtain more representative samples \({\boldsymbol{\gamma }}\) from the design parameters space. The CFD models include flow and mass transport governed by the Navier–Stokes equations under laminar flow conditions (Eq. (1)) and the convection–diffusion equation (Eq. (2)), respectively. There is a trade-off between computational efficiency and computational time because computational time can be reduced but it will lead to the decrease of computational efficiency as the increase of computational nodes. Our previous work39 indicated that it was a good choice for parallel solution of each 3D CFD model using 4 computational nodes in order to balance computational efficiency and computational time. The maximum computational scale can reach about 3840 nodes, and the total computational time for 960 3D multi-physics models is significantly reduced by more than 3000 times using the high throughput computing strategy.

The modified FNO-based method

FNO27 learns the mapping between two infinite dimensional spaces from the data of finite input-output pairs. In the event of the observations \({\{a{}_{j},{u}_{j}\}}_{j=1}^{N}\), where \({a}_{j}\sim \mu\) is an i.i.d. sequence from the probability measure µ supported on \({{A}}\), sequence \({u}_{j}={G}^{\dagger }({a}_{j})\) is possibly tampered with noise. There could be a parametric map constructed to build an approximation of \({G}^{\dagger }\), where \({G}^{\dagger}:{{A}}\to {{U}}\) is the nonlinear map and can be used to express the solution operator of a PDE,

$$ \begin{array}{c}G:{{A}}\times \it{\varTheta} \to {{U}},\\ {\rm{orequivalently}},\\ {G}_{\theta }:{{A}}\to {{U}},\theta \in \it{\varTheta} .\end{array}$$
(3)

By applying neural network minimize loss functions \({{C}}:{{U}}\times {{U}}\to {\mathbb{R}}\),

$$\mathop{\min}\limits_{\theta \in \it{\varTheta}}{{\mathbb{E}}}_{a\sim \mu}[{{C}}(G(a,\theta),{G}^{\dagger}(a))].$$
(4)

we can obtain the representation of the solution operator.

The framework of FNO proposed by Li et al. 27 firstly maps the input data to the high-dimensional space through the fully connected layer. Then it extracts the data through several Fourier layers, and finally reduces the dimensionality to the required output function space. However, the existence of fully connected layers makes FNO unable to deal with different sizes of data. Therefore, the fully connected layers are replaced with convolution layers whose convolution kernel is 1 × 1 × 1. The modified framework is shown in Fig. 1.

The transformation in the Fourier layer is defined as27,

$${v}_{{\rm{t}}+1}(x)=\sigma (W{v}_{{\rm{t}}}(x)+{\hat{{{\rm{F}}}}}^{-1}({R}_{{\rm{\varphi }}}\cdot \hat{{\rm{{F} }}}({v}_{{\rm{t}}}))(x)),\forall x\in D,$$
(5)

where \(\sigma\) denotes the activation function and GELU function is used as the activation function in this work. Moreover, the number of Fourier layers is set to be 4. A 1 × 1 × 1 convolution layer is used to transform input data \({\it{a}}({\it{x}})\) to \({\it{v}}({\it{x}})\), while the number of output channels is 32. The output data \({\it{u}}({\it{x}})\) is transformed by two 1 × 1 × 1 convolution layers. The number of output channels in the first convolution layer is 128, while the second convolution layer is related to the predicted model.

The loss function \({{{L} }}({u}_{{\rm{true}}},{u}_{{\rm{pred}}})\) is defined by27,

$${{{L}}}({u}_{{\rm{true}}},{u}_{{\rm{pred}}})=\mathop{\sum}\limits_{j=0}^{M}\frac{{\left|\left| {u}_{{\rm{true}},{{j}}}-{u}_{{\rm{pred}},{{j}}}\right|\right|}_{2}}{{\left|\left| {u}_{{\rm{true}},{{j}}}\right|\right|}_{2}},$$
(6)

where \({u}_{{\rm{pred}}}\), \({u}_{{\rm{true}}}\) denote the predicted result of the modified FNO-based method and the computational result of CFD, respectively. M represents the total number of data points in dataset. Since the concentration in this problem only changes over a small range in the z direction, fluid flow and mass transfer are trained by two different networks. The network of fluid flow predicts the distribution of the three components of velocity and pressure. And mass transfer predicts the distribution of concentration. Geometry, inlet velocity and coordinates serve as inputs to the fluid flow network, while the mass transfer network uses geometry, initial concentration, inlet velocity and coordinates as input. The loss functions of the two networks are set as,

$$\begin{array}{c}Loss={{{L}}}({u}_{{\rm{x}}},\widehat{{u}_{{\rm{x}}}})+{{{L}}}({u}_{{\rm{y}}},\widehat{{u}_{{\rm{y}}}})+{{{L} }}({u}_{{\rm{z}}},\widehat{{u}_{{\rm{z}}}})+{{{L}}}(p,\widehat{p}),\\ Loss=100\times {{{L} }}(c,\widehat{c}),\end{array}$$
(7)

where \(\hat{\bullet }\) denotes the predicted result of the modified FNO-based method, and the other denotes the computational result of CFD.

Data preparation details

In order to construct structured data input, linear interpolation is applied with the interval of \(\varDelta x=\frac{0.044}{200}\ \rm{m},\varDelta y=\frac{0.0078}{100}\ \rm{m},\varDelta z=\frac{0.001}{50}\ \rm{m}\) in the prediction of fluid flow. To reflect the data variation in a small range in the prediction of mass transfer, as shown in Supplementary Fig. S5(a), logarithmic intervals are used for interpolation in the z direction. Specifically, logarithmic intervals are intervals small at both ends and large in the middle, as shown in Supplementary Fig. S5(b). And the other two directions still use the same intervals as fluid flow.

Since the size of each model used in this work is quite different, the truncated modes in the Fourier layers are selected by the smallest size, and set as \(mod{e}_{{\rm{x}}}=16\), \(mod{e}_{{\rm{y}}}=6\), \(mod{e}_{{\rm{z}}}=8\).

Model train and inference speed

The Adaptive moment estimation (Adam) optimizer51 is applied to update the network parameters with the learning rate lr = 0.001, and the regularization coefficient \(\lambda ={10}^{-4}\). The learning rate is set to reduce to 50% every 50 iterations. The number of epochs in this work is set to 500.

When training the modified FNO-based model, Windows 10 is used as the operating system, and Intel(R) Xeon(R) CPU E5-2640 v4 and NVIDIA GeForce RTX 2080 Ti GPU are configured. Each epoch takes about 40 s. When predicting the modified FNO-based model, Windows 11 operating system is selected as the operating system, and Intel(R) Core (TM) i7-8700 CPU and NVIDIA GeForce GTX 1060 GPU are configured. The experimental results show that the average prediction time of each model is only 0.04 s, which fully demonstrates the efficient performance of the model.