Source code for matador.swaps.swaps

# coding: utf-8
# Distributed under the terms of the MIT License.

""" This file implements atomic swaps through the `AtomicSwapper` class. """

import re
from copy import deepcopy
from matador.utils.print_utils import print_success, print_warning
from matador.utils.chem_utils import get_periodic_table, get_stoich

[docs]class AtomicSwapper: """This class handles the creation of input files from database queries that have swapped atoms. """ def __init__( self, cursor, swap=None, uniq=False, top=None, maintain_num_species=True, debug=False, **kwargs ): """Initialise class with query cursor and arguments. Parameters: cursor (list): cursor of documents to swap. Keyword arguments: swap (str): specification of swaps to perform, e.g. "LiP:KSn" will swap all Li->P and all K->Sn in the cursor. uniq (bool/float): filter documents by similarity with the default sim_tol (True) or the value provided here. top (int): only swap from the first `top` structures in the cursor. maintain_num_species (bool): only perform swaps that maintain the number of species in the structure debug (bool): enable debug output kwargs (dict): dictionary of extra arguments that should be ignored. """ # define some swap macros self.periodic_table = get_periodic_table() self.maintain_num_species = maintain_num_species self.swap_dict_list = None self.swap_args = swap del self.periodic_table["X"] self.template_structure = None self.cursor = list(cursor) if top is not None: self.cursor = self.cursor[:top] if len(self.cursor) == 0: return self.swap_counter = 0 self.parse_swaps(self.swap_args) swap_cursor = [] for doc in self.cursor: docs, counter = self.atomic_swaps(doc) self.swap_counter += counter if counter > 0: swap_cursor.extend(docs) self.cursor = swap_cursor if self.swap_counter > 0: print_success("Performed {} swaps.".format(self.swap_counter)) else: print_warning("No swaps performed.") if uniq: from matador.utils.cursor_utils import filter_unique_structures print("Filtering for unique structures...") filtered_cursor = filter_unique_structures( self.cursor, debug=debug, sim_tol=uniq ) print( "Filtered {} down to {}".format(len(self.cursor), len(filtered_cursor)) ) self.cursor = filtered_cursor
[docs] def parse_swaps(self, swap_args=None): """Parse command line options into valid atomic species swaps. e.g. --swap LiP:NaAs ==> [[['Li'], ['P']], [['Na'], ['P']]. Handles multiple many-to-many swaps, macros for groups of the periodic table, and wildcards. Keyword arguments: swap_args (str): overrides command-line swap args. """ self.swap_pairs = [] if swap_args is None: swap_args = self.swap_args if swap_args is None: raise RuntimeError("No swap arguments passed.") if isinstance(swap_args, str): swap_args = [swap_args.strip()] if len(swap_args) > 1: raise RuntimeError( "Detected whitespace in your input clear it and try again." ) swap_list = swap_args[0].split(":") for swap in swap_list: if len(swap) <= 1: raise RuntimeError("Not enough arguments for swap!") # check is both options are groups if "][" in swap: tmp_list = [x for x in swap.split("][") if x != ""] # check if only first option is group elif swap[0] == "[": tmp_list = [x for x in swap.split("]") if x != ""] # check if only last option is group elif swap[-1] == "]": tmp_list = [x for x in swap.split("[") if x != ""] # check if no groups else: tmp_list = [x for x in re.split(r"([A-Z][a-z]*)", swap) if x != ""] for ind, tmp in enumerate(tmp_list): tmp_list[ind] = self._atoms_to_list(tmp) if len(tmp_list) != 2: raise RuntimeError( "Unable to parse swap! {} should contain only two entries".format( tmp_list ) ) self.swap_pairs.append(tmp_list) self.construct_swap_options()
def _atoms_to_list(self, atom_string): """For a given set of atoms in a string, parse any macros and return a list of options. e.g. '[V' -> [<all group V atoms>], and 'V' -> ['V']. Parameters: atom_string (str): formula string with macros. """ if "[" in atom_string or "]" in atom_string: group = atom_string.replace("[", "") group = group.replace("]", "") if group in self.periodic_table: atom_list = self.periodic_table[group] else: atom_list = group.split(",") else: return [atom_string] return [x.strip() for x in atom_list]
[docs] def construct_swap_options(self): """Iterate over possible combinations of multiple many-to-many swaps and create a dict for each swap. """ self.swap_dict_list = [] from itertools import product for branch in product(*([pair[1] for pair in self.swap_pairs])): self.swap_dict_list.append(dict()) for ind, pair in enumerate(self.swap_pairs): for swap_from in pair[0]: if swap_from != branch[ind]: self.swap_dict_list[-1][swap_from] = branch[ind]
[docs] def atomic_swaps(self, source_doc): """Swap atomic species according to parsed options. Parameters: source_doc (dict): matador doc to swap from. """ new_doc = deepcopy(source_doc) swapped_docs = [] unswapped_num_species = len(set(source_doc["atom_types"])) for swap in self.swap_dict_list: if any(key in source_doc["atom_types"] for key in swap): new_doc["atom_types"] = [ swap.get(atom, atom) for atom in source_doc["atom_types"] ] new_doc["_swapped_stoichiometry"] = get_stoich(source_doc["atom_types"]) new_doc["stoichiometry"] = get_stoich(new_doc["atom_types"]) new_doc["elems"] = set(new_doc["atom_types"]) new_doc["num_species"] = len(new_doc["elems"]) if ( not self.maintain_num_species or new_doc["num_species"] == unswapped_num_species ): swapped_doc = deepcopy(new_doc) swapped_docs.append(swapped_doc) return swapped_docs, len(swapped_docs)