Coverage for jaxquantum/core/dims.py: 100%

1"""dims.""" 

2 

3from typing import List, Tuple 

4from copy import deepcopy 

5from math import prod 

6from jax import Array 

7 

8from enum import Enum 

9 

10 

11DIMS_TYPE = List[List[int]] 

12 

13 

14def isket_dims(dims: DIMS_TYPE) -> bool: 

15 return prod(dims[1]) == 1 

16 

17 

18def isbra_dims(dims: DIMS_TYPE) -> bool: 

19 return prod(dims[0]) == 1 

20 

21 

22def isop_dims(dims: DIMS_TYPE) -> bool: 

23 return prod(dims[1]) == prod(dims[0]) 

24 

25 

26def ket_from_op_dims(dims: DIMS_TYPE) -> DIMS_TYPE: 

27 return (dims[0], tuple([1 for _ in dims[1]])) 

28 

29 

30def check_dims(dims: Tuple[Tuple[int]], bdims: Tuple[int], data_shape: Array) -> bool: 

31 if len(data_shape) == 1 and data_shape[0] == 0: 

32 # E.g. empty list of operators 

33 assert bdims == (0,) 

34 assert dims == ((), ()) 

35 return 

36 

37 if isop_dims(dims) and not (isket_dims(dims) or isbra_dims(dims)): 

38 # Operator: the last two axes hold the matrix dimensions. 

39 assert bdims == data_shape[:-2], ( 

40 "Data shape should be consistent with dimensions." 

41 ) 

42 assert data_shape[-2] == prod(dims[0]), ( 

43 "Data shape should be consistent with dimensions." 

44 ) 

45 assert data_shape[-1] == prod(dims[1]), ( 

46 "Data shape should be consistent with dimensions." 

47 ) 

48 else: 

49 # Ket/bra: the Hilbert space lives on a single trailing axis (no (N,1)). 

50 # The trivial side of dims has product 1, so N == prod(dims[0])*prod(dims[1]). 

51 N = prod(dims[0]) * prod(dims[1]) 

52 assert bdims == data_shape[:-1], ( 

53 "Data shape should be consistent with dimensions." 

54 ) 

55 assert data_shape[-1] == N, ( 

56 "Data shape should be consistent with dimensions." 

57 ) 

58 

59 

60class Qdims: 

61 def __init__(self, dims): 

62 self._dims = deepcopy(dims) 

63 self._dims = (tuple(self._dims[0]), tuple(self._dims[1])) 

64 self._qtype = Qtypes.from_dims(self._dims) 

65 

66 @property 

67 def dims(self): 

68 return self._dims 

69 

70 @property 

71 def from_(self): 

72 return self._dims[1] 

73 

74 @property 

75 def to_(self): 

76 return self._dims[0] 

77 

78 @property 

79 def qtype(self): 

80 return self._qtype 

81 

82 def __str__(self): 

83 return str(self.dims) 

84 

85 def __repr__(self): 

86 return self.__str__() 

87 

88 def __eq__(self, other): 

89 return (self.dims == other.dims) and (self.qtype == other.qtype) 

90 

91 def __ne__(self, other): 

92 return (self.dims != other.dims) or (self.qtype != other.qtype) 

93 

94 def __hash__(self): 

95 return hash(self.dims) 

96 

97 def __matmul__(self, other): 

98 if self.from_ != other.to_: 

99 raise TypeError(f"incompatible dimensions {self} and {other}") 

100 

101 new_dims = [self.to_, other.from_] 

102 return Qdims(new_dims) 

103 

104 

105class Qtypes(str, Enum): 

106 ket = "ket" 

107 bra = "bra" 

108 oper = "oper" 

109 

110 @classmethod 

111 def from_dims(cls, dims: Array): 

112 if isket_dims(dims): 

113 return cls.ket 

114 if isbra_dims(dims): 

115 return cls.bra 

116 if isop_dims(dims): 

117 return cls.oper 

118 raise ValueError("Invalid data shape") 

119 

120 @classmethod 

121 def from_str(cls, string: str): 

122 return cls(string) 

123 

124 def __str__(self): 

125 return self.value 

126 

127 def __repr__(self): 

128 return self.__str__() 

129 

130 def __eq__(self, other): 

131 return self.value == other.value 

132 

133 def __ne__(self, other): 

134 return self.value != other.value 

135 

136 def __hash__(self): 

137 return hash(self.value)