Various doc fixes (#2497)

- Add a few missing doc strings
- Fix some example sections coming before args/results
- Fix a few missing Args
- Change some type annotations to refer to exposed 'cirq.XXX' instead of local cirq.YYY.ZZZ.XXXX

cirq/circuits/text_diagram_drawer.py

@@ -55,6 +55,7 @@ class TextDiagramDrawer:
     """
 
     def __init__(self):
+        """Initializes an empty diagram drawer."""
         self.entries: Dict[Tuple[int, int], _DiagramText] = dict()
         self.vertical_lines: List[_VerticalLine] = []
         self.horizontal_lines: List[_HorizontalLine] = []

cirq/devices/noise_model.py

@@ -183,7 +183,11 @@ def _json_dict_(self):
 
 @value.value_equality
 class ConstantQubitNoiseModel(NoiseModel):
-    """Applies noise to each qubit individually at the start of every moment."""
+    """Applies noise to each qubit individually at the start of every moment.
+
+    This is the noise model that is wrapped around an operation when that
+    operation is given as "the noise to use" for a `NOISE_MODEL_LIKE` parameter.
+    """
 
     def __init__(self, qubit_noise_gate: 'cirq.Gate'):
         if qubit_noise_gate.num_qubits() != 1:

cirq/google/ops/sycamore_gate.py

@@ -25,8 +25,7 @@
 
 
 class SycamoreGate(ops.FSimGate):
-    """The Sycamore gate is a two-qubit operation equivalent to
-    FSimGate(π/2, π/6).
+    """The Sycamore gate is a two-qubit gate equivalent to FSimGate(π/2, π/6).
 
     The unitary of this gate is
 
@@ -38,7 +37,7 @@ class SycamoreGate(ops.FSimGate):
     This gate can be performed on the Google's Sycamore chip and
     is close to the gates that were used to demonstrate quantum
     supremacy used in this paper:
-    https://www.nature.com/articles/s41586-019-1666-5.
+    https://www.nature.com/articles/s41586-019-1666-5
     """
 
     def __init__(self):

cirq/linalg/transformations.py

@@ -388,10 +388,16 @@ def subwavefunction(wavefunction: np.ndarray,
     of kets like  $|\psi\rangle$ = $|x\rangle \otimes |y\rangle$, and
     $|x\rangle$ is defined over the subset `keep_indices` of k qubits, then
     this method will factor $|\psi\rangle$ into $|x\rangle$ and $|y\rangle$ and
-    return $|x\rangle$.  Note that $|x\rangle$ is not unique, because $(e^{i
-    \theta} |y\rangle) \otimes (|x\rangle) = (|y\rangle) \otimes (e^{i \theta}
-    |x\rangle)$ . This method randomizes the global phase of $|x\rangle$ in
-    order to avoid accidental reliance on it.
+    return $|x\rangle$. Note that $|x\rangle$ is not unique, because scalar
+    multiplication may be absorbed by any factor of a tensor product:
+
+    $$
+        (e^{i \theta} |y\rangle) \otimes (|x\rangle)
+        = (|y\rangle) \otimes (e^{i \theta} |x\rangle)
+    $$
+
+    This method randomizes the global phase of $|x\rangle$ in order to avoid
+    accidental reliance on the global phase being some specific value.
 
     If the provided wavefunction cannot be factored into a pure state over
     `keep_indices`, the method will fall back to return `default`. If `default`

cirq/study/flatten_expressions.py

@@ -42,7 +42,15 @@ def flatten(val: Any) -> Tuple[Any, 'ExpressionMap']:
     to a sweep over the flattened symbols e.g. a sweep over `sympy.Symbol('x')`
     to a sweep over `sympy.Symbol('<x + 1>')`.
 
-    Example:
+    Args:
+        val: The value to copy and substitute parameter expressions with
+        flattened symbols.
+
+    Returns:
+        The tuple (new value, expression map) where new value and expression map
+        are described above.
+
+    Examples:
         >>> qubit = cirq.LineQubit(0)
         >>> a = sympy.Symbol('a')
         >>> circuit = cirq.Circuit(
@@ -75,8 +83,9 @@ def flatten(val: Any) -> Tuple[Any, 'ExpressionMap']:
         {'<a/4>': 0.75, '<1 - a/2>': -0.5}
 
         >>> for params in sweep:  # Original
-        ...     print(circuit, '=>', end=' ')
-        ...     print(cirq.resolve_parameters(circuit, params))
+        ...     print(circuit,
+        ...           '=>',
+        ...           cirq.resolve_parameters(circuit, params))
         0: ───X^(a/4)───Y^(1 - a/2)─── => 0: ───X^0───Y───
         0: ───X^(a/4)───Y^(1 - a/2)─── => 0: ───X^0.25───Y^0.5───
         0: ───X^(a/4)───Y^(1 - a/2)─── => 0: ───X^0.5───Y^0───
@@ -89,14 +98,6 @@ def flatten(val: Any) -> Tuple[Any, 'ExpressionMap']:
         0: ───X^(<a/4>)───Y^(<1 - a/2>)─── => 0: ───X^0.25───Y^0.5───
         0: ───X^(<a/4>)───Y^(<1 - a/2>)─── => 0: ───X^0.5───Y^0───
         0: ───X^(<a/4>)───Y^(<1 - a/2>)─── => 0: ───X^0.75───Y^-0.5───
-
-    Args:
-        val: The value to copy and substitute parameter expressions with
-        flattened symbols.
-
-    Returns:
-        The tuple (new value, expression map) where new value and expression map
-        are described above.
     """
     flattener = _ParamFlattener()
     val_flat = flattener.flatten(val)

cirq/study/resolver.py

@@ -47,12 +47,13 @@ class ParamResolver(object):
             assigned value.
     """
 
-    def __new__(cls, param_dict: ParamResolverOrSimilarType = None):
+    def __new__(cls, param_dict: 'cirq.ParamResolverOrSimilarType' = None):
         if isinstance(param_dict, ParamResolver):
             return param_dict
         return super().__new__(cls)
 
-    def __init__(self, param_dict: ParamResolverOrSimilarType = None) -> None:
+    def __init__(self,
+                 param_dict: 'cirq.ParamResolverOrSimilarType' = None) -> None:
         if hasattr(self, 'param_dict'):
             return  # Already initialized. Got wrapped as part of the __new__.
 

cirq/testing/lin_alg_utils.py

@@ -29,6 +29,9 @@ def random_superposition(dim: int,
 
     Args:
         dim: The dimension of the vector.
+        random_state: A seed (int) or `np.random.RandomState` class to use when
+            generating random values. If not set, defaults to using the module
+            methods in `np.random`.
 
     Returns:
         The sampled unit-length vector.
@@ -46,11 +49,11 @@ def random_unitary(dim: int, *,
     """Returns a random unitary matrix distributed with Haar measure.
 
     Args:
-      dim: The width and height of the matrix.
-      random_state: A seed to use for random number generation.
+        dim: The width and height of the matrix.
+        random_state: A seed to use for random number generation.
 
     Returns:
-      The sampled unitary matrix.
+        The sampled unitary matrix.
 
     References:
         'How to generate random matrices from the classical compact groups'
@@ -69,10 +72,13 @@ def random_orthogonal(dim: int, *, random_state: value.RANDOM_STATE_LIKE = None
     """Returns a random orthogonal matrix distributed with Haar measure.
 
     Args:
-      dim: The width and height of the matrix.
+        dim: The width and height of the matrix.
+        random_state: A seed (int) or `np.random.RandomState` class to use when
+            generating random values. If not set, defaults to using the module
+            methods in `np.random`.
 
     Returns:
-      The sampled orthogonal matrix.
+        The sampled orthogonal matrix.
 
     References:
         'How to generate random matrices from the classical compact groups'
@@ -93,11 +99,13 @@ def random_special_unitary(dim: int,
     """Returns a random special unitary distributed with Haar measure.
 
     Args:
-      dim: The width and height of the matrix.
-      random_state: A seed to use for random number generation.
+        dim: The width and height of the matrix.
+        random_state: A seed (int) or `np.random.RandomState` class to use when
+            generating random values. If not set, defaults to using the module
+            methods in `np.random`.
 
     Returns:
-      The sampled special unitary.
+        The sampled special unitary.
     """
     r = random_unitary(dim, random_state=random_state)
     r[0, :] /= np.linalg.det(r)
@@ -111,10 +119,13 @@ def random_special_orthogonal(dim: int,
     """Returns a random special orthogonal matrix distributed with Haar measure.
 
     Args:
-      dim: The width and height of the matrix.
+        dim: The width and height of the matrix.
+        random_state: A seed (int) or `np.random.RandomState` class to use when
+            generating random values. If not set, defaults to using the module
+            methods in `np.random`.
 
     Returns:
-      The sampled special orthogonal matrix.
+        The sampled special orthogonal matrix.
     """
     m = random_orthogonal(dim, random_state=random_state)
     if np.linalg.det(m) < 0:

cirq/value/digits.py

@@ -18,6 +18,12 @@
 def big_endian_bits_to_int(bits: Iterable[Any]) -> int:
     """Returns the big-endian integer specified by the given bits.
 
+    Args:
+        bits: Descending bits of the integer, with the 1s bit at the end.
+
+    Returns:
+        The integer.
+
     Examples:
 
         >>> cirq.big_endian_bits_to_int([0, 1])
@@ -31,12 +37,6 @@ def big_endian_bits_to_int(bits: Iterable[Any]) -> int:
 
         >>> cirq.big_endian_bits_to_int([1, 0, 0, 1, 0])
         18
-
-    Args:
-        bits: Descending bits of the integer, with the 1s bit at the end.
-
-    Returns:
-        The integer.
     """
     result = 0
     for e in bits:
@@ -49,16 +49,6 @@ def big_endian_bits_to_int(bits: Iterable[Any]) -> int:
 def big_endian_int_to_bits(val: int, *, bit_count: int) -> List[int]:
     """Returns the big-endian bits of an integer.
 
-    Examples:
-        >>> cirq.big_endian_int_to_bits(19, bit_count=8)
-        [0, 0, 0, 1, 0, 0, 1, 1]
-
-        >>> cirq.big_endian_int_to_bits(19, bit_count=4)
-        [0, 0, 1, 1]
-
-        >>> cirq.big_endian_int_to_bits(-3, bit_count=4)
-        [1, 1, 0, 1]
-
     Args:
         val: The integer to get bits from. This integer is permitted to be
             larger than `2**bit_count` (in which case the high bits of the
@@ -68,6 +58,16 @@ def big_endian_int_to_bits(val: int, *, bit_count: int) -> List[int]:
 
     Returns:
         The bits.
+
+    Examples:
+        >>> cirq.big_endian_int_to_bits(19, bit_count=8)
+        [0, 0, 0, 1, 0, 0, 1, 1]
+
+        >>> cirq.big_endian_int_to_bits(19, bit_count=4)
+        [0, 0, 1, 1]
+
+        >>> cirq.big_endian_int_to_bits(-3, bit_count=4)
+        [1, 1, 0, 1]
     """
     return [(val >> i) & 1 for i in range(bit_count)[::-1]]
 
@@ -76,17 +76,6 @@ def big_endian_digits_to_int(digits: Iterable[int], *,
                              base: Union[int, Iterable[int]]) -> int:
     """Returns the big-endian integer specified by the given digits and base.
 
-    Examples:
-
-        >>> cirq.big_endian_digits_to_int([0, 1], base=10)
-        1
-
-        >>> cirq.big_endian_digits_to_int([1, 0], base=10)
-        10
-
-        >>> cirq.big_endian_digits_to_int([1, 2, 3], base=[2, 3, 4])
-        23
-
     Args:
         digits: Digits of the integer, with the least significant digit at the
             end.
@@ -104,6 +93,17 @@ def big_endian_digits_to_int(digits: Iterable[int], *,
             One of the digits is out of range for its base.
             The base was specified per-digit (as a list) but the length of the
                 bases list is different from the number of digits.
+
+    Examples:
+
+        >>> cirq.big_endian_digits_to_int([0, 1], base=10)
+        1
+
+        >>> cirq.big_endian_digits_to_int([1, 0], base=10)
+        10
+
+        >>> cirq.big_endian_digits_to_int([1, 2, 3], base=[2, 3, 4])
+        23
     """
     digits = tuple(digits)
     base = (base,) * len(digits) if isinstance(base, int) else tuple(base)
@@ -138,13 +138,6 @@ def big_endian_int_to_digits(val: int,
                              base: Union[int, Iterable[int]]) -> List[int]:
     """Separates an integer into big-endian digits.
 
-    Examples:
-        >>> cirq.big_endian_int_to_digits(11, digit_count=4, base=10)
-        [0, 0, 1, 1]
-
-        >>> cirq.big_endian_int_to_digits(11, base=[2, 3, 4])
-        [0, 2, 3]
-
     Args:
         val: The integer to get digits from. Must be non-negative and less than
             the maximum representable value, given the specified base(s) and
@@ -165,6 +158,13 @@ def big_endian_int_to_digits(val: int,
                 `digit_count` was not provided.
             Inconsistent digit count. The `base` was specified as a per-digit
                 list, and `digit_count` was also provided, but they disagree.
+
+    Examples:
+        >>> cirq.big_endian_int_to_digits(11, digit_count=4, base=10)
+        [0, 0, 1, 1]
+
+        >>> cirq.big_endian_int_to_digits(11, base=[2, 3, 4])
+        [0, 2, 3]
     """
     if isinstance(base, int):
         if digit_count is None:

cirq/vis/heatmap.py

@@ -195,7 +195,8 @@ def plot(self, ax: Optional[plt.Axes] = None, **pcolor_options: Any
                 plotted on, and shown.
             pcolor_options: keyword arguments passed to ax.pcolor().
 
-        Returns: a 3-tuple (ax, mesh, value_table)
+        Returns:
+            A 3-tuple (ax, mesh, value_table).
             ax: the `plt.Axes` that is plotted on.
             mesh: the collection of paths drawn and filled.
             value_table: the 2-D pandas DataFrame of values constructed from

cirq/work/pauli_sum_collector.py

@@ -13,20 +13,23 @@
 # limitations under the License.
 
 import collections
-from typing import cast, Dict, Optional, Union
+from typing import cast, Dict, Optional, Union, TYPE_CHECKING
 
 import numpy as np
 
-from cirq import circuits, study, ops
+from cirq import ops
 from cirq.work import collector
 
+if TYPE_CHECKING:
+    import cirq
+
 
 class PauliSumCollector(collector.Collector):
     """Estimates the energy of a linear combination of Pauli observables."""
 
     def __init__(self,
-                 circuit: circuits.Circuit,
-                 observable: ops.PauliSumLike,
+                 circuit: 'cirq.Circuit',
+                 observable: 'cirq.PauliSumLike',
                  *,
                  samples_per_term: int,
                  max_samples_per_job: int = 1000000):
@@ -61,7 +64,7 @@ def __init__(self,
         self._samples_per_term = samples_per_term
         self._total_samples_requested = 0
 
-    def next_job(self) -> Optional[collector.CircuitSampleJob]:
+    def next_job(self) -> Optional['cirq.CircuitSampleJob']:
         i = self._total_samples_requested // self._samples_per_term
         if i >= len(self._pauli_coef_terms):
             return None
@@ -76,8 +79,8 @@ def next_job(self) -> Optional[collector.CircuitSampleJob]:
             repetitions=amount_to_request,
             tag=pauli)
 
-    def on_job_result(self, job: collector.CircuitSampleJob,
-                      result: study.TrialResult):
+    def on_job_result(self, job: 'cirq.CircuitSampleJob',
+                      result: 'cirq.TrialResult'):
         job_id = cast(ops.PauliString, job.tag)
         parities = result.histogram(key='out',
                                     fold_func=lambda bits: np.sum(bits) % 2)
@@ -99,9 +102,9 @@ def estimated_energy(self) -> Union[float, complex]:
         return energy
 
 
-def _circuit_plus_pauli_string_measurements(circuit: circuits.Circuit,
-                                            pauli_string: ops.PauliString
-                                           ) -> circuits.Circuit:
+def _circuit_plus_pauli_string_measurements(circuit: 'cirq.Circuit',
+                                            pauli_string: 'cirq.PauliString'
+                                           ) -> 'cirq.Circuit':
     """A circuit measuring the given observable at the end of the given circuit.
     """
     assert pauli_string