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