impasse_algorithm.py

"""
Implementation of Rabier & Rheinboldt's (1994) impasse point algorithm.

Reference:
    P.J. Rabier and W.C. Rheinboldt, "On the Computation of Impasse Points
    of Quasi-Linear Differential-Algebraic Equations," Mathematics of
    Computation, Vol. 62, No. 205, pp. 133-154, January 1994.
    DOI: 10.1090/S0025-5718-1994-1208224-6
"""

import numpy as np
import scipy as sp
from numpy.typing import NDArray
from scikits.odes import dae
from abc import ABC, abstractmethod

# Newton iteration parameters
DEFAULT_NEWTON_TOLERANCE = 1e-10
MAX_NEWTON_ITERATIONS = 100

# Integrator parameters
DEFAULT_REL_TOL = 1e-3
DEFAULT_ABS_TOL = 1e-3
DEFAULT_MAX_STEP = 1e-3
MIN_S_STEP = 1e-14
MIN_TIME_STEP = 1e-14
MAX_RK_ATTEMPTS = 20
DEFAULT_TIME_STEP = 1e-4


class DAESystem(ABC):
    """
    Abstract base class for DAE system definition.

    The DAE system has the semi-explicit form of eq. (4.1) 
    from Rabier & Rheinboldt (1994):

        A1(x) @ x' = G1(x)     (r differential equations)
        G2(x) = 0              (p algebraic constraints)

    where:
        x in R_n         state vector (n = r + p)
        A1: R_n → R_rxn  differential equation matrix
        G1: R_n → R_r    differential equation RHS
        G2: R_n → R_p    algebraic constraints

    Note. Any non-autonomous system can be made autonomous
    by solving for time as unknown and adding the equation
    dt/ds = 1 to the system. However, it would be preferable
    to see what role non-autonomous dynamics play throughout
    the algorithm and generalize the later to handle non-autonomous
    systems directly without having to explicitely add time as a
    variable.
    """

    @abstractmethod
    def A1(self, x: NDArray) -> NDArray:
        """Differential equation matrix (r x n)"""
        pass

    @abstractmethod
    def G1(self, x: NDArray) -> NDArray:
        """Differential equation RHS (r,)"""
        pass

    @abstractmethod
    def G2(self, x: NDArray) -> NDArray:
        """Algebraic constraints (p,)"""
        pass

    @abstractmethod
    def DG2(self, x: NDArray) -> NDArray:
        """Jacobian of constraints (p x n)"""
        pass


def compute_normal_and_tangent_space_bases(
    x: NDArray,
    system: DAESystem
) -> tuple[NDArray, NDArray, NDArray]:
    """
    Computes the QR factorization of DG2(x)^T to get normal 
    and tangent space bases of the r-dimensional manifold

        N{x : G2(x) = 0}

    at x. See eq. (4.9).

    Args:
        x: State vector (n,)
        system: DAE system

    Returns:
        Q1: Normal space basis (n x p)
        Q2: Tangent space basis (n x r)
        R1: Upper triangular factor (p x p)

    Note. The paper assumes that rank(DG2(x)) = p everywhere in its
    domain of definition. Can the algorithm be generalized to admit 
    equations where DG2(x) doesn't have full rank everywhere? 
    In practice, if DG2(x) is very poorly conditionned, then it might
    be effectively singular from a numerical perspective.
    """
    DG2_x = system.DG2(x)
    p = DG2_x.shape[0]

    Q, R = sp.linalg.qr(DG2_x.T, mode='full')
    Q1 = Q[:, :p]  # Normal space basis (n x p)
    Q2 = Q[:, p:]  # Tangent space basis (n x r)
    R1 = R[:p]     # Upper triangular (p x p)
    return Q1, Q2, R1


