NetKet: A Machine Learning Toolkit for Many-Body Quantum Systems

We introduce NetKet, a comprehensive open source framework for the study of many-body quantum systems using machine learning techniques. The framework is built around a general and flexible implementation of neural-network quantum states, which are used as a variational ansatz for quantum wave functions. NetKet provides algorithms for several key tasks in quantum many-body physics and quantum technology, namely quantum state tomography, supervised learning from wave-function data, and ground state searches for a wide range of customizable lattice models. Our aim is to provide a common platform for open research and to stimulate the collaborative development of computational methods at the interface of machine learning and many-body physics.


Motivation and significance
Recent years have seen a tremendous activity around the development of physics-oriented numerical techniques based on machine learning (ML) tools [1].In the context of many-body quantum physics, one of the main goals of these approaches is to tackle complex quantum problems using compact representations of many-body states based on artificial neural networks.These representations, dubbed neural-network quantum states (NQS) [2], can be used for several applications.In the supervised learning setting, they can be used, e.g., to learn existing quantum states for which a non-NQS representation is available [3].In the unsupervised setting, they can be used to reconstruct complex quantum states from experimental measurements, a task known as quantum state tomography [4].Finally, in the context of purely variational applications, NQS can be used to find approximate ground-and excited-state solutions of the Schrödinger equation [2,[5][6][7][8][9], as well as to describe unitary [2,10,11] and dissipative [12][13][14][15] many-body dynamics.Despite the increasing methodological and theoretical interest in NQS and their applications, a set of comprehensive, easy-to-use tools for research applications is still lacking.This is particularly pressing as the complexity of NQS-related approaches and algorithms is expected to grow rapidly given these first successes, steepening the learning curve.
The goal of NetKet is to provide a set of primitives and flexible tools to ease the development of cutting-edge ML applications for quantum many-body physics.NetKet also wants to help bridge the gap between the latest and technically demanding developments in the field and those scholars and students who approach the subject for the first time.Pedagogical tutorials are provided to this aim.Serving as a common platform for future research, the NetKet project is meant to stimulate the open and easy-to-certify development of new methods and to provide a common set of tools to reproduce published results.
A central philosophy of the NetKet framework is to provide tools that are as simple as possible to use for the end user.Given the huge popularity of the Python programming language and of the many accompanying tools gravitating around the Python ecosystem, we have built NetKet as a full-fledged Python library.This simplicity of use however does not come at the expense of performance.With this efficiency requirement in mind, all critical routines and components of NetKet have been written in C++11.

Software description
We will first give a general overview of the structure of the code in Section 2.1 and then provide additional details on the functionality of NetKet in Section 2.2.

Software architecture
The core of NetKet is implemented in C++.For ease of use and in order to facilitate the integration with other frameworks, a Python interface is provided, which exposes all high-level functionality from the C++ core via pybind11 [16] bindings.Use of the Python interface is recommended for users building on the library for research purposes, while the C++ code should be modified for extending the NetKet library itself.
NetKet is divided into several submodules.The modules graph, hilbert, and operator contain the classes necessary for specifying the structure of the many-body Hilbert space, the Hamiltonian, and other observables of a quantum system.The exact module provides functions for exact diagonalization (ED) based on SciPy [18] and time propagation of the full quantum state, in order to allow for easy benchmarking and exploration of small systems within the NetKet framework.The NetKet operator classes implement the SciPy linear-operator interface and can also be converted to sparse and dense matrices, providing interoperability with Python code.In particular, the sparse and dense ED routines provided by the exact module are implemented as thin wrappers around SciPy functionality.
The dynamics module provides basic ODE solvers for exact time propagation.
The utility modules output, stats, and util contain some additional functionality for output and statistics that is used internally in other parts of NetKet.
An overview of the most important modules and their dependencies is given in Fig. 1.A more detailed description of the module contents will be given in the next section.
NetKet uses the Eigen 3 library [19] for linear algebra routines.In the Python interface, Eigen datatypes are transparently converted to and from NumPy [20] arrays by pybind11.The NetKet driver classes provide methods to directly write the simulation output to JSON files, which is done with the help of the nlohmann/json library for C++ [21].Parallelization is implemented based on the Message Passing Interface (MPI), allowing to substantially decrease running time.Specifically, the Monte Carlo sampling of expectation values implemented in the variational.Vmc class is parallelized, with each node drawing independent samples from the probability distribution which are averaged over all nodes.

Software functionalities
The core feature of NetKet is the variational representation of quantum states by artificial neural networks.Given a variational state, the task is to optimize its parameters with regard to a specified loss function, such as the total energy for ground state searches or the (negative) overlap with a given target state.In this section, we will discuss the models, types of variational wavefunctions, and learning schemes that are available in NetKet.

Model specification
NetKet currently supports lattice models with a finite Hilbert space of the form H = H ⊗N local where N denotes the number of lattice sites and H local denotes the local Hilbert space of each site.The system is defined on a graph with a set of N sites and a set of edges (also called bonds) between pairs of sites.This graph structure is used to help with the definition of operators on the lattice and to encode the spatial structure of the model, which is necessary, e.g., to work with convolutional neural networks (CNNs).NetKet provides the predefined Hypercube and Lattice graphs.Furthermore, CustomGraph supports arbitrary edge-colored graphs, where each edge is associated with an integer label called its color.This color can be used to describe different types of bonds.
General lattice spin models can be described straightforwardly in this manner.Bosonic lattice models can also be easily represented by truncating the local Hilbert space to only allow for occupations of up to N local bosons per site [22].NetKet currently provides pre-defined Hamiltonians for the transverse-field Ising, Heisenberg, and Bose-Hubbard models.Other observables and custom Hamiltonians can also be specified: a convenient option for common lattice models is to use the GraphOperator class, which allows to construct a Hamiltonian from a family of 2-local operators acting on each bond of a selected color and a family of 1-local operators acting on each site.It is also possible to specify general k-local operators (as well as their products and sums) using the LocalOperator class.
While fermionic Hamiltonians are not fully supported in the present version, they can be implemented using a custom Jordan-Wigner mapping and the LocalOperator class [23].

Variational quantum states
The purpose of variational states is to provide a compact and computationally efficient representation of quantum states.Since generally only a subset of the full many-body Hilbert space will be covered by a given variational ansatz, the aim is to use a parametrization that captures the relevant physical states for a given problem.
As activation functions, rectified linear units (Relu) [33], hyperbolic tangent (Tanh) [34], and the logarithm of the hyperbolic cosine (Lncosh) are provided.RBMs without visible bias can be represented as single-layer FFNNs with ln cosh activation, allowing for a generalization of these machines to multiple layers [5].
The machine module also provides more traditional variational wavefunctions, namely MPS with periodic boundary conditions (MPSPeriodic) and long-range Jastrow (Jastrow) wavefunctions, which allows for comparison of NQS with results obtained using these approaches.
Finally, NetKet also includes representations of mixed states in the machine.densitymatrixsubmodule, which most notably includes the real-valued NDO ansatz (NdmSpinPhase).For compatibility with the rest of the package, the vectorized representation of density matrices can be accessed through the same interface as NQS.
Custom wavefunctions may be provided by implementing subclasses of the AbstractMachine class in C++ or in Python by deriving netket.machine.CxxMachine.

Supervised learning
In supervised learning, a target wavefunction is given and the task is to optimize a chosen ansatz to represent it.This functionality is contained within the supervised module.Given a variational state |Ψ NN (α)⟩ depending on the parameters α ∈ C m and a target state |Ψ tar ⟩, the negative log overlap is taken as the loss function to be minimized.The loss is computed in a Monte Carlo fashion by direct sampling of the target wavefunction.To minimize the loss, the gradient ∇ α L of the loss function with respect to the parameters is calculated.This gradient is then used to update the parameters according to a specified gradient-based optimization scheme.For example, in stochastic gradient descent (SGD) the parameters are updated as where λ is the learning rate.The different update rules supported by NetKet are contained in the optimizer module.Various types of optimizers are available, including SGD, AdaGrad [35], AdaMax and AdaDelta [36], AMSGrad [37], and RMSProp.

Unsupervised learning
NetKet also allows to carry out unsupervised learning of unknown probability distributions, which in this context corresponds to quantum state tomography [38].Given an unknown quantum state, a neural network can be trained on projective measurement data to discover an approximate reconstruction of the state [4].In NetKet, this functionality is contained within the where π denotes the probability distribution The network parameters are updated according to the gradient of the log-likelihood L. This can be computed analytically, and it requires expectation values over both the training data points and the network distribution π(σ).While the first is trivial to compute, the latter should be approximated by a Monte Carlo average over configurations sampled from a Markov chain.

Variational Monte Carlo
Finally, NetKet supports ground state searches for a given many-body quantum Hamiltonian Ĥ.In this context, the task is to optimize the parameters of a variational wavefunction Ψ in order The energy of a wavefunction Ψ (σ) = ⟨σ|Ψ ⟩ can be estimated as where in the last line ⟨ • ⟩ σ denotes a stochastic expectation value taken over a sample of configurations {σ} drawn from the probability distribution corresponding to the variational wavefunction (4).This sampling is performed by classes from the sampler module, which generate Markov chains of configurations using the Metropolis algorithm [39] to ensure detailed balance.Parallel tempering [40] options are also available to improve sampling efficiency.
In order to optimize the parameters of a machine to minimize the energy, a gradient-based optimization scheme can be applied as discussed in the previous section.The energy gradient can be estimated at the same time as ⟨ Ĥ⟩ [2,25].This requires computing the partial derivatives of the wavefunction with respect to the variational parameters, which can be obtained analytically for the RBM [2] or via backpropagation [30,31,34] for multi-layer FFNNs.In this case, the steepest descent update according to Eq. ( 2) is also a form of SGD, because the energy is estimated using a subset of the full data available from the variational wavefunction.Alternatively, often more stable convergence can be achieved by using the stochastic reconfiguration (SR) method [41,42], which approximates the imaginary time evolution of the system on the submanifold of variational states.The SR approach is closely related to the natural gradient descent method used in machine learning [43].In the NetKet implementation, SR is performed using either an exact or an iterative linear solver, the latter being recommended when the number of variational parameters is large.
Information on the optimization run (sampler acceptance rates, energy, energy variance, expectation of additional observables, and the current variational parameters) for each iteration Listing 1: Example script for finding the ground state of the one-dimensional spin- 1  2 Heisenberg model using an RBM ansatz.
can be written to a log file in JSON format.Alternatively, they can be accessed directly inside the simulation loop in Python to allow for more flexible output.

Illustrative examples
NetKet is available as a Python package and can be obtained from the Python package index (PyPI) [44].Assuming a properly configured Python environment, NetKet can be installed via the shell command pip install netket which will download, compile, and install the package.A working MPI environment is required to run NetKet.In case multiple MPI installations are present on the system and in order to avoid potential conflicts, we recommend to run the installation command as CC=mpicc CXX=mpicxx pip install netket with the desired MPI environment loaded in order to perform the build with the correct compiler.After a successful installation, the NetKet module can be imported in Python scripts.
Alternatively to installing NetKet locally, NetKet also uses the deployment of BinderHub from mybinder.org[45] to build and deploy a stable version of the software, which can be found at https://mybinder.org/v2/gh/netket/netket/v.2.0.This allows users to run the tutorials or other small jobs without installing NetKet.transverse field Ising model with 10 sites from ED data, using an RBM with 20 hidden units.The blue line shows the overlap between the RBM wavefunction and the exact wavefunction for each iteration.

One-dimensional Heisenberg model
As a first example, we present a Python script for obtaining a variational RBM representation of the ground state of the spin-1 2 Heisenberg model on a one-dimensional chain with periodic boundary conditions.The code for this example is shown in Listing 1. Fig. 2 shows the evolution of the energy expectation value over the course of the optimization run.We see that for a small chain of 20 sites and an RBM with 20 hidden units, the energy converges to a relative error of the order 10 −5 within about 100 iteration steps.

Supervised learning
As a second example, we use the supervised learning module in NetKet to optimize an RBM to represent the ground state of the transverse field Ising model.The example script is shown in Listing 2. The exact ground state wavefunction is first obtained by exact diagonalization and then used for training the RBM state by minimizing the overlap loss (1).Fig. 3 shows the evolution of the overlap over the training iterations.

Impact
Given the flexibility of NetKet, we envision several potential applications of this library both in data-driven experimental research and in more theoretical, problem-driven research on interacting quantum many-body systems.For example, several important theoretical and practical questions concerning the expressibility of NQS, the learnability of experimental quantum states, and the efficiency at finding ground states of k-local Hamiltonians, can be directly addressed using the current functionality of the software.
Moreover, having an easy-to-extend set of tools to work with NQS-based applications can propel future research in the field, without researchers having to pay a significant cost of entry in terms of algorithm implementation and testing.Since its early release in April 2017, NetKet has already been used for research purposes by several groups worldwide [5,22,23,[46][47][48].We also hope that, building upon a common set of tools, practices like publishing accompanying codes to research papers, largely popular in the ML community, can become standard practice also for ML applications in quantum physics.
Finally, for a fast-growing community like ML for quantum science, it is also crucial to have pedagogical tools available that can be conveniently used by new generations of students and researchers.Benefiting from a growing set of tutorials and stepby-step explanations, NetKet can be comfortably used in schools and lectures.

Conclusions and future directions
We have introduced NetKet, a comprehensive open source framework for the study of many-body quantum systems using machine learning techniques.Central to this framework are variational parameterizations of many-body wavefunctions in the form of artificial neural networks.NetKet is a Python framework implemented in C++11, designed with efficiency as well as ease of use in mind.Several examples, tutorials, and notebooks are provided with our software in order to reduce the learning curve for newcomers.
The NetKet project is meant to continuously evolve in future releases, welcoming suggestions and contributions from its users.For example, future versions may provide a natural interface with general ML frameworks such as PyTorch [49] and Tensorflow [50].On the algorithmic side, future goals include the extension of NetKet to incorporate unitary dynamics [11,51], convenient Fermionic operators, as well as full support for density-matrix tomography [17].

Declaration of competing interest
The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.

Fig. 1 .
Fig. 1.The main submodules of the netket Python module and their dependencies from a user perspective (i.e., only dependencies in the public interface are shown).Below each submodule, examples of contained classes and functions are displayed.In a typical workflow, users will first define a quantum model, specify a variational representation of the wavefunction as well as the Monte Carlo sampling and optimization methods, and then run the simulation using one of the driver classes.A more detailed description of the software architecture and features is given in the main text.

1 2
systems as well as two other variants: the symmetric RBM (RbmSpinSymm) to capture lattice symmetries such as translation and inversion symmetries and the multi-valued RBM (RbmMultiVal) for systems with larger local Hilbert spaces (such as higher spins or bosonic systems).FFNNs represent a broad and flexible class of networks and are implemented by the FFNN class.They consist of a sequence of layers available from the layer submodule, each layer performing either an affine transformation to the input vector or applying a non-linear activation function.There are currently two types of affine maps available:• Dense fully-connected layers, which for an input x ∈ C n and output y ∈ C m have the form y = Wx + b where W ∈ C m×n and b ∈ C m are called the weight matrix and bias vector, respectively.

2 .
unsupervised.Qsr class.For some given target quantum state |Ψ tar ⟩, the training dataset D consists of a sequence of projective measurements σ b in different bases b, with underlying probability distribution P(σ b ) = |Ψ tar (σ b )| The quantum reconstruction of the target state translates into minimizing the statistical divergence between the distribution of the measurement outcomes and the distribution generated by the NQS.This corresponds, up to a constant dataset entropy contribution, to maximizing the loglikelihood of the network distribution over the measurement data

( 4 )
generated by the NQS wavefunction.Note that, for every training sample where the measurement basis differs from the reference basis |σ⟩ of the NQS, a unitary transformation Û should be applied to appropriately change the basis, Ψ NN (σ b ) = Ûb Ψ NN (σ).

2 .
Variational optimization of the restricted Boltzmann machine for the one-dimensional spin-1 2 Heisenberg model.The main plot shows the Monte Carlo energy estimate, which converges to the exact ground state energy up to a relative error |(E − E exact )/E exact | of 4.16 × 10 −5 within the 200 iteration steps shown.The inset shows the Monte Carlo estimate of the energy variance, which becomes zero in an exact eigenstate of the Hamiltonian.tominimize the energy ⟨ Ĥ⟩.The variational.Vmc driver class contains the main logic to optimize a variational wavefunction given a Hamiltonian, a sampler, and an optimizer.

Fig. 3 .
Fig. 3. Supervised learning of the ground state of the one-dimensional spin-1 2

Listing 2 :
g e t s = t a r g e t s )38 39 # Run the optimization f o r 2000 i t e r a t i o n s 40 spvsd .run ( n _ i t e r =2000 , o u t p u t _ p r e f i x = ' t e s t ' , 41 l o s s _ f u n c t i o n = " Overlap_phi " ) Example script for supervised learning.A RBM ansatz is optimized to represent the ground state of the one-dimensional spin-1  2 transverse field Ising model obtained by ED for this example.
[17]ented in the machine.densitymatrixsubmodule in the form of Neural Density Operators (NDO)[17].The variational, supervised, and unsupervised modules contain driver classes The core component of NetKet is the machine module, which provides different variational representations of the quantum wavefunction, particularly in the form of NQS.Encodings of mixed states, needed to describe dissipative quantum system, are im-