Discrete time burgers

In [ ]:

import os
import sys
import time
import warnings
from typing import Tuple, List, Optional
from functools import partial

import os import sys import time import warnings from typing import Tuple, List, Optional from functools import partial

In [ ]:

from tqdm import tqdm
import torch
import torch.nn as nn
import numpy as np
import scipy.io
import pandas as pd
from scipy.interpolate import griddata
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.gridspec as gridspec
from pyDOE import lhs

from tqdm import tqdm import torch import torch.nn as nn import numpy as np import scipy.io import pandas as pd from scipy.interpolate import griddata import matplotlib.pyplot as plt from mpl_toolkits.axes_grid1 import make_axes_locatable import matplotlib.gridspec as gridspec from pyDOE import lhs

In [ ]:

class NeuralNetwork(nn.Module):
    """Neural network model for the PINNs implementation."""
    
    def __init__(self, input_size: int, output_size: int, hidden_layers: int, 
                 hidden_units: int, activation_function: nn.Module):
        """
        Initialize the neural network.
        
        Args:
            input_size: Number of input features
            output_size: Number of output features
            hidden_layers: Number of hidden layers
            hidden_units: Number of units in each hidden layer
            activation_function: Activation function to use
        """
        super().__init__()
        self.linear_in = nn.Linear(input_size, hidden_units)
        self.linear_out = nn.Linear(hidden_units, output_size)
        self.layers = nn.ModuleList([
            nn.Linear(hidden_units, hidden_units) for _ in range(hidden_layers)
        ])
        self.act = activation_function

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the network."""
        x = self.linear_in(x)
        for layer in self.layers:
            x = self.act(layer(x))
        return self.linear_out(x)

class NeuralNetwork(nn.Module): """Neural network model for the PINNs implementation.""" def __init__(self, input_size: int, output_size: int, hidden_layers: int, hidden_units: int, activation_function: nn.Module): """ Initialize the neural network. Args: input_size: Number of input features output_size: Number of output features hidden_layers: Number of hidden layers hidden_units: Number of units in each hidden layer activation_function: Activation function to use """ super().__init__() self.linear_in = nn.Linear(input_size, hidden_units) self.linear_out = nn.Linear(hidden_units, output_size) self.layers = nn.ModuleList([ nn.Linear(hidden_units, hidden_units) for _ in range(hidden_layers) ]) self.act = activation_function def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward pass through the network.""" x = self.linear_in(x) for layer in self.layers: x = self.act(layer(x)) return self.linear_out(x)

In [ ]:

