pytorch实现:PINN 模拟抛射运动(基于时间)
1. 抛射运动(基于时间)
将问题置于抛射运动场景,在处于自由落体运动中的抛射体,它会受到重力和阻力的共同作用。对于这个问题,应该掌握一些基本的物理动力学知识,这些将为我们进一步探索 PINN 在模拟此类物理现象中的应用提供有力支撑。
抛射体在时间
t
t
t 的位置矢量(即其位置)有一个函数来定义:
s
⃗
=
(
s
x
,
s
y
)
=
f
(
t
)
\vec{s}=(s_x, s_y)=f(t)
s=(sx,sy)=f(t)
该位置矢量的一阶导数则表示速度矢量,具体定义为:
v
⃗
=
(
v
x
,
v
y
)
=
d
s
⃗
d
t
\vec{v}=(v_x, v_y)=\frac{d\vec{s}}{dt}
v=(vx,vy)=dtds
二阶导数表示加速度矢量,定义为:
a
⃗
=
(
a
x
,
a
y
)
=
d
2
s
⃗
d
t
2
\vec{a}=(a_x, a_y)=\frac{d^2\vec{s}}{dt^2}
a=(ax,ay)=dt2d2s
在阻尼介质中,如空气,抛射体同时受到重力和空气阻力的作用,抛射体的运动方程由牛顿第二定律给出,具体形式如下:
m
d
2
s
⃗
d
t
2
=
−
m
g
⃗
−
c
∥
d
s
⃗
d
t
∥
d
s
⃗
d
t
m\frac{d^2\vec{s}}{dt^2}=-m\vec{g}-c\|\frac{d\vec{s}}{dt}\|\frac{d\vec{s}}{dt}
mdt2d2s=−mg−c∥dtds∥dtds
其中, s ⃗ \vec{s} s 是位置矢量, g ⃗ \vec{g} g 是重力加速度矢量, m m m 为抛射体质量, c c c 为空气阻力系数。
要解决这个方程,通常要使用数值方法,例如欧拉方法或者龙格-库塔方法,因为解析解很难获得。在实际应用中,如炮弹或者导弹的轨迹模拟,考虑空气阻力是非常重要的,因为它影响射程和命中精度。
给定抛射体和环境的相关参数,其中抛射体质量
m
=
1.0
k
g
m=1.0 kg
m=1.0kg,重力加速度
g
=
9.81
m
/
s
2
g=9.81m/s^2
g=9.81m/s2,空气阻力系数
c
=
0.1
c = 0.1
c=0.1,初始速度
v
⃗
=
(
v
x
,
v
y
)
=
100
×
(
cos
π
4
,
sin
π
4
)
\vec{v}=(v_x, v_y)=100\times(\cos\frac{\pi}{4},\sin\frac{\pi}{4})
v=(vx,vy)=100×(cos4π,sin4π),即初始速度为
100
m
/
s
100m/s
100m/s,方向与
x
x
x 轴夹角
45
°
45\degree
45°。采用四阶龙格-库塔方法来模拟上述过程:
2. 理想和现实的差异
2.1. 理想完美情况
理想情况下,假设能获取整个区域内的理想数据,因此使用简单的监督学习训练神经网络便能快速学习到一个准确的轨迹模型,如下图所示:
2.2. 实际不完美数据
然而实际上,能够覆盖整个域的理想数据是十分有限的,并且我们能够获得的数据也是掺杂噪声的稀疏数据。在这种情况下,采用简单的监督学习方法会给我们带来不理想甚至完全错误的模型。
2.2.1. 过拟合
使用带有噪声、稀疏且不完整的数据训练神经网络时,网络会严格按照我们的要求工作,并基于我们提供的数据拟合一个函数。然而,这个函数对于其预期的用途(预测抛射体的位移)来说,并不是那么有用,换句话说,它过拟合了,如下图所示:
2.2.2 正则化
为防止网络过拟合,一种典型有效的方法是使用正则化,公式如下:
L
2
=
∑
i
=
1
n
w
i
2
L_{2} = \sum_{i=1}^n{w_i^2}
L2=i=1∑nwi2
L2 正则化将对网络中过大权重进行惩罚,从而防止模型对局部梯度的过渡依赖和特化,增强模型的泛化能力。如下图是经过 L2 正则化后的训练结果,通过 L2 正则化在带有噪声、稀疏且不完整的数据上训练神经网络,该模型在数据存在的区域更为有用,但无法在整个域内进行外推。
3. 物理信息神经网络
与 L2 正则化相似,PINN 不仅仅最小化数据损失,并且还使用额外的物理规律作为额外的正则化项,通俗来讲,不仅要拟合数据,而且要确保解决方案和给定的物理规律一致。
网络模型的目标是通过监督学习来学习位移矢量函数
f
(
t
)
f(t)
f(t),因此网络包含标准的数据损失,即回归目标和预测值之间的均方误差(MSE),如下:
L
d
a
t
a
=
1
n
∑
i
=
1
n
(
s
i
⃗
−
s
i
⃗
^
)
2
L_{data}=\frac{1}{n}\sum_{i=1}^n{(\vec{s_i}-\hat{\vec{s_i}})^2}
Ldata=n1i=1∑n(si−si^)2
前文提到的运动方程,
m
d
2
s
⃗
d
t
2
=
−
m
g
⃗
−
c
∥
d
s
⃗
d
t
∥
d
s
⃗
d
t
m\frac{d^2\vec{s}}{dt^2}=-m\vec{g}-c\|\frac{d\vec{s}}{dt}\|\frac{d\vec{s}}{dt}
mdt2d2s=−mg−c∥dtds∥dtds
可以变形为下式,将
m
、
c
m、c
m、c 两个环境变量合成一个超参数,让网络自发的去优化寻找最佳值:
d
2
s
⃗
d
t
2
=
−
g
⃗
−
c
m
∥
d
s
⃗
d
t
∥
d
s
⃗
d
t
\frac{d^2\vec{s}}{dt^2}=-\vec{g}-c_m\|\frac{d\vec{s}}{dt}\|\frac{d\vec{s}}{dt}
dt2d2s=−g−cm∥dtds∥dtds
因为我们要确保网络遵循上述物理方程,即
f
(
t
)
f(t)
f(t) 关于时间的二阶导数(加速度)应等于一阶导数(速度)的某个函数,换言之,我们希望两者的误差尽可能接近0,如下数学表达式:
−
g
⃗
−
c
m
∥
d
s
⃗
d
t
∥
d
s
⃗
d
t
−
d
2
s
⃗
d
t
2
=
0
-\vec{g}-c_m\|\frac{d\vec{s}}{dt}\|\frac{d\vec{s}}{dt}-\frac{d^2\vec{s}}{dt^2} = 0
−g−cm∥dtds∥dtds−dt2d2s=0
因此,物理方程损失项可以表示为:
L
p
h
y
s
i
c
s
=
1
n
∑
i
=
1
n
(
−
g
⃗
−
c
m
∥
d
s
⃗
d
t
∥
d
s
⃗
d
t
−
d
2
s
⃗
d
t
2
)
L_{physics}=\frac{1}{n}\sum_{i=1}^n{(-\vec{g}-c_m\|\frac{d\vec{s}}{dt}\|\frac{d\vec{s}}{dt}-\frac{d^2\vec{s}}{dt^2})}
Lphysics=n1i=1∑n(−g−cm∥dtds∥dtds−dt2d2s)
4. 代码实现
在本示例中,采用了两个全连接层,每层128个神经元的网络架构,并使用高斯误差线性单元(GELU)作为激活函数。由于损失函数包含网络中的梯度,因此要求所使用的激活函数必须在任何点可微,以确保获得连续的梯度,显然标准的ReLu函数不满足要求(ReLu是分段线性的)。
class PINN(nn.Module):
def __init__(self, input_size, output_size, hidden_size=None):
super().__init__()
if hidden_size is None:
hidden_size = [40, 40, 40, 40]
# 定义网络
layers = []
layer_sizes = [input_size] + hidden_size + [output_size]
self.net = nn.Sequential(
nn.GELU(),
nn.Linear(layer_sizes[0], layer_sizes[1]),
# nn.Tanh(),
nn.GELU(),
nn.Linear(layer_sizes[1], layer_sizes[2]),
# nn.Tanh(),
nn.GELU(),
nn.Linear(layer_sizes[2], layer_sizes[3]),
)
# 初始化参数
self.init_bias()
# 初始化参数
def init_bias(self):
for layer in self.net.children():
if isinstance(layer, nn.Linear):
# 权重采用Xavier正态初始化
nn.init.xavier_normal_(layer.weight, gain=5 / 3) # gain 根据激活函数选择, torch.nn.init.calculate_gain()查看
# 偏置置0
nn.init.constant_(layer.bias, 0.)
# 前向传播
def forward(self, x):
x = self.net(x)
return x
求解器类实现如下:
class PINNSolver:
def __init__(self, model, data, g, device, initial_lr=1e-3, decay=False):
self.model = model
self.g = g
self.data = data[0:90]
self.s_x = self.data[..., 0].reshape(-1, 1)
self.s_y = self.data[..., 1].reshape(-1, 1)
self.t = self.data[..., 2].reshape(-1, 1)
self.t_p = torch.linspace(min(data[:, 2]).item(), max(data[:, 2]).item(), 20000, device=device).reshape(-1, 1)
self.mc = torch.tensor(0.1, dtype=torch.float32, requires_grad=True, device=device)
# 要求追踪导数信息
self.s_x.requires_grad, self.s_y.requires_grad, self.t.requires_grad, self.t_p.requires_grad = True, True, True, True
# 优化器及优化策略
self.optimizer_Adam = torch.optim.Adamax(list(self.model.parameters()) + [self.mc], lr=initial_lr,
weight_decay=1e-3)
self.optimizer = torch.optim.LBFGS(
self.model.parameters(),
lr=initial_lr,
max_iter=50000,
max_eval=50000,
history_size=50,
tolerance_grad=1e-5,
tolerance_change=1.0 * np.finfo(float).eps,
line_search_fn="strong_wolfe"
)
self.iter = 0
self.decay = decay
if self.decay:
self.scheduler = optim.lr_scheduler.StepLR(self.optimizer_Adam, step_size=1000, gamma=0.9)
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer_Adam, factor=0.1, patience=1000, verbose=True)
# 历史损失
self.his_loss = []
def closure(self):
loss = self.loss_fn(self.model, self.s_x, self.s_y, self.t, self.t_p)
self.iter += 1
if self.iter % 100 == 0:
print(f'Loss: {loss.item():.6f}')
self.optimizer.zero_grad()
loss.backward()
return loss
# 计算 u 对 x 的导数, order 为导数阶数
def get_gard(self, u, x, order=1):
if order == 1:
return torch.autograd.grad(outputs=u, inputs=x, grad_outputs=torch.ones_like(u), retain_graph=True,
create_graph=True)[0]
else:
return self.get_gard(self.get_gard(u, x), x, order=order - 1)
# 损失函数
def loss_fn(self, model, x, y, t, t_p, loss=torch.nn.MSELoss()):
output = model(t)
s_x = output[:, 0].reshape(-1, 1)
s_y = output[:, 1].reshape(-1, 1)
output = model(t_p)
s_x_p = output[:, 0].reshape(-1, 1)
s_y_p = output[:, 1].reshape(-1, 1)
v_x = self.get_gard(s_x_p, t_p, order=1)
v_y = self.get_gard(s_y_p, t_p, order=1)
a_x = self.get_gard(v_x, t_p, order=1)
a_y = self.get_gard(v_y, t_p, order=1)
v = torch.cat([v_x, v_y], dim=1)
a = torch.cat([a_x, a_y], dim=1)
g = torch.tensor([[0., self.g]], device=device)
speed = torch.norm(v, p=2, dim=1, keepdim=True)
l1 = loss(s_x, x) + loss(s_y, y)
l2 = torch.mean(torch.square(-a-g-self.mc*speed*v))
l = 0.4*l1 + 0.6*l2
return l
def train(self, nIter):
# 开始训练时间
t0 = time()
for epoch in tqdm(range(nIter)):
loss = self.loss_fn(self.model, self.s_x, self.s_y, self.t, self.t_p)
self.optimizer_Adam.zero_grad()
loss.backward()
self.optimizer_Adam.step()
if self.decay:
self.scheduler.step()
if epoch % 50 == 0:
self.his_loss.append(loss.cpu().detach().numpy())
if epoch % 100 == 0:
tqdm.write(f"Iteration={epoch:04d}, loss={loss.cpu().detach().numpy():.5f}")
self.optimizer.step(closure=self.closure)
# 训练时间
print(f"TrainTime: {time() - t0:0.4f} seconds")
def plot_loss_history(self, ax=None):
if not ax:
fig = plt.figure(figsize=(7, 5))
ax = fig.add_subplot(111)
ax.semilogy(range(len(self.his_loss)), self.his_loss, 'k-')
ax.set_xlabel('$n_{epoch}$')
ax.set_ylabel('$\\phi^{n_{epoch}}$')
ax.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda x, pos: f'{int(x):,}'))
plt.show()
return ax
def predict(self, raw, true_trajectory, device):
true_s_x = true_trajectory[:, 0].reshape(-1, 1)
true_s_y = true_trajectory[:, 1].reshape(-1, 1)
true_t = true_trajectory[:, 2].reshape(-1, 1)
# pred_t = np.linspace(np.min(true_t), np.max(true_t), 1000, dtype=np.float32).reshape(-1, 1)
output = model(true_t)
s_x_pred = output[:, 0].reshape(-1, 1)
s_y_pred = output[:, 1].reshape(-1, 1)
print(s_x_pred)
print(self.mc)
# 误差
error = nn.MSELoss()(s_x_pred, torch.tensor(true_s_x, dtype=torch.float32).to(device)) + \
nn.MSELoss()(s_y_pred, torch.tensor(true_s_y, dtype=torch.float32).to(device))
print(f"error = {error.cpu().detach().numpy():.5f}")
# 可视化展示
plt.figure(figsize=(24, 8))
plt.scatter(self.s_x.cpu().detach().numpy(), self.s_y.cpu().detach().numpy(), label='Train Data', color='r')
plt.plot(s_x_pred.cpu().detach().numpy(), s_y_pred.cpu().detach().numpy(), label='NN Prediction', color='k')
plt.plot(true_s_x.cpu().detach().numpy(), true_s_y.cpu().detach().numpy(), label='True Data', color='g', linewidth=3)
plt.legend(loc='upper left', fontsize=18)
plt.title('抛射运动', fontname='FangSong', fontsize=22, fontweight='bold')
plt.xlabel('$s_x(m)$', fontsize=22, fontweight='bold')
plt.ylabel('$s_y(m)$', fontsize=22, fontweight='bold')
plt.grid(linestyle='--')
# 保存图片
# plt.savefig('./figures/PINN.png', dpi=1200)
plt.show()
5. 结果展示
通过使用使用物理正则化(PINN)在带有噪声、稀疏且不完整的数据上训练神经网络,物理损失使得网络能够通过数据点进行正则化,并以与已知物理规律一致的方式在训练数据之外进行外推。由于噪声数据量较少,这里无法发现真实解。但是随着数据量的增加或噪声的减少,网络能够以极高的准确度学习到真实解。
运行结果如图所示:
若有不足之处,欢迎批评指正!