Integrator Benchmark Example

Test Hiten integrators against SciPy on various ODE problems and provides

Testing Hiten integrators against SciPy on various ODE problems

#!/usr/bin/env python3
"""Integrator benchmark script.

Tests Hiten integrators against SciPy on various ODE problems and provides
detailed accuracy, performance, and error analysis.
"""

import os
import sys
import time
import warnings
from dataclasses import dataclass
from typing import List

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp

sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', "src"))

from hiten.algorithms.dynamics.rhs import create_rhs_system
from hiten.algorithms.integrators import (AdaptiveRK, RungeKutta)

warnings.filterwarnings('ignore', category=UserWarning)


@dataclass
class TestResult:
    """Store results from a single integrator test."""
    integrator_name: str
    problem_name: str
    final_error: float
    max_error: float
    relative_error: float
    computation_time: float
    n_steps: int
    energy_error: float = None
    converged: bool = True
    error_message: str = ""


# Test problems as simple functions
def harmonic_oscillator(t, y):
    """Harmonic oscillator: x'' + x = 0."""
    return np.array([y[1], -y[0]])

def van_der_pol(t, y):
    """Van der Pol oscillator: x'' - μ(1-x²)x' + x = 0."""
    mu = 1.0
    return np.array([y[1], mu * (1 - y[0]**2) * y[1] - y[0]])

def lorenz(t, y):
    """Lorenz system: chaotic attractor."""
    sigma, rho, beta = 10.0, 28.0, 8.0/3.0
    return np.array([
        sigma * (y[1] - y[0]),
        y[0] * (rho - y[2]) - y[1],
        y[0] * y[1] - beta * y[2]
    ])

def kepler(t, y):
    """Kepler problem: 2D central force motion."""
    x, y_pos, vx, vy = y
    mu = 1.0
    r = np.sqrt(x**2 + y_pos**2)
    r3 = r**3
    return np.array([vx, vy, -mu * x / r3, -mu * y_pos / r3])

def duffing(t, y):
    """Duffing oscillator: x'' + δx' + αx + βx³ = γcos(ωt)."""
    delta, alpha, beta = 0.1, -1.0, 1.0
    gamma, omega = 0.3, 1.2
    return np.array([
        y[1],
        -delta * y[1] - alpha * y[0] - beta * y[0]**3 + 
        gamma * np.cos(omega * t)
    ])

# Test problem configurations
TEST_PROBLEMS = [
    {
        "name": "Harmonic Oscillator",
        "rhs": harmonic_oscillator,
        "y0": np.array([1.0, 0.0]),
        "t_span": (0.0, 4*np.pi),
        "exact": lambda t: np.column_stack([np.cos(t), -np.sin(t)]),
        "energy": lambda y: 0.5 * (y[0]**2 + y[1]**2)
    },
    {
        "name": "Van der Pol",
        "rhs": van_der_pol,
        "y0": np.array([2.0, 0.0]),
        "t_span": (0.0, 20.0),
        "exact": None,
        "energy": None
    },
    {
        "name": "Lorenz",
        "rhs": lorenz,
        "y0": np.array([1.0, 1.0, 1.0]),
        "t_span": (0.0, 20.0),
        "exact": None,
        "energy": None
    },
    {
        "name": "Kepler",
        "rhs": kepler,
        "y0": np.array([1.0, 0.0, 0.0, 0.8]),
        "t_span": (0.0, 10.0),
        "exact": None,
        "energy": lambda y: 0.5 * (y[2]**2 + y[3]**2) - 1.0 / np.sqrt(y[0]**2 + y[1]**2)
    },
    {
        "name": "Duffing",
        "rhs": duffing,
        "y0": np.array([1.0, 0.0]),
        "t_span": (0.0, 20.0),
        "exact": None,
        "energy": None
    }
]


def run_integrator_test(integrator, system, problem_config, t_eval: np.ndarray) -> TestResult:
    """Run a single integrator test and return results."""
    
    # Warm-up JIT compilation (excluded from timing)
    try:
        if t_eval.size >= 2:
            dt = t_eval[1] - t_eval[0]
            warmup_t = np.array([t_eval[0], t_eval[0] + max(dt, 1e-8)], dtype=float)
        else:
            warmup_t = np.array([t_eval[0], t_eval[0] + 1e-8], dtype=float)
        _ = integrator.integrate(system, problem_config["y0"], warmup_t)
    except Exception:
        pass

    start_time = time.perf_counter()
    
    try:
        solution = integrator.integrate(system, problem_config["y0"], t_eval)
        computation_time = time.perf_counter() - start_time
        
        y_solution = solution.states
        
        # Calculate errors
        if problem_config["exact"] is not None:
            y_exact = problem_config["exact"](t_eval)
            errors = np.abs(y_solution - y_exact)
            final_error = np.linalg.norm(errors[-1])
            max_error = np.max(errors)
            relative_error = max_error / (np.max(np.abs(y_exact)) + 1e-16)
        else:
            # Use SciPy as reference
            scipy_sol = solve_ivp(
                system.rhs, problem_config["t_span"], problem_config["y0"], 
                t_eval=t_eval, rtol=1e-12, atol=1e-14
            )
            y_ref = scipy_sol.y.T
            errors = np.abs(y_solution - y_ref)
            final_error = np.linalg.norm(errors[-1])
            max_error = np.max(errors)
            relative_error = max_error / (np.max(np.abs(y_ref)) + 1e-16)
        
        # Calculate energy error if applicable
        energy_error = None
        if problem_config["energy"] is not None:
            initial_energy = problem_config["energy"](problem_config["y0"])
            final_energy = problem_config["energy"](y_solution[-1])
            energy_error = abs(final_energy - initial_energy) / abs(initial_energy)
        
        # Count steps (approximate for fixed-step methods)
        n_steps = len(t_eval) - 1
        
        return TestResult(
            integrator_name=str(integrator),
            problem_name=problem_config["name"],
            final_error=final_error,
            max_error=max_error,
            relative_error=relative_error,
            computation_time=computation_time,
            n_steps=n_steps,
            energy_error=energy_error,
            converged=True
        )
        
    except Exception as e:
        computation_time = time.perf_counter() - start_time
        return TestResult(
            integrator_name=str(integrator),
            problem_name=problem_config["name"],
            final_error=np.inf,
            max_error=np.inf,
            relative_error=np.inf,
            computation_time=computation_time,
            n_steps=0,
            converged=False,
            error_message=str(e)
        )


