Add random k(x) to algorithm and setup Project3

Project3/src/SystemGeneric.py

+'''
+Jan Habscheid
+Jan.Habscheid@rwth-aachen.de
+'''
+
+import matplotlib.pyplot as plt
+import numpy as np 
+
+class System:
+    '''
+    System to solve the (un)coupled transient diffusion-reaction equations
+    '''
+    def __init__(self, Nx:int, x0:float, x1:float, NTime:int, T_end=float):
+        '''
+        Constructor for the System class
+
+        All parameters are in dimensionless units
+
+        Parameters
+        ----------
+        Nx : int
+            Number of grid points
+        x0 : float
+            Left boundary of the domain
+        x1 : float
+            Right boundary of the domain
+        NTime : int
+            Number of time steps
+        T_end : float
+            End time of simulation
+        '''
+        # Store parameter
+        self.Nx = Nx
+        self.x0 = x0
+        self.x1 = x1
+        self.dx = (x1 - x0) / Nx
+        self.x = np.linspace(x0, x1, Nx)
+
+        self.NTime = NTime
+        self.T_end = T_end
+        self.dt = T_end / NTime
+
+    def get_x(self) -> np.ndarray:
+        '''
+        Creates the spatial grid
+
+        Returns
+        -------
+        np.ndarray
+            Spatial grid
+        '''
+        return self.x
+
+    def set_u0(self, u0:np.ndarray):
+        '''
+        Set the initial condition
+
+        Parameters
+        ----------
+        u0 : np.ndarray
+            Initial condition
+        '''
+        self.u0 = u0
+
+    def set_u_analatical(self, u_analytical:callable):
+        '''
+        Set the analytical solution
+
+        Parameters
+        ----------
+        u_analytical : callable
+            Analytical solution
+        '''
+        self.u_analytical = u_analytical
+
+    def get_u_analytical(self, t:float) -> np.ndarray:
+        '''
+        Get the analytical solution
+
+        Parameters
+        ----------
+        t : float
+            Time
+
+        Returns
+        -------
+        np.ndarray
+            Analytical solution at time t
+        '''
+        return self.u_analytical(self.x, t)
+
+    def plot_u0(self):
+        '''
+        Plot the initial condition
+        '''
+        fig, axs = plt.subplots()
+        axs.plot(self.x, self.u0, color='blue')
+        axs.grid()
+        axs.set_xlabel('x')
+        axs.set_ylabel('u')
+        fig.tight_layout()
+
+        return fig, axs
+
+    def set_f(self, f:callable):
+        '''
+        Set the source term
+
+        Parameters
+        ----------
+        f : np.ndarray
+            Source term
+        '''
+        self.f = f
+
+    def set_k(self, k:callable):
+        '''
+        Set the offset
+
+        Parameters
+        ----------
+        k : callable
+            Offset
+        '''
+        self.k = k
+
+    def plot_f(self, u):
+        '''
+        Plot the source term
+
+        Parameters
+        ----------
+        u : np.ndarray
+            Solution at time t
+        '''
+        fig, axs = plt.subplots()
+        axs.plot(u, self.f(u), color='red')
+        axs.grid()
+        axs.set_xlabel('x')
+        axs.set_ylabel('f(u)')
+        fig.tight_layout()
+
+        return fig, axs
+    
+    def plot_k(self):
+        '''
+        Plot the offset
+        '''
+        x = self.x
+
+        fig, axs = plt.subplots()
+        axs.plot(x, self.k(x), color='red')
+        axs.grid()
+        axs.set_xlabel('x')
+        axs.set_ylabel('k(x)')
+        fig.tight_layout()
+
+        return fig, axs
+    
+    def Upwind(self, u_old:np.array, M:float=1) -> tuple:
+        '''
+        Upwind scheme for the transient diffusion-reaction equation
+
+        Parameters
+        ----------
+        u_old : np.array
+            Solution at time t
+        M : float
+            Diffusion coefficient
+
+        Returns
+        -------
+        tuple
+            Fluxes at time t
+        '''
+        F = np.zeros_like(u_old)
+        F[:-1] = .5 * (
+            self.f(u_old[:-1]) + self.f(u_old[1:])
+            - M * (u_old[1:] - u_old[:-1])
+        )
+        F[-1] = self.f(u_old[-1])  # Ensure boundary condition correctness
+
+        return F[1:], F[:-1]
+    
+    def k_central(self) -> np.ndarray:
+        '''
+        Mean of two neighboring cells for k
+
+        Returns
+        -------
+        tuple
+            k_offset vector
+        '''
+        k = self.k(self.x)
+
+        return .5 * (k[1:] + k[:-1])
+
+    def check_cfl(self):
+        '''
+        Check the CFL condition
+        '''
+        a = np.max(np.abs(self.f(self.u0)))
+        cfl = a * self.dt / self.dx
+        return cfl, cfl < 1
+    
+    def solve(self, M:float=1):
+        '''
+        Solve the transient diffusion-reaction equation
+
+        Parameters
+        ----------
+        M : float
+            Diffusion coefficient
+
+        Returns
+        -------
+        np.ndarray
+            Solution at time t
+        '''
+        # Initial condition
+        u_vec = [self.u0.copy()]
+        u = self.u0.copy()
+        t = .0
+        t_vec = [t]
+
+        J = [self.f(self.u0)]
+
+        # Check stability
+        cfl, cfl_cond = self.check_cfl()
+        print(f'CFL Condition: {cfl_cond}, CFL number: {cfl}')
+        if not cfl_cond:
+            raise ValueError('CFL condition not satisfied, CFL number is: {}'.format(cfl))
+
+        # Time loop
+        for j in range(self.NTime):
+            # Store old values
+            u_old = u_vec[-1]
+            # Get Upwind
+            (upwind_n_1, upwind_n_2) = self.Upwind(u_old, M=M)
+            k_central_ = self.k_central()
+            # no exchange
+            u[1:] = u_old[1:] - self.dt * (
+                1 / self.dx * k_central_ * (upwind_n_1 - upwind_n_2)
+            )
+            J.append(self.f(u))
+
+            # Apply BCs
+            u[0] = u_old[0]
+
+            # Store
+            u_vec.append(u.copy())
+
+            t_vec.append(t)
+            t += self.dt
+
+        # Store solutions
+        self.t_vec = np.array(t_vec)
+        self.u_vec = np.array(u_vec)
+        self.J = np.array(J)
+
+    def plot_solution(self, u=np.linspace(0,1,101), analytical=False):
+        '''
+        Plot the solution
+
+        Parameters
+        ----------
+        u : np.ndarray
+            Solution at time t
+        analytical : bool
+            Plot analytical solution
+
+        Returns
+        -------
+        tuple
+            Figure and axis
+        '''
+        fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(15, 5))
+
+        fig.suptitle('Solution at time t = {}'.format(self.T_end))
+
+        # # Function u on first axis
+        # axs[0].set_title('Function f(u)')
+        # axs[0].plot(u, self.f(u), color='blue')
+        # axs[0].grid()
+        # axs[0].set_xlabel('u')
+        # axs[0].set_ylabel('f(u)')
+
+        # Initial Distribution and final distribution
+        axs[0].set_title('Solution of $u_t + f(u)_x = 0$')
+        axs[0].plot(self.x, self.u_vec[0], label='$u_0(x)$', color='red', ls='--')
+        axs[0].plot(self.x, self.u_vec[-1], label='u(x,t)', color='red', ls='-')
+        if analytical:
+            axs[0].plot(self.x, self.get_u_analytical(self.T_end), label='u(x,t) analytical', color
+            ='blue', ls=':')
+        axs[0].legend(loc='best')
+        axs[0].grid()
+        axs[0].set_xlabel('$x$')
+        axs[0].set_ylabel('$u(x)$')
+
+        # Flux function (initial and final)
+        axs[1].set_title('Flux $f(u(x,t))$')
+        axs[1].plot(self.x, self.J[0], label='$f(u_0(x))$', color='green', ls='--')
+        axs[1].plot(self.x, self.J[-1], label='$f(u(x,t))$', color='green', ls='-')
+        axs[1].legend()
+        axs[1].grid()
+        axs[1].set_xlabel('$u$')
+        axs[1].set_ylabel('$f(u)$')
+
+        fig.tight_layout()
+
+        return fig, axs
+    
+    def plot_solution_3D(self, analytical=False):
+        '''
+        Plot the solution in 3D
+
+        Parameters
+        ----------
+        analytical : bool
+            Plot analytical solution
+        
+        Returns
+        -------
+        tuple
+            Figure and axis
+        '''
+        fig = plt.figure(figsize=(15, 5))
+
+        # Initial Distribution and final distribution
+        ax0 = fig.add_subplot(131, projection='3d')
+        ax0.set_title('Numerical solution of $u_t + f(u)_x = 0$')
+        X, T = np.meshgrid(self.x, self.t_vec)
+        ax0.plot_surface(X, T, self.u_vec, cmap='coolwarm')
+        ax0.set_xlabel('$x$')
+        ax0.set_ylabel('$t$')
+        ax0.set_zlabel('$u(x)$')
+        ax0.grid()
+
+        # Analytical solution
+        if analytical:
+            ax1 = fig.add_subplot(132, projection='3d')
+            ax1.set_title('Analytical solution of $u_t + f(u)_x = 0$')
+            ax1.plot_surface(X, T, np.array([self.get_u_analytical(t) for t in self.t_vec]), cmap='coolwarm')
+            ax1.set_xlabel('$x$')
+            ax1.set_ylabel('$t$')
+            ax1.set_zlabel('$u(x)$')
+            ax1.grid()
+
+        # # Flux function (initial and final)
+        # ax2 = fig.add_subplot(133, projection='3d')
+        # ax2.set_title('Flux $f(u(x,t))$')
+        # ax2.plot_surface(X, T, self.J, cmap='coolwarm')
+        # ax2.set_xlabel('x')
+        # ax2.set_ylabel('t')
+        # ax2.set_zlabel('f(u)')
+        # ax2.grid()
+        
+        # Abs Error
+        u_numerical = np.array(self.u_vec)
+        u_analytical = np.array([self.get_u_analytical(t) for t in self.t_vec])
+        abs_error = np.abs(u_numerical - u_analytical)
+        title = 'Absolute error: $|u_{num} - u_{exact}|$'
+
+        ax2 = fig.add_subplot(133, projection='3d')
+        ax2.plot_surface(X, T, abs_error, cmap='coolwarm')
+        ax2.grid()
+        ax2.set_title(title)
+        ax2.set_xlabel('x')
+        ax2.set_ylabel('t')
+        ax2.set_zlabel('Error')
+
+        fig.tight_layout()
+
+        return fig, [ax0, ax1, ax2] if analytical else [ax0, ax2]
+    
+    def plot_error(self, type='rel'):
+        '''
+        Plot the error
+
+        Parameters
+        ----------
+        type : str
+            Type of error to plot
+            Choose from 'rel' or 'abs'
+
+        Returns
+        -------
+        tuple
+            Figure and axis
+        '''
+        if type not in ['rel', 'abs']:
+            raise ValueError('Type must be either "rel" or "abs"')
+        
+        X, T = np.meshgrid(self.x, self.t_vec)
+        u_numerical = np.array(self.u_vec)
+        u_analytical = np.array([self.get_u_analytical(t) for t in self.t_vec])
+        abs_error = np.abs(u_numerical - u_analytical)
+        rel_error = abs_error / np.abs(u_analytical)
+        error = rel_error if type == 'rel' else abs_error
+        title = 'Relative error' if type == 'rel' else 'Absolute error'
+
+        fig = plt.figure(figsize=(5, 5))
+        ax = fig.add_subplot(111, projection='3d')
+        ax.plot_surface(X, T, error, cmap='coolwarm')
+        ax.grid()
+        ax.set_title(title)
+        ax.set_xlabel('x')
+        ax.set_ylabel('t')
+        ax.set_zlabel('Error')
+
+        fig.tight_layout()
+
+        return fig, ax
\ No newline at end of file