eigenfunctions.py 文件源码

def return_real_part(to_return):
    """
    Check if the imaginary part of to_return vanishes
    and return the real part
    :param to_return:
    :return:
    """
    if not isinstance(to_return, (Number, list, np.ndarray)):
        raise TypeError
    if isinstance(to_return, (list, np.ndarray)):
        if not all([isinstance(num, Number) for num in to_return]):
            raise TypeError

    maybe_real = np.atleast_1d(np.real_if_close(to_return))

    if maybe_real.dtype == 'complex':
        raise ValueError("Something goes wrong, imaginary part does not vanish")
    else:
        if maybe_real.shape == (1,):
            maybe_real = maybe_real[0]
        return maybe_real