<commit_before>from collections import OrderedDict
import numpy as np
from scipy.spatial import distance, Delaunay, cKDTree
from scipy import sparse
import scipy.sparse.linalg
import functools
import numexpr as ne

def _memo(fn):
    """Helper decorator memoizes the given zero-argument function.
    """
    @functools.wraps(fn)
    def memofn(self):
        if id(fn) not in self._cache:
            self._cache[id(fn)] = fn(self)
        return self._cache[id(fn)]

    return memofn

class Surface(object):
    def __init__(self, pts, polys):
        self.pts = pts
        self.polys = polys

        self._cache = dict()
        self._rlfac_solvers = dict()
        self._nLC_solvers = dict()

    @property
    @_memo
    def ppts(self):
        """3D matrix of points in each face: n faces x 3 points per face x 3 coords per point.
        """
        return self.pts[self.polys]
    
    @property
    @_memo
    def connected(self):
        """Sparse matrix of vertex-face associations.
        """
        npt = len(self.pts)
        npoly = len(self.polys)
        return sparse.coo_matrix((np.ones((3*npoly,)), # data
                                  (np.hstack(self.polys.T), # row
                                   np.tile(range(npoly),(1,3)).squeeze())), # col
                                 (npt, npoly)).tocsr() # size
    @property
    @_memo
    def adj(self):
        """Sparse vertex adjacency matrix.
        """
        npt = len(self.pts)
        npoly = len(self.polys)
        adj1 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,0], self.polys[:,1])), (npt,npt))
        adj2 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,0], self.polys[:,2])), (npt,npt))
        adj3 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,1], self.polys[:,2])), (npt,npt))
        return (adj1 + adj2 + adj3).tocsr()
    
    @property
    @_memo
    def face_normals(self):
        """Normal vector for each face.
        """
        nnfnorms = np.cross(self.ppts[:,1] - self.ppts[:,0], self.ppts[:,2] - self.ppts[:,0])
        return nnfnorms / np.sqrt((nnfnorms**2).sum(1))[:,np.newaxis]

    @property
    @_memo
    def vertex_normals(self):
        """Normal vector for each vertex (average of normals for neighboring faces).
        """
        nnvnorms = np.nan_to_num(self.connected.dot(self.face_normals) / self.connected.sum(1)).A
        return nnvnorms / np.sqrt((nnvnorms**2).sum(1))[:,np.newaxis]

    @property
    @_memo
    def face_areas(self):
        """Area of each face.
        """
        nnfnorms = np.cross(self.ppts[:,1] - self.ppts[:,0], self.ppts[:,2] - self.ppts[:,0])
        return np.sqrt((nnfnorms**2).sum(-1))

    @property
    @_memo
    def cotangent_weights(self):
        """Cotangent of angle opposite each vertex in each face.
        """
        ppts = self.ppts
        cots1 = ((ppts[:,1]-ppts[:,0]) *
                 (ppts[:,2]-ppts[:,0])).sum(1) / np.sqrt((np.cross(ppts[:,1]-ppts[:,0],
                                                                   ppts[:,2]-ppts[:,0])**2).sum(1))
        cots2 = ((ppts[:,2]-ppts[:,1]) *
                 (ppts[:,0]-ppts[:,1])).sum(1) / np.sqrt((np.cross(ppts[:,2]-ppts[:,1],
                                                                   ppts[:,0]-ppts[:,1])**2).sum(1))
        cots3 = ((ppts[:,0]-ppts[:,2]) *
                 (ppts[:,1]-ppts[:,2])).sum(1) / np.sqrt((np.cross(ppts[:,0]-ppts[:,2],
                                                                   ppts[:,1]-ppts[:,2])**2).sum(1))

        # Then we have to sanitize the fuck out of everything..
        cots = np.vstack([cots1, cots2, cots3])
        cots[np.isinf(cots)] = 0
        cots[np.isnan(cots)] = 0
        # Clip to a reasonable range, these meshes are the worst
        cots = np.clip(cots, -100, 100)
        return cots

    @property
    @_memo
    def laplace_operator(self):
        """Laplace-Beltrami operator for this surface. A sparse adjacency matrix with
        edge weights determined by the cotangents of the angles opposite each edge.
        Returns a triplet (D,W,V) where D is the 'lumped mass matrix', W is the weighted
        adjacency matrix, and V is a diagonal matrix that normalizes the adjacencies.
        The 'stiffness matrix', A, can be computed as V - W.
        
        See 'Discrete Laplace-Beltrami operators for shape analysis and segmentation'
        by Reuter et al., 2009 for details.
        """
        ## Lumped mass matrix
        D = self.connected.dot(self.face_areas) / 3.0

        ## Stiffness matrix
        npt = len(self.pts)
        cots1, cots2, cots3 = self.cotangent_weights
        # W is weighted adjacency matrix
        W1 = sparse.coo_matrix((cots1, (self.polys[:,1], self.polys[:,2])), (npt, npt))
        W2 = sparse.coo_matrix((cots2, (self.polys[:,2], self.polys[:,0])), (npt, npt))
        W3 = sparse.coo_matrix((cots3, (self.polys[:,0], self.polys[:,1])), (npt, npt))
        W = (W1 + W1.T + W2 + W2.T + W3 + W3.T).tocsr() / 2.0
        
        # V is sum of each col
        V = sparse.dia_matrix((np.array(W.sum(0)).ravel(),[0]), (npt,npt))

        # A is stiffness matrix
        #A = W - V # negative operator -- more useful in practice

        # For FEM:
        Be1 = sparse.coo_matrix((self.face_areas, (self.polys[:,1], self.polys[:,2])), (npt, npt))
        Be2 = sparse.coo_matrix((self.face_areas, (self.polys[:,2], self.polys[:,0])), (npt, npt))
        Be3 = sparse.coo_matrix((self.face_areas, (self.polys[:,0], self.polys[:,1])), (npt, npt))
        Bd = self.connected.dot(self.face_areas) / 6
        dBd = scipy.sparse.dia_matrix((Bd,[0]), (len(D),len(D)))
        B = (Be1 + Be1.T + Be2 + Be2.T + Be3 + Be3.T)/12 + dBd
        return B, D, W, V

    @property
    @_memo
    def avg_edge_length(self):
        """Average length of all edges in the surface.
        """
        adj = self.laplace_operator[1] # use laplace operator as adjacency matrix
        tadj = sparse.triu(adj, 1) # only entries above main diagonal, in coo format
        edgelens = np.sqrt(((self.pts[tadj.row] - self.pts[tadj.col])**2).sum(1))
        return edgelens.mean()

    def surface_gradient(self, vertvals, at_verts=True):
        """Gradient of a function with values vertvals at each vertex along surface.
        If at_verts, returns values at each vertex. Otherwise, returns values at each
        face.
        """
        pu = vertvals[self.polys]
        fe12, fe23, fe31 = [f.T for f in self._facenorm_cross_edge]
        pu1, pu2, pu3 = pu.T
        fa = self.face_areas
        #gradu = ((fe12.T * pu[:,2] +
        #          fe23.T * pu[:,0] +
        #          fe31.T * pu[:,1]) / (2 * self.face_areas)).T
        gradu = ne.evaluate("(fe12 * pu3 + fe23 * pu1 + fe31 * pu2) / (2 * fa)").T
        #gradu = np.nan_to_num(gradu)

        if at_verts:
            return (self.connected.dot(gradu).T / self.connected.sum(1).A.squeeze()).T
        return gradu

    @property
    @_memo
    def _facenorm_cross_edge(self):
        ppts = self.ppts
        fnorms = self.face_normals
        fe12 = np.cross(fnorms, ppts[:,1] - ppts[:,0])
        fe23 = np.cross(fnorms, ppts[:,2] - ppts[:,1])
        fe31 = np.cross(fnorms, ppts[:,0] - ppts[:,2])
        return fe12, fe23, fe31

    def geodesic_distance(self, verts, m=1.0, fem=False):
        """Minimum mesh geodesic distance (in mm) from each vertex in surface to any
        vertex in the collection verts.

        Geodesic distance is estimated using heat-based method (see 'Geodesics in Heat',
        Crane et al, 2012). The diffusion of heat along the mesh is simulated and then
        used to infer geodesic distance. The duration of the simulation is controlled
        by the parameter m. Larger values of m will smooth & regularize the distance
        computation. Smaller values of m will roughen and will usually increase error
        in the distance computation. The default value of 1.0 is probably pretty good.

        This function caches some data (sparse LU factorizations of the laplace-beltrami
        operator and the weighted adjacency matrix), so it will be much faster on
        subsequent runs.
<commit_msg>added surface interpolation function to polyutils<commit_after>from collections import OrderedDict
import numpy as np
from scipy.spatial import distance, Delaunay, cKDTree
from scipy import sparse
import scipy.sparse.linalg
import functools
import numexpr as ne

def _memo(fn):
    """Helper decorator memoizes the given zero-argument function.
    """
    @functools.wraps(fn)
    def memofn(self):
        if id(fn) not in self._cache:
            self._cache[id(fn)] = fn(self)
        return self._cache[id(fn)]

    return memofn

class Surface(object):
    def __init__(self, pts, polys):
        self.pts = pts
        self.polys = polys

        self._cache = dict()
        self._rlfac_solvers = dict()
        self._nLC_solvers = dict()

    @property
    @_memo
    def ppts(self):
        """3D matrix of points in each face: n faces x 3 points per face x 3 coords per point.
        """
        return self.pts[self.polys]
    
    @property
    @_memo
    def connected(self):
        """Sparse matrix of vertex-face associations.
        """
        npt = len(self.pts)
        npoly = len(self.polys)
        return sparse.coo_matrix((np.ones((3*npoly,)), # data
                                  (np.hstack(self.polys.T), # row
                                   np.tile(range(npoly),(1,3)).squeeze())), # col
                                 (npt, npoly)).tocsr() # size
    @property
    @_memo
    def adj(self):
        """Sparse vertex adjacency matrix.
        """
        npt = len(self.pts)
        npoly = len(self.polys)
        adj1 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,0], self.polys[:,1])), (npt,npt))
        adj2 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,0], self.polys[:,2])), (npt,npt))
        adj3 = sparse.coo_matrix((np.ones((npoly,)),
                                  (self.polys[:,1], self.polys[:,2])), (npt,npt))
        return (adj1 + adj2 + adj3).tocsr()
    
    @property
    @_memo
    def face_normals(self):
        """Normal vector for each face.
        """
        nnfnorms = np.cross(self.ppts[:,1] - self.ppts[:,0], self.ppts[:,2] - self.ppts[:,0])
        return nnfnorms / np.sqrt((nnfnorms**2).sum(1))[:,np.newaxis]

    @property
    @_memo
    def vertex_normals(self):
        """Normal vector for each vertex (average of normals for neighboring faces).
        """
        nnvnorms = np.nan_to_num(self.connected.dot(self.face_normals) / self.connected.sum(1)).A
        return nnvnorms / np.sqrt((nnvnorms**2).sum(1))[:,np.newaxis]

    @property
    @_memo
    def face_areas(self):
        """Area of each face.
        """
        nnfnorms = np.cross(self.ppts[:,1] - self.ppts[:,0], self.ppts[:,2] - self.ppts[:,0])
        return np.sqrt((nnfnorms**2).sum(-1))

    @property
    @_memo
    def cotangent_weights(self):
        """Cotangent of angle opposite each vertex in each face.
        """
        ppts = self.ppts
        cots1 = ((ppts[:,1]-ppts[:,0]) *
                 (ppts[:,2]-ppts[:,0])).sum(1) / np.sqrt((np.cross(ppts[:,1]-ppts[:,0],
                                                                   ppts[:,2]-ppts[:,0])**2).sum(1))
        cots2 = ((ppts[:,2]-ppts[:,1]) *
                 (ppts[:,0]-ppts[:,1])).sum(1) / np.sqrt((np.cross(ppts[:,2]-ppts[:,1],
                                                                   ppts[:,0]-ppts[:,1])**2).sum(1))
        cots3 = ((ppts[:,0]-ppts[:,2]) *
                 (ppts[:,1]-ppts[:,2])).sum(1) / np.sqrt((np.cross(ppts[:,0]-ppts[:,2],
                                                                   ppts[:,1]-ppts[:,2])**2).sum(1))

        # Then we have to sanitize the fuck out of everything..
        cots = np.vstack([cots1, cots2, cots3])
        cots[np.isinf(cots)] = 0
        cots[np.isnan(cots)] = 0
        # Clip to a reasonable range, these meshes are the worst
        cots = np.clip(cots, -100, 100)
        return cots

    @property
    @_memo
    def laplace_operator(self):
        """Laplace-Beltrami operator for this surface. A sparse adjacency matrix with
        edge weights determined by the cotangents of the angles opposite each edge.
        Returns a 4-tuple (B,D,W,V) where D is the 'lumped mass matrix', W is the weighted
        adjacency matrix, and V is a diagonal matrix that normalizes the adjacencies.
        The 'stiffness matrix', A, can be computed as V - W.
        
        See 'Discrete Laplace-Beltrami operators for shape analysis and segmentation'
        by Reuter et al., 2009 for details.
        """
        ## Lumped mass matrix
        D = self.connected.dot(self.face_areas) / 3.0

        ## Stiffness matrix
        npt = len(self.pts)
        cots1, cots2, cots3 = self.cotangent_weights
        # W is weighted adjacency matrix
        W1 = sparse.coo_matrix((cots1, (self.polys[:,1], self.polys[:,2])), (npt, npt))
        W2 = sparse.coo_matrix((cots2, (self.polys[:,2], self.polys[:,0])), (npt, npt))
        W3 = sparse.coo_matrix((cots3, (self.polys[:,0], self.polys[:,1])), (npt, npt))
        W = (W1 + W1.T + W2 + W2.T + W3 + W3.T).tocsr() / 2.0
        
        # V is sum of each col
        V = sparse.dia_matrix((np.array(W.sum(0)).ravel(),[0]), (npt,npt))
        
        # A is stiffness matrix
        #A = W - V # negative operator -- more useful in practice

        # For FEM:
        Be1 = sparse.coo_matrix((self.face_areas, (self.polys[:,1], self.polys[:,2])), (npt, npt))
        Be2 = sparse.coo_matrix((self.face_areas, (self.polys[:,2], self.polys[:,0])), (npt, npt))
        Be3 = sparse.coo_matrix((self.face_areas, (self.polys[:,0], self.polys[:,1])), (npt, npt))
        Bd = self.connected.dot(self.face_areas) / 6
        dBd = scipy.sparse.dia_matrix((Bd,[0]), (len(D),len(D)))
        B = (Be1 + Be1.T + Be2 + Be2.T + Be3 + Be3.T)/12 + dBd
        return B, D, W, V

    @property
    @_memo
    def avg_edge_length(self):
        """Average length of all edges in the surface.
        """
        adj = self.laplace_operator[1] # use laplace operator as adjacency matrix
        tadj = sparse.triu(adj, 1) # only entries above main diagonal, in coo format
        edgelens = np.sqrt(((self.pts[tadj.row] - self.pts[tadj.col])**2).sum(1))
        return edgelens.mean()

    def surface_gradient(self, vertvals, at_verts=True):
        """Gradient of a function with values vertvals at each vertex along surface.
        If at_verts, returns values at each vertex. Otherwise, returns values at each
        face.
        """
        pu = vertvals[self.polys]
        fe12, fe23, fe31 = [f.T for f in self._facenorm_cross_edge]
        pu1, pu2, pu3 = pu.T
        fa = self.face_areas
        #gradu = ((fe12.T * pu[:,2] +
        #          fe23.T * pu[:,0] +
        #          fe31.T * pu[:,1]) / (2 * self.face_areas)).T
        gradu = ne.evaluate("(fe12 * pu3 + fe23 * pu1 + fe31 * pu2) / (2 * fa)").T
        #gradu = np.nan_to_num(gradu)

        if at_verts:
            return (self.connected.dot(gradu).T / self.connected.sum(1).A.squeeze()).T
        return gradu

    def interp(self, verts, vals):
        """Interpolates a function between N knot points `verts` with the values `vals`.
        `vals` can be a D x N array to interpolate multiple functions with the same
        knot points.
        """
        from scikits.sparse.cholmod import cholesky
        B, D, W, V = self.laplace_operator
        npt = len(D)
        r = np.zeros((npt,))
        r[verts] = vals
        g = np.nonzero(D > 0)[0]
        spDinv = sparse.dia_matrix((D**-1,[0]), (npt,npt)).tocsr()
        L = spDinv.dot((V-W))
        notp = np.setdiff1d(np.arange(npt)[g], verts)

        L2r = (V-W).dot(L)#[notp][:,notp]
        #L2rfac = sparse.linalg.dsolve.factorized(L2r) # 7.2 s
        L2rfac = cholesky(L2r[notp][:,notp]) # 832 ms stock BLAS/LAPACK; 245 ms w/ MKL    
        vr = spDinv.dot(L2r.dot(r))
        phi = L2rfac(-D[notp] * vr[notp])

        tphi = np.zeros((npt,))
        tphi[notp] = phi
        tphi[verts] = vals
        
        return tphi

    @property
    @_memo
    def _facenorm_cross_edge(self):
        ppts = self.ppts
        fnorms = self.face_normals
        fe12 = np.cross(fnorms, ppts[:,1] - ppts[:,0])
        fe23 = np.cross(fnorms, ppts[:,2] - ppts[:,1])
        fe31 = np.cross(fnorms, ppts[:,0] - ppts[:,2])
        return fe12, fe23, fe31

    def geodesic_distance(self, verts, m=1.0, fem=False):
        """Minimum mesh geodesic distance (in mm) from each vertex in surface to any
        vertex in the collection verts.

        Geodesic distance is estimated using heat-based method (see 'Geodesics in Heat',
        Crane et al, 2012). The diffusion of heat along the mesh is simulated and then
        used to infer geodesic distance. The duration of the simulation is controlled
        by the parameter m. Larger values of m will smooth & regularize the distance
        computation. Smaller values of m will roughen and will usually increase error
        in the distance computation. The default value of 1.0 is probably pretty good.

        This function caches some data (sparse LU factorizations of the laplace-beltrami
        operator and the weighted adjacency matrix), so it will be much faster on
        subsequent runs.

{'fraction_non_alphanumeric': 0.14320608952510794, 'fraction_numerical': 0.022778913883208362, 'mean_word_length': 2.066004876349704, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 2, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5352627', 'n_tokens_mistral': 6124, 'n_tokens_neox': 5953, 'n_words': 1499}