lagrangian.py

from enum import Enum
from typing import List, Tuple, Optional, Union, Callable
import numpy as np
from dataclasses import dataclass
from abc import ABC, abstractmethod


class Geometry(ABC):
    """Abstract base class for geometric configurations."""
    
    @property
    @abstractmethod
    def alpha(self) -> int:
        """Return the geometry parameter alpha."""
        pass
    
    @property
    @abstractmethod
    def geometric_constant(self) -> float:
        """Return the geometry-dependent constant (Cₐ)."""
        pass
    
    @abstractmethod
    def calculate_volume(self, r_inner: float, r_outer: float) -> float:
        """Calculate volume between two radial positions."""
        pass
    
    @abstractmethod
    def calculate_area(self, r: float) -> float:
        """Calculate area at a given radial position."""
        pass
    
    @abstractmethod
    def calculate_mass(self, r_inner: float, r_outer: float, rho: float) -> float:
        """Calculate mass between two radial positions with given density."""
        pass


class PlanarGeometry(Geometry):
    """Planar (Cartesian) geometry implementation."""
    
    @property
    def alpha(self) -> int:
        return 0
    
    @property
    def geometric_constant(self) -> float:
        return 1.0
    
    def calculate_volume(self, r_inner: float, r_outer: float) -> float:
        """Calculate volume (area in 1D) between two positions."""
        return r_outer - r_inner
    
    def calculate_area(self, r: float) -> float:
        """Calculate area (constant in planar geometry)."""
        return self.geometric_constant
    
    def calculate_mass(self, r_inner: float, r_outer: float, rho: float) -> float:
        """Calculate mass between two positions with given density."""
        return rho * self.calculate_volume(r_inner, r_outer)


class CylindricalGeometry(Geometry):
    """Cylindrical geometry implementation."""
    
    @property
    def alpha(self) -> int:
        return 1
    
    @property
    def geometric_constant(self) -> float:
        return 2.0 * np.pi
    
    def calculate_volume(self, r_inner: float, r_outer: float) -> float:
        """Calculate volume between two radial positions."""
        return self.geometric_constant * 0.5 * (r_outer**2 - r_inner**2)
    
    def calculate_area(self, r: float) -> float:
        """Calculate area at a given radius."""
        return self.geometric_constant * r
    
    def calculate_mass(self, r_inner: float, r_outer: float, rho: float) -> float:
        """Calculate mass between two radial positions with given density."""
        return rho * self.calculate_volume(r_inner, r_outer)


class SphericalGeometry(Geometry):
    """Spherical geometry implementation."""
    
    @property
    def alpha(self) -> int:
        return 2
    
    @property
    def geometric_constant(self) -> float:
        return 4.0 * np.pi
    
    def calculate_volume(self, r_inner: float, r_outer: float) -> float:
        """Calculate volume between two radial positions."""
        return self.geometric_constant * (r_outer**3 - r_inner**3) / 3.0
    
    def calculate_area(self, r: float) -> float:
        """Calculate area at a given radius."""
        return self.geometric_constant * r**2
    
    def calculate_mass(self, r_inner: float, r_outer: float, rho: float) -> float:
        """Calculate mass between two radial positions with given density."""
        return rho * self.calculate_volume(r_inner, r_outer)


@dataclass
class SimulationParameters:
    """Parameters for the simulation."""
    gamma: float = 1.4          # Ratio of specific heats for ideal gas
    cfl: float = 0.5            # CFL number for time step calculation
    t_end: float = 0.1          # End time for simulation
    output_frequency: int = 10  # How often to save results (in timesteps)
    use_artificial_viscosity: bool = True