def phi(
    y: NDArray,
    xc: NDArray,
    Q1c: NDArray,
    Uc: NDArray,
    R1c: NDArray,
    system: DAESystem,
    tol: float = DEFAULT_NEWTON_TOLERANCE,
    max_iter: int = MAX_NEWTON_ITERATIONS
) -> tuple[NDArray, bool]:
    """
    Evaluate the local tangential coordinate map.
    See algorithm Eval-phi on p.146, section 4.

    Args:
        y: Local coordinates (r,)
        xc: Center point (n,)
        Q1c: Normal space basis at center (n x p)
        Uc: Tangent space basis at center (n x r)
        R1c: QR factor at center (p x p)
        system: DAE system
        tol: Tolerance for Newton iteration
        max_iter: Maximum nb. of iterations

    Returns:
        x: Point on constraint manifold (n,)
        converged: True if Newton iteration converged

    Note. Using Newton chord for the projection allows us to
    reuse the QR decomposition that we already have of DG2(xc)
    without recomputing a new Jacobian at each iteration.
    But if G2(x) is highly nonlinear or if DG2(xc) is poorly
    conditionned, this could be replaced by a better root-finding
    algorithm for nonlinear algebraic equations.
    """
    x = xc + Uc @ y
    converged = False

    # Project point onto the constraint manifold using chord Newton
    for _ in range(max_iter):
        G2x = system.G2(x)  # Shape (p,)
        residual_norm = np.linalg.norm(G2x)

        if residual_norm < tol:
            converged = True
            break

        z = sp.linalg.solve_triangular(R1c.T, G2x, lower=True)  # Shape (p,)
        x = x - Q1c @ z  # (n,) - (n,p) @ (p,) = (n,) - (n,) = (n,)

    return x, converged


def Dphi(
    x: NDArray,
    Uc: NDArray,
    system: DAESystem
) -> NDArray:
    """
    Compute the Jacobian of the coordinate map. See eq. (4.11)

    Args:
        x: State vector (n,)
        Uc: Tangent space basis at center (n x r)
        system: DAE system

    Returns:
        Dphi: Derivative matrix (n x r)
    """
    _, rhs, _ = compute_normal_and_tangent_space_bases(x, system)
    lhs = Uc.T @ rhs
    return np.linalg.solve(lhs.T, rhs.T).T


def compute_augmentation_vectors(
    Bc: NDArray,
    Hc: NDArray,
    flip_orientation: bool = False
) -> tuple[NDArray, NDArray, NDArray]:
    """
    Compute augmentation vectors e, c1, c2 (case q=2).
    See equations (4.12)-(4.14) on p.146.

    Args:
        Bc: Reduced system matrix (r x r), i.e. A1c @ Uc
        Hc: Reduced RHS vector (r,), i.e. G1c
        flip_orientation: Whether z1, z2 should be flipped when computing gamma for
            scaling. Must match the flip_orientation used in the integration step.

    Returns:
        e: Augmentation vector (r,)
        c1: First normalization vector (r,)
        c2: Second normalization vector (r,)
    """
    r = Bc.shape[0]

    # SVD of (Bc, -Hc)
    # VL is r x r matrix
    # VR is (r+1) x (r+1) matrix
    VL, _, VR = sp.linalg.svd(np.hstack([Bc, -Hc[:, np.newaxis]]))

    # Helper function to find rightmost nonzero column
    def find_nonzero_column(matrix: NDArray, start_idx: int, end_idx: int,
                           name: str, threshold: float = 1e-12) -> tuple[NDArray, float, int]:
        """Find rightmost column with norm > threshold in range [end_idx, start_idx]."""
        for idx in range(start_idx, end_idx, -1):
            col = matrix[:, idx]
            norm = float(np.linalg.norm(col))
            if norm > threshold:
                return col, norm, idx
        raise RuntimeError(
            f"Cannot compute augmentation: norm of {name} is too small. "
            f"The matrix may be rank-deficient.")

    # Extract u1 and u2 from rightmost nonzero columns of VR
    # See eq. (4.13) on p.147
    u1, norm_u1, idx_u1 = find_nonzero_column(VR[:r, :], r, 0, "u1")
    u2, norm_u2, _ = find_nonzero_column(VR[:r, :], idx_u1 - 1, -1, "u2")

    # Extract e from rightmost nonzero column of VL
    # See eq. (4.14) on p.146
    # The paper forgets to mention that e also needs to be nonzero,
    # but we would get a zero column in eq. (4.12) otherwise.
    e, _, _ = find_nonzero_column(VL, r - 1, -1, "e")

    # Paper's original unit normalization
    c1 = u1 / norm_u1
    c2 = u2 / norm_u2

    return e, c1, c2


