assembly ¤

def assemble_gc_complex(
    y_guess: Array,
    component_groups: dict,
) -> tuple[Array, Array]:
    """Return separate complex G and C COO value arrays at the linearisation point.

    Complex counterpart of :func:`assemble_gc_real`.  The input state vector is
    in unrolled block format (real parts followed by imaginary parts, shape
    ``(2 * num_vars,)``).  Returns **complex128** COO value arrays aligned with
    the **N-sized** (not 2N) index arrays from ``_build_index_arrays(...,
    is_complex=False)``, suitable for building an N x N complex admittance
    matrix in the AC sweep.

    The complex Jacobian is recovered from the four real Jacobian blocks via
    the Wirtinger formula::

        J_complex = (1/2)(J_RR + J_II) + (j/2)(J_IR - J_RI)

    For holomorphic functions this simplifies to ``J_RR + j * J_IR``; for
    real-valued functions it reduces to ``J_RR + 0j``.

    """
    half = y_guess.shape[0] // 2
    y_real, y_imag = y_guess[:half], y_guess[half:]

    g_vals_list: list[Array] = []
    c_vals_list: list[Array] = []

    for k in sorted(component_groups.keys()):
        group = component_groups[k]
        n_entries = int(jnp.array(group.jac_rows).reshape(-1).shape[0])

        if _is_osdi(group):
            g_sep, c_sep = _assemble_osdi_gc_separate(y_real, group)
            g_vals_list.append(g_sep.reshape(-1).astype(jnp.complex128))
            c_vals_list.append(c_sep.reshape(-1).astype(jnp.complex128))
            continue

        if group.is_fdomain:
            g_vals_list.append(jnp.zeros(n_entries, dtype=jnp.complex128))
            c_vals_list.append(jnp.zeros(n_entries, dtype=jnp.complex128))
            continue

        v_r, v_i = y_real[group.var_indices], y_imag[group.var_indices]

        if group.combined_func is not None:
            _, _, df_l, dq_l = jax.vmap(lambda v, p: group.combined_func(v, p, 0.0))(v_r, group.params)
            g_vals_list.append(df_l.reshape(-1).astype(jnp.complex128))
            c_vals_list.append(dq_l.reshape(-1).astype(jnp.complex128))
            continue

        physics_split = functools.partial(_complex_physics, group=group, t1=0.0)

        (_, _, _, _), (dfr_r, dfi_r, dqr_r, dqi_r), (dfr_i, dfi_i, dqr_i, dqi_i) = jax.vmap(
            functools.partial(_primal_and_jac_complex, physics_split)
        )(v_r, v_i, group.params)

        g_complex = 0.5 * (dfr_r + dfi_i) + 0.5j * (dfi_r - dfr_i)
        c_complex = 0.5 * (dqr_r + dqi_i) + 0.5j * (dqi_r - dqr_i)

        g_vals_list.append(g_complex.reshape(-1))
        c_vals_list.append(c_complex.reshape(-1))

    return jnp.concatenate(g_vals_list), jnp.concatenate(c_vals_list)