class LagrangianMesh:
    """Representation of the computational mesh in mass coordinates."""
    
    def __init__(
        self, 
        r_initial: np.ndarray,
        rho_initial: np.ndarray,
        u_initial: np.ndarray,
        e_initial: np.ndarray,
        geometry: Geometry,
        gamma: float
    ) -> None:
        """
        Initialize the Lagrangian mesh.
        
        Args:
            r_initial: Initial positions of cell interfaces
            rho_initial: Initial densities in each cell
            u_initial: Initial velocities at cell interfaces
            e_initial: Initial specific internal energies in each cell
            geometry: Geometric configuration
            gamma: Ratio of specific heats for equation of state
        """
        self.geometry = geometry
        self.gamma = gamma
        
        # Number of cells and interfaces
        self.n_cells = len(rho_initial)
        self.n_interfaces = len(r_initial)
        
        # State variables
        self.r = r_initial.copy()  # Interface positions
        self.u = u_initial.copy()  # Interface velocities
        
        # Cell-centered quantities
        self.rho = rho_initial.copy()     # Density
        self.e = e_initial.copy()         # Specific internal energy
        self.V = np.zeros(self.n_cells)   # Specific volume (1/rho)
        self.p = np.zeros(self.n_cells)   # Pressure
        self.p_eff = np.zeros(self.n_cells)  # Effective pressure (with artificial viscosity)
        
        # Compute derived quantities
        self.compute_cell_masses()
        self.compute_volumes()
        self.compute_pressure()
    
    def compute_cell_masses(self) -> None:
        """Compute cell masses (constant in time for Lagrangian scheme)."""
        self.m = np.zeros(self.n_cells)
        
        for i in range(self.n_cells):
            self.m[i] = self.geometry.calculate_mass(
                self.r[i], self.r[i+1], self.rho[i]
            )
    
    def compute_volumes(self) -> None:
        """Compute cell volumes and specific volumes."""
        for i in range(self.n_cells):
            # Calculate volume using geometry
            volume = self.geometry.calculate_volume(self.r[i], self.r[i+1])
            
            # Specific volume (V = volume/mass)
            self.V[i] = volume / self.m[i]
        
        # Update density (ρ = 1/V)
        self.rho = 1.0 / self.V
    
    def compute_pressure(self) -> None:
        """Compute pressures using equation of state."""
        # Ideal gas equation of state: p = (γ-1)ρe = (γ-1)e/V
        self.p = (self.gamma - 1.0) * self.e / self.V
        
        # Initialize effective pressure (without artificial viscosity)
        self.p_eff = self.p.copy()


class ArtificialViscosity:
    """Artificial viscosity for shock capturing."""
    
    @staticmethod
    def calculate(mesh: LagrangianMesh) -> np.ndarray:
        """
        Calculate artificial viscosity terms for each cell.
        
        Args:
            mesh: The computational mesh
            
        Returns:
            Array of artificial viscosity terms
        """
        q = np.zeros(mesh.n_cells)
        
        for i in range(mesh.n_cells):
            # Velocity divergence (compression rate)
            div_u = (mesh.u[i+1] - mesh.u[i]) / (mesh.r[i+1] - mesh.r[i])
            
            # Only apply artificial viscosity in compression regions
            if div_u < 0:
                # Cell width as length scale
                l = mesh.r[i+1] - mesh.r[i]
                
                # Sound speed
                c = np.sqrt(mesh.gamma * mesh.p[i] * mesh.V[i])
                
                # Von Neumann-Richtmyer artificial viscosity with linear term
                q[i] = mesh.rho[i] * (
                    l**2 * div_u**2 +   # Quadratic term (von Neumann-Richtmyer)
                    l * c * abs(div_u)   # Linear term (adds stability)
                )
        
        return q


