Neural Network Differential Equations

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

# Neural network model
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(1, 50),
            nn.Tanh(),
            nn.Linear(50, 1)
        )

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

# Differential equation dy/dx = -y (loss function)
def ode_loss(x, model):
    y = model(x)
    dy_dx = torch.autograd.grad(y, x, grad_outputs=torch.ones_like(x), create_graph=True)[0]
    return torch.mean((dy_dx + y)**2)

x_train = torch.linspace(0, 2, 100).view(-1, 1).requires_grad_(True)
model = Net()
optimizer = optim.Adam(model.parameters(), lr=0.01)

for epoch in range(2000):
    optimizer.zero_grad()
    loss = ode_loss(x_train, model)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

# Plot result
with torch.no_grad():
    x_test = torch.linspace(0, 2, 100).view(-1, 1)
    y_pred = model(x_test)
    plt.plot(x_test, y_pred, label="NN solution")
    plt.plot(x_test, torch.exp(-x_test), label="Exact solution")
    plt.legend()
    plt.show()