def assemble_gc_complex_2n(
    y_guess: Array,
    component_groups: dict,
) -> tuple[list[Array], list[Array]]:
    """Return four-block G and C COO value arrays for the 2N×2N real-block system.

    Unlike :func:`assemble_gc_complex` which collapses the four Jacobian blocks
    into N×N complex matrices via the Wirtinger formula, this function keeps the
    blocks separate — returning ``[RR, RI, IR, II]`` for both G and C.  This is
    required for circuits containing non-holomorphic operations (e.g.
    ``jnp.real()``) where the Wirtinger N×N system does not fully capture the
    conjugate field response.

    Returns:
        ``(G_blocks, C_blocks)`` where each is a list of four JAX arrays
        ``[RR, RI, IR, II]``.  Each array contains the COO values aligned with
        the N-sized index arrays from ``_build_index_arrays(..., is_complex=False)``.

    """
    half = y_guess.shape[0] // 2
    y_real, y_imag = y_guess[:half], y_guess[half:]

    g_blocks: list[list[Array]] = [[], [], [], []]
    c_blocks: list[list[Array]] = [[], [], [], []]

    for k in sorted(component_groups.keys()):
        group = component_groups[k]
        n_entries = int(jnp.array(group.jac_rows).reshape(-1).shape[0])

        if _is_osdi(group):
            g_sep, c_sep = _assemble_osdi_gc_separate(y_real, group)
            flat_g = g_sep.reshape(-1)
            flat_c = c_sep.reshape(-1)
            zeros = jnp.zeros_like(flat_g)
            g_blocks[0].append(flat_g)
            g_blocks[1].append(zeros)
            g_blocks[2].append(zeros)
            g_blocks[3].append(zeros)
            c_blocks[0].append(flat_c)
            c_blocks[1].append(zeros)
            c_blocks[2].append(zeros)
            c_blocks[3].append(zeros)
            continue

        if group.is_fdomain:
            zeros = jnp.zeros(n_entries, dtype=jnp.float64)
            for blk in g_blocks:
                blk.append(zeros)
            for blk in c_blocks:
                blk.append(zeros)
            continue

        v_r, v_i = y_real[group.var_indices], y_imag[group.var_indices]

        if group.combined_func is not None:
            _, _, df_l, dq_l = jax.vmap(lambda v, p: group.combined_func(v, p, 0.0))(v_r, group.params)
            flat_g = df_l.reshape(-1)
            flat_c = dq_l.reshape(-1)
            zeros = jnp.zeros_like(flat_g)
            g_blocks[0].append(flat_g)
            g_blocks[1].append(zeros)
            g_blocks[2].append(zeros)
            g_blocks[3].append(zeros)
            c_blocks[0].append(flat_c)
            c_blocks[1].append(zeros)
            c_blocks[2].append(zeros)
            c_blocks[3].append(zeros)
            continue

        physics_split = functools.partial(_complex_physics, group=group, t1=0.0)

        (_, _, _, _), (dfr_r, dfi_r, dqr_r, dqi_r), (dfr_i, dfi_i, dqr_i, dqi_i) = jax.vmap(
            functools.partial(_primal_and_jac_complex, physics_split)
        )(v_r, v_i, group.params)

        g_blocks[0].append(dfr_r.reshape(-1))
        g_blocks[1].append(dfr_i.reshape(-1))
        g_blocks[2].append(dfi_r.reshape(-1))
        g_blocks[3].append(dfi_i.reshape(-1))
        c_blocks[0].append(dqr_r.reshape(-1))
        c_blocks[1].append(dqr_i.reshape(-1))
        c_blocks[2].append(dqi_r.reshape(-1))
        c_blocks[3].append(dqi_i.reshape(-1))

    return [jnp.concatenate(b) for b in g_blocks], [jnp.concatenate(b) for b in c_blocks]