class LagrangianSolver:
    """Solver for the Lagrangian formulation of Euler equations."""
    
    def __init__(
        self,
        mesh: LagrangianMesh,
        params: SimulationParameters
    ) -> None:
        """
        Initialize the Lagrangian solver.
        
        Args:
            mesh: The computational mesh
            params: Simulation parameters
        """
        self.mesh = mesh
        self.params = params
        self.time = 0.0
        self.step_count = 0
    
    def calculate_time_step(self) -> float:
        """
        Calculate stable time step using CFL condition.
        
        Returns:
            Time step size
        """
        dt_values = []
        
        for i in range(self.mesh.n_cells):
            # Sound speed
            c = np.sqrt(self.mesh.gamma * self.mesh.p[i] * self.mesh.V[i])
            
            # Cell width
            dx = self.mesh.r[i+1] - self.mesh.r[i]
            
            # CFL condition: dt ≤ dx / (|u| + c)
            u_max = max(abs(self.mesh.u[i]), abs(self.mesh.u[i+1]))
            dt_cell = self.params.cfl * dx / (u_max + c)
            dt_values.append(dt_cell)
        
        return min(dt_values)
    
    def apply_boundary_conditions(self) -> None:
        """Apply boundary conditions to velocities."""
        # Example: fixed wall boundaries (zero velocity)
        self.mesh.u[0] = 0.0
        self.mesh.u[-1] = 0.0
    
    def update_velocities(self, dt: float) -> None:
        """
        Update interface velocities using the momentum equation.
        
        Args:
            dt: Time step
        """
        # Copy current velocities
        u_new = self.mesh.u.copy()
        
        # Update interior interfaces
        for i in range(1, self.mesh.n_interfaces - 1):
            # Area at interface
            area = self.mesh.geometry.calculate_area(self.mesh.r[i])
            
            # Average mass for pressure gradient calculation
            # Interface i is between cells (i-1) and i
            m_avg = 0.5 * (self.mesh.m[i-1] + self.mesh.m[i])
            
            # Pressure gradient term
            dp_dm = (self.mesh.p_eff[i] - self.mesh.p_eff[i-1]) / m_avg
            
            # Update velocity: du/dt = -area * dp/dm
            u_new[i] = self.mesh.u[i] - dt * area * dp_dm
        
        # Apply boundary conditions
        self.mesh.u = u_new
        self.apply_boundary_conditions()
    
    def update_positions(self, dt: float) -> None:
        """
        Update interface positions using velocities.
        
        Args:
            dt: Time step
        """
        for i in range(self.mesh.n_interfaces):
            self.mesh.r[i] += dt * self.mesh.u[i]
    
    def update_energy(self, dt: float) -> None:
        """
        Update specific internal energy using the energy equation.
        
        Args:
            dt: Time step
        """
        # Store old specific volumes
        V_old = self.mesh.V.copy()
        
        # Update volumes with new positions
        self.mesh.compute_volumes()
        
        # Update energy: de/dt = -p * dV/dt
        for i in range(self.mesh.n_cells):
            dV = self.mesh.V[i] - V_old[i]
            self.mesh.e[i] -= self.mesh.p_eff[i] * dV
    
    def step(self) -> float:
        """
        Advance the solution by one time step.
        
        Returns:
            Size of time step taken
        """
        # Calculate time step
        dt = self.calculate_time_step()
        
        # Limit time step to not exceed end time
        if self.time + dt > self.params.t_end:
            dt = self.params.t_end - self.time
        
        # Calculate artificial viscosity if enabled
        if self.params.use_artificial_viscosity:
            q = ArtificialViscosity.calculate(self.mesh)
            self.mesh.p_eff = self.mesh.p + q
        else:
            self.mesh.p_eff = self.mesh.p.copy()
        
        # Update variables in sequence:
        
        # 1. Update velocities using momentum equation
        self.update_velocities(dt)
        
        # 2. Update positions using new velocities
        self.update_positions(dt)
        
        # 3. Update energy using energy equation
        self.update_energy(dt)
        
        # 4. Update pressure using equation of state
        self.mesh.compute_pressure()
        
        # Update simulation time and step count
        self.time += dt
        self.step_count += 1
        
        return dt
    
    def solve(self) -> Tuple[np.ndarray, List[np.ndarray]]:
        """
        Solve the Euler equations until the end time.
        
        Returns:
            Tuple of (time_history, solution_history)
            where solution_history is a list of (r, u, rho, e, p) at each saved time
        """
        # Initialize result storage
        time_history = [self.time]
        solution_history = [(
            self.mesh.r.copy(),
            self.mesh.u.copy(),
            self.mesh.rho.copy(),
            self.mesh.e.copy(),
            self.mesh.p.copy()
        )]
        
        # Main simulation loop
        while self.time < self.params.t_end:
            # Take a time step
            dt = self.step()
            
            # Save results at specified intervals
            if self.step_count % self.params.output_frequency == 0 or self.time >= self.params.t_end:
                time_history.append(self.time)
                solution_history.append((
                    self.mesh.r.copy(),
                    self.mesh.u.copy(),
                    self.mesh.rho.copy(),
                    self.mesh.e.copy(),
                    self.mesh.p.copy()
                ))
                
                # Print progress
                print(f"Time: {self.time:.6f}, Step: {self.step_count}, dt: {dt:.6f}")
        
        return np.array(time_history), solution_history


def create_geometry(geometry_type: str) -> Geometry:
    """
    Factory method to create geometry object.
    
    Args:
        geometry_type: Type of geometry ("planar", "cylindrical", "spherical")
        
    Returns:
        Geometry object
    """
    if geometry_type.lower() in ["planar", "cartesian", "plane"]:
        return PlanarGeometry()
    elif geometry_type.lower() in ["cylindrical", "cylinder"]:
        return CylindricalGeometry()
    elif geometry_type.lower() in ["spherical", "sphere"]:
        return SphericalGeometry()
    else:
        raise ValueError(f"Unsupported geometry type: {geometry_type}")