def run_benchmark():
    """Run comprehensive tests on all integrators and problems."""
    
    # Define integrators to test
    integrators = {
        # Fixed-step RK methods
        "RK4": RungeKutta(order=4),
        "RK6": RungeKutta(order=6),
        "RK8": RungeKutta(order=8),
        
        # Adaptive RK methods
        "RK45": AdaptiveRK(order=5, rtol=1e-8, atol=1e-10),
        "DOP853": AdaptiveRK(order=8, rtol=1e-8, atol=1e-10),
    }
    
    # SciPy reference
    scipy_integrators = {
        "SciPy-RK45": "RK45",
        "SciPy-DOP853": "DOP853",
        "SciPy-RK23": "RK23",
        "SciPy-BDF": "BDF",
        "SciPy-LSODA": "LSODA"
    }
    
    # Time grid for evaluation
    t_eval = np.linspace(0, 10, 1001)
    
    all_results = []
    
    print("Running comprehensive integrator tests...")
    print("=" * 60)
    
    for problem_config in TEST_PROBLEMS:
        print(f"\nTesting problem: {problem_config['name']}")
        print("-" * 40)
        # Build a single dynamical system instance per problem for all tests
        system = create_rhs_system(problem_config["rhs"], dim=len(problem_config["y0"]), name=problem_config["name"])
        
        # Test Hiten integrators
        for name, integrator in integrators.items():
            result = run_integrator_test(integrator, system, problem_config, t_eval)
            all_results.append(result)
            
            if result.converged:
                print(f"{name:15s}: "
                      f"Final Error: {result.final_error:.2e}, "
                      f"Max Error: {result.max_error:.2e}, "
                      f"Time: {result.computation_time:.4f}s")
                
                if result.energy_error is not None:
                    print(f"{'':15s} Energy Error: {result.energy_error:.2e}")
            else:
                print(f"{name:15s}: FAILED - {result.error_message}")
        
        # Test SciPy integrators
        for name, method in scipy_integrators.items():
            start_time = time.perf_counter()
            
            try:
                scipy_sol = solve_ivp(
                    system.rhs, problem_config["t_span"], problem_config["y0"],
                    t_eval=t_eval, method=method, rtol=1e-8, atol=1e-10
                )
                computation_time = time.perf_counter() - start_time
                
                if scipy_sol.success:
                    y_solution = scipy_sol.y.T
                    
                    # Calculate errors
                    if problem_config["exact"] is not None:
                        y_exact = problem_config["exact"](t_eval)
                        errors = np.abs(y_solution - y_exact)
                        final_error = np.linalg.norm(errors[-1])
                        max_error = np.max(errors)
                        relative_error = max_error / (np.max(np.abs(y_exact)) + 1e-16)
                    else:
                        # Use highest accuracy SciPy as reference
                        ref_sol = solve_ivp(
                            system.rhs, problem_config["t_span"], problem_config["y0"],
                            t_eval=t_eval, method="DOP853", rtol=1e-12, atol=1e-14
                        )
                        y_ref = ref_sol.y.T
                        errors = np.abs(y_solution - y_ref)
                        final_error = np.linalg.norm(errors[-1])
                        max_error = np.max(errors)
                        relative_error = max_error / (np.max(np.abs(y_ref)) + 1e-16)
                    
                    # Energy error
                    energy_error = None
                    if problem_config["energy"] is not None:
                        initial_energy = problem_config["energy"](problem_config["y0"])
                        final_energy = problem_config["energy"](y_solution[-1])
                        energy_error = abs(final_energy - initial_energy) / abs(initial_energy)
                    
                    result = TestResult(
                        integrator_name=name,
                        problem_name=problem_config["name"],
                        final_error=final_error,
                        max_error=max_error,
                        relative_error=relative_error,
                        computation_time=computation_time,
                        n_steps=len(scipy_sol.t) - 1,
                        energy_error=energy_error,
                        converged=True
                    )
                    all_results.append(result)
                    
                    print(f"{name:15s}: "
                          f"Final Error: {result.final_error:.2e}, "
                          f"Max Error: {result.max_error:.2e}, "
                          f"Time: {result.computation_time:.4f}s")
                    
                    if result.energy_error is not None:
                        print(f"{'':15s} Energy Error: {result.energy_error:.2e}")
                else:
                    print(f"{name:15s}: FAILED - {scipy_sol.message}")
                    
            except Exception as e:
                print(f"{name:15s}: ERROR - {e}")
    
    return all_results


