QuCumber: wavefunction reconstruction with neural networks

TL;DR

QuCumber is an open-source machine learning tool that reconstructs quantum states from measurement data using neural networks, enabling efficient analysis of large qubit systems and generation of new measurements.

Contribution

The paper introduces QuCumber, a novel software package employing restricted Boltzmann machines for quantum wavefunction reconstruction from experimental data.

Findings

Successfully reconstructs quantum states from measurement data.

Generates new measurements to access additional physical observables.

Scales efficiently to large qubit systems.

Abstract

As we enter a new era of quantum technology, it is increasingly important to develop methods to aid in the accurate preparation of quantum states for a variety of materials, matter, and devices. Computational techniques can be used to reconstruct a state from data, however the growing number of qubits demands ongoing algorithmic advances in order to keep pace with experiments. In this paper, we present an open-source software package called QuCumber that uses machine learning to reconstruct a quantum state consistent with a set of projective measurements. QuCumber uses a restricted Boltzmann machine to efficiently represent the quantum wavefunction for a large number of qubits. New measurements can be generated from the machine to obtain physical observables not easily accessible from the original data.

Figures

Equations

$${\psi }_{\mathbf{\lambda }}(\mathrm{x})\equiv \sqrt{p_{\mathbf{\lambda }}(\mathrm{x})}\simeq \Psi (\mathrm{x}).$$

$$\hat{H}=-J\sum _i{\hat{\sigma }}_i^z{\hat{\sigma }}_{i+1}^z-h\sum _i{\hat{\sigma }}_i^x,$$

$$\text{expectationvalue}\ast \hat{\mathcal{O}}=\frac{1}{{\sum }_{\mathbf{\sigma }}|{\psi }_{\mathbf{\lambda }}(\mathbf{\sigma }){|}^2}\sum _{\mathbf{\sigma }}{\mathcal{O}}_{\mathbf{\sigma }}|{\psi }_{\mathbf{\lambda }}(\mathbf{\sigma }){|}^2.$$

$$\text{expectationvalue}\ast \hat{\mathcal{O}}\approx \frac{1}{N_{\mathrm{MC}}}\sum _{k=1}^{N_{\mathrm{MC}}}{\mathcal{O}}_{{\mathbf{\sigma }}_k},$$

$$\text{expectationvalue}\ast \hat{\mathcal{O}}=\frac{1}{{\sum }_{\mathbf{\sigma }}|{\psi }_{\mathbf{\lambda }}(\mathbf{\sigma }){|}^2}\sum _{{\mathbf{\sigma \sigma }}^{\prime }}{\psi }_{\mathbf{\lambda }}^{\ast }(\mathbf{\sigma }){\psi }_{\mathbf{\lambda }}({\mathbf{\sigma }}^{\prime }){\mathcal{O}}_{{\mathbf{\sigma \sigma }}^{\prime }}.$$

$$\text{expectationvalue}\ast \hat{\mathcal{O}}\approx \frac{1}{N_{\mathrm{MC}}}\sum _{k=1}^{N_{\mathrm{MC}}}{\mathcal{O}}_{{\mathbf{\sigma }}_k}^{[L]}$$

$${\mathcal{O}}_{{\mathbf{\sigma }}_k}^{[L]}=\sum _{{\mathbf{\sigma }}^{\prime }}\frac{{\psi }_{\mathbf{\lambda }}({\mathbf{\sigma }}^{\prime })}{{\psi }_{\mathbf{\lambda }}({\mathbf{\sigma }}_k)}{\mathcal{O}}_{{\mathbf{\sigma }}_k{\mathbf{\sigma }}^{\prime }}.$$

$$\text{matrixelement}\mathbf{\sigma }{\hat{\sigma }}_j^x{\mathbf{\sigma }}^{\prime }={\delta }_{{\sigma }_j^{\prime },1-{\sigma }_j}\prod _{i\ne j}{\delta }_{{\sigma }_i^{\prime },{\sigma }_j}.$$

$${({\sigma }_j^x)}^{[L]}=\frac{{\psi }_{\mathbf{\lambda }}({\sigma }_1,\dots ,1-{\sigma }_j,\dots ,{\sigma }_N)}{{\psi }_{\mathbf{\lambda }}({\sigma }_1,\dots ,{\sigma }_j,\dots ,{\sigma }_N)}.$$

$$S_2=-\ln ⟨{\widehat{\text{Swap}}}_A⟩.$$

$${\widehat{\text{Swap}}}_A|{\mathbf{\sigma }}_A,{\mathbf{\sigma }}_{\bar{A}}{⟩}_1\otimes |{\mathbf{\sigma }}_A^{\prime },{\mathbf{\sigma }}_{\bar{A}}^{\prime }{⟩}_2=|{\mathbf{\sigma }}_A^{\prime },{\mathbf{\sigma }}_{\bar{A}}{⟩}_1\otimes |{\mathbf{\sigma }}_A,{\mathbf{\sigma }}_{\bar{A}}^{\prime }{⟩}_2.$$

$${\psi }_{\mathbf{\lambda }\mathbf{\mu }}(\mathbf{\sigma })=\sqrt{p_{\mathbf{\lambda }}(\mathbf{\sigma })}{e}^{i{\varphi }_{\mathbf{\mu }}(\mathbf{\sigma })/2}$$

$$\begin{array}{rl}\Psi ({\mathbf{\sigma }}^{\mathbf{b}})& =⟨{\mathbf{\sigma }}^{\mathbf{b}}|\Psi ⟩=\sum _{\mathbf{\sigma }}⟨{\mathbf{\sigma }}^{\mathbf{b}}|\mathbf{\sigma }⟩⟨\mathbf{\sigma }|\Psi ⟩\\ & =\sum _{\mathbf{\sigma }}\mathcal{U}({\mathbf{\sigma }}^{\mathbf{b}},\mathbf{\sigma })\Psi (\mathbf{\sigma }),\end{array}$$

