Overriding defaults in NetKet

In this section we will discuss how it is possible to extend NetKet or override default behaviour. NetKet is architectured in such a way that an user can override several parts of its default behaviour without having to edit its source code. You can simply fire up a jupyter instance, define a few methods overloading the default, and you’re all set. This is also valid if you want to define a custom object, such as a custom hilbert space or operator.

NetKet Architecture: Multiple Dispatch

Some parts of NetKet rely on multiple-dispatch in order to select the right implementation of a function, instead of the more limited single dispatch used in traditional Object-Oriented programming. Significant inspiration has been taken from the Julia programming language, and NetKet relies on an somewhat-faithful implementation of Julia’s dispatch into python, provided by the plum-dispatch package.

If you are already familiar with multiple-dispatch, you can skip this section. If you want to see how to use multiple-dispatch in python, we refer you to the plum documentation. Otherwise, read on!

The rationale behind multiple dispatch is the following: If you have a function acting on two objects, such as the expect function computing the expectation value of an operator over a variational state, how do you make sure that expect works for any combination of different variational states and operators?

The standard object-oriented approach is to use single-dispatch to dispatch on the first argument (or owner) of the method. Therefore if expect is a method of VariationalState, you would be writing an expect method for every different VariationalState. This method should work with any operator you throw at it.

Maybe you are capable of defining a very broad interface for Operators so that your implementation works for all the operators already defined, but how can an user define a new CrazyOperator that does not follow this convention?

Multiple Dispatch makes it possible to define a method of a function that is only used when all the types match a given signature, and as long as the return type of the function makes sense, it will work with the rest of NetKet. If my expect function is defined as:

@netket.utils.dispatch.dispatch
def expect(vstate : VariationalState, operator: MyOperator, chunk_size: Any):
    # a generic algorithm that works for any operator in NetKet
    return netket.stats.statistics(result)

Dispatch makes a clear distinction between positional and keyword arguments. Be careful in following the example precisely.

And this function works for anything already implemented in Netket and always returns a statistics object (netket.stats.Stat). However, imagine that my CrazyOperator is defined on a Crazy Hilbert space that always gives random results. This is clearly incompatible with the standard implementation. However, instead of figuring a way to make such CrazyOperator fit in the what already exists, we can just define a custom dispatch rule as follows:

class CrazyOperator:
    """ A really crazy implementation. """
    sigma = 0.1

@expect.dispatch
def expect(vstate : MCState, operator: CrazyOperator, chunk_size: Any):
    # A crazy implementation that only works for CrazyOperator
    return netket.stats.statistics(np.random.rand(100)*operator.sigma)

And everything will work as expected. Of course, to make everything really work, it would be best to have CrazyOperator subclass AbstractOperator, but that is not needed if you are carefull enough.

The types that determine the dispatch are picked up by the type hints.