Skip to content

channels

Functional direct-channel constructors and reusable kernels.

Channel(dims, apply, *, params=None, kraus=None, name='Channel')

Create a channel from a pure density-matrix kernel.

Source code in jaxquantum/circuits/channels.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def Channel(
    dims,
    apply: Callable,
    *,
    params=None,
    kraus=None,
    name="Channel",
):
    """Create a channel from a pure density-matrix kernel."""
    num_modes = 1 if isinstance(dims, int) else len(dims)
    return Gate.create(
        dims,
        name=name,
        params={} if params is None else params,
        gen_KM=_kraus_generator(kraus),
        channel_apply=apply,
        lazy_kraus=kraus is not None,
        num_modes=num_modes,
    )

ElementwiseChannel(dims, factor, *, params=None, kraus=None, name='ElementwiseChannel')

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

Source code in jaxquantum/circuits/channels.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def ElementwiseChannel(
    dims,
    factor,
    *,
    params=None,
    kraus=None,
    name="ElementwiseChannel",
):
    """Create ``rho[m,n] *= factor[m,n]`` channel."""
    params = dict(params or {})
    params["_factor"] = jnp.asarray(factor)
    return Channel(
        dims,
        apply_elementwise_channel,
        params=params,
        kraus=kraus,
        name=name,
    )

ShiftedChannel(dimension, coefficients, shifts, *, params=None, kraus=None, name='ShiftedChannel')

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

Source code in jaxquantum/circuits/channels.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def ShiftedChannel(
    dimension,
    coefficients,
    shifts: Sequence[int],
    *,
    params=None,
    kraus=None,
    name="ShiftedChannel",
):
    """Create a channel from output coefficients and input-index shifts."""
    coefficients = jnp.asarray(coefficients)
    shifts = jnp.asarray(shifts)
    if shifts.ndim != 1 or shifts.shape[0] == 0:
        raise ValueError("shifts must be a non-empty one-dimensional sequence")
    if coefficients.shape[-2:] != (shifts.shape[0], dimension):
        raise ValueError("coefficients must end in (num_shifts, dimension)")
    params = dict(params or {})
    params.update(
        {
            "_coefficients": coefficients,
            "_shifts": shifts,
        }
    )
    return Channel(
        dimension,
        apply_shifted_channel,
        params=params,
        kraus=kraus,
        name=name,
    )

apply_channel(channel, rho, axes=(-2, -1))

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

Source code in jaxquantum/circuits/channels.py 
69
70
71
72
73
74
75
76
77
78
79
80
def apply_channel(channel, rho, axes=(-2, -1)):
    """Apply a channel to density-matrix axes, with Kraus fallback."""
    input_ndim = rho.ndim
    rho = jnp.moveaxis(rho, axes, (-2, -1))
    if channel.channel_apply is not None:
        result = channel.channel_apply(rho[..., None, :, :], channel.params)
        result = jnp.squeeze(result, axis=-3)
    else:
        result = apply_kraus_map(channel.KM, rho)
    extra_dims = result.ndim - input_ndim
    output_axes = tuple(axis + extra_dims if axis >= 0 else axis for axis in axes)
    return jnp.moveaxis(result, (-2, -1), output_axes)

apply_elementwise_channel(rho, params)

Multiply trailing density-matrix axes by a channel factor.

Source code in jaxquantum/circuits/channels.py 
15
16
17
def apply_elementwise_channel(rho, params):
    """Multiply trailing density-matrix axes by a channel factor."""
    return rho * params["_factor"][..., None, :, :]

apply_kraus_map(kraus, rho)

Apply a dense Kraus stack whose leading axis indexes branches.

Source code in jaxquantum/circuits/channels.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def apply_kraus_map(kraus, rho):
    """Apply a dense Kraus stack whose leading axis indexes branches."""
    kraus = kraus.to_dense().data if isinstance(kraus, Qarray) else kraus
    if kraus.shape[0] == 0:
        return rho

    def branch(index):
        matrix = lax.dynamic_index_in_dim(kraus, index, 0, False)
        return matrix @ rho @ jnp.swapaxes(jnp.conj(matrix), -1, -2)

    if kraus.shape[0] == 1:
        return branch(0)
    return lax.fori_loop(
        1,
        kraus.shape[0],
        lambda index, total: total + branch(index),
        branch(0),
    )

apply_shifted_channel(rho, params)

Apply output-indexed shifted Kraus branches.

Source code in jaxquantum/circuits/channels.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def apply_shifted_channel(rho, params):
    """Apply output-indexed shifted Kraus branches."""
    coefficients = params["_coefficients"]
    shifts = params["_shifts"]
    indices = jnp.arange(rho.shape[-1])

    def branch(index):
        source = indices + shifts[index]
        valid = (source >= 0) & (source < rho.shape[-1])
        source = jnp.clip(source, 0, rho.shape[-1] - 1)
        shifted = rho[..., source[:, None], source[None, :]]
        coefficient = jnp.where(
            valid,
            coefficients[..., index, None, :],
            0,
        )
        return coefficient[..., :, None] * shifted * jnp.conj(coefficient[..., None, :])

    if jax.default_backend() != "cpu":
        return jax.vmap(branch)(jnp.arange(coefficients.shape[-2])).sum(axis=0)

    return lax.fori_loop(
        1,
        coefficients.shape[-2],
        lambda index, total: total + branch(index),
        branch(0),
    )