Architecture & Philosophy

Architecture & Philosophy#

NetKet is designed to empower researchers to both explore new ideas in quantum many-body physics and run established Neural Quantum State (NQS) algorithms with ease. Understanding NetKet’s modular architecture is key to effectively using and extending the library.

Tiered Modular Architecture#

NetKet is organized into three distinct tiers, each serving different needs and providing different levels of abstraction:

NetKet Architecture Overview#
Tier	Purpose	Usage	Characteristics
Tier 0	Core JAX utilities	Algorithm implementation	JAX-native, reusable tools. Few changes over the years.
Tier 1	VMC building blocks	Everyone, including those who want to write their own algorithms.	Extensible, composable and well-defined interfaces.
Tier 2	MCState interface & optimisation drivers	General users, beginners.	Opinionated, easy-to-use abstractions.

Tier 0: Core JAX Tools#

Purpose: Low-level utilities that make working with JAX more efficient, particularly for distributed computing and memory management.

These tools are JAX-native and could be used outside of NetKet. They are quite intricate and have seen few changes over the years, but are essential if you want to write code that efficiently runs across many GPUs while automatically maintaining a low memoryy consumption.

In general, if you want to write algorithms using NetKet (or not even) operators, you will have to use those tools.

Tier 0 Modules#
Module	Description
`netket.jax`	JAX utilities: chunking, sharding, distributed operations. `vmap_chunked()`, `vjp_chunked()` and `jacobian()`.
`netket.utils.struct`	Custom PyTree implementations and data structures. Virtually all NetKet classes are `Pytree`s.
`netket.utils.timing`	Performance profiling and timing utilities
`netket.utils.numbers`	Numerical utilities and special number types

Tier 1: VMC Building Blocks#

Fundamental components needed to build Variational Monte Carlo algorithms with well-defined, extensible interfaces.

This tier provides the mathematical abstractions and is designed for flexibility. Components work with bare Flax modules and can be combined to create custom VMC engines.

If you want to work with exotic systems you might need to implement some custom objects from this tier, such as custom operators if you are working in Nuclear Physics. This is in general not a trivial task, but it’s well supported and if you do everything else from NetKet will work fine with your custom objects.

Tier 1 Modules#
Module	Description
`netket.hilbert`	Hilbert space definitions and computational basis enumeration
`netket.operator`	Quantum operators, Hamiltonians, and expectation value computation
`netket.sampler`	Monte Carlo samplers and Markov chain transition rules
`netket.models`	Neural network architectures for quantum states
`netket.nn`	Neural network layers and building blocks

Tier 2: High-Level Workflows#

Opinionated, user-friendly abstractions that hide complexity and provide a streamlined interface for common tasks.

This tier is designed for ease of use and provides “batteries included” functionality for standard NQS optimization workflows.

It might be hard to fit your original ideas and custom training algorithms within the boundaries of this interface, so for algorithm research we expect that you might have to re-implement the login in the drivers yourself.

Tier 2 Modules#
Module	Description
`netket.vqs`	Variational state abstractions (`MCState`, `MCMixedState`)
`netket.driver`	Optimization drivers (`VMC`, `SteadyState`)
`netket.optimizer`	Advanced optimizers with QGT and Stochastic Reconfiguration
`netket.logging`	Data logging and serialization utilities
`netket.callbacks`	Optimization callbacks for monitoring and control

Integrated Workflows: Complete optimization loops with logging and checkpointing
Advanced Optimizers: Built-in Quantum Geometric Tensor and Stochastic Reconfiguration
Automatic Management: Handles sampling, parameter updates, and state management
Rich Logging: Comprehensive data collection and analysis tools

How to approach NetKet:#

Beginners:: Use Tier-2 functionality from MCState, VMC, and pre-built models from netket.models. This provides a complete, working NQS simulation with minimal complexity.

Algorithm development:: If you want to invent new optimization protocols or algorithms, you should re-implement some of the Tier-2 functionality yourself. We suggest to start by copy-pasting a driver such as the VMC_SR and editing it directly, such that you will still benefit from the loggers and more..

Simulating systems not built-in:: If you want to work on composite systems, or nuclear physics models, or complicated continuous space setups, you will have to implement some custom operators (Tier-1). You will not need, in principle, to re-implement other parts of NetKet.

Editing NetKet internals directly#

NetKet has been designed to allow users to customize most things without the need to edit NetKet’s files directly. While editing the source code of the installed NetKet version directly might seem a simple and fast way to achieve what you want, it has the downside that you will not be able to (i) easily update to new NetKet versions and (ii) it is very hard to share your code with others.

Instead, if you want to change something, you should be able to copy the relevant files, and then import them from anywhere else.