NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems

Filippo Vicentini1,2$\star$, Damian Hofmann3, Attila Szabó4,5, Dian Wu1,2,
Christopher Roth6, Clemens Giuliani1,2, Gabriel Pescia1,2, Jannes Nys1,2,
Vladimir Vargas-Calderón7,

Nikita Astrakhantsev8 and Giuseppe Carleo1,2

1 École Polytechnique Fédérale de Lausanne (EPFL), Institute of Physics,

CH-1015 Lausanne, Switzerland

2 Center for Quantum Science and Engineering, École Polytechnique Fédérale de
Lausanne (EPFL), CH-1015 Lausanne, Switzerland

3 Max Planck Institute for the Structure and Dynamics of Matter,

Center for Free-Electron Laser Science (CFEL),

Luruper Chaussee 149, 22761 Hamburg, Germany

4 Rudolf Peierls Centre for Theoretical Physics,

University of Oxford, Oxford OX1 3PU, United Kingdom

5 ISIS Facility, Rutherford Appleton Laboratory,

Harwell Campus, Didcot OX11 0QX, United Kingdom

6 Physics Department, University of Texas at Austin

7 Grupo de Superconductividad y Nanotecnología, Departamento de Física,

Universidad Nacional de Colombia, Bogotá, Colombia

8 Department of Physics, University of Zurich, CH-8057 Zurich, Switzerland

⋆<EMAIL_ADDRESS>
## Abstract

We introduce version 3 of NetKet, the machine learning toolbox for many-body
quantum physics. NetKet is built around neural quantum states and provides
efficient algorithms for their evaluation and optimization. This new version
is built on top of JAX, a differentiable programming and accelerated linear
algebra framework for the Python programming language. The most significant
new feature is the possibility to define arbitrary neural network ansätze in
pure Python code using the concise notation of machine-learning frameworks,
which allows for just-in-time compilation as well as the implicit generation
of gradients thanks to automatic differentiation. NetKet 3 also comes with
support for GPU and TPU accelerators, advanced support for discrete symmetry
groups, chunking to scale up to thousands of degrees of freedom, drivers for
quantum dynamics applications, and improved modularity, allowing users to use
only parts of the toolbox as a foundation for their own code.

###### Contents

  1. 1 Introduction
    1. 1.1 What’s new
    2. 1.2 Outline
    3. 1.3 Installing NetKet
  2. 2 Quantum-mechanical primitives
    1. 2.1 Hilbert spaces
    2. 2.2 Linear operators
  3. 3 Variational quantum states
    1. 3.1 Abstract interface
    2. 3.2 Defining the variational ansatz
    3. 3.3 Estimating observables
    4. 3.4 Monte Carlo samplers
    5. 3.5 Quantum geometric tensor
  4. 4 Algorithms for variational states
    1. 4.1 Ground-state search
    2. 4.2 Finding steady states
    3. 4.3 Time propagation
    4. 4.4 Implementing custom algorithms using NetKet
  5. 5 Symmetry-aware neural quantum states
    1. 5.1 Symmetry groups and representation theory
    2. 5.2 Using group convolutional neural networks (GCNNs)
  6. 6 Quantum systems with continuous degrees of freedom
    1. 6.1 Continuous Hilbert spaces
    2. 6.2 Linear operators
    3. 6.3 Samplers
    4. 6.4 Harmonic oscillators
    5. 6.5 Interacting system with continuous degrees of freedom
  7. 7 Example: Finding ground and excited states of a lattice model
    1. 7.1 Defining the lattice and the Hamiltonian
    2. 7.2 Defining and training a symmetric ansatz
    3. 7.3 Finding an excited state
  8. 8 Example: Fermions on a lattice
  9. 9 Example: Real-time dynamics
    1. 9.1 Unitary dynamics
    2. 9.2 Dissipative dynamics (Lindblad master equation)
  10. 10 Benchmarks
    1. 10.1 Variational Monte Carlo
    2. 10.2 MPI for NetKet
    3. 10.3 Comparison with jVMC
  11. 11 Discussion and conclusion
  12. A Details of group convolutions
    1. A.1 Group convolutions and equivariance
    2. A.2 Fast group convolutions using Fourier transforms
  13. B Implementation details of the quantum geometric tensor
    1. B.1 Jacobians and their products
    2. B.2 QGTJacobian
    3. B.3 QGTOnTheFly

##  1 Introduction

During the last two decades, we have witnessed tremendous advances in machine
learning (ML) algorithms which have been used to solve previously difficult
problems such as image recognition [1, 2] or natural language processing [3].
This has only been possible thanks to sustained hardware development: the last
decade alone has seen a 50-fold increase in available computing power [4].
However, unlocking the full computational potential of modern arithmetic
accelerators, such as GPUs, used to require significant technical skills,
hampering researchers in their efforts. The incredible pace of algorithmic
advances must therefore be attributed, at least in part, to the development of
frameworks allowing researchers to tap into the full potential of computer
clusters while writing high-level code [5, 6].

In the last few years, researchers in quantum physics have increasingly
utilized machine-learning techniques to develop novel algorithms or improve on
existing approaches [7]. In the context of variational methods for many-body
quantum physics in particular, the method of _neural quantum states_ (NQS) has
been developed [8]. NQS are based on the idea of using neural networks as an
efficient parametrization of the quantum wave function. They are of particular
interest because of their potential to represent highly entangled states in
more than one dimension with polynomial resources [9], which is a significant
challenge for more established families of variational states. NQS are also
flexible: they have been successfully used to determine variational ground
states of classical [10] and quantum Hamiltonians [11, 12, 13, 14, 15, 16, 17]
as well as excited states [13], to approximate Hamiltonian unitary dynamics
[8, 18, 19, 20, 21, 22, 23], and to solve the Lindblad master equation [24,
25, 26]. In particular, NQS are currently used in the study of frustrated
quantum systems [13, 15, 27, 17, 16, 28, 29, 30], which have so far been
challenging to optimize by established numerical techniques. They have also
been used to perform tomographic state reconstruction [31] and efficiently
approximate quantum circuits [32].

A complication often encountered when working with NQS is, however, that
standard ML frameworks like TensorFlow [33] or PyTorch [34] are not geared
towards these kind of quantum mechanical problems, and it often takes
considerable technical expertise to use them for such non-standard tasks.
Alternatively, researchers sometimes avoid those frameworks and implement
their routines from scratch, but this often leads to sub-optimal performance.
We believe that it is possible to foster research at this intersection of
quantum physics and ML by providing an easy-to-use interface exposing quantum
mechanical objects to ML frameworks.

We therefore introduce version 3 of the NetKet framework [35].111This
manuscript refers to NetKet v3.5, released in August 2022. NetKet 3 is an
open-source Python toolbox expressing several quantum mechanical primitives in
the differentiable programming framework jax [5, 36]. NetKet provides an easy-
to-use interface to high-performance variational techniques without the need
to delve into the details of their implementations, but customizability is not
sacrificed and advanced users can inspect, modify, and extend practically
every aspect of the package. Moreover, integration of our quantum object
primitives with the jax ecosystem allows users to easily define custom neural-
network architectures and compute a range of quantum mechanical quantities, as
well as their gradients, which are auto-generated through jax’s tracing-based
approach. jax provides the ability to write numerical code in pure Python
using NumPy-like calls for array operations, while still achieving high
performance through just-in-time compilation using xla, the accelerated linear
algebra compiler that underlies TensorFlow. We have also integrated jax and
MPI with the help of mpi4jax [37] to make NetKet scale to hundreds of
computing nodes.

###  1.1 What’s new

With the release of version 3, NetKet has moved from internally relying on a
custom C++ core to the jax framework, which allows models and algorithms to be
written in pure Python and just-in-time compiled for high performance on both
CPU and GPU platforms.222Google’s Tensor Processing Units (TPUs) are also, in
principle, supported. However, at the time of writing they only support half-
precision float16. Some modifications would be necessary to work-around loss
of precision and gradient underflow. By using only Python, the installation
process is greatly simplified and the barrier of entry for new contributors is
lowered.

iFrom a user perspective, the most important new feature is the possibility of
writing custom NQS wave functions using jax, which allows for quick
prototyping and deployment, frees users from having to manually implement
gradients due to jax’s support for automatic differentiation, and makes models
easily portable to GPU platforms. Other prominent new features are

  * •

support for (real and imaginary time) unitary and Markovian dissipative
dynamics;

  * •

support for continuous systems;

  * •

support for composite Hilbert spaces;

  * •

efficient implementations of the quantum geometric tensor and stochastic
reconfiguration, which scale to models with millions of parameters;

  * •

group-invariant and group-equivariant layers and architectures which support
arbitrary discrete symmetries.

A more advanced feature is an extension mechanism built around multiple
dispatch [38], which allows users to override algorithms used internally by
NetKet without editing the source itself. This can be used to make NetKet work
with custom objects and algorithms to study novel problems that do not easily
fit what is already available.

###  1.2 Outline

NetKet provides both an intuitive high-level interface with sensible defaults
to welcome beginners, as well as a complete set of options and lower-level
functions for flexible use by advanced users. The high-level interface is
built around quantum-mechanical objects such as Hilbert spaces (
netket.hilbert) and operators ( netket.operator), presented in Section 2.

The central object in NetKet 3 is the variational state, discussed in Section
3, which bring together the neural-network ansatz, its variational parameters,
and a Monte-Carlo sampler. In Section 3.2, we give an example on how to define
an arbitrary neural network using a NetKet/jax-compatible framework, while
Section 3.4 presents the new API of stochastic samplers. In Section 3.5, we
show how to compute the quantum geometric tensor (QGT) with NetKet, and
compare the different implementations.

Section 4 shows how to use the three built-in optimization drivers to perform
ground-state, steady-state, and dynamics calculations. Section 5 discusses
NetKet’s implementation of spatial symmetries and symmetric neural quantum
states, which can be exploited to lower the size of the variational manifold
and to target excited states in nontrivial symmetry sectors. In Section 6, we
also show how to study a system with continuous degrees of freedom, such as
interacting particles in one or more spatial dimensions.

The final sections present detailed workflow examples of some of the more
common use cases of NetKet. In Section 7, we show how to study the ground
state and the excited state of a lattice Hamiltonian. Section 9 gives examples
of both unitary and Lindbladian dynamical simulations.

To conclude, Section 10 presents scaling benchmarks of NetKet running across
multiple devices and a performance comparison with jvmc [39], another library
similar in scope to NetKet.

Readers who are already familiar with the previous version of NetKet might be
especially interested in the variational state interface described in Section
3.1, which replaces what was called machine in NetKet 2 [35], the QGT
interface described in Section 3.5, algorithms for dynamics (Section 4.3 and
Section 9), and symmetry-aware NQS (Section 5).

JAX fluency. Using NetKet’s high-level interface and built-in neural network
architectures does not require the user to be familiar with jax and concepts
such as just-in-time compilation and automatic differentiation. However, when
defining custom classes such as neural network architectures, operators, or
Monte Carlo samplers, some proficiency with writing jax-compatible code will
be required. We refrain from discussing jax in detail and instead point the
reader towards its documentation at jax.readthedocs.io.

###  1.3 Installing NetKet

NetKet is a package written in pure Python; it requires a recent Python
version, currently at least version 3.7. Even though NetKet itself is
platform-agnostic, jax, its main dependency, only works on MacOS and Linux at
the time of writing.333In principle, jax runs on Windows, but users must
compile it themselves, which is not an easy process. Installing NetKet is
straightforward and can be achieved running the following line inside a python
environment:

⬇

1pip install --upgrade netket

To enable GPU support, Linux with a recent CUDA version is required and a
special version of jax must be installed. As the appropriate installation
procedure can change between jax versions, we refer the reader to the official
documentation444https://github.com/google/jax#installation for detailed
instructions.

NetKet by default does not make use of multiple CPUs that might be available
to the user. Exploiting multiple processors, or even running across multiple
nodes, requires MPI dependencies, which can be installed using the command

⬇

1pip install --upgrade "netket[mpi]"

These dependencies, namely mpi4py and mpi4jax, can only be installed if a
working MPI distribution is already available.

Once NetKet is installed, it can be imported in a Python session or script and
its version can be checked as

⬇

1>>> import netket as nk

2>>> print(nk.__version__)

33.5.0

We recommend that users use an up-to-date version when starting a new project.
In code listings, we will often refer to the netket module as nk for brevity.

NetKet also comes with a set of so-called experimental functionalities which
are packaged into the netket.experimental submodule which mirrors the
structure of the standard netket module. Experimental APIs are marked as such
because they are relatively young and we might want to change the function
names or options keyword arguments without guaranteeing backward compatibility
as we do for the rest of NetKet. In general, we import the experimental
submodule as follows

⬇

1>>> from netket import experimental as nkx

and use nkx as a shorthand for it.

##  2 Quantum-mechanical primitives

In general, when working with NetKet, the workflow is the following: first,
one defines the Hilbert space of the system (Section 2.1) and the Hamiltonian
or super-operator of interest (Section 2.2). Then, one builds a variational
state (Section 3.1), usually combining a neural-network model and a stochastic
sampler. In this section, we describe the first step in this process, namely,
how to define a quantum-mechanical system to be modeled.

###  2.1 Hilbert spaces

Hilbert-space objects determine the state space of a quantum system and a
specific choice of basis. Functionality related to Hilbert spaces is contained
in the nk.hilbert module; for brevity, we will often leave out the prefix
nk.hilbert in this section.

All implementations of Hilbert spaces derive from the class AbstractHilbert
and fall into two classes:

  * •

discrete Hilbert spaces, which inherit from the abstract class DiscreteHilbert
and include spin ( Spin), qubit ( Qubit), Fock ( Fock) as well as fermionic
orbitals (dubbed SpinOrbitalFermions) Hilbert spaces. Discrete spaces are
typically used to describe lattice systems. The lattice structure itself is,
however, not part of the Hilbert space class and can be defined separately.

  * •

continuous Hilbert spaces, which inherit from the abstract class
ContinuousHilbert. Currently, the only concrete continuous space provided by
NetKet is Particle.

Continuous Hilbert spaces are discussed in Section 6. A general discrete space
with $N$ sites has the structure

$\displaystyle\mathcal{H}_{\mathrm{discrete}}=\operatorname{span}\\{\ket{s_{0}}\otimes\cdots\otimes\ket{s_{N-1}}\mid
s_{i}\in\mathcal{L}_{i},i\in\\{0,\ldots,N-1\\}\\},$ (1)

where $\mathcal{L}_{i}$ is the set of local quantum numbers at site $i$ (e.g.,
$\mathcal{L}=\\{0,1\\}$ for a qubit, $\mathcal{L}=\\{\pm 1\\}$ for a
spin-$1/2$ system in the $\sigma^{z}$ basis, or
$\mathcal{L}=\\{0,1,\ldots,N_{\mathrm{max}}\\}$ for a Fock space with up to
$N_{\mathrm{max}}$ particles per site). Constraints on the allowed quantum
numbers are supported, resulting in Hilbert spaces that are subspaces of Eq.
(1). For example, Spin(1/2, total_sz=0) creates a spin-$1/2$ space which only
includes configurations $\\{s_{i}\\}$ that satisfy $\sum_{i=1}^{N}s_{i}=0.$
The corresponding basis states $\ket{s}$ span the zero-magnetization subspace.
Similarly, constraints on the total population in Fock spaces are also
supported.

Different spaces can be composed to create coupled systems by using the
exponent operator ( ∗∗) and the multiplication operator ( ∗). For example, the
code below creates the Hilbert space of a bosonic cavity with a cutoff of 10
particles at each site, coupled to 6 spin$-\frac{1}{2}$ degrees of freedom.

