Add Concat to Sweeps (#6819)

* add Concat to sweeps

* json repr

* json testing

* add test for empty sweep

* comments

* coverage for keys property

Author: senecameeks

cirq-core/cirq/__init__.py

@@ -494,6 +494,7 @@
 )
 
 from cirq.study import (
+    Concat as Concat,
     dict_to_product_sweep as dict_to_product_sweep,
     dict_to_zip_sweep as dict_to_zip_sweep,
     ExpressionMap as ExpressionMap,

cirq-core/cirq/json_resolver_cache.py

@@ -120,6 +120,7 @@ def _symmetricalqidpair(qids):
         'CliffordState': cirq.CliffordState,
         'CliffordTableau': cirq.CliffordTableau,
         'CNotPowGate': cirq.CNotPowGate,
+        'Concat': cirq.Concat,
         'ConstantQubitNoiseModel': cirq.ConstantQubitNoiseModel,
         'ControlledGate': cirq.ControlledGate,
         'ControlledOperation': cirq.ControlledOperation,

cirq-core/cirq/protocols/json_test_data/Concat.json

@@ -0,0 +1,19 @@
+{
+    "cirq_type": "Concat",
+    "sweeps": [
+      {
+        "cirq_type": "Linspace",
+        "key": "a",
+        "start": 0,
+        "stop": 1,
+        "length": 2
+      },
+      {
+        "cirq_type": "Linspace",
+        "key": "a",
+        "start": 0,
+        "stop": 2,
+        "length": 4
+      }
+    ]
+  }

cirq-core/cirq/protocols/json_test_data/Concat.repr

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

cirq-core/cirq/study/__init__.py

@@ -36,6 +36,7 @@
 )
 
 from cirq.study.sweeps import (
+    Concat as Concat,
     Linspace as Linspace,
     ListSweep as ListSweep,
     Points as Points,

cirq-core/cirq/study/sweeps.py

@@ -276,6 +276,63 @@ def _from_json_dict_(cls, factors, **kwargs):
         return Product(*factors)
 
 
+class Concat(Sweep):
+    """Concatenates multiple to a new sweep.
+
+    All sweeps must share the same descriptors.
+
+    If one sweep assigns 'a' to the values 0, 1, 2, and another sweep assigns
+    'a' to the values 3, 4, 5, the concatenation produces a sweep assigning
+    'a' to the values 0, 1, 2, 3, 4, 5 in sequence.
+    """
+
+    def __init__(self, *sweeps: Sweep) -> None:
+        if not sweeps:
+            raise ValueError("Concat requires at least one sweep.")
+
+        # Validate consistency across sweeps
+        first_sweep = sweeps[0]
+        for sweep in sweeps[1:]:
+            if sweep.keys != first_sweep.keys:
+                raise ValueError("All sweeps must have the same descriptors.")
+
+        self.sweeps = sweeps
+
+    def __eq__(self, other):
+        if not isinstance(other, Concat):
+            return NotImplemented
+        return self.sweeps == other.sweeps
+
+    def __hash__(self):
+        return hash(tuple(self.sweeps))
+
+    @property
+    def keys(self) -> List['cirq.TParamKey']:
+        return self.sweeps[0].keys
+
+    def __len__(self) -> int:
+        return sum(len(sweep) for sweep in self.sweeps)
+
+    def param_tuples(self) -> Iterator[Params]:
+        for sweep in self.sweeps:
+            yield from sweep.param_tuples()
+
+    def __repr__(self) -> str:
+        sweeps_repr = ', '.join(repr(sweep) for sweep in self.sweeps)
+        return f'cirq.Concat({sweeps_repr})'
+
+    def __str__(self) -> str:
+        sweeps_repr = ', '.join(repr(s) for s in self.sweeps)
+        return f'Concat({sweeps_repr})'
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        return protocols.obj_to_dict_helper(self, ['sweeps'])
+
+    @classmethod
+    def _from_json_dict_(cls, sweeps, **kwargs):
+        return Concat(*sweeps)
+
+
 class Zip(Sweep):
     """Zip product (direct sum) of one or more sweeps.
 

cirq-core/cirq/study/sweeps_test.py

@@ -246,6 +246,8 @@ def test_equality():
     et.make_equality_group(lambda: cirq.Linspace('b', 0, 10, 11))
     et.make_equality_group(lambda: cirq.Points('a', list(range(11))))
     et.make_equality_group(lambda: cirq.Points('b', list(range(11))))
+    et.make_equality_group(lambda: cirq.Concat(cirq.Linspace('a', 0, 10, 11)))
+    et.make_equality_group(lambda: cirq.Concat(cirq.Linspace('b', 0, 10, 11)))
 
     # Product and Zip sweeps can also be equated.
     et.make_equality_group(lambda: cirq.Linspace('a', 0, 5, 6) * cirq.Linspace('b', 10, 15, 6))
@@ -373,3 +375,105 @@ def test_dict_to_zip_sweep():
     assert cirq.dict_to_zip_sweep({'t': [0, 1], 's': [2, 3], 'r': 4}) == (
         cirq.Zip(cirq.Points('t', [0, 1]), cirq.Points('s', [2, 3]), cirq.Points('r', [4]))
     )
+
+
+def test_concat_linspace():
+    sweep1 = cirq.Linspace('a', 0.34, 9.16, 4)
+    sweep2 = cirq.Linspace('a', 10, 20, 4)
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    assert len(concat_sweep) == 8
+    assert concat_sweep.keys == ['a']
+    params = list(concat_sweep.param_tuples())
+    assert len(params) == 8
+    assert params[0] == (('a', 0.34),)
+    assert params[3] == (('a', 9.16),)
+    assert params[4] == (('a', 10.0),)
+    assert params[7] == (('a', 20.0),)
+
+
+def test_concat_points():
+    sweep1 = cirq.Points('a', [1, 2])
+    sweep2 = cirq.Points('a', [3, 4, 5])
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    assert concat_sweep.keys == ['a']
+    assert len(concat_sweep) == 5
+    params = list(concat_sweep)
+    assert len(params) == 5
+    assert _values(concat_sweep, 'a') == [1, 2, 3, 4, 5]
+
+
+def test_concat_many_points():
+    sweep1 = cirq.Points('a', [1, 2])
+    sweep2 = cirq.Points('a', [3, 4, 5])
+    sweep3 = cirq.Points('a', [6, 7, 8])
+    concat_sweep = cirq.Concat(sweep1, sweep2, sweep3)
+
+    assert len(concat_sweep) == 8
+    params = list(concat_sweep)
+    assert len(params) == 8
+    assert _values(concat_sweep, 'a') == [1, 2, 3, 4, 5, 6, 7, 8]
+
+
+def test_concat_mixed():
+    sweep1 = cirq.Linspace('a', 0, 1, 3)
+    sweep2 = cirq.Points('a', [2, 3])
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    assert len(concat_sweep) == 5
+    assert _values(concat_sweep, 'a') == [0.0, 0.5, 1.0, 2, 3]
+
+
+def test_concat_inconsistent_keys():
+    sweep1 = cirq.Linspace('a', 0, 1, 3)
+    sweep2 = cirq.Points('b', [2, 3])
+
+    with pytest.raises(ValueError, match="All sweeps must have the same descriptors"):
+        cirq.Concat(sweep1, sweep2)
+
+
+def test_concat_sympy_symbol():
+    a = sympy.Symbol('a')
+    sweep1 = cirq.Linspace(a, 0, 1, 3)
+    sweep2 = cirq.Points(a, [2, 3])
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    assert len(concat_sweep) == 5
+    assert _values(concat_sweep, 'a') == [0.0, 0.5, 1.0, 2, 3]
+
+
+def test_concat_repr_and_str():
+    sweep1 = cirq.Linspace('a', 0, 1, 3)
+    sweep2 = cirq.Points('a', [2, 3])
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    expected_repr = (
+        "cirq.Concat(cirq.Linspace('a', start=0, stop=1, length=3), cirq.Points('a', [2, 3]))"
+    )
+    expected_str = "Concat(cirq.Linspace('a', start=0, stop=1, length=3), cirq.Points('a', [2, 3]))"
+
+    assert repr(concat_sweep) == expected_repr
+    assert str(concat_sweep) == expected_str
+
+
+def test_concat_large_sweep():
+    sweep1 = cirq.Points('a', list(range(101)))
+    sweep2 = cirq.Points('a', list(range(101, 202)))
+    concat_sweep = cirq.Concat(sweep1, sweep2)
+
+    assert len(concat_sweep) == 202
+    assert _values(concat_sweep, 'a') == list(range(101)) + list(range(101, 202))
+
+
+def test_concat_different_keys_raises():
+    sweep1 = cirq.Linspace('a', 0, 1, 3)
+    sweep2 = cirq.Points('b', [2, 3])
+
+    with pytest.raises(ValueError, match="All sweeps must have the same descriptors."):
+        _ = cirq.Concat(sweep1, sweep2)
+
+
+def test_concat_empty_sweep_raises():
+    with pytest.raises(ValueError, match="Concat requires at least one sweep."):
+        _ = cirq.Concat()