Code for project2 + remove ci

Project2/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 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 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 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 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)]
+
+        # 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)
+            # no exchange
+            u[1:] = u_old[1:] - self.dt * (
+                1 / self.dx * (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)):
+        '''
+        Plot the solution
+        '''
+        fig, axs = plt.subplots(ncols=3, nrows=1, figsize=(15, 5))
+
+        # 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[1].set_title('Solution of $u_t + f(u)_x = 0$')
+        axs[1].plot(self.x, self.u_vec[0], label='$u_0(x)$', color='red', ls='--')
+        axs[1].plot(self.x, self.u_vec[-1], label='u(x,t)', color='red', ls='-')
+        axs[1].legend()
+        axs[1].grid()
+        axs[1].set_xlabel('x')
+        axs[1].set_ylabel('U(x)')
+
+        # Flux function (initial and final)
+        axs[2].set_title('Flux $f(u(x,t))$')
+        axs[2].plot(self.x, self.J[0], label='$f(u_0(x))$', color='green', ls='--')
+        axs[2].plot(self.x, self.J[-1], label='$f(u(x,t))$', color='green', ls='-')
+        axs[2].legend()
+        axs[2].grid()
+        axs[2].set_xlabel('x')
+        axs[2].set_ylabel('f(u)')
+
+        fig.tight_layout()
+
+        return fig, axs
\ No newline at end of file