def assemble_gc_real(
    y_guess: Array,
    component_groups: dict,
) -> tuple[Array, Array]:
    """Return separate G and C COO value arrays at the linearisation point.

    Mirrors :func:`assemble_system_real` but returns ``df/dy`` (conductance) and
    ``dq/dy`` (capacitance) separately instead of combining them as
    ``G + C/dt``.  Frequency-domain groups contribute zero-filled blocks so that
    the returned arrays align with the static COO index arrays produced by
    ``_build_index_arrays`` (which includes all groups).

    Args:
        y_guess: Linearisation point (DC operating point), shape ``(num_vars,)``.
        component_groups: Compiled component groups from :func:`compile_netlist`.

    Returns:
        A two-tuple ``(G_vals, C_vals)`` of real-valued 1-D JAX arrays.  Both
        have the same length as the concatenated ``jac_rows``/``jac_cols`` COO
        index arrays from ``_build_index_arrays``.

    """
    g_vals_list = []
    c_vals_list = []

    for k in sorted(component_groups.keys()):
        group = component_groups[k]
        n_entries = int(jnp.array(group.jac_rows).reshape(-1).shape[0])

        if _is_osdi(group):
            g_sep, c_sep = _assemble_osdi_gc_separate(y_guess, group)
            g_vals_list.append(g_sep.reshape(-1).astype(y_guess.dtype))
            c_vals_list.append(c_sep.reshape(-1).astype(y_guess.dtype))
            continue

        if group.is_fdomain:
            # Fdomain groups are re-evaluated per-frequency in ac_sweep.
            # Emit zero blocks so COO alignment with _build_index_arrays is preserved.
            g_vals_list.append(jnp.zeros(n_entries, dtype=y_guess.dtype))
            c_vals_list.append(jnp.zeros(n_entries, dtype=y_guess.dtype))
            continue

        v_locs = y_guess[group.var_indices]

        # Direct combined bypass for VA components with combined_fn.
        if group.combined_func is not None:
            _, _, df_l, dq_l = jax.vmap(
                lambda v, p: group.combined_func(v, p, 0.0)
            )(v_locs, group.params)
            g_vals_list.append(df_l.reshape(-1))
            c_vals_list.append(dq_l.reshape(-1))
            continue

        physics_at_dc = functools.partial(_real_physics, group=group, t1=0.0)

        (_, _), (df_l, dq_l) = jax.vmap(functools.partial(_primal_and_jac_real, physics_at_dc))(v_locs, group.params)

        g_vals_list.append(df_l.reshape(-1))
        c_vals_list.append(dq_l.reshape(-1))

    return jnp.concatenate(g_vals_list), jnp.concatenate(c_vals_list)

def assemble_residual_only_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
) -> tuple[Array, Array]:
    """Assemble the residual vectors for an unrolled complex system, without computing the Jacobian.

    The complex counterpart of :func:`assemble_residual_only_real`. The state
    vector is expected in unrolled block format (real parts followed by imaginary
    parts) matching the layout used by :func:`assemble_system_complex`.

    Args:
        y_guess: Unrolled state vector of shape ``(2 * num_vars,)``.
        component_groups: Compiled component groups returned by
            :func:`compile_netlist`, keyed by group name.
        t1: Time at which the system is being evaluated.
        dt: Unused; present for signature symmetry with
            :func:`assemble_system_complex` so the two functions are
            interchangeable at call sites.

    Returns:
        A two-tuple ``(total_f, total_q)`` where both arrays have shape
        ``(2 * num_vars,)`` and ``dtype`` matching ``y_guess.dtype``.

    """
    sys_size = y_guess.shape[0]
    half_size = sys_size // 2
    y_real, y_imag = y_guess[:half_size], y_guess[half_size:]

    total_f = jnp.zeros(sys_size, dtype=y_guess.dtype)
    total_q = jnp.zeros(sys_size, dtype=y_guess.dtype)

    for k in sorted(component_groups.keys()):
        group = component_groups[k]

        if _is_osdi(group):
            f_l, q_l, _ = _assemble_osdi_group(
                y_guess[:half_size], group, alpha=1.0, dt=1.0, residual_only=True,
            )
            total_f = total_f.at[group.eq_indices].add(f_l)
            total_q = total_q.at[group.eq_indices].add(q_l)
            continue

        if group.is_fdomain:
            # F-domain groups have no time-domain physics; their contribution is
            # added directly in the frequency domain by the HB solver.
            continue

        v_r, v_i = y_real[group.var_indices], y_imag[group.var_indices]

        physics_split = functools.partial(_complex_physics, group=group, t1=t1)

        fr, fi, qr, qi = jax.vmap(physics_split)(v_r, v_i, group.params)

        idx_r, idx_i = group.eq_indices, group.eq_indices + half_size
        total_f = total_f.at[idx_r].add(fr).at[idx_i].add(fi)
        total_q = total_q.at[idx_r].add(qr).at[idx_i].add(qi)

    return total_f, total_q

