assembly ¤

Assembly functions for the transient circuit solver.

Provides functions for evaluating the residual vectors and effective Jacobian of the discretised circuit equations at each Newton iteration. Functions are provided in two variants:

Full assembly (:func:assemble_system_real, :func:assemble_system_complex) — evaluates both the residual and the forward-mode Jacobian via jax.jacfwd. Used once per timestep to assemble and factor the frozen Jacobian in :class:~circulax.solver.FactorizedTransientSolver.
Residual only (:func:assemble_residual_only_real, :func:assemble_residual_only_complex) — evaluates only the primal residual, with no Jacobian computation. Used inside the Newton loop where the Jacobian has already been factored and only needs to be applied.

Each pair has a real and a complex variant. The complex variant operates on state vectors in unrolled block format — real parts concatenated with imaginary parts — allowing complex circuit analyses to reuse real-valued sparse linear algebra kernels.

Functions:

Name	Description
`assemble_gc_complex`	Return separate complex G and C COO value arrays at the linearisation point.
`assemble_gc_complex_2n`	Return four-block G and C COO value arrays for the 2N×2N real-block system.
`assemble_gc_real`	Return separate G and C COO value arrays at the linearisation point.
`assemble_residual_only_complex`	Assemble the residual vectors for an unrolled complex system, without computing the Jacobian.
`assemble_residual_only_real`	Assemble the residual vectors for a real system, without computing the Jacobian.
`assemble_system_complex`	Assemble the residual vectors and effective Jacobian values for an unrolled complex system.
`assemble_system_real`	Assemble the residual vectors and effective Jacobian values for a real system.

assemble_gc_complex ¤

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.

Source code in circulax/solvers/assembly.py

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 or group.has_delay:
            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)

assemble_gc_complex_2n ¤

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:

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

Source code in circulax/solvers/assembly.py

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 or group.has_delay:
            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]

assemble_gc_real ¤

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).

Parameters:

Name	Type	Description	Default
`y_guess`	`Array`	Linearisation point (DC operating point), shape `(num_vars,)`.	required
`component_groups`	`dict`	Compiled component groups from :func:`compile_netlist`.	required

Returns:

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

Source code in circulax/solvers/assembly.py

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 or group.has_delay:
            # 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)

assemble_residual_only_complex ¤

assemble_residual_only_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
    delay_hist: dict[str, Array] | None = None,
) -> 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.

Parameters:

Name	Type	Description	Default
`y_guess`	`Array`	Unrolled state vector of shape `(2 * num_vars,)`.	required
`component_groups`	`dict`	Compiled component groups returned by :func:`compile_netlist`, keyed by group name.	required
`t1`	`float`	Time at which the system is being evaluated.	required
`dt`	`float`	Present for signature symmetry with :func:`assemble_system_complex`.	required
`hist_t`	`Array \| None`	Accepted-step sample times for delayed reads, or `None` if the circuit has no delayed components.	`None`
`hist_y`	`Array \| None`	Accepted-step state samples for delayed reads, or `None`.	`None`
`delay_hist`	`dict[str, Array] \| None`	Per-group periodic delayed samples supplied by HB, or `None` outside harmonic balance.	`None`

Returns:

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

Source code in circulax/solvers/assembly.py

