Merge branch 'reviewer-edits' of github.com:MunchLab/ect into reviewer-edits

Author: lizliz

src/ect/ect.py

@@ -8,6 +8,21 @@
 from .results import ECTResult
 
 
+def _thresholds_are_uniform(thresholds: np.ndarray) -> bool:
+    thresholds = np.asarray(thresholds, dtype=float)
+    if thresholds.ndim != 1:
+        raise ValueError("thresholds must be a 1-dimensional array")
+    n = thresholds.size
+    if n <= 1:
+        return True
+    diffs = np.diff(thresholds)
+    first = diffs[0]
+    if first == 0.0:
+        return bool(np.all(diffs == 0.0))
+    tol = 1e-12 * max(1.0, abs(first))
+    return bool(np.all(np.abs(diffs - first) <= tol))
+
+
 class ECT:
     """
     A class to calculate the Euler Characteristic Transform (ECT) from an input :class:`ect.embed_complex.EmbeddedComplex`,
@@ -55,6 +70,24 @@ def __init__(
         self.bound_radius = bound_radius
         self.thresholds = thresholds
         self.dtype = dtype
+        self._thresholds_validated = False
+        if self.thresholds is not None:
+            self.thresholds = np.asarray(self.thresholds, dtype=float)
+            if self.thresholds.ndim != 1:
+                raise ValueError("thresholds must be a 1-dimensional array")
+            self._thresholds_validated = True
+        if num_thresh is not None:
+            self.is_uniform = True
+        else:
+            self.is_uniform = False
+            if self.thresholds is None:
+                raise ValueError(
+                    "thresholds must be provided if num_thresh is not provided"
+                )
+            if not _thresholds_are_uniform(self.thresholds):
+                raise ValueError(
+                    "thresholds must be uniform if num_thresh is not provided"
+                )
 
     def _ensure_directions(self, graph_dim, theta=None):
         """Ensures directions is a valid Directions object of correct dimension"""
@@ -97,11 +130,14 @@ def _ensure_thresholds(self, graph, override_bound_radius=None):
                 or graph.get_bounding_radius()
             )
             self.thresholds = np.linspace(-radius, radius, self.num_thresh, dtype=float)
+            self.is_uniform = True
+            self._thresholds_validated = True
         else:
-            # validate and convert existing thresholds
-            self.thresholds = np.asarray(self.thresholds, dtype=float)
-            if self.thresholds.ndim != 1:
-                raise ValueError("thresholds must be a 1-dimensional array")
+            if not self._thresholds_validated:
+                self.thresholds = np.asarray(self.thresholds, dtype=float)
+                if self.thresholds.ndim != 1:
+                    raise ValueError("thresholds must be a 1-dimensional array")
+                self._thresholds_validated = True
 
     def calculate(
         self,
@@ -124,21 +160,32 @@ def _compute_ect(
         cell_vertex_pointers, cell_vertex_indices_flat, cell_euler_signs, N = (
             graph._build_incidence_csr()
         )
-        thresholds = np.asarray(thresholds, dtype=np.float64)
+        thresholds = np.asarray(thresholds, dtype=np.float32)
 
         V = directions.vectors
         X = graph.coord_matrix
         H = X @ V.T  # (N, m)
         H_T = np.ascontiguousarray(H.T)  # (m, N) for contiguous per-direction rows
 
-        out64 = _ect_all_dirs(
-            H_T,
-            cell_vertex_pointers,
-            cell_vertex_indices_flat,
-            cell_euler_signs,
-            thresholds,
-            N,
-        )
+        is_uniform = bool(self.is_uniform) and thresholds[0] != thresholds[-1]
+        if is_uniform:
+            out64 = _ect_all_dirs_uniform(
+                H_T,
+                cell_vertex_pointers,
+                cell_vertex_indices_flat,
+                cell_euler_signs,
+                thresholds,
+                N,
+            )
+        else:
+            out64 = _ect_all_dirs_search(
+                H_T,
+                cell_vertex_pointers,
+                cell_vertex_indices_flat,
+                cell_euler_signs,
+                thresholds,
+                N,
+            )
         if dtype == np.int32:
             return out64.astype(np.int32)
         return out64
@@ -246,3 +293,126 @@ def _ect_all_dirs(
             ect_values[dir_idx, thresh_idx] = euler_prefix[rank_cursor]
 
     return ect_values
+
+
+@njit(cache=True, parallel=True)
+def _ect_all_dirs_uniform(
+    heights_by_direction,
+    cell_vertex_pointers,
+    cell_vertex_indices_flat,
+    cell_euler_signs,
+    threshold_values,
+    num_vertices,
+):
+    num_directions = heights_by_direction.shape[0]
+    num_thresholds = threshold_values.shape[0]
+    t_min = threshold_values[0] if num_thresholds > 0 else 0.0
+    t_max = threshold_values[-1] if num_thresholds > 0 else 0.0
+    span = t_max - t_min
+    inv_span = 1.0 / span
+    n_minus_1 = num_thresholds - 1
+
+    ect_values = np.empty((num_directions, num_thresholds), dtype=np.int64)
+
+    for dir_idx in prange(num_directions):
+        heights = heights_by_direction[dir_idx]
+
+        diff = np.zeros(num_thresholds, dtype=np.int64)
+        vertex_thresh_index = np.empty(num_vertices, dtype=np.int64)
+
+        for v in range(num_vertices):
+            h = heights[v]
+            u = (h - t_min) * inv_span
+            idx = int(np.ceil(u * n_minus_1))
+            if idx < 0:
+                idx = 0
+            elif idx >= num_thresholds:
+                idx = num_thresholds
+
+            vertex_thresh_index[v] = idx
+            if idx < num_thresholds:
+                diff[idx] += 1
+
+        num_cells = cell_vertex_pointers.shape[0] - 1
+
+        for cell_idx in range(num_cells):
+            start = cell_vertex_pointers[cell_idx]
+            end = cell_vertex_pointers[cell_idx + 1]
+
+            entrance_idx = -1
+            for k in range(start, end):
+                v = cell_vertex_indices_flat[k]
+                t_idx = vertex_thresh_index[v]
+                if t_idx > entrance_idx:
+                    entrance_idx = t_idx
+
+            if 0 <= entrance_idx < num_thresholds:
+                diff[entrance_idx] += cell_euler_signs[cell_idx]
+
+        running = 0
+        for j in range(num_thresholds):
+            running += diff[j]
+            ect_values[dir_idx, j] = running
+
+    return ect_values
+
+
+@njit(cache=True, parallel=True)
+def _ect_all_dirs_search(
+    heights_by_direction,
+    cell_vertex_pointers,
+    cell_vertex_indices_flat,
+    cell_euler_signs,
+    threshold_values,
+    num_vertices,
+):
+    num_directions = heights_by_direction.shape[0]
+    num_thresholds = threshold_values.shape[0]
+
+    ect_values = np.empty((num_directions, num_thresholds), dtype=np.int64)
+
+    for dir_idx in prange(num_directions):
+        heights = heights_by_direction[dir_idx]
+
+        diff = np.zeros(num_thresholds, dtype=np.int64)
+        vertex_thresh_index = np.empty(num_vertices, dtype=np.int64)
+
+        for v in range(num_vertices):
+            h = heights[v]
+
+            left = 0
+            right = num_thresholds
+            while left < right:
+                mid = (left + right) // 2
+                if threshold_values[mid] >= h:
+                    right = mid
+                else:
+                    left = mid + 1
+            idx = left
+
+            vertex_thresh_index[v] = idx
+            if idx < num_thresholds:
+                diff[idx] += 1
+
+        num_cells = cell_vertex_pointers.shape[0] - 1
+
+        for cell_idx in range(num_cells):
+            start = cell_vertex_pointers[cell_idx]
+            end = cell_vertex_pointers[cell_idx + 1]
+
+            entrance_idx = -1
+            for k in range(start, end):
+                v = cell_vertex_indices_flat[k]
+                t_idx = vertex_thresh_index[v]
+                if t_idx > entrance_idx:
+                    entrance_idx = t_idx
+
+            if 0 <= entrance_idx < num_thresholds:
+                diff[entrance_idx] += cell_euler_signs[cell_idx]
+
+        running = 0
+        for j in range(num_thresholds):
+            running += diff[j]
+            ect_values[dir_idx, j] = running
+
+    return ect_values

src/ect/embed_complex.py

@@ -9,11 +9,6 @@
 from sklearn.decomposition import PCA
 
 from .utils.naming import next_vert_name
-from .utils.face_check import (
-    point_in_polygon,
-    validate_face_embedding,
-    validate_edge_embedding,
-)
 from .validation import EmbeddingValidator, ValidationRule
 
 
@@ -55,6 +50,7 @@ def __init__(self, validate_embedding=False, embedding_tol=1e-10):
         self._node_to_index = {}
         self._coord_matrix = None
         self.cells = defaultdict(list)
+        self._incidence_csr_cache = None
 
         self.validate_embedding = validate_embedding
         self.embedding_tol = embedding_tol
@@ -69,6 +65,12 @@ def edge_checker(v1_idx: int, v2_idx: int) -> bool:
 
         self._validator = EmbeddingValidator(embedding_tol, edge_checker)
 
+    def _invalidate_incidence_csr_cache(self) -> None:
+        self._incidence_csr_cache = None
+
+    def precompute_incidence_csr(self) -> tuple:
+        return self._build_incidence_csr()
+
     @property
     def coord_matrix(self):
         """
