Una guía paso a paso para resolver la ecuación de hamburguesas 1D con redes neuronales informadas por física (PINN): un enfoque de Pytorch utilizando métodos automáticos de diferenciación y colocación

En este tutorial, exploramos un enfoque innovador que combina aprendizaje profundo Con las leyes físicas aprovechando las redes neuronales informadas por física (PINN) para resolver la ecuación de hamburguesas unidimensionales. Usando Pytorch en Google Colab, demostramos cómo codificar la ecuación diferencial de gobierno directamente en la función de pérdida de la red neuronal, lo que permite que el modelo aprenda la solución 𝑢 (𝑥, 𝑡) que respeta inherentemente la física subyacente. Esta técnica reduce la dependencia de grandes conjuntos de datos etiquetados y ofrece una nueva perspectiva para resolver ecuaciones diferenciales parciales complejas y no lineales utilizando herramientas computacionales modernas.

!pip install torch matplotlib

Primero, instalamos las bibliotecas Pytorch y Matplotlib utilizando PIP, asegurando que tenga las herramientas necesarias para construir redes neuronales y visualizar los resultados en su entorno de Google Colab.

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


torch.set_default_dtype(torch.float32)

Importamos bibliotecas esenciales: Pytorch para el aprendizaje profundo, Numpy para operaciones numéricas y matplotlib para trazar. Establecimos el tipo de datos de tensor predeterminado en Float32 para una precisión numérica consistente a lo largo de sus cálculos.

x_min, x_max = -1.0, 1.0
t_min, t_max = 0.0, 1.0
nu = 0.01 / np.pi


N_f = 10000  
N_0 = 200    
N_b = 200    


X_f = np.random.rand(N_f, 2)
X_f[:, 0] = X_f[:, 0] * (x_max - x_min) + x_min  # x in [-1, 1]
X_f[:, 1] = X_f[:, 1] * (t_max - t_min) + t_min    # t in [0, 1]


x0 = np.linspace(x_min, x_max, N_0)[:, None]
t0 = np.zeros_like(x0)
u0 = -np.sin(np.pi * x0)


tb = np.linspace(t_min, t_max, N_b)[:, None]
xb_left = np.ones_like(tb) * x_min
xb_right = np.ones_like(tb) * x_max
ub_left = np.zeros_like(tb)
ub_right = np.zeros_like(tb)


X_f = torch.tensor(X_f, dtype=torch.float32, requires_grad=True)
x0 = torch.tensor(x0, dtype=torch.float32)
t0 = torch.tensor(t0, dtype=torch.float32)
u0 = torch.tensor(u0, dtype=torch.float32)
tb = torch.tensor(tb, dtype=torch.float32)
xb_left = torch.tensor(xb_left, dtype=torch.float32)
xb_right = torch.tensor(xb_right, dtype=torch.float32)
ub_left = torch.tensor(ub_left, dtype=torch.float32)
ub_right = torch.tensor(ub_right, dtype=torch.float32)

Establecemos el dominio de simulación para la ecuación de las hamburguesas definiendo límites espaciales y temporales, viscosidad y el número de puntos de colocación, iniciales y límites. Luego genera puntos de datos aleatorios y espaciados uniformemente para estas condiciones y los convierte en tensores de pytorch, lo que permite el cálculo de gradiente donde sea necesario.

class PINN(nn.Module):
    def __init__(self, layers):
        super(PINN, self).__init__()
        self.activation = nn.Tanh()
       
        layer_list = []
        for i in range(len(layers) - 1):
            layer_list.append(nn.Linear(layers[i], layers[i+1]))
        self.layers = nn.ModuleList(layer_list)
       
    def forward(self, x):
        for i, layer in enumerate(self.layers[:-1]):
            x = self.activation(layer(x))
        return self.layers[-1](x)


layers = [2, 50, 50, 50, 50, 1]
model = PINN(layers)
print(model)

Aquí, definimos una red neuronal informada por física personalizada (PINN) al extender el nn.module de Pytorch. La arquitectura de red se crea dinámicamente utilizando una lista de tamaños de capa, donde cada capa lineal es seguida por una activación de Tanh (excepto la capa de salida final). En este ejemplo, la red toma una entrada bidimensional, la pasa a través de cuatro capas ocultas (cada una con 50 neuronas) y genera un valor único. Finalmente, el modelo está instanciado con la arquitectura especificada, y su estructura está impresa.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

