PINNS: Adaptive weights¶
The network $\hat{u}(x)$ solves $\nu u'' - u = e^x$ on $(-1,1)$ with $u(-1)=1$, $u(1)=0$, and $\nu = 0.001$. Each training step draws $N=16$ random points $x_i$ in $(-1,1)$ and weights the residual loss and the boundary loss,
$$ \mathcal{L}(\theta,\lambda_{\text{pde}},\lambda_{\text{bc}}) = w(\lambda_{\text{pde}})\,\mathcal{L}_{\text{pde}}(\theta) + w(\lambda_{\text{bc}})\,\mathcal{L}_{\text{bc}}(\theta), $$
where
$$ \mathcal{L}_{\text{pde}} = \frac{1}{N}\sum_{i=1}^{N}\big(\nu\,\hat{u}''(x_i) - \hat{u}(x_i) - e^{x_i}\big)^2, \qquad \mathcal{L}_{\text{bc}} = \big(\hat{u}(-1)-1\big)^2 + \hat{u}(1)^2 . $$
Each weight is a scaled sigmoid of a trainable parameter $\lambda$,
$$ w(\lambda) = 100\,\sigma(\lambda) = \frac{100}{1+e^{-\lambda}}, $$
so the weight stays between 0 and 100. The starting value $\lambda=-4.6$ gives $w \approx 1$, and $\lambda=-2.2$ gives $w \approx 10$.
The derivative of the loss with respect to $\lambda_{\text{pde}}$ follows from $\sigma' = \sigma(1-\sigma)$,
$$ \begin{align*} \frac{\partial \mathcal{L}}{\partial \lambda_{\text{pde}}} &= \mathcal{L}_{\text{pde}}\,w'(\lambda_{\text{pde}}) \\ &= 100\,\mathcal{L}_{\text{pde}}\,\sigma(\lambda_{\text{pde}})\big(1-\sigma(\lambda_{\text{pde}})\big) \\ &= \frac{100\,\mathcal{L}_{\text{pde}}\,e^{-\lambda_{\text{pde}}}}{\big(1+e^{-\lambda_{\text{pde}}}\big)^2}, \end{align*} $$
and the same formula with $\mathcal{L}_{\text{bc}}$ gives $\partial\mathcal{L}/\partial\lambda_{\text{bc}}$.
Training solves the min-max problem $\min_{\theta}\max_{\lambda_{\text{pde}},\lambda_{\text{bc}}}\mathcal{L}$. Adam lowers $\mathcal{L}$ over $\theta$, and each $\lambda$ takes a gradient-ascent step,
$$ \lambda \leftarrow \lambda + \eta\,\frac{\partial \mathcal{L}}{\partial \lambda}, \qquad \eta = 0.01 . $$
The derivative is positive, so each $\lambda$ never decreases. The step for each $\lambda$ is proportional to its own loss term, so the weight on a loss term that stays large rises. The factor $\sigma(1-\sigma)$ shrinks once $\lambda$ passes 0, so the steps shrink and the weight never exceeds 100.
The code stores $\lambda_{\text{pde}}$ and $\lambda_{\text{bc}}$ as nn.Parameters.
Autograd computes the derivative above in loss.backward(), and optim.SGD(..., maximize=True) applies the ascent step.
The fixed-weight runs set requires_grad=False, so their $\lambda$ values stay at the starting values.
In [1]:
Copied!
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from dataclasses import dataclass
from typing import List, Tuple, Dict
from tqdm import tqdm
@dataclass
class ModelConfig:
"""Configuration for the neural network model."""
input_dim: int
hidden_layers: List[int]
output_dim: int
activation: nn.Module = nn.Tanh()
device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@dataclass
class TrainingConfig:
"""Configuration for the training process."""
epochs: int = 20000
batch_size: int = 16
initial_lr: float = 0.01
lr_decay_factor: float = 0.1
lr_decay_epochs: int = 4000
lambda_pde_init: float = -4.6 # Initial lambda for PDE residual (weight 100*sigmoid(-4.6) ≈ 1)
lambda_bc_init: float = -4.6 # Initial lambda for boundary conditions
adaptive_lambda: bool = False # True trains the lambdas, False keeps them fixed
adaptive_rate: float = 0.01 # SGD step size for the lambdas
class PINN(nn.Module):
"""Physics-Informed Neural Network for solving differential equations."""
def __init__(self, config: ModelConfig):
super().__init__()
self.layers = self._build_layers(config)
self.activation = config.activation
self.device = config.device
def _build_layers(self, config: ModelConfig) -> nn.ModuleList:
"""Construct neural network layers with Xavier initialization."""
layer_sizes = [config.input_dim] + config.hidden_layers + [config.output_dim]
layers = nn.ModuleList()
for i in range(len(layer_sizes) - 1):
layer = nn.Linear(layer_sizes[i], layer_sizes[i + 1])
nn.init.xavier_normal_(layer.weight)
nn.init.zeros_(layer.bias)
layers.append(layer)
return layers
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass through the network."""
for i, layer in enumerate(self.layers[:-1]):
x = self.activation(layer(x))
return self.layers[-1](x)
class DiffEqSolver:
"""Solver for differential equations using PINN."""
def __init__(self, model_config: ModelConfig, training_config: TrainingConfig):
self.model = PINN(model_config).to(model_config.device)
self.config = training_config
self.device = model_config.device
self.nu = 0.001 # Diffusion coefficient
# Loss weights w = 100 * sigmoid(lambda), with lambda a trainable parameter
self.lambda_pde = nn.Parameter(torch.tensor(training_config.lambda_pde_init, device=self.device),
requires_grad=training_config.adaptive_lambda)
self.lambda_bc = nn.Parameter(torch.tensor(training_config.lambda_bc_init, device=self.device),
requires_grad=training_config.adaptive_lambda)
def _adaptive_weight(self, lambda_val: torch.Tensor) -> torch.Tensor:
"""Map lambda to a loss weight between 0 and 100."""
return 100 * torch.sigmoid(lambda_val)
def _compute_pde_residual(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
"""Compute the PDE residual."""
du = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u),
create_graph=True)[0]
d2u = torch.autograd.grad(du, x, grad_outputs=torch.ones_like(du),
create_graph=True)[0]
return self.nu * d2u - u - torch.exp(x)
def _compute_loss(self, x: torch.Tensor) -> Tuple[torch.Tensor, float, float]:
"""Compute total loss with separate weights for PDE and BC."""
u = self.model(x)
pde_residual = self._compute_pde_residual(x, u)
pde_loss = torch.mean(pde_residual**2)
left = torch.tensor([-1.]).reshape(1, 1).to(self.device)
right = torch.tensor([1.]).reshape(1, 1).to(self.device)
bc_loss = torch.mean((self.model(left) - 1)**2 + (self.model(right) - 0)**2)
# Apply separate adaptive weights
w_pde = self._adaptive_weight(self.lambda_pde)
w_bc = self._adaptive_weight(self.lambda_bc)
total_loss = w_pde * pde_loss + w_bc * bc_loss
return total_loss, pde_loss.item(), bc_loss.item()
def train(self, desc: str = "Training") -> Tuple[List[float], List[float], List[float], List[float], List[float]]:
"""Train the model and return loss history."""
optimizer = optim.Adam(self.model.parameters(), lr=self.config.initial_lr)
# The network minimizes the weighted loss and the lambdas maximize it
lambda_optimizer = optim.SGD([self.lambda_pde, self.lambda_bc],
lr=self.config.adaptive_rate, maximize=True)
total_losses = []
pde_losses = []
bc_losses = []
lambda_pde_values = []
lambda_bc_values = []
pbar = tqdm(range(1, self.config.epochs + 1), desc=desc)
for epoch in pbar:
x = (2 * torch.rand(self.config.batch_size, 1) - 1).to(self.device)
x.requires_grad_(True)
optimizer.zero_grad()
lambda_optimizer.zero_grad()
loss, pde_loss, bc_loss = self._compute_loss(x)
loss.backward()
optimizer.step()
if self.config.adaptive_lambda:
lambda_optimizer.step()
pbar.set_description(f'{desc} (λ_pde={self.lambda_pde.item():.2f}, λ_bc={self.lambda_bc.item():.2f})')
if epoch % self.config.lr_decay_epochs == 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= self.config.lr_decay_factor
pbar.set_description(f'{desc} LR: {param_group["lr"]:.2e}')
if epoch % 100 == 0:
total_losses.append(loss.item())
pde_losses.append(pde_loss)
bc_losses.append(bc_loss)
lambda_pde_values.append(self.lambda_pde.item())
lambda_bc_values.append(self.lambda_bc.item())
pbar.set_postfix({
'total_loss': f'{loss.item():.6e}',
'pde_loss': f'{pde_loss:.6e}',
'bc_loss': f'{bc_loss:.6e}'
})
return total_losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values
def plot_detailed_losses(losses: Tuple[List[float], List[float], List[float], List[float], List[float]]):
"""Plot detailed loss components and lambda evolution."""
total_losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values = losses
total_losses = np.array(pde_losses) + np.array(bc_losses)
epochs = range(100, len(total_losses) * 100 + 1, 100)
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
# Plot losses
ax1.semilogy(epochs, total_losses, 'k-', label='Total Loss', linewidth=2)
ax1.semilogy(epochs, pde_losses, 'r--', label='PDE Residual Loss', linewidth=2)
ax1.semilogy(epochs, bc_losses, 'b--', label='BC Loss', linewidth=2)
ax1.set_xlabel('Iteration')
ax1.set_ylabel('Loss (log scale)')
ax1.set_title('Loss Components Evolution')
ax1.grid(True)
ax1.legend()
# Plot lambda evolution
ax2.plot(epochs, lambda_pde_values, 'r-', label='λ_PDE', linewidth=2)
ax2.plot(epochs, lambda_bc_values, 'b-', label='λ_BC', linewidth=2)
ax2.set_xlabel('Iteration')
ax2.set_ylabel('λ value')
ax2.set_title('Adaptive λ Evolution')
ax2.grid(True)
ax2.legend()
plt.tight_layout()
return fig
def exact_solution(x: np.ndarray) -> np.ndarray:
"""Compute the exact solution of the differential equation."""
c1 = ((1 + 1000/np.exp(1)/999)*np.exp((-10*np.sqrt(10))) -
np.exp(10*np.sqrt(10))*1000*np.exp(1)/999) / (np.exp(-20*np.sqrt(10)) -
np.exp(20*np.sqrt(10)))
c2 = (np.exp(-10*np.sqrt(10))*1000*np.exp(1)/999 -
np.exp(10*np.sqrt(10)) * (1+1000/np.exp(1)/999)) / (np.exp(-20*np.sqrt(10)) -
np.exp(20*np.sqrt(10)))
return (c1*np.exp(10*np.sqrt(10)*x) + c2*np.exp(-10*np.sqrt(10)*x) -
(1000/999)*np.exp(x))
def plot_comparison(solvers: Dict[str, DiffEqSolver], x_plot: np.ndarray,
y_exact: np.ndarray) -> Dict[str, float]:
"""Plot comparison of different methods and return their errors."""
plt.figure(figsize=(12, 8))
plt.plot(x_plot, y_exact, 'k-', label='Exact Solution', linewidth=2)
colors = ['r--', 'g--', 'b--']
errors = {}
for (name, solver), color in zip(solvers.items(), colors):
x_tensor = torch.tensor(x_plot.reshape(-1, 1), dtype=torch.float32).to(solver.device)
with torch.no_grad():
y_pred = solver.model(x_tensor).cpu().numpy()
errors[name] = np.linalg.norm(y_exact - y_pred.flatten()) / np.linalg.norm(y_exact)
plt.plot(x_plot, y_pred, color, label=name, linewidth=1.5)
plt.plot([-1, 1], [1, 0], 'r*', label='Boundary Conditions')
plt.xlabel('x')
plt.ylabel('u(x)')
plt.title('Comparison of Different Methods (ν = 0.001)')
plt.grid(True)
plt.legend()
# Add error text box
error_text = '\n'.join([f'{name}: {error:.2e}' for name, error in errors.items()])
plt.text(0.02, 0.02, f'L2 Relative Errors:\n{error_text}',
transform=plt.gca().transAxes,
bbox=dict(facecolor='white', alpha=0.8),
fontsize=8)
plt.tight_layout()
return errors
def plot_losses(all_losses: Dict[str, List[float]]):
"""Plot training loss history for all methods."""
plt.figure(figsize=(10, 6))
colors = ['r-o', 'g-o', 'b-o']
for (name, losses), color in zip(all_losses.items(), colors):
plt.semilogy(range(1000, len(losses) * 1000 + 1, 1000),
losses, color, label=name)
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Training Loss History')
plt.grid(True)
plt.legend()
plt.tight_layout()
def main():
# Set random seeds
torch.manual_seed(42)
np.random.seed(42)
# Base configurations
model_config = ModelConfig(
input_dim=1,
hidden_layers=[16, 16, 16],
output_dim=1
)
# Create solvers with different lambda strategies
solvers = {}
training_configs = {
'Fixed λ=1': TrainingConfig(lambda_pde_init=-4.6, lambda_bc_init=-4.6, adaptive_lambda=False),
'Fixed λ=10': TrainingConfig(lambda_pde_init=-2.2, lambda_bc_init=-2.2, adaptive_lambda=False),
'Adaptive λ': TrainingConfig(lambda_pde_init=-4.6, lambda_bc_init=-4.6, adaptive_lambda=True),
}
# Train all models
all_losses = {}
adaptive_losses = None
for name, config in training_configs.items():
solver = DiffEqSolver(model_config, config)
losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values = solver.train(desc=f'Training {name}')
if name == 'Adaptive λ':
adaptive_losses = (losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values)
solvers[name] = solver
all_losses[name] = np.array(pde_losses) + np.array(bc_losses) #losses
# Plot detailed losses
plot_detailed_losses(adaptive_losses)
plt.show()
# Generate comparison data
x_plot = np.linspace(-1, 1, 100)
y_exact = exact_solution(x_plot)
# Plot comparisons
errors = plot_comparison(solvers, x_plot, y_exact)
plt.show()
# Plot loss histories
plot_losses(all_losses)
plt.show()
# Print final errors
print('\nFinal L2 relative errors:')
for name, error in errors.items():
print(f'{name}: {error:.6e}')
if __name__ == "__main__":
main()
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from dataclasses import dataclass
from typing import List, Tuple, Dict
from tqdm import tqdm
@dataclass
class ModelConfig:
"""Configuration for the neural network model."""
input_dim: int
hidden_layers: List[int]
output_dim: int
activation: nn.Module = nn.Tanh()
device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@dataclass
class TrainingConfig:
"""Configuration for the training process."""
epochs: int = 20000
batch_size: int = 16
initial_lr: float = 0.01
lr_decay_factor: float = 0.1
lr_decay_epochs: int = 4000
lambda_pde_init: float = -4.6 # Initial lambda for PDE residual (weight 100*sigmoid(-4.6) ≈ 1)
lambda_bc_init: float = -4.6 # Initial lambda for boundary conditions
adaptive_lambda: bool = False # True trains the lambdas, False keeps them fixed
adaptive_rate: float = 0.01 # SGD step size for the lambdas
class PINN(nn.Module):
"""Physics-Informed Neural Network for solving differential equations."""
def __init__(self, config: ModelConfig):
super().__init__()
self.layers = self._build_layers(config)
self.activation = config.activation
self.device = config.device
def _build_layers(self, config: ModelConfig) -> nn.ModuleList:
"""Construct neural network layers with Xavier initialization."""
layer_sizes = [config.input_dim] + config.hidden_layers + [config.output_dim]
layers = nn.ModuleList()
for i in range(len(layer_sizes) - 1):
layer = nn.Linear(layer_sizes[i], layer_sizes[i + 1])
nn.init.xavier_normal_(layer.weight)
nn.init.zeros_(layer.bias)
layers.append(layer)
return layers
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass through the network."""
for i, layer in enumerate(self.layers[:-1]):
x = self.activation(layer(x))
return self.layers[-1](x)
class DiffEqSolver:
"""Solver for differential equations using PINN."""
def __init__(self, model_config: ModelConfig, training_config: TrainingConfig):
self.model = PINN(model_config).to(model_config.device)
self.config = training_config
self.device = model_config.device
self.nu = 0.001 # Diffusion coefficient
# Loss weights w = 100 * sigmoid(lambda), with lambda a trainable parameter
self.lambda_pde = nn.Parameter(torch.tensor(training_config.lambda_pde_init, device=self.device),
requires_grad=training_config.adaptive_lambda)
self.lambda_bc = nn.Parameter(torch.tensor(training_config.lambda_bc_init, device=self.device),
requires_grad=training_config.adaptive_lambda)
def _adaptive_weight(self, lambda_val: torch.Tensor) -> torch.Tensor:
"""Map lambda to a loss weight between 0 and 100."""
return 100 * torch.sigmoid(lambda_val)
def _compute_pde_residual(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
"""Compute the PDE residual."""
du = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u),
create_graph=True)[0]
d2u = torch.autograd.grad(du, x, grad_outputs=torch.ones_like(du),
create_graph=True)[0]
return self.nu * d2u - u - torch.exp(x)
def _compute_loss(self, x: torch.Tensor) -> Tuple[torch.Tensor, float, float]:
"""Compute total loss with separate weights for PDE and BC."""
u = self.model(x)
pde_residual = self._compute_pde_residual(x, u)
pde_loss = torch.mean(pde_residual**2)
left = torch.tensor([-1.]).reshape(1, 1).to(self.device)
right = torch.tensor([1.]).reshape(1, 1).to(self.device)
bc_loss = torch.mean((self.model(left) - 1)**2 + (self.model(right) - 0)**2)
# Apply separate adaptive weights
w_pde = self._adaptive_weight(self.lambda_pde)
w_bc = self._adaptive_weight(self.lambda_bc)
total_loss = w_pde * pde_loss + w_bc * bc_loss
return total_loss, pde_loss.item(), bc_loss.item()
def train(self, desc: str = "Training") -> Tuple[List[float], List[float], List[float], List[float], List[float]]:
"""Train the model and return loss history."""
optimizer = optim.Adam(self.model.parameters(), lr=self.config.initial_lr)
# The network minimizes the weighted loss and the lambdas maximize it
lambda_optimizer = optim.SGD([self.lambda_pde, self.lambda_bc],
lr=self.config.adaptive_rate, maximize=True)
total_losses = []
pde_losses = []
bc_losses = []
lambda_pde_values = []
lambda_bc_values = []
pbar = tqdm(range(1, self.config.epochs + 1), desc=desc)
for epoch in pbar:
x = (2 * torch.rand(self.config.batch_size, 1) - 1).to(self.device)
x.requires_grad_(True)
optimizer.zero_grad()
lambda_optimizer.zero_grad()
loss, pde_loss, bc_loss = self._compute_loss(x)
loss.backward()
optimizer.step()
if self.config.adaptive_lambda:
lambda_optimizer.step()
pbar.set_description(f'{desc} (λ_pde={self.lambda_pde.item():.2f}, λ_bc={self.lambda_bc.item():.2f})')
if epoch % self.config.lr_decay_epochs == 0:
for param_group in optimizer.param_groups:
param_group['lr'] *= self.config.lr_decay_factor
pbar.set_description(f'{desc} LR: {param_group["lr"]:.2e}')
if epoch % 100 == 0:
total_losses.append(loss.item())
pde_losses.append(pde_loss)
bc_losses.append(bc_loss)
lambda_pde_values.append(self.lambda_pde.item())
lambda_bc_values.append(self.lambda_bc.item())
pbar.set_postfix({
'total_loss': f'{loss.item():.6e}',
'pde_loss': f'{pde_loss:.6e}',
'bc_loss': f'{bc_loss:.6e}'
})
return total_losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values
def plot_detailed_losses(losses: Tuple[List[float], List[float], List[float], List[float], List[float]]):
"""Plot detailed loss components and lambda evolution."""
total_losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values = losses
total_losses = np.array(pde_losses) + np.array(bc_losses)
epochs = range(100, len(total_losses) * 100 + 1, 100)
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
# Plot losses
ax1.semilogy(epochs, total_losses, 'k-', label='Total Loss', linewidth=2)
ax1.semilogy(epochs, pde_losses, 'r--', label='PDE Residual Loss', linewidth=2)
ax1.semilogy(epochs, bc_losses, 'b--', label='BC Loss', linewidth=2)
ax1.set_xlabel('Iteration')
ax1.set_ylabel('Loss (log scale)')
ax1.set_title('Loss Components Evolution')
ax1.grid(True)
ax1.legend()
# Plot lambda evolution
ax2.plot(epochs, lambda_pde_values, 'r-', label='λ_PDE', linewidth=2)
ax2.plot(epochs, lambda_bc_values, 'b-', label='λ_BC', linewidth=2)
ax2.set_xlabel('Iteration')
ax2.set_ylabel('λ value')
ax2.set_title('Adaptive λ Evolution')
ax2.grid(True)
ax2.legend()
plt.tight_layout()
return fig
def exact_solution(x: np.ndarray) -> np.ndarray:
"""Compute the exact solution of the differential equation."""
c1 = ((1 + 1000/np.exp(1)/999)*np.exp((-10*np.sqrt(10))) -
np.exp(10*np.sqrt(10))*1000*np.exp(1)/999) / (np.exp(-20*np.sqrt(10)) -
np.exp(20*np.sqrt(10)))
c2 = (np.exp(-10*np.sqrt(10))*1000*np.exp(1)/999 -
np.exp(10*np.sqrt(10)) * (1+1000/np.exp(1)/999)) / (np.exp(-20*np.sqrt(10)) -
np.exp(20*np.sqrt(10)))
return (c1*np.exp(10*np.sqrt(10)*x) + c2*np.exp(-10*np.sqrt(10)*x) -
(1000/999)*np.exp(x))
def plot_comparison(solvers: Dict[str, DiffEqSolver], x_plot: np.ndarray,
y_exact: np.ndarray) -> Dict[str, float]:
"""Plot comparison of different methods and return their errors."""
plt.figure(figsize=(12, 8))
plt.plot(x_plot, y_exact, 'k-', label='Exact Solution', linewidth=2)
colors = ['r--', 'g--', 'b--']
errors = {}
for (name, solver), color in zip(solvers.items(), colors):
x_tensor = torch.tensor(x_plot.reshape(-1, 1), dtype=torch.float32).to(solver.device)
with torch.no_grad():
y_pred = solver.model(x_tensor).cpu().numpy()
errors[name] = np.linalg.norm(y_exact - y_pred.flatten()) / np.linalg.norm(y_exact)
plt.plot(x_plot, y_pred, color, label=name, linewidth=1.5)
plt.plot([-1, 1], [1, 0], 'r*', label='Boundary Conditions')
plt.xlabel('x')
plt.ylabel('u(x)')
plt.title('Comparison of Different Methods (ν = 0.001)')
plt.grid(True)
plt.legend()
# Add error text box
error_text = '\n'.join([f'{name}: {error:.2e}' for name, error in errors.items()])
plt.text(0.02, 0.02, f'L2 Relative Errors:\n{error_text}',
transform=plt.gca().transAxes,
bbox=dict(facecolor='white', alpha=0.8),
fontsize=8)
plt.tight_layout()
return errors
def plot_losses(all_losses: Dict[str, List[float]]):
"""Plot training loss history for all methods."""
plt.figure(figsize=(10, 6))
colors = ['r-o', 'g-o', 'b-o']
for (name, losses), color in zip(all_losses.items(), colors):
plt.semilogy(range(1000, len(losses) * 1000 + 1, 1000),
losses, color, label=name)
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Training Loss History')
plt.grid(True)
plt.legend()
plt.tight_layout()
def main():
# Set random seeds
torch.manual_seed(42)
np.random.seed(42)
# Base configurations
model_config = ModelConfig(
input_dim=1,
hidden_layers=[16, 16, 16],
output_dim=1
)
# Create solvers with different lambda strategies
solvers = {}
training_configs = {
'Fixed λ=1': TrainingConfig(lambda_pde_init=-4.6, lambda_bc_init=-4.6, adaptive_lambda=False),
'Fixed λ=10': TrainingConfig(lambda_pde_init=-2.2, lambda_bc_init=-2.2, adaptive_lambda=False),
'Adaptive λ': TrainingConfig(lambda_pde_init=-4.6, lambda_bc_init=-4.6, adaptive_lambda=True),
}
# Train all models
all_losses = {}
adaptive_losses = None
for name, config in training_configs.items():
solver = DiffEqSolver(model_config, config)
losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values = solver.train(desc=f'Training {name}')
if name == 'Adaptive λ':
adaptive_losses = (losses, pde_losses, bc_losses, lambda_pde_values, lambda_bc_values)
solvers[name] = solver
all_losses[name] = np.array(pde_losses) + np.array(bc_losses) #losses
# Plot detailed losses
plot_detailed_losses(adaptive_losses)
plt.show()
# Generate comparison data
x_plot = np.linspace(-1, 1, 100)
y_exact = exact_solution(x_plot)
# Plot comparisons
errors = plot_comparison(solvers, x_plot, y_exact)
plt.show()
# Plot loss histories
plot_losses(all_losses)
plt.show()
# Print final errors
print('\nFinal L2 relative errors:')
for name, error in errors.items():
print(f'{name}: {error:.6e}')
if __name__ == "__main__":
main()
Training Fixed λ=1 LR: 1.00e-07: 100%|██████████| 20000/20000 [00:11<00:00, 1668.11it/s, total_loss=1.186793e-04, pde_loss=1.060072e-04, bc_loss=1.324681e-05] Training Fixed λ=10 LR: 1.00e-07: 100%|██████████| 20000/20000 [00:11<00:00, 1674.57it/s, total_loss=1.679444e-03, pde_loss=1.573222e-04, bc_loss=1.104227e-05] Training Adaptive λ LR: 1.00e-07: 100%|██████████| 20000/20000 [00:21<00:00, 942.68it/s, total_loss=2.078834e-01, pde_loss=2.068892e-03, bc_loss=1.237454e-05]
Final L2 relative errors: Fixed λ=1: 1.032303e-02 Fixed λ=10: 1.008857e-02 Adaptive λ: 1.075754e-02