<jupyter_start><jupyter_text>$$ u_{tt}  - u_{xx} + u ^{3} = f(x,t) $$<jupyter_code>pi = torch.tensor(np.pi)
def ue(x,t):
    return x * torch.cos(5. * pi * t) + (x * t).pow(3) 

def ue_t(x,t):
    return - 5.0 * pi * x * torch.sin(5. * pi * t) + 3.0 * x * (x*t).pow(2) 

def ue_tt(x,t):
    return - (5.0 * pi).pow(2) * x * torch.cos(5. * pi * t) + 6.0 * x.pow(2) * (x*t)

def ue_x(x,t):
    return torch.cos(5. * pi * t) + 3.0 * t* (x * t).pow(2)

def ue_xx(x,t): 
    return 6.0 * t.pow(2) * (x * t)

def source(x,t):
    return ue_tt(x,t) - ue_xx(x,t) + ue(x,t).pow(3) 
kwargs ={"mean":0.50, "stdev":0.2887}  
domain = np.array([[0.,0.],[1.,1.]])
model  = Network(in_N=2,m=50,H_Layer=5,out_N=1,**kwargs)
model.to(device)
model.apply(init_weights)
print(model)
print(model.mu)
print(model.std)
def evaluate(model,domain):
    model.eval()
    x = np.linspace(domain[0][0],domain[1][0],64)
    t = np.linspace(domain[0][1],domain[1][1],64)

    x_star,t_star = np.meshgrid(x,t)
    x_test = torch.from_numpy(x_star.flatten()[:,None])
    t_test = torch.from_numpy(t_star.flatten()[:,None])

    u_star = ue(x_test,t_test)   
    u_pred = model(x_test,t_test)
    
    l2   = np.linalg.norm(u_star- u_pred.detach(), 2)/np.linalg.norm(u_star, 2)
    linf = max(abs(u_star- u_pred.detach().numpy())).item()
    return l2,linf
epochs          = 1500
disp            = 500
print_to_consol = True

model.apply(init_weights)

# update the optimizer
optimizer = LBFGS(model.parameters(),line_search_fn="strong_wolfe")

# initialize penalty parameter
mu           = torch.tensor(1.0)

# maximum penalty value for safeguarding
mu_max      = torch.tensor(1e4)

# l2 norm of constraints |C|_2
eta          = torch.tensor(0.0)

# penalty tolerance
epsilon      = torch.tensor(1e-8)

# number of collocation points in the domain
N_data     = 512

# collocation points
x_dm,t_dm  = fetch_interior_data(domain,N_data)
x_dm       = x_dm.requires_grad_(True)
t_dm       = t_dm.requires_grad_(True)


N_u        = 256
# boundary points 
x_bc,t_bc  = fetch_boundary_data(domain,N_u)

N_u0       = 128

# 1st initial condition points
x_ic,t_ic  = fetch_init_data(domain,N_u0)

N_ut       = 128

# 2nd intial condition points 
x_nic,t_nic = fetch_init_data(domain,N_ut)
x_nic       = x_nic.requires_grad_(True)
t_nic       = t_nic.requires_grad_(True)

# lagrange multiplier for the boundary constraints
Lambda_u  = torch.zeros_like(x_bc)

# lagrange multiplier for the initial constraints
Lambda_ut = torch.zeros_like(x_nic)

Lambda_u0 = torch.zeros_like(x_ic)

# starting to train neural network model
for epoch in range(epochs):
    def closure():
        if torch.is_grad_enabled():
            model.train()
            optimizer.zero_grad()

        pde_loss     = physics_loss(model,x_dm,t_dm)
        u_loss       = boundary_loss(model,x_bc,t_bc)
        u0_loss      = first_init_loss(model,x_ic,t_ic)
        ut_loss      = second_init_loss(model,x_nic,t_nic)
        

        penalty      = (u_loss.T@ u_loss + \
                           ut_loss.T@ut_loss + \
                           u0_loss.T@u0_loss).sum()

        loss         = pde_loss.sum() + \
                        (Lambda_u  * u_loss).sum() + \
                        (Lambda_ut * ut_loss).sum() + \
                        (Lambda_u0 * u0_loss).sum() +\
                        0.5 * mu   * penalty

        if loss.requires_grad:
            loss.backward()
        return loss

    def _closure():
        model.eval()
        u_loss       = boundary_loss(model,x_bc,t_bc)
        u0_loss      = first_init_loss(model,x_ic,t_ic)
        ut_loss      = second_init_loss(model,x_nic,t_nic)

        penalty      =  (u_loss.T@ u_loss + \
                        ut_loss.T@ut_loss + \
                        u0_loss.T@u0_loss).sum()

        return u_loss, ut_loss, u0_loss, penalty

    optimizer.step(closure)
    u_loss, ut_loss, u0_loss,penalty = _closure()
    with torch.no_grad():
        if (torch.sqrt(penalty) >= 0.25*eta) and (torch.sqrt(penalty) > epsilon):
            mu          = min(mu*2.0, mu_max)
            Lambda_u   += mu*u_loss
            Lambda_ut  += mu*ut_loss
            Lambda_u0  += mu*u0_loss

        eta  = torch.sqrt(penalty)
    if (epoch + 1)%disp == 0 and print_to_consol:
        print(f'epoch : {epoch+1:3d}, penalty loss: {penalty.item():2.3e}')


