Solve a Poisson equation

examples/pde_poisson_vqls.py

The problem

Solve a discretised PDE — a full variational application. The equation is the 1D Poisson boundary-value problem −u″(x) = f(x) on a rod with both ends held at zero: the steady-state heat profile under a uniform heat source. Discretised on a grid it becomes a linear system A·u = b, and that reduction is why PDEs fit the platform: a grid of 2ⁿ points lives in the amplitudes of n qubits.

The variational scheme minimises

C(θ) = 1 − ⟨ψ(θ)|M|ψ(θ)⟩ / ⟨ψ(θ)|A²|ψ(θ)⟩,   M = A·|b⟩⟨b|·A

which is zero exactly when |ψ⟩ ∝ A⁻¹·b — the solution.

Needs the qiskit extra: pip install 'qubitra-sdk[qiskit]'.

The walkthrough

The discretised system

# Two qubits carry a four-point grid. Small on purpose: the point is the method, every
# step of it is checkable against the classical solve at the end, and the whole run
# stays a few minutes of wall clock.
NUM_QUBITS = 2
GRID_POINTS = 2**NUM_QUBITS

# The discretised operator: the second difference with Dirichlet (held-at-zero) ends,
# the textbook tridiagonal matrix. The load f is uniform, so b is the uniform state.
A = 2 * np.eye(GRID_POINTS) - np.eye(GRID_POINTS, k=1) - np.eye(GRID_POINTS, k=-1)
b = np.full(GRID_POINTS, 1 / np.sqrt(GRID_POINTS))

Two qubits carry the four-point grid, small enough that every step is checkable against the classical solve at the end. A is the second-difference matrix and b the uniform load, normalised as a state.

The two cost operators

# Both cost operators, decomposed into Pauli terms numerically — no hand derivation.
# A² and A·|b⟩⟨b|·A are real symmetric, so every coefficient is real, which is all the
# platform's observable form measures.
A_SQUARED = SparsePauliOp.from_operator(A @ A).simplify()
M = SparsePauliOp.from_operator(A @ np.outer(b, b) @ A).simplify()

The cost is a ratio of two expectation values, so the scheme needs two operators: M (the numerator, built from A and the load) and A_SQUARED (the denominator). Both are Hermitian, so both are ordinary observables and every cost evaluation is an estimator call — the Hadamard tests a textbook VQLS uses are unnecessary here.

The ansatz

    # A real-amplitude ansatz: the solution of a real symmetric system is a real
    # vector, so there is nothing for complex phases to do.
    ansatz = real_amplitudes(NUM_QUBITS, reps=1)
    parameters = list(ansatz.parameters)
    theta = np.random.default_rng(7).normal(0, 0.3, len(parameters))

real_amplitudes prepares real-valued states, matching the real solution of a real symmetric system. The seeded initial parameters make the run reproducible.

One step’s evaluations, as one job

    def shift_rows(centre: np.ndarray) -> np.ndarray:
        """The centre point plus both parameter-shift rows per parameter — one step's
        every evaluation, batched into a single sweep."""
        rows = [centre]
        for index in range(len(centre)):
            for shift in (np.pi / 2, -np.pi / 2):
                row = centre.copy()
                row[index] += shift
                rows.append(row)
        return np.array(rows)

    def evaluate(rows: np.ndarray) -> np.ndarray:
        """⟨M⟩ and ⟨A²⟩ for every row — one estimator job for the whole batch. The
        observables are shaped (2, 1) so they broadcast against the row axis."""
        result = estimator.run([(ansatz, [[M], [A_SQUARED]], rows)]).result()
        return np.asarray(result[0].data.evs, dtype=float).reshape(2, -1)

shift_rows builds the centre point plus two shifted rows per parameter — every evaluation one optimisation step needs. evaluate reads both observables at every row in a single estimator job: the observables broadcast against the row axis, so a step is one platform round trip.

The optimiser

    m_avg = np.zeros_like(theta)
    v_avg = np.zeros_like(theta)
    cost = 1.0
    for step in range(1, STEPS + 1):
        values = evaluate(shift_rows(theta))
        m_c, a2_c = values[0, 0], values[1, 0]
        cost = 1.0 - m_c / a2_c
        d_m = (values[0, 1::2] - values[0, 2::2]) / 2
        d_a2 = (values[1, 1::2] - values[1, 2::2]) / 2
        gradient = -(d_m * a2_c - m_c * d_a2) / a2_c**2
        m_avg = 0.9 * m_avg + 0.1 * gradient
        v_avg = 0.999 * v_avg + 0.001 * gradient**2
        theta -= (
            LEARNING_RATE * (m_avg / (1 - 0.9**step)) / (np.sqrt(v_avg / (1 - 0.999**step)) + 1e-8)
        )
        if step % 10 == 0:
            print(f"step {step:3}: cost {cost:.4f}  (0 is a perfect solution)")
        if cost < 2e-3:
            print(f"step {step:3}: cost {cost:.4f} — converged")
            break

Adam over parameter-shift gradients. The centre row gives the cost; the shifted-row differences give each expectation’s own gradient, and the quotient rule combines them into the cost’s gradient — still nothing measured beyond the two observables.

Sample the solution

    # Read the solution the only way a quantum state allows: sample it. The counts
    # estimate |u_i|² — the solution profile up to sign and normalisation, which for
    # this positive solution is the profile itself, squared.
    measured = ansatz.assign_parameters(dict(zip(parameters, theta, strict=True)))
    measured.measure_all()
    counts = QubitraSamplerV2(backend).run([measured], shots=8192).result()[0].data.meas
    histogram = counts.get_counts()
    shots = sum(histogram.values())
    sampled = np.array(
        [histogram.get(format(i, f"0{NUM_QUBITS}b"), 0) / shots for i in range(GRID_POINTS)]
    )

The converged parameters are bound into the ansatz, measurements added, and the state sampled. Each basis state is one grid point, so the counts estimate the squared solution profile.

Compare with the classical solve

    # The classical answer, for the comparison that makes the run checkable at a glance.
    exact = np.linalg.solve(A, b)
    exact_profile = (exact / np.linalg.norm(exact)) ** 2

    print("\ngrid   classical  sampled")
    for i in range(GRID_POINTS):
        bar = "█" * round(40 * sampled[i])
        print(f"  u_{i}   {exact_profile[i]:.3f}      {sampled[i]:.3f}  {bar}")

    overlap = float(np.sqrt(exact_profile) @ np.sqrt(sampled))
    print(f"\noverlap with the classical solution: {overlap:.3f}  (1.000 is exact)")
    return 0 if overlap > 0.98 else 1

np.linalg.solve gives the exact answer on this small grid, and the overlap between the two profiles condenses the comparison to one number. The example exits non-zero below 0.98, which is what makes it a smoke test.

Running it

QUBITRA_API_KEY=qpk_... python examples/pde_poisson_vqls.py

Output from a live run, trimmed:

step  10: cost 0.6420  (0 is a perfect solution)
step  20: cost 0.1099  (0 is a perfect solution)
step  38: cost 0.0019 — converged

grid   classical  sampled
  u_0   0.154      0.144  ██████
  u_1   0.346      0.360  ██████████████
  u_2   0.346      0.349  ██████████████
  u_3   0.154      0.147  ██████

overlap with the classical solution: 1.000  (1.000 is exact)

The cost falls to convergence in 38 steps. The sampled profile — low at the held ends, high in the middle — matches the classical solve point for point within shot noise, and the overlap rounds to 1.000. The method carries to larger grids unchanged.

Where to go next

• Primitives and PUBs — the estimator jobs every step submits.
• Qiskit — the adapter the example builds on.