Fix bug in mesh refinement with multiple points

discretize/_extensions/tree_ext.pyx

@@ -1164,7 +1164,11 @@ cdef class _TreeMesh:
         Parameters
         ----------
         points : (N, dim) array_like
-        levels : (N) array_like of int
+            Array with coordinates of points on which cells will be inserted.
+        levels : int or (N) array_like of int
+            Integer indicating the refinement level for every point in ``points``,
+            or array of integers specifying the refinement level for each one of the
+            points.
         finalize : bool, optional
             Whether to finalize after inserting point(s)
         diagonal_balance : bool or None, optional
@@ -1191,6 +1195,20 @@ cdef class _TreeMesh:
         levels = np.require(np.atleast_1d(levels), dtype=np.int32, requirements='C')
         cdef int_t n_points = _check_first_dim_broadcast(points=points, levels=levels)
 
+        if levels.ndim != 1:
+            msg = (
+                f"Invalid levels argument with '{levels.ndim}' dimensions. "
+                "It must be a single integer or a 1D array."
+            )
+            raise ValueError(msg)
+        if levels.size > points.shape[0]:
+            msg = (
+                f"Invalid levels argument with '{levels.size}' elements. "
+                "It must be a single integer or an arry with the same amount of "
+                f"elements as points ('{points.shape[0]}')."
+            )
+            raise ValueError(msg)
+
         cdef double[:, :] cs = points
         cdef int[:] ls = levels
 
@@ -1200,11 +1218,12 @@ cdef class _TreeMesh:
             diagonal_balance = self._diagonal_balance
         cdef bool diag_balance = diagonal_balance
 
+        cdef int_t n_iters = max(ls.size, cs.shape[0])
         cdef int_t p_step = cs.shape[0] > 1
         cdef int_t l_step = ls.shape[0] > 1
         cdef int_t i_p=0, i_l=0
 
-        for i in range(ls.shape[0]):
+        for i in range(n_iters):
             l = _wrap_levels(ls[i_l], max_level)
             with self._tree_modify_lock:
                 self.tree.insert_cell(&cs[i_p, 0], l, diagonal_balance)

tests/tree/test_refine.py

@@ -239,6 +239,20 @@ def test_insert_errors():
     with pytest.raises(ValueError):
         mesh.insert_cells(x0s2d, [1, 1, 3], finalize=False)
 
+    # Incorrect dimension of levels
+    levels_nd = np.array([[1, 2]])
+    points = np.array([[0.1, 0.1], [0.5, 0.5]])
+    msg = re.escape("Invalid levels argument with '2' dimensions")
+    with pytest.raises(ValueError, match=msg):
+        mesh.insert_cells(points, levels_nd, finalize=False)
+
+    # Multiple levels on a single point
+    levels_large = np.array([3, 2, 1])
+    points = np.array([[0.1, 0.1]])
+    msg = re.escape("Invalid levels argument with '3' elements")
+    with pytest.raises(ValueError, match=msg):
+        mesh.insert_cells(points, levels_large, finalize=False)
+
 
 def test_refine_triang_prism():
     xyz = np.array(