linear ¤

def assume_full_rank(self) -> bool:
    """Indicate if the solver assumes the operator is full rank."""
    return False

def compute(self, state: Any, vector: jax.Array, options: Any) -> lx.Solution:
    """Satisfies the lineax API; call ``_solve_impl`` directly for internal use."""
    msg = "Directly call _solve_impl for internal use."
    raise NotImplementedError(msg)

def init(self, operator: Any, options: Any) -> Any:  # noqa: ARG002
    """Initialize the solver state (no-op for stateless solvers)."""
    return None

def solve_dc(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC operating point via damped Newton-Raphson.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess vector (shape ``[N]`` or ``[2N]``).

    Returns:
        Converged solution vector.

    """
    y, _ = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)
    return y

def solve_dc_auto(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    g_start: float = 1e-2,
    n_gmin: int = 10,
    n_source: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point with automatic homotopy fallback.

    Attempts a direct Newton solve first.  If it fails to converge, falls
    back to GMIN stepping followed by source stepping — all inside a single
    JIT-compiled kernel via ``jax.lax.cond``.

    Strategy:
    1.  ``_run_newton`` — plain damped Newton from ``y_guess``.
    2.  On failure: ``solve_dc_gmin`` (GMIN stepping) starting from
        ``y_guess``, then ``solve_dc_source`` (source stepping) from the
        GMIN result.

    Because ``jax.lax.cond`` evaluates both branches at *trace* time but
    only one at *runtime*, this compiles to a single kernel with no Python-
    level branching.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        g_start: Starting leakage for GMIN stepping (rescue branch).
        n_gmin: Number of GMIN steps in the rescue branch.
        n_source: Number of source steps in the rescue branch.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector.

    """
    y_direct, converged = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    def rescue(_: None) -> jax.Array:
        y_gmin = self.solve_dc_gmin(
            component_groups,
            y_guess,
            g_start=g_start,
            n_steps=n_gmin,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )
        return self.solve_dc_source(
            component_groups,
            y_gmin,
            n_steps=n_source,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )

    return jax.lax.cond(converged, lambda _: y_direct, rescue, None)

def solve_dc_checked(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> tuple[jax.Array, jax.Array]:
    """DC Operating Point with convergence status.

    Identical to :meth:`solve_dc` but additionally returns a boolean JAX
    scalar indicating whether the Newton-Raphson fixed-point iteration
    reported success.  Because the flag is a JAX array (not a Python bool)
    it can be consumed inside compiled programs:

    .. code-block:: python

        y, converged = solver.solve_dc_checked(groups, y0)
        # Outside JIT — inspect in Python:
        if not converged:
            y = solver.solve_dc_gmin(groups, y0)
        # Inside JIT — branch without Python-level control flow:
        y = jax.lax.cond(converged, lambda: y, lambda: solver.solve_dc_gmin(groups, y0))

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess vector (shape ``[N]`` or ``[2N]``).
        rtol: Relative tolerance for ``optx.fixed_point``.
        atol: Absolute tolerance for ``optx.fixed_point``.
        max_steps: Maximum Newton iterations.

    Returns:
        ``(y, converged)`` — solution vector and boolean success flag.

    """
    return self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

def solve_dc_gmin(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    g_start: float = 1e-2,
    n_steps: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point via GMIN stepping (homotopy rescue).

    Steps the diagonal regularisation conductance ``g_leak`` logarithmically
    from ``g_start`` down to ``self.g_leak``, using each converged solution
    as the warm start for the next step.  The large initial ``g_leak``
    linearises highly nonlinear components (diodes above threshold, lasers)
    that would otherwise cause Newton to diverge from a flat 0V start.

    Implemented with ``jax.lax.scan`` — fully JIT/grad/vmap-compatible.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        g_start: Starting leakage conductance (large value, e.g. ``1e-2``).
        n_steps: Number of log-uniform steps from ``g_start`` to
            ``self.g_leak``.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector after the full GMIN schedule.

    """
    g_values = jnp.logspace(jnp.log10(g_start), jnp.log10(self.g_leak), n_steps)

    def step(y: jax.Array, g_leak_val: jax.Array) -> tuple[jax.Array, None]:
        stepped_solver = eqx.tree_at(lambda s: s.g_leak, self, g_leak_val)
        y_new, _ = stepped_solver._run_newton(  # noqa: SLF001
            component_groups, y, rtol=rtol, atol=atol, max_steps=max_steps
        )
        return y_new, None

    y_final, _ = jax.lax.scan(step, y_guess, g_values)
    return y_final

def solve_dc_source(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    n_steps: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point via source stepping (homotopy rescue).

    Ramps all source amplitudes (components tagged with ``amplitude_param``)
    from 10 % to 100 % of their netlist values, using each converged
    solution as the warm start for the next step.  This guides Newton
    through the nonlinear region without the large initial step from 0V to
    full excitation.

    Implemented with ``jax.lax.scan`` — fully JIT/grad/vmap-compatible.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        n_steps: Number of uniformly-spaced steps from 0.1 to 1.0.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector at full source amplitude.

    """
    scales = jnp.linspace(0.1, 1.0, n_steps)

    def step(y: jax.Array, scale: jax.Array) -> tuple[jax.Array, None]:
        y_new, _ = self._run_newton(component_groups, y, source_scale=scale, rtol=rtol, atol=atol, max_steps=max_steps)
        return y_new, None

    y_final, _ = jax.lax.scan(step, y_guess, scales)
    return y_final

def assume_full_rank(self) -> bool:
    """Indicate if the solver assumes the operator is full rank."""
    return False

def compute(self, state: Any, vector: jax.Array, options: Any) -> lx.Solution:
    """Satisfies the lineax API; call ``_solve_impl`` directly for internal use."""
    msg = "Directly call _solve_impl for internal use."
    raise NotImplementedError(msg)

@classmethod
def from_component_groups(
    cls, component_groups: dict[str, Any], num_vars: int, *, is_complex: bool = False, g_leak: float = 1e-9
) -> "DenseSolver":
    """Factory method to pre-calculate indices for the dense matrix."""
    rows, cols, ground_idxs, sys_size = _build_index_arrays(component_groups, num_vars, is_complex)
    return cls(
        static_rows=jnp.array(rows),
        static_cols=jnp.array(cols),
        sys_size=sys_size,
        ground_indices=jnp.array(ground_idxs),
        is_complex=is_complex,
        g_leak=g_leak,
    )

def init(self, operator: Any, options: Any) -> Any:  # noqa: ARG002
    """Initialize the solver state (no-op for stateless solvers)."""
    return None

def solve_dc(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC operating point via damped Newton-Raphson.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess vector (shape ``[N]`` or ``[2N]``).

    Returns:
        Converged solution vector.

    """
    y, _ = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)
    return y

def solve_dc_auto(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    g_start: float = 1e-2,
    n_gmin: int = 10,
    n_source: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point with automatic homotopy fallback.

    Attempts a direct Newton solve first.  If it fails to converge, falls
    back to GMIN stepping followed by source stepping — all inside a single
    JIT-compiled kernel via ``jax.lax.cond``.

    Strategy:
    1.  ``_run_newton`` — plain damped Newton from ``y_guess``.
    2.  On failure: ``solve_dc_gmin`` (GMIN stepping) starting from
        ``y_guess``, then ``solve_dc_source`` (source stepping) from the
        GMIN result.

    Because ``jax.lax.cond`` evaluates both branches at *trace* time but
    only one at *runtime*, this compiles to a single kernel with no Python-
    level branching.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        g_start: Starting leakage for GMIN stepping (rescue branch).
        n_gmin: Number of GMIN steps in the rescue branch.
        n_source: Number of source steps in the rescue branch.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector.

    """
    y_direct, converged = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    def rescue(_: None) -> jax.Array:
        y_gmin = self.solve_dc_gmin(
            component_groups,
            y_guess,
            g_start=g_start,
            n_steps=n_gmin,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )
        return self.solve_dc_source(
            component_groups,
            y_gmin,
            n_steps=n_source,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )

    return jax.lax.cond(converged, lambda _: y_direct, rescue, None)

def solve_dc_checked(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> tuple[jax.Array, jax.Array]:
    """DC Operating Point with convergence status.

    Identical to :meth:`solve_dc` but additionally returns a boolean JAX
    scalar indicating whether the Newton-Raphson fixed-point iteration
    reported success.  Because the flag is a JAX array (not a Python bool)
    it can be consumed inside compiled programs:

    .. code-block:: python

        y, converged = solver.solve_dc_checked(groups, y0)
        # Outside JIT — inspect in Python:
        if not converged:
            y = solver.solve_dc_gmin(groups, y0)
        # Inside JIT — branch without Python-level control flow:
        y = jax.lax.cond(converged, lambda: y, lambda: solver.solve_dc_gmin(groups, y0))

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess vector (shape ``[N]`` or ``[2N]``).
        rtol: Relative tolerance for ``optx.fixed_point``.
        atol: Absolute tolerance for ``optx.fixed_point``.
        max_steps: Maximum Newton iterations.

    Returns:
        ``(y, converged)`` — solution vector and boolean success flag.

    """
    return self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

def solve_dc_gmin(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    g_start: float = 1e-2,
    n_steps: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point via GMIN stepping (homotopy rescue).

    Steps the diagonal regularisation conductance ``g_leak`` logarithmically
    from ``g_start`` down to ``self.g_leak``, using each converged solution
    as the warm start for the next step.  The large initial ``g_leak``
    linearises highly nonlinear components (diodes above threshold, lasers)
    that would otherwise cause Newton to diverge from a flat 0V start.

    Implemented with ``jax.lax.scan`` — fully JIT/grad/vmap-compatible.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        g_start: Starting leakage conductance (large value, e.g. ``1e-2``).
        n_steps: Number of log-uniform steps from ``g_start`` to
            ``self.g_leak``.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector after the full GMIN schedule.

    """
    g_values = jnp.logspace(jnp.log10(g_start), jnp.log10(self.g_leak), n_steps)

    def step(y: jax.Array, g_leak_val: jax.Array) -> tuple[jax.Array, None]:
        stepped_solver = eqx.tree_at(lambda s: s.g_leak, self, g_leak_val)
        y_new, _ = stepped_solver._run_newton(  # noqa: SLF001
            component_groups, y, rtol=rtol, atol=atol, max_steps=max_steps
        )
        return y_new, None

    y_final, _ = jax.lax.scan(step, y_guess, g_values)
    return y_final

def solve_dc_source(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    n_steps: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point via source stepping (homotopy rescue).

    Ramps all source amplitudes (components tagged with ``amplitude_param``)
    from 10 % to 100 % of their netlist values, using each converged
    solution as the warm start for the next step.  This guides Newton
    through the nonlinear region without the large initial step from 0V to
    full excitation.

    Implemented with ``jax.lax.scan`` — fully JIT/grad/vmap-compatible.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        n_steps: Number of uniformly-spaced steps from 0.1 to 1.0.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector at full source amplitude.

    """
    scales = jnp.linspace(0.1, 1.0, n_steps)

    def step(y: jax.Array, scale: jax.Array) -> tuple[jax.Array, None]:
        y_new, _ = self._run_newton(component_groups, y, source_scale=scale, rtol=rtol, atol=atol, max_steps=max_steps)
        return y_new, None

    y_final, _ = jax.lax.scan(step, y_guess, scales)
    return y_final

def assume_full_rank(self) -> bool:
    """Indicate if the solver assumes the operator is full rank."""
    return False

def compute(self, state: Any, vector: jax.Array, options: Any) -> lx.Solution:
    """Satisfies the lineax API; call ``_solve_impl`` directly for internal use."""
    msg = "Directly call _solve_impl for internal use."
    raise NotImplementedError(msg)

@classmethod
def from_component_groups(
    cls, component_groups: dict[str, Any], num_vars: int, *, is_complex: bool = False, g_leak: float = 1e-9
) -> "KLUSolver":
    """Factory method to pre-hash indices for sparse coalescence."""
    rows, cols, ground_idxs, sys_size = _build_index_arrays(component_groups, num_vars, is_complex)
    u_rows, u_cols, map_idx, n_unique = _klu_deduplicate(rows, cols, ground_idxs, sys_size)
    return cls(
        u_rows=jnp.array(u_rows),
        u_cols=jnp.array(u_cols),
        map_idx=jnp.array(map_idx),
        n_unique=n_unique,
        ground_indices=jnp.array(ground_idxs),
        sys_size=sys_size,
        is_complex=is_complex,
        g_leak=g_leak,
    )

def init(self, operator: Any, options: Any) -> Any:  # noqa: ARG002
    """Initialize the solver state (no-op for stateless solvers)."""
    return None

def solve_dc(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC operating point via damped Newton-Raphson.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess vector (shape ``[N]`` or ``[2N]``).

    Returns:
        Converged solution vector.

    """
    y, _ = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)
    return y

def solve_dc_auto(
    self,
    component_groups: dict[str, Any],
    y_guess: jax.Array,
    g_start: float = 1e-2,
    n_gmin: int = 10,
    n_source: int = 10,
    rtol: float = 1e-6,
    atol: float = 1e-6,
    max_steps: int = 100,
) -> jax.Array:
    """DC Operating Point with automatic homotopy fallback.

    Attempts a direct Newton solve first.  If it fails to converge, falls
    back to GMIN stepping followed by source stepping — all inside a single
    JIT-compiled kernel via ``jax.lax.cond``.

    Strategy:
    1.  ``_run_newton`` — plain damped Newton from ``y_guess``.
    2.  On failure: ``solve_dc_gmin`` (GMIN stepping) starting from
        ``y_guess``, then ``solve_dc_source`` (source stepping) from the
        GMIN result.

    Because ``jax.lax.cond`` evaluates both branches at *trace* time but
    only one at *runtime*, this compiles to a single kernel with no Python-
    level branching.

    Args:
        component_groups: Compiled circuit components.
        y_guess: Initial guess (typically ``jnp.zeros(sys_size)``).
        g_start: Starting leakage for GMIN stepping (rescue branch).
        n_gmin: Number of GMIN steps in the rescue branch.
        n_source: Number of source steps in the rescue branch.
        rtol: Relative tolerance for each inner Newton solve.
        atol: Absolute tolerance for each inner Newton solve.
        max_steps: Max Newton iterations per step.

    Returns:
        Converged solution vector.

    """
    y_direct, converged = self._run_newton(component_groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    def rescue(_: None) -> jax.Array:
        y_gmin = self.solve_dc_gmin(
            component_groups,
            y_guess,
            g_start=g_start,
            n_steps=n_gmin,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )
        return self.solve_dc_source(
            component_groups,
            y_gmin,
            n_steps=n_source,
            rtol=rtol,
            atol=atol,
            max_steps=max_steps,
        )

    return jax.lax.cond(converged, lambda _: y_direct, rescue, None)