nkDSL 0.1.2

Symbolic operator DSL for NetKet discrete JAX operators

nkDSL

nkDSL is a standalone symbolic operator DSL which provides a compact, declarative way to define discrete operators in NetKet and compile them into executable JAX-backed operator kernels.

[!NOTE] nkDSL is currently in beta. The core DSL and compiler workflow are ready for broader testing, while some advanced APIs may still evolve before a fully stable release.

What it does

With nkDSL, you describe an operator in terms of

• where to iterate, such as sites, pairs, or user-defined index sets,
• when a branch is active, through symbolic predicates,
• how to rewrite a configuration, through update programs,
• what matrix element to assign to each emitted branch.

The result is a symbolic operator that can be compiled into a NetKet-compatible executable operator.

Example: 1D transverse-field Ising model

The example below builds $H = -J \sum_{\langle i,j \rangle} \sigma_i^z \sigma_j^z - h \sum_i \sigma_i^x$ for a periodic chain. The diagonal Ising interaction is emitted with an identity update, while the transverse-field term flips one spin at a time.

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

L = 8

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

ising = (
    SymbolicDiscreteJaxOperator(hilbert, "ising_tfim", hermitian=True)
    .for_each(("i", "j"), over=graph.edges())
    .emit(
        matrix_element=-1.0 * site("i").value * site("j").value,
    )
    .for_each_site("i")
    .emit(
        write("i", -site("i").value),
        matrix_element=-0.5,
    )
    .compile()
)

# now, `ising` is a NetKet DiscreteJaxOperator subclass, and can be used in NetKet

Installation

Install directly from PyPI:

python -m pip install --upgrade pip
python -m pip install nkdsl

Status

nkdsl is currently a small, focused package for the recent DSL operator work developed in neuraLQX. Right now the goal is clarity and usability rather than a broad, frozen API surface.

License

This project is licensed under Apache License 2.0. See LICENSE for the full text.