def solve_augmented_system(
    B: NDArray,
    H: NDArray,
    e: NDArray,
    c1: NDArray,
    c2: NDArray,
    flip_orientation: bool = False
) -> tuple[NDArray, NDArray, float, float, float, float]:
    """
    Solve augmented system (2.22) for q=2 case.

    Args:
        B: Reduced system matrix (r x r)
        H: Reduced RHS vector (r,)
        e: Augmentation vector (r,)
        c1: First normalization vector (r,)
        c2: Second normalization vector (r,)
        flip_orientation: Whether to flip z1, z2 after crossing singularity

    Returns:
        v1: First solution vector (r,)
        v2: Second solution vector (r,)
        w1: First scalar coefficient
        w2: Second scalar coefficient
        z1: First normalization scalar
        z2: Second normalization scalar
    """
    r = B.shape[0]

    # Assemble augmented matrix (2.22)
    rows1 = np.hstack([B, -H[:, np.newaxis], e[:, np.newaxis]])
    rows2 = np.hstack([c1[np.newaxis, :], np.zeros((1, 2))])
    rows3 = np.hstack([c2[np.newaxis, :], np.zeros((1, 2))])
    B_aug = np.vstack([rows1,
                       rows2,
                       rows3])

    rhs = np.vstack([np.zeros((r, 2)),
                     np.eye(2)])

    # Use LU to solve and get determinant sign
    lu, piv = sp.linalg.lu_factor(B_aug)
    sol = sp.linalg.lu_solve((lu, piv), rhs)

    v1 = sol[:r, 0]
    v2 = sol[:r, 1]
    w1 = sol[r, 0]
    w2 = sol[r, 1]
    z1 = sol[r+1, 0]
    z2 = sol[r+1, 1]

    # I am too lazy right now to use the prevously computed
    # decomposition to compute the sign of the determinant
    det_B_aug = np.linalg.det(B_aug)
    sign_det_B_aug = np.sign(det_B_aug)

    # Flip z1, z2 if det(B_aug) < 0, see p.148 in sec. 4 just
    # before the Eval-v algorithm
    if sign_det_B_aug < 0:
        z1 = -z1
        z2 = -z2

    # Flip z1, z2 again if we've crossed the singularity
    # The paper suggests on p.149 to change the sign of the integration step h
    # after crossing. Instead, we flip z1, z2 to reverse the vector field v(y).
    # I need to check if it is always always mathematically equivalent,
    # but flipping z1, z2 is prefered for industrial solvers if it generalizes
    # to nicely in eq. (4.14). It would make it much easier to implement in
    # already existing solvers.
    if flip_orientation:
        z1 = -z1
        z2 = -z2

    return v1, v2, w1, w2, z1, z2


class ButcherTableau:
    """Butcher tableau for explicit Runge-Kutta methods."""

    def __init__(self, A, b, c, b_hat=None, ctrl_order=None, name="RK"):
        self.A = np.array(A)      # Coefficient matrix (q x q)
        self.b = np.array(b)      # Weights for solution (q,)
        self.c = np.array(c)      # Node points (q,)
        # Weights for error estimate
        self.b_hat = np.array(b_hat) if b_hat is not None else None
        self.ctrl_order = ctrl_order # min(order of estimator, order of auxiliary)
        self.name = name


# Classical RK4
RK4 = ButcherTableau(
    A=[[0, 0, 0, 0],
       [0.5, 0, 0, 0],
       [0, 0.5, 0, 0],
       [0, 0, 1, 0]],
    b=[1/6, 1/3, 1/3, 1/6],
    c=[0, 0.5, 0.5, 1],
    b_hat=[1/2, 0, 0, 1/2],  # Heun's method (RK2): uses k1 (c=0) and k4 (c=1)
    ctrl_order=2, # this is probably not true
    name="RK4"
)

