Skip to content

Math

bosonic

per(mtx, column, selected, prod, output=False)

Row expansion for the permanent of matrix mtx. The counter column is the current column, selected is a list of indices of selected rows, and prod accumulates the current product.

Source code in src/squint/math/bosonic.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def per(mtx, column, selected, prod, output=False):
    """
    Row expansion for the permanent of matrix mtx.
    The counter column is the current column,
    selected is a list of indices of selected rows,
    and prod accumulates the current product.
    """
    if column == mtx.shape[1]:
        if output:
            print(selected, prod)
        return prod
    else:
        result = 0
        for row in range(mtx.shape[0]):
            if not row in selected:
                result = result + per(
                    mtx, column + 1, selected + [row], prod * mtx[row, column]
                )
        return result

permanent(mat)

Returns the permanent of the matrix mat.

Source code in src/squint/math/bosonic.py 
41
42
43
44
45
def permanent(mat):
    """
    Returns the permanent of the matrix mat.
    """
    return per(mat, 0, [], 1)

get_fixed_sum_tuples(length, total)

Generate all tuples of a given length that sum to a specified total.

Source code in src/squint/math/bosonic.py 
48
49
50
51
52
53
54
55
56
def get_fixed_sum_tuples(length, total):
    """Generate all tuples of a given length that sum to a specified total."""
    if length == 1:
        yield (total,)
        return

    for i in range(total + 1):
        for t in get_fixed_sum_tuples(length - 1, total - i):
            yield (i,) + t

compile_Aij_indices(i_s: jnp.array, j_s: jnp.array, m: int, n: int)

Compile all indices for generating the \(A_{ij}\) matrices for all i and j combinations.

Source code in src/squint/math/bosonic.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def compile_Aij_indices(i_s: jnp.array, j_s: jnp.array, m: int, n: int):
    """Compile all indices for generating the $A_{ij}$ matrices for all i and j combinations."""
    # checkify.check(
    #     jnp.all(i_s.sum(axis=1) == n), f"Some input bases do not have n={n} photons."
    # )
    # checkify.check(
    #     jnp.all(j_s.sum(axis=1) == n), f"Some output bases do not have n={n} photons."
    # )

    unitary_inds = jnp.indices((m, m))

    def repeated_indices(i_basis: jnp.array, j_basis: jnp.array):
        rectangular = jnp.concat(
            [
                einops.repeat(
                    unitary_inds[:, :, i : i + 1],
                    "ind row col -> ind row (rep col)",
                    rep=i_basis[i],
                )
                for i in range(m)
            ],
            axis=2,
        )

        square = jnp.concat(
            [
                einops.repeat(
                    rectangular[:, i : i + 1, :],
                    "ind row col -> ind (rep row) col",
                    rep=j_basis[i],
                )
                for i in range(m)
            ],
            axis=1,
        )

        return square

    transition_inds = jnp.array(
        [[repeated_indices(i_basis, j_basis) for j_basis in j_s] for i_basis in i_s]
    )
    return transition_inds

compute_transition_amplitudes(unitary: jnp.array, transition_inds: jnp.array)

Calculates all i -> j transition amplitudes in a jit-able manner.

Source code in src/squint/math/bosonic.py 
103
104
105
106
107
108
109
110
111
112
@jax.jit
def compute_transition_amplitudes(unitary: jnp.array, transition_inds: jnp.array):
    """Calculates all i -> j transition amplitudes in a jit-able manner."""
    a_ijs = unitary[transition_inds[:, :, 0, :, :], transition_inds[:, :, 1, :, :]]

    # swapping axes required when using recursive permanent function
    a_ijs_swapaxes = einops.rearrange(a_ijs, "i o a b -> a b i o")
    transition_amplitudes = permanent(a_ijs_swapaxes)  # fastest after jit of the three

    return transition_amplitudes

gellmann

The code for the gellman function is adapted from the PySME project, which is licensed under the MIT license.

