source: inundation/ga/storm_surge/pyvolution/util.py @ 1697

"""This module contains various auxiliary function used by pyvolution.

It is also a clearing house for functions that may later earn a module
of their own.
"""

#FIXME: Move polygon stuff out

def angle(v):
    """Compute angle between e1 (the unit vector in the x-direction)
    and the specified vector
    """

    from math import acos, pi, sqrt
    from Numeric import sum, array

    l = sqrt( sum (array(v)**2))
    v1 = v[0]/l
    v2 = v[1]/l


    theta = acos(v1)
    
    #try:
    #   theta = acos(v1)
    #except ValueError, e:
    #    print 'WARNING (util.py): Angle acos(%s) failed: %s' %(str(v1), e)
    #
    #    #FIXME, hack to avoid acos(1.0) Value error
    #    # why is it happening?
    #    # is it catching something we should avoid?
    #    #FIXME (Ole): When did this happen? We need a unit test to
    #    #reveal this condition or otherwise remove the hack
    #    
    #    s = 1e-6
    #    if (v1+s >  1.0) and (v1-s < 1.0) :
    #        theta = 0.0
    #    elif (v1+s >  -1.0) and (v1-s < -1.0):
    #        theta = 3.1415926535897931
    #    print 'WARNING (util.py): angle v1 is %f, setting acos to %f '\
    #          %(v1, theta)

    if v2 < 0:
        #Quadrant 3 or 4
        theta = 2*pi-theta

    return theta


def anglediff(v0, v1):
    """Compute difference between angle of vector x0, y0 and x1, y1.
    This is used for determining the ordering of vertices,
    e.g. for checking if they are counter clockwise.

    Always return a positive value
    """

    from math import pi

    a0 = angle(v0)
    a1 = angle(v1)

    #Ensure that difference will be positive
    if a0 < a1:
        a0 += 2*pi

    return a0-a1


def mean(x):
    from Numeric import sum
    return sum(x)/len(x)


def point_on_line(x, y, x0, y0, x1, y1):
    """Determine whether a point is on a line segment

    Input: x, y, x0, x0, x1, y1: where
        point is given by x, y
        line is given by (x0, y0) and (x1, y1)

    """

    from Numeric import array, dot, allclose
    from math import sqrt

    a = array([x - x0, y - y0])
    a_normal = array([a[1], -a[0]])

    b = array([x1 - x0, y1 - y0])

    if dot(a_normal, b) == 0:
        #Point is somewhere on the infinite extension of the line

        len_a = sqrt(sum(a**2))
        len_b = sqrt(sum(b**2))
        if dot(a, b) >= 0 and len_a <= len_b:
           return True
        else:
           return False
    else:
      return False

def ensure_numeric(A, typecode = None):
    """Ensure that sequence is a Numeric array.
    Inputs:
        A: Sequance. If A is already a Numeric array it will be returned
                     unaltered
                     If not, an attempt is made to convert it to a Numeric
                     array
        typecode: Numeric type. If specified, use this in the conversion.
                                If not, let Numeric decide

    This function is necessary as array(A) can cause memory overflow.
    """

    from Numeric import ArrayType, array

    if typecode is None:
        if type(A) == ArrayType:
            return A
        else:
            return array(A)
    else:
        if type(A) == ArrayType:
            if A.typecode == typecode:
                return array(A)
            else:
                return A.astype(typecode)
        else:
            return array(A).astype(typecode)



def file_function(filename,
                  domain = None,
                  quantities = None,
                  interpolation_points = None, verbose = False):
    """If quantities is not specified, derive them from domain
    (if that is specified)


    see get_netcdf_file_function
    """


    #FIXME (OLE): Should check origin of domain against that of file
    #In fact, this is where origin should be converted to that of domain
    #Also, check that file covers domain fully.

    assert type(filename) == type(''),\
               'First argument to File_function must be a string'

    try:
        fid = open(filename)
    except Exception, e:
        msg = 'File "%s" could not be opened: Error="%s"'\
                  %(filename, e)
        raise msg

    line = fid.readline()
    fid.close()

    if quantities is None: 
        if domain is not None:
            quantities = domain.conserved_quantities



    if line[:3] == 'CDF':
        return get_netcdf_file_function(filename, domain, quantities,
                                        interpolation_points, verbose = verbose)
    else:
        raise 'Must be a NetCDF File'