def print_summary_table(results: List[TestResult]):
    """Print a summary table of all results."""
    
    print("\n" + "=" * 100)
    print("SUMMARY TABLE")
    print("=" * 100)
    print(f"{'Integrator':<15} {'Problem':<20} {'Final Error':<12} {'Max Error':<12} {'Time (s)':<10} {'Energy Error':<12}")
    print("-" * 100)
    
    # Group by problem for better readability
    problems = {}
    for result in results:
        if result.problem_name not in problems:
            problems[result.problem_name] = []
        problems[result.problem_name].append(result)
    
    for problem_name in sorted(problems.keys()):
        problem_results = problems[problem_name]
        for result in sorted(problem_results, key=lambda x: x.final_error):
            energy_str = f"{result.energy_error:.2e}" if result.energy_error is not None else "N/A"
            print(f"{result.integrator_name:<15} {result.problem_name:<20} "
                  f"{result.final_error:<12.2e} {result.max_error:<12.2e} "
                  f"{result.computation_time:<10.4f} {energy_str:<12}")


def create_performance_plots(results: List[TestResult]):
    """Create performance comparison plots."""
    
    # Create results directory
    os.makedirs('_debug/results/plots', exist_ok=True)
    
    # Extract data for plotting
    integrators = list(set(r.integrator_name for r in results if r.converged))
    problems = list(set(r.problem_name for r in results if r.converged))
    
    # Create accuracy vs speed plot
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    colors = plt.cm.tab10(np.linspace(0, 1, len(integrators)))
    integrator_colors = dict(zip(integrators, colors))
    
    # Plot 1: Final Error vs Time
    for integrator in integrators:
        integrator_results = [r for r in results if r.integrator_name == integrator and r.converged]
        if integrator_results:
            times = [r.computation_time for r in integrator_results]
            errors = [r.final_error for r in integrator_results]
            ax1.loglog(times, errors, 'o', color=integrator_colors[integrator], 
                      label=integrator, markersize=8)
    
    ax1.set_xlabel('Computation Time (s)')
    ax1.set_ylabel('Final Error')
    ax1.set_title('Accuracy vs Speed')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Max Error vs Time
    for integrator in integrators:
        integrator_results = [r for r in results if r.integrator_name == integrator and r.converged]
        if integrator_results:
            times = [r.computation_time for r in integrator_results]
            errors = [r.max_error for r in integrator_results]
            ax2.loglog(times, errors, 'o', color=integrator_colors[integrator], 
                      label=integrator, markersize=8)
    
    ax2.set_xlabel('Computation Time (s)')
    ax2.set_ylabel('Max Error')
    ax2.set_title('Max Error vs Speed')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('_debug/results/plots/performance_comparison.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print("Performance plots saved to _debug/results/plots/")


def main():
    """Main function to run all tests."""
    
    print("Hiten Integrators vs SciPy Comprehensive Benchmark")
    print("=" * 60)
    
    # Run comprehensive tests
    results = run_benchmark()
    
    # Print summary
    print_summary_table(results)
    
    # Create plots
    create_performance_plots(results)
    
    print("\nBenchmark completed successfully!")


if __name__ == "__main__":
    main()

Event Integrator Benchmark Example

Tests Hiten’s event-capable integrator (DOP853) against SciPy solvers on simple problems with known event times. Reports detection accuracy and speed, and saves a comparison plot.

Testing Hiten’s event-capable integrator (DOP853) against SciPy solvers on simple problems with known event times

"""Event-handling integrator benchmark.

Compares Hiten's event-capable integrator (DOP853) against SciPy solvers on
simple problems with known event times. Reports detection accuracy and speed,
and saves a comparison plot.
"""

import os
import sys
import time
import warnings
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from numba import njit, types
from scipy.integrate import solve_ivp

# Make project src importable when running from repository root
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'src'))

from hiten.algorithms.dynamics.rhs import create_rhs_system
from hiten.algorithms.integrators import RungeKutta
from hiten.algorithms.integrators.configs import _EventConfig

warnings.filterwarnings('ignore', category=UserWarning)


@dataclass
class EventBenchmarkResult:
    integrator_name: str
    problem_name: str
    t_event: Optional[float]
    y_event: Optional[np.ndarray]
    t_error: float
    y_error: Optional[float]
    computation_time: float
    converged: bool = True
    error_message: str = ""


# --- Precompiled global RHS and event functions (stable identities) ---
@njit(types.float64(types.float64, types.float64[:]), cache=True, fastmath=True)
def g_y1_jit(t: float, y: np.ndarray) -> float:
    return y[0] - 1.0


@njit(types.float64(types.float64, types.float64[:]), cache=True, fastmath=True)
def g_x0_jit(t: float, y: np.ndarray) -> float:
    return y[0]


@njit(types.float64[:](types.float64, types.float64[:]), cache=True, fastmath=True)
def rhs_inc_jit(t: float, y: np.ndarray) -> np.ndarray:
    out = np.empty(1, dtype=np.float64)
    out[0] = 1.0
    return out


@njit(types.float64[:](types.float64, types.float64[:]), cache=True, fastmath=True)
def rhs_dec_jit(t: float, y: np.ndarray) -> np.ndarray:
    out = np.empty(1, dtype=np.float64)
    out[0] = -1.0
    return out


@njit(types.float64[:](types.float64, types.float64[:]), cache=True, fastmath=True)
def rhs_ho_jit(t: float, y: np.ndarray) -> np.ndarray:
    out = np.empty(2, dtype=np.float64)
    out[0] = y[1]
    out[1] = -y[0]
    return out


# --- Event test problems using precompiled functions ---
def _problems() -> List[Dict]:
    problems: List[Dict] = []

    problems.append({
        "name": "y' = 1, hit y=1 (inc)",
        "rhs": rhs_inc_jit,
        "y0": np.array([0.2], dtype=float),
        "t_span": (0.0, 5.0),
        "event": g_y1_jit,
        "direction": +1,
        "t_expected": lambda y0: 1.0 - float(y0[0]),
        "dim": 1,
    })

    problems.append({
        "name": "y' = -1, hit y=1 (dec)",
        "rhs": rhs_dec_jit,
        "y0": np.array([2.0], dtype=float),
        "t_span": (0.0, 5.0),
        "event": g_y1_jit,
        "direction": -1,
        "t_expected": lambda y0: float(y0[0]) - 1.0,
        "dim": 1,
    })

    problems.append({
        "name": "Harmonic oscillator, hit x=0",
        "rhs": rhs_ho_jit,
        "y0": np.array([1.0, 0.0], dtype=float),
        "t_span": (0.0, 5.0),
        "event": g_x0_jit,
        "direction": -1,
        "t_expected": lambda y0: 0.5 * np.pi,
        "dim": 2,
    })

    return problems


