<commit_before>import numpy as np
import itertools


class Structure:

    def __init__(self,
                 positions=None,
                 scaled_positions=None,
                 masses=None,
                 cell=None,
                 force_sets=None,
                 force_constants=None,
                 atomic_numbers=None,
                 atomic_elements=None,
                 atom_type_index=None,
                 primitive_matrix=None):

        """
        :param positions: atoms cartesian positions (array Ndim x Natoms)
        :param scaled_positions: atom positions scaled to 1 (array Ndim x Natoms)
        :param masses: masses of the atoms (vector NAtoms)
        :param cell: Numpy array containing the unit cell (lattice vectors in rows)
        :param force_sets: force constants: Harmonic force constants
        :param force_constants: atomic numbers vector (1x Natoms):
        :param atomic_numbers: number of total atoms in the crystal:
        :param atomic_elements: atomic names of each element (ex: H, Be, Si,..) (vector Natoms):
        :param atom_type_index: index vector that contains the number of different types of atoms in crystal (vector NdiferentAtoms)
        :param primitive_matrix: matrix that defines the primitive cell respect to the unicell
        """

        self._cell = np.array(cell, dtype='double')
        self._masses = np.array(masses, dtype='double')
        self._atomic_numbers = np.array(atomic_numbers, dtype='double')

        self._force_constants = force_constants
        self._force_sets = force_sets
        self._atomic_elements = atomic_elements
        self._atom_type_index = atom_type_index
        self._scaled_positions = scaled_positions
        self._positions = positions
        self._primitive_matrix = primitive_matrix

        self._primitive_cell = None
        self._supercell_matrix = None
        self._supercell_phonon = None
        self._supercell_phonon_renormalized = None
        self._number_of_cell_atoms = None
        self._number_of_atoms = None
        self._number_of_atom_types = None
        self._number_of_primitive_atoms = None

        # Get normalized cell from cell
        #self._cell_normalized = (cell.T / np.linalg.norm(cell.T, axis=-1)[:, np.newaxis]).T

        # Get atomic types from masses if available
        if atomic_elements is None and masses is not None:
            self._atomic_elements = []
            for i in masses:
                for j in atom_data:
                    if "{0:.1f}".format(i) == "{0:.1f}".format(j[3]):
                        self._atomic_elements.append(j[1])

        # Get atomic numbers and masses from available information
        if atomic_numbers is None and self._atomic_elements is not None:
            self._atomic_numbers =  np.array([symbol_map[i] for i in self._atomic_elements])
        else:
            self._atomic_numbers = np.array(atomic_numbers)

        if masses is None and self._atomic_numbers is not None:
            self._masses = np.array([atom_data[i][3] for i in self._atomic_numbers])
        else:
            self._masses = masses

    # -------- Methods start here -----------

    # Getting data
    def get_data_from_dict(self, data_dictionary):
        for data in self.__dict__:
            try:
                self.__dict__[data] = data_dictionary[data]
            except KeyError:
                continue

    # Cell related methods
    def set_cell(self, cell):
        self._cell = np.array(cell, dtype='double')

    def get_cell(self, supercell=None):

        if supercell is not None:
            return np.dot(self._cell.T, np.diag(supercell))

        return self._cell

    def get_supercell_phonon(self):

        if self.get_force_constants() is not None:
            supercell_phonon = self.get_force_constants().get_supercell()
        elif self.get_force_sets() is not None:
            supercell_phonon = self.get_force_sets().get_supercell()
        else:
            supercell_phonon = np.identity(self.get_number_of_dimensions(), dtype=int)
        return supercell_phonon

    def set_supercell_matrix(self, supercell_matrix):
        self._supercell_matrix = supercell_matrix

    def get_supercell_matrix(self):
        if self._supercell_matrix is None:
            self._supercell_matrix = np.array(self.get_number_of_dimensions() * [1], dtype=int)
        return self._supercell_matrix

    def get_primitive_cell(self):
        if self._primitive_cell is None:
            self._primitive_cell = (np.dot(self.get_cell().T, self.get_primitive_matrix()))
        return self._primitive_cell

    # Cell matrix related methods
    def set_primitive_matrix(self,primitive_matrix):
        self._primitive_matrix = primitive_matrix
        self._number_of_atom_types = None
        self._number_of_primitive_atoms = None
        self._atom_type_index = None
        self._primitive_cell = None

    def get_primitive_matrix(self):
        if self._primitive_matrix is None:
            if self._primitive_cell is None:
                print('Warning: No primitive matrix defined! Using unit cell as primitive')
                self._primitive_matrix = np.identity(self.get_number_of_dimensions())
            else:
                self._primitive_matrix = np.dot(np.linalg.inv(self.get_cell()), self._primitive_cell)
        return  self._primitive_matrix

    # Positions related methods
    def set_positions(self, cart_positions):
        self._scaled_positions = np.dot(cart_positions, np.linalg.inv(self.get_cell()))

    def get_positions(self, supercell=None):
        if self._positions is None:
            if self._scaled_positions is None:
                print('No positions provided')
                exit()
            else:
                self._positions = np.dot(self._scaled_positions, self.get_cell())

        if supercell is None:
            supercell = self.get_number_of_dimensions() * [1]

        position_supercell = []
        for k in range(self._positions.shape[0]):
            for r in itertools.product(*[range (i) for i in supercell[::-1]]):
                position_supercell.append(self._positions[k,:] + np.dot(np.array(r[::-1]), self.get_cell()))
        position_supercell = np.array(position_supercell)

        return position_supercell

    def get_scaled_positions(self, supercell=None):

        if self._scaled_positions is None:
            self._scaled_positions = np.dot(self.get_positions(), np.linalg.inv(self.get_cell()))

        if supercell is not None:
            cell = (np.dot(self.get_cell().T, np.diag(supercell)).T)
            scaled_positions = np.dot(self.get_positions(supercell), np.linalg.inv(cell))
            return scaled_positions

        return self._scaled_positions

    # Force related methods
    def forces_available(self):
        if self.get_force_constants() is not None or self.get_force_sets() is not None:
            return True
        else:
            return False

    def set_force_constants(self, force_constants):
        self._force_constants = force_constants

    def get_force_constants(self):
        return self._force_constants

    def set_force_set(self, force_set):
        self._force_sets = force_set

    def get_force_sets(self):
