gistfile1.txt

"""
1D Finite Difference Solver

This module provides a high-performance, flexible implementation of the 1D finite difference
method for solving partial differential equations. It uses advanced numerical techniques
and adheres to best practices in scientific computing and Python development.

Features:
- Supports various boundary conditions (Dirichlet, Neumann, periodic)
- Implements multiple time-stepping schemes (explicit, implicit, Crank-Nicolson)
- Utilizes sparse matrices for efficient memory usage and computation
- Provides vectorized operations for improved performance
- Includes comprehensive error checking and informative error messages
- Offers an intuitive API with type hinting for better IDE support

Author: [Your Name]
License: MIT
"""

from __future__ import annotations

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from typing import Callable, Literal, Optional, Union
from dataclasses import dataclass
from enum import Enum, auto


class BoundaryCondition(Enum):
    """Enumeration of supported boundary condition types."""
    DIRICHLET = auto()
    NEUMANN = auto()
    PERIODIC = auto()


class TimeScheme(Enum):
    """Enumeration of supported time-stepping schemes."""
    EXPLICIT = auto()
    IMPLICIT = auto()
    CRANK_NICOLSON = auto()


@dataclass
class GridParameters:
    """Dataclass to store grid parameters."""
    x_min: float
    x_max: float
    nx: int
    t_max: float
    nt: int

    @property
    def dx(self) -> float:
        """Calculate and return the spatial step size."""
        return (self.x_max - self.x_min) / (self.nx - 1)

    @property
    def dt(self) -> float:
        """Calculate and return the time step size."""
        return self.t_max / self.nt

    @property
    def x(self) -> np.ndarray:
        """Generate and return the spatial grid."""
        return np.linspace(self.x_min, self.x_max, self.nx)