# Dormand-Prince 5(4) - DOPRI5
DOPRI5 = ButcherTableau(
    A=[[0, 0, 0, 0, 0, 0, 0],
       [1/5, 0, 0, 0, 0, 0, 0],
       [3/40, 9/40, 0, 0, 0, 0, 0],
       [44/45, -56/15, 32/9, 0, 0, 0, 0],
       [19372/6561, -25360/2187, 64448/6561, -212/729, 0, 0, 0],
       [9017/3168, -355/33, 46732/5247, 49/176, -5103/18656, 0, 0],
       [35/384, 0, 500/1113, 125/192, -2187/6784, 11/84, 0]],
    b=[35/384, 0, 500/1113, 125/192, -2187/6784, 11/84, 0],
    c=[0, 1/5, 3/10, 4/5, 8/9, 1, 1],
    b_hat=[5179/57600, 0, 7571/16695, 393/640, -92097/339200, 187/2100, 1/40],
    ctrl_order=4, 
    name="DOPRI5"
)


def evaluate_velocity(
    y: NDArray,
    xc: NDArray,
    Q1c: NDArray,
    Uc: NDArray,
    R1c: NDArray,
    e: NDArray,
    c1: NDArray,
    c2: NDArray,
    system: DAESystem,
    flip_orientation: bool = False
) -> tuple[NDArray, float, bool]:
    """
    Eval-v algorithm from page 148.

    Args:
        y: Local coordinates (r,)
        xc: Center point (n,)
        Q1c: Normal space basis at xc (n x p)
        Uc: Tangent basis at xc (n x r)
        R1c: QR factor at xc (p x p)
        e: Augmentation vector (r,)
        c1: First normalization vector (r,)
        c2: Second normalization vector (r,)
        system: DAE system
        flip_orientation: Whether to flip z1, z2 after crossing singularity

    Returns:
        v: Velocity vector (r,)
        gamma: Time scaling factor (scalar)
        success: True if coordinate map evaluation succeeded
    """
    # The following steps match the steps given in the paper.

    # Step 1: compute x = phi(y)
    x, converged = phi(y, xc, Q1c, Uc, R1c, system)
    if not converged:
        return np.zeros_like(y), 0.0, False

    # Step 2: compute Dphi(y) and B(y) = A1(x)Dphi(y)
    Dphi_y = Dphi(x, Uc, system)
    B = system.A1(x) @ Dphi_y
    H = system.G1(x)

    # Steps 3-4: solve augmented system (2.22)
    v1, v2, w1, w2, z1, z2 = solve_augmented_system(
        B, H, e, c1, c2, flip_orientation)

    # Step 5: compute v(y) and gamma(y)
    v = z2 * v1 - z1 * v2
    gamma = z2 * w1 - z1 * w2
    return v, gamma, True


