Как обучить модель для оценки коэффициентов связанных ОДУ?

Рассмотрите следующую систему связных ОДУ (уравнения Лотки-Вольтерра):

$$
\frac{dx}{dt} = \alpha x – \beta x y, \\
\frac{dy}{dt} = – \gamma y + \delta x y ,
$$

Как я могу обучить модель для оценивания положительных параметров $\alpha$, $\beta$, $\delta$, $\gamma$, имея данные о траекториях $x(t)$ и $y(t)$?

Я демонстрирую этот подход в своем ответе:

1. Создание тренировочного набора

   • Случайным образом выберите значения параметров и начальные значения (или выберите диапазоны, релевантные вашей задаче)
   • Решите уравнения для траекторий $x(t)$ и $y(t)$ каждого набора параметров, используя решатель ОДУ
   • Каждая пара $\large[x$ и $y$ траектория$,$ 4 параметра$\large]$ становится образцом в тренировочном наборе.

2. Обучите модель многорегрессии, чтобы предсказывать 4 параметра для каждой двумерной последовательности. Так как данные являются сильно последовательными, последовательная модель будет хорошим кандидатом (например, LSTM/GRU, 1D CNNs).

   • На вход поступает пакет двумерных последовательностей (пакет $x(t)$ и $y(t)$), и целью являются соответствующие пакеты значений 4D параметров.

Используя LSTM, я получаю 10-20% ошибку предсказания по всем четырем параметрам. Я обучил её на относительно широком диапазоне параметров и начальных значений и ожидаю, что производительность будет лучше, если вы ограничите диапазон. Вы также можете дополнительно настроить модель и/или обучить её дольше.

Размер модели: 2304 обучаемых параметра
эпоха  1/100  [ЗНА обучение:  0.213 | мАПОИ [54 56 55 57]] [ЗНА валидации:  0.210 | мАПОИ [57 55 55 54]]
...
эпоха 100/100 [ЗНА обучение:  0.023 | мАПОИ [12 17 17 11]] [ЗНА валидации:  0.025 | мАПОИ [12 18 19 11]]

Я случайно выбираю наборы параметров, где каждый параметр выбирается из диапазона [0.1, 0.9), и начальные значения также выбраны случайно между [1, 50). Один такой образец:

Если вы используете различные диапазоны для параметров, стоит использовать взвешивание ошибок, чтобы дать меньшим параметрам шанс на улучшение (вы могли бы также нормализовать цели для их привязки к одной шкале).

Сначала я использую solve_ivp() из SciPy для решения траекторий $x(t), y(t)$. Затем я обучаю однослойный LSTM для уменьшения MSE между предсказаниями и целями. Я обеспечиваю положительные значения для каждого параметра, используя слой Softplus().

Тренировочный набор данных состоял примерно из 3800 последовательностей, где каждая последовательность покрывала временной диапазон (0, 100) в 200 шагов. Другие гиперпараметры обучения: размер пакета 64, 100 эпох, однослойная LSTM с размером скрытого слоя 64, и оптимизатор NAdam.

Я отслеживал MAPE для каждого параметра. Кажется, что $\beta$ и $\delta$ было сложнее предсказать правильно, потому что их MAPE были ближе к 20%, в то время как параметры $\alpha$ и $\gamma$ имели значение около 10%. Финальная валидационная ошибка составляла 0.025.

Воспроизводимый пример

Синтезировать данные и разделить:

import numpy as np
from matplotlib import pyplot as plt
from scipy.integrate import solve_ivp

np.random.seed(0)

#
# Данные для тестирования
#
n_sequences = 5_000
n_timesteps = 200

t_span = (0, 100)
alpha_span = (0.1, 0.9)
beta_span = (0.1, 0.9)
delta_span = (0.1, 0.9)
gamma_span = (0.1, 0.9)
x0_span = (1, 50)
y0_span = (1, 50)

