Consolidate quantum state representations in cirq

[viathor]

Is your design idea/issue related to a use case or problem? Please describe.

There is a zoo of quantum states and special representations they admit:

• computational basis states (which can be specified by an int),
• pure states (which can be described by a 2^n-dimensional array or a (2, 2, ..., 2)-shaped tensor),
• mixed states (which can be represented as a 2^n-by-2^n positive matrix with unit trace),
• stabilizer states (which can be described by generators of their stabilizer group),
• graph states (stabilizer states whose stabilizer generators can be described by a graph),
• MPS states (which can be described by a simple tensor network),
• PEPS states (which can be described by a little more complex tensor network),
• ...

In various places in cirq we have the need to represent states in some of the sets above (e.g. mixed states in cirq.DensityMatrixSimulator and MPS states in cirq.contrib.quimb.mps_simulator).

A while back we discussed the question whether we should have a class to encapsulate a generic quantum state to hide the zoo behind a consistent interface. We decided not to do that and instead allow quantum states to be specified by simple primitive and numpy types (int, np.ndarray). This made sense at the time since most code manipulating quantum states was particular to the specific use-case and given a small subset of the zoo that we needed to support it was easy to infer the nature of the state based on the type and shape information. We also defined a union type cirq.STATE_VECTOR_LIKE to describe the subset of the zoo we found relevant to us. This consisted of only four cases.

By now, we have added a class cirq.ProductState to cirq.STATE_VECTOR_LIKE and also added another union type cirq.QUANTUM_STATE_LIKE which encompasses cirq.STATE_VECTOR_LIKE and some additional options including another class for representing quantum states called cirq.QuantumState. We also have a mixin called cirq.StateVectorMixin...

In summary, we had considered a choice between a lightweight class-free approach of a simple union whose members were either primitive or maintained elsewhere (numpy) and a more heavy-weight approach with a single class to encapsulate the zoo. We ended up with the worst of both worlds :-) We now have two union types each with a class encapsulating a subset of the states. This means that on one hand code handling quantum states needs to deal with the diversity of the zoo and we lack a common place to extend quantum states with new functionality and on the other hand we have three classes to maintain.

Describe your design idea/issue

Perhaps, it is time to re-evaluate. One approach merges cirq.ProductState, cirq.QuantumState and cirq.StateVectorMixin and adds cirq.state(...) which can take as arguments any of the present members of the two union types. For example, we could say things like this

cirq.state(0)                                   # |0>
cirq.state(np.diag([1, np.exp(1j*np.pi / 4)]))  # magic T state
cirq.state(np.eye(2) / 2, n_qubits=1)           # maximally mixed state of one qubit
cirq.state(np.eye(2) / np.sqrt(2), n_qubits=2)  # Bell state
cirq.state(stabilizers={cirq.Z})                # |0>

The factory cirq.state would continue to allow us to use simple ways to specify quantum states such as specifying computational basis states using an int while also providing us with a single place to add general state-related functionality. In particular, the new class would be a natural place for code that

• converts between state representations (e.g. producing density matrices of state vectors or stabilizer tableau from computational basis index etc),
• interprets a numpy array (together with optional hints) as a state (which can be done in a few different ways depending on Hilbert space dimension and qubit/qudit count),
• checks whether a state fulfills some basic condition (e.g. is pure, entangled, separable),
• generates human-readable notation (e.g. Dirac notation).

Filing to start discussion. Context: #3517.

[viathor]

One question this raises is to what extent we wish to use mypy and to what extent runtime checks to see whether a given state belongs to one of the sets of quantum states listed earlier (e.g. is pure, is separable, etc) .

I suspect that fully committing to the use of type system would lead to very complicated type relationships. Partially committing to it would lead to a situation where some of the checks would be done by mypy and others at runtime. On the other hand, having all such checks done at runtime is a uniform and easily extensible approach.

[viathor]

