Coverage for jaxquantum/circuits/channels.py: 91%

1"""Functional direct-channel constructors and reusable kernels.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable, Sequence 

6 

7import jax 

8import jax.numpy as jnp 

9from jax import lax 

10 

11from jaxquantum.circuits.gates import Gate 

12from jaxquantum.core.qarray import Qarray 

13 

14 

15def apply_elementwise_channel(rho, params): 

16 """Multiply trailing density-matrix axes by a channel factor.""" 

17 return rho * params["_factor"][..., None, :, :] 

18 

19 

20def apply_shifted_channel(rho, params): 

21 """Apply output-indexed shifted Kraus branches.""" 

22 coefficients = params["_coefficients"] 

23 shifts = params["_shifts"] 

24 indices = jnp.arange(rho.shape[-1]) 

25 

26 def branch(index): 

27 source = indices + shifts[index] 

28 valid = (source >= 0) & (source < rho.shape[-1]) 

29 source = jnp.clip(source, 0, rho.shape[-1] - 1) 

30 shifted = rho[..., source[:, None], source[None, :]] 

31 coefficient = jnp.where( 

32 valid, 

33 coefficients[..., index, None, :], 

34 0, 

35 ) 

36 return coefficient[..., :, None] * shifted * jnp.conj(coefficient[..., None, :]) 

37 

38 if jax.default_backend() != "cpu": 

39 return jax.vmap(branch)(jnp.arange(coefficients.shape[-2])).sum(axis=0) 

40 

41 return lax.fori_loop( 

42 1, 

43 coefficients.shape[-2], 

44 lambda index, total: total + branch(index), 

45 branch(0), 

46 ) 

47 

48 

49def apply_kraus_map(kraus, rho): 

50 """Apply a dense Kraus stack whose leading axis indexes branches.""" 

51 kraus = kraus.to_dense().data if isinstance(kraus, Qarray) else kraus 

52 if kraus.shape[0] == 0: 

53 return rho 

54 

55 def branch(index): 

56 matrix = lax.dynamic_index_in_dim(kraus, index, 0, False) 

57 return matrix @ rho @ jnp.swapaxes(jnp.conj(matrix), -1, -2) 

58 

59 if kraus.shape[0] == 1: 

60 return branch(0) 

61 return lax.fori_loop( 

62 1, 

63 kraus.shape[0], 

64 lambda index, total: total + branch(index), 

65 branch(0), 

66 ) 

67 

68 

69def apply_channel(channel, rho, axes=(-2, -1)): 

70 """Apply a channel to density-matrix axes, with Kraus fallback.""" 

71 input_ndim = rho.ndim 

72 rho = jnp.moveaxis(rho, axes, (-2, -1)) 

73 if channel.channel_apply is not None: 

74 result = channel.channel_apply(rho[..., None, :, :], channel.params) 

75 result = jnp.squeeze(result, axis=-3) 

76 else: 

77 result = apply_kraus_map(channel.KM, rho) 

78 extra_dims = result.ndim - input_ndim 

79 output_axes = tuple(axis + extra_dims if axis >= 0 else axis for axis in axes) 

80 return jnp.moveaxis(result, (-2, -1), output_axes) 

81 

82 

83def _kraus_generator(kraus): 

84 if kraus is None: 

85 return None 

86 if callable(kraus): 

87 

88 def generate(params): 

89 result = kraus(params) 

90 return result if isinstance(result, Qarray) else Qarray.from_list(result) 

91 

92 return generate 

93 kraus = kraus if isinstance(kraus, Qarray) else Qarray.from_list(kraus) 

94 return lambda _: kraus 

95 

96 

97def Channel( 

98 dims, 

99 apply: Callable, 

100 *, 

101 params=None, 

102 kraus=None, 

103 name="Channel", 

104): 

105 """Create a channel from a pure density-matrix kernel.""" 

106 num_modes = 1 if isinstance(dims, int) else len(dims) 

107 return Gate.create( 

108 dims, 

109 name=name, 

110 params={} if params is None else params, 

111 gen_KM=_kraus_generator(kraus), 

112 channel_apply=apply, 

113 lazy_kraus=kraus is not None, 

114 num_modes=num_modes, 

115 ) 

116 

117 

118def ElementwiseChannel( 

119 dims, 

120 factor, 

121 *, 

122 params=None, 

123 kraus=None, 

124 name="ElementwiseChannel", 

125): 

126 """Create ``rho[m,n] *= factor[m,n]`` channel.""" 

127 params = dict(params or {}) 

128 params["_factor"] = jnp.asarray(factor) 

129 return Channel( 

130 dims, 

131 apply_elementwise_channel, 

132 params=params, 

133 kraus=kraus, 

134 name=name, 

135 ) 

136 

137 

138def ShiftedChannel( 

139 dimension, 

140 coefficients, 

141 shifts: Sequence[int], 

142 *, 

143 params=None, 

144 kraus=None, 

145 name="ShiftedChannel", 

146): 

147 """Create a channel from output coefficients and input-index shifts.""" 

148 coefficients = jnp.asarray(coefficients) 

149 shifts = jnp.asarray(shifts) 

150 if shifts.ndim != 1 or shifts.shape[0] == 0: 

151 raise ValueError("shifts must be a non-empty one-dimensional sequence") 

152 if coefficients.shape[-2:] != (shifts.shape[0], dimension): 

153 raise ValueError("coefficients must end in (num_shifts, dimension)") 

154 params = dict(params or {}) 

155 params.update( 

156 { 

157 "_coefficients": coefficients, 

158 "_shifts": shifts, 

159 } 

160 ) 

161 return Channel( 

162 dimension, 

163 apply_shifted_channel, 

164 params=params, 

165 kraus=kraus, 

166 name=name, 

167 )