Scientific Machine Learning:
Fourier Neural Operator
Learning Objectives¶
After studying this section, you should be able to:
- Explain why the response of a linear time-invariant system to a general input is a convolution with its impulse response, and generalise it to a kernel integral operator with a learnable two-argument kernel.
- Explain why a local convolution cannot represent such an operator, and why the passage to the frequency domain, where a convolution is a multiplication, makes the kernel a small set of learnable numbers.
- Describe the Fourier layer, a spectral path that multiplies the retained low modes by learned complex weights plus a pointwise linear path and a nonlinearity, and the lifting and projection layers around a stack of them.
- State what mode truncation buys (a parameter count independent of the grid and a smoothness prior) and what it costs (sharp features).
- Read the PyTorch implementation,
spectral_conv1dwithtorch.fft.rfft, the complex weight multiplication andirfft, and theFNO1dmodule with its liftingP, spectral weightsR, pointwise pathsWand projectionQ, say which line realises which idea, and extend it to two dimensions withrfft2. - Build and train a one-dimensional FNO on the cantilever beam of Deep Neural Operators, with a coordinate channel and padding because the domain is not periodic, and evaluate it on loads it never saw.
- Hand a scalar such as the time to an operator network as a constant input channel, impose fixed-end boundary values exactly through an envelope, and choose the number of modes from the physics of the solution rather than from the input.
- Learn the operator of a vibrating string with the time as an input channel, and show that a plain time channel learns an oscillating solution slowly while Fourier features of the time let the same network follow it.
- Add the wave-equation residual, formed by finite differences in time and in space from the FNO's output field, to the data loss of that operator, and show that, given a record that stops at a quarter of a transit time, the physics carries the solution forward where the data alone cannot.
- Carry the time-dependent heat problem to a square plate with a two-dimensional FNO, and say which four things change in the code and why nothing else does.
Where we are going. Deep Neural Operators learned the map from an input function to a solution with a branch network on sampled sensors and a trunk network on the output coordinate. This notebook builds an operator learner on a different principle: the solution operators of physics are integral operators, an integral operator with a translation-invariant kernel is a convolution, and a convolution is a multiplication in the frequency domain, so the kernel can be learned as a handful of complex numbers on the low Fourier modes. Section 1 traces that chain from the impulse response of a linear system to the full architecture; Sections 2 to 4 build the PyTorch implementation on the cantilever beam, the cooling rod and the vibrating string of Deep Neural Operators, the last with the wave equation added to the loss, so that the problems are familiar and their answers known; Section 5 takes the time-dependent heat problem to a square plate, the first two-dimensional problem and the gentlest. Fourier Neural Operator: Labs then applies the same architecture to engineering problems whose data come from a solver.
1. Fourier Neural Operator¶
Overview
Deep Neural Operators set up the problem this notebook continues. We are given a family of physical situations indexed by an input function — a coefficient field, an initial condition, a forcing term — and for each of them a numerical solver can produce the corresponding solution field. Solving is expensive and must be repeated from scratch for every new input. The proposal of operator learning is to spend that expense once, offline, generating a dataset of input-output function pairs, and then to fit a single map
$$ \mathcal{G}_\theta : a(\cdot) \;\longmapsto\; u(\cdot) $$
between function spaces. Training is done once. Afterwards a new input function is evaluated in a single forward pass. DeepONet, the construction of that notebook, is assumed here and will not be repeated.
The Fourier Neural Operator pursues the same goal with a different construction. FNO represents the operator as a composition of layers, each of which performs a global convolution over the domain, implemented as a multiplication in the Fourier domain. Nothing is queried pointwise; the whole output field is produced at once from the whole input field.
The route to that construction is a chain of three steps, and Sections 1.1 through 1.6 follow it. First, for a linear time-invariant system, the response to an arbitrary input is the convolution of that input with the impulse response — one function characterizes the entire system. Second, dropping the time-invariance assumption replaces convolution by a general kernel integral operator, which is the right general form for a solution operator but is expensive to evaluate and awkward to parameterize. Third, restricting attention back to the convolution case and moving to the frequency domain turns the integral into a multiplication, and the object to be learned from a kernel function into a short list of complex numbers. Sections 1.7 through 1.9 assemble the resulting layer into a network.
1.1. Motivation: Response to a General Input¶
A central question in engineering and physics is: given a system and an arbitrary input, what is the output? Posed that baldly the question has no general answer, because "arbitrary input" ranges over an infinite-dimensional space of functions and the system may respond to each of them in an unrelated way. For an important class of systems, however, the question has a complete and finite answer, and understanding why is the beginning of everything that follows. The class is defined by two properties: linearity, meaning that the response to a sum of inputs is the sum of the responses, and time-invariance, meaning that delaying the input delays the response by the same amount without changing its shape.
Consider a simple first-order ODE as a representative example:
$$ \dot{y} + 5y = x(t), \quad y(0) = 0 $$
The system takes a forcing function $x(t)$ and returns the resulting $y(t)$. It is linear because the equation is linear in $y$ and $x$, and time-invariant because the coefficient $5$ does not depend on $t$.
For a Linear Time-Invariant (LTI) system, the response to any general input $x(t)$ can be computed from the impulse response $h(t)$ alone:
$$ y(t) = h(t) * x(t) $$
The impulse response is the output produced by a unit impulse $\delta(t)$ at the input, and for the example above it can be computed by hand. Setting $x(t) = \delta(t)$, the impulse deposits a unit of "content" at $t = 0$ and vanishes afterwards; integrating the equation across the instant $t = 0$ gives $y(0^+) = 1$, and for $t > 0$ the equation reduces to the homogeneous $\dot{y} + 5y = 0$, whose solution starting from $1$ is the decaying exponential. Hence
$$ h(t) = e^{-5t}, \qquad t \ge 0, $$
and $h(t) = 0$ for $t < 0$, since a system that responds before it is excited would violate causality. The convolution formula then states that for any forcing whatsoever,
$$ y(t) = \int_{0}^{t} e^{-5(t-\tau)}\, x(\tau)\, d\tau, $$
which you may verify directly by differentiating under the integral sign. The lower limit is $0$ and the upper limit is $t$ because $h$ vanishes for negative arguments and $x$ is applied from $t = 0$ onward.
This is a remarkably powerful result, and it deserves a moment of reflection before we generalize it. A single experiment — apply an impulse, record the response — determines the output for every input the system will ever see. An infinite-dimensional map has been compressed into one function. The LTI system is therefore a mapping from one function to another; it is, in the precise mathematical sense, an operator, and the impulse response is a complete description of it.
This perspective motivates the central question of neural operator learning: can we learn such an operator from data, without assuming linearity or time-invariance? The answer developed over the next several subsections is that we can keep the form of the LTI answer — an integral against a kernel — while relaxing the assumptions that produced it, and that the resulting form is exactly what a neural network can be built around.
1.2. Convolution as an Operator¶
The convolution formula was quoted above; it is worth deriving, because the derivation shows exactly which assumption does which piece of work, and because the same reasoning reappears when we generalize.
To understand why convolution arises naturally, consider decomposing any continuous signal $x(t)$ into a superposition of scaled and shifted impulses. A continuous signal can be approximated as:
$$ x(t) = \lim_{\Delta \to 0} \sum_k x(k\Delta)\, p(t - k\Delta)\, \Delta $$
where $p(t)$ is a rectangular pulse of width $\Delta$ and height $1/\Delta$. Read the right-hand side as a staircase: on the $k$-th interval the pulse $p(t - k\Delta)$ is the only nonzero term, it has height $1/\Delta$, and multiplying by $x(k\Delta)\Delta$ makes the staircase equal to $x(k\Delta)$ there. As $\Delta \to 0$ the staircase converges to $x$, the pulse $p(t)$ concentrates into the unit impulse $\delta(t)$, and the sum becomes an integral:
$$ x(t) = \int_{-\infty}^{\infty} x(\tau)\, \delta(t - \tau)\, d\tau $$
This identity — the sifting property — says something modest: any signal is a continuous superposition of impulses, weighted by its own values. Its usefulness comes from what the two assumptions now allow us to do with it.
By linearity, the response to a superposition is the superposition of the responses, so it suffices to know the response to each individual term $x(\tau)\delta(t - \tau)$. That term is an impulse located at $\tau$ and scaled by the number $x(\tau)$, and linearity says its response is the response to $\delta(t - \tau)$ scaled by the same number. By time-invariance, the response to the delayed impulse $\delta(t-\tau)$ is the delayed impulse response $h(t - \tau)$: the shape does not change, only its position on the time axis. Each term therefore contributes $x(\tau)h(t - \tau)$, and summing all contributions gives:
$$ y(t) = \int_{-\infty}^{\infty} x(\tau)\, h(t - \tau)\, d\tau = h(t) * x(t) $$
Notice how cleanly the labour divides. Linearity is what lets us decompose the input and reassemble the output. Time-invariance is what collapses the two-argument object "response at time $t$ to an impulse at time $\tau$" into the one-argument object $h(t - \tau)$. Remove time-invariance and the first half of the argument still stands; only the collapse fails. That observation is the whole content of the next subsection.
This is the convolution integral. The LTI system is therefore a mapping in functional space: it takes the input function $x(t)$ and returns the output function $y(t)$. It can be viewed as an operator $\mathcal{G}$ such that $y = \mathcal{G}(x)$, where $\mathcal{G}$ is completely characterized by its kernel $h$.
1.3. Kernel Integral Operator¶
The convolution operator is powerful but restrictive: it applies only to LTI systems. Many physical systems of interest are nonlinear, spatially inhomogeneous, or time-varying, and convolution cannot represent them exactly. A composite plate whose stiffness varies from point to point is not translation-invariant in any useful sense: pushing on the stiff region and pushing on the soft region produce responses of different shape, not merely shifted copies of one another.
The natural generalization is exactly the one the previous derivation suggests. Keep linearity, which gave us superposition, and discard time-invariance, which gave us the collapse to a single argument. Replace the fixed impulse response $h(t - \tau)$ with a flexible two-argument kernel $k(t, \tau)$:
$$ y(t) = \int_{-\infty}^{\infty} h(t-\tau)\, x(\tau)\, d\tau \quad \longrightarrow \quad y(t) = \int_{-\infty}^{\infty} k(t, \tau)\, x(\tau)\, d\tau $$
The two-argument kernel $k(t, \tau)$ relaxes the LTI constraint entirely: the response at time $t$ can now depend on the source at $\tau$ in an arbitrary, non-shift-invariant way. The standard convolution is recovered as the special case $k(t, \tau) = h(t - \tau)$.
Using spatial notation, the kernel integral operator is defined as:
$$ v'(x) = \int_\Omega k(x, y)\, v(y)\, dy $$
It maps an input function $v : \Omega \to \mathbb{R}$ to an output function $v' : \Omega \to \mathbb{R}$. The kernel $k(x, y)$ determines how the value at location $y$ influences the output at location $x$. When $k(x, y) = k(x - y)$, this reduces to convolution. For a general $k(x, y)$, the operator can represent non-local, non-translation-invariant mappings.
This is not an arbitrary functional form chosen for convenience. It is the form that solution operators of linear PDEs actually take. For the Poisson problem $-\Delta u = f$ on a domain $\Omega$ with $u = 0$ on the boundary, the solution is
$$ u(x) = \int_\Omega G(x, y)\, f(y)\, dy, $$
where $G$ is the Green's function of the problem — precisely a two-argument kernel, and precisely the spatial analogue of the impulse response, since $G(\cdot, y)$ is the response to a point source at $y$. For a domain with no special symmetry, $G(x,y)$ genuinely depends on both arguments separately, not on their difference.
The key insight for machine learning is that $k(t, \tau)$ can be learned from data. If we parameterize $k$ by a neural network, we can train it to approximate any target operator from input-output pairs. The Green's function is rarely available in closed form for a domain of engineering interest, and it does not exist at all for nonlinear problems; but pairs $(a, u)$ of input and solution can always be manufactured by running a solver.
Why a Local Convolution Is Not Enough
The single most important structural fact about a solution operator is that it is non-local: changing the input at one point changes the solution everywhere. The Green's function above makes this concrete. For an elliptic problem, $G(x,y)$ is nonzero for essentially every pair of points in the domain; it decays with distance, but it does not vanish outside a small neighbourhood. Press on one corner of an elastic plate and the whole plate deforms. Heat one edge of a conducting slab and the steady-state temperature changes at the opposite edge. Equilibrium is a global condition, and the operator that produces it is a global operator.
This is exactly the property that an ordinary convolution layer of the kind used in Convolutional Neural Networks (CNN) does not have. A $3 \times 3$ kernel reads only the eight neighbours of each pixel, so a single layer has a receptive field of radius one. Stacking layers grows the radius by one per layer, so for a point at the centre of a $64 \times 64$ grid to see the whole domain requires on the order of thirty layers, and on a $256 \times 256$ grid, on the order of a hundred and thirty. Depth is not free: deep stacks are harder to train, and information passing through many nonlinearities is attenuated and distorted along the way.
The other way to widen the receptive field is to widen the kernel. A kernel large enough to span the domain is a global convolution by construction, and this is the honest option — but its parameter count is the number of grid points it spans, so it scales with the resolution. Train on a $64 \times 64$ grid and the learned kernel is a $64 \times 64$ array of numbers that has no meaning on any other grid. The parameter count grows quadratically with the linear resolution in two dimensions, and the whole model becomes tied to the discretization it was trained on.
So the requirement and the obstacle are both clear. We want a layer whose receptive field is the entire domain, and we want its parameter count not to grow with the grid. Section 1.5 finds a way to have both, and the mechanism is the convolution theorem.
1.4. Multi-layer Kernel Integral Operators¶
A single kernel integral operator is a linear map between function spaces. To represent highly nonlinear mappings — such as the solution operator of a nonlinear PDE — a single linear layer is insufficient, just as a single linear layer is insufficient for learning nonlinear functions between finite-dimensional spaces.
The nonlinearity is not an exotic requirement. Even when a PDE is linear in its solution, the dependence of the solution on a coefficient field is nonlinear. In a composite plate whose stiffness varies from cell to cell, the equilibrium equations are linear in the displacement, but the map from the stiffness layout to the resulting stress field is not: doubling the stiffness everywhere does not double the stress, and the stress produced by a layout that superposes two microstructures is not the sum of the stresses each produces separately. An operator learner that is linear in its input can never represent such a map, however many parameters its kernel has.
The standard remedy is to introduce a pointwise nonlinear activation $\sigma(\cdot)$ after the integral step:
$$ v'(x) = \sigma\!\left( \int_\Omega k(x, y)\, v(y)\, dy \right) $$
and to stack multiple such layers in sequence:
$$ x(t) \;\longrightarrow\; \text{Operator}_1 \;\longrightarrow\; \text{Operator}_2 \;\longrightarrow\; \text{Operator}_3 \;\longrightarrow\; y(t) $$
Each layer applies a kernel integral followed by a pointwise nonlinearity. Note the division of labour, which is the same as in an ordinary feedforward network: the integral is where information moves between different locations, and the activation is where the value at a location is transformed nonlinearly without moving anywhere. Neither alone is enough. A stack of integrals with no activation collapses to a single integral, since the composition of linear operators is linear. A stack of activations with no integral is a pointwise function, unable to let any location influence any other.
The composition of multiple such layers gives the network the expressiveness needed to approximate general nonlinear operators between function spaces. This is directly analogous to how deep neural networks approximate nonlinear functions between finite-dimensional spaces by composing linear maps with nonlinearities.
The theoretical justification for this approach is the universal approximation theorem for operators, which guarantees that a sufficiently deep and wide network of this form can approximate any continuous operator to arbitrary accuracy. As with the universal approximation theorem for ordinary networks, discussed in Artificial Neural Networks (ANN), this is an existence statement and not a construction: it promises that adequate parameters exist, not that training will find them, and not that the number required is small.
1.5. Bridge to the Frequency Domain¶
While the kernel integral operator framework is mathematically elegant, directly computing $\int k(x, y) v(y) dy$ at every spatial location $x$ is computationally expensive. For a discrete grid of $n$ points, naively evaluating all pairwise interactions requires $O(n^2)$ operations, which becomes prohibitive for large $n$. On a $64 \times 64$ grid that is $n = 4096$ points and roughly seventeen million kernel evaluations for a single layer applied to a single sample, and the kernel itself is a function of two spatial arguments that must be represented somehow.
The key computational insight comes from the convolution theorem: convolution in the spatial domain corresponds to pointwise multiplication in the frequency domain.
The Fourier transform decomposes a signal into sinusoidal components:
$$ X(j\omega) = \int_{-\infty}^{\infty} x(t)\, e^{-j\omega t}\, dt \qquad \text{(analysis)} $$
$$ x(t) = \frac{1}{2\pi} \int_{-\infty}^{\infty} X(j\omega)\, e^{j\omega t}\, d\omega \qquad \text{(synthesis)} $$
The convolution theorem states:
$$ x(t) * h(t) \quad \leftrightarrow \quad X(j\omega)\, H(j\omega) $$
The proof is one change of variable and is worth writing out, because the step where translation-invariance is used is the step that explains why the theorem applies to convolutions and not to general kernel integrals. Take the transform of the convolution and exchange the order of integration:
$$ \int_{-\infty}^{\infty}\!\left(\int_{-\infty}^{\infty} x(\tau)\, h(t-\tau)\, d\tau \right) e^{-j\omega t}\, dt = \int_{-\infty}^{\infty} x(\tau) \left( \int_{-\infty}^{\infty} h(t-\tau)\, e^{-j\omega t}\, dt \right) d\tau . $$
In the inner integral substitute $u = t - \tau$, so that $t = u + \tau$ and $dt = du$. The exponential factors as $e^{-j\omega t} = e^{-j\omega u} e^{-j\omega \tau}$, and the factor $e^{-j\omega\tau}$ does not involve $u$, so it comes outside:
$$ \int_{-\infty}^{\infty} x(\tau)\, e^{-j\omega \tau} \left( \int_{-\infty}^{\infty} h(u)\, e^{-j\omega u}\, du \right) d\tau = H(j\omega) \int_{-\infty}^{\infty} x(\tau)\, e^{-j\omega \tau}\, d\tau = H(j\omega)\, X(j\omega). $$
The substitution is legitimate only because $h$ depends on $t$ and $\tau$ through their difference. For a general kernel $k(t,\tau)$ the inner integral would still depend on $\tau$ and nothing would factor. Shift-invariance is precisely what the theorem consumes.
This is a profound simplification: the complex operation of convolution in the time domain becomes simple pointwise multiplication in the frequency domain. Furthermore, the Fast Fourier Transform (FFT) computes this transform in $O(n \log n)$ operations, making frequency-domain computation far more efficient than direct spatial convolution. A global convolution therefore costs a transform, a multiplication, and an inverse transform — quasi-linear in the number of grid points rather than quadratic.
The Fourier basis functions $e^{j\omega t}$ are the eigenfunctions of LTI systems: if $e^{j\omega t}$ is the input, the output is $H(j\omega) e^{j\omega t}$, where $H(j\omega)$ is the eigenvalue. This eigenfunction property is what makes the frequency domain so natural for analyzing and learning LTI-like operators. It is the functional-space version of a fact from linear algebra: a matrix acts on its eigenvectors by simple scaling, so in the eigenbasis a matrix-vector product is a list of independent multiplications. The Fourier basis diagonalizes every shift-invariant operator at once, which is why the same basis serves for all of them.
1.6. Learning the Kernel in the Frequency Domain¶
Returning to the kernel integral operator, the convolution special case can be written as:
$$ v'(x) = \mathcal{F}^{-1}\!\left(\mathcal{F}(k) \cdot \mathcal{F}(v)\right)(x) $$
where $\mathcal{F}(k)$ is the Fourier transform of the kernel $k$. So far this is only a restatement of the convolution theorem: it is a cheaper way to compute the same integral, given the kernel. The step that turns it into an architecture is a change of what is treated as unknown.
Nothing in the training problem requires us to represent the kernel in space. The kernel is not observed; it is a parameter to be fitted, and the network only ever uses it through its transform. So instead of learning the kernel $k$ in the spatial domain and then transforming it, we can directly learn a trainable weight matrix $R$ in the frequency domain:
$$ \mathcal{F}(k) \;\longrightarrow\; R $$
so the operator becomes:
$$ v'(x) = \mathcal{F}^{-1}\!\left(R \cdot \mathcal{F}(v)\right)(x) $$
$R$ is a complex-valued matrix learned by the neural network, acting as a frequency-domain filter. Each entry of $R$ controls how much of a particular frequency component is passed through or amplified. Standard neural network training (backpropagation through the FFT and inverse FFT) can learn $R$ directly from data. The FFT is a linear operation with an exact adjoint, so gradients flow through it without any special treatment; the transform is just another differentiable layer.
This is the central idea of the Fourier Neural Operator, and it is worth stating in one sentence: the parameters of the layer are the multipliers attached to frequency modes, and the layer they define is a convolution over the entire domain. The receptive field problem of Section 1.3 disappears — every mode is supported everywhere, so every output point sees every input point — and it disappears at the cost of a transform rather than at the cost of depth.
An important practical choice is frequency truncation: rather than learning $R$ for all $n$ frequency components, we retain only the lowest $k_{\max}$ modes and set the rest to zero. This truncation serves two purposes. First, it dramatically reduces the number of trainable parameters. Second, it acts as a regularizer that prevents overfitting to high-frequency noise. The truncated operator can still capture the smooth, large-scale features of most PDE solutions, which tend to be dominated by low-frequency content.
The parameter-count consequence deserves emphasis because it is the one that distinguishes this construction from a large spatial kernel. If $R$ has one entry per retained mode per channel pair, and the number of retained modes is fixed at $k_{\max}$ in advance, then the number of parameters does not depend on $n$ at all. A finer grid produces more Fourier coefficients, but the layer simply ignores the additional ones. The same weights are meaningful on a $64 \times 64$ grid and on a $256 \times 256$ grid.
What Truncation Costs
Truncation is usually presented as a free lunch — fewer parameters, built-in regularization — and it is worth being blunt about the other side of the ledger, because it determines which problems FNO is good at.
Retaining only the lowest $k_{\max}$ modes means the spectral path is a low-pass filter. Whatever comes out of it is a band-limited function: a superposition of a fixed, finite set of smooth sinusoids. Such a function cannot have a jump, cannot have a corner, and cannot have a feature narrower than roughly one period of the highest retained mode. Truncation is therefore not merely a computational economy; it is a smoothness prior, a hard statement that the operator is not allowed to depend on fine oscillations in its input and is not allowed to produce them in its output.
For many PDE solutions this prior is close to true. Diffusion smooths; elliptic problems have solutions smoother than their data; the large-scale structure of a flow carries most of the energy. For others it is badly wrong. Shocks in compressible flow, cracks and delamination fronts, contact boundaries, and stress concentrations at sharp material interfaces are exactly the features that live in the high modes, and they are exactly what the truncation suppresses. A band-limited approximation to a discontinuity does not merely blur it; it also oscillates around it, the Gibbs phenomenon familiar from Fourier series.
This matters for any problem whose input is discontinuous, a two-phase composite for instance: a piecewise-constant stiffness field whose value jumps between two materials across sharp internal boundaries. The stress field it produces has its most interesting and most engineering-relevant behaviour precisely at those boundaries, where stiff and soft phases meet and stress concentrates. A low-mode FNO reproduces the overall load path of such a plate convincingly and softens the peaks at the interfaces, which is the wrong place to be optimistic, since failure is governed by the peaks. Two mitigations are standard: retain more modes, at a proportional cost in parameters, and add a path through each layer that acts locally in space, which is the subject of Section 1.7.
The figure is the standard schematic of the construction. Read it from left to right as three stages. The input function is first lifted pointwise into a higher-dimensional channel space; a sequence of Fourier layers then transforms it, each performing the transform-multiply-invert operation just described in parallel with a path that stays in physical space; and a final pointwise map projects the result down to the output quantity. The rest of Section 1 fills in each stage in turn, beginning with the layer in the middle, which is the only one that moves information across the domain.
1.7. The Fourier Layer¶
The purely spectral path described above captures global, long-range interactions through the frequency domain. However, physical systems also exhibit local behavior — the value at a point is influenced strongly by its immediate neighbors — and, as the previous subsection argued, the truncated spectral path is structurally incapable of representing anything sharp. To handle both scales simultaneously, each Fourier layer adds a local path $Wv(x)$ alongside the global spectral path:
$$ v'(x) = \sigma\!\left( \int_\Omega k(x,y)\, v(y)\, dy + W v(x) \right) $$
In the Fourier domain, with the convolution kernel replaced by the learnable weight $R$, this becomes:
$$ v'(x) = \sigma\!\left( \mathcal{F}^{-1}(R \cdot \mathcal{F}(v))(x) + W v(x) \right) $$
The two paths play complementary roles:
- Global path $\mathcal{F}^{-1}(R \cdot \mathcal{F}(v))$: captures long-range dependencies and global periodic patterns by operating in the frequency domain. Each frequency mode interacts with every spatial location simultaneously.
- Local path $Wv(x)$: a pointwise linear transformation applied directly in the spatial domain. It preserves and passes local spatial information that might be lost or attenuated by the low-pass truncation in the spectral path.
The local path is worth understanding precisely, because its name invites two misreadings. It is not a small convolution and it does not look at neighbouring points at all: $W$ is a single matrix applied to the channel vector at each location independently, mixing channels but never mixing positions. Its receptive field is a single point. That is exactly why it complements the spectral path. A pointwise map in physical space is, in the frequency domain, an operation that touches every mode including the ones the spectral path discarded. Whatever information the truncation threw away, the local path carries forward unmodified in the sum. Without it, high-frequency content in the input would be annihilated in the first layer and could never be recovered; with it, the sharp structure survives the stack even though only the low modes are mixed globally.
The nonlinearity $\sigma$ then acts pointwise on the sum. It has one further effect that is easy to overlook: applying a nonlinear function to a band-limited signal creates new frequency content. The high modes suppressed by truncation in one layer are partly regenerated by the activation, only to be truncated again in the next layer. The stack is therefore not a single low-pass filter, and its effective bandwidth is larger than $k_{\max}$ suggests, though there is no clean formula for how much larger.
The combination of these two paths ensures that the Fourier layer is sensitive to both global structure and local detail. This is the complete Fourier layer, the fundamental building block of FNO.
1.8. Lifting and Projection Layers¶
The input function $a(x)$ to a PDE operator (e.g., a coefficient field or initial condition) typically has a small number of channels — often just one scalar value per spatial location. The Fourier layers, however, operate in a much higher-dimensional latent space of $d_v$ channels, which is necessary to give the network sufficient representational capacity.
The reason a single channel is not enough is the same reason a hidden layer is needed in an ordinary network. Within a Fourier layer, the spectral weight $R$ mixes channels: the output in channel $i$ at a given mode is a linear combination of the inputs in all channels at that mode. With one channel there is nothing to mix and the entire layer reduces to a single scalar multiplier per mode. The channel dimension is where the layer's expressive capacity lives, and lifting is what creates it.
Two additional fully connected layers are therefore placed at the beginning and end of the network:
Lifting layer $P$ maps the low-dimensional input $a(x) \in \mathbb{R}^{d_a}$ to the high-dimensional latent representation $v_0(x) \in \mathbb{R}^{d_v}$:
$$ v_0(x) = P(a(x)) $$
where $P$ is a pointwise fully connected network (typically a single linear layer). This operation is applied independently at each spatial location, so it does not mix information across space — it only lifts the channel dimension.
Projection layer $Q$ maps the final latent representation $v_T(x) \in \mathbb{R}^{d_v}$ back to the target output $u_\theta(x) \in \mathbb{R}^{d_u}$:
$$ u_\theta(x) = Q(v_T(x)) $$
where $Q$ is another pointwise fully connected network, typically a two-layer MLP with a nonlinear activation between the layers. The nonlinearity in $Q$ provides additional expressive power at the output stage, which is important because the Fourier layers themselves are linear in the spectral domain (the nonlinearity in each Fourier layer is applied after the spectral and local paths are combined).
Because $P$ and $Q$ act at each location independently, they are indifferent to how many locations there are. Together with the mode-indexed spectral weights, this is what leaves the whole network free of any parameter tied to a particular grid.
The full pipeline is:
$$ a(x) \;\xrightarrow{P}\; v_0(x) \;\xrightarrow{\text{Fourier layers}}\; v_T(x) \;\xrightarrow{Q}\; u_\theta(x) $$
1.9. Full FNO Architecture¶
Assembling all components, the complete FNO consists of:
- A lifting layer $P$ that maps the input function to a high-dimensional latent space.
- A sequence of $T$ Fourier layers, each consisting of a global spectral path and a local linear path, combined and passed through a nonlinear activation.
- A projection layer $Q$ that maps the final latent representation back to the output function.
$$ a(x) \;\xrightarrow{P}\; v_0(x) \;\xrightarrow{L_1}\; v_1(x) \;\xrightarrow{L_2}\; \cdots \;\xrightarrow{L_T}\; v_T(x) \;\xrightarrow{Q}\; u_\theta(x) $$
where each Fourier layer $L_t$ computes:
$$ v_{t+1}(x) = \sigma\!\left( \mathcal{F}^{-1}(R_t \cdot \mathcal{F}(v_t))(x) + W_t\, v_t(x) \right) $$
The trainable parameters are the spectral weight matrices $\{R_t\}$, the local weight matrices $\{W_t\}$, and the parameters of $P$ and $Q$. Each layer has its own $R_t$ and $W_t$; the modes are not shared across depth, so different layers may attend to different parts of the spectrum.
The full architecture learns the solution operator:
$$ \mathcal{G}_\theta : a(x) \mapsto u(x) $$
mapping an input function (e.g., a PDE coefficient field, initial condition, or forcing term) to the corresponding output function (e.g., the PDE solution field), without re-solving the PDE for each new input. Once trained, a single forward pass through the FNO produces the solution for a new input in milliseconds, compared to the seconds or minutes required by traditional numerical solvers.
One assumption is built into the construction so deeply that it is easy to miss. The discrete Fourier transform treats the domain as periodic: it represents the field as a sum of sinusoids that wrap around from one edge to the opposite one. On a genuinely periodic domain, such as a flow on a periodic box, that is exactly right. On a domain with ordinary boundary conditions — a clamped edge, a loaded edge — it is false, and the transform sees an artificial discontinuity where the field wraps. The usual remedy is to pad the field with a margin before transforming and to discard the margin afterwards, so that the artificial jump lands in the padding rather than in the region of interest. The beam laboratory of Lab 1 does exactly this, and the padding is there for this reason and no other.
In summary, FNO operates in the frequency domain, where convolution becomes multiplication. Everything the architecture does follows from that single exchange. To make the arrangement concrete, the FNO model is built as follows.
The diagram splits the network into the three stages already described. Taken in order:
(1) Lifting layer. The input function $a(x)$ — a load along a beam or an initial temperature profile, say — is expanded from its few physical channels into a wider latent representation. The transformation is pointwise: a fully connected map applied independently at every spatial location, mixing channels without moving information across space.
(2) Fourier layers. Each layer performs a global convolution by transforming to the frequency domain, multiplying a fixed number of retained low-frequency modes by learned complex weights, and transforming back to physical space. In parallel, a pointwise linear path carries local information forward; the two are summed and passed through a nonlinearity. Stacking several such layers gives the depth needed for a nonlinear operator.
(3) Projection layer. The final latent representation is mapped back down to the physical quantity of interest, again pointwise, by a small fully connected network. This is where the high-dimensional internal representation becomes a predicted deflection or temperature field.
One practical detail appears in the diagram and in every implementation: the spatial coordinates are supplied to the network as extra input channels alongside the input function itself. Without them, the lifting layer sees only the field value at each point and the network is very nearly translation-equivariant, which is wrong whenever the boundary conditions single out particular edges of the domain. Appending the coordinates tells the network where in the domain each point sits, so it can learn that the clamped end and the free end of a beam behave differently.
2. Lab 1: Euler Beam¶
Overview
The purpose of this laboratory is to learn the code, on a problem whose answer is known. The cantilever beam of Deep Neural Operators, Lab 2, is clamped at $x = 0$ and free at $x = 1$; its loads are a small family of smooth functions, its deflections are found by integrating the beam equation, and its solution operator is linear, so nothing about the answer is in doubt and every line of the implementation can be checked against Section 1 rather than taken on trust. Everything the later laboratories need appears here in its one-dimensional form: the data layout, the spectral convolution, the lifting and projection maps, the four Fourier layers with their pointwise paths, the training loop and the evaluation. The two-dimensional version of Lab 4 changes four things and nothing else.
2.1. Problem Setup¶
The cantilever is clamped at $x = 0$ and free at $x = 1$, and carries a distributed load $q(x)$. Its deflection $y(x)$ obeys the beam equation
$$ y'''' = -q(x), \qquad y(0) = y'(0) = 0, \qquad y''(1) = y'''(1) = 0 $$
The loads are the family of Deep Neural Operators, a constant and the first three sine and cosine modes,
$$ q(x) = a_0 + \sum_{k=1}^{3}\big[a_k \sin(k\pi x) + b_k \cos(k\pi x)\big], \qquad a_k, b_k \sim U(-1, 1) $$
and the deflection is found by the integrations of that notebook, carried out numerically. The one change is in the data layout: the load and the deflection are both sampled on the same $64$-point grid, including both ends, because an FNO reads and writes the whole field on one grid. Two hundred loads are drawn for training and a hundred for testing.
What is being learned
The operator maps the load to the deflection, $G: q(\cdot) \mapsto y(\cdot)$, and it is linear, so a single linear layer could represent it in principle.
What the Network Takes In and Gives Out
A DeepONet had two inputs of two kinds: the condition, sampled at sensors, went into the branch network, and the coordinate at which the answer was wanted went into the trunk network, so that a training sample was one condition paired with one point. An FNO has one network and no place to ask for a point. Its input is a field on the grid, an array with one or more numbers at every grid point, and its output is another field on the same grid: one sample is the whole beam, and there is no pairing of conditions with points. Whatever the network is to know must therefore be written onto the grid as a channel. Here the grid has $n = 64$ points $x_1, \dots, x_{64}$ on $[0, 1]$, and at each point the network is given two numbers:
$$ \text{input at } x_i = \big(\, q(x_i), \;\; x_i \,\big) $$
The first channel is the condition, the load at that point; the grid itself plays the part of the sensors, one value per point rather than a vector handed over whole. The second is the coordinate. A Fourier layer treats every point alike, so the network cannot tell by itself where a point lies, and for a beam whose two ends are held differently that matters: the coordinate channel tells the lifting map where it is. Stacked over the grid, one sample is an array of shape $(64, 2)$, a batch has shape $(\text{batch}, 64, 2)$, and the output comes back as $(\text{batch}, 64, 1)$, the deflection at every grid point, to be compared with the integrated deflection on the same grid.
One more thing the beam needs that a periodic problem would not: the domain is not periodic, so the field is padded before the transforms, as Section 1.9 described.
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import cumulative_trapezoid
import torch
import torch.nn as nn
import torch.nn.functional as F
from matplotlib import animation
from IPython.display import HTML
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Generate Data
load evaluates the load family on any grid, one load per row, and deflection integrates the beam equation on a fine grid with the cumulative trapezoidal rule and reads the result off at the grid points.
n = 64
x = np.linspace(0, 1, n) # both ends included: the domain is not periodic
k = np.arange(1, 4)
def load(coef, x): # q(x) = a0 + sum a_k sin(k pi x) + b_k cos(k pi x)
modes = np.column_stack([np.ones_like(x), np.sin(np.pi*np.outer(x, k)), np.cos(np.pi*np.outer(x, k))])
return coef @ modes.T # one load per row
xi = np.linspace(0, 1, 20*(n - 1) + 1) # fine grid for integration; every 20th point is a grid point
def integrate_from_right(f): # int_x^1 f(s) ds
cum = cumulative_trapezoid(f, xi, initial = 0)
return cum[:, -1:] - cum
def integrate_from_left(f): # int_0^x f(s) ds
return cumulative_trapezoid(f, xi, initial = 0)
def deflection(coef):
q = load(coef, xi)
y3 = integrate_from_right(q) # y''' , using y'''(1) = 0
y2 = -integrate_from_right(y3) # y'' , using y''(1) = 0
y1 = integrate_from_left(y2) # y' , using y'(0) = 0
y = integrate_from_left(y1) # y , using y(0) = 0
return y[:, ::20] # (N, 64)
np.random.seed(0)
coef_train = np.random.uniform(-1, 1, (200, 7))
coef_test = np.random.uniform(-1, 1, (100, 7))
train_q = load(coef_train, x)
test_q = load(coef_test, x)
train_y = deflection(coef_train)
test_y = deflection(coef_test)
print('train loads shape :', train_q.shape)
print('train deflections shape :', train_y.shape)
print('test loads shape :', test_q.shape)
print('test deflections shape :', test_y.shape)
Three training pairs: a load on the left and the deflection it produces on the right. The deflection is far smoother than the load and grows toward the free end, where nothing holds the beam.
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
for j, c in enumerate(['k', 'b', 'r']):
plt.plot(x, train_q[j], c)
plt.xlabel('x')
plt.ylabel('q(x)')
plt.title('Three Training Loads')
plt.grid(True)
plt.subplot(1, 2, 2)
for j, c in enumerate(['k', 'b', 'r']):
plt.plot(x, train_y[j], c)
plt.xlabel('x')
plt.ylabel('y(x)')
plt.title('Their Deflections')
plt.grid(True)
plt.tight_layout()
plt.show()
Tensors
The two channels of Section 2.1 are assembled: the coordinate is appended to every load once, here, giving A of shape (200, 64, 2), and GT holds the deflections as (200, 64, 1). get_batch draws a random subset of the solved cases at each step; a sample is a whole field.
channel_x = np.repeat(x[None, :, None], len(train_q), axis = 0) # (200, 64, 1)
A = torch.tensor(np.concatenate([train_q[:, :, None], channel_x], -1), dtype = torch.float32).to(device) # (200, 64, 2)
GT = torch.tensor(train_y, dtype = torch.float32).reshape(-1, n, 1).to(device) # (200, 64, 1)
A_test = torch.tensor(np.concatenate([test_q[:, :, None], channel_x[:len(test_q)]], -1), dtype = torch.float32).to(device) # (100, 64, 2)
def get_batch(batch_size = 64):
i = torch.randint(0, len(A), (batch_size,))
return A[i], GT[i]
Define FNO Structure
This is Section 1 written out, with the names of Sections 1.7 to 1.9. spectral_conv1d is the global path $\mathcal{F}^{-1}(R \cdot \mathcal{F}(v))$: it transforms the field with rfft, multiplies the lowest modes coefficients by the learned complex weights R, leaves every other coefficient at zero, and transforms back with irfft. The einsum is the channel mixing of Section 1.8: for each retained mode, output channel $o$ is a linear combination of the input channels $i$ with the complex weights $R_{io}$ of that mode. The weights are complex because a shift is a phase: a real multiplier can only scale a mode, and corresponds to an even kernel that cannot move anything, so an operator that transports its input, advection for instance, is out of reach of real weights. FNO1d holds the lifting map P, four spectral weights R and four pointwise paths W, one pair per Fourier layer, and the projection Q; its forward is the pipeline $a \to P \to L_1 \cdots L_4 \to Q \to u$, each layer computing $\sigma\!\left(\mathcal{F}^{-1}(R \cdot \mathcal{F}(v)) + Wv\right)$, which in the code is spectral_conv1d(v, R) + W(v) followed by gelu, with the activation omitted after the last layer.
The number of retained modes is read off the shape of R. The two keyword arguments in_channels and padding are the two devices of this problem: the input width is two, for the load and the coordinate, and the field is padded with zeros before the transforms and the padding is stripped afterwards, so that the artificial jump the periodic transform sees between $x = 1$ and $x = 0$ falls in the margin rather than in the beam. Two details are worth noticing now, because they change in Lab 4. The transform of a real signal of length $n$ has only $n/2 + 1$ independent coefficients, so rfft returns them and the lowest modes are a single block [:modes]. And every pointwise map in the model, the lifting P, the paths W and the projection Q, is an nn.Linear, which acts on the last axis of its input, the channels, at every grid point independently: it mixes channels but never neighbours, exactly as Section 1.7 said it should. The tensor keeps the layout (batch, n, channels) from the input to the output; the only operation that touches the grid axis is the transform, and rfft is told which axis that is with dim = 1.
def spectral_conv1d(v, R): # F^-1 (R . F(v)) along the grid axis, only the lowest modes are kept
modes = R.shape[-1]
v_ft = torch.fft.rfft(v, dim = 1) # v: (batch, n, width) -> (batch, n//2 + 1, width)
out_ft = torch.zeros_like(v_ft)
out_ft[:, :modes] = torch.einsum('bxi,iox->bxo', v_ft[:, :modes], R)
return torch.fft.irfft(out_ft, n = v.shape[1], dim = 1)
class FNO1d(nn.Module):
def __init__(self, width, modes = 16, in_channels = 1, padding = 0):
super().__init__()
self.padding = padding
self.P = nn.Linear(in_channels, width) # lifting: input channels -> width channels
self.R = nn.ParameterList([nn.Parameter(torch.rand(width, width, modes, dtype = torch.cfloat)/width)
for _ in range(4)]) # spectral weights, one per layer
self.W = nn.ModuleList([nn.Linear(width, width) for _ in range(4)]) # pointwise paths, one per layer
self.Q = nn.Sequential(nn.Linear(width, 128), nn.GELU(), nn.Linear(128, 1)) # projection
def forward(self, a): # a: (batch, n, in_channels)
v = self.P(a) # (batch, n, width)
v = F.pad(v, [0, 0, 0, self.padding]) # pad the grid axis; zero for a periodic domain
for k in range(4):
v = spectral_conv1d(v, self.R[k]) + self.W[k](v)
if k < 3:
v = F.gelu(v)
v = v[:, :v.shape[1] - self.padding]
return self.Q(v) # (batch, n, 1)
A Training Function
Sixteen modes are kept out of the thirty-three that a $64$-point real transform provides, and the latent width is $32$ channels; both are generous for this problem, and they are the settings of the later laboratories so that the code reads the same. The training loop is wrapped in a function, so that the models of this and the following laboratories are trained by one call: a fresh batch at each step, and one Adam step on the mean squared error between the predicted and exact fields.
width = 32
def train(model, epochs = 2000):
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
for epoch in range(1, epochs + 1):
a, gt = get_batch()
loss = torch.mean((model(a) - gt)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print('epoch %4d | loss %.4e' % (epoch, loss.item()))
model.eval()
return model
Train
The model has two input channels, for the load and the coordinate, and a padding of eight. Two thousand steps of Adam on the mean squared error between the predicted and integrated deflections, with a fresh batch at each step; the loss should fall by several orders of magnitude, since the operator is linear and smooth.
torch.manual_seed(0)
model = train(FNO1d(width, in_channels = 2, padding = 8).to(device))
Evaluation
The trained operator is applied to the hundred test loads in one pass, and its relative $L^2$ error against the integrated deflections is printed. One case is plotted: the load, the exact deflection and the prediction.
with torch.no_grad():
pred = model(A_test).cpu().numpy()[:, :, 0] # (100, 64)
print('relative L2 error on the test set : %.4f' % (np.linalg.norm(pred - test_y)/np.linalg.norm(test_y)))
j = 0
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
plt.plot(x, test_q[j], 'k')
plt.xlabel('x')
plt.ylabel('q(x)')
plt.title('Test Load')
plt.grid(True)
plt.subplot(1, 2, 2)
plt.plot(x, test_y[j], 'k', alpha = 0.3, linewidth = 4, label = 'Exact')
plt.plot(x, pred[j], 'r--', label = 'FNO')
plt.xlabel('x')
plt.ylabel('y(x)')
plt.title('Deflection')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
What the Experiment Shows
A single forward pass returns the deflection of a load the network never saw, from the clamped end, where the coordinate channel told the network the beam is held, to the free end, where it is not; the prediction lies on the integrated curve. The whole of the implementation has now been seen once: a field in and a field out on one grid, the spectral convolution with its complex weights on the retained modes, the pointwise path beside it, the lifting and projection maps around the four layers, and a training loop that is the ordinary supervised one.
The beam exercised the code and the two devices of a non-periodic domain, and its operator took one function and returned one. The laboratories that follow make the time an input, change the solution from one that decays to one that oscillates, add the governing equation to the loss, and go to two dimensions, one step at a time.
3. Lab 2: Transient Heat Conduction¶
Overview
Every operator so far took one function and returned one function. The beam of Lab 1 was steady, and a time-dependent problem solved the same way would have fixed one time $T$ and needed a new operator and a new training run for any other. This laboratory makes the time an input. The operator takes an initial condition and a time, and returns the temperature at that time, so that one trained network answers at every instant of an interval. It keeps the fixed ends, so the coordinate channel and the padding of Lab 1 return, and one device is added: the boundary values are imposed exactly by an envelope rather than learned. The exact solution is known. The problem is the transient conduction of Deep Neural Operators, Lab 3, with the same family of initial conditions, now solved by an FNO.
3.1. Problem Setup¶
Problem Setup
Picture a thin metal rod whose two ends are clamped to large heat sinks held at a fixed reference temperature. Along the rod the temperature has been shaped by a heater, unevenly, into some profile $a(x)$. At $t = 0$ the heater is switched off. What is the temperature along the rod a moment later, and a minute later? The ends cannot change, since the sinks hold them, so heat has nowhere to go but out through the ends, and the profile relaxes: sharp bumps smooth out first, the broad shape decays more slowly, and eventually the rod is at the sink temperature everywhere.
This is the question a thermal engineer asks over and over, for a chip after the load drops, for a brake disc after a stop, for a weld cooling between two clamps, and each time with a different starting profile. A classical solver answers one profile at a time. An operator network answers for every profile and every time at once, which is exactly the setting in which neural operators earn their keep. The same equation also governs the spread of a dye in a still channel between two clean reservoirs, the drying of a wet sheet held between two dry edges, and the decay of a voltage along a leaky cable grounded at both ends; whatever the physical variable, the mathematics below is identical.
Measured in units of the rod length $L$ and the diffusion time $L^2/\alpha$, where $\alpha$ is the thermal diffusivity, the problem has no parameters at all. For a $10$ cm copper rod the diffusion time is about $90$ seconds, so the dimensionless interval studied here, up to $T = 0.5$, is three quarters of a minute; for the same rod in steel it is about seven minutes. One trained network serves both; only the clock is read differently.
Consider, then, the dimensionless heat equation on the unit interval with both ends held at zero,
$$ \frac{\partial u}{\partial t} = \frac{\partial^2 u}{\partial x^2}, \qquad x \in (0, 1), \quad t \in (0, T], \qquad u(0, t) = u(1, t) = 0, \qquad u(x, 0) = a(x) $$
with $T = 0.5$. The operator to be learned is
$$ G: (a, t) \;\longmapsto\; u(\cdot, t) $$
With fixed ends the natural basis is the sine series, since every $\sin(k\pi x)$ vanishes at both ends. The initial conditions are drawn as
$$ a(x) = \sum_{k=1}^{20} c_k \sin(k\pi x), \qquad c_k = \frac{z_k}{k^{1.5}}, \quad z_k \sim \mathcal{N}(0, 1) $$
where the factor $k^{-1.5}$ makes the initial conditions smooth but not band-limited: twenty modes are present, with slowly decreasing weight. Substituting one mode into the equation gives $\dot c_k = -(k\pi)^2 c_k$, so the exact solution is
$$ u(x, t) = \sum_{k=1}^{20} c_k\, e^{-k^2 \pi^2 t}\, \sin(k\pi x) $$
and the data need no solver. Each training sample is an initial condition together with a time in $(0, T]$, and the exact solution at that time; the network never sees the same initial condition at two times, and it is asked at test time for times it never saw.
What the Network Takes In and Gives Out
An FNO does not take a number or a vector; it takes a function sampled on a grid, with one or more values at every grid point. In this example the grid has $n = 128$ points $x_1, \dots, x_{128}$ on $[0, 1]$, and at each point the network is given three numbers:
$$ \text{input at } x_i = \big(\, a(x_i), \;\; t/T, \;\; x_i \,\big) $$
The first is the initial temperature at that point. The second is the time at which the answer is wanted, scaled to $[0, 1]$ and written down at every point, so that it is the same number in all 128 places. The third is the coordinate itself. Stacked over the grid, one sample is an array of shape $(128, 3)$: 128 points, 3 channels. A batch of such samples has shape $(\text{batch}, 128, 3)$, and this is the tensor A built by make_input and handed to model(a).
The output is one number per grid point, the temperature at that point at time $t$, so one sample comes back with shape $(128, 1)$ and a batch with $(\text{batch}, 128, 1)$. This is compared with GT, the exact $u(x_i, t)$ on the same grid.
Two of the three channels deserve a word. The time channel is constant along $x$, which looks wasteful, but a Fourier layer acts on fields, not scalars, and the only way to show it a scalar is to spread the scalar over the field. The lifting map $P$ then sees at each point the pair $(a(x_i), t/T)$ and can combine them; the spectral layers mix the result across the grid. This is the standard way to hand a parameter to an operator network. It also fixes what a training sample is: one initial condition together with one time. A DeepONet could pair every initial condition with every query point in one matrix product; an FNO cannot, so a data set is a list of pairs of an initial condition and a time, each sample carrying its own time channel. Here each initial condition is drawn with one random time; in Lab 3 every shape will be paired with every recorded time. The coordinate channel is the one from Lab 1: the domain is not periodic, and the coordinate lets the network know where the ends are.
The fixed ends are handled in two places. The domain is not periodic, so the field is padded before the transforms, as in Lab 1. And the boundary values are imposed exactly rather than learned: the network's output is multiplied by $x(1 - x)$, which vanishes at $x = 0$ and $x = 1$ whatever the network does. This is the hard constraint of Embedding Physical Knowledge into Neural Networks, and it removes the ends from the loss altogether.
How many modes
Eight modes are kept, and the reason can be read off the solution formula. The initial condition has twenty modes, but mode $k$ is reduced by $e^{-k^2\pi^2 t}$, which at $t = 0.1$ is $0.37$ for $k = 1$, $0.02$ for $k = 2$ and essentially zero for $k \geq 3$, so that by then only the first mode survives; by $T = 0.5$ even that has shrunk to $e^{-\pi^2/2} \approx 0.007$ of its initial size. The output is band-limited almost at once, however many modes the input carries, and eight is already generous for every time but the very earliest. The mode count is a property of the solution family, not of the network, and here the family tells us what to choose.
3.2. Solve the Transient Heat Problem using FNO¶
Generate Data
The grid has $128$ points and includes both ends. initial evaluates the sine series for a set of coefficients, and solution evaluates it at a time $t$ with every coefficient multiplied by its decay factor; both are matrix products with the same basis. The three input channels, $a(x)$, $t/T$ and $x$, are assembled by make_input.
The training times are not drawn uniformly. Diffusion does almost all of its work early, so a uniform draw from $(0, 0.5]$ would place only a handful of samples below $t = 0.01$, where the solution changes fastest. Squaring a uniform number puts half of the samples below $t = T/4$ and a quarter below $t = T/16$; the late times, where the solution is nearly flat, need few samples.
n = 128
x = np.linspace(0, 1, n) # both ends included: fixed ends, not periodic
T = 0.5
k = np.arange(1, 21) # modes in the initial condition
basis = np.sin(np.pi*np.outer(k, x)) # (20, 128)
def initial(coef): # a(x) = sum c_k sin(k pi x)
return coef @ basis
def solution(coef, t): # u(x, t) = sum c_k exp(-k^2 pi^2 t) sin(k pi x)
decay = np.exp(-(np.pi*k)**2*t[:, None]) # one decay factor per mode and per sample
return (coef*decay) @ basis
def make_input(a, t): # three channels at every grid point: a(x), t/T, x
channel_t = np.repeat(t[:, None]/T, n, axis = 1)
channel_x = np.repeat(x[None, :], len(a), axis = 0)
return np.stack([a, channel_t, channel_x], axis = -1)
np.random.seed(0)
coef_train = np.random.normal(size = (1600, 20))/k**1.5 # c_k = z_k / k^1.5
coef_test = np.random.normal(size = (400, 20))/k**1.5
t_train = T*np.random.uniform(0, 1, 1600)**2 # squared: more samples at early times
train_a = initial(coef_train) # (1600, 128)
train_gt = solution(coef_train, t_train) # (1600, 128)
train_in = make_input(train_a, t_train) # (1600, 128, 3)
print('train input shape :', train_in.shape)
print('train gt shape :', train_gt.shape)
Three initial conditions are plotted on the left, and on the right one of them is followed through time. The wiggles disappear first, the single arch that remains at $t = 0.1$ is the first mode alone, and by $t = 0.5$ almost nothing is left.
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
for j, c in enumerate(['k', 'b', 'r']):
plt.plot(x, train_a[j], c)
plt.xlabel('x')
plt.ylabel('a(x)')
plt.title('Three Initial Conditions')
plt.grid(True)
plt.subplot(1, 2, 2)
times = np.array([0, 0.002, 0.01, 0.03, 0.1, 0.3, 0.5])
coef = np.repeat(coef_train[0:1], len(times), axis = 0) # the same coefficients, once per time
u = solution(coef, times)
for j in range(len(times)):
plt.plot(x, u[j], label = 't = %g' % times[j])
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('One Initial Condition Diffusing')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
The same initial condition can be watched diffusing. The frames are spaced in the same way as the training times, densely at the beginning and sparsely at the end, so that the fast early change is not over in the first frame.
frames = T*np.linspace(0, 1, 60)**2 # 60 times, dense at the beginning
coef = np.repeat(coef_train[0:1], len(frames), axis = 0)
u = solution(coef, frames)
fig = plt.figure(figsize = (6, 4))
line = plt.plot(x, u[0], 'k')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(u.min() - 0.1, u.max() + 0.1)
plt.grid(True)
plt.close() # the frames are shown by the animation, not here
def update(i): # redraw the line and the title for frame i
line.set_ydata(u[i])
title.set_text('t = %.3f' % frames[i])
return line, title
anim = animation.FuncAnimation(fig, update, frames = len(frames), interval = 100)
HTML(anim.to_jshtml())
Tensors
The three channels of Section 3.1 as A, the exact temperatures as GT, and the envelope on the grid.
A = torch.tensor(train_in, dtype = torch.float32).to(device) # (1600, 128, 3)
GT = torch.tensor(train_gt, dtype = torch.float32).reshape(-1, n, 1).to(device) # (1600, 128, 1)
envelope = torch.tensor(x*(1 - x), dtype = torch.float32).reshape(1, n, 1).to(device) # vanishes at both ends
def get_batch(batch_size = 64):
i = torch.randint(0, len(A), (batch_size,))
return A[i], GT[i]
Define FNO Structure
The model is the FNO1d of Lab 1 with three input channels and a padding of eight. The hard constraint on the ends is applied outside the model, in predict, by multiplying the output with the envelope $x(1 - x)$.
def predict(model, a):
return model(a)*envelope # u(0, t) = u(1, t) = 0 exactly
width = 32
torch.manual_seed(0)
model = FNO1d(width, modes = 8, in_channels = 3, padding = 8).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
print('number of parameters :', sum(w.numel() for w in model.parameters()))
Train
The loss is the mean squared error over the grid, at whatever time each sample carries, and it is formed through predict so that the envelope is applied; the loop is written out rather than calling train of Lab 1 for that reason.
for epoch in range(1, 6001):
a, gt = get_batch()
loss = torch.mean((predict(model, a) - gt)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('epoch %4d | loss %.4e' % (epoch, loss.item()))
One Initial Condition through Time
One test initial condition, the exact solution and the prediction at six times, from $t = 0.002$, where the wiggles of the initial condition are still present, to $t = 0.5$, where the solution has almost vanished. The same network produces all six curves; only the number in the second channel changes. The relative error is printed for each time as well.
plot_times = [0.002, 0.01, 0.03, 0.1, 0.3, 0.5]
colors = ['b', 'g', 'r', 'm', 'c', 'y']
test_a = initial(coef_test) # (400, 128)
j = 0 # the test sample to plot
plt.figure(figsize = (6, 4))
plt.plot(x, test_a[j], 'k', linewidth = 1, label = 'a(x)')
for t, c in zip(plot_times, colors):
t_test = np.full(len(coef_test), t) # the same time for every test sample
exact = solution(coef_test, t_test)
a = torch.tensor(make_input(test_a, t_test), dtype = torch.float32).to(device)
with torch.no_grad():
pred = predict(model, a).cpu().numpy()[:, :, 0] # (400, 128)
error = np.linalg.norm(pred - exact, axis = 1)/np.linalg.norm(exact, axis = 1)
print('t = %-5g | relative L2 error %.4f' % (t, error.mean()))
plt.plot(x, exact[j], c, alpha = 0.4, linewidth = 3)
plt.plot(x, pred[j], c + '--', label = 'FNO, t = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('Exact (solid) and FNO (dashed)')
plt.legend()
plt.grid(True)
plt.show()
Continuous in Time
Because the time is an input channel, the prediction can be asked for at as many times as one likes. The animation moves through sixty times, none of them in the training set, with the exact solution as a solid line and the FNO as a dashed one.
frames = T*np.linspace(0, 1, 60)**2 # 60 times, dense at the beginning
exact = solution(np.repeat(coef_test[j:j + 1], len(frames), axis = 0), frames) # (60, 128)
def predict_at(t): # the test sample j on the grid at time t
a = torch.tensor(make_input(test_a[j:j + 1], np.array([t])), dtype = torch.float32).to(device)
with torch.no_grad():
return predict(model, a).cpu().numpy()[0, :, 0]
pred = np.stack([predict_at(t) for t in frames]) # (60, 128)
fig = plt.figure(figsize = (6, 4))
line_exact = plt.plot(x, exact[0], 'k', alpha = 0.3, linewidth = 4, label = 'Exact')[0]
line_pred = plt.plot(x, pred[0], 'r--', label = 'FNO')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(exact.min() - 0.1, exact.max() + 0.1)
plt.legend()
plt.grid(True)
plt.close()
def update(i):
line_exact.set_ydata(exact[i])
line_pred.set_ydata(pred[i])
title.set_text('t = %.3f' % frames[i])
return line_exact, line_pred, title
anim = animation.FuncAnimation(fig, update, frames = len(frames), interval = 100)
HTML(anim.to_jshtml())
What the Experiment Shows
A single FNO with 8 modes, trained once, predicts $u(x, t)$ at every time in $[0, T]$. Time is simply one more input channel; nothing in the architecture had to change. The figure shows one network following one initial condition from a sharp profile at $t = 0.002$ to a smooth arch at $t = 0.1$ and on to the nearly flat line at $t = 0.5$, and the animation follows it through sixty times, none of them seen in training.
The lesson is the one anticipated in the problem setup: the number of modes should be chosen from the physics of the solution, not from the complexity of the input. Diffusion kills high modes almost immediately, and 8 modes are enough for the output even though the input has 20.
4. Lab 3: A Vibrating String¶
Overview
The cooling rod settled; this problem rings. It is the vibrating string of Deep Neural Operators, Lab 4, with the same family of initial shapes and the same training data, now solved by an FNO. The solution oscillates in time instead of decaying, and a network given $t$ as a plain input builds the slow part of the dependence first and the fast part only after many more steps, if at all, the spectral bias of PINN: Limitations and Overcoming Strategies along the time axis. The question of this laboratory is whether an FNO, which reads the time as a constant input channel, has this difficulty, and whether the remedy of that notebook, cosines of $t$ in place of $t$, works here as well.
4.1. Problem Setup¶
Problem Setup
The dimensionless wave equation on the unit interval with both ends fixed,
$$ \frac{\partial^2 u}{\partial t^2} = \frac{\partial^2 u}{\partial x^2}, \qquad x \in (0, 1), \quad t \in (0, T], \qquad u(0, t) = u(1, t) = 0, \qquad u(x, 0) = a(x), \quad \frac{\partial u}{\partial t}(x, 0) = 0 $$
with $T = 1$, one transit time. The initial shapes are drawn from the sine series with ten modes, $a(x) = \sum_{k=1}^{10} c_k \sin(k\pi x)$, $c_k = z_k/k^{1.5}$, and the exact solution is
$$ u(x, t) = \sum_{k=1}^{10} c_k\, \cos(k\pi t)\, \sin(k\pi x) $$
The factor $e^{-k^2\pi^2 t}$ of the rod has become $\cos(k\pi t)$: nothing decays, the solution at $t = 1$ is the initial shape mirrored and upside down, and the $k$th mode completes $k/2$ oscillations in the interval. The operator is $G: (a, t) \mapsto u(\cdot, t)$, as in Lab 2.
What the Network Takes In and Gives Out
The input is a field on the $128$-point grid with the initial shape $a(x_i)$ and the coordinate $x_i$ in two of its channels, as in Lab 2, and the time in the rest. Two ways of writing the time are compared.
(1) Plain time channel. The time is written as in Lab 2: one channel holding $t/T$ at every grid point. At $x_i$ the network reads three numbers,
$$ \text{input at } x_i = \big(\, a(x_i), \;\; x_i, \;\; t/T \,\big) $$
so one sample has shape $(128, 3)$. Everything the network knows about the time is this one constant, and whatever oscillation in $t$ the solution requires has to be built out of it by the lifting map, the pointwise paths and their activations.
(2) Fourier features in time. The time is written as the cosines of its first $M$ harmonics, the Fourier feature embedding of Deep Neural Operators, Lab 4, with $M = 10$: $M$ channels, each again a constant over the grid. At $x_i$ the network reads $2 + M = 12$ numbers,
$$ \text{input at } x_i = \big(\, a(x_i), \;\; x_i, \;\; \cos(\pi t), \; \cos(2\pi t), \; \ldots, \; \cos(M\pi t) \,\big) $$
so one sample has shape $(128, 12)$. The oscillations the solution requires, $\cos(k\pi t)$ for $k = 1, \ldots, 10$, are now among the inputs themselves, and what remains for the network is to pair each of them with the right spatial mode, a smooth combination. Only cosines are used, and not sines, because the string is released from rest: the exact solution contains only $\cos(k\pi t)$, and a network that reads $t$ through even functions alone produces an output that is even in $t$, so that $\partial u/\partial t = 0$ at $t = 0$ holds whatever the weights. The initial velocity is built in, as the fixed ends are built in by the envelope.
The two models differ in nothing but these time channels: in both cases the lifting map reads the time at every point next to the initial shape, the spectral layers mix the result across the grid, and the output is one number per grid point, multiplied by the envelope $x(1 - x)$ so that the ends are fixed exactly. In the code the two are time_plain and time_fourier, handed to the same make_input.
In the frequency domain the solution operator is simple to describe: the coefficient of $\sin(k\pi x)$ is multiplied by $\cos(k\pi t)$. The multiplier is a function of $t$ as well as of the mode, which is why the time has to be an input at all, and it oscillates in $t$ with frequencies up to $10\pi$. The spatial modes are a different matter. All ten are present at every time, since nothing decays, so sixteen retained modes are kept rather than the eight of the rod: here the mode count is set by the solution family at every time, not only at the start.
4.2. Solve the String Problem using FNO¶
Generate Data
The data are those of Deep Neural Operators, Lab 4, arranged as the FNO reads them: forty training shapes at thirty-two uniformly spaced times on $(0, T]$, so that every pair of a shape and a time is one sample of the field, $1280$ samples in all, and ten test shapes evaluated at whatever times are asked for. make_input assembles the channels and takes the time representation as an argument, so that the same function serves both models.
n = 128
x = np.linspace(0, 1, n) # both ends included: fixed ends, not periodic
T = 1.0
k = np.arange(1, 11) # modes in the initial shape
basis = np.sin(np.pi*np.outer(k, x)) # (10, 128)
M = 10 # Fourier features in time
def initial(coef): # a(x) = sum c_k sin(k pi x)
return coef @ basis
def solution(coef, t): # u(x, t) = sum c_k cos(k pi t) sin(k pi x), t: (N,)
return (coef*np.cos(np.pi*k*t[:, None])) @ basis
def time_plain(t): # one channel: t/T
return (t/T)[:, None]
def time_fourier(t): # M channels: cos(m pi t), even in t since the string starts from rest
m = np.arange(1, M + 1)
return np.cos(np.pi*np.outer(t, m))
def make_input(a, t, time_features): # channels at every grid point: a(x), x, and the time features
N = len(a)
channel_x = np.repeat(x[None, :, None], N, axis = 0) # (N, 128, 1)
channel_t = np.repeat(time_features(t)[:, None, :], n, axis = 1) # (N, 128, 1 or M), constant over the grid
return np.concatenate([a[:, :, None], channel_x, channel_t], axis = -1)
t_grid = np.linspace(0, T, 33)[1:] # 32 times, evenly spaced on (0, T]
np.random.seed(0)
coef_train = np.random.normal(size = (40, 10))/k**1.5 # c_k = z_k / k^1.5
coef_test = np.random.normal(size = (10, 10))/k**1.5
coef_pairs = np.repeat(coef_train, len(t_grid), axis = 0) # every shape at every time: (1280, 10)
t_pairs = np.tile(t_grid, len(coef_train)) # (1280,)
train_a = initial(coef_pairs) # (1280, 128)
train_gt = solution(coef_pairs, t_pairs) # (1280, 128)
test_a = initial(coef_test) # (10, 128)
print('train condition shape :', train_a.shape)
print('train gt shape :', train_gt.shape)
Three initial shapes, and one of them followed through time: the shape splits into two travelling halves, reflects at the supports, and at $t = 1$ is its own mirror image upside down, as large as it started.
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
for j, c in enumerate(['k', 'b', 'r']):
plt.plot(x, initial(coef_train[j:j + 1])[0], c)
plt.xlabel('x')
plt.ylabel('a(x)')
plt.title('Three Initial Shapes')
plt.grid(True)
plt.subplot(1, 2, 2)
times = np.array([0, 0.25, 0.5, 0.75, 1.0])
u = solution(np.repeat(coef_train[0:1], len(times), axis = 0), times)
for j in range(len(times)):
plt.plot(x, u[j], label = 't = %g' % times[j])
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('One Initial Shape Vibrating')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
Tensors and Models
Two input tensors are built from the same shapes and times, one with the plain time channel and one with the Fourier features, and two models of the same width, mode count and padding are trained on them from the same seed. train_string is the loop of Lab 2, with the envelope applied through predict; it takes the input tensor and the model, so that the two runs differ in nothing but the time channels.
A_plain = torch.tensor(make_input(train_a, t_pairs, time_plain), dtype = torch.float32).to(device) # (1280, 128, 3)
A_fourier = torch.tensor(make_input(train_a, t_pairs, time_fourier), dtype = torch.float32).to(device) # (1280, 128, 2 + M)
GT = torch.tensor(train_gt, dtype = torch.float32).reshape(-1, n, 1).to(device) # (1280, 128, 1)
envelope = torch.tensor(x*(1 - x), dtype = torch.float32).reshape(1, n, 1).to(device) # vanishes at both ends
def predict(model, a):
return model(a)*envelope # u(0, t) = u(1, t) = 0 exactly
def train_string(A, model, epochs = 10000, batch_size = 64):
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 2500, gamma = 0.5)
for epoch in range(1, epochs + 1):
i = torch.randint(0, len(A), (batch_size,))
loss = torch.mean((predict(model, A[i]) - GT[i])**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 1000 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
model.eval()
return model
width = 32
torch.manual_seed(0)
model_plain = train_string(A_plain, FNO1d(width, modes = 16, in_channels = 3, padding = 8).to(device))
torch.manual_seed(0)
model_fourier = train_string(A_fourier, FNO1d(width, modes = 16, in_channels = 2 + M, padding = 8).to(device))
One Initial Shape through Time
One test shape, the exact solution and the two predictions at six times across the interval, ending at $t = 0.99$ where the string is nearly the inverted mirror image of its start; none of the six times is in t_grid. The relative error over the ten test shapes is printed for each time and each model.
def predict_at(model, time_features, t): # every test shape on the grid at time t, (10, 128)
t_test = np.full(len(coef_test), t)
a = torch.tensor(make_input(test_a, t_test, time_features), dtype = torch.float32).to(device)
with torch.no_grad():
return predict(model, a).cpu().numpy()[:, :, 0]
def relative_error(model, time_features, t): # mean relative L2 error over the test shapes
exact = solution(coef_test, np.full(len(coef_test), t))
pred = predict_at(model, time_features, t)
return np.mean(np.linalg.norm(pred - exact, axis = 1)/np.linalg.norm(exact, axis = 1))
plot_times = [0.05, 0.2, 0.4, 0.6, 0.8, 0.99]
colors = ['b', 'g', 'r', 'm', 'c', 'y']
j = 0 # the test shape to plot
print(' t plain time Fourier features')
for t in plot_times:
print('%5.2f %.4f %.4f' % (t, relative_error(model_plain, time_plain, t), relative_error(model_fourier, time_fourier, t)))
plt.figure(figsize = (6, 4))
plt.plot(x, test_a[j], 'k', linewidth = 1, label = 'a(x)')
for t, c in zip(plot_times, colors):
plt.plot(x, solution(coef_test[j:j + 1], np.array([t]))[0], c, alpha = 0.4, linewidth = 3)
plt.plot(x, predict_at(model_plain, time_plain, t)[j], c + '--')
plt.plot(x, predict_at(model_fourier, time_fourier, t)[j], c + ':', linewidth = 2, label = 't = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('Exact (solid), plain time (dashed), Fourier features (dotted)')
plt.legend()
plt.grid(True)
plt.show()
Continuous in Time
The Fourier-feature model and the exact solution for the same test shape through sixty times, none of them in the training set.
frames = np.linspace(0.001, T, 60)
exact = np.stack([solution(coef_test[j:j + 1], np.array([t]))[0] for t in frames]) # (60, 128)
pred = np.stack([predict_at(model_fourier, time_fourier, t)[j] for t in frames]) # (60, 128)
fig = plt.figure(figsize = (6, 4))
line_exact = plt.plot(x, exact[0], 'k', alpha = 0.4, linewidth = 3, label = 'Exact')[0]
line_pred = plt.plot(x, pred[0], 'r--', label = 'FNO, Fourier features')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(exact.min() - 0.1, exact.max() + 0.1)
plt.legend()
plt.grid(True)
plt.close()
def update(i):
line_exact.set_ydata(exact[i])
line_pred.set_ydata(pred[i])
title.set_text('t = %.3f' % frames[i])
return line_exact, line_pred, title
anim = animation.FuncAnimation(fig, update, frames = len(frames), interval = 100)
HTML(anim.to_jshtml())
What the Experiment Shows
The two models read the same shapes at the same times and differ only in how the time is written into the input. With the Fourier features the FNO follows the string through the travelling halves, the reflections and the inverted shape at the end of the interval, at every time alike, since nothing decays and no part of the interval is written off by the loss. With the plain time channel it does less well: its loss falls more slowly, and in the same number of steps it reaches about twice the error, missing most at the steep fronts of the travelling halves, whose shape is set by the highest modes, the ones that oscillate fastest in time. The difference is not that one model works and the other fails; it is that one is given the oscillation and the other has to build it, and in the same training time builds it only partly.
This is spectral bias along the time axis. The time is a channel that the lifting map and the pointwise paths read at every grid point, so the dependence on $t$ has to be built by those layers out of their activations, and a network builds slow dependence before fast. The Fourier layers do not help, because they act along $x$, not along $t$: the time enters as a constant over the grid, and a constant is the zeroth mode, which the spectral path can only scale. When the solution oscillates in an input that reaches the network this way, giving the network the oscillation at the input is the remedy.
4.3. Adding the Physics: Physics-informed FNO¶
A Short Record
The model of Section 4.2 was given the solution at thirty-two times across the whole interval and asked only to interpolate between them. Suppose instead that the string is recorded for a quarter of a transit time, nine times in $[0, 0.25]$, and that the question is what it does over the next quarter, up to $t = 0.5$. A model trained on the record alone has nothing to go on beyond $t = 0.25$; the wave equation holds there as well as inside the record, and an operator that is made to satisfy it over the whole rectangle $[0, 1] \times [0, 0.5]$ can be carried forward by it. The loss is the hybrid one of Physics-informed Neural Networks (PINN): the data term on the record, and the residual
$$ r(a; x, t) = \frac{\partial^2 u}{\partial t^2} - \frac{\partial^2 u}{\partial x^2} $$
at random pairs of a training shape and a random time in $[0, 0.5]$, over the whole grid, divided by $(5\pi)^2$, the size of $u_{tt}$ for a middle mode, so that it is of order one. The rest of the problem is built in rather than penalised: the initial displacement is in the data, since the record starts at $t = 0$, the fixed ends are in the envelope, and the zero initial velocity is in the time features, which are even in $t$.
Derivatives of an FNO Output
A PINN obtains its derivatives by automatic differentiation with respect to its input coordinates. An FNO offers no such coordinate: its output is the whole field on the grid, a function of the input field and of the time channels, not of a coordinate it could be differentiated with respect to. The derivatives are therefore formed by finite differences, as in PINN: Derivatives by Automatic Differentiation and Finite Differences. In space the three-point stencil on the grid gives $u_{xx}$ at every interior point from one output field. In time the field is evaluated three times, at $t - \Delta t$, $t$ and $t + \Delta t$, which is three forward passes with three sets of time channels and the same shape, and the same stencil gives $u_{tt}$. Both steps are small against the shortest period and wavelength in the solution, $0.2$ for the tenth mode, so the truncation error of the stencils is well below the residual the network is asked to reach.
Two models are trained on the record from the same initial weights, one on the data alone and one with the physics term added, and both are tested up to $t = 0.5$.
The Record
The data of Section 4.2 cut at $t = 0.25$, with the initial shape added at $t = 0$: the same forty shapes at nine times, $360$ samples in all.
t_record = np.linspace(0, 0.25, 9) # 9 times in [0, 0.25]: the recorded quarter, initial shape included
coef_record = np.repeat(coef_train, len(t_record), axis = 0) # every shape at every recorded time: (360, 10)
t_pairs_record = np.tile(t_record, len(coef_train)) # (360,)
A_record = torch.tensor(make_input(initial(coef_record), t_pairs_record, time_fourier), dtype = torch.float32).to(device) # (360, 128, 12)
GT_record = torch.tensor(solution(coef_record, t_pairs_record), dtype = torch.float32).reshape(-1, n, 1).to(device) # (360, 128, 1)
print('record input shape :', tuple(A_record.shape))
print('record gt shape :', tuple(GT_record.shape))
The Physics Loss
u_at evaluates a model for a batch of shapes at a batch of times, with the channels assembled in PyTorch so that the time can be shifted by $\Delta t$ without leaving the GPU. physics_loss draws $256$ pairs of a training shape and a random time in $[0, 0.5]$ and forms the scaled residual from three evaluations in time and the stencil in space.
A_shapes = torch.tensor(initial(coef_train), dtype = torch.float32).to(device) # the 40 training shapes, (40, 128)
x_grid = torch.tensor(x, dtype = torch.float32).to(device)
m_grid = torch.arange(1, M + 1, dtype = torch.float32).to(device)
def make_input_torch(a, t): # a: (N, 128), t: (N,) -> (N, 128, 2 + M)
N = len(a)
channel_x = x_grid.reshape(1, n, 1).expand(N, n, 1)
features = torch.cos(np.pi*torch.outer(t, m_grid))
channel_t = features.reshape(N, 1, M).expand(N, n, M)
return torch.cat([a.reshape(N, n, 1), channel_x, channel_t], dim = -1)
def u_at(model, a, t): # the predicted string for shapes a at times t, (N, 128)
return predict(model, make_input_torch(a, t))[:, :, 0]
dt = 0.01
h = x[1] - x[0]
T_end = 0.5
def physics_loss(model, n_pairs = 256):
i = torch.randint(0, len(A_shapes), (n_pairs,), device = device)
a = A_shapes[i]
t = T_end*torch.rand(n_pairs, device = device)
u_minus = u_at(model, a, t - dt)
u = u_at(model, a, t)
u_plus = u_at(model, a, t + dt)
u_tt = (u_plus - 2*u + u_minus)/dt**2
u_xx = (u[:, 2:] - 2*u[:, 1:-1] + u[:, :-2])/h**2
r = (u_tt[:, 1:-1] - u_xx)/(5*np.pi)**2 # wave equation residual at the interior points, scaled to order one
return torch.mean(r**2)
Train on the Record: Data Only, and Data and Physics
train_record is the loop of Section 4.2 on the record, with the physics loss added at weight w_phys; a weight of zero gives the data-only model. The two losses are printed separately, since they are measured in different units.
def train_record(model, w_phys, epochs = 20000, batch_size = 64):
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5000, gamma = 0.5)
for epoch in range(1, epochs + 1):
i = torch.randint(0, len(A_record), (batch_size,))
loss_data = torch.mean((predict(model, A_record[i]) - GT_record[i])**2)
loss_phys = physics_loss(model) if w_phys > 0 else torch.zeros(())
loss = loss_data + w_phys*loss_phys
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 2000 == 0:
print('epoch %5d | data %.4e | physics %.4e' % (epoch, loss_data.item(), loss_phys.item()))
model.eval()
return model
torch.manual_seed(0)
model_data = train_record(FNO1d(width, modes = 16, in_channels = 2 + M, padding = 8).to(device), w_phys = 0.0)
torch.manual_seed(0)
model_pi = train_record(FNO1d(width, modes = 16, in_channels = 2 + M, padding = 8).to(device), w_phys = 10.0)
The Recorded Quarter and the Quarter Beyond
The relative error over the test shapes is printed for both models at six times up to $t = 0.5$, of which the first three lie inside the record; the same test shape is drawn at three times beyond the record, with the exact solution as a solid line, the data-only prediction dashed and the physics-informed prediction dotted; and the two errors are drawn as curves over $(0, 0.5]$, with the end of the record marked.
print(' t data only data and physics')
for t in [0.05, 0.15, 0.25, 0.3, 0.4, 0.5]:
print('%5.2f %.4f %.4f%s' % (t, relative_error(model_data, time_fourier, t), relative_error(model_pi, time_fourier, t),
'' if t <= 0.25 else ' (beyond the record)'))
plt.figure(figsize = (6, 4))
for t, c in zip([0.3, 0.4, 0.5], ['m', 'c', 'y']):
plt.plot(x, solution(coef_test[j:j + 1], np.array([t]))[0], c, alpha = 0.4, linewidth = 3)
plt.plot(x, predict_at(model_data, time_fourier, t)[j], c + '--')
plt.plot(x, predict_at(model_pi, time_fourier, t)[j], c + ':', linewidth = 2, label = 't = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('Exact (solid), data only (dashed), data and physics (dotted)')
plt.legend()
plt.grid(True)
plt.show()
t_curve = np.linspace(0.02, T_end, 50)
plt.figure(figsize = (6, 4))
plt.plot(t_curve, [relative_error(model_data, time_fourier, t) for t in t_curve], 'r', label = 'data only')
plt.plot(t_curve, [relative_error(model_pi, time_fourier, t) for t in t_curve], 'b', label = 'data and physics')
plt.axvline(0.25, color = 'gray', linestyle = '--')
plt.xlabel('t')
plt.ylabel('relative L2 error')
plt.legend()
plt.grid(True, alpha = 0.3)
plt.show()
Beyond the Record, Continuously
The two predictions and the exact solution for the same test shape, through sixty times up to $t = 0.5$; the second half of the animation is the quarter the data never saw.
frames = np.linspace(0.001, T_end, 60)
exact = np.stack([solution(coef_test[j:j + 1], np.array([t]))[0] for t in frames]) # (60, 128)
pred_data = np.stack([predict_at(model_data, time_fourier, t)[j] for t in frames]) # (60, 128)
pred_pi = np.stack([predict_at(model_pi, time_fourier, t)[j] for t in frames]) # (60, 128)
fig = plt.figure(figsize = (6, 4))
line_exact = plt.plot(x, exact[0], 'k', alpha = 0.4, linewidth = 3, label = 'Exact')[0]
line_data = plt.plot(x, pred_data[0], 'r--', label = 'Data only')[0]
line_pi = plt.plot(x, pred_pi[0], 'b:', linewidth = 2, label = 'Data and physics')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(exact.min() - 0.1, exact.max() + 0.1)
plt.legend(loc = 'upper right')
plt.grid(True)
plt.close()
def update(i):
line_exact.set_ydata(exact[i])
line_data.set_ydata(pred_data[i])
line_pi.set_ydata(pred_pi[i])
title.set_text('t = %.3f (%s)' % (frames[i], 'inside the record' if frames[i] <= 0.25 else 'beyond the record'))
return line_exact, line_data, line_pi, title
anim = animation.FuncAnimation(fig, update, frames = len(frames), interval = 100)
HTML(anim.to_jshtml())
What the Experiment Shows
Inside the record the two models are close. Beyond it they part company: the data-driven model has nothing to go on and its prediction drifts away from the string at once, while the physics-informed model carries the motion forward, following the travelling halves and their reflections although none of these were in its data, with an error that grows with the distance from the record. That growth is the mark of an equation enforced only approximately, and here it has a second source besides the residual the optimizer leaves: the stencils themselves, whose truncation error a smaller $\Delta t$ and a finer grid would reduce.
The division of labour between the two terms is this. The data fix the coefficients of the operator, which the equation cannot do, since every combination of the modes $\cos(k\pi t)\sin(k\pi x)$ satisfies the wave equation, the fixed ends and the zero initial velocity whatever its coefficients; the equation fixes how those coefficients evolve, which the data cannot do beyond the last recorded time. The residual is formed from the output field by differences, in space between neighbouring grid values and in time between evaluations at neighbouring times, and the physics enters through it with the same gain as in a PINN: an operator that answers where the data do not reach.
5. Lab 4: Heat Conduction on a Square Plate¶
Overview
The purpose of this laboratory is the step into two dimensions, taken on the easiest problem available. Lab 2 made the time an input: an initial temperature profile and a time went in, and the temperature at that time came out. This laboratory repeats the construction on a square plate. The purpose is to see how little changes. The input gains a dimension and a coordinate channel, the transform becomes rfft2, the envelope gains a factor, and nothing else in the notebook is new. The exact solution is still a sine series, so the data still need no solver and the error can still be read at any time.
It is also where spectral_conv2d and FNO2d are defined, so the two-dimensional code is met on a problem whose answer is known.
5.1. Problem Setup¶
Problem Setup
Picture a thin square metal plate whose four edges are clamped to large heat sinks held at a fixed reference temperature. Across the plate the temperature has been shaped by a heater, unevenly, into some profile $a(x, y)$. At $t = 0$ the heater is switched off. What is the temperature across the plate a moment later, and a minute later? The edges cannot change, since the sinks hold them, so heat has nowhere to go but out through the edges, and the profile relaxes: small hot spots smooth out first, the broad shape decays more slowly, and eventually the plate is at the sink temperature everywhere.
This is the rod of Lab 2 with one more dimension, and it is closer to the questions thermal engineers actually ask: a chip die after the load drops, a brake disc after a stop, a plate cooling in a frame. The same equation also governs a dye spreading on a shallow tray with clean edges and the pressure in a square porous block drained at its boundary; whatever the physical variable, the mathematics below is identical.
Measured in units of the plate side $L$ and the diffusion time $L^2/\alpha$, where $\alpha$ is the thermal diffusivity, the problem has no parameters at all. For a $10$ cm copper plate the diffusion time is about $90$ seconds, so the dimensionless interval studied here, up to $T = 0.25$, is about twenty seconds; for the same plate in steel it is about three and a half minutes. One trained network serves both; only the clock is read differently.
Consider, then, the dimensionless heat equation on the unit square with all four edges held at zero,
$$ \frac{\partial u}{\partial t} = \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}, \qquad (x, y) \in (0, 1)^2, \quad t \in (0, T], \qquad u = 0 \text{ on the edges}, \qquad u(x, y, 0) = a(x, y) $$
with $T = 0.25$. The operator to be learned is
$$ G: (a, t) \;\longmapsto\; u(\cdot, \cdot, t) $$
With all edges fixed the natural basis is the double sine series, since every product $\sin(k\pi x)\sin(l\pi y)$ vanishes on all four edges. The initial conditions are drawn as
$$ a(x, y) = \sum_{k=1}^{8} \sum_{l=1}^{8} c_{kl} \sin(k\pi x)\sin(l\pi y), \qquad c_{kl} = \frac{z_{kl}}{(k^2 + l^2)^{3/4}}, \quad z_{kl} \sim \mathcal{N}(0, 1) $$
where the factor $(k^2 + l^2)^{-3/4}$ plays the role of $k^{-1.5}$ in one dimension: it makes the initial conditions smooth but not band-limited, with sixty-four modes present at slowly decreasing weight, and it treats the two directions alike. Substituting one mode into the equation gives $\dot c_{kl} = -(k^2 + l^2)\pi^2 c_{kl}$, so the exact solution is
$$ u(x, y, t) = \sum_{k=1}^{8} \sum_{l=1}^{8} c_{kl}\, e^{-(k^2 + l^2) \pi^2 t}\, \sin(k\pi x)\sin(l\pi y) $$
and the data need no solver. The slowest mode, $k = l = 1$, decays at the rate $2\pi^2$, so by $T = 0.25$ it has shrunk to $e^{-\pi^2/2} \approx 0.007$ of its initial size, the same endpoint as $T = 0.5$ in one dimension. Each training sample is an initial condition together with a time in $(0, T]$, and the exact solution at that time; the network never sees the same initial condition at two times, and it is asked at test time for times it never saw.
What the Network Takes In and Gives Out
An FNO takes a function sampled on a grid, with one or more values at every grid point. In this example the grid has $64 \times 64$ points on the unit square, and at each point $(x_i, y_j)$ the network is given four numbers:
$$ \text{input at } (x_i, y_j) = \big(\, a(x_i, y_j), \;\; t/T, \;\; x_i, \;\; y_j \,\big) $$
The first is the initial temperature at that point. The second is the time at which the answer is wanted, scaled to $[0, 1]$ and written down at every point, so that it is the same number in all $4096$ places. The last two are the coordinates. Stacked over the grid, one sample is an array of shape $(64, 64, 4)$, and a batch has shape $(\text{batch}, 64, 64, 4)$; this is the tensor A built by make_input and handed to model(a). The output is one number per grid point, the temperature there at time $t$, so a batch comes back with shape $(\text{batch}, 64, 64, 1)$ and is compared with GT, the exact $u(x_i, y_j, t)$ on the same grid.
The time channel is constant over the plate, for the same reason as in one dimension: a Fourier layer acts on fields, and the only way to show it a scalar is to spread the scalar over the field. The two coordinate channels are the coordinate channel of Lab 1 in two directions: the domain is not periodic, and the coordinates let the network know where the edges are.
The fixed edges are handled in two places. The domain is not periodic, so the field is padded in both directions before the transforms. And the boundary values are imposed exactly rather than learned: the network's output is multiplied by $x(1 - x)\,y(1 - y)$, which vanishes on all four edges whatever the network does. This is the hard constraint of Embedding Physical Knowledge into Neural Networks, and it removes the edges from the loss altogether.
How many modes
Eight modes are kept in each direction, and the reason can be read off the solution formula. The initial condition has sixty-four modes, but mode $(k, l)$ is reduced by $e^{-(k^2 + l^2)\pi^2 t}$, which at $t = 0.01$ is $0.82$ for $(1, 1)$, $0.61$ for $(2, 1)$, $0.04$ for $(4, 4)$ and about $10^{-6}$ for $(8, 8)$; at $t = 0.1$ only the $(1, 1)$ mode survives at all. The output is band-limited almost at once, however many modes the input carries, and eight is already generous for every time but the very earliest. The mode count is a property of the solution family, not of the network, and here the family tells us what to choose.
Generate Data
The grid has $64 \times 64$ points and includes all four edges. The double sine series is two matrix products: initial multiplies the coefficient array by the sine basis in $x$ from the left and by the sine basis in $y$ from the right, and solution does the same after every coefficient has been multiplied by its decay factor. The four input channels, $a(x, y)$, $t/T$, $x$ and $y$, are assembled by make_input.
The training times are not drawn uniformly. Diffusion does almost all of its work early, so a uniform draw from $(0, 0.25]$ would place only a handful of samples below $t = 0.005$, where the solution changes fastest. Squaring a uniform number puts half of the samples below $t = T/4$ and a quarter below $t = T/16$; the late times, where the solution is nearly flat, need few samples.
n = 64
x = np.linspace(0, 1, n) # both ends included: fixed edges, not periodic
y = np.linspace(0, 1, n)
T = 0.25
k = np.arange(1, 9) # modes in each direction
basis_x = np.sin(np.pi*np.outer(k, x)) # (8, 64)
basis_y = np.sin(np.pi*np.outer(k, y)) # (8, 64)
kk, ll = np.meshgrid(k, k, indexing = 'ij') # (8, 8): the mode numbers (k, l)
rate = (np.pi*kk)**2 + (np.pi*ll)**2 # (8, 8): decay rate of mode (k, l)
def initial(coef): # a(x, y) = sum c_kl sin(k pi x) sin(l pi y), coef: (N, 8, 8)
return basis_x.T @ coef @ basis_y # (N, 64, 64)
def solution(coef, t): # one decay factor per mode and per sample, t: (N,)
decay = np.exp(-rate*t[:, None, None])
return initial(coef*decay)
X, Y = np.meshgrid(x, y, indexing = 'ij') # (64, 64) each: the coordinates of every grid point
def make_input(a, t): # four channels at every grid point: a(x, y), t/T, x, y
N = len(a)
channel_t = np.broadcast_to((t/T)[:, None, None], (N, n, n))
channel_x = np.broadcast_to(X, (N, n, n))
channel_y = np.broadcast_to(Y, (N, n, n))
return np.stack([a, channel_t, channel_x, channel_y], axis = -1)
np.random.seed(0)
coef_train = np.random.normal(size = (1000, 8, 8))/(kk**2 + ll**2)**0.75 # c_kl = z_kl / (k^2 + l^2)^(3/4)
coef_test = np.random.normal(size = (200, 8, 8))/(kk**2 + ll**2)**0.75
t_train = T*np.random.uniform(0, 1, 1000)**2 # squared: more samples at early times
train_a = initial(coef_train) # (1000, 64, 64)
train_gt = solution(coef_train, t_train) # (1000, 64, 64)
train_in = make_input(train_a, t_train) # (1000, 64, 64, 4)
print('train input shape :', train_in.shape)
print('train gt shape :', train_gt.shape)
Three initial conditions are shown in the top row, and in the bottom row one of them is followed through time. The small hot and cold spots disappear first, the single hump that remains at $t = 0.03$ is already almost the $(1, 1)$ mode alone, and by $t = 0.25$ almost nothing is left. All images share one color scale so that the decay is visible.
def show(field, title, vmax): # one temperature field as an image
plt.imshow(field.T, origin = 'lower', extent = [0, 1, 0, 1], cmap = 'RdBu_r', vmin = -vmax, vmax = vmax)
plt.colorbar()
plt.xlabel('x')
plt.ylabel('y')
plt.title(title)
vmax = np.abs(train_a[0]).max()
plt.figure(figsize = (9, 5))
for j in range(3):
plt.subplot(2, 3, j + 1)
show(train_a[j], 'Initial condition %d' % j, vmax)
times = np.array([0.002, 0.03, 0.25])
coef = np.repeat(coef_train[0:1], len(times), axis = 0) # the same coefficients, once per time
u = solution(coef, times)
for j in range(3):
plt.subplot(2, 3, j + 4)
show(u[j], 'Initial condition 0 at t = %g' % times[j], vmax)
plt.tight_layout()
plt.show()
The same initial condition can be watched diffusing. The frames are spaced like the training times, densely at the beginning and sparsely at the end, so that the fast early change is not over in the first frame.
frames = T*np.linspace(0, 1, 60)**2 # 60 times, dense at the beginning
coef = np.repeat(coef_train[0:1], len(frames), axis = 0)
u = solution(coef, frames)
fig = plt.figure(figsize = (6, 5))
image = plt.imshow(u[0].T, origin = 'lower', extent = [0, 1, 0, 1], cmap = 'RdBu_r', vmin = -vmax, vmax = vmax)
plt.colorbar()
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('y')
plt.close() # the frames are shown by the animation, not here
def update(i): # redraw the image and the title for frame i
image.set_array(u[i].T)
title.set_text('t = %.3f' % frames[i])
return image, title
anim = animation.FuncAnimation(fig, update, frames = len(frames), interval = 100)
HTML(anim.to_jshtml())
Tensors
The four channels of Section 5.1 as A, the exact temperatures as GT, and the envelope on the grid.
A = torch.tensor(train_in, dtype = torch.float32).to(device) # (1000, 64, 64, 4)
GT = torch.tensor(train_gt, dtype = torch.float32).reshape(-1, n, n, 1).to(device) # (1000, 64, 64, 1)
envelope = torch.tensor(X*(1 - X)*Y*(1 - Y), dtype = torch.float32).reshape(1, n, n, 1).to(device) # vanishes on all edges
def get_batch(batch_size = 32):
i = torch.randint(0, len(A), (batch_size,))
return A[i], GT[i]
Define FNO Structure
spectral_conv2d and FNO2d are the two-dimensional versions of the one-dimensional code, with rfft2 and irfft2 in place of rfft and irfft. The transform now acts on the two grid axes, dim = (1, 2), and the real transform halves only the second of them, so along the first axis both the lowest positive and the lowest negative modes have to be kept, which is why R carries two blocks of weights. The class takes the number of input and output channels and the padding as arguments; here it takes four input channels, returns one, and is padded by eight in both directions. The hard constraint on the edges is applied outside the model, in predict, by multiplying the output with the envelope $x(1 - x)\,y(1 - y)$.
def spectral_conv2d(v, R): # F^-1 (R . F(v)) along the two grid axes, only the lowest modes are kept
modes = R.shape[-1]
v_ft = torch.fft.rfft2(v, dim = (1, 2)) # v: (batch, n, n, width) -> (batch, n, n//2 + 1, width)
out_ft = torch.zeros_like(v_ft)
out_ft[:, :modes, :modes] = torch.einsum('bxyi,ioxy->bxyo', v_ft[:, :modes, :modes], R[0])
out_ft[:, -modes:, :modes] = torch.einsum('bxyi,ioxy->bxyo', v_ft[:, -modes:, :modes], R[1])
return torch.fft.irfft2(out_ft, s = v.shape[1:3], dim = (1, 2))
class FNO2d(nn.Module):
def __init__(self, width, modes = 12, in_channels = 1, out_channels = 1, padding = 0):
super().__init__()
self.padding = padding
self.P = nn.Linear(in_channels, width) # lifting: input channels -> width channels
self.R = nn.ParameterList([nn.Parameter(torch.rand(2, width, width, modes, modes, dtype = torch.cfloat)/width)
for _ in range(4)]) # spectral weights, one per layer
self.W = nn.ModuleList([nn.Linear(width, width) for _ in range(4)]) # pointwise paths, one per layer
self.Q = nn.Sequential(nn.Linear(width, 128), nn.GELU(), nn.Linear(128, out_channels)) # projection
def forward(self, a): # a: (batch, n, n, in_channels)
v = self.P(a) # (batch, n, n, width)
v = F.pad(v, [0, 0, 0, self.padding, 0, self.padding]) # pad both grid axes
for k in range(4):
v = spectral_conv2d(v, self.R[k]) + self.W[k](v)
if k < 3:
v = F.gelu(v)
v = v[:, :v.shape[1] - self.padding, :v.shape[2] - self.padding]
return self.Q(v) # (batch, n, n, out_channels)
def predict(model, a):
return model(a)*envelope # u = 0 on all four edges exactly
width = 32
torch.manual_seed(0)
model = FNO2d(width, modes = 8, in_channels = 4, padding = 8).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
print('number of parameters :', sum(w.numel() for w in model.parameters()))
Train
The loss is the mean squared error over the grid, at whatever time each sample carries, formed through predict as in Lab 2.
for epoch in range(1, 5001):
a, gt = get_batch()
loss = torch.mean((predict(model, a) - gt)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('epoch %4d | loss %.4e' % (epoch, loss.item()))
One Initial Condition through Time
One test initial condition, the exact solution, the prediction and their difference at four times, from $t = 0.002$, where the spots of the initial condition are still present, to $t = 0.1$, where only the first mode is left. Each row shares one color scale, and the difference is drawn on one fixed scale, so that it can be compared from row to row. The same network produces every row; only the number in the second channel changes. The relative error over all $200$ test initial conditions is printed for each time, with $t = 0.25$ added to the list.
plot_times = [0.002, 0.01, 0.03, 0.1, 0.25]
test_a = initial(coef_test) # (200, 64, 64)
j = 0 # the test sample to plot
for t in plot_times:
t_test = np.full(len(coef_test), t) # the same time for every test sample
exact = solution(coef_test, t_test)
a = torch.tensor(make_input(test_a, t_test), dtype = torch.float32).to(device)
with torch.no_grad():
pred = predict(model, a).cpu().numpy()[:, :, :, 0] # (200, 64, 64)
error = np.linalg.norm((pred - exact).reshape(len(exact), -1), axis = 1)/np.linalg.norm(exact.reshape(len(exact), -1), axis = 1)
print('t = %-5g | relative L2 error %.4f' % (t, error.mean()))
if t < 0.25:
vmax = np.abs(exact[j]).max()
plt.figure(figsize = (9, 2.4))
plt.subplot(1, 3, 1)
show(exact[j], 'Exact, t = %g' % t, vmax)
plt.subplot(1, 3, 2)
show(pred[j], 'FNO, t = %g' % t, vmax)
plt.subplot(1, 3, 3)
show(pred[j] - exact[j], 'Difference', 0.1)
plt.tight_layout()
plt.show()
What the Experiment Shows
A single FNO with $8 \times 8$ modes, trained once, predicts $u(x, y, t)$ at every time in $[0, T]$. The step from the rod to the plate changed four things in the code, and it is worth listing them because they are all there is: the data have one more axis, the input has one more coordinate channel, the transform is rfft2 with two blocks of spectral weights, and the envelope has one more factor. The training loop, the loss and the evaluation are the one-dimensional code unchanged. The figures follow one initial condition from the spots of the first instants to the single hump of the first mode.
The lesson is the one of Lab 2, now in two dimensions: the number of modes should be chosen from the physics of the solution, not from the complexity of the input. Diffusion kills high modes almost immediately, in every direction at once, and $8 \times 8$ modes are enough for the output even though the input has sixty-four and the grid has four thousand points.
Summary¶
Key takeaways
- A linear time-invariant system responds to a general input by convolution with its impulse response; the generalisation is a kernel integral operator $v \mapsto \int k(x, y)\, v(y)\, dy$ with a learnable kernel, and a stack of such operators with nonlinearities between them is the operator-learning counterpart of a deep network. A local convolution cannot represent it, and a global spatial kernel costs as many parameters as the grid has points.
- In the frequency domain a convolution is a multiplication, so the kernel becomes a set of complex weights, one per retained mode per channel pair. Keeping only the lowest modes makes the parameter count independent of the grid, and imposes a smoothness prior, at the cost of sharp features that the truncated spectrum cannot carry.
- A Fourier layer is
rfft, multiplication of the retained modes by learned complex weights,irfft, added to a pointwise linear path and followed by an activation; in the code it isspectral_conv1d(v, R) + W(v), andFNO1dlifts the input channels withP, stacks four such layers and projects back withQ; the two-dimensional version replaces the transforms byrfft2andirfft2and keeps two blocks of spectral weights. Coordinates are appended as input channels when the operator is not translation-equivariant, and the field is padded when the domain is not periodic. - The beam lab trains a one-dimensional FNO on the cantilever of Deep Neural Operators in a minute, with a coordinate channel because the two ends are held differently and padding because the domain is not periodic, and returns the deflection of loads it never saw in one pass.
- With the time supplied as a second input channel, the same one-dimensional FNO predicts the fixed-end heat equation at every instant of an interval from one training run; the number of modes is set by how fast diffusion smooths the solution, not by how rough the initial condition is.
- On the vibrating string, whose solution oscillates instead of decaying, the same one-dimensional FNO with the time as a plain channel follows the motion more slowly and less closely, spectral bias along the time axis; given $\cos(m\pi t)$ as channels instead, it follows the motion through the whole interval. The Fourier layers act along $x$ and cannot help along $t$, which enters as a constant over the grid.
- A physics term carries the FNO where the data do not reach: on the string, with a record that stops at a quarter of a transit time, the wave-equation residual, formed by finite differences in time and in space from the output field, carries the solution to half a transit time, while the data-only model fails beyond the record. The residual of a field-to-field operator is a difference, not an autograd derivative; the data fix the coefficients, the equation fixes how they evolve.
- The plate lab takes the time-dependent heat problem to two dimensions and shows how little changes: one more axis in the data, one more coordinate channel,
rfft2with two blocks of spectral weights, and one more factor in the envelope, with the training loop and the evaluation untouched. Looking ahead. Fourier Neural Operator: Labs applies the architecture built here to engineering problems whose data come from a solver, a nozzle flow, a building with a tuned mass damper and a plate on an elastic foundation, and uses a trained operator as a surrogate inside a design optimization.