source: branches/inundation-numpy-branch/geospatial_data/geospatial_data.py @ 6689

"""Class Geospatial_data - Manipulation of locations on the planet and 
associated attributes.

"""


from anuga.utilities.numerical_tools import ensure_numeric

from numpy import concatenate

from coordinate_transforms.geo_reference import Geo_reference
        
class Geospatial_data:

    def __init__(self,
                 data_points,
                 attributes = None,
                 geo_reference = None,
                 default_attribute_name = None):

        """Create instance from data points and associated attributes


        data_points: x,y coordinates in meters. Type must be either a
        sequence of 2-tuples or an Mx2 Numeric array of floats.

        attributes: Associated values for each data point. The type
        must be either a list or an array of length M or a dictionary
        of lists (or arrays) of length M. In the latter case the keys
        in the dictionary represent the attribute names, in the former
        the attribute will get the default name 'attribute'.
        
        geo_reference: Object representing the origin of the data
        points. It contains UTM zone, easting and northing and data
        points are assumed to be relative to this origin.
        If geo_reference is None, the default geo ref object is used

        default_attribute_name: Name of default attribute to be used with
        get_attribute_values. The idea is that the dataset can be
        equipped with information about which attribute to return.
        If None, the default is the 'first'
        
        """

        self.data_points = ensure_numeric(data_points)
        self.set_attributes(attributes)
        self.set_geo_reference(geo_reference)
        self.set_default_attribute_name(default_attribute_name)
        



    def set_attributes(self, attributes):
        """Check and assign attributes dictionary
        """
        
        from types import DictType
       
        if attributes is None:
            self.attributes = None
            return
        
        if not isinstance(attributes, DictType):
            #Convert single attribute into dictionary
            attributes = {'attribute': attributes}

        #Check input attributes    
        for key in attributes.keys():
            try:
                attributes[key] = ensure_numeric(attributes[key])
            except:
                msg = 'Attribute %s could not be converted' %key
                msg += 'to a numeric vector'
                raise msg

        self.attributes = attributes    


    def set_geo_reference(self, geo_reference):

        from coordinate_transforms.geo_reference import Geo_reference


        if geo_reference is None:
            self.geo_reference = Geo_reference() # Use default
        elif isinstance(geo_reference, Geo_reference):
            self.geo_reference = geo_reference
        else:
            msg = 'Argument geo_reference must be a valid Geo_reference \n'
            msg += 'object or None.'
            raise msg


    def set_default_attribute_name(self, default_attribute_name):
        self.default_attribute_name = default_attribute_name


    def get_geo_reference(self):
        return self.geo_reference
       
    def get_data_points(self, absolute = True):
        if absolute is True:
            xll = self.geo_reference.xllcorner
            yll = self.geo_reference.yllcorner
            return self.data_points + [xll, yll]
            
        else: 
            return self.data_points
    
    def get_attributes(self, attribute_name = None):
        """Return values for one named attribute.

        If attribute_name is None, default_attribute_name is used
        """

        if attribute_name is None:
            if self.default_attribute_name is not None:
                attribute_name = self.default_attribute_name
            else:
                attribute_name = self.attributes.keys()[0] 
                # above line takes the first one from keys
                

        msg = 'Attribute name %s does not exist in data set' %attribute_name
        assert self.attributes.has_key(attribute_name), msg

        return self.attributes[attribute_name]

    def __add__(self,other):
        """ 
        returns the addition of 2 geospatical objects,
        objects are concatencated to the end of each other
            
        NOTE: doesn't add if objects contain different 
        attributes  
        """

        # sets xll and yll as the smallest from self and other
        if self.geo_reference.xllcorner <= other.geo_reference.xllcorner:
            xll = self.geo_reference.xllcorner
        else:
            xll = other.geo_reference.xllcorner

        if self.geo_reference.yllcorner <= other.geo_reference.yllcorner:
            yll = self.geo_reference.yllcorner
        else:
            yll = other.geo_reference.yllcorner

        # find objects zone and checks if the same
        geo_ref1 = self.get_geo_reference()
        zone1 = geo_ref1.get_zone()
        
        geo_ref2 = other.get_geo_reference()
        zone2 = geo_ref2.get_zone()
        
        geo_ref = Geo_reference(zone1, xll, yll)
        
        if zone1 == zone2:  
            a_rel_points = self.get_data_points()
            b_rel_points = other.get_data_points()
            
            # subtracts xll and yll from self's and other's
            # reletive data point and concatenates
            c_points = concatenate((a_rel_points-[xll, yll],
                                    b_rel_points-[xll, yll]), axis = 0)

            new_dictionary = {}
            for x in self.attributes.keys():
                if other.attributes.has_key(x):

                    a_attrib = self.attributes[x]
                    b_attrib = other.attributes[x]
                    new_dictionary[x] = concatenate((a_attrib, b_attrib))

                else:
                    msg = 'Both geospatial_data objects must have the same \n'
                    msg += 'attributes to allow addition.'
                    raise msg
 
            return Geospatial_data(c_points, new_dictionary, geo_ref)
            
        else:
            msg = 'Both geospatial_data objects must be in the same \
            ZONE to allow addition.'
            raise msg
      


def geospatial_data2points_dictionary(geospatial_data):
    """Convert geospatial data to points_dictionary
    """

    points_dictionary = {}
    points_dictionary['pointlist'] = geospatial_data.data_points

    points_dictionary['attributelist'] = {}

    for attribute_name in geospatial_data.attributes.keys():
        val = geospatial_data.attributes[attribute_name]
        points_dictionary['attributelist'][attribute_name] = val

    points_dictionary['geo_reference'] = geospatial_data.geo_reference

    return points_dictionary

    
    




def points_dictionary2geospatial_data(points_dictionary):
    """Convert points_dictionary to geospatial data object
    """

    msg = 'Points dictionary must have key pointlist' 
    assert points_dictionary.has_key('pointlist'), msg

    msg = 'Points dictionary must have key attributelist'     
    assert points_dictionary.has_key('attributelist'), msg        

    if points_dictionary.has_key('geo_reference'):
        geo = points_dictionary['geo_reference']
    else:
        geo = None
    
    return Geospatial_data(points_dictionary['pointlist'],
                           points_dictionary['attributelist'],
                           geo_reference = geo)