基于神经网络求解二阶ODE：代码适配与实现疑问

问题描述 #

本人正在学习基于机器学习的常微分方程（ODE）数值解法，现有一套用于求解一阶ODE的TensorFlow代码，希望将其适配以求解二阶ODE：$y'' = -10y$（初值条件$y(0)=0$，$y'(0)=1$）。存在以下疑问：

1. 如何将该二阶ODE拆分为两个一阶ODE？是否需要修改代码中的f(x)函数？
2. 如何在损失函数中实现二阶ODE的约束？
3. GradientTape是否需要调整？

本人机器学习基础薄弱，尝试适配时卡在损失函数、f(x)及g(x)的修改环节，希望获取书籍、文章、代码示例或相关思路指导。

核心思路与解答 #

1. 二阶ODE拆分与函数修改 #

对于二阶ODE $y'' = -10y$，引入新变量$z = y'$，将其转化为一阶方程组：

• $y' = z$
• $z' = -10y$

对应初值条件更新为：$y(0)=0$，$z(0)=1$。

此时需要调整代码核心函数：

• 神经网络输出维度改为2，分别对应$y(x)$和$z(x)$
• 原g(x)从单输出改为返回二维向量，同时要满足初值条件
• 原f(x)替换为方程组的右端逻辑，即输入$y,z$返回$z$和$-10y$

2. 损失函数的二阶约束实现 #

损失函数需包含两部分：

• 初值损失：强制$y(0)=0$、$z(0)=1$
• ODE残差损失：验证$y'=z$和$z'=-10y$，通过有限差分计算导数代入残差

具体逻辑：
对每个采样点$x$，用有限差分计算$y'(x) \approx \frac{g_y(x+inf_s)-g_y(x)}{inf_s}$、$z'(x) \approx \frac{g_z(x+inf_s)-g_z(x)}{inf_s}$，然后计算残差的平方和，与初值损失相加得到总损失。

3. GradientTape的调整 #

无需大幅调整，只需保证损失函数的所有计算过程在GradientTape作用域内即可，原有的梯度计算与优化逻辑可复用。

修改后的完整代码 #

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(123)
tf.random.set_seed(123)

# 初值条件：y(0)=0，y'(0)=1 → z(0)=1
y0 = 0.0
z0 = 1.0

# 微小量用于有限差分计算导数
inf_s = np.sqrt(np.finfo(np.float32).eps)

# 训练参数
learning_rate = 0.001
training_steps = 10000
display_step = training_steps // 10

# 网络参数：输出改为2维（对应y和z）
n_input = 1 
n_hidden_1 = 32
n_hidden_2 = 32
n_output = 2  # 输出y和z两个值

weights = {
    'h1': tf.Variable(tf.random.normal([n_input, n_hidden_1])),
    'h2': tf.Variable(tf.random.normal([n_hidden_1, n_hidden_2])),
    'out': tf.Variable(tf.random.normal([n_hidden_2, n_output]))
}
biases = {
    'b1': tf.Variable(tf.random.normal([n_hidden_1])),
    'b2': tf.Variable(tf.random.normal([n_hidden_2])),
    'out': tf.Variable(tf.random.normal([n_output]))
}

# 优化器改用Adam，收敛效果比SGD更好
optimizer = tf.optimizers.Adam(learning_rate)

# 多层感知机模型：输出二维向量[y, z]
def multilayer_perceptron(x):
    x = tf.convert_to_tensor([[x]], dtype=tf.float32)
    # 隐藏层1
    layer_1 = tf.add(tf.matmul(x, weights['h1']), biases['b1'])
    layer_1 = tf.nn.tanh(layer_1)  # tanh比sigmoid更适合这类问题
    # 隐藏层2
    layer_2 = tf.add(tf.matmul(layer_1, weights['h2']), biases['b2'])
    layer_2 = tf.nn.tanh(layer_2)
    # 输出层
    output = tf.matmul(layer_2, weights['out']) + biases['out']
    return output[0]  # 返回[ y, z ]

# 近似函数：满足初值条件
def g(x):
    # 构造方式：确保x=0时，y=y0，z=z0
    nn_output = multilayer_perceptron(x)
    y = x * nn_output[0] + y0
    z = x * nn_output[1] + z0
    return y, z

# ODE方程组的右端项
def ode_rhs(y, z):
    return z, -10 * y

# 自定义损失函数
def custom_loss():
    loss = 0.0
    # 初值条件损失
    y_0, z_0 = g(0.0)
    loss += tf.square(y_0 - y0) + tf.square(z_0 - z0)
    
    # ODE残差损失：采样多个点计算
    for x in np.linspace(0, 2, 20):  # 求解区间设为[0,2]，对应解析解周期
        y_x, z_x = g(x)
        # 用有限差分计算y'和z'
        y_next, z_next = g(x + inf_s)
        y_prime = (y_next - y_x) / inf_s
        z_prime = (z_next - z_x) / inf_s
        
        # 计算残差：y'应该等于z，z'应该等于-10y
        rhs_y, rhs_z = ode_rhs(y_x, z_x)
        loss += tf.square(y_prime - rhs_y) + tf.square(z_prime - rhs_z)
    
    return loss

# 训练步骤
def train_step():
    with tf.GradientTape() as tape:
        loss = custom_loss()
    trainable_vars = list(weights.values()) + list(biases.values())
    gradients = tape.gradient(loss, trainable_vars)
    optimizer.apply_gradients(zip(gradients, trainable_vars))

# 开始训练
for i in range(training_steps):
    train_step()
    if i % display_step == 0:
        print(f"Step {i}, Loss: {custom_loss().numpy():.6f}")

# 可视化结果
figure(figsize=(10, 6))

# 解析解：y(x) = (1/√10) sin(√10 x)
def true_solution(x):
    return (1 / np.sqrt(10)) * np.sin(np.sqrt(10) * x)

X = np.linspace(0, 2, 100)
nn_y = []
true_y = true_solution(X)

for x in X:
    y, _ = g(x)
    nn_y.append(y.numpy())

plt.plot(X, true_y, label="解析解", linewidth=2)
plt.plot(X, nn_y, label="神经网络近似解", linestyle='--', linewidth=2)
plt.xlabel("x")
plt.ylabel("y(x)")
plt.legend(fontsize=12)
plt.grid(True)
plt.show()

学习资源推荐 #

• 书籍：《Physics-Informed Neural Networks》（物理信息神经网络专著，详细讲解用神经网络求解ODE/PDE）
• 核心思路：搜索"Physics-Informed Neural Networks for ODEs"，重点关注二阶ODE转一阶方程组的实现细节
• 代码参考：查看TensorFlow官方文档中自动微分与ODE结合的示例，理解损失函数构造逻辑

内容的提问来源于stack exchange，提问作者Leonardo Cardoso Belintani