(The classes we have for product and general quantum states demonstrate the temptation (encouraged by OOP-inspired programming languages) to build taxonomies. The reason taxonomies lead to clunky code is that the most appropriate classification of our objects is generally a local concern at the place of use. Using a type system (via user-defined classes) to force a particular taxonomy globally on the codebase leads to the use of a classification that is a poor fit in at least some of it.

A nice alternative to a globally imposed taxonomy is provided by pattern matching wherein each place of use can choose a property to look at for classification at that location. For example, code computing fidelity could classify states by whether they are pure or mixed (to improve performance in the former case) and code checking separability could look at the number of qubits (to exploit the sufficiency of the PPT criterion for two qubits). The good news is that pattern matching is coming to python in version 3.10, see PEP 636. Perhaps we should keep that in mind when considering how to represent quantum states in cirq.)

[tanujkhattar]

• This sounds hard but the current situation is getting bad and would get worse.
• Need a design doc and we should put it in the roadmap.
• Adam: Any code that consume a quantum state needs to inspect the union. Since the union types are expanding, older code doesn't have a way to consume the objects safely.

Tl;Dr - We should think about this and consider adding to the roadmap / have a propose a design.

[viathor]

(Re-added kind/design-issue label since a "kind" label is required and this does seem like the most appropriate, at least until we agree on enough detail about how to proceed to change it to kind/roadmap-item.)

[daxfohl]

I'd like to do this, but whatever we do I'd like to make it compatible for use in the simulator framework. What "level" are you most interested in here?

1. We could have states that represent high-level state and include things like qubit ids (this is somewhat already covered by ActOnArgs -- we just have to rename that to e.g. QuantumState and perhaps move it to cirq/qis)
2. We could have states that represent low-level state and work with axes instead of qubits, and a member apply_gate(gate, axes) (I'd not want apply_operation on these because some operations are too high-level, and the state probably won't know what to do with them).
3. Similar state representation as (2) but no apply_gate. Instead StateVector would have apply_unitary, DensityMatix would have apply_channel, etc.
4. We could do all of the above, where (1) contains a (2) which contains a (3). (I like this option but then you get into the problem of "naming things", not to mention to get the state vector array from the top level you have to do x.state.state.state.state.state_vector or write wrapper methods at each level)

