Scientific Machine Learning:
Deep Neural Operators
Physics-informed Neural Networks (PINN) ended in a good place and a narrow one. A neural network was trained until it satisfied a differential equation, its boundary conditions, and whatever measurements were available, and the result was a solution field obtained without a mesh. That is a real achievement. But look closely at what was trained. Physics-informed Neural Networks (PINN) trained one network per problem instance. One beam, one load, one set of supports, one network. Change the load from a uniform distribution to a triangular one, change the fixed end to a pinned end, change the thermal conductivity of the material, and the trained weights are worthless. You start again, and you pay the full training cost again.
For a single simulation this is merely inconvenient. For engineering design it is fatal. Design is not the act of solving one problem; it is the act of sweeping over many. An engineer sizing a bracket wants the deflection under fifty load cases. An engineer choosing a cooling strategy wants the temperature field for a range of heat fluxes. An optimizer wants a thousand evaluations, and a controller wants an answer in milliseconds while the plant is running. A method that costs a full training run per query cannot serve any of these, no matter how elegant the single answer it produces. The honest summary of Physics-informed Neural Networks (PINN) is that it produced a very sophisticated way to solve one problem, and design needs a way to solve a family.
This notebook changes the object being approximated, and that change is the whole point. Instead of learning a function — the solution of one problem, a map from a spatial coordinate to a temperature or a displacement — we learn an operator: the map from an input function to the output function it produces. The input function is whatever varies across the family. It may be a distributed load, a heat source, a boundary temperature profile, an inlet velocity profile, or a spatially varying material coefficient. The output function is the solution field that this input generates. An operator takes a whole function and returns a whole function, and an operator learned once covers the entire family at once. Training is paid a single time, offline. Afterwards each new case is a forward pass.
That is a genuine change of ambition, and it raises an immediate difficulty that this notebook has to answer honestly. A neural network consumes a finite vector of numbers. A function is an infinite-dimensional object. Something has to give, and what gives is that the input function is represented by its values at a fixed, finite set of sensor locations. Section 2 examines what that discretisation buys and what it costs. The architecture that results, the Deep Operator Network or DeepONet, splits into two networks whose outputs are combined by an inner product, and the reason for that peculiar split is the most instructive idea in the notebook: the trunk network learns a set of basis functions over the domain, and the branch network learns the coefficients of the particular solution in that basis.
Six top-level sections follow. Section 1 states the limitation of instance-wise solvers precisely, defines what an operator is with examples you already know from calculus and signal processing, and states the universal approximation theorem for operators that licenses the enterprise. Section 2 builds DeepONet: the sensor representation of an input function, the branch and trunk networks, the inner product that joins them, the reading of that inner product as coefficients times basis functions, and the physics-informed variant PI-DeepONet. Section 3 is Lab 1, a one-dimensional steady heat conduction problem whose operator is so simple that you can write it down by hand, which makes it the right place to see the machinery work. Section 4 is Lab 2, an Euler beam under a distributed load, the first case where the input is a genuine function sampled at many points rather than a single number. Section 5 is Lab 3, transient heat conduction in a rod with fixed ends, where the operator takes an initial condition and returns the temperature at every later time, so that the time becomes a coordinate of the trunk. Section 6 is Lab 4, the same rod set vibrating instead of cooling, where the solution oscillates in time rather than decaying, so that the time enters the trunk as Fourier features, and where the wave equation itself is then added to the loss, as a first physics-informed DeepONet, to carry the solution beyond the times that were recorded.
The prerequisite is Physics-informed Neural Networks (PINN), and it is a strict one. Artificial Neural Networks (ANN) supplies the fully connected network and the training loop, and Stochastic and Mini-batch Gradient Descent supplies the optimizer. Looking forward, Fourier Neural Operator builds an operator learner on a completely different principle, and it will be easier to appreciate once the DeepONet construction is familiar.
Learning Objectives¶
After studying this section, you should be able to:
- State the limitation of a PINN, one network per problem instance, and explain what changes when the object being learned is an operator, a map from functions to functions, rather than a function.
- Describe the DeepONet architecture, a branch network on the sampled input function and a trunk network on the output coordinate combined by a dot product, and read it as a learned basis with input-dependent coefficients.
- Explain why the input function must be sampled at fixed sensors while the output can be queried anywhere, and what the universal approximation theorem for operators does and does not guarantee.
- Distinguish a data-driven DeepONet, trained on solved instances, from a physics-informed DeepONet, trained on the PDE residual and boundary conditions at collocation points, and say what each costs.
- Implement a DeepONet in PyTorch, in the matrix form that evaluates every condition at every evaluation point in one product, and train it on the one-dimensional conduction and cantilever problems; predict from the rank of the solution operator how easy each fit will be and what the test measures.
- Learn a time-dependent operator with a DeepONet by giving the trunk the pair $(x, t)$, impose fixed-end boundary values exactly through an envelope, and evaluate the trained network at times that never appeared in training.
- Learn the operator of a vibrating string, whose solution oscillates in time rather than decaying, by giving the trunk Fourier features of the time coordinate, and say why a plain coordinate would not do.
- Add the wave-equation residual, evaluated pair by pair through the trunk coordinates, 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 to half a transit time, where the data alone cannot.
- Explain the offline cost and online saving of operator learning, and when the trade is worth making.
Where we are going. Physics-informed Neural Networks (PINN) solved one problem at a time: change the load, the boundary temperature or the viscosity and the network must be trained again. This notebook changes the object being learned from a function to an operator, the map from a problem's input function to its solution, so that a new instance costs a forward pass instead of a training run. The DeepONet does it with two networks, a branch on the sampled input and a trunk on the output coordinate, combined by a dot product that is a learned basis with input-dependent coefficients. Four laboratories, all in PyTorch, take it from a rank-one conduction operator through the cantilever beam and a cooling rod to a vibrating string, where the solution oscillates in time and the trunk is given the oscillation as Fourier features. Fourier Neural Operator builds an operator learner on a different principle, in frequency space.
1. Deep Neural Operators¶
Overview
This section makes the case for a change of target. It begins by naming the limitation of the physics-informed neural network as sharply as possible, because the limitation is not a defect of the implementation and cannot be repaired by a better optimizer or a larger network. It is structural: a PINN represents a single solution field, so a single solution field is what it can give you.
The section then introduces the mathematical object that replaces the function, namely the operator, and shows that operators are entirely familiar even if the word is not. Differentiation is an operator. Integration is an operator. The Fourier transform is an operator. You have been applying operators since your first calculus course; what is new is treating one as an object to be learned rather than one to be applied. Finally the section states the theoretical result that makes learning an operator plausible at all, and points at the two papers that opened the field.
1.1. Limitations of PINN¶
PINNs have gained significant attention for their ability to solve forward and inverse problems governed by partial differential equations. Their success has been demonstrated across fluid dynamics, heat transfer, and materials science, and the reason is easy to state: by embedding the governing equation directly in the loss, a PINN turns physical law into supervision, and supervision that costs nothing to generate is available at every point of the domain. That is what makes a PINN data-efficient in a way that a purely data-driven regressor is not.
Now consider what the trained weights of a PINN actually encode. During training, the residual of one particular differential equation was evaluated at collocation points inside one particular domain, and penalties were applied for violating one particular set of initial and boundary conditions. Every one of those choices entered the loss function, and therefore every one of them is written into the converged weights. The network is not a solver. It is a solution.
Two consequences follow, and they are the reason this notebook exists.
Single instance solver. A trained PINN model solves a specific instance of a PDE problem. It predicts accurately for the initial conditions, boundary conditions, and equation parameters that were used during training, and it makes no claim at all about any other instance. If you trained a PINN for a beam under a uniform load of $q = 1$, the network has no representation of what happens under $q = 2$. It does not contain a mechanism; it contains one answer.
Need for re-training. If the initial or boundary conditions change, the entire model must be re-trained from scratch. There is no cheap update, no fine-tuning shortcut that is obviously correct, and no way to interpolate between two trained networks. Real-time prediction is therefore impossible, and parametric studies are prohibitively expensive: a sweep over fifty load cases costs fifty full training runs.
Put the cost side by side with the alternative it was meant to replace. A conventional finite element solver also solves one instance at a time, but it solves that instance in seconds and requires no training. A PINN that takes an hour to produce one field is competitive only if it offers something the solver does not, such as handling an inverse problem or assimilating scattered measurements. For the forward parametric sweep, which is the daily work of design, the instance-wise PINN is strictly worse than the classical method it was supposed to accelerate. This is worth saying plainly, because enthusiasm for physics-informed learning often skips it.
This raises the critical question of the notebook: how can we obtain real-time predictions for varying conditions without retraining for each new scenario?
Transition to Deep Neural Operators
Deep Neural Operators answer that question by changing what is learned. Rather than solving individual instances, they learn the solution mapping itself: given the conditions, return the solution. Where a PINN approximates the solution of a fixed problem setup, a DNO approximates a function-to-function mapping from input conditions — initial conditions, boundary conditions, forcing terms, material parameters — to the corresponding solution field. One trained model then covers a whole range of conditions, which is precisely what makes real-time and multi-instance prediction possible.
The economics are worth stating explicitly, because they are the practical argument for the entire field. Operator learning does not make the total computation smaller; if anything it makes it larger, since generating a training set of solved instances means running a conventional solver many times. What it does is move the cost. All of the expense is paid once, offline, before anyone needs an answer. Afterwards every query costs one forward pass through two small networks, which is microseconds. For a design sweep, an optimization loop, or a real-time controller, an expensive offline stage followed by a negligible online stage is exactly the right shape of cost, and it is the shape that instance-wise methods cannot offer.
By shifting from instance-specific learning to operator learning, Deep Neural Operators open up parametric PDEs to rapid and flexible prediction across varying physical conditions.
1.2. What is an Operator?¶
The word "operator" sounds forbidding and refers to something entirely ordinary. In mathematics and functional analysis, an operator is a mapping between functions, typically defined over infinite-dimensional spaces. An ordinary function maps numbers to numbers: give it $x$, it returns $f(x)$. An operator maps functions to functions: give it the whole function $f$, and it returns another whole function. You have been using operators since your first calculus course without calling them that.
The distinction is easier to feel with a concrete contrast. The function $f(x) = x^2$ accepts the number $3$ and returns the number $9$. The derivative operator accepts the entire function $x \mapsto x^2$ and returns the entire function $x \mapsto 2x$. You cannot hand the derivative operator a single number and expect an answer, because a derivative at a point depends on the behaviour of the function in a neighbourhood of that point. The input to an operator is irreducibly a function.
Examples of Operators
(1) Derivative Operator: maps a function to its derivative.
$$f(x) \to \frac{d}{dx} f(x)$$
If $f(x) = x^2$, the operator returns $f'(x) = 2x$. Note that both the input and the output are functions on the same domain, which will also be true of the solution operators we learn in this notebook.
(2) Integral Operator: maps a function to its integral.
$$f(x) \to \int f(x)\, dx$$
If $f(x) = 2x$, the operator yields $F(x) = x^2 + C$. Integration is the operator most closely analogous to a PDE solution map, because solving a differential equation is, at bottom, a sophisticated form of integration: the load is integrated, subject to boundary conditions, into a displacement.
(3) Projection Operator: maps a function or a vector onto a lower-dimensional subspace. Orthogonal projection in a Hilbert space is the standard example, and it appears throughout signal processing and numerical methods. A projection operator makes an important point about operators in general: they need not be invertible, and they routinely destroy information.
(4) Fourier Transform and Laplace Transform: map a function from the time or spatial domain to the frequency domain.
$$ f(x) \quad \to \quad F(j\omega) = \int_{-\infty}^{\infty} f(x) e^{-j 2\pi \omega x}\, dx$$
$$f(t) \quad \to \quad F(s) = \int_0^\infty f(t) e^{-st}\, dt$$
These transforms are used everywhere in signal processing, differential equations, and control theory. They also matter to this course specifically: Fourier Neural Operator that follows this one builds its operator out of exactly this transform, parameterising the map in frequency space instead of learning a basis from scratch.
The operator we actually care about
The four examples above are operators whose formulas we already know. The operator this notebook wants is one whose formula we do not know in closed form: the solution operator of a boundary value problem. Write the problem abstractly as
$$\mathcal{N}[s; u] = 0 \quad \text{on } \Omega, \qquad \mathcal{B}[s; u] = 0 \quad \text{on } \partial\Omega$$
where $u$ is the input data of the problem — a source term, a boundary profile, a coefficient field — and $s$ is the solution. For each admissible $u$ there is a unique solution $s$, and the map
$$G: u \; \longmapsto \; s$$
is the solution operator.
The figure shows the idea for the beam of Section 4. In both panels the input $u$ is a distributed load, a whole function along the beam, and the output $G(u)$ is the deflection curve, another whole function; the two loads differ in shape, and so do the two deflections. The operator takes the upper function to the lower one. Asking for the deflection at one position $y$ gives one number, $G(u)(y)$, and that number is what the network will actually compute, one query at a time.
For the heat conduction problem of Section 3 it maps a boundary temperature to a temperature profile. For the beam of Section 4 it maps a distributed load to a deflection curve. For the flow around a cylinder of the next notebook it maps a viscosity to a velocity and pressure field. In every case $G$ exists, is determined by the physics, and has no elementary formula. It is exactly the kind of object a neural network is supposed to be good at approximating, provided we can say what it means to feed a function in and get a function out.
Can Neural Networks Learn Operators?
Deep Neural Operators extend the idea of operator to something learned rather than derived. The theoretical justification is the universal approximation theorem for operators, which guarantees that networks can approximate function-to-function mappings given sufficient capacity.
Universal Approximation Theorem for Operators
The theorem states that a neural network can approximate any continuous operator that maps one function space to another. Formally, let $\mathcal{A}: \mathcal{X} \to \mathcal{Y}$ be an operator taking functions in a Banach space $\mathcal{X}$ to functions in another Banach space $\mathcal{Y}$. Then a neural network with sufficiently many parameters and a suitable architecture can approximate $\mathcal{A}$ to arbitrary accuracy. Neural networks are therefore capable of learning entire function-to-function mappings, not merely function values at specific points.
Be as careful with this theorem as Physics-informed Neural Networks (PINN) was with its scalar counterpart. It is an existence result about a specific architecture, and it is silent about everything that determines whether the method works in practice. It does not say how many parameters are needed, and the number may be astronomical. It does not say that gradient descent will find the approximating weights; it only says they exist. It does not say how many training examples are required to identify them. It says nothing about accuracy outside the class of input functions the approximation was constructed for, which is the limitation that will matter most in the laboratories. And it requires continuity of the operator, which excludes problems where an arbitrarily small change in the input produces a large change in the solution — shock formation, bifurcation, buckling past a critical load. What the theorem does is remove the objection of principle. It tells us the target is not impossible, and then leaves the engineering entirely to us.
Both approaches are valuable, and this course covers both: DeepONet here, FNO in the notebook that follows. The two differ in where the structure comes from. DeepONet learns a basis for the output function from scratch and represents the input function by point samples. FNO instead assumes the input and output are given on a grid, moves to the frequency domain with a fast Fourier transform, applies a learned filter to a truncated set of modes, and transforms back. One learns its basis; the other adopts a basis that is already known to be well suited to the smooth, translation-invariant behaviour of many physical fields. That is the essential contrast, and it is enough to carry forward; the details belong to Fourier Neural Operator.
We begin with DeepONet, for three reasons.
Its architecture is intuitive. DeepONet explicitly separates the input function, handled by the branch network, from the query location, handled by the trunk network, and that separation makes the design readable rather than mysterious once you know what to look for. There is a single formula at its centre and the whole method follows from it.
It connects directly to PINNs. Because the trunk network takes a spatial coordinate as its input, the output of a DeepONet is differentiable with respect to position, exactly as a PINN's output is. Everything Physics-informed Neural Networks (PINN) did with automatic differentiation therefore transfers unchanged, and the physics-informed variant of Section 2.4 follows almost immediately. This continuity from the previous notebook is not a coincidence of curriculum design; it is a property of the architecture.
It rests on the operator approximation theorem quoted above, which was proved for precisely this branch-and-trunk structure rather than assumed for it after the fact.
By concentrating on DeepONet we can build a solid foundation for operator learning while keeping the applicability to scientific and engineering problems in view throughout.
2. Deep Operator Network (DeepONet)¶
Motivation
Section 1 argued for learning an operator and left one obstacle standing. A neural network accepts a fixed-length vector of real numbers. An operator accepts a function, which is an infinite-dimensional object: to specify $q(x)$ on $[0,1]$ completely you must give uncountably many values. These two facts are not compatible, and DeepONet's response to the incompatibility is the first thing to understand about it.
The response is the obvious engineering one. Represent the input function by its values at a fixed, finite set of locations, called sensors, and treat the resulting vector as the network's input. If the load $q$ is measured at twenty equally spaced stations along a beam, then $(u_1, u_2, \ldots, u_{20})$ with $u_j = q(x_j)$ is a perfectly good finite stand-in for $q$. The name "sensor" is chosen deliberately: this is exactly what an instrumented experiment gives you, and it is exactly what a simulation gives you when you read the input field off its mesh.
This section develops the architecture that grows out of that decision, starting from a concrete flow problem and generalising in three steps: one scalar condition, several condition values, and finally an arbitrary discretised input function. It then gives the interpretation that makes the design intelligible, and closes with the physics-informed variant and an honest account of what the method cannot do.
A concrete starting point
Consider a flow field in two dimensions. Suppose we want the pressure $P$ at the location $(y_1, y_2)$, given that the viscosity of the fluid is $u_1$. Our goal is to learn efficiently the mapping from viscosity and spatial coordinate to the corresponding pressure field.
Since $u_1$, $y_1$ and $y_2$ together determine the pressure, an obvious first candidate is an ordinary fully connected network taking all three as input features and returning the pressure. Nothing is wrong with this in principle. The physical dependence really is a function of three real numbers, and a multilayer perceptron approximates functions of three real numbers.
In practice, however, the approximation improves when the viscosity $u_1$ and the spatial coordinates $(y_1, y_2)$ are processed separately through two distinct networks, and the two latent representations are then combined by an element-wise multiplication followed by a sum. The final output therefore reflects the combined influence of the condition and of the evaluation point, and empirically this factored form approximates complex operators far better than the fused network does.
It is worth pausing on why the split helps, because "it works better" is an unsatisfying reason to adopt an architecture. The two inputs play fundamentally different roles. The viscosity indexes which problem we are solving; the coordinate indexes where in the answer we are looking. A fused network must discover that distinction from data, and nothing in its structure encourages it to. The factored network has the distinction built in: whatever the condition network computes cannot depend on position, and whatever the position network computes cannot depend on the condition. That is a strong architectural prior, and it happens to be the correct one, for a reason Section 2.3 makes precise.
Training
To train the two networks we need solution samples, obtained either from a numerical solver or from real-world measurement. Those samples supply the supervision, and the model is fitted by ordinary supervised learning so that the networks together reproduce the mapping from input condition to pressure field.
This is the point at which the practical difficulty of operator learning appears, and it is the one students most often overlook. A data-driven operator network needs many solved instances. Not many measurements of one flow, but many different flows, each solved to completion. In practice that means running a conventional CFD or finite element solver once per training condition, which is expensive, and it means the operator network can only be as good as the solver that generated its labels. The compensating fact, and the reason anyone accepts this bargain, is that the expense is incurred once, before deployment. After training, evaluating a new condition costs two forward passes, so a design sweep that would have required a hundred solver runs requires none. The offline cost is real and the online saving is enormous; whether the trade is worth making depends entirely on how many queries you expect to issue.
A physics-informed approach can reduce or eliminate the need for solved examples, and we return to it in Section 2.4 and use it in Section 6 and in the next notebook. For now the treatment is purely data-driven.
Evaluation
Once training is complete, we take an unseen condition $u_1$ and sweep the spatial coordinates $(y_1, y_2)$ across the domain, evaluating the network at each. This produces the whole pressure field for that condition, and comparing it against a withheld solver run is how we assess whether the model has generalised.
Note the shape of that evaluation loop, because it is characteristic. The condition is fixed and the coordinate varies. Each query is independent of the others, so the field can be evaluated at as many or as few points as you like, in any order, at any resolution. There is no mesh, no assembly, and no linear system.
From one condition value to many
In the example above the viscosity $u_1$ is a single number, and the whole condition is captured by that one value. Many problems are not so obliging. An inlet velocity distribution is not one number but a profile, and to characterise it we need several values — say $(u_1, u_2, \ldots, u_5)$, the profile sampled at five stations across the inlet. This multi-dimensional representation captures variation in the flow properties along the boundary, and it is what makes the model general enough to be interesting.
This is the discretisation decision announced at the start of the section, now stated concretely. The input function is replaced by the vector of its values at $m$ fixed sensor locations, and the branch network takes that vector.
What does the choice buy? Three things. It makes the input finite, so an ordinary network can consume it. It makes the input format uniform, so that every training example and every query has the same shape regardless of how complicated the underlying function is. And it makes no assumption about the form of the input function: any load, any temperature profile, any coefficient field can be sampled, so the same architecture handles cases that would require different parameterisations if we tried to describe them analytically.
What does it cost? Also three things, and they are permanent. First, information is lost: two different input functions that agree at all $m$ sensors are indistinguishable to the model, so any variation finer than the sensor spacing is invisible. Second, the sensor locations are fixed at training time and cannot be changed afterwards. A trained DeepONet demands its inputs at exactly those $m$ places, in exactly that order; if a new experiment reports the load at a different set of stations, the values must be interpolated onto the training grid before the model can be used. Third, the sensor count $m$ fixes the width of the branch network's input layer, so refining the representation of the input function means retraining. Notice the asymmetry that results and keep it in mind for the rest of the notebook: the input function is tied to a fixed grid, while the output function is not tied to a grid at all.
The network structure must then be widened to accept the whole vector of condition values as its branch input, while the trunk input is unchanged.
(1) Branch network. It encodes the input function values $(u_1, u_2, \cdots, u_m)$ — the initial or boundary conditions, the forcing term, or the material coefficients of the PDE, sampled at the $m$ sensor locations — and outputs a latent vector $b(u) \in \mathbb{R}^p$ that captures the condition. Its input dimension is the number of sensors and its output dimension is $p$, the width of the representation, which is a design choice.
(2) Trunk network. It takes the spatial or parametric coordinate $(y_1, y_2, \cdots, y_d)$ at which we want the answer, and outputs a second latent vector $t(y) \in \mathbb{R}^p$. Its input dimension is the dimension of the domain — one for a rod or a beam, two for a plate or a planar flow — and its output dimension is the same $p$ as the branch.
(3) Final prediction. The two latent vectors are combined by a dot product:
$$G(u)(y) = \sum_{i=1}^{p} b_i(u)\, t_i(y)$$
where $G(u)(y)$ is the predicted value of the solution function for input $u$ at location $y$, the $b_i(u)$ are the branch outputs, and the $t_i(y)$ are the trunk outputs.
Read the notation carefully, since it carries the whole idea. $G$ is the operator; $G(u)$ is a function, the solution corresponding to input $u$; and $G(u)(y)$ is a number, the value of that function at $y$. The network never produces $G(u)$ as an object. It produces $G(u)(y)$ one query at a time, and the function is whatever you get by sweeping $y$ over the domain.
Two properties follow immediately from this formula and are worth naming.
The output can be queried at any location. The trunk network is an ordinary function of the coordinate, defined everywhere in the domain, so $y$ may be any point — between sensors, between training collocation points, on a grid ten times finer than anything used in training. The sensor grid constrains the input, not the output. This is a genuine advantage over grid-to-grid methods, and it is the property that makes a DeepONet a mesh-free representation of the solution in exactly the sense Physics-informed Neural Networks (PINN) used that phrase.
The output is differentiable in $y$. Because $y$ enters through a smooth network, $\partial G(u)(y)/\partial y$ and higher derivatives are available by automatic differentiation at any point. This is what will let us put a PDE residual into the loss in Section 2.4, and it is why the trunk must take the coordinate as an input rather than emitting a fixed vector of nodal values.
2.2. From Function to Operator¶
The distinguishing feature of DeepONet is the separation of the network that processes the input function from the network that processes the evaluation point. Ordinary networks approximate functions; DeepONet is built to approximate operators, that is, mappings between functions. That raises a fair question: how can a neural network approximate an operator at all, and does it need a fundamentally different kind of architecture to do so?
Function: mapping from vectors to vectors
A standard feedforward network, a multilayer perceptron, approximates functions that map a vector input to a vector output. This is the classical function-approximation setting, and it is what every previous notebook in this course has used. The input lives in $\mathbb{R}^n$, the output in $\mathbb{R}^k$, and both are finite.
Operator: mapping from functions to functions
An operator maps a function to another function. The route to approximating one with a neural network is to make both ends finite. On the input side, represent the function by sampled values or by its coefficients in some set of basis functions. On the output side, reconstruct the result from output basis functions. What sits between the two is then an ordinary finite-dimensional map from input coefficients to output coefficients, which is precisely what a multilayer perceptron does. Given suitable bases at both ends, an operator becomes a function, and a network can learn it.
This construction is old and long predates deep learning. Classical model order reduction does exactly this: choose a basis for the input space and a basis for the solution space, project, learn or compute the map between coefficient vectors, and reconstruct. The bases might be Fourier modes, Chebyshev polynomials, finite element shape functions, or the leading modes of a proper orthogonal decomposition of a set of snapshots. All of these are hand-crafted in the sense that the basis is chosen before the data is seen.
The key innovation of DeepONet is that it does not rely on hand-crafted basis functions. It learns both the latent representation of the input function and the output basis functions directly from data. That is a more flexible arrangement, and it means the basis is adapted to the operator being learned rather than to a general notion of smoothness.
2.3. DeepONet as Basis and Coefficients Representation¶
At first sight the architecture looks arbitrary. Why two networks? Why a dot product, of all the ways one could combine two vectors? The design is in fact rooted in the universal approximation theorem for operators, which extends the classical theorem from functions to operators and which was proved for this very structure. The branch network processes the input function into a latent code; the trunk network processes the evaluation point into a second set of features; and the dot product joins them.
There is an interpretation of this structure that makes it not merely defensible but inevitable, and it is the single most useful idea in the notebook. Looked at correctly, the trunk network learns a set of basis functions over the domain, and the branch network learns the coefficients of the solution in that basis.
Mathematical interpretation
Given an operator $G$ mapping an input function $u$ to an output function $G(u)$, DeepONet approximates the mapping as
$$G(u)(y) = \sum_{i=1}^{p} b_i(u)\, t_i(y)$$
Now fix the trained network and hold the input function $u$ fixed. Then each $b_i(u)$ is a number, and each $t_i$ is a function of $y$ alone. The formula therefore reads
$$G(u)(\cdot) = b_1(u)\, t_1(\cdot) + b_2(u)\, t_2(\cdot) + \cdots + b_p(u)\, t_p(\cdot)$$
which is a linear combination of $p$ fixed functions with $p$ scalar coefficients. The functions $t_1, \ldots, t_p$ do not depend on $u$ at all: they are the same for every input the model will ever see. Only the coefficients change from one input to the next. In other words, the trained model has decided on a $p$-dimensional space of candidate solution fields, spanned by the trunk outputs, and its entire response to a new input function is to choose a point in that space.
This is a familiar way of representing a function, and you have met it repeatedly. A Fourier series writes $f(x) = \sum_k c_k e^{ikx}$: fixed basis functions, varying coefficients. A polynomial expansion writes $f(x) = \sum_k a_k x^k$. A finite element solution writes $u(x) = \sum_j u_j \phi_j(x)$ with the shape functions $\phi_j$ fixed by the mesh and the nodal values $u_j$ determined by the load. A principal component reconstruction writes a snapshot as a mean plus a weighted sum of modes. In every case the basis is fixed in advance and the coefficients carry the case-specific information.
DeepONet has exactly this form with one difference: the basis is not chosen in advance, it is learned. The trunk network is trained to discover, from the data, which $p$ functions span the solution family well. The branch network is trained to compute the right coefficients in that discovered basis for any given input. The dot product is not an arbitrary combination rule; it is the only thing it could be, because expansion in a basis is a dot product between a coefficient vector and a basis-function vector.
Two consequences follow, and they explain a great deal of what happens in the laboratories.
First, $p$ is the dimension of the solution space the model can represent, and it should be compared with the intrinsic dimension of the family of solutions. If the true solution family lives in a two-dimensional space, a model with $p = 16$ has ample room and will train easily. If the family is genuinely high-dimensional, $p$ must grow, and a small $p$ imposes an error floor that no amount of further training can remove. Section 4 gives a case where the intrinsic dimension can be computed exactly, which is a rare and instructive situation.
Second, the nonlinearity of the operator is carried entirely by the branch network. For a fixed basis, the map from coefficients to output function is linear; if the operator $G$ is nonlinear in $u$, that nonlinearity must appear in how the coefficients depend on $u$, which is to say in $b(u)$. This is a useful diagnostic. A linear PDE with a linear dependence on its data has a linear solution operator, so the branch network only has to learn a linear map, which is why the first two laboratories converge so readily. The Navier-Stokes problem of the next notebook is nonlinear, and there the branch network has real work to do.
Comparison with classical representations
The figure shows the same idea pictorially: DeepONet decomposes the solution into learned coefficients and learned basis functions, one factor supplied by each network.
2.4. Physics-informed DeepONet (PI-DeepONet)¶
Everything so far has been data-driven. The DeepONet as described relies on solution samples, from a numerical solver or from measurement, to learn the operator, and Section 2.1 was honest about what that costs: a solver run for every training condition.
Physics-informed DeepONet incorporates the governing physics directly into the learning process instead. This is precisely the idea of Physics-informed Neural Networks (PINN), lifted one level up. There, a network $u_\theta(x)$ was differentiated with respect to $x$ by automatic differentiation, the derivatives were substituted into the differential equation, and the residual was penalised at collocation points. Here the network is $G_\theta(u)(y)$, the derivatives are taken with respect to $y$ through the trunk network, and the residual is penalised at collocation points for every input function in the training family. By embedding conservation laws and PDEs into the training pipeline, PI-DeepONet leverages both data and physics, which improves accuracy and generalisation and allows it to cope with sparse or noisy data while remaining physically consistent.
It is worth being precise about what changes and what does not, because the answer is pleasingly narrow.
What does not change: the architecture. The branch network, the trunk network, the dot product, the sensor discretisation of the input function, and the ability to query any $y$ are all identical. A PI-DeepONet is not a different model; it is the same model trained against a different objective.
What does not change either: the need to sample input functions. Physics tells you what a valid solution looks like for a given input, but it does not tell you which inputs you care about. You still choose a family of input functions and still draw training conditions from it, and the trained operator is still only trustworthy on that family.
What changes: the loss. In the data-driven case the loss is a mean squared error between predicted and reference values at points where a solver has told us the answer. In the physics-informed case the loss is built from the PDE residual evaluated at collocation points, plus terms enforcing each boundary condition, exactly as in Physics-informed Neural Networks (PINN) — and in the hybrid case, plus a data term wherever reference values happen to exist.
What changes as a consequence: how much solved data is required. This is the practical payoff. Because the residual supplies supervision at every collocation point for every input function, far fewer solved examples are needed, and in the pure form none at all. Section 6 uses the hybrid form in its simplest setting, a data term on a short record of a vibrating string and the wave equation everywhere else. The next notebook trains a PI-DeepONet for a thin plate using nothing but the governing equations and the boundary conditions, with no reference solution anywhere in the objective, and its flow laboratory uses the hybrid form again.
What also changes: the cost of a training step, and the difficulty of the optimization. Each residual evaluation requires second derivatives of the network output with respect to the coordinates, obtained by nested automatic differentiation, and the loss now has several terms of differing magnitude that must be balanced against one another. Both of these were already true of PINNs; neither is made easier by the operator setting, and the training runs of the next notebook's laboratories are long enough that it supplies pretrained weights rather than asking you to reproduce them.
In principle a PI-DeepONet can be trained purely from collocation points and governing equations, with no data at all. In practice a combination of observed data and physical law is usually the better choice. The hybrid objective takes the strengths of both: the data anchors the solution where measurements exist, and the physics constrains it everywhere else, which is exactly the arrangement one wants when data is sparse or noisy.
Honest limits
Before turning to the laboratories, three limitations deserve to be stated plainly rather than discovered later.
Accuracy degrades outside the training distribution of input functions. The model has learned a $p$-dimensional basis fitted to the family of inputs it was shown, and the coefficient map $b(\cdot)$ was fitted on the same family. Present it with an input function of a shape it has never seen — a discontinuous load when it was trained on smooth ones, a viscosity outside the sampled range, a boundary profile with sharper features than anything in the training set — and there is no mechanism that keeps the prediction close to the truth. This is ordinary out-of-distribution failure, and it is not softened by the presence of physics in the loss, because the physics was also enforced only over the training family. Note that both laboratories in Sections 3 and 4 test at a condition that lies inside the range of the training conditions, so they measure interpolation and say nothing about extrapolation.
The sensor grid is fixed. The number and location of the sensors are frozen when the branch network's input layer is sized. Changing them means retraining. A conventional solver has no analogous constraint: you may refine the mesh, change the quadrature, or supply the load in whatever form you have it.
Error control is not comparable to a classical solver's. This is the most important of the three. A finite element analysis comes with a convergence theory: refine the mesh and the error decreases at a known rate, and a posteriori estimators can tell you where the error lives. A trained operator network offers nothing of the kind. Its error on a new input can be measured only by comparing against a reference solution, which is the very thing the surrogate was built to avoid computing. The residual can be evaluated cheaply and is a useful sanity check, but a small residual does not bound the solution error. In practice this means an operator surrogate belongs in the exploratory and screening stages of a design workflow — where thousands of cheap approximate answers are worth more than ten exact ones — with a classical solver retained for the final verification of whatever design the sweep selects.
The four laboratories
Each of the four laboratories that follow has one job, and the code in them is nearly identical from one to the next; read each for its mapping, what the input function is, what the output function is and where the training data come from, rather than for its implementation details. Lab 1 is about the code: a problem so simple that nothing can go wrong with the physics, so that the data layout, the two networks and the training loop can be learned without distraction. Lab 2 is the data-driven DeepONet in its most general form: a genuine input function, sampled at sensors, mapped to a solution curve from a set of solved examples, which is how an operator is learned from a simulator or an experiment. Lab 3 keeps the data-driven setting and adds time: the output is a function of position and time, and the time enters the trunk as a second coordinate, which is all that a time-dependent problem asks of a DeepONet. Lab 4 keeps the rod and sets it vibrating: the solution oscillates in time instead of decaying, so the time enters the trunk as Fourier features, a lesson about what a trunk finds easy to represent; it ends by adding the wave-equation residual to the data loss, the hybrid physics-informed DeepONet of the next notebook, and showing that the equation carries the solution beyond the end of a short record.
3. Lab 1: Heat Transfer¶
Overview
The purpose of this laboratory is to learn the code. It uses the simplest DeepONet that can be built: one neuron of input to the branch network, representing the condition, and one neuron of input to the trunk network, representing the evaluation location. The architecture is otherwise complete, so the essential structure of the method is preserved while the physics is stripped to nothing.
That is the right way to start, and it is worth defending. The problem has a closed-form solution that a first-year student can derive in two lines, so nothing about the answer is in doubt. Everything you observe can therefore be attributed to the machinery rather than to the physics, and when the same machinery is applied to a plate or a flow in later sections you will already know what a healthy training run looks like. Use this laboratory to fix the data layout in your mind: which array holds conditions, which holds evaluation points, and how a pair of the two indexes one training target.
3.1. Problem Setting¶
Consider steady-state heat conduction in a one-dimensional rod of unit length, with the temperature fixed at both ends:
- Left boundary: $T(0) = 0$
- Right boundary: $T(1) = T_{BC}$
With no internal heat generation and constant conductivity, the governing equation is
$$\frac{d^2 T}{dy^2} = 0$$
Analytical solution. Integrating twice gives $T(y) = c_1 y + c_2$. The condition $T(0) = 0$ forces $c_2 = 0$, and $T(1) = T_{BC}$ then forces $c_1 = T_{BC}$. Hence
$$T(y) = T_{BC} \cdot y$$
a linear profile between the two fixed ends, which is what the physics of steady conduction in a uniform rod requires: with no source and no accumulation, the heat flux $-k\,dT/dy$ must be the same at every cross-section, so the gradient must be constant.
What is being learned here
The operator to be learned maps the boundary condition to the temperature field:
$$G: T_{BC} \; \longmapsto \; T(\cdot), \qquad G(T_{BC})(y) = T_{BC}\, y$$
The input function is degenerate in this example. The condition is a single number rather than a profile, so the branch network takes one input, $u = T_{BC}$, and the sensor discretisation of Section 2 collapses to the trivial case $m = 1$. The output function is the temperature distribution $T(y)$ along the rod, and the trunk network takes the single coordinate $y$. The model learns $T(y)$ from data.
Written in the DeepONet form, $G(u)(y) = \sum_{i=1}^{p} b_i(u)\, t_i(y)$ must reproduce $u \cdot y$, a product of one function of $u$ and one function of $y$, so a single basis function would suffice. The network of Section 3.2 uses $p = 16$, far more than needed; the point of this laboratory is the code, not efficiency.
3.2. Solve the Heat Transfer 1D Problem using DeepONet¶
The network for this problem is the one in the diagram: one input to the branch, one input to the trunk, two independent stacks of fully connected layers, and a dot product at the end.
Setup
NumPy, Matplotlib and PyTorch. The models in this notebook are small, and the first two laboratories train in seconds on a CPU; the physics-informed loop of Section 6 and those of the next notebook run for many thousands of iterations and benefit from a GPU.
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from scipy.integrate import cumulative_trapezoid
from matplotlib import animation
from IPython.display import HTML
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Define the desired output locations
The trunk network needs coordinates at which to evaluate the solution. Twenty-five equally spaced points on $[0, 1]$ are used here. These are the locations at which ground truth will be supplied during training, and it is worth being clear that they are a property of the training data, not of the model: after training, the trunk network can be evaluated anywhere on the interval, at any resolution, because it is an ordinary function of $y$.
y_eval = np.linspace(0, 1, 25).reshape(-1, 1)
plt.figure(figsize = (6, 2))
plt.plot(y_eval, np.zeros(25), 'ko', markersize = 4)
plt.yticks([])
plt.xlabel('y')
plt.title('Evaluation Points')
plt.show()
Boundary condition $u$ and its ground truth
Ten different values of $T_{\text{BC}}$ are used for training, and one held-out value, $T_{\text{BC}} = 60$, is used for testing.
The test value lies between the training values, so the test is an interpolation test, as it will be in every laboratory.
train_u = np.array([10, 20, 30, 40, 50, 70, 80, 90, 100, 110]).reshape(-1, 1) # T_BC for training
test_u = np.array([[60]]) # T_BC for testing
train_gt = train_u * y_eval.T # (10, 25), T(y) = T_BC * y
test_gt = test_u * y_eval.T # (1, 25)
print('Conditions for train :', train_u.shape)
print('Ground truth for train :', train_gt.shape)
print('Conditions for test :', test_u.shape)
print('Ground truth for test :', test_gt.shape)
The ground truth is generated analytically rather than by a solver, which is possible only because this problem has a closed-form solution. In every later laboratory the labels come from a numerical simulation or from the physics itself, which is the realistic situation. The printed shapes are the ones to memorise, since the same layout recurs in all four laboratories: conditions indexed by case, evaluation points indexed by location, and ground truth indexed by both.
plt.figure(figsize = (6, 4))
plt.plot(y_eval, train_gt.T, 'b-o', markersize = 4, alpha = 0.3)
plt.plot(y_eval, test_gt.T, 'r-o', markersize = 4, label = 'test')
plt.xlabel('y')
plt.ylabel('Temperature')
plt.title('Temperature Distribution in the Rod')
plt.legend()
plt.grid(True)
plt.show()
Tensors
The three arrays are moved to PyTorch. Keep their roles in mind, because the same three recur in every laboratory: U holds one condition per row, Y one evaluation point per row, and GT the solution value for each condition at each point.
U = torch.tensor(train_u, dtype = torch.float32).to(device) # conditions
Y = torch.tensor(y_eval, dtype = torch.float32).to(device) # evaluation points
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # ground truth
Define the DeepONet
The two subnetworks are built by the same function and then combined. branch reads the input function, here the single boundary temperature; trunk reads the spatial location where a prediction is wanted; and deeponet combines them.
The one number to notice is the width of the final layer of each subnetwork, which is $p$, the number of learned basis functions. Everything else about the two stacks — their depth, their width, their activation — is ordinary network design of the kind Artificial Neural Networks (ANN) covered, and nothing in this section depends on those choices. The requirement that binds the two subnetworks together is only that their output dimensions agree, since otherwise the dot product is undefined.
deeponet evaluates every condition at every point in one line. branch(U) is a $10 \times p$ matrix of coefficients, one row per condition, and trunk(Y) is a $25 \times p$ matrix of basis function values, one row per point; their product branch(U) @ trunk(Y).T is the $10 \times 25$ matrix whose entry $(i, j)$ is $\sum_k b_k(u_i)\, t_k(y_j) = G(u_i)(y_j)$. The dot product of Section 2.1 is this matrix product read one entry at a time, and the basis-and-coefficients picture of Section 2.3 is this matrix product read one row at a time.
def new_net(n_in, n_out):
return nn.Sequential(nn.Linear(n_in, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, n_out)).to(device)
p = 16 # number of learned basis functions
torch.manual_seed(0)
branch = new_net(1, p)
trunk = new_net(1, p)
def deeponet(U, Y):
return branch(U) @ trunk(Y).T # (conditions, points)
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 5e-4)
Train
A plain supervised loop follows: predict all ten conditions at all twenty-five points, compare with the analytical ground truth by mean squared error, and step. The data set is small enough to be used whole at every iteration; drawing batches from it becomes necessary only when the table grows too large to use whole, as it does in the next notebook. There is no physics in this objective; it is the data-driven case of Section 2.
for epoch in range(1, 2001):
loss = torch.mean((deeponet(U, Y) - GT)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print('epoch %4d | loss %.4e' % (epoch, loss.item()))
Evaluation
The trained model is now queried at the held-out condition $T_{\text{BC}} = 60$, sweeping the twenty-five evaluation points and comparing the result with the exact solution $T(y) = 60y$.
Notice the shape of the evaluation: one condition in, one row of predictions out, which is the pattern described in Section 2. This is the whole cost of answering a new design question, and it is what replaces a full retraining run in the PINN setting.
To see that the output is not tied to a grid, replace Y by a hundred points, or by points between the training locations; the curve stays smooth and correct because the trunk is a function of $y$ defined everywhere.
U_test = torch.tensor(test_u, dtype = torch.float32).to(device)
with torch.no_grad():
pred = deeponet(U_test, Y).cpu().numpy() # (1, 25)
plt.figure(figsize = (6, 4))
plt.plot(y_eval, test_gt.T, 'k', lw = 4, label = 'Exact', alpha = 0.3)
plt.plot(y_eval, pred.T, 'r--', label = 'DeepONet')
plt.xlabel('y')
plt.ylabel('Temperature')
plt.title('Test Result for Heat Conduction Problem')
plt.legend()
plt.grid(True)
plt.show()
4. Lab 2: Euler Beam¶
Overview
The purpose of this laboratory is the data-driven DeepONet in its most general form: an input function, sampled at sensors, and a set of solved examples. In the heat conduction laboratory the condition was a single number, so the sensor discretisation of Section 2 never had to do any work. This laboratory takes the step that makes DeepONet an operator network in the full sense: the condition is now a distributed load $q(x)$, an actual function defined along the beam, and it enters the branch network as its values at twenty equally spaced stations.
Everything else is deliberately held constant. The trunk network still takes a single coordinate, the combination is still a dot product, the training loop is still plain supervised regression, and the code is nearly a copy of the previous laboratory. The one substantive change is the width of the branch network's input layer, from one to twenty, and that change is worth dwelling on: it is the moment where the fixed sensor grid of Section 2 becomes a concrete architectural commitment. A model trained on twenty stations requires twenty values forever after, at those twenty positions.
The cantilever is fixed at $x = 0$ and free at $x = l$, and carries a distributed load $q(x)$, in N/m, along its length. The right-hand side shows what makes this an operator problem rather than a single beam problem: the load is not one function but a family of them, each a different shape along the beam, and one trained network has to return the deflection for every member of the family.
Problem properties
$$E = 1 \operatorname{pa}, \quad I = 1 \operatorname{kg\cdot m^2}, \quad l = 1 \operatorname{m}$$
The Euler-Bernoulli beam equation relates the transverse deflection $y(x)$ to the distributed load $q(x)$ through the fourth derivative,
$$EI \frac{d^4 y}{dx^4} = -q(x)$$
with the sign convention implied by the exact solution. Because $E$ and $I$ are both set to unity, the flexural rigidity drops out and the equation reduces to $y^{(4)} = -q$.
The training dataset consists of two hundred different loading conditions. The load distribution $q(x)$ is represented using twenty equally spaced points, denoted $u_1, u_2, \cdots, u_{20}$ — these are the sensors of Section 2, and the vector $(u_1, \ldots, u_{20})$ is what the branch network receives.
The external loading $q(x)$ could be arbitrary. To have a family that is varied but easy to generate, the loads are drawn from 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)$$
so that a loading condition is fixed by seven numbers even though it is presented to the network as twenty. That redundancy is intentional and instructive: the branch network is given a twenty-dimensional input whose admissible values lie on a seven-dimensional subspace, and part of what it must learn is to recognise that structure.
The exact solution is elementary. For $q(x) = 1$,
$$y(x) = -{1 \over 24}x^4 + {1 \over 6}x^3 - {1 \over 4}x^2$$
The desired output locations are twenty-five equally spaced evaluation points in $[0, 1]$.
The rank of this solution family
Section 2.3 argued that $p$, the number of learned basis functions, should be compared with the dimension of the space in which the solutions actually live. For this problem that dimension can be computed exactly. The deflection under any load is found by integrating $y'''' = -q$ four times, using the two conditions at the free end, $y''(1) = y'''(1) = 0$, and the two at the clamped end, $y(0) = y'(0) = 0$:
$$ y'''(x) = \int_x^1 q(s)\, ds, \qquad y''(x) = -\int_x^1 y'''(s)\, ds, \qquad y'(x) = \int_0^x y''(s)\, ds, \qquad y(x) = \int_0^x y'(s)\, ds $$
The code of Section 4.2 carries out these four integrations numerically, so that the same few lines serve every load in the family. Integration is linear, so the deflection of a sum of loads is the sum of their deflections. Every load is a combination of seven fixed functions, $1$, $\sin(k\pi x)$ and $\cos(k\pi x)$, so every deflection is the same combination of seven fixed deflection curves, the responses to the seven basis loads:
$$G(u)(y) = \sum_{i=1}^{7} b_i(u)\, t_i(y), \qquad b_i = \text{the seven load coefficients}, \qquad t_i = \text{the seven deflection curves}$$
The solution operator for this family has rank exactly seven. Seven basis functions suffice, and the network is given sixteen, which is why the fit will be easy. The branch, for its part, has only a linear map to learn, from the twenty sampled load values to the seven coefficients; this is the observation of Section 2.3 that a linear PDE with linear dependence on its data has a linear solution operator.
4.2. Solve the Euler Beam Problem with DeepONet¶
The corresponding DeepONet structure is shown in the figure: twenty inputs to the branch network, one to the trunk, and a dot product of two sixteen-dimensional latent vectors.
Sensor and evaluation points
The load is sampled at twenty equally spaced stations, and the deflection is evaluated at twenty-five equally spaced points in $[0, 1]$.
Keep the two grids of this laboratory distinct in your mind, because they are different objects that happen to lie on the same interval. The twenty sensor points are where the input load is sampled and are frozen into the branch network's input layer. The twenty-five evaluation points are where the output deflection is supplied during training, and they constrain nothing about the trained model, which can be queried anywhere. The two grids need not agree in number or in position, and here they do not.
x_sensor = np.linspace(0, 1, 20) # where the load is sampled
x_eval = np.linspace(0, 1, 25).reshape(-1, 1) # where the deflection is evaluated
Load $q(x)$ and its ground truth
Two hundred loads are drawn from the seven-term family for training and one more for testing. Ten loads were enough in Section 3, where the condition was one number; here the branch network has to learn a map from twenty inputs, and with only ten examples it can fit them perfectly in many ways that do not generalise, so more examples are needed. Since the data are generated rather than solved for, more are free. The deflections are computed by the four integrations of Section 4.1, carried out on a fine grid with the cumulative trapezoidal rule, and then read off at the twenty-five evaluation points.
The test load is a member of the same family that the network never sees. As in Section 3, the test is an interpolation test.
def load(coef, x): # q(x) = a0 + sum a_k sin(k pi x) + b_k cos(k pi x)
k = np.arange(1, 4)
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
x_fine = np.linspace(0, 1, 1201) # fine grid for integration; every 50th point is an evaluation point
def integrate_from_right(f): # int_x^1 f(s) ds
F = cumulative_trapezoid(f, x_fine, initial = 0)
return F[:, -1:] - F
def integrate_from_left(f): # int_0^x f(s) ds
return cumulative_trapezoid(f, x_fine, initial = 0)
def deflection(coef):
q = load(coef, x_fine)
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[:, ::50] # (n, 25)
np.random.seed(0)
coef_train = np.random.uniform(-1, 1, (200, 7))
coef_test = np.random.uniform(-1, 1, (1, 7))
train_q = load(coef_train, x_sensor)
test_q = load(coef_test, x_sensor)
train_gt = deflection(coef_train)
test_gt = deflection(coef_test)
print('Conditions for train :', train_q.shape)
print('Ground truth for train :', train_gt.shape)
print('Conditions for test :', test_q.shape)
print('Ground truth for test :', test_gt.shape)
The printed shapes tell the whole story of the data layout, and they differ from the previous laboratory in exactly one place. The loading conditions now have twenty columns rather than one, because each condition is a sampled function rather than a scalar. The ground truth still has one row per condition and one column per evaluation point.
This is also where the practical question of Section 2 becomes concrete: where did these ground truth arrays come from? For a cantilever they can be produced by integrating the beam equation, as the cell above does: two hundred sets of coefficients are drawn at random, the loads are sampled at the twenty stations, and the deflections are integrated at the twenty-five output points. For any problem interesting enough to need operator learning they could not, and each row of the ground truth array would be one run of a finite element solver. That is the offline cost, and the whole argument for the method is that it is paid once.
External load $q(x)$
The load distribution $q(x)$ is represented using twenty equally spaced points $u_1, u_2, \cdots, u_{20}$.
The plot shows ten of the two hundred training loads and the test load as functions of position. The curves are smooth but of varied shape and sign, and the test load is one more member of the same family.
plt.figure(figsize = (6, 4))
plt.plot(x_sensor, train_q[:10].T, 'b-o', markersize = 4, alpha = 0.3) # 10 of 200 shown
plt.plot(x_sensor, test_q.T, 'r-o', markersize = 4, label = 'test')
plt.xlabel('x')
plt.ylabel('q(x)')
plt.title('External Loading Conditions')
plt.legend()
plt.grid(True)
plt.show()
The test load is another member of the seven-term family, so this is again an interpolation test. A load of a shape outside the family, a point load or a discontinuous one, is the case the model has not been prepared for, and this laboratory does not probe it.
Corresponding ground truth for bending displacement
plt.figure(figsize = (6, 4))
plt.plot(x_eval, train_gt[:10].T, 'b-o', markersize = 4, alpha = 0.3) # 10 of 200 shown
plt.plot(x_eval, test_gt.T, 'r-o', markersize = 4, label = 'test')
plt.xlabel('x')
plt.ylabel('y(x)')
plt.title('Deflection of Euler Beam')
plt.legend()
plt.grid(True)
plt.show()
Tensors
The same three tensors as in Section 3: conditions, evaluation points and ground truth.
U = torch.tensor(train_q, dtype = torch.float32).to(device) # conditions
Y = torch.tensor(x_eval, dtype = torch.float32).to(device) # evaluation points
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # ground truth
Define the DeepONet
branch reads the input function, here the discretised loading condition; trunk reads the spatial location; and branch(U) @ trunk(Y).T is the prediction for every load at every point.
The only difference from Section 3 is the branch input width, which is now the sensor count. That number is the architectural embodiment of the fixed sensor grid: it is chosen before training and cannot be altered afterwards without discarding the first layer's weights, which in practice means retraining.
torch.manual_seed(0)
branch = new_net(20, p)
trunk = new_net(1, p)
def deeponet(U, Y):
return branch(U) @ trunk(Y).T # (conditions, points)
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 1e-3)
Train
The objective is again a plain mean squared error against solved deflections, with no physics involved. The run takes seconds, and the rank-seven argument says why the fit is easy.
for epoch in range(1, 5001):
loss = torch.mean((deeponet(U, Y) - GT)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('epoch %4d | loss %.4e' % (epoch, loss.item()))
Evaluation
The trained operator is queried at the held-out load and swept across the twenty-five evaluation points, then compared against the integrated solution.
U_test = torch.tensor(test_q, dtype = torch.float32).to(device)
with torch.no_grad():
pred = deeponet(U_test, Y).cpu().numpy() # (1, 25)
plt.figure(figsize = (6, 4))
plt.plot(x_eval, test_gt.T, 'k', lw = 4, label = 'Exact', alpha = 0.3)
plt.plot(x_eval, pred.T, 'r--', label = 'DeepONet')
plt.xlabel('x')
plt.ylabel('Deflection')
plt.title('Test Result for Euler Beam')
plt.legend()
plt.grid(True)
plt.show()
What the Lab Showed
The condition became a function, and the only change in the code was the width of the branch input, from one to twenty. The prediction lies on the integrated solution for a load the network never saw, because the family of deflections is spanned by seven curves, the trunk has learned a basis for that space, and the branch has learned the map from twenty sampled load values to the coefficients. That division of labour, basis from the trunk and coefficients from the branch, is what every later laboratory relies on.
5. Lab 3: Transient Heat Conduction¶
Overview
The purpose of this laboratory is a time-dependent operator. The laboratories so far asked for the solution of a steady problem, a field over space that does not change. This laboratory adds time. 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. In a DeepONet this needs no new machinery at all: the trunk already takes the coordinate at which the answer is wanted, and the time is simply one more coordinate.
The data-driven setting of Lab 2 is kept, and so is its matrix form: every initial condition is paired with the same set of query points, and one matrix product trains everything. What is new is only that a query point is now a pair, a position and a time. The exact solution is known, so the data need no solver, and the error can be read at any time, including times the network never saw.
5.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 \;\longmapsto\; u, \qquad u = u(x, t) $$
one input function, the initial condition, and one output function of two variables. 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.
What the Network Takes In and Gives Out
The figure shows the arrangement for this laboratory, and it is the architecture of Section 2.1 with nothing added: the branch on the left reads one initial condition $a(x)$, the trunk reads one query point, and their inner product is the temperature there. A DeepONet has two inputs of two different kinds, and keeping them apart is the whole of the design. The branch takes the input function: the initial condition $a(x)$, sampled at $128$ sensor points on the rod, so one initial condition is a vector of $128$ numbers. The trunk takes the coordinate at which the answer is wanted. In the steady laboratories that coordinate was a position; here it is a position and a time together, the pair $(x, t)$. The output is one number, the temperature at that position and that time:
$$ u(x, t) \;\approx\; \sum_{k=1}^{p} b_k(a)\, \tau_k(x, t) $$
The branch output $b_k(a)$ depends on the initial condition only, and the trunk output $\tau_k(x, t)$ on the coordinate only. The trunk therefore learns $p$ basis functions of space and time, and the branch learns the coefficients of a particular solution in that basis. Nothing about the time is special: it is a second column in the trunk's input.
Two consequences follow, and both are what make this architecture comfortable with time. First, the time is not tied to a grid. The trunk accepts any real number, so the trained network can be asked for the temperature at a time that appeared nowhere in the training data, and at a position between the sensors. Second, the network is trained on the same set of query points for every initial condition, so that the matrix form of the first two laboratories applies unchanged: branch(U) @ trunk(Y).T gives, for every initial condition in U and every query point in Y, the predicted temperature, in one matrix product.
The fixed ends are imposed exactly rather than learned: the 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.
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 query points are all pairs $(x_i, t_m)$ of the $128$ grid points with $32$ times, $4096$ points in all, and the ground truth for one initial condition is the exact solution at all of them, laid out time by time as one row of train_gt. The training times are not spaced uniformly. Diffusion does almost all of its work early, so uniform spacing on $(0, 0.5]$ would place only a few times below $t = 0.01$, where the solution changes fastest. Squaring an evenly spaced number puts half of the times below $T/4$ and a quarter below $T/16$; the late times, where the solution is nearly flat, need few.
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), one time t
decay = np.exp(-(np.pi*k)**2*t)
return (coef*decay) @ basis
t_grid = T*np.linspace(0, 1, 32)**2 # 32 times, dense at the beginning
X_mesh, T_mesh = np.meshgrid(x, t_grid) # every (x, t) pair, (32, 128) each
query = np.column_stack([X_mesh.ravel(), T_mesh.ravel()]) # (4096, 2): one row per query point
np.random.seed(0)
coef_train = np.random.normal(size = (100, 20))/k**1.5 # c_k = z_k / k^1.5
coef_test = np.random.normal(size = (25, 20))/k**1.5
train_a = initial(coef_train) # (100, 128)
train_gt = np.concatenate([solution(coef_train, t) for t in t_grid], axis = 1) # (100, 4096), same order as query
print('train condition shape :', train_a.shape)
print('query points shape :', query.shape)
print('train gt shape :', train_gt.shape)
Three initial conditions, and one of them followed through time: the fine structure disappears first and the broad shape last.
plt.figure(figsize = (12, 4))
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)
for t in [0, 0.002, 0.01, 0.03, 0.1, 0.3, 0.5]:
plt.plot(x, solution(coef_train[0], t), label = 't = %g' % t)
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 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
u_frames = np.stack([solution(coef_train[0], t) for t in frames]) # (60, 128)
fig = plt.figure(figsize = (6, 4))
line = plt.plot(x, u_frames[0], 'k')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(u_frames.min() - 0.1, u_frames.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_frames[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
U holds the initial conditions, one per row, and is the branch input. Y holds the query points, one $(x, t)$ pair per row, and is the trunk input. GT is the exact temperature for every initial condition at every query point, so its shape is that of branch(U) @ trunk(Y).T.
U = torch.tensor(train_a, dtype = torch.float32).to(device) # (100, 128)
Y = torch.tensor(query, dtype = torch.float32).to(device) # (4096, 2)
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # (100, 4096)
Define DeepONet Structure
The two networks are built by new_net as in the earlier laboratories, redefined here with $64$ units per layer. The branch takes the $128$ sensor values and the trunk takes the two coordinates; both return $p = 64$ numbers, and the inner product of the two is the prediction. They are wider than before, $64$ units instead of $32$, because the trunk has the harder job here: at early times the solution still carries twenty sine modes, and a narrow network produces the high ones poorly. The envelope $x(1 - x)$ is computed from the first column of the query points, so that deeponet can be evaluated on any set of query points, not only the training set. deeponet is a plain function because it has no parameters of its own; everything learnable lives in branch and trunk. One small indirection is built in for later: the trunk does not read the query points directly but through features, which here returns them unchanged. Lab 4 will replace that one function and nothing else.
def new_net(n_in, n_out):
return nn.Sequential(nn.Linear(n_in, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, n_out)).to(device)
p = 64
def features(Y): # what the trunk reads: here the query point (x, t) itself
return Y
def envelope(Y): # x (1 - x) at every query point, (1, points)
return (Y[:, 0]*(1 - Y[:, 0])).reshape(1, -1)
torch.manual_seed(0)
branch = new_net(n, p) # initial condition at 128 sensors -> p coefficients
trunk = new_net(2, p) # (x, t) -> p basis values
def deeponet(U, Y):
return (branch(U) @ trunk(features(Y)).T)*envelope(Y) # (conditions, points), u = 0 at both ends exactly
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5000, gamma = 0.5) # halve the learning rate every 5000 epochs
Train
Full-batch Adam, as in the beam laboratory: every initial condition and every query point enter each step, and the loss is the mean squared error over the whole $100 \times 4096$ table. Twenty thousand epochs, with the learning rate halved every five thousand; the early-time error keeps falling long after the loss looks flat, and the decay lets the last epochs settle rather than wander.
for epoch in range(1, 20001):
loss = torch.mean((deeponet(U, Y) - GT)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 2000 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
One Initial Condition through Time
predict evaluates a model, given its two networks, for every test initial condition on the grid at one time, and relative_error compares the result with the exact solution; both are used again in Lab 4. Then 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. None of the six times is in t_grid; the trunk is simply handed a new column of times. The relative error over the $25$ test initial conditions is printed for each time as well.
test_a = initial(coef_test) # (25, 128)
U_test = torch.tensor(test_a, dtype = torch.float32).to(device)
def predict(branch_net, trunk_net, t): # every test condition on the grid at time t, (test conditions, 128)
query_t = np.column_stack([x, np.full(n, t)]) # (128, 2): the grid at one time t
Y_t = torch.tensor(query_t, dtype = torch.float32).to(device)
with torch.no_grad():
return ((branch_net(U_test) @ trunk_net(features(Y_t)).T)*envelope(Y_t)).cpu().numpy()
def relative_error(branch_net, trunk_net, t): # mean relative L2 error over the test conditions
exact = solution(coef_test, t)
pred = predict(branch_net, trunk_net, t)
return np.mean(np.linalg.norm(pred - exact, axis = 1)/np.linalg.norm(exact, axis = 1))
plot_times = [0.002, 0.01, 0.03, 0.1, 0.3, 0.5]
colors = ['b', 'g', 'r', 'm', 'c', 'y']
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):
print('t = %-5g | relative L2 error %.4f' % (t, relative_error(branch, trunk, t)))
plt.plot(x, solution(coef_test[j], t), c, alpha = 0.4, linewidth = 3)
plt.plot(x, predict(branch, trunk, t)[j], c + '--', label = 'DeepONet, t = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('Exact (solid) and DeepONet (dashed)')
plt.legend()
plt.grid(True)
plt.show()
Continuous in Time
Because the time is a trunk coordinate, the prediction can be asked for at as many times as one likes. The animation moves through $60$ times, none of them in the training set, with the exact solution as a solid line and the DeepONet as a dashed one.
frames = T*np.linspace(0.001, 1, 60)**2
exact = np.stack([solution(coef_test[j], t) for t in frames]) # (60, 128)
pred = np.stack([predict(branch, trunk, 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 = 'DeepONet')[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 DeepONet, trained once, predicts $u(x, t)$ at every time in $[0, T]$. The time entered as a second coordinate of the trunk, next to the position, and nothing else in the construction changed: the branch still reads the initial condition at its sensors, the inner product still joins the two, and the matrix form of the earlier laboratories still trains everything in one product. The six plotted times and the sixty animated ones were never seen in training, and the network answers at them because the trunk is a function of a real variable, not a lookup over a grid.
The error is small while the solution is large and grows at late times, where the solution has decayed to a small fraction of its initial size: the absolute error there is no larger than before, but the relative error divides it by a much smaller number.
The lesson is about where the inputs go. An operator has an input function and an output function; whatever indexes the input function, here the initial condition, belongs to the branch, and whatever indexes the output function, here position and time, belongs to the trunk. Once that division is made, a time-dependent problem is no harder for a DeepONet than a steady one.
6. Lab 4: A Vibrating String¶
Overview
The purpose of this laboratory is a time-dependent operator whose solution does not settle but rings. Lab 3 made the time a coordinate of the trunk, and its solution relaxed: every mode decayed, the fast ones first, and by the end of the interval almost nothing was left. This laboratory keeps the rod, the sine series, the matrix form and the training loop of Lab 3, and changes the physics from diffusion to vibration. Every mode now oscillates at its own frequency instead of decaying, and the operator has to carry the phase of every one of them across the whole interval.
That changes one thing in the model. A trunk that reads $t$ as a plain coordinate learns an oscillation slowly, for the reason PINN: Limitations and Overcoming Strategies called spectral bias, so the time is not handed to the trunk as a coordinate but as sines and cosines, the Fourier feature embedding of that notebook. It is the one place where this laboratory differs from Lab 3, and it is decided before training rather than discovered after it. The laboratory then goes one step further than any before it: the wave equation is added to the loss, as a residual evaluated through the trunk coordinates, and the physics-informed operator is shown to carry the solution beyond the end of a short record, where the data-driven one cannot go.
6.1. Problem Setup¶
Problem Setup
Picture a string stretched between two fixed supports, a guitar string or a cable. It has been pulled into some shape $a(x)$ and is released from rest at $t = 0$. What is its displacement a moment later, and a few periods later? The ends cannot move, so the disturbance travels along the string, reflects at the supports, and travels back. Unlike heat, nothing is lost: the shape does not relax but keeps moving, and after one full period it is back where it started. This is the problem of every plucked instrument, of a cable between two towers after a gust, and of a transmission line between two terminations; whatever the physical variable, the mathematics below is identical.
Measured in units of the length $L$ and the transit time $L/c$, where $c$ is the wave speed, the problem has no parameters at all. The dimensionless interval studied here, up to $T = 1$, is the time a disturbance takes to travel once from one end to the other: for a $65$ cm guitar string tuned to $110$ Hz it is about $4.5$ milliseconds, and for a $100$ m cable with a wave speed of $50$ m/s it is two seconds. One trained network serves both; only the clock is read differently.
Consider, then, 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$. The operator to be learned is
$$ G: a \;\longmapsto\; u, \qquad u = u(x, t) $$
one input function, the initial shape, and one output function of two variables, exactly as in Lab 3. The initial shapes are drawn from the same sine series,
$$ a(x) = \sum_{k=1}^{10} c_k \sin(k\pi x), \qquad c_k = \frac{z_k}{k^{1.5}}, \quad z_k \sim \mathcal{N}(0, 1) $$
with ten modes rather than twenty, because every mode now has a frequency to be tracked rather than a decay to be forgotten. Substituting one mode into the equation gives $\ddot c_k = -(k\pi)^2 c_k$, so each coefficient oscillates instead of decaying, and with zero initial velocity the exact solution is
$$ u(x, t) = \sum_{k=1}^{10} c_k\, \cos(k\pi t)\, \sin(k\pi x) $$
and the data need no solver. Compare this with the solution of Lab 3: the factor $e^{-k^2\pi^2 t}$ has become $\cos(k\pi t)$, and that is the whole difference. Two features of the new factor decide everything that follows. Its size never shrinks, so the solution at $t = 1$ is as large as at $t = 0$ and a loss on $u$ sees every time alike. And it oscillates, with the $k$th mode completing $k/2$ oscillations in the interval, so the trunk's basis functions have to oscillate in $t$ with it.
The formula also says what the string does, and it is worth knowing before looking at the pictures. Since $\sin(k\pi(1 - x)) = -(-1)^k \sin(k\pi x)$, the solution at $t = 1$ is $u(x, 1) = -a(1 - x)$: after one transit time the string is the mirror image of its initial shape, upside down. At $t = 2$ it is back to $a(x)$, and the motion repeats with period $2$. In between, the initial shape splits into two halves that travel in opposite directions, reflect at the supports with a change of sign, and recombine. This is d'Alembert's description of the string, and the sine series is the same motion written mode by mode.
What the Network Takes In and Gives Out
The figure shows the arrangement: three initial shapes on the left, the branch and trunk in the middle, and on the right one initial shape followed through a full period, where the string passes through its inverted mirror image at $t = 1$ and returns to its starting shape, the motion the trunk has to represent. The branch takes the input function, the initial shape $a(x)$, sampled at $128$ sensor points. The trunk takes the coordinate at which the answer is wanted, the pair $(x, t)$, and the output is one number, the displacement at that position and that time:
$$ u(x, t) \;\approx\; \sum_{k=1}^{p} b_k(a)\, \tau_k(x, t) $$
The trunk learns $p$ basis functions of space and time together, and the branch learns the coefficients of a particular solution in that basis. All of this is as in Lab 3, and so is the matrix form: every initial condition is paired with the same set of query points, and branch(U) @ trunk(Y).T trains everything in one product. The fixed ends are imposed exactly by the envelope $x(1 - x)$, as before.
Fourier Features in Time
What is new is the shape of the dependence on $t$. Along $t$ the solution oscillates, with frequencies from $\pi$ to $10\pi$, and a $\tanh$ network given $t$ as a coordinate builds the slow part of such a dependence first and the fast part only after many more steps, if at all; this is the spectral bias of PINN: Limitations and Overcoming Strategies, now along the time axis of an operator. The remedy is the one that notebook used for a high-frequency solution: give the network the oscillations instead of asking it to build them. The trunk's input is therefore not $(x, t)$ but
$$ \gamma(x, t) = \bigl[\,x,\ \cos(\pi t),\ \sin(\pi t),\ \cos(2\pi t),\ \sin(2\pi t),\ \ldots,\ \cos(M\pi t),\ \sin(M\pi t)\,\bigr] $$
with $M = 10$. The position enters as it is, because the basis in space is a smooth sine and a $\tanh$ network handles it; only the time is embedded. The trunk then has $2M + 1 = 21$ inputs, and any oscillation in $t$ up to the frequency $M\pi$ is a slowly varying combination of inputs that already oscillate. For the string the modal frequencies $k\pi$ happen to be exactly the feature frequencies, which makes the trunk's job especially easy; in general, as on a square membrane whose frequencies are $\pi\sqrt{k^2 + l^2}$, the features need only cover the range of frequencies in the solution, and the network combines them.
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 oscillation factor $\cos(k\pi t)$; both are matrix products with the same basis, and the only change from Lab 3 is in that factor.
The query points are all pairs $(x_i, t_m)$ of the $128$ grid points with $32$ times, $4096$ points in all, and the ground truth for one initial condition is the exact solution at all of them, laid out time by time as one row of train_gt. The training times are spaced uniformly on $(0, T]$. In Lab 3 they were crowded at the beginning, because diffusion does almost all of its work early; a string does the same amount of work at every instant, so no part of the interval deserves more points than another.
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 condition
basis = np.sin(np.pi*np.outer(k, x)) # (10, 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 cos(k pi t) sin(k pi x), one time t
return (coef*np.cos(np.pi*k*t)) @ basis
t_grid = np.linspace(0, T, 33)[1:] # 32 times, evenly spaced on (0, T]
X_mesh, T_mesh = np.meshgrid(x, t_grid) # every (x, t) pair, (32, 128) each
query = np.column_stack([X_mesh.ravel(), T_mesh.ravel()]) # (4096, 2): one row per query point
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
train_a = initial(coef_train) # (40, 128)
train_gt = np.concatenate([solution(coef_train, t) for t in t_grid], axis = 1) # (40, 4096), same order as query
print('train condition shape :', train_a.shape)
print('query points shape :', query.shape)
print('train gt shape :', train_gt.shape)
Three initial shapes are plotted on the left, and on the right one of them is followed through time. Nothing fades: the shape at $t = 0.25$ and $t = 0.5$ is the initial shape split into two travelling halves, at $t = 1$ it is the initial shape mirrored and upside down, exactly as the formula says, and at every time the displacement is as large as at the start.
plt.figure(figsize = (12, 4))
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 Shapes')
plt.grid(True)
plt.subplot(1, 2, 2)
for t in [0, 0.25, 0.5, 0.75, 1.0]:
plt.plot(x, solution(coef_train[0], t), label = 't = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('One Initial Shape Vibrating')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
The same initial shape can be watched vibrating. The frames are spaced uniformly in time, since no part of the interval is calmer than any other.
frames = np.linspace(0, T, 60)
u_frames = np.stack([solution(coef_train[0], t) for t in frames]) # (60, 128)
fig = plt.figure(figsize = (6, 4))
line = plt.plot(x, u_frames[0], 'k')[0]
title = plt.title('t = 0')
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.ylim(u_frames.min() - 0.1, u_frames.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_frames[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
U holds the initial shapes, one per row, and is the branch input. Y holds the query points, one $(x, t)$ pair per row. GT is the exact displacement for every initial shape at every query point, so its shape is that of branch(U) @ trunk(Y).T. U_test holds the test shapes for predict.
U = torch.tensor(train_a, dtype = torch.float32).to(device) # (40, 128)
Y = torch.tensor(query, dtype = torch.float32).to(device) # (4096, 2)
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # (40, 4096)
test_a = initial(coef_test) # (10, 128)
U_test = torch.tensor(test_a, dtype = torch.float32).to(device)
Define DeepONet Structure
Everything is that of Lab 3, new_net, envelope, deeponet and predict, with one function replaced. features now turns each query point $(x, t)$ into the $2M + 1 = 21$ numbers $\gamma(x, t)$ of the problem setup, the position and the sines and cosines of the time, and the trunk reads those, so its input width is $21$ rather than $2$. New networks are built from the same seed.
M = 10
def features(Y): # what the trunk reads: (x, cos(m pi t), sin(m pi t)), m = 1, ..., M
t = Y[:, 1:2]
m = torch.arange(1, M + 1, device = Y.device, dtype = Y.dtype).reshape(1, -1)
return torch.cat([Y[:, 0:1], torch.cos(m*np.pi*t), torch.sin(m*np.pi*t)], dim = 1)
torch.manual_seed(0)
branch = new_net(n, p) # initial shape at 128 sensors -> p coefficients
trunk = new_net(2*M + 1, p) # features of (x, t) -> p basis values
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5000, gamma = 0.5) # halve the learning rate every 5000 epochs
Train
Full-batch Adam, as in Lab 3: every initial shape and every query point enter each step, and the loss is the mean squared error over the whole $40 \times 4096$ table. Twenty thousand epochs, with the learning rate halved every five thousand.
for epoch in range(1, 20001):
loss = torch.mean((deeponet(U, Y) - GT)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 2000 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
One Initial Shape through Time
One test initial shape, the exact solution and the prediction at six times across the interval, ending at $t = 1$ where the string is the inverted mirror image of its start. None of the six times is in t_grid; the trunk is simply handed a new column of times. The relative error over the $10$ test initial shapes is printed for each time as well, with predict and relative_error of Lab 3.
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 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):
print('t = %-5g | relative L2 error %.4f' % (t, relative_error(branch, trunk, t)))
plt.plot(x, solution(coef_test[j], t), c, alpha = 0.4, linewidth = 3)
plt.plot(x, predict(branch, trunk, t)[j], c + '--', label = 'DeepONet, t = %g' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.title('Exact (solid) and DeepONet (dashed)')
plt.legend()
plt.grid(True)
plt.show()
Continuous in Time
Because the time is a trunk coordinate, embedded or not, the prediction can be asked for at as many times as one likes. The animation moves through $60$ times, none of them in the training set, with the exact solution as a solid line and the DeepONet as a dashed one.
frames = np.linspace(0.001, T, 60)
exact = np.stack([solution(coef_test[j], t) for t in frames]) # (60, 128)
pred = np.stack([predict(branch, trunk, 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 = 'DeepONet')[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 DeepONet, trained once, predicts $u(x, t)$ at every time in $[0, T]$: the dashed curves follow the solid ones through the two travelling halves, the reflections at the supports and the inverted shape at the end of the interval, and the animation shows the same through sixty times the network never saw. The error is larger than for the cooling rod, of the order of ten to twenty percent, largest where the travelling fronts are steepest, and nearly the same at every time, since nothing decays and no part of the interval is written off by the loss; a longer run or more training shapes would reduce it.
The lesson is the same as in Lab 3 with one addition. Whatever indexes the output function belongs to the trunk, and time is no exception; but the trunk is a network like any other, and it learns slow dependence before fast. A solution that oscillates in time is therefore given the oscillation at the input, as sines and cosines of $t$, and the trunk's job becomes the smooth one of combining them. Nothing else in the construction changed: the same branch, the same inner product, the same matrix form, and the same training loop as for the cooling rod.
Remark (what the Fourier features buy). The same model with the trunk reading $(x, t)$ as plain coordinates,
featuresof Lab 3 andtrunk = new_net(2, p), is worth training once for comparison. Its loss falls far more slowly, and what it reaches in the same twenty thousand epochs has the spatial shape roughly right and the phases wrong: the right shape at the wrong moment of its motion. The loss is not at fault, since the string is as large at $t = 1$ as at $t = 0$ and every time is weighed alike; the representation is. A sum of basis functions that vary slowly in $t$ cannot oscillate, whatever coefficients the branch chooses.
6.3. Adding the Physics: PI-DeepONet¶
A Short Record
The data-driven model of Section 6.2 was given the solution at thirty-two times across the whole interval, and asked only to interpolate between them. Measurements are rarely so generous. Suppose instead that the string is recorded for a quarter of a transit time, eight times in $(0, 0.25]$, and that the record then stops, and that the question is what the string 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 trunk is a function of $t$, so it will return something there, but nothing in its training says what. The solution beyond the record is not unknown, however. It obeys the wave equation, and the wave equation can be evaluated at any $(x, t)$, inside the record or beyond it, by differentiating the network twice with respect to its trunk coordinates. This is the physics-informed DeepONet of Section 2.4, used here in its hybrid form: the data term of Section 6.2 on the record, and the residual
$$ r_\theta(a; x, t) = \frac{\partial^2 u_\theta}{\partial t^2} - \frac{\partial^2 u_\theta}{\partial x^2} $$
at random pairs of a training initial shape $a$ and a random point $(x, t)$ of the whole rectangle $[0, 1] \times [0, 0.5]$, record and beyond alike. The residual carries the factor $(k\pi)^2$ of the fastest mode, up to $(10\pi)^2 \approx 1000$, so it is divided by $(10\pi)^2$ before being squared, as the residual of PINN: Limitations and Overcoming Strategies was divided by $\omega^2$; this brings it to order one and makes the weight between the two terms a number near one. The initial velocity $\partial u/\partial t = 0$ at $t = 0$ is a second physics term, evaluated at random positions on the line $t = 0$. The initial displacement is in the data, and the fixed ends are in the envelope.
Two models are trained on the record, from the same initial weights and with the same schedule as before. The first is data-driven and sees nothing but the eight early times. The second adds the physics terms. Both are tested up to $t = 0.5$, and the question of this section is whether the equation is enough to carry the solution from the quarter that was recorded to the quarter that was not.
The Pair Form
The matrix form branch(U) @ trunk(Y).T cannot be differentiated condition by condition, because every column of Y serves every row of U. For the residual the prediction is therefore written for pairs, one initial shape with one point, torch.sum(branch(u)*trunk(features(y)), dim = 1), so that torch.autograd.grad returns the derivative of each prediction with respect to its own point. The two forms are the same inner product arranged differently, and the same two networks serve both: the matrix form for the data loss over the record, the pair form for the physics loss over a batch of collocation pairs. physics_loss draws $4096$ pairs, a random training shape with a random point of the rectangle, forms the scaled residual by two calls of derivative in each coordinate, and adds the initial-velocity term on $4096$ further points of the line $t = 0$.
def deeponet_pairs(u, y): # one prediction per (condition, point) pair
return torch.sum(branch(u)*trunk(features(y)), dim = 1, keepdim = True)*envelope(y).T
def derivative(f, y, column):
return torch.autograd.grad(f, y, torch.ones_like(f), create_graph = True)[0][:, column:column + 1]
T_end = 0.5
def physics_loss(n_pairs = 4096):
i = torch.randint(0, len(U), (n_pairs,), device = device)
y = torch.rand(n_pairs, 2, device = device)*torch.tensor([1.0, T_end], device = device)
y.requires_grad_(True)
u = deeponet_pairs(U[i], y)
u_x = derivative(u, y, 0)
u_t = derivative(u, y, 1)
r = (derivative(u_t, y, 1) - derivative(u_x, y, 0))/(10*np.pi)**2 # wave equation residual, scaled to order one
y0 = torch.column_stack([torch.rand(n_pairs, 1, device = device), torch.zeros(n_pairs, 1, device = device)])
y0.requires_grad_(True)
v0 = derivative(deeponet_pairs(U[i], y0), y0, 1) # initial velocity, should be zero
return torch.mean(r**2) + torch.mean(v0**2)
The Record
The data table is that of Section 6.2 cut at $t = 0.25$: the same grid, the same spacing of the times, and a quarter as many of them.
t_record = np.linspace(0, 0.25, 9)[1:] # 8 times in (0, 0.25]: the recorded quarter
X_mesh, T_mesh = np.meshgrid(x, t_record)
query_record = np.column_stack([X_mesh.ravel(), T_mesh.ravel()]) # (1024, 2)
Y_record = torch.tensor(query_record, dtype = torch.float32).to(device)
GT_record = torch.tensor(np.concatenate([solution(coef_train, t) for t in t_record], axis = 1), dtype = torch.float32).to(device)
print('query points shape :', query_record.shape)
print('train gt shape :', tuple(GT_record.shape))
Train on the Record: Data Only
torch.manual_seed(0)
branch = new_net(n, p)
trunk = new_net(2*M + 1, p)
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5000, gamma = 0.5)
for epoch in range(1, 20001):
loss = torch.mean((deeponet(U, Y_record) - GT_record)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 2000 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
branch_data, trunk_data = branch, trunk # data-driven, trained on the record t <= 0.25
Train on the Record: Data and Physics
The same eight times in the data term, and the physics terms over the rectangle up to $t = 0.5$, with weight w_phys. The two losses are printed separately, since they are measured in different units.
torch.manual_seed(0)
branch = new_net(n, p)
trunk = new_net(2*M + 1, p)
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5000, gamma = 0.5)
w_phys = 10.0
for epoch in range(1, 20001):
loss_data = torch.mean((deeponet(U, Y_record) - GT_record)**2)
loss_phys = physics_loss()
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()))
branch_pi, trunk_pi = branch, trunk # physics-informed, data on t <= 0.25 and the equation up to t = 0.5
The Recorded Quarter and the Quarter Beyond
The relative error over the test shapes, by relative_error of Lab 3 with each model's two networks, 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-driven 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(branch_data, trunk_data, t), relative_error(branch_pi, trunk_pi, 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], t), c, alpha = 0.4, linewidth = 3)
plt.plot(x, predict(branch_data, trunk_data, t)[j], c + '--')
plt.plot(x, predict(branch_pi, trunk_pi, 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(branch_data, trunk_data, t) for t in t_curve], 'r', label = 'data only')
plt.plot(t_curve, [relative_error(branch_pi, trunk_pi, 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 physics-informed prediction and the exact solution for the same test shape, through sixty times up to $t = 0.5$. The title reports the time and says whether it lies inside the record or beyond it; 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], t) for t in frames]) # (60, 128)
pred_data = np.stack([predict(branch_data, trunk_data, t)[j] for t in frames]) # (60, 128)
pred_pi = np.stack([predict(branch_pi, trunk_pi, 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, and the data-driven one is slightly better, since it has nothing else to satisfy. Beyond the record they part company at once. The data-driven model has nothing to go on, and its error soon exceeds that of predicting nothing at all; the physics-informed model carries the solution forward, following the travelling halves and their reflections although none of these were in its data, with an error that grows only slowly with the distance from the record. The animation shows the same through the whole interval. The growth is the mark of an equation enforced only approximately: a small residual integrated over a longer time gives a larger error, as a time-stepping scheme accumulates its truncation error step by step, and a longer record or a smaller residual would push the usable horizon further.
This is the division of labour that makes the hybrid worth having. The data fix the coefficients of the operator, which the equation alone cannot do, because 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 then fixes how those coefficients evolve, which the data alone cannot do beyond the last time they were recorded. A physics term is worth its cost where the data leave gaps, in time, in space, or in the number of solved cases, and the equation has to fill them; where the data are dense the equation can only repeat what they already say.
The construction is the one the next notebook uses throughout. Whether the physics term supplements the data, as here and in the cylinder flow, or replaces it altogether, as for the thin plate, the residual is evaluated pair by pair through the trunk coordinates, exactly as a PINN evaluates it through its input, and the branch and trunk are trained on it with the same optimizer as before.
Summary¶
Key takeaways
- A PINN approximates one solution; an operator network approximates the solution operator $\mathcal{G}: u \mapsto s$ over a family of inputs, and is trained once. The offline cost is a set of solved (or physics-constrained) instances, the online saving is that every new condition is a forward pass; whether the trade pays depends on how many queries will be issued.
- The DeepONet evaluates $\mathcal{G}(u)(y) \approx \sum_{k=1}^p b_k(u)\, t_k(y)$: the branch network maps the input function, sampled at $m$ fixed sensors, to $p$ coefficients, and the trunk network maps the output coordinate $y$ to $p$ basis functions. The sensors are frozen into the branch's first layer; the output is a function of $y$ and can be queried anywhere at any resolution. The universal approximation theorem for operators guarantees existence, not size, trainability or extrapolation.
- Reading the operator as basis and coefficients predicts the training runs. Conduction, $T = T_{BC}\,y$, has rank one; the cantilever under loads drawn from a seven-term family has rank seven; both fit in seconds with $p = 16$, and both tests are interpolation inside the training family and say nothing about loads of a new shape.
- In PyTorch the whole model is two
nn.Sequentialnetworks and one product:branch(U) @ trunk(Y).Tgives every condition at every point when the data set is small enough to use whole, andtorch.sum(b*t, dim = 1)gives one value per condition-point pair when each pair has to be handled on its own, as it must when the physics residual is differentiated through the trunk; the training loop is the ordinary supervised one. - A time-dependent problem needs nothing new: the transient-conduction lab gives the trunk the pair $(x, t)$, keeps the matrix form of the beam lab, and after one training run answers at times it never saw.
- On the vibrating string the solution oscillates instead of decaying, and a trunk that reads $t$ as a plain coordinate cannot follow it, the spectral bias of a network along the time axis; given $\cos(m\pi t)$ and $\sin(m\pi t)$ instead, with nothing else changed from the cooling rod, the same matrix-form training follows the motion through the whole interval.
- A physics term carries the solution where the data do not reach. On the string the wave-equation residual is evaluated pair by pair, one initial shape with one random point, by
torch.autograd.gradthrough the trunk coordinates, and added to the data term. Given a record that stops early, the data-driven operator fails beyond it, while the physics-informed one follows the motion on: the data fix the coefficients, the equation fixes how they evolve. This hybrid, data where it exists and the equation everywhere else, is the construction of the next notebook. - A physics-informed DeepONet, developed in the next notebook, keeps the architecture and replaces the data term with the PDE residual at collocation points plus boundary terms, each obtained by
torch.autograd.gradthrough the trunk coordinate exactly as in a PINN, so that no solved examples are needed; the price is a long, delicately balanced optimisation, and the input family still bounds what the operator can be trusted on.
Looking ahead. Deep Operator Networks: Labs carries the DeepONet into seven laboratories, among them the cooling plate whose decaying solution hides the late times from the loss, and ending with a physics-informed thin plate and a hybrid cylinder flow; Fourier Neural Operator learns the same kind of operator by a different route, parameterising the map in frequency space so that it is resolution-independent by construction; Embedding Prior Knowledge into Neural Networks steps back to survey every way, from the loss to the architecture to the data, in which what is already known can be built into a network.