cirq-core 1.2.0.dev20230717225858__py3-none-any.whl → 1.3.0__py3-none-any.whl

cirq/transformers/transformer_primitives.py

@@ -104,6 +104,134 @@ def map_moments(
     )
 
 
+def _map_operations_impl(
+    circuit: CIRCUIT_TYPE,
+    map_func: Callable[[ops.Operation, int], ops.OP_TREE],
+    *,
+    deep: bool = False,
+    raise_if_add_qubits=True,
+    tags_to_ignore: Sequence[Hashable] = (),
+    wrap_in_circuit_op: bool = True,
+) -> CIRCUIT_TYPE:
+    """Applies local transformations, by calling `map_func(op, moment_index)` for each operation.
+
+    This method provides a fast, iterative implementation for the two `map_operations_*` variants
+    exposed as public transformer primitives. The high level idea for the iterative implementation
+    is to
+        1) For each operation `op`, find the corresponding mapped operation(s) `mapped_ops`. The
+            set of mapped operations can be either wrapped in a circuit operation or not, depending
+            on the value of flag `wrap_in_circuit_op` and whether the mapped operations will end up
+            occupying more than one moment or not.
+        2) Use the `get_earliest_accommodating_moment_index` infrastructure built for `cirq.Circuit`
+            construction to determine the index at which the mapped operations should be inserted.
+            This step takes care of the nuances that arise due to (a) preserving moment structure
+            and (b) mapped operations spanning across multiple moments (these both are trivial when
+            `op` is mapped to a single `mapped_op` that acts on the same set of qubits).
+
+    By default, the function assumes `issubset(qubit_set(map_func(op, moment_index)), op.qubits)` is
+    True.
+
+    Args:
+        circuit: Input circuit to apply the transformations on. The input circuit is not mutated.
+        map_func: Mapping function from (cirq.Operation, moment_index) to a cirq.OP_TREE. If the
+            resulting optree spans more than 1 moment, it's either wrapped in a tagged circuit
+            operation and inserted in-place in the same moment (if  `wrap_in_circuit_op` is True)
+            OR the mapped operations are inserted directly in the circuit, preserving moment
+            strucutre. The effect is equivalent to (but much faster) a two-step approach of first
+            wrapping the operations in a circuit operation and then calling `cirq.unroll_circuit_op`
+            to unroll the corresponding circuit ops.
+        deep: If true, `map_func` will be recursively applied to circuits wrapped inside
+            any circuit operations contained within `circuit`.
+        raise_if_add_qubits: Set to True by default. If True, raises ValueError if
+            `map_func(op, idx)` adds operations on qubits outside of `op.qubits`.
+        tags_to_ignore: Sequence of tags which should be ignored while applying `map_func` on
+            tagged operations -- i.e. `map_func(op, idx)` will be called only for operations that
+            satisfy `set(op.tags).isdisjoint(tags_to_ignore)`.
+        wrap_in_circuit_op: If True, the mapped operations will be wrapped in a tagged circuit
+        operation and inserted in-place if they occupy more than one moment.
+
+    Raises:
+          ValueError if `issubset(qubit_set(map_func(op, idx)), op.qubits) is False` and
+            `raise_if_add_qubits is True`.
+
+    Returns:
+        Copy of input circuit with mapped operations.
+    """
+    tags_to_ignore_set = set(tags_to_ignore)
+
+    def apply_map_func(op: 'cirq.Operation', idx: int) -> List['cirq.Operation']:
+        if tags_to_ignore_set.intersection(op.tags):
+            return [op]
+        if deep and isinstance(op.untagged, circuits.CircuitOperation):
+            op = op.untagged.replace(
+                circuit=_map_operations_impl(
+                    op.untagged.circuit,
+                    map_func,
+                    deep=deep,
+                    raise_if_add_qubits=raise_if_add_qubits,
+                    tags_to_ignore=tags_to_ignore,
+                    wrap_in_circuit_op=wrap_in_circuit_op,
+                )
+            ).with_tags(*op.tags)
+        mapped_ops = [*ops.flatten_to_ops(map_func(op, idx))]
+        op_qubits = set(op.qubits)
+        mapped_ops_qubits: Set['cirq.Qid'] = set()
+        has_overlapping_ops = False
+        for mapped_op in mapped_ops:
+            if raise_if_add_qubits and not op_qubits.issuperset(mapped_op.qubits):
+                raise ValueError(
+                    f"Mapped operations {mapped_ops} should act on a subset "
+                    f"of qubits of the original operation {op}"
+                )
+            if mapped_ops_qubits.intersection(mapped_op.qubits):
+                has_overlapping_ops = True
+            mapped_ops_qubits = mapped_ops_qubits.union(mapped_op.qubits)
+        if wrap_in_circuit_op and has_overlapping_ops:
+            # Mapped operations should be wrapped in a `CircuitOperation` only iff they occupy more
+            # than one moment, i.e. there are at least two operations that share a qubit.
+            mapped_ops = [
+                circuits.CircuitOperation(circuits.FrozenCircuit(mapped_ops)).with_tags(
+                    MAPPED_CIRCUIT_OP_TAG
+                )
+            ]
+        return mapped_ops
+
+    new_moments: List[List['cirq.Operation']] = []
+
+    # Keep track of the latest time index for each qubit, measurement key, and control key.
+    qubit_time_index: Dict['cirq.Qid', int] = {}
+    measurement_time_index: Dict['cirq.MeasurementKey', int] = {}
+    control_time_index: Dict['cirq.MeasurementKey', int] = {}
+
+    # New mapped operations in the current moment should be inserted after `last_moment_time_index`.
+    last_moment_time_index = -1
+
+    for idx, moment in enumerate(circuit):
+        if wrap_in_circuit_op:
+            new_moments.append([])
+        for op in moment:
+            mapped_ops = apply_map_func(op, idx)
+
+            for mapped_op in mapped_ops:
+                # Identify the earliest moment that can accommodate this op.
+                placement_index = circuits.circuit.get_earliest_accommodating_moment_index(
+                    mapped_op, qubit_time_index, measurement_time_index, control_time_index
+                )
+                placement_index = max(placement_index, last_moment_time_index + 1)
+                new_moments.extend([[] for _ in range(placement_index - len(new_moments) + 1)])
+                new_moments[placement_index].append(mapped_op)
+                for qubit in mapped_op.qubits:
+                    qubit_time_index[qubit] = placement_index
+                for key in protocols.measurement_key_objs(mapped_op):
+                    measurement_time_index[key] = placement_index
+                for key in protocols.control_keys(mapped_op):
+                    control_time_index[key] = placement_index
+
+        last_moment_time_index = len(new_moments) - 1
+
+    return _create_target_circuit_type([circuits.Moment(moment) for moment in new_moments], circuit)
+
+
 def map_operations(
     circuit: CIRCUIT_TYPE,
     map_func: Callable[[ops.Operation, int], ops.OP_TREE],
@@ -139,29 +267,13 @@ def map_operations(
     Returns:
         Copy of input circuit with mapped operations (wrapped in a tagged CircuitOperation).
     """
-
-    def apply_map(op: ops.Operation, idx: int) -> ops.OP_TREE:
-        if not set(op.tags).isdisjoint(tags_to_ignore):
-            return op
-        c = circuits.FrozenCircuit(map_func(op, idx))
-        if raise_if_add_qubits and not c.all_qubits().issubset(op.qubits):
-            raise ValueError(
-                f"Mapped operations {c.all_operations()} should act on a subset "
-                f"of qubits of the original operation {op}"
-            )
-        if len(c) <= 1:
-            # Either empty circuit or all operations act in the same moment;
-            # So, we don't need to wrap them in a circuit_op.
-            return c[0].operations if c else []
-        circuit_op = circuits.CircuitOperation(c).with_tags(MAPPED_CIRCUIT_OP_TAG)
-        return circuit_op
-
-    return map_moments(
+    return _map_operations_impl(
         circuit,
-        lambda m, i: circuits.Circuit(apply_map(op, i) for op in m.operations).moments
-        or [circuits.Moment()],
+        map_func,
         deep=deep,
+        raise_if_add_qubits=raise_if_add_qubits,
         tags_to_ignore=tags_to_ignore,
+        wrap_in_circuit_op=True,
     )
 
 
@@ -191,15 +303,13 @@ def map_operations_and_unroll(
     Returns:
         Copy of input circuit with mapped operations, unrolled in a moment preserving way.
     """
-    return unroll_circuit_op(
-        map_operations(
-            circuit,
-            map_func,
-            deep=deep,
-            raise_if_add_qubits=raise_if_add_qubits,
-            tags_to_ignore=tags_to_ignore,
-        ),
+    return _map_operations_impl(
+        circuit,
+        map_func,
         deep=deep,
+        raise_if_add_qubits=raise_if_add_qubits,
+        tags_to_ignore=tags_to_ignore,
+        wrap_in_circuit_op=False,
     )
 
 

cirq/value/abc_alt_test.py

@@ -140,7 +140,7 @@ def test_classcell_in_namespace():
     class _(metaclass=ABCMetaImplementAnyOneOf):
         def other_method(self):
             # Triggers __classcell__ to be added to the class namespace
-            super()  # coverage: ignore
+            super()  # pragma: no cover
 
 
 def test_two_alternatives():
@@ -170,17 +170,17 @@ def test_two_alternatives():
             return 'alt1'
 
         def alt2(self) -> NoReturn:
-            raise RuntimeError  # coverage: ignore
+            raise RuntimeError  # pragma: no cover
 
     class TwoAlternativesOverride(TwoAlternatives):
         def my_method(self, arg, kw=99) -> str:
             return 'override'
 
         def alt1(self) -> NoReturn:
-            raise RuntimeError  # coverage: ignore
+            raise RuntimeError  # pragma: no cover
 
         def alt2(self) -> NoReturn:
-            raise RuntimeError  # coverage: ignore
+            raise RuntimeError  # pragma: no cover
 
     class TwoAlternativesForceSecond(TwoAlternatives):
         def _do_alt1_with_my_method(self):

cirq/value/duration.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 """A typed time delta that supports picosecond accuracy."""
 
-from typing import AbstractSet, Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
+from typing import AbstractSet, Any, Dict, Optional, Tuple, TYPE_CHECKING, Union, List
 import datetime
 
 import sympy
 import numpy as np
 
 from cirq import protocols
-from cirq._compat import proper_repr
+from cirq._compat import proper_repr, cached_method
 from cirq._doc import document
 
 if TYPE_CHECKING:
@@ -79,48 +79,53 @@ class Duration:
             >>> print(cirq.Duration(micros=1.5 * sympy.Symbol('t')))
             (1500.0*t) ns
         """
+        self._time_vals: List[_NUMERIC_INPUT_TYPE] = [0, 0, 0, 0]
+        self._multipliers = [1, 1000, 1000_000, 1000_000_000]
         if value is not None and value != 0:
             if isinstance(value, datetime.timedelta):
                 # timedelta has microsecond resolution.
-                micros += int(value / datetime.timedelta(microseconds=1))
+                self._time_vals[2] = int(value / datetime.timedelta(microseconds=1))
             elif isinstance(value, Duration):
-                picos += value._picos
+                self._time_vals = value._time_vals
             else:
                 raise TypeError(f'Not a `cirq.DURATION_LIKE`: {repr(value)}.')
-
-        val = picos + nanos * 1000 + micros * 1000_000 + millis * 1000_000_000
-        self._picos: _NUMERIC_OUTPUT_TYPE = float(val) if isinstance(val, np.number) else val
+        input_vals = [picos, nanos, micros, millis]
+        self._time_vals = _add_time_vals(self._time_vals, input_vals)
 
     def _is_parameterized_(self) -> bool:
-        return protocols.is_parameterized(self._picos)
+        return protocols.is_parameterized(self._time_vals)
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return protocols.parameter_names(self._picos)
+        return protocols.parameter_names(self._time_vals)
 
     def _resolve_parameters_(self, resolver: 'cirq.ParamResolver', recursive: bool) -> 'Duration':
-        return Duration(picos=protocols.resolve_parameters(self._picos, resolver, recursive))
+        return _duration_from_time_vals(
+            protocols.resolve_parameters(self._time_vals, resolver, recursive)
+        )
 
+    @cached_method
     def total_picos(self) -> _NUMERIC_OUTPUT_TYPE:
         """Returns the number of picoseconds that the duration spans."""
-        return self._picos
+        val = sum(a * b for a, b in zip(self._time_vals, self._multipliers))
+        return float(val) if isinstance(val, np.number) else val
 
     def total_nanos(self) -> _NUMERIC_OUTPUT_TYPE:
         """Returns the number of nanoseconds that the duration spans."""
-        return self._picos / 1000
+        return self.total_picos() / 1000
 
     def total_micros(self) -> _NUMERIC_OUTPUT_TYPE:
         """Returns the number of microseconds that the duration spans."""
-        return self._picos / 1000_000
+        return self.total_picos() / 1000_000
 
     def total_millis(self) -> _NUMERIC_OUTPUT_TYPE:
         """Returns the number of milliseconds that the duration spans."""
-        return self._picos / 1000_000_000
+        return self.total_picos() / 1000_000_000
 
     def __add__(self, other) -> 'Duration':
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return Duration(picos=self._picos + other._picos)
+        return _duration_from_time_vals(_add_time_vals(self._time_vals, other._time_vals))
 
     def __radd__(self, other) -> 'Duration':
         return self.__add__(other)
@@ -129,29 +134,36 @@ class Duration:
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return Duration(picos=self._picos - other._picos)
+        return _duration_from_time_vals(
+            _add_time_vals(self._time_vals, [-x for x in other._time_vals])
+        )
 
     def __rsub__(self, other) -> 'Duration':
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return Duration(picos=other._picos - self._picos)
+        return _duration_from_time_vals(
+            _add_time_vals(other._time_vals, [-x for x in self._time_vals])
+        )
 
     def __mul__(self, other) -> 'Duration':
         if not isinstance(other, (int, float, sympy.Expr)):
             return NotImplemented
-        return Duration(picos=self._picos * other)
+        if other == 0:
+            return _duration_from_time_vals([0] * 4)
+        return _duration_from_time_vals([x * other for x in self._time_vals])
 
     def __rmul__(self, other) -> 'Duration':
         return self.__mul__(other)
 
     def __truediv__(self, other) -> Union['Duration', float]:
         if isinstance(other, (int, float, sympy.Expr)):
-            return Duration(picos=self._picos / other)
+            new_time_vals = [x / other for x in self._time_vals]
+            return _duration_from_time_vals(new_time_vals)
 
         other_duration = _attempt_duration_like_to_duration(other)
         if other_duration is not None:
-            return self._picos / other_duration._picos
+            return self.total_picos() / other_duration.total_picos()
 
         return NotImplemented
 
@@ -159,56 +171,57 @@ class Duration:
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos == other._picos
+        return self.total_picos() == other.total_picos()
 
     def __ne__(self, other):
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos != other._picos
+        return self.total_picos() != other.total_picos()
 
     def __gt__(self, other):
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos > other._picos
+        return self.total_picos() > other.total_picos()
 
     def __lt__(self, other):
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos < other._picos
+        return self.total_picos() < other.total_picos()
 
     def __ge__(self, other):
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos >= other._picos
+        return self.total_picos() >= other.total_picos()
 
     def __le__(self, other):
         other = _attempt_duration_like_to_duration(other)
         if other is None:
             return NotImplemented
-        return self._picos <= other._picos
+        return self.total_picos() <= other.total_picos()
 
     def __bool__(self):
-        return bool(self._picos)
+        return bool(self.total_picos())
 
     def __hash__(self):
-        if isinstance(self._picos, (int, float)) and self._picos % 1000000 == 0:
-            return hash(datetime.timedelta(microseconds=self._picos / 1000000))
-        return hash((Duration, self._picos))
+        if isinstance(self.total_picos(), (int, float)) and self.total_picos() % 1000000 == 0:
+            return hash(datetime.timedelta(microseconds=self.total_picos() / 1000000))
+        return hash((Duration, self.total_picos()))
 
     def _decompose_into_amount_unit_suffix(self) -> Tuple[int, str, str]:
+        picos = self.total_picos()
         if (
-            isinstance(self._picos, sympy.Mul)
-            and len(self._picos.args) == 2
-            and isinstance(self._picos.args[0], (sympy.Integer, sympy.Float))
+            isinstance(picos, sympy.Mul)
+            and len(picos.args) == 2
+            and isinstance(picos.args[0], (sympy.Integer, sympy.Float))
         ):
-            scale = self._picos.args[0]
-            rest = self._picos.args[1]
+            scale = picos.args[0]
+            rest = picos.args[1]
         else:
-            scale = self._picos
+            scale = picos
             rest = 1
 
         if scale % 1000_000_000 == 0:
@@ -234,7 +247,7 @@ class Duration:
         return amount * rest, unit, suffix
 
     def __str__(self) -> str:
-        if self._picos == 0:
+        if self.total_picos() == 0:
             return 'Duration(0)'
         amount, _, suffix = self._decompose_into_amount_unit_suffix()
         if not isinstance(amount, (int, float, sympy.Symbol)):
@@ -257,3 +270,21 @@ def _attempt_duration_like_to_duration(value: Any) -> Optional[Duration]:
     if isinstance(value, (int, float)) and value == 0:
         return Duration()
     return None
+
+
+def _add_time_vals(
+    val1: List[_NUMERIC_INPUT_TYPE], val2: List[_NUMERIC_INPUT_TYPE]
+) -> List[_NUMERIC_INPUT_TYPE]:
+    ret: List[_NUMERIC_INPUT_TYPE] = []
+    for i in range(4):
+        if val1[i] and val2[i]:
+            ret.append(val1[i] + val2[i])
+        else:
+            ret.append(val1[i] or val2[i])
+    return ret
+
+
+def _duration_from_time_vals(time_vals: List[_NUMERIC_INPUT_TYPE]):
+    ret = Duration()
+    ret._time_vals = time_vals
+    return ret

cirq/value/duration_test.py

@@ -168,9 +168,11 @@ def test_sub():
 def test_mul():
     assert Duration(picos=2) * 3 == Duration(picos=6)
     assert 4 * Duration(picos=3) == Duration(picos=12)
+    assert 0 * Duration(picos=10) == Duration()
 
     t = sympy.Symbol('t')
     assert t * Duration(picos=3) == Duration(picos=3 * t)
+    assert 0 * Duration(picos=t) == Duration(picos=0)
 
     with pytest.raises(TypeError):
         _ = Duration() * Duration()

cirq/value/measurement_key_test.py

@@ -43,7 +43,7 @@ def test_eq_and_hash():
             self.some_str = some_str
 
         def __str__(self):
-            return self.some_str  # coverage: ignore
+            return self.some_str  # pragma: no cover
 
     mkey = cirq.MeasurementKey('key')
     assert mkey == 'key'

cirq/value/product_state.py

@@ -61,8 +61,7 @@ class ProductState:
 
     def __init__(self, states=None):
         if states is None:
-            # coverage: ignore
-            states = dict()
+            states = dict()  # pragma: no cover
 
         object.__setattr__(self, 'states', states)
 
@@ -200,8 +199,7 @@ class _XEigenState(_PauliEigenState):
             return np.array([1, 1]) / np.sqrt(2)
         elif self.eigenvalue == -1:
             return np.array([1, -1]) / np.sqrt(2)
-        # coverage: ignore
-        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")
+        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")  # pragma: no cover
 
     def stabilized_by(self) -> Tuple[int, 'cirq.Pauli']:
         # Prevent circular import from `value.value_equality`
@@ -218,8 +216,7 @@ class _YEigenState(_PauliEigenState):
             return np.array([1, 1j]) / np.sqrt(2)
         elif self.eigenvalue == -1:
             return np.array([1, -1j]) / np.sqrt(2)
-        # coverage: ignore
-        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")
+        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")  # pragma: no cover
 
     def stabilized_by(self) -> Tuple[int, 'cirq.Pauli']:
         from cirq import ops
@@ -235,8 +232,7 @@ class _ZEigenState(_PauliEigenState):
             return np.array([1, 0])
         elif self.eigenvalue == -1:
             return np.array([0, 1])
-        # coverage: ignore
-        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")
+        raise ValueError(f"Bad eigenvalue: {self.eigenvalue}")  # pragma: no cover
 
     def stabilized_by(self) -> Tuple[int, 'cirq.Pauli']:
         from cirq import ops

cirq/value/value_equality_attr.py

@@ -17,7 +17,7 @@ from typing import Any, Callable, Optional, overload, Union
 
 from typing_extensions import Protocol
 
-from cirq import protocols
+from cirq import protocols, _compat
 
 
 class _SupportsValueEquality(Protocol):
@@ -50,8 +50,7 @@ class _SupportsValueEquality(Protocol):
         Returns:
             Any type supported by `cirq.approx_eq()`.
         """
-        # coverage: ignore
-        return self._value_equality_values_()
+        return self._value_equality_values_()  # pragma: no cover
 
     def _value_equality_values_cls_(self) -> Any:
         """Automatically implemented by the `cirq.value_equality` decorator.
@@ -222,13 +221,21 @@ def value_equality(
             )
     else:
         setattr(cls, '_value_equality_values_cls_', lambda self: cls)
-    setattr(cls, '__hash__', None if unhashable else _value_equality_hash)
+    cached_values_getter = values_getter if unhashable else _compat.cached_method(values_getter)
+    setattr(cls, '_value_equality_values_', cached_values_getter)
+    setattr(cls, '__hash__', None if unhashable else _compat.cached_method(_value_equality_hash))
     setattr(cls, '__eq__', _value_equality_eq)
     setattr(cls, '__ne__', _value_equality_ne)
 
     if approximate:
         if not hasattr(cls, '_value_equality_approximate_values_'):
-            setattr(cls, '_value_equality_approximate_values_', values_getter)
+            setattr(cls, '_value_equality_approximate_values_', cached_values_getter)
+        else:
+            approx_values_getter = getattr(cls, '_value_equality_approximate_values_')
+            cached_approx_values_getter = (
+                approx_values_getter if unhashable else _compat.cached_method(approx_values_getter)
+            )
+            setattr(cls, '_value_equality_approximate_values_', cached_approx_values_getter)
         setattr(cls, '_approx_eq_', _value_equality_approx_eq)
 
     return cls

cirq/vis/heatmap.py

@@ -15,6 +15,7 @@ import copy
 from dataclasses import astuple, dataclass
 from typing import (
     Any,
+    cast,
     Dict,
     List,
     Mapping,
@@ -217,7 +218,7 @@ class Heatmap:
         )
         position = self._config['colorbar_position']
         orien = 'vertical' if position in ('left', 'right') else 'horizontal'
-        colorbar = ax.figure.colorbar(
+        colorbar = cast(plt.Figure, ax.figure).colorbar(
             mappable, colorbar_ax, ax, orientation=orien, **self._config.get("colorbar_options", {})
         )
         colorbar_ax.tick_params(axis='y', direction='out')
@@ -230,15 +231,15 @@ class Heatmap:
         ax: plt.Axes,
     ) -> None:
         """Writes annotations to the center of cells. Internal."""
-        for (center, annotation), facecolor in zip(centers_and_annot, collection.get_facecolors()):
+        for (center, annotation), facecolor in zip(centers_and_annot, collection.get_facecolor()):
             # Calculate the center of the cell, assuming that it is a square
             # centered at (x=col, y=row).
             if not annotation:
                 continue
             x, y = center
-            face_luminance = vis_utils.relative_luminance(facecolor)
+            face_luminance = vis_utils.relative_luminance(facecolor)  # type: ignore
             text_color = 'black' if face_luminance > 0.4 else 'white'
-            text_kwargs = dict(color=text_color, ha="center", va="center")
+            text_kwargs: Dict[str, Any] = dict(color=text_color, ha="center", va="center")
             text_kwargs.update(self._config.get('annotation_text_kwargs', {}))
             ax.text(x, y, annotation, **text_kwargs)
 
@@ -295,6 +296,7 @@ class Heatmap:
         show_plot = not ax
         if not ax:
             fig, ax = plt.subplots(figsize=(8, 8))
+            ax = cast(plt.Axes, ax)
         original_config = copy.deepcopy(self._config)
         self.update_config(**kwargs)
         collection = self._plot_on_axis(ax)
@@ -381,6 +383,7 @@ class TwoQubitInteractionHeatmap(Heatmap):
         show_plot = not ax
         if not ax:
             fig, ax = plt.subplots(figsize=(8, 8))
+            ax = cast(plt.Axes, ax)
         original_config = copy.deepcopy(self._config)
         self.update_config(**kwargs)
         qubits = set([q for qubits in self._value_map.keys() for q in qubits])

cirq/vis/heatmap_test.py

@@ -14,6 +14,7 @@
 """Tests for Heatmap."""
 
 import pathlib
+import shutil
 import string
 from tempfile import mkdtemp
 
@@ -33,6 +34,14 @@ def ax():
     return figure.add_subplot(111)
 
 
+def test_default_ax():
+    row_col_list = ((0, 5), (8, 1), (7, 0), (13, 5), (1, 6), (3, 2), (2, 8))
+    test_value_map = {
+        grid_qubit.GridQubit(row, col): np.random.random() for (row, col) in row_col_list
+    }
+    _, _ = heatmap.Heatmap(test_value_map).plot()
+
+
 @pytest.mark.parametrize('tuple_keys', [True, False])
 def test_cells_positions(ax, tuple_keys):
     row_col_list = ((0, 5), (8, 1), (7, 0), (13, 5), (1, 6), (3, 2), (2, 8))
@@ -60,6 +69,8 @@ def test_two_qubit_heatmap(ax):
     title = "Two Qubit Interaction Heatmap"
     heatmap.TwoQubitInteractionHeatmap(value_map, title=title).plot(ax)
     assert ax.get_title() == title
+    # Test default axis
+    heatmap.TwoQubitInteractionHeatmap(value_map, title=title).plot()
 
 
 def test_invalid_args():
@@ -104,10 +115,8 @@ def test_cell_colors(ax, colormap_name):
         col = int(round(np.mean([v[0] for v in vertices])))
         value = test_row_col_map[(row, col)]
         color_scale = (value - vmin) / (vmax - vmin)
-        if color_scale < 0.0:
-            color_scale = 0.0
-        if color_scale > 1.0:
-            color_scale = 1.0
+        color_scale = max(color_scale, 0.0)
+        color_scale = min(color_scale, 1.0)
         expected_color = np.array(colormap(color_scale))
         assert np.all(np.isclose(facecolor, expected_color))
 
@@ -309,6 +318,7 @@ def test_colorbar(ax, position, size, pad):
 
     plt.close(fig1)
     plt.close(fig2)
+    shutil.rmtree(tmp_dir)
 
 
 @pytest.mark.usefixtures('closefigures')