# --- Warm-up helpers (avoid JIT bias in timing) ---
def _make_warmup_times(t0: float) -> np.ndarray:
    return np.array([t0, t0 + 1.0e1], dtype=float)


def _is_fixed_step(integrator) -> bool:
    try:
        return hasattr(integrator, "_integrate_fixed_rk")
    except Exception:
        return False


def _warmup_hiten_event(integrator, system, y0: np.ndarray, event_fn: Callable[[float, np.ndarray], float], direction: int) -> None:
    try:
        warmup_t = _make_warmup_times(0.0)
        if _is_fixed_step(integrator):
            # Use a small grid to trigger step-based event path and refinement
            t_grid = np.linspace(warmup_t[0], warmup_t[-1], 64)
            _ = integrator.integrate(
                system,
                y0,
                t_grid,
                event_fn=event_fn,
                event_cfg=_EventConfig(direction=direction, terminal=True),
            )
        else:
            _ = integrator.integrate(
                system,
                y0,
                warmup_t,
                event_fn=event_fn,
                event_cfg=_EventConfig(direction=direction, terminal=True),
            )
    except Exception:
        pass


def _warmup_scipy_event(method: str, rhs: Callable[[float, np.ndarray], np.ndarray], y0: np.ndarray, t0: float, event_fn: Callable[[float, np.ndarray], float], direction: int) -> None:
    try:
        def ev(t, y):
            return event_fn(t, y)
        ev.terminal = True
        ev.direction = float(direction)
        warmup_t = _make_warmup_times(t0)
        _ = solve_ivp(rhs, (warmup_t[0], warmup_t[-1]), y0, method=method, events=ev, rtol=1e-6, atol=1e-8)
    except Exception:
        pass


def run_hiten_event(
    integrator,
    system,
    y0: np.ndarray,
    t_span: Tuple[float, float],
    event_fn: Callable[[float, np.ndarray], float],
    direction: int,
    t_expected: float,
    problem_name: str,
) -> EventBenchmarkResult:
    _warmup_hiten_event(integrator, system, y0, event_fn, direction)
    start = time.perf_counter()
    try:
        # Adaptive integrators only need endpoints; fixed-step benefits from a grid
        if _is_fixed_step(integrator):
            # Use a moderately fine grid for fixed-step methods
            n_grid = 4097
            t_eval = np.linspace(t_span[0], t_span[1], n_grid, dtype=float)
        else:
            t_eval = np.array([t_span[0], t_span[1]], dtype=float)
        sol = integrator.integrate(
            system,
            y0,
            t_eval,
            event_fn=event_fn,
            event_cfg=_EventConfig(direction=direction, terminal=True),
        )
        dt = time.perf_counter() - start
        # Event path returns two nodes [t0, t_hit] or [t0, tmax] when no hit
        t_event = float(sol.times[-1])
        y_event = sol.states[-1].copy()
        t_err = abs(t_event - t_expected)
        y_err = float(np.linalg.norm(y_event - y_event))  # always 0 vs itself (placeholder)
        return EventBenchmarkResult(
            integrator_name=str(integrator),
            problem_name=problem_name,
            t_event=t_event,
            y_event=y_event,
            t_error=t_err,
            y_error=y_err,
            computation_time=dt,
            converged=True,
        )
    except Exception as e:
        dt = time.perf_counter() - start
        return EventBenchmarkResult(
            integrator_name=str(integrator),
            problem_name=problem_name,
            t_event=None,
            y_event=None,
            t_error=np.inf,
            y_error=None,
            computation_time=dt,
            converged=False,
            error_message=str(e),
        )


def run_scipy_event(
    name: str,
    method: str,
    rhs: Callable[[float, np.ndarray], np.ndarray],
    y0: np.ndarray,
    t_span: Tuple[float, float],
    event_fn: Callable[[float, np.ndarray], float],
    direction: int,
    t_expected: float,
    problem_name: str,
) -> EventBenchmarkResult:
    _warmup_scipy_event(method, rhs, y0, t_span[0], event_fn, direction)
    def ev(t, y):
        return event_fn(t, y)
    ev.terminal = True
    ev.direction = float(direction)
    start = time.perf_counter()
    try:
        sol = solve_ivp(
            rhs,
            t_span,
            y0,
            method=method,
            events=ev,
            rtol=1e-8,
            atol=1e-10,
        )
        dt = time.perf_counter() - start
        if sol.t_events and len(sol.t_events[0]) > 0:
            t_event = float(sol.t_events[0][0])
            y_event = sol.y_events[0][0].copy() if hasattr(sol, 'y_events') and sol.y_events and len(sol.y_events[0]) > 0 else None
            t_err = abs(t_event - t_expected)
            y_err = float(np.linalg.norm(y_event - y_event)) if y_event is not None else None
            return EventBenchmarkResult(
                integrator_name=name,
                problem_name=problem_name,
                t_event=t_event,
                y_event=y_event,
                t_error=t_err,
                y_error=y_err,
                computation_time=dt,
                converged=True,
            )
        else:
            return EventBenchmarkResult(
                integrator_name=name,
                problem_name=problem_name,
                t_event=None,
                y_event=None,
                t_error=np.inf,
                y_error=None,
                computation_time=dt,
                converged=False,
                error_message=str(sol.message),
            )
    except Exception as e:
        dt = time.perf_counter() - start
        return EventBenchmarkResult(
            integrator_name=name,
            problem_name=problem_name,
            t_event=None,
            y_event=None,
            t_error=np.inf,
            y_error=None,
            computation_time=dt,
            converged=False,
            error_message=str(e),
        )


