Fix issues with multi-circuit size in Engine

cirq-google/cirq_google/engine/abstract_job.py

@@ -141,9 +141,16 @@ def failure(self) -> tuple[str, str] | None:
         """Return failure code and message of the job if present."""
 
     @abc.abstractmethod
-    def get_repetitions_and_sweeps(self) -> tuple[int, list[cirq.Sweep]]:
+    def get_repetitions_and_sweeps(
+        self, circuit_num: int | None = None
+    ) -> tuple[int, list[cirq.Sweep]]:
         """Returns the repetitions and sweeps for the job.
 
+        Args:
+            circuit_num: if this is a batch job, the index of the circuit
+                to return the sweeps for.  This argument is zero-indexed.
+                Negative values index from the end of the list.
+
         Returns:
             A tuple of the repetition count and list of sweeps.
         """
@@ -159,11 +166,11 @@ def get_calibration(self) -> calibration.Calibration | None:
         one was captured, else None."""
 
     @abc.abstractmethod
-    def get_circuit(self, program_num: int | None = None) -> cirq.Circuit:
+    def get_circuit(self, circuit_num: int | None = None) -> cirq.Circuit:
         """Returns the cirq Circuit for the job.
 
         Args:
-            program_num: if this is a multi-circuit job, the index of the circuit
+            circuit_num: if this is a multi-circuit job, the index of the circuit
                 to return.  This argument is zero-indexed. Negative values
                 index from the end of the list.  Ignored if not multi-circuit.
 

cirq-google/cirq_google/engine/abstract_job_test.py

@@ -80,7 +80,7 @@ def get_processor(self):
     def get_calibration(self):
         pass
 
-    def get_circuit(self, program_num: int | None = None) -> cirq.Circuit:
+    def get_circuit(self, circuit_num: int | None = None) -> cirq.Circuit:
         return cirq.Circuit()
 
     def cancel(self) -> None:

cirq-google/cirq_google/engine/abstract_local_job.py

@@ -157,13 +157,48 @@ def processor_ids(self) -> list[str]:
         """Returns the processor ids provided when the job was created."""
         return [self._processor_id]
 
-    def get_repetitions_and_sweeps(self) -> tuple[int, list[cirq.Sweep]]:
+    def get_repetitions_and_sweeps(
+        self, circuit_num: int | None = None
+    ) -> tuple[int, list[cirq.Sweep]]:
         """Returns the repetitions and sweeps for the job.
 
+        Args:
+            circuit_num: if this is a batch job, the index of the circuit
+                to return the sweeps for.  This argument is zero-indexed.
+                Negative values index from the end of the list.
+
         Returns:
             A tuple of the repetition count and list of sweeps.
         """
-        return (self._repetitions, self._sweeps)
+        is_batch = self.program().is_batch()
+        batch_size = self.program().batch_size() if is_batch else 1
+
+        is_mapped = is_batch and len(self._sweeps) == batch_size and len(self._sweeps) > 1
+
+        if circuit_num is None:
+            if is_mapped:
+                raise ValueError(
+                    f"This is a batch job with {len(self._sweeps)} mapped sweeps. "
+                    "Please specify `circuit_num` to get sweeps for a specific circuit."
+                )
+            return (self._repetitions, self._sweeps)
+
+        if not is_batch:
+            if circuit_num != 0 and circuit_num != -1:
+                raise IndexError(f"Job is not a batch job, cannot index {circuit_num}")
+            return (self._repetitions, self._sweeps)
+
+        if not is_mapped:
+            # Shared sweeps in a batch job, return all of them
+            return (self._repetitions, self._sweeps)
+
+        # Mapped sweeps in a batch job
+        try:
+            return (self._repetitions, [self._sweeps[circuit_num]])
+        except IndexError:
+            raise IndexError(
+                f"Index {circuit_num} out of range for sweeps of size {len(self._sweeps)}."
+            )
 
     def get_processor(self) -> AbstractLocalProcessor:
         """Returns the AbstractProcessor for the processor the job is/was run on,
@@ -175,15 +210,15 @@ def get_calibration(self) -> calibration.Calibration | None:
         from the parent Engine object."""
         return self.get_processor().get_latest_calibration(int(self._create_time.timestamp()))
 
-    def get_circuit(self, program_num: int | None = None) -> cirq.Circuit:
+    def get_circuit(self, circuit_num: int | None = None) -> cirq.Circuit:
         """Returns the cirq Circuit for the job.
 
         Args:
-            program_num: if this is a multi-circuit job, the index of the circuit
+            circuit_num: if this is a multi-circuit job, the index of the circuit
                 to return.  This argument is zero-indexed. Negative values
                 index from the end of the list.  Ignored if not multi-circuit.
 
         Returns:
             The job's cirq Circuit.
         """
-        return self.program().get_circuit(program_num)
+        return self.program().get_circuit(circuit_num)

cirq-google/cirq_google/engine/abstract_local_job_test.py

@@ -21,6 +21,8 @@
 from typing import TYPE_CHECKING
 from unittest import mock
 
