transient ¤

def free_numeric(handle, dependency=None):  # noqa: ANN001, ANN201, ARG001
    """Dispatch free to the correct backend based on handle type.

    ``klujax.KLUHandleManager`` exposes ``.close()``, which calls the backend's FFI free
    function.  Duck-typing is used so handles are freed correctly without ``isinstance``
    checks against concrete backend types.
    """
    if hasattr(handle, "close"):
        handle.close()
        return
    if _klujax_free_numeric is not None:
        _klujax_free_numeric(handle)

def setup_transient(
    groups: list, linear_strategy: CircuitLinearSolver, transient_solver: AbstractSolver = None
) -> Callable[..., diffrax.Solution]:
    """Configures and returns a function for executing transient analysis.

    This function acts as a factory, preparing a transient solver that is
    pre-configured with the circuit's linear strategy. It returns a callable
    that executes the time-domain simulation using `diffrax.diffeqsolve`.

    Args:
        groups (list): A list of component groups that define the circuit.
        linear_strategy (CircuitLinearSolver): The configured linear solver
            strategy, typically obtained from `analyze_circuit`.
        transient_solver (optional): The transient solver class to use.
            If None, `BDF2VectorizedTransientSolver` will be used.

    Returns:
        Callable[..., Any]: A function that executes the transient analysis.
        This returned function accepts the following arguments:

            t0 (float): The start time of the simulation.
            t1 (float): The end time of the simulation.
            dt0 (float): The initial time step for the solver.
            y0 (ArrayLike): The initial state vector of the system.
            saveat (diffrax.SaveAt, optional): Specifies time points at which
                to save the solution. Defaults to None.
            max_steps (int, optional): The maximum number of steps the solver
                can take. Defaults to 100000.
            throw (bool, optional): If True, the solver will raise an error on
                failure. Defaults to False.
            term (diffrax.AbstractTerm, optional): The term defining the ODE.
                Defaults to a zero-value ODETerm.
            stepsize_controller (diffrax.AbstractStepSizeController, optional):
                The step size controller. Defaults to `ConstantStepSize()`.
                Delays shorter than an accepted step are interpolated against
                the current Newton trial with the corresponding Jacobian, so
                adaptive controllers do not need a delay-based step cap.
            **kwargs: Additional keyword arguments to pass directly to
                `diffrax.diffeqsolve`.

    """
    fdomain_names = [g.name for g in groups.values() if getattr(g, "is_fdomain", False)]
    if fdomain_names:
        msg = (
            "Frequency-domain components cannot be used in transient simulation "
            "(time-domain convolution is not supported). "
            f"Offending groups: {fdomain_names}. "
            "Use setup_harmonic_balance() instead."
        )
        raise RuntimeError(msg)

    # Fixed time-delay support: gate the delay-history buffer behind whether
    # any group actually has a delayed component, so undelayed circuits pay
    # no cost. Delayed reads inside the current step retain their dependence
    # on the Newton trial, so neither fixed nor adaptive stepping needs to be
    # artificially capped by the smallest delay.
    has_delay = any(getattr(g, "has_delay", False) for g in groups.values())

    if transient_solver is None:
        # Pick the best Trap variant the linear solver supports.
        # Trap (2nd order, A-stable, zero numerical damping) is the right
        # pick for oscillators, LC tanks, and harmonic balance.  For circuits
        # with sharp digital edges that cause trapezoidal ringing, pass
        # BDF2*TransientSolver explicitly.
        if hasattr(linear_strategy, "refactor_jacobian"):
            transient_solver = TrapRefactoringTransientSolver
        elif hasattr(linear_strategy, "factor_jacobian"):
            transient_solver = TrapFactorizedTransientSolver
        else:
            transient_solver = TrapVectorizedTransientSolver

    import inspect

    tsolver = transient_solver(linear_solver=linear_strategy) if inspect.isclass(transient_solver) else transient_solver

    sys_size = linear_strategy.sys_size // 2 if linear_strategy.is_complex else linear_strategy.sys_size

    def _execute_transient(
        *,
        t0: float,
        t1: float,
        dt0: float,
        y0: ArrayLike,
        saveat: diffrax.SaveAt = None,
        max_steps: int = 100000,
        throw: bool = False,
        **kwargs: Any,
    ) -> diffrax.Solution:
        """Executes the transient simulation for the pre-configured circuit."""
        term = kwargs.pop("term", diffrax.ODETerm(lambda t, y, args: jnp.zeros_like(y)))
        solver = kwargs.pop("solver", tsolver)
        args = kwargs.pop("args", (groups, sys_size))
        stepsize_controller = kwargs.pop("stepsize_controller", ConstantStepSize())
        checkpoints = kwargs.pop("checkpoints", None)

        sol = circuit_diffeqsolve(
            terms=term,
            solver=solver,
            t0=t0,
            t1=t1,
            dt0=dt0,
            y0=y0,
            args=args,
            saveat=saveat,
            max_steps=max_steps,
            throw=throw,
            stepsize_controller=stepsize_controller,
            checkpoints=checkpoints,
            record_history=has_delay,
        )

        return sol

    return _execute_transient