Diffusion-Reaction Class Implementation

# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt
from pymrm import construct_grad, construct_div, newton, construct_coefficient_matrix, NumJac, non_uniform_grid

class DiffusionReaction:
    """
    A class to model diffusion-reaction systems.
    
    This class encapsulates the parameters, initialization, and solution methods for a diffusion-reaction model.
    It supports multi-component systems and allows for modular and reusable code.
    """
    def __init__(self, D=[[1.0, 1.0, 1.0]], nu=2, c_b=[[1.0, 1.0, 0.0]], k=1.0, dt=np.inf, freq_out=1):
        """
        Initialize the diffusion-reaction model.
        
        Parameters:
        - D (list): Diffusion coefficients for each component.
        - nu (int): Geometric factor for divergence operator.
        - c_b (list): Boundary conditions for each component.
        - k (float): Reaction rate constant.
        - dt (float): Time step size.
        - freq_out (int): Frequency of output during simulation.
        """
        # Initialize model parameters
        self.D = np.asarray(D)  # Diffusion coefficients
        self.nu = nu  # Geometric factor for divergence operator
        self.R = 1.0  # Domain length
        self.num_r = 30  # Number of grid points
        self.num_c = 3  # Number of components (e.g., species)
        self.bc = ({'a': 1, 'b': 0, 'd': 0},  # Left boundary conditions
                   {'a': 0, 'b': 1, 'd': c_b})  # Right boundary conditions
        self.k = k  # Reaction rate constant
        self.dt = dt  # Time step size
        self.freq_out = freq_out  # Frequency of output (e.g., every 10 steps)

        # Generate a non-uniform grid for cell-face positions
        dr_large = 0.1 * self.R
        self.r_f = non_uniform_grid(0, self.R, self.num_r + 1, dr_large, 0.75)  # Cell-face positions
        self.r_c = 0.5 * (self.r_f[:-1] + self.r_f[1:])  # Cell-centered grid points

        # Initialize the concentration field and Jacobian matrix
        self.init_field()
        self.init_jac()

    def init_field(self, c0=0.0):
        """
        Initialize the concentration field with a uniform value.
        
        Parameters:
        - c0 (float or array): Initial concentration value(s).
        """
        c = np.asarray(c0)
        shape = (1,) * (2 - c.ndim) + c.shape  # Ensure c is 2D
        c = c.reshape(shape)  # Reshape c to 2D
        self.c = np.broadcast_to(c, (self.num_r, self.num_c)).copy()  # Broadcast c to the correct shape

    def set_diff_field(self):
        """
        Define a position-dependent diffusion field.
        
        Returns:
        - diff_field (array): Diffusion field values.
        """
        diff_field = (1.0 + 4.0 * (self.r_f > 0.5 * self.R)).reshape((-1, 1)) * np.asarray(self.D)
        return diff_field

    def init_jac(self):
        """
        Construct the Jacobian matrix and constant terms for the system.
        """
        # Construct gradient and divergence operators
        grad_mat, grad_bc = construct_grad(self.c.shape, self.r_f, self.r_c, self.bc, axis=0)
        div_mat = construct_div(self.c.shape, self.r_f, nu=self.nu, axis=0)

        # Construct diffusion coefficient matrix
        diff_mat = construct_coefficient_matrix(self.D, shape=(self.num_r, self.num_c), axis=0)

        # Compute flux terms
        flux_mat = -diff_mat @ grad_mat
        flux_bc = -diff_mat @ grad_bc

        # Compute Jacobian of diffusion term and boundary forcing term
        jac_diff = div_mat @ flux_mat
        f_diff_bc = div_mat @ flux_bc

        # Accumulation term for time-dependent problems
        jac_accum = construct_coefficient_matrix(1.0 / self.dt, shape=(self.num_r, self.num_c))

        # Precompute constant terms for the residual and Jacobian
        self.g_const = f_diff_bc
        self.jac_const = jac_accum + jac_diff

        # Apparent reaction rate terms
        self.r_apparent_mat = -((self.nu + 1) / self.R) * flux_mat[-self.num_c:, :]
        self.r_apparent_bc = -((self.nu + 1) / self.R) * flux_bc[-self.num_c:, 0].reshape((-1, 1))

        # Numerical Jacobian for reaction terms
        self.numjac = NumJac((self.num_r, self.num_c))

    def set_reaction_rate(self, k):
        """
        Update the reaction rate constant.
        
        Parameters:
        - k (float): New reaction rate constant.
        """
        self.k = k

    def set_c_b(self, c_b):
        """
        Update the boundary conditions for the system.
        
        Parameters:
        - c_b (list): New boundary condition values.
        """
        self.bc[1]['d'] = c_b
        self.init_jac()  # Reinitialize the Jacobian matrix with the new boundary condition

    def reaction_kinetics(self, c):
        """
        Compute the reaction kinetics for the given concentration.
        
        Parameters:
        - c (array): Concentration values.
        
        Returns:
        - array: Reaction rates for each component.
        """
        r = self.k * c[..., [0]] * c[..., [1]]
        return r * np.array([[-1.0, -1.0, 1.0]])

    def residual(self, c, c_old):
        """
        Compute the residual vector and Jacobian matrix for the current time step.
        
        Parameters:
        - c (array): Current concentration values.
        - c_old (array): Concentration values from the previous time step.
        
        Returns:
        - tuple: Residual vector and Jacobian matrix.
        """
        g_react, jac_react = self.numjac(self.reaction_kinetics, c)
        g = self.g_const + self.jac_const @ c.reshape((-1, 1)) - c_old.reshape((-1, 1)) / self.dt - g_react.reshape((-1, 1))
        jac = self.jac_const - jac_react
        return g, jac

    def solve(self, num_timesteps=1, callback=None):
        """
        Solve the system iteratively over a specified number of time steps.
        
        Parameters:
        - num_timesteps (int): Number of time steps to simulate.
        - callback (function): Optional callback function for custom operations (e.g., plotting).
        """
        for i in range(num_timesteps):
            c_old = self.c.copy()

            # Solve the system using Newton's method
            result = newton(lambda c: self.residual(c, c_old), c_old, maxfev=10)
            self.c[...] = result.x.reshape((self.c.shape))

            # Invoke the callback function if provided
            if i % self.freq_out == 0 and callback is not None:
                callback(i, self)

    def compute_apparent_reaction_rate(self):
        """
        Compute the apparent reaction rate based on the current concentration field.
        
        Returns:
        - array: Apparent reaction rates for each component.
        """
        r_apparent = -(self.r_apparent_mat @ self.c.reshape((-1, 1)) + self.r_apparent_bc)
        return r_apparent.reshape((self.num_c,))

