Building a Physics-Informed Neural Network (PINN) Using PyTorch From Scratch

Solving the 1D Heat Equation

Image created by the author using a generative AI tool

Physics-Informed Neural Networks (PINNs) represent a breakthrough in applying deep learning to solve differential equations. These neural networks learn solutions that adhere to physical laws rather than just fitting data, which makes them especially useful in scientific and engineering applications.

In this article, I’ll build a PINN to approximate the solution to the 1D heat equation, a widely-used partial differential equation (PDE) describing the diffusion of heat over time. I’ll explain each part of the code, understand its purpose, and visualize the results using Python and PyTorch.

Why Use a PINN?

Traditional numerical methods for solving PDEs, like finite element methods, require discretization of the domain and suffer from high computational costs as the dimensionality of the problem grows. PINNs offer a mesh-free alternative, making them computationally efficient and suitable for high-dimensional problems.

Understanding the 1D Heat Equation

The 1D heat equation is a PDE that models heat distribution over time:

here:

• u(x,t) is the temperature at position x and time t.
• α is the thermal diffusivity constant, determining the rate of heat transfer.

Our goal is to approximate the temperature u(x,t) using a neural network trained to satisfy this equation.

Steps to Build and Train the PINN

I’ll create a neural network that will:

1. Predict the temperature u(x,t) at any point x and time t.
2. Learn solutions that satisfy the 1D heat equation by minimizing the residual error.

Let’s go through the code in detail.

Import Libraries:

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

I’ll use PyTorch for building the neural network, calculating derivatives, and optimizing the model. Additionally, **matplotlib **will help us visualize the results.

Define the Physics-Informed Neural Network (PINN):

import torch
import torch.nn as nn

class HeatPINN(nn.Module):
    """
    A neural network model for solving the 1D heat equation using the Physics-Informed Neural Network (PINN) approach.

    Attributes:
        model (nn.Sequential): The neural network architecture for the PINN.
    """
    def __init__(self):
        super(HeatPINN, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(2, 64),  # Input layer: 2 inputs (x, t) to 64 neurons
            nn.Tanh(),
            nn.Linear(64, 64), # Hidden layer
            nn.Tanh(),
            nn.Linear(64, 1)   # Output layer: 1 output (temperature u)
        )

    def forward(self, x, t):
        """
        Forward pass through the neural network.

        Args:
            x (torch.Tensor): The spatial coordinate tensor.
            t (torch.Tensor): The time coordinate tensor.

        Returns:
            torch.Tensor: Predicted temperature `u(x, t)` tensor.
        """
        inputs = torch.cat((x, t), dim=1)
        u = self.model(inputs)
        return u

This neural network takes two inputs, **x (position) and `t** (**time**), and outputs the predicted temperatureu(x, t)`.

The model is simple, with three fully connected layers and **Tanh activation functions. I used `torch.cat** to combinex` and **t** as a single input for the network. This setup is flexible enough to learn complex, smooth functions that satisfy the PDE requirements.

Define the Loss Function Based on the Heat Equation:

The core of a PINN lies in its loss function, which incorporates the physics of the problem, in this case, the heat equation. I define a loss function that penalizes the model if it produces solutions that do not satisfy the equation. PyTorch’s **autograd **feature allows us to compute derivatives required by the heat equation.

def heat_loss_fn(model, x, t, alpha):
    """
    Compute the physics-informed loss based on the heat equation residual.

    Args:
        model (HeatPINN): The PINN model.
        x (torch.Tensor): The spatial coordinate tensor.
        t (torch.Tensor): The time coordinate tensor.
        alpha (float): The thermal diffusivity constant.

    Returns:
        torch.Tensor: Mean squared loss of the heat equation residual.
    """
    u = model(x, t)
    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    residual = u_t - alpha * u_xx
    return torch.mean(residual ** 2)

Train the Model:

def train_pinn(model, optimizer, x_train, t_train, alpha, epochs=1000):
    """
    Train the PINN model to solve the heat equation.

    Args:
        model (HeatPINN): Instance of the PINN model.
        optimizer (torch.optim.Optimizer): Optimizer to update the model weights.
        x_train (torch.Tensor): Spatial training coordinates.
        t_train (torch.Tensor): Temporal training coordinates.
        alpha (float): Thermal diffusivity constant.
        epochs (int): Number of training epochs.

    Returns:
        list: Training losses recorded at each epoch.
    """
    losses = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        loss = heat_loss_fn(model, x_train, t_train, alpha)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        if epoch % 100 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item()}")
    return losses

The train_pinn function iterates through training epochs, calculating the physics-informed loss at each step. The loss is backpropagated to adjust model weights. I’ll record and print the loss every 100 epochs to monitor training progress.

Visualization Function:

def plot_results(x_train, t_train, u_pred, losses):
    plt.figure(figsize=(10, 4))
    plt.plot(losses, label="Training Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.title("Training Loss Over Epochs")
    plt.legend()
    plt.show()

    plt.figure(figsize=(10, 4))
    plt.scatter(t_train.detach().numpy(), u_pred.detach().numpy(), color='blue', label="Predicted u(x, t)")
    plt.xlabel("Time (t)")
    plt.ylabel("Temperature u(x, t)")
    plt.title("Predicted Temperature over Time")
    plt.legend()
    plt.show()

Training Loss Plot shows how the physics-informed loss decreases over time, indicating how well the model is learning. And the Predicted Solution plots the predicted temperature values for ***u(x, t)***, giving a visual of the heat distribution over time. This helps validate the model’s performance.

Run the Complete Model:

# Define parameters and synthetic data
alpha = 0.1  # thermal diffusivity
x_train = torch.rand(100, 1, requires_grad=True)
t_train = torch.rand(100, 1, requires_grad=True)

# Initialize model and optimizer
model = HeatPINN()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Train the model
losses = train_pinn(model, optimizer, x_train, t_train, alpha)

# Get predictions for visualization
u_pred = model(x_train, t_train)

# Plot the results
plot_results(x_train, t_train, u_pred, losses)

I have generated random **x_train and `t_train**values to act as our training points. These points are sufficient for the network to learn an approximate solution. The model is trained using thetrain_pinn` function. After training, I have visualized both the training loss and the predicted temperature values, allowing us to assess the solution’s accuracy.

Outputs:

In this article, I’ve built a simple but effective Physics-Informed Neural Network (PINN) to solve the 1D heat equation. Through automatic differentiation and physics-informed loss, the network learned a solution that approximates the heat equation’s behaviour without needing explicit temperature data. PINNs are a powerful tool in scientific computing, offering flexible solutions for complex physical systems. While this example focused on a basic 1D problem, PINNs can be extended to higher dimensions and more complex equations, opening up new possibilities in fields like fluid dynamics, electromagnetics, and more.