Experimental noise model and simulation module by bramathon · Pull Request #1848 · rigetti/pyquil

bramathon · 2026-05-19T10:25:59Z

Description

This PR adds new noise modeling and simulation utilities. It is intended to replace the existing reference simulators and the existing NoiseModel. The proposed migration path is to add these utilities as experimental private modules, while marking the existing simulation and noise utilities as deprecated. In the next major version of pyquil, the deprecated utilities will be replaced by the new utilities.

Dependencies

This PR adds the new dependency rigetti-quax. This package provides a wide set of utilities for representing quantum objects including states, gates, measurements and noise channels. It's based on jax, which delivers high performance but adds jax as a transitive dependency to pyquil. jax is a relatively mature package but somewhat larger than existing dependencies (Around 80MB compared to scipy at 40MB). So this is worth some discussion.

Noise

Below, I discuss the new noise modeling classes.

First, the noise.py file is promoted to a module. This is non-breaking as the same functions continue to be exported from pyquil.noise. Inside the module are several new private files. When these are promoted to the public api, they will be exported from pyquil.noise.

New Noise Model System (`pyquil/noise/_noise_model.py`)

We introduce a frozen-dataclass-based NoiseModel container that collects per-instruction noise channels. The role of the noise model is to store a collection of channels, which together make up a device noise model. Channels are associated with instructions in the program and we can get the channel for a particular instruction via NoiseModel.get_channel(inst).

NoiseModelLike — Protocol defining the get_channel interface for custom implementations.
NoiseModel — The canonical use case. Accepts an Iterable of channels (list, tuple, set, generator), stores them as an immutable Tuple. Supports + for composition and get_channel(inst) for lookup.
DepolarizingNoiseModel — Convenience model returning a depolarizing channel for any gate.
CompositeNoiseModel — Chains multiple NoiseModelLike objects, returning the first non-None channel.

We can also construct noise models from the instruction set architecture, giving users a straightforward path to a device-realistic noise model.

Quax-Backed Noise Channels (`pyquil/noise/_channels.py`)

Four frozen dataclasses representing physical noise processes, each backed by a quax operator (SuperOp, KrausMap, or QuantumInstrument): The main role of the channel is to associate a superoperator with a particular instruction. For example, the instruction CZ 0 1 indicates that qubits 0 and 1 undergo the unitary CZ operation (via the quil spec). The channel associates that instruction to a superoperator, which is a higher dimensional representation of the operation which includes the effects of noise.

Class	Purpose
`Channel`	Noisy gate instruction
`MeasurementChannel`	Noisy measurement (`QuantumInstrument`)
`ResetChannel`	Noisy reset operation
`CycleChannel`	Noise for an entire instruction cycle (multi-gate group)

Simulation

A noise model is not very interesting on it's own, we want a simulator which can simulate it's effects on programs. The existing pyquil simulators are limited in various ways. The numpy reference simulators have poor performance, while the QVM has limited options for representing noise. Here we attempt to solve both problems by replacing the reference numpy simulators with a jax-accelerated, highly flexible simulation framework.

Three New Simulators (`pyquil/simulation/_simulator.py`)

A simulator object is constructed from a program. We have a simulation hierarchy based on the capabilities and scale.

Simulator	Use Case	JAX-Compatible	Scale	Noise	Measurements
`PureStateVectorSimulator`	Gate-only programs (no noise/measurements/resets)	`jit` + `grad`	<26 qubits	No	No
`DensityMatrixSimulator`	Any program, optionally with noise	`jit` + `grad`	<13 qubits	Yes	No
`TrajectorySimulator`	Monte Carlo trajectory simulation	`jit` (per-batch)	<26 qubits	Yes	Yes

Linearizer / DAG / Resolver / Compressor Pipeline (`pyquil/simulation/_resolver.py`)

The reason that the simulator is an object rather than a simple function is because efficient simulation requires the construction of several related closures. We call these the linearizer, the resolver, the compressor an the calculator.

Linearizer — Converts MemoryMap → flat JAX parameter vector.
Resolver — Converts parameter vector → (operator, subsystem) pairs, consulting the noise model.
Compressor — Greedy edge contraction on the program DAG, merging adjacent operators up to max_subsystem_size qubits.
Calculator — Compute the final state from the resolved and compressed operators.

The functions are all tightly coupled, each closure is constructed based on the program structure. They can also all be compiled with jax.jit. This is important to achieving good performance. The role of the class is therefore to combine all these objects into a logical construct.

Usage

from pyquil.gates import CNOT, H

