Scientific Machine Learning:
Embedding Prior Knowledge into Neural Networks: Labs
These laboratory problems accompany the discussion of embedding prior knowledge into neural networks. Three compact examples are treated, each one an application in which the required behavior is known before any data is seen.
| Lab | Application | Prior knowledge | Embedding route |
|---|---|---|---|
| 1 | Drug concentration | Positive, non-increasing, known initial value | Loss penalty + hard parameterization |
| 2 | Epidemic compartments | Positive fractions and total population = 1 | Loss penalty + projection + Softmax |
| 3 | Pairwise interaction | Translation, rotation, and permutation symmetry | Augmentation + pooling + invariant feature |
All three follow the same workflow: state the prior knowledge, train alternative models that embed it in different places, apply the required transformations, and verify the behavior directly from the outputs.
Learning Objectives¶
After studying this notebook, you should be able to
- name the places in a neural-network pipeline where prior knowledge can enter — the data, the input features, the loss, the output transformation, a post-processing step, and differentiation of the output — and say what each one guarantees;
- impose a requirement softly by adding its violation to the loss, and state why such a penalty discourages a violation without preventing it;
- impose the same requirement exactly by construction, writing a positive, monotone decreasing curve as an exponential of an integrated positive rate, and nonnegative fractions summing to one as a Softmax output;
- correct a linear equality after prediction by projection, and explain why projection fixes the sum but not the sign;
- compare a hard construction against a soft penalty on the same data and read the resulting trade-off between guaranteed behavior and data error;
- encode a symmetry by data augmentation, by averaging predictions over a finite transformation group, and by replacing the coordinates with an invariant feature, and distinguish invariance from predictive accuracy;
- obtain an equivariant vector field by differentiating an invariant scalar, and verify the relation numerically.
Question: can a neural network fit noisy measurements while preserving physically valid concentration behavior?
| Item | Definition |
|---|---|
| Input | Time $t$ |
| Output | Drug concentration $C(t)$ |
| Prior knowledge | $C(t)\ge0$, $dC/dt\le0$, and $C(0)=C_0$ |
| Comparison | Baseline vs. soft penalty vs. hard parameterization |
After a single dose, concentration should remain nonnegative, decrease over time, and begin from the known initial value:
$$ C(t)\ge0, \qquad \frac{dC}{dt}\le0, \qquad C(0)=C_0. $$
Setup
The cell below loads the libraries and fixes the random seeds. Run it once before any of the sections that follow. Everything else is written out where it is used: each section builds its own network with nn.Sequential, writes its own training loop, and draws its own figures, so that no part of a laboratory problem is hidden in a helper defined elsewhere.
In [ ]:
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
np.random.seed(7)
torch.manual_seed(7)
Out[ ]:
In [ ]:
t = np.linspace(0, 8, 121).reshape(-1, 1) # time grid
C0 = 4.0 # initial concentration
C_true = C0*np.exp(-0.55*t) # reference curve
# 12 noisy measurements taken at these grid points
idx = np.array([0, 4, 9, 16, 24, 33, 45, 58, 73, 90, 106, 120])
t_obs = t[idx]
C_obs = C_true[idx] + 0.10*np.random.randn(len(idx), 1)
C_obs[5] = C_obs[5] + 0.35 # one outlier
C_obs[-1] = -0.05 # one negative measurement
plt.figure(figsize = (6, 4))
plt.plot(t, C_true, 'k--', alpha = 0.3)
plt.plot(t_obs, C_obs, 'ko', markersize = 4)
plt.xlim([0, 8])
plt.xlabel('Time', fontsize = 12)
plt.ylabel('Drug concentration', fontsize = 12)
plt.legend(['True', 'Measured'], fontsize = 12)
plt.grid(True)
plt.show()
In [ ]:
# the network input is the time rescaled to [-1, 1]
x = (t - 4)/4
x_all = torch.tensor(x, dtype = torch.float32)
x_obs = torch.tensor(x[idx], dtype = torch.float32)
y_obs = torch.tensor(C_obs, dtype = torch.float32)
The baseline minimizes data error only. The soft model additionally penalizes negative concentration, positive slope, and initial-condition error:
$$ \mathcal{L}_{\mathrm{soft}} =\mathcal{L}_{\mathrm{data}} +\lambda_{+}\operatorname{ReLU}(-C)^2 +\lambda_m\operatorname{ReLU}\!\left(\frac{dC}{dt}\right)^2 +\lambda_0\left(C(0)-C_0\right)^2. $$
Soft enforcement discourages violations but does not make them impossible.
In [ ]:
# baseline: data error only
base = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1))
optm = torch.optim.Adam(base.parameters(), lr = 0.002)
for epoch in range(1800):
loss = torch.mean((base(x_obs) - y_obs)**2)
optm.zero_grad()
loss.backward()
optm.step()
print('baseline data loss = %.3e' % loss.item())
C_base = base(x_all).detach().numpy()
In [ ]:
# soft: data error + three penalties
soft = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1))
optm = torch.optim.Adam(soft.parameters(), lr = 0.002)
x_check = x_all.clone().requires_grad_(True) # where the constraints are checked
x_zero = torch.tensor([[-1.0]]) # input that corresponds to t = 0
for epoch in range(1800):
C = soft(x_check)
dC = torch.autograd.grad(C, x_check, torch.ones_like(C), create_graph = True)[0]
data_loss = torch.mean((soft(x_obs) - y_obs)**2)
positive_loss = torch.mean(torch.relu(-C)**2) # active when C < 0
decrease_loss = torch.mean(torch.relu(dC)**2) # active when C increases
initial_loss = torch.mean((soft(x_zero) - C0)**2) # active when C(0) != C0
loss = data_loss + 20*positive_loss + 20*decrease_loss + 10*initial_loss
optm.zero_grad()
loss.backward()
optm.step()
print('soft total loss = %.3e' % loss.item())
C_soft = soft(x_all).detach().numpy()
The network predicts a positive elimination rate and integrates it into the concentration:
$$ r_\theta(t)=\operatorname{Softplus}(z_\theta(t))>0, $$
$$ C(t)=C_0\exp\left[-\int_0^t r_\theta(\tau)\,d\tau\right]. $$
This form guarantees positivity, monotone decrease, and the initial value. A small rate-smoothness penalty prevents unrealistically sharp decay, producing a hybrid constraint: hard validity with soft regularization.
In [ ]:
# hard: the network predicts a positive rate, the curve is built from it
hard = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1))
optm = torch.optim.Adam(hard.parameters(), lr = 0.002)
t_torch = torch.tensor(t, dtype = torch.float32)
dt = t_torch[1:] - t_torch[:-1]
zero = torch.zeros(1, 1)
for epoch in range(1800):
rate = torch.log(1 + torch.exp(hard(x_all))) # softplus keeps the rate positive
area = 0.5*(rate[1:] + rate[:-1])*dt # trapezoid area of each interval
integral = torch.cat([zero, torch.cumsum(area, 0)]) # integral of the rate up to each time
C = C0*torch.exp(-integral) # positive and decreasing by construction
data_loss = torch.mean((C[idx] - y_obs)**2)
smooth_loss = torch.mean((rate[1:] - rate[:-1])**2) # keeps the rate from changing too fast
loss = data_loss + 100*smooth_loss
optm.zero_grad()
loss.backward()
optm.step()
print('hard total loss = %.3e' % loss.item())
C_hard = C.detach().numpy()
In [ ]:
plt.figure(figsize = (11, 4))
plt.subplot(1, 2, 1)
plt.plot(t, C_true, 'k--')
plt.plot(t_obs, C_obs, 'ko')
plt.plot(t, C_base, 'b')
plt.plot(t, C_soft, 'r')
plt.plot(t, C_hard, 'g')
plt.xlabel('Time', fontsize = 12)
plt.ylabel('Drug concentration', fontsize = 12)
plt.legend(['True', 'Measured', 'Baseline', 'Soft', 'Hard'], fontsize = 10)
plt.grid(True)
plt.subplot(1, 2, 2) # zoom on the tail
plt.plot(t, C_true, 'k--')
plt.plot(t, C_base, 'b')
plt.plot(t, C_soft, 'r')
plt.plot(t, C_hard, 'g')
plt.plot(t, 0*t, 'k')
plt.xlim(5.5, 8)
plt.ylim(-0.2, 0.35)
plt.xlabel('Time', fontsize = 12)
plt.ylabel('Drug concentration', fontsize = 12)
plt.title('Tail: below the black line the curve is negative', fontsize = 12)
plt.grid(True)
plt.show()
1.4.1. Numerical Check¶
In [ ]:
print('model negative increasing IC error RMSE')
for name, C in [('baseline', C_base), ('soft', C_soft), ('hard', C_hard)]:
negative = np.sum(C < 0) # how many points fall below zero
increasing = np.sum(np.diff(C[:, 0]) > 1e-4) # how many steps go up
ic_error = abs(C[0, 0] - C0) # error at t = 0
rmse = np.sqrt(np.mean((C - C_true)**2)) # error against the true curve
print('%-10s %8d %11d %9.3f %5.3f' % (name, negative, increasing, ic_error, rmse))
1.4.2. Interpretation¶
- Baseline: follows the observations without any validity guarantee.
- Soft: reduces violations, but a small negative prediction can remain.
- Hard: stays nonnegative, decreases monotonically, and satisfies $C(0)=C_0$ by construction.
Key point: hard constraints guarantee the prescribed behavior, not the smallest data error. If the prior knowledge is incorrect or too restrictive, prediction accuracy can decrease.
Question: how can a neural network keep population fractions valid at every time?
| Item | Definition |
|---|---|
| Input | Time $t$ |
| Output | Susceptible $S$, infected $I$, recovered $R$ |
| Prior knowledge | Nonnegative fractions and $S+I+R=1$ |
| Comparison | Soft penalty vs. projection vs. Softmax |
Every person belongs to one compartment, so the constraints must hold locally at every time:
$$ S(t),I(t),R(t)\ge0, \qquad S(t)+I(t)+R(t)=1. $$
In [ ]:
beta = 0.9 # infection rate
gamma = 0.25 # recovery rate
dt = 0.01
n = 1801
S = np.zeros(n)
I = np.zeros(n)
R = np.zeros(n)
S[0] = 0.98
I[0] = 0.02
R[0] = 0.00
for k in range(n - 1): # explicit Euler steps
infection = beta*S[k]*I[k]
recovery = gamma*I[k]
S[k+1] = S[k] - dt*infection
I[k+1] = I[k] + dt*(infection - recovery)
R[k+1] = R[k] + dt*recovery
sub = np.arange(0, n, 12) # keep every 12th point
t2 = dt*sub
Y = np.column_stack([S[sub], I[sub], R[sub]])
obs = np.arange(0, 151, 9) # 17 noisy measurements
Y_obs = Y[obs] + 0.012*np.random.randn(len(obs), 3)
plt.figure(figsize = (6, 4))
plt.plot(t2, Y[:, 0], 'b', t2, Y[:, 1], 'r', t2, Y[:, 2], 'g', alpha = 0.3)
plt.plot(t2[obs], Y_obs[:, 0], 'bo', t2[obs], Y_obs[:, 1], 'ro', t2[obs], Y_obs[:, 2], 'go', markersize = 4)
plt.title('SIR Model', fontsize = 12)
plt.xlabel('Time', fontsize = 12)
plt.ylabel('Population fraction', fontsize = 12)
plt.legend(['S', 'I', 'R'], fontsize = 12)
plt.grid(True)
plt.show()
In [ ]:
x2 = torch.tensor((t2.reshape(-1, 1) - 9)/9, dtype = torch.float32)
x2_obs = x2[obs]
y2_obs = torch.tensor(Y_obs, dtype = torch.float32)
| Method | Implementation | Guaranteed behavior |
|---|---|---|
| Soft | Add violations to the loss | Approximate positivity and conservation |
| Projection | Correct the raw output after prediction | Exact sum only |
| Hard | Apply Softmax to network logits | Exact positivity and sum |
2.2.1. Soft: Loss Penalty¶
$$ \mathcal{L}=\mathcal{L}_{data} +\lambda_c(\hat S+\hat I+\hat R-1)^2 +\lambda_p\sum_{k\in\{S,I,R\}}\operatorname{ReLU}(-\hat y_k)^2. $$
2.2.2. Projection: Output Correction¶
$$ \hat{\mathbf y}_{proj} =\hat{\mathbf y} +\frac{1-(\hat S+\hat I+\hat R)}{3}[1,1,1]. $$
Projection requires no additional training and guarantees $\hat S_{proj}+\hat I_{proj}+\hat R_{proj}=1$, but negative fractions may remain.
2.2.3. Hard: Softmax Output¶
$$ \hat y_k=\frac{\exp(z_k)}{\exp(z_S)+\exp(z_I)+\exp(z_R)}. $$
Softmax guarantees $\hat S,\hat I,\hat R>0$ and $\hat S+\hat I+\hat R=1$ simultaneously.
In [ ]:
# baseline: data error only
base2 = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 3))
optm = torch.optim.Adam(base2.parameters(), lr = 0.002)
for epoch in range(1800):
loss = torch.mean((base2(x2_obs) - y2_obs)**2)
optm.zero_grad()
loss.backward()
optm.step()
print('baseline loss = %.3e' % loss.item())
In [ ]:
# soft: data error + conservation penalty + positivity penalty
soft2 = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 3))
optm = torch.optim.Adam(soft2.parameters(), lr = 0.002)
for epoch in range(1800):
P = soft2(x2)
data_loss = torch.mean((soft2(x2_obs) - y2_obs)**2)
sum_loss = torch.mean((torch.sum(P, 1) - 1)**2) # active when S + I + R != 1
positive_loss = torch.mean(torch.relu(-P)**2) # active when a fraction is negative
loss = data_loss + 10*sum_loss + 10*positive_loss
optm.zero_grad()
loss.backward()
optm.step()
print('soft loss = %.3e' % loss.item())
In [ ]:
# hard: the three outputs pass through a softmax
hard2 = nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 3))
optm = torch.optim.Adam(hard2.parameters(), lr = 0.002)
for epoch in range(1800):
P = torch.softmax(hard2(x2_obs), 1)
loss = torch.mean((P - y2_obs)**2)
optm.zero_grad()
loss.backward()
optm.step()
print('hard loss = %.3e' % loss.item())
In [ ]:
with torch.no_grad():
P_base = base2(x2).numpy()
P_soft = soft2(x2).numpy()
P_hard = torch.softmax(hard2(x2), 1).numpy()
# projection: shift the three raw outputs so that they sum to one
shift = (1 - np.sum(P_base, 1))/3
P_proj = P_base + shift.reshape(-1, 1)
print('sum before projection at t = 0: %.4f' % np.sum(P_base[0]))
print('sum after projection at t = 0: %.4f' % np.sum(P_proj[0]))
In [ ]:
names = ['Soft', 'Projection', 'Hard']
predictions = [P_soft, P_proj, P_hard]
plt.figure(figsize = (6, 10))
for k in range(3):
plt.subplot(3, 1, k+1)
plt.plot(t2, Y[:, 0], 'b', t2, Y[:, 1], 'r', t2, Y[:, 2], 'g', alpha = 0.3)
plt.plot(t2, predictions[k][:, 0], 'b--')
plt.plot(t2, predictions[k][:, 1], 'r--')
plt.plot(t2, predictions[k][:, 2], 'g--')
plt.title(names[k], fontsize = 12)
plt.ylabel('Population fraction', fontsize = 12)
plt.grid(True)
plt.show()
In [ ]:
plt.figure(figsize = (6, 4))
plt.plot(t2, np.sum(P_base, 1), 'b')
plt.plot(t2, np.sum(P_soft, 1), 'r')
plt.plot(t2, np.sum(P_hard, 1), 'g')
plt.plot(t2, np.ones(len(t2)), 'k--', alpha = 0.3)
plt.xlabel('Time', fontsize = 12)
plt.ylabel('S + I + R', fontsize = 12)
plt.legend(['Baseline', 'Soft', 'Hard', 'Target'], fontsize = 12)
plt.grid(True)
plt.show()
2.5.1. Numerical Check¶
In [ ]:
print('model max |sum - 1| negative values')
for name, P in [('baseline', P_base), ('soft', P_soft), ('projection', P_proj), ('hard', P_hard)]:
sum_error = np.max(np.abs(np.sum(P, 1) - 1))
negative = np.sum(P < 0)
print('%-11s %13.2e %14d' % (name, sum_error, negative))
2.5.2. Interpretation¶
- Soft: remains close to the conservation target, but small sum or positivity violations can remain.
- Projection: corrects the sum exactly after prediction, but does not guarantee nonnegative fractions.
- Hard: Softmax guarantees both positivity and unit sum at every time.
Key point: projection is effective for a linear equality, while a suitable output transformation can satisfy multiple compatible constraints at once.
Question: should the predicted interaction change when the same particle pair is moved, rotated, or reordered?
| Item | Definition |
|---|---|
| Input | Two particle coordinates $X=[\mathbf{x}_1,\mathbf{x}_2]$ |
| Output | Interaction energy $U$ |
| Prior knowledge | Translation, rotation, and permutation invariance |
| Derived quantity | Equivariant force $\mathbf F=-\nabla_XU$ |
For this example, energy depends only on pairwise distance:
$$ U(TX)=U(X), \qquad r=\left\|\mathbf{x}_2-\mathbf{x}_1\right\|_2, \qquad U=\frac{1}{r}. $$
In [ ]:
n = 90
r = np.random.uniform(0.7, 2.0, n) # distance between the two particles
a = np.random.uniform(-0.25, 0.25, n) # orientation: a narrow range only
cx = np.random.uniform(-0.15, 0.15, n) # center of the pair
cy = np.random.uniform(-0.15, 0.15, n)
dx = r*np.cos(a)
dy = r*np.sin(a)
X = np.zeros((n, 4)) # [x1, y1, x2, y2]
X[:, 0] = cx - dx/2
X[:, 1] = cy - dy/2
X[:, 2] = cx + dx/2
X[:, 3] = cy + dy/2
U = (1/r).reshape(-1, 1) # interaction energy
plt.figure(figsize = (6, 4))
plt.scatter(X[:, 2] - X[:, 0], X[:, 3] - X[:, 1], c = U[:, 0])
plt.colorbar()
plt.axis('equal')
plt.xlabel('Relative x', fontsize = 12)
plt.ylabel('Relative y', fontsize = 12)
plt.title('Training pairs: limited orientations', fontsize = 12)
plt.grid(True)
plt.show()
| Strategy | Where symmetry enters | Expected behavior |
|---|---|---|
| Raw coordinates | Nowhere | Sensitive to representation |
| Data augmentation | Training data | Approximate invariance |
| Augmented + pooling | Output aggregation | More consistent predictions |
| Invariant feature | Model input | Invariance by construction |
Pooling averages predictions over a finite transformation set $G$:
$$ \bar U(X)=\frac{1}{|G|}\sum_{g\in G}f_\theta(gX). $$
It is an inference-time operation, not a separately trained model. By contrast, the invariant-feature model receives only
$$ r=\left\|\mathbf{x}_2-\mathbf{x}_1\right\|_2, $$
which is unchanged by translation, rotation, or particle swapping.
In [ ]:
# augmented data: each pair is rotated, translated, and sometimes swapped
X_aug = np.zeros((6*n, 4))
U_aug = np.zeros((6*n, 1))
for m in range(6):
a = np.random.uniform(-np.pi, np.pi)
sx = np.random.uniform(-2, 2)
sy = np.random.uniform(-2, 2)
x1 = np.cos(a)*X[:, 0] - np.sin(a)*X[:, 1] + sx
y1 = np.sin(a)*X[:, 0] + np.cos(a)*X[:, 1] + sy
x2 = np.cos(a)*X[:, 2] - np.sin(a)*X[:, 3] + sx
y2 = np.sin(a)*X[:, 2] + np.cos(a)*X[:, 3] + sy
rows = np.arange(m*n, (m+1)*n)
if m % 2 == 0:
X_aug[rows] = np.column_stack([x1, y1, x2, y2])
else:
X_aug[rows] = np.column_stack([x2, y2, x1, y1]) # the two particles are swapped
U_aug[rows] = U
r_train = np.sqrt((X[:, 2] - X[:, 0])**2 + (X[:, 3] - X[:, 1])**2).reshape(-1, 1)
In [ ]:
def train(net, inputs, targets):
xx = torch.tensor(inputs, dtype = torch.float32)
yy = torch.tensor(targets, dtype = torch.float32)
optm = torch.optim.Adam(net.parameters(), lr = 0.002)
for epoch in range(1800):
loss = torch.mean((net(xx) - yy)**2)
optm.zero_grad()
loss.backward()
optm.step()
return loss.item()
net_raw = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
net_aug = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
net_inv = nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
print('raw loss = %.3e' % train(net_raw, X, U))
print('augmented loss = %.3e' % train(net_aug, X_aug, U_aug))
print('invariant loss = %.3e' % train(net_inv, r_train, U))
3.3.1. Transformation Pooling at Inference¶
In [ ]:
def move_pair(Xt, a, swap):
p = np.zeros(Xt.shape)
cx = (Xt[:, 0] + Xt[:, 2])/2 # center the pair first
cy = (Xt[:, 1] + Xt[:, 3])/2
x1 = Xt[:, 0] - cx
y1 = Xt[:, 1] - cy
x2 = Xt[:, 2] - cx
y2 = Xt[:, 3] - cy
p[:, 0] = np.cos(a)*x1 - np.sin(a)*y1
p[:, 1] = np.sin(a)*x1 + np.cos(a)*y1
p[:, 2] = np.cos(a)*x2 - np.sin(a)*y2
p[:, 3] = np.sin(a)*x2 + np.cos(a)*y2
if swap:
p = p[:, [2, 3, 0, 1]]
return p
def pool_predict(net, Xt): # average over 4 rotations x 2 orderings
total = np.zeros((len(Xt), 1))
for a in [0, np.pi/2, np.pi, 3*np.pi/2]:
for swap in [False, True]:
p = move_pair(Xt, a, swap)
with torch.no_grad():
total = total + net(torch.tensor(p, dtype = torch.float32)).numpy()
return total/8
In [ ]:
def predict_all(Xt):
r = np.sqrt((Xt[:, 2] - Xt[:, 0])**2 + (Xt[:, 3] - Xt[:, 1])**2).reshape(-1, 1)
with torch.no_grad():
p_raw = net_raw(torch.tensor(Xt, dtype = torch.float32)).numpy()
p_aug = net_aug(torch.tensor(Xt, dtype = torch.float32)).numpy()
p_inv = net_inv(torch.tensor(r, dtype = torch.float32)).numpy()
p_pool = pool_predict(net_aug, Xt)
return p_raw[:, 0], p_aug[:, 0], p_pool[:, 0], p_inv[:, 0]
pair = np.array([-0.6, 0.0, 0.6, 0.0]) # the reference pair, distance 1.2
U_ref = 1/1.2
print('reference energy = %.4f' % U_ref)
In [ ]:
# rotation test: the same pair turned through a full circle
angles = np.linspace(0, 2*np.pi, 121)
X_rot = np.zeros((121, 4))
for k in range(121):
X_rot[k] = move_pair(pair.reshape(1, 4), angles[k], False)
rot_raw, rot_aug, rot_pool, rot_inv = predict_all(X_rot)
# translation test: the same pair moved along a line
shifts = np.linspace(-2, 2, 121)
X_tr = np.zeros((121, 4))
for k in range(121):
X_tr[k] = pair + np.array([shifts[k], 0.5*shifts[k], shifts[k], 0.5*shifts[k]])
tr_raw, tr_aug, tr_pool, tr_inv = predict_all(X_tr)
# particle-order test: the same pair with the two particles exchanged
X_sw = np.zeros((2, 4))
X_sw[0] = pair
X_sw[1] = pair[[2, 3, 0, 1]]
sw_raw, sw_aug, sw_pool, sw_inv = predict_all(X_sw)
In [ ]:
plt.figure(figsize = (6, 8))
plt.subplot(2, 1, 1)
plt.plot(angles*180/np.pi, rot_raw, 'k')
plt.plot(angles*180/np.pi, rot_aug, 'b')
plt.plot(angles*180/np.pi, rot_pool, 'r')
plt.plot(angles*180/np.pi, rot_inv, 'g')
plt.plot(angles*180/np.pi, U_ref*np.ones(121), 'k--')
plt.xlabel('Rotation angle (degree)', fontsize = 12)
plt.ylabel('Predicted energy', fontsize = 12)
plt.legend(['Raw', 'Augmented', 'Augmented + pooling', 'Invariant feature', 'True'], fontsize = 9)
plt.grid(True)
plt.subplot(2, 1, 2)
plt.plot(shifts, tr_raw, 'k')
plt.plot(shifts, tr_aug, 'b')
plt.plot(shifts, tr_pool, 'r')
plt.plot(shifts, tr_inv, 'g')
plt.plot(shifts, U_ref*np.ones(121), 'k--', alpha = 0.3)
plt.xlabel('Translation', fontsize = 12)
plt.ylabel('Predicted energy', fontsize = 12)
plt.grid(True)
plt.show()
print('method original swapped')
print('raw %9.4f %9.4f' % (sw_raw[0], sw_raw[1]))
print('augmented %9.4f %9.4f' % (sw_aug[0], sw_aug[1]))
print('augmented + pooling %9.4f %9.4f' % (sw_pool[0], sw_pool[1]))
print('invariant feature %9.4f %9.4f' % (sw_inv[0], sw_inv[1]))
3.4.1. Interpretation¶
The reference pair has $r=1.2$, so every equivalent configuration should produce $U=1/1.2\approx0.833$.
| Method | Rotation | Translation | Particle order | What the plots show |
|---|---|---|---|---|
| Raw | Sensitive | Sensitive | Sensitive | Coordinate representation changes the prediction |
| Augmented | Approximate | Approximate | Approximate | Additional data reduces but does not eliminate sensitivity |
| Augmented + pooling | Discrete approximation | Exact after centering | Exact after averaging | Consistent output can still be inaccurate |
| Invariant feature | Exact | Exact | Exact | The green prediction overlaps the reference |
Key point: invariance and predictive accuracy are different properties. Pooling can stabilize a biased prediction, while an appropriate invariant feature removes irrelevant variation before learning.
Force is obtained from the gradient of the learned scalar energy:
$$ \mathbf F_2=-\nabla_{\mathbf{x}_2}U. $$
If $U$ depends only on invariant distance, rotating the input preserves energy and rotates the force in the same way:
$$ U(RX)=U(X), \qquad \mathbf F(RX)=R\mathbf F(X). $$
The following cell checks this relation numerically. It verifies the equivariance structure, not the physical accuracy of the force magnitude.
In [ ]:
def energy_and_force(net, p):
q = torch.tensor(p.reshape(1, 4), dtype = torch.float32, requires_grad = True)
r = torch.sqrt((q[:, 2] - q[:, 0])**2 + (q[:, 3] - q[:, 1])**2).reshape(1, 1)
energy = net(r)
gradient = torch.autograd.grad(energy, q)[0]
force = -gradient[0, 2:].numpy() # force acting on particle 2
return energy.item(), force
a = 55*np.pi/180 # rotate the pair by 55 degrees
pair_rot = np.zeros(4)
pair_rot[0] = np.cos(a)*pair[0] - np.sin(a)*pair[1]
pair_rot[1] = np.sin(a)*pair[0] + np.cos(a)*pair[1]
pair_rot[2] = np.cos(a)*pair[2] - np.sin(a)*pair[3]
pair_rot[3] = np.sin(a)*pair[2] + np.cos(a)*pair[3]
U1, F1 = energy_and_force(net_inv, pair)
U2, F2 = energy_and_force(net_inv, pair_rot)
# the force of the original pair, rotated by the same angle
F1_rot = np.array([np.cos(a)*F1[0] - np.sin(a)*F1[1],
np.sin(a)*F1[0] + np.cos(a)*F1[1]])
print('energy before rotation = %.6f' % U1)
print('energy after rotation = %.6f' % U2)
print('error between the rotated force and the predicted force = %.2e' % np.linalg.norm(F2 - F1_rot))
plt.figure(figsize = (10, 5))
plt.subplot(1, 2, 1)
plt.plot([pair[0], pair[2]], [pair[1], pair[3]], 'ko')
plt.quiver(pair[2], pair[3], F1[0], F1[1], angles = 'xy', scale_units = 'xy', scale = 1, color = 'r')
plt.gca().set_aspect('equal')
plt.xlim(-1.5, 1.5)
plt.ylim(-1.5, 1.5)
plt.title('Original', fontsize = 12)
plt.xlabel('x', fontsize = 12)
plt.ylabel('y', fontsize = 12)
plt.grid(True)
plt.subplot(1, 2, 2)
plt.plot([pair_rot[0], pair_rot[2]], [pair_rot[1], pair_rot[3]], 'ko')
plt.quiver(pair_rot[2], pair_rot[3], F2[0], F2[1], angles = 'xy', scale_units = 'xy', scale = 1, color = 'r')
plt.gca().set_aspect('equal')
plt.xlim(-1.5, 1.5)
plt.ylim(-1.5, 1.5)
plt.title('Rotated', fontsize = 12)
plt.xlabel('x', fontsize = 12)
plt.ylabel('y', fontsize = 12)
plt.grid(True)
plt.show()
3.5.1. Interpretation¶
- The energy is unchanged before and after rotation.
- The force arrow rotates with the particle pair.
- The small rotated-force error confirms $\mathbf F(RX)\approx R\mathbf F(X)$.
Key point: an invariant scalar potential can generate an equivariant vector field through automatic differentiation.
Key takeaways
- Prior knowledge can enter a network at six places, the data, the input features, the loss, the output transformation, the post-processing and the differentiation, and the three labs used all of them.
- A loss penalty discourages a violation but does not forbid it; an output transformation or a projection makes the constraint exact by construction, and the two can be combined when a quantity must satisfy several conditions at once.
- Symmetry is cheapest to enforce early: a distance feature removes the irrelevant coordinates before the network sees them, whereas augmentation only teaches the symmetry approximately and pooling has to correct for it afterwards.
Where was prior knowledge embedded?
| Pipeline location | Example in the labs | Effect |
|---|---|---|
| Data | Rotated, translated, and swapped particle pairs | Symmetry learned approximately |
| Input feature | Pairwise distance | Irrelevant coordinates removed |
| Loss | Positivity, monotonicity, IC, conservation, smoothness | Violations discouraged |
| Output transformation | Cumulative exponential and Softmax | Constraints exact by construction |
| Post-processing | Projection and transformation pooling | Output corrected or symmetrized |
| Differentiation | Force from invariant energy | Equivariant vector output |
Design rule: enforce what must hold, penalize what should hold, and encode symmetry as early as practical.
Looking ahead. Every constraint in these labs was a statement about the output, a sign, a monotone trend, a sum, a symmetry. The next notebook, Physics-informed Neural Networks (PINN), adds the strongest prior of all, the governing differential equation itself, and puts it into the loss through the same mechanism used here for the penalties, with the output transformations of Lab 1 reappearing as the hard constraints of a PINN.