#Выбираем `n_sequences` наборов параметров из вышеуказанных параметров
alphas, betas, deltas, gammas = [
    np.random.uniform(*span, size=n_sequences)
    for span in [alpha_span, beta_span, delta_span, gamma_span]
]
parameters = np.column_stack([alphas, betas, deltas, gammas])

initial_values = np.column_stack([
    np.random.randint(*x0_span, size=n_sequences),
    np.random.randint(*y0_span, size=n_sequences)
])

#
# Решение для x/y траекторий каждого набора параметров
#
def solve_lotka_volterra(parameters, xy0, t_span, n_timesteps):
    def lokta_volterra(t, xy, alpha, beta, delta, gamma):
        x, y = xy.T

        dx_dt = alpha*x - beta*x*y
        dy_dt = delta*x*y - gamma*y
        return np.column_stack([dx_dt, dy_dt])
    #\def lotka_volterra

    t_eval = np.linspace(*t_span, num=n_timesteps)
    t_span = t_eval.min(), t_eval.max()

    solver_results = solve_ivp(
        lokta_volterra, t_span, xy0, t_eval=t_eval, args=parameters
    )

    if not solver_results['success']:
        print('[!] solve_ivp не удалось |', solver_results['message'])

    return solver_results['y'].T #возвращает (n_timesteps, 2)

#Для каждого набора параметров решаем траекторию размера (n_timesteps, 2)
# Стекаем в (n_sequences, n_timesteps, 2)
sequences = np.stack(
    [solve_lotka_volterra(params_i, xy0_i, t_span, n_timesteps)
     for params_i, xy0_i in zip(parameters, initial_values)
     ],
     axis=0
)

#
# Разделение на тренировочную и валидационную выборку
#
from sklearn.model_selection import train_test_split

splits = train_test_split(sequences, parameters, train_size=3/4)
sequences_trn, sequences_val, parameters_trn, parameters_val = splits

Визуализация случайно выбранного образца:

#
#Визуализация случайного образца из тренировочного и валидационного набора
#
ix = np.random.randint(n_sequences)
x, y = sequences[ix].T
params = parameters[ix]

x, y = sequences[np.random.randint(n_sequences)].T
t = np.linspace(*t_span, n_timesteps)

plt.suptitle(r'$\alpha, \beta, \delta, \gamma=$' + str(params.round(3)))
plt.subplot(121)
plt.plot(t, x, t, y)
plt.xlabel('t')
plt.ylabel('x, y')

plt.subplot(122)
plt.scatter(x, y, marker=".", c=t, s=8, edgecolors="none")
plt.xlabel('x')
plt.ylabel('y')

plt.gcf().set_size_inches(9, 4)
plt.tight_layout()

Преобразуйте в тензоры PyTorch и оберните в загрузчик пакетов:

import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

n_features = sequences.shape[-1] #x, y
output_size = parameters.shape[1] #alpha, beta, delta, gamma

#
# Масштабирование датасета
#
from sklearn.preprocessing import StandardScaler, MinMaxScaler

scaler = StandardScaler().fit(sequences_trn.reshape(-1, n_features))
X_trn = scaler.transform(sequences_trn.reshape(-1, n_features)).reshape(sequences_trn.shape)
X_val = scaler.transform(sequences_val.reshape(-1, n_features)).reshape(sequences_val.shape)

#
# В тензорный датасет
#
trn_dataset = TensorDataset(torch.tensor(X_trn).float(), torch.tensor(parameters_trn).float())
val_dataset = TensorDataset(torch.tensor(X_val).float(), torch.tensor(parameters_val).float())

#Разбиение на пакеты
batch_size = 64

trn_loader = DataLoader(trn_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)

Определите и обучите модель, сообщая о прогрессе:

torch.manual_seed(0)
np.random.seed(0)

class LambdaLayer (nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)

model = nn.Sequential(
    #B L C

    nn.LSTM(input_size=n_features, hidden_size=64, num_layers=1, batch_first=True, proj_size=output_size),
    #output: B L D*out, h_n: D*layers B out, c_n: D*layers B cell

    LambdaLayer(lambda output_hncn: output_hncn[0][:, -1]),

    #Параметры должны быть положительными
    nn.Softplus(),
)

