per_vertex_normal contribution adjusted

src/gpytoolbox/per_vertex_normals.py

@@ -2,9 +2,10 @@
 from .per_face_normals import per_face_normals
 from .doublearea import doublearea
 from scipy.sparse import csr_matrix
+from scipy.spatial import KDTree
 
 
-def compute_angles(V,F):
+def compute_angles(V, F):
     """Computes the angle at each vertex in a triangle mesh.
 
     Parameters
@@ -61,8 +62,9 @@ def compute_angles(V,F):
     return angles
 
 
-def per_vertex_normals(V,F,weights='area'):
-    """Compute per-vertex normal vectors for a triangle mesh with different weighting options.
+def per_vertex_normals(V, F, weights='area', correct_invalid_normals=True):
+    """Compute per-vertex normal vectors for a triangle mesh with different weighting options,
+    with an option to correct invalid normals (True) or omit them (False).
 
     Parameters
     ----------
@@ -71,16 +73,20 @@ def per_vertex_normals(V,F,weights='area'):
     F : (m,d) numpy int array
         face index list of a triangle mesh
     weights : str, optional
-        The type of weighting to use ('area', 'angles', 'uniform'), default is 'area'
+        the type of weighting to use ('area', 'angle', 'uniform'), default is 'area'
+    correct_invalid_normals : bool, optional
+        if True, invalid normals (NaN, inf, zero vectors) will be corrected, either by calculating
+        them based on nearby valid normals, or -in extreme cases- replacing them with a default normal vector
+        if False, the invalid normals are omitted
 
     Returns
     -------
     N : (n,d) numpy array
-        Matrix of per-vertex normals
+        matrix of per-vertex normals
 
     See Also
     --------
-    per_vertex_normals, per_face_normals.
+    per_face_normals
 
     Examples
     --------
@@ -90,10 +96,16 @@ def per_vertex_normals(V,F,weights='area'):
     n = per_vertex_normals(v,f)
     ```
     """
+    # Ensure vertices are in float64
+    V = V.astype(np.float64)  
+    # Ensure faces are in int32
+    F = F.astype(np.int32)  
+
+    # Compute face normals
+    face_normals = per_face_normals(V, F, unit_norm=True)
+
     if weights=="area":
         dim = V.shape[1]
-        # First compute face normals
-        face_normals = per_face_normals(V,F,unit_norm=True)
         # We blur these normals onto vertices, weighing by area
         areas = doublearea(V,F)
         vals = np.concatenate([areas for _ in range(dim)])
@@ -104,47 +116,67 @@ def per_vertex_normals(V,F,weights='area'):
         # I = np.concatenate((F[:,0],F[:,1],F[:,2]))
 
         weight_mat = csr_matrix((vals,(I,J)),shape=(V.shape[0],F.shape[0]))
-
-        vertex_normals = weight_mat @ face_normals
-        # Now, normalize
-        N = vertex_normals/np.tile(np.linalg.norm(vertex_normals,axis=1)[:,None],(1,dim))
-
-
-    elif weights=="angles":
-        # Ensure vertices are in float64
-        V = V.astype(np.float64)  
-        # Ensure faces are in int32
-        F = F.astype(np.int32)    
-
-        # Compute face normals
-        normals = per_face_normals(V,F)
+        weighted_normals = weight_mat @ face_normals
+
+    elif weights=="angle":
         # Compute angles at each vertex
         angles = compute_angles(V,F)
         # Weight the face normals by the angles at each vertex
         weighted_normals = np.zeros_like(V)
         for i in range(3):
-            np.add.at(weighted_normals, F[:, i], normals * angles[:, i, np.newaxis])
-
-        # Calculate norms
-        norms = np.linalg.norm(weighted_normals, axis=1, keepdims=True)
-        # Avoid division by zero by setting zero norms to a very small number
-        norms[norms == 0] = np.finfo(float).eps
-        # Normalize the vertex normals
-        N = weighted_normals / norms
-
-
+            np.add.at(weighted_normals, F[:, i], face_normals * angles[:, i, np.newaxis])
+
     elif weights=="uniform":
-        # Compute face normals
-        face_normals = per_face_normals(V, F)
-        # Initialize vertex normals
-        vertex_normals = np.zeros_like(V)
+        # Compute (non-)weighted normals uniformly
+        weighted_normals = np.zeros_like(V)
         # Accumulate face normals to vertices
         for i in range(3):
-            np.add.at(vertex_normals, F[:, i], face_normals)
+            np.add.at(weighted_normals, F[:, i], face_normals)
+
+    # Calculate norms
+    norms = np.linalg.norm(weighted_normals, axis=1, keepdims=True)
+    # Identify indices with problematic normals (nan, inf, negative norms)
+    problematic_indices = np.where(np.isnan(norms) | np.isinf(norms) | (norms <= 0))[0]
+    # Identify indices with valid normals
+    valid_indices = np.where(np.isfinite(norms) & (norms > 0))[0]
+
+    # Normalize valid normals
+    N = np.where(norms > 0, weighted_normals / norms, np.zeros_like(weighted_normals))
+
+    # Ensure all normals are unit vectors (handling edge cases)
+    if problematic_indices.size > 0 and correct_invalid_normals:
+        print("Handling problematic indices")
 
-        # Normalize
-        norms = np.linalg.norm(vertex_normals, axis=1, keepdims=True)
-        norms[norms == 0] = np.finfo(float).eps
-        N = vertex_normals / norms
+        # Build KDTree using only valid vertices
+        if valid_indices.size > 0:
+            valid_tree = KDTree(V[valid_indices])
+
+        for idx in problematic_indices:
+            print(f"Handling vertex {idx}")
+
+            # Use nearest valid normal for invalid normals
+            if valid_indices.size > 0:
+                # Query the nearest valid normal
+                dist, pos = valid_tree.query(V[idx], k=1)  
+                nearest_valid_idx = valid_indices[pos]  
+                # Assign the nearest valid normal to the current problematic normal
+                if np.isfinite(dist) and dist > 0:
+                    N[idx] = N[nearest_valid_idx]
+                    norms[idx] = np.linalg.norm(N[nearest_valid_idx], keepdims=True)
+                    print(f"Updated normal for vertex {idx} from valid neighbor.")
+                # Else assign a default normal
+                else:
+                    N[idx] = np.array([1, 0, 0])
+                    norms[idx] = 1.0
+                    print(f"Assigned default normal due to lack of valid neighbors.")
+
+            # Else assign a default normal
+            else:
+                N[idx] = np.array([1, 0, 0])
+                norms[idx] = 1.0
+                print(f"No valid vertices available, default normal assigned.")
+
+        # Normalize only the previously problematic normals
+        N[problematic_indices] = N[problematic_indices] / np.maximum(norms[problematic_indices], np.finfo(float).eps)
 
     return N