Add Circulax transmon optimization notebook with harmonic-balance and transient analysis

[Copilot]

Adds an example notebook demonstrating Circulax (differentiable JAX-based circuit simulator) integration with qpdk for transmon qubit design and optimization.

Notebook content

• Josephson junction component — Custom @component with sinusoidal current-phase relation in flux formulation
• Layout→circuit mapping — Converts qpdk double_pad_transmon geometry (pad area, gap, JJ area) to circuit parameters (C_s, I_c, E_J, E_C)
• Harmonic-balance analysis — Periodic steady-state of a driven transmon via Circulax's HB solver
• Gradient-based optimization — jax.grad + Optax to optimize Ic/Cs toward target f₀₁ = 5 GHz and α = −300 MHz
• Transient crosstalk simulation — Two coupled qubits with parasitic C_m, pulse on Q1, measure induced voltage on Q2
• Differentiable crosstalk metric — Gradient of peak victim voltage w.r.t. coupling capacitance for layout optimization

@component(ports=("p1", "p2"), states=("phi",))
def JosephsonJunction(signals, s, Ic=50e-9, R_sub=1e6):
    i_jj = Ic * jnp.sin(s.phi)
    v_drop = signals.p1 - signals.p2
    return (
        {"p1": i_jj + v_drop/R_sub, "p2": -(i_jj + v_drop/R_sub), "phi": v_drop},
        {"phi": -FLUX_PER_RAD * s.phi},
    )

Other changes

• pyproject.toml — circulax = ["circulax>=0.1.7", "optax"] optional dependency group
• docs/notebooks.rst — New "Differentiable circuit simulation" section
• docs/conf.py — Added to nb_execution_excludepatterns (external dep, like HFSS/MATLAB notebooks)
• uv.lock — Regenerated with circulax resolution

Summary by Sourcery

Add a Circulax-based differentiable transmon circuit simulation and optimization example notebook, along with supporting script, and wire it into the docs and dependency configuration.

New Features:

• Provide a new Jupyter notebook and Python script demonstrating Circulax integration with qpdk for transmon qubit design, harmonic-balance analysis, and transient crosstalk optimization.

Enhancements:

• Adjust existing JAX jit usage in generic LC resonator models for improved compatibility with functools.partial.

Build:

• Introduce an optional 'circulax' dependency group with Circulax and Optax in the project configuration.

Documentation:

• Document the new differentiable circuit simulation example in the notebooks index and exclude the Circulax notebook from automated execution in the Sphinx config.

[sourcery-ai]

Reviewer's Guide

Adds a new Circulax-based transmon optimization example (notebook + jupytext .py) integrating qpdk layout with differentiable circuit simulation, plus small core-library and docs updates to support the example.

Sequence diagram for harmonic-balance transmon optimization workflow

sequenceDiagram
    actor User
    participant NB as Notebook
    participant QPDK as QPDK
    participant Cx as Circulax
    participant JAX as JAX
    participant Optax as Optax

    User->>NB: Run notebook cells (Part 1)

    NB->>QPDK: PDK.activate()
    QPDK-->>NB: Activated PDK

    NB->>QPDK: double_pad_transmon parameters
    QPDK-->>NB: pad_size, pad_gap

    NB->>QPDK: plate_capacitor_capacitance_analytical
    QPDK-->>NB: Cs_initial

    NB->>NB: layout_to_circuit_params
    NB-->>NB: Ic_initial, Cs_initial, EJ2_ratio

    NB->>Cx: build_transmon_netlist(Cs, Ic, EJ2_ratio)
    Cx-->>NB: netlist_dict

    NB->>Cx: compile_netlist(netlist_dict, models)
    Cx-->>NB: groups, num_vars, port_map

    NB->>Cx: analyze_circuit(groups, num_vars)
    Cx-->>NB: solver

    NB->>Cx: solver.solve_dc(initial_guess)
    Cx-->>NB: y_dc

    NB->>Cx: setup_harmonic_balance(groups, num_vars, freq, num_harmonics)
    Cx-->>NB: run_hb callable

    NB->>Cx: run_hb(y_dc)
    Cx-->>NB: y_time, y_freq

    NB->>NB: define hb_loss_fn(params_vec)
    NB->>JAX: grad = jax.grad(hb_loss_fn)
    JAX-->>NB: grad_fn

    NB->>Optax: optimizer = optax.adam()
    Optax-->>NB: opt_state

    loop Optimization steps
        NB->>JAX: grads = grad_fn(params_opt)
        JAX-->>NB: grads

        NB->>Optax: updates, opt_state = optimizer.update(grads, opt_state)
        Optax-->>NB: updates, opt_state

        NB->>NB: params_opt = apply_updates(params_opt, updates)
        NB->>NB: evaluate hb_loss_fn(params_opt)
    end

    NB-->>User: Optimized Ic, Cs, f01, alpha

Loading

Sequence diagram for transient crosstalk simulation and gradient computation