$$P_{\bm{b}}(\bm{\sigma}^{\bm{b}})=\bigg{|}\sum_{\bm{\sigma}}\bm{\mathcal{U}}(\bm{\sigma}^{\bm{b}},\bm{\sigma})\Psi(\bm{\sigma})\bigg{|}^{2}\>.$$

$$|\Psi ⟩=\sum _{{\sigma }_0,{\sigma }_1}{\Phi }_{{\sigma }_0{\sigma }_1}{e}^{i{\theta }_{{\sigma }_0{\sigma }_1}}|{\sigma }_0{\sigma }_1⟩,$$

$$\hat{\mathrm{H}}=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}1& 1\\ 1& -1\end{array}\right]$$

$$P_{\mathbf{b}}({\sigma }_0^x,{\sigma }_1)=\frac{{\Phi }_{0{\sigma }_1}^2+{\Phi }_{1{\sigma }_1}^2}{4}+\frac{1-2{\sigma }_0^x}{2}{\Phi }_{0{\sigma }_1}{\Phi }_{1{\sigma }_1}\cos (\Delta \theta ),$$

$$\hat{\mathrm{K}}=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}1& -i\\ 1& i\end{array}\right].$$

$$(b_0,b_1)=(\mathrm{z},\mathrm{z}),(\mathrm{x},\mathrm{z}),(\mathrm{z},\mathrm{x}),(\mathrm{y},\mathrm{z}),(\mathrm{z},\mathrm{y})$$

$$E_{\mathbf{\lambda }}(\mathbf{v},\mathbf{h})=-\sum _{j=1}^{n_v}{b}_j{v}_j-\sum _{i=1}^{n_h}{c}_i{h}_i-\sum _{ij}{h}_i{W}_{ij}{v}_j.$$

$${\mathcal{E}}_{\mathbf{\lambda }}(\mathbf{v})=-\sum _{j=1}^{n_v}{b}_j{v}_j-\sum _{i=1}^{n_h}\log \left[1+\exp \left(\sum _j^{n_v}{W}_{ij}{v}_j+c_i\right)\right].$$

$$p_{\mathbf{\lambda }}(\mathbf{v},\mathbf{h})=\frac{1}{Z_{\mathbf{\lambda }}}{e}^{-E_{\mathbf{\lambda }}(\mathbf{v},\mathbf{h})}.$$

$$\mathrm{KL}(P||Q)=\sum _{\mathbf{v}}P(\mathbf{v})\log \frac{P(\mathbf{v})}{Q(\mathbf{v})}.$$

$$p_{\mathbf{\lambda }}(\mathbf{v})=\frac{1}{Z_{\mathbf{\lambda }}}\sum _{\mathbf{h}}{e}^{-E_{\mathbf{\lambda }}(\mathbf{v},\mathbf{h})}=\frac{1}{Z_{\mathbf{\lambda }}}{e}^{-{\mathcal{E}}_{\mathbf{\lambda }}(\mathbf{v})}.$$

$$Z_{\mathbf{\lambda }}=\sum _{\mathbf{v},\mathbf{h}}{e}^{-E_{\mathbf{\lambda }}(\mathbf{v},\mathbf{h})}.$$

Full text

QuCumber: wavefunction reconstruction with neural networks

Matthew J. S. Beach1,2, Isaac De Vlugt2, Anna Golubeva1,2, Patrick Huembeli1,3, Bohdan Kulchytskyy1,2, Xiuzhe Luo2, Roger G. Melko1,2*, Ejaaz Merali2, Giacomo Torlai1,2,4

${\bf 1}$ Perimeter Institute for Theoretical Physics, Waterloo, Ontario N2L 2Y5, Canada

${\bf 2}$ Department of Physics and Astronomy, University of Waterloo, Ontario N2L 3G1, Canada

${\bf 3}$ ICFO-Institut de Ciencies Fotoniques, Barcelona Institute of Science and Technology, 08860 Castelldefels (Barcelona), Spain

${\bf 4}$ Center for Computational Quantum Physics, Flatiron Institute, 162 5th Avenue, New York, NY 10010, USA

• rgmelko@uwaterloo.ca

Abstract

**As we enter a new era of quantum technology, it is increasingly important to develop methods to aid in the accurate preparation of quantum states for a variety of materials, matter, and devices. Computational techniques can be used to reconstruct a state from data, however the growing number of qubits demands ongoing algorithmic advances in order to keep pace with experiments. In this paper, we present an open-source software package called QuCumber that uses machine learning to reconstruct a quantum state consistent with a set of projective measurements. QuCumber uses a restricted Boltzmann machine to efficiently represent the quantum wavefunction for a large number of qubits. New measurements can be generated from the machine to obtain physical observables not easily accessible from the original data. **

Contents

1. 1 Introduction

2. 2 Positive wavefunctions

3. 2.1 Setup

4. 2.2 Training

5. 2.3 Reconstruction of physical observables

6. 2.3.1 Diagonal observables

7. 2.3.2 Off-diagonal observables

8. 2.3.3 Entanglement entropy

9. 3 Complex wavefunctions

10. 3.1 Setup

11. 3.2 Training

12. 4 Conclusion

13. A Glossary

1 Introduction

