<jupyter_start><jupyter_text>Define neural network<jupyter_code># Model
import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):

    def __init__(self, nof=50, dim=2, num_layers=4):
        super().__init__()
        self.dim = dim
        self.input_layer = nn.Linear(dim, nof)
        
        def f():
          return nn.Sequential(nn.Linear(nof, nof), nn.Tanh())
        self.hidden_layer = nn.Sequential(
            *[f() for _ in range(num_layers)],
        )
        self.output_layer = nn.Linear(nof, 1)


    def forward(self, x, y):
        input = torch.cat((x, y), dim=-1)
        input = torch.tanh(self.input_layer(input))
        input = self.hidden_layer(input)
        input = self.output_layer(input)
        return input<jupyter_output><empty_output><jupyter_text>Initialize data and forcing function<jupyter_code># Initialize data
gridsize = 100
x = torch.linspace(0, 1, gridsize, device=device)
X, Y = torch.meshgrid(x, x)
X = X.to(device)
Y = Y.to(device)
f = torch.zeros_like(X)

# load f data into torch
X = X.reshape((-1, gridsize, 1))
Y = Y.reshape((-1, gridsize, 1))
f = f.reshape((-1, gridsize, 1)).to(device)
Xc = X.clone().to(device)
Xc.requires_grad = True
Yc = Y.clone().to(device)
Yc.requires_grad = True

# load BC data into torch
bound_left = 0.0
bound_right = 0.0
bound_bottom = 0.0
bound_top = 0.0
bound_slope_top = np.pi * (np.exp(np.pi) + np.exp(-np.pi)) * torch.sin(np.pi * x)

def plot(*args):
  return map(lambda x : x.cpu().detach().numpy(), args)<jupyter_output><empty_output><jupyter_text>Loss functions for different pdes<jupyter_code># loss functions
def pde_loss_poisson(pred, x, y, f, lambda_pde):
    du_dx = torch.autograd.grad(pred.sum(), x, create_graph=True)[0]
    du2_dx2 = torch.autograd.grad(du_dx.sum(), x, create_graph=True)[0]
    du_dy = torch.autograd.grad(pred.sum(), y, create_graph=True)[0]
    du2_dy2 = torch.autograd.grad(du_dy.sum(), y, create_graph=True)[0]
    loss = torch.sum(torch.square(du2_dx2 + du2_dy2 - f))
    return lambda_pde * loss / (gridsize * gridsize)

def boundary_loss_dirichlet_neumann_mixed(
    pred, x, y,
    bound_left, bound_right, 
    bound_bottom, bound_slope_top, 
    lambda_bound=1.0):
    du_dx = torch.autograd.grad(pred.sum(), x, create_graph=True)[0]
    du_dy = torch.autograd.grad(pred.sum(), y, create_graph=True)[0]
    loss_dirichlet = torch.sum(
        (pred[0, :, 0] - bound_left) ** 2 + (pred[:, 0, 0] - bound_bottom) ** 2 + (pred[-1, :, 0] - bound_right) ** 2)
    loss_neumann = torch.sum((du_dy[:, -1, 0] - bound_slope_top) ** 2)
    loss = loss_dirichlet + loss_neumann
    return lambda_bound * loss / (4.0 * gridsize)<jupyter_output><empty_output><jupyter_text>Define analytical solutions for comparison<jupyter_code>def generate_ana_sol(X, Y):
    ana_sol = 0.0 * X
    ana_sol = (torch.exp(np.pi * Y) - torch.exp(-np.pi * Y)) * torch.sin(np.pi * X)
    return ana_sol

# define network
net = Net(nof=50, num_layers=4).to(device)

# define optimizer
learning_rate = 0.005
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
# train
num_epochs = 10000
lambda_pde = 1.0
lambda_bound = 5.0
log = {"loss_pde" : [],
       "loss_bound" : [],
       "loss_total" : []}