class FiniteDifference1D:
    """
    A class for solving 1D partial differential equations using the finite difference method.
    """

    def __init__(
        self,
        grid_params: GridParameters,
        diffusion_coefficient: Union[float, Callable[[np.ndarray, float], np.ndarray]],
        source_term: Optional[Callable[[np.ndarray, float], np.ndarray]] = None,
        left_bc: tuple[BoundaryCondition, Union[float, Callable[[float], float]]],
        right_bc: tuple[BoundaryCondition, Union[float, Callable[[float], float]]],
        time_scheme: TimeScheme = TimeScheme.CRANK_NICOLSON,
    ):
        """
        Initialize the FiniteDifference1D solver.

        Args:
            grid_params (GridParameters): Parameters defining the computational grid.
            diffusion_coefficient (Union[float, Callable]): Diffusion coefficient (constant or function).
            source_term (Optional[Callable]): Source term function.
            left_bc (tuple): Left boundary condition (type and value/function).
            right_bc (tuple): Right boundary condition (type and value/function).
            time_scheme (TimeScheme): Time-stepping scheme to use.
        """
        self.grid = grid_params
        self.diffusion_coefficient = diffusion_coefficient
        self.source_term = source_term
        self.left_bc = left_bc
        self.right_bc = right_bc
        self.time_scheme = time_scheme

        self._validate_inputs()
        self._setup_matrices()

    def _validate_inputs(self) -> None:
        """Validate input parameters and raise informative errors if invalid."""
        if self.grid.nx < 3:
            raise ValueError("Number of spatial points must be at least 3.")
        if self.grid.nt < 1:
            raise ValueError("Number of time steps must be at least 1.")
        if self.grid.x_max <= self.grid.x_min:
            raise ValueError("x_max must be greater than x_min.")
        if self.grid.t_max <= 0:
            raise ValueError("t_max must be positive.")

        if isinstance(self.diffusion_coefficient, (int, float)) and self.diffusion_coefficient <= 0:
            raise ValueError("Diffusion coefficient must be positive.")

        if self.left_bc[0] == self.right_bc[0] == BoundaryCondition.PERIODIC:
            if self.left_bc[1] != self.right_bc[1]:
                raise ValueError("Periodic boundary conditions must have matching values.")

    def _setup_matrices(self) -> None:
        """Set up the finite difference matrices based on the chosen scheme and boundary conditions."""
        dx = self.grid.dx
        dt = self.grid.dt
        nx = self.grid.nx

        # Set up the base tridiagonal matrix
        main_diag = np.ones(nx)
        off_diag = -0.5 * np.ones(nx - 1)
        self.A = sp.diags([off_diag, main_diag, off_diag], [-1, 0, 1], format="csr")

        # Adjust for boundary conditions
        self._apply_boundary_conditions()

        # Set up the right-hand side matrix
        self.B = sp.eye(nx, format="csr")

        # Adjust matrices based on the time-stepping scheme
        if self.time_scheme == TimeScheme.EXPLICIT:
            self.A = self.B
        elif self.time_scheme == TimeScheme.IMPLICIT:
            pass  # A and B are already set up correctly
        elif self.time_scheme == TimeScheme.CRANK_NICOLSON:
            self.A = 0.5 * (self.B + self.A)
            self.B = 0.5 * (self.B - self.A)

        # Scale matrices by diffusion coefficient and time step
        if isinstance(self.diffusion_coefficient, (int, float)):
            coeff = self.diffusion_coefficient * dt / (dx ** 2)
            self.A *= coeff
            self.B *= coeff
        else:
            # For spatially varying diffusion coefficient, we'll update in the solve method
            pass

    def _apply_boundary_conditions(self) -> None:
        """Apply the specified boundary conditions to the finite difference matrix."""
        nx = self.grid.nx

        # Left boundary
        if self.left_bc[0] == BoundaryCondition.DIRICHLET:
            self.A[0, :2] = [1, 0]
        elif self.left_bc[0] == BoundaryCondition.NEUMANN:
            self.A[0, :2] = [-1, 1]

        # Right boundary
        if self.right_bc[0] == BoundaryCondition.DIRICHLET:
            self.A[-1, -2:] = [0, 1]
        elif self.right_bc[0] == BoundaryCondition.NEUMANN:
            self.A[-1, -2:] = [1, -1]

        # Periodic boundary conditions
        if self.left_bc[0] == self.right_bc[0] == BoundaryCondition.PERIODIC:
            self.A[0, -1] = self.A[-1, 0] = -0.5

    def solve(self, initial_condition: np.ndarray) -> np.ndarray:
        """
        Solve the PDE using the finite difference method.

        Args:
            initial_condition (np.ndarray): Initial condition values.

        Returns:
            np.ndarray: Solution array with shape (nt + 1, nx).
        """
        nx, nt = self.grid.nx, self.grid.nt
        dt, dx = self.grid.dt, self.grid.dx
        x = self.grid.x

        if len(initial_condition) != nx:
            raise ValueError(f"Initial condition must have {nx} elements.")

        # Initialize solution array
        u = np.zeros((nt + 1, nx))
        u[0] = initial_condition

        for n in range(nt):
            t = n * dt

            # Update matrices for spatially varying diffusion coefficient
            if callable(self.diffusion_coefficient):
                D = self.diffusion_coefficient(x, t)
                A = self.A.multiply(D[:, np.newaxis]) * dt / (dx ** 2)
                B = self.B.multiply(D[:, np.newaxis]) * dt / (dx ** 2)
            else:
                A, B = self.A, self.B

            # Compute right-hand side
            rhs = B @ u[n]

            # Add source term if present
            if self.source_term:
                rhs += self.source_term(x, t) * dt

            # Apply boundary conditions
            self._apply_time_dependent_bc(rhs, t)

            # Solve the linear system
            u[n + 1] = spla.spsolve(sp.eye(nx) - A, rhs)

        return u

    def _apply_time_dependent_bc(self, rhs: np.ndarray, t: float) -> None:
        """Apply time-dependent boundary conditions to the right-hand side."""
        if self.left_bc[0] == BoundaryCondition.DIRICHLET:
            rhs[0] = self.left_bc[1](t) if callable(self.left_bc[1]) else self.left_bc[1]
        elif self.left_bc[0] == BoundaryCondition.NEUMANN:
            rhs[0] += self.grid.dx * (self.left_bc[1](t) if callable(self.left_bc[1]) else self.left_bc[1])

        if self.right_bc[0] == BoundaryCondition.DIRICHLET:
            rhs[-1] = self.right_bc[1](t) if callable(self.right_bc[1]) else self.right_bc[1]
        elif self.right_bc[0] == BoundaryCondition.NEUMANN:
            rhs[-1] -= self.grid.dx * (self.right_bc[1](t) if callable(self.right_bc[1]) else self.right_bc[1])

# Example usage
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    # Set up problem parameters
    L = 1.0  # Domain length
    T = 0.1  # Total time
    nx = 101  # Number of spatial points
    nt = 1000  # Number of time steps

    grid_params = GridParameters(0, L, nx, T, nt)

    # Define problem-specific functions
    def initial_condition(x):
        return np.sin(np.pi * x / L)

    def analytical_solution(x, t):
        return np.exp(-np.pi**2 * t) * np.sin(np.pi * x / L)

    # Create and solve the problem
    fd_solver = FiniteDifference1D(
        grid_params,
        diffusion_coefficient=1.0,
        left_bc=(BoundaryCondition.DIRICHLET, 0),
        right_bc=(BoundaryCondition.DIRICHLET, 0),
        time_scheme=TimeScheme.CRANK_NICOLSON,
    )

    solution = fd_solver.solve(initial_condition(grid_params.x))

    # Plot results
    plt.figure(figsize=(10, 6))
    plt.plot(grid_params.x, solution[-1], 'b-', label='Numerical')
    plt.plot(grid_params.x, analytical_solution(grid_params.x, T), 'r--', label='Analytical')
    plt.xlabel('x')
    plt.ylabel('u')
    plt.title(f'1D Heat Equation Solution at t = {T}')
    plt.legend()
    plt.grid(True)
    plt.show()

    # Compute and print error
    error = np.abs(solution[-1] - analytical_solution(grid_params.x, T)).max()
    print(f"Maximum absolute error: {error:.6e}")