Aquí, verificamos si hay una GPU habilitada para CUDA disponible, establecer el dispositivo en consecuencia y mover el modelo a ese dispositivo para obtener un cálculo acelerado durante el entrenamiento y la inferencia.

def pde_residual(model, X):
    x = X[:, 0:1]
    t = X[:, 1:2]
    u = model(torch.cat([x, t], dim=1))
   
    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True, retain_graph=True)[0]
    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True, retain_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True, retain_graph=True)[0]
   
    f = u_t + u * u_x - nu * u_xx
    return f


def loss_func(model):
    f_pred = pde_residual(model, X_f.to(device))
    loss_f = torch.mean(f_pred**2)
   
    u0_pred = model(torch.cat([x0.to(device), t0.to(device)], dim=1))
    loss_0 = torch.mean((u0_pred - u0.to(device))**2)
   
    u_left_pred = model(torch.cat([xb_left.to(device), tb.to(device)], dim=1))
    u_right_pred = model(torch.cat([xb_right.to(device), tb.to(device)], dim=1))
    loss_b = torch.mean(u_left_pred**2) + torch.mean(u_right_pred**2)
   
    loss = loss_f + loss_0 + loss_b
    return loss

Ahora, calculamos el residuo de la ecuación de hamburguesas en los puntos de colocación calculando las derivadas requeridas a través de la diferenciación automática. Luego, definimos una función de pérdida que agregue la pérdida residual de PDE, el error de la condición inicial y los errores de las condiciones de contorno. Esta pérdida combinada guía a la red para aprender una solución que satisfaga tanto la ley física como las condiciones impuestas.

optimizer = optim.Adam(model.parameters(), lr=1e-3)
num_epochs = 5000


for epoch in range(num_epochs):
    optimizer.zero_grad()
    loss = loss_func(model)
    loss.backward()
    optimizer.step()
   
    if (epoch+1) % 500 == 0:
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.5e}')
       
print("Training complete!")

Aquí, configuramos el bucle de entrenamiento de PinN usando el Adam Optimizer con una tasa de aprendizaje de 1 × 10⁻³. Más de 5000 épocas, calcula repetidamente la pérdida (que incluye los errores de condición residuales, iniciales y de límite de PDE), retrocede los gradientes y actualiza los parámetros del modelo. Cada 500 épocas, imprime la época actual y la pérdida para monitorear el progreso y finalmente anuncia cuándo se completa la capacitación.

N_x, N_t = 256, 100
x = np.linspace(x_min, x_max, N_x)
t = np.linspace(t_min, t_max, N_t)
X, T = np.meshgrid(x, t)
XT = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
XT_tensor = torch.tensor(XT, dtype=torch.float32).to(device)


model.eval()
with torch.no_grad():
    u_pred = model(XT_tensor).cpu().numpy().reshape(N_t, N_x)


plt.figure(figsize=(8, 5))
plt.contourf(X, T, u_pred, levels=100, cmap='viridis')
plt.colorbar(label="u(x,t)")
plt.xlabel('x')
plt.ylabel('t')
plt.title("Predicted solution u(x,t) via PINN")
plt.show()

Finalmente, creamos una cuadrícula de puntos sobre el dominio espacial definido (𝑥) y temporal (𝑡), alimentamos estos puntos al modelo entrenado para predecir la solución 𝑢 (𝑥, 𝑡) y remodelar la salida en una matriz 2D. Además, visualiza la solución predicha como una gráfica de contorno usando matplotlib, completo con una barra de colores, etiquetas del eje y un título, lo que le permite observar cómo el PINN ha aproximado la dinámica de la ecuación de las hamburguesas.

En conclusión, este tutorial ha mostrado cómo se pueden implementar efectivamente PinN para resolver la ecuación de hamburguesas 1D incorporando la física del problema en el proceso de capacitación. A través de la construcción cuidadosa de la red neuronal, la generación de datos de colocación y límites, y la diferenciación automática, logramos un modelo que aprende una solución consistente con el PDE y las condiciones prescritas. Esta fusión de aprendizaje automático Y la física tradicional allana el camino para abordar problemas más desafiantes en la ciencia y la ingeniería computacionales, invitando a una mayor exploración a sistemas de mayor dimensión y arquitecturas neuronales más sofisticadas.

Aquí está el Cuaderno de colab. Además, no olvides seguirnos Gorjeo y únete a nuestro Canal de telegrama y LinkedIn GRsalpicar. No olvides unirte a nuestro 85k+ ml de subreddit.