Source code for netket.operator._pauli_strings.jax

# Copyright 2023 The NetKet Authors - All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import partial, wraps
from typing import Optional, Union, TYPE_CHECKING

import numpy as np

import jax
from jax import numpy as jnp
from jax.tree_util import register_pytree_node_class

from netket.hilbert import AbstractHilbert, HomogeneousHilbert
from netket.errors import concrete_or_error, JaxOperatorSetupDuringTracingError
from netket.utils.types import DType
from netket.utils import HashableArray

from .._discrete_operator_jax import DiscreteJaxOperator

from .base import PauliStringsBase

    from .numba import PauliStrings

# pauli-strings operator written in nJax
# the general idea is the following:

# observe that Y = -i Z X
# therefore for every pauli in a given operator we
# (1.) apply X  (flip site)                  -- if it is X or Y
# (2.) absorb -i to the weights              -- if it is Y
# (3.) apply Z (pick up -1. if site is up/1) -- if it is Y or Z

# TODO special case for the diagonal
# TODO eventually also implement _ising_conn_states_jax with indexing instead of mask
# TODO eventually add version with sparse jax arrays (achieving the same as indexing)

# duplicated from ising
def _ising_conn_states_jax(x, cond, local_states):
    was_state_0 = x == local_states[0]
    state_0 = jnp.asarray(local_states[0], dtype=x.dtype)
    state_1 = jnp.asarray(local_states[1], dtype=x.dtype)
    return jnp.where(cond ^ was_state_0, state_0, state_1)

def pack_internals(operators, weights, cutoff=0):
    # here we group together operators with same final state
    # The final state is determined by the sites we flip,
    # we store this in the keys of `acting`
    # Each value of `acting` contains a list, with an entry for of the opeartors
    # which does the same flips (given by the key). Each entry contains
    # 1. a weight
    # 2. list of sites we apply the Z on (those are the sites we
    #    will need to check in the operator to determine if the sign is flipped or not)

    acting = {}

    def find_char(s, ch):
        return [i for i, ltr in enumerate(s) if ltr == ch]

    def append(key, k):
        # convert list to tuple
        key = tuple(sorted(key))  # order of X and Y does not matter
        if key in acting:
            acting[key] = [k]

    for i, op in enumerate(operators):
        b_to_change = []  # list of all the sites we will need to act on with X
        b_z_check = []  # list of all the sites we act on with Z
        b_weight = weights[i]

        if abs(b_weight) <= cutoff:

        x_ops = find_char(op, "X")  # find sites we act on with X w/ op
        if len(x_ops):
            b_to_change += x_ops

        y_ops = find_char(op, "Y")  # find sites we act on with Y w/ op
        if len(y_ops):
            b_to_change += y_ops
            b_weight *= (-1.0j) ** (len(y_ops))  # absorb the -i into weights
            b_z_check += y_ops

        z_ops = find_char(op, "Z")  # find sites we act on with Z w/ op
        if len(z_ops):
            b_z_check += z_ops

        # by appending we concat all (b_weights, b_z_check) which have the same b_to_change
        # i.e. the one which we act on with X on the same sites (also X coming from Y obviously)

        # sort b_z_check in ascending order, for better locality
        b_z_check = list(sorted(b_z_check))

        # If there is an even number of Y in a string, the weight should be real
        if np.isreal(b_weight):
            b_weight = b_weight.real

        append(b_to_change, (b_weight, b_z_check))
    return acting

def sites_to_mask(sites, n_sites, dtype=bool):
    mask = np.zeros(n_sites, dtype=dtype)
    mask[(sites,)] = 1
    return mask

_available_modes = ["index", "mask"]