def run_event_benchmark() -> List[EventBenchmarkResult]:
    results: List[EventBenchmarkResult] = []

    # Hiten integrators (adaptive and fixed-step)
    hiten_integrators = {
        "HITEN-RK45": RungeKutta(order=45, rtol=1e-8, atol=1e-10),
        "HITEN-DOP853": RungeKutta(order=853, rtol=1e-8, atol=1e-10),
        "HITEN-RK4-fixed": RungeKutta(order=4),
        "HITEN-RK6-fixed": RungeKutta(order=6),
        "HITEN-RK8-fixed": RungeKutta(order=8),
    }

    # SciPy integrators with event support
    scipy_integrators = {
        "SciPy-RK45": "RK45",
        "SciPy-DOP853": "DOP853",
        "SciPy-Radau": "Radau",
        "SciPy-BDF": "BDF",
        "SciPy-LSODA": "LSODA",
    }

    print("Event Integrator Benchmark (Hiten vs SciPy)")
    print("=" * 60)

    for prob in _problems():
        print(f"\nProblem: {prob['name']}")
        print("-" * 40)
        system = create_rhs_system(prob["rhs"], dim=prob["dim"], name=prob["name"])  
        t_expected = float(prob["t_expected"](prob["y0"]))

        # Hiten
        for name, integrator in hiten_integrators.items():
            res = run_hiten_event(
                integrator=integrator,
                system=system,
                y0=prob["y0"],
                t_span=prob["t_span"],
                event_fn=prob["event"],
                direction=prob["direction"],
                t_expected=t_expected,
                problem_name=prob["name"],
            )
            results.append(res)
            if res.converged:
                print(f"{str(integrator):15s}: t_hit={res.t_event:.10f}, |dt|={res.t_error:.2e}, time={res.computation_time:.4f}s")
            else:
                print(f"{str(integrator):15s}: FAILED - {res.error_message}")

        # SciPy
        for name, method in scipy_integrators.items():
            res = run_scipy_event(
                name=name,
                method=method,
                rhs=prob["rhs"],
                y0=prob["y0"],
                t_span=prob["t_span"],
                event_fn=prob["event"],
                direction=prob["direction"],
                t_expected=t_expected,
                problem_name=prob["name"],
            )
            results.append(res)
            if res.converged:
                print(f"{name:15s}: t_hit={res.t_event:.10f}, |dt|={res.t_error:.2e}, time={res.computation_time:.4f}s")
            else:
                print(f"{name:15s}: FAILED - {res.error_message}")

    return results


def print_summary_table(results: List[EventBenchmarkResult]) -> None:
    print("\n" + "=" * 100)
    print("SUMMARY TABLE (Events)")
    print("=" * 100)
    print(f"{'Integrator':<15} {'Problem':<30} {'t_hit':<18} {'|dt|':<12} {'Time (s)':<10}")
    print("-" * 100)

    groups: Dict[str, List[EventBenchmarkResult]] = {}
    for r in results:
        groups.setdefault(r.problem_name, []).append(r)

    for problem_name in sorted(groups.keys()):
        for r in sorted(groups[problem_name], key=lambda x: (not x.converged, x.t_error)):
            t_hit_str = f"{r.t_event:.10f}" if r.t_event is not None else "N/A"
            dt_str = f"{r.t_error:.2e}" if np.isfinite(r.t_error) else "inf"
            print(f"{r.integrator_name:<15} {r.problem_name:<30} {t_hit_str:<18} {dt_str:<12} {r.computation_time:<10.4f}")


def create_performance_plot(results: List[EventBenchmarkResult]) -> None:
    os.makedirs('_debug/results/plots', exist_ok=True)

    # Plot |dt| vs time, grouped by integrator names
    names = list({r.integrator_name for r in results})
    colors = plt.cm.tab10(np.linspace(0, 1, max(1, len(names))))
    cmap: Dict[str, np.ndarray] = dict(zip(names, colors))

    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    for name in names:
        subset = [r for r in results if r.integrator_name == name and r.converged]
        if not subset:
            continue
        times = [r.computation_time for r in subset]
        terrs = [r.t_error for r in subset]
        ax.loglog(times, terrs, 'o', label=name, color=cmap[name], markersize=8)

    ax.set_xlabel('Computation Time (s)')
    ax.set_ylabel('Absolute Event Time Error |dt|')
    ax.set_title('Event Detection: Accuracy vs Speed')
    ax.grid(True, which='both', alpha=0.3)
    ax.legend()

    plt.tight_layout()
    out_path = '_debug/results/plots/events_performance_comparison.png'
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Performance plot saved to {out_path}")


def main() -> None:
    results = run_event_benchmark()
    print_summary_table(results)
    create_performance_plot(results)
    print("\nEvent benchmark completed.")


if __name__ == "__main__":
    main()

Symplectic Event Integrator Benchmark Example

Tests Hiten’s extended symplectic integrators with event detection enabled against SciPy solvers on simple problems with known event times. Reports detection accuracy and speed, and saves a comparison plot.

Testing Hiten’s extended symplectic integrators with event detection enabled against SciPy solvers on simple problems with known event times

