"""Defines 2D and 3D data-structures storing vectors of integers.
This module defines several classes to store pairs (2D) or 3-tuples (3D) of
integers representing coordinates. They are used all over the code base to avoid
any coordinate system mess and to explicitly annotate each coordinate with its
significance. Basically, instead of having
.. code-block:: python
coords = (0, 4) # Is it (y, x) or (x, y)?
# 0 used below, but if the coords tuple was obtained from a call to
# `.shape` from a numpy array it should rather be a 1.
x = coords[0]
we have
.. code-block:: python
coords = Position2D(0, 4)
x = coords.x
This is particularly useful when we, as humans, are mostly used to always have
(x, y) coordinates but some libraries (such as numpy) reverse that order for
indexing.
"""
from __future__ import annotations
from dataclasses import astuple, dataclass
from enum import Enum
import numpy as np
import numpy.typing as npt
@dataclass(frozen=True, order=True)
class Vec2D:
x: int
y: int
@dataclass(frozen=True, order=True)
class Vec3D:
x: int
y: int
z: int
class Position2D(Vec2D):
"""Represents a position on a 2-dimensional plane.
Warning:
This class represents a position without any knowledge of the coordinate
system being used. As such, it should only be used when the coordinate
system is meaningless or in localised places where the coordinate system
is obvious. In particular, this class should be avoided in interfaces.
"""
def with_block_coordinate_system(self) -> BlockPosition2D:
return BlockPosition2D(self.x, self.y)
[docs]
class PhysicalQubitPosition2D(Position2D):
"""Represents the position of a physical qubit on a 2-dimensional plane."""
[docs]
class PlaquettePosition2D(Position2D):
"""Represents the position of a plaquette on a 2-dimensional plane."""
[docs]
def get_origin_position(self, displacement: Shift2D) -> PhysicalQubitPosition2D:
"""Returns the position of the plaquette origin."""
return PhysicalQubitPosition2D(displacement.x * self.x, displacement.y * self.y)
[docs]
class BlockPosition2D(Position2D):
"""Represents the position of a block on a 2-dimensional plane."""
[docs]
def get_top_left_plaquette_position(
self, block_shape: Shape2D
) -> PlaquettePosition2D:
"""Returns the position of the top-left plaquette of the block."""
return PlaquettePosition2D(block_shape.x * self.x, block_shape.y * self.y)
[docs]
class Shape2D(Vec2D):
[docs]
def to_numpy_shape(self) -> tuple[int, int]:
"""Returns the shape according to numpy indexing.
In the coordinate system used in this library, numpy indexes
arrays using (y, x) coordinates. This method is here to
translate a Shape instance to a numpy shape transparently for
the user.
"""
return (self.y, self.x)
[docs]
class Shift2D(Vec2D):
def __mul__(self, factor: int) -> Shift2D:
return Shift2D(factor * self.x, factor * self.y)
def __rmul__(self, factor: int) -> Shift2D:
return self.__mul__(factor)
[docs]
class Position3D(Vec3D):
"""A 3D integer position."""
x: int
y: int
z: int
[docs]
def shift_by(self, dx: int = 0, dy: int = 0, dz: int = 0) -> Position3D:
"""Shift the position by the given offset."""
return Position3D(self.x + dx, self.y + dy, self.z + dz)
[docs]
def shift_in_direction(self, direction: Direction3D, shift: int) -> Position3D:
"""Shift the position in the given direction by the given shift."""
if direction == Direction3D.X:
return self.shift_by(dx=shift)
elif direction == Direction3D.Y:
return self.shift_by(dy=shift)
else:
return self.shift_by(dz=shift)
[docs]
def is_neighbour(self, other: Position3D) -> bool:
"""Check if the other position is near to this position, i.e. Manhattan
distance is 1."""
return (
abs(self.x - other.x) + abs(self.y - other.y) + abs(self.z - other.z) == 1
)
[docs]
def as_tuple(self) -> tuple[int, int, int]:
"""Return the position as a tuple."""
return astuple(self)
def __str__(self) -> str:
return f"({self.x},{self.y},{self.z})"
[docs]
def as_2d(self) -> Position2D:
"""Return the position as a 2D position."""
return Position2D(self.x, self.y)
[docs]
class Direction3D(Enum):
"""Axis directions in the 3D spacetime diagram."""
X = 0
Y = 1
Z = 2
[docs]
@staticmethod
def all_directions() -> list[Direction3D]:
"""Return all the directions."""
return [Direction3D.X, Direction3D.Y, Direction3D.Z]
def __str__(self) -> str:
return self.name
[docs]
@dataclass(frozen=True)
class SignedDirection3D:
"""Signed directions in the 3D spacetime diagram."""
direction: Direction3D
towards_positive: bool
def __neg__(self) -> SignedDirection3D:
return SignedDirection3D(self.direction, not self.towards_positive)
def __str__(self) -> str:
return f"{self.direction}{'+' if self.towards_positive else '-'}"
@dataclass(frozen=True, order=True)
class FloatPosition3D:
"""A 3D float position."""
x: float
y: float
z: float
def shift_by(self, dx: float = 0, dy: float = 0, dz: float = 0) -> FloatPosition3D:
"""Shift the position by the given offset."""
return FloatPosition3D(self.x + dx, self.y + dy, self.z + dz)
def shift_in_direction(
self, direction: Direction3D, shift: float
) -> FloatPosition3D:
"""Shift the position in the given direction by the given shift."""
if direction == Direction3D.X:
return self.shift_by(dx=shift)
elif direction == Direction3D.Y:
return self.shift_by(dy=shift)
else:
return self.shift_by(dz=shift)
def as_array(self) -> npt.NDArray[np.float32]:
"""Return the position as a numpy array."""
return np.asarray(astuple(self), dtype=np.float32)