Fourier Neural Operator: Labs
Learning Objectives¶
After studying this notebook, you should be able to
- assemble the input channels of an FNO for a problem whose input is a shape, a scalar design parameter or a pair of fields: the physical fields, the coordinates, and a mask where part of the grid is solid;
- train an FNO on data that come from a numerical solver rather than from a formula, and judge it on held-out cases against a constant predictor and against the discretization error of the labels themselves;
- use an FNO on a grid whose axes are not both spatial, the degree of freedom of a structure along one axis and time along the other, and choose the number of retained modes per axis from how the solution varies along it;
- use a trained operator as a surrogate inside a design optimization, and validate the resulting design once against the original solver;
- evaluate a trained FNO on a finer grid than it was trained on, and say what in the architecture makes that possible;
- state the limits of a purely data-driven operator: it inherits the discretization error of its labels and satisfies no governing equation unless one is put into the loss.
Where we are going. Fourier Neural Operator built the architecture on problems whose answers were known in closed form. This notebook is the practice on problems whose answers are not: every label here comes from a numerical solver, and every input is something an engineer actually varies, a wall shape, a damper design, a pressure and a foundation stiffness. Three labs share the one
FNO2dof the previous notebook. What changes from lab to lab is what is written into the input channels, and that is the skill the notebook teaches.
1. Operator Learning on Engineering Data¶
1.1. What Changes from the Previous Notebook¶
In Fourier Neural Operator the beam, the rod, the string and the plate all had exact solutions, so the data cost nothing and the error could be read exactly. The labs below give up both comforts. The labels are produced by a solver, a Laplace solve on a carved grid, a Newmark time integration, a finite-difference plate solver, and they carry that solver's discretization error, which the operator inherits and can never beat. The error against the labels is therefore the quantity reported, and what it means is stated once in each lab.
The other change is in the input. An FNO reads a field on a grid with one or more values at every point, and the previous notebook always had an obvious field to give it: the load, the initial temperature, the initial shape. Here the input is a shape, a set of design numbers, or two fields at once, and each has to be written onto the grid before the network can read it. A shape becomes a signed distance and a mask; a scalar becomes a constant channel; two fields become two channels. The coordinates are appended as before, so that the network knows where on the grid each point sits. Designing these channels is the part of each lab that is specific to its problem; everything after the lifting layer is the same code.
1.2. The Three Labs¶
| Problem | Input channels | Output | Labels from | |
|---|---|---|---|---|
| Lab 1 | potential flow through a nozzle | wall gap, $x$, $y$, fluid mask | velocity $(u, v)$ | Laplace solve on the carved grid |
| Lab 2 | five-story building with a tuned mass damper | ground motion, $m_d$, $k_d$, $c_d$, DOF, $t$ | response $Z(\mathrm{DOF}, t)$ | Newmark time integration |
| Lab 3 | plate on an elastic foundation | pressure $p$, stiffness $k_f$, $x$, $y$ | deflection $w$ | finite-difference plate solver |
Lab 1 is the field-to-field case with a geometry as input: the wall shape is written as a field, part of the grid is solid, and the network has to learn to leave it empty. Lab 2 departs furthest from a field problem. The grid has the degree of freedom of the structure along one axis and time along the other, the damper design enters as three constant channels, and the trained operator is not the end product but a surrogate inside a design optimization, which is where a fast approximate solver earns its keep. Lab 3 is the plainest case, two input fields to one output field on one square grid, and it is where a trained operator is evaluated on a finer grid than it was trained on.
2. Lab 1: Potential Flow Through a Nozzle¶
Overview
The purpose of this laboratory is to learn an operator whose input is a geometry. A converging-diverging nozzle is described by the shape of its walls, and the quantity wanted is the velocity field of the flow through it. The question is how a shape is handed to a network that reads fields on a grid, and how well the network learns to respect the solid wall it is told about.
2.1. Problem Setup¶
The flow is incompressible, inviscid and irrotational, so the stream function satisfies Laplace's equation and the velocity is its rotated gradient,
$$ \nabla^2\psi = 0, \qquad u = \frac{\partial\psi}{\partial y}, \qquad v = -\frac{\partial\psi}{\partial x} $$
The lower wall is $\psi = 0$ and the upper wall is $\psi = 1$, so the flux through every section is one and the speed rises where the duct narrows.
The labels are computed on a Cartesian grid of $64 \times 64$ points carved by the walls, so the wall becomes a staircase. Against a $256^2$ solve the $64^2$ labels carry about four to five percent discretization error, and that is the accuracy floor of the data: an operator that reproduced the labels exactly would still be that far from the flow.
What is being learned
The operator maps the wall shape to the velocity field, $G: a(x, y) \mapsto (u, v)(x, y)$. It is nonlinear in the shape, since narrowing a duct does not change the speed in proportion, and its output is two fields rather than one.
What the Network Takes In and Gives Out
The shape has to be written onto the grid. At each of the $64 \times 64$ points the network is given four numbers,
$$ \text{input at } (x_i, y_j) = \big(\, \text{gap}(x_i, y_j), \;\; x_i, \;\; y_j, \;\; \text{mask}(x_i, y_j) \,\big) $$
The gap is the signed distance of the point from the nearest wall, positive inside the duct and negative in the solid, normalized by its mean and standard deviation over the training set; it is the shape, read at every point. The coordinates are the channels of every lab, since the inlet and the outlet are held differently and the network cannot tell by itself where a point lies. The mask is one in the fluid and zero in the solid, and it is stored with the labels rather than recomputed from the gap, because the labels leave the inlet and outlet columns empty while the gap does not; recomputing it would make the input disagree with the output on exactly those columns. One sample has shape $(64, 64, 4)$ and its label $(64, 64, 2)$, the two velocity components, zero in the solid.
2.2. Solve the Nozzle Problem using FNO¶
Import Library
The same libraries as Fourier Neural Operator, and scipy.optimize for the design search of Lab 2.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from scipy.optimize import minimize
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Load Data
The nozzle shapes and their Laplace-solver labels are generated separately, in L3_FNO_Nozzle_data_generation.ipynb, and loaded from fno_nozzle_dataset.npz: fifteen hundred shapes for training and two hundred for testing, each with its four input channels and two output fields on the $64 \times 64$ grid.
!gdown 1Wlh-IMBHiRHoYhFwmVSApF7NAQYRDBzm -O /content/fno_nozzle_dataset.npz
data = np.load('/content/fno_nozzle_dataset.npz')
L = 2.0 # nozzle length
train_a = data['x_train_full'] # (1500, 64, 64, 4): gap, x, y, mask
train_gt = data['y_train_full'] # (1500, 64, 64, 2): u, v
test_a = data['x_test_full'] # (200, 64, 64, 4)
test_gt = data['y_test_full'] # (200, 64, 64, 2)
n = train_a.shape[1]
print('train condition shape :', train_a.shape)
print('train gt shape :', train_gt.shape)
print('test condition shape :', test_a.shape)
print('test gt shape :', test_gt.shape)
One nozzle, solved on a fine $257 \times 257$ grid, shows what the labels contain: the speed field, highest at the throat, and the speed along the centreline against the one-dimensional estimate $1/h(x)$ from the flux. The two agree except near the throat, where the flow turns and the one-dimensional picture is too crude; the operator is being asked for the two-dimensional answer.
X_257, Y_257, mask_257, uv_257 = data['X_257'], data['Y_257'], data['mask_257'], data['uv_257']
y_up_257, y_lo_257 = data['y_up_257'], data['y_lo_257']
x_257 = np.linspace(0, L, 257)
speed_257 = np.linalg.norm(uv_257, axis = -1)
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
plt.pcolormesh(X_257, Y_257, np.where(mask_257, speed_257, np.nan), cmap = 'rainbow', shading = 'auto')
plt.colorbar(label = '|V|')
plt.plot(x_257, y_up_257, 'k', x_257, y_lo_257, 'k')
plt.xlim(0, L)
plt.ylim(-0.65, 0.65)
plt.gca().set_aspect('equal')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Exact Velocity Magnitude')
plt.subplot(1, 2, 2)
plt.plot(x_257, uv_257[:, 257//2, 0], 'k', alpha = 0.3, linewidth = 4, label = '2D exact')
plt.plot(x_257, 1/(y_up_257 - y_lo_257), 'r--', label = '1D theory, 1/h(x)')
plt.xlabel('x')
plt.ylabel('u')
plt.title('Centreline Speed')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
Tensors
The four channels as A and the two velocity fields as GT; get_batch draws a random subset of the solved shapes at each step, a sample being a whole nozzle.
A = torch.tensor(train_a, dtype = torch.float32).to(device) # (1500, 64, 64, 4)
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # (1500, 64, 64, 2)
A_test = torch.tensor(test_a, dtype = torch.float32).to(device) # (200, 64, 64, 4)
def get_batch(batch_size = 16):
i = torch.randint(0, len(A), (batch_size,))
return A[i], GT[i]
Define FNO Structure
spectral_conv2d and FNO2d are those of Fourier Neural Operator, Lab 4, with one generalization: the number of retained modes may differ between the two axes, so modes is a pair, and R has shape (2, width, width, modes[0], modes[1]). Lab 2 needs this, since its two axes are of different kinds. The class takes the number of input and output channels and the padding as arguments; here it reads four channels, returns two, and is padded by nine in both directions, because the nozzle is not periodic. predict evaluates a model on a large set in chunks, and relative_error is the mean over the samples of the relative $L^2$ error of each field, the measure reported in every lab.
def spectral_conv2d(v, R): # F^-1 (R . F(v)) along the two grid axes, only the lowest modes are kept
m1, m2 = R.shape[-2:]
v_ft = torch.fft.rfft2(v, dim = (1, 2)) # v: (batch, n1, n2, width) -> (batch, n1, n2//2 + 1, width)
out_ft = torch.zeros_like(v_ft)
out_ft[:, :m1, :m2] = torch.einsum('bxyi,ioxy->bxyo', v_ft[:, :m1, :m2], R[0])
out_ft[:, -m1:, :m2] = torch.einsum('bxyi,ioxy->bxyo', v_ft[:, -m1:, :m2], R[1])
return torch.fft.irfft2(out_ft, s = v.shape[1:3], dim = (1, 2))
class FNO2d(nn.Module):
def __init__(self, width, modes = (12, 12), in_channels = 1, out_channels = 1, padding = 0):
super().__init__()
self.padding = padding
self.P = nn.Linear(in_channels, width) # lifting: input channels -> width channels
self.R = nn.ParameterList([nn.Parameter(torch.rand(2, width, width, *modes, dtype = torch.cfloat)/width)
for _ in range(4)]) # spectral weights, one per layer
self.W = nn.ModuleList([nn.Linear(width, width) for _ in range(4)]) # pointwise paths, one per layer
self.Q = nn.Sequential(nn.Linear(width, 128), nn.GELU(), nn.Linear(128, out_channels)) # projection
def forward(self, a): # a: (batch, n1, n2, in_channels)
v = self.P(a) # (batch, n1, n2, width)
v = F.pad(v, [0, 0, 0, self.padding, 0, self.padding]) # pad both grid axes
for k in range(4):
v = spectral_conv2d(v, self.R[k]) + self.W[k](v)
if k < 3:
v = F.gelu(v)
v = v[:, :v.shape[1] - self.padding, :v.shape[2] - self.padding]
return self.Q(v) # (batch, n1, n2, out_channels)
def predict(model, a, chunk = 50): # the whole set, chunk by chunk
with torch.no_grad():
return torch.cat([model(a[i:i + chunk]) for i in range(0, len(a), chunk)]).cpu().numpy()
def relative_error(pred, gt): # mean over the samples of the relative L2 error of each field
pred, gt = pred.reshape(len(pred), -1), gt.reshape(len(gt), -1)
return np.mean(np.linalg.norm(pred - gt, axis = 1)/np.linalg.norm(gt, axis = 1))
A Training Function
Twelve modes in each direction and a latent width of $32$, the settings of the previous notebook. The loop is the ordinary supervised one, a fresh batch at each step and one Adam step on the mean squared error against the solver's fields; it is wrapped in a function so that Labs 1 and 2 are trained by one call.
width = 32
def train(model, epochs = 5000, batch_size = 16):
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
for epoch in range(1, epochs + 1):
a, gt = get_batch(batch_size)
loss = torch.mean((model(a) - gt)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
model.eval()
return model
Train
Five thousand steps with batches of sixteen on the fifteen hundred training shapes.
torch.manual_seed(0)
model = train(FNO2d(width, modes = (12, 12), in_channels = 4, out_channels = 2, padding = 9).to(device))
Evaluation
The trained operator is applied to the two hundred test shapes, and its error is compared with that of a constant predictor, the mean training velocity field returned for every input. The constant predictor is the right baseline here, because every nozzle shares the same inlet, the same outlet and roughly the same flow, so a good part of any label is predictable without looking at the shape; what the operator has to earn is the rest.
pred_test = predict(model, A_test) # (200, 64, 64, 2)
constant = np.broadcast_to(train_gt.mean(0), test_gt.shape) # the mean training field for every input
print('relative L2 error, constant predictor : %.4f' % relative_error(constant, test_gt))
print('relative L2 error, FNO : %.4f' % relative_error(pred_test, test_gt))
Example Prediction
The worked nozzle on the $64 \times 64$ grid: its geometry, the exact and predicted speed on one color scale, and the pointwise velocity error. The error map is the informative panel, since it shows where the operator is weak rather than how weak it is on average.
X_64, Y_64, mask_64, uv_64 = data['X_64'], data['Y_64'], data['mask_64'], data['uv_64']
y_up_64, y_lo_64 = data['y_up_64'], data['y_lo_64']
x_64 = np.linspace(0, L, n)
a_example = torch.tensor(data['example_input_64_full'], dtype = torch.float32).to(device)
pred_example = predict(model, a_example)[0] # (64, 64, 2)
print('relative L2 error of the example : %.4f' % relative_error(pred_example[None], uv_64[None]))
speed_exact = np.linalg.norm(uv_64, axis = -1)
speed_fno = np.linalg.norm(pred_example, axis = -1)
error = np.linalg.norm(pred_example - uv_64, axis = -1)
fields = [(mask_64.astype(float), 'Nozzle Geometry', 'gray', 1), (speed_exact, 'Exact Speed', 'rainbow', speed_exact.max()),
(speed_fno, 'FNO Speed', 'rainbow', speed_exact.max()), (error, 'Velocity Error', 'magma', error.max())]
plt.figure(figsize = (16, 3.5))
for j, (field, title, cmap, vmax) in enumerate(fields):
plt.subplot(1, 4, j + 1)
plt.pcolormesh(X_64, Y_64, field, cmap = cmap, shading = 'auto', vmin = 0, vmax = vmax)
plt.colorbar()
plt.plot(x_64, y_up_64, 'k', x_64, y_lo_64, 'k', linewidth = 1)
plt.xlim(0, L)
plt.ylim(-0.65, 0.65)
plt.gca().set_aspect('equal')
plt.xlabel('x')
plt.ylabel('y')
plt.title(title)
plt.tight_layout()
plt.show()
2.3. What This Lab Showed¶
A geometry became an operator input by being written onto the grid, as a signed distance, a mask and the coordinates, and the network learned the map from shape to velocity field well below the discretization error of the labels themselves, from a constant predictor that already explains most of each field. The solid region is part of what was learned: the network was told where the wall is through the mask and returns zero velocity there, which is the property a purely data-driven operator has to acquire from the data, since nothing in the architecture knows what a wall is.
The reported error is an error against a staircase approximation of the true flow, and no amount of training lowers the staircase. An operator trained on solver output is a surrogate for that solver, at that resolution, and inherits its error as a floor.
3. Lab 2: Vibration Reduction with a Tuned Mass Damper¶
Overview
The purpose of this laboratory is twofold. The grid the FNO works on is not a spatial field at all: one axis is the degree of freedom of a five-story building and the other is time, so a $6 \times 512$ image is six response histories stacked, and the damper design enters as three numbers spread over the grid as constant channels. And the trained operator is not the end product. It is used as a surrogate inside an optimization over the damper design, with the original solver called once at the end to check the answer, which is the pattern in which a fast approximate solver earns its keep.
3.1. Problem Setup¶
The primary structure is a five-degree-of-freedom shear building with unit masses, $m_1 = \cdots = m_5 = 1$, story stiffnesses $k_1 = \cdots = k_5 = 10$ and five percent Rayleigh damping. A tuned mass damper of mass $m_d$, stiffness $k_d$ and damping $c_d$ is attached to the fifth floor, and the whole is a six-degree-of-freedom system under ground acceleration $\ddot y_g(t)$,
$$ M_6\,\ddot{\mathbf z} + C_6\,\dot{\mathbf z} + K_6\,\mathbf z = -M_6\,\mathbf r\,\ddot y_g(t), \qquad \mathbf z = [x_1, x_2, x_3, x_4, x_5, x_d]^T $$
The ground motions are a decaying-envelope sum of three sinusoids with random amplitudes, frequency ratios and phases,
$$ \ddot y_g(t) = \left(1 - e^{-0.6t}\right) \sum_{i=1}^{3} A_i \sin(r_i\,\omega_i t + \phi_i) $$
so that the family is broad enough that the operator cannot succeed by memorizing one signal. The responses are computed by Newmark time integration at $512$ instants.
What is being learned
The operator maps a ground motion and a damper design to the response of all six degrees of freedom over the record,
$$ G: \big[\ddot y_g(t),\; m_d,\; k_d,\; c_d\big] \;\longmapsto\; Z(\mathrm{DOF}, t) \in \mathbb{R}^{6 \times 512} $$
What the Network Takes In and Gives Out
The grid has the six degrees of freedom along one axis and the $512$ instants along the other, and at each point $(\mathrm{DOF}, t)$ the network is given six numbers,
$$ \text{input at } (\mathrm{DOF}, t) = \big(\, \ddot y_g(t), \;\; m_d, \;\; k_d, \;\; c_d, \;\; \mathrm{DOF}, \;\; t \,\big) $$
The ground motion is the same along the DOF axis, since every floor feels the same ground. The three damper parameters are constant over the whole grid, the device of Fourier Neural Operator for handing a scalar to a field-to-field network. The last two are the coordinates of the grid, the degree of freedom and the time. One sample has shape $(6, 512, 6)$ and its label $(6, 512, 1)$, the displacement of every degree of freedom at every instant.
How many modes
The retained modes differ by axis. Along the six degrees of freedom the response varies slowly, from floor to floor, and three modes are kept; along the $512$ instants it oscillates at the natural frequencies of the building and sixteen are kept. This is the first place where the two axes of the grid are of different kinds, and the mode counts say so.
3.2. Solve the Problem using FNO¶
Load Data
The data-generation notebook stores the building matrices, the training and test cases and the grid in FNO_TMD_data.npz: five hundred cases for training and a hundred for testing, each a ground motion with its damper parameters and the response it produces. The natural frequencies and effective modal masses of the building without the damper are printed first. The first mode carries most of the effective mass, which is why the damper will be tuned to it.
!gdown 1_PZady7iIiJHR-P3m5bC9GXTBqDf17Yf -O /content/FNO_TMD_data.npz
dataset = np.load('/content/FNO_TMD_data.npz')
M, K, C = dataset['M'], dataset['K'], dataset['C'] # the 5-DOF building
omega_n = dataset['omega_n'] # natural frequencies (rad/s)
eig_vec = dataset['eig_vec'] # mode shapes
effective_mass_ratio = dataset['effective_mass_ratio']
n_floor = len(omega_n)
n_dof = n_floor + 1 # five floors and the damper
for i in range(n_floor):
print('mode %d : %.4f rad/s | effective mass ratio %.4f' % (i + 1, omega_n[i], effective_mass_ratio[i]))
plt.figure(figsize = (6, 4))
for i in range(n_floor):
plt.plot(eig_vec[:, i], np.arange(1, n_floor + 1), '-o', label = 'Mode %d' % (i + 1))
plt.xlabel('Normalized mode shape')
plt.ylabel('Floor')
plt.title('Mode Shapes of the 5-DOF Structure')
plt.legend()
plt.grid(True)
plt.show()
train_a, train_y, train_gt = dataset['train_a'], dataset['train_y'], dataset['train_gt'] # conditions, coordinates, responses
test_a, test_y, test_gt = dataset['test_a'], dataset['test_y'], dataset['test_gt']
ag_test = dataset['ag_test'] # the test ground motions, (100, 512)
mesh_pts = dataset['mesh_pts'] # (6, 512, 2): DOF and time at every grid point
record_time = dataset['record_time'] # (512,)
n_time = train_a.shape[2]
print('train condition shape :', train_a.shape)
print('train location shape :', train_y.shape)
print('train gt shape :', train_gt.shape)
print('test condition shape :', test_a.shape)
print('test gt shape :', test_gt.shape)
Ten of the training ground motions and the test motion that will be used for the design, each a sum of three sinusoids growing out of rest.
plt.figure(figsize = (6, 4))
for i in range(10):
plt.plot(record_time, train_a[i, 0, :, 0], 'k', alpha = 0.3)
plt.plot(record_time, test_a[10, 0, :, 0], 'r', label = 'test case 10')
plt.xlabel('t')
plt.ylabel('ground acceleration')
plt.title('Base Excitations')
plt.legend()
plt.grid(True)
plt.show()
Tensors
The four condition channels and the two coordinate channels are concatenated into the six-channel input A; GT holds the responses. A batch is sixteen complete cases.
A = torch.tensor(np.concatenate([train_a, train_y], -1), dtype = torch.float32).to(device) # (500, 6, 512, 6)
GT = torch.tensor(train_gt, dtype = torch.float32).to(device) # (500, 6, 512, 1)
A_test = torch.tensor(np.concatenate([test_a, test_y], -1), dtype = torch.float32).to(device) # (100, 6, 512, 6)
Train
The FNO2d of Lab 1 with six input channels, one output channel, three modes along the degrees of freedom and sixteen along time; fifteen thousand steps with batches of sixteen on the mean squared error against the integrated responses. No physics term appears: the operator is learned from solver output alone, a surrogate for the simulation and not a substitute for the equations.
torch.manual_seed(0)
model = train(FNO2d(width, modes = (3, 16), in_channels = 6, out_channels = 1, padding = 9).to(device), epochs = 15000)
Result
The error over the hundred unseen cases, and one of them in full: the ground acceleration, the predicted $6 \times 512$ response field and the integrated one on the same color scale.
pred_test = predict(model, A_test) # (100, 6, 512, 1)
print('relative L2 error on the test set : %.4f' % relative_error(pred_test, test_gt))
j = 10 # the test case to show
labels = ['Floor 1', 'Floor 2', 'Floor 3', 'Floor 4', 'Floor 5', 'TMD']
vmin, vmax = test_gt[j].min(), test_gt[j].max()
plt.figure(figsize = (12, 3))
for k, (field, title, cmap) in enumerate([(test_a[j, :, :, 0], 'Ground Acceleration', 'gray'),
(pred_test[j, :, :, 0], 'FNO', 'jet'), (test_gt[j, :, :, 0], 'Numerical', 'jet')]):
plt.subplot(1, 3, k + 1)
plt.imshow(field, aspect = 'auto', cmap = cmap, vmin = None if k == 0 else vmin, vmax = None if k == 0 else vmax)
plt.colorbar()
plt.yticks(np.arange(n_dof), labels)
plt.xlabel('time step')
plt.title(title)
plt.tight_layout()
plt.show()
Top-Floor Response
The fifth floor is where the design will be judged, so this is the curve that matters. The prediction has to keep its phase over the whole record, which is harder than matching the envelope.
plt.figure(figsize = (6, 4))
plt.plot(record_time, test_gt[j, n_floor - 1, :, 0], 'k', alpha = 0.3, linewidth = 4, label = 'Numerical')
plt.plot(record_time, pred_test[j, n_floor - 1, :, 0], 'r--', label = 'FNO')
plt.xlabel('t')
plt.ylabel('$x_5(t)$')
plt.title('5th Floor Displacement')
plt.legend()
plt.grid(True)
plt.show()
3.3. TMD Design Optimization¶
The ground motion is fixed to the test case above, and the damper is designed for it. The design variables are the damper mass $m_d$, the tuning ratio $r = \omega_d/\omega_1$ of the damper frequency to the first natural frequency, and the damping ratio $\zeta_d$, from which
$$ k_d = m_d\,(r\,\omega_1)^2, \qquad c_d = 2\,\zeta_d\,m_d\,(r\,\omega_1) $$
The objective is the root mean square of the fifth-floor displacement over the second half of the record, once the transient has passed,
$$ \min_{m_d,\, r,\, \zeta_d} \; \mathrm{RMS}\big(x_5(t)\big), \qquad t \in [30, 60]\ \mathrm{s} $$
The Objective Function
Read this function as the point of the laboratory. It writes a candidate design into the three constant channels, calls the trained network, reads the fifth-floor response off the output field and returns its RMS: a complete structural evaluation in one forward pass.
ag_opt = ag_test[j] # the ground motion to design for
bounds = [(0.03, 0.12), (0.85, 1.15), (0.04, 0.20)] # m_d, r, zeta_d
def tmd_input(m_d, k_d, c_d): # the six-channel field of one design, (1, 6, 512, 6)
a = np.zeros((1, n_dof, n_time, 4), dtype = np.float32)
a[0, :, :, 0] = ag_opt
a[0, :, :, 1], a[0, :, :, 2], a[0, :, :, 3] = m_d, k_d, c_d
return torch.tensor(np.concatenate([a, mesh_pts[None]], -1), dtype = torch.float32).to(device)
def tmd_objective(z):
m_d, r, zeta_d = z
omega_d = r*omega_n[0]
k_d, c_d = m_d*omega_d**2, 2*zeta_d*m_d*omega_d
x = predict(model, tmd_input(m_d, k_d, c_d))[0, :, :, 0] # (6, 512)
return np.sqrt(np.mean(x[n_floor - 1, n_time//2:]**2)) # RMS of the 5th floor over the second half
Search
Nelder–Mead needs no derivatives and makes many objective calls, which is exactly the situation in which a surrogate pays for itself; each call would otherwise be a six-degree-of-freedom time integration.
result = minimize(tmd_objective, x0 = [0.08, 1.00, 0.10], method = 'Nelder-Mead', bounds = bounds, options = {'maxiter': 80})
best_m, best_r, best_zeta = result.x
best_omega = best_r*omega_n[0]
best_k, best_c = best_m*best_omega**2, 2*best_zeta*best_m*best_omega
print('optimal TMD mass : %.4f' % best_m)
print('optimal TMD stiffness : %.4f' % best_k)
print('optimal TMD damping : %.4f' % best_c)
print('tuning ratio : %.4f' % best_r)
print('damping ratio : %.4f' % best_zeta)
print('objective (RMS x5) : %.4f' % result.fun)
3.4. Numerical Validation¶
The optimizer worked entirely inside the surrogate, so the winning design is checked once against the original solver: the six-degree-of-freedom system with the optimal damper is integrated by the Newmark method under the same ground motion, the fifth-floor response is compared with the FNO's, and the reduction in RMS relative to the building without a damper is computed from the solver's answer. A fast search followed by one trustworthy evaluation is the shape of every surrogate-based design loop.
def newmark(M_sys, C_sys, K_sys, ag): # average-acceleration Newmark integration under base motion
beta, gamma = 0.25, 0.5
dt = record_time[1] - record_time[0]
a0, a1, a2 = 1/(beta*dt**2), gamma/(beta*dt), 1/(beta*dt)
a3, a4, a5 = 1/(2*beta) - 1, gamma/beta - 1, dt*(gamma/(2*beta) - 1)
K_eff = K_sys + a1*C_sys + a0*M_sys
r = np.ones(len(M_sys))
x, v, acc = np.zeros((3, len(M_sys), n_time))
acc[:, 0] = np.linalg.solve(M_sys, -M_sys @ r*ag[0])
for i in range(n_time - 1):
p_eff = -M_sys @ r*ag[i + 1] + M_sys @ (a0*x[:, i] + a2*v[:, i] + a3*acc[:, i]) + C_sys @ (a1*x[:, i] + a4*v[:, i] + a5*acc[:, i])
x[:, i + 1] = np.linalg.solve(K_eff, p_eff)
acc[:, i + 1] = a0*(x[:, i + 1] - x[:, i]) - a2*v[:, i] - a3*acc[:, i]
v[:, i + 1] = v[:, i] + dt*((1 - gamma)*acc[:, i] + gamma*acc[:, i + 1])
return x
def build_6dof(m_d, k_d, c_d): # the building with the damper on the 5th floor
M6, C6, K6 = np.zeros((3, n_dof, n_dof))
M6[:n_floor, :n_floor], C6[:n_floor, :n_floor], K6[:n_floor, :n_floor] = M, C, K
M6[-1, -1] = m_d
for S, s in [(C6, c_d), (K6, k_d)]:
S[n_floor - 1, n_floor - 1] += s
S[-1, -1] += s
S[n_floor - 1, -1] -= s
S[-1, n_floor - 1] -= s
return M6, C6, K6
x_fno = predict(model, tmd_input(best_m, best_k, best_c))[0, :, :, 0]
x_num = newmark(*build_6dof(best_m, best_k, best_c), ag_opt)
x_base = newmark(M, C, K, ag_opt) # the building without a damper
rms_base = np.sqrt(np.mean(x_base[-1, n_time//2:]**2))
rms_best = np.sqrt(np.mean(x_num[n_floor - 1, n_time//2:]**2))
print('RMS reduction of the 5th floor, by the solver : %.2f %%' % (100*(1 - rms_best/rms_base)))
plt.figure(figsize = (9, 3))
plt.subplot(1, 2, 1)
plt.plot(record_time, x_num[n_floor - 1], 'k', alpha = 0.3, linewidth = 4, label = 'Numerical')
plt.plot(record_time, x_fno[n_floor - 1], 'r--', label = 'FNO')
plt.xlabel('t')
plt.ylabel('$x_5(t)$')
plt.title('Optimal TMD: FNO against the Solver')
plt.legend()
plt.grid(True)
plt.subplot(1, 2, 2)
plt.plot(record_time, x_base[-1], 'k', alpha = 0.3, linewidth = 4, label = 'without TMD')
plt.plot(record_time, x_num[n_floor - 1], 'r--', label = 'optimal TMD')
plt.xlabel('t')
plt.ylabel('$x_5(t)$')
plt.title('Vibration Reduction')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
3.5. What This Lab Showed¶
An FNO does not care what its two axes mean. A grid of degrees of freedom against time is read exactly like a grid of $x$ against $y$, the damper design entered as three constant channels, and the mode count per axis was chosen from how the response varies along it, slowly across floors and quickly in time. The same FNO2d served, with its arguments changed.
The second half of the laboratory is the reason to train such a surrogate at all. Eighty objective evaluations inside a derivative-free search cost eighty forward passes instead of eighty time integrations, and the design they found was then checked, once, against the integrator. The surrogate's error is visible in that check, as the gap between the dashed and the solid curve, and it is the solver's reduction figure, not the surrogate's objective value, that is the result. A surrogate is for searching; the solver is for deciding.
4. Lab 3: Plate on a Heterogeneous Elastic Foundation¶
Overview
The purpose of this laboratory is the plainest field-to-field case, two input fields in and one output field out on one square grid, and a check that a trained operator can be evaluated on a finer grid than it was trained on. The pressure on a plate and the stiffness of the foundation under it both vary from point to point, and the deflection depends on both at once.
4.1. Problem Setup¶
A simply supported thin plate rests on a Winkler foundation whose stiffness varies over the plate, and carries a spatially varying pressure. Its deflection $w$ satisfies
$$ D\,\nabla^4 w(x, y) + k_f(x, y)\, w(x, y) = p(x, y), \qquad D = \frac{E h^3}{12(1 - \nu^2)} $$
where $p$ is the applied pressure, $k_f$ the local foundation stiffness, so that $k_f w$ is the reaction of the foundation, and $D$ the flexural rigidity. The plate is one metre square, $E = 70\ \mathrm{GPa}$, $\nu = 0.3$ and $h = 10\ \mathrm{mm}$, the same for every sample; only $p$ and $k_f$ vary, and the labels come from a finite-difference solution of the plate equation on a $64 \times 64$ grid.
What is being learned
The operator maps the two fields to the deflection, $G: [p(x, y), k_f(x, y)] \mapsto w(x, y)$. For a fixed foundation it is linear in the pressure, but it is not linear in the stiffness, since a stiffer foundation reduces the deflection in a way that depends on the pressure it carries, so the two inputs interact.
What the Network Takes In and Gives Out
At each of the $64 \times 64$ points the network is given four numbers,
$$ \text{input at } (x_i, y_j) = \big(\, p(x_i, y_j), \;\; k_f(x_i, y_j), \;\; x_i, \;\; y_j \,\big) $$
Pressure, stiffness and deflection are physically unrelated quantities with very different magnitudes, so each is divided by its maximum over the training set before it reaches the network or the loss; otherwise a single squared-error loss would be dominated by whichever happens to be largest. The output is the scaled deflection at every grid point, one sample $(64, 64, 1)$, converted back to metres for the figures.
The plate is simply supported, not periodic, so the field is padded before the transforms, by one eighth of the grid: eight points on the $64$-point grid, twelve on the $96$-point grid of Section 4.3. Writing the padding as a fraction of the grid rather than a fixed number is what lets the same model move to another resolution.
4.2. Solve the Plate Problem using FNO¶
Load Data
The dataset is generated separately and loaded from fno_plate_k_dataset.npz: eight hundred cases for training, two hundred for testing, each a pressure field, a stiffness field, the grid coordinates and the deflection the solver computed. Nothing in this notebook solves the plate equation.
!gdown 1wQXtSO6O8MklzJUcMHTmOuqPk5gnGgJ2 -O /content/fno_plate_k_dataset.npz
data = np.load('/content/fno_plate_k_dataset.npz')
train_p = np.concatenate([data['train_p'], data['val_p']]) # (800, 64, 64): pressure (N/m^2)
train_k = np.concatenate([data['train_k'], data['val_k']]) # (800, 64, 64): foundation stiffness (N/m^3)
train_xy = np.concatenate([data['train_xy'], data['val_xy']]) # (800, 64, 64, 2): coordinates
train_w = np.concatenate([data['train_w'], data['val_w']]) # (800, 64, 64): deflection (m)
test_p, test_k, test_xy, test_w = data['test_p'], data['test_k'], data['test_xy'], data['test_w']
Lx = Ly = 1.0
n = train_p.shape[1]
print('train pressure shape :', train_p.shape)
print('test pressure shape :', test_p.shape)
Tensors
Each field is divided by its maximum over the training set; make_input assembles the four channels, and is used again at the other resolution in Section 4.3.
p_scale, k_scale, w_scale = train_p.max(), train_k.max(), train_w.max()
def make_input(p, kf, xy): # four channels at every grid point: p, k_f, x, y, all of order one
return np.concatenate([(p/p_scale)[..., None], (kf/k_scale)[..., None], xy], axis = -1)
A = torch.tensor(make_input(train_p, train_k, train_xy), dtype = torch.float32).to(device) # (800, 64, 64, 4)
GT = torch.tensor(train_w/w_scale, dtype = torch.float32).reshape(-1, n, n, 1).to(device) # (800, 64, 64, 1)
def get_batch(batch_size = 20):
i = torch.randint(0, len(A), (batch_size,))
return A[i], GT[i]
Train
The FNO2d of Lab 1 with four input channels, one output and a padding of eight. Four thousand steps with batches of twenty, AdamW with a small weight decay and a cosine learning-rate schedule, on the mean squared error of the scaled deflection.
torch.manual_seed(0)
model = FNO2d(width, modes = (12, 12), in_channels = 4, out_channels = 1, padding = n//8).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr = 1e-3, weight_decay = 1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 4000, eta_min = 1e-5)
history = []
for epoch in range(1, 4001):
a, gt = get_batch()
loss = torch.mean((model(a) - gt)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
history.append(loss.item())
if epoch % 500 == 0:
print('epoch %5d | loss %.4e' % (epoch, loss.item()))
model.eval();
plt.figure(figsize = (6, 4))
plt.semilogy(history, 'k', alpha = 0.5)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.title('Training History')
plt.grid(True, alpha = 0.3)
plt.show()
Evaluation
The error over the two hundred unseen pairs of fields, and the mean absolute error of the deflection in millimetres.
A_test = torch.tensor(make_input(test_p, test_k, test_xy), dtype = torch.float32).to(device)
pred_test = predict(model, A_test)[:, :, :, 0]*w_scale # (200, 64, 64), in metres
print('relative L2 error on the test set : %.4f' % relative_error(pred_test, test_w))
print('mean absolute error : %.4f mm' % (1e3*np.mean(np.abs(pred_test - test_w))))
Example Prediction
The median case of the test set, chosen at the median of the error rather than at the best: the two input fields, the reference and predicted deflections on one color scale, and the absolute error.
errors = np.linalg.norm((pred_test - test_w).reshape(len(test_w), -1), axis = 1)/np.linalg.norm(test_w.reshape(len(test_w), -1), axis = 1)
j = int(np.argsort(errors)[len(errors)//2]) # the median case
print('relative L2 error of the example : %.4f' % errors[j])
def show_plate(p, kf, w_ref, w_pred):
fields = [(p/1e3, 'Pressure', 'kPa', 'inferno', None), (kf/1e6, 'Foundation Stiffness', 'MN/m$^3$', 'viridis', None),
(w_ref*1e3, 'Reference Deflection', 'mm', 'magma', (w_ref.min()*1e3, w_ref.max()*1e3)),
(w_pred*1e3, 'FNO Deflection', 'mm', 'magma', (w_ref.min()*1e3, w_ref.max()*1e3)),
(np.abs(w_pred - w_ref)*1e3, 'Absolute Error', 'mm', 'cividis', None)]
plt.figure(figsize = (18, 3.5))
for k, (field, title, unit, cmap, lim) in enumerate(fields):
plt.subplot(1, 5, k + 1)
plt.imshow(field, origin = 'lower', extent = [0, Lx, 0, Ly], cmap = cmap, vmin = None if lim is None else lim[0], vmax = None if lim is None else lim[1])
plt.colorbar(label = unit)
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.title(title)
plt.tight_layout()
plt.show()
show_plate(test_p[j], test_k[j], test_w[j], pred_test[j])
4.3. Resolution Transfer¶
The spectral weights of an FNO are indexed by mode, not by grid point, and the lifting and projection maps act pointwise, so nothing in the trained model refers to the $64 \times 64$ grid it was trained on. The model is therefore evaluated, without retraining, on a second set of cases solved on $64 \times 64$ and on $96 \times 96$ grids. The one thing that has to follow the grid is the padding, which is set to one eighth of the new grid before the finer fields are passed in.
def evaluate(p, kf, xy, w): # the trained model on fields of any resolution
model.padding = p.shape[1]//8 # the padding follows the grid
pred = predict(model, torch.tensor(make_input(p, kf, xy), dtype = torch.float32).to(device))[:, :, :, 0]*w_scale
return pred, np.linalg.norm((pred - w).reshape(len(w), -1), axis = 1)/np.linalg.norm(w.reshape(len(w), -1), axis = 1)
res64 = [data['res64_p'], data['res64_k'], data['res64_xy'], data['res64_w']]
res96 = [data['res96_p'], data['res96_k'], data['res96_xy'], data['res96_w']]
pred_64, err_64 = evaluate(*res64)
pred_96, err_96 = evaluate(*res96)
model.padding = n//8
for size, err in [(64, err_64), (96, err_96)]:
print('%d x %d grid | relative L2 error: mean %.4f | median %.4f' % (size, size, err.mean(), np.median(err)))
The first case of the $96 \times 96$ set, with the same weights and no retraining.
show_plate(res96[0][0], res96[1][0], res96[3][0], pred_96[0])
4.4. What This Lab Showed¶
Two input fields became two channels, and the operator learned the deflection they produce together, for pressures and foundations it had not seen. The error against the finite-difference labels is of the order of a few percent, and, as in Lab 1, it is an error against a solver on a grid, which the operator cannot improve on.
The same weights then answered on a finer grid, at an error of the same order. That is what the architecture of Fourier Neural Operator promised: the parameters live on the modes and on the channels, not on the points, so a finer grid produces more Fourier coefficients and the layer simply ignores the extra ones. The only line that had to change was the padding, which is a length in grid points and so must follow the grid. A finer grid does not make the answer more accurate, since the model has not learned anything new about the plate; it makes the same answer available where the user wants it.
Summary¶
Key takeaways
| Lab 1 | Lab 2 | Lab 3 | |
|---|---|---|---|
| Input | wall shape | ground motion and damper design | pressure and foundation stiffness |
| Channels | gap, $x$, $y$, mask | $\ddot y_g$, $m_d$, $k_d$, $c_d$, DOF, $t$ | $p$, $k_f$, $x$, $y$ |
| Grid | $64 \times 64$ in space | $6$ DOF $\times$ $512$ instants | $64 \times 64$ in space |
| Modes | $12 \times 12$ | $3 \times 16$ | $12 \times 12$ |
| Output | velocity $(u, v)$ | response $Z$ | deflection $w$ |
| Labels | Laplace solve | Newmark integration | finite differences |
- The input channels are where the engineering goes. A shape became a signed distance and a mask, a design became three constant channels, two fields became two channels, and the coordinates were appended every time. After the lifting layer the three labs ran the same code.
- An FNO does not care what its axes mean. Degrees of freedom against time was read like $x$ against $y$, and the mode count per axis was chosen from how the solution varies along it.
- A data-driven operator is a surrogate for its solver. Every label here carried a discretization error, and the operator inherits it as a floor; the reported error is an error against the solver, not against the physics. Nothing in the loss enforces a governing equation.
- A surrogate is for searching; the solver is for deciding. The damper design was found by eighty forward passes and confirmed by one time integration, and it is the solver's figure that is the result.
- The same weights serve another grid. The parameters live on the modes and the channels, so the plate model trained at $64 \times 64$ answered at $96 \times 96$ with no retraining; only the padding, a length in grid points, had to follow the grid.