def runge_kutta_step(
    tableau: ButcherTableau,
    y0: NDArray,
    ds: float,
    xc: NDArray,
    Q1c: NDArray,
    Uc: NDArray,
    R1c: NDArray,
    e: NDArray,
    c1: NDArray,
    c2: NDArray,
    system: DAESystem,
    flip_orientation: bool
) -> tuple[NDArray, NDArray, float, float, float, bool]:
    """
    Perform one Runge-Kutta step for dy/ds = v(y) and dt/ds = gamma(y).

    Args:
        tableau: Butcher tableau defining the RK method
        y0: Current y value (r,)
        ds: Step size in pseudo-time
        xc, Q1c, Uc, R1c, e, c1, c2: Reduction/augmentation data
        system: DAE system
        flip_orientation: Whether orientation is flipped

    Returns:
        y_next: Next y value using primary weights (r,)
        y_error: Error estimate y_next - y_hat (r,)
        dt: Real time increment
        gamma0: gamma at y0
        gamma_next: gamma at y_next
        success: True if all velocity evaluations succeeded
    """
    q = len(tableau.b)
    k = np.zeros((q, len(y0)))
    gamma_stages = np.zeros(q)

    # Compute stages
    for i in range(q):
        y_stage = y0.copy()
        for j in range(i):
            y_stage += ds * tableau.A[i, j] * k[j]

        k[i], gamma_stages[i], success = evaluate_velocity(
            y_stage, xc, Q1c, Uc, R1c, e, c1, c2, system, flip_orientation)

        if not success:
            return np.zeros_like(y0), np.zeros_like(y0), 0.0, 0.0, 0.0, False

    # Compute solution
    y_next = y0 + ds * np.sum(tableau.b[:, np.newaxis] * k, axis=0)
    dt = ds * np.sum(tableau.b * gamma_stages)

    # Gamma at endpoints
    gamma0 = gamma_stages[0]
    _, gamma_next, success = evaluate_velocity(
        y_next, xc, Q1c, Uc, R1c, e, c1, c2, system, flip_orientation)

    if not success:
        return np.zeros_like(y0), np.zeros_like(y0), 0.0, gamma0, 0.0, False

    # Error estimate
    if tableau.b_hat is not None:
        y_hat = y0 + ds * np.sum(tableau.b_hat[:, np.newaxis] * k, axis=0)
        y_error = y_next - y_hat
    else:
        y_error = np.zeros_like(y_next)

    return y_next, y_error, dt, gamma0, gamma_next, True


