Source code for graphix.fundamentals

"""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