"""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
import re
from dataclasses import astuple, dataclass
from enum import Enum
import numpy as np
import numpy.typing as npt
from typing_extensions import Self
from tqec.utils.exceptions import TQECException
@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)
def is_neighbour(self, other: Position2D) -> 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) == 1
def to_3d(self, z: int = 0) -> Position3D:
return Position3D(self.x, self.y, z)
[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, shift: Shift2D) -> PhysicalQubitPosition2D:
"""Returns the position of the plaquette origin."""
return PhysicalQubitPosition2D(shift.x * self.x, shift.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: PlaquetteShape2D
) -> PlaquettePosition2D:
"""Returns the position of the top-left plaquette of the block."""
return PlaquettePosition2D(block_shape.x * self.x, block_shape.y * self.y)
class Shape2D(Vec2D):
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)
class PlaquetteShape2D(Shape2D):
"""Represents a 2-dimensional shape using plaquette coordinate system."""
class PhysicalQubitShape2D(Shape2D):
"""Represents a 2-dimensional shape using physical qubit coordinate system."""
[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) -> Self:
"""Shift the position by the given offset."""
return self.__class__(self.x + dx, self.y + dy, self.z + dz)
[docs]
def shift_in_direction(self, direction: Direction3D, shift: int) -> Self:
"""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)
class BlockPosition3D(Position3D):
"""Represents the position of a block in 3D space."""
def as_2d(self) -> BlockPosition2D:
return BlockPosition2D(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]
[docs]
@staticmethod
def spatial_directions() -> list[Direction3D]:
return [Direction3D.X, Direction3D.Y]
[docs]
@staticmethod
def temporal_directions() -> list[Direction3D]:
return [Direction3D.Z]
def __str__(self) -> str:
return self.name
[docs]
@staticmethod
def from_neighbouring_positions(
source: Position3D, sink: Position3D
) -> Direction3D:
assert source.is_neighbour(sink)
for direction, (source_coord, sink_coord) in zip(
Direction3D.all_directions(), zip(source.as_tuple(), sink.as_tuple())
):
if source_coord != sink_coord:
return direction
raise TQECException(
"Could not find the direction from two neighbouring positions "
f"{source:=} and {sink:=}."
)
@property
def orthogonal_directions(self) -> tuple[Direction3D, Direction3D]:
i = self.value
return Direction3D((i + 1) % 3), Direction3D((i + 2) % 3)
[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"{'+' if self.towards_positive else '-'}{self.direction}"
[docs]
@staticmethod
def from_string(s: str) -> SignedDirection3D:
"""Return the signed direction from a string.
Args:
s: The string representation of the signed direction. The string
should have the format "<sign><direction>", where "<direction>"
is one of "X", "Y", "Z" and "<sign>" is either "+" or "-".
Returns:
The signed direction.
Raises:
TQECException: If the string does not match the expected format.
"""
match = re.match(r"([+-])([XYZ])", s)
if match is None:
raise TQECException(
f"Invalid signed direction: {s}, expected format: [+/-][XYZ]"
)
sign, direction = match.groups()
return SignedDirection3D(Direction3D("XYZ".index(direction)), sign == "+")
@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)