def step_singDAE(
    xk: NDArray,
    dsk: float,
    system: DAESystem,
    reltol: float = DEFAULT_REL_TOL,
    abstol: float = DEFAULT_ABS_TOL,
    ds_min: float = MIN_S_STEP,
    flip_orientation: bool = False,
    rk_method: ButcherTableau = DOPRI5
) -> tuple[NDArray, float, bool, float, float]:
    """
    SingDAE algorithm from page 148-149.
    Performs one step of the integration in pseudo-time s.

    Args:
        xk: Current point (n,)
        dsk: Suggested pseudo-time step size
        system: DAE system
        reltol: Relative tolerance for error control
        abstol: Absolute tolerance for error control
        ds_min: Minimal pseudo-time step size
        flip_orientation: Whether to flip z1, z2 after crossing singularity
        rk_method: Runge-Kutta method to use (default: DOPRI5)

    Returns:
        xk_next: Next point (n,)
        dsk_next: Suggested next pseudo-time step size
        gamma_sign_changed: Boolean indicating if gamma changed sign
        dt: Real time increment, integrated using selected RK method
        gamma_next: Gamma value at the endpoint of this step
    """
    # The following steps match the steps given in the paper.

    # Step 1: Set xc = xk and compute QR factorization
    xc = xk
    Q1c, Uc, R1c = compute_normal_and_tangent_space_bases(xc, system)

    # Step 2: Evaluate Bc = A1(xc)Uc
    Bc = system.A1(xc) @ Uc
    Hc = system.G1(xc)

    # Step 3: Compute augmentation vectors
    e, c1, c2 = compute_augmentation_vectors(Bc, Hc,
                                              flip_orientation=flip_orientation)

    # Step 4: Take Runge-Kutta step for (2.12) from y0 = 0 in pseudo-time s
    r = Bc.shape[0]  # Derive r from matrix shape
    y0 = np.zeros(r)
    ds = dsk

    y1: NDArray | None = None
    gamma0: float | None = None
    gamma1: float | None = None
    dt: float | None = None
    dsk_next: float | None = None
    xk_next: NDArray | None = None

    # Step size control parameters
    safety = 0.9
    max_factor = 5.0 # Don't increase too aggressively
    min_factor = 0.2 # Don't decrease too agressively

    step_accepted = False
    for _ in range(MAX_RK_ATTEMPTS):
        y1, y_error, dt, gamma0, gamma1, rk_success = runge_kutta_step(
            rk_method, y0, ds, xc, Q1c, Uc, R1c, e, c1, c2,
            system, flip_orientation)

        if not rk_success:
            # Coordinate map evaluation failed during RK stages, reduce step size
            ds *= 0.5
            if ds < ds_min:
                raise RuntimeError(
                    f"Coordinate map evaluation failed and step size {ds:.2e} "
                    f"reached minimum {ds_min:.2e}. Cannot continue integration.")
            continue

        # RMS norm of scaled error
        sc = reltol * np.maximum(np.abs(y0), np.abs(y1)) + abstol
        error = np.sqrt(np.mean((y_error / sc) ** 2))

        # Compute new step size based on error
        if rk_method.ctrl_order is not None and error > 0:
            exponent = 1.0 / (rk_method.ctrl_order + 1)
            factor = safety * (1.0 / error) ** exponent
            factor = min(factor, max_factor) # Limit growth
            factor = max(factor, min_factor) # Limit reduction
            ds_new = ds * factor
        else:
            ds_new = ds  # Keep same step size if no order info

        if error < 1.0:
            # Step passes error control, now check final coordinate map evaluation
            # Step 6: Compute xk+1 = phi(y1)
            xk_next, converged = phi(y1, xc, Q1c, Uc, R1c, system)

            if not converged:
                # Final coordinate map evaluation failed, reduce step size and retry
                ds *= 0.5
                if ds < ds_min:
                    raise RuntimeError(
                        "Final coordinate map evaluation failed and "
                        f"step size {ds:.2e} reached minimum {ds_min:.2e}. "
                        "Cannot continue integration.")
                continue

            # Step fully accepted
            dsk_next = ds_new
            if rk_method.ctrl_order is None:
                dsk_next *= 1.2  # Gradually increase when no error control
            step_accepted = True
            break

        # Reject step due to error control, reduce step size and retry
        ds = ds_new
        if rk_method.ctrl_order is None:
            ds *= 0.5
        if ds < ds_min:
            raise RuntimeError(
                f"Error control failed and step size {ds:.2e} "
                f"reached minimum {ds_min:.2e}. Cannot continue integration.")

    # Check if we exhausted all attempts
    if not step_accepted:
        raise RuntimeError(
            f"Failed to accept step after {MAX_RK_ATTEMPTS} attempts. "
            f"Integration cannot continue.")

    # Step 7: Check for sign change in gamma
    # At this point, all variables must be defined since step_accepted is True
    assert gamma0 is not None and gamma1 is not None
    assert dt is not None and dsk_next is not None and xk_next is not None

    gamma_sign_changed = (np.sign(gamma0) != np.sign(gamma1)
                          if gamma0 != 0 and gamma1 != 0 else False)

    # Step 8: Output
    return xk_next, dsk_next, gamma_sign_changed, dt, gamma1


# ============================================================================
# Integration and visualization
# ============================================================================