Whichever we choose, I agree with @tanujkhattar that a union will be hard to maintain, and a type hierarchy would be preferred. To integrate these into simulation, the base "interface" would need to support (params depending on which layer we're at):

• apply_whichever(op/gate/matrix) # no return value
• measure(qubits/axes) -> List[int]
• sample(qubits/axes, reps) -> List[List[int]]
• deep_copy() -> TSelf
• kron(TSelf other) -> TSelf
• factor(qubits/axes) -> Tuple[TSelf, TSelf]
• transpose(qubits/axes) -> TSelf

Which I think is a good baseline for any quantum state representation to start from. The last three are a required part of the "interface", though it's fine if they throw NotImplemented; the simulator framework can work around that.

[daxfohl]

FWIW my preference is (2) above. I think (1) is great but probably higher level than what's useful from a QIS standpoint. (3) is nice in that it avoids bringing in the notion of a "gate", which is a cirq.ops thing, into qis. But the downside of (3) is that the "whatever" in the "apply_whatever" function is different for every subtype, so that wouldn't be part of the base type, and thus just leaves a bunch of combinators in the base interface, severely limiting the usefulness of the abstraction.

So I like (2) in that it's low-level, but useful as an abstraction. I think this could be extremely useful in combination with #4632, where we could have these classes be generic on the types of gates they support in apply_gate(TGateType), where TGateType is one of UnitaryGate, StabilizerGate, etc. But that's not a prereq for doing this.

For stabilizers we would need to rework those protocols a bit first. Right now the only way to apply a stabilizer effect is via act_on, which is high-level and requires qubits rather than axes. I wouldn't expect this to be too much of a challenge, and would make the stabilizer implementation more robust than what we currently have, but needs to be called out.

The other thing that we'd have to do in order for this to be used in simulators, is to make sure all non-gate operations implement _acton. Right now the only missing one I see (outside of artificial tests) is PauliStringPhasor.

Edit: ...And actually the thing that occurred to me later is that this change would also break any third-party users that are implementing operations without gates, which maybe is nobody, or maybe is a huge breaking change. (I believe we could come up with a deprecation path, however it could be a painful one if users have dozens of gateless operation classes). That said, I still think this is the correct design for a low-level abstraction around quantum state regardless of whether we can use it in simulators. We just wouldn't be able to use it in simulators if it's considered too big of a change, which is a bummer, but shouldn't block the feature in its own right.

[daxfohl]

I tried this and (2) ends up being not particularly clean. The reason is that state vector simulator records which mixture option was chosen. This means apply_gate would have to take the measurement log and a prng. Also axes might not be enough if we eventually do qudit subdimensions, so apply_gate would need a subdimensions param as well.

At that point we're passing in so many args from ActOnArgs._act_on_fallback_ to QuantumState.apply_gate that it no longer makes sense to do this in a separate class.

[daxfohl]

Option (3) that I mentioned two comments above is done as of #5065. That PR created an interface qis.QuantumStateRepresentation: https://github.com/quantumlib/Cirq/blob/fe7fe4e1d29e6b5fcc1643c0c73c26f8b37adae1/cirq-core/cirq/qis/clifford_tableau.py#L28

Implementations are expected to be bare bones. The interface operates on indices, not qubits. A prng must be explicitly passed in where needed. Implementations don't contain their own prng. There's also no notion of classical data here. It's just the raw data required to represent the state.

This has implementations

1. sim._BufferedStateVector
2. sim._BufferedDensityMatrix
3. contrib._MpsHandler
4. qis.CliffordTableau
5. sim.CliffordChForm

It has the following interface:

1. Combinator methods: copy, kron, factor, transpose.
2. Measurement methods: measure, sample.
3. Informational properties: supports_factor, can_represent_mixed_states.

Notably it doesn't have any abstract apply method. To add one was more or less impossible due to the rationale I spelled out above: apply has lots of complexities (measurements, subcircuits, classical controls, channel measurements, decompositions, etc), so being able to apply generically requires use of the act_on protocol, and can't be done in a single method here. Here, each subclass provides its own set of apply_X methods that are used by the simulators, but there is not a way to abstract one out into the base interface; that's what act_on is for.

Anyway, perhaps all that's left for this issue is

1. Move all the implementations to qis and make them public.
2. Make the proposed cirq.state(...) function that analyzes the input and returns the corresponding state object.

[daxfohl]

After thinking about this for a bit, I believe the root problem is not going to be solved by a new quantum state representation (though the representation I added worked great for simulators), I think the core issue is that there are always going to be multiple incompatible ways to represent quantum states. We'll have some that are useful in simulators, some that are useful in low-level code, some that are useful in analysis, some that come from third parties, etc., and they may have no relationship to each other. If we want to support them all, then the STATE_VECTOR_LIKE union and friends will continue to grow as we add more supported types, and we'll have to revisit each function that uses these parameters every time they change, which will be ugly.

What I'd propose here is that we get rid of STATE_VECTOR_LIKE etc unions entirely, and instead create corresponding protocols, i.e. protocols.state_vector(...), protocols.density_matrix(...), etc. Those should mostly be lifted from the existing to_valid_state_vector etc methods, but then add a final obj._state_vector_() check if none of the previous things succeed.

This allows us to get rid of all the mandatory isinstance checks for each possible type in STATE_VECTOR_LIKE and just rely on the protocol. For example, we could change def fidelity(state1: 'cirq.QUANTUM_STATE_LIKE', state2: 'cirq.QUANTUM_STATE_LIKE', ...) to be def fidelity(state1: Any, state2: Any, ...), have it explicitly handle the couple cases where it can run faster (like if they're both ints), but use the state_vector protocol to retrieve the state vector and calculate fidelities in the general case.

This also allows us to instrument a wider range of things as state-vector-like, just by adding the _state_vector_, _density_matrix_, etc handlers to them. For instance our SimulationTrialResult contains a state vector somewhere under the hood, so it could be useful to allow users to calculate the fidelity of two of those, without having to dig for the function that gets the state vector.