<commit_msg>clean code<commit_after>import numpy as np
import itertools


class Structure:

    def __init__(self,
                 positions=None,
                 scaled_positions=None,
                 masses=None,
                 cell=None,
                 force_sets=None,
                 force_constants=None,
                 atomic_numbers=None,
                 atomic_elements=None,
                 atom_type_index=None,
                 primitive_matrix=None):

        """
        :param positions: atoms cartesian positions (array Ndim x Natoms)
        :param scaled_positions: atom positions scaled to 1 (array Ndim x Natoms)
        :param masses: masses of the atoms (vector NAtoms)
        :param cell: Numpy array containing the unit cell (lattice vectors in rows)
        :param force_sets: force constants: Harmonic force constants
        :param force_constants: atomic numbers vector (1x Natoms):
        :param atomic_numbers: number of total atoms in the crystal:
        :param atomic_elements: atomic names of each element (ex: H, Be, Si,..) (vector Natoms):
        :param atom_type_index: index vector that contains the number of different types of atoms in crystal (vector NdiferentAtoms)
        :param primitive_matrix: matrix that defines the primitive cell respect to the unicell
        """

        self._cell = np.array(cell, dtype='double')
        self._masses = np.array(masses, dtype='double')
        self._atomic_numbers = np.array(atomic_numbers, dtype='double')

        self._force_constants = force_constants
        self._force_sets = force_sets
        self._atomic_elements = atomic_elements
        self._atom_type_index = atom_type_index
        self._scaled_positions = scaled_positions
        self._positions = positions
        self._primitive_matrix = primitive_matrix

        self._primitive_cell = None
        self._supercell_matrix = None
        self._supercell_phonon = None
        self._supercell_phonon_renormalized = None
        self._number_of_cell_atoms = None
        self._number_of_atoms = None
        self._number_of_atom_types = None
        self._number_of_primitive_atoms = None

        # Get normalized cell from cell
        #self._cell_normalized = (cell.T / np.linalg.norm(cell.T, axis=-1)[:, np.newaxis]).T

        # Get atomic types from masses if available
        if atomic_elements is None and masses is not None:
            self._atomic_elements = []
            for i in masses:
                for j in atom_data:
                    if "{0:.1f}".format(i) == "{0:.1f}".format(j[3]):
                        self._atomic_elements.append(j[1])

        # Get atomic numbers and masses from available information
        if atomic_numbers is None and self._atomic_elements is not None:
            self._atomic_numbers =  np.array([symbol_map[i] for i in self._atomic_elements])
        else:
            self._atomic_numbers = np.array(atomic_numbers)

        if masses is None and self._atomic_numbers is not None:
            self._masses = np.array([atom_data[i][3] for i in self._atomic_numbers])
        else:
            self._masses = masses

    # -------- Methods start here -----------

    # Getting data
    def get_data_from_dict(self, data_dictionary):
        for data in self.__dict__:
            try:
                self.__dict__[data] = data_dictionary[data]
            except KeyError:
                continue

    # Cell related methods
    def set_cell(self, cell):
        self._cell = np.array(cell, dtype='double')

    def get_cell(self, supercell=None):

        if supercell is not None:
            return np.dot(self._cell.T, np.diag(supercell))

        return self._cell

    def get_supercell_phonon(self):

        if self.get_force_constants() is not None:
            supercell_phonon = self.get_force_constants().get_supercell()
        elif self.get_force_sets() is not None:
            supercell_phonon = self.get_force_sets().get_supercell()
        else:
            supercell_phonon = np.identity(self.get_number_of_dimensions(), dtype=int)
        return supercell_phonon

    def set_supercell_matrix(self, supercell_matrix):
        self._supercell_matrix = supercell_matrix

    def get_supercell_matrix(self):
        if self._supercell_matrix is None:
            self._supercell_matrix = np.array(self.get_number_of_dimensions() * [1], dtype=int)
        return self._supercell_matrix

    def get_primitive_cell(self):
        if self._primitive_cell is None:
            self._primitive_cell = (np.dot(self.get_cell().T, self.get_primitive_matrix()))
        return self._primitive_cell

    # Cell matrix related methods
    def set_primitive_matrix(self,primitive_matrix):
        self._primitive_matrix = primitive_matrix
        self._number_of_atom_types = None
        self._number_of_primitive_atoms = None
        self._atom_type_index = None
        self._primitive_cell = None

    def get_primitive_matrix(self):
        if self._primitive_matrix is None:
            if self._primitive_cell is None:
                print('Warning: No primitive matrix defined! Using unit cell as primitive')
                self._primitive_matrix = np.identity(self.get_number_of_dimensions())
            else:
                self._primitive_matrix = np.dot(np.linalg.inv(self.get_cell()), self._primitive_cell)
        return  self._primitive_matrix

    # Positions related methods
    def set_positions(self, cart_positions):
        self._scaled_positions = np.dot(cart_positions, np.linalg.inv(self.get_cell()))

    def get_positions(self, supercell=None):
        if self._positions is None:
            if self._scaled_positions is None:
                print('No positions provided')
                exit()
            else:
                self._positions = np.dot(self._scaled_positions, self.get_cell())

        if supercell is None:
            supercell = self.get_number_of_dimensions() * [1]

        position_supercell = []
        for k in range(self._positions.shape[0]):
            for r in itertools.product(*[range (i) for i in supercell[::-1]]):
                position_supercell.append(self._positions[k,:] + np.dot(np.array(r[::-1]), self.get_cell()))
        position_supercell = np.array(position_supercell)

        return position_supercell

    def get_scaled_positions(self, supercell=None):

        if self._scaled_positions is None:
            self._scaled_positions = np.dot(self.get_positions(), np.linalg.inv(self.get_cell()))

        if supercell is not None:
            cell = np.dot(self.get_cell().T, np.diag(supercell))
            scaled_positions = np.dot(self.get_positions(supercell), np.linalg.inv(cell))
            return scaled_positions

        return self._scaled_positions

    # Force related methods
    def forces_available(self):
        if self.get_force_constants() is not None or self.get_force_sets() is not None:
            return True
        else:
            return False

    def set_force_constants(self, force_constants):
        self._force_constants = force_constants

    def get_force_constants(self):
        return self._force_constants

    def set_force_set(self, force_set):
        self._force_sets = force_set

    def get_force_sets(self):

{'fraction_non_alphanumeric': 0.08451963413795475, 'fraction_numerical': 0.0020631318341242004, 'mean_word_length': 2.393698949824971, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2844968', 'n_tokens_mistral': 4217, 'n_tokens_neox': 4042, 'n_words': 997}