def assemble_residual_only_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
) -> tuple[Array, Array]:
    """Assemble the residual vectors for a real system, without computing the Jacobian.

    Cheaper than :func:`assemble_system_real` as it performs only primal
    evaluations. Used inside the frozen-Jacobian Newton loop where the
    Jacobian has already been factored and only the residual needs to be
    recomputed at each iteration.

    Args:
        y_guess: Current state vector of shape ``(sys_size,)``.
        component_groups: Compiled component groups returned by
            :func:`compile_netlist`, keyed by group name.
        t1: Time at which the system is being evaluated.
        dt: Unused; present for signature symmetry with
            :func:`assemble_system_real` so the two functions are
            interchangeable at call sites.

    Returns:
        A two-tuple ``(total_f, total_q)`` where both arrays have shape
        ``(sys_size,)`` and ``dtype`` matching ``y_guess.dtype``.

    """
    sys_size = y_guess.shape[0]
    total_f = jnp.zeros(sys_size, dtype=y_guess.dtype)
    total_q = jnp.zeros(sys_size, dtype=y_guess.dtype)

    for k in sorted(component_groups.keys()):
        group = component_groups[k]

        if _is_osdi(group):
            f_l, q_l, _ = _assemble_osdi_group(
                y_guess, group, alpha=1.0, dt=1.0, residual_only=True,
            )
            total_f = total_f.at[group.eq_indices].add(f_l)
            total_q = total_q.at[group.eq_indices].add(q_l)
            continue

        if group.is_fdomain:
            # F-domain groups have no time-domain physics; their contribution is
            # added directly in the frequency domain by the HB solver.
            continue

        v = y_guess[group.var_indices]

        physics_at_t1 = functools.partial(_real_physics, group=group, t1=t1)

        f_l, q_l = jax.vmap(physics_at_t1)(v, group.params)

        total_f = total_f.at[group.eq_indices].add(f_l)
        total_q = total_q.at[group.eq_indices].add(q_l)

    return total_f, total_q