#!/usr/bin/env python3
"""Symplectic event-handling benchmark and example.

This script mirrors ``event_integrator_benchmark.py`` but focuses on
Hiten's extended symplectic integrators with event detection enabled.
It constructs a truncated pendulum Hamiltonian, integrates it with
event-enabled symplectic schemes, compares the detected event times
against SciPy solvers, and saves an accuracy vs speed plot.
"""

from __future__ import annotations

import os
import sys
import time
import warnings
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import matplotlib.pyplot as plt
import numpy as np
from numba import njit, types
from numba.typed import List as NumbaList
from scipy.integrate import solve_ivp

# Make project src importable when running from repository root
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'src'))

from hiten.algorithms.dynamics.hamiltonian import create_hamiltonian_system
from hiten.algorithms.integrators import ExtendedSymplectic
from hiten.algorithms.integrators.configs import _EventConfig
from hiten.algorithms.integrators.symplectic import (
    N_SYMPLECTIC_DOF,
    N_VARS_POLY,
    P_POLY_INDICES,
    Q_POLY_INDICES,
)
from hiten.algorithms.polynomial.base import (
    _create_encode_dict_from_clmo,
    _encode_multiindex,
    _init_index_tables,
)

warnings.filterwarnings('ignore', category=UserWarning)


@dataclass
class EventResult:
    solver_name: str
    problem_name: str
    t_event: Optional[float]
    y_event: Optional[np.ndarray]
    t_error: float
    y_error: Optional[float]
    computation_time: float
    converged: bool = True
    error_message: str = ""


@dataclass
class SymplecticEventProblem:
    name: str
    description: str
    y0_extended: np.ndarray
    y0_canonical: np.ndarray
    t_span: Tuple[float, float]
    grid_size: int
    direction: int


@njit(types.float64(types.float64, types.float64[:]), cache=True, fastmath=True)
def event_q1_jit(t: float, y: np.ndarray) -> float:
    """Event function detecting q1 = 0 crossings."""
    return y[0]


def truncated_pendulum_rhs(_t: float, y: np.ndarray) -> List[float]:
    """Return dynamics for the 2D truncated pendulum system."""
    q, p = y
    sin_taylor = q - (q ** 3) / 6.0 + (q ** 5) / 120.0
    return [p, -sin_taylor]


def embed_state(qp_state: np.ndarray) -> np.ndarray:
    """Embed a 2D (q, p) state into the 6D symplectic phase space."""
    state = np.zeros(2 * N_SYMPLECTIC_DOF, dtype=np.float64)
    state[0] = qp_state[0]
    state[N_SYMPLECTIC_DOF] = qp_state[1]
    return state


def build_truncated_pendulum_system(max_degree: int = 6):
    """Construct a polynomial Hamiltonian system for a truncated pendulum."""
    psi_tables, clmo_tables = _init_index_tables(max_degree)
    encode_dict_list = _create_encode_dict_from_clmo(clmo_tables)

    H_blocks = [
        np.zeros(psi_tables[N_VARS_POLY, deg], dtype=np.complex128)
        for deg in range(max_degree + 1)
    ]

    idx_p1 = P_POLY_INDICES[0]
    idx_q1 = Q_POLY_INDICES[0]
    mono = np.zeros(N_VARS_POLY, dtype=np.int64)

    # Kinetic term: p1^2 / 2 (degree 2)
    mono[:] = 0
    mono[idx_p1] = 2
    encoded = _encode_multiindex(mono, 2, encode_dict_list)
    if encoded != -1:
        H_blocks[2][encoded] = 0.5

    # Potential offset: -1 (degree 0)
    mono[:] = 0
    encoded = _encode_multiindex(mono, 0, encode_dict_list)
    if encoded != -1:
        H_blocks[0][encoded] = -1.0

    # Quadratic potential: +q1^2 / 2 (degree 2)
    mono[:] = 0
    mono[idx_q1] = 2
    encoded = _encode_multiindex(mono, 2, encode_dict_list)
    if encoded != -1:
        H_blocks[2][encoded] += 0.5

    # Quartic correction: -q1^4 / 24 (degree 4)
    if max_degree >= 4:
        mono[:] = 0
        mono[idx_q1] = 4
        encoded = _encode_multiindex(mono, 4, encode_dict_list)
        if encoded != -1:
            H_blocks[4][encoded] = -1.0 / 24.0

    # Sextic correction: +q1^6 / 720 (degree 6)
    if max_degree >= 6:
        mono[:] = 0
        mono[idx_q1] = 6
        encoded = _encode_multiindex(mono, 6, encode_dict_list)
        if encoded != -1:
            H_blocks[6][encoded] = 1.0 / 720.0

    H_blocks_typed = NumbaList()
    for arr in H_blocks:
        H_blocks_typed.append(arr.copy())

    system = create_hamiltonian_system(
        H_blocks=H_blocks_typed,
        degree=max_degree,
        psi_table=psi_tables,
        clmo_table=clmo_tables,
        encode_dict_list=encode_dict_list,
        n_dof=N_SYMPLECTIC_DOF,
        name="Truncated Pendulum Hamiltonian",
    )

    return system


def make_event_problems() -> List[SymplecticEventProblem]:
    """Configure example problems for event detection."""
    angle = np.deg2rad(45.0)
    base_state = np.zeros(2 * N_SYMPLECTIC_DOF, dtype=np.float64)
    base_state_neg = base_state.copy()

    # Positive release, expect crossing with negative direction
    base_state[0] = angle
    problem_pos = SymplecticEventProblem(
        name="Pendulum release (+45 deg)",
        description="Pendulum released from +45 degrees, expect q1 -> 0 with negative crossing.",
        y0_extended=base_state.copy(),
        y0_canonical=np.array([angle, 0.0], dtype=np.float64),
        t_span=(0.0, 5.0),
        grid_size=4097,
        direction=-1,
    )

    # Negative release, expect crossing with positive direction
    base_state_neg[0] = -angle
    problem_neg = SymplecticEventProblem(
        name="Pendulum release (-45 deg)",
        description="Pendulum released from -45 degrees, expect q1 -> 0 with positive crossing.",
        y0_extended=base_state_neg.copy(),
        y0_canonical=np.array([-angle, 0.0], dtype=np.float64),
        t_span=(0.0, 5.0),
        grid_size=4097,
        direction=+1,
    )

    return [problem_pos, problem_neg]


