2025 USA-NA-AIO Round 2, Problem 1, Part 12

AI Olympiads 2025 USA-NA-AIO Round 2
1

Part 12 (10 points, coding task)

In this part, you are asked to do the following tasks to test the effectiveness of our PINN model.

  1. Generate a dataset \left\{ \left( t , x \right) \in \left\{ 0, 0.01, \cdots, 1 \right\}^2 \right\}. Save the dataset as a tensor with name tx_test and shape (101,2).

  2. For each data point, compute u \left( t , x \right) whose formula is given in Part 1. Save the result as a tensor with name u_test and shape (101,2).

  3. For each data point, use our trained PINN model to compute the predicted value U \left( t, x \mid \mathbf{\theta} \right). Save the result as a tensor with name U_test and shape (101,2).

  4. Print the mean squared error between u_test and U_test.

  5. Generate two 2-dim scatter plots for \left( t, x \right) by using the above data points.

    • In Figure 1, the value on each position is the ground-truth temperature u \left( t , x \right).

    • In Figure 2, the value on each position is the predicted temperature U \left( t, x \mid \mathbf{\theta} \right).

    • In each plot,

      • Set c as the values on those scattered positions

      • Set cmap='viridis'

      • Add plt.colorbar(label='Value').

2 
### WRITE YOUR SOLUTION HERE ###

model.to('cpu')
data_1dim = torch.linspace(0, 1, 101)
data_grid = torch.meshgrid(data_1dim, data_1dim)
tx_test = torch.stack([data_grid[0].reshape(-1), data_grid[1].reshape(-1)], dim=1)
u_test = torch.exp(- alpha * torch.pi**2 * tx_test[:,0]) * torch.sin(torch.pi * tx_test[:,1])
U_test = model(tx_test).reshape(-1).detach()

mse = torch.mean((u_test - U_test)**2).item()
print(mse)

plt.figure(1)
plt.scatter(tx_test[:,0], tx_test[:,1], c=u_test, cmap='viridis')
plt.xlabel('t')
plt.ylabel('x')
plt.colorbar(label='Value')
plt.title('Ground-truth')

plt.show()

plt.figure(2)
plt.scatter(tx_test[:,0], tx_test[:,1], c=U_test, cmap='viridis')
plt.xlabel('t')
plt.ylabel('x')
plt.colorbar(label='Value')
plt.title('Predicted')

plt.show()


""" END OF THIS PART """