<jupyter_start><jupyter_text>Solving a stiff differential equation with a deep neural netWritten by [<NAME>](https://sites.google.com/site/mahdiebrahimikahou/about-me)  The problem Consider the following first order differential equation:$\frac{d y(t)}{dt} = g(y,t) = -21 y(t)+ e^{-t}$$y(0) = 0.0$The example is taken from these [notes](https://ece.uwaterloo.ca/~dwharder/NumericalAnalysis/14IVPs/stiff/complete.html). Note that this problem has an analytical solution of the form  $y(t) = A e^{-st}+ y_p(t)$, where $y_p(t)$ is the particular solution. So let's solve it, after a bit of algebra we get:$y(t) = 0.05 \big( e^{-t} - e^{-21 t}\big)$  Let's plot the theoretical solution<jupyter_code>import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from matplotlib import cm
fontsize= 18
ticksize = 14
figsize = (7, 4.5)
params = {"text.usetex": True,
    'font.family':'serif',
    "figure.figsize":figsize, 
    'figure.dpi': 80,
    'figure.edgecolor': 'k',
    'font.size': fontsize, 
    'axes.labelsize': fontsize,
    'axes.titlesize': fontsize,
    'xtick.labelsize': ticksize,
    'ytick.labelsize': ticksize
}
plt.rcParams.update(params) 
class params:
    def __init__(self,
                 α = -21,
                 β = -1.0,
                 y_0 = 0.0,
                 T = 4,
                ):
        self.α = α
        self.β = β
        self.y_0 = y_0
        self.T = T
def theory_sol(time):
    α = params().α
    β = params().β
    return 0.05*(-torch.exp(α*time)+ torch.exp(β*time))
time_mesh = torch.arange(0.0, params().T , 0.005)
dim = time_mesh.shape[0]
time_mesh = time_mesh.reshape([dim,1])
y_theory = theory_sol(time_mesh)
plt.plot(time_mesh,y_theory,label= r"Analytical Solution")
plt.xlabel(r"Time(t)")
plt.tight_layout()
plt.legend(loc='best')
plt.show()

from scipy.integrate import odeint
def g_1(y,t):
    α = params().α
    β = params().β
    dydt = α * y + np.exp(β*t)
    return dydt

y0 = 0.0

# time points
time_1 = np.linspace(0,params().T)

# solve ODE
y_1 = odeint(g_1,y0,time_1)

# plot results
plt.plot(time_1,y_1, label = r"ODEINT Solution")
plt.plot(time_mesh,y_theory,label= r"Analytical Solution")
plt.title(r"$\dot{y} = -21y+\exp(-t)$")
plt.xlabel('Time(t)')
plt.ylabel('y(t)')
plt.legend(loc='best')

plt.show()
print("ODEINT uses",time_1.shape[0], "data points")<jupyter_output>ODEINT uses 50 data points<jupyter_text>Solving the problem with a DNN<jupyter_code>from torch.utils.data import Dataset, DataLoader
class Collocation_data:
    def __init__(self,
                 mesh_step = 0.08,
                 batch_size = 4,
                ):
        
        
        self.T = params().T
        self.mesh_step = mesh_step
        self.batch_size = batch_size
    
        self.time_range = torch.arange(start = 0.0, end = self.T , step = self.mesh_step)
        dim = self.time_range.shape[0]
        self.time = self.time_range.reshape([dim,1])<jupyter_output><empty_output><jupyter_text>Writing the dataloader<jupyter_code>class Data_label(Dataset):
    
    def __init__(self,data):
        self.data = data
        self.time = self.data[:,[0]]
        self.n_samples = self.data.shape[0]
    
    def __getitem__(self,index):
        #first element is time, second element is y
            return self.time[index]
        
    def __len__(self):
        return self.n_samples
train_data = Collocation_data().time
train_labeled = Data_label(train_data)
batch_len = Collocation_data().batch_size
train = DataLoader(dataset = train_labeled, batch_size = batch_len , shuffle = True )
epsilon = 1e-8
np.sqrt(epsilon)<jupyter_output><empty_output><jupyter_text>Here I am going to use an old trick introduced by [Lagaris, Likas, and Fotiadis](https://arxiv.org/pdf/physics/9705023.pdf), equation 11. The trick is essentially representing the solution by :$y(t) = y_0 + t\tilde{y}(t)$Note that:(1) $y(0) = y_0$(2) $\tilde{y}$ is going to be approximated by a neural net  Setting up the neuran net and some useful functions<jupyter_code>def y(model,t):
    y_0 = params().y_0
    return y_0+t*model(t)
def derivative_forward(model,t): #forward differencing: dy/t
    sqrt_eps = np.sqrt(epsilon)
    return (y(model,t+sqrt_eps)-y(model,t))/sqrt_eps
def derivative_back(model,t): #backward differencing : dy/dt
    sqrt_eps = np.sqrt(epsilon)
    return (y(model,t)-y(model,t-sqrt_eps))/sqrt_eps

def derivative_center(model,t): #centered differencing : This sucks
    sqrt_eps = np.sqrt(epsilon)
    return (y(model,t+sqrt_eps)-y(model,t-sqrt_eps))/2*sqrt_eps

def f(model,t): # f(y,t)
    α = params().α
    β = params().β
    dfdt = α * y(model,t) + torch.exp(β*t)
    return dfdt
def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']<jupyter_output><empty_output><jupyter_text>Defining the neural network<jupyter_code>class NN(nn.Module):
    def __init__(self, 
                 dim_hidden = 128,
                layers = 4,
                hidden_bias = True):
        super().__init__()
        self.dim_hidden= dim_hidden
        self.layers = layers
        self.hidden_bias = hidden_bias
        
        torch.manual_seed(123)
        module = []
        module.append(nn.Linear(1,self.dim_hidden, bias = self.hidden_bias))
        module.append(nn.Tanh())
        
        for i in range(self.layers-1):
            module.append(nn.Linear(self.dim_hidden,self.dim_hidden, bias = self.hidden_bias))
            module.append(nn.Tanh())  
            
        module.append(nn.Linear(self.dim_hidden,1))
        module.append(nn.Softplus(beta = 1.0))
        
        self.y = nn.Sequential(*module)


    def forward(self, x):
        out = self.y(x)
        return  out<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>model= NN()
learning_rate = 1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.8)
num_epochs = 2501
for epoch in range(num_epochs):
    for i, time in enumerate(train):
        res =  derivative_back(model,time) - f(model,time)
        
        loss = res.pow(2).mean()
                
        optimizer.zero_grad()
        loss.backward()
        
        optimizer.step()
    scheduler.step()
    if epoch == 0:
         print('epoch' , ',' , 'loss' , ',' , 'lr_rate')
    if epoch % 100 == 0:
          print(epoch,',',"{:.2e}".format(loss.item()),',', "{:.2e}".format(get_lr(optimizer)) )  
        
model.eval()
time_tensor = Collocation_data(mesh_step = 0.005).time
y_approx = y(model,time_tensor).detach()
plt.plot(time_tensor, y_approx, color = 'r' ,label= r"Approximate Solution")
plt.plot(time_mesh,y_theory, label= r"Analytical Solution")
#plt.plot(time_1,y_1, label = r"ODEINT Solution")
plt.title(r"$\dot{y} = -21y+\exp(-t)$")

plt.xlabel(r"Time(t)")
plt.tight_layout()
plt.legend(loc='best')
plt.show()<jupyter_output><empty_output><jupyter_text>Plotting the GIF:<jupyter_code>model= NN()
learning_rate = 1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.8)
num_epochs = 2500
gif_results = torch.zeros([time_tensor.shape[0],250])
j = 0
for epoch in range(num_epochs):
    if epoch % 10 ==0:
        gif_results[:,j] = y(model,time_tensor).detach().squeeze()
        j = j+1
    for i, time in enumerate(train):
        res =  derivative_back(model,time) - f(model,time)
        
        loss = res.pow(2).mean()
                
        optimizer.zero_grad()
        loss.backward()
        
        optimizer.step()
    scheduler.step()
    if epoch == 0:
         print('epoch' , ',' , 'loss' , ',' , 'lr_rate')
    if epoch % 100 == 0:
          print(epoch,',',"{:.2e}".format(loss.item()),',', "{:.2e}".format(get_lr(optimizer)) )  
        
