<jupyter_start><jupyter_text>Implementing an Physics-informed neural network for the 1D Schrodinger equation using the PINN framework  Including necessary libaries<jupyter_code>import os
#print(os.getcwd())
import sys
sys.path.append('..') # examples
sys.path.append('../..') # PINNFramework etc.
from PINNFramework.PINN import Interface
from PINNFramework.models.mlp import MLP
from torch.autograd import grad
import torch
import numpy as np
import torch.nn as nn
import scipy.io
from pyDOE import lhs
import torch.optim as optim<jupyter_output><empty_output><jupyter_text>Underlying PDE$f:=i h_{t}+0.5 h_{x x}+|h|^{2} h$  Implementing needed functions<jupyter_code>class SchrodingerPINN(Interface):
    def __init__(self, model, input_d, output_d, lb, ub):
        super().__init__(model,input_d,output_d)
        self.lb = lb
        self.ub = ub
        
    def pde(self, x, u, derivatives):
        u_xx = derivatives[:,0]
        v_xx = derivatives[:,1]
        _u = u[:,0]
        _v = u[:,1]
        real_part = - 0.5 * v_xx - (_u**2 - _v**2)*_v
        imaginary_part= 0.5 * u_xx + (_u**2 + _v**2)*_u 
        result = torch.stack([real_part,imaginary_part],1)
        return result
        
    def derivatives(self, u, x):
        grads= torch.ones(x.shape[0])
        pred_u = u[:,0]
        pred_v = u[:,1]
        J_u = grad(pred_u, x, create_graph=True, grad_outputs=grads)[0]
        J_v = grad(pred_v, x, create_graph=True, grad_outputs=grads)[0]
        
        #calculate first order derivatives
        u_x = J_u[:,0]
        u_t = J_u[:,1]

        v_x = J_v[:,0]
        v_t = J_v[:,1]
        
        # calculate second order derivatives
        J_u_x = grad(u_x, x, create_graph=True, grad_outputs=grads)[0]
        J_v_x = grad(v_x, x, create_graph=True, grad_outputs=grads)[0]

        u_xx = J_u_x[:,0]
        v_xx = J_v_x[:,0]
        pred_derivatives = torch.stack([u_xx,v_xx,u_t,v_t],1)
        return pred_derivatives
    
        
    def input_normalization(self,x):
        """
        Implementation of min-max scaling in range of [-1,1]
        """
        return 2.0 * (x - self.lb) / (self.ub - self.lb) - 1.0<jupyter_output><empty_output><jupyter_text>Creating a model with the sequential API from torch<jupyter_code>#pinn_model = nn.Sequential(
#          nn.Linear(2,100),
#          nn.Tanh(),
#          nn.Linear(100,100),
#          nn.Tanh(),
#          nn.Linear(100,2)
#        )

pinn_model = MLP(input_size=2, output_size=2, num_hidden=3, hidden_size=100, activation=torch.tanh)
lb = torch.tensor([-5.0, 0.0])
ub = torch.tensor([[5.0, np.pi / 2]])
model = SchrodingerPINN(model = pinn_model, input_d = 2, output_d = 2, lb = lb, ub= ub)<jupyter_output><empty_output><jupyter_text>Testing forward function of the model (testing normalization)<jupyter_code>sample_x = torch.randn(100,2)
sample_pred = model(sample_x)
sample_pred.shape<jupyter_output><empty_output><jupyter_text>Preparing data (implement your dataset here)<jupyter_code>noise = 0.0
np.random.seed(1234)

# Doman bounds
lb = np.array([-5.0, 0.0])  # lower bound consists of [lower bound of x, lower bound of t]
ub = np.array([5.0, np.pi / 2])  # upper bound follows from lower bound

# defines the sizes of the neural network
N0 = 50
N_b = 50
N_f = 20000

data = scipy.io.loadmat('NLS.mat')


t = data['tt'].flatten()[:, None]  # get timestamps
x = data['x'].flatten()[:, None]  # get x positions
Exact = data['uu']
# definie labels
Exact_u = np.real(Exact)
Exact_v = np.imag(Exact)
Exact_h = np.sqrt(Exact_u ** 2 + Exact_v ** 2)

X, T = np.meshgrid(x, t)

X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))  # concats the arrays
u_star = Exact_u.T.flatten()[:, None]  #
v_star = Exact_v.T.flatten()[:, None]
h_star = Exact_h.T.flatten()[:, None]

