This PR implements most of the features discussed in #703 pertaining to symmorphic space groups.

New stuff and changes

Groups

• New module netket.utils.group to replace semigroup.py: it takes care of everything related to symmetry groups
  • contains SemiGroup, identity, Element, PermutationGroup, Permutation with unchanged API
  • some new features, e.g. Permutation objects can carry an arbitrary name
• New class Group: base class for all group-like objects that are guaranteed to satisfy group axioms
  • methods to calculate inverse mapping, times table, conjugacy classes, character tables (using Burnside's algorithm) for a generic group
  • equality checking via the function _canonical() provided by subclasses: this must return an integer array for all group elements such that equal arrays imply equal group elements (this is a property of the specific group classes to allow them to handle Identity as they see fit)
• New class PGSymmetry: represents a point group symmetry around the origin, specified by a transformation matrix
  • autogenerated name describes transformation in human-readable form
• New class PointGroup: stores PGSymmetry objects
• All crystallographically relevant point groups provided in submodules planar (2D), axial, cubic (3D)
• New class SpaceGroupBuilder (in netket.graph)
  • translates PointGroups into PermutationGroups acting on a particular Lattice
  • generates the translation group of a Lattice as PermutationGroups (with sensible names attached)
  • hence generates space groups as PermutationGroups
  • helps calculate the character table of the space group intuitively (i.e., calculates irreps consistent with a given wave vector)

Graphs

• The custom rotation etc. groups of Lattice are removed in favour of using the above machinery

• ~NetworkX.automorphisms() changed to only return an array of permutation indices; the PermutationGroup is made by a free-floating method in symmetry. (SpaceGroupBuilder needs to use Lattice, which means that having a reference to anything in symmetry within Lattice produces a circular import via symmetry/__init__py.) It is also deprecated and should eventually be replaced by a hidden method that feeds into symmetry.automorphism_group(): it makes little sense to have this single piece of symmetry functionality outside symmetry. Alternatively, SpaceGroupBuilder could live in the graph module and be blended into the functionality of Lattice, similar to how automorphisms() behaves now.~

• Lattice is given several new methods:

  • space_group_builder() returns a SpaceGroupBuilder object (see above) corresponding to the translations of the lattice and the supplied point group. The Lattice constructor also takes a PointGroup argument that is cached as a default point group.
  • point_group() returns the representation of its PointGroup argument or the default point group as a PermutationGroup
  • rotation_group() picks out the rotations (determinant of rotation matrix is +1)
  • translation_group() returns the group of lattice translations as a PermutationGroup. It takes an optional argument to specify the axes along which to translate
  • space_group() is the semidirect product translation_group() @ point_group().

  All of these are convenience wrappers around methods of space_group_builder().

• The "hashing logic" in Lattice is tidied up and extended to wave vectors (needed in SpaceGroupBuilder). It now honours periodic and open BCs.

• The Grid class is removed and replaced by functions of the same calling sequence that return Lattices. (The space-group functionality is built around Lattices, so it is better to focus on improving that one API rather than developing several independent ones.)

  • This breaks Grid's ability to colour its edges by direction. A more flexible constructor for Lattice will solve this problem.
  • planar_rotation() and axis_reflection() are dropped as they were in Lattice
  • The name space_group() was used incorrectly: instead of that and lattice_group(), we have Lattice.point_group() and Lattice.space_group(). Deprecation would be hard, since one of the names is reused in a different meaning.
  • point_group() only returns symmetries that leave the origin in place. This is different from the original behaviour for open BC axes (could be fixed by allowing nonsymmorphic point groups, but I need more reason than this to implement those).

• Specialised constructors for triangle, honeycomb and kagome lattices are added.

Odds and ends

• The "hashing logic" (used both in PointGroup and Lattice) is moved into netket.utils.float~_utils~ and fine-tuned. It also supports a mix of periodic and open boundary conditions. I've also added
  • a function that prunes nearly-zero real and imaginary parts from an array;
  • a function that checks whether elements of an array are nearly integers.
• netket.jax.logsumexp is added, which extends the functionality of JAX logsumexp to handle complex numbers well (i.e., it forces the output to be complex and doesn't error out on complex inputs/outputs)
• Functions to project the outputs of DenseSymm and DenseEquivariant onto irreps using their characters. The default flavour uses logsumexp, but there is one with plain sums too.

Typical workflows

A simple workflow, without character tables

We just want to generate the space group of a Lattice given a point group we know it's invariant under:

from netket.utils import group
from netket.graph import Lattice

graph = Lattice(basis_vectors = [[1,0],[0.5,0.75**0.5]], extent = (6,6)) # triangle lattice
space_group = graph.space_group(group.planar.D(6))

The resulting space_group is a PermutationGroup that can be used directly in a GCNN, for instance. Alternatively, we can use the premade triangular lattice that is loaded with the D_6 group:

from netket.graph import TriangularLattice

graph = TriangularLattice([6,6])
space_group = graph.space_group()

Using character tables

For this, one needs a basic appreciation of how crystallographic character tables are constructed. They can be described in terms of a wave vector (or rather a star of symmetry-related wave vectors) and the irreps of the corresponding little group (the subgroup of the point group that leaves the wave vector unchanged). The latter can be read off from a human-readable character table one can generate in an interactive session:

from netket.utils import group
from netket.graph import TriangularLattice
from math import pi

graph = TriangularLattice([6,6]) 
sgb = graph.space_group_builder()

k = [4*pi/3,0] # corner of the hexagonal BZ
sgb.little_group(k) 

> PointGroup(elems=[Id(), Rot(120°), Rot(-120°), Refl(0°), Refl(-60°), Refl(60°)], ndim=2)

sgb.little_group(k).character_table_readable()

> (['1xId()', '2xRot(120°)', '3xRefl(0°)'], 
array([[ 1.,  1.,  1.],
       [ 1.,  1., -1.],
       [ 2., -1.,  0.]]))

The first output confirms that the little group of D_6 at the corner of the Brillouin zone is D_3, whose known character table is generated by the second command; given the labels in the first part of the output, they are easy to match to the characters in these tables, so we can look up their physical/geometrical meaning. Any of these can be turned into an irrep of the full space group using SpaceGroupBuilder: in fact, it generates all of them as a 2D array (in the same order as the irreps printed above), so we'd write something like

chi = sgb.space_group_irreps(k)[2] # [2] selects the "E" irrep
# ...
# in the definition of the GCNN
return irrep_project_logsumexp(output, chi)

To do

• ~Writing tests. I tested most of the stuff manually and it seems to work in all cases, but of course it has to be more systematic.~
• Writing docs. Probably the best place for the kind of workflow docs you see above would be in @chrisrothUT's tutorial on GCNNs, and #700 will be updated with how the abstract stuff gets implemented here.
• ~Checking if the stuff that got caught up in an earlier git-rebase (see first 2 commits) affects the behaviour of struct.dataclass. Everything seems to work fine, so I'm not too worried, but @PhilipVinc could you perhaps check and suggest what I should do?~
• Non-symmorphic groups? I've given some thought to it, PointGroup wouldn't be too hard to extend, the main question is whether the automatic construction of character tables generalises nicely. I have a hunch that it does, but I would need some downtime with a group theory textbook to make sure. Probably left for another PR
• Extending Lattice so it can have further-neighbour and coloured edges (this is functionality that is lost from the new Grid for instance). I just flag this up, but it can wait.

An Easter egg

The NetworkX algorithm really looks for all automorphisms:

lattice = nk.graph.Square(4)
len(symmetry.automorphism_group(lattice))
> 384
len(symmetry.space_group(lattice, symmetry.planar.D(4)))
> 128

It turns out that a 4x4 square lattice with PBC is isomorphic to a 2^4 hypercube, which has many more symmetries. E.g., you can check that this maps nearest neighbours to nearest neighbours without making any geometrical sense:

 0,  1,  5,  4
 3,  2,  6,  7
15, 14, 10, 11
12, 13,  9,  8

This is an interesting caveat for using NetworkX graph matching. PS. The 3×3 triangle lattice turns out to have 1296 isomorphisms, of which only 108 are space-group symmetries!

I put this here in case anybody wants to see the result of our discussion the other day. It implements the design described in #649.

I'm not done yet but shouldn't be too far off. I think the only thing left to do for me is to rewrite the nk.optimizer.SR in order to detect if using the old api/new api, convert to new api and print deprecation warnings. +tests

Summary for version 2.0

This PR introduces the first beta version for NetKet 2.0.

Major changes in version 2.0

1. NetKet now fully exposes its internal types and classes, thus becoming a full-fledged python library built on a monolithic c++ core
2. Python Bindings are provided using pybind11, thus fully addressing issue #8
3. As a result of this transformation, JSON input is not accepted anymore and the netket executable does not exist anymore
4. A great deal of flexibility in the usage of NetKet comes with these changes, removing the majority of constraints coming from the JSON-based input and usage pattern
5. Bindings are provided for the great majority of NetKet classes, with the exception of some internals, that don't need python exposure.

Reasons behind these changes

1. The most common current usage pattern of NetKet is not to write directly the JSON file, but use the python scripts we have in the Tutorials/ folder. That's why moving to a python library is a desirable goal, improving the user experience and opening the way to a number of applications previously hard or impossible to perform
2. With a full-fledged python library it is also much easier to submit multiple jobs with different parameters, and have more flexibility on the output
3. There is a substantial overhead for maintaining the JSON support. Supporting all the possible ways of putting together the NetKet classes, in a consistent way, does not scale well given the increasing size of this project.
4. It becomes easier to interface to a whole spectrum of Python-only libraries, including advanced visualization tools and state-of-the-art machine learning libraries (pytorch, tensorflow etc)
5. New generations of students (unfortunately?) are not very familiar with c++, and strive to have a pure Python library

Installing, Tutorials, Examples

To install the latest development version do

pip install . 

The directory /Tutorials/ contains pynotebook tutorials. More will be added. Several example codes showcasing a specific application are contained in /Examples.

Other remarks

• Version 2.0 removes all "glue" classes, such as Graph,Hamiltonian etc. The corresponding abstract interfaces remain, with little changes.

• We introduce the concept of Operator, which generalizes Hamiltonian,Observable etc. Hamiltonian and Observable disappear.

• We introduce a LocalOperator, and all associated overloaded operations, for example

hi=nk.Spin(s=0.5,graph=g)
X=[[0,1],[1,0]]
o=nk.LocalOperator(hi,X,[0])*(nk.LocalOperator(hi,X,[1]))

defines a quantum operator X(0)*X(1), automatically performing the tensor product. Other operations such as product times a scalar and addition of local operators are also supported.

TLDR

I'd like some feedback on this proposal. It's not yet done, though a big part is complete (I need to put a few @jit calls in the right place). Download it with pip install git+https://github.com/netket/[email protected] and play around a bit.

If you want to review the PR, since it's huge, I'd suggest to look at it commit by commit. Every chunk of changes is in a separate commit. Skip the first commit which is not really important anymore.

see how VMC becomes easy and more generic, or this gists for an example of the api

Netket v3.0 master plan

For the so-called, unreleased version of netket v3.0, until now, we have transitioned all of our infrastructure to python and removed all C++ code. This was done by temporarily using numpy and later converting most code to jax.

Still, the API has largely remained unchanged (except few changes to the construction of hilbert spaces and graphs, mostly aestetical).

v3.0 is still unreleased because me and @gcarleo wanted to get the API right before we commit to it, however, months have passed (almost an year) and this is still unreleased.

The big issues that I would like to address are the following:

1. #437 Add an extensible VariationalState and use it in the drivers, making it easier to develop new applications.
2. #525 Make it easy and intuitive to write custom Metropolis-Hastings transition rules and use them.
3. Support non homogeneous hilbert spaces.
4. Make it easier to define new arbitrary networks, with real or complex weights.
5. Actually support arbitrary networks that mix real and complex layers (right now they aren't)
6. Support models with a state that can change (but should not trigger recompilation)
7. Remove legacy code.

In particular, while giving some lectures on netket I recently noticed Points 2 and 4, that is, it's quite messy to define those things.

Points 3/5 can be resolved by moving to flax, where (as I will argue below) defining models is much more compact and intuitive.

The following is my proposal for netket v3.0 API

A somwhat central point of this API stems from the discussion in #480, namely my proposal, which seemed somewhat accepted, to remove numpy and torch backend and convert netket to a pure jax package. This was already quite well received, and it seems to me that a) jax has now complete support from Google and it's a stable project and b) it allows us to writ more compact, simpler code.

This is my proposal for netket v3.0:

Roughly:

• add a new requirement: jax and flax. ~This will also bump the minimum required version to python 3.7. This should not be an issue as also jax is discussing dropping python 3.6.~ EDIT: it's possible to still support 3.6.

Add to netket the following sub-packages:

• netket.nn : re-exporting flax.nn but wrapping some functions in order to make them work better with complex numbers.

  • Jax/Flax already supports complex numbers, but some activation functions do not for different reasons and the kwargs necessary to use complex weights in a layer are a bit complicated to use. The main reason to re-export and wrap is to make our exported api of flax work out of the box with complex. Slowly, I hope to get some PRs merged in flax itself so we can drop our own code.

  • Example:

    >> import netket as nk
    >> import flax
    >> from jax import numpy as jnp
    
    >> x = jax.random.normal(jax.random.PRNGKey(0), (4,4), dtype=jax.numpy.complex64)
    # flax does not work
    >> flax.nn.activation.softplus(x)
    TypeError: add requires arguments to have the same dtypes, got complex64, float32.
    # netket.nn works
    >> nk.nn.activation.softplus(x)
    DeviceArray([[1.9650044+0.14860448j, 0.362084 -0.14605658j],
                 [0.6844594+0.72896177j, 0.4604799-0.02623011j]],            dtype=complex64)
    
    # To use flax, we need to define our own complex init function
    def complex_kernel_init(rng, shape):
      fan_in = np.prod(shape) // shape[-1]
      x = random.normal(random.PRNGKey(0), shape) + 1j * random.normal(random.PRNGKey(0), shape)
      return x * (2 * fan_in) ** -0.5
    
    complex_bias_init = lambda _, shape: jnp.zeros(shape, jnp.complex64)
    
    # complex-valued dense flax version
    m = flax.nn.Dense(features=3, dtype=jax.numpy.complex64, kernel_init=complex_kernel_init, bias_init=complex_bias_init)
    
    # nk-version
    m = netket.nn.Dense(features=3, dtype=jax.numpy.complex64)
    

    I hope the above convinces you that having flax working out of the box with complex values is handy.

  • Since people might still have jax machines around, we provide a very simple function wrapping a jax module into a flax one so that everything works out of the box (nk.nn.wrap_jax)

• netket.jax : wrapping some functions in order to support complex numbers and functions that may one of R->R, R->C, C->C with the same syntax. Jax does not and will not support this out of the box. Notably, this will have our own version of jax.vjp and jax.grad based on the code we have already in jax_utils, plus a few over utilities.

• netket.optim: (What was netket.optimizer).

  • I propose to change name (with a slow deprecation so not to break code) because optim is the default in jax and pyro world. tensorflfow uses optimisers, and we use optimiser. Let's pick one and be consistent.
  • the optimisers are simply re-exported from flax. No code.
  • We also export SR
    • Sr is rewritten with a new interface. see below.

Remove AbstractMachine and its implementations

• Functionally replaced by pure flax modules, which will be objects that contain no state (parameters) but only two fucntions: the init_params, returning the pytree of params and apply to compute the forward pass.

  • Also support pure jax modules and make it easy to support other jax frameworks (optax, for example).

  • While those modules do not contain the hilbert space, provide an easy to use constructor that accepts an hilbert space and extract the size.

  • By default use np.float32 and not np.float64, but depending on what the user uses, anything is supported

  • An advantage of this is that we can now copy-paste any machine written for jax/flax and they will work out of the box, regardless of what they do! (modulo replacing flax.nn -> netket.nn for complex-number compatibility until things are fixed upstream).

  • See for example how to define a Convnet or a RBM with spin and phase: it's very easy. compare it with our old jax code for a RBMModPhase which

  import netket as nk
  from netket import nn
  
  class RBMModPhase(nn.Module):
      dtype : Any = np.float32
      activation : Any = nknn.logcosh
      alpha : Union[float, int] = 1
      use_bias : bool = True
  
      @nn.compact
      def __call__(self, x):
          re = nknn.Dense(features=self.alpha*x.shape[-1], dtype=self.dtype, use_bias=self.use_bias)(x)
          re = self.activation(re)
          re = jnp.sum(re, axis=-1)
  
          im = nknn.Dense(features=self.alpha*x.shape[-1], dtype=self.dtype, use_bias=self.use_bias)(x)
          im = self.activation(im)
          im = jnp.sum(im, axis=-1)
  
          return mod + 1j * im
  
  
  class CNN(nn.Module):
    @nn.compact
    def __call__(self, x):
      x = nn.Conv(features=32, kernel_size=(3, 3))(x)
      x = nn.relu(x)
      x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
      x = x.reshape((x.shape[0], -1))  # flatten
      x = nn.Dense(features=256)(x)
      x = nn.relu(x)
      x = nn.log_softmax(x)
      return x
  
  machine = CNN(hilbert)
  W = machine.init_params(rng_key) # rng_key either an int/uint or jax.random.PRNGkey
  

  • We wish to support non-differentiable variables in machines (for RNNs, batchnorm or other things.). To do so, we adopt the flax standard: the parameters pytree has the shape {'params':params_pytree, **other_state}.

Rewrite the samplers to be functional. Mainly:

• Samplers are now datastructures only containing the parameters of the sampler, and no state, nor machine.

  • Using datastructures allow us to pass those objects straight to jitted functions without issues.
  • The state of a sampler is stored in a separate struct.
    • The state mainly carries the current rng state, plus additional stuff if needed.
  • The api will be comprised of the following functions:
    • netket.sampler.init_state(sampler, machine, params) or sampler.init_state(), creating the state
    • netket.sampler.sample(sampler, machine, params, chain_length=XXX, state=None) or sampler.sample(...)`, sampling n_samples
      • If state is None, then init_state is used to create a new_state
      • Returns the modified state and sampled values.
      • I chose chain_length instead of the old n_samples because that's technically what it is.

• Example:

  >>>import netket as nk
  >>>from netket import nn
  >>>sampler = nk.sampler.ExactSampler(hilbert, seed=0)
  ExactSampler(
    hilbert = Spin(s=1/2, N=4),
    seed = [         0 2385058908],
    n_batches = 8,
    machine_power = 2)
  # notice the seed
  
  # create the state for the sampler:
  state = nk.sampler.init_state(sampler, machine, params)
  # reset the chain (could also have passed state = None, and it would create it)
  state = nk.sampler.reset(sampler, machine, params, state)
  
  samples, state = nk.sampler.sample(sampler, machine, params, chain_length = 1000, state=state)
  >>> samples
  DeviceArray([[[-1., -1., -1., -1.],
                [-1.,  1.,  1.,  1.],
                [-1., -1., -1., -1.],
                ...,
  >>> state
  ExactSamplerState(pdf=DeviceArray([3.8351530e-01, 2.0624077e-02, 7.7783165e-04, 7.2943498e-03,
               4.4462629e-02, 1.9844480e-04, 9.5614446e-03, 3.3565965e-02,
               3.3565965e-02, 9.5614446e-03, 1.9844480e-04, 4.4462629e-02,
               7.2943498e-03, 7.7783165e-04, 2.0624077e-02, 3.8351530e-01], dtype=float32), 
              rng=DeviceArray([1255698341, 3859703708], dtype=uint32))
              # notice the rng
  

  • We store in the sampler the seed, so that if you reuse the same samplere, without carrying with you the state, you will get the same samples again. I had a brief discussion about this with @gcarleo and it seems the most sensible thing to do.

  # If state is not passed in, the state is automatically creaeted...
  samples, state = nk.sampler.sample(sampler, machine, params, chain_length = 1000)
  >>> samples
  DeviceArray([[[-1., -1., -1., -1.],
                [-1.,  1.,  1.,  1.],
                [-1., -1., -1., -1.],
                ...,
  >>> state.rng
  DeviceArray([1255698341, 3859703708], dtype=uint32)                
  # notice the rng: it's the same as before
  

  • While the interface is fully functional, you can also call those methods as my_sampler.sample(...) and so on.
  • This is all pure jax, even the sampler themselvees, so everything can be jitted through.
    • If you create a new sampler, the function is not recompiled. The only thing triggering recompilation is changing the n_batches, the hilbert or other things declared static.

### Implement a variationalstate

• A variationalState has the following interface:

  • vs.parameters : (or params?) returns the PyTree of the variational parameters one may want to optimise.

  • vs.expect(operator) : computes expectation value of operator

  • vs.expect_and_grad(operator, is_hermitian=auto/True/False) computes the expectation value of operator, and the gradient of it.

    • is_hermitian can be used to decide whever to use the simpler form we use for the gradient energy now, or the more standard formula.

  • vs.QGT() -> Callable[Grad, Grad] : returns the quantum geometric tensor/ S matrix.

    • The returned object should be a lazy object (a la scipy.functor) that takes as input a gradient and returns another gradient.

  • vs.reset() : resets the internal state among iterations

  • save/load

  • For a machine/sampler, the ClassicalVariationalState will be thee first implementation of this interfacee.

  • It will provide some functionality similar to a sampler/machine from before, with a bunch of extra tricks:

    • A ClassicalVariational State is constructed by taking a
      • hilbert,
      • Machine/Module
      • Sampler
      • optional SR object? (~maybe~ probably)
      • some configuration data
        • chain_length
      • the api is the one deescribed above plus:
        • vs.sample(chain_length) : sample and store the samples internally until a reset() is callede
        • vs.samples : access the samplse. if it was resetted, resample.
        • vs.model_state : returns the PyTree of any other parameters that might change but we don't want to differentiate against (think batchnorm, rnn...). Might also put it in the common api. I'm not sure.

Minor changes:

• the file netket/utils.py has been moved into a subfolder and utils is now a full fledges submodule. I moved here our logic for detecting if MPI is installed, deprecation warnings, and so on.
• I more carefully check that the discoverable names in every module are things that are relevant. For example, before we had that in every module like netket.hilbert there was both Spin (the class) and spin (the module generated by file). There is a small utility in netket.utils to hide all file-generated modules that we don't want and I use it extensively.
• netket.hilbert.Boson has been renamed to netket.hilbert.Fock, as I find it more accurate. However, there is a deprecated constructor called Boson that forwards to Fock.
• Maybe we should rename all Samplers from MetropolisSampler, ExactSampler to Metropolis, Exact ? As we usually use them from their module (netket.sampler) it might make sense, and makes the API lighter.
• The old API to create samplers is deprecated (I'd like to remove it) but it is still there. If you try to construct a MetropolisSampler with a machine, it will give you the same objects as before.
• I'd like to remove netket.random. There is nothing of interest there anymore.

People

@inailuig If you have some time I'd like to know if this can solve your problem of recompiling. I think it should.

Dear all,

I wrote a wrapper for using the IETL Lanczos code (http://alps.comp-phys.org/static/doc2.1.0/ietl.html). The ground state energy computation in NetKet would now be done with this library. Probably we should discuss some more detailed user interface for the algorithm (how to set parameters, like precision for the Lanczos iteration. There are also some test, which could be extended.

Best, Alex

I have implemented the SR algorithm with precomputed gradients. This is less elegant than the lazy vjp/jvp-based implementation in LazySMatrix, but has practical advantages:

• It has minimal memory overhead: precomputing the gradients requires the same amount of memory as a single pass of vjp since the different samples on which the neural network is run are independent, so we only need to compute one of them to get a row of the Jacobian. (I.e., instead of looping through vjp with vmap as jacrev does, we can loop through grad.) Storing the matrix of gradients is guaranteed to take up less memory: backpropagation has to store a lot of internal information about the forward pass, which takes up many times the memory needed for the gradients in a deep network (In my experiments, they were 40 MB vs 1.2 GB). The bottom line is that if there is enough memory for a single vjp, there is enough memory for this too.
• It yields a massive speedup: the gradient matrix can again be calculated in the same asymptotic time as a single vjp (it will of course be a bit slower because vmap is used, but not by large factors); afterwards, we only need matrix-vector multiplications, which are way faster than a full backpropagation of a complex neural network. In my experiments (on 20 CPUs), I could reduce the time of an SR step from 8x that with Adam to about 1.2x.
• It allows for regularising the S matrix in a scale-invariant way by factoring out the magnitude of diagonal elements, as described in Becca & Sorella, p. 143. This has minimal overhead and can be crucial for heterogeneous networks that may have very different gradients in different parts.

Unfortunately, calculating the full Jacobian cannot be done as dtype-agnostically as a VJP, so some information about the network structure needs to be passed. This is done through the jacobian parameter, which is implemented for the values "R2R" (both the network entries and the wave function are real), "R2C" (real entries, complex wave function), and "holomorphic" (holomorphic function of complex parameters), None (uses a LazySMatrix).

• The parameter could use a better name, but jacobian also signifies it is a switch for using this code vs. the original one
• It may be possible to automate this choice, although it's not trivial (e.g., R2R is the right choice even if the wave function has negative entries, which makes the output complex...)
• A few more cases might be implemented, e.g. non-holomorphic C2C, although that is just sugar for R2C...

The boolean parameter rescale_shiftspecifies whether the scale-invariant regularisation should be used.

This is an implementation of discrete fermionic degrees of freedom by @imi-hub and myself. It contains the hilbert space and an operator. The operator handles the minus signs coming from the exchange symmetry, while the hilbert space stores occupation numbers only. This follows e.g. the way qiskit and openfermion would implement things. For the operator, we follow the implementation of openfermion and added an 'from_openfermion' function to create operators. This allows us to use all their hamiltonians, such as Fermi-Hubbard. The hilbert space is not optimized for fixed number of fermions (since it inherits from HomogeneousHilbert), for that we might need to add another Hilbert implementation.

FYI: formatting will fail because flake8 and flakehell are giving me hell.

This implements translations, rotations and reflections on graphs generated by Lattice.py.

Let’s define a square grid to show how we use this

graph = nk.graph.Lattice(basis_vectors = [[1,0],[0,1]], extent=[5,5])

We can generate a SymmGroup with translations by the basis vectors as follows

graph.basis_translations()

In order to generate rotations we need to specify the period. This will generate the C4 rotational symmetry group

graph.planar_rotations(period=4,axes=(0,1))

If we specify a period that doesn’t map the lattice to itself we get a value error. Finally we can generate the reflections about a plane

graph.reflections(axis=0)

This should work for arbitrary dimensions (I only tested a cube). @femtobit @PhilipVinc @attila-i-szabo can you help me bug hunt?

`

@attila-i-szabo noticed that jax.nn.selu loses it's self normalizing property with complex numbers. Namely, it applies selu(x)when it should be applying selu(Re(x)) + 1j*selu(Im(x)). Making this small change improves performance tremendously. I introduce the correct non-linearity as piecewise_selu

Me and @gcarleo were recently discussing the fact that we should also redesign the interface used to access the S matrix and to perform Stochastic Reconfiguration Natural Gradient.

The aim is to achieve easier extendability (so that anyone can write his own version of the S matrix if he wants) and composability (play well with solvers from jax/scipy and others, possibly without requiring to wrap them as it's needed now).

This Issue wants to discuss two items:

• Whever we should do this redesign for v3.0, when to deprecate the old interface, or if we should put it on hold for a v3.1 or 3.2 later.
• The design of the new system

To recap, the current interface is the following:

• A SR object holds settings on the type of S matrix representation used (dense, lazy-onthefly, lazy-jacobian...) and the algorithm used to solve Sx=F. Every object should correspond to only 1 set of choices (lazy-onthefly+cg, lazy-onthefly+gmres, ecc...) . In netket v3.0b1 there is only one possible representation of the S matrix so that was not a big issue.

  • As shown in the PR #648, this design has issues, as adding a new type of S matrix representation (lazy-jacobian in the PR) requires either duplicating all the types corresponding to solvers, which is inconvenient
  • The natural thing to do would be to split the S matrix representation from the solver used.

• The SR object can build the S matrix itself if it's given a variational state S = sr.create_S(state).

• The S matrix keeps a reference to the SR object, so that you can do S.solve(F) and it will use the parameters from the sr object.

sr = nk.optimizer.sr.LazyCG(diag_shift=0.01, maxiters=100)
S = vstate.quantum_geometric_tensor(sr) # equivalent to S = sr.create(vstate)
_, F = vstate.expect_and_grad(ham)

x, info = S.solve(F) # uses an inner field S.sr for the parameters and solve function

While the system described above, where all configuration is stored in a single structure, it makes it rather easy to use SR inside of a driver, as we can simply pass

gs = nk.VMC(ham, optim, variational_state=vstate, sr=sr)

and inside the driver will use this object like shown above.

--

Tentative new design:

Goal: be able to use scipy/jax solvers out of the box.

• Every kind of SR matrix has it's own constructor. SMatrixOnTheFly(vstate), SMatrixJacbian(vstate) etcetera that can be used explicitly.

• An helper function SMatrix(SMatrixType, vstate) = SMatrixType(vstate) is provided. If SMatrixType is not passed some sensible default representation that always works is used.

• vstate.quantum_geometric_tensor() will now take as input the SMatrix type and relay the call to the type. If no type is passed, a sensible default is used.

• All S matrix must have the methods __matmul__ and __call__ supporting both PyTrees and dense vectors, so that the matrix can be passed to sparse solvers.

With this, what is below should work.

S = SMatrixOnTheFly(vstate)
x, info = jax.scipy.sparse.cg(S, F, maxiter=100)

A question arises: we often regularise the S matrix to have a shift in the diagonal. To support it under this API, we should also ask all implementation of the S matrix to support ___add__(self, x:Number) and keep the diagonal shift in memory (if it's a lazy representation) or simply add it to the dense matrix if it's a dense representation.

Some implementations might even support __mul__ or other conditionings.

We will then be able to do

S = SMatrixOnTheFly(vstate)
S = S + 0.01 # equivalent to S.diag_shift+=0.01
x, info = jax.scipy.sparse.cg(S, F, maxiter=100)

(Note: this is inconsistent with numpy api, where adding a number to a matrix adds it to all the entries in the matrix, but is consistent with our implementation of LocalOperators where adding a number to a local operator only adds it to the diagonal).

We could even get this to work with scipy (not jax) sparse solvers if we also implement S.shape to report the number of parameters.

S = SMatrixOnTheFly(vstate)
# assuming S.shape = (vstate.n_parameters, vstate.n_parameters)
S = S + 0.01 # equivalent to S.diag_shift+=0.01
F_dense, F_unravel = nk.jax.ravel(F)
x, info = scipy.sparse.cg(S, F_dense, maxiter=100)

or even

S = SMatrixOnTheFly(vstate)
S = S + 0.01 # equivalent to S.diag_shift+=0.01
Sm1 = np.linalg.pinv(S.to_dense)
x = [email protected]_dense

So all seems great! The only thing we need to think about is how to make all this play with the Driver API.

How to support this? we could accept two kwargs in the drivers:

 - S_type: Optional[SMatrixType] = The type of the S Matrix you want to use, that should support doing `S = S_type(vstate)` 
 - SR_solver: Optional[Callable] = The function to solve the linear system Sx=F. It must have signature SR_solver(S:SMatrix, F:PyTree, **kwargs) -> Tuple[x:PyTree, info:Any] 

If both are not declared (None) sr is not used. If one of the two is passed we use SR and the unspecified kwarg goes to a default. Internally we could do something like

def __init__(S_type=None, SR_solver=None):
   if S_type is None and Sr_solver is None:
     self.use_sr = False
   else:
     self.use_sr = True

   if use_sr = True and S_type is None:
     S_type = default
   ...


def _forward_and_backward(self):
    """
    Performs a number of VMC optimization steps.
    """

    self.state.reset()

    # Compute the local energy estimator and average Energy
    self._loss_stats, self._loss_grad = self.state.expect_and_grad(self._ham)

    if self.sr is not None:
        self._S = self.S_type(self.vstate)

        # use the previous solution as an initial guess to speed up the solution of the linear system
        x0 = self._dp if self.sr_restart is False else None
        self._dp, self._sr_info = self.SR_solver(self._S, self._loss_grad, x0=x0)

For the user, to specify kwargs of the solver like we do now, he would need to consult the docs of that solver and specify it with a functools.partial. Example:

from functools import partial

SR_solver = partial(jax.scipy.sparse.gmres, maxiter=300, restart=10)

# use default S_matrix type
gs = nk.VMC(ham, optim, variational_state=vstate, SR_solver= SR_solver) 
# or
gs = nk.VMC(ham, optim, variational_state=vstate, SR_solver= SR_solver, S_matrix=nk.optimizer.sr.SJacobian) 

However how to include the diagonal shift? One would have to do

def srsolver(S,F,**kwargs):
   return jax.scipy.sparse.gmres(S+0.01, F, **kwargs)

SR_solver = partial(srsolver, maxiter=300, restart=10)

which is not too clean...

This PR allows specifying both a scale-invariant and a constant diagonal offset in QGTJacobian* at the same time.

@chrisrothUT and I discovered that adding a diagonal shift of the form diag_shift + diag_scale * S_ii is much more stable than pure scale-invariant shifting but remains faster and better converging than simply using diag_shift. To accommodate both, this PR deprecates rescale_shift and introduces diag_scale, the coefficient of S_ii in the shift.

Defaults:

• If nothing is specified, diag_shift=0.01, diag_scale=0.0 to recover the original behaviour
• If only diag_scale is specified, diag_shift=0.0
• If rescale_shift is specified, it behaves as it used to, but there is a deprecation warning
• Specifying rescale_shift and diag_scale together leads to an error

Internally, diag_scale is implemented the same way rescale_shift once was; diag_shift is added to it by adding offset=diag_shift/diag_scale to the scale factors used to rescale rows/columns of the S matrix.

There are some bits to iron out, most importantly whether the defaults above are good and diag_scale is a good name for the new parameter.

Bumps sphinx from 4.5.0 to 6.1.1.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

Updates the requirements on nbsphinx to permit the latest version.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

Hi, first of all I'm sorry if this actually belongs in the discussions section, I've never worked with github before and don't know what problems are Issues-worthy.

I noticed that after the new implementation of the masked GCNNs, passing masks to the GCNN does not have any effect, e.g. for a model with only 2nd nearest neighbors convolutions the number of parameters does not change, so

# Define system
L = 8
lattice = nk.graph.Square(length=L)
hi = nk.hilbert.Spin(s=1 / 2, total_sz=0, N=lattice.n_nodes)

# Define Metropolis-Hastings sampler
sampler = nk.sampler.MetropolisExchange(hilbert=hi, graph=lattice)

# Define the model: masked GCNN
input_mask = np.zeros([L, L])
for i in range(-1, 2):
    for j in range(-1, 2):
        input_mask[i][j] = 1
input_mask = input_mask.ravel()
hidden_mask = np.repeat(np.expand_dims(input_mask, 1), repeats=8, axis=1).ravel()

machine = nk.models.GCNN(symmetries=lattice, layers=2, features=(4, 2), param_dtype=jnp.complex128,
                         input_mask=input_mask, hidden_mask=hidden_mask)
vstate = nk.vqs.MCState(sampler=sampler, model=machine)
print("number of parameters:", vstate.n_parameters)

returns number of parameters: 4358 instead of the desired number of parameters: 618

I think the error is that in netket.models.equivariant , the masks are not passed on from the general GCNN method to the constructors of the different modes (FFT, irreps).

Also, if I use the DenseSymm layer with a mask I get an error, so changing the model in the above code to machine = nk.nn.DenseSymm(symmetries=lattice, mode="matrix", mask=HashableArray(input_mask), features=2) gives the error TypeError: nonzero requires ndarray or scalar arguments, got <class 'netket.utils.array.HashableArray'> at position 0.

It should work if self.kernel_indices = jnp.nonzero(self.mask)[0] in the DenseSymmMatrix class is changed to (self.kernel_indices,) = np.nonzero(self.mask), as in the other modes.

Simple example of computing σ^2 for a single site.

>>> hi = nk.hilbert.Spin(s=0.5, total_sz = 0, N=4)
>>> heisenberg = np.array([[1, 0, 0, 0],[0, -1, 2, 0],[0, 2, -1, 0],[0, 0, 0, 1]])
>>> ha = nk.operator.LocalOperator(hilbert=hi,operators=[heisenberg],acting_on=[(0,0)])
>>> vstate.expect(ha)
1

The answer should be 3 of course but Netket returns 1 because it doesn't account for SxSx and SySy returning the state back to itself.

Updates the requirements on jax to permit the latest version.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

Bumps flake8 from 5.0.4 to 6.0.0.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.