Source: https://pysme.readthedocs.io/en/latest/_modules/gellmann.html .. module:: gellmann.py :synopsis: Generate generalized Gell-Mann matrices .. moduleauthor:: Jonathan Gross jarthurgross@gmail.com

Functions to generate the generalized Pauli (i.e., Gell-Mann matrices)

gellmann(j, k, d)

Returns a generalized Gell-Mann matrix of dimension d. According to the convention in Bloch Vectors for Qubits by Bertlmann and Krammer (2008), returns :math:\Lambda^j for :math:1\leq j=k\leq d-1, :math:\Lambda^{kj}_s for :math:1\leq k<j\leq d, :math:\Lambda^{jk}_a for :math:1\leq j<k\leq d, and :math:I for :math:j=k=d.

:param j: First index for generalized Gell-Mann matrix :type j: positive integer :param k: Second index for generalized Gell-Mann matrix :type k: positive integer :param d: Dimension of the generalized Gell-Mann matrix :type d: positive integer :returns: A genereralized Gell-Mann matrix. :rtype: numpy.array

Source code in src/squint/math/gellmann.py 
19
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def gellmann(j, k, d):
    r"""Returns a generalized Gell-Mann matrix of dimension d. According to the
    convention in *Bloch Vectors for Qubits* by Bertlmann and Krammer (2008),
    returns :math:`\Lambda^j` for :math:`1\leq j=k\leq d-1`,
    :math:`\Lambda^{kj}_s` for :math:`1\leq k<j\leq d`,
    :math:`\Lambda^{jk}_a` for :math:`1\leq j<k\leq d`, and
    :math:`I` for :math:`j=k=d`.

    :param j: First index for generalized Gell-Mann matrix
    :type j:  positive integer
    :param k: Second index for generalized Gell-Mann matrix
    :type k:  positive integer
    :param d: Dimension of the generalized Gell-Mann matrix
    :type d:  positive integer
    :returns: A genereralized Gell-Mann matrix.
    :rtype:   numpy.array

    """

    if j > k:
        gjkd = jnp.zeros((d, d), dtype=jnp.complex64)
        gjkd = gjkd.at[j - 1, k - 1].set(1.0)
        gjkd = gjkd.at[k - 1, j - 1].set(1.0)
    elif k > j:
        gjkd = jnp.zeros((d, d), dtype=jnp.complex64)
        gjkd = gjkd.at[j - 1, k - 1].set(-1.0j)
        gjkd = gjkd.at[k - 1, j - 1].set(1.0j)
    elif j == k and j < d:
        gjkd = jnp.sqrt(2 / (j * (j + 1))) * jnp.diag(
            jnp.array(
                [
                    1 + 0.0j if n <= j else (-j + 0.0j if n == (j + 1) else 0 + 0.0j)
                    for n in range(1, d + 1)
                ],
                dtype=jnp.complex64,
            )
        )
    else:
        gjkd = jnp.diag(
            jnp.array([1 + 0.0j for n in range(1, d + 1)], dtype=jnp.complex64)
        )

    return gjkd

information_matrices

qfim(psi: Array, dspi: Array)

Computes the quantum Fisher information matrix from the already computed arrays representing the probability amplitudes and their gradients.

Parameters:

Name Type Description Default
psi Array

Quantum amplitudes.

 required
dpsi Array

Gradients of the quantum amplitudes.

 required

Returns:

Name Type Description
qfim ndarray

Quantum Fisher information matrix.
Source code in src/squint/math/information_matrices.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def qfim(
    psi: Array,
    dspi: Array,
):
    """
    Computes the quantum Fisher information matrix from the already computed arrays representing
    the probability amplitudes and their gradients.

    Args:
        psi (Array): Quantum amplitudes.
        dpsi (Array): Gradients of the quantum amplitudes.

    Returns:
        qfim (jnp.ndarray): Quantum Fisher information matrix.
    """
    dpsi_conj = jnp.conjugate(dspi)
    return 4 * jnp.real(
        jnp.real(jnp.einsum("i..., j... -> ij", dpsi_conj, dspi))
        + jnp.einsum(
            "i,j->ij",
            jnp.einsum("i..., ... -> i", dpsi_conj, psi),
            jnp.einsum("j..., ... -> j", dpsi_conj, psi),
        )
    )

