"""Fundamental components related to quantum mechanics."""
from __future__ import annotations
import cmath
import enum
from abc import ABC, ABCMeta, abstractmethod
from enum import Enum, EnumMeta
from math import pi
from typing import TYPE_CHECKING, Literal, SupportsComplex, SupportsFloat, SupportsIndex, overload
import typing_extensions
# override introduced in Python 3.12
from typing_extensions import override
from graphix.parameter import Expression, cos_sin
from graphix.repr_mixins import EnumReprMixin
if TYPE_CHECKING:
from typing import TypeAlias
from graphix.parameter import ExpressionOrFloat
Angle: TypeAlias = float
"""In Graphix, angles are represented as floats and expressed in units of π."""
ParameterizedAngle: TypeAlias = Expression | Angle
ANGLE_PI: Angle = 1
"""The constant ``ANGLE_PI = 1`` is defined for convenience to make expressions involving angles more readable."""
@overload
def rad_to_angle(angle: float) -> Angle: ...
@overload
def rad_to_angle(angle: Expression) -> Expression: ...
def rad_to_angle(angle: ExpressionOrFloat) -> ParameterizedAngle:
"""Convert an angle expressed in radians to a Graphix angle.
In Graphix, angles are expressed in units of π.
"""
return angle / pi
@overload
def angle_to_rad(angle: Angle) -> float: ...
@overload
def angle_to_rad(angle: Expression) -> Expression: ...
def angle_to_rad(angle: ParameterizedAngle) -> ExpressionOrFloat:
"""Convert a Graphix angle to radians.
In Graphix, angles are expressed in units of π.
"""
return angle * pi
SupportsComplexCtor = SupportsComplex | SupportsFloat | SupportsIndex | complex
[docs]
class Sign(EnumReprMixin, Enum):
"""Sign, plus or minus."""
PLUS = 1
MINUS = -1
def __str__(self) -> str:
"""Return `+` or `-`."""
if self == Sign.PLUS:
return "+"
return "-"
[docs]
@staticmethod
def plus_if(b: bool) -> Sign:
"""Return *+* if *b* is *True*, *-* otherwise."""
if b:
return Sign.PLUS
return Sign.MINUS
[docs]
@staticmethod
def minus_if(b: bool) -> Sign:
"""Return *-* if *b* is *True*, *+* otherwise."""
if b:
return Sign.MINUS
return Sign.PLUS
def __neg__(self) -> Sign:
"""Swap the sign."""
return Sign.minus_if(self == Sign.PLUS)
@overload
def __mul__(self, other: Sign) -> Sign: ...
@overload
def __mul__(self, other: int) -> int: ...
@overload
def __mul__(self, other: float) -> float: ...
@overload
def __mul__(self, other: complex) -> complex: ...
def __mul__(self, other: Sign | complex) -> Sign | int | float | complex:
"""Multiply the sign with another sign or a number."""
if isinstance(other, Sign):
return Sign.plus_if(self == other)
if isinstance(other, int):
return int(self) * other
if isinstance(other, float):
return float(self) * other
if isinstance(other, complex):
return complex(self) * other
return NotImplemented
@overload
def __rmul__(self, other: int) -> int: ...
@overload
def __rmul__(self, other: float) -> float: ...
@overload
def __rmul__(self, other: complex) -> complex: ...
def __rmul__(self, other: complex) -> int | float | complex:
"""Multiply the sign with a number."""
if isinstance(other, (int, float, complex)):
return self.__mul__(other)
return NotImplemented
def __int__(self) -> int:
"""Return `1` for `+` and `-1` for `-`."""
# mypy does not infer the return type correctly
return self.value # type: ignore[no-any-return]
def __float__(self) -> float:
"""Return `1.0` for `+` and `-1.0` for `-`."""
return float(self.value)
def __complex__(self) -> complex:
"""Return `1.0 + 0j` for `+` and `-1.0 + 0j` for `-`."""
return complex(self.value)
[docs]
class ComplexUnit(EnumReprMixin, Enum):
"""
Complex unit: 1, -1, j, -j.
Complex units can be multiplied with other complex units,
with Python constants 1, -1, 1j, -1j, and can be negated.
"""
# HACK: complex(u) == (1j) ** u.value for all u in ComplexUnit.
ONE = 0
J = 1
MINUS_ONE = 2
MINUS_J = 3
[docs]
@staticmethod
def try_from(
value: ComplexUnit | SupportsComplexCtor, rel_tol: float = 1e-09, abs_tol: float = 0.0
) -> ComplexUnit | None:
"""Return the ComplexUnit instance if the value is compatible, None otherwise.
Parameters
----------
value : ComplexUnit | SupportsComplexCtor
Complex value to convert.
rel_tol : float, optional
Relative tolerance for comparing values, passed to :func:`math.isclose`. Default is ``1e-9``.
abs_tol : float, optional
Absolute tolerance for comparing values, passed to :func:`math.isclose`. Default is ``0.0``.
Returns
-------
ComplexUnit | None
Complex unit close to value, or ``None`` otherwise.
"""
if isinstance(value, ComplexUnit):
return value
value = complex(value)
for reference, result in (
(1, ComplexUnit.ONE),
(-1, ComplexUnit.MINUS_ONE),
(1j, ComplexUnit.J),
(-1j, ComplexUnit.MINUS_J),
):
if cmath.isclose(value, reference, rel_tol=rel_tol, abs_tol=abs_tol):
return result
return None
[docs]
@staticmethod
def from_properties(*, sign: Sign = Sign.PLUS, is_imag: bool = False) -> ComplexUnit:
"""Construct ComplexUnit from its properties."""
osign = 0 if sign == Sign.PLUS else 2
oimag = 1 if is_imag else 0
return ComplexUnit(osign + oimag)
@property
def sign(self) -> Sign:
"""Return the sign."""
return Sign.plus_if(self.value < 2)
@property
def is_imag(self) -> bool:
"""Return *True* if *j* or *-j*."""
return bool(self.value % 2)
def __complex__(self) -> complex:
"""Return the unit as complex number."""
ret: complex = 1j**self.value
return ret
def __str__(self) -> str:
"""Return a human-readable representation of the unit."""
result = "1j" if self.is_imag else "1"
if self.sign == Sign.MINUS:
result = "-" + result
return result
def __mul__(self, other: ComplexUnit | SupportsComplexCtor) -> ComplexUnit:
"""Multiply the complex unit with a number."""
if isinstance(other, ComplexUnit):
return ComplexUnit((self.value + other.value) % 4)
if isinstance(
other,
(SupportsComplex, SupportsFloat, SupportsIndex, complex),
) and (other_ := ComplexUnit.try_from(other)):
return self.__mul__(other_)
return NotImplemented
def __rmul__(self, other: SupportsComplexCtor) -> ComplexUnit:
"""Multiply the complex unit with a number."""
return self.__mul__(other)
def __neg__(self) -> ComplexUnit:
"""Return the opposite of the complex unit."""
return ComplexUnit((self.value + 2) % 4)
class CustomMeta(ABCMeta, EnumMeta):
"""Custom metaclass to allow multiple inheritance from `Enum` and `ABC`."""
class AbstractMeasurement(ABC):
"""Abstract base class for measurement objects.
Measurement objects are:
- :class:`graphix.measurements.Measurement`.
- :class:`graphix.fundamentals.Plane`.
- :class:`graphix.fundamentals.Axis`.
"""
@abstractmethod
def to_plane_or_axis(self) -> Plane | Axis:
"""Return the plane or axis of a measurement object.
Returns
-------
Plane | Axis
"""
# The parameters `rel_tol` and `abs_tol` are not used in the base
# implementation, but can be used in overrides.
def isclose(self, other: AbstractMeasurement, rel_tol: float = 1e-09, abs_tol: float = 0.0) -> bool: # noqa: ARG002
"""Determine whether this measurement is close to another.
Subclasses should implement a notion of “closeness” between two measurements, comparing measurement-specific attributes. The default comparison for ``float`` values involves checking equality within given relative or absolute tolerances.
Parameters
----------
other : AbstractMeasurement
The measurement to compare against.
rel_tol : float, optional
Relative tolerance for determining closeness. Relevant for comparing angles in the `Measurement` subclass. Default is ``1e-9``.
abs_tol : float, optional
Absolute tolerance for determining closeness. Relevant for comparing angles in the `Measurement` subclass. Default is ``0.0``.
Returns
-------
bool
``True`` if this measurement is considered close to ``other`` according
to the subclass's comparison rules; ``False`` otherwise.
"""
return self == other
class AbstractPlanarMeasurement(AbstractMeasurement):
"""Abstract base class for planar measurement objects.
Planar measurement objects are:
- :class:`graphix.measurements.Measurement`.
- :class:`graphix.fundamentals.Plane`.
"""
@abstractmethod
def to_plane(self) -> Plane:
"""Return the plane of a measurement object.
Returns
-------
Plane
"""
@override
def to_plane_or_axis(self) -> Plane:
return self.to_plane()
[docs]
class Axis(AbstractMeasurement, EnumReprMixin, Enum, metaclass=CustomMeta):
"""Axis: *X*, *Y* or *Z*."""
X = enum.auto()
Y = enum.auto()
Z = enum.auto()
[docs]
@override
def to_plane_or_axis(self) -> Axis:
return self
class SingletonI(Enum):
"""Singleton I."""
I = enum.auto()
I = SingletonI.I
IXYZ: TypeAlias = Literal[SingletonI.I] | Axis
IXYZ_VALUES: tuple[IXYZ, ...] = (I, Axis.X, Axis.Y, Axis.Z)
[docs]
class Plane(AbstractPlanarMeasurement, EnumReprMixin, Enum, metaclass=CustomMeta):
# TODO: Refactor using match
"""Plane: *XY*, *YZ* or *XZ*."""
XY = enum.auto()
YZ = enum.auto()
XZ = enum.auto()
@property
def axes(self) -> tuple[Axis, Axis]:
"""Return the pair of axes that carry the plane."""
match self:
case Plane.XY:
return (Axis.X, Axis.Y)
case Plane.YZ:
return (Axis.Y, Axis.Z)
case Plane.XZ:
return (Axis.X, Axis.Z)
case _:
typing_extensions.assert_never(self)
@property
def orth(self) -> Axis:
"""Return the axis orthogonal to the plane."""
match self:
case Plane.XY:
return Axis.Z
case Plane.YZ:
return Axis.X
case Plane.XZ:
return Axis.Y
case _:
typing_extensions.assert_never(self)
@property
def cos(self) -> Axis:
"""Return the axis of the plane that conventionally carries the cos."""
match self:
case Plane.XY:
return Axis.X
case Plane.YZ:
return Axis.Z # former convention was Y
case Plane.XZ:
return Axis.Z # former convention was X
case _:
typing_extensions.assert_never(self)
@property
def sin(self) -> Axis:
"""Return the axis of the plane that conventionally carries the sin."""
match self:
case Plane.XY:
return Axis.Y
case Plane.YZ:
return Axis.Y # former convention was Z
case Plane.XZ:
return Axis.X # former convention was Z
case _:
typing_extensions.assert_never(self)
@overload
def polar(self, angle: Angle) -> tuple[float, float, float]: ...
@overload
def polar(self, angle: Expression) -> tuple[Expression, Expression, Expression]: ...
[docs]
def polar(
self, angle: ParameterizedAngle
) -> tuple[float, float, float] | tuple[ExpressionOrFloat, ExpressionOrFloat, ExpressionOrFloat]:
"""Return the Cartesian coordinates of the point of module 1 at the given angle, following the conventional orientation for cos and sin."""
pp = (self.cos, self.sin)
# Angles are in units of π whereas `cos_sin` expects radians.
cos, sin = cos_sin(angle_to_rad(angle))
match pp:
case (Axis.X, Axis.Y):
return (cos, sin, 0)
case (Axis.Z, Axis.Y):
return (0, sin, cos)
case (Axis.Z, Axis.X):
return (sin, 0, cos)
case _:
raise RuntimeError("Unreachable.") # pragma: no cover
[docs]
@staticmethod
def from_axes(a: Axis, b: Axis) -> Plane:
"""Return the plane carried by the given axes."""
ab = {a, b}
if ab == {Axis.X, Axis.Y}:
return Plane.XY
if ab == {Axis.Y, Axis.Z}:
return Plane.YZ
if ab == {Axis.X, Axis.Z}:
return Plane.XZ
assert a == b
raise ValueError(f"Cannot make a plane giving the same axis {a} twice.")
[docs]
@override
def to_plane_or_axis(self) -> Plane:
"""Return the plane.
Returns
-------
Plane
"""
return self
[docs]
@override
def to_plane(self) -> Plane:
"""Return the plane.
Returns
-------
Plane
"""
return self