Repository: quantumlib/Cirq

Issue:
Add support for indexing circuits and moments by qubits
It would be convenient if we could slice circuits and moments by qubit instead o fjust by index. For example:

```python
# returns an empty moment if no operation on qubit, otherwise a moment containing the single operation
moment[qubit]  

# returns Circuit([moment[qubit] for moment in circuit])
circuit[qubit] 

circuit[ [qubit1, qubit2] ]

circuit[0:10, qubit1)

# range support
line_qubit1:line_qubit2 has obvious interpretation
grid_qubit_1:grid_qubit_2 could be corners of rectangle?
```

In general the goal is for the indexing to behave like existing python slicing, and e.g. numpy slicing (probably we can't make mutable views). Any implementation should go over the relevant cases.


Discussion:
I'm worried about relying on type to distinguish moment index from qubit index.  Because of numpy, I would be surprised if `object[indices0, indices1] == object[indices1, indices0]` even if `indices1` is a list of qubits.

If this is implemented, I think it's more important to maintain the invariant `obj[x, y] == obj[x][y]` no matter the types of `x` and `y`.

The API I would expect is this:
```python
# Returns a new circuit containing the first 10 moments
circuit[0:10]

# Returns a circuit (with the same number of moments) containing only operations that touch qubit.
# (Or maybe return just a list of operations to better agree with moment[qubit].  Pad with identity gates?)
circuit[:, qubit]

# Returns a circuit (with the same number of moments) containing only operations that touch any of [qubit1, qubit2].
circuit[:, [qubit1, qubit2]]

assert circuit[0:10, qubits] == circuit[0:10][:, qubits] == circuit[:, qubits][0:10] == circuit[:, qubits][0:10, :]
assert circuit[5, qubit] == circuit[5][qubit]  # A single operation (or KeyError)
assert circuit[5, [qubit1, qubit2]] == circuit[5][[qubit1, qubit2]]  # A moment

circuit[qubit]  # TypeError or IndexError

# Returns the **operation** that touches qubit or KeyError.
moment[qubit]

# Returns a moment containing only operations touching any of [qubit1, qubit2]
moment[[qubit1, qubit2]]

# Ranges
line_qubit1:line_qubit2:step_size
grid_qubit_1:grid_qubit_2:int_or_pair_of_int?
```
Hi. I would like to implement this. Can someone please assign this to me?
@Strilanc @cduck 
@cduck I agree with you. But I don't understand the part about ranges. What does the integer or pair of integers in the middle of the slice mean?
For LineQubit it's the optional step size.  For GridQubit it's the optional x and y step sizes.  I don't know how useful this is but it's part of the slice API.
Which is more intuitive for you?  Momemt[qubit] is type Moment or type Operation?
Moment[qubit] should be an Operation.

Ah so you meant to write `line_qubit1:line_qubit2:integer` instead of having `integer` in the middle.
Yes.  My mistake.

Resolving patch:
diff --git a/cirq/circuits/circuit.py b/cirq/circuits/circuit.py
--- a/cirq/circuits/circuit.py
+++ b/cirq/circuits/circuit.py
@@ -18,7 +18,6 @@
 Operations. Each Operation is a Gate that acts on some Qubits, for a given
 Moment the Operations must all act on distinct Qubits.
 """
-
 from collections import defaultdict
 from fractions import Fraction
 from itertools import groupby
@@ -85,6 +84,13 @@ class Circuit:
     and sliced,
         circuit[1:3] is a new Circuit made up of two moments, the first being
             circuit[1] and the second being circuit[2];
+        circuit[:, qubit] is a new Circuit with the same moments, but with only
+            those operations which act on the given Qubit;
+        circuit[:, qubits], where 'qubits' is list of Qubits, is a new Circuit
+            with the same moments, but only with those operations which touch
+            any of the given qubits;
+        circuit[1:3, qubit] is equivalent to circuit[1:3][:, qubit];
+        circuit[1:3, qubits] is equivalent to circuit[1:3][:, qubits];
     and concatenated,
         circuit1 + circuit2 is a new Circuit made up of the moments in circuit1
             followed by the moments in circuit2;
@@ -211,15 +217,50 @@ def __getitem__(self, key: slice) -> 'Circuit':
     def __getitem__(self, key: int) -> 'cirq.Moment':
         pass
 
+    @overload
+    def __getitem__(self, key: Tuple[int, 'cirq.Qid']) -> 'cirq.Operation':
+        pass
+
+    @overload
+    def __getitem__(self,
+                    key: Tuple[int, Iterable['cirq.Qid']]) -> 'cirq.Moment':
+        pass
+
+    @overload
+    def __getitem__(self, key: Tuple[slice, 'cirq.Qid']) -> 'cirq.Circuit':
+        pass
+
+    @overload
+    def __getitem__(self,
+                    key: Tuple[slice, Iterable['cirq.Qid']]) -> 'cirq.Circuit':
+        pass
+
     def __getitem__(self, key):
         if isinstance(key, slice):
             sliced_circuit = Circuit(device=self.device)
             sliced_circuit._moments = self._moments[key]
             return sliced_circuit
-        if isinstance(key, int):
+        if hasattr(key, '__index__'):
             return self._moments[key]
-
-        raise TypeError('__getitem__ called with key not of type slice or int.')
+        if isinstance(key, tuple):
+            if len(key) != 2:
+                raise ValueError('If key is tuple, it must be a pair.')
+            moment_idx, qubit_idx = key
+            # moment_idx - int or slice; qubit_idx - Qid or Iterable[Qid].
+            selected_moments = self._moments[moment_idx]
+            # selected_moments - Moment or list[Moment].
+            if isinstance(selected_moments, list):
+                if isinstance(qubit_idx, cirq.Qid):
+                    qubit_idx = [qubit_idx]
+                new_circuit = Circuit(device=self.device)
+                new_circuit._moments = [
+                    moment[qubit_idx] for moment in selected_moments
+                ]
+                return new_circuit
+            return selected_moments[qubit_idx]
+
+        raise TypeError(
+            '__getitem__ called with key not of type slice, int or tuple.')
 
     @overload
     def __setitem__(self, key: int, value: 'cirq.Moment'):
diff --git a/cirq/ops/moment.py b/cirq/ops/moment.py
--- a/cirq/ops/moment.py
+++ b/cirq/ops/moment.py
@@ -15,7 +15,7 @@
 """A simplified time-slice of operations within a sequenced circuit."""
 
 from typing import (Any, Callable, Iterable, Sequence, TypeVar, Union, Tuple,
-                    FrozenSet, TYPE_CHECKING, Iterator)
+                    FrozenSet, TYPE_CHECKING, Iterator, overload)
 from cirq import protocols
 from cirq.ops import raw_types
 
@@ -33,6 +33,13 @@ class Moment:
     moment should execute at the same time (to the extent possible; not all
     operations have the same duration) and it is expected that all operations
     in a moment should be completed before beginning the next moment.
+
+    Moment can be indexed by qubit or list of qubits:
+        moment[qubit] returns the Operation in the moment which touches the
+            given qubit, or throws KeyError if there is no such operation.
+        moment[qubits] returns another Moment which consists only of those
+            operations which touch at least one of the given qubits. If there
+            are no such operations, returns an empty Moment.
     """
 
     def __init__(self, operations: Iterable[raw_types.Operation] = ()) -> None:
@@ -117,6 +124,18 @@ def without_operations_touching(self, qubits: Iterable[raw_types.Qid]):
             operation for operation in self.operations
             if qubits.isdisjoint(frozenset(operation.qubits)))
 
