# -*- coding: utf-8 -*- """正演代理模型网络结构。 ForwardSurrogate 输入标准化后的物理参数特征和可选的流量制度编码，输出固定长度 拼接曲线：log_pressure、log_derivative 和 slope。模型采用“参数分支 + 流量制度 分支 + 融合主干 + 多输出头”的结构，便于分别学习静态地层信息、动态制度信息以及 二者共同决定的曲线形态。 压力和导数输出被拆成 level 与 shape 两部分：level 学习整条曲线的纵向偏移，shape 学习去均值后的局部形态，从结构上减少整体幅值与局部形状之间的相互干扰。 """ from __future__ import annotations import torch import torch.nn as nn def build_mlp( in_dim: int, hidden_dims: list[int], out_dim: int, dropout: float = 0.0, ) -> nn.Sequential: """按隐藏层列表搭建 Linear-ReLU-Dropout 组成的多层感知机。""" layers: list[nn.Module] = [] prev = in_dim for h in hidden_dims: layers.append(nn.Linear(prev, h)) layers.append(nn.ReLU()) if dropout > 0: layers.append(nn.Dropout(dropout)) prev = h layers.append(nn.Linear(prev, out_dim)) return nn.Sequential(*layers) class ScheduleEncoder(nn.Module): """神经网络中的流量制度分支，把固定长度制度向量编码为隐层特征。""" def __init__(self, schedule_dim: int, hidden_dim: int, dropout: float = 0.0): """按流量制度向量维度构建两层编码网络。""" super().__init__() self.net = nn.Sequential( nn.Linear(schedule_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout) if dropout > 0 else nn.Identity(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), ) def forward(self, x: torch.Tensor) -> torch.Tensor: """把流量制度统计特征映射到与参数分支同宽度的隐藏表示。""" # 该分支只处理制度向量，便于后续与地层参数特征拼接融合。 return self.net(x) class ParamEncoder(nn.Module): """神经网络中的参数分支，把变换后的物理参数编码为隐层特征。""" def __init__(self, param_dim: int, hidden_dim: int, dropout: float = 0.0): """按物理参数特征维度构建两层编码网络。""" super().__init__() self.net = nn.Sequential( nn.Linear(param_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout) if dropout > 0 else nn.Identity(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), ) def forward(self, x: torch.Tensor) -> torch.Tensor: """把变换后的地层和井筒参数映射为隐藏表示。""" # 参数特征通常来自 log/asinh 等尺度变换，先编码再与制度分支融合。 return self.net(x) class ForwardSurrogate(nn.Module): """完整曲线正演代理模型。 输入: params_x: 标准化后的物理参数特征，形状 [B, param_dim]。 schedule_x: 标准化后的流量制度向量，形状 [B, schedule_dim]；当 use_schedule=False 时该输入可为空。 输出: curve_pred: 形状 [B, curve_dim]，按 log_pressure、log_derivative、slope 三段 顺序拼接。curve_dim 必须能被 3 整除，以便每段拥有相同时间点数。 """ def __init__( self, param_dim: int, schedule_dim: int, curve_dim: int, hidden_dim: int = 128, fusion_hidden_dims: list[int] | None = None, dropout: float = 0.0, use_schedule: bool = True, ): """构建参数分支、可选流量制度分支、融合主干和三组曲线输出头。""" super().__init__() if curve_dim % 3 != 0: raise ValueError(f"curve_dim={curve_dim} 不能被 3 整除；期望为 pressure/derivative/slope 三段") if fusion_hidden_dims is None: fusion_hidden_dims = [256, 256] self.curve_dim = curve_dim self.part_dim = curve_dim // 3 self.use_schedule = bool(use_schedule) # 参数和流量制度的物理含义与尺度差异较大，因此采用两个分支分别编码。 self.param_encoder = ParamEncoder(param_dim, hidden_dim, dropout=dropout) if self.use_schedule: self.schedule_encoder = ScheduleEncoder(schedule_dim, hidden_dim, dropout=dropout) trunk_in_dim = hidden_dim * 2 else: self.schedule_encoder = None trunk_in_dim = hidden_dim trunk_out_dim = fusion_hidden_dims[-1] self.trunk = build_mlp( in_dim=trunk_in_dim, hidden_dims=fusion_hidden_dims, out_dim=trunk_out_dim, dropout=dropout, ) # 压力曲线拆成 level + centered shape： # level 学习整体纵向偏移，shape 学习局部曲线形态。 self.pressure_level_head = build_mlp( in_dim=trunk_out_dim, hidden_dims=[128], out_dim=1, dropout=dropout, ) self.pressure_shape_head = build_mlp( in_dim=trunk_out_dim, hidden_dims=[128], out_dim=self.part_dim, dropout=dropout, ) # 导数曲线同样拆分为 level + shape，因为平台、谷值和过渡段 # 对自动拟合筛选非常重要。 self.derivative_level_head = build_mlp( in_dim=trunk_out_dim, hidden_dims=[128], out_dim=1, dropout=dropout, ) self.derivative_shape_head = build_mlp( in_dim=trunk_out_dim, hidden_dims=[128], out_dim=self.part_dim, dropout=dropout, ) # slope 是辅助输出，主要用于保持数据布局兼容。 self.slope_head = build_mlp( in_dim=trunk_out_dim, hidden_dims=[128], out_dim=self.part_dim, dropout=dropout, ) @staticmethod def center_shape(x: torch.Tensor) -> torch.Tensor: """去除每个样本 shape 分支的均值，让 level 分支专门学习整体偏移。""" return x - x.mean(dim=1, keepdim=True) def forward(self, params_x: torch.Tensor, schedule_x: torch.Tensor | None = None) -> torch.Tensor: """执行一次前向预测。 参数分支和流量制度分支先分别编码，再在隐空间拼接。融合主干提取共同特征后， 压力和导数各自通过 level + centered shape 两个输出头生成；slope 作为辅助通道 直接由单独输出头预测。返回值仍保持预处理阶段约定的曲线拼接布局。 """ p = self.param_encoder(params_x) if self.use_schedule: if schedule_x is None: raise ValueError("use_schedule=True，但 forward 没有传入 schedule_x") s = self.schedule_encoder(schedule_x) # 两个分支在隐藏空间拼接，避免直接混合量纲差异很大的原始特征。 fused = torch.cat([p, s], dim=-1) else: fused = p # trunk 负责学习参数-制度共同决定的曲线整体形态。 trunk_feat = self.trunk(fused) pressure_level = self.pressure_level_head(trunk_feat) # [B, 1] pressure_shape = self.pressure_shape_head(trunk_feat) # [B, T] # shape 去均值后只表达相对形态，纵向偏移交给 level 分支学习。 pressure_shape = self.center_shape(pressure_shape) pressure_pred = pressure_level + pressure_shape derivative_level = self.derivative_level_head(trunk_feat) # [B, 1] derivative_shape = self.derivative_shape_head(trunk_feat) # [B, T] # 导数也采用 level + shape，减少平台值和局部过渡段之间的相互牵制。 derivative_shape = self.center_shape(derivative_shape) derivative_pred = derivative_level + derivative_shape slope_pred = self.slope_head(trunk_feat) # [B, T] curve_pred = torch.cat([pressure_pred, derivative_pred, slope_pred], dim=1) return curve_pred