Differential Equations

import numpy as np

def euler(f, y0, t_span, h=0.01):
    t = np.arange(t_span[0], t_span[1], h)
    y = np.zeros(len(t))
    y[0] = y0
    for i in range(1, len(t)):
        y[i] = y[i-1] + h * f(t[i-1], y[i-1])
    return t, y

# Example: dy/dt = -y, y(0) = 1
f = lambda t, y: -y
t, y = euler(f, 1.0, [0, 5])

def rk4(f, y0, t_span, h=0.01):
    t = np.arange(t_span[0], t_span[1], h)
    y = np.zeros(len(t))
    y[0] = y0
    for i in range(1, len(t)):
        k1 = h * f(t[i-1], y[i-1])
        k2 = h * f(t[i-1] + h/2, y[i-1] + k1/2)
        k3 = h * f(t[i-1] + h/2, y[i-1] + k2/2)
        k4 = h * f(t[i-1] + h, y[i-1] + k3)
        y[i] = y[i-1] + (k1 + 2*k2 + 2*k3 + k4) / 6
    return t, y

t, y = rk4(lambda t, y: -y, 1.0, [0, 5])

from scipy.integrate import solve_ivp
import numpy as np
import matplotlib.pyplot as plt

# Define the ODE: dy/dt = -2y + sin(t)
def ode(t, y):
    return -2*y + np.sin(t)

# Solve with default RK45
sol = solve_ivp(ode, [0, 10], [1.0], dense_output=True, max_step=0.1)

# Evaluate solution on a fine grid
t_eval = np.linspace(0, 10, 500)
y_eval = sol.sol(t_eval)[0]

# Compare with Euler for visualization
def euler(f, y0, t_span, h=0.01):
    t = np.arange(t_span[0], t_span[1], h)
    y = np.zeros(len(t))
    y[0] = y0
    for i in range(1, len(t)):
        y[i] = y[i-1] + h * f(t[i-1], y[i-1])
    return t, y

t_euler, y_euler = euler(ode, 1.0, [0, 10], h=0.05)

plt.plot(t_eval, y_eval, label='RK45 (scipy)', linewidth=2)
plt.plot(t_euler, y_euler, '--', label='Euler (h=0.05)', alpha=0.7)
plt.xlabel('t'); plt.ylabel('y')
plt.legend(); plt.title('Comparison of Numerical Methods')
plt.show()

import torch
import torch.nn as nn
from torchdiffeq import odeint  # pip install torchdiffeq

class NeuralODEFunc(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, input_dim)
        )

    def forward(self, t, y):
        return self.net(y)

# Usage
func = NeuralODEFunc(input_dim=2, hidden_dim=16)
y0 = torch.tensor([1.0, 0.0])
t_span = torch.linspace(0, 5.0, 100)
y_traj = odeint(func, y0, t_span)  # Shape: (100, 2)

# Simplified PINN for dy/dt = -y, y(0) = 1
class PINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32), nn.Tanh(),
            nn.Linear(32, 32), nn.Tanh(),
            nn.Linear(32, 1)
        )

    def forward(self, t):
        return self.net(t)

model = PINN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(5000):
    t = torch.rand(100, 1) * 5  # Collocation points
    t.requires_grad_(True)
    y = model(t)

    # ODE residual: dy/dt + y = 0
    dy_dt = torch.autograd.grad(y, t, grad_outputs=torch.ones_like(y),
                                 create_graph=True)[0]
    ode_loss = torch.mean((dy_dt + y)**2)

    # Initial condition: y(0) = 1
    ic_loss = (model(torch.zeros(1, 1)) - 1.0)**2

    loss = ode_loss + 10 * ic_loss
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()