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

AI Olympiads 2025 USA-NA-AIO Round 2
1

The high level idea of PINN is as follows:

  1. (Neural network) We design a neural network (functional mapping) U \left( \cdot, \cdot \mid \mathbf{\theta} \right): \left[ 0 , 1 \right]^2 \rightarrow \Bbb R, such that

    • \mathbf{\theta}: learnable parameters in U.

    • Inputs are time t and position x.

    • Output is the predicted temperature.

  2. (Training data) To train U (equivalently, to learn \mathbf{\theta}), we use the following three groups of temporal-spacial data \left( t, x \right):

    • (Training data for PDE) \left( t, x \right) are randomly sampled from \left[ 0 , 1 \right]^2.

      Denote by \mathcal D_{PDE} the set of these data points.

    • (Training data for IC) \left( 0, x \right) with x that are evenly distributed on \left[ 0 , 1 \right].

      Denote by \mathcal D_{IC} the set of these data points.

    • (Training data for BC) \left( t, 0 \right) and \left( t, 1 \right) with t that are evenly distributed on \left[ 0 , 1 \right].

      Denote by \mathcal D_{BC} the set of these data points.

  3. (Loss function in training)

    L_{total} = L_{PDE} + L_{IC} + L_{BC} ,

    where

    • Residual loss in PDE:
    L_{PDE} = \frac{1}{| \mathcal D_{PDE} |} \sum_{\left( t, x \right) \in \mathcal D_{PDE}} \left( \frac{\partial U \left( t, x \mid \mathbf{\theta} \right)}{\partial t} - \alpha \frac{\partial^2 U \left( t, x \mid \mathbf{\theta} \right)}{\partial x^2} \right)^2 .
    • IC loss:
    L_{IC} = \frac{1}{| \mathcal D_{IC} |} \sum_{\left( t, x \right) \in \mathcal D_{IC}} \left( U \left( t, x \mid \mathbf{\theta} \right) - u \left( t, x \right) \right)^2 .
    • BC loss:
    L_{BC} = \frac{1}{| \mathcal D_{BC} |} \sum_{\left( t, x \right) \in \mathcal D_{BC}} \left( U \left( t, x \mid \mathbf{\theta} \right) - u \left( t, x \right) \right)^2 .

Part 2 (10 points, coding task)

In this part, you are asked to build a deep neural network that is used to output PDE solutions.

  1. The class name is HeatPINN.

    It subclasses nn.Module.

  2. The model consists of the following layers that are sequentially conntected:

    (1) Fully connected layer with out_features = 64 (you need to determine in_features taken from the input).

    (2) Activation layer with tanh function.

    (3) Fully connected layer with in_features = 64 and out_features = 64.

    (4) Activation layer with tanh function.

    (5) Fully connected layer with in_features = 64 (you need to determine out_features as the output of the entire model).

  3. Construct a model who is an object of this class and is called as model.

2 
### WRITE YOUR SOLUTION HERE ###

class HeatPINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )

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

model = HeatPINN()

""" END OF THIS PART """