def compute_reference_event(problem: SymplecticEventProblem) -> Tuple[float, np.ndarray]:
    """High-accuracy reference using SciPy Radau."""
    event = lambda t, y: y[0]
    event.terminal = True
    event.direction = float(problem.direction)

    sol = solve_ivp(
        truncated_pendulum_rhs,
        problem.t_span,
        problem.y0_canonical,
        method="Radau",
        events=event,
        rtol=1e-12,
        atol=1e-14,
    )

    if sol.status == 1 and sol.t_events and len(sol.t_events[0]) > 0:
        t_event = float(sol.t_events[0][0])
        y_event = sol.y_events[0][0].copy()
        return t_event, y_event

    raise RuntimeError(f"Reference solver failed to detect event for problem '{problem.name}'")


def _warmup_symplectic_event(integrator, system, problem: SymplecticEventProblem) -> None:
    """Trigger compilation overhead outside timed region."""
    try:
        warmup_grid = np.linspace(
            problem.t_span[0],
            problem.t_span[0] + 1.0e-2,
            8,
            dtype=np.float64,
        )
        integrator.integrate(
            system,
            problem.y0_extended,
            warmup_grid,
            event_fn=event_q1_jit,
            event_cfg=_EventConfig(direction=problem.direction, terminal=True),
        )
    except Exception:
        pass


def _warmup_scipy_event(method: str, problem: SymplecticEventProblem) -> None:
    """Trigger SciPy solver overhead outside timed region."""
    try:
        event = lambda t, y: y[0]
        event.terminal = True
        event.direction = float(problem.direction)
        warmup_t_span = (problem.t_span[0], problem.t_span[0] + 1.0e-2)
        solve_ivp(
            truncated_pendulum_rhs,
            warmup_t_span,
            problem.y0_canonical,
            method=method,
            events=event,
            rtol=1e-6,
            atol=1e-8,
        )
    except Exception:
        pass


def run_symplectic_event(
    solver_name: str,
    integrator,
    system,
    problem: SymplecticEventProblem,
    t_expected: float,
    y_expected: np.ndarray,
) -> EventResult:
    """Execute symplectic integration with event detection."""
    time_grid = np.linspace(
        problem.t_span[0], problem.t_span[1], problem.grid_size, dtype=np.float64
    )
    event_cfg = _EventConfig(direction=problem.direction, terminal=True, xtol=1.0e-12, gtol=1.0e-12)

    _warmup_symplectic_event(integrator, system, problem)

    start = time.perf_counter()
    try:
        sol = integrator.integrate(
            system,
            problem.y0_extended,
            time_grid,
            event_fn=event_q1_jit,
            event_cfg=event_cfg,
        )
        elapsed = time.perf_counter() - start

        t_event = float(sol.times[-1])
        y_event = sol.states[-1].copy()
        t_error = abs(t_event - t_expected)
        y_error = float(np.linalg.norm(y_event - y_expected)) if y_expected is not None else None

        return EventResult(
            solver_name=solver_name,
            problem_name=problem.name,
            t_event=t_event,
            y_event=y_event,
            t_error=t_error,
            y_error=y_error,
            computation_time=elapsed,
            converged=True,
        )
    except Exception as exc:
        elapsed = time.perf_counter() - start
        return EventResult(
            solver_name=solver_name,
            problem_name=problem.name,
            t_event=None,
            y_event=None,
            t_error=np.inf,
            y_error=None,
            computation_time=elapsed,
            converged=False,
            error_message=str(exc),
        )


def run_scipy_event(
    solver_name: str,
    method: str,
    solver_opts: Dict,
    problem: SymplecticEventProblem,
    t_expected: float,
    y_expected: np.ndarray,
) -> EventResult:
    """Run SciPy solve_ivp with event detection on the 2D system."""
    event = lambda t, y: y[0]
    event.terminal = True
    event.direction = float(problem.direction)

    _warmup_scipy_event(method, problem)

    start = time.perf_counter()
    try:
        sol = solve_ivp(
            truncated_pendulum_rhs,
            problem.t_span,
            problem.y0_canonical,
            method=method,
            events=event,
            **solver_opts,
        )
        elapsed = time.perf_counter() - start

        if sol.status == 1 and sol.t_events and len(sol.t_events[0]) > 0:
            t_event = float(sol.t_events[0][0])
            qp_event = sol.y_events[0][0].copy()
            y_event = embed_state(qp_event)
            t_error = abs(t_event - t_expected)
            y_error = float(np.linalg.norm(y_event - y_expected)) if y_expected is not None else None
            return EventResult(
                solver_name=solver_name,
                problem_name=problem.name,
                t_event=t_event,
                y_event=y_event,
                t_error=t_error,
                y_error=y_error,
                computation_time=elapsed,
                converged=True,
            )

        message = sol.message if hasattr(sol, 'message') else "event not detected"
        return EventResult(
            solver_name=solver_name,
            problem_name=problem.name,
            t_event=None,
            y_event=None,
            t_error=np.inf,
            y_error=None,
            computation_time=elapsed,
            converged=False,
            error_message=message,
        )
    except Exception as exc:
        elapsed = time.perf_counter() - start
        return EventResult(
            solver_name=solver_name,
            problem_name=problem.name,
            t_event=None,
            y_event=None,
            t_error=np.inf,
            y_error=None,
            computation_time=elapsed,
            converged=False,
            error_message=str(exc),
        )


