python 代码实现了一个基于物理信息神经网络（PINN）的模型，用于解决特定的流体力学问题

import torch
import torch.nn as nn
import numpy as np
import math
import scipy.io
from fun import func  # 确保fun模块及其func类已正确定义
import matplotlib.pyplot as plt
import time
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/pinn_experiment')

# 检查GPU是否可用，并设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {
     device}')

# 设置随机种子以确保结果可复现
np.random.seed(42)
torch.manual_seed(42)
if device.type == 'cuda':
    torch.cuda.manual_seed_all(42)

# 定义基本物理参数
L = 37.2  # 管长
D = 0.0221  # 管道直径
a = 1319  # 波速
C0 = 0.0019  # 阀门流量系数
TC = 0.07  # 关阀时间
H0 = 26.5  # 上游水位
H1 = 26.295  # 下游水位
vv = 1.007 * 10 ** (-6)  # 水在20℃运动粘度系数
k = 1.5 * 10 ** (-3)  # 管壁绝对粗糙度
Delta = k / D  # 相对粗糙度
g = 9.8  # 重力加速度
A = math.pi * D ** 2 / 4  # 主管面积
Q01 = 6.3324 * 10 ** (-5)  # 定义流量Q01
V = Q01 / A  # 平均流速

# 计算雷诺数和相关物理量
fu = func()
Re = fu.Reynold_number(Q01, D, vv)  # 雷诺数
k3 = fu.f_k(Re)  # 柯尔布鲁克参数
f1 = fu.cole_brook_white_2(Re, Delta)  # Colebrook-White方程结果

class Pinn:
    def __init__(self, X, T, v, h, normalization_params, lambda_reg=1e-6):
        """
        初始化PINN模型
        :param X: 空间变量 (NT x 1)
        :param T: 时间变量 (NT x 1)
        :param v: 流量数据 (NT x 1)
        :param h: 水头数据 (NT x 1)
        :param normalization_params: 字典，包含归一化参数
        :param lambda_reg: L2正则化参数
        """
        # 存储归一化参数
        self.min_v = normalization_params['min_v']
        self.max_v = normalization_params['max_v']
        self.min_h = normalization_params['min_h']
        self.max_h = normalization_params['max_h']
        self.min_x = normalization_params['min_x']
        self.max_x = normalization_params['max_x']
        self.min_t = normalization_params['min_t']
        self.max_t = normalization_params['max_t']
        self.Num = normalization_params['Num']
        self.Tim = normalization_params['Tim']

        # 将数据转换为张量并移动到指定设备
        self.x = torch.tensor(X, dtype=torch.float32, requires_grad=True).to(device)
        self.t = torch.tensor(T, dtype=torch.float32, requires_grad=True).to(device)
        self.v = torch.tensor(v, dtype=torch.float32).to(device)
        self.h = torch.tensor(h, dtype=torch.float32).to(device)

        # 定义一个零张量用于PDE损失
        self.null = torch.zeros((self.t.shape[0], 1), device=device)

        # 初始化神经网络
        self.network()

        # 将网络移动到指定设备
        self.net.to(device)

        # 设置L2正则化参数
        self.lambda_reg = lambda_reg

        # 设置优化器和学习率调度器
        self.optimizer_adam = torch.optim.Adam(self.net.parameters(), lr=0.01)
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer_adam, patience=100, factor=0.5, verbose=True)
        # 定义LBFGS优化器
        self.optimizer_lbfgs = torch.optim.LBFGS(
            self.net.parameters(),
            lr=0.001,
            max_iter=500,
            max_eval=500,
            history_size=50,
            tolerance_grad=1e-07,
            tolerance_change=1e-09,
            line_search_fn="strong_wolfe"
        )
        # 定义损失函数
        self.mse = nn.MSELoss().to(device)
        self.ls = 0  # 当前损失值
        self.iter = 0  # 当前迭代次数

        # 历史损失记录
        self.total_loss_history = []
        self.pde_loss_history = []
        self.data_loss_history = []

        # 预测结果初始化
        self.v_pred = None
        self.h_pred = None

    def network(self):
        """
        定义神经网络结构
        """
        self.net = nn.Sequential(
            nn.Linear(2, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 2)  # 输出两个值：v和h
        )
        # 初始化网络权重
        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zer