quantum_fisher_information_matrix(_forward_amplitudes: Callable, _grad_amplitudes: Callable, *params: PyTree)

Performs the forward pass to compute quantum amplitudes and their gradients, and then calculates the quantum Fisher information matrix. Args: _forward_amplitudes (Callable): Function to compute quantum amplitudes. _grad_amplitudes (Callable): Function to compute gradients of quantum amplitudes. *params (list[PyTree]): Parameters for the quantum circuit, partitioned via eqx.partition. The argnum is already defined in the callables Returns: qfim (jnp.ndarray): Quantum Fisher information matrix.

Source code in src/squint/math/information_matrices.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def quantum_fisher_information_matrix(
    _forward_amplitudes: Callable,
    _grad_amplitudes: Callable,
    # get: Callable,
    *params: PyTree,
):
    """
    Performs the forward pass to compute quantum amplitudes and their gradients,
    and then calculates the quantum Fisher information matrix.
    Args:
        _forward_amplitudes (Callable): Function to compute quantum amplitudes.
        _grad_amplitudes (Callable): Function to compute gradients of quantum amplitudes.
        *params (list[PyTree]): Parameters for the quantum circuit, partitioned via `eqx.partition`.
            The argnum is already defined in the callables
    Returns:
        qfim (jnp.ndarray): Quantum Fisher information matrix."""
    amplitudes = _forward_amplitudes(*params)
    grads, _ = jax.tree.flatten(_grad_amplitudes(*params))
    grads = jnp.stack(grads, axis=0)
    return qfim(amplitudes, grads)

cfim(p: Array, dp: Array)

Computes the classical Fisher information matrix from the already computed arrays representing the probabilities and their gradients. Args: p (Array): Classical probabilities. dp (Array): Gradients of the classical probabilities. Returns: cfim (jnp.ndarray): Classical Fisher information matrix.

Source code in src/squint/math/information_matrices.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def cfim(
    p: Array,
    dp: Array,
):
    """
    Computes the classical Fisher information matrix from the already computed arrays representing
    the probabilities and their gradients.
    Args:
        p (Array): Classical probabilities.
        dp (Array): Gradients of the classical probabilities.
    Returns:
        cfim (jnp.ndarray): Classical Fisher information matrix.
    """

    return jnp.einsum(
        "i..., j..., ... -> ij",
        dp,
        dp,
        1
        / (p[None, ...] + 1e-14),  # add a small constant to avoid division by zero
    )

classical_fisher_information_matrix(_forward_prob: Callable, _grad_prob: Callable, *params: PyTree)

Performs the forward pass to compute classical probabilities and their gradients, and then calculates the classical Fisher information matrix. Args: _forward_prob (Callable): Function to compute classical probabilities. _grad_prob (Callable): Function to compute gradients of classical probabilities. *params (list[PyTree]): Parameters for the quantum circuit, partitioned via eqx.partition. The argnum is already defined in the callables Returns: cfim (jnp.ndarray): Classical Fisher information matrix.

Source code in src/squint/math/information_matrices.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def classical_fisher_information_matrix(
    _forward_prob: Callable,
    _grad_prob: Callable,
    # get: Callable,
    *params: PyTree,
):
    """
    Performs the forward pass to compute classical probabilities and their gradients,
    and then calculates the classical Fisher information matrix.
    Args:
        _forward_prob (Callable): Function to compute classical probabilities.
        _grad_prob (Callable): Function to compute gradients of classical probabilities.
        *params (list[PyTree]): Parameters for the quantum circuit, partitioned via `eqx.partition`.
            The argnum is already defined in the callables
    Returns:
        cfim (jnp.ndarray): Classical Fisher information matrix.
    """
    probs = _forward_prob(*params)
    grads, _ = jax.tree.flatten(_grad_prob(*params))
    grads = jnp.stack(grads, axis=0)
    return cfim(probs, grads)