def run_benchmark() -> List[EventResult]:
    """Execute the symplectic event benchmark."""
    system = build_truncated_pendulum_system()
    problems = make_event_problems()

    symplectic_schemes = {
        "Symplectic4": ExtendedSymplectic(order=4, c_omega_heuristic=15.0),
        "Symplectic6": ExtendedSymplectic(order=6, c_omega_heuristic=20.0),
        "Symplectic8": ExtendedSymplectic(order=8, c_omega_heuristic=25.0),
    }

    scipy_solvers = {
        "SciPy-RK45": ("RK45", {"rtol": 1.0e-8, "atol": 1.0e-10}),
        "SciPy-DOP853": ("DOP853", {"rtol": 1.0e-9, "atol": 1.0e-11}),
        "SciPy-BDF": ("BDF", {"rtol": 1.0e-8, "atol": 1.0e-10}),
    }

    results: List[EventResult] = []

    print("Symplectic Event Integrator Benchmark")
    print("=" * 70)

    for problem in problems:
        print(f"\nProblem: {problem.name}")
        print(problem.description)
        print("-" * 70)

        t_ref, y_ref_qp = compute_reference_event(problem)
        y_ref_ext = embed_state(y_ref_qp)

        print(f"Reference event time (Radau): {t_ref:.10f} s")

        reference_result = EventResult(
            solver_name="SciPy-Radau (reference)",
            problem_name=problem.name,
            t_event=t_ref,
            y_event=y_ref_ext,
            t_error=0.0,
            y_error=0.0,
            computation_time=0.0,
            converged=True,
        )
        results.append(reference_result)

        for solver_name, integrator in symplectic_schemes.items():
            res = run_symplectic_event(
                solver_name,
                integrator,
                system,
                problem,
                t_ref,
                y_ref_ext,
            )
            results.append(res)
            if res.converged:
                print(
                    f"{solver_name:20s}: t_hit={res.t_event:.10f}, |dt|={res.t_error:.2e}, "
                    f"time={res.computation_time:.4f}s"
                )
            else:
                print(f"{solver_name:20s}: FAILED - {res.error_message}")

        for solver_name, (method, opts) in scipy_solvers.items():
            res = run_scipy_event(
                solver_name,
                method,
                opts,
                problem,
                t_ref,
                y_ref_ext,
            )
            results.append(res)
            if res.converged:
                print(
                    f"{solver_name:20s}: t_hit={res.t_event:.10f}, |dt|={res.t_error:.2e}, "
                    f"time={res.computation_time:.4f}s"
                )
            else:
                print(f"{solver_name:20s}: FAILED - {res.error_message}")

    return results


def print_summary_table(results: List[EventResult]) -> None:
    """Print formatted summary of all runs."""
    print("\n" + "=" * 110)
    print("SUMMARY TABLE (Symplectic Events)")
    print("=" * 110)
    print(f"{'Integrator':<22} {'Problem':<32} {'t_hit':<18} {'|dt|':<12} {'|dy|':<12} {'Time (s)':<10}")
    print("-" * 110)

    grouped: Dict[str, List[EventResult]] = {}
    for res in results:
        grouped.setdefault(res.problem_name, []).append(res)

    for problem_name in sorted(grouped):
        for res in sorted(grouped[problem_name], key=lambda r: (not r.converged, r.t_error)):
            if res.converged:
                t_hit_str = f"{res.t_event:.10f}" if res.t_event is not None else "N/A"
                dt_str = f"{res.t_error:.2e}" if np.isfinite(res.t_error) else "inf"
                dy_str = (
                    f"{res.y_error:.2e}"
                    if res.y_error is not None and np.isfinite(res.y_error)
                    else "n/a"
                )
                print(
                    f"{res.solver_name:<22} {res.problem_name:<32} {t_hit_str:<18} "
                    f"{dt_str:<12} {dy_str:<12} {res.computation_time:<10.4f}"
                )
            else:
                print(
                    f"{res.solver_name:<22} {res.problem_name:<32} {'FAILED':<18} "
                    f"{'inf':<12} {'n/a':<12} {res.computation_time:<10.4f}"
                )


def create_performance_plot(results: List[EventResult]) -> None:
    """Save log-log plot of |dt| vs computation time."""
    os.makedirs(os.path.join('_debug', 'results', 'plots'), exist_ok=True)

    names = sorted({res.solver_name for res in results if res.converged and res.t_error >= 0.0})
    if not names:
        print("No successful runs to plot.")
        return

    colors = plt.cm.tab10(np.linspace(0, 1, max(1, len(names))))
    cmap = dict(zip(names, colors))

    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    for name in names:
        subset = [r for r in results if r.solver_name == name and r.converged and np.isfinite(r.t_error)]
        if not subset:
            continue
        times = [r.computation_time for r in subset]
        terrs = [r.t_error for r in subset]
        ax.loglog(times, terrs, 'o', label=name, color=cmap[name], markersize=8)

    ax.set_xlabel('Computation Time (s)')
    ax.set_ylabel('Absolute Event Time Error |dt|')
    ax.set_title('Symplectic Event Detection: Accuracy vs Speed')
    ax.grid(True, which='both', alpha=0.3)
    ax.legend()

    plt.tight_layout()
    out_path = os.path.join('_debug', 'results', 'plots', 'symplectic_events_performance.png')
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Performance plot saved to {out_path}")


def main() -> None:
    results = run_benchmark()
    print_summary_table(results)
    create_performance_plot(results)
    print("\nSymplectic event benchmark completed.")


if __name__ == "__main__":
    main()