def _forward_and_backward(self):
"""
Performs a number of VMC optimization steps.
Args:
n_steps (int): Number of steps to perform.
"""
self._sampler.reset()
# Burnout phase
self._sampler.generate_samples(self._n_discard)
# Generate samples and store them
self._samples = self._sampler.generate_samples(self._n_samples_node,
samples=self._samples)
# Compute the local energy estimator and average Energy
eloc, self._loss_stats = self._get_mc_stats(self._ham)
# Center the local energy
eloc -= _mean(eloc)
samples_r = self._samples.reshape((-1, self._samples.shape[-1]))
eloc_r = eloc.reshape(-1, 1)
# Perform update
if self._sr:
# When using the SR (Natural gradient) we need to have the full jacobian
self._grads, self._jac = self._machine.vector_jacobian_prod(
samples_r,
eloc_r / self._n_samples,
self._grads,
return_jacobian=True)
self._grads = tree_map(_sum_inplace, self._grads)
self._dp = self._sr.compute_update(self._jac, self._grads,
self._dp)
else:
# Computing updates using the simple gradient
self._grads = self._machine.vector_jacobian_prod(
samples_r, eloc_r / self._n_samples, self._grads)
self._grads = tree_map(_sum_inplace, self._grads)
# if Real pars but complex gradient, take only real part
# not necessary for SR because sr already does it.
if not self._machine.has_complex_parameters:
self._dp = tree_map(lambda x: x.real, self._grads)
else:
self._dp = self._grads
return self._dp
def estimate_expectations(ops,
sampler,
n_samples,
n_discard=None,
compute_gradients=False):
"""
For a sequence of linear operators, computes a statistical estimate of the
respective expectation values, variances, and optionally gradients of the
expectation values with respect to the variational parameters.
The estimate is based on `n_samples` configurations
obtained from `sampler`.
Args:
ops: pytree of linear operators
sampler: A NetKet sampler
n_samples: Number of MC samples used to estimate expectation values
n_discard: Number of MC samples dropped from the start of the
chain (burn-in). Defaults to `n_samples //10`.
compute_gradients: Whether to compute the gradients of the
observables.
Returns:
Either `stats` or, if `der_logs` is passed, a tuple of `stats` and `grad`:
stats: A sequence of Stats object containing mean, variance,
and MC diagonstics for each operator in `ops`.
grad: A sequence of gradients of the expectation value of `op`,
as ndarray of shape `(psi.n_par,)`, for each `op` in `ops`.
"""
from netket.operator import local_values as _local_values
from ._C_netket import stats as nst
psi = sampler.machine
if not n_discard:
n_discard = n_samples // 10
# Burnout phase
sampler.generate_samples(n_discard)
# Generate samples
samples = sampler.generate_samples(n_samples)
if compute_gradients:
der_logs = psi.der_log(samples)
def estimate(op):
lvs = _local_values(op, psi, samples)
stats = nst.statistics(lvs)
if compute_gradients:
grad = nst.covariance_sv(lvs, der_logs)
return stats, grad
else:
return stats
return tree_map(estimate, ops)
def estimate(self, observables):
"""
Return MCMC statistics for the expectation value of observables in the
current state of the driver.
Args:
observables: A pytree of operators for which statistics should be computed.
Returns:
A pytree of the same structure as the input, containing MCMC statistics
for the corresponding operators as leaves.
"""
return tree_map(self._estimate_stats, observables)
def test_tree_map():
# Structure: name -> (input, expected)
test_trees = {
"None": (None, None),
"leaf": (0, 1),
"list": ([0, 1], [1, 2]),
"tuple1": ((0,), (1,)),
"tuple3": ((0, 1, 0), (1, 2, 1)),
"list of list": ([[0, 1], [1], [0, 1, 0]], [[1, 2], [2], [1, 2, 1]]),
"dict": ({"a": 0, "b": 1}, {"a": 1, "b": 2}),
"mixed list": ([(0, 1), [1], 1, {"a": 0}], [(1, 2), [2], 2, {"a": 1}]),
}
for name, (inp, expected) in test_trees.items():
info = "{}, input={}".format(name, inp)
assert tree_map(lambda x: x + 1, inp) == expected, info
def run(
self,
output_prefix,
n_iter,
obs=None,
save_params_every=50,
write_every=50,
step_size=1,
show_progress=True,
):
"""
Executes the Monte Carlo Variational optimization, updating the weights of the network
stored in this driver for `n_iter` steps and dumping values of the observables `obs`
in the output. The output is a json file at `output_prefix`, overwriting files with
the same prefix.
Args:
:output_prefix: The prefix at which json output should be stored (ignored if logger
is provided).
:n_iter: the total number of iterations
:obs: An iterable containing all observables that should be computed
:save_params_every: Every how many steps the parameters of the network should be
serialized to disk (ignored if logger is provided)
:write_every: Every how many steps the json data should be flushed to disk (ignored if
logger is provided)
:step_size: Every how many steps should observables be logged to disk (default=1)
:show_progress: If true displays a progress bar (default=True)
"""
if obs is None:
# TODO remove the first case after deprecation of self._obs in 3.0
if len(self._obs) != 0:
obs = self._obs
else:
obs = {}
logger = _JsonLog(output_prefix, save_params_every, write_every)
# Don't log on non-root nodes
if self._mynode != 0:
logger = None
with tqdm(
self.iter(n_iter, step_size), total=n_iter, disable=not show_progress
) as itr:
for step in itr:
# if the cost-function is defined then report it in the progress bar
if self._loss_stats is not None:
itr.set_postfix_str(self._loss_name + "=" + str(self._loss_stats))
obs_data = self.estimate(obs)
if self._loss_stats is not None:
obs_data[self._loss_name] = self._loss_stats
log_data = tree_map(_obs_stat_to_dict, obs_data)
if logger is not None:
logger(step, log_data, self.machine)
# flush at the end of the evolution so that final values are saved to
# file
if logger is not None:
logger.flush(self.machine)
def run(
self,
n_iter,
out=None,
obs=None,
show_progress=True,
save_params_every=50, # for default logger
write_every=50, # for default logger
step_size=1, # for default logger
output_prefix=None, # TODO: deprecated
):
"""
Executes the Monte Carlo Variational optimization, updating the weights of the network
stored in this driver for `n_iter` steps and dumping values of the observables `obs`
in the output `logger`. If no logger is specified, creates a json file at `output_prefix`,
overwriting files with the same prefix.
!! Compatibility v2.1
Before v2.1 the order of the first two arguments, `n_iter` and `output_prefix` was
reversed. The reversed ordering will still be supported until v3.0, but is deprecated.
Args:
:n_iter: the total number of iterations
:out: A logger object to be used to store simulation log and data.
If this argument is a string, it will be used as output prefix for the standard JSON logger.
:obs: An iterable containing all observables that should be computed
:save_params_every: Every how many steps the parameters of the network should be
serialized to disk (ignored if logger is provided)
:write_every: Every how many steps the json data should be flushed to disk (ignored if
logger is provided)
:step_size: Every how many steps should observables be logged to disk (default=1)
:show_progress: If true displays a progress bar (default=True)
:output_prefix: (Deprecated) The prefix at which json output should be stored (ignored if out
is provided).
"""
# TODO Remove this deprecated code in v3.0
# manage deprecated where argument names are not specified, and
# prefix is passed as the first positional argument and the number
# of iterations as a second argument.
if type(n_iter) is str and type(out) is int:
n_iter, out = out, n_iter
warn_deprecation(
"The positional syntax run(output_prefix, n_iter, **args) is deprecated, use run(n_iter, output_prefix, **args) instead."
)
if obs is None:
# TODO remove the first case after deprecation of self._obs in 3.0
if len(self._obs) != 0:
obs = self._obs
else:
obs = {}
# output_prefix is deprecated. out should be used and takes over
# error out if both are passed
# TODO: remove in v3.0
if out is not None and output_prefix is not None:
raise ValueError(
"Invalid out and output_prefix arguments. Only one of the two can be passed. Note that output_prefix is deprecated and you should use out."
)
elif out is None:
warn_deprecation(
"The output_prefix argument is deprecated. Use out instead.")
out = output_prefix
# Log only non-root nodes
if self._mynode == 0:
# if out is a path, create an overwriting Json Log for output
if isinstance(out, str):
logger = _JsonLog(out, "w", save_params_every, write_every)
else:
logger = out
else:
logger = None
with tqdm(self.iter(n_iter, step_size),
total=n_iter,
disable=not show_progress) as itr:
for step in itr:
# if the cost-function is defined then report it in the progress bar
if self._loss_stats is not None:
itr.set_postfix_str(self._loss_name + "=" +
str(self._loss_stats))
obs_data = self.estimate(obs)
if self._loss_stats is not None:
obs_data[self._loss_name] = self._loss_stats
log_data = tree_map(_obs_stat_to_dict, obs_data)
if logger is not None:
logger(step, log_data, self.machine)
# flush at the end of the evolution so that final values are saved to
# file
if logger is not None:
logger.flush(self.machine)
def _forward_and_backward(self):
"""
Performs a number of VMC optimization steps.
Args:
n_steps (int): Number of steps to perform.
"""
self._sampler.reset()
self._obs_samples_valid = False
# Burnout phase
self._sampler.generate_samples(self._n_discard)
# Generate samples and store them
self._samples = self._sampler.generate_samples(self._n_samples_node,
samples=self._samples)
# Estimate C^[loc] (local energy) and LdagL
self._lloc, self._loss_stats = self._get_mc_superop_stats(self._lind)
# Flatten chain dimension
samples_r = self._samples.reshape((-1, self._samples.shape[-1]))
lloc_r = self._lloc.reshape(-1, 1)
# Compute Log derivatives
self._der_logs = self._machine.der_log(samples_r)
# Compute statistical average (also across nodes)
self._der_logs_ave = tree_map(_mean, self._der_logs, axis=0)
# Compute \nabla C^[Loc]
self._der_loc_vals = _der_local_values(self._lind,
self._machine,
samples_r,
center_derivative=False)
# apply this function to every element of the gradient.
n_samples_node = lloc_r.shape[0]
def gradfun(der_loc_vals, der_logs_ave):
par_dims = der_loc_vals.ndim - 1
_lloc_r = lloc_r.reshape((n_samples_node, ) +
tuple(1 for i in range(par_dims)))
grad = _mean(der_loc_vals.conjugate() * _lloc_r, axis=0) - (
der_logs_ave.conjugate() * self._loss_stats.mean)
return grad
self._grads = trees2_map(gradfun, self._der_loc_vals,
self._der_logs_ave)
# Perform update
if self._sr:
# Center the log derivatives
self._jac = trees2_map(lambda x, y: x - y, self._der_logs,
self._der_logs_ave)
self._dp = self._sr.compute_update(self._jac, self._grads,
self._dp)
else:
# if Real pars but complex gradient, take only real part
# not necessary for SR because sr already does it.
if not self._machine.has_complex_parameters:
self._dp = tree_map(lambda x: x.real, self._grads)
else:
self._dp = self._grads
return self._dp
