Source code for netket.sampler.metropolis

# Copyright 2021 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
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import partial
from typing import Any, Callable, Optional, Union
from textwrap import dedent

import numpy as np

import jax
from flax import linen as nn
from flax import serialization
from jax import numpy as jnp

from netket.hilbert import AbstractHilbert, ContinuousHilbert

from netket.utils import mpi, wrap_afun
from netket.utils.types import PyTree, DType

from netket.utils.deprecation import deprecated, warn_deprecation
from netket.utils import struct

from netket.utils.config_flags import config
from netket.jax.sharding import (
    extract_replicated,
    gather,
    distribute_to_devices_along_axis,
    device_count,
    with_samples_sharding_constraint,
)

from .base import Sampler, SamplerState
from .rules import MetropolisRule


[docs] class MetropolisSamplerState(SamplerState): """ State for a Metropolis sampler. Contains the current configuration, the RNG state and the (optional) state of the transition rule. """ σ: jnp.ndarray """Current batch of configurations in the Markov chain.""" rng: jnp.ndarray """State of the random number generator (key, in jax terms).""" rule_state: Optional[Any] """Optional state of the transition rule.""" n_steps_proc: int = struct.field(default_factory=lambda: jnp.zeros((), dtype=int)) """Number of moves performed along the chains in this process since the last reset.""" n_accepted_proc: jnp.ndarray """Number of accepted transitions among the chains in this process since the last reset.""" def __init__(self, σ: jnp.ndarray, rng: jnp.ndarray, rule_state: Optional[Any]): self.σ = σ self.rng = rng self.rule_state = rule_state self.n_accepted_proc = jnp.zeros(σ.shape[0], dtype=int) self.n_steps_proc = jnp.zeros((), dtype=int) super().__init__() @property def acceptance(self) -> float: """The fraction of accepted moves across all chains and MPI processes. The rate is computed since the last reset of the sampler. Will return None if no sampling has been performed since then. """ if self.n_steps == 0: return None return self.n_accepted / self.n_steps @property def n_steps(self) -> int: """Total number of moves performed across all processes since the last reset.""" return self.n_steps_proc * mpi.n_nodes @property def n_accepted(self) -> int: """Total number of moves accepted across all processes since the last reset.""" # jit sum for gda res, _ = mpi.mpi_sum_jax(jax.jit(jnp.sum)(self.n_accepted_proc)) return res def __repr__(self): if self.n_steps > 0: acc_string = "# accepted = {}/{} ({}%), ".format( self.n_accepted, self.n_steps, self.acceptance * 100 ) else: acc_string = "" return f"{type(self).__name__}({acc_string}rng state={self.rng})"
# serialization when sharded def serialize_MetropolisSamplerState_sharding(sampler_state): state_dict = MetropolisSamplerState._to_flax_state_dict( MetropolisSamplerState._pytree__static_fields, sampler_state ) for prop in ["σ", "n_accepted_proc"]: x = state_dict.get(prop, None) if x is not None and isinstance(x, jax.Array) and len(x.devices()) > 1: state_dict[prop] = gather(x) state_dict = extract_replicated(state_dict) return state_dict def deserialize_MetropolisSamplerState_sharding(sampler_state, state_dict): for prop in ["σ", "n_accepted_proc"]: x = state_dict[prop] if x is not None: state_dict[prop] = distribute_to_devices_along_axis(x) return MetropolisSamplerState._from_flax_state_dict( MetropolisSamplerState._pytree__static_fields, sampler_state, state_dict ) if config.netket_experimental_sharding: # when running on multiple jax processes the σ and n_accepted_proc are not fully addressable # however, when serializing they need to be so here we register custom handlers which # gather all the data to every process. # when deserializing we distribute the samples again to all availale devices # this way it is enough to serialize on process 0, and we can restart the simulation # also on a different number of devices, provided the number of samples is still # divisible by the new number of devices serialization.register_serialization_state( MetropolisSamplerState, serialize_MetropolisSamplerState_sharding, deserialize_MetropolisSamplerState_sharding, override=True, ) def _assert_good_sample_shape(samples, shape, dtype, obj=""): canonical_dtype = jax.dtypes.canonicalize_dtype(dtype) if samples.shape != shape or samples.dtype != canonical_dtype: raise ValueError( dedent( f""" The samples returned by the {obj} have `shape={samples.shape}` and `dtype={samples.dtype}`, but the sampler requires `shape={shape} and `dtype={canonical_dtype}` (canonicalized from {dtype}). If you are using a custom transition rule, check that it returns the correct shape and dtype. If you are using a built-in transition rule, there might be a mismatch between hilbert spaces, or it's a bug in NetKet. """ ) ) def _assert_good_log_prob_shape(log_prob, n_chains_per_rank, machine): if log_prob.shape != (n_chains_per_rank,): raise ValueError( dedent( f""" The output of the model {machine} has `shape={log_prob.shape}`, but `shape=({n_chains_per_rank},)` was expected. This might be because of an hilbert space mismatch or because your model is ill-configured. """ ) ) def _round_n_chains_to_next_multiple( n_chains, n_chains_per_whatever, n_whatever, whatever_str ): # small helper function to round the number of chains to the next multiple of [whatever] # here [whatever] can be e.g. mpi ranks or jax devices # if n_chains is None and n_chains_per_whatever is None: # n_chains_per_whatever = default if n_chains is not None and n_chains_per_whatever is not None: raise ValueError( f"Cannot specify both `n_chains` and `n_chains_per_{whatever_str}`" ) elif n_chains is not None: n_chains_per_whatever = max(int(np.ceil(n_chains / n_whatever)), 1) if n_chains_per_whatever * n_whatever != n_chains: if mpi.rank == 0: import warnings warnings.warn( f"Using {n_chains_per_whatever} chains per {whatever_str} among {n_whatever} {whatever_str}s " f"(total={n_chains_per_whatever * n_whatever} instead of n_chains={n_chains}). " f"To directly control the number of chains on every {whatever_str}, specify " f"`n_chains_per_{whatever_str}` when constructing the sampler. " f"To silence this warning, either use `n_chains_per_{whatever_str}` or use `n_chains` " f"that is a multiple of the number of {whatever_str}s", category=UserWarning, stacklevel=2, ) return n_chains_per_whatever * n_whatever
[docs] class MetropolisSampler(Sampler): r""" Metropolis-Hastings sampler for a Hilbert space according to a specific transition rule. The transition rule is used to generate a proposed state :math:`s^\prime`, starting from the current state :math:`s`. The move is accepted with probability .. math:: A(s \rightarrow s^\prime) = \mathrm{min} \left( 1,\frac{P(s^\prime)}{P(s)} e^{L(s,s^\prime)} \right) , where the probability being sampled from is :math:`P(s)=|M(s)|^p`. Here :math:`M(s)` is a user-provided function (the machine), :math:`p` is also user-provided with default value :math:`p=2`, and :math:`L(s,s^\prime)` is a suitable correcting factor computed by the transition kernel. The dtype of the sampled states can be chosen. """ rule: MetropolisRule = None """The Metropolis transition rule.""" sweep_size: int = struct.field(pytree_node=False, default=None) """Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space.""" n_chains: int = struct.field(pytree_node=False) """Total number of independent chains across all MPI ranks and/or devices.""" reset_chains: bool = struct.field(pytree_node=False, default=False) """If True, resets the chain state when `reset` is called on every new sampling."""
[docs] def __init__( self, hilbert: AbstractHilbert, rule: MetropolisRule, *, n_sweeps: int = None, sweep_size: int = None, reset_chains: bool = False, n_chains: Optional[int] = None, n_chains_per_rank: Optional[int] = None, machine_pow: int = 2, dtype: DType = float, ): """ Constructs a Metropolis Sampler. Args: hilbert: The Hilbert space to sample. rule: A `MetropolisRule` to generate random transitions from a given state as well as uniform random states. n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. If MPI is disabled, the two are equivalent; if MPI is enabled and `n_chains` is specified, then every MPI rank will run `n_chains/mpi.n_nodes` chains. In general, we recommend specifying `n_chains_per_rank` as it is more portable. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). If netket_experimental_sharding is enabled this is interpreted as the number of independent chains on every jax device, and the n_chains_per_rank property of the sampler will return the total number of chains on all devices. sweep_size: Number of sweeps for each step along the chain. This is equivalent to subsampling the Markov chain. (Defaults to the number of sites in the Hilbert space.) reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). """ # Validate the inputs if not isinstance(rule, MetropolisRule): raise TypeError( f"The second positional argument, rule, must be a MetropolisRule but " f"`type(rule)={type(rule)}`." ) if not isinstance(reset_chains, bool): raise TypeError("reset_chains must be a boolean.") if n_sweeps is not None: warn_deprecation( "Specifying `n_sweeps` when constructing sampler is deprecated. Please use `sweep_size` instead." ) if sweep_size is not None: raise ValueError("Cannot specify both `sweep_size` and `n_sweeps`") sweep_size = n_sweeps if sweep_size is None: sweep_size = hilbert.size # Default n_chains per rank, if unset if n_chains is None and n_chains_per_rank is None: # TODO set it to a few hundred if on GPU? n_chains_per_rank = 16 n_chains = _round_n_chains_to_next_multiple( n_chains, n_chains_per_rank, device_count(), "rank", ) super().__init__( hilbert=hilbert, machine_pow=machine_pow, dtype=dtype, ) self.n_chains = n_chains self.reset_chains = reset_chains self.rule = rule self.sweep_size = sweep_size
@property def n_sweeps(self): warn_deprecation( "`MetropolisSampler.n_sweeps` is deprecated. Please use `MetropolisSampler.sweep_size` instead." ) return self.sweep_size
[docs] def sample_next( sampler, machine: Union[Callable, nn.Module], parameters: PyTree, state: Optional[SamplerState] = None, ) -> tuple[SamplerState, jnp.ndarray]: """ Samples the next state in the Markov chain. Args: machine: A Flax module or callable with the forward pass of the log-pdf. If it is a callable, it should have the signature :code:`f(parameters, σ) -> jnp.ndarray`. parameters: The PyTree of parameters of the model. state: The current state of the sampler. If not specified, then initialize and reset it. Returns: state: The new state of the sampler. σ: The next batch of samples. Note: The return order is inverted wrt `sample` because when called inside of a scan function the first returned argument should be the state. """ if state is None: state = sampler.reset(machine, parameters) return sampler._sample_next(wrap_afun(machine), parameters, state)
@partial(jax.jit, static_argnums=1) def _init_state(sampler, machine, parameters, key): key_state, key_rule = jax.random.split(key) rule_state = sampler.rule.init_state(sampler, machine, parameters, key_rule) σ = jnp.zeros((sampler.n_batches, sampler.hilbert.size), dtype=sampler.dtype) σ = with_samples_sharding_constraint(σ) state = MetropolisSamplerState(σ=σ, rng=key_state, rule_state=rule_state) # If we don't reset the chain at every sampling iteration, then reset it # now. if not sampler.reset_chains: key_state, rng = jax.jit(jax.random.split)(key_state) σ = sampler.rule.random_state(sampler, machine, parameters, state, rng) _assert_good_sample_shape( σ, (sampler.n_batches, sampler.hilbert.size), sampler.dtype, f"{sampler.rule}.random_state", ) σ = with_samples_sharding_constraint(σ) state = state.replace(σ=σ, rng=key_state) return state @partial(jax.jit, static_argnums=1) def _reset(sampler, machine, parameters, state): rng = state.rng if sampler.reset_chains: rng, key = jax.random.split(state.rng) σ = sampler.rule.random_state(sampler, machine, parameters, state, rng) _assert_good_sample_shape( σ, (sampler.n_batches, sampler.hilbert.size), sampler.dtype, f"{sampler.rule}.random_state", ) σ = with_samples_sharding_constraint(σ) else: σ = state.σ rule_state = sampler.rule.reset(sampler, machine, parameters, state) return state.replace( σ=σ, rng=rng, rule_state=rule_state, n_steps_proc=jnp.zeros_like(state.n_steps_proc), n_accepted_proc=jnp.zeros_like(state.n_accepted_proc), ) def _sample_next(sampler, machine, parameters, state): """ Implementation of `sample_next` for subclasses of `MetropolisSampler`. If you subclass `MetropolisSampler`, you should override this and not `sample_next` itself, because `sample_next` contains some common logic. """ def loop_body(i, s): # 1 to propagate for next iteration, 1 for uniform rng and n_chains for transition kernel s["key"], key1, key2 = jax.random.split(s["key"], 3) σp, log_prob_correction = sampler.rule.transition( sampler, machine, parameters, state, key1, s["σ"] ) _assert_good_sample_shape( σp, (sampler.n_batches, sampler.hilbert.size), sampler.dtype, f"{sampler.rule}.transition", ) proposal_log_prob = sampler.machine_pow * machine.apply(parameters, σp).real _assert_good_log_prob_shape(proposal_log_prob, sampler.n_batches, machine) uniform = jax.random.uniform(key2, shape=(sampler.n_batches,)) if log_prob_correction is not None: do_accept = uniform < jnp.exp( proposal_log_prob - s["log_prob"] + log_prob_correction ) else: do_accept = uniform < jnp.exp(proposal_log_prob - s["log_prob"]) # do_accept must match ndim of proposal and state (which is 2) s["σ"] = jnp.where(do_accept.reshape(-1, 1), σp, s["σ"]) s["accepted"] += do_accept s["log_prob"] = jax.numpy.where( do_accept.reshape(-1), proposal_log_prob, s["log_prob"] ) return s new_rng, rng = jax.random.split(state.rng) s = { "key": rng, "σ": state.σ, "log_prob": sampler.machine_pow * machine.apply(parameters, state.σ).real, # for logging "accepted": state.n_accepted_proc, } s = jax.lax.fori_loop(0, sampler.sweep_size, loop_body, s) new_state = state.replace( rng=new_rng, σ=s["σ"], n_accepted_proc=s["accepted"], n_steps_proc=state.n_steps_proc + sampler.sweep_size * sampler.n_batches, ) return new_state, new_state.σ @partial(jax.jit, static_argnums=(1, 4)) def _sample_chain(sampler, machine, parameters, state, chain_length): """ Samples `chain_length` batches of samples along the chains. Internal method used for jitting calls. Arguments: sampler: The Monte Carlo sampler. machine: A Flax module with the forward pass of the log-pdf. parameters: The PyTree of parameters of the model. state: The current state of the sampler. chain_length: The length of the chains. Returns: σ: The next batch of samples. state: The new state of the sampler """ state, samples = jax.lax.scan( lambda state, _: sampler.sample_next(machine, parameters, state), state, xs=None, length=chain_length, ) # make it (n_chains, n_samples_per_chain) as expected by netket.stats.statistics samples = jnp.swapaxes(samples, 0, 1) return samples, state def __repr__(sampler): return ( f"{type(sampler).__name__}(" + f"\n hilbert = {sampler.hilbert}," + f"\n rule = {sampler.rule}," + f"\n n_chains = {sampler.n_chains}," + f"\n sweep_size = {sampler.sweep_size}," + f"\n reset_chains = {sampler.reset_chains}," + f"\n machine_power = {sampler.machine_pow}," + f"\n dtype = {sampler.dtype}" + ")" ) def __str__(sampler): return ( f"{type(sampler).__name__}(" + f"rule = {sampler.rule}, " + f"n_chains = {sampler.n_chains}, " + f"sweep_size = {sampler.sweep_size}, " + f"reset_chains = {sampler.reset_chains}, " + f"machine_power = {sampler.machine_pow}, " + f"dtype = {sampler.dtype})" )
@deprecated( "The module function `sample_next` is deprecated in favor of the class method `sample_next`." ) def sample_next( sampler: MetropolisSampler, machine: Union[Callable, nn.Module], parameters: PyTree, state: Optional[SamplerState] = None, ) -> tuple[SamplerState, jnp.ndarray]: """ Samples the next state in the Markov chain. Args: sampler: The Metropolis sampler. machine: A Flax module or callable with the forward pass of the log-pdf. If it is a callable, it should have the signature :code:`f(parameters, σ) -> jnp.ndarray`. parameters: The PyTree of parameters of the model. state: The current state of the sampler. If not specified, then initialize and reset it. Returns: state: The new state of the sampler. σ: The next batch of samples. """ return sampler.sample_next(machine, parameters, state)
[docs] def MetropolisLocal(hilbert, **kwargs) -> MetropolisSampler: r""" Sampler acting on one local degree of freedom. This sampler acts locally only on one local degree of freedom :math:`s_i`, and proposes a new state: :math:`s_1 \dots s^\prime_i \dots s_N`, where :math:`s^\prime_i \neq s_i`. The transition probability associated to this sampler can be decomposed into two steps: 1. One of the site indices :math:`i = 1\dots N` is chosen with uniform probability. 2. Among all the possible (:math:`m - 1`) values that :math:`s^\prime_i` can take, one of them is chosen with uniform probability. For example, in the case of spin :math:`1/2` particles, :math:`m=2` and the possible local values are :math:`s_i = -1,+1`. In this case then :class:`MetropolisLocal` is equivalent to flipping a random spin. In the case of bosons, with occupation numbers :math:`s_i = 0, 1, \dots n_{\mathrm{max}}`, :class:`MetropolisLocal` would pick a random local occupation number uniformly between :math:`0` and :math:`n_{\mathrm{max}}` except the current :math:`s_i`. Args: hilbert: The Hilbert space to sample. n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). sweep_size: Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space. This is equivalent to subsampling the Markov chain. reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). """ from .rules import LocalRule return MetropolisSampler(hilbert, LocalRule(), **kwargs)
[docs] def MetropolisExchange( hilbert, *, clusters=None, graph=None, d_max=1, **kwargs ) -> MetropolisSampler: r""" This sampler acts locally only on two local degree of freedom :math:`s_i` and :math:`s_j`, and proposes a new state: :math:`s_1 \dots s^\prime_i \dots s^\prime_j \dots s_N`, where in general :math:`s^\prime_i \neq s_i` and :math:`s^\prime_j \neq s_j`. The sites :math:`i` and :math:`j` are also chosen to be within a maximum graph distance of :math:`d_{\mathrm{max}}`. The transition probability associated to this sampler can be decomposed into two steps: 1. A pair of indices :math:`i,j = 1\dots N`, and such that :math:`\mathrm{dist}(i,j) \leq d_{\mathrm{max}}`, is chosen with uniform probability. 2. The sites are exchanged, i.e. :math:`s^\prime_i = s_j` and :math:`s^\prime_j = s_i`. Notice that this sampling method generates random permutations of the quantum numbers, thus global quantities such as the sum of the local quantum numbers are conserved during the sampling. This scheme should be used then only when sampling in a region where :math:`\sum_i s_i = \mathrm{constant}` is needed, otherwise the sampling would be strongly not ergodic. Args: hilbert: The Hilbert space to sample. d_max: The maximum graph distance allowed for exchanges. n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). sweep_size: Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space. This is equivalent to subsampling the Markov chain. reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). Examples: Sampling from a RBM machine in a 1D lattice of spin 1/2, using nearest-neighbor exchanges. >>> import netket as nk >>> >>> g=nk.graph.Hypercube(length=10,n_dim=2,pbc=True) >>> hi=nk.hilbert.Spin(s=0.5, N=g.n_nodes) >>> >>> # Construct a MetropolisExchange Sampler >>> sa = nk.sampler.MetropolisExchange(hi, graph=g) >>> print(sa) MetropolisSampler(rule = ExchangeRule(# of clusters: 200), n_chains = 16, sweep_size = 100, reset_chains = False, machine_power = 2, dtype = <class 'float'>) """ from .rules import ExchangeRule rule = ExchangeRule(clusters=clusters, graph=graph, d_max=d_max) return MetropolisSampler(hilbert, rule, **kwargs)
[docs] def MetropolisHamiltonian(hilbert, hamiltonian, **kwargs) -> MetropolisSampler: r""" Sampling based on the off-diagonal elements of a Hamiltonian (or a generic 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}|), where :math:`\theta(x)` is the Heaviside step function, and :math:`\mathcal{N}(\mathbf{s})` is a state-dependent normalization. The effect of this transition probability is then to connect (with uniform probability) a given state :math:`\mathbf{s}` to all those states :math:`\mathbf{s}^\prime` for which the Hamiltonian has finite matrix elements. Notice that this sampler preserves by construction all the symmetries of the Hamiltonian. This is in generally not true for the local samplers instead. This sampler only works on the CPU. To use the Hamiltonian sampler with GPUs, you should use :class:`netket.sampler.MetropolisHamiltonianNumpy` Args: hilbert: The Hilbert space to sample. hamiltonian: The operator used to perform off-diagonal transition. n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). sweep_size: Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space. This is equivalent to subsampling the Markov chain. reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). Examples: Sampling from a RBM machine in a 1D lattice of spin 1/2 >>> import netket as nk >>> >>> g=nk.graph.Hypercube(length=10,n_dim=2,pbc=True) >>> hi=nk.hilbert.Spin(s=0.5, N=g.n_nodes) >>> >>> # Transverse-field Ising Hamiltonian >>> ha = nk.operator.Ising(hilbert=hi, h=1.0, graph=g) >>> >>> # Construct a MetropolisHamiltonian Sampler >>> sa = nk.sampler.MetropolisHamiltonian(hi, hamiltonian=ha) >>> print(sa) MetropolisSampler(rule = HamiltonianRuleNumba(Ising(J=1.0, h=1.0; dim=100)), n_chains = 16, sweep_size = 100, reset_chains = False, machine_power = 2, dtype = <class 'float'>) """ from .rules import HamiltonianRule rule = HamiltonianRule(hamiltonian) return MetropolisSampler(hilbert, rule, **kwargs)
[docs] def MetropolisGaussian(hilbert, sigma=1.0, **kwargs) -> MetropolisSampler: """This sampler acts on all particle positions simultaneously and proposes a new state according to a Gaussian distribution with width `sigma`. Args: hilbert: The continuous Hilbert space to sample. sigma: The width of the Gaussian proposal distribution (default = 1.0). n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). sweep_size: Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space. This is equivalent to subsampling the Markov chain. reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). """ if not isinstance(hilbert, ContinuousHilbert): raise ValueError("This sampler only works for Continuous Hilbert spaces.") from .rules import GaussianRule rule = GaussianRule(sigma) return MetropolisSampler(hilbert, rule, **kwargs)
[docs] def MetropolisAdjustedLangevin( hilbert, dt=0.001, chunk_size=None, **kwargs ) -> MetropolisSampler: r"""This sampler acts on all particle positions simultaneously and takes a Langevin step [1]: .. math:: x_{t+dt} = x_t + dt \nabla_x \log p(x) \vert_{x=x_t} + \sqrt{2 dt}\eta, where :math:`\eta` is normal distributed noise :math:`\eta \sim \mathcal{N}(0,1)`. This sampler only works for continuous Hilbert spaces. [1]: https://en.wikipedia.org/wiki/Metropolis-adjusted_Langevin_algorithm Args: hilbert: The continuous Hilbert space to sample. dt: Time step size for the Langevin dynamics (noise with variance 2*dt). chunk_size: Chunk size to compute the gradients of the log probability. n_chains: The total number of independent Markov chains across all MPI ranks. Either specify this or `n_chains_per_rank`. n_chains_per_rank: Number of independent chains on every MPI rank (default = 16). sweep_size: Number of sweeps for each step along the chain. Defaults to the number of sites in the Hilbert space. This is equivalent to subsampling the Markov chain. reset_chains: If True, resets the chain state when `reset` is called on every new sampling (default = False). machine_pow: The power to which the machine should be exponentiated to generate the pdf (default = 2). dtype: The dtype of the states sampled (default = np.float64). """ if not isinstance(hilbert, ContinuousHilbert): raise ValueError("This sampler only works for Continuous Hilbert spaces.") from .rules import LangevinRule rule = LangevinRule(dt=dt, chunk_size=chunk_size) return MetropolisSampler(hilbert, rule, **kwargs)