Symbolic discrete operators for NetKet

[waleed-sh]

Hi everyone,

As a show and tell, I wanted to share a small project I've been working on which shows the power and flexibility of NetKet's modular API and also might be useful to the broader NetKet community. It's a way to build NetKet discrete operators using [d]eclarative [s]ymbolic [l]anguage (DSL), and it’s called nkDSL.

Basically you describe how the operator acts and then it gets "compiled" to a DiscreteJaxOperator compatible with NetKet. This is really possible thanks to NetKet’s API, which makes integrations like this easy to do. For example it makes use of NetKet’s machinery to add and multiply discrete operators to make things more friendly, so doing A @ B and A + B works thanks to NetKet’s ProductOperator and SumOperator.

For example, an Ising Hamiltonian can be written as:

import netket as nk
from nkdsl import SymbolicDiscreteJaxOperator, site, write

L = 6
J = 1.2
h = 0.7

hi = nk.hilbert.Spin(s=0.5, N=L)
g = nk.graph.Chain(length=L, pbc=True)

H = (
    SymbolicDiscreteJaxOperator(hi, “Ising”, hermitian=True)
    .for_each(("i", "j"), over=g.edges())  # iterator over nearest-neighbor bonds (i, j)
    .emit(
        matrix_element=J * site("i").value * site("j").value,  # diagonal contribution: J * s_i * s_j
    )
    .for_each_site("i")  # iterator over all sites i
    .emit(
        write("i", -site("i").value),  # update rule: flip spin at site i (non-diagonal)
        matrix_element=-h,  # off-diagonal matrix element
    )
    .compile()  # "compile" to a usable DiscreteJaxOperator
)

And now H is a subclass of DiscreteJaxOperator with a working get_conn_padded function and so it can be used in NetKet's VMC and co.

A bosonic number operator for example would simply be

import netket as nk
from nkdsl import SymbolicDiscreteJaxOperator, AmplitudeExpr

hi = nk.hilbert.Fock(n_max=4, N=6)

k = 1  # act on site 1

n_k = (
    SymbolicDiscreteJaxOperator(hi, f"n_{k}", hermitian=True)
    .globally()  # act on the entire basis element x, not a specific basis label
    .emit(matrix_element=AmplitudeExpr.static_index(k))  # matrix element is x[k]
    .compile()
)

Right now, it works for the DiscreteJaxOperator, and therefore any Hilbert space compatible with it.

Not all operators should be constructed using this DSL... It cannot handle really complicated operators yet, but I tried to make the ability to add to the DSL easy so it can be expanded if needed. It’s intended for operators which are too difficult/memory consuming to write as local operators, and would be too much work to write as a discrete operator by hand using JAX and have to worry about batch dimensions, indexing, etc.

This started as a personal project and is now public since I thought someone might benefit from it. It is however an alpha release, I am sure not everything is thought through still, and it has to be really tested out, but I do verify results against native NetKet implementations and exact diagonalisation as a correctness baseline.

If this sounds interesting, please have a look! I would really appreciate any feedback, suggestions and help!

The documentation is here, the repo is here, and here are some tutorials.