如何在Darts中实现含静态协变量的Lorenz系统ODE自定义损失函数

问题背景 #

你当前用Darts的TFT模型预测3D Lorenz混沌时间序列，训练数据是分组合成的200组模拟数据，每组包含固定的静态协变量sigma、rho、beta，数据加载代码如下：

train_series = TimeSeries.from_group_dataframe(training_df, 
                                group_cols = 'simNum',
                                time_col = 'time_idx',
                                value_cols = ['x','y','z'],
                                static_cols = ['sigma','rho','beta'])

所有数值列已做Scaler缩放，time_idx为每组独立重置的时间整数。目标是将Lorenz系统的3个ODE作为残差损失，与MSE数据损失结合训练模型，但遇到核心问题：

• 不知道如何在配点（训练时间步）生成模型预测值用于ODE残差计算
• 不清楚如何将每组的静态协变量融入ODE损失的计算逻辑
• 纠结是否需要随机生成sigma/rho/beta来整合预测与ODE计算

核心解决方案 #

1. 配点预测的获取：直接利用模型forward输出 #

不需要额外调用model.predict，在训练循环中，TFT的LightningModule会接收batch输入，直接拿到模型对当前输入序列的预测值（配点输出）。重写TFT的training_step方法即可实现：

from darts.models import TFTModel
import torch

class ODEIntegratedTFT(TFTModel):
    def __init__(self, *args, ode_loss_weight=0.1, **kwargs):
        super().__init__(*args, **kwargs)
        self.ode_loss_weight = ode_loss_weight
        self.mse_loss = torch.nn.MSELoss()

    def lorenz_ode(self, x, y, z, sigma, rho, beta):
        # 适配缩放后数据的Lorenz ODE实现
        dx_dt = sigma * (y - x)
        dy_dt = x * (rho - z) - y
        dz_dt = x * y - beta * z
        return dx_dt, dy_dt, dz_dt

    def training_step(self, batch, batch_idx):
        # 提取batch中的输入与目标序列
        x, y = batch
        # 获取模型对输入序列的配点预测
        pred = self(x)
        # 计算数据损失：预测值与目标的MSE
        data_loss = self.mse_loss(pred, y)

        # 提取当前batch的静态协变量：每组的sigma/rho/beta
        static_covs = x['static_covariates']
        sigma = static_covs[:, 0]
        rho = static_covs[:, 1]
        beta = static_covs[:, 2]

        # 从预测值中拆分x/y/z分量，用差分近似连续导数（适配离散时间序列）
        pred_x = pred[:, :, 0].squeeze()
        pred_y = pred[:, :, 1].squeeze()
        pred_z = pred[:, :, 2].squeeze()
        # 用后向差分近似导数（步长对应time_idx的间隔）
        dx_dt_pred = torch.diff(pred_x, dim=1)
        dy_dt_pred = torch.diff(pred_y, dim=1)
        dz_dt_pred = torch.diff(pred_z, dim=1)

        # 用ODE计算理论导数（对应差分后的时间步）
        dx_dt_ode, dy_dt_ode, dz_dt_ode = self.lorenz_ode(
            pred_x[:, :-1], pred_y[:, :-1], pred_z[:, :-1], sigma, rho, beta
        )

        # 计算ODE残差损失
        ode_loss = self.mse_loss(
            torch.stack([dx_dt_pred, dy_dt_pred, dz_dt_pred], dim=-1),
            torch.stack([dx_dt_ode, dy_dt_ode, dz_dt_ode], dim=-1)
        )

        # 加权合并总损失
        total_loss = data_loss + self.ode_loss_weight * ode_loss

        # 记录训练日志
        self.log("train_data_loss", data_loss)
        self.log("train_ode_loss", ode_loss)
        self.log("train_total_loss", total_loss)
        return total_loss

2. 静态协变量的融入：直接从batch提取，无需随机生成 #

训练数据中每组的sigma/rho/beta是已知的静态协变量，Darts加载分组数据时已将这些信息绑定到样本中。在训练循环的batch里，静态协变量会作为输入的一部分传递给模型，直接从x['static_covariates']提取即可用于ODE计算，完全不需要随机生成参数——随机生成仅适用于无监督参数估计场景，你这里训练数据的参数是已知的，应该用真实值约束模型。

3. 损失函数的整合：在training_step中加权结合 #

不要直接定义loss_fn = MSELoss() + odeLoss()，这种方式无法传递静态协变量和配点预测的上下文。正确方式是在training_step中分别计算两种损失，再根据需求加权求和，可灵活调整权重避免某一种损失主导训练。

注意事项 #

• 确保Lorenz ODE实现与数据缩放后的数值范围匹配，必要时对ODE输出做反向缩放
• 导数近似方式可根据时间步长调整，比如用中心差分或数值积分提升精度
• 若需在验证阶段加入ODE损失，只需重写validation_step方法，逻辑与training_step一致

内容的提问来源于stack exchange，提问作者Farooq Zahid