SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left\| u - \sum_ {i, j} ^ {r} a _ {i j} f _ {i} g _ {j} \right\| _ {L ^ {2} (X \times Y)} (16) \\ \leq \left\| u - \sum_ {i, j} ^ {r} a _ {i j} \phi_ {i} \psi_ {j} \right\| _ {L ^ {2} (X \times Y)} + \left\| \sum_ {i, j} ^ {r} a _ {i j} \phi_ {i} \psi_ {j} - \sum_ {i, j} ^ {r} a _ {i j} f _ {i} g _ {j} \right\| _ {L ^ {2} (X \times Y)} (17) \\ \equiv I + I I (18) \\ \end{array}
$$
Since  $|I| < \varepsilon / 2$  from (14), it is enough to estimate  $II$ . For this, we consider
$$
\begin{array}{l} \sum_ {i, j} ^ {r} a _ {i j} \phi_ {i} \psi_ {j} - \sum_ {i, j} ^ {r} a _ {i j} f _ {i} g _ {j} = \sum_ {i, j} ^ {r} a _ {i j} \phi_ {i} (\psi_ {j} - g _ {j}) + \sum_ {i, j} ^ {r} a _ {i j} (\phi_ {i} - f _ {i}) g _ {j} (19) \\ \equiv I I _ {1} + I I _ {2}. (20) \\ \end{array}
$$
We first compute  $II_{1}$ :
$$
\begin{array}{l} \left\| I I _ {1} \right\| _ {L ^ {2} (X \times Y)} ^ {2} = \int_ {X \times Y} \left\{\sum_ {i, j} ^ {r} a _ {i j} \phi_ {i} (\psi_ {j} - f _ {j}) \right\} ^ {2} d x d y (21) \\ = \int_ {X \times Y} \left\{\sum_ {j} ^ {r} \left(\sum_ {i} ^ {r} a _ {i j} \phi_ {i}\right) \left(\psi_ {j} - f _ {j}\right) \right\} ^ {2} d x d y (22) \\ \end{array}
$$
$$
A _ {j} (x) = \sum_ {i} ^ {r} a _ {i j} \phi_ {i} (x), \quad B _ {j} (y) = \psi_ {j} (y) - g _ {j} (y) \tag {23}
$$
to write  $II_{1}$  as
$$
\left\| I I _ {1} \right\| _ {L ^ {2} (X \times Y)} ^ {2} = \int_ {X \times Y} \left\{\sum_ {j} ^ {r} A _ {j} (x) B _ {j} (y) \right\} ^ {2} d x d y. \tag {24}
$$
We then apply Cauchy-Scharwz inequality to get
$$
\begin{array}{l} \left\| I I _ {1} \right\| _ {L ^ {2} (X \times Y)} ^ {2} \leq \int_ {X \times Y} \left(\sum_ {j} ^ {r} | A _ {j} (x) | ^ {2}\right) \left(\sum_ {j} ^ {r} | B _ {j} (y) | ^ {2}\right) d x d y (25) \\ = \left(\int_ {X} \sum_ {j} ^ {r} | A _ {j} (x) | ^ {2} d x\right) \left(\int_ {Y} \sum_ {j} ^ {r} | B _ {j} (y) | ^ {2} d y\right). (26) \\ \end{array}
$$
$$
\begin{array}{l} \int_ {X} \sum_ {j} ^ {r} | A _ {j} (x) | ^ {2} d x = \int_ {X} \sum_ {j} ^ {r} \left(\sum_ {i} ^ {r} a _ {i j} \phi_ {i}\right) ^ {2} d x (27) \\ = \sum_ {j} ^ {r} \left\| \sum_ {i} ^ {r} a _ {i j} \phi_ {i} \right\| _ {L (X)} ^ {2}. (28) \\ \end{array}
$$
Since  $\{\phi_i\}$  is an orthonormal basis, we see that
$$
\left\| \sum_ {i} ^ {r} a _ {i j} \phi_ {i} \right\| _ {L (X)} \leq \sum_ {i} ^ {r} | a _ {i j} | \left\| \phi_ {i} \right\| _ {L (X)} = \sum_ {i} ^ {r} | a _ {i j} |. \tag {29}
$$
$$
\begin{array}{l} \int_ {X} \sum_ {j} ^ {r} | A _ {j} (x) | ^ {2} d x = \int_ {X} \sum_ {j} ^ {r} \left(\sum_ {i} ^ {r} a _ {i j} \phi_ {i}\right) ^ {2} d x (30) \\ = \sum_ {j} ^ {r} \int_ {X} \left(\sum_ {i} ^ {r} a _ {i j} \phi_ {i}\right) ^ {2} d x (31) \\ = \sum_ {j} ^ {r} \left\| \sum_ {i} ^ {r} a _ {i j} \phi_ {i} \right\| _ {L (X)} ^ {2} (32) \\ = \sum_ {j} ^ {r} \left(\sum_ {i} ^ {r} \left| a _ {i j} \right|\right) ^ {2}. (33) \\ \end{array}
$$
Finally, we recall
$$
\left(\sum_ {i} ^ {r} \left| a _ {i j} \right|\right) ^ {2} \leq 2 \sum_ {i} ^ {r} \left| a _ {i j} \right| ^ {2} \tag {34}
$$
to conclude
$$
\int_ {X} \sum_ {j} ^ {r} | A _ {j} (x) | ^ {2} d x \leq 2 \sum_ {j} ^ {r} \sum_ {i} ^ {r} \left| a _ {i j} \right| ^ {2} <   2 \sum_ {i, j} ^ {\infty} \left| a _ {i j} \right| ^ {2} = 2 \| u \| _ {L ^ {2} (X \times Y)} ^ {2}. \tag {35}
$$
On the other hand, we have from (15)
$$
\begin{array}{l} \int_ {Y} \sum_ {j} ^ {r} | B _ {j} (y) | ^ {2} d y = \sum_ {j} ^ {r} \int_ {Y} | \psi_ {j} (y) - g _ {j} (y) | ^ {2} d y (36) \\ = \sum_ {j} ^ {r} \left\| \psi_ {j} (y) - g _ {j} (y) \right\| _ {L ^ {2} (Y)} ^ {2} (37) \\ \leq \sum_ {j} ^ {r} \frac {\varepsilon^ {2}}{9 ^ {j} \| u \| _ {L ^ {2} (X \times Y)} ^ {2}} (38) \\ <   \frac {\varepsilon^ {2}}{8 \| u \| _ {L ^ {2} (X \times Y)} ^ {2}}. (39) \\ \end{array}
$$
$$
\left\| I I _ {1} \right\| _ {L ^ {2} (X \times Y)} ^ {2} <   \frac {\varepsilon^ {2}}{8 \| u \| _ {L ^ {2} (X \times Y)} ^ {2}} 2 \| u \| _ {L ^ {2} (X \times Y)} ^ {2} = \frac {\varepsilon^ {2}}{4}. \tag {40}
$$
$$
\left\| I I _ {2} \right\| _ {L ^ {2} (X \times Y)} ^ {2} <   \frac {\varepsilon^ {2}}{4}. \tag {41}
$$
$$
\| I I \| _ {L ^ {2} (X \times Y)} ^ {2} <   \frac {\varepsilon^ {2}}{2}. \tag {42}
$$
We go back to (16) with the estimates (14) and (42) to derive
$$
\left\| u - \sum_ {i, j} ^ {r} a _ {i j} f _ {i} g _ {j} \right\| _ {L ^ {2} (X \times Y)} <   \frac {\varepsilon}{2} + \frac {\varepsilon}{2} = \varepsilon . \tag {43}
$$
Finally, we reorder the index to rewrite
$$
\begin{array}{l} \sum_ {i, j} ^ {r} a _ {i j} f _ {i} g _ {j} = \sum_ {i} ^ {\tilde {r}} b _ {i} \tilde {f} _ {i} \tilde {g} _ {i} (44) \\ = \sum_ {i} ^ {\tilde {r}} \left\{b _ {i} \tilde {f} _ {i} \right\} \tilde {g} _ {i} (45) \\ = \sum_ {i} ^ {\tilde {r}} k _ {i} \tilde {g} _ {i} (46) \\ \end{array}
$$
Without loss of generality, we rewrite  $\tilde{r}, k_i, \tilde{g}_i$  into  $r, f_i, g_i$  respectively, to complete the proof.
B Training with Physics-Informed Loss
After SPINN predicts an output function with the methods described above, the rest of the training procedure follows the same process used in conventional PINN training [15], except we use forward-mode AD to compute PDE residuals (standard back-propagation, a.k.a. reverse-mode AD for parameter updates). With the slight abuse of notation, our predicted solution function is denoted as  $\hat{u}^{(\theta)}(x,t)$  from onwards, explicitly expressing time coordinates. Given an underlying PDE (or ODE), the initial, and the boundary conditions, SPINN is trained with a 'physics-informed' loss function:
$$
\min  _ {\theta} \mathcal {L} \left(\hat {u} ^ {(\theta)} (x, t)\right) = \min  _ {\theta} \lambda_ {\mathrm {p d e}} \mathcal {L} _ {\mathrm {p d e}} + \lambda_ {\mathrm {i c}} \mathcal {L} _ {\mathrm {i c}} + \lambda_ {\mathrm {b c}} \mathcal {L} _ {\mathrm {b c}}, \tag {47}
$$
$$
\mathcal {L} _ {\mathrm {p d e}} = \int_ {\Gamma} \int_ {\Omega} \| \mathcal {N} [ \hat {u} ^ {(\theta)} ] (x, t) \| ^ {2} d x d t, \tag {48}
$$
$$
\mathcal {L} _ {\mathrm {i c}} = \int_ {\Omega} \| \hat {u} ^ {(\theta)} (x, 0) - u _ {\mathrm {i c}} (x) \| ^ {2} d x, \tag {49}
$$
$$
\mathcal {L} _ {\mathrm {b c}} = \int_ {\Gamma} \int_ {\partial \Omega} \| \mathcal {B} [ \hat {u} ^ {(\theta)} ] (x, t) - u _ {\mathrm {b c}} (x, t) \| ^ {2} d x d t, \tag {50}
$$
where  $\Omega$  is an input domain,  $\mathcal{N},\mathcal{B}$  are generic differential operators and  $u_{\mathrm{ic}},u_{\mathrm{bc}}$  are initial, boundary conditions, respectively.  $\lambda$  are weighting factors for each loss term. When calculating the PDE loss  $(\mathcal{L}_{\mathrm{pde}})$  with Monte-Carlo integral approximation, we sampled collocation points from factorized coordinates and used forward-mode AD. The remaining  $\mathcal{L}_{\mathrm{ic}}$  and  $\mathcal{L}_{\mathrm{bc}}$  are then computed with initial and boundary coordinates to regress the given conditions. By minimizing the objective loss in Eq. 47, the model output is enforced to satisfy the given equation, the initial, and the boundary conditions.
C FLOPs Estimation
The FLOPs for evaluating the derivatives can be systematically calculated by disassembling the computational graph into elementary operations such as additions and multiplications. Given a computational graph of forward pass for computing the primals, AD augments each elementary operation into other elementary operations. The FLOPs in the forward pass can be precisely calculated since it consists of a series of matrix multiplications and additions. We used the method described in [6] to estimate FLOPs for evaluating the derivatives. Table 1 shows the number of additions (ADDS) and multiplications (MULTS) in each evaluation process. Note that FLOPs is a summation of ADDS and MULTS by definition.
One thing to note here is that this is a theoretical estimation. Theoretically, the number of JVP evaluations for computing the gradient with respect to the input coordinates is  $Nd$ , when  $N$  is the number of coordinates for each axis and  $d$  is the input dimension (see section 4.3 in the main paper). However, our actual implementation of gradient calculation involves re-computing the feature representations  $f^{(\theta_i)}$ , which makes the complexity of network propagations from  $\mathcal{O}(Nd)$  to  $\mathcal{O}(Nd^2)$ . Ideally, these feature vectors can be computed only once and stored to be used later for gradient computations. Although it is still significantly more efficient than the conventional PINN's complexity ( $Nd^2 \ll N^d$ ), there is a room to bridge the gap between theoretical FLOPs and actual training runtime by further software optimization.
Table 1: The number of elementary operations for evaluating forward pass, first and second-order derivatives. The calculation is based on  $64^{3}$  collocation points in a 3-d system and the vanilla MLP settings used for diffusion, Helmholtz, and Klein-Gordon equations. We assumed that each derivative is evaluated on every coordinate axis.
SPINN (ours)	PINN (baseline)
ADDS (×106)	MULTS (×106)	ADDS (×106)	MULTS (×106)
forward pass	20	20	21,609	21,609
1st-order derivative	40	40	86,638	43,419
2nd-order derivative	80	80	130,057	87,040
MFLOPs (total)	280	390,370
D Experimental Details and Results
In this section, we provide experimental details, numerical results, and visualizations for each experiment in the main paper. Below, we show the list of examples where the solution aligns with the variable separation form or not.
separable form: 3-d Helmholtz,  $(2 + 1)$ -d Klein-Gordon,  $(3 + 1)$ -d Klein-Gordon,  $(3 + 1)$ -d Navier-Stokes,  $(5 + 1)$ -d diffusion
non-separable form:  $(2 + 1)$ -d diffusion,  $(2 + 1)$ -d Navier-Stokes,  $(2 + 1)$ -d flow mixing, 2-d L-shaped Poisson
D.1 Diffusion Equation
The diffusion equation is one of the most representative parabolic PDEs, often used for modeling the heat diffusion process. We especially choose a nonlinear diffusion equation where it can be written
$$
\partial_ {t} u = \alpha \left(\| \nabla u \| ^ {2} + u \Delta u\right), \quad x \in \Omega , t \in \Gamma , \tag {51}
$$
$$
u (x, 0) = u _ {\mathrm {i c}} (x), \quad x \in \Omega , \tag {52}
$$
$$
u (x, t) = 0, \quad x \in \partial \Omega , t \in \Gamma . \tag {53}
$$
We used diffusivity  $\alpha = 0.05$ , spatial domain  $\Omega = [-1,1]^2$ , temporal domain  $\Gamma = [0,1]$  and used superposition of three Gaussian functions for the initial condition  $u_{\mathrm{ic}}$ . We obtained the reference solution  $(101\times 101\times 101$  resolution) through a widely-used PDE solver platform FEniCS [12]. Note that FEniCS is a FEM-based solver. We particularly set the initial condition to be a superposition of three gaussian functions:
$$
\begin{array}{l} u _ {\mathrm {i c}} (x, y) = 0. 2 5 \exp \left[ - 1 0 \{(x - 0. 2) ^ {2} + (y - 0. 3) ^ {2} \} \right] + 0. 4 \exp \left[ - 1 5 \{(x + 0. 1) ^ {2} + (y + 0. 5) ^ {2} \} \right] \\ + 0. 3 \exp \left[ - 2 0 \left\{\left(x + 0. 5\right) ^ {2} + y ^ {2} \right\} \right]. \tag {54} \\ \end{array}
$$
For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors  $\lambda$  in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 3 shows the numerical results, and the visualized solutions are provided in Fig. 3.
D.2 Helmholtz Equation
The Helmholtz equation is a time-independent wave equation that takes the form:
$$
\Delta u + k ^ {2} u = q, \quad x \in \Omega , \tag {55}
$$
$$
u (x) = 0, \quad x \in \partial \Omega , \tag {56}
$$
where the spatial domain is  $\Omega = [-1,1]^3$ . For a given source term  $q = -(a_{1}\pi)^{2}u - (a_{2}\pi)^{2}u - (a_{3}\pi)^{2}u + k^{2}u$ , we devised a manufactured solution  $u = \sin(a_{1}\pi x_{1})\sin(a_{2}\pi x_{2})\sin(a_{3}\pi x_{3})$ , where we take  $k = 1, a_{1} = 4, a_{2} = 4, a_{3} = 3$ .
For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors  $\lambda$  in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 4 shows the numerical results and the visualized solutions are provided in Fig. 4.
D.3 Klein-Gordon Equation
The Klein-Gordon equation is a nonlinear hyperbolic PDE, which arises in diverse applied physics for modeling relativistic wave propagation. The inhomogeneous Klein-Gordon equation is given by
$$
\partial_ {t t} u - \Delta u + u ^ {2} = f, \quad x \in \Omega , t \in \Gamma , \tag {57}
$$
$$
u (x, 0) = x _ {1} + x _ {2}, \quad x \in \Omega , \tag {58}
$$
$$
u (x, t) = u _ {\mathrm {b c}} (x), \quad x \in \partial \Omega , t \in \Gamma , \tag {59}
$$
where we chose the spatial/temporal domain to be  $\Omega = [-1,1]^2$  and  $\Gamma = [0,10]$ , respectively. For error measurement, we used a manufactured solution  $u = (x_{1} + x_{2})\cos (2t) + x_{1}x_{2}\sin (2t)$  and  $f$ ,  $u_{bc}$  are extracted from this exact solution.
For our model, we used three body networks of 4 hidden layers with 64/32 hidden feature/output size each. For the baseline model, we used a single MLP of 5 hidden layers with 128 hidden feature sizes. We used Adam [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations for every experiment. All weight factors  $\lambda$  in the loss function in Eq. 47 are set to 1. The final reported errors are extracted where the total loss was minimum across the entire training iteration. We also resampled the input points every 100 epochs. Tab. 5 shows the numerical results.
We used the same settings used in  $(2 + 1)$ -d Klein-Gordon experiment for the  $(3 + 1)$ -d experiment except the manufactured solution was chosen as:
$$
u = \left(x _ {1} + x _ {2} + x _ {3}\right) \cos (t) + x _ {1} x _ {2} x _ {3} \sin (t), \tag {60}
$$
where  $f$ ,  $u_{bc}$  are extracted from this exact solution. Tab. 6 shows the numerical results. The number of collocation points of  $23^4$ ,  $18^4$  were the maximum value for PINN and PINN with modified MLP, respectively.
D.4 (2+1)-d Navier-Stokes Equation
Navier-Stokes equation is a nonlinear time-dependent PDE that describes the motion of a viscous fluid. Various engineering fields rely on this equation, such as modeling the weather, airflow, or ocean currents. The vorticity form for incompressible fluid can be written as below:
$$
\partial_ {t} \omega + u \cdot \nabla \omega = \nu \Delta \omega , \quad x \in \Omega , t \in \Gamma , \tag {61}
$$
$$
\nabla \cdot u = 0, \quad x \in \Omega , t \in \Gamma , \tag {62}
$$
$$
\omega (x, 0) = \omega_ {0} (x), \quad x \in \Omega , \tag {63}
$$
where  $u \in \mathbb{R}^2$  is the velocity field,  $\omega = \nabla \times u$  is the vorticity,  $\omega_0$  is the initial vorticity, and  $\nu$  is the viscosity. We used the viscosity 0.01 and made the spatial/temporal domain  $\Omega = [0,2\pi]^2$  and  $\Gamma = [0,1]$ , respectively. Note that Eq. 61 models decaying turbulence since there is no forcing term and Eq. 62 is the incompressible fluid condition. The reference solution is generated by JAX-CFD solver [9] which specifically used the pseudo-spectral method. The initial condition was generated using the gaussian random field with a maximum velocity of 5. The resolution of the obtained solution is  $100 \times 128 \times 128 \left( N_t \times N_x \times N_y \right)$ , and we tested our model on this data.
For our model, we used three body networks (modified MLP) of 3 hidden layers with 128/256 hidden feature/output sizes each. We divided the temporal domain into ten time windows to adopt the time marching method [10, 18]. We used Adam [8] optimizer with a learning rate of 0.002 and each time window is trained for 100,000 iterations. Followed by causal PINN, the PDE (residual) loss and initial condition loss function are written as follows.
$$
\mathcal {L} _ {\mathrm {p d e}} = \frac {\lambda_ {w}}{N _ {c}} \sum^ {N _ {c}} \left| \partial_ {t} w + u _ {x} \partial_ {x} w + u _ {y} \partial_ {y} w - \nu \left(\partial_ {x x} w + \partial_ {y y} w\right) \right| ^ {2} + \frac {\lambda_ {c}}{N _ {c}} \sum^ {N _ {c}} \left| \partial_ {x} u _ {x} + \partial_ {y} u _ {y} \right| ^ {2}, \tag {64}
$$
$$
\mathcal {L} _ {\mathrm {i c}} = \frac {\lambda_ {i c}}{N _ {i c}} \sum^ {N _ {i c}} \left(| u _ {x} - u _ {x 0} | ^ {2} + | u _ {y} - u _ {y 0} | ^ {2} + | w - w _ {0} | ^ {2}\right), \tag {65}
$$
where  $u_x, u_y$  are  $x, y$  components of predicted velocity,  $w = \partial_x u_y - \partial_y u_x$ ,  $N_c$  is the number of collocation points,  $N_{ic}$  is the number of coordinates for initial condition and  $u_{x0}, u_{y0}, u_{w0}$  are the initial conditions. We chose the weighting factors  $\lambda_w = 1$ ,  $\lambda_c = 5,000$ , and  $\lambda_{ic} = 10,000$ . We also resampled the input points every 100 epochs. The periodic boundary condition can be explicitly enforced by positional encoding [3], and we specifically used the following encoding function only for the spatial input coordinates.
$$
\gamma (x) = [ 1, \sin (x), \sin (2 x), \sin (3 x), \sin (4 x), \sin (5 x), \cos (x), \cos (2 x), \cos (3 x), \cos (4 x), \cos (5 x) ] ^ {\top}. \tag {66}
$$
Unlike other experiments, the solution of the Navier-Stokes equation is a 2-dimensional vector-valued function  $u: \mathbb{R}^3 \to \mathbb{R}^2$ . We can rewrite the feature merging equation Eq. 5 in the main paper to construct SPINN into a 2-dimensional vector function:
$$
u _ {1} = \sum_ {j = 1} ^ {r} \prod_ {i = 1} ^ {d} f _ {j} ^ {\left(\theta_ {i}\right)} \left(x _ {i}\right), \tag {67}
$$
$$
u _ {2} = \sum_ {j = r + 1} ^ {2 r} \prod_ {i = 1} ^ {d} f _ {j} ^ {\left(\theta_ {i}\right)} \left(x _ {i}\right). \tag {68}
$$
For example in our  $(2 + 1)$ -d Navier-Stokes equation setting,  $r = 128$  since the network output feature size is 256. This can be applied to any  $m$ -dimensional vector function if we use a larger output feature
size. More formally, if we want to construct an  $m$ -dimensional vector function with SPINN of rank  $r$ , the  $k$ -th element of the function output can be written as
$$
u _ {k} = \sum_ {j = (k - 1) r + 1} ^ {k r} \prod_ {i = 1} ^ {d} f _ {j} ^ {\left(\theta_ {i}\right)} \left(x _ {i}\right), \tag {69}
$$
where the output feature size of each body network is  $mr$ .
D.5 (3+1)-d Navier-Stokes Equation
The vorticity form of  $(3 + 1)$ -d Navier-Stokes equation is given as:
$$
\partial_ {t} \omega + (u \cdot \nabla) \omega = (\omega \cdot \nabla) u + \nu \Delta \omega + F, \quad x \in \Omega , t \in \Gamma , \tag {70}
$$
$$
\nabla \cdot u = 0, \quad x \in \Omega , t \in \Gamma , \tag {71}
$$
$$
\omega (x, 0) = \omega_ {0} (x), \quad x \in \Omega . \tag {72}
$$
We constructed the spatial/temporal domain to be  $\Omega = [0,2\pi ]^3$  and  $\Gamma = [0,5]$ , respectively. We constructed the analytic solution for  $(3 + 1)$ -d Navier-Stokes equation introduced in Taylor et al. [17]. The manufactured velocity and vorticity are
$$
u _ {x} = 2 e ^ {- 9 \nu t} \cos (2 x) \sin (2 y) \sin (z), \tag {73}
$$
$$
u _ {y} = - e ^ {- 9 \nu t} \sin (2 x) \cos (2 y) \sin (z), \tag {74}
$$
$$
u _ {z} = - 2 e ^ {- 9 \nu t} \sin (2 x) \sin (2 y) \cos (z), \tag {75}
$$
$$
\omega_ {x} = - 3 e ^ {- 9 \nu t} \sin (2 x) \cos (2 y) \cos (z), \tag {76}
$$
$$
\omega_ {y} = 6 e ^ {- 9 \nu t} \cos (2 x) \sin (2 y) \cos (z), \tag {77}
$$
$$
\omega_ {z} = - 6 e ^ {- 9 \nu t} \cos (2 x) \cos (2 y) \sin (z), \tag {78}
$$
where we chose the viscosity to be  $\nu = 0.05$ . Each forcing term  $F$  in the Eq. 70 is then given as
$$
F _ {x} = - 6 e ^ {- 1 8 \nu t} \sin (4 y) \sin (2 z), \tag {79}
$$
$$
F _ {y} = - 6 e ^ {- 1 8 \nu t} \sin (4 x) \sin (2 z), \tag {80}
$$
$$
F _ {z} = 6 e ^ {- 1 8 \nu t} \sin (4 x) \sin (4 y). \tag {81}
$$
We constructed SPINN to be four body networks (modified MLP) of 5 hidden layers with 64/384 hidden feature/output sizes each. We used Adam [8] optimizer with a learning rate of 0.001 and trained for 50,000 iterations. The weight factors in the loss function in Eq. 48 are chosen as  $\lambda_{\mathrm{pde}} = 1$ ,  $\lambda_{\mathrm{ic}} = 10$ , and  $\lambda_{\mathrm{bc}} = 1$ . We also weighted the incompressibility loss (Eq. 71) with 100. The visualized solution vector field is shown in Fig. 6.
E.1 (5+1)-d diffusion Equation
We tested our model on the  $(5 + 1)$ -d diffusion equation to verify the effectiveness of our model on higher dimensional PDE:
$$
\frac {\partial u (t , x)}{\partial t} = \Delta u (t, x), \quad x \in [ - 1, 1 ] ^ {5}, t \in [ 0, 1 ], \tag {82}
$$
where the manufactured solution is chosen to be  $\| x\| ^2 +10t$ . When trained with  $8^{6}$  collocation points, SPINN achieved a relative  $L_{2}$  error of 0.0074 within 2 minutes.
E.2 (2+1)-d Flow Mixing Problem
This is a time dependent PDE which describes the behavior of two fluids being mixed at the interface in a 2-d environment. Following the settings in CAN-PINNs [1], the equation is
$$
\frac {\partial u (t , x , y)}{\partial t} + a \frac {\partial u (t , x , y)}{\partial x} + b \frac {\partial u (t , x , y)}{\partial y} = 0, \quad t \in [ 0, 4 ], x \in [ - 4, 4 ], y \in [ - 4, 4 ], \tag {83}
$$
$$
a (x, y) = - \frac {v _ {t}}{v _ {t , \max }} \frac {y}{r}, \tag {84}
$$
$$
b (x, y) = \frac {v _ {t}}{v _ {t , \max }} \frac {x}{r}, \tag {85}
$$
$$
v _ {t} = \operatorname {s e c h} ^ {2} (r) \tanh  (r), \tag {86}
$$
$$
r = \sqrt {x ^ {2} + y ^ {2}}, \tag {87}
$$
$$
v _ {t, \max } = 0. 3 8 5. \tag {88}
$$
It has an analytic solution [16]:
$$
u (t, x, y) = - \tanh  \left(\frac {y}{2} \cos (\omega t) - \frac {x}{2} \sin (\omega t)\right), \tag {89}
$$
where  $\omega = \frac{1}{r}\frac{v_t}{v_{t,\max}}$ . The initial and boundary conditions are extracted from the exact solution Eq. 89. When trained with  $256^3$  collocation points, SPINN achieved a relative  $L_{2}$  error of 0.0029 in 6 minutes. Figure 1 shows the visualized result.
Figure 1: Visualized solution of flow mixing problem obtained by SPINN.
E.3 2-d Poisson equation on an L-shaped domain
Besides the coordinate or mesh transformations [5, 11] mentioned in the limitations section 6, we would like to present yet another simpler way to handle arbitrary domain shapes with SPINN. For boundary points, we train in a non-separable way, just as normal PINN training (non-factorizable coordinates). For collocation points inside the domain, we train in a separable way (factorizable coordinates) and ignore the points outside the domain. This can be easily achieved by finding a tight bounding box of the input domain and masking out the PDE loss when the coordinates are outside.
To show the effectiveness, we tested SPINN on the Poisson equation on an L-shaped domain following the settings used in DeepXDE [13]:
$$
- \Delta u (x, y) = 1, \quad (x, y) \in \Omega , \tag {90}
$$
$$
u (x, y) = 0, \quad (x, y) \in \partial \Omega , \tag {91}
$$
where  $\Omega = [-1,1]^2\setminus [0,1]^2$ . The reference solution was obtained by the spectral method. SPINN achieved a relative  $L_{2}$  error of 0.0322, while PINNs with modified MLP achieved 0.0392. Figure 2 shows the visualized result. For practical reasons, we split the L-shaped domain into two rectangular domains instead of masking out the top-right part of the domain. The L-shaped domain we showed is just an example case, and we believe that this method can be applied to any arbitrary input domain.
Figure 2: Visualized solution of Poisson equation obtained by SPINN and PINNs with modified MLP.
E.4 Fine Tuning with L-BFGS