def assemble_system_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
) -> tuple[Array, Array, Array]:
    """Assemble the residual vectors and effective Jacobian values for an unrolled complex system.

    The complex state vector is stored in unrolled (block) format: the first
    half of ``y_guess`` holds the real parts of all node voltages/states, the
    second half holds the imaginary parts. This avoids JAX's limited support
    for complex-valued sparse linear solvers by keeping all arithmetic real.

    The Jacobian is split into four real blocks — RR, RI, IR, II — representing
    the partial derivatives of the real and imaginary residual components with
    respect to the real and imaginary state components respectively. The blocks
    are concatenated in RR→RI→IR→II order to match the sparsity index layout
    produced during compilation.

    Args:
        y_guess: Unrolled state vector of shape ``(2 * num_vars,)``, where
            ``y_guess[:num_vars]`` are real parts and ``y_guess[num_vars:]``
            are imaginary parts.
        component_groups: Compiled component groups returned by
            :func:`compile_netlist`, keyed by group name.
        t1: Time at which the system is being evaluated.
        dt: Timestep duration, used to scale the reactive Jacobian blocks.
        source_scale: Multiplicative scale applied to source amplitudes
            (components whose ``amplitude_param`` is set).  Use ``1.0``
            for a standard evaluation and values in ``(0, 1)`` during
            DC homotopy source stepping.
        alpha: Jacobian scaling factor for the reactive blocks.  Use ``1.0``
            for Backward Euler, the variable-step BDF2 ``α₀`` coefficient for
            BDF2, or ``1/γ`` for SDIRK3 stages.

    Returns:
        A three-tuple ``(total_f, total_q, jac_vals)`` where:

        - **total_f** — assembled resistive residual in unrolled format,
            shape ``(2 * num_vars,)``.
        - **total_q** — assembled reactive residual in unrolled format,
            shape ``(2 * num_vars,)``.
        - **jac_vals** — concatenated non-zero values of the four effective
            Jacobian blocks (RR, RI, IR, II) in group-sorted order.

    """
    sys_size = y_guess.shape[0]
    half_size = sys_size // 2
    y_real, y_imag = y_guess[:half_size], y_guess[half_size:]

    total_f = jnp.zeros(sys_size, dtype=jnp.float64)
    total_q = jnp.zeros(sys_size, dtype=jnp.float64)

    vals_blocks: list[list[Array]] = [[], [], [], []]

    for k in sorted(component_groups.keys()):
        group = component_groups[k]

        if _is_osdi(group):
            f_l, q_l, j_eff = _assemble_osdi_group(y_guess[:half_size], group, alpha, dt)
            total_f = total_f.at[group.eq_indices].add(f_l)
            total_q = total_q.at[group.eq_indices].add(q_l)
            vals_blocks[0].append(j_eff.reshape(-1))
            vals_blocks[1].append(jnp.zeros(j_eff.size, dtype=jnp.float64))
            vals_blocks[2].append(jnp.zeros(j_eff.size, dtype=jnp.float64))
            vals_blocks[3].append(jnp.zeros(j_eff.size, dtype=jnp.float64))
            continue

        if group.is_fdomain:
            # F-domain component: evaluate admittance at f=0 (DC) — complex circuit path.
            v_r, v_i = y_real[group.var_indices], y_imag[group.var_indices]
            v_c = v_r + 1j * v_i  # (N, n_ports) complex
            Y_mats = jax.vmap(lambda p: group.physics_func(0.0, p))(group.params)
            i_c = jnp.einsum("nij,nj->ni", Y_mats, v_c)  # (N, n_ports) complex
            idx_r, idx_i = group.eq_indices, group.eq_indices + half_size
            total_f = total_f.at[idx_r].add(i_c.real).at[idx_i].add(i_c.imag)
            # Jacobian blocks: dI/dVr = Y.real, dI/dVi = -Y.imag (by Cauchy-Riemann)
            # For general complex Y: dIr/dVr = Yr, dIr/dVi = -Yi, dIi/dVr = Yi, dIi/dVi = Yr
            Yr = Y_mats.real  # (N, n_ports, n_ports)
            Yi = Y_mats.imag
            vals_blocks[0].append(Yr.reshape(-1))  # RR: dIr/dVr
            vals_blocks[1].append((-Yi).reshape(-1))  # RI: dIr/dVi
            vals_blocks[2].append(Yi.reshape(-1))  # IR: dIi/dVr
            vals_blocks[3].append(Yr.reshape(-1))  # II: dIi/dVi
            continue

        v_r, v_i = y_real[group.var_indices], y_imag[group.var_indices]

        ap = group.amplitude_param
        params = (
            eqx.tree_at(lambda p, _ap=ap: getattr(p, _ap), group.params, getattr(group.params, ap) * source_scale)
            if ap
            else group.params
        )

        physics_split = functools.partial(_complex_physics, group=group, t1=t1)

        (fr, fi, qr, qi), (dfr_r, dfi_r, dqr_r, dqi_r), (dfr_i, dfi_i, dqr_i, dqi_i) = jax.vmap(
            functools.partial(_primal_and_jac_complex, physics_split)
        )(v_r, v_i, params)

        idx_r, idx_i = group.eq_indices, group.eq_indices + half_size
        total_f = total_f.at[idx_r].add(fr).at[idx_i].add(fi)
        total_q = total_q.at[idx_r].add(qr).at[idx_i].add(qi)

        vals_blocks[0].append((dfr_r + (alpha / dt) * dqr_r).reshape(-1))  # RR
        vals_blocks[1].append((dfr_i + (alpha / dt) * dqr_i).reshape(-1))  # RI
        vals_blocks[2].append((dfi_r + (alpha / dt) * dqi_r).reshape(-1))  # IR
        vals_blocks[3].append((dfi_i + (alpha / dt) * dqi_i).reshape(-1))  # II

    all_vals = jnp.concatenate([jnp.concatenate(b) for b in vals_blocks])
    return total_f, total_q, all_vals