def get_netcdf_file_function(filename, domain=None, quantity_names=None,
                             interpolation_points=None, verbose = False):
    """Read time history of spatial data from NetCDF sww file and
    return a callable object f(t,x,y)
    which will return interpolated values based on the input file.

    If domain is specified, model time (domain.starttime)
    will be checked and possibly modified
    
    All times are assumed to be in UTC

    See Interpolation function for further documetation
    

    """
    
    #verbose = True
    
    #FIXME: Check that model origin is the same as file's origin
    #(both in UTM coordinates)
    #If not - modify those from file to match domain
    #Take this code from e.g. dem2pts in data_manager.py
    #FIXME: Use geo_reference to read and write xllcorner...
        

    import time, calendar, types
    from config import time_format
    from Scientific.IO.NetCDF import NetCDFFile
    from Numeric import array, zeros, Float, alltrue, concatenate, reshape        

    #Open NetCDF file
    if verbose: print 'Reading', filename
    fid = NetCDFFile(filename, 'r')

    if type(quantity_names) == types.StringType:
        quantity_names = [quantity_names]        
    
    if quantity_names is None or len(quantity_names) < 1:
        #Get quantities from file
        quantity_names = []
        for name in fid.variables.keys():
            if name not in ['x', 'y', 'time']:
                quantity_names.append(name)
                
        #msg = 'ERROR: At least one independent value must be specified'
        #raise msg

    missing = []
        
    for quantity in ['time'] + quantity_names:
        if not fid.variables.has_key(quantity):
            missing.append(quantity)

    if len(missing) > 0:
        msg = 'Quantities %s could not be found in file %s'\
              %(str(missing), filename)
        fid.close()
        raise msg

    #Decide whether this data has a spatial dimension
    #FIXME: Let us split this up into two functions: One that reads
    #sww and one that reads tms (time series only)
    spatial = True
    for quantity in ['x', 'y']:
        if not fid.variables.has_key(quantity):
            spatial = False



    #Get first timestep
    try:
        starttime = fid.starttime[0]
    except ValueError:
        msg = 'Could not read starttime from file %s' %filename
        raise msg

    #Get variables
    if verbose: print 'Get variables'    
    time = fid.variables['time'][:]    

    if spatial:
        #Get origin
        xllcorner = fid.xllcorner[0]
        yllcorner = fid.yllcorner[0]

        x = fid.variables['x'][:]
        y = fid.variables['y'][:]
        triangles = fid.variables['volumes'][:]

        x = reshape(x, (len(x),1))
        y = reshape(y, (len(y),1))
        vertex_coordinates = concatenate((x,y), axis=1) #m x 2 array


    if domain is not None:
        #Update domain.startime if it is *earlier* than starttime
        if starttime > domain.starttime:
            msg = 'WARNING: Start time as specified in domain (%f)'\
                  %domain.starttime
            msg += ' is earlier than the starttime of file %s (%f).'\
                     %(filename, starttime)
            msg += ' Modifying domain starttime accordingly.'
            
            if verbose: print msg
            domain.starttime = starttime #Modifying model time
            if verbose: print 'Domain starttime is now set to %f'\
               %domain.starttime


        #If domain.startime is *later* than starttime,
        #move time back - relative to domain's time
        if domain.starttime > starttime:
            time = time - domain.starttime + starttime


    if verbose:
        print 'File_function data obtained from: %s' %filename
        print '  References:'
        #print '    Datum ....' #FIXME
        if spatial:
            print '    Lower left corner: [%f, %f]'\
                  %(xllcorner, yllcorner)
        print '    Start time:   %f' %starttime                
        
    
    #Produce values for desired data points at
    #each timestep

    quantities = {}
    for i, name in enumerate(quantity_names):
        quantities[name] = fid.variables[name][:]
    fid.close()
    

    from least_squares import Interpolation_function

    if not spatial:
        vertex_coordinates = triangles = interpolation_points = None         
        
    return Interpolation_function(time,
                                  quantities,
                                  quantity_names,
                                  vertex_coordinates,
                                  triangles,
                                  interpolation_points,
                                  verbose = verbose)
    