sequenceDiagram
    actor User
    participant NB as Notebook
    participant Cx as Circulax
    participant Diffrax as Diffrax
    participant JAX as JAX

    User->>NB: Run notebook cells (Part 2)

    NB->>NB: Define CouplingCapacitor component
    NB->>Cx: coupled_netlist with JJ1, JJ2, Cs1, Cs2, Cm, Vpulse
    Cx-->>NB: netlist_dict

    NB->>Cx: compile_netlist(netlist_dict, models_coupled)
    Cx-->>NB: groups_c, num_vars_c, port_map_c

    NB->>Cx: analyze_circuit(groups_c, num_vars_c)
    Cx-->>NB: solver_c

    NB->>Cx: solver_c.solve_dc(initial_guess)
    Cx-->>NB: y_dc_c

    NB->>Cx: sim = setup_transient(groups_c, solver_c)
    Cx-->>NB: sim callable

    NB->>Diffrax: sim(t0, t1, dt0, y_dc_c, SaveAt(ts))
    Diffrax-->>NB: sol.ys over time

    NB->>NB: Extract v_q1, v_q2 from sol.ys
    NB-->>User: Crosstalk metrics

    NB->>NB: Define crosstalk_metric(log_Cm)
    NB->>JAX: grad_fn = jax.grad(crosstalk_metric)
    JAX-->>NB: grad_fn

    NB->>JAX: grad_crosstalk = grad_fn(log_Cm_init)
    JAX-->>NB: grad_crosstalk

    NB-->>User: ∂(crosstalk)/∂(log Cm) for layout tuning

Loading

Flow diagram for layout-to-circuit mapping and optimization loop

flowchart LR
    subgraph Layout_to_Circuit
        L_Params["Layout parameters\n(pad_width_um, pad_height_um, pad_gap_um, jj_area_um2)"]
        L_Map["layout_to_circuit_params"]
        C_Params["Circuit parameters\n(Cs, Ic, EJ2_ratio)"]
    end

    subgraph Optimization_Loop
        P_Init["Initial params_vec = [log(Ic), log(Cs)]"]
        LossFn["hb_loss_fn(params_vec)"]
        Net["build_transmon_netlist(Cs, Ic)"]
        Compile["compile_netlist"]
        HB["setup_harmonic_balance + run_hb"]
        Spectrum["Extract V_1st_harmonic at JJ1"]
        Anharm["transmon_anharmonicity"]
        LossVal["resonance_loss + 0.1*alpha_error"]
        Grad["jax.grad(hb_loss_fn)"]
        Opt["optax.adam update"]
        P_Update["Updated params_vec"]
    end

    L_Params --> L_Map --> C_Params
    C_Params --> P_Init

    P_Init --> LossFn
    LossFn --> Net --> Compile --> HB --> Spectrum --> LossVal
    LossFn --> Anharm --> LossVal

    P_Init --> Grad --> Opt --> P_Update --> LossFn

Loading

File-Level Changes

Change	Details	Files
Introduce a Circulax-based, differentiable transmon design and crosstalk analysis example notebook (with jupytext .py mirror) that maps qpdk transmon layout parameters to circuit models, runs harmonic-balance and transient simulations, and performs gradient-based optimization using JAX/Optax.	Define a custom JosephsonJunction Circulax component using a flux-based formulation with optional second-harmonic term and resistive shunt for stability. Add layout_to_circuit_params helper to convert qpdk transmon geometry (pad sizes, gaps, JJ area, critical current density) to circuit parameters (Cs, Ic, EJ2_ratio) using analytical capacitor models. Build a driven transmon netlist using the custom JJ, capacitor, resistor, and voltage source components and connect it to Circulax harmonic-balance solver to compute periodic steady state and spectra. Implement a differentiable harmonic-balance loss function over log(Ic) and log(Cs), combining resonance amplitude at target f01 and anharmonicity penalty, then optimize with Optax/Adam using jax.grad. Add a second workflow building a two-qubit coupled netlist with a custom CouplingCapacitor and SmoothPulse sources, running transient simulation with Diffrax-based solver to analyze crosstalk waveforms. Define a differentiable crosstalk_metric that rebuilds the circuit as a function of coupling capacitance, runs a transient solve, measures victim qubit peak voltage, and differentiates w.r.t. log(Cm) to obtain gradients for layout optimization. Mirror all notebook logic into a jupytext-compatible Python script in notebooks/src for reproducible execution and version control friendliness.	notebooks/circulax_transmon_optimization.ipynb notebooks/src/circulax_transmon_optimization.py
Enable Circulax as an optional dependency and ensure docs build infrastructure is aware of the new notebook and its execution constraints.	Add a circulax optional extra group declaring Circulax and Optax dependencies. Register the new notebook in the Sphinx notebooks index under a "Differentiable circuit simulation" section (or equivalent) so it shows up in rendered docs. Exclude the Circulax notebooks from automated doc notebook execution to avoid failures when optional dependencies are missing. Regenerate the uv.lock lockfile to capture the new optional dependency resolution.	pyproject.toml docs/notebooks.rst docs/conf.py uv.lock
Adjust existing generic LC resonator model JAX compilation to use functools.partial for jit configuration, improving compatibility with JAX static_argnames usage.	Replace bare @jax.jit(static_argnames=["grounded"]) decorators on lc_resonator and lc_resonator_coupled with @partial(jax.jit, static_argnames=["grounded"]).	qpdk/models/generic.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[nikosavola]

@copilot resolve the merge conflicts in this pull request

[Copilot]

@copilot resolve the merge conflicts in this pull request

Resolved the merge conflict in tests/test_models_inductor.py (took main's rel=1e-2 tolerance). Merge commit: 15511e3.