+import pytest
+
 import cirq
 from cirq_google.cloud import quantum
 from cirq_google.engine.abstract_local_job import AbstractLocalJob
@@ -75,14 +77,52 @@ def test_description_and_labels():
 
 
 def test_reps_and_sweeps():
+    # Single program (non-batch)
+    mock_program = mock.Mock()
+    mock_program.is_batch.return_value = False
     job = NothingJob(
         job_id='test',
         processor_id='grill',
-        parent_program=None,
+        parent_program=mock_program,
         repetitions=100,
         sweeps=[cirq.Linspace('t', 0, 10, 0.1)],
     )
     assert job.get_repetitions_and_sweeps() == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+    assert job.get_repetitions_and_sweeps(0) == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+    with pytest.raises(IndexError, match="Job is not a batch job"):
+        _ = job.get_repetitions_and_sweeps(1)
+
+    # Batch program, shared sweep
+    mock_program_batch = mock.Mock()
+    mock_program_batch.is_batch.return_value = True
+    mock_program_batch.batch_size.return_value = 2
+    job_batch_shared = NothingJob(
+        job_id='test',
+        processor_id='grill',
+        parent_program=mock_program_batch,
+        repetitions=100,
+        sweeps=[cirq.Linspace('t', 0, 10, 0.1)],
+    )
+    # Shared sweep, works with None or any index
+    assert job_batch_shared.get_repetitions_and_sweeps() == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+    assert job_batch_shared.get_repetitions_and_sweeps(0) == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+    assert job_batch_shared.get_repetitions_and_sweeps(1) == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+
+    # Batch program, mapped sweeps
+    job_batch_mapped = NothingJob(
+        job_id='test',
+        processor_id='grill',
+        parent_program=mock_program_batch,
+        repetitions=100,
+        sweeps=[cirq.Linspace('t', 0, 10, 0.1), cirq.Linspace('u', 0, 5, 0.5)],
+    )
+    # Mapped sweeps, requires index
+    with pytest.raises(ValueError, match="mapped sweeps"):
+        _ = job_batch_mapped.get_repetitions_and_sweeps()
+    assert job_batch_mapped.get_repetitions_and_sweeps(0) == (100, [cirq.Linspace('t', 0, 10, 0.1)])
+    assert job_batch_mapped.get_repetitions_and_sweeps(1) == (100, [cirq.Linspace('u', 0, 5, 0.5)])
+    with pytest.raises(IndexError, match="Index 2 out of range"):
+        _ = job_batch_mapped.get_repetitions_and_sweeps(2)
 
 
 def test_create_update_time():

cirq-google/cirq_google/engine/abstract_local_program.py

@@ -190,22 +190,47 @@ def remove_labels(self, keys: list[str]) -> AbstractProgram:
             del self._labels[key]
         return self
 
-    def get_circuit(self, program_num: int | None = None) -> cirq.Circuit:
+    def get_circuit(self, circuit_num: int | None = None) -> cirq.Circuit:
         """Returns the cirq Circuit for the program. This is only
         supported if the program was created with the V2 protos.
 
         Args:
-            program_num: if this is a multi-circuit program, the index of the circuit
+            circuit_num: if this is a multi-circuit program, the index of the circuit
                 to return.  This argument is zero-indexed. Negative values
                 indexing from the end of the list.
 
         Returns:
             The program's cirq Circuit.
         """
-        if program_num is not None:
-            return self._circuits[program_num]
-        return self._circuits[0]
+        if circuit_num is None:
+            if self.is_batch():
+                raise ValueError(
+                    f"This program is a batch program containing {len(self._circuits)} circuits. "
+                    "Please specify `circuit_num` to get a specific circuit, "
+                    "or use `get_circuits()` to get all of them."
+                )
+            return self._circuits[0]
+        try:
+            return self._circuits[circuit_num]
+        except IndexError:
+            raise IndexError(
+                f"Index {circuit_num} out of range for batch program of size {len(self._circuits)}."
+            )
+
+    def get_circuits(self) -> list[cirq.Circuit]:
+        """Returns all the cirq Circuits for the program."""
+        return self._circuits
+
+    def is_batch(self) -> bool:
+        """Returns True if the program is a batch program."""
+        return len(self._circuits) > 1
 
     def batch_size(self) -> int:
-        """Returns the number of programs in a batch program."""
+        """Returns the number of programs in a batch program.
+
+        Raises:
+            ValueError: if the program created was not a batch program.
+        """
+        if not self.is_batch():
+            raise ValueError("This program is not a batch program.")
         return len(self._circuits)

cirq-google/cirq_google/engine/abstract_local_program_test.py

@@ -149,12 +149,25 @@ def test_description_and_labels():
 def test_circuit():
     circuit1 = cirq.Circuit(cirq.X(cirq.LineQubit(1)))
     circuit2 = cirq.Circuit(cirq.Y(cirq.LineQubit(2)))
+
+    # Single circuit, non-batch
     program = NothingProgram([circuit1], None)
-    assert program.batch_size() == 1
+    assert not program.is_batch()
+    with pytest.raises(ValueError, match="not a batch program"):
+        _ = program.batch_size()
     assert program.get_circuit() == circuit1
     assert program.get_circuit(0) == circuit1
