derived_array adjusted to work with irregular Monte Carlo chains.

Additional test required for verification.
This commit is contained in:
Fabian Joswig 2021-11-11 10:45:19 +00:00
parent 0f14aba083
commit 4d3b00eb48

View file

@ -1,23 +1,25 @@
import numpy as np import numpy as np
from autograd import jacobian from autograd import jacobian
import autograd.numpy as anp # Thinly-wrapped numpy import autograd.numpy as anp # Thinly-wrapped numpy
from .obs import derived_observable, CObs, Obs from .obs import derived_observable, CObs, Obs, merge_idx, expand_deltas_for_merge, filter_zeroes
from functools import partial from functools import partial
from autograd.extend import defvjp from autograd.extend import defvjp
def derived_array(func, data, **kwargs): def derived_array(func, data, **kwargs):
"""Construct a derived Obs according to func(data, **kwargs) of matrix value data """Construct a derived Obs for a matrix valued function according to func(data, **kwargs) using automatic differentiation.
using automatic differentiation.
Parameters Parameters
---------- ----------
func -- arbitrary function of the form func(data, **kwargs). For the func : object
arbitrary function of the form func(data, **kwargs). For the
automatic differentiation to work, all numpy functions have to have automatic differentiation to work, all numpy functions have to have
the autograd wrapper (use 'import autograd.numpy as anp'). the autograd wrapper (use 'import autograd.numpy as anp').
data -- list of Obs, e.g. [obs1, obs2, obs3]. data : list
man_grad -- manually supply a list or an array which contains the jacobian list of Obs, e.g. [obs1, obs2, obs3].
man_grad : list
manually supply a list or an array which contains the jacobian
of func. Use cautiously, supplying the wrong derivative will of func. Use cautiously, supplying the wrong derivative will
not be intercepted. not be intercepted.
""" """
@ -30,25 +32,29 @@ def derived_array(func, data, **kwargs):
if isinstance(i_data, Obs): if isinstance(i_data, Obs):
first_name = i_data.names[0] first_name = i_data.names[0]
first_shape = i_data.shape[first_name] first_shape = i_data.shape[first_name]
first_idl = i_data.idl[first_name]
break break
for i in range(len(raveled_data)): for i in range(len(raveled_data)):
if isinstance(raveled_data[i], (int, float)): if isinstance(raveled_data[i], (int, float)):
raveled_data[i] = Obs([raveled_data[i] + np.zeros(first_shape)], [first_name]) raveled_data[i] = Obs([raveled_data[i] + np.zeros(first_shape)], [first_name], idl=[first_idl])
n_obs = len(raveled_data) n_obs = len(raveled_data)
new_names = sorted(set([y for x in [o.names for o in raveled_data] for y in x])) new_names = sorted(set([y for x in [o.names for o in raveled_data] for y in x]))
new_shape = {} is_merged = len(list(filter(lambda o: o.is_merged is True, raveled_data))) > 0
for i_data in raveled_data: reweighted = len(list(filter(lambda o: o.reweighted is True, raveled_data))) > 0
new_idl_d = {}
for name in new_names: for name in new_names:
tmp = i_data.shape.get(name) idl = []
for i_data in raveled_data:
tmp = i_data.idl.get(name)
if tmp is not None: if tmp is not None:
if new_shape.get(name) is None: idl.append(tmp)
new_shape[name] = tmp new_idl_d[name] = merge_idx(idl)
else: if not is_merged:
if new_shape[name] != tmp: is_merged = (1 != len(set([len(idx) for idx in [*idl, new_idl_d[name]]])))
raise Exception('Shapes of ensemble', name, 'do not match.')
if data.ndim == 1: if data.ndim == 1:
values = np.array([o.value for o in data]) values = np.array([o.value for o in data])
else: else:
@ -71,8 +77,6 @@ def derived_array(func, data, **kwargs):
deriv = np.asarray(kwargs.get('man_grad')) deriv = np.asarray(kwargs.get('man_grad'))
if new_values.shape + data.shape != deriv.shape: if new_values.shape + data.shape != deriv.shape:
raise Exception('Manual derivative does not have correct shape.') raise Exception('Manual derivative does not have correct shape.')
elif kwargs.get('num_grad') is True:
raise Exception('Multi mode currently not supported for numerical derivative')
else: else:
deriv = jacobian(func)(values, **kwargs) deriv = jacobian(func)(values, **kwargs)
@ -83,7 +87,7 @@ def derived_array(func, data, **kwargs):
d_extracted[name] = [] d_extracted[name] = []
for i_dat, dat in enumerate(data): for i_dat, dat in enumerate(data):
ens_length = dat.ravel()[0].shape[name] ens_length = dat.ravel()[0].shape[name]
d_extracted[name].append(np.array([o.deltas[name] for o in dat.reshape(np.prod(dat.shape))]).reshape(dat.shape + (ens_length, ))) d_extracted[name].append(np.array([expand_deltas_for_merge(o.deltas[name], o.idl[name], o.shape[name], new_idl_d[name]) for o in dat.reshape(np.prod(dat.shape))]).reshape(dat.shape + (ens_length, )))
for i_val, new_val in np.ndenumerate(new_values): for i_val, new_val in np.ndenumerate(new_values):
new_deltas = {} new_deltas = {}
@ -95,12 +99,21 @@ def derived_array(func, data, **kwargs):
new_samples = [] new_samples = []
new_means = [] new_means = []
for name in new_names: new_idl = []
new_samples.append(new_deltas[name]) if is_merged:
filtered_names, filtered_deltas, filtered_idl_d = filter_zeroes(new_names, new_deltas, new_idl_d)
else:
filtered_names = new_names
filtered_deltas = new_deltas
filtered_idl_d = new_idl_d
for name in filtered_names:
new_samples.append(filtered_deltas[name])
new_means.append(new_r_values[name][i_val]) new_means.append(new_r_values[name][i_val])
new_idl.append(filtered_idl_d[name])
final_result[i_val] = Obs(new_samples, new_names, means=new_means) final_result[i_val] = Obs(new_samples, filtered_names, means=new_means, idl=new_idl)
final_result[i_val]._value = new_val final_result[i_val]._value = new_val
final_result[i_val].is_merged = is_merged
final_result[i_val].reweighted = reweighted
return final_result return final_result