def assemble_system_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
) -> tuple[Array, Array, Array]:
    """Assemble the residual vectors and effective Jacobian values for a real system.

    For each component group, evaluates the physics at ``t1`` and computes the
    forward-mode Jacobian via ``jax.jacfwd``. The effective Jacobian combines
    the resistive and reactive contributions as ``J_eff = df/dy + (alpha/dt) * dq/dy``,
    where ``alpha=1`` recovers Backward Euler and ``alpha=3/2`` (uniform step) gives BDF2.

    Components are processed in sorted key order to ensure a deterministic
    non-zero layout in the sparse Jacobian, which is required for the
    factorisation step.

    Args:
        y_guess: Current state vector of shape ``(sys_size,)``.
        component_groups: Compiled component groups returned by
            :func:`compile_netlist`, keyed by group name.
        t1: Time at which the system is being evaluated.
        dt: Timestep duration, used to scale the reactive Jacobian block.
        source_scale: Multiplicative scale applied to source amplitudes
            (components whose ``amplitude_param`` is set).  Use ``1.0``
            for a standard evaluation and values in ``(0, 1)`` during
            DC homotopy source stepping.
        alpha: Jacobian scaling factor for the reactive block.  Use ``1.0``
            for Backward Euler, the variable-step BDF2 ``α₀`` coefficient for
            BDF2, or ``1/γ`` for SDIRK3 stages.

    Returns:
        A three-tuple ``(total_f, total_q, jac_vals)`` where:

        - **total_f** — assembled resistive residual, shape ``(sys_size,)``.
        - **total_q** — assembled reactive residual, shape ``(sys_size,)``.
        - **jac_vals** — concatenated non-zero values of the effective Jacobian
            in group-sorted order, ready to be passed to the sparse linear solver.

    """
    sys_size = y_guess.shape[0]
    total_f = jnp.zeros(sys_size, dtype=y_guess.dtype)
    total_q = jnp.zeros(sys_size, dtype=y_guess.dtype)
    vals_list = []

    for k in sorted(component_groups.keys()):
        group = component_groups[k]

        if _is_osdi(group):
            f_l, q_l, j_eff = _assemble_osdi_group(y_guess, group, alpha, dt)
            total_f = total_f.at[group.eq_indices].add(f_l)
            total_q = total_q.at[group.eq_indices].add(q_l)
            vals_list.append(j_eff.reshape(-1))
            continue

        if group.is_fdomain:
            # F-domain component: evaluate admittance at f=0 (DC).
            v_locs = y_guess[group.var_indices]
            Y_mats = jax.vmap(lambda p: group.physics_func(0.0, p))(group.params)
            Y_real = Y_mats.real  # (N, n_ports, n_ports)
            f_l = jnp.einsum("nij,nj->ni", Y_real, v_locs)  # (N, n_ports)
            total_f = total_f.at[group.eq_indices].add(f_l)
            vals_list.append(Y_real.reshape(-1))  # Jacobian = Y at DC
            continue

        v_locs = y_guess[group.var_indices]

        ap = group.amplitude_param
        params = (
            eqx.tree_at(lambda p, _ap=ap: getattr(p, _ap), group.params, getattr(group.params, ap) * source_scale)
            if ap
            else group.params
        )

        # Direct combined bypass — calls _fast_combined ONCE per device per Newton
        # iteration instead of routing through vmap(jvp(fast_physics, v, eye[i]))
        # which fires the custom JVP n times.  Only active when the VA emitter
        # produced a combined_fn (i.e. group.combined_func is not None).
        if group.combined_func is not None:
            f_l, q_l, df_l, dq_l = jax.vmap(
                lambda v, p: group.combined_func(v, p, t1)
            )(v_locs, params)
            total_f = total_f.at[group.eq_indices].add(f_l)
            total_q = total_q.at[group.eq_indices].add(q_l)
            j_eff = df_l + (alpha / dt) * dq_l
            vals_list.append(j_eff.reshape(-1))
            continue

        physics_at_t1 = functools.partial(_real_physics, group=group, t1=t1)

        (f_l, q_l), (df_l, dq_l) = jax.vmap(functools.partial(_primal_and_jac_real, physics_at_t1))(v_locs, params)

        total_f = total_f.at[group.eq_indices].add(f_l)
        total_q = total_q.at[group.eq_indices].add(q_l)
        j_eff = df_l + (alpha / dt) * dq_l
        vals_list.append(j_eff.reshape(-1))

    return total_f, total_q, jnp.concatenate(vals_list)