def read_xya(filename, format = 'netcdf'):
    """Read simple xya file, possibly with a header in the first line, just
    x y [attributes]
    separated by whitespace

    Format can be either ASCII or NetCDF

    Return list of points, list of attributes and
    attribute names if present otherwise None
    """
    #FIXME: Probably obsoleted by similar function in load_ASCII
    #FIXME: Name read_data_points (or see 'load' in pylab)

    from Scientific.IO.NetCDF import NetCDFFile

    if format.lower() == 'netcdf':
        #Get NetCDF

        fid = NetCDFFile(filename, 'r')

        # Get the variables
        points = fid.variables['points']
        keys = fid.variables.keys()
        attributes = {}
        for key in keys:
            if key != 'points':
                attributes[key] = fid.variables[key]
        #Don't close - arrays are needed outside this function,
        #alternatively take a copy here
    else:
        #Read as ASCII file assuming that it is separated by whitespace
        fid = open(filename)
        lines = fid.readlines()
        fid.close()

        #Check if there is a header line
        fields = lines[0].strip().split()
        try:
            float(fields[0])
        except:
            #This must be a header line
            attribute_names = fields
            lines = lines[1:]
        else:
            attribute_names = ['elevation']  #HACK

        attributes = {}
        for key in attribute_names:
            attributes[key] = []

        points = []
        for line in lines:
            fields = line.strip().split()
            points.append( (float(fields[0]), float(fields[1])) )
            for i, key in enumerate(attribute_names):
                attributes[key].append( float(fields[2+i]) )

    return points, attributes


#####################################
#POLYGON STUFF
#
#FIXME: All these should be put into new module polygon.py



#These have been redefined to use separate_by_polygon which will
#put all inside indices in the first part of the indices array and
#outside indices in the last

def inside_polygon(points, polygon, closed = True, verbose = False):
    """See separate_points_by_polygon for documentation
    """

    from Numeric import array, Float, reshape

    if verbose: print 'Checking input to inside_polygon'
    try:
        points = ensure_numeric(points, Float)
    except:
        msg = 'Points could not be converted to Numeric array'
        raise msg

    try:
        polygon = ensure_numeric(polygon, Float)
    except:
        msg = 'Polygon could not be converted to Numeric array'
        raise msg



    if len(points.shape) == 1:
        one_point = True
        points = reshape(points, (1,2))
    else:
        one_point = False

    indices, count = separate_points_by_polygon(points, polygon,
                                                closed, verbose)

    if one_point:
        return count == 1
    else:
        return indices[:count]

def outside_polygon(points, polygon, closed = True, verbose = False):
    """See separate_points_by_polygon for documentation
    """

    from Numeric import array, Float, reshape

    if verbose: print 'Checking input to outside_polygon'
    try:
        points = ensure_numeric(points, Float)
    except:
        msg = 'Points could not be converted to Numeric array'
        raise msg

    try:
        polygon = ensure_numeric(polygon, Float)
    except:
        msg = 'Polygon could not be converted to Numeric array'
        raise msg



    if len(points.shape) == 1:
        one_point = True
        points = reshape(points, (1,2))
    else:
        one_point = False

    indices, count = separate_points_by_polygon(points, polygon,
                                                closed, verbose)

    if one_point:
        return count != 1
    else:
        return indices[count:][::-1]  #return reversed


def separate_points_by_polygon(points, polygon,
                               closed = True, verbose = False):
    """Determine whether points are inside or outside a polygon

    Input:
       points - Tuple of (x, y) coordinates, or list of tuples
       polygon - list of vertices of polygon
       closed - (optional) determine whether points on boundary should be
       regarded as belonging to the polygon (closed = True)
       or not (closed = False)

    Outputs:
       indices: array of same length as points with indices of points falling
       inside the polygon listed from the beginning and indices of points
       falling outside listed from the end.

       count: count of points falling inside the polygon

       The indices of points inside are obtained as indices[:count]
       The indices of points outside are obtained as indices[count:]


    Examples:
       U = [[0,0], [1,0], [1,1], [0,1]] #Unit square

       separate_points_by_polygon( [[0.5, 0.5], [1, -0.5], [0.3, 0.2]], U)
       will return the indices [0, 2, 1] and count == 2 as only the first
       and the last point are inside the unit square

    Remarks:
       The vertices may be listed clockwise or counterclockwise and
       the first point may optionally be repeated.
       Polygons do not need to be convex.
       Polygons can have holes in them and points inside a hole is
       regarded as being outside the polygon.

    Algorithm is based on work by Darel Finley,
    http://www.alienryderflex.com/polygon/

    """

    from Numeric import array, Float, reshape, Int, zeros


    #Input checks
    try:
        points = ensure_numeric(points, Float)
    except:
        msg = 'Points could not be converted to Numeric array'
        raise msg

    try:
        polygon = ensure_numeric(polygon, Float)
    except:
        msg = 'Polygon could not be converted to Numeric array'
        raise msg

    assert len(polygon.shape) == 2,\
       'Polygon array must be a 2d array of vertices'

    assert polygon.shape[1] == 2,\
       'Polygon array must have two columns'

    assert len(points.shape) == 2,\
       'Points array must be a 2d array'

    assert points.shape[1] == 2,\
       'Points array must have two columns'

    N = polygon.shape[0] #Number of vertices in polygon
    M = points.shape[0]  #Number of points

    px = polygon[:,0]
    py = polygon[:,1]

    #Used for an optimisation when points are far away from polygon
    minpx = min(px); maxpx = max(px)
    minpy = min(py); maxpy = max(py)


    #Begin main loop
    indices = zeros(M, Int)

    inside_index = 0    #Keep track of points inside
    outside_index = M-1 #Keep track of points outside (starting from end)

    for k in range(M):

        if verbose:
            if k %((M+10)/10)==0: print 'Doing %d of %d' %(k, M)

        #for each point
        x = points[k, 0]
        y = points[k, 1]

        inside = False

        if not x > maxpx or x < minpx or y > maxpy or y < minpy:
            #Check polygon
            for i in range(N):
                j = (i+1)%N

                #Check for case where point is contained in line segment
                if point_on_line(x, y, px[i], py[i], px[j], py[j]):
                    if closed:
                        inside = True
                    else:
                        inside = False
                    break
                else:
                    #Check if truly inside polygon
                    if py[i] < y and py[j] >= y or\
                       py[j] < y and py[i] >= y:
                        if px[i] + (y-py[i])/(py[j]-py[i])*(px[j]-px[i]) < x:
                            inside = not inside

        if inside:
            indices[inside_index] = k
            inside_index += 1
        else:
            indices[outside_index] = k
            outside_index -= 1

    return indices, inside_index