@@ -192,6 +194,7 @@ def add_node(self, node_id, coord):
         self._node_list.append(node_id)
         self._node_to_index[node_id] = len(self._node_list) - 1
         super().add_node(node_id)
+        self._invalidate_incidence_csr_cache()
 
     def add_nodes_from_dict(self, nodes_with_coords: Dict[Union[str, int], np.ndarray]):
         """Add multiple vertices to the complex.
@@ -232,6 +235,7 @@ def add_edge(self, node_id1, node_id2):
             raise ValueError(node_result.message)
 
         super().add_edge(node_id1, node_id2)
+        self._invalidate_incidence_csr_cache()
 
     def add_cell(
         self,
@@ -302,6 +306,7 @@ def add_cell(
             self.add_edge(cell_vertices[0], cell_vertices[1])
 
         self.cells[dim].append(cell_indices)
+        self._invalidate_incidence_csr_cache()
 
     def enable_embedding_validation(self, tol: float = 1e-10):
         """
@@ -434,6 +439,7 @@ def add_cycle(self, coord_matrix):
         new_names = next_vert_name(self._node_list[-1] if self._node_list else 0, n)
         self.add_nodes_from(zip(new_names, coord_matrix))
         self.add_edges_from([(new_names[i], new_names[(i + 1) % n]) for i in range(n)])
+        self.precompute_incidence_csr()
 
     def get_center(self, method: str = "bounding_box") -> np.ndarray:
         """
@@ -967,6 +973,9 @@ def _build_incidence_csr(self) -> tuple:
         Example: takes the complex [(1,3),(2,4),(1,2,3)] and returns [(0,2,4,7),(1,3,2,4,1,2,3),(-1,-1,1),4]
 
         """
+        if self._incidence_csr_cache is not None:
+            return self._incidence_csr_cache
+
         n_vertices = len(self.node_list)
 
         cells_by_dimension = {}
@@ -1010,12 +1019,14 @@ def _build_incidence_csr(self) -> tuple:
                 cell_vertex_pointers[cell_index] = len(cell_vertex_indices_flat)
 
         cell_vertex_indices_flat = np.asarray(cell_vertex_indices_flat, dtype=np.int32)
-        return (
+        out = (
             cell_vertex_pointers,
             cell_vertex_indices_flat,
             cell_euler_signs,
             n_vertices,
         )
+        self._incidence_csr_cache = out
+        return out
 
 
 EmbeddedGraph = EmbeddedComplex