###########################

idx_x = np.random.choice(x.shape[0], N0, replace=False)
idx_x = np.sort(idx_x)

x0 = x[idx_x, :]
u0 = Exact_u[idx_x, 0:1]
v0 = Exact_v[idx_x, 0:1]

idx_t = np.random.choice(t.shape[0], N_b, replace=False)
idx_t = np.sort(idx_t)
tb = t[idx_t, :]

x_f = lb + (ub - lb) * lhs(2, N_f) # determine sampling points 
t0 = torch.zeros([x0.shape[0],1])

X_lb = np.concatenate((0 * tb + lb[0], tb), 1)  # (lb[0], tb)
X_ub = np.concatenate((0 * tb + ub[0], tb), 1)  # (ub[0], tb)

x_b = np.vstack((X_lb,X_ub)) # [x,t]
x_0 = np.concatenate([x0,t0],1)
u_b = np.zeros(x_b.shape)
u_0 = np.concatenate([u0,v0],1)<jupyter_output><empty_output><jupyter_text>Create input_data dictionary and transfer data to torch<jupyter_code>x = {"x_0": torch.tensor(x_0).float(), "x_b": torch.tensor(x_b).float(), "x_f":torch.tensor(x_f).float()}
u_0 = torch.tensor(u_0).float()
u_b = torch.tensor(u_b).float()<jupyter_output><empty_output><jupyter_text>Preparing Optimizer for training<jupyter_code>optimizer = optim.Adam(model.parameters(),lr=1e-3)<jupyter_output><empty_output><jupyter_text>Training_loop<jupyter_code>num_epochs = 100
for epoch in range(num_epochs):
    optimizer.zero_grad()
    loss = model.pinn_loss(x, u_0, u_b,interpolation_criterion=nn.MSELoss(), boundary_criterion=nn.MSELoss(), pde_norm=nn.MSELoss())
    loss.backward()
    print("Epoch %d Loss %.10f:"%(epoch + 1, loss.item()))
    optimizer.step()<jupyter_output>Epoch 1 Loss 0.4857752025:
Epoch 2 Loss 0.3936398029:
Epoch 3 Loss 0.4228673279:
Epoch 4 Loss 0.3832141459:
Epoch 5 Loss 0.3486938775:
Epoch 6 Loss 0.3396714032:
Epoch 7 Loss 0.3420429230:
Epoch 8 Loss 0.3440891206:
Epoch 9 Loss 0.3419007659:
Epoch 10 Loss 0.3351892829:
Epoch 11 Loss 0.3272061944:
Epoch 12 Loss 0.3239412010:
Epoch 13 Loss 0.3278098702:
Epoch 14 Loss 0.3337433934:
Epoch 15 Loss 0.3347472250:
Epoch 16 Loss 0.3296834230:
Epoch 17 Loss 0.3229019344:
Epoch 18 Loss 0.3192938268:
Epoch 19 Loss 0.3199869990:
Epoch 20 Loss 0.3218687773:
Epoch 21 Loss 0.3218515217:
Epoch 22 Loss 0.3196863234:
Epoch 23 Loss 0.3166038096:
Epoch 24 Loss 0.3138881624:
Epoch 25 Loss 0.3122970760:
Epoch 26 Loss 0.3116075993:
Epoch 27 Loss 0.3112331033:
Epoch 28 Loss 0.3108410537:
Epoch 29 Loss 0.3098304868:
Epoch 30 Loss 0.3077088296:
Epoch 31 Loss 0.3051316440:
Epoch 32 Loss 0.3030409515:
Epoch 33 Loss 0.3015531003:
Epoch 34 Loss 0.3001136184:
Epoch 35 Loss 0.2977615893:
Epoch 36 Loss 0.2941000462:
E[...]
{'fraction_non_alphanumeric': 0.14781491002570693, 'fraction_numerical': 0.09607969151670952, 'mean_word_length': 3.86328125, 'pattern_counts': {'":': 3, '<': 32, '<?xml version=': 0, '>': 32, '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': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '333564', 'license': 'MIT', 'n_tokens_mistral': 2904, 'n_tokens_neox': 2540, 'n_words': 562, 'path': 'examples/1D_Schroedinger/SchrodingerExample.ipynb', 'repo_name': 's0vereign/NeuralSolvers', 'size': '6224'}