Implement own np.take and remove 'reflected lists'

Author: FranzBangar

src/classy_blocks/optimize/iteration.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import List
 
-from classy_blocks.util.constants import VBIG, VSMALL
+from classy_blocks.util.constants import TOL, VBIG, VSMALL
 from classy_blocks.util.tools import report
 
 
@@ -129,4 +129,8 @@ def converged(self) -> bool:
             print("Tolerance reached, stopping optimization.")
             return True
 
+        if self.iterations[-1].final_quality < TOL:
+            print("Nothing left to optimize.")
+            return True
+
         return False

src/classy_blocks/optimize/junction.py

@@ -1,6 +1,8 @@
 import dataclasses
 from typing import List, Optional, Set
 
+import numpy as np
+
 from classy_blocks.base.exceptions import ClampExistsError
 from classy_blocks.cbtyping import NPPointListType, NPPointType
 from classy_blocks.optimize.cell import CellBase, HexCell
@@ -60,4 +62,4 @@ def quality(self) -> float:
         else:
             quality_function = get_quad_quality
 
-        return sum(quality_function(self.points, cell.indexes) for cell in self.cells)
+        return sum(quality_function(self.points, np.array(cell.indexes, dtype=np.int32)) for cell in self.cells)

src/classy_blocks/optimize/quality.py

@@ -1,11 +1,14 @@
-from typing import List, Tuple
+from typing import Tuple
 
 import numba  # type:ignore
 import numpy as np
+from nptyping import Int32, NDArray, Shape
 
 from classy_blocks.cbtyping import NPPointListType, NPPointType, NPVectorType
 from classy_blocks.util.constants import VSMALL
 
+NPIndexType = NDArray[Shape["*, 1"], Int32]
+
 
 @numba.jit(nopython=True, cache=True)
 def scale_quality(base: float, exponent: float, factor: float, value: float) -> float:
@@ -27,6 +30,19 @@ def scale_aspect(ratio: float) -> float:
     return scale_quality(3, 2.5, 3, np.log10(ratio))
 
 
+@numba.jit(nopython=True, cache=True)
+def take(points: NPPointListType, indexes: NPIndexType):
+    n_points = len(indexes)
+    dim = points.shape[1]
+    result = np.empty((n_points, dim), dtype=points.dtype)
+
+    for i in range(n_points):
+        for j in range(dim):
+            result[i, j] = points[indexes[i], j]
+
+    return result
+
+
 @numba.jit(nopython=True, cache=True)
 def get_center_point(points: NPPointListType) -> NPPointType:
     return np.sum(points, axis=0) / len(points)
@@ -103,8 +119,8 @@ def get_quad_angle_quality(quad_points: NPPointListType) -> float:
 
 
 @numba.jit(nopython=True, cache=True)
-def get_quad_quality(grid_points: NPPointListType, cell_indexes: List[int]) -> float:
-    cell_points = np.take(grid_points, cell_indexes, axis=0)
+def get_quad_quality(grid_points: NPPointListType, cell_indexes: NPIndexType) -> float:
+    cell_points = take(grid_points, cell_indexes)
     cell_center, cell_normal, cell_aspect = get_quad_normal(cell_points)
 
     # non-ortho
@@ -119,8 +135,8 @@ def get_quad_quality(grid_points: NPPointListType, cell_indexes: List[int]) -> f
 
 
 @numba.jit(nopython=True, cache=True)
-def get_hex_quality(grid_points: NPPointListType, cell_indexes: List[int]) -> float:
-    cell_points = np.take(grid_points, cell_indexes, axis=0)
+def get_hex_quality(grid_points: NPPointListType, cell_indexes: NPIndexType) -> float:
+    cell_points = take(grid_points, cell_indexes)
     cell_center = get_center_point(cell_points)
 
     side_indexes = np.array([[0, 1, 2, 3], [7, 6, 5, 4], [4, 0, 3, 7], [6, 2, 1, 5], [0, 4, 5, 1], [7, 3, 2, 6]])
@@ -136,7 +152,7 @@ def get_hex_quality(grid_points: NPPointListType, cell_indexes: List[int]) -> fl
         # For this kind of optimization it is quite sufficient to take
         # only the latter as it's not much different and we'll optimize
         # other cells too.
-        side_points = np.take(cell_points, side, axis=0)
+        side_points = take(cell_points, side)
         side_center, side_normal, side_aspect = get_quad_normal(side_points)
         center_vector = cell_center - side_center