Skip to main content
`Ctrl`+`K`

#

# Source code for ase.io.cfg

    # fmt: off

    import numpy as np

    import ase
    from ase.data import chemical_symbols
    from ase.utils import reader, writer

    cfg_default_fields = np.array(['positions', 'momenta', 'numbers', 'magmoms'])

    @writer
    def write_cfg(fd, atoms):
        """Write atomic configuration to a CFG-file (native AtomEye format).
           See: http://mt.seas.upenn.edu/Archive/Graphics/A/
        """

        fd.write('Number of particles = %i\n' % len(atoms))
        fd.write('A = 1.0 Angstrom\n')
        cell = atoms.get_cell(complete=True)
        for i in range(3):
            for j in range(3):
                fd.write('H0(%1.1i,%1.1i) = %f A\n' % (i + 1, j + 1, cell[i, j]))

        entry_count = 3
        for x in atoms.arrays.keys():
            if x not in cfg_default_fields:
                if len(atoms.get_array(x).shape) == 1:
                    entry_count += 1
                else:
                    entry_count += atoms.get_array(x).shape[1]

        vels = atoms.get_velocities()
        if isinstance(vels, np.ndarray):
            entry_count += 3
        else:
            fd.write('.NO_VELOCITY.\n')

        fd.write('entry_count = %i\n' % entry_count)

        i = 0
        for name, aux in atoms.arrays.items():
            if name not in cfg_default_fields:
                if len(aux.shape) == 1:
                    fd.write('auxiliary[%i] = %s [a.u.]\n' % (i, name))
                    i += 1
                else:
                    if aux.shape[1] == 3:
                        for j in range(3):
                            fd.write('auxiliary[%i] = %s_%s [a.u.]\n' %
                                     (i, name, chr(ord('x') + j)))
                            i += 1

                    else:
                        for j in range(aux.shape[1]):
                            fd.write('auxiliary[%i] = %s_%1.1i [a.u.]\n' %
                                     (i, name, j))
                            i += 1

        # Distinct elements
        spos = atoms.get_scaled_positions()
        for i in atoms:
            el = i.symbol

            fd.write('%f\n' % ase.data.atomic_masses[chemical_symbols.index(el)])
            fd.write(f'{el}\n')

            x, y, z = spos[i.index, :]
            s = f'{x:e} {y:e} {z:e} '

            if isinstance(vels, np.ndarray):
                vx, vy, vz = vels[i.index, :]
                s = s + f' {vx:e} {vy:e} {vz:e} '

            for name, aux in atoms.arrays.items():
                if name not in cfg_default_fields:
                    if len(aux.shape) == 1:
                        s += ' %e' % aux[i.index]
                    else:
                        s += (aux.shape[1] * ' %e') % tuple(aux[i.index].tolist())

            fd.write(f'{s}\n')

    default_color = {
        'H': [0.800, 0.800, 0.800],
        'C': [0.350, 0.350, 0.350],
        'O': [0.800, 0.200, 0.200]}

    default_radius = {'H': 0.435, 'C': 0.655, 'O': 0.730}

    def write_clr(fd, atoms):
        """Write extra color and radius code to a CLR-file (for use with AtomEye).
           Hit F12 in AtomEye to use.
           See: http://mt.seas.upenn.edu/Archive/Graphics/A/
        """
        color = None
        radius = None
        if atoms.has('color'):
            color = atoms.get_array('color')
        if atoms.has('radius'):
            radius = atoms.get_array('radius')

        if color is None:
            color = np.zeros([len(atoms), 3], dtype=float)
            for a in atoms:
                color[a.index, :] = default_color[a.symbol]

        if radius is None:
            radius = np.zeros(len(atoms), dtype=float)
            for a in atoms:
                radius[a.index] = default_radius[a.symbol]

        radius.shape = (-1, 1)

        for c1, c2, c3, r in np.append(color, radius, axis=1):
            fd.write(f'{c1:f} {c2:f} {c3:f} {r:f}\n')

    @reader
    def read_cfg(fd):
        """Read atomic configuration from a CFG-file (native AtomEye format).
           See: http://mt.seas.upenn.edu/Archive/Graphics/A/
        """
        nat = None
        naux = 0
        aux = None
        auxstrs = None

        cell = np.zeros([3, 3])
        transform = np.eye(3)
        eta = np.zeros([3, 3])

        current_atom = 0
        current_symbol = None
        current_mass = None

        L = fd.readline()
        while L:
            L = L.strip()
            if len(L) != 0 and not L.startswith('#'):
                if L == '.NO_VELOCITY.':
                    vels = None
                    naux += 3
                else:
                    s = L.split('=')
                    if len(s) == 2:
                        key, value = s
                        key = key.strip()
                        value = [x.strip() for x in value.split()]
                        if key == 'Number of particles':
                            nat = int(value[0])
                            spos = np.zeros([nat, 3])
                            masses = np.zeros(nat)
                            syms = [''] * nat
                            vels = np.zeros([nat, 3])
                            if naux > 0:
                                aux = np.zeros([nat, naux])
                        elif key == 'A':
                            pass  # unit = float(value[0])
                        elif key == 'entry_count':
                            naux += int(value[0]) - 6
                            auxstrs = [''] * naux
                            if nat is not None:
                                aux = np.zeros([nat, naux])
                        elif key.startswith('H0('):
                            i, j = (int(x) for x in key[3:-1].split(','))
                            cell[i - 1, j - 1] = float(value[0])
                        elif key.startswith('Transform('):
                            i, j = (int(x) for x in key[10:-1].split(','))
                            transform[i - 1, j - 1] = float(value[0])
                        elif key.startswith('eta('):
                            i, j = (int(x) for x in key[4:-1].split(','))
                            eta[i - 1, j - 1] = float(value[0])
                        elif key.startswith('auxiliary['):
                            i = int(key[10:-1])
                            auxstrs[i] = value[0]
                    else:
                        # Everything else must be particle data.
                        # First check if current line contains an element mass or
                        # name. Then we have an extended XYZ format.
                        s = [x.strip() for x in L.split()]
                        if len(s) == 1:
                            if L in chemical_symbols:
                                current_symbol = L
                            else:
                                current_mass = float(L)
                        elif current_symbol is None and current_mass is None:
                            # Standard CFG format
                            masses[current_atom] = float(s[0])
                            syms[current_atom] = s[1]
                            spos[current_atom, :] = [float(x) for x in s[2:5]]
                            vels[current_atom, :] = [float(x) for x in s[5:8]]
                            current_atom += 1
                        elif (current_symbol is not None and
                              current_mass is not None):
                            # Extended CFG format
                            masses[current_atom] = current_mass
                            syms[current_atom] = current_symbol
                            props = [float(x) for x in s]
                            spos[current_atom, :] = props[0:3]
                            off = 3
                            if vels is not None:
                                off = 6
                                vels[current_atom, :] = props[3:6]
                            aux[current_atom, :] = props[off:]
                            current_atom += 1
            L = fd.readline()

        # Sanity check
        if current_atom != nat:
            raise RuntimeError('Number of atoms reported for CFG file (={}) and '
                               'number of atoms actually read (={}) differ.'
                               .format(nat, current_atom))

        if np.any(eta != 0):
            raise NotImplementedError('eta != 0 not yet implemented for CFG '
                                      'reader.')
        cell = np.dot(cell, transform)

        if vels is None:
            a = ase.Atoms(
                symbols=syms,
                masses=masses,
                scaled_positions=spos,
                cell=cell,
                pbc=True)
        else:
            a = ase.Atoms(
                symbols=syms,
                masses=masses,
                scaled_positions=spos,
                momenta=masses.reshape(-1, 1) * vels,
                cell=cell,
                pbc=True)

        i = 0
        while i < naux:
            auxstr = auxstrs[i]
            if auxstr[-2:] == '_x':
                a.set_array(auxstr[:-2], aux[:, i:i + 3])
                i += 3
            else:
                a.set_array(auxstr, aux[:, i])
                i += 1

        return a