⬇

1>>> hi = nk.hilbert.Fock(10) ∗ nk.hilbert.Spin(1/2)∗∗6

2>>> print("Size of the hilbert space: ", hi.n_states)

3Size of the hilbert space: 704

4>>> print("Size of the basis: ", hi.size)

5Size of the basis: 7

6>>> hi.random_state(jax.random.PRNGKey(0), (2,))

7DeviceArray([[10., -1., 1., -1., 1., -1., -1.],

8 [ 9., 1., -1., -1., 1., -1., 1.]], dtype=float32)

All Hilbert objects can generate random basis elements through the function
random_state(rng_key, shape, dtype), which has the same signature as standard
random number generators in jax. The first argument is a jax random-generator
state as returned by jax.random.PRNGKey, while the other arguments specify the
number of output states and optionally the jax data type. In this example, an
array with two state vectors has been returned. The first entry of each
corresponds to the Fock space and is thus an integer in $\\{0,1,\ldots,10\\}$,
while the rest contains the spin quantum numbers.

Custom Hilbert spaces can be constructed by defining a class inheriting either
from ContinuousHilbert for continuous spaces or DiscreteHilbert for discrete
spaces. In the rest of the paper, we will always be working with discrete
Hilbert spaces unless stated otherwise.

NetKet also supports working with super-operators, such as the Liouvillian
used to define open quantum systems, and variational mixed states. The density
matrix is an element of the space of linear operators acting on a Hilbert
space, $\mathbb{B}(\mathcal{H})$. NetKet represents this space using the
Choi–Jamilkowski isomorphism [40, 41] convention
$\mathbb{B}(\mathcal{H})\sim\mathcal{H}\otimes\mathcal{H}$; this “doubled”
Hilbert space is implemented as DoubledHilbert. Doubled Hilbert spaces behave
largely similarly to standard Hilbert spaces, but their bases have double the
number of degrees of freedom; for example, super-operators can be defined
straightforwardly as operators acting on them.

###  2.2 Linear operators

NetKet is designed to allow users to work with large systems, beyond the
typically small system sizes that are accessible through exact diagonalization
techniques. In order to compute expectation values $\langle\hat{O}\rangle$ on
such large spaces, we must be able to efficiently represent the operators
$\hat{O}$ and work with their matrix elements $\bra{\sigma}\hat{O}\ket{\eta}$
without storing them in memory.

NetKet provides different implementations for the operators, tailored for
different use cases, which are available in the netket.operator submodule.
NetKet operators are always defined relative to a specific underlying Hilbert
space object and inherit from one of the abstract classes DiscreteOperator or
ContinuousOperator, depending on the classes of supported Hilbert spaces. We
defer the discussion of operators acting on a continuous space to Section 6
and focus on discrete-space operators in the remainder of this section.

An operator acting on a discrete space can be represented as a matrix with
some matrix elements $\bra{\sigma}\hat{O}\ket{\eta}$. As most of those
elements are zero in physical systems, a standard approach is to store the
operator as sparse matrices, a format that lowers the memory cost by only
storing non-zero entries. However, the number of non-zero matrix elements
still scales exponentially with the number of degrees of freedom, so sparse
matrices cannot scale to the thousands of lattice sites that we want to
support, either. For this reason, NetKet uses one of three custom formats to
represent operators:

  * •

LocalOperator

is an implementation that can efficiently represent sums of $K$-local
operators, that is, operators that only act nontrivially on a set of $K$
sites. The memory cost of this format grows linearly with the number of
operator terms and the number of degrees of freedom, but it scales
exponentially in $K$.

  * •

PauliStrings

is an implementation that efficiently represents a product of Pauli $X,Y,Z$
operators acting on the whole system. This format only works with qubit-like
Hilbert spaces, but it is extremely efficient and has negligible memory cost.

  * •

FermionOperator2nd

is an efficient implementation of second-quantized fermionic operators built
out of the on-site creation and annihilation operators
$f_{i}^{\dagger},f_{i}$. It works together with SpinOrbitalFermions and the
equivalent Fock spaces.

  * •

Special implementations like Ising, which hard-code the matrix elements of the
operator. Those are the most efficient, though they cannot be customized at
all.

The nk.operator submodule also contains ready-made implementations of commonly
used operators, such as Pauli matrices, bosonic ladder or projection
operators, and common Hamiltonians such as the Heisenberg, or the Bose–Hubbard
models.

####  2.2.1 Manipulating operators

Operators can be manipulated similarly to standard matrices: they can be
added, subtracted, and multiplied using standard Python operators. In the
example below we show how to construct the operator

$\displaystyle\hat{O}=\left(\hat{\sigma}^{x}_{0}+\hat{\sigma}^{x}_{1}\right)^{2}=2(\hat{\sigma}^{x}_{0}\hat{\sigma}^{x}_{1}+1).$
(2)

starting from the Pauli $X$ operator acting on the $i$-th site,
$\sigma^{x}_{i}$, given by the clearly-named function
nk.operator.spin.sigmax(hi, i):

⬇

1>>> hi = nk.hilbert.Spin(1/2)∗∗2

2>>> op = nk.operator.spin.sigmax(hi,0) + nk.operator.spin.sigmax(hi,1)

3>>> op = op ∗ op

4>>> op

5LocalOperator(dim=2, acting_on=[[0], [0, 1], [1]], constant=0, dtype=float64)

6>>> op.to_dense()

7array([[2., 0., 0., 2.],

8 [0., 2., 2., 0.],

9 [0., 2., 2., 0.],

10 [2., 0., 0., 2.]])

Note that each operator requires the Hilbert space object hi as well as the
specific sites it acts on as constructor arguments. In the last step (line 6),
we convert the operator into a dense matrix using the to_dense() method; it is
also possible to convert an operator into a SciPy sparse matrix using
to_sparse().

While it is possible to inspect those operators and (if the Hilbert space is
small enough) to convert them to dense matrices, NetKet’s operators are built
in order to support efficient row indexing, similar to row-sparse (CSR)
matrices. Given a basis vector $|\sigma\rangle$ in a Hilbert space, one can
efficiently query the list of basis states $\ket{\eta}$ and matrix elements
$O(\sigma,\eta)$ such that

$O(\sigma,\eta)=\bra{\sigma}\hat{O}\ket{\eta}\neq 0$ (3)

using the function operator.get_conn(sigma), which returns both the vector of
non-zero matrix elements and the corresponding list of indices $\ket{\eta}$,
stored as a matrix.555This querying is currently performed in Python code,
just-in-time compiled using numba [42], which runs on the CPU. If you run your
computations on a GPU with a small number of samples, this might introduce a
considerable slowdown. We are aware of this issue and plan to adapt our
operators to be indexed directly on the GPU in the future.

##  3 Variational quantum states

In this section, we first introduce the general interface of variational
states, which can be used to represent both pure states (vectors in the
Hilbert space) and mixed states (positive-definite density operators). We then
present how to define variational ansätze and the stochastic samplers needed
that generate Monte Carlo states.

###  3.1 Abstract interface

A variational state describes a parametrized quantum state that depends on a
(possibly large) set of variational parameters $\theta$. The quantum state can
be either pure (denoted as $\ket{\psi_{\theta}}$) or mixed (written as a
density matrix $\hat{\rho}_{\theta}$). NetKet defines an abstract interface,
netket.vqs.VariationalState, for such objects; all classes that implement this
interface will automatically work with all the high-level drivers (e.g.,
ground-state optimization or time-dependent variational dynamics) discussed in
Section 4. The VariationalState interface is relatively simple, as it has only
four requirements:

  * •

The parameters $\theta$ of the variational state are exposed through the
attribute parameters and should be stored as an array or a nested dictionary
of arrays.

  * •

The expectation value $\langle\hat{A}\rangle_{\theta}$ of an operator
$\hat{A}$ can be computed or estimated by the method expect.

  * •

The gradient of an expectation value with respect to the variational
parameters, $\partial\langle\hat{A}\rangle_{\theta}/\partial\theta_{j}$, is
computed by the method expect_and_grad666For complex-valued parameters
$\theta_{j}\in\mathbb{C}$, expect_and_grad returns the conjugate gradient
$\partial\langle\hat{A}\rangle_{\theta}/\partial\theta_{j}^{*}$ instead. This
is done because the conjugate gradient corresponds to the direction of
steepest ascent when optimizing a real-valued function of complex variables
[43]..

  * •

The quantum geometric tensor (Section 3.5) of a variational state can be
constructed with the method quantum_geometric_tensor.

At the time of writing, NetKet exposes three types of variational state:

  * •

nk.vqs.ExactState

represents a variational pure state $|\psi_{\theta}\rangle$ and computes
expectation values, gradients and the geometric tensor by performing exact
summation over the full Hilbert space.

  * •

nk.vqs.MCState

(short for Monte Carlo state) represents a variational pure state and computes
expectation values, gradients and the geometric tensor by performing Markov
chain Monte Carlo (MCMC) sampling over the Hilbert space.

  * •

nk.vqs.MCMixedState

represents a variational mixed state and computes expectation values by
sampling diagonal entries of the density matrix.

Variational states based on Monte Carlo sampling are the main tools that we
expose to users, together with a wide variety of high-performance Monte Carlo
samplers. More details about stochastic estimates and Monte Carlo sampling
will be discussed in Section 3.3 and Section 3.4.

##### Dispatch and algorithm selection.

With three different types of variational state and several different
operators supported, it is hard to write a well-performing algorithm that
works with all possible combinations of types that users might require. In
order not to sacrifice performance for generic algorithms, NetKet uses the
approach of multiple dispatch based on the plum module [38]. Combined with
jax’s just-in-time compilation, this solution bears a strong resemblance to
the approach commonly used in the Julia language [6].

Every time the user calls VariationalState.expect or .expect_and_grad, the
types of the variational state and the operator are used to select the most
specific algorithm that applies to those two types. This allows NetKet to
provide generic algorithms that work for all operators, but keeps it easy to
supply custom algorithms for specific operator types if desired.

This mechanism is also exposed to users: it is possible to override the
algorithms used by NetKet to compute expectation values and gradients without
modifying the source code of NetKet but simply by defining new dispatch rules
using the syntax shown below.

⬇

<EMAIL_ADDRESS>
2def expect(vstate: MCState, operator: Ising):

3 # more efficient implementation than default one

4 #

5 # expectation_value = ...

6 #

7 return expectation_value

###  3.2 Defining the variational ansatz

The main feature defining a variational state is the parameter-dependent
mapping of an input configuration to the corresponding probability amplitude
or, in other words, the quantum wave function (for pure states)

$\displaystyle(\theta,s)$
$\displaystyle\mapsto\psi_{\theta}(s)=\innerproduct{s}{\psi_{\theta}}$ (4) or
quantum density matrix (for mixed states) $\displaystyle(\theta,s,s^{\prime})$
$\displaystyle\mapsto\rho_{\theta}(s,s^{\prime})=\bra{s}\hat{\rho}_{\theta}\ket{s^{\prime}}.$
(5)

In NetKet, this mapping is called a _model_ (of the quantum state).777The
notion of “model” in NetKet 3 is related to the “machine” classes in NetKet 2
[35]. However, while NetKet 2 machines both define the mapping (4) and store
the current parameters, this has been decoupled in NetKet 3\. The model only
specifies the mapping, while the parameters are stored in the variational
state classes. In the case of NQS, the model is given by a neural network. For
defining models, NetKet primarily relies on flax [44], a jax-based neural-
network library888While our primary focus has been the support of flax, NetKet
can in principle be used with any jax-compatible neural network model. For
example, NetKet currently includes a compatibility layer which ensures that
models defined using the Haiku framework by DeepMind [45] will work
automatically as well. Furthermore, any model represented by a pair of init
and apply functions (as used, e.g., in the stax framework included with jax)
is also supported.. Ansätze are implemented as flax modules that map input
configurations (the structure of which is determined by the Hilbert space) to
the corresponding log-probability amplitudes. For example, pure quantum states
are evaluated as

$\displaystyle\ln\psi_{\theta}(s)=\definecolor{tcbcolback}{rgb}{0.95,0.95,0.95}\definecolor{tcbcolframe}{rgb}{1,1,1}\definecolor{tcbcol@origin}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\leavevmode\hbox
to202.22pt{\vbox to17.06pt{\pgfpicture\makeatletter\hbox{\hskip 0.0pt\lower
0.0pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{{}{}{}{}\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{1,1,1}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{0.0pt}{4.26788pt}\pgfsys@lineto{0.0pt}{12.79318pt}\pgfsys@curveto{0.0pt}{15.1503pt}{1.91077pt}{17.06107pt}{4.26788pt}{17.06107pt}\pgfsys@lineto{197.95413pt}{17.06107pt}\pgfsys@curveto{200.31125pt}{17.06107pt}{202.22202pt}{15.1503pt}{202.22202pt}{12.79318pt}\pgfsys@lineto{202.22202pt}{4.26788pt}\pgfsys@curveto{202.22202pt}{1.91077pt}{200.31125pt}{0.0pt}{197.95413pt}{0.0pt}\pgfsys@lineto{4.26788pt}{0.0pt}\pgfsys@curveto{1.91077pt}{0.0pt}{0.0pt}{1.91077pt}{0.0pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{0.95,0.95,0.95}\pgfsys@color@rgb@fill{0.95}{0.95}{0.95}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{1.42262pt}{4.26788pt}\pgfsys@lineto{1.42262pt}{12.79318pt}\pgfsys@curveto{1.42262pt}{14.3646pt}{2.69647pt}{15.63844pt}{4.26788pt}{15.63844pt}\pgfsys@lineto{197.95413pt}{15.63844pt}\pgfsys@curveto{199.52554pt}{15.63844pt}{200.7994pt}{14.3646pt}{200.7994pt}{12.79318pt}\pgfsys@lineto{200.7994pt}{4.26788pt}\pgfsys@curveto{200.7994pt}{2.69647pt}{199.52554pt}{1.42262pt}{197.95413pt}{1.42262pt}\pgfsys@lineto{4.26788pt}{1.42262pt}\pgfsys@curveto{2.69647pt}{1.42262pt}{1.42262pt}{2.69647pt}{1.42262pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{{}}{{}}{{}}{{}}{{}}{{}}{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{4.42262pt}{6.24835pt}\pgfsys@invoke{
}\hbox{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\hbox{\set@color{\ignorespaces\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language{\@listingGroup{ltx_lst_identifier}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}module}}.{\@listingGroup{ltx_lst_identifier}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}apply}}}}}}}}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}(\theta,s).$ (6)

The use of log-amplitudes has the benefit that the log-derivatives
$\partial\ln\psi_{\theta}(s)/\partial\theta_{j}$, often needed in variational
optimization algorithms, are directly available through automatic
differentiation of the model. It also makes it easier for the model to learn
amplitudes with absolute values ranging over several orders of magnitude,
which is common for many types of quantum states.

Real and complex amplitudes. NetKet supports both real-valued and complex-
valued model outputs. However, since model outputs correspond to log-
amplitudes, real-valued networks can only represent states that have
exclusively non-negative amplitudes,
$\ln\psi_{\theta}(s)\in\mathbb{R}\Rightarrow\psi_{\theta}(s)\geq 0.$

Since the input configurations $s$ are real, in many pre-defined NetKet models
the data type of the network parameters
($\theta\in\mathbb{R}^{{N_{\mathrm{p}}}}$ or
$\theta\in\mathbb{C}^{{N_{\mathrm{p}}}}$) determines whether an ansatz
represents a general or a real non-negative state. This should be kept in mind
in particular when optimizing Hamiltonians with ground states that can have
negative amplitudes.

