Repository: quantumlib/Cirq

Issue:
Add __getitem__[start:stop] slicing support to parameter sweeps
If we have a sweepable that is really big it is useful to break it up into smaller sweeps.  We should add functionality which does this slicing.


Discussion:

Resolving patch:
diff --git a/cirq/study/sweeps.py b/cirq/study/sweeps.py
--- a/cirq/study/sweeps.py
+++ b/cirq/study/sweeps.py
@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import cast, Iterable, Iterator, List, Sequence, Tuple, Union
+from typing import (cast, Dict, Iterable, Iterator, List, overload, Sequence,
+                    Tuple, Union)
 
 import abc
 import collections
@@ -101,6 +102,40 @@ def __iter__(self) -> Iterator[resolver.ParamResolver]:
         for params in self.param_tuples():
             yield resolver.ParamResolver(collections.OrderedDict(params))
 
+    # pylint: disable=function-redefined
+    @overload
+    def __getitem__(self, val: int) -> resolver.ParamResolver:
+        pass
+
+    @overload
+    def __getitem__(self, val: slice) -> 'Sweep':
+        pass
+
+    def __getitem__(self, val):
+        n = len(self)
+        if isinstance(val, int):
+            if val < -n or val >= n:
+                raise IndexError(f'sweep index out of range: {val}')
+            if val < 0:
+                val += n
+            return next(itertools.islice(self, val, val + 1))
+        if not isinstance(val, slice):
+            raise TypeError(
+                'Sweep indices must be either int or slices, not {}'.format(
+                    type(val)))
+
+        inds_map: Dict[int, int] = {
+            sweep_i: slice_i for slice_i, sweep_i in enumerate(range(n)[val])
+        }
+        results = [resolver.ParamResolver()] * len(inds_map)
+        for i, item in enumerate(self):
+            if i in inds_map:
+                results[inds_map[i]] = item
+
+        return ListSweep(results)
+
+    # pylint: enable=function-redefined
+
     @abc.abstractmethod
     def param_tuples(self) -> Iterator[Params]:
         """An iterator over (key, value) pairs assigning Symbol key to value."""


Test patch:
diff --git a/cirq/study/sweeps_test.py b/cirq/study/sweeps_test.py
--- a/cirq/study/sweeps_test.py
+++ b/cirq/study/sweeps_test.py
@@ -74,6 +74,56 @@ def test_product():
     assert _values(sweep, 'b') == [4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7]
 
 
+def test_slice_access_error():
+    sweep = cirq.Points('a', [1, 2, 3])
+    with pytest.raises(TypeError, match='<class \'str\'>'):
+        _ = sweep['junk']
+
+    with pytest.raises(IndexError):
+        _ = sweep[4]
+
+    with pytest.raises(IndexError):
+        _ = sweep[-4]
+
+
+def test_slice_sweep():
+    sweep = cirq.Points('a', [1, 2, 3]) * cirq.Points('b', [4, 5, 6, 7])
+
+    first_two = sweep[:2]
+    assert list(first_two.param_tuples())[0] == (('a', 1), ('b', 4))
+    assert list(first_two.param_tuples())[1] == (('a', 1), ('b', 5))
+    assert len(list(first_two)) == 2
+
+    middle_three = sweep[5:8]
+    assert list(middle_three.param_tuples())[0] == (('a', 2), ('b', 5))
+    assert list(middle_three.param_tuples())[1] == (('a', 2), ('b', 6))
+    assert list(middle_three.param_tuples())[2] == (('a', 2), ('b', 7))
+    assert len(list(middle_three.param_tuples())) == 3
+
+    odd_elems = sweep[6:1:-2]
+    assert list(odd_elems.param_tuples())[2] == (('a', 1), ('b', 6))
+    assert list(odd_elems.param_tuples())[1] == (('a', 2), ('b', 4))
+    assert list(odd_elems.param_tuples())[0] == (('a', 2), ('b', 6))
+    assert len(list(odd_elems.param_tuples())) == 3
+
+    sweep_reversed = sweep[::-1]
+    assert list(sweep) == list(reversed(list(sweep_reversed)))
+
+    single_sweep = sweep[5:6]
+    assert list(single_sweep.param_tuples())[0] == (('a', 2), ('b', 5))
+    assert len(list(single_sweep.param_tuples())) == 1
+
+
+def test_access_sweep():
+    sweep = cirq.Points('a', [1, 2, 3]) * cirq.Points('b', [4, 5, 6, 7])
+
+    first_elem = sweep[-12]
+    assert first_elem == cirq.ParamResolver({'a': 1, 'b': 4})
+
+    sixth_elem = sweep[5]
+    assert sixth_elem == cirq.ParamResolver({'a': 2, 'b': 5})
+
+
 @pytest.mark.parametrize('r_list', [
     [{
         'a': a,