pyerrors.linalg

  1import numpy as np
  2import autograd.numpy as anp  # Thinly-wrapped numpy
  3from .obs import derived_observable, CObs, Obs, import_jackknife
  4
  5
  6def matmul(*operands):
  7    """Matrix multiply all operands.
  8
  9    Parameters
 10    ----------
 11    operands : numpy.ndarray
 12        Arbitrary number of 2d-numpy arrays which can be real or complex
 13        Obs valued.
 14
 15    This implementation is faster compared to standard multiplication via the @ operator.
 16    """
 17    if any(isinstance(o[0, 0], CObs) for o in operands):
 18        extended_operands = []
 19        for op in operands:
 20            tmp = np.vectorize(lambda x: (np.real(x), np.imag(x)))(op)
 21            extended_operands.append(tmp[0])
 22            extended_operands.append(tmp[1])
 23
 24        def multi_dot(operands, part):
 25            stack_r = operands[0]
 26            stack_i = operands[1]
 27            for op_r, op_i in zip(operands[2::2], operands[3::2]):
 28                tmp_r = stack_r @ op_r - stack_i @ op_i
 29                tmp_i = stack_r @ op_i + stack_i @ op_r
 30
 31                stack_r = tmp_r
 32                stack_i = tmp_i
 33
 34            if part == 'Real':
 35                return stack_r
 36            else:
 37                return stack_i
 38
 39        def multi_dot_r(operands):
 40            return multi_dot(operands, 'Real')
 41
 42        def multi_dot_i(operands):
 43            return multi_dot(operands, 'Imag')
 44
 45        Nr = derived_observable(multi_dot_r, extended_operands, array_mode=True)
 46        Ni = derived_observable(multi_dot_i, extended_operands, array_mode=True)
 47
 48        res = np.empty_like(Nr)
 49        for (n, m), entry in np.ndenumerate(Nr):
 50            res[n, m] = CObs(Nr[n, m], Ni[n, m])
 51
 52        return res
 53    else:
 54        def multi_dot(operands):
 55            stack = operands[0]
 56            for op in operands[1:]:
 57                stack = stack @ op
 58            return stack
 59        return derived_observable(multi_dot, operands, array_mode=True)
 60
 61
 62def jack_matmul(*operands):
 63    """Matrix multiply both operands making use of the jackknife approximation.
 64
 65    Parameters
 66    ----------
 67    operands : numpy.ndarray
 68        Arbitrary number of 2d-numpy arrays which can be real or complex
 69        Obs valued.
 70
 71    For large matrices this is considerably faster compared to matmul.
 72    """
 73
 74    def _exp_to_jack(matrix):
 75        base_matrix = np.empty_like(matrix)
 76        for index, entry in np.ndenumerate(matrix):
 77            base_matrix[index] = entry.export_jackknife()
 78        return base_matrix
 79
 80    def _imp_from_jack(matrix, name, idl):
 81        base_matrix = np.empty_like(matrix)
 82        for index, entry in np.ndenumerate(matrix):
 83            base_matrix[index] = import_jackknife(entry, name, [idl])
 84        return base_matrix
 85
 86    def _exp_to_jack_c(matrix):
 87        base_matrix = np.empty_like(matrix)
 88        for index, entry in np.ndenumerate(matrix):
 89            base_matrix[index] = entry.real.export_jackknife() + 1j * entry.imag.export_jackknife()
 90        return base_matrix
 91
 92    def _imp_from_jack_c(matrix, name, idl):
 93        base_matrix = np.empty_like(matrix)
 94        for index, entry in np.ndenumerate(matrix):
 95            base_matrix[index] = CObs(import_jackknife(entry.real, name, [idl]),
 96                                      import_jackknife(entry.imag, name, [idl]))
 97        return base_matrix
 98
 99    if any(isinstance(o.flat[0], CObs) for o in operands):
