lattice/pyerrors
1
0
Fork 0
mirror of https://github.com/fjosw/pyerrors.git synced 2025-03-15 14:50:25 +01:00
Code Issues Projects Releases Packages Wiki Activity

feat: first working version of array_mode in dervived_observable

Browse source
This commit is contained in:
Fabian Joswig 2021-12-02 12:50:08 +00:00
parent ed47d50286
commit 147bc6b24b
2 changed files with 22 additions and 120 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

124
pyerrors/linalg.py
View file

@ -1,123 +1,11 @@
import numpy as np
from autograd import jacobian
import autograd.numpy as anp # Thinly-wrapped numpy
from .obs import derived_observable, CObs, Obs, _merge_idx, _expand_deltas_for_merge, _filter_zeroes, import_jackknife
from .obs import derived_observable, CObs, Obs, import_jackknife
from functools import partial
from autograd.extend import defvjp
def derived_array(func, data, **kwargs):
"""Construct a derived Obs for a matrix valued function according to func(data, **kwargs) using automatic differentiation.
Parameters
----------
func : object
arbitrary function of the form func(data, **kwargs). For the
automatic differentiation to work, all numpy functions have to have
the autograd wrapper (use 'import autograd.numpy as anp').
data : list
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
not be intercepted.
"""
data = np.asarray(data)
raveled_data = data.ravel()
# Workaround for matrix operations containing non Obs data
for i_data in raveled_data:
if isinstance(i_data, Obs):
first_name = i_data.names[0]
first_shape = i_data.shape[first_name]
first_idl = i_data.idl[first_name]
break
for i in range(len(raveled_data)):
if isinstance(raveled_data[i], (int, float)):
raveled_data[i] = Obs([raveled_data[i] + np.zeros(first_shape)], [first_name], idl=[first_idl])
n_obs = len(raveled_data)
new_names = sorted(set([y for x in [o.names for o in raveled_data] for y in x]))
is_merged = {name: (len(list(filter(lambda o: o.is_merged.get(name, False) is True, raveled_data))) > 0) for name in new_names}
reweighted = len(list(filter(lambda o: o.reweighted is True, raveled_data))) > 0
new_idl_d = {}
for name in new_names:
idl = []
for i_data in raveled_data:
tmp = i_data.idl.get(name)
if tmp is not None:
idl.append(tmp)
new_idl_d[name] = _merge_idx(idl)
if not is_merged[name]:
is_merged[name] = (1 != len(set([len(idx) for idx in [*idl, new_idl_d[name]]])))
if data.ndim == 1:
values = np.array([o.value for o in data])
else:
values = np.vectorize(lambda x: x.value)(data)
new_values = func(values, **kwargs)
new_r_values = {}
for name in new_names:
tmp_values = np.zeros(n_obs)
for i, item in enumerate(raveled_data):
tmp = item.r_values.get(name)
if tmp is None:
tmp = item.value
tmp_values[i] = tmp
tmp_values = np.array(tmp_values).reshape(data.shape)
new_r_values[name] = func(tmp_values, **kwargs)
if 'man_grad' in kwargs:
deriv = np.asarray(kwargs.get('man_grad'))
if new_values.shape + data.shape != deriv.shape:
raise Exception('Manual derivative does not have correct shape.')
else:
deriv = jacobian(func)(values, **kwargs)
final_result = np.zeros(new_values.shape, dtype=object)
d_extracted = {}
for name in new_names:
d_extracted[name] = []
for i_dat, dat in enumerate(data):
ens_length = len(new_idl_d[name])
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):
new_deltas = {}
for name in new_names:
ens_length = d_extracted[name][0].shape[-1]
new_deltas[name] = np.zeros(ens_length)
for i_dat, dat in enumerate(d_extracted[name]):
new_deltas[name] += np.tensordot(deriv[i_val + (i_dat, )], dat)
new_samples = []
new_means = []
new_idl = []
for name in new_names:
if is_merged[name]:
filtered_deltas, filtered_idl_d = _filter_zeroes(new_deltas[name], new_idl_d[name])
else:
filtered_deltas = new_deltas[name]
filtered_idl_d = new_idl_d[name]
new_samples.append(filtered_deltas)
new_idl.append(filtered_idl_d)
new_means.append(new_r_values[name][i_val])
final_result[i_val] = Obs(new_samples, new_names, means=new_means, idl=new_idl)
final_result[i_val]._value = new_val
final_result[i_val].is_merged = is_merged
final_result[i_val].reweighted = reweighted
return final_result
def matmul(*operands):
"""Matrix multiply all operands.
@ -157,8 +45,8 @@ def matmul(*operands):
def multi_dot_i(operands):
return multi_dot(operands, 'Imag')
Nr = derived_array(multi_dot_r, extended_operands)
Ni = derived_array(multi_dot_i, extended_operands)
Nr = derived_observable(multi_dot_r, extended_operands, array_mode=True)
Ni = derived_observable(multi_dot_i, extended_operands, array_mode=True)
res = np.empty_like(Nr)
for (n, m), entry in np.ndenumerate(Nr):
@ -171,7 +59,7 @@ def matmul(*operands):
for op in operands[1:]:
stack = stack @ op
return stack
return derived_array(multi_dot, operands)
return derived_observable(multi_dot, operands, array_mode=True)
def jack_matmul(*operands):
@ -360,7 +248,7 @@ def _mat_mat_op(op, obs, **kwargs):
if kwargs.get('num_grad') is True:
op_big_matrix = _num_diff_mat_mat_op(op, big_matrix, **kwargs)
else:
op_big_matrix = derived_array(lambda x, **kwargs: op(x), [big_matrix])[0]
op_big_matrix = derived_observable(lambda x, **kwargs: op(x), [big_matrix], array_mode=True)[0]
dim = op_big_matrix.shape[0]
op_A = op_big_matrix[0: dim // 2, 0: dim // 2]
op_B = op_big_matrix[dim // 2:, 0: dim // 2]
@ -371,7 +259,7 @@ def _mat_mat_op(op, obs, **kwargs):
else:
if kwargs.get('num_grad') is True:
return _num_diff_mat_mat_op(op, obs, **kwargs)
return derived_array(lambda x, **kwargs: op(x), [obs])[0]
return derived_observable(lambda x, **kwargs: op(x), [obs], array_mode=True)[0]
def eigh(obs, **kwargs):

18
pyerrors/obs.py
View file

@ -1015,7 +1015,7 @@ def _filter_zeroes(deltas, idx, eps=Obs.filter_eps):
return deltas, idx
def derived_observable(func, data, **kwargs):
def derived_observable(func, data, array_mode=False, **kwargs):
"""Construct a derived Obs according to func(data, **kwargs) using automatic differentiation.
Parameters
@ -1138,14 +1138,28 @@ def derived_observable(func, data, **kwargs):
final_result = np.zeros(new_values.shape, dtype=object)
if array_mode is True:
d_extracted = {}
for name in new_names:
d_extracted[name] = []
for i_dat, dat in enumerate(data):
ens_length = len(new_idl_d[name])
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):
new_deltas = {}
new_grad = {}
if array_mode is True:
for name in new_names:
ens_length = d_extracted[name][0].shape[-1]
new_deltas[name] = np.zeros(ens_length)
for i_dat, dat in enumerate(d_extracted[name]):
new_deltas[name] += np.tensordot(deriv[i_val + (i_dat, )], dat)
for j_obs, obs in np.ndenumerate(data):
for name in obs.names:
if name in obs.cov_names:
new_grad[name] = new_grad.get(name, 0) + deriv[i_val + j_obs] * obs.covobs[name].grad
else:
elif array_mode is False:
new_deltas[name] = new_deltas.get(name, 0) + deriv[i_val + j_obs] * _expand_deltas_for_merge(obs.deltas[name], obs.idl[name], obs.shape[name], new_idl_d[name])
new_covobs = {name: Covobs(0, allcov[name], name, grad=new_grad[name]) for name in new_grad}