Source code for fcdmft.utils.arraymath

import scipy.linalg as sla
import numpy as np


def _is_contiguous(arr):
    return arr.flags['C_CONTIGUOUS'] or arr.flags['F_CONTIGUOUS']




def array_scale(arr, alpha):
    """Scale an array in place by a scalar with BLAS if possible.

    Parameters
    ----------
    arr : array_like
        Array to be scaled.
    alpha : scalar
        Scale factor.

    Returns
    -------
    arr : ndarray
        Scaled array equal to alpha * arr.
    """
    if _is_contiguous(arr):
        scal = sla.get_blas_funcs('scal', (arr,))
        scal(a=alpha, x=arr.reshape(-1))
    else:
        arr *= alpha





def diag_array_view(arr):
    """
    Returns a view of the diagonal of a N-d array.

    Parameters
    ----------
    arr : (..., M, M) array_like
        Input array.

    Returns
    -------
    diag : (..., M) ndarray
        View of the diagonal of the input array.
    """
    if not (arr.ndim >= 2 and arr.shape[-1] == arr.shape[-2]):
        raise ValueError('Input must be a stack of square arrays.')
    diagstride = arr.strides[-1] + arr.strides[-2]
    diagshape = list(arr.shape[:-1])
    diagstrides = list(arr.strides[:-1])
    diagstrides[-1] = diagstride
    return np.lib.stride_tricks.as_strided(arr, shape=diagshape, strides=diagstrides)





def get_id_minus_pi(Pi):
    """Calculate I - Pi in place.

    Parameters
    ----------
    Pi : (M, M) array_like
        Input matrix. Must be C-contiguous.

    Returns
    -------
    id_minus_pi : (M, M) ndarray
        (I - Pi)
    """
    if Pi.flags.f_contiguous:
        Pi_c = Pi.T
    else:
        Pi_c = Pi
    assert Pi_c.flags.c_contiguous
    Pi_diag = diag_array_view(Pi_c)
    Pi_diag -= 1.0
    array_scale(Pi_c, -1.0)
    return Pi





def get_pi_inv(Pi, overwrite_input=False):
    """Calculate (I - Pi)^-1 - I.

    Parameters
    ----------
    Pi : (M, M) array_like
        Input matrix. Must be C-contiguous.

    Returns
    -------
    Pi_inv : (M, M) ndarray
        (I - Pi)^-1 - I.
    """
    assert Pi.flags.c_contiguous
    Pi_diag = diag_array_view(Pi)
    Pi = get_id_minus_pi(Pi)
    Pi_inv = sla.inv(Pi.T, overwrite_a=overwrite_input, check_finite=False).T
    Pi_inv_diag = diag_array_view(Pi_inv)
    Pi_inv_diag -= 1.0
    return Pi_inv




def mkslice(l):
    """
    Try to make a slice from a list of integers.
    If this is not possible, return the input unchanged.

    Parameters
    ----------
    l : slice | list | range | ndarray
        Various ways of representing a list of integer indices

    Returns
    -------
    slice | list | ndarray
        slice if possible, otherwise returns the input unchanged
    """
    # If l is already a slice, return it
    if isinstance(l, slice):
        return l

    if l is None or l == ():
        return slice(None)

    # range to slice is easy
    if isinstance(l, range):
        return slice(l.start, l.stop, l.step)

    if len(l) < 2:
        return l

    strides = np.diff(l)
    if not np.all(strides == strides[0]):
        return l
    else:
        start = l[0]
        stop = l[-1] + strides[0]
        step = strides[0]
        return slice(start, stop, step)