pyerrors.linalg

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