def separate_points_by_polygon_c(points, polygon,
                                 closed = True, verbose = False):
    """Determine whether points are inside or outside a polygon

    C-wrapper
    """

    from Numeric import array, Float, reshape, zeros, Int


    if verbose: print 'Checking input to separate_points_by_polygon'
    #Input checks
    try:
        points = ensure_numeric(points, Float)
    except:
        msg = 'Points could not be converted to Numeric array'
        raise msg

    #if verbose: print 'Checking input to separate_points_by_polygon 2'
    try:
        polygon = ensure_numeric(polygon, Float)
    except:
        msg = 'Polygon could not be converted to Numeric array'
        raise msg

    if verbose: print 'check'

    assert len(polygon.shape) == 2,\
       'Polygon array must be a 2d array of vertices'

    assert polygon.shape[1] == 2,\
       'Polygon array must have two columns'

    assert len(points.shape) == 2,\
       'Points array must be a 2d array'

    assert points.shape[1] == 2,\
       'Points array must have two columns'

    N = polygon.shape[0] #Number of vertices in polygon
    M = points.shape[0]  #Number of points

    from util_ext import separate_points_by_polygon

    if verbose: print 'Allocating array for indices'

    indices = zeros( M, Int )

    if verbose: print 'Calling C-version of inside poly'
    count = separate_points_by_polygon(points, polygon, indices,
                                       int(closed), int(verbose))

    return indices, count



class Polygon_function:
    """Create callable object f: x,y -> z, where a,y,z are vectors and
    where f will return different values depending on whether x,y belongs
    to specified polygons.

    To instantiate:

       Polygon_function(polygons)

    where polygons is a list of tuples of the form

      [ (P0, v0), (P1, v1), ...]

      with Pi being lists of vertices defining polygons and vi either
      constants or functions of x,y to be applied to points with the polygon.

    The function takes an optional argument, default which is the value
    (or function) to used for points not belonging to any polygon.
    For example:

       Polygon_function(polygons, default = 0.03)

    If omitted the default value will be 0.0

    Note: If two polygons overlap, the one last in the list takes precedence

    FIXME? : Currently, coordinates specified here are assumed to be relative to the origin (georeference) used by domain. This function is more general than domain so perhaps it'd be an idea to allow the optional argument georeference???

    """

    def __init__(self, regions, default = 0.0):

        try:
            len(regions)
        except:
            msg = 'Polygon_function takes a list of pairs (polygon, value). Got %s' %polygons
            raise msg


        T = regions[0]
        try:
            a = len(T)
        except:
            msg = 'Polygon_function takes a list of pairs (polygon, value). Got %s' %polygons
            raise msg

        assert a == 2, 'Must have two component each: %s' %T

        self.regions = regions
        self.default = default


    def __call__(self, x, y):
        from util import inside_polygon
        from Numeric import ones, Float, concatenate, array, reshape, choose

        x = array(x).astype(Float)
        y = array(y).astype(Float)

        N = len(x)
        assert len(y) == N

        points = concatenate( (reshape(x, (N, 1)),
                               reshape(y, (N, 1))), axis=1 )

        if callable(self.default):
            z = self.default(x,y)
        else:
            z = ones(N, Float) * self.default

        for polygon, value in self.regions:
            indices = inside_polygon(points, polygon)

            #FIXME: This needs to be vectorised
            if callable(value):
                for i in indices:
                    xx = array([x[i]])
                    yy = array([y[i]])
                    z[i] = value(xx, yy)[0]
            else:
                for i in indices:
                    z[i] = value

        return z