def run_sedov_blast_wave(
    n_cells: int = 100,
    geometry_type: str = "spherical",
    gamma: float = 1.4,
    t_end: float = 0.1,
    cfl: float = 0.3
) -> Tuple[np.ndarray, List[np.ndarray]]:
    """
    Run a Sedov blast wave problem.
    
    Args:
        n_cells: Number of cells
        geometry_type: Type of geometry ("planar", "cylindrical", "spherical")
        gamma: Ratio of specific heats
        t_end: End time for simulation
        cfl: CFL number
        
    Returns:
        Tuple of (time_history, solution_history)
    """
    # Create geometry
    geometry = create_geometry(geometry_type)
    
    # Domain size
    r_max = 1.0
    
    # Setup mesh
    r = np.linspace(0.0, r_max, n_cells + 1)  # Interface positions
    
    # Initial conditions
    rho = np.ones(n_cells)  # Uniform density
    u = np.zeros(n_cells + 1)  # Zero initial velocity
    e = np.zeros(n_cells)  # Zero initial energy
    
    # Add energy to central region (point blast)
    total_energy = 1.0
    
    # Put all energy in first cell
    volume_first_cell = geometry.calculate_volume(r[0], r[1])
    e[0] = total_energy / (rho[0] * volume_first_cell)
    
    # Create mesh and solver
    mesh = LagrangianMesh(r, rho, u, e, geometry, gamma)
    
    params = SimulationParameters(
        gamma=gamma,
        cfl=cfl,
        t_end=t_end,
        output_frequency=10,
        use_artificial_viscosity=True
    )
    
    solver = LagrangianSolver(mesh, params)
    
    # Solve
    return solver.solve()


def plot_results(
    times: np.ndarray,
    solutions: List[np.ndarray],
    geometry_type: str,
    output_file: Optional[str] = None
) -> None:
    """
    Plot simulation results.
    
    Args:
        times: Array of simulation times
        solutions: List of solution snapshots (r, u, rho, e, p)
        geometry_type: Type of geometry used
        output_file: Optional file to save plot to
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.gridspec import GridSpec
    except ImportError:
        print("Matplotlib not available. Skipping plotting.")
        return
    
    # Set up figure
    fig = plt.figure(figsize=(12, 10))
    gs = GridSpec(2, 2, figure=fig)
    
    # Extract final solution
    r_final, u_final, rho_final, e_final, p_final = solutions[-1]
    
    # Plot density
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(0.5*(r_final[:-1] + r_final[1:]), rho_final, 'b-', linewidth=2)
    ax1.set_title('Density')
    ax1.set_ylabel('Density')
    ax1.grid(True)
    
    # Plot velocity
    ax2 = fig.add_subplot(gs[0, 1])
    ax2.plot(r_final, u_final, 'r-', linewidth=2)
    ax2.set_title('Velocity')
    ax2.set_ylabel('Velocity')
    ax2.grid(True)
    
    # Plot pressure
    ax3 = fig.add_subplot(gs[1, 0])
    ax3.plot(0.5*(r_final[:-1] + r_final[1:]), p_final, 'g-', linewidth=2)
    ax3.set_title('Pressure')
    ax3.set_xlabel('Radius')
    ax3.set_ylabel('Pressure')
    ax3.grid(True)
    
    # Plot internal energy
    ax4 = fig.add_subplot(gs[1, 1])
    ax4.plot(0.5*(r_final[:-1] + r_final[1:]), e_final, 'm-', linewidth=2)
    ax4.set_title('Specific Internal Energy')
    ax4.set_xlabel('Radius')
    ax4.set_ylabel('Energy')
    ax4.grid(True)
    
    plt.tight_layout()
    
    # Add main title
    plt.suptitle(f'Sedov Blast Wave ({geometry_type} geometry) - t = {times[-1]:.3f}', fontsize=16)
    plt.subplots_adjust(top=0.93)
    
    # Save or show
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"Plot saved to {output_file}")
    plt.show()


if __name__ == "__main__":
    # Example use case: Sedov blast wave in spherical geometry
    print("Running Sedov blast wave simulation...")
    
    # Simulation parameters
    n_cells = 100
    geometry_type = "spherical"  # Options: "planar", "cylindrical", "spherical"
    gamma = 1.4
    t_end = 0.1
    cfl = 0.3
    
    # Run simulation
    times, solutions = run_sedov_blast_wave(
        n_cells=n_cells,
        geometry_type=geometry_type,
        gamma=gamma,
        t_end=t_end,
        cfl=cfl
    )
    
    print(f"Simulation completed.")
    print(f"Number of time steps: {len(times)}")
    print(f"Final time: {times[-1]:.6f}")
    
    # Optional: Plot results if matplotlib is available
    plot_results(times, solutions, geometry_type)
    
    # Optional: Calculate convergence by tracking total energy
    initial_energy = sum([solutions[0][3][i] * solutions[0][2][i] * 
                         (solutions[0][0][i+1] - solutions[0][0][i]) 
                         for i in range(len(solutions[0][3]))])
    
    final_energy = sum([solutions[-1][3][i] * solutions[-1][2][i] * 
                        (solutions[-1][0][i+1] - solutions[-1][0][i]) 
                        for i in range(len(solutions[-1][3]))])
    
    energy_conservation = final_energy / initial_energy
    
    print(f"Energy conservation: {energy_conservation:.6f} (1.0 is perfect)")