Fix small issues in Shor tutorial (quantumlib#5639)

* Fix small issues in Shor tutorial

- use cirq.X instead of cirq.ops.X
- Add a circuit_diagram_info for better printing
- Add a few clarifications here and there.

Also, this changes ArithmeticOperation to use a raw
docstring so it will format better.

cirq-core/cirq/ops/arithmetic_operation.py

@@ -31,7 +31,7 @@
 
 @deprecated_class(deadline='v0.16', fix='Use cirq.ArithmeticGate')
 class ArithmeticOperation(Operation, metaclass=abc.ABCMeta):
-    """A helper class for implementing reversible classical arithmetic.
+    r"""A helper class for implementing reversible classical arithmetic.
 
     Child classes must override the `registers`, `with_registers`, and `apply`
     methods.

docs/experiments/shor.ipynb

@@ -375,7 +375,11 @@
     "        return Adder(*new_registers)\n",
     "    \n",
     "    def apply(self, target_value, input_value):\n",
-    "        return target_value + input_value"
+    "        return target_value + input_value\n",
+    "\n",
+    "    def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs):\n",
+    "        wire_symbols = [' + ' for _ in range(len(self.input_register)+len(self.target_register))]\n",
+    "        return cirq.CircuitDiagramInfo(wire_symbols=tuple(wire_symbols))"
    ]
   },
   {
@@ -404,8 +408,8 @@
     "\n",
     "# Define the circuit.\n",
     "circ = cirq.Circuit(\n",
-    "    cirq.ops.X.on(qreg1[0]),\n",
-    "    cirq.ops.X.on(qreg2[1]),\n",
+    "    cirq.X.on(qreg1[0]),\n",
+    "    cirq.X.on(qreg2[1]),\n",
     "    Adder(input_register=qreg1, target_register=qreg2),\n",
     "    cirq.measure_each(*qreg1),\n",
     "    cirq.measure_each(*qreg2)\n",
@@ -531,6 +535,7 @@
     "            self,\n",
     "            *new_registers: Union[int, Sequence['cirq.Qid']],\n",
     "    ) -> cirq.ArithmeticOperation:\n",
+    "        \"\"\"Returns a new ModularExp object with new registers.\"\"\"\n",
     "        if len(new_registers) != 4:\n",
     "            raise ValueError(f'Expected 4 registers (target, exponent, base, '\n",
     "                             f'modulus), but got {len(new_registers)}')\n",
@@ -547,6 +552,18 @@
     "        return ModularExp(target, exponent, base, modulus)\n",
     "\n",
     "    def apply(self, *register_values: int) -> int:\n",
+    "        \"\"\"Applies modular exponentiation to the registers.\n",
+    "\n",
+    "        Four values should be passed in.  They are, in order:\n",
+    "          - the target\n",
+    "          - the exponent\n",
+    "          - the base\n",
+    "          - the modulus\n",
+    "\n",
+    "        Note that the target and exponent should be qubit\n",
+    "        registers, while the base and modulus should be\n",
+    "        constant parameters that control the resulting unitary.\n",
+    "        \"\"\"\n",
     "        assert len(register_values) == 4\n",
     "        target, exponent, base, modulus = register_values\n",
     "        if target >= modulus:\n",
@@ -557,6 +574,12 @@
     "            self,\n",
     "            args: cirq.CircuitDiagramInfoArgs,\n",
     "    ) -> cirq.CircuitDiagramInfo:\n",
+    "        \"\"\"Returns a 'CircuitDiagramInfo' object for printing circuits.\n",
+    "\n",
+    "        This function just returns information on how to print this operation\n",
+    "        out in a circuit diagram so that the registers are labeled\n",
+    "        appropriately as exponent ('e') and target ('t').\n",
+    "        \"\"\"\n",
     "        assert args.known_qubits is not None\n",
     "        wire_symbols: List[str] = []\n",
     "        t, e = 0, 0\n",
@@ -676,7 +699,7 @@
     "    \"\"\"Returns quantum circuit which computes the order of x modulo n.\n",
     "\n",
     "    The circuit uses Quantum Phase Estimation to compute an eigenvalue of\n",
-    "    the unitary\n",
+    "    the following unitary:\n",
     "\n",
     "        U|y⟩ = |y * x mod n⟩      0 <= y < n\n",
     "        U|y⟩ = |y⟩                n <= y\n",
@@ -1096,6 +1119,7 @@
  ],
  "metadata": {
   "colab": {
+   "collapsed_sections": [],
    "name": "shor.ipynb",
    "toc_visible": true
   },