<commit_before>#!/usr/bin/env python

'''
This module implements various transmuter classes.
Transmuters are essentially classes that generate TransformedStructures from
various data sources. They enable the high-throughput generation of new
structures and input files.

It also includes the helper function, batch_write_vasp_input to generate an 
entire directory of vasp input files for running.
'''

from __future__ import division

__author__ = "<NAME>, <NAME>"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "0.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__date__ = "Mar 4, 2012"

import os
import re
import collections

from pymatgen.alchemy.materials import TransformedStructure


class StandardTransmuter(object):
    """
    An example of a Transmuter object, which performs a sequence of
    transformations on many structures to generate TransformedStructures.
    """

    def __init__(self, transformed_structures, transformations=[],
                 extend_collection=0):
        """
        Args:
            transformed_structures:
                Input transformed structures
            transformations:
                New transformations to be applied to all structures
            extend_collection:
                Whether to use more than one output structure from one-to-many
                transformations. extend_collection can be a number, which
                determines the maximum branching for each transformation.
        """
        self._transformed_structures = transformed_structures
        for trans in transformations:
            self.append_transformation(trans, extend_collection=extend_collection)

    def get_transformed_structures(self):
        """
        Returns all TransformedStructures.
        """
        return self._transformed_structures

    def __getitem__(self, index):
        return self._transformed_structures[index]

    def __getattr__(self, name):
        return [getattr(x, name) for x in self._transformed_structures]

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.
        
        Raises:
            IndexError if already at the oldest change.
        """
        for x in self._transformed_structures:
            x.undo_last_transformation()

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.
        
        Raises:
            IndexError if already at the latest change.
        """
        for x in self._transformed_structures:
            x.redo_next_transformation()

    def __len__(self):
        return len(self._transformed_structures)

    def append_transformation(self, transformation, extend_collection=False,
                              clear_redo=True):
        """
        Appends a transformation to all TransformedStructures.
        
        Args:
            transformation:
                Transformation to append
            extend_collection:
                Whether to use more than one output structure from one-to-many
                transformations. extend_collection can be a number, which
                determines the maximum branching for each transformation.
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the 
                redo list should not be cleared to allow multiple redos.
                
        Returns:
            List of booleans corresponding to initial transformed structures
            each boolean describes whether the transformation altered the 
            structure
        """
        new_structures = []

        for x in self._transformed_structures:
            new = x.append_transformation(transformation, return_alternatives=extend_collection)
            if new is not None:
                new_structures.extend(new)

        self._transformed_structures.extend(new_structures)

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.
        
        Args:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def write_vasp_input(self, vasp_input_set, output_dir,
                         create_directory=True, subfolder=None):
        """
        Batch write vasp input for a sequence of transformed structures to 
        output_dir, following the format output_dir/{formula}_{number}.
        
        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
            subfolder:
                function to create subdirectory name from transformed_structure.
                eg. lambda x: x.other_parameters['tags'][0] to use the first tag
        """
        batch_write_vasp_input(self._transformed_structures, vasp_input_set, output_dir, create_directory, subfolder)

    def set_parameter(self, key, value):
        for x in self._transformed_structures:
            x.set_parameter(key, value)

    def __str__(self):
        output = ["Current structures"]
        output.append("------------")
        for x in self._transformed_structures:
            output.append(str(x._structures[-1]))
        return "\n".join(output)



class CifTransmuter(StandardTransmuter):

    def __init__(self, cif_string, transformations=[], primitive=True,
                  extend_collection=False):
        '''
        Generates a Transmuter from a cif string, possibly
        containing multiple structures.
        
        Args:
            cif_string:<commit_msg>Added + method to transmuter<commit_after>#!/usr/bin/env python

'''
This module implements various transmuter classes.
Transmuters are essentially classes that generate TransformedStructures from
various data sources. They enable the high-throughput generation of new
structures and input files.

