source: inundation/ga/storm_surge/pyvolution/least_squares.py @ 613

"""Least squares smooting and interpolation.

   Implements a penalised least-squares fit and associated interpolations.

   The penalty term (or smoothing term) is controlled by the smoothing
   parameter alpha.
   With a value of alpha=0, the fit function will attempt
   to interpolate as closely as possible in the least-squares sense.
   With values alpha > 0, a certain amount of smoothing will be applied.
   A positive alpha is essential in cases where there are too few
   data points.
   A negative alpha is not allowed.
   A typical value of alpha is 1.0e-6
          

   Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
   Geoscience Australia, 2004.   
"""


#FIXME (Ole): Currently datapoints outside the triangular mesh are ignored.
#             Is there a clean way of including them?



import exceptions
class ShapeError(exceptions.Exception): pass

from general_mesh import General_mesh
from Numeric import zeros, array, Float, Int, dot, transpose
from LinearAlgebra import solve_linear_equations
from sparse import Sparse, Sparse_CSR
from cg_solve import conjugate_gradient, VectorShapeError

try:
    from util import gradient
except ImportError, e: 
    #FIXME reduce the dependency of modules in pyvolution
    # Have util in a dir, working like load_mesh, and get rid of this
    def gradient(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


DEFAULT_ALPHA = 0.001
    
def fit_to_mesh_file(mesh_file, point_file, mesh_output_file, alpha=DEFAULT_ALPHA):
    """
    Given a mesh file (tsh) and a point attribute file (xya), fit
    point attributes to the mesh and write a mesh file with the
    results.
    """
    from load_mesh.loadASCII import mesh_file_to_mesh_dictionary, \
                 load_xya_file, export_trianglulation_file
    # load in the .tsh file
    mesh_dict = mesh_file_to_mesh_dictionary(mesh_file)
    vertex_coordinates = mesh_dict['generatedpointlist']
    triangles = mesh_dict['generatedtrianglelist']
    
    old_point_attributes = mesh_dict['generatedpointattributelist'] 
    old_title_list = mesh_dict['generatedpointattributetitlelist']

   
    # load in the .xya file
    try:
        point_dict = load_xya_file(point_file)
    except SyntaxError,e:
        point_dict = load_xya_file(point_file,delimiter = ' ')
    point_coordinates = point_dict['pointlist']
    point_attributes = point_dict['pointattributelist']
    title_string = point_dict['title']
    title_list = title_string.split(',') #FIXME iffy! Hard coding title delimiter  
    for i in range(len(title_list)):
        title_list[i] = title_list[i].strip() 
    #print "title_list stripped", title_list    
    f = fit_to_mesh(vertex_coordinates,
                    triangles,
                    point_coordinates,
                    point_attributes,
                    alpha = alpha)
    
    # convert array to list of lists
    new_point_attributes = f.tolist()

    #FIXME have this overwrite attributes with the same title - DSG
    #Put the newer attributes last
    if old_title_list <> []:
        old_title_list.extend(title_list)
        #FIXME can this be done a faster way? - DSG
        for i in range(len(old_point_attributes)):
            old_point_attributes[i].extend(new_point_attributes[i])
        mesh_dict['generatedpointattributelist'] = old_point_attributes
        mesh_dict['generatedpointattributetitlelist'] = old_title_list
    else:
        mesh_dict['generatedpointattributelist'] = new_point_attributes
        mesh_dict['generatedpointattributetitlelist'] = title_list
    
    export_trianglulation_file(mesh_output_file, mesh_dict)
        

def fit_to_mesh(vertex_coordinates,
                triangles,
                point_coordinates,
                point_attributes,
                alpha = DEFAULT_ALPHA):
    """
    Fit a smooth surface to a trianglulation,
    given data points with attributes.

          
        Inputs:
        
          vertex_coordinates: List of coordinate pairs [xi, eta] of points
          constituting mesh (or a an m x 2 Numeric array)
        
          triangles: List of 3-tuples (or a Numeric array) of
          integers representing indices of all vertices in the mesh.

          point_coordinates: List of coordinate pairs [x, y] of data points
          (or an nx2 Numeric array)

          alpha: Smoothing parameter.

          point_attributes: Vector or array of data at the point_coordinates.
    """
    interp = Interpolation(vertex_coordinates,
                           triangles,
                           point_coordinates,
                           alpha = alpha)
    
    vertex_attributes = interp.fit_points(point_attributes)
    return vertex_attributes


class Interpolation:

    def __init__(self,
                 vertex_coordinates,
                 triangles,
                 point_coordinates = None,
                 alpha = DEFAULT_ALPHA):

        
        """ Build interpolation matrix mapping from
        function values at vertices to function values at data points

        Inputs:
        
          vertex_coordinates: List of coordinate pairs [xi, eta] of 
          points constituting mesh (or a an m x 2 Numeric array)
        
          triangles: List of 3-tuples (or a Numeric array) of
          integers representing indices of all vertices in the mesh.

          point_coordinates: List of coordinate pairs [x, y] of 
          data points (or an nx2 Numeric array)
          If point_coordinates is absent, only smoothing matrix will 
          be built 

          alpha: Smoothing parameter
          
        """


        #Convert input to Numeric arrays
        vertex_coordinates = array(vertex_coordinates).astype(Float)
        triangles = array(triangles).astype(Int)                
        
        #Build underlying mesh
        self.mesh = General_mesh(vertex_coordinates, triangles)

        #Smoothing parameter
        self.alpha = alpha

        #Build coefficient matrices
        self.build_coefficient_matrix_B(point_coordinates)    

    def set_point_coordinates(self, point_coordinates):
        """
        A public interface to setting the point co-ordinates.
        """
        self.build_coefficient_matrix_B(point_coordinates)
        
    def build_coefficient_matrix_B(self, point_coordinates=None):
        """Build final coefficient matrix"""
        

        if self.alpha <> 0:
            self.build_smoothing_matrix_D()
        
        if point_coordinates:

            self.build_interpolation_matrix_A(point_coordinates)

            if self.alpha <> 0:
                self.B = self.AtA + self.alpha*self.D
            else:
                self.B = self.AtA

            #Convert self.B matrix to CSR format for faster matrix vector
            self.B = Sparse_CSR(self.B)
        
    def build_interpolation_matrix_A(self, point_coordinates):
        """Build n x m interpolation matrix, where
        n is the number of data points and
        m is the number of basis functions phi_k (one per vertex)

        This algorithm uses a quad tree data structure for fast binning of data points
        """

        from quad import build_quadtree
        
        #Convert input to Numeric arrays
        point_coordinates = array(point_coordinates).astype(Float)
        
        #Build n x m interpolation matrix        
        m = self.mesh.coordinates.shape[0] #Nbr of basis functions (1/vertex)
        n = point_coordinates.shape[0]     #Nbr of data points
        
        self.A = Sparse(n,m)
        self.AtA = Sparse(m,m)

        #Build quad tree of vertices (FIXME: Is this the right spot for that?)
        root = build_quadtree(self.mesh)

        #Compute matrix elements
        for i in range(n):
            #For each data_coordinate point

            #print 'Doing %d of %d' %(i, n)

            x = point_coordinates[i]

            #Find vertices near x
            candidate_vertices = root.search(x[0], x[1])

            is_more_elements = True
            if candidate_vertices == []:
                # The point isn't even within the root cell!
                is_more_elements = False
                element_found = False
            else:
                element_found, sigma0, sigma1, sigma2, k = \
                    self.search_triangles_of_vertices(candidate_vertices, x)

            while not element_found and is_more_elements: 
                candidate_vertices = root.expand_search()
                if candidate_vertices == []:
                    # All the triangles have been searched.
                    is_more_elements = False
                else:
                    element_found, sigma0, sigma1, sigma2, k = \
                      self.search_triangles_of_vertices(candidate_vertices, x)
                    
                
            
            #Update interpolation matrix A if necessary     
            if element_found is True:        
                #Assign values to matrix A

                j0 = self.mesh.triangles[k,0] #Global vertex id
                #self.A[i, j0] = sigma0

                j1 = self.mesh.triangles[k,1] #Global vertex id
                #self.A[i, j1] = sigma1

                j2 = self.mesh.triangles[k,2] #Global vertex id
                #self.A[i, j2] = sigma2

                sigmas = {j0:sigma0, j1:sigma1, j2:sigma2}
                js     = [j0,j1,j2]

                for j in js:
                    self.A[i,j] = sigmas[j]
                    for k in js:
                        self.AtA[j,k] += sigmas[j]*sigmas[k]
            else:
                pass
                #Ok if there is no triangle for datapoint 
                #(as in brute force version)
                #raise 'Could not find triangle for point', x


    def search_triangles_of_vertices(self, candidate_vertices, x):
            #Find triangle containing x:
            element_found = False

            # This will be returned if element_found = False
            sigma2 = -10.0
            sigma0 = -10.0
            sigma1 = -10.0

            #For all vertices in same cell as point x
            for v in candidate_vertices:
            
                #for each triangle id (k) which has v as a vertex 
                for k, _ in self.mesh.vertexlist[v]:
                    
                    #Get the three vertex_points of candidate triangle
                    xi0 = self.mesh.get_vertex_coordinate(k, 0)
                    xi1 = self.mesh.get_vertex_coordinate(k, 1)
                    xi2 = self.mesh.get_vertex_coordinate(k, 2)     

                    #print "PDSG - k", k
                    #print "PDSG - xi0", xi0
                    #print "PDSG - xi1", xi1
                    #print "PDSG - xi2", xi2
                    #print "PDSG element %i verts((%f, %f),(%f, %f),(%f, %f))" \
                    #      % (k, xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1])
                    
                    #Get the three normals
                    n0 = self.mesh.get_normal(k, 0)
                    n1 = self.mesh.get_normal(k, 1)
                    n2 = self.mesh.get_normal(k, 2)                

                    

                    #Compute interpolation
                    sigma2 = dot((x-xi0), n2)/dot((xi2-xi0), n2)
                    sigma0 = dot((x-xi1), n0)/dot((xi0-xi1), n0)
                    sigma1 = dot((x-xi2), n1)/dot((xi1-xi2), n1)

                    #print "PDSG - sigma0", sigma0
                    #print "PDSG - sigma1", sigma1
                    #print "PDSG - sigma2", sigma2
                    
                    #FIXME: Maybe move out to test or something
                    epsilon = 1.0e-6
                    assert abs(sigma0 + sigma1 + sigma2 - 1.0) < epsilon
                    
                    #Check that this triangle contains the data point
                    
                    #Sigmas can get negative within
                    #machine precision on some machines (e.g nautilus)
                    #Hence the small eps                    
                    eps = 1.0e-15
                    if sigma0 >= -eps and sigma1 >= -eps and sigma2 >= -eps:
                        element_found = True
                        break

                if element_found is True:
                    #Don't look for any other triangle
                    break
            return element_found, sigma0, sigma1, sigma2, k     
                    

        
    def build_interpolation_matrix_A_brute(self, point_coordinates):
        """Build n x m interpolation matrix, where
        n is the number of data points and
        m is the number of basis functions phi_k (one per vertex)

        This is the brute force which is too slow for large problems, 
        but could be used for testing
        """


        
        #Convert input to Numeric arrays
        point_coordinates = array(point_coordinates).astype(Float)
        
        #Build n x m interpolation matrix        
        m = self.mesh.coordinates.shape[0] #Nbr of basis functions (1/vertex)
        n = point_coordinates.shape[0]     #Nbr of data points
        
        self.A = Sparse(n,m)
        self.AtA = Sparse(m,m)

        #Compute matrix elements
        for i in range(n):
            #For each data_coordinate point

            x = point_coordinates[i]
            element_found = False
            k = 0
            while not element_found and k < len(self.mesh):
                #For each triangle (brute force)
                #FIXME: Real algorithm should only visit relevant triangles

                #Get the three vertex_points
                xi0 = self.mesh.get_vertex_coordinate(k, 0)
                xi1 = self.mesh.get_vertex_coordinate(k, 1)
                xi2 = self.mesh.get_vertex_coordinate(k, 2)                 

                #Get the three normals
                n0 = self.mesh.get_normal(k, 0)
                n1 = self.mesh.get_normal(k, 1)
                n2 = self.mesh.get_normal(k, 2)                

                #Compute interpolation
                sigma2 = dot((x-xi0), n2)/dot((xi2-xi0), n2)
                sigma0 = dot((x-xi1), n0)/dot((xi0-xi1), n0)
                sigma1 = dot((x-xi2), n1)/dot((xi1-xi2), n1)

                #FIXME: Maybe move out to test or something
                epsilon = 1.0e-6
                assert abs(sigma0 + sigma1 + sigma2 - 1.0) < epsilon

                #Check that this triangle contains data point
                if sigma0 >= 0 and sigma1 >= 0 and sigma2 >= 0:
                    element_found = True
                    #Assign values to matrix A

                    j0 = self.mesh.triangles[k,0] #Global vertex id
                    #self.A[i, j0] = sigma0

                    j1 = self.mesh.triangles[k,1] #Global vertex id
                    #self.A[i, j1] = sigma1

                    j2 = self.mesh.triangles[k,2] #Global vertex id
                    #self.A[i, j2] = sigma2

                    sigmas = {j0:sigma0, j1:sigma1, j2:sigma2}
                    js     = [j0,j1,j2]

                    for j in js:
                        self.A[i,j] = sigmas[j]
                        for k in js:
                            self.AtA[j,k] += sigmas[j]*sigmas[k]
                k = k+1
        

        
    def get_A(self):
        return self.A.todense() 

    def get_B(self):
        return self.B.todense()
    
    def get_D(self):
        return self.D.todense()
    
        #FIXME: Remember to re-introduce the 1/n factor in the
        #interpolation term
        
    def build_smoothing_matrix_D(self):
        """Build m x m smoothing matrix, where
        m is the number of basis functions phi_k (one per vertex)

        The smoothing matrix is defined as

        D = D1 + D2

        where

        [D1]_{k,l} = \int_\Omega
           \frac{\partial \phi_k}{\partial x}
           \frac{\partial \phi_l}{\partial x}\,
           dx dy

        [D2]_{k,l} = \int_\Omega
           \frac{\partial \phi_k}{\partial y}
           \frac{\partial \phi_l}{\partial y}\,
           dx dy


        The derivatives \frac{\partial \phi_k}{\partial x},
        \frac{\partial \phi_k}{\partial x} for a particular triangle
        are obtained by computing the gradient a_k, b_k for basis function k
        """

        #FIXME: algorithm might be optimised by computing local 9x9
        #"element stiffness matrices:

        m = self.mesh.coordinates.shape[0] #Nbr of basis functions (1/vertex)

        self.D = Sparse(m,m)

        #For each triangle compute contributions to D = D1+D2        
        for i in range(len(self.mesh)):

            #Get area
            area = self.mesh.areas[i]

            #Get global vertex indices
            v0 = self.mesh.triangles[i,0]
            v1 = self.mesh.triangles[i,1]
            v2 = self.mesh.triangles[i,2]
            
            #Get the three vertex_points
            xi0 = self.mesh.get_vertex_coordinate(i, 0)
            xi1 = self.mesh.get_vertex_coordinate(i, 1)
            xi2 = self.mesh.get_vertex_coordinate(i, 2)                 

            #Compute gradients for each vertex
            a0, b0 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              1, 0, 0)

            a1, b1 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              0, 1, 0)

            a2, b2 = gradient(xi0[0], xi0[1], xi1[0], xi1[1], xi2[0], xi2[1],
                              0, 0, 1)            

            #Compute diagonal contributions
            self.D[v0,v0] += (a0*a0 + b0*b0)*area
            self.D[v1,v1] += (a1*a1 + b1*b1)*area
            self.D[v2,v2] += (a2*a2 + b2*b2)*area            

            #Compute contributions for basis functions sharing edges
            e01 = (a0*a1 + b0*b1)*area
            self.D[v0,v1] += e01
            self.D[v1,v0] += e01

            e12 = (a1*a2 + b1*b2)*area
            self.D[v1,v2] += e12
            self.D[v2,v1] += e12

            e20 = (a2*a0 + b2*b0)*area
            self.D[v2,v0] += e20
            self.D[v0,v2] += e20             

            
    def fit(self, z):
        """Fit a smooth surface to given 1d array of data points z.

        The smooth surface is computed at each vertex in the underlying
        mesh using the formula given in the module doc string. 

        Pre Condition:
          self.A, self.At and self.B have been initialised
          
        Inputs:
          z: Single 1d vector or array of data at the point_coordinates.
        """

        #Convert input to Numeric arrays
        z = array(z).astype(Float)


        if len(z.shape) > 1 :
            raise VectorShapeError, 'Can only deal with 1d data vector'
        
        #Compute right hand side based on data
        Atz = self.A.trans_mult(z)

        
        #Check sanity
        n, m = self.A.shape
        if n<m and self.alpha == 0.0:
            msg = 'ERROR (least_squares): Too few data points\n'
            msg += 'There only %d data points. Need at least %d\n' %(n,m)
            msg += 'Alternatively, increase smoothing parameter alpha' 
            raise msg



        return conjugate_gradient(self.B, Atz, Atz,imax=2*len(Atz) )
        #FIXME: Should we store the result here for later use? (ON)        

            
    def fit_points(self, z):
        """Like fit, but more robust when each point has two or more attributes
        FIXME (Ole): The name fit_points doesn't carry any meaning 
        for me. How about something like fit_multiple or fit_columns?
        """
        
        try:
            return self.fit(z)
        except VectorShapeError, e:
            # broadcasting is not supported.

            #Convert input to Numeric arrays
            z = array(z).astype(Float)
            
            #Build n x m interpolation matrix        
            m = self.mesh.coordinates.shape[0] #Number of vertices
            n = z.shape[1]               #Number of data points         

            f = zeros((m,n), Float) #Resulting columns
            
            for i in range(z.shape[1]):
                f[:,i] = self.fit(z[:,i])
                
            return f
            
        
    def interpolate(self, f):
        """Evaluate smooth surface f at data points implied in self.A. 

        The mesh values representing a smooth surface are
        assumed to be specified in f. This argument could,
        for example have been obtained from the method self.fit()
        
        Pre Condition:
          self.A has been initialised
        
        Inputs:
          f: Vector or array of data at the mesh vertices.
          If f is an array, interpolation will be done for each column
        """

        return self.A * f
        
            
#-------------------------------------------------------------
if __name__ == "__main__":
    """
    Load in a mesh and data points with attributes.
    Fit the attributes to the mesh.
    Save a new mesh file.
    """
    import os, sys
    usage = "usage: %s mesh_input.tsh point.xya mesh_output.tsh alpha"\
            %os.path.basename(sys.argv[0])

    if len(sys.argv) < 4:
        print usage
    else:
        mesh_file = sys.argv[1]
        point_file = sys.argv[2]
        mesh_output_file = sys.argv[3]
        if len(sys.argv) > 4:
            alpha = sys.argv[4]
        else:
            alpha = DEFAULT_ALPHA
        fit_to_mesh_file(mesh_file, point_file, mesh_output_file, alpha)