Source code for index_mapper

from itertools import count
from typing import Iterable, Dict, List

[docs]class IndexMapper: """ This class maps between the different atoms objects in a MAZE-sim project. It is essential for keeping track of the various """ id = 0
[docs] @staticmethod def get_id(): """ Get a unique id :return: """ += 1 return str(
[docs] @staticmethod def get_unique_name(name: str): """ Get a unique name :param name: name :return: name + _ + unique id number """ return name + '_' + IndexMapper.get_id()
[docs] def __init__(self, atom_indices: Iterable) -> None: """ :param atom_indices: A list of atom indices from a Zeotype """ self.main_index = {} self.i_max = 0 self.names = ['parent'] for main_index, atom_index in zip(count(), atom_indices): self.main_index[main_index] = {'parent': atom_index} self.i_max = main_index
[docs] def get_reverse_main_index(self, name: str) -> Dict[int, int]: """ Reverses an index so that the name indices are used as a key for the main index :param name: name of the item to make the key :return: a reverse index map where the indices of name are the key """ return self._reverse_main_index(name)
[docs] def _reverse_main_index(self, name: str) -> Dict[int, int]: """ Internal function for making a reverse index map :param name: name of the item to make the key :return:a reverse index map where the indices of name are the key """ name_to_main_dict = {} for main_index, value in self.main_index.items(): try: name_index = value[name] except KeyError: print(f'name not found at index {main_index}') continue if name_index is None: # ignore none indices continue name_to_main_dict[name_index] = main_index return name_to_main_dict
[docs] def get_name1_to_name2_map(self, name1: str, name2: str) -> Dict[int, int]: """ Gets a map between the indices of name1 and the indices of name2 :param name1: name whose indices are the key in the map :param name2: name whose indices are the value in the map :return: name1.index -> name2.index map """ name1_to_name2_map = {} for row in self.main_index.values(): if row[name1] is not None: name1_to_name2_map[row[name1]] = row[name2] return name1_to_name2_map
[docs] def register_with_main(self, new_name: str, main_to_new_map: Dict[int, int]) -> None: """ Register a new object by using the main index in mapping. :param new_name: name of the new object being registered :param main_to_new_map: :return: """ self.names.append(new_name) for main_i in self.main_index.keys(): self.main_index[main_i][new_name] = main_to_new_map.get(main_i, None)
[docs] def register(self, old_name: str, new_name: str, old_to_new_map: Dict[int, int]) -> None: """ Register a new object with the indexer :param old_name: name of object known by the indexer :param new_name: name of new object being registered :param old_to_new_map: a index mapping between the old map and the new map note that additional atoms in new will be added to the index :return: """ already_registered = f'Error: {new_name} has already been registered' not_registered = f'Error: {old_name} has not been registered' if new_name in self.names: print(f'old name is {old_name}, new name is {new_name}') raise ValueError(already_registered) assert old_name in self.names, not_registered self.names.append(new_name) old_name_to_main_dict = self._reverse_main_index(old_name) main_to_new_name_dict = {} for old_ind, main_ind in old_name_to_main_dict.items(): main_to_new_name_dict[main_ind] = old_to_new_map.get(old_ind, None) for i in self.main_index.keys(): self.main_index[i][new_name] = main_to_new_name_dict.get(i, None)
[docs] def _make_none_dict(self) -> Dict[str, None]: """ Get a dictionary full of None for each name :return: A dictionary full of None for each known name """ none_dict = {} for name in self.names: none_dict[name] = None return none_dict
[docs] def add_atoms(self, name: str, new_atom_indices: Iterable[int]) -> None: """ Add new atoms to the dictionary :param name: name of object with new atoms :param new_atom_indices: indices of the new atoms :return: """ assert name not in self.names, 'name already exists' for index, value in self.main_index.items(): value[name] = None self.names.append(name) for index in new_atom_indices: none_dict = self._make_none_dict() # could be slow none_dict[name] = index self.i_max += 1 self.main_index[self.i_max] = none_dict
[docs] def pop(self, name: str, atom_index_to_delete: int) -> int: """ Deletes a single atom index :param name: name of Zeotype to delete atom :type name: str :param atom_index_to_delete: index to delete (using own mapping) :type atom_index_to_delete: int :return: index deleted :rtype: int """ name_to_main_dict = self._reverse_main_index(name) name_indices = list(name_to_main_dict.keys()) name_indices.sort() for i in name_indices: if i == atom_index_to_delete: self.main_index[name_to_main_dict[i]][name] = None elif i > atom_index_to_delete: self.main_index[name_to_main_dict[i]][name] -= 1 return atom_index_to_delete
[docs] def extend(self, name: str, new_atom_indices: Iterable[int]) -> None: """ This adds additional atom indices to the zeolite :param name: name to add indices to :type name: str :param new_atom_indices: list of indices :type new_atom_indices: Iterable[int] :return: None :rtype: None """ assert name in self.names, 'name not in index mapper' for index in new_atom_indices: none_dict = self._make_none_dict() # could be slow none_dict[name] = index self.i_max += 1 self.main_index[self.i_max] = none_dict
[docs] def delete_atoms(self, name: str, atom_indices_to_delete: Iterable[int]) -> None: """ :param name: name of with indices to delete :param atom_indices_to_delete: indices to delete :return: None """ name_to_main_dict = self._reverse_main_index(name) for i in atom_indices_to_delete: self.main_index[name_to_main_dict[i]][name] = None
[docs] def get_index(self, sender_name: str, receiver_name: str, sender_index: int) -> int: """ get the index of another object :param sender_name: name of the sender zeolite :param receiver_name: the name of the receiving zeolite :param sender_index: the index of the sender :return: the receiving index """ for name_dict in self.main_index.values(): my_name_dict = name_dict[sender_name] my_sender_index = sender_index if name_dict[sender_name] == sender_index: return name_dict[receiver_name]
[docs] def delete_name(self, name: str) -> None: """ Delete zeolite from the index :param name: The name of the zeolite to delete from the index :return: None """ try: self.names.remove(name) except: print(name) # TODO: Write custom error message return None for index, old_row in self.main_index.items(): new_row = self._make_none_dict() for key in new_row.keys(): new_row[key] = old_row[key] self.main_index[index] = new_row
[docs] def get_overlap(self, name1: str, name2: str) -> List[int]: """ Get the list of names that overlap :param name1: name of object 1 :param name2: name of object 2 :return: overlapping indices """ overlap_indices_name1 = [] for name_dict in self.main_index.values(): if name_dict[name1] is not None and name_dict[name2] is not None: overlap_indices_name1.append(name_dict[name1]) return overlap_indices_name1
[docs] def reregister_parent(self, main_to_parent_map) -> None: """ Register a Zeolite as the parent using a main_to_parent_map This is useful when building a Perfect Zeolite from a file. :param main_to_parent_map: dict that maps between parent and the main indices :type main_to_parent_map: Dict[int, int] :return: None :rtype: None """ for key, value in self.main_index.items(): if key in main_to_parent_map: value['parent'] = main_to_parent_map[key] else: value['parent'] = None