def CP2KOutputReader(fh, module=None, type_map=None, kind_map=None, format=None): # mapping from run type to (default module index, list of available module) run_types = { 'QS': ['QUICKSTEP'], 'QMMM': ['FIST', 'QM/MM', 'QUICKSTEP'], 'MM': ['FIST'] } filename, lines = read_text_file(fh) run_type = cp2k_run_type(cp2k_output=lines) if type_map is None: type_map = {} if kind_map is None: kind_map = {} try: available_modules = run_types[run_type] except KeyError: raise ValueError('Unknown CP2K run type %s' % run_type) if module is None: module = available_modules[0] try: cell_index = available_modules.index(module) except ValueError: raise ValueError("Don't know how to read module %s from file %s" % (module, filename)) cell_lines = [ i for i, line in enumerate(lines) if line.startswith(" CELL| Vector a") ] if cell_lines == []: raise ValueError("Cannot find cell in file %s" % filename) try: cell_line = cell_lines[cell_index] except IndexError: raise ValueError( "Cannot find cell with index %d in file %s for module %s" % (cell_index, filename, module)) lattice = fzeros((3, 3)) for i in [0, 1, 2]: lattice[:, i + 1] = [float(c) for c in lines[cell_line + i].split()[4:7]] try: start_line = lines.index( " MODULE %s: ATOMIC COORDINATES IN angstrom\n" % module) except ValueError: raise ValueError( "Cannot find atomic positions for module %s in file %s" % (module, filename)) kinds = [] species = [] Zs = [] pos = [] masses = [] Zeffs = [] types = [] qeffs = [] for line in lines[start_line + 4:]: if line.strip() == '': break if module == 'FIST': atom, kind, typ, x, y, z, qeff, mass = line.split() types.append(typ) Z = type_map.get(typ, 0) kind = int(kind) if Z == 0: Z = kind_map.get(kind, 0) Zs.append(Z) qeffs.append(float(qeff)) else: atom, kind, sp, Z, x, y, z, Zeff, mass = line.split() species.append(sp) Zs.append(int(Z)) Zeffs.append(float(Zeff)) kinds.append(int(kind)) pos.append([float(x), float(y), float(z)]) masses.append(float(mass)) at = Atoms(n=len(kinds), lattice=lattice) at.pos[...] = farray(pos).T at.set_atoms(Zs) at.add_property('mass', farray(masses) * MASSCONVERT) at.add_property('kind', kinds) if module == 'FIST': at.add_property('type', ' ' * TABLE_STRING_LENGTH) at.add_property('qm', False) at.qm[:] = (at.type.stripstrings() == '_QM_') | (at.type.stripstrings() == '_LNK') at.type[...] = s2a(types, TABLE_STRING_LENGTH) at.add_property('qeff', qeffs) else: at.species[...] = s2a(species, TABLE_STRING_LENGTH) at.add_property('zeff', Zeffs) yield at
def add_property(self, name, value, n_cols=None, overwrite=None, property_type=None): """ Add a new property to this Atoms object. `name` is the name of the new property and `value` should be either a scalar or an array representing the value, which should be either integer, real, logical or string. If a scalar is given for `value` it is copied to every element in the new property. `n_cols` can be specified to create a 2D property from a scalar initial value - the default is 1 which creates a 1D property. If an array is given for `value` it should either have shape (self.n,) for a 1D property or (n_cols,self.n) for a 2D property. In this case `n_cols` is inferred from the shape of the `value` and shouldn't be passed as an argument. If `property_type` is present, then no attempt is made to infer the type from `value`. This is necessary to resolve ambiguity between integer and logical types. If property with the same type is already present then no error occurs.If `overwrite` is true, the value will be overwritten with that given in `value`, otherwise the old value is retained. Here are some examples:: a = Atoms(n=10, lattice=10.0*fidentity(3)) a.add_property('mark', 1) # Scalar integer a.add_property('bool', False) # Scalar logical a.add_property('local_energy', 0.0) # Scalar real a.add_property('force', 0.0, n_cols=3) # Vector real a.add_property('label', '') # Scalar string a.add_property('count', [1,2,3,4,5,6,7,8,9,10]) # From list a.add_property('norm_pos', a.pos.norm()) # From 1D array a.add_property('pos', new_pos) # Overwrite positions with array new_pos # which should have shape (3,10) """ kwargs = {} if n_cols is not None: kwargs['n_cols'] = n_cols if overwrite is not None: kwargs['overwrite'] = overwrite if (isinstance(value, np.ndarray) and value.dtype.kind in ['O', 'S'] and value.shape != (len(self), TABLE_STRING_LENGTH)): value = s2a(value.astype('str'), TABLE_STRING_LENGTH).T if property_type is None: _atoms.Atoms.add_property(self, name, value, **kwargs) else: # override value_ref if property_type is specified new_property = not self.has_property(name) type_to_value_ref = { T_INTEGER_A: 0, T_REAL_A: 0.0, T_CHAR_A: " " * TABLE_STRING_LENGTH, T_LOGICAL_A: False, T_INTEGER_A2: 0, T_REAL_A2: 0.0 } try: value_ref = type_to_value_ref[property_type] except KeyError: raise ValueError('Unknown property_type %d' % property_type) if (hasattr(value, 'shape') and len(value.shape) == 2 and property_type != T_CHAR_A and n_cols is None): kwargs['n_cols'] = value.shape[0] _atoms.Atoms.add_property(self, name, value_ref, **kwargs) if new_property or overwrite: getattr(self, name.lower())[:] = value