100        name = operands[0].flat[0].real.names[0]
101        idl = operands[0].flat[0].real.idl[name]
102
103        r = _exp_to_jack_c(operands[0])
104        for op in operands[1:]:
105            if isinstance(op.flat[0], CObs):
106                r = r @ _exp_to_jack_c(op)
107            else:
108                r = r @ op
109        return _imp_from_jack_c(r, name, idl)
110    else:
111        name = operands[0].flat[0].names[0]
112        idl = operands[0].flat[0].idl[name]
113
114        r = _exp_to_jack(operands[0])
115        for op in operands[1:]:
116            if isinstance(op.flat[0], Obs):
117                r = r @ _exp_to_jack(op)
118            else:
119                r = r @ op
120        return _imp_from_jack(r, name, idl)
121
122
123def einsum(subscripts, *operands):
124    """Wrapper for numpy.einsum
125
126    Parameters
127    ----------
128    subscripts : str
129        Subscripts for summation (see numpy documentation for details)
130    operands : numpy.ndarray
131        Arbitrary number of 2d-numpy arrays which can be real or complex
132        Obs valued.
133    """
134
135    def _exp_to_jack(matrix):
136        base_matrix = []
137        for index, entry in np.ndenumerate(matrix):
138            base_matrix.append(entry.export_jackknife())
139        return np.asarray(base_matrix).reshape(matrix.shape + base_matrix[0].shape)
140
141    def _exp_to_jack_c(matrix):
142        base_matrix = []
143        for index, entry in np.ndenumerate(matrix):
144            base_matrix.append(entry.real.export_jackknife() + 1j * entry.imag.export_jackknife())
145        return np.asarray(base_matrix).reshape(matrix.shape + base_matrix[0].shape)
146
147    def _imp_from_jack(matrix, name, idl):
148        base_matrix = np.empty(shape=matrix.shape[:-1], dtype=object)
149        for index in np.ndindex(matrix.shape[:-1]):
150            base_matrix[index] = import_jackknife(matrix[index], name, [idl])
151        return base_matrix
152
153    def _imp_from_jack_c(matrix, name, idl):
154        base_matrix = np.empty(shape=matrix.shape[:-1], dtype=object)
155        for index in np.ndindex(matrix.shape[:-1]):
156            base_matrix[index] = CObs(import_jackknife(matrix[index].real, name, [idl]),
157                                      import_jackknife(matrix[index].imag, name, [idl]))
158        return base_matrix
159
160    for op in operands:
161        if isinstance(op.flat[0], CObs):
162            name = op.flat[0].real.names[0]
163            idl = op.flat[0].real.idl[name]
164            break
165        elif isinstance(op.flat[0], Obs):
166            name = op.flat[0].names[0]
167            idl = op.flat[0].idl[name]
168            break
169
170    conv_operands = []
171    for op in operands:
172        if isinstance(op.flat[0], CObs):
173            conv_operands.append(_exp_to_jack_c(op))
174        elif isinstance(op.flat[0], Obs):
175            conv_operands.append(_exp_to_jack(op))
176        else:
177            conv_operands.append(op)
178
179    tmp_subscripts = ','.join([o + '...' for o in subscripts.split(',')])
180    extended_subscripts = '->'.join([o + '...' for o in tmp_subscripts.split('->')[:-1]] + [tmp_subscripts.split('->')[-1]])
181    einsum_path = np.einsum_path(extended_subscripts, *conv_operands, optimize='optimal')[0]
182    jack_einsum = np.einsum(extended_subscripts, *conv_operands, optimize=einsum_path)
183
184    if jack_einsum.dtype == complex:
185        result = _imp_from_jack_c(jack_einsum, name, idl)
186    elif jack_einsum.dtype == float:
187        result = _imp_from_jack(jack_einsum, name, idl)
188    else:
189        raise Exception("Result has unexpected datatype")
190
191    if result.shape == ():
192        return result.flat[0]
193    else:
194        return result
195
196
197def inv(x):
198    """Inverse of Obs or CObs valued matrices."""
199    return _mat_mat_op(anp.linalg.inv, x)
200
201
202def cholesky(x):
203    """Cholesky decomposition of Obs valued matrices."""
204    if any(isinstance(o, CObs) for o in x.ravel()):
205        raise Exception("Cholesky decomposition is not implemented for CObs.")
206    return _mat_mat_op(anp.linalg.cholesky, x)
207
208
209def det(x):
210    """Determinant of Obs valued matrices."""
211    return _scalar_mat_op(anp.linalg.det, x)
212
213
214def _scalar_mat_op(op, obs, **kwargs):
215    """Computes the matrix to scalar operation op to a given matrix of Obs."""
216    def _mat(x, **kwargs):
217        dim = int(np.sqrt(len(x)))
218
219        mat = []
220        for i in range(dim):
221            row = []
222            for j in range(dim):
223                row.append(x[j + dim * i])
224            mat.append(row)
225
226        return op(anp.array(mat))
227
228    if isinstance(obs, np.ndarray):
229        raveled_obs = (1 * (obs.ravel())).tolist()
230    else:
231        raise TypeError('Unproper type of input.')
232    return derived_observable(_mat, raveled_obs, **kwargs)
233
234
235def _mat_mat_op(op, obs, **kwargs):
236    """Computes the matrix to matrix operation op to a given matrix of Obs."""
237    # Use real representation to calculate matrix operations for complex matrices
238    if any(isinstance(o, CObs) for o in obs.ravel()):
239        A = np.empty_like(obs)
240        B = np.empty_like(obs)
241        for (n, m), entry in np.ndenumerate(obs):
242            if hasattr(entry, 'real') and hasattr(entry, 'imag'):
243                A[n, m] = entry.real
244                B[n, m] = entry.imag
245            else:
246                A[n, m] = entry
247                B[n, m] = 0.0
248        big_matrix = np.block([[A, -B], [B, A]])
249        op_big_matrix = derived_observable(lambda x, **kwargs: op(x), [big_matrix], array_mode=True)[0]
250        dim = op_big_matrix.shape[0]
251        op_A = op_big_matrix[0: dim // 2, 0: dim // 2]
252        op_B = op_big_matrix[dim // 2:, 0: dim // 2]
253        res = np.empty_like(op_A)
254        for (n, m), entry in np.ndenumerate(op_A):
255            res[n, m] = CObs(op_A[n, m], op_B[n, m])
256        return res
257    else:
258        return derived_observable(lambda x, **kwargs: op(x), [obs], array_mode=True)[0]
259
260
261def eigh(obs, **kwargs):
262    """Computes the eigenvalues and eigenvectors of a given hermitian matrix of Obs according to np.linalg.eigh."""
263    w = derived_observable(lambda x, **kwargs: anp.linalg.eigh(x)[0], obs)
264    v = derived_observable(lambda x, **kwargs: anp.linalg.eigh(x)[1], obs)
265    return w, v
266
267
268def eig(obs, **kwargs):
269    """Computes the eigenvalues of a given matrix of Obs according to np.linalg.eig."""
270    w = derived_observable(lambda x, **kwargs: anp.real(anp.linalg.eig(x)[0]), obs)
271    return w
272
273
274def eigv(obs, **kwargs):
275    """Computes the eigenvectors of a given hermitian matrix of Obs according to np.linalg.eigh."""
276    v = derived_observable(lambda x, **kwargs: anp.linalg.eigh(x)[1], obs)
277    return v
278
279
280def pinv(obs, **kwargs):
281    """Computes the Moore-Penrose pseudoinverse of a matrix of Obs."""
282    return derived_observable(lambda x, **kwargs: anp.linalg.pinv(x), obs)
283
284
285def svd(obs, **kwargs):
286    """Computes the singular value decomposition of a matrix of Obs."""
287    u = derived_observable(lambda x, **kwargs: anp.linalg.svd(x, full_matrices=False)[0], obs)
288    s = derived_observable(lambda x, **kwargs: anp.linalg.svd(x, full_matrices=False)[1], obs)
289    vh = derived_observable(lambda x, **kwargs: anp.linalg.svd(x, full_matrices=False)[2], obs)
290    return (u, s, vh)