def assemble_residual_only_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
    delay_hist: dict[str, Array] | None = None,
) -> 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: Present for signature symmetry with :func:`assemble_system_complex`.
        hist_t: Accepted-step sample times for delayed reads, or ``None``
            if the circuit has no delayed components.
        hist_y: Accepted-step state samples for delayed reads, or ``None``.
        delay_hist: Per-group periodic delayed samples supplied by HB, or
            ``None`` outside harmonic balance.

    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]

        if group.has_delay and delay_hist is not None:
            hist_locs = delay_hist[group.name]
            physics_split = functools.partial(_complex_physics_hist, group=group, t1=t1)
            fr, fi, qr, qi = jax.vmap(physics_split)(v_r, v_i, group.params, hist_locs)
        elif group.has_delay and hist_t is not None:
            current = v_r + 1j * v_i
            hist_locs = _complex_hist_locs(group, group.params, current, t1, hist_t, hist_y, half_size)
            physics_split = functools.partial(_complex_physics_hist, group=group, t1=t1)
            fr, fi, qr, qi = jax.vmap(physics_split)(v_r, v_i, group.params, hist_locs)
        elif group.has_delay:
            physics_split = functools.partial(_complex_physics_hist, group=group, t1=t1)
            fr, fi, qr, qi = jax.vmap(physics_split)(v_r, v_i, group.params, v_r + 1j * v_i)
        else:
            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

assemble_residual_only_real ¤

assemble_residual_only_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
    delay_hist: dict[str, Array] | None = None,
) -> 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.

Parameters:

Name	Type	Description	Default
`y_guess`	`Array`	Current state vector of shape `(sys_size,)`.	required
`component_groups`	`dict`	Compiled component groups returned by :func:`compile_netlist`, keyed by group name.	required
`t1`	`float`	Time at which the system is being evaluated.	required
`dt`	`float`	Present for signature symmetry with :func:`assemble_system_real`.	required
`hist_t`	`Array \| None`	Accepted-step sample times for delayed reads, or `None` if the circuit has no delayed components.	`None`
`hist_y`	`Array \| None`	Accepted-step state samples for delayed reads, or `None`.	`None`
`delay_hist`	`dict[str, Array] \| None`	Per-group periodic delayed samples supplied by HB, or `None` outside harmonic balance.	`None`

Returns:

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

Source code in circulax/solvers/assembly.py

def assemble_residual_only_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
    delay_hist: dict[str, Array] | None = None,
) -> 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: Present for signature symmetry with :func:`assemble_system_real`.
        hist_t: Accepted-step sample times for delayed reads, or ``None``
            if the circuit has no delayed components.
        hist_y: Accepted-step state samples for delayed reads, or ``None``.
        delay_hist: Per-group periodic delayed samples supplied by HB, or
            ``None`` outside harmonic balance.

    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]

        if group.has_delay and delay_hist is not None:
            hist_locs = delay_hist[group.name]
            f_l, q_l = jax.vmap(functools.partial(_real_physics_hist, group=group, t1=t1))(v, group.params, hist_locs)
        elif group.has_delay and hist_t is not None:
            hist_locs = _real_hist_locs(group, group.params, v, t1, hist_t, hist_y)
            f_l, q_l = jax.vmap(functools.partial(_real_physics_hist, group=group, t1=t1))(v, group.params, hist_locs)
        elif group.has_delay:
            f_l, q_l = jax.vmap(functools.partial(_real_physics_hist, group=group, t1=t1))(v, group.params, v)
        else:
            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

assemble_system_complex ¤

assemble_system_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
) -> 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.

Parameters:

Name	Type	Description	Default
`y_guess`	`Array`	Unrolled state vector of shape `(2 * num_vars,)`, where `y_guess[:num_vars]` are real parts and `y_guess[num_vars:]` are imaginary parts.	required
`component_groups`	`dict`	Compiled component groups returned by :func:`compile_netlist`, keyed by group name.	required
`t1`	`float`	Time at which the system is being evaluated.	required
`dt`	`float`	Timestep duration, used to scale the reactive Jacobian blocks.	required
`source_scale`	`float`	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.	`1.0`
`alpha`	`float`	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.	`1.0`

Returns:

Type	Description
`Array`	A three-tuple `(total_f, total_q, jac_vals)` where:
`Array`	total_f — assembled resistive residual in unrolled format, shape `(2 * num_vars,)`.
`Array`	total_q — assembled reactive residual in unrolled format, shape `(2 * num_vars,)`.
`tuple[Array, Array, Array]`	jac_vals — concatenated non-zero values of the four effective Jacobian blocks (RR, RI, IR, II) in group-sorted order.