+    def _operation_touching(self, qubit: raw_types.Qid) -> 'cirq.Operation':
+        """Returns the operation touching given qubit.
+        Args:
+            qubit: Operations that touch this qubit will be returned.
+        Returns:
+            The operation which touches `qubit`.
+        """
+        for op in self.operations:
+            if qubit in op.qubits:
+                return op
+        raise KeyError("Moment doesn't act on given qubit")
+
     def __copy__(self):
         return type(self)(self.operations)
 
@@ -198,6 +217,25 @@ def __add__(self, other):
             return self.with_operation(other)
         return NotImplemented
 
+    # pylint: disable=function-redefined
+    @overload
+    def __getitem__(self, key: raw_types.Qid) -> 'cirq.Operation':
+        pass
+
+    @overload
+    def __getitem__(self, key: Iterable[raw_types.Qid]) -> 'cirq.Moment':
+        pass
+
+    def __getitem__(self, key):
+        if isinstance(key, raw_types.Qid):
+            return self._operation_touching(key)
+        elif isinstance(key, Iterable):
+            qubits_to_keep = frozenset(key)
+            ops_to_keep = tuple(
+                op for op in self.operations
+                if not qubits_to_keep.isdisjoint(frozenset(op.qubits)))
+            return Moment(ops_to_keep)
+
 
 def _list_repr_with_indented_item_lines(items: Sequence[Any]) -> str:
     block = '\n'.join([repr(op) + ',' for op in items])


