Source code for pyretis.inout.formats.cp2k

# -*- coding: utf-8 -*-
# Copyright (c) 2021, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Some method for interacting with input/output from CP2K.

Important methods defined here

update_cp2k_input (:py:func:`.update_cp2k_input`)
    A method for updating a CP2K input file and creating a new one.

read_cp2k_input (:py:func:`.read_cp2k_input`)
    A method to read a CP2K input file.
import logging
import numpy as np
from import box_matrix_to_list, box_vector_angles
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

[docs]class SectionNode: """A class representing a section in the CP2K input. Attributes ---------- title : string The title of the section parent : string The parent section if this node represents a sub-section. settings : list of strings The setting(s) for this particular node. data : string A section of settings if the node defines several settings. children : set of objects like :py:class:`.SectionNode` A set with the sub-sections of this section. level : integer An integer to remember how far down this node is. E.g. if the level is 2, this node is a sub-sub-section. This is used for printing. parents : list of strings or None A list representing the path from the node to the top section. """
[docs] def __init__(self, title, parent, settings, data=None): """Initialise a node. Parameters ---------- title : string The title of the section. parent : object like :py:class:`.SectionNode` The parent if this section is a sub-section. settings : list of strings The settings defined in this section. data : list of strings, optional A section of settings. """ self.title = title self.parent = parent self.settings = settings if data: = list(data) else: = [] self.children = set() self.level = 0 self.parents = None
[docs] def add_child(self, child): """Add a sub-section to the current section.""" self.children.add(child)
[docs] def get_all_parents(self): """Find the path to the top of the tree.""" parents = [self.title] prev = self.parent while prev is not None: parents.append(prev.title) prev = prev.parent self.parents = parents[::-1]
[docs]def dfs_print(node, visited): """Walk through the nodes and print out text. Parameters ---------- node : object like :py:class:`.SectionNode` The object representing a C2PK section. visited : set of objects like :py:class:`.SectionNode` The set contains the nodes we have already visited. Returns ------- out : list of strings These strings represent the CP2K input file. """ out = [] pre = ' ' * (2 * node.level) if not node.settings: out.append(f'{pre}&{node.title}') else: out.append(f'{pre}&{node.title} {" ".join(node.settings)}') for lines in out.append(f'{pre} {lines}') visited.add(node) for child in node.children: if child not in visited: for lines in dfs_print(child, visited): out.append(lines) out.append(f'{pre}&END {node.title}') return out
[docs]def set_parents(listofnodes): """Set parents for all nodes.""" node_ref = {} def dfs_set(node, vis): """DFS traverse the nodes.""" if node.parents is None: node.get_all_parents() par = '->'.join(node.parents) if par in node_ref: prev = node_ref.pop(par) par1 = f'{par}->{" ".join(prev.settings)}' par2 = f'{par}->{" ".join(node.settings)}' node_ref[par1] = prev node_ref[par2] = node else: node_ref[par] = node vis.add(node) for child in node.children: if child not in visited: dfs_set(child, vis) for nodes in listofnodes: visited = set() dfs_set(nodes, visited) return node_ref
[docs]def read_cp2k_input(filename): """Read a CP2K input file. Parameters ---------- filename : string The file to open and read. Returns ------- nodes : list of objects like :py:class:`.SectionNode` The root section nodes found in the file. """ nodes = [] current_node = None with open(filename, 'r', encoding="utf-8") as infile: for lines in infile: lstrip = lines.strip() if not lstrip: # skip empty lines continue if lstrip.startswith('&'): strip = lstrip[1:].split() if lstrip[1:].lower().startswith('end'): current_node = current_node.parent else: if len(strip) > 1: setts = strip[1:] else: setts = [] new_node = SectionNode(strip[0].upper(), current_node, setts) if current_node is None: nodes.append(new_node) else: new_node.level = current_node.level + 1 current_node.add_child(new_node) current_node = new_node else: if current_node is not None: return nodes
[docs]def _add_node(target, settings, data, nodes, node_ref): """Just add a new node.""" # check if this is a root node: root = target.find('->') == -1 if root: new_node = SectionNode(target, None, settings, data=data) nodes.append(new_node) else: parents = target.split('->') title = parents[-1] par = '->'.join(parents[:-1]) if par not in node_ref: _add_node(par, None, None, nodes, node_ref) parent = node_ref['->'.join(parents[:-1])] new_node = SectionNode(title, parent, settings, data=data) new_node.level = parent.level + 1 parent.add_child(new_node) node_ref[target] = new_node
[docs]def update_node(target, settings, data, node_ref, nodes, replace=False): """Update the given target node. If the node does not exist, it will be created. Parameters ---------- target : string The target node, on form root->section->subsection settings : list of strings The settings for the node. data : list of strings or dict The data for the node. node_ref : dict of :py:class:`.SectionNode` A dict of all nodes in the tree. nodes : list of :py:class:`.SectionNode` The root nodes. replace : boolean, optional If this is True and if the nodes have some data, the already existing data will be ignored. We also assume that the data is already formatted. """ if target not in node_ref: # add node # TODO: remove decommented try-except construction later # try: _add_node(target, settings, data, nodes, node_ref) # except KeyError: # pass return None node = node_ref[target] new_data = [] done = set() if not replace: for line in key = line.split()[0] if key in data: new_data.append(f'{key} {data[key]}') done.add(key) else: new_data.append(line) for key in data: if key in done: continue if data[key] is None: new_data.append(str(key)) else: new_data.append(f'{key} {data[key]}') = list(new_data) else: = list(data) if settings is not None: if replace: node.settings = list(settings) else: node.settings += settings return node
[docs]def remove_node(target, node_ref, root_nodes): """Remove a node (and it's children) from the tree. Parameters ---------- target : string The target node, on form root->section->subcection. node_ref : dict A dict with all the nodes. root_nodes : list of objects like :py:class:`.SectionNode` The root nodes. """ to_del = node_ref.pop(target, None) if to_del is None: pass else: # remove all it's children: visited = set() nodes = [to_del] while nodes: node = nodes.pop() if node not in visited: visited.add(node) for i in node.children: nodes.append(i) # remove the reference to this node from the parent parent = to_del.parent if parent is None: # This is a root node. root_nodes.remove(to_del) else: parent.children.remove(to_del) del to_del for key in visited: _ = node_ref.pop(key, None)
[docs]def update_cp2k_input(template, output, update=None, remove=None): """Read a template input and create a new CP2K input. Parameters ---------- template : string The CP2K input file we use as a template. output : string The CP2K input file we will create. update : dict, optional The settings we will update. remove : list of strings, optional The nodes we will remove. """ nodes = read_cp2k_input(template) node_ref = set_parents(nodes) if update is not None: for target in update: value = update[target] settings = value.get('settings', None) replace = value.get('replace', False) data = value.get('data', []) update_node(target, settings, data, node_ref, nodes, replace=replace) if remove is not None: for nodei in remove: remove_node(nodei, node_ref, nodes) with open(output, 'w', encoding='utf-8') as outf: for i, nodei in enumerate(nodes): vis = set() if i > 0: outf.write('\n') outf.write('\n'.join(dfs_print(nodei, vis))) outf.write('\n')
[docs]def read_box_data(box_data): """Read the box data. Parameters ---------- box_data : list of strings The settings for the SUBSYS->CELL section. Returns ------- out[0] : numpy.array, 1D The box vectors, in the correct order. out[1] : list of booleans The periodic boundary setting for each dimension. """ to_read = {'A': 'vec', 'B': 'vec', 'C': 'vec', 'PERIODIC': 'string', 'ABC': 'vec', 'ALPHA_BETA_GAMMA': 'vec'} data = {} for lines in box_data: for key, val in to_read.items(): keyword = f'{key} ' if lines.startswith(keyword): if val == 'vec': data[key] = [float(i) for i in lines.split()[1:]] elif val == 'string': data[key] = ' '.join(lines.split()[1:]) if all(('A' in data, 'B' in data, 'C' in data)): box_matrix = np.zeros((3, 3)) box_matrix[:, 0] = data['A'] box_matrix[:, 1] = data['B'] box_matrix[:, 2] = data['C'] box = box_matrix_to_list(box_matrix) elif 'ABC' in data: if 'ALPHA_BETA_GAMMA' in data: box_matrix = box_vector_angles( data['ABC'], data['ALPHA_BETA_GAMMA'][0], data['ALPHA_BETA_GAMMA'][1], data['ALPHA_BETA_GAMMA'][2], ) box = box_matrix_to_list(box_matrix) else: box = np.array(data['ABC']) else: box = None periodic = [] periodic_setting = data.get('PERIODIC', 'XYZ') for val in ('X', 'Y', 'Z'): periodic.append(val in periodic_setting.upper()) return box, periodic
[docs]def read_cp2k_energy(energy_file): """Read and return CP2K energies. Parameters ---------- energy_file : string The input file to read. Returns ------- out : dict This dict contains the energy terms read from the CP2K energy file. """ data = np.loadtxt(energy_file) energy = {} for i, key in ((1, 'time'), (2, 'ekin'), (3, 'temp'), (4, 'vpot')): try: energy[key] = data[:, i] except IndexError: logger.warning('Could not read energy term %s from CP2kfile %s', key, energy_file) if 'ekin' in energy and 'vpot' in energy: energy['etot'] = energy['ekin'] + energy['vpot'] return energy
[docs]def read_cp2k_restart(restart_file): """Read some info from a CP2K restart file. Parameters ---------- restart_file : string The file to read. Returns ------- pos : numpy.array The positions. vel : numpy.array The velocities. box_size : numpy.array The box vectors. periodic : list of booleans For each dimension, the list entry is True if periodic boundaries should be applied. """ nodes = read_cp2k_input(restart_file) node_ref = set_parents(nodes) velocity = 'FORCE_EVAL->SUBSYS->VELOCITY' coord = 'FORCE_EVAL->SUBSYS->COORD' cell = 'FORCE_EVAL->SUBSYS->CELL' atoms, pos, vel = [], [], [] for posi, veli in zip(node_ref[coord].data, node_ref[velocity].data): pos_split = posi.split() atoms.append(pos_split[0]) pos.append([float(i) for i in pos_split[1:4]]) vel.append([float(i) for i in veli.split()]) pos = np.array(pos) vel = np.array(vel) box, periodic = read_box_data(node_ref[cell].data) return atoms, pos, vel, box, periodic
[docs]def read_cp2k_box(inputfile): """Read the box from a CP2K file. Parameters ---------- inputfile : string The file we will read from. Returns ------- out[0] : numpy.array The box vectors. out[1] : list of booleans For each dimension, the list entry is True if periodic boundaries should be applied. """ nodes = read_cp2k_input(inputfile) node_ref = set_parents(nodes) try: box, periodic = read_box_data( node_ref['FORCE_EVAL->SUBSYS->CELL'].data ) except KeyError: logger.warning('No CELL found in CP2K file "%s"', inputfile) box = np.array([[100., 0., 0.], [0., 100., 0.], [0., 0., 100.]]) periodic = [True, True, True] return box, periodic