feat: finished jitting the easily jittable functions

src/qibojit/custom_operators/quantum_info.py

@@ -8,39 +8,6 @@
 
 ENGINE = qinfo.ENGINE  # this should be numpy
 
-SIGNATURES = {
-    # "_vectorization_row": (
-    #    ["c16[:,::1](c16[:,::1], i8)", "c16[:,::1](c16[:,:,::1], i8)"],
-    #    {"parallel": True, "cache": True},
-    # ),
-    # "_unvectorization_row": (
-    #    ["c16[:,:,::1](c16[:,::1], i8)", "c16[:,:,::1](c16[:,:,::1], i8)"],
-    #    {"parallel": True, "cache": True},
-    # ),
-    # "_random_hermitian": ("c16[:,:](i8)", {"parallel": True, "cache": True})
-    # "_vectorize_pauli_basis_row": ("c16[:,::1](i8, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1], f8)", {"parallel": True, "cache": True}),
-    # "_vectorize_pauli_basis_column": ("c16[:,::1](i8, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1], f8)", {"parallel": True, "cache": True}),
-}
-
-
-def jit_function(signature, function, **kwargs):
-    if isinstance(function, str):
-        function = getattr(qinfo, function)
-    return njit(signature, **kwargs)(function)
-
-
-class QinfoNumba:
-    pass
-
-
-QINFO = QinfoNumba()
-
-for function, signature in SIGNATURES.items():
-    print(function)
-    jitted = jit_function(signature[0], function, **signature[1])
-    globals()[function] = jitted
-    setattr(QINFO, function, jitted)
-
 
 @njit("c16[:,:,::1](i8, c16[:,:,:,::1])", parallel=True, cache=True)
 def _pauli_basis_inner(
@@ -49,9 +16,9 @@ def _pauli_basis_inner(
 ):
     dim = 2**nqubits
     basis = ENGINE.empty((len(prod), dim, dim), dtype=ENGINE.complex128)
-    for i in prange(len(prod)):  # pylint: disable=not-an-iterable
+    for i in prange(len(prod)):
         elem = prod[i][0]
-        for j in prange(1, len(prod[i])):  # pylint: disable=not-an-iterable
+        for j in prange(1, len(prod[i])):
             elem = ENGINE.kron(elem, prod[i][j])
         basis[i] = elem
     return basis
@@ -131,6 +98,8 @@ def _vectorization_column(state, dim):
 
 
 # dynamic tuple creation is not possible in numba
+# this might be jittable if we passed the shape
+# dim = (2,) * 2 * nqubits as inputs
 @njit
 def _vectorization_system(state, dim=0):
     nqubits = int(ENGINE.log2(state.shape[-1]))
@@ -163,6 +132,19 @@ def _unvectorization_column(state, dim):
     return state.T
 
 
+@njit("c16[:,:](c16[:,:], i8, i8)", parallel=True, cache=True)
+def _reshuffling(super_op, ax1: int, ax2: int):
+    dim = int(ENGINE.sqrt(super_op.shape[0]))
+    super_op = ENGINE.reshape(ENGINE.ascontiguousarray(super_op), (dim, dim, dim, dim))
+    axes = ENGINE.arange(super_op.ndim)
+    tmp = axes[ax1]
+    axes[ax1] = axes[ax2]
+    axes[ax2] = tmp
+    # axes[[ax1, ax2]] = axes[[ax2, ax1]]
+    super_op = ENGINE.transpose(super_op, to_fixed_tuple(axes, 4))
+    return ENGINE.reshape(ENGINE.ascontiguousarray(super_op), (dim**2, dim**2))
+
+
 @njit(
     [
         nbt.complex128[:](
@@ -290,6 +272,30 @@ def _pauli_to_comp_basis_sparse_row(
     return ENGINE.ascontiguousarray(unitary).reshape(unitary.shape[0], -1), nonzero[1]
 
 
+@njit(
+    nbt.Tuple((nbt.complex128[:, ::1], nbt.int64[:]))(
+        nbt.int64,
+        nbt.complex128[:, ::1],
+        nbt.complex128[:, ::1],
+        nbt.complex128[:, ::1],
+        nbt.complex128[:, ::1],
+        nbt.float64,
+    ),
+    parallel=True,
+    cache=True,
+)
+def _pauli_to_comp_basis_sparse_column(
+    nqubits: int, pauli_0, pauli_1, pauli_2, pauli_3, normalization: float = 1.0
+):
+    unitary = _vectorize_pauli_basis_column(
+        nqubits, pauli_0, pauli_1, pauli_2, pauli_3, normalization
+    )
+    unitary = unitary.T
+    nonzero = ENGINE.nonzero(unitary)
+    unitary = _array_at_2d_indices(unitary, nonzero)
+    return ENGINE.ascontiguousarray(unitary).reshape(unitary.shape[0], -1), nonzero[1]
+
+
 @njit(
     nbt.Tuple((nbt.complex128[:, :], nbt.complex128[:, :], nbt.float64[:, :, ::1]))(
         nbt.complex128[:, :]
@@ -445,6 +451,9 @@ def expm(A):
 """
 
 
+# if we can implement the expm in pure numba
+# we will be able to completely jit random unitary
+# and the other functions that depend on it
 def _random_unitary(dims: int):
     H = _random_hermitian(dims)
     return expm(-1.0j * H / 2)
@@ -461,6 +470,9 @@ def _random_density_matrix_bures_inner(
     return state / ENGINE.trace(state)
 
 
+# not entirely jittable because depends on random unitary
+
+
 def _random_density_matrix_bures(dims: int, rank: int, mean: float, stddev: float):
     unitary = _random_unitary(dims)
     return _random_density_matrix_bures_inner(unitary, dims, rank, mean, stddev)
@@ -612,6 +624,36 @@ def _super_op_from_bcsz_measure_column(dims: int, rank: int):
     return operator @ super_op @ operator
 
 
+@njit("c16[:,:](i8)", parallel=True, cache=True)
+def _super_op_from_haar_measure_row(dims: int):
+    super_op = _random_unitary_haar(dims)
+    super_op = _vectorization_row(super_op, dims)
+    return ENGINE.outer(super_op, ENGINE.conj(super_op))
+
+
+@njit("c16[:,:](i8)", parallel=True, cache=True)
+def _super_op_from_haar_measure_column(dims: int):
+    super_op = _random_unitary_haar(dims)
+    super_op = _vectorization_column(super_op, dims)
+    return ENGINE.outer(super_op, ENGINE.conj(super_op))
+
+
+# these two can't be jitted (at least globally) because of
+# random unitary
+
+
+def _super_op_from_hermitian_measure_row(dims: int):
+    super_op = _random_unitary(dims)
+    super_op = _vectorization_row(super_op, dims)
+    return ENGINE.outer(super_op, ENGINE.conj(super_op))
+
+
+def _super_op_from_hermitian_measure_column(dims: int):
+    super_op = _random_unitary(dims)
+    super_op = _vectorization_column(super_op, dims)
+    return ENGINE.outer(super_op, ENGINE.conj(super_op))
+
+
 @njit("c16[:,:](c16[:,:,::1], c16[:], i8)", parallel=True, cache=True)
 def _kraus_to_stinespring(kraus_ops, initial_state_env, dim_env: int):
     alphas = ENGINE.zeros((dim_env, dim_env, dim_env), dtype=initial_state_env.dtype)
@@ -645,18 +687,97 @@ def _stinespring_to_kraus(stinespring, initial_state_env, dim: int, dim_env: int
     return kraus.reshape(dim, dim_env, dim_env)
 
 
+@njit("c16[:,:](c16[:,:])", parallel=True, cache=True)
+def _to_choi_row(channel):
+    channel = _vectorization_row(channel, channel.shape[-1])
+    return ENGINE.outer(channel, ENGINE.conj(channel))
+
+
+@njit("c16[:,:](c16[:,:])", parallel=True, cache=True)
+def _to_choi_column(channel):
+    channel = _vectorization_column(channel, channel.shape[-1])
+    return ENGINE.outer(channel, ENGINE.conj(channel))
+
+
+@njit("c16[:,:](c16[:,:])", parallel=True, cache=True)
+def _to_liouville_row(channel):
+    channel = _to_choi_row(channel)
+    return _reshuffling(channel, 1, 2)
+
+
+@njit("c16[:,:](c16[:,:])", parallel=True, cache=True)
+def _to_liouville_column(channel):
+    channel = _to_choi_column(channel)
+    return _reshuffling(channel, 0, 3)
+
+
+@njit(
+    nbt.Tuple((nbt.complex128[:, :, :], nbt.float64[:]))(
+        nbt.float64[:], nbt.complex128[:, :], nbt.float64
+    ),
+    parallel=True,
+    cache=True,
+)
+def _choi_to_kraus_cp_row(eigenvalues, eigenvectors, precision: float):
+    eigv_gt_tol = ENGINE.abs(eigenvalues) > precision
+    coefficients = ENGINE.sqrt(eigenvalues[eigv_gt_tol])
+    eigenvectors = eigenvectors[eigv_gt_tol]
+    dim = int(ENGINE.sqrt(eigenvectors.shape[-1]))
+    kraus_ops = coefficients.reshape(-1, 1, 1) * _unvectorization_row(eigenvectors, dim)
+    return kraus_ops, coefficients
+
+
+@njit(
+    nbt.Tuple((nbt.complex128[:, :, :], nbt.float64[:]))(
+        nbt.float64[:], nbt.complex128[:, :], nbt.float64
+    ),
+    parallel=True,
+    cache=True,
+)
+def _choi_to_kraus_cp_column(eigenvalues, eigenvectors, precision: float):
+    eigv_gt_tol = ENGINE.abs(eigenvalues) > precision
+    coefficients = ENGINE.sqrt(eigenvalues[eigv_gt_tol])
+    eigenvectors = eigenvectors[eigv_gt_tol]
+    dim = int(ENGINE.sqrt(eigenvectors.shape[-1]))
+    kraus_ops = coefficients.reshape(-1, 1, 1) * _unvectorization_column(
+        eigenvectors, dim
+    )
+    return kraus_ops, coefficients
+
+
+@njit("c16[:,:](c16[:,:,:])", parallel=True, cache=True)
+def _kraus_to_choi_row(kraus_ops):
+    kraus_ops = _vectorization_row(kraus_ops, kraus_ops.shape[-1])
+    return kraus_ops.T @ ENGINE.conj(kraus_ops)
+
+
+@njit("c16[:,:](c16[:,:,:])", parallel=True, cache=True)
+def _kraus_to_choi_column(kraus_ops):
+    kraus_ops = _vectorization_column(kraus_ops, kraus_ops.shape[-1])
+    return kraus_ops.T @ ENGINE.conj(kraus_ops)
+
+
+class QinfoNumba:
+    pass
+
+
+QINFO = QinfoNumba()
+
+
 for function in (
     _pauli_basis,
     _vectorization_row,
     _vectorization_column,
     _unvectorization_row,
     _unvectorization_column,
+    _reshuffling,
     _post_sparse_pauli_basis_vectorization,
     _vectorize_pauli_basis_row,
     _vectorize_pauli_basis_column,
     _vectorize_sparse_pauli_basis_row,
     _vectorize_sparse_pauli_basis_column,
     _pauli_to_comp_basis_sparse_row,
+    _pauli_to_comp_basis_sparse_column,
     _choi_to_kraus_row,
     _choi_to_kraus_column,
     _random_statevector,
@@ -672,7 +793,42 @@ def _stinespring_to_kraus(stinespring, initial_state_env, dim: int, dim_env: int
     _gamma_delta_matrices,
     _super_op_from_bcsz_measure_row,
     _super_op_from_bcsz_measure_column,
+    _super_op_from_haar_measure_row,
+    _super_op_from_haar_measure_column,
+    _super_op_from_hermitian_measure_row,
+    _super_op_from_hermitian_measure_column,
     _kraus_to_stinespring,
     _stinespring_to_kraus,
+    _to_choi_row,
+    _to_choi_column,
+    _to_liouville_row,
+    _to_liouville_column,
+    _choi_to_kraus_cp_row,
+    _choi_to_kraus_cp_column,
+    _kraus_to_choi_row,
+    _kraus_to_choi_column,
 ):
     setattr(QINFO, function.__name__, function)
+
+
+# it would be quite cool and spare us a lot of code repetition if
+# we could make a recursive approach like the one below working
+
+SIGNATURES = {
+    # "_random_hermitian": ("c16[:,:](i8)", {"parallel": True, "cache": True})
+    # "_vectorize_pauli_basis_row": ("c16[:,::1](i8, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1], f8)", {"parallel": True, "cache": True}),
+    # "_vectorize_pauli_basis_column": ("c16[:,::1](i8, c16[:,::1], c16[:,::1], c16[:,::1], c16[:,::1], f8)", {"parallel": True, "cache": True}),
+}
+
+
+def jit_function(signature, function, **kwargs):
+    if isinstance(function, str):
+        function = getattr(qinfo, function)
+    return njit(signature, **kwargs)(function)
+
+
+for function, signature in SIGNATURES.items():
+    print(function)
+    jitted = jit_function(signature[0], function, **signature[1])
+    globals()[function] = jitted
+    setattr(QINFO, function, jitted)