Test patch:
diff --git a/cirq/circuits/circuit_test.py b/cirq/circuits/circuit_test.py
--- a/cirq/circuits/circuit_test.py
+++ b/cirq/circuits/circuit_test.py
@@ -3615,3 +3615,112 @@ def test_transform_qubits():
     assert c.transform_qubits(lambda q: q).device is cg.Foxtail
     assert c.transform_qubits(lambda q: q, new_device=cg.Bristlecone
                              ).device is cg.Bristlecone
+
+
+def test_indexing_by_pair():
+    # 0: ───H───@───X───@───
+    #           │       │
+    # 1: ───────H───@───@───
+    #               │   │
+    # 2: ───────────H───X───
+    q = cirq.LineQubit.range(3)
+    c = cirq.Circuit([
+        cirq.H(q[0]),
+        cirq.H(q[1]).controlled_by(q[0]),
+        cirq.H(q[2]).controlled_by(q[1]),
+        cirq.X(q[0]),
+        cirq.CCNOT(*q),
+    ])
+
+    # Indexing by single moment and qubit.
+    assert c[0, q[0]] == c[0][q[0]] == cirq.H(q[0])
+    assert c[1, q[0]] == c[1, q[1]] == cirq.H(q[1]).controlled_by(q[0])
+    assert c[2, q[0]] == c[2][q[0]] == cirq.X(q[0])
+    assert c[2, q[1]] == c[2, q[2]] == cirq.H(q[2]).controlled_by(q[1])
+    assert c[3, q[0]] == c[3, q[1]] == c[3, q[2]] == cirq.CCNOT(*q)
+
+    # Indexing by moment and qubit - throws if there is no operation.
+    with pytest.raises(KeyError, match="Moment doesn't act on given qubit"):
+        _ = c[0, q[1]]
+
+    # Indexing by single moment and multiple qubits.
+    assert c[0, q] == c[0]
+    assert c[1, q] == c[1]
+    assert c[2, q] == c[2]
+    assert c[3, q] == c[3]
+    assert c[0, q[0:2]] == c[0]
+    assert c[0, q[1:3]] == cirq.Moment([])
+    assert c[1, q[1:2]] == c[1]
+    assert c[2, [q[0]]] == cirq.Moment([cirq.X(q[0])])
+    assert c[2, q[1:3]] == cirq.Moment([cirq.H(q[2]).controlled_by(q[1])])
+    assert c[np.int64(2), q[0:2]] == c[2]
+
+    # Indexing by single qubit.
+    assert c[:, q[0]] == cirq.Circuit([
+        cirq.Moment([cirq.H(q[0])]),
+        cirq.Moment([cirq.H(q[1]).controlled_by(q[0])]),
+        cirq.Moment([cirq.X(q[0])]),
+        cirq.Moment([cirq.CCNOT(q[0], q[1], q[2])]),
+    ])
+    assert c[:, q[1]] == cirq.Circuit([
+        cirq.Moment([]),
+        cirq.Moment([cirq.H(q[1]).controlled_by(q[0])]),
+        cirq.Moment([cirq.H(q[2]).controlled_by(q[1])]),
+        cirq.Moment([cirq.CCNOT(q[0], q[1], q[2])]),
+    ])
+    assert c[:, q[2]] == cirq.Circuit([
+        cirq.Moment([]),
+        cirq.Moment([]),
+        cirq.Moment([cirq.H(q[2]).controlled_by(q[1])]),
+        cirq.Moment([cirq.CCNOT(q[0], q[1], q[2])]),
+    ])
+
+    # Indexing by several qubits.
+    assert c[:, q] == c[:, q[0:2]] == c[:, [q[0], q[2]]] == c
+    assert c[:, q[1:3]] == cirq.Circuit([
+        cirq.Moment([]),
+        cirq.Moment([cirq.H(q[1]).controlled_by(q[0])]),
+        cirq.Moment([cirq.H(q[2]).controlled_by(q[1])]),
+        cirq.Moment([cirq.CCNOT(q[0], q[1], q[2])]),
+    ])
+
+    # Indexing by several moments and one qubit.
+    assert c[1:3, q[0]] == cirq.Circuit([
+        cirq.H(q[1]).controlled_by(q[0]),
+        cirq.X(q[0]),
+    ])
+    assert c[1::2, q[2]] == cirq.Circuit([
+        cirq.Moment([]),
+        cirq.Moment([cirq.CCNOT(*q)]),
+    ])
+
+    # Indexing by several moments and several qubits.
+    assert c[0:2, q[1:3]] == cirq.Circuit([
+        cirq.Moment([]),
+        cirq.Moment([cirq.H(q[1]).controlled_by(q[0])]),
+    ])
+    assert c[::2, q[0:2]] == cirq.Circuit([
+        cirq.Moment([cirq.H(q[0])]),
+        cirq.Moment([cirq.H(q[2]).controlled_by(q[1]),
+                     cirq.X(q[0])]),
+    ])
+
+    # Equivalent ways of indexing.
+    assert c[0:2, q[1:3]] == c[0:2][:, q[1:3]] == c[:, q[1:3]][0:2]
+
+    # Passing more than 2 items is forbidden.
+    with pytest.raises(ValueError, match='If key is tuple, it must be a pair.'):
+        _ = c[0, q[1], 0]
+
+    # Can't swap indices.
+    with pytest.raises(TypeError,
+                       match='list indices must be integers or slices'):
+        _ = c[q[1], 0]
+
+
+def test_indexing_by_numpy_integer():
+    q = cirq.NamedQubit('q')
+    c = cirq.Circuit(cirq.X(q), cirq.Y(q))
+
+    assert c[np.int32(1)] == cirq.Moment([cirq.Y(q)])
+    assert c[np.int64(1)] == cirq.Moment([cirq.Y(q)])
diff --git a/cirq/ops/moment_test.py b/cirq/ops/moment_test.py
--- a/cirq/ops/moment_test.py
+++ b/cirq/ops/moment_test.py
@@ -274,3 +274,35 @@ def test_add():
     circuit2 = cirq.Circuit(cirq.CNOT(a, b), cirq.Y(b))
     circuit2[1] += cirq.X(a)
     assert circuit2 == expected_circuit