def _check_mode(mode):
    if mode not in _available_modes:
        raise ValueError(
            f"unknown mode {mode}. Available modes are {_available_modes}."

def pack_internals_jax(
    Take the internal lazy representation of a paulistrings operator and returns the
    arrays needed for the jax implementation.

    This takes as input a numpy array of strings (`operators`) of equal lengths
    containing only the characters 'X', 'Y', 'Z' or 'I',  and an equal length
    vector of real or complex coefficients (`weights`).

    `weights` must be a numpy-array with statically known values otherwise an error
    is thrown.

    Returns a dictionary with all the data fields
    # index_dtype needs to be signed (we use -1 for padding)


    # group together operators with same final state (i.e. those which flip the same sites)
    # see code of pack_internals for more details
    acting = pack_internals(operators, weights)

    n_sites = len(operators[0])

    # now group those ones which have the same number of operators for a given final state
    acting_by_num_ops = {}
    for k, v in acting.items():
        num = len(v)
        acting_by_num_ops[num] = acting_by_num_ops.get(num, []) + [(k, v)]

    x_flip_masks = {}
    weights = {}
    z_sign_masks = {}
    z_sign_indices = {}
    z_sign_indices_masks = {}

    for l, ops in acting_by_num_ops.items():
        x_flip_masks_l = []
        weights_l = []
        z_sign_masks_l = []  # if we decide to use masks
        z_sign_indices_l = []  # if we decide to use indexing

        for sites_to_flip, rest in ops:
            w = [r[0] for r in rest]
            sites_for_sgn = [r[1] for r in rest]
            x_flip_mask = sites_to_mask(sites_to_flip, n_sites)
            z_sign_mask = list(
                map(partial(sites_to_mask, n_sites=n_sites), sites_for_sgn)


        # turn into arrays
        x_flip_masks[l] = jnp.array(x_flip_masks_l, dtype=mask_dtype)
        if weight_dtype is not None:
            weights[l] = jnp.array(weights_l, dtype=weight_dtype)
            weights[l] = jnp.array(weights_l)

        z_sign_masks[l] = jnp.array(z_sign_masks_l, dtype=mask_dtype)

        # prepare index arrays if we are indexing
        num_z = [len(y) for x in z_sign_indices_l for y in x]
        maxlen = max(num_z, default=0)
        pad = maxlen != min(num_z, default=0)
        if pad:
            # arbitrarily fill with -1
            tmp1 = np.full((len(z_sign_indices_l), l, maxlen), -1, dtype=index_dtype)
            tmp2 = np.full((len(z_sign_indices_l), l, maxlen), False, dtype=mask_dtype)
            for i, inds in enumerate(z_sign_indices_l):
                for j, ind in enumerate(inds):
                    tmp1[i, j, : len(ind)] = ind
                    tmp2[i, j, : len(ind)] = True
            z_sign_indices[l] = jnp.array(tmp1)
            z_sign_indices_masks[l] = jnp.array(tmp2)
        else:  # no padding needed, all have the same length
            z_sign_indices[l] = jnp.array(z_sign_indices_l, dtype=index_dtype)
            z_sign_indices_masks[l] = None

    # transform the arrays into lists so that we have consistent ordering
    # as we will concatenate the results in the operator
    keys = sorted(x_flip_masks.keys())
    x_flip_masks = [x_flip_masks[k] for k in keys]
    weights = [weights[k] for k in keys]

    # TODO here would be the place we could decide wether to use index or mask
    # depending on how much we padded, use a hybrid scheme etc
    z_sign_masks = [z_sign_masks[k] if (mode == "mask") else None for k in keys]
    z_sign_indices = [z_sign_indices[k] if (mode == "index") else None for k in keys]
    z_sign_indices_masks = [
        z_sign_indices_masks[k] if (mode == "index") else None for k in keys

    x_flip_masks_stacked = jnp.concatenate(x_flip_masks, axis=0)
    z_data = (weights, z_sign_masks, z_sign_indices, z_sign_indices_masks)
    return x_flip_masks_stacked, z_data

def _pauli_strings_mels_jax(local_states, z_data, x):
    # supports both masks and indexing (can be padded, so also with a mask but smaller)
    # which path is taken is flexible, and can be fully determined by z_data
    # z_data: a list of tuples weights, z_sign_mask, z_sign_indices, z_sign_indexmask
    state1 = local_states[1]
    was_state_1 = x == state1
    mels = []
    for w, z_sign_mask, z_sign_indices, z_sign_indexmask in zip(*z_data):
        if z_sign_mask is not None:  # use masks
            was_state = was_state_1[..., None, None, :]
            mask = z_sign_mask
        else:  # use indexing
            assert z_sign_indices is not None
            was_state = x[..., z_sign_indices] == state1
            if z_sign_indexmask is None:  # no padding, so no mask necessary
                mask = 1
                mask = z_sign_indexmask
        # sgn = (-1)**(was_state*mask).sum(axis=-1)
        sgn = (1 - 2 * (was_state * mask).astype(np.int8)).prod(
            axis=-1, promote_integers=False
        mels.append(jnp.einsum("...ab,ab->...a", sgn, w))
    return jnp.concatenate(mels, axis=-1)

def _pauli_strings_kernel_jax(local_states, x_flip_masks_all, z_data, x, cutoff=None):
    mels = _pauli_strings_mels_jax(local_states, z_data, x)
    if cutoff is not None:
        nonzero_mels_mask = jnp.abs(mels) > cutoff
        mels =, mels, jnp.zeros_like(mels))
        # only flip if corresponding mel is nonzero
        x_flip_masks_all = x_flip_masks_all & nonzero_mels_mask[..., None]
        nonzero_mels_mask = None
    # can re-use function from Ising
    x_prime = _ising_conn_states_jax(x[..., None, :], x_flip_masks_all, local_states)
    # TODO we could optionally move nonzeros to front here
    return x_prime, mels, nonzero_mels_mask

def _pauli_strings_n_conn_jax(local_states, x_flip_masks_all, z_data, x, cutoff=None):
    _, _, nonzero_mels_mask = _pauli_strings_kernel_jax(
        local_states, x_flip_masks_all, z_data, x, cutoff
    if nonzero_mels_mask is not None:
        return nonzero_mels_mask.sum(axis=-1, dtype=np.int32)
        max_conn_size = x_flip_masks_all.shape[0]
        return jnp.full(x.shape[:-1], max_conn_size, dtype=np.int32)

[docs] @register_pytree_node_class class PauliStringsJax(PauliStringsBase, DiscreteJaxOperator): """ Jax-compatible version of :class:`netket.operator.PauliStrings`. """ @wraps(PauliStringsBase.__init__) def __init__( self, hilbert: AbstractHilbert, operators: Union[None, str, list[str]] = None, weights: Union[None, float, complex, list[Union[float, complex]]] = None, *, cutoff: float = 1.0e-10, dtype: Optional[DType] = None, _mode: str = "index", ): super().__init__(hilbert, operators, weights, cutoff=cutoff, dtype=dtype) if len(self.hilbert.local_states) != 2: raise ValueError( "PauliStringsJax only supports Hamiltonians with two local states" ) # check that it is homogeneous, throw error if it's not if not isinstance(self.hilbert, HomogeneousHilbert): local_states = self.hilbert.states_at_index(0) if not all( np.allclose(local_states, self.hilbert.states_at_index(i)) for i in range(self.hilbert.size) ): raise ValueError( "Hilbert spaces with non homogeneous local_states are not " "yet supported by PauliStrings." ) # private variable for setting the mode # currently there are two modes: # index: indexes into the vector to flip qubits and compute the sign # faster if the strings act only on a few qubits # mask: uses masks to flip qubits and compute the sign # faster if the strings act on many of the qubits # (and possibly on gpu) # By adapting pack_internals_jax hybrid approaches are also possible. # depending on performance tests we might expose or remove it self._hi_local_states = tuple(self.hilbert.local_states) self._initialized = False self._mode = _mode @property def _mode(self): """ (Internal) Indexing mode of the operator. Valid values are "index" or "mask". 'Index' uses the standard LocalOperator-like indexing of changed points, while the latter uses constant-size masks. The latter does not really need recompilation for paulistrings with different values, and this could be changed in the future. """ return self._mode_attr @_mode.setter def _mode(self, mode): _check_mode(mode) self._mode_attr = mode self._reset_caches() @property def max_conn_size(self) -> int: """The maximum number of non zero ⟨x|O|x'⟩ for every x.""" self._setup() return self._x_flip_masks_stacked.shape[0] def _setup(self, force=False): if force or not self._initialized: weights = concrete_or_error( np.asarray, self.weights, JaxOperatorSetupDuringTracingError, self, ) # Necessary for the tree_flatten in jax.jit, because # metadata must be hashable and comparable. We don't # want to re-hash it at every unpacking so we do it # once in here. if self._mode == "index": self._operators_hashable = HashableArray(self.operators) else: self._operators_hashable = None x_flip_masks_stacked, z_data = pack_internals_jax( self.operators, weights, weight_dtype=self.dtype, mode=self._mode ) self._x_flip_masks_stacked = x_flip_masks_stacked self._z_data = z_data self._initialized = True def _reset_caches(self): super()._reset_caches() self._initialized = False
[docs] def n_conn(self, x): self._setup() return _pauli_strings_n_conn_jax( self._hi_local_states, self._x_flip_masks_stacked, self._z_data, x, cutoff=self._cutoff, )
[docs] def get_conn_padded(self, x): self._setup() xp, mels, _ = _pauli_strings_kernel_jax( self._hi_local_states, self._x_flip_masks_stacked, self._z_data, x, cutoff=self._cutoff, ) return xp, mels
[docs] def tree_flatten(self): self._setup() data = (self.weights, self._x_flip_masks_stacked, self._z_data) metadata = { "hilbert": self.hilbert, "operators": self._operators_hashable, "dtype": self.dtype, "mode": self._mode, } return data, metadata
[docs] @classmethod def tree_unflatten(cls, metadata, data): (weights, xm, zd) = data hi = metadata["hilbert"] operators_hashable = metadata["operators"] dtype = metadata["dtype"] mode = metadata["mode"] op = cls(hi, dtype=dtype, _mode=mode) op._operators = ( operators_hashable.wrapped if operators_hashable is not None else None ) op._operators_hashable = operators_hashable op._weights = weights op._x_flip_masks_stacked = xm op._z_data = zd op._initialized = True return op
[docs] def to_numba_operator(self) -> "PauliStrings": # noqa: F821 """ Returns the standard numba version of this operator, which is an instance of :class:`netket.operator.PauliStrings`. """ from .numba import PauliStrings return PauliStrings( self.hilbert, self.operators, self.weights, dtype=self.dtype, cutoff=self._cutoff, )