def integrate(
    x0: NDArray,
    t_final: float,
    ds0: float,
    system: DAESystem,
    reltol: float = DEFAULT_REL_TOL,
    abstol: float = DEFAULT_ABS_TOL,
    max_ds: float = DEFAULT_MAX_STEP,
    flip_after_crossing: bool = False,
    initial_flip: bool = False,
    rk_method: ButcherTableau = DOPRI5,
    verbose: bool = False
) -> tuple[NDArray, NDArray]:
    """
    Integrate the DAE from initial point x0 until final real time t_final.
    Integration is performed in pseudo-time s, with dt/ds = gamma(y(s)).

    Args:
        x0: Initial point on constraint manifold N (n,)
        t_final: Final real time t
        ds0: Initial pseudo-time step size (must be positive)
        reltol: Relative tolerance for RK step
        abstol: Absolute tolerance for RK step
        max_ds: Maximum reparametrized stepsize
        system: DAE system
        flip_after_crossing: If True, flip orientation when crossing
            singularity
        initial_flip: If True, start with flipped orientation (reverses
            initial direction)
        rk_method: Runge-Kutta method (default: DOPRI5)
        verbose: If True, print integration progress information

    Returns:
        trajectory: Array of solution points, shape (num_points, n)
        time: Array of real time t values, shape (num_points,)
    """
    sol = [x0.copy()]
    time = [0.0]
    flipped = initial_flip
    num_crossings = 0

    xk = x0
    dsk = ds0
    t = 0.0
    step_count = 0

    if verbose:
        print(f"Starting integration: t_final={t_final}, ds0={ds0}")
        print(f"Initial state: x0={x0}")

    while t < t_final:
        try:
            xk_next, dsk_next, gamma_sign_changed, dt, gamma_next = step_singDAE(
                xk, dsk, system, reltol=reltol, abstol=abstol,
                flip_orientation=flipped, rk_method=rk_method)

        except RuntimeError as e:
            # Integration failed due to exceptional circumstance
            # (step size too small, coordinate map failed, etc.)
            # Return partial solution computed so far
            if verbose:
                print(f"Integration stopped at t={t:.6e}: {e}")
            break

        # Real time increment: use |dt| since we want t to always increase
        # (gamma can be negative, but physical time should always go forward)
        t += abs(dt)
        sol.append(xk_next.copy())
        time.append(t)

        step_count += 1

        if verbose and step_count % 100 == 0:
            print(f"Step {step_count}: t={t:.6e}, ds={dsk:.6e}, gamma={gamma_next:.6e}, crossings={num_crossings}")

        if gamma_sign_changed:
            num_crossings += 1
            if verbose:
                print(f"  Singularity crossing detected at step {step_count}, t={t:.6e}")
            if flip_after_crossing:
                flipped = not flipped
                if verbose:
                    print(f"  Orientation flipped")

        xk = xk_next
        dsk = min(dsk_next, max_ds)

    if verbose:
        print(f"Integration complete: {step_count} steps, {num_crossings} crossings")
        print(f"Final time: {t:.6e} (target: {t_final})")

    return np.array(sol), np.array(time)


def integrate_IDA(
    x0: NDArray,
    t_final: float,
    system: DAESystem,
    xdot0: NDArray | None = None,
    dt: float = DEFAULT_TIME_STEP,
    algebraic_vars_idx: list[int] | None = None
) -> tuple[NDArray, NDArray]:
    """
    Integrate a DAE system using scikits.odes (SUNDIALS IDA) in real time t.

    Args:
        x0: Initial state vector (n,)
        t_final: Final real time t
        system: DAE system to integrate
        xdot0: Initial time derivative (n,). If None, IDA will compute it.
        dt: Real time step size for output sampling
        algebraic_vars_idx: List of indices for algebraic variables (for IDA)

    Returns:
        y: Solution array of shape (len(t_eval), n)
        t: Time points corresponding to solution

    Note:
        This solver will likely fail near impasse points. Use integrate()
        with the impasse point algorithm for systems with singularities.
    """
    def residual(_t, x, xdot, result):
        # _t is intentionally unused
        f1 = system.G1(x)
        f2 = system.G2(x)

        r = f1.shape[0]  # Number of differential equations

        # Residual for differential equations: A1(x) @ xdot - G1(x) = 0
        A1_mat = system.A1(x)
        result[:r] = A1_mat @ xdot - f1

        # Residual for algebraic constraints: G2(x) = 0
        result[r:] = f2

    # If xdot0 not provided
    if xdot0 is None:
        # Provide a zero initial guess and let IDA refine it
        xdot0 = np.zeros_like(x0)

    # Configure IDA solver with appropriate options
    solver = dae('ida', residual,
                 algebraic_vars_idx=algebraic_vars_idx,
                 exclude_algvar_from_error=True,  # Don't use algebraic vars in error control
                 max_steps=500000)

    # Let IDA compute consistent initial conditions
    solver.set_options(compute_initcond='yp0')  # Compute consistent yp0

    t_eval = np.arange(0, t_final + dt, dt)
    sol = solver.solve(t_eval, x0, xdot0)

    return sol.values.y, sol.values.t