def read_polygon(filename):
    """Read points assumed to form a polygon
       There must be exactly two numbers in each line
    """

    #Get polygon
    fid = open(filename)
    lines = fid.readlines()
    fid.close()
    polygon = []
    for line in lines:
        fields = line.split(',')
        polygon.append( [float(fields[0]), float(fields[1])] )

    return polygon

def populate_polygon(polygon, number_of_points, seed = None):
    """Populate given polygon with uniformly distributed points.

    Input:
       polygon - list of vertices of polygon
       number_of_points - (optional) number of points
       seed - seed for random number generator (default=None)

    Output:
       points - list of points inside polygon

    Examples:
       populate_polygon( [[0,0], [1,0], [1,1], [0,1]], 5 )
       will return five randomly selected points inside the unit square
    """

    from random import uniform, seed

    seed(seed)

    points = []

    #Find outer extent of polygon
    max_x = min_x = polygon[0][0]
    max_y = min_y = polygon[0][1]
    for point in polygon[1:]:
        x = point[0]
        if x > max_x: max_x = x
        if x < min_x: min_x = x
        y = point[1]
        if y > max_y: max_y = y
        if y < min_y: min_y = y


    while len(points) < number_of_points:
        x = uniform(min_x, max_x)
        y = uniform(min_y, max_y)

        if inside_polygon( [x,y], polygon ):
            points.append([x,y])

    return points

def tsh2sww(filename, verbose=True):
    """
    to check if a tsh/msh file 'looks' good.
    """

    #FIXME Move to data_manager
    from shallow_water import Domain
    from pmesh2domain import pmesh_to_domain_instance
    import time, os
    from data_manager import get_dataobject
    from os import sep, path
    #from util import mean

    if verbose == True:print 'Creating domain from', filename
    domain = pmesh_to_domain_instance(filename, Domain)
    if verbose == True:print "Number of triangles = ", len(domain)

    domain.smooth = True
    domain.format = 'sww'   #Native netcdf visualisation format
    file_path, filename = path.split(filename)
    filename, ext = path.splitext(filename)
    domain.filename = filename
    domain.reduction = mean
    if verbose == True:print "file_path",file_path
    if file_path == "":file_path = "."
    domain.set_datadir(file_path)

    if verbose == True:
        print "Output written to " + domain.get_datadir() + sep + \
              domain.filename + "." + domain.format
    sww = get_dataobject(domain)
    sww.store_connectivity()
    sww.store_timestep('stage')

####################################################################
#Python versions of function that are also implemented in util_gateway.c
#

def gradient_python(x0, y0, x1, y1, x2, y2, q0, q1, q2):
    """
    """

    det = (y2-y0)*(x1-x0) - (y1-y0)*(x2-x0)
    a = (y2-y0)*(q1-q0) - (y1-y0)*(q2-q0)
    a /= det

    b = (x1-x0)*(q2-q0) - (x2-x0)*(q1-q0)
    b /= det

    return a, b


def gradient2_python(x0, y0, x1, y1, q0, q1):
    """Compute radient based on two points and enforce zero gradient
    in the direction orthogonal to (x1-x0), (y1-y0) 
    """

    #Old code
    #det = x0*y1 - x1*y0
    #if det != 0.0:
    #    a = (y1*q0 - y0*q1)/det
    #    b = (x0*q1 - x1*q0)/det

    #Correct code (ON)
    det = (x1-x0)**2 + (y1-y0)**2
    if det != 0.0:
        a = (x1-x0)*(q1-q0)/det
        b = (y1-y0)*(q1-q0)/det
        
    return a, b        


##############################################
#Initialise module

import compile
if compile.can_use_C_extension('util_ext.c'):
    from util_ext import gradient, gradient2, point_on_line
    separate_points_by_polygon = separate_points_by_polygon_c
else:
    gradient = gradient_python
    gradient2 = gradient2_python    


if __name__ == "__main__":
    pass