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

"""Least squares smooting and interpolation.

   Implements a penalised least-squares fit and associated interpolations.

   The panalty 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: Current implementation uses full matrices and a general solver.
#Later on we may consider using sparse techniques

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 scipy import sparse
from sparse import Sparse
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
    
    point_dict = load_xya_file(point_file)
    point_coordinates = point_dict['pointlist']
    point_attributes = point_dict['pointattributelist']
    title_string = point_dict['title']
    title_list = title_string.split(',') #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)

          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 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

        
    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)"""
        
        #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 = zeros((n,m), Float)
        self.A = Sparse(n,m)
        self.AtA = Sparse(m,m)

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

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

            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
        
##         self.A   = (self.A).tocsc()
##         self.AtA = (self.AtA).tocsc()
##         self.At  =  self.A.transp()


        
    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 = zeros((m,m), Float)
        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             

        #self.D = (self.D).tocsc()
            
    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)


        #print 'fit: Atz',Atz
        
        #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


        #Solve and return
        #return solve_linear_equations(self.get_B(), Atz)

        # caused errors
        #return sparse.solve(self.B,Atz)

        return conjugate_gradient(self.B, Atz, 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
        """
        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)
            #f = sparse.dok_matrix() # even though it wont be sparse?
            
            for i in range(z.shape[1]):
                f[:,i] = self.fit(z[:,i])
            return f
            
        
    def interpolate(self, f):
        """Compute predicted values 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
        """

        try:
            return self.A * f
        except ValueError:
            # We are here, so it probalby means that f is 2 dimensional
            # and so will do each column separately due to problems in
            # sparse matrix, 2d array multiplication
            (N , M) = self.A.shape
            f = array(f).astype(Float)
            (m , n) = f.shape
            #print "m",m
            #print "M",M
            if m != M :
                #print "!!self.A",self.A
                #print "!!self.A.todense() ",self.A.todense() 
                #print "self.get_A()",self.get_A()
                #return dot(self.get_A(),f)
                raise VectorShapeError, 'Mismatch between A and f dimensions'
            
            y = zeros( (N,n), Float)

            for i in range(n):
                y[:,i] = self.A * f[:,i]

            #print "!!self.A.todense() ",self.A.todense() 
            return y
     
    #FIXME: We will need a method 'evaluate(self):' that will interpolate
    #a computed surface living on the mesh onto a collection of
    #arbitrary data points
    #
    #Precondition: self.fit(z) has stored its result in self.f.
    #
    #Input: data_points
    #Algorithm:
    #  1 Build a new temporary A matrix based on mesh and new data points
    #  2 Apply it to self.f (return A*self.f)
    #
    # ON



#-------------------------------------------------------------
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)