""" Extended XYZ support Read/write files in "extended" XYZ format, storing additional per-configuration information as key-value pairs on the XYZ comment line, and additional per-atom properties as extra columns. See http://jrkermode.co.uk/quippy/io.html#extendedxyz for a full description of the Extended XYZ file format. Contributed by James Kermode """ from __future__ import print_function import re import numpy as np from ase.atoms import Atoms from ase.calculators.calculator import all_properties, Calculator from ase.calculators.singlepoint import SinglePointCalculator from ase.parallel import paropen from ase.utils import basestring __all__ = ['read_xyz', 'write_xyz'] PROPERTY_NAME_MAP = {'positions': 'pos', 'numbers': 'Z', 'charges': 'charge', 'symbols': 'species'} REV_PROPERTY_NAME_MAP = dict(zip(PROPERTY_NAME_MAP.values(), PROPERTY_NAME_MAP.keys())) KEY_QUOTED_VALUE = re.compile(r'([A-Za-z_]+[A-Za-z0-9_]*)' + r'\s*=\s*["\{\}]([^"\{\}]+)["\{\}e+-]\s*') KEY_VALUE = re.compile(r'([A-Za-z_]+[A-Za-z0-9_]*)\s*=' + r'\s*([-0-9A-Za-z_.:\[\]()e+-/]+)\s*') KEY_RE = re.compile(r'([A-Za-z_]+[A-Za-z0-9_]*)\s*') UNPROCESSED_KEYS = ['uid'] def key_val_str_to_dict(s): """ Parse strings in the form 'key1=value1 key2="quoted value"' """ d = {} s = s.strip() while True: # Match quoted string first, then fall through to plain key=value m = KEY_QUOTED_VALUE.match(s) if m is None: m = KEY_VALUE.match(s) if m is not None: s = KEY_VALUE.sub('', s, 1) else: # Just a key with no value m = KEY_RE.match(s) if m is not None: s = KEY_RE.sub('', s, 1) else: s = KEY_QUOTED_VALUE.sub('', s, 1) if m is None: break # No more matches key = m.group(1) try: value = m.group(2) except IndexError: # default value is 'T' (True) value = 'T' if key.lower() not in UNPROCESSED_KEYS: # Try to convert to (arrays of) floats, ints try: numvalue = [] for x in value.split(): if x.find('.') == -1: numvalue.append(int(float(x))) else: numvalue.append(float(x)) if len(numvalue) == 1: numvalue = numvalue[0] # Only one number elif len(numvalue) == 9: # special case: 3x3 matrix, fortran ordering numvalue = np.array(numvalue).reshape((3, 3), order='F') else: numvalue = np.array(numvalue) # vector value = numvalue except (ValueError, OverflowError): pass # Parse boolean values: 'T' -> True, 'F' -> False, # 'T T F' -> [True, True, False] if isinstance(value, str): str_to_bool = {'T': True, 'F': False} if len(value.split()) > 1: if all([x in str_to_bool.keys() for x in value.split()]): value = [str_to_bool[x] for x in value.split()] elif value in str_to_bool: value = str_to_bool[value] d[key] = value return d def key_val_dict_to_str(d, sep=' '): """ Convert atoms.info dictionary to extended XYZ string representation """ if len(d) == 0: return '' s = '' type_val_map = {(bool, True): 'T', (bool, False): 'F', (np.bool_, True): 'T', (np.bool_, False): 'F'} s = '' for key in d.keys(): val = d[key] if hasattr(val, '__iter__'): val = np.array(val) val = ' '.join(str(type_val_map.get((type(x), x), x)) for x in val.reshape(val.size, order='F')) val.replace('[', '') val.replace(']', '') else: val = type_val_map.get((type(val), val), val) if val is None: s = s + '%s%s' % (key, sep) elif isinstance(val, basestring) and ' ' in val: s = s + '%s="%s"%s' % (key, val, sep) else: s = s + '%s=%s%s' % (key, str(val), sep) return s.strip() def parse_properties(prop_str): """ Parse extended XYZ properties format string Format is "[NAME:TYPE:NCOLS]...]", e.g. "species:S:1:pos:R:3". NAME is the name of the property. TYPE is one of R, I, S, L for real, integer, string and logical. NCOLS is number of columns for that property. """ properties = {} properties_list = [] dtypes = [] converters = [] fields = prop_str.split(':') def parse_bool(x): """ Parse bool to string """ return {'T': True, 'F': False, 'True': True, 'False': False}.get(x) fmt_map = {'R': ('d', float), 'I': ('i', int), 'S': (object, str), 'L': ('bool', parse_bool)} for name, ptype, cols in zip(fields[::3], fields[1::3], [int(x) for x in fields[2::3]]): if ptype not in ('R', 'I', 'S', 'L'): raise ValueError('Unknown property type: ' + ptype) ase_name = REV_PROPERTY_NAME_MAP.get(name, name) dtype, converter = fmt_map[ptype] if cols == 1: dtypes.append((name, dtype)) converters.append(converter) else: for c in range(cols): dtypes.append((name + str(c), dtype)) converters.append(converter) properties[name] = (ase_name, cols) properties_list.append(name) dtype = np.dtype(dtypes) return properties, properties_list, dtype, converters def read_xyz(fileobj, index=-1): """ Read from a file in Extended XYZ format index is the frame to read, default is last frame (index=-1). """ if isinstance(fileobj, str): fileobj = open(fileobj) if not isinstance(index, int) and not isinstance(index, slice): raise TypeError('Index argument is neither slice nor integer!') # If possible, build a partial index up to the last frame required last_frame = None if isinstance(index, int) and index >= 0: last_frame = index elif isinstance(index, slice): if index.stop is not None and index.stop >= 0: last_frame = index.stop # scan through file to find where the frames start fileobj.seek(0) frames = [] while fileobj: frame_pos = fileobj.tell() line = fileobj.readline() if line.strip() == '': break natoms = int(line) frames.append((frame_pos, natoms)) if last_frame is not None and len(frames) > last_frame: break fileobj.readline() # read comment line for i in range(natoms): fileobj.readline() if isinstance(index, int): if index < 0: tmpsnp = len(frames) + index trbl = range(tmpsnp, tmpsnp + 1, 1) else: trbl = range(index, index + 1, 1) elif isinstance(index, slice): start = index.start stop = index.stop step = index.step if start is None: start = 0 elif start < 0: start = len(frames) + start if step is None: step = 1 if stop is None: stop = len(frames) elif stop < 0: stop = len(frames) + stop trbl = range(start, stop, step) if step < 0: trbl.reverse() for index in trbl: frame_pos, natoms = frames[index] fileobj.seek(frame_pos) # check for consistency with frame index table assert int(fileobj.readline()) == natoms # comment line line = fileobj.readline() info = key_val_str_to_dict(line) pbc = None if 'pbc' in info: pbc = info['pbc'] del info['pbc'] elif 'Lattice' in info: # default pbc for extxyz file containing Lattice # is True in all directions pbc = [True, True, True] cell = None if 'Lattice' in info: # NB: ASE cell is transpose of extended XYZ lattice cell = info['Lattice'].T del info['Lattice'] if 'Properties' not in info: # Default set of properties is atomic symbols and positions only info['Properties'] = 'species:S:1:pos:R:3' properties, names, dtype, convs = parse_properties(info['Properties']) del info['Properties'] data = [] for ln in range(natoms): line = fileobj.readline() vals = line.split() row = tuple([conv(val) for conv, val in zip(convs, vals)]) data.append(row) try: data = np.array(data, dtype) except TypeError: raise IOError('Badly formatted data, ' + 'or end of file reached before end of frame') arrays = {} for name in names: ase_name, cols = properties[name] if cols == 1: value = data[name] else: value = np.vstack([data[name + str(c)] for c in range(cols)]).T arrays[ase_name] = value symbols = None if 'symbols' in arrays: symbols = arrays['symbols'] del arrays['symbols'] numbers = None duplicate_numbers = None if 'numbers' in arrays: if symbols is None: numbers = arrays['numbers'] else: duplicate_numbers = arrays['numbers'] del arrays['numbers'] positions = None if 'positions' in arrays: positions = arrays['positions'] del arrays['positions'] atoms = Atoms(symbols=symbols, positions=positions, numbers=numbers, cell=cell, pbc=pbc, info=info) for name, array in arrays.items(): atoms.new_array(name, array) if duplicate_numbers is not None: atoms.set_atomic_numbers(duplicate_numbers) # Load results of previous calculations into SinglePointCalculator results = {} for key in list(atoms.info.keys()): if key in all_properties: results[key] = atoms.info[key] # special case for stress- convert to Voigt 6-element form if key.startswith('stress') and results[key].shape == (3, 3): stress = results[key] stress = np.array([stress[0, 0], stress[1, 1], stress[2, 2], stress[1, 2], stress[0, 2], stress[0, 1]]) results[key] = stress del atoms.info[key] for key in list(atoms.arrays.keys()): if key in all_properties: results[key] = atoms.arrays[key] del atoms.arrays[key] if results != {}: calculator = SinglePointCalculator(atoms, **results) atoms.set_calculator(calculator) yield atoms def output_column_format(atoms, columns, arrays, write_info=True, results=None): """ Helper function to build extended XYZ comment line """ fmt_map = {'d': ('R', '%16.8f '), 'f': ('R', '%16.8f '), 'i': ('I', '%8d '), 'O': ('S', '%s'), 'S': ('S', '%s'), 'U': ('S', '%s'), 'b': ('L', ' %.1s ')} # NB: Lattice is stored as tranpose of ASE cell, # with Fortran array ordering lattice_str = ('Lattice="' + ' '.join([str(x) for x in np.reshape(atoms.cell.T, 9, order='F')]) + '"') property_names = [] property_types = [] property_ncols = [] dtypes = [] formats = [] for column in columns: array = arrays[column] dtype = array.dtype property_name = PROPERTY_NAME_MAP.get(column, column) property_type, fmt = fmt_map[dtype.kind] property_names.append(property_name) property_types.append(property_type) if (len(array.shape) == 1 or (len(array.shape) == 2 and array.shape[1] == 1)): ncol = 1 dtypes.append((column, dtype)) else: ncol = array.shape[1] for c in range(ncol): dtypes.append((column + str(c), dtype)) formats.extend([fmt] * ncol) property_ncols.append(ncol) props_str = ':'.join([':'.join(x) for x in zip(property_names, property_types, [str(nc) for nc in property_ncols])]) comment_str = lattice_str + ' Properties=' + props_str info = {} if write_info: info.update(atoms.info) if results is not None: info.update(results) info['pbc'] = atoms.get_pbc() # always save periodic boundary conditions comment_str += ' ' + key_val_dict_to_str(info) dtype = np.dtype(dtypes) fmt = ''.join(formats) + '\n' return comment_str, property_ncols, dtype, fmt def write_xyz(fileobj, images, comment='', columns=None, write_info=True, write_results=True, append=False): """ Write output in extended XYZ format Optionally, specify which columns (arrays) to include in output, and whether to write the contents of the Atoms.info dict to the XYZ comment line (default is True) and the results of any calculator attached to this Atoms. """ if isinstance(fileobj, str): fileobj = paropen(fileobj, 'w') if not isinstance(images, (list, tuple)): images = [images] for atoms in images: natoms = len(atoms) if columns is None: fr_cols = None else: fr_cols = columns[:] if fr_cols is None: fr_cols = (['symbols', 'positions'] + [key for key in atoms.arrays.keys() if key not in ['symbols', 'positions', 'species', 'pos']]) per_frame_results = {} per_atom_results = {} if write_results: calculator = atoms.get_calculator() if (calculator is not None and isinstance(calculator, Calculator)): for key in all_properties: value = calculator.results.get(key, None) if value is None: # skip missing calculator results continue if (isinstance(value, np.ndarray) and value.shape[0] == len(atoms)): # per-atom quantities (forces, energies, stresses) per_atom_results[key] = value else: # per-frame quantities (energy, stress) # special case for stress, which should be converted # to 3x3 matrices before writing if key.startswith('stress'): xx, yy, zz, yz, xz, xy = value value = np.array([(xx, xy, xz), (xy, yy, yz), (xz, yz, zz)]) per_frame_results[key] = value # Move symbols and positions to first two properties if 'symbols' in fr_cols: i = fr_cols.index('symbols') fr_cols[0], fr_cols[i] = fr_cols[i], fr_cols[0] if 'positions' in fr_cols: i = fr_cols.index('positions') fr_cols[1], fr_cols[i] = fr_cols[i], fr_cols[1] # Check first column "looks like" atomic symbols if fr_cols[0] in atoms.arrays: symbols = atoms.arrays[fr_cols[0]] else: symbols = atoms.get_chemical_symbols() if not isinstance(symbols[0], basestring): raise ValueError('First column must be symbols-like') # Check second column "looks like" atomic positions pos = atoms.arrays[fr_cols[1]] if pos.shape != (natoms, 3) or pos.dtype.kind != 'f': raise ValueError('Second column must be position-like') # Collect data to be written out arrays = {} for column in fr_cols: if column in atoms.arrays: arrays[column] = atoms.arrays[column] elif column == 'symbols': arrays[column] = np.array(symbols) else: raise ValueError('Missing array "%s"' % column) if write_results: fr_cols += per_atom_results.keys() arrays.update(per_atom_results) comm, ncols, dtype, fmt = output_column_format(atoms, fr_cols, arrays, write_info, per_frame_results) if comment != '': # override key/value pairs with user-speficied comment string comm = comment # Pack fr_cols into record array data = np.zeros(natoms, dtype) for column, ncol in zip(fr_cols, ncols): value = arrays[column] if ncol == 1: data[column] = np.squeeze(value) else: for c in range(ncol): data[column + str(c)] = value[:, c] # Write the output fileobj.write('%d\n' % natoms) fileobj.write('%s\n' % comm) for i in range(natoms): fileobj.write(fmt % tuple(data[i])) # create aliases for read/write functions read_extxyz = read_xyz write_extxyz = write_xyz