Source code in circulax/solvers/assembly.py

def assemble_system_complex(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
) -> 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
        )

        if group.has_delay and hist_t is not None:
            tau = _group_tau(group, params)

            def _delayed_physics(vr, vi, p, tau_i, idx_i, _group=group):
                def _with_current(vr_current, vi_current, p_current):
                    h_r = _interp_delayed(
                        hist_t,
                        hist_y[:, :half_size],
                        tau_i[None],
                        idx_i[None],
                        t1,
                        vr_current[None],
                    )[0]
                    h_i = _interp_delayed(
                        hist_t,
                        hist_y[:, half_size:],
                        tau_i[None],
                        idx_i[None],
                        t1,
                        vi_current[None],
                    )[0]
                    h = h_r + 1j * h_i
                    return _complex_physics_hist(vr_current, vi_current, p_current, h, group=_group, t1=t1)

                return _primal_and_jac_complex(_with_current, vr, vi, p)

            (fr, fi, qr, qi), (dfr_r, dfi_r, dqr_r, dqi_r), (dfr_i, dfi_i, dqr_i, dqi_i) = jax.vmap(_delayed_physics)(
                v_r, v_i, params, tau, group.var_indices
            )
        elif group.has_delay:

            def physics_split(vr, vi, p, _group=group):
                return _complex_physics_hist(vr, vi, p, vr + 1j * vi, 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)
        else:
            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

assemble_system_real ¤

assemble_system_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
) -> 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.

Parameters:

Name	Type	Description	Default
`y_guess`	`Array`	Current state vector of shape `(sys_size,)`.	required
`component_groups`	`dict`	Compiled component groups returned by :func:`compile_netlist`, keyed by group name.	required
`t1`	`float`	Time at which the system is being evaluated.	required
`dt`	`float`	Timestep duration, used to scale the reactive Jacobian block.	required
`source_scale`	`float`	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.	`1.0`
`alpha`	`float`	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.	`1.0`

Returns:

Type	Description
`Array`	A three-tuple `(total_f, total_q, jac_vals)` where:
`Array`	total_f — assembled resistive residual, shape `(sys_size,)`.
`Array`	total_q — assembled reactive residual, shape `(sys_size,)`.
`tuple[Array, Array, Array]`	jac_vals — concatenated non-zero values of the effective Jacobian in group-sorted order, ready to be passed to the sparse linear solver.

Source code in circulax/solvers/assembly.py

def assemble_system_real(
    y_guess: Array,
    component_groups: dict,
    t1: float,
    dt: float,
    source_scale: float = 1.0,
    alpha: float = 1.0,
    hist_t: Array | None = None,
    hist_y: Array | None = None,
) -> 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

        if group.has_delay and hist_t is not None:
            tau = _group_tau(group, params)

            def _delayed_physics(v, p, tau_i, idx_i, _group=group):
                def _with_current(v_current, p_current):
                    h = _interp_delayed(hist_t, hist_y, tau_i[None], idx_i[None], t1, v_current[None])[0]
                    return _group.physics_func(y=v_current, args=p_current, t=t1, hist=h)

                return _primal_and_jac_real(_with_current, v, p)

            (f_l, q_l), (df_l, dq_l) = jax.vmap(_delayed_physics)(v_locs, params, tau, group.var_indices)
        elif group.has_delay:
            # At DC every finite fixed delay is the identity. Differentiating
            # through hist=v includes both the direct and delayed dependence.
            def physics_at_t1(v, p, _group=group):
                return _group.physics_func(y=v, args=p, t=t1, hist=v)

            (f_l, q_l), (df_l, dq_l) = jax.vmap(functools.partial(_primal_and_jac_real, physics_at_t1))(v_locs, params)
        else:
            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)