Add AnalogTrajectory and FreqMap

cirq-google/cirq_google/experimental/analog_experiments/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2019 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.
+
+"""Folder for Running Analog experiments."""
+
+from cirq_google.experimental.analog_experiments.analog_trajectory_util import (
+    FreqMap as FreqMap,
+    AnalogTrajectory as AnalogTrajectory,
+)

cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py

@@ -0,0 +1,197 @@
+# Copyright 2019 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 typing import AbstractSet, TYPE_CHECKING
+
+import attrs
+import matplotlib.pyplot as plt
+import numpy as np
+import tunits as tu
+
+import cirq
+from cirq_google.experimental.analog_experiments import symbol_util as su
+
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+
+
+@attrs.mutable
+class FreqMap:
+    """Object containing information about the step to a new analog Hamiltonian.
+
+    Attributes:
+        duration: duration of step
+        qubit_freqs: dict describing qubit frequencies at end of step (None if idle)
+        couplings: dict describing coupling rates at end of step
+    """
+
+    duration: su.ValueOrSymbol
+    qubit_freqs: dict[str, su.ValueOrSymbol | None]
+    couplings: dict[tuple[str, str], su.ValueOrSymbol]
+
+    def _is_parameterized_(self) -> bool:
+        return (
+            cirq.is_parameterized(self.duration)
+            or su.is_parameterized_dict(self.qubit_freqs)
+            or su.is_parameterized_dict(self.couplings)
+        )
+
+    def _parameter_names_(self) -> AbstractSet[str]:
+        return (
+            cirq.parameter_names(self.duration)
+            | su.dict_param_name(self.qubit_freqs)
+            | su.dict_param_name(self.couplings)
+        )
+
+    def _resolve_parameters_(
+        self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
+    ) -> FreqMap:
+        resolver_ = cirq.ParamResolver(resolver)
+        return FreqMap(
+            duration=su.direct_symbol_replacement(self.duration, resolver_),
+            qubit_freqs={
+                k: su.direct_symbol_replacement(v, resolver_) for k, v in self.qubit_freqs.items()
+            },
+            couplings={
+                k: su.direct_symbol_replacement(v, resolver_) for k, v in self.couplings.items()
+            },
+        )
+
+
+class AnalogTrajectory:
+    """Class for handling qubit frequency and coupling trajectories that
+    define analog experiments. The class is defined using a sparse_trajectory,
+    which contains time durations of each Hamiltonian ramp element and the
+    corresponding qubit frequencies and couplings (unassigned qubits and/or
+    couplers are left unchanged).
+    """
+
+    def __init__(
+        self, *, full_trajectory: list[FreqMap], qubits: list[str], pairs: list[tuple[str, str]]
+    ):
+        self.full_trajectory = full_trajectory
+        self.qubits = qubits
+        self.pairs = pairs
+
+    @classmethod
+    def from_sparse_trajectory(
+        cls,
+        sparse_trajectory: list[
+            tuple[
+                tu.Value,
+                dict[str, su.ValueOrSymbol | None],
+                dict[tuple[str, str], su.ValueOrSymbol],
+            ],
+        ],
+        qubits: list[str] | None = None,
+        pairs: list[tuple[str, str]] | None = None,
+    ):
+        """Construct AnalogTrajectory from sparse trajectory.
+
+        Args:
+            sparse_trajectory: A list of tuples, where each tuple defines a segment of `FreqMap`
+                and contains three elements: (duration, qubit_freqs, coupling_strengths).
+                `duration` is a tunits value, `qubit_freqs` is a dictionary mapping qubit strings
+                to detuning frequencies, and `coupling_strengths` is a dictionary mapping qubit
+                pairs to their coupling strength. This format is considered "sparse" because each
+                tuple does not need to fully specify all qubits and coupling pairs; any missing
+                detuning frequency or coupling strength will be set to the same value as the
+                previous value in the list.
+            qubits: The qubits in interest. If not provided, automatically parsed from trajectory.
+            pairs: The pairs in interest. If not provided, automatically parsed from trajectory.
+        """
+        if qubits is None or pairs is None:
+            qubits_in_traj: list[str] = []
+            pairs_in_traj: list[tuple[str, str]] = []
+            for _, q, p in sparse_trajectory:
+                qubits_in_traj.extend(q.keys())
+                pairs_in_traj.extend(p.keys())
+            qubits = list(set(qubits_in_traj))
+            pairs = list(set(pairs_in_traj))
+
+        full_trajectory: list[FreqMap] = []
+        init_qubit_freq_dict: dict[str, tu.Value | None] = {q: None for q in qubits}
+        init_g_dict: dict[tuple[str, str], tu.Value] = {p: 0 * tu.MHz for p in pairs}
+        full_trajectory.append(FreqMap(0 * tu.ns, init_qubit_freq_dict, init_g_dict))
+
+        for dt, qubit_freq_dict, g_dict in sparse_trajectory:
+            # If no freq provided, set equal to previous
+            new_qubit_freq_dict = {
+                q: qubit_freq_dict.get(q, full_trajectory[-1].qubit_freqs.get(q)) for q in qubits
+            }
+            # If no g provided, set equal to previous
+            new_g_dict: dict[tuple[str, str], tu.Value] = {
+                p: g_dict.get(p, full_trajectory[-1].couplings.get(p)) for p in pairs  # type: ignore[misc]
+            }
+
+            full_trajectory.append(FreqMap(dt, new_qubit_freq_dict, new_g_dict))
+        return cls(full_trajectory=full_trajectory, qubits=qubits, pairs=pairs)
+
+    def get_full_trajectory_with_resolved_idles(
+        self, idle_freq_map: dict[str, tu.Value]
+    ) -> list[FreqMap]:
+        """Insert idle frequencies instead of None in trajectory."""
+
+        resolved_trajectory: list[FreqMap] = []
+        for freq_map in self.full_trajectory:
+            resolved_qubit_freqs = {
+                q: idle_freq_map[q] if f is None else f for q, f in freq_map.qubit_freqs.items()
+            }
+            resolved_trajectory.append(attrs.evolve(freq_map, qubit_freqs=resolved_qubit_freqs))
+        return resolved_trajectory
+
+    def plot(
+        self,
+        idle_freq_map: dict[str, tu.Value] | None = None,
+        default_idle_freq: tu.Value = 6.5 * tu.GHz,
+        resolver: cirq.ParamResolverOrSimilarType | None = None,
+        axes: tuple[Axes, Axes] | None = None,
+    ) -> tuple[Axes, Axes]:
+        if idle_freq_map is None:
+            idle_freq_map = {q: default_idle_freq for q in self.qubits}
+        full_trajectory_resolved = cirq.resolve_parameters(
+            self.get_full_trajectory_with_resolved_idles(idle_freq_map), resolver
+        )
+        unresolved_param_names = set().union(
+            *[cirq.parameter_names(freq_map) for freq_map in full_trajectory_resolved]
+        )
+        if unresolved_param_names:
+            raise ValueError(f"There are some parameters {unresolved_param_names} not resolved.")
+
+        times = np.cumsum([step.duration[tu.ns] for step in full_trajectory_resolved])
+
+        if axes is None:
+            _, axes = plt.subplots(1, 2, figsize=(10, 4))
+
+        for qubit_agent in self.qubits:
+            axes[0].plot(
+                times,
+                [step.qubit_freqs[qubit_agent][tu.GHz] for step in full_trajectory_resolved],  # type: ignore[index]
+                label=qubit_agent,
+            )
+        for pair_agent in self.pairs:
+            axes[1].plot(
+                times,
+                [step.couplings[pair_agent][tu.MHz] for step in full_trajectory_resolved],
+                label=pair_agent,
+            )
+
+        for ax, ylabel in zip(axes, ["Qubit freq. (GHz)", "Coupling (MHz)"]):
+            ax.set_xlabel("Time (ns)")
+            ax.set_ylabel(ylabel)
+            ax.legend()
+        plt.tight_layout()
+        return axes

cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py

@@ -0,0 +1,131 @@
+# Copyright 2019 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 pytest
+import sympy
+import tunits as tu
+
+import cirq
+from cirq_google.experimental.analog_experiments import analog_trajectory_util as atu
+
+
+@pytest.fixture
+def freq_map() -> atu.FreqMap:
+    return atu.FreqMap(
+        10 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
+        {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): sympy.Symbol("g_q0_1_q0_2")},
+    )
+
+
+def test_freq_map_param_names(freq_map: atu.FreqMap) -> None:
+    assert cirq.is_parameterized(freq_map)
+    assert cirq.parameter_names(freq_map) == {"f_q0_2", "g_q0_1_q0_2"}
+
+
+def test_freq_map_resolve(freq_map: atu.FreqMap) -> None:
+    resolved_freq_map = cirq.resolve_parameters(
+        freq_map, {"f_q0_2": 6 * tu.GHz, "g_q0_1_q0_2": 7 * tu.MHz}
+    )
+    assert resolved_freq_map == atu.FreqMap(
+        10 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 6 * tu.GHz},
+        {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 7 * tu.MHz},
+    )
+
+
+FreqMapType = tuple[tu.Value, dict[str, tu.Value | None], dict[tuple[str, str], tu.Value]]
+
+
+@pytest.fixture
+def sparse_trajectory() -> list[FreqMapType]:
+    traj1: FreqMapType = (20 * tu.ns, {"q0_1": 5 * tu.GHz}, {})
+    traj2: FreqMapType = (30 * tu.ns, {"q0_2": 8 * tu.GHz}, {})
+    traj3: FreqMapType = (
+        40 * tu.ns,
+        {"q0_0": 8 * tu.GHz, "q0_1": None, "q0_2": None},
+        {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
+    )
+    return [traj1, traj2, traj3]
+
+
+def test_full_traj(sparse_trajectory: list[FreqMapType]) -> None:
+    analog_traj = atu.AnalogTrajectory.from_sparse_trajectory(sparse_trajectory)
+    assert len(analog_traj.full_trajectory) == 4
+    assert analog_traj.full_trajectory[0] == atu.FreqMap(
+        0 * tu.ns,
+        {"q0_0": None, "q0_1": None, "q0_2": None},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert analog_traj.full_trajectory[1] == atu.FreqMap(
+        20 * tu.ns,
+        {"q0_0": None, "q0_1": 5 * tu.GHz, "q0_2": None},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert analog_traj.full_trajectory[2] == atu.FreqMap(
+        30 * tu.ns,
+        {"q0_0": None, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert analog_traj.full_trajectory[3] == atu.FreqMap(
+        40 * tu.ns,
+        {"q0_0": 8 * tu.GHz, "q0_1": None, "q0_2": None},
+        {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
+    )
+
+
+def test_get_full_trajectory_with_resolved_idles(sparse_trajectory: list[FreqMapType]) -> None:
+
+    analog_traj = atu.AnalogTrajectory.from_sparse_trajectory(sparse_trajectory)
+    resolved_full_traj = analog_traj.get_full_trajectory_with_resolved_idles(
+        {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz}
+    )
+
+    assert len(resolved_full_traj) == 4
+    assert resolved_full_traj[0] == atu.FreqMap(
+        0 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert resolved_full_traj[1] == atu.FreqMap(
+        20 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 5 * tu.GHz, "q0_2": 7 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert resolved_full_traj[2] == atu.FreqMap(
+        30 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+    )
+    assert resolved_full_traj[3] == atu.FreqMap(
+        40 * tu.ns,
+        {"q0_0": 8 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz},
+        {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
+    )
+
+
+def test_plot_with_unresolved_parameters():
+    traj1: FreqMapType = (20 * tu.ns, {"q0_1": sympy.Symbol("qf")}, {})
+    traj2: FreqMapType = (sympy.Symbol("t"), {"q0_2": 8 * tu.GHz}, {})
+    analog_traj = atu.AnalogTrajectory.from_sparse_trajectory([traj1, traj2])
+
+    with pytest.raises(ValueError):
+        analog_traj.plot()
+
+
+def test_analog_traj_plot():
+    traj1: FreqMapType = (20 * tu.ns, {"q0_1": sympy.Symbol("qf")}, {})
+    traj2: FreqMapType = (sympy.Symbol("t"), {"q0_2": 8 * tu.GHz}, {})
+    analog_traj = atu.AnalogTrajectory.from_sparse_trajectory([traj1, traj2])
+    analog_traj.plot(resolver={"t": 10 * tu.ns, "qf": 5 * tu.GHz})