print(
    'Размер модели:',
    sum(p.numel() for p in model.parameters() if p.requires_grad),
    'обучаемых параметров'
)

optimizer = torch.optim.NAdam(model.parameters())

n_epochs = 100

#Балансировать ошибки относительно самого большого значения цели
target_weights = torch.tensor(
    parameters_trn.mean(axis=0).max() / parameters_trn.mean(axis=0)
).float().ravel()
# target_weights = torch.tensor([1, 1, 1, 1]) #без балансировки

@torch.no_grad()
def calc_metrics(model, loader, target_weights=None):
    model.eval()

    if target_weights is None:
        target_weights =  torch.ones(output_size)

    cum_loss = 0
    maes = []
    mapes = []

    for X_mb, target_mb in loader:
        preds = model(X_mb)
        cum_loss += preds.sub(target_mb).pow(2).mul(target_weights).sum()

        mae_pertarget = preds.sub(target_mb).abs().mean(dim=0)
        mape_pertarget = mae_pertarget.div(target_mb).mul(100)

        maes.append(mae_pertarget)
        mapes.append(mape_pertarget)

    loss = cum_loss / len(loader.dataset)
    maes = np.row_stack(maes)
    mapes = np.row_stack(mapes)

    return loss, maes.mean(axis=0), mapes.mean(axis=0)

#
# Цикл обучения
#
for epoch in range(n_epochs):
    model.train()

    for (X_mb, target_mb) in trn_loader:
        preds = model(X_mb)

        loss = preds.sub(target_mb).pow(2).mul(target_weights).mean()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    #/конец эпохи

    # if (epoch % 5) != 0:
    #     continue
    model.eval()
    with torch.no_grad():
        trn_loss, _, trn_mapes = calc_metrics(model, trn_loader, target_weights)
        val_loss, _, val_mapes = calc_metrics(model, val_loader, target_weights)

        print(
            f'эпоха {epoch + 1:2d}/{n_epochs}',
            f'[ЗНА обучение: {trn_loss:6.3f} | мАПОИ {trn_mapes.round()}]',
            f'[ЗНА валидация: {val_loss:6.3f} | мАПОИ {val_mapes.round()}]',
        )

Визуализируйте цели и предсказанные значения:

#
# Визуализировать предсказанные vs целевые параметры для каждого параметра
#
model.eval()

with torch.no_grad():
    trn_preds = torch.concatenate([model(X_mb) for X_mb, _ in trn_loader], dim=0).numpy()
    val_preds = torch.concatenate([model(X_mb) for X_mb, _ in val_loader], dim=0).numpy()

trn_loss, trn_maes, trn_mapes = calc_metrics(model, trn_loader, target_weights)
val_loss, val_maes, val_mapes = calc_metrics(model, val_loader, target_weights)

viz_val = True

if viz_val:
    use_preds = val_preds
    use_mapes = val_mapes
else:
    use_preds = trn_preds
    use_mapes = trn_mapes

pred_alpha, pred_beta, pred_delta, pred_gamma = use_preds.T

f, axs = plt.subplots(ncols=output_size, figsize=(12, 3.5), layout="tight")

for ax, pred, target, name, mape in zip(
    axs,
    use_preds.T,
    parameters_val.T,
    r'$\alpha$ $\beta$ $\delta$ $\gamma$'.split(),
    use_mapes
):
    ax.scatter(target, pred, marker=".", alpha=0.3, edgecolors="none")
    ax.set_title(f'{name}\nMAPE={mape:.0f}%')
    ax.set(xlabel="цель", ylabel="прогноз")
    ax.axline([0, 0], slope=1, linestyle="--", color="black")

axs[0].plot([np.nan], [np.nan], linestyle="--", color="black", label="линия идентичности")
f.legend(loc="lower left", shadow=True, bbox_to_anchor=(0.05, -0.1))
f.suptitle(f'Предсказанные vs целевые параметры для {"валидационная" if viz_val else "тренировочная"} выборка', weight="bold");