####  3.2.1 Custom models using Flax

The recommended way to define a custom module is to subclass flax.linen.Module
and to provide a custom implementation of the __call__ method. As an example,
we define a simple one-layer NQS with a wave function of the form

$\displaystyle\ln\psi(s)=\sum_{j=1}^{M}\tanh[Ws+b]_{j}.$ (7)

with the number of visible units $N$ matching the number of physical sites, a
number of hidden units $M$, and complex parameters $W\in\mathbb{C}^{M\times
N}$ (the weight matrix) and $b\in\mathbb{C}^{M}$ (the bias vector). Using
NetKet and flax, this ansatz can be implemented as follows:

⬇

1import netket as nk

2import jax.numpy as jnp

3import flax

4import flax.linen as nn

5

6class OneLayerNQS(nn.Module):

7 # Module hyperparameter:

8 n_hidden_units: int

9

10 @nn.compact

11 def __call__(self, s):

12 n_visible_units = s.shape[-1]

13 # define parameters

14 # the arguments are: name, initializer, shape, dtype

15 W = self.param(

16 "weights",

17 nn.initializers.normal(),

18 (self.n_hidden_units, n_visible_units),

19 jnp.complex128,

20 )

21 b = self.param(

22 "bias",

23 nn.initializers.normal(),

24 (self.n_hidden_units,),

25 jnp.complex128,

26 )

27

28 # multiply with weight matrix over last dimension of s

29 y = jnp.einsum("ij,...j", W, s)

30 # add bias

31 y += b

32 # apply tanh activation and sum

33 y = jnp.sum(jnp.tanh(y), axis=-1)

34

35 return y

The decorator flax.linen.compact used on __call__ (line 10) makes it possible
to define the network parameters directly in the body of the call function via
self.param as done above (lines 15 and 21). For performance reasons, the input
to the module is batched. This means that, instead of passing a single array
of quantum numbers $s$ of size $N$, a batch of multiple state vectors is
passed as a matrix of shape (batch_size, N). Therefore, operations like the
sum over all feature indices in the example above need to be explicitly
performed over the last axis. Just-in-time compilation. Note that the network
will be just-in-time (JIT) compiled to efficient machine code for the target
device (CPU, GPU, or TPU) using jax.jit, which means that all code inside the
__call__ method needs to written in a way compatible with jax.jit.

In particular, users should use jax.numpy for NumPy calls that need to happen
at runtime; explicit Python control flow, such as for loops and if statements,
should also be avoided, unless one explicitly wants to have them evaluate once
at compile time. We refer users to the jax documentation for further
information on how to write efficient JIT-compatible code.

The module defined above can be used by first initializing the parameters
using module.init and then computing log-amplitudes through module.apply:

⬇

1>>> module = OneLayerNQS(n_hidden_units=16)

2# init takes two arguments, a PRNG key for random initialization

3# and a dummy array used to determine the input shape

4# (here with a batch size of one):

5>>> params = module.init(nk.jax.PRNGKey(0), jnp.zeros((1, 8)))

6>>> module.apply(params, jnp.array([[-1, 1, -1, 1, -1, 1, -1, 1]]))

7DeviceArray([-0.00047843+0.07939122j], dtype=complex128)

####  3.2.2 Network parametrization and pytrees

##### Parameter data types.

NetKet supports models with both real-valued and complex-valued network
parameters. The data type of the parameters does not determine the output
type. It is possible to define a model with real parameters that produces
complex output. A simple example is the sum of two real-valued and real-
parameter networks, representing real and imaginary part of the log-amplitudes
(and thus phase and absolute value of the wave function) [31, 29]:

$\displaystyle\ln\psi_{(\theta,\eta)}(s)=f_{\theta}(s)+ig_{\eta}(s)\quad\Leftrightarrow\quad|\psi_{(\theta,\eta)}(s)|=\exp[f_{\theta}(s)],\quad\operatorname{arg}\psi_{(\theta,\eta)}(s)=g_{\eta}(s)$
(8)

(where all $\theta_{i},\eta_{i}\in\mathbb{R}$ and $f(s),g(s)\in\mathbb{R}$).

More generally, any model with ${N_{\mathrm{p}}}$ complex parameters
$\theta=\alpha+i\beta$ can be represented by a model with $2{N_{\mathrm{p}}}$
real parameters $(\alpha,\beta)$. While these parametrizations are formally
equivalent, the choice of complex parameters can be particularly useful in the
case where the variational mapping is holomorphic or, equivalently, complex
differentiable with respect to $\theta$. This is the case for many standard
network architectures such as RBMs or feed-forward networks, since both linear
transformations and typical activation functions are holomorphic (such as
$\tanh$, $\cosh$, and their logarithms) or piecewise holomorphic (such as
ReLU), which is sufficient in practice. Note, however, that there are also
common architectures, such as autoregressive networks, that are not
holomorphic. In the holomorphic case, the computational cost of
differentiating the model, e.g., to compute the quantum geometric tensor
(Section 3.5), can be reduced by exploiting the Cauchy–Riemann equations [46],

$\displaystyle
i\nabla_{\alpha}\psi_{(\alpha,\beta)}(s)=\nabla_{\beta}\psi_{(\alpha,\beta)}(s).$
(9)

Note that NetKet generally supports models with _arbitrary parametrizations_
(i.e., real and both holomorphic and non-holomorphic complex
parametrizations). The default assumption is that models with complex weights
are non-holomorphic, but some objects (most notably the quantum geometric
tensor) accept a flag holomorphic=True to enable a more efficient code path
for holomorphic networks. It is the user’s responsibility to only set
holomorphic=True for models that are, in fact, holomorphic. If this is
incorrectly specified, NetKet code may give incorrect results. To check
whether a specific architecture is holomorphic, one can verify the condition

$\displaystyle\partial\psi_{\theta}(s)/\partial\theta_{j}^{*}=(\partial/\partial\alpha_{j}+i\partial/\partial\beta_{j})\psi_{\theta}(s)=0,$
(10)

which is equivalent to Eq. (9).

##### Pytrees.

In NetKet, model parameters do not need to be stored as a contiguous vector.
Instead, models can support any collection of parameters that forms a so-
called _pytree_. Pytree is jax terminology for collections of numerical
tensors stored as the leaf nodes inside layers of nested standard Python
containers (such as lists, tuples, and dictionaries).999See
https://github.com/google/jax/blob/jax-v0.2.28/docs/pytrees.md for a detailed
introduction of pytrees. Any object that is not itself a pytree, in particular
NumPy or jax arrays, is referred to as a _leaf_. Networks defined as flax
modules store their parameters in a (potentially nested) dictionary, which
provides name-based access to the network parameters.101010Specifically, flax
stores networks parameters in an immutable FrozenDict object, which otherwise
has the same semantics as a standard Python dictionary and, in particular, is
also a valid pytree. The parameters can be modified by converting to a
standard mutable dict via flax.core.unfreeze(params). For the OneLayerNQS
defined above, the parameter pytree has the structure:

⬇

1# For readability, the actual array data has been replaced with ... below.

2>>> print(params)

3FrozenDict({

4 params: {

5 weights: DeviceArray(..., dtype=complex128),

6 bias: DeviceArray(..., dtype=complex128),

7 },

8})

The names of the entries in the parameter dictionary correspond to those given
in the param call when defining the model. NetKet functions often work
directly with both plain arrays and pytrees of arrays. Furthermore, any Python
function can be applied to the leaves of a pytree using jax.tree_map. For
example, the following code prints a pytree containing the shape of each leaf
of params, preserving the nested dictionary structure:

⬇

1>>> print(jax.tree_map(jnp.shape, params))

2FrozenDict({

3 params: {

4 weights: (16, 8),

5 bias: (16,),

6 },

7})

Functions accepting multiple leaves as arguments can be mapped over the
corresponding number of pytrees (with compatible structure) using
jax.tree_map.

For example, the difference of two parameter pytrees of the same model can be
computed using delta = jax.tree_map(lambda a, b: a - b, params1, params2).
NetKet provides an additional set of utility functions to perform linear
algebra operations on such pytrees in the nk.jax submodule.

####  3.2.3 Pre-defined ansätze included with NetKet

Name |  NetKet class | References
---|---|---
Jastrow ansatz |  Jastrow |  [47, 48]
Restricted Boltzmann machine (RBM) |  RBM , RBMMultiVal, RBMModPhase |  [8]
Symmetric RBM |  RBMSymm |  [8]
Group-Equivariant Convolutional Neural Network |  GCNN |  Section 5 [49]
Autoregressive Neural Network |  ARNNDense , ARNNConv1D, ARNNConv2D, FastARNNConv1D, FastARNNConv2D |  [50]
Neural Density Matrix |  NDM |  [51, 24]
|  |
Table 1: List of models included in NetKet’s nk.models submodule, together
with relevant references.

NetKet provides a collection of pre-defined modules under nk.models, which
allow quick access to many commonly used NQS architectures (Table 1):

  * •

Jastrow: The Jastrow ansatz [47, 48] is an extremely simple yet effective
many-body ansatz that can capture some inter-particle correlations. The log-
wavefunction is the linear function
$\log\psi(\mathbf{\sigma})=\sum_{i}\sigma_{i}W_{i,j}\sigma_{j}$. Evaluation of
this ansatz is very fast but it is also the least powerful model implemented
in NetKet;

  * •

