[Pasqal updates] New qubit classes (#3133)

This the first child-PR of #3126 .

- Introduces new Pasqal qubit classes: `ThreeDQubit`and `TwoDQubit` (a subclass of `ThreeDQubit`)
- These new classes are intended to replace the original `ThreeDGridQubit`, which will be deleted on the next PR (where other classes' dependencies will be resolved).
- In comparison with `ThreeDGridQubit`, these new classes allow for arbitrary placement of qubits in 3D or 2D space, instead of restricting them to a grid. This more closely resembles the capabilities of Pasqal devices.

cirq/pasqal/__init__.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 
 from cirq.pasqal.pasqal_qubits import (
-    ThreeDGridQubit,)
+    ThreeDGridQubit,
+    ThreeDQubit,
+    TwoDQubit,
+)
 
 from cirq.pasqal.pasqal_device import (
     PasqalDevice,)

cirq/pasqal/pasqal_qubits.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
 from numpy import sqrt, isclose
+import numpy as np
 
 import cirq
 
@@ -53,7 +54,7 @@ def is_adjacent(self, other: cirq.ops.Qid) -> bool:
         return isclose(self.distance(other), 1)
 
     def distance(self, other: cirq.ops.Qid) -> float:
-        """Returns the distance between two qubits in a 3D grid."""
+        """Returns the distance between two qubits in a 3d grid."""
         if not isinstance(other, ThreeDGridQubit):
             raise TypeError(
                 "Can compute distance to another ThreeDGridQubit, but {}".
@@ -192,3 +193,159 @@ def __rsub__(self, other: Tuple[int, int, int]) -> 'ThreeDGridQubit':
 
     def __neg__(self) -> 'ThreeDGridQubit':
         return ThreeDGridQubit(row=-self.row, col=-self.col, lay=-self.lay)
+
+
+class ThreeDQubit(cirq.ops.Qid):
+    """A qubit in 3d.
+
+    ThreeDQubits use z-y-x ordering:
+
+        ThreeDQubit(0, 0, 0) < ThreeDQubit(1, 0, 0)
+        < ThreeDQubit(0, 1, 0) < ThreeDQubit(1, 1, 0)
+        < ThreeDQubit(0, 0, 1) < ThreeDQubit(1, 0, 1)
+        < ThreeDQubit(0, 1, 1) < ThreeDQubit(1, 1, 1)
+    """
+
+    def __init__(self, x: float, y: float, z: float):
+        self.x = x
+        self.y = y
+        self.z = z
+
+    def _comparison_key(self):
+        return round(self.z, 15), round(self.y, 15), round(self.x, 15)
+
+    @property
+    def dimension(self) -> int:
+        return 2
+
+    def distance(self, other: cirq.ops.Qid) -> float:
+        """Returns the distance between two qubits in 3d."""
+        if not isinstance(other, ThreeDQubit):
+            raise TypeError(
+                "Can compute distance to another ThreeDQubit, but {}".format(
+                    other))
+        return sqrt((self.x - other.x)**2 + (self.y - other.y)**2 +
+                    (self.z - other.z)**2)
+
+    @staticmethod
+    def cube(diameter: int, x0: float = 0, y0: float = 0,
+             z0: float = 0) -> List['ThreeDQubit']:
+        """Returns a cube of ThreeDQubits.
+
+        Args:
+            diameter: Length of a side of the square
+            x0: x-coordinate of the first qubit
+            y0: y-coordinate of the first qubit
+            z0: z-coordinate of the first qubit
+
+        Returns:
+            A list of ThreeDQubits filling in a square grid
+        """
+        return ThreeDQubit.parallelep(diameter,
+                                      diameter,
+                                      diameter,
+                                      x0=x0,
+                                      y0=y0,
+                                      z0=z0)
+
+    @staticmethod
+    def parallelep(rows: int,
+                   cols: int,
+                   lays: int,
+                   x0: float = 0,
+                   y0: float = 0,
+                   z0: float = 0) -> List['ThreeDQubit']:
+        """Returns a parallelepiped of ThreeDQubits.
+
+        Args:
+            rows: Number of rows in the rectangle
+            cols: Number of columns in the rectangle
+            x0: x-coordinate of the first qubit
+            y0: y-coordinate of the first qubit
+            z0: z-coordinate of the first qubit
+
+        Returns:
+            A list of ThreeDQubits filling in a 3d grid
+        """
+        return [
+            ThreeDQubit(x0 + x, y0 + y, z0 + z) for z in range(lays)
+            for y in range(cols) for x in range(rows)
+        ]
+
+    def __repr__(self):
+        return 'pasqal.ThreeDQubit({}, {}, {})'.format(self.x, self.y, self.z)
+
+    def __str__(self):
+        return '({}, {}, {})'.format(self.x, self.y, self.z)
+
+    def _json_dict_(self):
+        return cirq.protocols.obj_to_dict_helper(self, ['x', 'y', 'z'])
+
+
+class TwoDQubit(ThreeDQubit):
+    """A qubit in 2d."""
+
+    def __init__(self, x: float, y: float):
+        super().__init__(x, y, z=0)
+
+    @staticmethod
+    def square(diameter: int, x0: float = 0,
+               y0: float = 0) -> List['TwoDQubit']:
+        """Returns a square of TwoDQubit.
+
+        Args:
+            diameter: Length of a side of the square
+            x0: x-coordinate of the first qubit
+            y0: y-coordinate of the first qubit
+
+        Returns:
+            A list of TwoDQubits filling in a square grid
+        """
+        return TwoDQubit.rect(diameter, diameter, x0=x0, y0=y0)
+
+    @staticmethod
+    def rect(rows: int, cols: int, x0: float = 0,
+             y0: float = 0) -> List['TwoDQubit']:
+        """Returns a rectangle of TwoDQubit.
+
+        Args:
+            rows: Number of rows in the rectangle
+            cols: Number of columns in the rectangle
+            x0: x-coordinate of the first qubit
+            y0: y-coordinate of the first qubit
+
+        Returns:
+            A list of TwoDQubits filling in a rectangular grid
+        """
+        return [
+            TwoDQubit(x0 + x, y0 + y) for y in range(cols) for x in range(rows)
+        ]
+
+    @staticmethod
+    def triangular_lattice(l: int, x0: float = 0, y0: float = 0):
+        """Returns a triangular lattice of TwoDQubits.
+
+        Args:
+            l: Number of qubits along one direction
+            x0: x-coordinate of the first qubit
+            y0: y-coordinate of the first qubit
+
+        Returns:
+            A list of TwoDQubits filling in a triangular lattice.
+        """
+        coords = np.array([[x, y] for x in range(l + 1) for y in range(l + 1)],
+                          dtype=float)
+        coords[:, 0] += 0.5 * np.mod(coords[:, 1], 2)
+        coords[:, 1] *= np.sqrt(3) / 2
+        coords += [x0, y0]
+
+        return [TwoDQubit(coord[0], coord[1]) for coord in coords]
+
+    def __repr__(self):
+        return 'pasqal.TwoDQubit({}, {})'.format(self.x, self.y)
+
+    def __str__(self):
+        return '({}, {})'.format(self.x, self.y)
+
+    def _json_dict_(self):
+        return cirq.protocols.obj_to_dict_helper(self, ['x', 'y'])

cirq/pasqal/pasqal_qubits_test.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 import cirq
-from cirq.pasqal import ThreeDGridQubit
+from cirq.pasqal import ThreeDGridQubit, ThreeDQubit, TwoDQubit
 
 
 def test_pasqal_qubit_init():
@@ -219,3 +219,165 @@ def test_to_json():
         'col': 1,
         'lay': 1,
     }
+
+
+def test_pasqal_qubit_init_3D():
+    q = ThreeDQubit(3, 4, 5)
+    assert q.x == 3
+    assert q.y == 4
+    assert q.z == 5
+
+
+def test_comparison_key_3D():
+    assert ThreeDQubit(3, 4, 5)._comparison_key() == (5, 4, 3)
+    coords = (np.cos(np.pi / 2), np.sin(np.pi / 2), 0)
+    assert ThreeDQubit(*coords) == ThreeDQubit(0, 1, 0)
+
+
+def test_pasqal_qubit_ordering_3D():
+    assert ThreeDQubit(0, 0, 1) >= ThreeDQubit(1, 0, 0)
+    assert ThreeDQubit(0, 0, 1) >= ThreeDQubit(0, 1, 0)
+    assert ThreeDQubit(0, 1, 0) >= ThreeDQubit(1, 0, 0)
+    for i in range(8):
+        v = [int(x) for x in bin(i)[2:].zfill(3)]
+
+        assert ThreeDQubit(0, 0, 0) <= ThreeDQubit(v[0], v[1], v[2])
+        assert ThreeDQubit(1, 1, 1) >= ThreeDQubit(v[0], v[1], v[2])
+
+        if i >= 1:
+            assert ThreeDQubit(0, 0, 0) < ThreeDQubit(v[0], v[1], v[2])
+        if i < 7:
+            assert ThreeDQubit(1, 1, 1) > ThreeDQubit(v[0], v[1], v[2])
+
+
+def test_distance_3D():
+    with pytest.raises(TypeError):
+        _ = ThreeDQubit(0, 0, 0).distance(cirq.GridQubit(0, 0))
+
+    for x in np.arange(-2, 3):
+        for y in np.arange(-2, 3):
+            for z in np.arange(-2, 3):
+                assert ThreeDQubit(0, 0, 0).distance(ThreeDQubit(
+                    x, y, z)) == np.sqrt(x**2 + y**2 + z**2)
+
+
+def test_grid_qubit_eq_3D():
+    eq = cirq.testing.EqualsTester()
+    eq.make_equality_group(lambda: ThreeDQubit(0, 0, 0))
+    eq.make_equality_group(lambda: ThreeDQubit(1, 0, 0))
+    eq.make_equality_group(lambda: ThreeDQubit(0, 1, 0))
+    eq.make_equality_group(lambda: ThreeDQubit(50, 25, 25))
+
+
+def test_cube_3D():
+    assert ThreeDQubit.cube(2, x0=1, y0=1, z0=1) == [
+        ThreeDQubit(1, 1, 1),
+        ThreeDQubit(2, 1, 1),
+        ThreeDQubit(1, 2, 1),
+        ThreeDQubit(2, 2, 1),
+        ThreeDQubit(1, 1, 2),
+        ThreeDQubit(2, 1, 2),
+        ThreeDQubit(1, 2, 2),
+        ThreeDQubit(2, 2, 2)
+    ]
+    assert ThreeDQubit.cube(2) == [
+        ThreeDQubit(0, 0, 0),
+        ThreeDQubit(1, 0, 0),
+        ThreeDQubit(0, 1, 0),
+        ThreeDQubit(1, 1, 0),
+        ThreeDQubit(0, 0, 1),
+        ThreeDQubit(1, 0, 1),
+        ThreeDQubit(0, 1, 1),
+        ThreeDQubit(1, 1, 1),
+    ]
+
+
+def test_parrallelep_3D():
+    assert ThreeDQubit.parallelep(1, 2, 2, x0=5, y0=6, z0=7) == [
+        ThreeDQubit(5, 6, 7),
+        ThreeDQubit(5, 7, 7),
+        ThreeDQubit(5, 6, 8),
+        ThreeDQubit(5, 7, 8),
+    ]
+
+    assert ThreeDQubit.parallelep(2, 2, 2) == [
+        ThreeDQubit(0, 0, 0),
+        ThreeDQubit(1, 0, 0),
+        ThreeDQubit(0, 1, 0),
+        ThreeDQubit(1, 1, 0),
+        ThreeDQubit(0, 0, 1),
+        ThreeDQubit(1, 0, 1),
+        ThreeDQubit(0, 1, 1),
+        ThreeDQubit(1, 1, 1)
+    ]
+
+
+def test_square_2D():
+    assert TwoDQubit.square(2, x0=1, y0=1) == [
+        TwoDQubit(1, 1),
+        TwoDQubit(2, 1),
+        TwoDQubit(1, 2),
+        TwoDQubit(2, 2)
+    ]
+    assert TwoDQubit.square(2) == [
+        TwoDQubit(0, 0),
+        TwoDQubit(1, 0),
+        TwoDQubit(0, 1),
+        TwoDQubit(1, 1)
+    ]
+
+
+def test_rec_2D():
+    assert TwoDQubit.rect(1, 2, x0=5,
+                          y0=6) == [TwoDQubit(5, 6),
+                                    TwoDQubit(5, 7)]
+    assert TwoDQubit.rect(2, 2) == [
+        TwoDQubit(0, 0),
+        TwoDQubit(1, 0),
+        TwoDQubit(0, 1),
+        TwoDQubit(1, 1)
+    ]
+
+
+def test_triangular_2D():
+    assert TwoDQubit.triangular_lattice(1) == [
+        TwoDQubit(0.0, 0.0),
+        TwoDQubit(0.5, 0.8660254037844386),
+        TwoDQubit(1.0, 0.0),
+        TwoDQubit(1.5, 0.8660254037844386)
+    ]
+
+    assert TwoDQubit.triangular_lattice(1, x0=5., y0=6.1) == [
+        TwoDQubit(5.0, 6.1),
+        TwoDQubit(5.5, 6.966025403784438),
+        TwoDQubit(6.0, 6.1),
+        TwoDQubit(6.5, 6.966025403784438)
+    ]
+
+
+def test_repr_():
+    assert repr(ThreeDQubit(4, -25, 109)) == 'pasqal.ThreeDQubit(4, -25, 109)'
+    assert repr(TwoDQubit(4, -25)) == 'pasqal.TwoDQubit(4, -25)'
+
+
+def test_str_():
+    assert str(ThreeDQubit(4, -25, 109)) == '(4, -25, 109)'
+    assert str(TwoDQubit(4, -25)) == '(4, -25)'
+
+
+def test_to_json_():
+    q = ThreeDQubit(1.3, 1, 1)
+    d = q._json_dict_()
+    assert d == {
+        'cirq_type': 'ThreeDQubit',
+        'x': 1.3,
+        'y': 1,
+        'z': 1,
+    }
+    q = TwoDQubit(1.3, 1)
+    d = q._json_dict_()
+    assert d == {
+        'cirq_type': 'TwoDQubit',
+        'x': 1.3,
+        'y': 1,
+    }

cirq/protocols/json_serialization.py

@@ -148,7 +148,9 @@ def two_qubit_matrix_gate(matrix):
                 'SycamoreGate': cirq.google.SycamoreGate,
                 'TaggedOperation': cirq.TaggedOperation,
                 'ThreeDGridQubit': cirq.pasqal.ThreeDGridQubit,
+                'ThreeDQubit': cirq.pasqal.ThreeDQubit,
                 'TrialResult': cirq.TrialResult,
+                'TwoDQubit': cirq.pasqal.TwoDQubit,
                 'TwoQubitMatrixGate': two_qubit_matrix_gate,
                 'TwoQubitDiagonalGate': cirq.TwoQubitDiagonalGate,
                 '_UnconstrainedDevice':

cirq/protocols/json_test_data/ThreeDQubit.json

@@ -0,0 +1,6 @@
+{
+  "cirq_type": "ThreeDQubit",
+  "x": 10,
+  "y": 11,
+  "z": 12
+}

cirq/protocols/json_test_data/ThreeDQubit.repr

@@ -0,0 +1 @@
+cirq.pasqal.ThreeDQubit(10, 11, 12)

cirq/protocols/json_test_data/TwoDQubit.json

@@ -0,0 +1,5 @@
+{
+  "cirq_type": "TwoDQubit",
+  "x": 10,
+  "y": 11
+}

cirq/protocols/json_test_data/TwoDQubit.repr

@@ -0,0 +1 @@
+cirq.pasqal.TwoDQubit(10, 11)