Source code for edrixs.iostream

__all__ = ['write_tensor', 'write_emat', 'write_umat', 'write_config',
           'read_poles_from_file', 'dump_poles', 'load_poles']

import numpy as np
import json


def write_tensor_1(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}',
                   fmt_float='{:.15f}'):
    (n1, ) = tensor.shape
    is_cmplx = False
    if tensor.dtype in (complex, np.complex128):
        is_cmplx = True
    space = "    "
    f = open(fname, 'w')
    for i in range(n1):
        if only_nonzeros and abs(tensor[i]) < tol:
            continue
        if is_cmplx:
            fmt_string = fmt_int + space + (fmt_float + space) * 2 + '\n'
            f.write(fmt_string.format(i + 1, tensor[i].real, tensor[i].imag))
        else:
            fmt_string = fmt_int + space + fmt_float + space + '\n'
            f.write(fmt_string.format(i + 1, tensor[i]))
    f.close()


def write_tensor_2(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}',
                   fmt_float='{:.15f}'):
    (n1, n2) = tensor.shape
    is_cmplx = False
    if tensor.dtype in (complex, np.complex128):
        is_cmplx = True
    space = "    "
    f = open(fname, 'w')
    for i in range(n1):
        for j in range(n2):
            if only_nonzeros and abs(tensor[i, j]) < tol:
                continue
            if is_cmplx:
                fmt_string = (fmt_int + space) * 2 + (fmt_float + space) * 2 + '\n'
                f.write(fmt_string.format(i + 1, j + 1, tensor[i, j].real, tensor[i, j].imag))
            else:
                fmt_string = (fmt_int + space) * 2 + fmt_float + space + '\n'
                f.write(fmt_string.format(i + 1, j + 1, tensor[i, j]))
    f.close()


def write_tensor_3(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}',
                   fmt_float='{:.15f}'):
    (n1, n2, n3) = tensor.shape
    is_cmplx = False
    if tensor.dtype in (complex, np.complex128):
        is_cmplx = True
    space = "    "
    f = open(fname, 'w')
    for i in range(n1):
        for j in range(n2):
            for k in range(n3):
                if only_nonzeros and abs(tensor[i, j, k]) < tol:
                    continue
                if is_cmplx:
                    fmt_string = (fmt_int + space) * 3 + (fmt_float + space) * 2 + '\n'
                    f.write(fmt_string.format(i + 1, j + 1, k + 1, tensor[i, j, k].real,
                                              tensor[i, j, k].imag))
                else:
                    fmt_string = (fmt_int + space) * 3 + fmt_float + space + '\n'
                    f.write(fmt_string.format(i + 1, j + 1, k + 1, tensor[i, j, k]))
    f.close()


def write_tensor_4(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}',
                   fmt_float='{:.15f}'):
    (n1, n2, n3, n4) = tensor.shape
    is_cmplx = False
    if tensor.dtype in (complex, np.complex128):
        is_cmplx = True
    space = "    "
    f = open(fname, 'w')
    for i in range(n1):
        for j in range(n2):
            for k in range(n3):
                for m in range(n4):
                    if only_nonzeros and abs(tensor[i, j, k, m]) < tol:
                        continue
                    if is_cmplx:
                        fmt_string = (fmt_int + space) * 4 + (fmt_float + space) * 2 + '\n'
                        f.write(fmt_string.format(i + 1, j + 1, k + 1, m + 1,
                                tensor[i, j, k, m].real, tensor[i, j, k, m].imag))
                    else:
                        fmt_string = (fmt_int + space) * 4 + fmt_float + space + '\n'
                        f.write(fmt_string.format(i + 1, j + 1, k + 1, m + 1, tensor[i, j, k, m]))
    f.close()


