From e8bcf8de6faa56f2147e64d333b76b7948aa4042 Mon Sep 17 00:00:00 2001 From: Fabian Joswig Date: Tue, 7 Dec 2021 08:09:38 +0000 Subject: [PATCH] fix: array mode now also works with covobs with N>1 --- pyerrors/obs.py | 14 ++++++++------ tests/linalg_test.py | 4 ++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/pyerrors/obs.py b/pyerrors/obs.py index f7029a63..bc60a3a2 100644 --- a/pyerrors/obs.py +++ b/pyerrors/obs.py @@ -1148,23 +1148,25 @@ def derived_observable(func, data, array_mode=False, **kwargs): if array_mode is True: - class _Zero_grad(): - def __init__(self): - self.grad = 0 + new_covobs_lengths = dict(set([y for x in [[(n, o.covobs[n].N) for n in o.cov_names] for o in raveled_data] for y in x])) - zero_grad = _Zero_grad() + class _Zero_grad(): + def __init__(self, N): + # self.grad = np.zeros(N) + self.grad = np.zeros((N, 1)) d_extracted = {} g_extracted = {} for name in new_sample_names: d_extracted[name] = [] + ens_length = len(new_idl_d[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.get(name, np.zeros(ens_length)), o.idl.get(name, new_idl_d[name]), o.shape.get(name, ens_length), new_idl_d[name]) for o in dat.reshape(np.prod(dat.shape))]).reshape(dat.shape + (ens_length, ))) for name in new_cov_names: g_extracted[name] = [] + zero_grad = _Zero_grad(new_covobs_lengths[name]) for i_dat, dat in enumerate(data): - g_extracted[name].append(np.array([o.covobs.get(name, zero_grad).grad for o in dat.reshape(np.prod(dat.shape))]).reshape(dat.shape + (1, ))) + g_extracted[name].append(np.array([o.covobs.get(name, zero_grad).grad for o in dat.reshape(np.prod(dat.shape))]).reshape(dat.shape + (new_covobs_lengths[name], 1))) for i_val, new_val in np.ndenumerate(new_values): new_deltas = {} diff --git a/tests/linalg_test.py b/tests/linalg_test.py index 73bcf5bb..55e65e30 100644 --- a/tests/linalg_test.py +++ b/tests/linalg_test.py @@ -195,7 +195,7 @@ def test_matmul_irregular_histories(): irregular_array = [] for i in range(dim ** 2): irregular_array.append(pe.Obs([np.random.normal(1.1, 0.2, len(idl))], ['ens1'], idl=[idl])) - irregular_matrix = np.array(irregular_array).reshape((dim, dim)) + irregular_matrix = np.array(irregular_array).reshape((dim, dim)) * pe.cov_Obs([1.0, 1.0], [[0.001,0.0001], [0.0001, 0.002]], 'norm')[0] t1 = standard_matrix @ irregular_matrix t2 = pe.linalg.matmul(standard_matrix, irregular_matrix) @@ -213,7 +213,7 @@ def test_irregular_matrix_inverse(): irregular_array = [] for i in range(dim ** 2): irregular_array.append(pe.Obs([np.random.normal(1.1, 0.2, len(idl)), np.random.normal(0.25, 0.1, 10)], ['ens1', 'ens2'], idl=[idl, range(1, 11)])) - irregular_matrix = np.array(irregular_array).reshape((dim, dim)) + irregular_matrix = np.array(irregular_array).reshape((dim, dim)) * pe.cov_Obs(1.0, 0.002, 'cov') * pe.pseudo_Obs(1.0, 0.002, 'ens2|r23') invertible_irregular_matrix = np.identity(dim) + irregular_matrix @ irregular_matrix.T