# checkpointing the model 
torch.save(model.state_dict(),f"klein_gordon.pt")
# evaluation
l2,linf = evaluate(model,domain)
# print the l2 norms
print(f"relative l2 error :{l2:2.3e}, linf error : {linf :2.3e}")
# https://joseph-long.com/writing/colorbars/
def colorbar(mappable,min_val,max_val):
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    import matplotlib.pyplot as plt
    last_axes = plt.gca()
    ax = mappable.axes
    fig = ax.figure
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.1)
    ticks = np.linspace(min_val, max_val, 4, endpoint=True)
    cbar = fig.colorbar(mappable, cax=cax,ticks=ticks)
    cbar.formatter.set_powerlimits((0, 0))
    plt.sca(last_axes)
    return cbar

params = {
    'text.latex.preamble': '\\usepackage{gensymb}',
    'image.origin': 'lower',
    'image.interpolation': 'nearest',
    'image.cmap': 'gray',
    'axes.grid': False,
    'savefig.dpi': 150,  # to adjust notebook inline plot size
    'axes.labelsize': 16, # fontsize for x and y labels
    'axes.titlesize': 16,
    'font.size': 16, 
    'legend.fontsize': 16, 
    'xtick.labelsize': 16,
    'ytick.labelsize': 16,
    'text.usetex': False,
    'figure.figsize': [20, 3],
    'font.family': 'serif',
}
plt.rcParams.update(params)
model.eval()
x = np.linspace(domain[0][0],domain[1][0],100)
t = np.linspace(domain[0][1],domain[1][1],100)

x_star,t_star = np.meshgrid(x,t)
x_test = torch.from_numpy(x_star.flatten()[:,None])
t_test = torch.from_numpy(t_star.flatten()[:,None])

u_star = ue(x_test,t_test)   
u_pred = model(x_test,t_test)
# ------------
# 2D figures
# --------------
cmap_list = ['jet','YlGnBu','coolwarm','rainbow','magma','plasma','inferno','Spectral','RdBu']
cmap = cmap_list[4]
gs = gridspec.GridSpec(1, 3)
gs.update(wspace=0.3)

points      = np.concatenate((x_star.flatten()[:,None],t_star.flatten()[:,None]),axis=1)
u_pred_plot = griddata(points, u_pred.detach().flatten(), (x_star,t_star), method='cubic')
u_star_plot = griddata(points, u_star.flatten(), (x_star,t_star), method='cubic')


save = True
#################################### Predicted Solution #####################################
ax = plt.subplot(gs[0,0])
min_val = np.min(u_star_plot)
max_val = np.amax(u_star_plot)
img = ax.pcolormesh(t_star,x_star,u_pred_plot, cmap = cmap,vmin=min_val,vmax=max_val,shading='gouraud')
# ax.set_title('$\hat{u}(x,y)$')
ax.set_xlabel('$x$')
ax.set_ylabel('$t$')
ax.axis('square')
cbar = colorbar(img,min_val,max_val)
cbar.formatter.set_powerlimits((-1, -1))
ax.axis('square')
ax.set_xticks([domain[0][0],(domain[1][0]+ domain[0][0])/2,domain[1][0]])
ax.set_yticks([domain[0][1],(domain[1][1]+domain[0][1])/2,domain[1][1]])

    
# #################################### Exact Solution #########################################
ax = plt.subplot(gs[0,1])
img = ax.pcolormesh(t_star,x_star,u_star_plot, cmap = cmap,vmin=min_val,vmax=max_val,shading='gouraud')
# ax.set_title('$u(x,y)$')
ax.set_xlabel('$x$')
ax.set_ylabel('$t$')
cbar = colorbar(img,min_val,max_val)
cbar.formatter.set_powerlimits((-1, -1))
ax.axis('square')
ax.set_xticks([domain[0][0],(domain[1][0]+ domain[0][0])/2,domain[1][0]])
ax.set_yticks([domain[0][1],(domain[1][1]+domain[0][1])/2,domain[1][1]])


# #################################### Absolute Error #########################################
ax = plt.subplot(gs[0,2])
img = ax.pcolormesh(t_star,x_star,np.abs(u_star_plot - u_pred_plot), cmap = cmap,shading='gouraud')
# ax.set_title('$|u - \hat{u}|$')
ax.set_xlabel('$x$')
ax.set_ylabel('$t$')

min_val = np.amin(np.abs(u_star_plot - u_pred_plot))
max_val = np.amax(np.abs(u_star_plot - u_pred_plot))
cbar = colorbar(img,min_val,max_val)
cbar.formatter.set_powerlimits((0, 0))
ax.axis('square')
ax.set_xticks([domain[0][0],(domain[1][0]+ domain[0][0])/2,domain[1][0]])
ax.set_yticks([domain[0][1],(domain[1][1]+domain[0][1])/2,domain[1][1]])


filename="pecanns_klein_gordon"
plt.figtext(0.229, -0.25,'(a)' ,wrap=True, horizontalalignment='center',fontsize=20)
plt.figtext(0.508, -0.25,'(b)', wrap=True, horizontalalignment='center', fontsize=20)
plt.figtext(0.790, -0.25,'(c)', wrap=True, horizontalalignment='center', fontsize=20)
plt.savefig('{}.png'.format(filename), bbox_inches='tight', pad_inches=0.02)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19914470642625984, 'fraction_numerical': 0.030744336569579287, 'mean_word_length': 3.6124733475479744, 'pattern_counts': {'":': 2, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 1, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '606462', 'license': 'MIT', 'n_tokens_mistral': 3648, 'n_tokens_neox': 3517, 'n_words': 525, 'path': 'Klein_Gordon/pecann_klein_gordon.ipynb', 'repo_name': 'shamsbasir/PECANN', 'size': '8652'}