for epoch in range(num_epochs):
    if epoch % 1000 == 0:
        learning_rate *= 0.99

    for param_group in optimizer.param_groups:
        param_group['lr'] = learning_rate

    net.train()

    upred = net.forward(Xc, Yc)

    loss_bound = boundary_loss_dirichlet_neumann_mixed(
        upred, Xc, Yc, bound_left, bound_right, bound_bottom, bound_slope_top, lambda_bound=lambda_bound)
    loss_pde = pde_loss_poisson(upred, Xc, Yc, f, lambda_pde)
    loss = loss_pde + loss_bound
    log["loss_pde"].append(loss_pde.item())
    log["loss_bound"].append(loss_bound.item())
    log["loss_total"].append(loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 1000 == 0:
        print(f"Epoch {epoch} --- Loss : {loss}, pde_loss : {loss_pde}, bound_loss : {loss_bound}")

print(loss)
plt.semilogy(log["loss_pde"], 'r', label = 'loss_pde')
plt.semilogy(log["loss_bound"], 'g', label = 'loss_bound')
plt.semilogy(log["loss_total"], 'b', label = 'loss_total')
plt.legend()
# Saving log file because plotly doesnt work on colab for some reason
import pickle

with open('loss_laplace_rect.pickle', 'wb') as handle:
    pickle.dump(log, handle)<jupyter_output><empty_output><jupyter_text>Compare results<jupyter_code>ana_sol = generate_ana_sol(X, Y)
pinn_sol = net.forward(Xc, Yc)
fig = plt.figure(figsize=(6,8))
ax1 = plt.subplot(211)
cntr1 = ax1.contourf(*plot(X[:, :, 0], Y[:, :, 0], pinn_sol[:, :, 0]), levels=100)
fig.colorbar(cntr1, ax=ax1)
ax2 = plt.subplot(212)
cntr2 = ax2.contourf(*plot(X[:, :, 0], Y[:, :, 0], ana_sol[:, :, 0]), levels=100)
fig.colorbar(cntr2, ax=ax2)

plt.figure()
error = torch.abs(pinn_sol - ana_sol) / torch.max(torch.abs(pinn_sol))
plt.contourf(*plot(X[:,:,0], Y[:,:,0], (error[:,:,0]) + 1e-12), cmap="inferno", levels=100)
plt.colorbar()

plt.figure()
plt.plot(*plot(x, pinn_sol[gridsize // 2, :, 0]), '--', label="x")
plt.plot(*plot(x, pinn_sol[:, gridsize // 2, 0]), '--', label="y")
plt.plot(*plot(x, ana_sol[gridsize // 2, :, 0]), label="anal_x")
plt.plot(*plot(x, ana_sol[:, gridsize // 2, 0]), label="anal_y")
plt.legend()<jupyter_output><empty_output><jupyter_text>Generate solution and error contours<jupyter_code>ana_sol = generate_ana_sol(X, Y)
pinn_sol = net.forward(Xc, Yc)
error = torch.abs(pinn_sol - ana_sol) / torch.max(torch.abs(ana_sol)) * 100.0

import numpy.linalg as la
error_l2 = la.norm((error / 100.0).cpu().detach().numpy()) / np.sqrt(gridsize * gridsize)

import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman'] + plt.rcParams['font.serif']

import matplotlib.ticker as ticker
# specify your own formatter for colorbar labels in return select desired format
def myfmt(x, pos):
    return '{0:.1f}'.format(x)

fig = plt.figure(figsize=(6,8))
ax1 = plt.subplot(211)
cntr1 = ax1.contourf(*plot(X[:, :, 0], Y[:, :, 0], pinn_sol[:, :, 0]), cmap="Spectral", levels=100)
cbar1 = fig.colorbar(cntr1, ax=ax1, label='u(x, y)', format=ticker.FuncFormatter(myfmt))
# ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_title('Solution')
# This is the fix for the white lines between contour levels
for c in cntr1.collections:
    c.set_edgecolor("face")

ax2 = plt.subplot(212)
cntr2 = ax2.contourf(*plot(X[:,:,0], Y[:,:,0], (error[:,:,0]) + 1e-12), cmap="inferno", levels=100)
cbar2 = fig.colorbar(cntr2, ax=ax2, label='Error (%)', format=ticker.FuncFormatter(myfmt))
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_title('Error')
for c in cntr2.collections:
    c.set_edgecolor("face")

plt.tight_layout()

error_l2
fig.savefig('laplace_2d_rect.eps', dpi=300)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17290059671081356, 'fraction_numerical': 0.027943530781545628, 'mean_word_length': 4.389803921568627, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '167488', 'license': 'MIT', 'n_tokens_mistral': 2863, 'n_tokens_neox': 2738, 'n_words': 536, 'path': 'cases/laplace_2d.ipynb', 'repo_name': 'fankiat/CS498DL-project', 'size': '6871'}