# Create an instance of the diffusion-reaction model
c_b = np.array([2.0, 1.0, 0.0])  # Boundary condition for species B
mrm_model = DiffusionReaction(k=10, c_b=c_b, D=[[1.0, 1.0, 1.0]], nu=3)  # Initialize the model with parameters

# Define a callback function for plotting the concentration profile
def plot_callback(step, model):
    plt.plot(model.r_c, model.c[:, 0], 'r-', label='A')
    plt.plot(model.r_c, model.c[:, 1], 'b-', label='B')
    if model.num_c >= 3:
        plt.plot(model.r_c, model.c[:, 2], 'g-', label='C')
    plt.title(f'Time step: {step * model.dt:.2f}')
    plt.xlabel('Position')
    plt.ylabel('Concentration')
    plt.legend()

# Solve the problem and use the callback function for plotting
mrm_model.solve(callback=plot_callback)

# Compute and print intrinsic and apparent reaction rates
r_intrinsic = mrm_model.reaction_kinetics(c_b).ravel()
print(f"Intrinsic reaction rate at c_A = {c_b[0]:>7.3f} and c_B = {c_b[1]:>7.3f}: r_A = {r_intrinsic[0]:>7.3f}, r_B = {r_intrinsic[1]:>7.3f}, r_C = {r_intrinsic[2]:>7.3f}")
r_apparent = mrm_model.compute_apparent_reaction_rate().ravel()
print(f"Apparent  reaction rate at c_A = {c_b[0]:>7.3f} and c_B = {c_b[1]:>7.3f}: r_A = {r_apparent[0]:>7.3f}, r_B = {r_apparent[1]:>7.3f}, r_C = {r_apparent[2]:>7.3f}")