# Source code for netket.sampler.rules.hamiltonian_numpy

# Copyright 2021 The NetKet Authors - All rights reserved.
#
# 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
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and

import math

from numba import jit

import numpy as np

from netket.operator import DiscreteOperator
from netket.utils import struct

from .base import MetropolisRule

@struct.dataclass
class HamiltonianRuleState:
sections: np.ndarray
"""Preallocated array for sections"""

[docs]
class HamiltonianRuleNumpy(MetropolisRule):
r"""
Rule for Numpy sampler backend proposing moves according to the terms in an operator.

In this case, the transition matrix is taken to be:

.. math::

T( \mathbf{s} \rightarrow \mathbf{s}^\prime) =
\frac{1}{\mathcal{N}(\mathbf{s})}\theta(|H_{\mathbf{s},\mathbf{s}^\prime}|),

"""

operator: DiscreteOperator = struct.field(pytree_node=False)
"""The (hermitian) operator giving the transition amplitudes."""

[docs]
def __init__(self, operator: DiscreteOperator):
"""
Constructs the Hamiltonian sampling rule for the
:class:netket.sampler.MetropolisNumpy sampler.

If you are using the standard jax sampler, look for
:class:netket.sampler.rules.HamiltonianRule .

Args:
operator: The hermitian operator to be used to generate
configurations.
"""
if not isinstance(operator, DiscreteOperator):
raise TypeError(
"Argument to HamiltonianRule must be a valid operator, "
f"but operator is a {type(operator)}."
)
self.operator = operator

[docs]
def init_state(rule, sampler, machine, params, key):
if sampler.hilbert != rule.operator.hilbert:
raise ValueError(
f"""
The hilbert space of the sampler ({sampler.hilbert}) and the hilbert space
of the operator ({rule.operator.hilbert}) for HamiltonianRule must be the same.
"""
)

return HamiltonianRuleState(
sections=np.empty(sampler.n_batches, dtype=np.int32)
)

[docs]
def transition(rule, sampler, machine, parameters, state, rng, Ïƒ):
Ïƒ = state.Ïƒ
Ïƒ1 = state.Ïƒ1
log_prob_corr = state.log_prob_corr

sections = state.rule_state.sections
Ïƒp = rule.operator.get_conn_flattened(Ïƒ, sections)[0]

rand_vec = rng.uniform(0, 1, size=Ïƒ.shape[0])

_choose(Ïƒp, sections, Ïƒ1, log_prob_corr, rand_vec)
rule.operator.n_conn(Ïƒ1, sections)
log_prob_corr -= np.log(sections)

def __repr__(self):
return f"HamiltonianRuleNumpy({self.operator})"

@jit(nopython=True)
def _choose(states, sections, out, w, rand_vec):
low_range = 0
for i, s in enumerate(sections):
n_rand = low_range + int(np.floor(rand_vec[i] * (s - low_range)))
out[i] = states[n_rand]
w[i] = math.log(s - low_range)
low_range = s