from pyquil.quil import Program
from pyquil.simulation._simulator import PureStateVectorSimulator

program = Program(H(0), CNOT(0, 1))
sim = PureStateVectorSimulator(program)
psi = sim.compute()

windsurf-bot

1 file skipped due to size limits:

poetry.lock

Looks good to me 🤙

💡 To request another review, post a new comment with "/windsurf-review".

mhodson-rigetti

I have not yet reviewed simulation/*, transform.py, or the tests (besides one thing I saw while cross-checking the code). Partial review submitted!

mhodson-rigetti · 2026-05-20T05:59:08Z

 [tool.poetry.dependencies]
 # TODO(#1816): Loosen this bound once we've resolved support for Python 3.13+.
-python = "^3.9,<3.13"
+python = ">=3.11, <3.13"


Nominally this is a breaking change but the versioning strategy here has been that retiring unsupported Python versions is allowed without a major version increment. Python 3.8 support was removed in 4.11 (June 2024). Python 3.9 is already EOL, so it should be dropped now. However, Python 3.10 is still in maintenance until October 2026. Can you do some due diligence to see if 3.10 can be accepted as a minimum version here?

The reason to upgrade should also be captured in a separate issue, even if its implementation is linked to this PR.

The python bump is because quax is >=3.11

quax is >= 3.11 because it requires jax >= 0.8.2

Jax 0.8.2 is >=3.11 because it follows SPEC-0

That seems reasonable to me.

mhodson-rigetti · 2026-05-20T06:03:47Z

+CustomGateMap = Dict[str, Union[qx.Unitary, Callable[..., qx.Unitary]]]
+
+
+def _parse_quil_instruction(quil_str: str) -> Gate | Measurement | Reset:


I was wondering what code paths would require string parsing for individual textual Quil instructions. I can't find any call sites for this method. Since it is private, it is also not meant to be used outside. Remove?

Left from a previous iteration, will remove

Correction it's used in the serialization of the noise model.

mhodson-rigetti · 2026-05-20T06:19:24Z

+            if len(pauli) != num_qubits:
+                raise ValueError(f"Pauli term '{pauli}' has length {len(pauli)}, expected {num_qubits}.")
+
+        all_pauli_terms = list(map(lambda term: "".join(term), itertools.product(*["IXYZ" for _ in range(num_qubits)])))


Why force a list here? I think it can be left as a generator and use O(1) storage given the single usage site.

mhodson-rigetti · 2026-05-20T07:17:45Z

+            return False
+        return bool(jnp.isclose(float(qx.process_fidelity(self.process, other.process)), 1.0, atol=1e-9))
+
+    def __hash__(self) -> int:


Well, spent at least an hour here.

This is inconsistent with __eq__. The hash and the equality check should agree on the semantics.

But I know what it really happening. You want the frozen dataclass. But, Jax arrays are not hashable. If you didn't override the default behavior, any attempt to use the channel as a key would fail with an unhashable type exception. It seems like it should, since Jax arrays are immutable, but they are not hashable, which is a design decision they made related to the JIT compiler. And equality checks on channels (including their process) does make general sense.

Still, the non-standard treatment is buried in this method and not documented to the user. If the user came along with some Python set() of CZ(0, 1) channels which they had prepared to sweep over some wide range of noise processes, the Python set would have hash collisions on every single insert and lookup, degrading the container performance from O(log N) insert and O(N.log N) full traversal to O(N.log N) and O(N^2.log N) respectively. One day they might notice and get quite grumpy.

The codified way to document that you are not hashing the process is to use field(hash=False), and probably say why in the property documentation. However, they still need to read the docs to not fall foul of the aforementioned scenario.

I would personally hash the Jax array. This means you do need to override this dunder method, but you should take inspiration maybe from the below to get some (one-time) binary blob as a key.

https://github.com/netket/netket/blob/a36129bbde5ea2152ebb7bbefc9aa824612e27e8/netket/utils/array.py#L22

You should also include the target process (unitary) which should be resolved prior. Mimic the expected behavior of a dataclass but overcome the non-hashability issue minimally. I don't think it will impact you as you are not keying on this anyway.

I decided to remove hash, it's not necessary.

mhodson-rigetti · 2026-05-20T07:43:22Z

+        """A depolarizing channel on a qutrit gate mixes the state."""
+        inst = Gate("TX", [], [0])
+        channel = Channel.from_gate_fidelity(inst=inst, fidelity=0.8)
+        noise_model = NoiseModel(channels=frozenset([channel]))
