Add ZipLongest to cirq_google

cirq-google/cirq_google/__init__.py

@@ -107,6 +107,8 @@
     Serializer,
 )
 
+from cirq_google.study import ZipLongest
+
 from cirq_google.workflow import (
     ExecutableSpec,
     KeyValueExecutableSpec,

cirq-google/cirq_google/json_resolver_cache.py

@@ -84,4 +84,5 @@ def _old_xmon(*args, **kwargs):
         'cirq.google.EngineResult': cirq_google.EngineResult,
         'cirq.google.GridDevice': cirq_google.GridDevice,
         'cirq.google.GoogleCZTargetGateset': cirq_google.GoogleCZTargetGateset,
+        'cirq.google.ZipLongest': cirq_google.ZipLongest,
     }

cirq-google/cirq_google/json_test_data/cirq.google.ZipLongest.json

@@ -0,0 +1,19 @@
+{
+  "cirq_type": "cirq.google.ZipLongest",
+  "sweeps": [
+    {
+      "cirq_type": "Linspace",
+      "key": "a",
+      "start": 0,
+      "stop": 1,
+      "length": 2
+    },
+    {
+      "cirq_type": "Linspace",
+      "key": "b",
+      "start": 0,
+      "stop": 2,
+      "length": 4
+    }
+  ]
+}

cirq-google/cirq_google/json_test_data/cirq.google.ZipLongest.repr

@@ -0,0 +1 @@
+cirq_google.ZipLongest(cirq.Linspace('a', start=0, stop=1, length=2), cirq.Linspace('b', start=0, stop=2, length=4))

cirq-google/cirq_google/json_test_data/spec.py

@@ -63,6 +63,7 @@
             'EngineResult',
             'GridDevice',
             'GoogleCZTargetGateset',
+            'ZipLongest',
         ]
     },
     resolver_cache=_class_resolver_dictionary(),

cirq-google/cirq_google/study/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2023 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 cirq_google.study.zip_longest import ZipLongest

cirq-google/cirq_google/study/zip_longest.py

@@ -0,0 +1,76 @@
+# Copyright 2023 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 Any, Dict, Iterator, List
+
+import itertools
+import cirq
+
+
+class ZipLongest(cirq.Zip):
+    """Iterate over constituent sweeps in parallel
+
+    Analogous to itertools.zip_longest.
+    Note that we iterate until all sweeps terminate,
+    so if the sweeps are different lengths, the
+    shorter sweeps will be filled by repeating their last value
+    until all sweeps have equal length.
+    This is different from itertools.zip_longest, which uses a fixed fill value.
+    """
+
+    def __init__(self, *sweeps: cirq.Sweep) -> None:
+        self.sweeps = sweeps
+
+    def __eq__(self, other):
+        if not isinstance(other, ZipLongest):
+            return NotImplemented
+        return self.sweeps == other.sweeps
+
+    def __hash__(self) -> int:
+        return hash(tuple(self.sweeps))
+
+    @property
+    def keys(self) -> List['cirq.TParamKey']:
+        return sum((sweep.keys for sweep in self.sweeps), [])
+
+    def __len__(self) -> int:
+        if not self.sweeps:
+            return 0
+        return max(len(sweep) for sweep in self.sweeps)
+
+    def __repr__(self) -> str:
+        sweeps_repr = ', '.join(repr(s) for s in self.sweeps)
+        return f'cirq_google.ZipLongest({sweeps_repr})'
+
+    def __str__(self) -> str:
+        sweeps_repr = ', '.join(repr(s) for s in self.sweeps)
+        return f'ZipLongest({sweeps_repr})'
+
+    def param_tuples(self) -> Iterator[cirq.study.sweeps.Params]:
+        iters = [
+            itertools.chain(sweep.param_tuples(), itertools.repeat(list(sweep.param_tuples())[-1]))
+            for sweep in self.sweeps
+        ]
+        for vals in itertools.islice(zip(*iters), len(self)):
+            yield sum(vals, ())
+
+    @classmethod
+    def _json_namespace_(cls) -> str:
+        return 'cirq.google'
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        return cirq.obj_to_dict_helper(self, ['sweeps'])
+
+    @classmethod
+    def _from_json_dict_(cls, sweeps, **kwargs):
+        return ZipLongest(*sweeps)

cirq-google/cirq_google/study/zip_longest_test.py

@@ -0,0 +1,52 @@
+# Copyright 2023 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import cirq
+import cirq_google as cg
+
+
+def test_zip_longest():
+    sweep = cg.ZipLongest(cirq.Points('a', [1, 2, 3]), cirq.Points('b', [4, 5, 6, 7]))
+    assert len(sweep) == 4
+    assert tuple(sweep.param_tuples()) == (
+        (('a', 1), ('b', 4)),
+        (('a', 2), ('b', 5)),
+        (('a', 3), ('b', 6)),
+        (('a', 3), ('b', 7)),
+    )
+    assert sweep.keys == ['a', 'b']
+    assert (
+        str(sweep) == 'ZipLongest(cirq.Points(\'a\', [1, 2, 3]), cirq.Points(\'b\', [4, 5, 6, 7]))'
+    )
+    assert (
+        repr(sweep)
+        == 'cirq_google.ZipLongest(cirq.Points(\'a\', [1, 2, 3]), cirq.Points(\'b\', [4, 5, 6, 7]))'
+    )
+
+
+def test_empty_zip():
+    assert len(cg.ZipLongest()) == 0
+
+
+def test_zip_eq():
+    sweep1 = cg.ZipLongest(cirq.Points('a', [1, 2, 3]), cirq.Points('b', [4, 5, 6, 7]))
+    sweep2 = cg.ZipLongest(cirq.Points('a', [1, 2, 3]), cirq.Points('b', [4, 5, 6, 7]))
+    sweep3 = cg.ZipLongest(cirq.Points('a', [1, 2]), cirq.Points('b', [4, 5, 6, 7]))
+    sweep4 = cirq.Zip(cirq.Points('a', [1, 2]), cirq.Points('b', [4, 5, 6, 7]))
+
+    assert sweep1 == sweep2
+    assert hash(sweep1) == hash(sweep2)
+    assert sweep2 != sweep3
+    assert hash(sweep2) != hash(sweep3)
+    assert sweep1 != sweep4
+    assert hash(sweep1) != hash(sweep4)