Current advances in fabricating quantum technologies, as well as in reliable control of synthetic quantum matter, are leading to a new era of quantum hardware where highly pure quantum states are routinely prepared in laboratories. With the growing number of controlled quantum degrees of freedom, such as superconducting qubits, trapped ions, and ultracold atoms [1, 2, 3, 4], reliable and scalable classical algorithms are required for the analysis and verification of experimentally prepared quantum states. Efficient algorithms can aid in extracting physical observables otherwise inaccessible from experimental measurements, as well as in identifying sources of noise to provide direct feedback for improving experimental hardware. However, traditional approaches for reconstructing unknown quantum states from a set of measurements, such as quantum state tomography, often suffer the exponential overhead that is typical of quantum many-body systems.

Recently, an alternative path to quantum state reconstruction was put forward, based on modern machine learning (ML) techniques [5, 6, 7, 8, 9, 10]. The most common approach relies on a powerful generative model called a restricted Boltzmann machine (RBM) [11], a stochastic neural network with two layers of binary units. A visible layer $\bm{v}$ describes the physical degrees of freedom, while a hidden layer $\bm{h}$ is used to capture high-order correlations between the visible units. Given a set of neural network parameters $\bm{\lambda}$, the RBM defines a probabilistic model described by the parametric distribution $p_{\bm{\lambda}}(\bm{v})$. RBMs have been widely used in the ML community for the pre-training of deep neural networks [12], for compressing high-dimensional data into lower-dimensional representations [13], and more [14]. More recently, RBMs have been adopted by the physics community in the context of representing both classical and quantum many-body states [15, 16]. They are currently being investigated for their representational power [17, 18, 19], their relationship with tensor networks and the renormalization group [20, 21, 22, 23, 24], and in other contexts in quantum many-body physics [25, 26, 27].

In this post, we present QuCumber: a quantum calculator used for many-body eigenstate reconstruction. QuCumber is an open-source Python package that implements neural-network quantum state reconstruction of many-body wavefunctions from projective measurement data. Examples of data to which QuCumber could be applied might be magnetic spin projections, orbital occupation number, polarization of photons, or the logical state of qubits. Given a training set of such measurements, QuCumber discovers the most likely compatible quantum state by finding the optimal set of parameters $\bm{\lambda}$ of an RBM. A properly trained RBM is an approximation of the unknown quantum state underlying the data. It can be used to calculate various physical observables of interest, including measurements that may not be possible in the original experiment.

This post is organized as follows. In Section 2, we introduce the reconstruction technique for the case where all coefficients of the wavefunction are real and positive. We discuss the required format for input data, as well as training of the RBM and the reconstruction of typical observables. In Section 3, we consider the more general case of a complex-valued wavefunction. We illustrate a general strategy to extract the phase structure from data by performing appropriate unitary rotations on the state before measurements. We then demonstrate a practical reconstruction of an entangled state of two qubits. Note, the detailed theory underlying the reconstruction methods used by QuCumber can be found in the original references [5, 6] and a recent review [28]. A list of useful terms and equations can be found at the end of this post in the Glossary.

2 Positive wavefunctions

We begin by presenting the application of QuCumber to reconstruct many-body quantum states described by wavefunctions $\ket{\Psi}$ with positive coefficients $\Psi(\bm{\mathrm{x}})=\innerproduct{\bm{\mathrm{x}}}{\Psi}\geq 0$, where $\ket{\bm{\mathrm{x}}}=\ket{\mathrm{x}_{1},\dots,\mathrm{x}_{N}}$ is a reference basis for the Hilbert space of $N$ quantum degrees of freedom. The neural-network quantum state reconstruction requires raw data $\mathcal{D}=(\bm{\mathrm{x}}_{1},\bm{\mathrm{x}}_{2},\dots)$ generated through projective measurements of the state $\ket{\Psi}$ in the reference basis. These measurements adhere to the probability distribution given by the Born rule, $P(\bm{\mathrm{x}})=|\Psi(\bm{\mathrm{x}})|^{2}$. Since the wavefunction is strictly positive, the quantum state is completely characterized by the measurement distribution, i.e. $\Psi(\bm{\mathrm{x}})=\sqrt{P(\bm{\mathrm{x}})}$.

The positivity of the wavefunction allows a simple and natural connection between quantum states and classical probabilistic models. QuCumber employs the probability distribution $p_{\bm{\lambda}}(\bm{\mathrm{x}})$ of an RBM (see Eq. 24 of the Glossary) to approximate the distribution $P(\bm{\mathrm{x}})$ underlying the measurement data. Using contrastive divergence (CD) [12], QuCumber trains the RBM to discover an optimal set of parameters ${\bm{\lambda}}$ that minimize the Kullback-Leibler (KL) divergence between the two distributions (see Eq. 23). Upon successful training ($p_{\bm{\lambda}}(\bm{\mathrm{x}})\sim P(\bm{\mathrm{x}})$), we obtain an approximate representation of the target quantum state,

[TABLE]

Note, the precise mathematical form of the marginal distribution $p_{\bm{\lambda}}(\bm{\mathrm{x}})$ defined in terms of an effective energy over the parameters of the RBM is defined in the Glossary.

In the following, we demonstrate the application of QuCumber for the reconstruction of the ground-state wavefunction of the one-dimensional transverse-field Ising model (TFIM). The Hamiltonian is

[TABLE]

where $\hat{\sigma}^{x/z}_{i}$ are spin-1/2 Pauli operators acting on site $i$, and we assume open boundary conditions. For this example, we consider a chain with $N=10$ spins at the quantum critical point $J=h=1$.

2.1 Setup

