sampler_test.py

# Copyright 2021 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 __future__ import annotations

from unittest import mock

import pandas as pd
import sympy as sp

import cirq
import cirq_ionq as ionq


def test_sampler_qpu():
    mock_service = mock.MagicMock()
    job_dict = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'qpu',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }

    job = ionq.Job(client=mock_service, job_dict=job_dict)
    mock_service.create_job.return_value = job
    mock_service.get_results.return_value = {'0': '0.25', '1': '0.75'}

    sampler = ionq.Sampler(service=mock_service, target='qpu')
    q0 = cirq.LineQubit(0)
    circuit = cirq.Circuit(cirq.X(q0), cirq.measure(q0, key='a'))
    results = sampler.sample(program=circuit, repetitions=4)
    pd.testing.assert_frame_equal(
        results, pd.DataFrame(columns=['a'], index=[0, 1, 2, 3], data=[[0], [1], [1], [1]])
    )
    mock_service.create_job.assert_called_once_with(circuit=circuit, repetitions=4, target='qpu')


def test_sampler_simulator():
    mock_service = mock.MagicMock()
    job_dict = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'simulator',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }

    job = ionq.Job(client=mock_service, job_dict=job_dict)
    mock_service.create_job.return_value = job
    mock_service.get_results.return_value = {'0': '0.25', '1': '0.75'}

    sampler = ionq.Sampler(service=mock_service, target='simulator', seed=10)
    q0 = cirq.LineQubit(0)
    circuit = cirq.Circuit(cirq.X(q0), cirq.measure(q0, key='a'))
    results = sampler.sample(program=circuit, repetitions=4)
    pd.testing.assert_frame_equal(
        results, pd.DataFrame(columns=['a'], index=[0, 1, 2, 3], data=[[1], [0], [1], [1]])
    )
    mock_service.create_job.assert_called_once_with(
        circuit=circuit, repetitions=4, target='simulator'
    )


def test_sampler_multiple_jobs():
    mock_service = mock.MagicMock()
    job_dict0 = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'qpu',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }
    job_dict1 = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'qpu',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }

    job0 = ionq.Job(client=mock_service, job_dict=job_dict0)
    job1 = ionq.Job(client=mock_service, job_dict=job_dict1)
    mock_service.create_job.side_effect = [job0, job1]
    mock_service.get_results.side_effect = [{'0': '0.25', '1': '0.75'}, {'0': '0.5', '1': '0.5'}]

    sampler = ionq.Sampler(service=mock_service, timeout_seconds=10, target='qpu')
    q0 = cirq.LineQubit(0)
    x = sp.Symbol('x')
    circuit = cirq.Circuit(cirq.X(q0) ** x, cirq.measure(q0, key='a'))
    results = sampler.sample(
        program=circuit,
        repetitions=4,
        params=[cirq.ParamResolver({x: 0.5}), cirq.ParamResolver({x: 0.6})],
    )
    pd.testing.assert_frame_equal(
        results,
        pd.DataFrame(
            columns=['x', 'a'],
            index=[0, 1, 2, 3] * 2,
            data=[[0.5, 0], [0.5, 1], [0.5, 1], [0.5, 1], [0.6, 0], [0.6, 0], [0.6, 1], [0.6, 1]],
        ),
    )
    circuit0 = cirq.Circuit(cirq.X(q0) ** 0.5, cirq.measure(q0, key='a'))
    circuit1 = cirq.Circuit(cirq.X(q0) ** 0.6, cirq.measure(q0, key='a'))
    mock_service.create_job.assert_has_calls(
        [
            mock.call(circuit=circuit0, repetitions=4, target='qpu'),
            mock.call(circuit=circuit1, repetitions=4, target='qpu'),
        ]
    )
    assert mock_service.create_job.call_count == 2


def test_sampler_run_sweep():
    mock_service = mock.MagicMock()
    job_dict = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'qpu',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }

    job = ionq.Job(client=mock_service, job_dict=job_dict)
    mock_service.create_job.return_value = job
    mock_service.get_results.return_value = {'0': '0.25', '1': '0.75'}

    sampler = ionq.Sampler(service=mock_service, target='qpu')
    q0 = cirq.LineQubit(0)
    circuit = cirq.Circuit(cirq.X(q0) ** sp.Symbol('theta'), cirq.measure(q0, key='a'))

    sweep = cirq.Linspace(key='theta', start=0.0, stop=2.0, length=5)

    results = sampler.run_sweep(program=circuit, params=sweep, repetitions=4)
    assert len(results) == 5  # Confirm that 5 result objects were created
    assert all([result.repetitions == 4 for result in results])  # Confirm all repetitions were 4

    # Assert that create_job was called 5 times
    assert mock_service.create_job.call_count == 5
    # For each call, assert that it was called with the correct repetitions and target
    for call in mock_service.create_job.call_args_list:
        _, kwargs = call
        assert kwargs["repetitions"] == 4
        assert kwargs["target"] == 'qpu'


def test_sampler_run_sweep_batched_job_results():
    mock_service = mock.MagicMock()
    job_dict = {
        'id': '1',
        'status': 'completed',
        'stats': {'qubits': '1'},
        'backend': 'qpu',
        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
    }

    job = ionq.Job(client=mock_service, job_dict=job_dict)
    mock_service.create_job.return_value = job

    # Create dummy results
    result1 = ionq.QPUResult(counts={0: 4}, num_qubits=1, measurement_dict={'a': [0]})
    result2 = ionq.QPUResult(counts={1: 4}, num_qubits=1, measurement_dict={'a': [0]})

    # Mock job.results to return a list
    with mock.patch.object(ionq.Job, 'results', return_value=[result1, result2]):
        sampler = ionq.Sampler(service=mock_service, target='qpu')
        q0 = cirq.LineQubit(0)
        circuit = cirq.Circuit(cirq.X(q0), cirq.measure(q0, key='a'))

        # We pass 1 resolver, but job returns 2 results (list).
        # This triggers `flattened_job_results.extend(res)`.
        results = sampler.run_sweep(program=circuit, params=[cirq.ParamResolver({})], repetitions=4)

        # Even though we got 2 results from the job, we only have 1 resolver,
        # so zip() truncates to 1 result.
        assert len(results) == 1
        # result1 counts={0: 4}, so we expect four 0s.
        # measurements is (repetitions, qubits), so [[0], [0], [0], [0]]
        assert results[0].measurements['a'].tolist() == [[0], [0], [0], [0]]