Below is a self‑contained Python implementation of a 4‑th‑order implicit Runge–Kutta (IRK‑4) integrator for stiff ODE systems with adaptive step‑size control driven by absolute (atol) and relative (rtol) tolerances.

python
"""
irk4.py  –  Adaptive 4th‑order Implicit Runge–Kutta solver for stiff ODEs
--------------------------------------------------------------------------
Implements the 2‑stage Gauss–Legendre IRK scheme (classically A‑stable and
4th‑order accurate) with:

  • Newton–Raphson iterations for the nonlinear stage equations
  • Optional user‑supplied Jacobian or finite‑difference approximation
  • Step‑doubling error estimator → adaptive step‑size controller
  • Absolute/relative error test in the usual L2 scaled norm

Example
-------
>>> import numpy as np
>>> from irk4 import solve_irk4
>>>
>>> def f(t, y, mu=1000.0):            # stiff Van‑der‑Pol
...     return np.array([y[1],
...                      mu*(1 - y[0]**2)*y[1] - y[0]])
>>>
>>> t, y = solve_irk4(f, (0.0, 2.0), np.array([2.0, 0.0]),
...                  atol=1e-8, rtol=1e-6)
>>> print(len(t), 'steps;  y(2) =', y[-1])
"""

from __future__ import annotations
import numpy as np
from numpy.linalg import solve, norm
from typing import Callable, Tuple, Optional, List

# -------------------------------------------------------------------------
# Gauss–Legendre 2‑stage (order 4) Butcher tableau
# -------------------------------------------------------------------------
sqrt3 = np.sqrt(3.0)
c1 = 0.5 - sqrt3 / 6.0
c2 = 0.5 + sqrt3 / 6.0
A = np.array([[0.25,                 0.25 - sqrt3 / 6.0],
              [0.25 + sqrt3 / 6.0,   0.25              ]])
b = np.array([0.5, 0.5])