Given the small size of the system, the ground state $\ket{\Psi}$ can be found with exact diagonalization. The training dataset $\mathcal{D}$ is generated by sampling the distribution $P(\bm{\sigma}^{z})=|\Psi(\bm{\sigma}^{z})|^{2}$, obtaining a sequence of $N_{S}=10^{5}$ independent spin projections in the reference basis $\bm{\mathrm{x}}=\bm{\sigma}^{z}$.111 The training dataset can be downloaded from [https://github.com/PIQuIL/QuCumber/blob/master/

examples/Tutorial1_TrainPosRealWaveFunction/tfim1d_data.txt](https://github.com/PIQuIL/QuCumber/blob/master/examples/Tutorial1_TrainPosRealWaveFunction/tfim1d_data.txt)

Each data point in $\mathcal{D}$ consists of an array $\bm{\sigma}^{z}_{j}=(\sigma^{z}_{1},\dots,\sigma^{z}_{N})$ with shape (N,) and should be passed to QuCumber as a numpy array or torch tensor. For example, $\bm{\sigma}^{z}_{j}=$ np.array([1,0,1,1,0,1,0,0,0,1]), where we use $\sigma_{j}^{z}=0,1$ to represent a spin-down and spin-up state respectively. Therefore, the entire input data set is contained in an array with shape (N_S, N).

Aside from the training data, QuCumber also allows us to import an exact wavefunction. This can be useful for monitoring the quality of the reconstruction during training. In our example, we evaluate the fidelity between the reconstructed state $\psi_{\bm{\lambda}}(\bm{\mathrm{x}})$ and the exact wavefunction $\Psi(\bm{\mathrm{x}})$. The training dataset, train_data, and the exact ground state, true_psi, are loaded with the data loading utility as follows:

⬇

train_path = "tfim1d_data.txt"

psi_path = "tfim1d_psi.txt"

train_data, true_psi = data.load_data(train_path, psi_path)

If psi_path is not provided, QuCumber will load only the training data.

Next, we initialize an RBM quantum state $\psi_{\bm{\lambda}}(\bf{x})$ with random weights and zero biases using the constructor PositiveWaveFunction:

⬇

state = PositiveWaveFunction(num_visible=10, num_hidden=10)

The number of visible units (num_visible) must be equal to the number of physical spins $N$, while the number of hidden units (num_hidden) can be adjusted to systematically increase the representational power of the RBM.

The quality of the reconstruction will depend on the structure underlying the specific quantum state and the ratio of visible to hidden units, $\alpha=\verb|num_hidden|/\verb|num_visible|$. In practice, we find that $\alpha=1$ often leads to good approximations of positive wavefunctions [6]. However, in the general case, the value of $\alpha$ required for a given wavefunction should be explored and adjusted by the user.

2.2 Training

Once an appropriate representation of the quantum state has been defined, QuCumber trains the RBM through the function PositiveWaveFunction.fit. Several input parameters need to be provided aside from the training dataset (train_data). These include the number of training iterations (epochs), the number of samples used for the positive/negative phase of CD (pos_batch_size/neg_batch_size), the learning rate (lr) and the number of sampling steps in the negative phase of CD (k). The last argument (callbacks) allows the user to pass a set of additional functions to be evaluated during training.

As an example of a callback, we show the MetricEvaluator, which evaluates a function log_every epochs during training. Given the small system size and the knowledge of the true ground state, we can evaluate the fidelity between the RBM state and the true ground-state wavefunction (true_psi). Similarly, we can calculate the KL divergence between the RBM distribution $p_{\bm{\lambda}}(\bm{\mathrm{x}})$, and the data distribution $P(\bm{\mathrm{x}})$, which should approach zero for a properly trained RBM. For the current example, we monitor the fidelity and KL divergence (defined in qucumber.utils.training_statistics):

⬇

import qucumber.utils.training_statistics as ts

log_every = 10

space = state.generate_hilbert_space(10)

callbacks = [

MetricEvaluator(

    log_every,

    {"Fidelity": ts.fidelity, "KL": ts.KL},

    target_psi=true_psi,

    space=space,

    verbose=True

)

]

With verbose=True, the program will print the epoch number and all callbacks every log_every epochs. For the current example, we monitor the fidelity and KL divergence (see Glossary). Note that the KL divergence is only tractable for small systems. The MetricEvaluator will compute the KL exactly when provided with a list of all states in the Hilbert space. For convenience these can be generated with

space = state.generate_hilbert_space(10).

Now that the metrics to monitor during training have been chosen, we can invoke the optimization with the fit function of PositiveWaveFunction.

⬇

train_data,

epochs=500,

pos_batch_size=100,

neg_batch_size=100,

lr=0.01,

k=5,

callbacks=callbacks,

)

Figure 1 shows the convergence of the fidelity and KL divergence during training. The convergence time will, in general, depend on the choice of hyperparameters. Finally, the network parameters $\bm{\lambda}$, together with the MetricEvaluator’s data, can be saved (or loaded) to a file:

⬇

"filename.pt",

metadata={

    "fidelity": callbacks[0].Fidelity,

    "KL": callbacks[0].KL

},

)

state.load("filename.pt")

With this we have demonstrated the most basic aspects of QuCumber regarding training a model and verifying its accuracy. We note that in this example the evaluation utilized the knowledge of the exact ground state and the calculation of the KL divergence, which we reemphasize is tractable only for small system sizes. However, we point out that QuCumber is capable of carrying out the reconstruction of much larger systems. In such cases, users must rely on other estimators to evaluate the training, such as expectation values of physical observables (magnetization, energy, etc). In the following, we show how to compute diagonal and off-diagonal observables in QuCumber.

2.3 Reconstruction of physical observables

In this section, we discuss how to calculate the average value of a generic physical observable $\hat{\mathcal{O}}$ from a trained RBM. We start with the case of observables that are diagonal in the reference basis where the RBM was trained. We then discuss the more general cases of off-diagonal observables and entanglement entropies.

2.3.1 Diagonal observables