It also includes the helper function, batch_write_vasp_input to generate an 
entire directory of vasp input files for running.
'''

from __future__ import division

__author__ = "<NAME>, <NAME>"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "0.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__date__ = "Mar 4, 2012"

import os
import re
import collections

from pymatgen.alchemy.materials import TransformedStructure


class StandardTransmuter(object):
    """
    An example of a Transmuter object, which performs a sequence of
    transformations on many structures to generate TransformedStructures.
    """

    def __init__(self, transformed_structures, transformations=[],
                 extend_collection=0):
        """
        Args:
            transformed_structures:
                Input transformed structures
            transformations:
                New transformations to be applied to all structures
            extend_collection:
                Whether to use more than one output structure from one-to-many
                transformations. extend_collection can be a number, which
                determines the maximum branching for each transformation.
        """
        self._transformed_structures = transformed_structures
        for trans in transformations:
            self.append_transformation(trans, extend_collection=extend_collection)

    def get_transformed_structures(self):
        """
        Returns all TransformedStructures.
        """
        return self._transformed_structures

    def __getitem__(self, index):
        return self._transformed_structures[index]

    def __getattr__(self, name):
        return [getattr(x, name) for x in self._transformed_structures]

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.
        
        Raises:
            IndexError if already at the oldest change.
        """
        for x in self._transformed_structures:
            x.undo_last_transformation()

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.
        
        Raises:
            IndexError if already at the latest change.
        """
        for x in self._transformed_structures:
            x.redo_next_transformation()

    def __len__(self):
        return len(self._transformed_structures)

    def append_transformation(self, transformation, extend_collection=False,
                              clear_redo=True):
        """
        Appends a transformation to all TransformedStructures.
        
        Args:
            transformation:
                Transformation to append
            extend_collection:
                Whether to use more than one output structure from one-to-many
                transformations. extend_collection can be a number, which
                determines the maximum branching for each transformation.
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the 
                redo list should not be cleared to allow multiple redos.
                
        Returns:
            List of booleans corresponding to initial transformed structures
            each boolean describes whether the transformation altered the 
            structure
        """
        new_structures = []

        for x in self._transformed_structures:
            new = x.append_transformation(transformation, return_alternatives=extend_collection)
            if new is not None:
                new_structures.extend(new)

        self._transformed_structures.extend(new_structures)

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.
        
        Args:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def write_vasp_input(self, vasp_input_set, output_dir,
                         create_directory=True, subfolder=None):
        """
        Batch write vasp input for a sequence of transformed structures to 
        output_dir, following the format output_dir/{formula}_{number}.
        
        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
            subfolder:
                function to create subdirectory name from transformed_structure.
                eg. lambda x: x.other_parameters['tags'][0] to use the first tag
        """
        batch_write_vasp_input(self._transformed_structures, vasp_input_set, output_dir, create_directory, subfolder)

    def set_parameter(self, key, value):
        for x in self._transformed_structures:
            x.set_parameter(key, value)

    def __str__(self):
        output = ["Current structures"]
        output.append("------------")
        for x in self._transformed_structures:
            output.append(str(x._structures[-1]))
        return "\n".join(output)
    
    def __add__(self, b):
        transformed_structures = self._transformed_structures
        transformed_structures.extend(b._transformed_structures)
        return StandardTransmuter(transformed_structures,[])



class CifTransmuter(StandardTransmuter):

    def __init__(self, cif_string, transformations=[], primitive=True,
                  extend_collection=False):
        '''
        Generates a Transmuter from a cif string, possibly
        containing multiple structures.
        
        Args:
            cif_string:
{'fraction_non_alphanumeric': 0.05878552971576227, 'fraction_numerical': 0.0022609819121447027, 'mean_word_length': 2.2540725170782974, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, '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': '4715993', 'n_tokens_mistral': 3057, 'n_tokens_neox': 2840, 'n_words': 960}