import os
import matplotlib.pyplot as plt
import imageio
filenames = []
for i in range(gif_results.shape[1]):
    # plot the line chart
    plt.plot(time_tensor,gif_results[:,i], color= 'b',label= r"Approximate Solution" )
    plt.plot(time_mesh,y_theory, "--", color= 'r', label= r"Analytical Solution")
    plt.title(r"$\dot{y} = -21y+\exp(-t)$")
    plt.ylabel(r"$y(t)$")
    plt.xlabel(r"Time(t)")
    plt.legend(loc='upper right')
    plt.ylim(0.0,0.06)
    plt.xlim(-0.05,4)
    plt.tight_layout()
    # create file name and append it to a list
    filename = f'{i}.png'
    filenames.append(filename)
    
    # save frame
    plt.savefig(filename)
    plt.close()# build gif
with imageio.get_writer('stiff.gif', mode='I') as writer:
    for filename in filenames:
        image = imageio.imread(filename)
        writer.append_data(image)
        
# Remove files
for filename in set(filenames):
    os.remove(filename)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1472420432665081, 'fraction_numerical': 0.019141465624646054, 'mean_word_length': 3.141651031894934, 'pattern_counts': {'":': 2, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 3, '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': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '116201', 'license': 'MIT', 'n_tokens_mistral': 3345, 'n_tokens_neox': 3169, 'n_words': 642, 'path': 'codes/stiff differential equations/7.1 Solving Stiff ODEs DNN, No Boundary In The Loss.ipynb', 'repo_name': 'Mekahou/Fun-Stuff', 'size': '8836'}