+
+
+def test_indexes_by_qubit():
+    a, b, c = cirq.LineQubit.range(3)
+    moment = cirq.Moment([cirq.H(a), cirq.CNOT(b, c)])
+
+    assert moment[a] == cirq.H(a)
+    assert moment[b] == cirq.CNOT(b, c)
+    assert moment[c] == cirq.CNOT(b, c)
+
+
+def test_throws_when_indexed_by_unused_qubit():
+    a, b = cirq.LineQubit.range(2)
+    moment = cirq.Moment([cirq.H(a)])
+
+    with pytest.raises(KeyError, match="Moment doesn't act on given qubit"):
+        _ = moment[b]
+
+
+def test_indexes_by_list_of_qubits():
+    q = cirq.LineQubit.range(4)
+    moment = cirq.Moment([cirq.Z(q[0]), cirq.CNOT(q[1], q[2])])
+
+    assert moment[[q[0]]] == Moment([cirq.Z(q[0])])
+    assert moment[[q[1]]] == Moment([cirq.CNOT(q[1], q[2])])
+    assert moment[[q[2]]] == Moment([cirq.CNOT(q[1], q[2])])
+    assert moment[[q[3]]] == Moment([])
+    assert moment[q[0:2]] == moment
+    assert moment[q[1:3]] == Moment([cirq.CNOT(q[1], q[2])])
+    assert moment[q[2:4]] == Moment([cirq.CNOT(q[1], q[2])])
+    assert moment[[q[0], q[3]]] == Moment([cirq.Z(q[0])])
+    assert moment[q] == moment