RBM: The restricted Boltzmann machine (RBM) ansatz is composed by a dense
layer followed by a nonlinearity. If the Hilbert space has $N$ degrees of
freedom of size $d$, RBM has $\alpha N$ features in the dense layer. This
ansatz requires param_dtype=complex to represent states that are non-positive
valued. RBMMultiVal is a one-hot encoding layer followed by an RBM with
$\alpha dN$ features in its dense layer. Finally, RBMModPhase consists of two
real-valued RBMs that encode respectively the modulus and phase of the
wavefunction as
$\log\psi(\mathbf{\sigma})=\definecolor{tcbcolback}{rgb}{0.95,0.95,0.95}\definecolor{tcbcolframe}{rgb}{1,1,1}\definecolor{tcbcol@origin}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\leavevmode\hbox
to105.84pt{\vbox to17.06pt{\pgfpicture\makeatletter\hbox{\hskip 0.0pt\lower
0.0pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{{}{}{}{}\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{1,1,1}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{0.0pt}{4.26788pt}\pgfsys@lineto{0.0pt}{12.79318pt}\pgfsys@curveto{0.0pt}{15.1503pt}{1.91077pt}{17.06107pt}{4.26788pt}{17.06107pt}\pgfsys@lineto{101.57007pt}{17.06107pt}\pgfsys@curveto{103.92719pt}{17.06107pt}{105.83795pt}{15.1503pt}{105.83795pt}{12.79318pt}\pgfsys@lineto{105.83795pt}{4.26788pt}\pgfsys@curveto{105.83795pt}{1.91077pt}{103.92719pt}{0.0pt}{101.57007pt}{0.0pt}\pgfsys@lineto{4.26788pt}{0.0pt}\pgfsys@curveto{1.91077pt}{0.0pt}{0.0pt}{1.91077pt}{0.0pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{0.95,0.95,0.95}\pgfsys@color@rgb@fill{0.95}{0.95}{0.95}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{1.42262pt}{4.26788pt}\pgfsys@lineto{1.42262pt}{12.79318pt}\pgfsys@curveto{1.42262pt}{14.3646pt}{2.69647pt}{15.63844pt}{4.26788pt}{15.63844pt}\pgfsys@lineto{101.57007pt}{15.63844pt}\pgfsys@curveto{103.14148pt}{15.63844pt}{104.41533pt}{14.3646pt}{104.41533pt}{12.79318pt}\pgfsys@lineto{104.41533pt}{4.26788pt}\pgfsys@curveto{104.41533pt}{2.69647pt}{103.14148pt}{1.42262pt}{101.57007pt}{1.42262pt}\pgfsys@lineto{4.26788pt}{1.42262pt}\pgfsys@curveto{2.69647pt}{1.42262pt}{1.42262pt}{2.69647pt}{1.42262pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{{}}{{}}{{}}{{}}{{}}{{}}{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{4.42262pt}{6.24835pt}\pgfsys@invoke{
}\hbox{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\hbox{\set@color{\ignorespaces\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language{\@listingGroup{ltx_lst_keywords3}{\color[rgb]{0.25,0.3828125,0.734375}\definecolor[named]{pgfstrokecolor}{rgb}{0.25,0.3828125,0.734375}RBM}}}}}}}}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}(\mathbf{\sigma})+i\definecolor{tcbcolback}{rgb}{0.95,0.95,0.95}\definecolor{tcbcolframe}{rgb}{1,1,1}\definecolor{tcbcol@origin}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\definecolor{.}{rgb}{0,0,0}\leavevmode\hbox
to105.84pt{\vbox to17.06pt{\pgfpicture\makeatletter\hbox{\hskip 0.0pt\lower
0.0pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{{}{}{}{}\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{1,1,1}\pgfsys@color@gray@fill{1}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{0.0pt}{4.26788pt}\pgfsys@lineto{0.0pt}{12.79318pt}\pgfsys@curveto{0.0pt}{15.1503pt}{1.91077pt}{17.06107pt}{4.26788pt}{17.06107pt}\pgfsys@lineto{101.57007pt}{17.06107pt}\pgfsys@curveto{103.92719pt}{17.06107pt}{105.83795pt}{15.1503pt}{105.83795pt}{12.79318pt}\pgfsys@lineto{105.83795pt}{4.26788pt}\pgfsys@curveto{105.83795pt}{1.91077pt}{103.92719pt}{0.0pt}{101.57007pt}{0.0pt}\pgfsys@lineto{4.26788pt}{0.0pt}\pgfsys@curveto{1.91077pt}{0.0pt}{0.0pt}{1.91077pt}{0.0pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}{}{}{}{}{}{}{}{}\definecolor[named]{pgffillcolor}{rgb}{0.95,0.95,0.95}\pgfsys@color@rgb@fill{0.95}{0.95}{0.95}\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}{{}{}{{}}}{{}{}{{}}}{}{}\pgfsys@moveto{1.42262pt}{4.26788pt}\pgfsys@lineto{1.42262pt}{12.79318pt}\pgfsys@curveto{1.42262pt}{14.3646pt}{2.69647pt}{15.63844pt}{4.26788pt}{15.63844pt}\pgfsys@lineto{101.57007pt}{15.63844pt}\pgfsys@curveto{103.14148pt}{15.63844pt}{104.41533pt}{14.3646pt}{104.41533pt}{12.79318pt}\pgfsys@lineto{104.41533pt}{4.26788pt}\pgfsys@curveto{104.41533pt}{2.69647pt}{103.14148pt}{1.42262pt}{101.57007pt}{1.42262pt}\pgfsys@lineto{4.26788pt}{1.42262pt}\pgfsys@curveto{2.69647pt}{1.42262pt}{1.42262pt}{2.69647pt}{1.42262pt}{4.26788pt}\pgfsys@closepath\pgfsys@fill\pgfsys@invoke{
}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@fill@opacity{1.0}\pgfsys@invoke{
}{{{}}{{}}{{}}{{}}{{}}{{}}{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{4.42262pt}{6.24835pt}\pgfsys@invoke{
}\hbox{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\hbox{\set@color{\ignorespaces\leavevmode\lstinline{{\lst@@@set@language\lst@@@set@numbers\lst@@@set@frame\lst@@@set@rulecolor\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language\lst@@@set@rulecolor\lst@@@set@frame\lst@@@set@numbers\lst@@@set@language{\@listingGroup{ltx_lst_keywords3}{\color[rgb]{0.25,0.3828125,0.734375}\definecolor[named]{pgfstrokecolor}{rgb}{0.25,0.3828125,0.734375}RBM}}}}}}}}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}(\mathbf{\sigma})$.
This ansatz only supports real parameters. If considering Hilbert spaces with
local dimension $d>2$, plain RBMs usually require a very large feature density
$\alpha$ and RBMMultiVals perform better.

  * •

RBMSymm: A symmetry-invariant RBM. Only symmetry groups that can be
represented as permutations of the computational basis are supported (see
Section 5). This architecture has fewer parameters than an RBM, but it is more
expensive to evaluate. It requires param_dtype=complex to represent states
that are non-positive valued.

  * •

GCNN: A symmetry-equivariant feed-forward network (see Section 5.2). Only
symmetry groups that can be represented as permutations of the computational
basis are supported. This model is much more complex and computationally
intensive than RBMSymm, but can also lead to more accurate results. It can
also be used to target an excited state of a lattice Hamiltonian. When working
with states that are real but non-positive, one can use real parameters
together with complex_output=True. If the states are to have a complex phase,
param_dtype=complex is required.

  * •

Autoregressive networks: ARNNs are models that can produce uncorrelated
samples when sampled with nk.sampler.ARNNSampler. Those architectures can be
efficiently sampled on GPUs, but they are much more expensive than traditional
RBMs.

  * •

NDM: A positive-semidefinite density matrix ansatz, comprised of a component
describing the pure part and one describing the mixed part of the state. The
pure part is equivalent to an RBM with feature density $\alpha$, while the
mixed part is an RBM with feature density $\beta$. This network only supports
real parameters.

####  3.2.4 Custom layers included with Flax and NetKet

The nk.nn submodule contains generic modules such as masked dense, masked
convolutional and symmetric layers to be used as building blocks for custom
neural networks. Those layers are complementary to those provided by flax and
can be combined together to develop novel neural-network
architectures.111111In the past flax had minor issues with complex numbers and
therefore NetKet included versions of some standard layers, such as Dense and
Conv, that handle complex numbers properly. Starting with flax version 0.5,
released in May 2022, those issues have been addressed and we now recommend
the use of flax layers also with complex numbers.

As an example, a multi-layer NQS with two convolutional and one final dense
layer acting as a weighted sum can be defined as follows:

⬇

1class MultiLayerCNN(nn.Module):

2 features1: int

3 features2: int

4 kernel_size: int

5

6 @nn.compact

7 def __call__(self, s):

8 # define layers

9 layer1 = nn.Conv(

10 features=self.features1,

11 kernel_size=self.kernel_size,

12 )

13 layer2 = nn.Conv(

14 features=self.features2,

15 kernel_size=self.kernel_size,

16 )

17 weighted_sum = nn.Dense(features=1)

18

19 # apply layers and tanh activations

20 y = jnp.tanh(layer1(s))

21 y = jnp.tanh(layer2(y))

22 y = weighted_sum(y)

23 # last axis only has one entry, so we just return that

24 # but keep the batch dimension

25 return y[..., 0]

flax network layers are available from the flax.linen submodule (imported as
nn in the example above), NetKet layers from netket.nn.

###  3.3 Estimating observables

For any variational ansatz, it is crucial to also have efficient algorithms
for computing quantities of interest, in particular observables and their
gradients. Since evaluating the wave function on all configurations is
infeasible for larger Hilbert spaces, NQS approaches rely on Monte Carlo
sampling of quantum expectation values.

##### Pure states.

The quantum expectation value of an operator $\hat{A}$ on a non-normalized
pure state $|\psi\rangle$ can be written as a classical expectation value over
the Born distribution $p(s)\propto|\psi(s)|^{2}$ using the identity

$\displaystyle\langle\hat{A}\rangle=\frac{\langle\psi|\hat{A}|\psi\rangle}{\innerproduct{\psi}{\psi}}=\sum_{s}\frac{|\psi(s)|^{2}}{\innerproduct{\psi}{\psi}}\tilde{A}(s)=\sum_{s}p(s)\tilde{A}(s)=\mathbb{E}[\tilde{A}],$
(11)

where $\tilde{A}$ is the _local estimator_

$\displaystyle\tilde{A}(s)=\frac{\langle
s|\hat{A}|\psi\rangle}{\innerproduct{s}{\psi}}=\sum_{s^{\prime}}\frac{\psi(s^{\prime})}{\psi(s)}\langle
s|\hat{A}|s^{\prime}\rangle,$ (12)

also known as the _local energy_ when $\hat{A}$ is the Hamiltonian [52]. Even
though the sum in Eq. 12 runs over the full Hilbert space basis, the local
estimator can be efficiently computed if the operator is sufficiently sparse
in the given basis, i.e., all but a tractable number of matrix elements
$\langle s|\hat{A}|s^{\prime}\rangle$ are zero. Thus, an efficient algorithm
is required that, given $s$, yields all connected configurations $s^{\prime}$
together with their respective matrix elements, as described in Section 2.2.
Given the derivatives of the log-amplitudes

$\displaystyle
O_{i}(s)=\frac{\partial\ln\psi_{\theta}(s)}{\partial\theta_{i}},$ (13)

gradients of expectation values can also be evaluated. Define the _force
vector_ as the covariance

$\displaystyle\tilde{f}_{i}=\operatorname{Cov}[O_{i},\tilde{A}]=\mathbb{E}[O_{i}^{*}(\tilde{A}-\mathbb{E}[\tilde{A}])].$
(14)

Then, if $\theta_{i}\in\mathbb{R}$ is a real-valued parameter,

$\displaystyle\frac{\partial\langle\hat{A}\rangle}{\partial\theta_{i}}=2\real[\tilde{f}_{i}].$
(15)

If $\theta_{i}\in\mathbb{C}$ and the mapping
$\theta_{i}\mapsto\psi_{\theta}(s)$ is complex differentiable (holomorphic),

$\displaystyle\frac{\partial\langle\hat{A}\rangle}{\partial\theta_{i}^{*}}=\tilde{f}_{i}.$
(16)

In case of a non-holomorphic mapping, $\real[\theta_{i}]$ and
$\imaginary[\theta_{i}]$ can be treated as two independent real parameters and
Eq. (15) applies to each.

The required classical expectation values are then estimated by averaging over
a sequence $\\{s_{i}\\}_{i=1}^{N_{\mathrm{s}}}$ of configurations distributed
according to the Born distribution $p(s)\propto|\psi(s)|^{2}$; e.g., Eq. 11
becomes

$\displaystyle\mathbb{E}[\tilde{A}]\approx\frac{1}{N_{\mathrm{s}}}\sum_{i=1}^{N_{\mathrm{s}}}\tilde{A}(s).$
(17)

For some models, in particular autoregressive neural networks [50], one can
efficiently draw samples from the Born distribution directly. For a general
ansatz, however, this is not possible and Markov-chain Monte Carlo (MCMC)
sampling methods [52] must be used: these generate a sequence (Markov chain)
of samples that asymptotically follows the Born distribution. Such a chain can
be generated using the Metropolis–Hastings algorithm [53], which is
implemented in NetKet’s sampler interface, described in the next section.

##### Mixed states.

When evaluating observables for mixed states, it is possible to exploit a
slightly different identity,

$\displaystyle\langle\hat{A}\rangle=\frac{\Tr[\hat{\rho}\hat{A}]}{\Tr[\hat{\rho}]}=\sum_{s}\frac{\rho(s,s)}{\Tr[\hat{\rho}]}\tilde{A}_{\rho}(s)=\mathbb{E}[\tilde{A}_{\rho}],$
(18)

which rewrites the quantum expectation value as a classical expectation over
the probability distribution defined by the diagonal of the density matrix
$p(s)\propto\rho(s,s)$. Here, $\tilde{A}_{\rho}$ denotes the local estimator
of the observable over a mixed state,

$\displaystyle\tilde{A}_{\rho}(s)=\frac{\langle
s|\hat{\rho}\hat{A}|s\rangle}{\bra{s}\hat{\rho}\ket{s}}=\sum_{s^{\prime}}\frac{\rho(s,s^{\prime})}{\rho(s,s)}\langle
s^{\prime}|\hat{A}|s\rangle.$ (19)

It is then possible to follow the same procedures detailed in the previous
paragraph for pure states to compute the gradient of an expectation value of
an operator over a mixed state by replacing the probability distribution over
which the average is computed and the local estimator.

####  3.3.1 Reducing memory usage with chunking

The number of variational state evaluations required to compute the local
estimators (12) typically scales superlinearly121212The exact scaling depends
on the sparsity of the observable in the computational basis (which in a
lattice model primarily depends on the locality of the operator and the
dimension of the lattice).  in the number of sites $N$. For optimal
performance, NetKet by default performs those evaluations in a single call
using batched inputs. However, for large Hilbert spaces or very deep models it
might be impossible to fit all required intermediate buffers into the
available memory, leading to out-of-memory errors. This is encountered
particularly often in calculations on GPUs, which have more limited memory.

To avoid those errors, NetKet’s nk.vqs.VariationalState exposes an attribute
called chunk_size, which controls the maximum number of configurations for
which a model is evaluated at the same time131313The chunk size can be
specified at model construction and freely changed later. Chunking can also be
disabled at any time by setting VariationalState.chunk_size = None. . The
chunk size effectively bounds the maximum amount of memory required to
evaluate the variational function at the expense of an increased computational
cost in some operations involving the derivatives of the model. For this
reason, we suggest using the largest chunk size that fits in memory.

Chunking is supported for the majority of operations, such as computing
expectation values and their gradients, as well as the evaluation the quantum
geometric tensor. If a chunk size is specified but an operation does not
support it, NetKet will print a warning and attempt to perform the operation
without chunking.

###  3.4 Monte Carlo samplers

The sampling algorithm used to obtain a sequence of configurations from the
probability distribution defined by the variational ansatz is specified by
sampler classes inheriting from nk.sampler.AbstractSampler. Following the
purely functional design of jax, we define the sampler to be a stateless
collection of settings and parameters, while storing all mutable state such as
the PRNG key and the statistics of acceptances in an immutable sampler state
object. Both the sampler and the sampler state are stored in the variational
state, but they can be used independently, as they are decoupled from the rest
of NetKet.

The Metropolis–Hastings algorithm is used to generate samples from an
arbitrary probability distribution. In each step, it suggests a transition
from the current configuration $s$ to a proposed configuration $s^{\prime}$.
The proposal is accepted with probability

$P_{\text{acc}}(s\rightarrow
s^{\prime})=\min\left(1,\frac{P(s^{\prime})}{P(s)}\frac{g(s\mid
s^{\prime})}{g(s^{\prime}\mid s)}\right),$ (20)

where $P$ is the distribution being sampled from and $g(s^{\prime}\mid s)$ is
the conditional probability of proposing $s^{\prime}$ given the current $s$.
We use $L(s,s^{\prime})=\log[g(s\mid s^{\prime})/g(s^{\prime}\mid s)]$ to
denote the correcting factor to the log probability due to the transition
kernel. This factor is needed for asymmetric kernels that might propose one
move with higher probability than its reverse. Simple kernels, such as a local
spin flip or exchange, are symmetric, therefore
$L(s,s^{\prime})=L(s^{\prime},s)=1$, but other proposals, such as Hamiltonian
sampling, are not necessarily symmetric and need this factor.

At the time of writing, NetKet exposes four types of rules to use with the
Metropolis sampler: MetropolisLocal, which changes one discrete local degree
of freedom in each transition; MetropolisExchange, which exchanges two local
degrees of freedom respecting a conserved quantity (e.g., total particle
number or magnetization); MetropolisHamiltonian, which transitions the
configuration according to the off-diagonal elements of the Hamiltonian; and
MetropolisGaussian, which moves a configuration with continuous degrees of
freedom according to a Gaussian distribution.

The different transition kernels in these samplers are represented by
MetropolisRule objects. To define a Metropolis sampling algorithm with a new
transition kernel, one only needs to subclass MetropolisRule and implement the
transition method, which gives $s^{\prime}$ and $L(s,s^{\prime})$ in each
transition. For example, the following transition rule changes the local
degree of freedom on two sites at a time:

⬇

1from netket.hilbert.random import flip_state

2from netket.sampler import MetropolisRule

3from netket.utils.struct import dataclass

4

5# To be jax-compatible, it must be a dataclass

6@dataclass

7class TwoLocalRule(MetropolisRule):

8 def transition(rule, sampler, machine, parameters, state, key,
$\mathtt{\sigma}$):

9 # Deduce the number of MCMC chains from input shape

10 n_chains = $\mathtt{\sigma}$.shape[0]

11 # Load the Hilbert space of the sampler

12 hilb = sampler.hilbert

13 # Split the rng key into 2: one for each random operation

14 key_indx, key_flip = jax.random.split(key, 2)

15 # Pick two random sites on every chain

16 indxs = jax.random.randint(

17 key_indx, shape=(n_chains, 2), minval=0, maxval=hilb.size

18 )

19 # flip those sites

20 $\mathtt{\sigma}$p, _ = flip_state(hilb, key_flip, $\mathtt{\sigma}$,
indxs)

21

22 # If this transition had a correcting factor L, it’s possible

23 # to return it as a vector in the second value

24 return $\mathtt{\sigma}$p, None

Once a custom rule is defined, a MCMC sampler using such rule can be
constructed with the command sampler = MetropolisSampler(hilbert,
TwoLocalRule()). Besides Metropolis algorithms, more advanced Markov chain
algorithms can also be implemented as NetKet samplers. Currently, parallel
tempering is provided as an experimental feature.

Some models allow us to directly generate samples that are exactly distributed
according to the desired probability, without the use of Markov chains and the
issue of autocorrelation, which often leads to more efficient sampling. In
this case, direct samplers can be implemented with an interface similar to
Markov chain samplers. Currently NetKet has implemented ARDirectSampler to be
used with ARNNs. For benchmarking purposes, NetKet also provides ExactSampler,
which allows direct sampling from any model by computing the full Born
distribution $|\psi(s)|^{2}$ for all $s$. Table 2 is a list of all the
samplers.

Type | Name |  Usage
---|---|---
MCMC (Metropolis) | MetropolisLocal |  discrete Hilbert spaces
2-3 | MetropolisExchange |  permutations of local states, conserving total magnetization in spin systems
2-3 | MetropolisHamiltonian |  preserving symmetries of the Hamiltonian
2-3 | MetropolisGaussian |  continuous Hilbert spaces
Direct | ExactSampler |  small Hilbert spaces, performs MC sampling from the exact distribution
2-3 | ARDirectSampler |  autoregressive models
|  |
Table 2: List of samplers in NetKet with their class names and a description.

###  3.5 Quantum geometric tensor

The quantum geometric tensor (QGT) [54] of a pure state is the metric tensor
induced by the Fubini–Study distance [55, 56]

$\displaystyle
d(\psi,\phi)=\cos^{-1}\sqrt{\frac{\langle\psi|\phi\rangle\langle\phi|\psi\rangle}{\langle\psi|\psi\rangle\langle\phi|\phi\rangle}},$
(21)

which is the natural and gauge-invariant distance between two pure quantum
states $\ket{\psi}$ and $\ket{\phi}$. The QGT is commonly used for time
evolution (see Section 4.3) and for quantum natural gradient descent [57],
which was originally developed in the VMC community under the name of
_stochastic reconfiguration_ (SR) [58]. Quantum natural gradient descent is
directly related to the natural gradient descent developed in the machine
learning community [59].

From now on, we assume that the state $\ket{\psi_{\theta}}$ is parametrized by
a set of parameters $\theta$. Assuming further that
$\ket{\phi}=\ket{\psi_{\theta+\delta\theta}}$, the distance (21) can be
expanded to second order in the infinitesimal parameter change $\delta\theta$
as
$d(\psi_{\theta},\psi_{\theta+\delta\theta})^{2}=(\delta\theta)^{\dagger}G(\delta\theta)$,
where $G$ is the quantum geometric tensor. For a holomorphic mapping
$\theta\mapsto\ket{\psi_{\theta}}$, the QGT is given by

$\displaystyle
G_{ij}(\theta)=\frac{\innerproduct{\partial_{\theta_{i}}{\psi_{\theta}}}{\partial_{\theta_{j}}\psi_{\theta}}}{\innerproduct{\psi}{\psi}}-\frac{\innerproduct{\partial_{\theta_{i\vphantom{j}}}\psi_{\theta}}{\psi_{\theta}}\innerproduct{\psi_{\theta}}{\partial_{\theta_{j}}\psi_{\theta}}}{\innerproduct{\psi}{\psi}^{2}}$
(22)

where the indices $i,j$ label the parameters and
$\innerproduct{x}{\partial_{\theta_{i}}\psi_{\theta}}=\partial_{\theta_{i}}\innerproduct{x}{\psi_{\theta}}$.
Similar to expectation values and their gradients, Eq. (22) can be rewritten
as a classical covariance with respect to the Born distribution
$\propto|\psi(s)|^{2}$:

$\displaystyle
G_{ij}(\theta)=\operatorname{Cov}[O_{i},O_{j}]=\mathbb{E}\left[O_{i}^{*}(O_{j}-\mathbb{E}[O_{j}])\right],$
(23)

where $O_{i}$ are the log-derivatives (13) of the ansatz.141414Strictly
speaking, this estimator is only correct if
$\psi_{\theta}(s)=0\implies\partial_{\theta_{k}}\psi_{\theta}(s)=0$. This is
because we multiplied and divided by $\psi_{\theta}(s)$ in the derivation of
the estimator, which is only valid if $\psi_{\theta}(s)\neq 0$. This allows
the quantum geometric tensor to be estimated using the same sampling procedure
used to obtain expectation values and gradients. The QGT or its stochastic
estimate is also commonly known as the _S matrix_ [52] or _quantum Fisher
matrix_ (QFM) in analogy to the classical Fisher information matrix [59, 60,
57].

For applications such as quantum natural gradient descent or time-evolution it
is usually not necessary to access the full, dense matrix representation
$G_{ij}(\theta)$ of the quantum geometric tensor, but only to compute its
product with a vector, $\tilde{v}_{i}=\sum_{j}G_{ij}(\theta)\,v_{j}$. When the
variational ansatz $\psi_{\theta}$ has millions of parameters, the QGT can
indeed be too large to be stored in memory. Exploiting the Gram matrix
structure of the geometric tensor [61], we can directly compute its action on
a vector without ever calculating the full matrix, trading memory requirements
for an increased computational cost.

Given a variational state vs, a QGT object can be obtained by calling:

⬇

1>>> qgt = vs.quantum_geometric_tensor()

This qgt object does not store the full matrix, but can still be applied to a
vector with the same shape as the parameters:

⬇

1>>> vec = jax.tree_map(jnp.ones_like, vs.parameters)

2>>> qgt_times_vec = qgt @ vec

It can be converted to a dense matrix by calling to_dense:

⬇

1# get the matrix (2d array) of the qgt

2>>> qgt_dense = qgt.to_dense()

3# flatten vec into a 1d Array

4>>> grad_dense, unravel = nk.jax.tree_ravel(vec)

5>>> qgt_times_vec = unravel(qgt_dense @ vec_dense)

The QGT can then be used together with a direct solver, such as
jnp.linalg.eigh, jnp.linalg.svd, or jnp.linalg.qr.

##### Mixed states.

When working with mixed states, which are encoded in a density matrix, it is
necessary to pick a suitable metric to induce the QGT. Even if the most
physical distances for density matrices are the spectral norm or other trace-
based norms [62, 63], it is generally hard to use them to define an expression
for the QGT that can be efficiently sampled and computed at polynomial cost.
While this might be regarded as a barbaric choice, it leads to an expression
equivalent to Eq. 23, where the expectation value is over the joint-
distribution of the of row and column labels $(s,s^{\prime})$ of the squared
density matrix $\propto\absolutevalue{\rho(s,s^{\prime})}^{2}$. Therefore,
when working with mixed states, we resort to the $L_{2}$ norm, which is
equivalent to treating the density matrix as a vector (“pure state”) in an
enlarged Hilbert space.

####  3.5.1 Implementation differences

There is some freedom in the way one can calculate the QGT, each different
implementation taking a different tradeoff between computational and memory
cost. In the example above, we have relied on NetKet to automatically select
the best implementation according to some internal heuristics, but if one
wants to push variational methods to the limits, it is useful to understand
the two different implementations on offer:

  * •

QGTJacobian

, which computes and stores the Jacobian $O_{i}(s)$ of the model (cf. Section
B.1) at initialization (using reverse-mode automatic differentiation) and can
be applied to a vector by performing two matrix-vector multiplications;

  * •

QGTOnTheFly

, which lazily computes the application of the quantum geometric tensor on a
vector through automatic differentiation.

##### QGTOnTheFly

is the most flexible and can be scaled to arbitrarily large systems. It is
based on the observation that the QGT is the Jacobian of the model multiplied
with its conjugate, which means that its action can be calculated by combining
forward and reverse-mode automatic differentiation. At initialization, it only
computes the linearization (forward pass), and then it effectively recomputes
the gradients every time it is applied to a vector. However, since it never
has to store these gradients, it is not limited by the available memory, which
also makes it perform well for shallow neural-network models like RBMs. This
method works with both holomorphic and non-holomorphic ansätze with no
difference in performance.

##### QGTJacobian.

For deep networks with ill-conditioned151515The number of steps required to
find a solution with an iterative linear solver grows with the condition
number of the matrix. Therefore, an ill-conditioned matrix requires many steps
of iterative solver. For a discussion on this issue, see the paragraph on
Linear Systems of this section. quantum geometric tensors, recomputing the
gradients at every step in an iterative solver might be very costly.
QGTJacobian can therefore achieve better performance at the cost of
considerably higher memory requirements because it precomputes the Jacobian at
construction and stores it. The downside is that it has to store a matrix of
shape $N_{\mathrm{samples}}\times N_{\mathrm{parameters}}$, which might not
fit in the memory of a GPU. We note that there are two different
implementations of QGTJacobian: QGTJacobianDense and QGTJacobianPyTree. The
difference among the two is that in the former the Jacobian is stored
contiguously in memory, leading to a better throughput on GPUs, while the
latter stores them in the same structure as the parameters (so each parameter
block is separated from the others). Converting from the non-contiguous (
QGTJacobianPyTree) to the contiguous ( QGTJacobianDense) format has, however,
a computational and memory cost which might shadow its benefit. Moreover, the
dense format does not work with non-homogeneous parameter data types. The
basic QGTJacobian algorithm supports both holomorphic and non-holomorphic NQS,
but a better performing algorithm for holomorphic ansätze can be accessed
instantiating it with the option holomorphic=True.

Key differences between the different QGT implementations are summarized in
Table 3. Implementations can be selected, and options passed to them, as shown
below:

⬇

1>>> from netket.optimizer.qgt import (

2... QGTJacobianPyTree, QGTJacobianDense, QGTOnTheFly

3... )

4>>> qgt1 = vs.quantum_geometric_tensor(QGTOnTheFly)

5>>> qgt2 = vs.quantum_geometric_tensor(QGTJacobianPyTree(holomorphic=True))

6>>> qgt3 = vs.quantum_geometric_tensor(QGTJacobianDense)

Implementation | Extra arguments |  Use-cases |  Limitations
---|---|---|---
QGTOnTheFly | None |  shallow networks, large numbers of parameters and samples, few solver steps |  Might be more computationally expensive for deep networks compared to QGTJacobian.
QGTJacobianDense |  mode holomorphic rescale_shift |  deep networks with narrow layers |  requires homogeneous parameter types, memory bound
1-13-4 QGTJacobianPyTree |  deep networks with heterogeneous parameters |  memory bound
|  |
Table 3:  Overview of the three QGT implementations currently provided by
NetKet with their respective options and limitations.

##### Holomorphicity.

When performing time evolution or natural gradient descent, one does not
always need the full quantum geometric tensor: for ansätze with real
parameters, as well as in the case of non-holomorphic wave
functions,161616Non-holomorphic functions of complex parameters are internally
handled by both jax and NetKet as real-parameter functions that take the real
and imaginary parts of the “complex parameters” separately. only the real part
of the QGT is used. The real and imaginary parts of the QGT are only required
when working with a holomorphic ansatz. (An in-depth discussion of why this is
the case can be found at [64, Table 1].) For this reason, NetKet’s QGT
implementations return the full geometric tensor only for holomorphic complex-
parameter ansätze, and its real part in all other cases.

####  3.5.2 Solving linear systems

For most applications involving the QGT, a linear system of equations of the
kind

$\sum_{j}G_{ij}\delta_{j}=f_{i}$ (24)

needs to be solved, where $G_{i,j}$ is the quantum geometric tensor of a NQS,
and $f_{i}$ is a gradient. This can be done using the standard jax/NumPy
functions, assuming $f$ is a pytree with the same structure as the variational
parameters:

⬇

1>>> # iterative solver

2>>> x, info = jax.scipy.sparse.linalg.cg(qgt, f)

3>>> # direct solver, acting on the dense matrix

4>>> x, info = jax.numpy.linalg.cholesky(qgt.to_dense(), f)

However, we recommend that users call the solve method on the QGT object,
which allows some additional optimization that may improve performance:

⬇

1>>> x, info = qgt.solve(jax.scipy.sparse.linalg.cg, f)

2>>> x, info = qgt.solve(nk.optimizer.solver.cholesky, f)

While this works with any of the representations it is advisable to only use
Jacobian based implementations ( QGTJacobianPyTree or QGTJacobianDense) with
direct solvers, since constructing the QGT matrix from QGTOnTheFly requires
multiplication with all basis vectors, which is not as efficient. Finally, we
highlight the fact that users can write their own functions to solve the
linear system (24) using advanced regularization schemes (see for instance
ref. [21]) and use them together with qgt.solve, as long as they respect the
standard NumPy solver interface.

When working with iterative solvers such as cg, gmres or minres, the number of
steps required to find a solution grows with the condition number of the
matrix. Therefore, an ill-conditioned geometric tensor requires many steps of
iterative solver, increasing the computational cost. Even then or when using
non-iterative methods such as singular value decomposition (SVD), the high
condition number can cause instabilities by amplifying noise in the right-hand
side of the linear equation [20, 21, 23]. This is especially true for NQS,
which typically feature QGTs with a spectrum spanning many orders of magnitude
[60], often making QGT-based algorithms challenging to stabilize [21, 23, 15].

To counter that, there is empirical evidence that in some situations,
increasing the number of samples used to estimate the QGT and gradients helps
to stabilize the solution [23]. Furthermore, it is possible to apply various
regularization techniques to the equation. A standard option is to add a small
diagonal shift $\epsilon$ to the QGT matrix before inverting it, thus solving
the linear equation

$\sum_{j}(G_{i,j}+\epsilon)\delta_{j}^{\prime}=f_{i}.$ (25)

When $\epsilon$ is small, the solution $\delta^{\prime}$ will be close to the
desired solution. Otherwise it is biased towards the plain force $f$, which is
still acceptable in gradient-based optimization. To add this diagonal shift in
NetKet, one of the following approaches can be used:

⬇

1>>> qgt_1 = vs.quantum_geometric_tensor(QGTOnTheFly(diag_shift=0.001))

2QGTOnTheFly(diag_shift=0.001)

3>>> qgt_2 = qgt_1.replace(diag_shift=0.005)

4QGTOnTheFly(diag_shift=0.005)

5>>> qgt_3 = qgt_2 + 0.005

6QGTOnTheFly(diag_shift=0.01)

Regularizing the QGT with a diagonal shift is an effective technique that can
be used when performing SR/natural gradient descent for ground state search
(see Section 4.1). Note, however, that since the diagonal shift biases the
solution of the linear equation towards the plain gradient, it may bias the
evolution of the system away from the physical trajectory in cases such as
real-time evolution. In those cases, non-iterative solvers such as those based
on SVD can be used, the stability of which can be controlled by suppressing
smaller singular values. It has also been suggested in the literature to
improve stability by suppressing particularly noisy gradient components [21,
39]. This is not currently implemented in NetKet, but planned for a future
release. SVD-based regularization also comes at the cost of potentially
suppressing physically relevant dynamics [23], making it necessary to find the
right balance between stabilization and physical accuracy, and increased
computational time as SVD is usually less efficient than iterative solvers.

##  4 Algorithms for variational states

The main use case of NetKet is variational optimization of wave function
ansätze. In the current version NetKet, three algorithms are provided out of
the box via high-level driver classes: variational Monte Carlo (VMC) for
finding ground states of (Hermitian) Hamiltonians, time-dependent variational
Monte Carlo (t-VMC) for real- and imaginary-time evolution, and steady-state
search of Liouvillian open-system dynamics.

These drivers are part of the nk.driver module but we also export them from
the nk namespace. They are constructed from the relevant physical model (e.g.,
a Hamiltonian), a variational state, and other objects used by the
optimization method. They all support the run method, which performs a number
of optimization steps and logs their progress (e.g., variational energies and
network parameters) in a variety of output formats.

We highlight that these drivers are built on top of the functionalities
described in Sections 2 and 3, and users are free to implement their own
drivers or optimization loops, as demonstrated in Section 4.4.

###  4.1 Ground-state search

NetKet provides the variational driver nk.VMC for searching for minimal-energy
states using VMC [52]. In the simplest case, the VMC constructor takes three
arguments: the Hamiltonian, an optimizer and the variational state (see
Section 3.1). NetKet makes use of optimizers provided by the jax-based optax
library [65],171717The nk.optimizer submodule includes a few optimizers for
ease of use and backward compatibility: these are simply re-exports from
optax. which can be directly passed to VMC, allowing the user to build complex
training schedules or custom optimizers. In each optimization step, new
samples of the variational state are drawn and used to estimate the gradient
of the Hamiltonian with respect to the parameters $\theta$ of the ansatz [52]
based on the force vector [compare Eq. (14)]

$\tilde{f}_{i}=\operatorname{Cov}[O_{i},\tilde{H}]=\mathbb{E}[O_{i}^{*}(\tilde{H}-\mathbb{E}[\tilde{H}])],$
(26)

where $\tilde{H}$ is the local estimator (12) of the Hamiltonian, known as the
local energy, and $O_{i}$ is the log-derivative (13) of the wave function. All
expectation values in Eq. 26 are evaluated over the Born distribution
$\propto|\psi({}\cdot{})|^{2}$ and can therefore be estimated by averaging
over the Monte Carlo samples. Given the vector $\tilde{f}$, the direction of
steepest descent is given by the energy gradient

$\displaystyle f\equiv\nabla_{\theta}\langle\hat{H}\rangle$
$\displaystyle=2\real[\tilde{f}]\quad\text{(real)}$ (27) or complex co-
gradient [43] $\displaystyle f\equiv\nabla_{\theta^{*}}\langle\hat{H}\rangle$
$\displaystyle=\tilde{f}\quad\text{(complex holomorphic)}.$ (28)

Here we have distinguished the case of i) real parameters and ii) complex
parameters with a variational mapping that is holomorphic with respect to
$\theta.$ For non-holomorphic ansätze (cf. Section 3.2.2), complex parameters
can be treated pairs of separate real-valued parameters (real and imaginary
part) in the sense of eq. (27). Therefore, this case can be considered
equivalent to the real parameter case.

The gradients are then passed on to the optax optimizer, which may transform
them (using, e.g., Adam) further before updating the parameters $\theta$.
Using the simple stochastic gradient descent optimizer optax.sgd (alias
nk.optimizer.Sgd), the update rule is

$\theta_{i}\mapsto\theta_{i}-\eta f_{i},$ (29)

where the gradient $f$ is given by Eq. 27 or (28) as appropriate.

Below we give a short snippet showing how to use the VMC driver to find the
ground-state of the Ising Hamiltonian.

⬇

1# Define the geometry of the lattice

2g = nk.graph.Hypercube(length=10, n_dim=1, pbc=False)

3# Hilbert space of spins on the graph

4hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)

5# Construct the Hamiltonian

6hamiltonian = nk.operator.Ising(hi, graph=g, h=0.5)

7

8# define a variational state with a Metropolis Sampler

9sa = nk.sampler.MetropolisLocal(hi)

10vstate = nk.vqs.MCState(sa, nk.models.RBM())

11

12# Construct the VMC driver

13vmc = nk.VMC(hamiltonian,

14 nk.optimizer.Sgd(learning_rate=0.1),

15 variational_state=vstate)

16

17# run the optimisation for 300 steps

18output = vmc.run(300)

To improve on plain stochastic gradient descent, the VMC interface allows
passing a keyword argument preconditioner. This must be a function that maps a
variational state and the gradient vector $f_{i}$ to the vector $\delta_{i}$
to be passed to the optimizer as gradients instead of $f_{i}$. An important
use case is stochastic reconfiguration [52], where the gradient is
preconditioned by solving the linear system of equations

$\displaystyle\sum_{j}\real[G_{ij}]\delta_{j}$
$\displaystyle=f_{i}=2\real[\tilde{f}_{i}]\quad\text{(real)}$ (30) or
$\displaystyle\sum_{j}G_{ij}\delta_{j}$
$\displaystyle=f_{i}\quad\text{(complex holomorphic)}.$ (31)

The corresponding preconditioner can be created from a QGT class and a jax-
compatible linear solver (the default is jax.scipy.sparse.linalg.cg) using
nk.optimizer.SR:

⬇

1# Construct the SR object with the chosen algorithm

2sr = nk.optimizer.SR(

3qgt = nk.optimizer.qgt.QGTOnTheFly,

4 solver=jax.scipy.sparse.linalg.bicgstab,

5 diag_shift=0.01,

6)

7

8# Construct the VMC driver

9vmc = nk.VMC(

10 hamiltonian, # The Hamiltonian to optimize

11 nk.optimizer.sgd(learning_rate=0.1), # The optimizer

12 variational_state=vstate, # The variational state

13 preconditioner=sr, # The preconditioner

14)

###  4.2 Finding steady states

In order to study open quantum systems, NetKet provides the nk.SteadyState
variational driver for determining the variational steady-state
$\hat{\rho}_{ss}$ defined as the stationary point of an arbitrary super-
operator $\mathcal{L}$,

$0=\frac{d\hat{\rho}}{dt}=\mathcal{L}\hat{\rho}.$ (32)

The search is performed by minimizing the Frobenius norm of the time-
derivative [24], which defines the cost function

$\mathcal{C}(\theta)=\frac{\norm{\mathcal{L}\hat{\rho}}_{2}^{2}}{\norm{\hat{\rho}}_{2}^{2}}=\frac{\Tr\left[\hat{\rho}^{\dagger}\mathcal{L}^{\dagger}\mathcal{L}\hat{\rho}\right]}{\Tr\left[\hat{\rho}^{\dagger}\hat{\rho}\right]},$
(33)

which has a global minimum for the steady state. The stochastic gradient is
estimated over the probability distribution of the entries of the vectorized
density matrix according to the formula:

$f_{i}\equiv\frac{\partial}{\partial\theta_{i}^{*}}\frac{\norm{\mathcal{L}\hat{\rho}}_{2}^{2}}{\norm{\hat{\rho}}_{2}^{2}}=\mathbb{E}\left[\tilde{\mathcal{L}}\nabla_{i}^{*}\tilde{\mathcal{L}}\right]-\mathbb{E}[O_{i}^{*}\tilde{\mathcal{L}}^{2}],$
(34)

where
$\tilde{\mathcal{L}}(s,s^{\prime})=\sum_{m,m^{\prime}}\mathcal{L}(s,s^{\prime};m,m^{\prime})\rho(m,m^{\prime})/\rho(s,s^{\prime})$
is the local estimator proposed in [24], and the expectation values are taken
with respect to the “Born distribution” of the vectorized density matrix,
$p(s,s^{\prime})\propto|\rho(s,s^{\prime})|^{2}$. The optimization works like
the ground-state optimization provided by nk.VMC: the gradient is passed to an
optax optimizer, which may transform it further before updating the parameters
$\theta$. The simplest optimizer, optax.sgd, would update the parameters
according to the equation

$\theta_{i}\rightarrow\theta_{i}-\eta f_{i}.$ (35)

To improve the performance of the optimization, it is possible to pass the
keyword argument preconditioner to specify a gradient preconditioner, such as
stochastic reconfiguration that uses the quantum geometric tensor to transform
the gradient. The geometric tensor is computed according to the $L_{2}$ norm
of the vectorized density matrix (see Section 3.5).

As an example, we provide a snippet to study the steady state of a transverse-
field Ising chain with 10 spins and spin relaxation corresponding to the
Lindblad master equation

$\mathcal{L}\hat{\rho}=-i\commutator{\hat{H}}{\hat{\rho}}+\sum_{i}\hat{\sigma}_{i}^{-}\hat{\rho}\hat{\sigma}_{i}^{+}-\frac{1}{2}\anticommutator{\hat{\sigma}_{i}^{+}\hat{\sigma}_{i}^{-}}{\hat{\rho}}.$
(36)

We first define the Hamiltonian and a list of jump operators, which are stored
in a LocalLiouvillian object, which is a lazy representation of the super-
operator $\mathcal{L}$. Next, a variational mixed state is built by defining a
sampler over the doubled Hilbert space and optionally a different sampler for
the diagonal distribution $p(s)\propto\rho(s,s)$, which is used to estimate
expectation values of operators. The number of samples used to estimate super-
operators and operators can be specified separately, as shown in the example
by specifying n_samples and n_samples_diag.

⬇

1# Define the geometry of the lattice

2g = nk.graph.Hypercube(length=10, n_dim=1, pbc=False)

3# Hilbert space of spins on the graph

4hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)

5

6# Construct the Liouvillian Master Equation

7ha = nk.operator.Ising(hi, graph=g, h=0.5)

8j_ops = [nk.operator.spin.sigmam(hi, i) for i in range(g.n_nodes)]

9# Create the Liouvillian with Hamiltonian and jump operators

10lind = nk.operator.LocalLiouvillian(ha, j_ops)

11

12# Observable

13sz = sum([nk.operator.spin.sigmam(hi, i) for i in range(g.n_nodes)])

14

15# Neural Density Matrix

16sa = nk.sampler.MetropolisLocal(lind.hilbert)

17vs = nk.vqs.MCMixedState(

18 sa, nk.models.NDM(beta=1), n_samples=2000, n_samples_diag=500

19)

20# Optimizer

21op = nk.optimizer.Sgd(0.01)

22sr = nk.optimizer.SR(diag_shift=0.01)

23

24ss = nk.SteadyState(lind, op, variational_state=vs, preconditioner=sr)

25out = ss.run(n_iter=300, obs={"Sz": sz})

###  4.3 Time propagation

Time propagation of variational states can be performed by incorporating the
time dependence in the variational parameters and deriving an equation of
motion that gives a trajectory in parameters space $\theta(t)$ that
approximates the desired quantum dynamics. For real-time dynamics of pure and
mixed NQS, such an equation of motion can be derived from the time-dependent
variational principle (TDVP) [66, 67, 64]. When combined with VMC sampling to
estimate the equation of motion (EOM), this is known as time-dependent
variational Monte Carlo (t-VMC) [68, 52] and is the primary approach currently
used in NQS literature [8, 18, 19, 21, 22, 23]. For complex holomorphic
parametrizations181818The TDVP can be implemented for real-parameter
wavefunctions by taking real parts of the right-hand side and QGT similar to
VMC (Section 4.1) [39, 46]. This is not yet available in the current version
of NetKet, but will be added in a future release. , the TDVP equation of
motion is

$\displaystyle\sum_{j}G_{ij}(\theta)\,\dot{\theta}_{j}$ $\displaystyle=\gamma
f_{i}(\theta,t),$ (37)

with the QGT $G$ and force vector $f$ defined in Sections 4.1 and 4.2. After
solving Eq. 37, the resulting parameter derivative $\dot{\theta}$ can be
passed to an ODE solver. The factor $\gamma$ determines the type of evolution:

  * •

For $\gamma=-i$, the EOM approximates the real-time Schrödinger equation on
the variational manifold, the simulation of which is the main use case for the
t-VMC implementation provided by NetKet.

  * •

For $\gamma=-1$, the EOM approximates the imaginary-time Schrödinger equation
on the variational manifold. When solved using the first-order Euler scheme
$\theta(t+dt)=\theta(t)+\dot{\theta}\,dt$, this EOM is equivalent to
stochastic reconfiguration with learning rate $dt$. Imaginary-time propagation
with higher-order ODE solvers can therefore also be used for ground state
search as an alternative to VMC. This can result in improved convergence in
some cases [17].

  * •

For $\gamma=1$ and with the Lindbladian super-operator taking the place of the
Hamiltonian in the definition of the force $f$, this ansatz yields the
dissipative real-time evolution according to the Gorini-Kossakowski-Lindblad-
Sudarshan master equation [69]. Our implementation uses the QGT induced by the
vector norm [25] as discussed in the last paragraph of Section 3.5.

The current version of NetKet provides a set of Runge–Kutta (RK) solvers based
on jax and a driver class TDVP implementing the t-VMC algorithm for the three
use cases listed above. At the time of writing, these features are provided as
a preview version in the netket.experimental namespace as their API is still
subject to ongoing development. The ODE solvers are located in the submodule
netket.experimental.dynamics, the driver under netket.experimental.TDVP.

Runge-Kutta solvers implement the propagation scheme [70]

$\displaystyle\theta(t+dt)$
$\displaystyle=\theta(t)+dt{\sum\nolimits_{l=1}^{L}}b_{l}k_{l}$ (38)

using a linear combination of slopes

$\displaystyle k_{l}$
$\displaystyle=F\left(\theta(t)+{\sum\nolimits_{m=1}^{l-1}}a_{lm}k_{m},\;t+c_{l}\,dt\right),$
(39)

each determined by the solution $F(\theta,t)=\dot{\theta}$ of the equation of
motion (37) at an intermediate time step. The coefficients $\\{a_{lm}\\}$,
$\\{b_{l}\\}$, and $\\{c_{l}\\}$ determine the specific RK scheme and its
order. NetKet further supports step size control when using adaptive RK
schemes. In this case, the step size is dynamically adjusted based on an
embedded error estimate that can be computed with little overhead during the
RK step (38) [70]. Step size control requires a norm on the parameters space
in order to estimate the magnitude of the error. Typically, the Euclidean norm
$\|\delta\|=\sqrt{\sum_{i}|\delta_{i}|^{2}}$ is used. However, since different
directions in parameters space influence the quantum state to different
degrees, it can be beneficial to use the norm
$\|\delta\|_{G}=\sqrt{\sum_{i}\delta_{i}^{*}G_{ij}\delta_{j}}$ induced by the
QGT $G$ as suggested in Ref. [21], which takes this curvature into account and
is also provided as an option with the NetKet time-evolution driver.

An example demonstrating the use of NetKet’s time evolution functionality is
provided in Sec. 9.

###  4.4 Implementing custom algorithms using NetKet

While key algorithms for energy optimization, steady states, and time
propagation are provided out of the box in the current NetKet version, there
are many more applications of NQS. While we wish to provide new high-level
driver classes for additional use cases, such as quantum state tomography [31]
or general overlap optimization [32], it is already possible and encouraged
for users to implement their own algorithms on top of NetKet. For this reason,
we provide the core building blocks of NQS algorithms in a composable fashion.

For example, it is possible to write a simple loop that solves the TDVP
equation of motion (37) for a holomorphic variational ansatz and using a
first-order Euler scheme [i.e., $\theta(t+dt)=\theta(t)+\dot{\theta}(t)\,dt$]
very concisely, making use of the elementary building blocks provided by the
VariationalState class:

⬇

1def custom_simple_tdvp(

2 hamiltonian: AbstractOperator, # Hamiltonian

3 vstate: VariationalState, # variational state

4 t0: float, # initial time

5 dt: float, # time step

6 t_end: float, # end time

7):

8 t = t0

9 while t < t_end:

10 # compute the energy gradient f

11 energy, f = vstate.expect_and_grad(hamiltonian)

12 G = vstate.quantum_geometric_tensor()

13 # multiply the gradient by -1.0j for unitary dynamics

14 gamma_f = jax.tree_map(lambda x: -1.0j ∗ x, f)

15 # Solve the linear system using any solver, such as CG

16 # (or write your own regularization scheme)

17 dtheta, _ = G.solve(jax.scipy.sparse.linalg.cg, gamma_f)

18 # update the parameters (theta = theta + dt ∗ dtheta)

19 vs.parameters = jax.tree_map(

20 lambda x, y: x + dt ∗ y, vs.parameters, dtheta

21 )

22 t = t + dt

While the included TDVP driver (Section 4.3) provides many additional features
(such as error handling, step size control, or higher-order integrators) and
makes use of jax’s just-in-time compilation, this simple implementation
already provides basic functionality and shows how NetKet can be used for
quick prototyping.

##  5 Symmetry-aware neural quantum states

NetKet includes a powerful set of utilities for implementing NQS ansätze that
are symmetric or transform correctly under the action of certain discrete
symmetry groups. Only groups that are isomorphic to a set of permutations of
the computational basis are supported. This is useful for modeling symmetric
(e.g., lattice) Hamiltonians, whose eigenstates transform under irreducible
representations of their symmetry groups. Restricting the Hilbert space to
individual symmetry sectors can improve the convergence of variational
optimization [71] and the accuracy of its result [72, 49, 14, 16].
Additionally, symmetry restrictions can be used to find excited states [13,
30, 16], provided they are the lowest energy level in a particular symmetry
sector.

While there is a growing interest for other symmetry groups, such as
continuous ones like $SU(2)$ or $SO(3)$, they cannot be compactly represented
in the computational basis and therefore the approach described in this
chapter cannot be used. Finding efficient encodings for continuous groups is
still an open research problem and it’s not yet clear which strategy will work
best [16].

NetKet uses _group convolutional neural networks_ (GCNNs) to build wave
functions that are symmetric over a finite group $G$. GCNNs generalize
convolutional neural networks, invariant under the Abelian translation group,
to general symmetry groups $G$ which may contain non-commuting elements [73].
GCNNs were originally designed to be invariant, but they can be modified to
transform under an arbitrary irreducible representation (irrep) of $G$, using
the projection operator [74]

$|\psi_{\chi}\rangle=\frac{d_{\chi}}{|G|}\sum_{g\in G}\chi_{g}^{*}\
g|\psi\rangle,$ (40)

where $g$ runs over all symmetry operations in $G$, with corresponding
characters $\chi_{g}$. Under the trivial irrep, where all characters are
unity, the invariant model is recovered.

NetKet can infer the full space group of a lattice, defined as a set of
permutations of lattice sites, starting from a geometric description of its
point group. It can also generate nontrivial irrep characters [to be used in
(40) for states with nonzero wave vectors or transforming nontrivially under
point-group symmetries] using a convenient interface that approximates
standard crystallographic formalism [75]. In addition, NetKet provides
powerful group-theoretic algorithms for arbitrary permutation groups of
lattice sites, allowing new symmetry elements to be easily defined.

Pre-built GCNNs are then provided in the nk.models submodule, which can be
constructed by specifying few parameters, such as the number of features in
each layer, and the lattice or permutation group under which the network
should transform. Symmetric RBMs [8] are also implemented as one-layer GCNNs
that aggregate convolutional features using a product rather than a sum. These
pre-built network architectures are made up of individual layers found in the
nk.nn submodule, which can be used directly to build custom symmetric ansätze.

Section 5.1 describes the NetKet interface for constructing space groups of
lattices and their irreps. Usage of GCNNs is described in Section 5.2, while
appendix A provides mathematical and implementation details.

###  5.1 Symmetry groups and representation theory

NetKet supports symmetry groups that act on a discrete Hilbert space defined
on a lattice. On such a Hilbert space, space-group symmetries act by permuting
sites; most generally, therefore, arbitrary subgroups of the symmetric group
$S_{N}$ of a lattice of $N$ sites are supported. A symmetry group can be
specified directly as a list of permutations, as in the following example,
which enforces the symmetry
$\psi(s_{0},s_{1},s_{2},s_{3})=\psi(s_{3},s_{1},s_{2},s_{0})$ for all four-
spin configurations $s=(s_{0},\ldots,s_{3})$, $s_{i}=\pm 1$:

⬇

1hi = nk.hilbert.Spin(1/2, N=4)

2symms = [

3 [0, 1, 2, 3], # identity element

4 [3, 1, 2, 0], # swap first and last site

5]

6model = nk.models.RBMSymm(symms, alpha=1)

The listed permutations are required to form a group and, in particular, the
identity operation $e:s\mapsto s$ must always be included as the first
element.

It is inconvenient and error-prone to specify all space-group symmetries of a
large lattice by their indices. Therefore, NetKet provides support for
abstract representations of permutation and point groups through the
nk.utils.group module, complete with algorithms to compute irreducible
representations [76, 77, 78]. The module also contains a library of two- and
three-dimensional point groups, which can be turned into lattice-site
permutation groups using the graph class nk.graph.Lattice (but not general
Graph objects, for they carry no geometric information about the system):

⬇

1from netket.utils.group.planar import D

2from netket.graph import Lattice

3

4# construct a centred rectangular lattice

5lattice = Lattice(

6 basis_vectors = [[2,0], [0,1]], # each row is a lattice vector

7 extent = (5,5),

8 site_offsets = [[0,0], [1,0.5]], # each row is the position of a site in the
unit cell

9 point_group = D(2) # the point group of the lattice, here Z_2 x Z_2

10)

NetKet contains specialized constructors for some lattices (e.g., Square or
Pyrochlore), which come with a default point group; however, these can be
overridden in methods like Lattice.space_group:

⬇

1from netket.utils.group.planar import rotation, reflection_group, D

2from netket.utils.group import PointGroup, Identity

3from netket.graph import Honeycomb

4

5# construct the D_6 point group of the honeycomb lattice by hand

6cyclic_6 = PointGroup(

7 [Identity()] + [rotation(360 / n ∗ i) for i in range(1, n)],

8 ndim=2,

9)

10# the @ operator returns the Cartesian product of groups

11# but doesn’t check for group structure

12dihedral_6 = reflection_group(angle=0) @ cyclic_6

13

14assert dihedral_6 == D(6)

15

16lattice = Honeycomb([6,6])

17

18# returns the full space group of ‘lattice‘ as a PermutationGroup

19space_group = lattice.space_group()

20# the space group is spanned by 6^2 translations and 12 point-group
symmetries

21assert len(space_group) == 12 ∗ 6 ∗ 6

22

23# do this if the Hamiltonian breaks reflection symmetry

24# can also be used for generic Lattices that have no default point group

25space_group = lattice.space_group(cyclic_6)

Irreducible representation (irrep) matrices can be computed for any point or
permutation group object using the method irrep_matrices(). Characters (the
traces of these matrices) are returned by the method character_table() as a
matrix, each row of which lists the characters of all group elements.
Character tables closer to the format familiar from quantum chemistry texts
are produced by character_table_readable(). Irrep matrices and character
tables are calculated using adaptations of Dixon’s [76] and Burnside’s [77]
algorithms, respectively.

It would, however, be impractical to inspect irreps of a large space group
directly to specify the symmetry sector on which to project a GCNN wave
function. Exploiting the semidirect-product structure of space groups [78],
space-group irreps are usually191919Representation theory for wave vectors on
the surface of the Brillouin zone in a nonsymmorphic space group is much more
complicated [78] and is not currently implemented in NetKet. described in
terms of a set of symmetry-related wave vectors (known as a star) and an irrep
of the subgroup of the point group that leaves the same invariant (known as
the little group) [75]. Irreps can be constructed in this paradigm using
SpaceGroupBuilder objects returned by Lattice.space_group_builder():

⬇

1from netket.graph import Triangular

2

3lattice = Triangular([6,6])

4momentum = [0,0]

5# space_group_builder() takes an optional PointGroup argument

6sgb = lattice.space_group_builder()

7

8# choosing a representation

9# this one corresponds to the B_2 irrep at the Gamma point

10chi = sgb.space_group_irreps(momentum)[3]

The irrep chi, generated using the little group, is equivalent to one of the
irreps in the table Lattice.space_group().character_table() and can thus be
used for symmetry-projecting GCNN ansätze. The order in which irreps of the
little group are returned can readily be checked in an interactive session:

⬇

1>>> sgb.little_group(momentum).character_table_readable()

2([’1xId()’, ’2xRot(60)’, ’2xRot(120)’, ’1xRot(180)’, ’3xRefl(0)’,
’3xRefl(-30)’],

3array([[ 1., 1., 1., 1., 1., 1.], # this is irrep A1

4 [ 1., 1., 1., 1., -1., -1.], # A2

5 [ 1., -1., 1., -1., 1., -1.], # B1

6 [ 1., -1., 1., -1., -1., 1.], # B2

7 [ 2., 1., -1., -2., 0., 0.], # E1

8 [ 2., -1., -1., 2., 0., 0.]])) # E2

The main caveat in using this machinery is that the point groups predefined in
NetKet all leave the origin invariant (except for cubic.Fd3m which represents
the “nonsymmorphic point group” of the diamond/pyrochlore lattice) and thus
only work well with lattices in which the origin has full point-group
symmetry. This behaviour can be changed (see the definition of cubic.Fd3m for
an example), but it is generally safer to define lattices using the proper
Wyckoff positions [75], of which the origin is usually maximally symmetric.

###  5.2 Using group convolutional neural networks (GCNNs)

NetKet uses GCNNs [73, 49] to create NQS ansätze that are symmetric under
space groups of lattices. These networks consist of alternating group
convolutional layers and pointwise nonlinearities. The former can be thought
of as a generalization of convolutional layers to a generic finite group $G$.
They are equivariant, that is, if their inputs are transformed by some space-
group symmetry, features in all subsequent layers are transformed accordingly.
As a result, the output of a GCNN can be understood as amplitudes of the wave
functions $g|\psi\rangle$ for all $g\in G$, which can be combined into a
symmetric wave function using the projection operator (40). Further details
about equivariance and group convolutions are given in Section A.1.

|  mode="irreps" |  mode="fft"
---|---|---
Can be used for |  any group |  only space groups
Symmetries can be specified by |  $\bullet$ Lattice
$\bullet$ PermutationGroup
$\bullet$ Symmetry permutations and irrep matrices |  $\bullet$ Lattice
$\bullet$ PermutationGroup and shape of translation group
$\bullet$ Symmetry permutations, product table, and shape of translation group
Kernel memory footprint per layer |  $O(f_{\mathrm{in}}f_{\mathrm{out}}|G|)$ |  $O(f_{\mathrm{in}}f_{\mathrm{out}}|G||P|)$
Evaluation time per layer per sample |  $O[(f_{\mathrm{in}}+f_{\mathrm{out}})|G|^{2}+f_{\mathrm{in}}f_{\mathrm{out}}\sum_{a}d_{a}^{3}]$ |  $O[(f_{\mathrm{in}}+f_{\mathrm{out}})|G|\log|T|+f_{\mathrm{in}}f_{\mathrm{out}}|G||P|]$
Preferable for |  $\bullet$ large point groups
$\bullet$ if expanded "fft" kernels don’t fit in memory |  $\bullet$ small point groups
$\bullet$ very large batches (see App. A.2)
|  |
Table 4: Comparison of GCNN implementations. $f_{\mathrm{in,out}}$ stands for
the number of input and output features, $|G|,|P|,|T|$ for the sizes of the
space group, point group, and translation group, respectively. $d_{a}$ are the
dimensions of irreps of $G$; in a large space group, $\sum_{a}d_{a}^{3}$
scales as $|G||P|$.

GCNNs are constructed by the function nk.models.GCNN. Symmetries are specified
either as a PermutationGroup or a Lattice. In the latter case, the symmetry
group is given by space_group(); an optional point_group argument to GCNN can
be used to override the default point group. By default, output transforms
under the trivial irrep $\chi_{g}\equiv 1$, that is, all output features are
averaged together to obtain a wave function that is fully symmetric under the
whole space group. Other irreps can be specified through the characters
argument, which takes a vector of the same size as the space group.

⬇

1from netket.graph import Triangular

2from netket.models import GCNN

3from netket.utils.group.planar import C

4

5lattice = Triangular([6,6])

6momentum = [0,0]

7sgb = lattice.space_group_builder()

8chi = sgb.space_group_irreps(momentum)[3]

9

10# This transforms as the trivial irrep Gamma A_1

11gcnn1 = GCNN(lattice, layers = 4, features=4)

12

13# This transforms as Gamma B_2

14gcnn2 = GCNN(lattice.space_group(), characters=chi, layers=4, features=4)

15

16# This does not enforce reflection symmetry

17gcnn3 = GCNN(lattice, point_group=C(6), layers=4, features=4)

NetKet currently supports two implementations of GCNNs, one based on group
Fourier transforms ( mode="irreps"), the other using fast Fourier transforms
on each coset of the translation group ( mode="fft"): these are discussed in
more detail in Section A.2. Their behavior is equivalent, but their
performance and calling sequence is different, as explained in Table 4. A
default mode="auto" is also available. For spin models, parity symmetry
(taking $\sigma^{z}$ to $-\sigma^{z}$) is a useful extension of the U(1) spin
symmetry group enforced by fixing magnetization along the $\sigma^{z}$ axis.
Parity-enforcing GCNNs can be constructed using the parity argument, which can
be set to $\pm 1$.

In addition to deep GCNNs, fully symmetric RBMs [8] are implemented in
nk.models.
RBMSymm as a single-layer GCNN from which the wave function is computed as

$\psi=\prod_{i,g\in G}\cosh f^{(i)}_{g}\implies\log\psi=\sum_{i,g\in
G}\ln\cosh f^{(i)}_{g}.$ (41)

Due to the products (rather than sums) used, this ansatz only supports wave
functions that transform under the trivial irrep. An RBM-like structure closer
to that of ref. [72] can be achieved using a single-layer GCNN:

⬇

1from netket.models import GCNN, RBMSymm

2from netket.nn import logcosh

3

4# fully symmetric RBM

5rbm1 = RBMSymm(group, alpha=4)

6

7# symmetrized RBM similar to (Nomura, 2021)

8rbm2 = GCNN(group, layers=1, features=4, output_activation=logcosh)

##  6 Quantum systems with continuous degrees of freedom

In this section we will introduce the tools provided by NetKet to study
systems with continuous degrees of freedom. The interface is very similar to
the one introduced in Section 2 for discrete degrees of freedom.

###  6.1 Continuous Hilbert spaces

Similar to the discrete Hilbert spaces, the bosonic Hilbert space of $N$
particles in continuous space has the structure

$\displaystyle\mathcal{H}_{\mathrm{continuous}}=\operatorname{span}\\{\ket{x_{0}}\otimes\cdots\otimes\ket{x_{N-1}}:x_{i}\in\mathcal{L}_{i},i\in\\{0,\ldots,N-1\\}\\}$
(42)

where $\mathcal{L}_{i}$ is the space available to each individual boson: for
example, $\mathcal{L}_{i}$ is $\mathbb{R}^{d}$ for a free particle in $d$
spatial dimensions, and $[0,L]^{d}$ for particles confined to a
$d$-dimensional box of side length $L$. In the case of finite simulation
cells, the boundaries can be equipped with periodic boundary conditions.

In the following snippet, we define the Hilbert space of five bosons in two
spatial dimensions, confined to $[0,10]^{2}$ with periodic boundary
conditions:

⬇

1>>> hilb = nk.hilbert.Particle(N=5, L=(10.0, 10.0), pbc=True)

2>>> print("Size of the basis: ", hilb.size)

3Size of the basis: 10

4>>> hilb.random_state(nk.jax.PRNGKey(0), (2,))

5[[0.02952452 0.21660899 2.836163 3.5628846 4.5622005 5.9473248

6 6.104126 8.14864 9.163713 9.263418 ]

7 [9.85617 0.4667458 2.211265 4.1587596 4.250165 6.69916

8 6.5165453 7.3764215 8.508119 0.08060825]]

As we discussed in Section 2.1, the Hilbert objects only define the
computational basis. For that reason, the flag pbc=True only affects what
configurations can be generated by samplers and how to compute the distance
between two different sets of positions. This option does not enforce any
boundary condition for the wave-function, which would have to be accounted for
into the variational ansatz.

###  6.2 Linear operators

Similar to the discrete-variable case, expectation values of operators can be
estimated as classical averages of the local estimator

$\tilde{O}(x)=\frac{\bra{x}\hat{O}\ket{\psi}}{\bra{x}\ket{\psi}}$ (43)

over the (continuous) Born distribution $p(x)\propto|\psi(x)|^{2}$. NetKet
provides the base class ContinuousOperator to write custom (local) operators
and readily implements Hamiltonians of the form ($\hbar=1$ in our units)

$\hat{H}=-\frac{1}{2}\sum_{i}\frac{1}{m_{i}}\nabla_{i}^{2}+\hat{V}\big{(}\\{\mathbf{x}_{i}\\}\big{)}$
(44)

using the predefined operators KineticEnergy and PotentialEnergy. For example,
a harmonically confined system described by
$\hat{H}=-\frac{1}{2}\sum_{i}\nabla_{i}^{2}+\frac{1}{2}\sum_{i}\mathbf{\hat{x}}_{i}^{2}$
can be implemented as

⬇

1# This function takes a single vector and returns a scalar

2def v(x):

3 return 0.5∗jnp.linalg.norm(x) ∗∗ 2

4

5# Construct the Kinetic energy term with unit mass

6H_kin = nk.operator.KineticEnergy(hilb, mass=1.0)

7# Construct the Potential energy term using the potential defined above

8H_pot = nk.operator.PotentialEnergy(hilb, v)

9

10# Sum the two objects into a single Operator

11H = H_kin + H_pot

Operators defined on continuous Hilbert spaces cannot be converted to a matrix
form or used in exact diagonalization, in contrast to those defined on
discrete Hilbert spaces. Continuous operators can still be used to compute
expectation values and their gradients with a variational state.

###  6.3 Samplers

Out of the built-in samplers in the current version of NetKet (Section 3.4),
only the Markov chain Monte Carlo sampler MetropolisSampler supports
continuous degrees of freedom, as both ExactSampler and the autoregressive
ARNNSampler rely on the sampled basis being countable. For continuous spaces,
we provide the transition rule sampler.rules.GaussianRule which proposes new
states by adding a random shift to every degree of freedom sampled from a
Gaussian of customizable width. More complex transition rules can be defined
following the instructions provided in Section 3.4.

###  6.4 Harmonic oscillators

---
Figure 1: VMC energy estimate as a function of the optimization step for a
continuous-space system of $N=10$ particles in $d=3$ spatial dimensions
subject to a harmonic confinement.

As a complete example of how to use continuous-space Hilbert spaces,
operators, variational states, and the VMC driver together, consider 10
particles in three-dimensional space, confined by a harmonic potential
$V(x)=x^{2}/2$. The exact ground-state energy of this system is known to be
$E_{0}=15$. We use the multivariate Gaussian ansatz
$\log\psi(x)=-x^{\mathrm{T}}\Sigma^{-1}x$, where $\Sigma=TT^{\intercal}$ and
$T$ is randomly initialized using a Gaussian with zero mean and variance one.
Note that the form of $\Sigma$ ensures that it is positive definite.

⬇

1import netket as nk

2import jax.numpy as jnp

3

4def v(x):

5 return 0.5∗jnp.linalg.norm(x) ∗∗ 2

6

7hilb = nk.hilbert.Particle(N=10, L=(jnp.inf,jnp.inf,jnp.inf), pbc=False)

8ekin = nk.operator.KineticEnergy(hilb, mass=1.0)

9pot = nk.operator.PotentialEnergy(hilb, v)

10ha = ekin + pot

11

12sa = nk.sampler.MetropolisGaussian(hilb, sigma=0.1, n_chains=16,
n_sweeps=32)

13model = nk.models.Gaussian(param_dtype=float)

14vs = nk.vqs.MCState(sa, model, n_samples=10 ∗∗ 4, n_discard=2000)

15

16op = nk.optimizer.Sgd(0.05)

17sr = nk.optimizer.SR(diag_shift=0.01)

18

19gs = nk.VMC(ha, op, sa, variational_state=vs, preconditioner=sr)

20gs.run(n_iter=100, out="HO_10_3d")

We show the training curve of above snippet in Fig. 1; exact ground-state
energy is recovered to a very high accuracy.

###  6.5 Interacting system with continuous degrees of freedom

In this example we want to tackle an interacting system of bosonic Helium
particles in one continuous spatial dimension. The two-body interaction is
given by the Aziz potential which qualitatively resembles a Lennard-Jones
potential [79, 80, 81]. The Hamiltonian reads

$\displaystyle
H=-\frac{\hbar^{2}}{2m}\sum_{i}\nabla_{i}^{2}+\sum_{i<j}V_{\mathrm{Aziz}}(r_{ij})$
(45)

We will examine the system at a density
$\rho=\frac{N}{L}=0.3$\mathrm{\SIUnitSymbolAngstrom}$^{-1}$ with $N=10$
particles in units where $\hbar=m=k_{b}=1$. To confine the system it is placed
in a box of size $L$ equipped with periodic boundary conditions. The Hilbert
space and sampler are initialized as shown above ($r_{m}$ is the length-scaled
defined in the Aziz potential):

⬇

1import netket as nk

2

3N = 10

4d = 0.3 # 1/Angstrom

5rm = 2.9673 # Angstrom

6L = N / (0.3 ∗ rm)

7hilb = nk.hilbert.Particle(N=N, L=L, pbc=True)

8sab = nk.sampler.MetropolisGaussian(hilb, sigma=0.05, n_chains=16,
n_sweeps=32)

####  6.5.1 Defining the Hamiltonian

We can define the Hamiltonian through the action of the interaction-potential
on a sample of positions $x$, and combine it with the predefined kinetic
energy operator. Since we are using periodic boundary conditions, we will use
the Minimum Image Convention (MIC) to compute distances between particles. In
the following snippet the Aziz potential (in the units above) is defined and
the Hamiltonian is instantiated:

⬇

1

2def minimum_distance(x, sdim):

3 """Computes distances between particles using mimimum image convention"""

4 n_particles = x.shape[0] // sdim

5 x = x.reshape(-1, sdim)

6

7 distances = (-x[jnp.newaxis, :, :] + x[:, jnp.newaxis, :])[

8 jnp.triu_indices(n_particles, 1)

9 ]

10 distances = jnp.remainder(distances + L / 2.0, L) - L / 2.0

11

12 return jnp.linalg.norm(distances, axis=1)

13

14def potential(x, sdim):

15 """Compute Aziz potential for single sample x"""

16 eps = 7.846373

17 A = 0.544850 ∗ 10 ∗∗ 6

18 alpha = 13.353384

19 c6 = 1.37332412

20 c8 = 0.4253785

21 c10 = 0.178100

22 D = 1.241314

23

24 dis = minimum_distance(x, sdim)

25 return jnp.sum(

26 eps

27 ∗ (

28 A ∗ jnp.exp(-alpha ∗ dis)

29 - (c6 / dis ∗∗ 6 + c8 / dis ∗∗ 8 + c10 / dis ∗∗ 10)

30 ∗ jnp.where(dis < D, jnp.exp(-((D / dis - 1) ∗∗ 2)), 1.0)

31 )

32 )

33

34ekin = nk.operator.KineticEnergy(hilb, mass=1.0)

35pot = nk.operator.PotentialEnergy(hilb, lambda x: potential(x, 1))

36ha = ekin + pot

####  6.5.2 Defining and training the variational Ansatz

There are two properties that the variational Ansatz for this system must
obey:

  1. 1.

It must be invariant with respect to the permutations of its particles,
because they are bosons;

  2. 2.

As the interaction resembles a Lennard-Jones potential we have a strong
divergence in the potential energy when particles get close to each other.
This divergence must be compensated by the kinetic energy.

We satisfy the permutation-invariance by using a neural network architecture
called DeepSets. DeepSets exploit the fact that any function
$f(x_{1},...,x_{N})$ which is invariant under permutations of its inputs can
be decomposed as [82]:

$\displaystyle f(x_{1},...,x_{N})=\rho\left(\sum_{i}\phi(x_{i})\right)$ (46)

where $\rho$ and $\phi$ are arbitrary functions. In this specific example,
$x_{i}$ denotes a single-particle position and $\rho$ and $\phi$ are
parameterized with dense feed forward neural networks.

The second requirement is fulfilled by using Kato’s cusp condition which
states that [83]

$\displaystyle\lim_{r\to
0}\left(\frac{\nabla^{2}\psi_{c}(r)}{\psi_{c}(r)}+V(r)\right)<\infty$ (47)

where $r$ denotes the distance between the particles and $\psi_{c}$ is the
cusp wave-function. For the case of a Lennard-Jones potential ($\propto
r^{-12}$-divergence), we have

$\psi_{c}(r)=\exp\left[-\frac{1}{2}\left(\frac{b}{r}\right)^{5}\right],$ (48)

where $b$ is a variational parameter.

We also need to handle the periodic conditions, making sure that the wave-
function does not exhibit divergent behaviour at the edges of the (periodic)
box. To this end we will use a surrogate distance function for the minimum
image convention, namely
$d_{\sin}(x_{i},x_{j})=\frac{L}{2}\sin(\frac{\pi}{L}r_{ij})$ in the
variational Ansatz. Additionally we replace the single-particle coordinates in
Eq. 46 by the two-particle distances $d_{\sin}(x_{i},x_{j})^{2}$, such that
all in all our variational Ansatz reads

$\displaystyle\psi(x_{1},...,x_{N})=\exp\left[\rho\left(\sum_{i<j}\phi(d_{\sin}(x_{i},x_{j})^{2})\right)\right]\cdot\exp\left[-\frac{1}{2}\left(\frac{b}{d_{\sin}(x_{i},x_{j})}\right)^{5}\right]$
(49)

The variational ansatz we described here is implemented in NetKet as
DeepSeetRelDistance, and a more in-depth discussion can be found in this
reference [84].

---
Figure 2: VMC energy estimate as a function of the optimization step for a
continuous-space system of $N=10$ particles in $d=1$ spatial dimensions
subject to a LJ-like interaction potential placed within a periodic box. The
green dashed line is the result given in the supplementary material of [80].

Having defined the ansatz, we run the VMC driver with the given variatianal
Ansatz to find an estimation of the ground-state energy of the system. This is
done as follows:

⬇

1model = nk.models.DeepSetRelDistance(

2 hilbert=hilb,

3 cusp_exponent=5,

4 layers_phi=2,

5 layers_rho=3,

6 features_phi=(16, 16),

7 features_rho=(16, 16, 1),

8)

9vs = nk.vqs.MCState(sab, model, n_samples=4096, n_discard_per_chain=128)

10

11op = nk.optimizer.Sgd(0.001)

12sr = nk.optimizer.SR(diag_shift=0.01)

13

14gs = nk.VMC(ha, op, sab, variational_state=vs, preconditioner=sr)

15gs.run(n_iter=1000, out="Helium_10_1d")

The result of this optimization and a comparison to literature results is
displayed in Fig. 2.

##  7 Example: Finding ground and excited states of a lattice model

In this example, we define the $J_{1}$–$J_{2}$ Heisenberg Hamiltonian

$H=J_{1}\sum_{\langle
ij\rangle}\vec{\sigma}_{i}\cdot\vec{\sigma}_{j}+J_{2}\sum_{\langle\langle
ij\rangle\rangle}\vec{\sigma}_{i}\cdot\vec{\sigma}_{j},$ (50)

on a $10\times 10$ square lattice and use the VMC code introduced in Section
4.1 to find a variational approximation of its ground state. This model gives
rise to several phases of matter, including magnetically ordered states, a
valence bond solid, and a quantum spin liquid. Here, we set
$J_{1}=1,J_{2}=0.5$, inside the spin liquid phase [30, 85].

Our example is optimized to run on a single GPU with 16 GB of memory. We will
make note of what should be changed when running the simulation on CPUs.

###  7.1 Defining the lattice and the Hamiltonian

We use the Lattice class to define the square lattice and generate its space-
group symmetries. By passing max_neighbor_order=2 to the constructor, we
generate graph edges for both nearest and next-nearest neighbours. The pre-
defined Heisenberg class supports passing different coupling constants for
both types of edge.

⬇

1import netket as nk

2import numpy as np

3import json

4from math import pi

5

6L = 10

7# Build square lattice with nearest and next-nearest neighbor edges

8lattice = nk.graph.Square(L, max_neighbor_order=2)

9hi = nk.hilbert.Spin(s=1 / 2, total_sz=0, N=lattice.n_nodes)

10# Heisenberg with coupling J=1.0 for nearest neighbors

11# and J=0.5 for next-nearest neighbors

12H = nk.operator.Heisenberg(hilbert=hi, graph=lattice, J=[1.0, 0.5])

###  7.2 Defining and training a symmetric ansatz

To enforce all spatial symmetries of (50), we use the GCNN ansatz described in
Section 5. By default, the GCNN projects onto the symmetric representation,
which contains the ground state for this geometry. We select singlet states by
only sampling basis states with $\sum_{i}S_{i}^{z}=0$ and setting spin parity
using parity=1. We use a model with four layers, each containing four feature
vectors (i.e., four hidden units for each of the $8L^{2}$ space-group
symmetries). To exploit the high degree of parallelism of GPUs, we set
sampler.n_chains equal to vstate.n_samples202020Note that this results in a
somewhat non-standard MC scheme where, instead of an ensemble of chains with
generally non-zero internal autocorrelation, the sampler produces an ensemble
of independently drawn single configurations ( n_samples_per_chain==1). Since
the sampler state of the previous VMC step is used as initial state for the
next step, there can still be a residual autocorrelation with the previous
samples, which is, however, alleviated by the sampler performing
$N_{\mathrm{sites}}$ intermediate updates before yielding the single requested
sample with the settings used here.. When using CPUs, n_chains should be set
to a smaller value.

⬇

1machine = GCNN(

2 symmetries=lattice,

3 parity=1,

4 layers=4,

5 features=4,

6 param_dtype=np.complex128,

7)

8sampler = nk.sampler.MetropolisExchange(

9 hilbert=hi,

10 n_chains=1024,

11 graph=lattice,

12 d_max=2,

13)

14opt = nk.optimizer.Sgd(learning_rate=0.02)

15sr = nk.optimizer.SR(diag_shift=0.01)

16vstate = nk.vqs.MCState(

17 sampler=sampler,

18 model=machine,

19 n_samples=1024,

20 n_discard_per_chain=0,

21 chunk_size=4096,

22)

23gs = nk.VMC(H, opt, variational_state=vstate, preconditioner=sr)

24gs.run(n_iter=200, out="ground_state")

25

26data = json.load(open("ground_state.log"))

27print(np.mean(data["Energy"]["Mean"]["real"][-20:])/400)

28# Output: -0.49562531096409457

29print(np.std(data["Energy"]["Mean"]["real"][-20:])/400)

30# Output: 0.0002492304032505182

We specify chunk_size=4096 in the variational state in order to reduce memory
consumption. As we have $L^{2}=100$ sites, at every VMC step we will need to
evaluate the network for
$\mathcal{O}(N_{\text{samples}}L^{2})=\mathcal{O}(10^{3}\cdot 10^{2})$
different configurations, but the memory available on commercial GPUs will not
be enough to perform this computation in a single pass. Instead, by setting
chunk_size NetKet will split the calculation in many smaller sub-calculations
(see Section 3.3.1 for more details).

This calculation, which takes about 30 minutes on an NVIDIA A100 GPU, already
delivers a fairly accurate variational energy. The evolution of variational
energy during the training procedure is shown in Fig. 3.

We note that a typical initialization of a GCNN gives rise to ferromagnetic
correlations, which can make training an antiferromagnetic Hamiltonian
unstable [49, 15]. Therefore, it is often good practice to pre-optimize the
phases by restricting all amplitudes to unity by setting equal_amplitudes=True
switch and training only the phases of the network. These parameters can then
be loaded into a model with equal_amplitudes=False.

Figure 3: Energy evolution of variational ground states (green) and excited
states after transfer learning (blue), compared to the lowest known
variational energy for the $10\times 10$ square-lattice $J_{1}$–$J_{2}$ model
[30] (black). The variational energies can be further improved by allowing
more training steps, Monte Carlo samples, etc. The plot on the right zooms in
on the lowest energies.

###  7.3 Finding an excited state

We can also find low-lying excited states using this procedure, by projecting
the wavefunction onto a different irrep. Here, we consider the first gapless
mode in the Anderson tower of states of the Néel antiferromagnet [86], a
triplet at wave vector $(\pi,\pi)$ that transforms trivially under all point-
group symmetries. This mode is still gapless in the quantum spin liquid; we
project out spin-singlets by focusing on parity odd states.