We begin by considering an observable with only diagonal matrix elements, $\bra{\bm{\sigma}}\hat{\mathcal{O}}\ket{\bm{\sigma}^{\prime}}=\mathcal{O}_{\bm{\sigma}}\delta_{\bm{\sigma\sigma}^{\prime}}$ where for convenience we denote the reference basis ${\bf x}=\bm{\sigma}^{z}$ as $\bm{\sigma}$ unless otherwise stated. The expectation value of $\hat{\mathcal{O}}$ is given by

[TABLE]

The expectation value can be approximated by a Monte Carlo estimator,

[TABLE]

where the spin configurations $\bm{\sigma}_{k}$ are sampled from the RBM distribution $p_{\bm{\lambda}}(\bm{\sigma})$. This process is particularly efficient given the bipartite structure of the network which allows the use of block Gibbs sampling.

A simple example for the TFIM is the average longitudinal magnetization per spin, $\expectationvalue{\hat{\sigma}^{z}}=\sum_{j}\expectationvalue*{\hat{\sigma}^{z}_{j}}/N$, which can be calculated directly on the spin configuration sampled by the RBM (i.e., the state of the visible layer). The visible samples are obtained with the sample function of the RBM state object:

which takes the total number of samples (num_samples) and the number of block Gibbs steps (k) as input. Once these samples are obtained, the magnetization can be calculated simply as

where we converted the binary samples of the RBM back into $\pm 1$ spins before taking the mean.

2.3.2 Off-diagonal observables

We turn now to the case of off-diagonal observables, where the expectation value assumes the following form

[TABLE]

This expression can once again be approximated with a Monte Carlo estimator

[TABLE]

of the so-called local estimator of the observable:

[TABLE]

As long as the matrix representation $\mathcal{O}_{\bm{\sigma\sigma}^{\prime}}$ is sufficiently sparse in the reference basis, the summation can be evaluated efficiently.

As an example, we consider the specific case of the transverse magnetization for the $j$-th spin, $\expectationvalue*{\hat{\sigma}^{x}_{j}}$, with matrix elements

[TABLE]

Therefore, the expectation values reduces to the Monte Carlo average of the local observable

[TABLE]

evaluated on spin configurations $\bm{\sigma}_{k}$ sampled from the RBM distribution $p_{\bm{\lambda}}(\bm{\sigma})$.

