bjmorgan/bsym

from bsym import Configuration, SymmetryOperation
import numpy as np

class ColourOperation( SymmetryOperation ):
    """
    This class subclasses `SymmetryOperation`.
    It defines a class of object for performing a compound operation on a
    configuration. 
    First, a matrix transform applied to the configuration
    vector space, equivlant to a `SymmetryOperation`. 
    Second, a colour mapping, that allows objects to be replaced in the
    configuration vector space.
    """

    def __init__( self, matrix, colour_mapping, label=None ):
        """
        Initialise a `ColourOperation` object.
        A `ColourOperation` object behaves similarly to a `SymmetryOperation`, but has
        an additional `colour_mapping` attribute, which is a `list` of ``dict`s,
        describing a per-site mapping between objects.

        Args:
            matrix (numpy.matrix|numpy.ndarray|list): 1D vector as either a
            `numpy.matrix`, `numpy.ndarray`, or `list` containing the site mappings
            for this symmetry operation.
            colour_mapping (list[dict]): A `list` of `dict`s, that describe per-site
            object mappings.
            label (default=None) (str): optional string label for this `SymmetryOperation` object.

        Returns:
            None

        Example:

            >>> matrix = np.array( [[1, 0], [0, 1]] )
            >>> colour_mapping = [ { 0: 1, 1: 0 }, { 0: 0, 1: 1 } ]
            >>> ColourOperation( matrix, colour_mapping )
         
        """
        super().__init__( matrix, label )
        self.colour_mapping = colour_mapping
         
    @classmethod
    def from_vector( cls, vector, colour_mapping, count_from_zero=False, label=None ):
        """
        Initialise a ColourOperation from a vector of site mappings.

        Args:
            vector (list): vector of integers defining a symmetry operation mapping.
            colour_mapping (list[dict]): A `list` of `dict`s, that describe per-site object mappings.
            count_from_zero (default = False) (bool): set to True if the site index counts from zero.
            label (default=None) (str): optional string label for this `SymmetryOperation` object.
   
        Returns:
            a new SymmetryOperation object
        """
        if not count_from_zero:
            vector = [ x - 1 for x in vector ]
        dim = len( vector )
        matrix = np.zeros( ( dim, dim ) )
        for index, element in enumerate( vector ):
            matrix[ element, index ] = 1
        new_colour_operation = cls( matrix, colour_mapping=colour_mapping, label=label )
        return new_colour_operation
    
    def operate_on( self, configuration ):
        """
        Return the Configuration generated by appliying this colour operation

        Args:
            configuration (Configuration): the configuration / occupation vector to operate on.

        Returns:
            (Configuration): the new configuration.
        """
        if not isinstance( configuration, Configuration ):
            raise TypeError
        new_configuration = Configuration( self.matrix.dot( configuration.vector ) )
        return Configuration( [ d[value] for value, d in zip( new_configuration.vector, self.colour_mapping ) ] )
    
    def __mul__( self, other ):
        """
        Operate on another object with this `ColourOperation`.

        Args:
            other (ColourOperation, SymmetryOperation, Configuration): the other object (colour operation, symmetry operation, configuration, or matrix).

        Returns:
            (ColourOperation): a new `ColourOperation` instance with the resultant matrix and colour_mapping..
            (Configuration): if `other` is a `Configuration`.
        """
        if isinstance( other, ColourOperation ):
            new_matrix = self.matrix.dot( other.matrix )
            new_mapping = [ {} for d in self.colour_mapping ]
            for i, (this_mapping, other_mapping) in enumerate( zip( self.colour_mapping, other.colour_mapping ) ):
                for key in this_mapping.keys():
                    new_mapping[i][key] = this_mapping[ other_mapping[ key ] ]
            return ColourOperation( new_matrix, colour_mapping=new_mapping ) 
        elif isinstance( other, SymmetryOperation ):
            return ColourOperation( self.matrix.dot( other.matrix ), colour_mapping=self.colour_mapping )
        elif isinstance( other, Configuration ):
            return self.operate_on( other )
        else:
            print( self.__class__, other.__class__ )
            raise TypeError
           
    def invert( self, label=None ):
        # TODO
        raise NotImplementedError
 
    def __repr__( self ):
        label = self.label if self.label else '---'
        return 'ColourOperation\nlabel(' + label + ")\n" + "\n".join( [ row.__str__() + ' ' + mapping.__repr__() for row, mapping in zip( self.matrix, self.colour_mapping ) ] )