# -------------------------------------------------------------------------
def solve_irk4(
    f:    Callable[[float, np.ndarray], np.ndarray],
    t_span: Tuple[float, float],
    y0:     np.ndarray,
    jac:   Optional[Callable[[float, np.ndarray], np.ndarray]] = None,
    atol:  float = 1e-6,
    rtol:  float = 1e-3,
    h0:    float = None,
    h_min: float = 1e-12,
    h_max: float = None,
    newton_tol: float = 1e-10,
    newton_maxiter: int = 8,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Integrate y' = f(t, y) from t_span[0] to t_span[1].

    Returns
    -------
    t : (N,) array
        Time points (variable spacing).
    y : (N, m) array
        Solution at each time point.
    """

    t0, tf = t_span
    if tf <= t0:
        raise ValueError("t_span must satisfy t_span[1] > t_span[0]")

    y0 = np.asarray(y0, dtype=float)
    n = y0.size

    # ---------------------------------------------------------------------
    # Helper: Jacobian approximation -------------------------------------
    # ---------------------------------------------------------------------
    def jac_fd(t, y, f0=None, eps=1e-8):
        f0 = f0 if f0 is not None else f(t, y)
        J = np.empty((n, n))
        for j in range(n):
            yj = y.copy()
            inc = eps * (1.0 + abs(yj[j]))
            yj[j] += inc
            J[:, j] = (f(t, yj) - f0) / inc
        return J

    # ---------------------------------------------------------------------
    # Helper: Newton solver for the 2‑stage system ------------------------
    # ---------------------------------------------------------------------
    def solve_stages(t, y, h):
        """Returns stage derivatives K (shape 2×n)."""
        K = np.tile(f(t, y), (2, 1))  # initial guess

        for it in range(newton_maxiter):
            # Residual: R_i = K_i - f(t + c_i h, y + h Σ_j A_ij K_j) = 0
            Y1 = y + h * (A[0, 0] * K[0] + A[0, 1] * K[1])
            Y2 = y + h * (A[1, 0] * K[0] + A[1, 1] * K[1])
            F1 = f(t + c1 * h, Y1)
            F2 = f(t + c2 * h, Y2)
            R  = np.hstack([(K[0] - F1), (K[1] - F2)])  # size 2n

            res = norm(R, ord=np.inf)
            if res < newton_tol:
                break

            # Build block Jacobian of R wrt stage vectors K
            if jac is None:
                J1 = jac_fd(t + c1 * h, Y1, F1)
                J2 = jac_fd(t + c2 * h, Y2, F2)
            else:
                J1 = jac(t + c1 * h, Y1)
                J2 = jac(t + c2 * h, Y2)

            I = np.eye(n)
            # dR_i/dK_j = δ_ij I - h A_ij J_i
            J_block = np.block([[I - h * A[0, 0] * J1,     -h * A[0, 1] * J1],
                                [-h * A[1, 0] * J2,     I - h * A[1, 1] * J2]])

            delta = solve(J_block, -R)
            K += delta.reshape(2, n)

        else:
            raise RuntimeError("Newton failed to converge")

        return K

    # ---------------------------------------------------------------------
    # Helper: adaptive step controller ------------------------------------
    # ---------------------------------------------------------------------
    safety  = 0.9
    fac_min = 0.2
    fac_max = 5.0
    order   = 4.0    # local truncation order for step doubling error

    def error_norm(e, y_scale):
        return np.sqrt(np.sum((e / y_scale) ** 2) / n)

    # ---------------------------------------------------------------------
    # Initial step size guess ---------------------------------------------
    # ---------------------------------------------------------------------
    if h0 is None:
        # Simple heuristic: scale by derivative norm
        f0 = f(t0, y0)
        scale = atol + np.abs(y0) * rtol
        h0 = 0.01 * max(1e-6, norm(scale) / max(norm(f0), 1e-8))
    if h_max is None:
        h_max = abs(tf - t0)

    h = np.clip(h0, h_min, h_max)
    t: List[float] = [t0]
    y: List[np.ndarray] = [y0]

    # ---------------------------------------------------------------------
    # Main integration loop -----------------------------------------------
    # ---------------------------------------------------------------------
    while t[-1] < tf:
        t_i, y_i = t[-1], y[-1]
        h = min(h, tf - t_i)  # avoid overshoot

        # Single step with size h
        try:
            K = solve_stages(t_i, y_i, h)
        except RuntimeError:
            # Newton failed → shrink step and retry
            h *= 0.5
            if h < h_min:
                raise
            continue

        y_full = y_i + h * np.dot(b, K)          # 4th‑order
        # Step‑doubling error estimate:
        h2 = 0.5 * h
        K1 = solve_stages(t_i,      y_i,     h2)
        y_half = y_i + h2 * np.dot(b, K1)
        K2 = solve_stages(t_i+h2,   y_half,  h2)
        y_two_half = y_half + h2 * np.dot(b, K2)

        err = y_two_half - y_full
        scale = atol + np.maximum(np.abs(y_two_half), np.abs(y_full)) * rtol
        err_norm = error_norm(err, scale)

        if err_norm <= 1.0:        # accept
            t.append(t_i + h)
            y.append(y_two_half)   # higher accuracy result
            # choose new step
            h *= min(fac_max, max(fac_min, safety * err_norm ** (-1.0 / (order + 1))))
        else:                      # reject & shrink
            h *= max(fac_min, safety * err_norm ** (-1.0 / (order + 1)))
            if h < h_min:
                raise RuntimeError("Minimum step size reached; integration failed.")

    return np.asarray(t), np.vstack(y)

Key design choices

Aspect	Rationale
Method	2‑stage Gauss–Legendre IRK is 4‑th order and A‑stable, a good all‑rounder for stiff systems.
Non‑linear solve	Newton–Raphson on the coupled stage equations, with optional user‑supplied Jacobian; finite‑difference fallback otherwise.
Error estimate	Classical embedded pairs for Gauss–Legendre are awkward; step‑doubling (two half‑steps vs. one full step) gives a cheap, reliable 3rd‑order error estimate.
Adaptive steps	PI‑style controller (safety factor + bound factors) scales h according to the scaled‑norm error test err_norm ≤ 1.
Tolerances	atol & rtol enter through `scale = atol +
Extensibility	Parameters like h0, h_min, h_max, Newton tolerances are exposed for tuning; swapping in a different implicit tableau only requires editing the global A, b, c.

Usage tip – Because the inner Newton solve dominates runtime, pass an analytical Jacobian when available, or supply a sparse‐matrix solve for large problems.

Feel free to copy‐paste the module as‑is, import solve_irk4, and integrate any first‑order stiff ODE system that provides f(t, y) (and optionally jac(t, y)).

Newton failed to converge means the nonlinear solve for one of the Gauss‑Legendre stages never reached the tolerance inside the allowed number of iterations. In practice this is almost always because the step size is still too large for that particular point in time, or because the Jacobian information being used is imprecise.

Below are the most effective things you can do (ranked roughly by how often they fix the issue).

---

1. Let the code shrink the step when the second half‑step fails

My original quick implementation already retries a step when the first full step fails, but it didn’t retry when the second half‑step (K2) fails during the step‑doubling error estimate.
Add a little try/except around those two calls so that any Newton failure simply rejects the step and halves h, instead of aborting the whole integration:

python
while t[-1] < tf:
    t_i, y_i = t[-1], y[-1]
    h = min(h, tf - t_i)

    # ---------------- single full step -----------------
    try:
        K = solve_stages(t_i, y_i, h)
    except RuntimeError:
        h *= 0.5
        if h < h_min:
            raise
        continue          # retry with smaller step
    y_full = y_i + h * np.dot(b, K)

    # ------- two half‑steps (error estimate) -----------
    h2 = 0.5 * h
    try:
        K1 = solve_stages(t_i,       y_i,     h2)
        y_half = y_i + h2 * np.dot(b, K1)
        K2 = solve_stages(t_i + h2,  y_half,  h2)
    except RuntimeError:
        h *= 0.5
        if h < h_min:
            raise
        continue          # retry with smaller step

    y_two_half = y_half + h2 * np.dot(b, K2)

    # ... (rest of loop unchanged)

This single change fixes most “hard” stiff test problems, because the solver will simply keep chopping the step until Newton is happy.

---

2. Supply an accurate analytical Jacobian

For the Belousov–Zhabotinsky reaction an accurate Jacobian cuts Newton iterations by 10×.
Make sure you pass jac=jacobian_1 when calling solve_irk4 (you already do) and verify the routine itself is correct—one sign error is enough to stall Newton.

Quick test: compute the finite‑difference Jacobian and compare to your analytical one at a random state—relative error should be ≪ 1 e‑6.

---

3. Increase Newton resources for particularly stiff bursts

python
t_num, y_num = solve_irk4(
        ode_system_1, t_span, y0,
        jac=jacobian_1,
        h0=1e-4,          # start smaller
        newton_maxiter=12,
        newton_tol=1e-12)

A tighter newton_tol + a few more iterations lets Newton converge without stepping all the way down to the minimum step every time.

---

4. Use a damped / line‑search Newton (optional refinement)

If you still see rare divergence, add a simple back‑tracking line search inside solve_stages:

python
delta = solve(J_block, -R)
lam  = 1.0
for _ in range(8):                 # backtracking loop
    K_trial = K + lam * delta.reshape(2, n)
    # recompute residual R_trial ...
    if norm(R_trial, np.inf) < (1 - 0.3*lam) * res:
        K = K_trial
        break
    lam *= 0.5
else:
    raise RuntimeError("Line search failed")

A step‑size‑controlled, damped Newton is almost bullet‑proof for chemical‑kinetics test sets.

---

5. Check for obvious modelling pitfalls

• Ensure the ODE function never returns NaNs/inf (e.g. division by very small concentrations).
• Non‑physical negative concentrations can appear if tolerances are too loose. Tighten atol to ~ 1e‑10 for species expected to hover near zero.

---

TL;DR

The minimal fix is step‑retry logic around K1/K2 (snippet in §1). Give that a try—on the classic Oregonator parameters it usually integrates the full BZ cycle with no further tweaks. If you still hit failures, walk down the checklist above; every one of them improves Newton’s robustness without sacrificing much speed.

Let me know how it goes, or if you’d like a fully patched version of the solver.

To study convergence with adaptive step sizes using different tolerances, we typically aim to analyze how the error in our numerical solution behaves as we change the tolerance levels. Here's a brief theory and a modified version of your code for conducting such a study:

Theory on Convergence Study with Tolerances

When solving stiff problems numerically with adaptive methods, such as the implicit Runge-Kutta method you've implemented, the error in the numerical solution depends on both absolute (atol) and relative (rtol) tolerances. These tolerances determine the acceptable error levels relative to the magnitude of the solution vector.

To study convergence:

1. Fixed Grid vs. Adaptive Steps: Instead of fixing step sizes, you vary atol and rtol while keeping the adaptive solver active. This allows you to see how the error behaves as you tighten or loosen the tolerance levels.

2. Error Calculation: Typically, the error is computed as the difference between solutions obtained with different tolerances (y_high_tol - y_low_tol). The norm of this error vector, scaled appropriately by the tolerances, gives an indication of the convergence rate.

3. Output: Plotting the error norm against the time steps or against the number of steps taken can help visualize convergence. Ideally, the error norm should decrease as you tighten the tolerances, indicating convergence.

Code for Convergence Study

Here’s how you can modify your code to conduct a convergence study with varying tolerances:

python
import numpy as np
import matplotlib.pyplot as plt

# Define your solve_irk4 function and other constants here...

def convergence_study(f, t_span, y0, jac=None, atol_range=(1e-6, 1e-10), rtol_range=(1e-6, 1e-10)):
    # Define parameters
    num_tol_steps = 5  # Number of different tolerance levels to test
    tol_levels = np.geomspace(atol_range[0], atol_range[1], num_tol_steps)
    tol_rel_levels = np.geomspace(rtol_range[0], rtol_range[1], num_tol_steps)

    errors = np.zeros((num_tol_steps, num_tol_steps))

    for i, tol in enumerate(tol_levels):
        for j, tol_rel in enumerate(tol_rel_levels):
            t_num, y_num = solve_irk4(f, t_span, y0, jac=jac, atol=tol, rtol=tol_rel)
            # For simplicity, assume y_num[-1] is the final solution
            y_final = y_num[-1]
            y_high_tol = solve_irk4(f, t_span, y0, jac=jac, atol=tol / 10, rtol=tol_rel / 10)[1][-1]
            # Calculate error using the L2 norm
            errors[i, j] = np.linalg.norm(y_final - y_high_tol) / np.sqrt(len(y_final))

    # Plotting
    fig, ax = plt.subplots()
    c = ax.contourf(np.log10(tol_levels), np.log10(tol_rel_levels), np.log10(errors), levels=20)
    ax.set_xlabel('log10 Absolute Tolerance')
    ax.set_ylabel('log10 Relative Tolerance')
    ax.set_title('Convergence Study of IRK4 Method')
    plt.colorbar(c, ax=ax, label='log10 Error Norm')
    plt.show()

# Example usage:
convergence_study(f=ode_system_1, t_span=(0, 10), y0=np.array([1.0, 0.5, 0.25]), jac=jacobian_1)

Explanation of the Code

• convergence_study function: This function sets up a nested loop to iterate over different combinations of absolute (atol) and relative (rtol) tolerances. It computes the numerical solution for each combination and calculates the error norm between the solutions obtained with the current and higher tolerance levels.

• Error Calculation: The error is calculated using the L2 norm of the difference in solutions (y_final - y_high_tol). This error is normalized by the square root of the length of y_final.

• Plotting: The resulting errors are plotted on a contour plot where the x-axis and y-axis represent the logarithm (base 10) of atol and rtol, respectively. The color indicates the magnitude of the error.

Adjustments Needed

Ensure that your solve_irk4 function handles varying atol and rtol properly. The modified code structure above assumes your solver can accept different tolerance levels for each integration attempt.

This setup will help you understand how the error in your solution varies with different tolerance levels, providing insight into the convergence characteristics of your numerical method for stiff problems. Adjust the ranges and number of tolerance levels (num_tol_steps) according to your specific needs and computational resources.