QuCumber provides an interface for sampling off-diagonal observables in the ObservableBase class. Thorough examples are available in the tutorial section in the documentation.​​ 222The observables tutorial is available at [https://qucumber.readthedocs.io/en/stable/_examples

/Tutorial3_DataGeneration_CalculateObservables/tutorial_sampling_observables.html](https://qucumber.readthedocs.io/en/stable/_examples/Tutorial3_DataGeneration_CalculateObservables/tutorial_sampling_observables.html)

As an example, $\sigma^{x}$ can be written as an observable class with

⬇

from qucumber.utils import cplx

from qucumber.observables import ObservableBase

class SigmaX(ObservableBase):

def apply(self, nn_state, samples):

    psi = nn_state.psi(samples)

    psi_ratio_sum = torch.zeros_like(psi)

    for i in range(samples.shape[-1]):  # sum over spin sites

        flip_spin(i, samples)  # flip the spin at site i

        # add ratio psi_(-i) / psi to the running sum

        psi_flip = nn_state.psi(samples)

        psi_ratio = cplx.elementwise_division(psi_flip, psi)

        psi_ratio_sum.add_(psi_ratio)

        flip_spin(i, samples)  # flip it back

    # take real part and divide by number of spins

    return psi_ratio_sum[0].div_(samples.shape[-1])

The value of the observable can be estimated from a set of samples with:

which produces a dictionary containing the mean, variance, and standard error of the observable. Similarly, the user can define other observables like the energy.

The reconstruction of two magnetic observables for the TFIM is shown in Fig. 2, where a different RBM was trained for each value of the transverse field $h$. In the left plot, we show the average longitudinal magnetization per site, which can be calculated directly from the configurations sampled by the RBM. In the right plot, we show the off-diagonal observable of transverse magnetization. In both cases, QuCumber successfully discovers an optimal set of parameters $\bm{\lambda}$ that accurately approximate the ground-state wavefunction underlying the data.

2.3.3 Entanglement entropy

A quantity of significant interest in quantum many-body systems is the degree of entanglement between a subregion $A$ and its complement $\bar{A}$. Numerically, measurement of bipartite entanglement entropy is commonly accessed through the computation of the second Rényi entropy $S_{2}=-\ln{\rm Tr}(\rho_{A}^{2})$. When one has access to a pure state wavefunction $\psi_{\bm{\lambda}}(\bm{\mathrm{x}})$, Rényi entropies can be calculated as an expectation value of the “Swap” operator [29],

[TABLE]

It is essentially an off-diagonal observable that acts on an extended product space consisting of two independent copies of the wavefunction, $\psi_{\bm{\lambda}}(\bm{\mathrm{x}})\otimes\psi_{\bm{\lambda}}(\bm{\mathrm{x}})$, referred to as “replicas”. As the name suggests, the action of the Swap operator is to swap the spin configurations in region $A$ between the replicas,

[TABLE]

Here the subcript of the ket indicates the replica index, while the two labels inside a ket, such as ${\bm{\sigma}}_{A},{\bm{\sigma}}_{\bar{A}}$, describe the spins configurations within the subregion and its complement.

In QuCumber, the Swap operator is implemented as a routine within the entanglement observable unit,

⬇

_s = s1[:, A].clone()

s1[:, A] = s2[:, A]

s2[:, A] = _s

return s1, s2

where s1 and s2 are batches of samples produced from each replica, and A is a list containing the indices of the sites in subregion $A$. While ideally those samples should be entirely independent, in order to save computational costs, QuCumber just splits a given batch into two equal parts and treats them as if they were independent samples. This is implemented within the SWAP observable,

⬇

def __init__(self, A):

    self.A = A

def apply(self, nn_state, samples):

    _ns = samples.shape[0] // 2

    samples1 = samples[:_ns, :]

    samples2 = samples[_ns : _ns * 2, :]

    psi_ket1 = nn_state.psi(samples1)

    psi_ket2 = nn_state.psi(samples2)

    psi_ket = cplx.elementwise_mult(psi_ket1, psi_ket2)

    psi_ket_star = cplx.conjugate(psi_ket)

    samples1_, samples2_ = swap(samples1, samples2, self.A)

    psi_bra1 = nn_state.psi(samples1_)

    psi_bra2 = nn_state.psi(samples2_)

    psi_bra = cplx.elementwise_mult(psi_bra1, psi_bra2)

    psi_bra_star = cplx.conjugate(psi_bra)

    return cplx.real(

        cplx.elementwise_division(psi_bra_star, psi_ket_star)

    )

Note the similarity in the implementation to that for the transverse magnetization observable from the last section, once the amplitude of a sample is substituted with the product of amplitudes drawn from each replica.

Using this observable, we can estimate the Rényi entropy of the region containing the first 5 sites in the chain using Eq. 10,

⬇

swap_ = SWAP(A)

swap_stats = swap_.statistics_from_samples(state, samples)

S_2 = -np.log(swap_stats["mean"])

We apply this measurement procedure to a TFIM chain with results shown in Fig. 3. As was the case with the magnetization observables, the trained RBM gives a good approximation to the second Rényi entropy for different subregion $A$ sizes. Being a basis-independent observable, this constitutes a useful test on the ability of QuCumber to capture the full wavefunction from the information contained in a single-basis dataset for TFIM.

3 Complex wavefunctions

For positive wavefunctions, the probability distribution underlying the outcomes of projective measurements in the reference basis contains all possible information about the unknown quantum state. However, in the more general case of a wavefunction with a non-trivial sign or phase structure, this is not the case. In this section, we consider a target quantum state where the wavefunction coefficients in the reference basis can be complex-valued, $\Psi(\bm{\sigma})=\Phi(\bm{\sigma})e^{i\theta(\bm{\sigma})}$. We continue to choose the reference basis as $\bm{\sigma}=\bm{\sigma}^{z}$. We first need to generalize the RBM representation of the quantum state to capture a generic complex wavefunction. To this end, we introduce an additional RBM with marginalized distribution $p_{\bm{\mu}}(\bm{\sigma})$ parameterized by a new set of network weights and biases $\bm{\mu}$. We use this to define the quantum state as:

[TABLE]

where $\phi_{\bm{\mu}}(\bm{\sigma})=\log(p_{\bm{\mu}}(\bm{\sigma}))$ [6]. In this case, the reconstruction requires a different type of measurement setting. It is easy to see that projective measurements in the reference basis do not convey any information on the phases $\theta(\bm{\sigma})$, since $P(\bm{\sigma})=|\Psi(\bm{\sigma})|^{2}=\Phi^{2}(\bm{\sigma})$.

The general strategy to learn a phase structure is to apply a unitary transformation $\bm{\mathcal{U}}$ to the state $\ket{\Psi}$ before the measurements, such that the resulting measurement distribution $P^{\>\prime}(\bm{\sigma})=|\Psi^{\prime}(\bm{\sigma})|^{2}$ of the rotated state $\Psi^{\prime}(\bm{\sigma})=\bra{\bm{\sigma}}\>\bm{\mathcal{U}}\>\ket{\Psi}$ contains fingerprints of the phases $\theta(\bm{\sigma})$ (Fig. 4). In general, different rotations must be independently applied to gain full information on the phase structure. We make the assumption of a tensor product structure of the rotations, $\bm{\mathcal{U}}=\bigotimes_{j=1}^{N}\hat{\mathcal{U}}_{j}$. This is equivalent to a local change of basis from $\ket{\bm{\sigma}}$ to $\{|\bm{\sigma}^{\bm{b}}\rangle=|\sigma_{1}^{b_{1}},\dots,\sigma_{N}^{b_{N}}\rangle\}$, where the vector $\bm{b}$ identifies the local basis $b_{j}$ for each site $j$. The target wavefunction in the new basis is given by

[TABLE]

and the resulting measurement distribution is

[TABLE]

To clarify the procedure, let us consider the simple example of a quantum state of two qubits:

[TABLE]

and rotation $\bm{\mathcal{U}}=\hat{\mathrm{H}}_{0}\otimes\hat{\mathcal{I}}_{1}$, where $\hat{\mathcal{I}}$ is the identity operator and

[TABLE]

is called the Hadamard gate. This transformation is equivalent to rotating the qubit $\sigma_{0}$ from the reference $\sigma_{0}^{z}$ basis the the $\sigma_{0}^{x}$ basis. A straightforward calculation leads to the following probability distribution of the projective measurement in the new basis $|\sigma_{0}^{x},\sigma_{1}\rangle$:

[TABLE]

where $\Delta\theta=\theta_{0\sigma_{1}}-\theta_{1\sigma_{1}}$. Therefore, the statistics collected by measuring in this basis implicitly contains partial information on the phases. To obtain the full phases structure, additional transformations are required, one example being the rotation from the reference basis to the $\sigma^{y}_{j}$ local basis, realized by the elementary gate

[TABLE]

3.1 Setup

We now proceed to use QuCumber to reconstruct a complex-valued wavefunction. For simplicity, we restrict ourselves to two qubits and consider the general case of a quantum state with random amplitudes $\Phi_{\sigma_{0}\sigma_{1}}$ and random phases $\theta_{\sigma_{0}\sigma_{1}}$. This example is available in the online tutorial. 333The tutorial for complex wavefunctions can be found at [https://qucumber.readthedocs.io/en/

stable/_examples/Tutorial2_TrainComplexWaveFunction/tutorial_qubits.html](https://qucumber.readthedocs.io/en/stable/_examples/Tutorial2_TrainComplexWaveFunction/tutorial_qubits.html)

We begin by importing the required packages:

⬇

import qucumber.utils.unitaries as unitaries

import qucumber.utils.cplx as cplx

Since we are dealing with a complex wavefunction, we load the corresponding module ComplexWaveFunction to build the RBM quantum state $\psi_{\bm{\lambda\mu}}(\bm{\sigma})$. Furthermore, the following additional utility modules are required: the utils.cplx backend for complex algebra, and the utils.unitaries module which contains a set of elementary local rotations. By default, the set of unitaries include local rotations to the $\sigma^{x}$ and $\sigma^{y}$ basis, implemented by the $\hat{\mathrm{H}}$ and $\hat{\mathrm{K}}$ gates respectively.

We continue by loading the data444 The training dataset can be downloaded from [https://github.com/PIQuIL/QuCumber/blob/master/

examples/Tutorial2_TrainComplexWaveFunction/](https://github.com/PIQuIL/QuCumber/blob/master/examples/Tutorial2_TrainComplexWaveFunction/) into QuCumber, which is done using the load_data function of the data utility:

⬇

train_bases_path = "qubits_train_bases.txt"

psi_path = "qubits_psi.txt"

bases_path = "qubits_bases.txt"

train_samples, true_psi, train_bases, bases = data.load_data(

train_path, psi_path, train_bases_path, bases_path

)

As before, we may load the true target wavefunction from qubits_psi.txt, which can be used to calculate the fidelity and KL divergence. In contrast with the positive case, we now have measurements performed in different bases. Therefore, the training data consists of an array of qubits projections $(\sigma_{0}^{b_{0}},\sigma_{1}^{b_{1}}$) in qubits_train_samples.txt, together with the corresponding bases $(b_{0},b_{1})$ where the measurement was taken, in qubits_train_bases.txt. Finally, QuCumber loads the set of all the bases appearing the in training dataset, stored in qubits_bases.txt. This is required to properly configure the various elementary unitary rotations that need to be applied to the RBM state during the training. For this example, we generated measurements in the following bases:

[TABLE]

Finally, before the training, we initialize the set of unitary rotations and create the RBM state object. In the case of the provided dataset, the unitaries are the $\hat{\mathrm{H}}$ and $\hat{\mathrm{K}}$ gates. The required dictionary can be created with unitaries.create_dict(). By default, when unitaries.create_dict() is called, it will contain the identity, the $\hat{\mathrm{H}}$ gate, and the $\hat{\mathrm{K}}$ gate, with the keys Z, X, and Y, respectively. It is possible to add additional gates by specifying them as

⬇

unitary_dict = unitaries.create_dict(<unitary_name>=U)

where re_part, im_part, and unitary_name are to be specified by the user.

We then initialize the complex RBM object with

⬇

num_visible=2, num_hidden=2, unitary_dict=unitary_dict

)

The key difference between positive and complex wavefunction reconstruction is the requirement of additional measurements in different basis. Despite this, loading the data, initializing models, and training the RBMs are all very similar to the positive case, as we now discuss.

3.2 Training

Like in the case of a positive wavefunction, for the complex case QuCumber optimizes the network parameters to minimize the KL divergence between the data and the RBM distribution. When measuring in multiple bases, the optimization now runs over the set of parameters $(\bm{\lambda},\bm{\mu})$ and minimizes the sum of KL divergences between the data distribution $P(\bm{\sigma}^{\bm{b}})$ and the RBM distribution $|\psi_{\bm{\lambda\mu}}(\bm{\sigma}^{\bm{b}})|^{2}$ for each bases $\bm{b}$ appearing in the training dataset [6]. For example, if a given training sample is measured in the basis $(\mathrm{x},\mathrm{z})$, QuCumber applies the appropriate unitary rotation $\bm{\mathcal{U}}=\hat{\mathrm{H}}_{0}\otimes\hat{\mathcal{I}}_{1}$ to the RBM state before collecting the gradient signal.

Similar to the case of positive wavefunction, we generate the Hilbert space (to compute fidelity and KL divergence) and initialize the callbacks

⬇

callbacks = [

MetricEvaluator(

log_every,

{"Fidelity": ts.fidelity, "KL": ts.KL},

target_psi=true_psi,

bases=bases,

verbose=True,

space=state.space,

)

]

The training is carried out by calling the fit function of ComplexWaveFunction, given the set of hyperparameters

⬇

train_samples,

epochs=100,

pos_batch_size=10,

neg_batch_size=10,

lr=0.05,

k=5,

input_bases=train_bases,

callbacks=callbacks,

)

In Fig. 5 we show the total KL divergence and the fidelity with the true two-qubit state during training. After successfully training a QuCumber model, we can once again compute expectation values of physical observables, as discussed in Section 2.3.

4 Conclusion

We have introduced the open source software package QuCumber, a quantum calculator used for many-body eigenstate reconstruction. QuCumber is capable of taking input data representing projective measurements of a quantum wavefunction, and reconstructing the wavefunction using a restricted Boltzmann machine (RBM).

Once properly trained, QuCumber can produce a new set of measurements, sampled stochastically from the RBM. These samples, generated in the reference basis, can be used to verify the training of the RBM against the original data set. More importantly, they can be used to calculate expectation values of many physical observables. In fact, any expectation value typically estimated by conventional Monte Carlo methods can be implemented as an estimator in QuCumber. Such estimators may be inaccessible in the reference basis, for example. Or, they may be difficult or impossible to implement in the setup for which the original data was obtained. This is particularly relevant for experiments, where it is easy to imagine many possible observables that are inaccessible, due to fundamental or technical challenges.

Future versions of QuCumber, as well as the next generation of quantum state reconstruction software, may explore different generative models, such as variational autoencoders, generative adversarial networks, or recurrent neural networks. The techniques described in this paper can also be extended to reconstruct mixed states, via the purification technique described in Reference [7]. In addition, future techniques may include hybridization between machine learning and other well-established methods in computational quantum many-body physics, such as variational Monte Carlo and tensor networks [8].

Acknowledgements

We acknowledge M. Albergo, G. Carleo, J. Carrasquilla, D. Sehayek, and L. Hayward Sierens for stimulating discussions. We thank the Perimeter Institute for Theoretical Physics for the continuing support of PIQuIL.

Author contributions

Authors are listed alphabetically. For an updated record of individual contributions, consult the repository at https://github.com/PIQuIL/QuCumber/graphs/contributors.

Funding information

This research was supported by the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada Research Chair program, and the Perimeter Institute for Theoretical Physics. We also gratefully acknowledge the support of NVIDIA Corporation with the donation of the Titan Xp GPU used in this work. Research at Perimeter Institute is supported by the Government of Canada through Industry Canada and by the Province of Ontario through the Ministry of Research & Innovation. P. H. acknowledges support from ICFOstepstone, funded by the Marie Sklodowska-Curie Co-funding of regional, national and international programmes (GA665884) of the European Commission, as well as by the Severo Ochoa 2016–2019’ program at ICFO (SEV–2015–0522), funded by the Spanish Ministry of Economy, Industry, and Competitiveness (MINECO).

Appendix A Glossary

This section contains an overview of terms discussed in the document which are relevant for RBMs. For more detail we refer the reader to the code documentation on https://qucumber.readthedocs.io/en/stable/, and References [12, 30].

• •

Batch: A subset of data upon which the gradient is computed and the network parameters are adjusted accordingly. A smaller batch size often results in a more stochastic trajectory, while a large batch size more closely approximates the exact gradient and has less variance.

• •

Biases: Adjustable parameters in an RBM, denoted by $b_{j}$ and $c_{i}$ in Eq. (20).

• •

Contrastive divergence: An approximate maximum-likelihood learning algorithm for RBMs [12]. CD estimates the gradient of the effective energy (20) with respect to model parameters by using Gibbs sampling to compare the generated and target distributions.

• •

Energy: The energy of the joint configuration $(\bm{v},\bm{h})$ of a RBM is defined as follows:

[TABLE]

• •

Effective energy: Obtained from the energy by tracing out the hidden units $\bm{h}$; often called the “free energy” in machine learning literature.

[TABLE]

• •

Epoch: A single pass through an entire training set. For example, with a training set of 1,000 samples and a batch size of 100, one epoch consists of 10 updates of the parameters.

• •

Gibbs sampling: A Monte Carlo algorithm that samples from the conditional distribution of one variable, given the state of other variables. In an RBM, the restricted weight connectivity allows Gibbs sampling between the visible “block”, conditioned on the hidden “block”, and vice versa.

• •

Hidden units: There are $n_{h}$ units in the hidden layer of the RBM, denoted by the vector $\bm{h}=(h_{1},\ldots,h_{n_{h}})$. The number of hidden units can be adjusted to increase the representational capacity of the RBM.

• •

Hyperparameters: A set of parameters that are not adjusted by a neural network during training. Examples include the learning rate, number of hidden units, batch size, and number of training epochs.

• •

Joint distribution: The RBM assigns a probability to each joint configuration $(\bm{v},\bm{h})$ according to the Boltzmann distribution:

[TABLE]

• •

KL divergence: The Kullback-Leibler divergence, or relative entropy, is a measure of the “distance” between two probability distributions $P$ and $Q$, defined as:

[TABLE]

The KL divergence between two identical distributions is zero. Note that it is not symmetric between $P$ and $Q$.

• •

Learning rate: The step size used in the gradient descent algorithm for the optimization of the network parameters. A small learning rate may result in better optimization but will take more time to converge. If the learning rate is too high, training might not converge or will find a poor optimum.

• •

Marginal distribution: Obtained by marginalizing out the hidden layer from the joint distribution via

[TABLE]

• •

QuCumber: A quantum calculator used for many-body eigenstate reconstruction.

• •

Parameters: The set of weights and biases $\bm{\lambda}=\{\bm{W},\bm{b},\bm{c}\}$ characterizing the RBM energy function. These are adjusted during training.

• •

Partition function: The normalizing constant of the Boltzmann distribution. It is obtained by tracing over all possible pairs of visible and hidden vectors:

[TABLE]

• •

Restricted Boltzmann Machine: A two-layer network with bidirectionally connected stochastic processing units. “Restricted” refers to the connections (or weights) between the visible and hidden units: each visible unit is connected with each hidden unit, but there are no intra-layer connections.

• •

Visible units: There are $n_{v}$ units in the visible layer of the RBM, denoted by the vector $\bm{v}=(v_{1},\dots,v_{n_{v}})$. These units correspond to the physical degrees of freedom. In the cases considered in this paper, the number of visible units is equal to the number of spins $N$.

• •

Weights: $W_{ij}$ is the symmetric connection or interaction between the visible unit $v_{j}$ and the hidden unit $h_{i}$.