<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_x = 100
gridsize_y = 100
x = torch.linspace(0, 1, gridsize_x, device=device)
y = torch.linspace(0, 1, gridsize_y, device=device)
X, Y = torch.meshgrid(x, y)
X = X.to(device)
Y = Y.to(device)
num_modes = 3
# coef = np.random.rand(num_modes) * 10
coef = np.array([10.0, 1.0, 5.0])
f = torch.zeros_like(X)
for i in range(num_modes):
    f += coef[i] * torch.sin((i+1) * np.pi * X) * torch.sin((i+4) * np.pi * Y)

# load f data into torch
X = X.reshape((gridsize_x, gridsize_y, 1))
Y = Y.reshape((gridsize_x, gridsize_y, 1))
f = f.reshape((gridsize_x, gridsize_y, 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
u0 = torch.sin(2.0 * np.pi * y)

pe = 100

def plot(*args):
  return map(lambda x : x.cpu().detach().numpy(), args)

plt.plot(*plot(y, u0))<jupyter_output>torch.Size([100, 100, 1])<jupyter_text>Loss functions for different pdes<jupyter_code># loss functions
def pde_loss_poisson(pred, x, y, pe, 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(du_dx + pred * du_dy - du2_dy2 / pe))
    return lambda_pde * loss / (gridsize_x * gridsize_y)

def boundary_loss_dirichlet(pred, bound_left, bound_right, u0, lambda_bound=1.0):
    loss = torch.sum((pred[:, 0, 0] - bound_left) ** 2 / gridsize_x + (pred[:, -1, 0] - bound_right) ** 2 / gridsize_x)
    loss += torch.sum((pred[0, :, 0] - u0) ** 2 / gridsize_y)
    return lambda_bound * loss
# define network
net = Net(nof=50, num_layers=4).to(device)

# define optimizer
learning_rate = 0.0025
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(upred, bound_left, bound_right, u0, lambda_bound=lambda_bound)
    loss_pde = pde_loss_poisson(upred, Xc, Yc, pe, 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 colab doesnt have plotly shit
import pickle

with open('loss_burgers_shock.pickle', 'wb') as handle:
    pickle.dump(log, handle)<jupyter_output><empty_output><jupyter_text>Preview results<jupyter_code>pinn_sol = net.forward(X, Y)
plt.contourf(*plot(X[:, :, 0], Y[:, :, 0], pinn_sol[:, :, 0]), levels=100)
plt.colorbar()

freq = 5
plt.figure()
for i in range(10):
  plt.plot(*plot(y, pinn_sol[i * 5, :, 0]))<jupyter_output><empty_output><jupyter_text>Compute FDM solution for comparison<jupyter_code>def fdm_sol_stationary_viscous_burgers(u0, pe, final_time):
    """
    FDM solution for viscous burgers on domain [0, 1]
    assumes u = 0.0 at x = 0.0, 1.0
    """
    from scipy.sparse.linalg import spsolve
    import scipy.sparse as spp
    import numpy.linalg as la
    
    fdm_grid_size = u0.shape[0] - 2 # takes 2 end pts in account
    dx = 1.0 / (fdm_grid_size + 1) # takes 2 end pts in account
    gradient = (spp.diags([-0.5, 0.5], [-1, 1], shape=(fdm_grid_size, fdm_grid_size), 
                          format="csr") / dx)
    laplacian = (spp.diags([1.0, -2.0, 1.0], [-1, 0, 1],
                           shape=(fdm_grid_size, fdm_grid_size), format="csr") / (dx ** 2))

    fdm_sol = u0[1:-1].copy()
    time = 0.0
    while time < final_time:
        dt = 0.01 * dx
        if time + dt > final_time:
            dt = final_time - time
        fdm_sol[...] += dt * (-fdm_sol * (gradient @ fdm_sol) + laplacian * fdm_sol / pe)
        time += dt
    
    final_fdm_sol = np.zeros_like(u0)
    final_fdm_sol[1:-1] = fdm_sol
    return final_fdm_sol

import numpy as np
import matplotlib.pyplot as plt

pe = 100.0
fdm_grid_size = gridsize_y
x_fdm = np.linspace(0.0, 1.0, fdm_grid_size)
u0 = np.sin(2 * np.pi * x_fdm)
final_time = np.linspace(0.0, 0.5, 10)
# plt.plot(x_fdm, u0, '--')

fdm_sol = 0 * pinn_sol
final_time = np.linspace(0.0, 1.0, gridsize_x)
for i, t in enumerate(final_time):
    fdm_sol[i, ...] = torch.from_numpy(fdm_sol_stationary_viscous_burgers(u0, pe, t)).unsqueeze(-1)
    # plt.plot(x_fdm, fdm_sol)

plt.contourf(*plot(X[:, :, 0], Y[:, :, 0], fdm_sol[:, :, 0]), levels=100)
ana_sol = fdm_sol
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((pinn_sol - ana_sol).cpu().detach().numpy()) / np.sqrt(gridsize_x * gridsize_y)

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=20)
cbar1 = fig.colorbar(cntr1, ax=ax1, label='u(t, 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('t')
ax2.set_ylabel('y')
ax2.set_title('Error')
for c in cntr2.collections:
    c.set_edgecolor("face")

plt.tight_layout()

error_l2
fig.savefig('burgers_shock.eps', dpi=300)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16336633663366337, 'fraction_numerical': 0.03224511640353225, 'mean_word_length': 3.8288113695090438, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '199867', 'license': 'MIT', 'n_tokens_mistral': 3169, 'n_tokens_neox': 3024, 'n_words': 631, 'path': 'cases/1d_burgers_time_shock.ipynb', 'repo_name': 'fankiat/CS498DL-project', 'size': '7474'}