import torch
import torch.nn as nn

# 数据生成 (x1, x2 -> y)
x1 = torch.rand(100, 1) * 4 - 2  # [-2, 2]
x2 = torch.rand(100, 1) * 3 - 1  # [-1, 2]
y_true = 2.7*x1 + 3.0*x2 + 5.0*x1*x2 - 1.5*x1**2*x2  # 真实多项式

# 特征工程
def make_features(x1, x2, degree=3):
    features = [torch.ones_like(x1)]
    for d in range(1, degree+1):
        for i in range(d+1):
            features.append((x1**i) * (x2**(d-i)))
    return torch.cat(features, dim=1)
x_poly = make_features(x1, x2, degree=3)

# 模型训练
model = nn.Sequential(nn.Linear(10, 1))  # 3次多项式共10项
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.1)
for epoch in range(2000):
    y_pred = model(x_poly)
    loss = nn.MSELoss()(y_pred, y_true)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 500 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}')