class DiscreteTimePINNs:
    """Implementation of Physics-Informed Neural Networks for discrete time problems."""
    
    def __init__(self, config: dict):
        """
        Initialize the DiscreteTimePINNs solver.
        
        Args:
            config: Dictionary containing configuration parameters
        """
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.setup_directories()
        self.setup_data()
        self.setup_model()
        self.results = []
        self.iter = 0

    def setup_directories(self):
        """Create necessary directories for output."""
        directories = ['results']
        for directory in directories:
            os.makedirs(directory, exist_ok=True)

    def setup_data(self):
        """Load and prepare data for training."""
        # Load Burgers equation data
        data = scipy.io.loadmat(self.config['data_path'])
        self.t = data['t'].flatten()[:, None]
        self.x = data['x'].flatten()[:, None]
        self.Exact = np.real(data['usol']).T.astype(np.float32)
        
        # Setup time steps and initial conditions
        self.setup_time_steps()
        self.setup_initial_conditions()
        self.setup_irk_weights()
        
        # Convert data to tensors
        self.prepare_tensors()

    def setup_time_steps(self):
        """Setup time step information."""
        self.idx_t0 = self.config['idx_t0']
        self.idx_t1 = self.config['idx_t1']
        dt_value = self.t[self.idx_t1] - self.t[self.idx_t0]
        self.dt = torch.tensor(dt_value, dtype=torch.float32)

    def setup_initial_conditions(self):
        """Setup initial and boundary conditions."""
        N = self.config['N']
        idx_x = np.random.choice(self.Exact.shape[1], N, replace=False)
        self.x0 = self.x[idx_x, :]
        self.u0 = self.Exact[self.idx_t0:self.idx_t0+1, idx_x].T
        
        # Add noise if specified
        if self.config.get('noise_u0', 0.0) > 0:
            self.u0 += (self.config['noise_u0'] * np.std(self.u0) * 
                       np.random.randn(*self.u0.shape))
        
        # Setup boundary conditions
        self.x1 = np.vstack((self.config['lb'], self.config['ub']))
        self.x_star = self.x

    def setup_irk_weights(self):
        """Load and setup IRK weights."""
        tmp = np.loadtxt(
            f"./IRK_weights/Butcher_IRK{self.config['q']}.txt", 
            ndmin=2
        ).astype(np.float32)
        weights = np.reshape(tmp[0:self.config['q']**2 + self.config['q']], 
                           (self.config['q']+1, self.config['q']))
        self.IRK_weights = torch.tensor(weights, dtype=torch.float32).T

    def prepare_tensors(self):
        """Convert numpy arrays to PyTorch tensors."""
        # Dictionary of numpy arrays that need to be converted to tensors
        numpy_arrays = {
            'x0': self.x0,
            'x1': self.x1,
            'u0': self.u0,
            'x_star': self.x_star
        }
        
        # Convert numpy arrays to tensors
        for name, array in numpy_arrays.items():
            tensor = torch.from_numpy(array).float().to(self.device)
            if name in ['x0', 'x1', 'x_star']:
                tensor.requires_grad = True
            setattr(self, name, tensor)
            
        # Handle dt and IRK_weights separately as they're already tensors
        self.dt = self.dt.to(self.device)
        self.IRK_weights = self.IRK_weights.to(self.device)

    def setup_model(self):
        """Initialize the neural network model."""
        self.model = NeuralNetwork(
            input_size=1,
            output_size=self.config['q'] + 1,
            hidden_layers=4,
            hidden_units=50,
            activation_function=nn.Tanh()
        ).float().to(self.device)
        
        self.model.apply(self.init_weights)

    @staticmethod
    def init_weights(m):
        """Initialize network weights using Xavier initialization."""
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_normal_(m.weight)
            m.bias.data.fill_(0.0)

    @staticmethod
    def set_seed(seed: int = 42):
        """Set random seed for reproducibility."""
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = True
        np.random.seed(seed)

    def f(self, x: torch.Tensor, x_1: torch.Tensor, 
          dt: torch.Tensor, IRK_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute the physics-informed neural network function.
        
        Args:
            x: Input tensor
            x_1: Boundary conditions tensor
            dt: Time step
            IRK_weights: IRK weights tensor
            
        Returns:
            Tuple of tensors (U0, U1)
        """
        nu = 0.01/torch.pi
        U1 = self.model(x)
        U = U1[:, :-1]
        
        # Compute derivatives
        U_x = self.fwd_gradients_0(U, x)
        U_xx = self.fwd_gradients_0(U_x, x)
        
        # Compute physics terms
        F = -U*U_x + nu*U_xx
        U0 = U1 - dt * torch.matmul(F, IRK_weights)
        U1 = self.model(x_1)
        
        return U0, U1

    def compute_loss(self, x: torch.Tensor, x_1: torch.Tensor, 
                    dt: torch.Tensor, IRK_weights: torch.Tensor, 
                    U0_real: torch.Tensor) -> torch.Tensor:
        """Compute the total loss for training."""
        U0, U1 = self.f(x, x_1, dt, IRK_weights)
        return torch.sum((U0_real - U0) ** 2) + torch.sum(U1 ** 2)

    def closure(self, optimizer: torch.optim.Optimizer, x: torch.Tensor, 
                x_1: torch.Tensor, x_star: torch.Tensor, dt: torch.Tensor, 
                IRK_weights: torch.Tensor, U0_real: torch.Tensor) -> torch.Tensor:
        """Closure function for LBFGS optimizer."""
        optimizer.zero_grad()
        loss = self.compute_loss(x, x_1, dt, IRK_weights, U0_real)
        loss.backward(retain_graph=True)
        
        # Update iteration count and compute error
        self.iter += 1
        U1_pred = self.model(x_star)
        pred = U1_pred[:, -1].detach().cpu().numpy()
        error = np.linalg.norm(pred - self.Exact[self.idx_t1, :], 2) / \
                np.linalg.norm(self.Exact[self.idx_t1, :], 2)
        
        # Store results and save model if needed
        self.results.append([self.iter, loss.item(), error])    
        return loss

    def train(self, num_iter: int):
        """Train using L-BFGS optimizer."""
        optimizer = torch.optim.LBFGS(
            self.model.parameters(),
            lr=1.0,
            max_iter=num_iter,
            max_eval=num_iter,
            tolerance_grad=1e-7,
            tolerance_change=1.0 * np.finfo(float).eps,
            history_size=100,
            line_search_fn='strong_wolfe'
        )
        
        start_time = time.time()
        closure_fn = partial(self.closure, optimizer, self.x0, self.x1, 
                            self.x_star, self.dt, self.IRK_weights, self.u0)
        
        with tqdm(total=num_iter, desc="LBFGS Training", unit="iter") as pbar:
            def closure_with_tqdm():
                loss = closure_fn()
                pbar.set_postfix(loss=loss.item())
                pbar.update(1)
                return loss

            optimizer.step(closure_with_tqdm)
            
        self.lbfgs_training_time = time.time() - start_time
        # Save final results
        self.save_results()


    def save_results(self):
        """Save training results and model."""
        # Save training summary
        total_time = getattr(self, 'lbfgs_training_time', 0)
        
        with open('results/burgers_discrete_time_training_summary.txt', 'w') as f:
            f.write(f"Total training time: {total_time:.6e} seconds\n")
            f.write(f"Total iterations: {self.iter}\n")
            f.write(f"Final Loss: {self.results[-1][1]:.6e}\n")
            f.write(f"Final L2: {self.results[-1][2]:.6e}\n")
        
        # Save training data
        results_array = np.array(self.results)
        np.savetxt("results/burgers_discrete_time_training_data.csv", 
                   results_array, 
                   delimiter=",", 
                   header="Iter,Loss,L2",
                   comments="")
        
        # Save final model
        self.save_model('burgers_discrete_time.pt')

    def save_model(self, path: str):
        """Save model state dict."""
        torch.save(self.model.state_dict(), path)

    def load_model(self, path: str):
        """Load model state dict."""
        self.model.load_state_dict(torch.load(path))
        self.model.eval()

    def plot_results(self):
        """Generate plots for the results."""
        # Create figure for training curves
        self.plot_training_curves()
        
        # Create figure for solution visualization
        self.plot_solution()

    def plot_training_curves(self):
        """Plot training loss and L2 error curves."""
        data = pd.read_csv('results/burgers_discrete_time_training_data.csv')
        
        fig, axarr = plt.subplots(1, 2, figsize=(12, 4))
        
        # Loss plot
        axarr[0].semilogy(data['Iter'], data['Loss'], 
                        label='Loss', color='gray', linewidth=1)
        axarr[0].set_xlabel('Iteration')
        axarr[0].set_ylabel('Loss')
        
        # L2 error plot
        axarr[1].semilogy(data['Iter'], data['L2'], 
                        label='L2 Error', color='gray', linewidth=1)
        axarr[1].set_xlabel('Iteration')
        axarr[1].set_ylabel(r'$\mathrm{L}_2$')
        
        plt.tight_layout()
        plt.savefig('results/burgers_discrete_time_training_curves.pdf')
        plt.close()

    def plot_solution(self):
        """Plot the solution, including exact solution and predictions."""
        fig = plt.figure(figsize=(12, 10))
        fsize = 12
        # Set up GridSpec
        gs0 = gridspec.GridSpec(1, 2)
        gs0.update(top=1-0.06, bottom=1-1/2 + 0.1, left=0.15, right=0.85, wspace=0)
        gs1 = gridspec.GridSpec(1, 2)
        gs1.update(top=1-1/2-0.05, bottom=0.15, left=0.15, right=0.85, wspace=0.5)
        
        # Plot exact solution
        ax = plt.subplot(gs0[:, :])
        h = ax.imshow(self.Exact.T, interpolation='nearest', cmap='rainbow',
                        extent=[self.t.min(), self.t.max(), self.x.min(), self.x.max()],
                        origin='lower', aspect='auto')
        
        # Add colorbar
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(h, cax=cax)
        
        # Add time slice indicators
        line = np.linspace(self.x.min(), self.x.max(), 2)[:, None]
        # Fix scalar conversion by using item()
        t0_scalar = self.t[self.idx_t0].item()
        t1_scalar = self.t[self.idx_t1].item()
        ax.plot(t0_scalar * np.ones((2,1)), line, 'w-', linewidth=1)
        ax.plot(t1_scalar * np.ones((2,1)), line, 'w-', linewidth=1)
        
        ax.set_xlabel('$t$')
        ax.set_ylabel('$x$')
        ax.set_title('$u(t,x)$', fontsize=fsize)
        
        # Plot initial condition
        ax = plt.subplot(gs1[0, 0])
        ax.plot(self.x, self.Exact[self.idx_t0,:], 'b-', linewidth=2)
        ax.plot(self.x0.cpu().detach().numpy(), self.u0, 'rx', linewidth=2, label='Data')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$u(t,x)$')
        ax.set_title('$t = %.2f$' % t0_scalar, fontsize=fsize)
        ax.set_xlim([self.config['lb']-0.1, self.config['ub']+0.1])
        ax.legend(loc='upper center', bbox_to_anchor=(0.3, -0.25), ncol=2, frameon=False)
        
        # Plot prediction vs exact solution
        ax = plt.subplot(gs1[0, 1])
        U1_pred = self.model(self.x_star)
        U1_pred = U1_pred.cpu().detach().numpy()
        
        ax.plot(self.x, self.Exact[self.idx_t1,:], 'b-', linewidth=2, label='Exact')
        ax.plot(self.x_star.cpu().detach().numpy(), U1_pred[:,-1], 'r--', 
                linewidth=2, label='Prediction')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$u(t,x)$')
        ax.set_title('$t = %.2f$' % t1_scalar, fontsize=fsize)
        ax.set_xlim([self.config['lb']-0.1, self.config['ub']+0.1])
        ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.25), ncol=2, frameon=False)
        
        plt.savefig('results/burgers_discrete_time.pdf')
        plt.show()
        plt.close()

    @staticmethod
    def fwd_gradients_0(dy: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """Compute forward gradients."""
        z = torch.ones(dy.shape, device=dy.device).requires_grad_()
        g = torch.autograd.grad(dy, x, grad_outputs=z, create_graph=True)[0]
        ones = torch.ones(g.shape, device=g.device)
        return torch.autograd.grad(g, z, grad_outputs=ones, create_graph=True)[0]

class DiscreteTimePINNs: """Implementation of Physics-Informed Neural Networks for discrete time problems.""" def __init__(self, config: dict): """ Initialize the DiscreteTimePINNs solver. Args: config: Dictionary containing configuration parameters """ self.config = config self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.setup_directories() self.setup_data() self.setup_model() self.results = [] self.iter = 0 def setup_directories(self): """Create necessary directories for output.""" directories = ['results'] for directory in directories: os.makedirs(directory, exist_ok=True) def setup_data(self): """Load and prepare data for training.""" # Load Burgers equation data data = scipy.io.loadmat(self.config['data_path']) self.t = data['t'].flatten()[:, None] self.x = data['x'].flatten()[:, None] self.Exact = np.real(data['usol']).T.astype(np.float32) # Setup time steps and initial conditions self.setup_time_steps() self.setup_initial_conditions() self.setup_irk_weights() # Convert data to tensors self.prepare_tensors() def setup_time_steps(self): """Setup time step information.""" self.idx_t0 = self.config['idx_t0'] self.idx_t1 = self.config['idx_t1'] dt_value = self.t[self.idx_t1] - self.t[self.idx_t0] self.dt = torch.tensor(dt_value, dtype=torch.float32) def setup_initial_conditions(self): """Setup initial and boundary conditions.""" N = self.config['N'] idx_x = np.random.choice(self.Exact.shape[1], N, replace=False) self.x0 = self.x[idx_x, :] self.u0 = self.Exact[self.idx_t0:self.idx_t0+1, idx_x].T # Add noise if specified if self.config.get('noise_u0', 0.0) > 0: self.u0 += (self.config['noise_u0'] * np.std(self.u0) * np.random.randn(*self.u0.shape)) # Setup boundary conditions self.x1 = np.vstack((self.config['lb'], self.config['ub'])) self.x_star = self.x def setup_irk_weights(self): """Load and setup IRK weights.""" tmp = np.loadtxt( f"./IRK_weights/Butcher_IRK{self.config['q']}.txt", ndmin=2 ).astype(np.float32) weights = np.reshape(tmp[0:self.config['q']**2 + self.config['q']], (self.config['q']+1, self.config['q'])) self.IRK_weights = torch.tensor(weights, dtype=torch.float32).T def prepare_tensors(self): """Convert numpy arrays to PyTorch tensors.""" # Dictionary of numpy arrays that need to be converted to tensors numpy_arrays = { 'x0': self.x0, 'x1': self.x1, 'u0': self.u0, 'x_star': self.x_star } # Convert numpy arrays to tensors for name, array in numpy_arrays.items(): tensor = torch.from_numpy(array).float().to(self.device) if name in ['x0', 'x1', 'x_star']: tensor.requires_grad = True setattr(self, name, tensor) # Handle dt and IRK_weights separately as they're already tensors self.dt = self.dt.to(self.device) self.IRK_weights = self.IRK_weights.to(self.device) def setup_model(self): """Initialize the neural network model.""" self.model = NeuralNetwork( input_size=1, output_size=self.config['q'] + 1, hidden_layers=4, hidden_units=50, activation_function=nn.Tanh() ).float().to(self.device) self.model.apply(self.init_weights) @staticmethod def init_weights(m): """Initialize network weights using Xavier initialization.""" if isinstance(m, nn.Linear): torch.nn.init.xavier_normal_(m.weight) m.bias.data.fill_(0.0) @staticmethod def set_seed(seed: int = 42): """Set random seed for reproducibility.""" torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = True np.random.seed(seed) def f(self, x: torch.Tensor, x_1: torch.Tensor, dt: torch.Tensor, IRK_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute the physics-informed neural network function. Args: x: Input tensor x_1: Boundary conditions tensor dt: Time step IRK_weights: IRK weights tensor Returns: Tuple of tensors (U0, U1) """ nu = 0.01/torch.pi U1 = self.model(x) U = U1[:, :-1] # Compute derivatives U_x = self.fwd_gradients_0(U, x) U_xx = self.fwd_gradients_0(U_x, x) # Compute physics terms F = -U*U_x + nu*U_xx U0 = U1 - dt * torch.matmul(F, IRK_weights) U1 = self.model(x_1) return U0, U1 def compute_loss(self, x: torch.Tensor, x_1: torch.Tensor, dt: torch.Tensor, IRK_weights: torch.Tensor, U0_real: torch.Tensor) -> torch.Tensor: """Compute the total loss for training.""" U0, U1 = self.f(x, x_1, dt, IRK_weights) return torch.sum((U0_real - U0) ** 2) + torch.sum(U1 ** 2) def closure(self, optimizer: torch.optim.Optimizer, x: torch.Tensor, x_1: torch.Tensor, x_star: torch.Tensor, dt: torch.Tensor, IRK_weights: torch.Tensor, U0_real: torch.Tensor) -> torch.Tensor: """Closure function for LBFGS optimizer.""" optimizer.zero_grad() loss = self.compute_loss(x, x_1, dt, IRK_weights, U0_real) loss.backward(retain_graph=True) # Update iteration count and compute error self.iter += 1 U1_pred = self.model(x_star) pred = U1_pred[:, -1].detach().cpu().numpy() error = np.linalg.norm(pred - self.Exact[self.idx_t1, :], 2) / \ np.linalg.norm(self.Exact[self.idx_t1, :], 2) # Store results and save model if needed self.results.append([self.iter, loss.item(), error]) return loss def train(self, num_iter: int): """Train using L-BFGS optimizer.""" optimizer = torch.optim.LBFGS( self.model.parameters(), lr=1.0, max_iter=num_iter, max_eval=num_iter, tolerance_grad=1e-7, tolerance_change=1.0 * np.finfo(float).eps, history_size=100, line_search_fn='strong_wolfe' ) start_time = time.time() closure_fn = partial(self.closure, optimizer, self.x0, self.x1, self.x_star, self.dt, self.IRK_weights, self.u0) with tqdm(total=num_iter, desc="LBFGS Training", unit="iter") as pbar: def closure_with_tqdm(): loss = closure_fn() pbar.set_postfix(loss=loss.item()) pbar.update(1) return loss optimizer.step(closure_with_tqdm) self.lbfgs_training_time = time.time() - start_time # Save final results self.save_results() def save_results(self): """Save training results and model.""" # Save training summary total_time = getattr(self, 'lbfgs_training_time', 0) with open('results/burgers_discrete_time_training_summary.txt', 'w') as f: f.write(f"Total training time: {total_time:.6e} seconds\n") f.write(f"Total iterations: {self.iter}\n") f.write(f"Final Loss: {self.results[-1][1]:.6e}\n") f.write(f"Final L2: {self.results[-1][2]:.6e}\n") # Save training data results_array = np.array(self.results) np.savetxt("results/burgers_discrete_time_training_data.csv", results_array, delimiter=",", header="Iter,Loss,L2", comments="") # Save final model self.save_model('burgers_discrete_time.pt') def save_model(self, path: str): """Save model state dict.""" torch.save(self.model.state_dict(), path) def load_model(self, path: str): """Load model state dict.""" self.model.load_state_dict(torch.load(path)) self.model.eval() def plot_results(self): """Generate plots for the results.""" # Create figure for training curves self.plot_training_curves() # Create figure for solution visualization self.plot_solution() def plot_training_curves(self): """Plot training loss and L2 error curves.""" data = pd.read_csv('results/burgers_discrete_time_training_data.csv') fig, axarr = plt.subplots(1, 2, figsize=(12, 4)) # Loss plot axarr[0].semilogy(data['Iter'], data['Loss'], label='Loss', color='gray', linewidth=1) axarr[0].set_xlabel('Iteration') axarr[0].set_ylabel('Loss') # L2 error plot axarr[1].semilogy(data['Iter'], data['L2'], label='L2 Error', color='gray', linewidth=1) axarr[1].set_xlabel('Iteration') axarr[1].set_ylabel(r'$\mathrm{L}_2$') plt.tight_layout() plt.savefig('results/burgers_discrete_time_training_curves.pdf') plt.close() def plot_solution(self): """Plot the solution, including exact solution and predictions.""" fig = plt.figure(figsize=(12, 10)) fsize = 12 # Set up GridSpec gs0 = gridspec.GridSpec(1, 2) gs0.update(top=1-0.06, bottom=1-1/2 + 0.1, left=0.15, right=0.85, wspace=0) gs1 = gridspec.GridSpec(1, 2) gs1.update(top=1-1/2-0.05, bottom=0.15, left=0.15, right=0.85, wspace=0.5) # Plot exact solution ax = plt.subplot(gs0[:, :]) h = ax.imshow(self.Exact.T, interpolation='nearest', cmap='rainbow', extent=[self.t.min(), self.t.max(), self.x.min(), self.x.max()], origin='lower', aspect='auto') # Add colorbar divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) fig.colorbar(h, cax=cax) # Add time slice indicators line = np.linspace(self.x.min(), self.x.max(), 2)[:, None] # Fix scalar conversion by using item() t0_scalar = self.t[self.idx_t0].item() t1_scalar = self.t[self.idx_t1].item() ax.plot(t0_scalar * np.ones((2,1)), line, 'w-', linewidth=1) ax.plot(t1_scalar * np.ones((2,1)), line, 'w-', linewidth=1) ax.set_xlabel('$t$') ax.set_ylabel('$x$') ax.set_title('$u(t,x)$', fontsize=fsize) # Plot initial condition ax = plt.subplot(gs1[0, 0]) ax.plot(self.x, self.Exact[self.idx_t0,:], 'b-', linewidth=2) ax.plot(self.x0.cpu().detach().numpy(), self.u0, 'rx', linewidth=2, label='Data') ax.set_xlabel('$x$') ax.set_ylabel('$u(t,x)$') ax.set_title('$t = %.2f$' % t0_scalar, fontsize=fsize) ax.set_xlim([self.config['lb']-0.1, self.config['ub']+0.1]) ax.legend(loc='upper center', bbox_to_anchor=(0.3, -0.25), ncol=2, frameon=False) # Plot prediction vs exact solution ax = plt.subplot(gs1[0, 1]) U1_pred = self.model(self.x_star) U1_pred = U1_pred.cpu().detach().numpy() ax.plot(self.x, self.Exact[self.idx_t1,:], 'b-', linewidth=2, label='Exact') ax.plot(self.x_star.cpu().detach().numpy(), U1_pred[:,-1], 'r--', linewidth=2, label='Prediction') ax.set_xlabel('$x$') ax.set_ylabel('$u(t,x)$') ax.set_title('$t = %.2f$' % t1_scalar, fontsize=fsize) ax.set_xlim([self.config['lb']-0.1, self.config['ub']+0.1]) ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.25), ncol=2, frameon=False) plt.savefig('results/burgers_discrete_time.pdf') plt.show() plt.close() @staticmethod def fwd_gradients_0(dy: torch.Tensor, x: torch.Tensor) -> torch.Tensor: """Compute forward gradients.""" z = torch.ones(dy.shape, device=dy.device).requires_grad_() g = torch.autograd.grad(dy, x, grad_outputs=z, create_graph=True)[0] ones = torch.ones(g.shape, device=g.device) return torch.autograd.grad(g, z, grad_outputs=ones, create_graph=True)[0]

In [ ]:

def main():
    """Main function to run the DiscreteTimePINNs solver."""
    # Configuration dictionary
    config = {
        'data_path': './burgers_shock.mat',
        'q': 500,
        'N': 250,
        'lb': np.array([-1.0]),
        'ub': np.array([1.0]),
        'idx_t0': 10,
        'idx_t1': 90,
        'noise_u0': 0.0
    }
    
    # Set random seed
    DiscreteTimePINNs.set_seed(42)
    
    # Initialize solver
    solver = DiscreteTimePINNs(config)
    
    # Train the model
    solver.train(num_iter=10000)
    
    # Plot results
    solver.plot_results()
    
    # Print final error
    U1_pred = solver.model(solver.x_star)
    U1_pred = U1_pred.cpu().detach().numpy()
    error = np.linalg.norm(U1_pred[:,-1] - solver.Exact[solver.idx_t1,:], 2) / \
            np.linalg.norm(solver.Exact[solver.idx_t1,:], 2)
    print('Final Error: %e' % (error))

def main(): """Main function to run the DiscreteTimePINNs solver.""" # Configuration dictionary config = { 'data_path': './burgers_shock.mat', 'q': 500, 'N': 250, 'lb': np.array([-1.0]), 'ub': np.array([1.0]), 'idx_t0': 10, 'idx_t1': 90, 'noise_u0': 0.0 } # Set random seed DiscreteTimePINNs.set_seed(42) # Initialize solver solver = DiscreteTimePINNs(config) # Train the model solver.train(num_iter=10000) # Plot results solver.plot_results() # Print final error U1_pred = solver.model(solver.x_star) U1_pred = U1_pred.cpu().detach().numpy() error = np.linalg.norm(U1_pred[:,-1] - solver.Exact[solver.idx_t1,:], 2) / \ np.linalg.norm(solver.Exact[solver.idx_t1,:], 2) print('Final Error: %e' % (error))

In [ ]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()