-    assert program.batch_size() == 1
+    assert program.get_circuits() == [circuit1]
+
+    # Multi circuit (always batch)
     program = NothingProgram([circuit1, circuit2], None)
+    assert program.is_batch()
     assert program.batch_size() == 2
+    with pytest.raises(ValueError, match="batch program containing 2 circuits"):
+        _ = program.get_circuit()
     assert program.get_circuit(0) == circuit1
     assert program.get_circuit(1) == circuit2
+    assert program.get_circuits() == [circuit1, circuit2]
+
+    with pytest.raises(IndexError):
+        _ = program.get_circuit(2)

cirq-google/cirq_google/engine/abstract_program.py

@@ -158,19 +158,31 @@ def remove_labels(self, keys: list[str]) -> AbstractProgram:
         """
 
     @abc.abstractmethod
-    def get_circuit(self, program_num: int | None = None) -> cirq.Circuit:
+    def get_circuit(self, circuit_num: int | None = None) -> cirq.Circuit:
         """Returns the cirq Circuit for the program. This is only
         supported if the program was created with the V2 protos.
 
         Args:
-            program_num: if this is a multi-circuit program, the index of the circuit
+            circuit_num: if this is a multi-circuit program, the index of the circuit
                 to return.  This argument is zero-indexed. Negative values
                 indexing from the end of the list.
 
         Returns:
             The program's cirq Circuit.
         """
 
+    @abc.abstractmethod
+    def get_circuits(self) -> list[cirq.Circuit]:
+        """Returns all the cirq Circuits for the program.
+
+        Returns:
+            A list of the program's cirq Circuits.
+        """
+
+    @abc.abstractmethod
+    def is_batch(self) -> bool:
+        """Returns True if the program is a batch program."""
+
     @abc.abstractmethod
     def batch_size(self) -> int:
         """Returns the number of programs in a batch program.

cirq-google/cirq_google/engine/engine_job.py

@@ -226,15 +226,52 @@ def failure(self) -> tuple[str, str] | None:
             return (failure.error_code.name, failure.error_message)
         return None
 
-    def get_repetitions_and_sweeps(self) -> tuple[int, list[cirq.Sweep]]:
+    def get_repetitions_and_sweeps(
+        self, circuit_num: int | None = None
+    ) -> tuple[int, list[cirq.Sweep]]:
         """Returns the repetitions and sweeps for the Quantum Engine job.
 
+        Args:
+            circuit_num: if this is a batch job, the index of the circuit
+                to return the sweeps for.  This argument is zero-indexed.
+                Negative values index from the end of the list.
+
         Returns:
             A tuple of the repetition count and list of sweeps.
         """
         if self._job is None or self._job.run_context is None:
             self._job = self._get_job(return_run_context=True)
-        return _deserialize_run_context(self._job.run_context)
+        reps, sweeps = _deserialize_run_context(self._job.run_context)
+
+        is_batch = self.program().is_batch()
+        batch_size = self.program().batch_size() if is_batch else 1
+
+        is_mapped = is_batch and len(sweeps) == batch_size and len(sweeps) > 1
+
+        if circuit_num is None:
+            if is_mapped:
+                raise ValueError(
+                    f"This is a batch job with {len(sweeps)} mapped sweeps. "
+                    "Please specify `circuit_num` to get sweeps for a specific circuit."
+                )
+            return (reps, sweeps)
+
+        if not is_batch:
+            if circuit_num != 0 and circuit_num != -1:
+                raise IndexError(f"Job is not a batch job, cannot index {circuit_num}")
+            return (reps, sweeps)
+
+        if not is_mapped:
+            # Shared sweeps in a batch job, return all of them
+            return (reps, sweeps)
+
+        # Mapped sweeps in a batch job
+        try:
+            return (reps, [sweeps[circuit_num]])
+        except IndexError:
+            raise IndexError(
+                f"Index {circuit_num} out of range for batch job sweeps of size {len(sweeps)}."
+            )
 
     def get_processor(self) -> engine_processor.EngineProcessor | None:
         """Returns the EngineProcessor for the processor the job is/was run on,
@@ -258,18 +295,18 @@ def get_calibration(self) -> calibration.Calibration | None:
         metrics = v2.metrics_pb2.MetricsSnapshot.FromString(response.data.value)
         return calibration.Calibration(metrics)
 
-    async def get_circuit_async(self, program_num: int | None = None) -> cirq.Circuit:
+    async def get_circuit_async(self, circuit_num: int | None = None) -> cirq.Circuit:
         """Returns the cirq Circuit for the Quantum Engine job.
 
         Args:
-            program_num: if this is a multi-circuit job, the index of the circuit
+            circuit_num: if this is a multi-circuit job, the index of the circuit
                 to return.  This argument is zero-indexed. Negative values
                 indexing from the end of the list.
 
         Returns:
             The job's cirq Circuit.
         """
-        return await self.program().get_circuit_async(program_num)
+        return await self.program().get_circuit_async(circuit_num)
 
     get_circuit = duet.sync(get_circuit_async)