def write_tensor_5(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}',
                   fmt_float='{:.15f}'):
    (n1, n2, n3, n4, n5) = tensor.shape
    is_cmplx = False
    if tensor.dtype in (complex, np.complex128):
        is_cmplx = True
    space = "    "
    f = open(fname, 'w')
    for i in range(n1):
        for j in range(n2):
            for k in range(n3):
                for r in range(n4):
                    for m in range(n5):
                        if only_nonzeros and abs(tensor[i, j, k, r, m]) < tol:
                            continue
                        if is_cmplx:
                            fmt_string = (fmt_int + space) * 5 + (fmt_float + space) * 2 + '\n'
                            f.write(fmt_string.format(i + 1, j + 1, k + 1, r + 1, m + 1,
                                    tensor[i, j, k, r, m].real, tensor[i, j, k, r, m].imag))
                        else:
                            fmt_string = (fmt_int + space) * 5 + fmt_float + space + '\n'
                            f.write(fmt_string.format(i + 1, j + 1, k + 1, r + 1, m + 1,
                                    tensor[i, j, k, r, m]))
    f.close()


[docs]def write_tensor(tensor, fname, only_nonzeros=False, tol=1E-10, fmt_int='{:10d}', fmt_float='{:.15f}'): """ Write :math:`n` -dimension numpy array to file, currently, :math:`n` can be 1, 2, 3, 4, 5. Parameters ---------- tensor: :math:`n` d float or complex array The array needs to be written. fname: str File name. only_nonzeros: logical (default: False) Only write nonzero elements. tol: float (default: 1E-10) Only write the elements when their absolute value are larger than tol and only_nonzeros=True. fmt_int: str (default: '{:10d}') The format for printing integer numbers. fmt_float: str (default: '{:.15f}') The format for printing float numbers. """ ndim = tensor.ndim if ndim == 1: write_tensor_1(tensor, fname, only_nonzeros=only_nonzeros, tol=tol, fmt_int=fmt_int, fmt_float=fmt_float) elif ndim == 2: write_tensor_2(tensor, fname, only_nonzeros=only_nonzeros, tol=tol, fmt_int=fmt_int, fmt_float=fmt_float) elif ndim == 3: write_tensor_3(tensor, fname, only_nonzeros=only_nonzeros, tol=tol, fmt_int=fmt_int, fmt_float=fmt_float) elif ndim == 4: write_tensor_4(tensor, fname, only_nonzeros=only_nonzeros, tol=tol, fmt_int=fmt_int, fmt_float=fmt_float) elif ndim == 5: write_tensor_5(tensor, fname, only_nonzeros=only_nonzeros, tol=tol, fmt_int=fmt_int, fmt_float=fmt_float) else: raise Exception("error in write_tensor: ndim >5, not implemented !")
[docs]def write_emat(emat, fname, tol=1E-12, fmt_int='{:10d}', fmt_float='{:.15f}'): """ Write the nonzeros of the rank-2 hopping matrices to file. The first line is the number of nonzeros, and the following lines are the nonzero elements. This file will be read by ed.x, xas.x or rixs.x. Parameters ---------- emat: 2d complex array The array to be written. fname: str File name. tol: float Precision. fmt_int: str (default: '{:10d}') Format for printing integer numbers. fmt_float: str (default: '{:.15f}') Format for printing float numbers. """ a1, a2 = np.nonzero(abs(emat) > tol) nonzero = np.stack((a1, a2), axis=-1) space = " " fmt_string = (fmt_int + space) * 2 + (fmt_float + space) * 2 + '\n' f = open(fname, 'w') if len(nonzero) == 0: f.write("{:10d}\n".format(1)) f.write(fmt_string.format(1, 1, 0.0, 0.0)) else: f.write("{:20d}\n".format(len(nonzero))) for i, j in nonzero: f.write(fmt_string.format(i + 1, j + 1, emat[i, j].real, emat[i, j].imag)) f.close()
[docs]def write_umat(umat, fname, tol=1E-12, fmt_int='{:10d}', fmt_float='{:.15f}'): """ Write the nonzeros of the rank-4 Coulomb U tensor to file. The first line is the number of nonzeros, and the following lines are the nonzero elements. This file will be read by ed.x, xas.x or rixs.x. Parameters ---------- umat: 4d complex array The array to be written. fname: str File name. tol: float (default: 1E-12) Precision. fmt_int: str (default: '{:10d}') Format for printing integer numbers. fmt_float: str (default: '{:.15f}') Format for printing float numbers. """ a1, a2, a3, a4 = np.nonzero(abs(umat) > tol) nonzero = np.stack((a1, a2, a3, a4), axis=-1) space = " " fmt_string = (fmt_int + space) * 4 + (fmt_float + space) * 2 + '\n' f = open(fname, 'w') if len(nonzero) == 0: f.write("{:10d}\n".format(1)) f.write(fmt_string.format(1, 1, 1, 1, 0.0, 0.0)) else: f.write("{:20d}\n".format(len(nonzero))) for i, j, k, l in nonzero: f.write(fmt_string.format(i + 1, j + 1, k + 1, l + 1, umat[i, j, k, l].real, umat[i, j, k, l].imag)) f.close()
[docs]def write_config( directory='.', ed_solver=1, num_val_orbs=2, num_core_orbs=2, neval=1, nvector=1, ncv=1, idump=True, num_gs=1, maxiter=500, linsys_max=1000, min_ndim=1000, nkryl=500, eigval_tol=1e-8, linsys_tol=1e-10, omega_in=0.0, gamma_in=0.1 ): """ Write control parameters in config.in file for ed_fsolver. """ if idump: dump_vector = '.true.' else: dump_vector = '.false.' config = [ "&control", "ed_solver=" + str(ed_solver), "num_val_orbs=" + str(num_val_orbs), "num_core_orbs=" + str(num_core_orbs), "neval=" + str(neval), "nvector=" + str(nvector), "ncv=" + str(ncv), "idump=" + str(dump_vector), "num_gs=" + str(num_gs), "maxiter=" + str(maxiter), "linsys_max=" + str(linsys_max), "min_ndim=" + str(min_ndim), "nkryl=" + str(nkryl), "eigval_tol=" + str(eigval_tol), "linsys_tol=" + str(linsys_tol), "omega_in=" + str(omega_in), "gamma_in=" + str(gamma_in), "&end" ] f = open(directory + '/config.in', 'w') for item in config: f.write(item + "\n") f.close()
[docs]def read_poles_from_file(file_list): """ Read informations in files xas_poles.n or rixs_poles.n to a dict. Parameters ---------- file_list: list of strings Names of pole files. pole_dict: dict A dict containing information of poles. """ pole_dict = { 'npoles': [], 'eigval': [], 'norm': [], 'alpha': [], 'beta': [] } for fname in file_list: f = open(fname, 'r') line = f.readline() neff = int(line.strip().split()[1]) pole_dict['npoles'].append(neff) line = f.readline() eigval = float(line.strip().split()[1]) pole_dict['eigval'].append(eigval) line = f.readline() norm = float(line.strip().split()[1]) pole_dict['norm'].append(norm) alpha = [] beta = [] for i in range(neff): line = f.readline() line = line.strip().split() alpha.append(float(line[1])) beta.append(float(line[2])) pole_dict['alpha'].append(alpha) pole_dict['beta'].append(beta) f.close() return pole_dict
[docs]def dump_poles(obj, file_name="poles"): """ Dump the objects of poles returned from XAS or RIXS calculations to file for later plotting. Parameters ---------- obj: Python object Object of poles, a dict or a list of dicts. file_name: string File name. """ with open(file_name+'.json', 'w') as f: json.dump(obj, f, indent=2)
[docs]def load_poles(file_name='poles'): """ Load the objects of poles from file. Parameters ---------- file_name: string Returns ------- obj: Python objects Poles object. """ with open(file_name+'.json', 'r') as f: obj = json.load(f) return obj