fix implementation and add tests

Author: jcapriot

discretize/_extensions/tree.cpp

@@ -418,46 +418,42 @@ void Cell::refine_func(node_map_t& nodes, function test_func, double *xs, double
     }
 }
 
-void Cell::refine_image(node_map_t& nodes, double* image, int_t *shape_cells, double *xs, double*ys, double *zs, bool diagonal_balance){
+void Cell::refine_image(node_map_t& nodes, double* image, int_t *shape_cells, double *xs, double*ys, double *zs, bool diag_balance){
     // early exit if my level is higher than or equal to target
     if (level == max_level){
         return;
     }
-    int_t start_ix = points[0].location_ind[0]/2;
-    int_t start_iy = points[0].location_ind[1]/2;
-    int_t start_iz = n_dim == 2 ? 0 : points[0].location_ind[2]/2;
-    int_t nx = shape_cells[0]
-    int_t ny = shape_cells[1]
-    int_t nz = shape_cells[2]
-
-    // same as 2**(max_level - level), but quicker since power of 2 and integers
-    int_t span = 1<<(max_level - level);
-    int_t span_z = n_dim == 2 ? 1 : span;
-    int_t i_image = (nz * ny) * start_ix + nz * start_iy + start_iz;
-    double val_start = image[0];
-    bool subdivide = false;
+    int_t start_ix = points[0]->location_ind[0]/2;
+    int_t start_iy = points[0]->location_ind[1]/2;
+    int_t start_iz = n_dim == 2 ? 0 : points[0]->location_ind[2]/2;
+    int_t end_ix = points[3]->location_ind[0]/2;
+    int_t end_iy = points[3]->location_ind[1]/2;
+    int_t end_iz = n_dim == 2? 1 : points[7]->location_ind[2]/2;
+    int_t nx = shape_cells[0];
+    int_t ny = shape_cells[1];
+    int_t nz = shape_cells[2];
+
+    int_t i_image = (nx * ny) * start_iz + nx * start_iy + start_ix;
+    double val_start = image[i_image];
+    bool all_unique = true;
 
     // if any of the image data contained in the cell are different, subdivide myself
-    for(int_t ix=0; ix<span && !subdivide; ++ix){
-        for(int_t iy=0; iy<span && !subdivide; ++iy){
-            for(int_t iz=0; iz<span_z && !subdivide; ++iz){
-                subdivide = image[i_image] != val_start;
-                i_image += 1;
+    for(int_t iz=start_iz; iz<end_iz && all_unique; ++iz)
+        for(int_t iy=start_iy; iy<end_iy && all_unique; ++iy)
+            for(int_t ix=start_ix; ix<end_ix && all_unique; ++ix){
+                i_image = (nx * ny) * iz + nx * iy + ix;
+                all_unique = image[i_image] == val_start;
             }
-            i_image += nz;
-        }
-        i_image += nz * ny;
-    }
-    if(subdivide){
+
+    if(!all_unique){
         if(is_leaf()){
             divide(nodes, xs, ys, zs, true, diag_balance);
         }
         // recurse into children
         for(int_t i = 0; i < (1<<n_dim); ++i){
-            children[i]->refine_image(nodes, geom, p_level, xs, ys, zs, diag_balance);
+            children[i]->refine_image(nodes, image, shape_cells, xs, ys, zs, diag_balance);
         }
     }
-
 }
 
 void Cell::divide(node_map_t& nodes, double* xs, double* ys, double* zs, bool balance, bool diag_balance){
@@ -938,15 +934,15 @@ void Tree::refine_function(function test_func, bool diagonal_balance){
                 roots[iz][iy][ix]->refine_func(nodes, test_func, xs, ys, zs, diagonal_balance);
 };
 
-void Tree:refine_image(double *image, bool diagonal_balance){
-    int_t[3] shape_cells;
+void Tree::refine_image(double *image, bool diagonal_balance){
+    int_t shape_cells[3];
     shape_cells[0] = nx/2;
     shape_cells[1] = ny/2;
     shape_cells[2] = nz/2;
     for(int_t iz=0; iz<nz_roots; ++iz)
         for(int_t iy=0; iy<ny_roots; ++iy)
             for(int_t ix=0; ix<nx_roots; ++ix)
-                roots[iz][iy][ix]->refine_image(nodes, image, xs, ys, zs, diagonal_balance);
+                roots[iz][iy][ix]->refine_image(nodes, image, shape_cells, xs, ys, zs, diagonal_balance);
 }
 
 

discretize/_extensions/tree_ext.pyx

@@ -7,7 +7,7 @@ cimport numpy as np
 from libc.stdlib cimport malloc, free
 from libcpp.vector cimport vector
 from libcpp cimport bool
-from libc.math cimport INFINITYs
+from libc.math cimport INFINITY
 
 from .tree cimport int_t, Tree as c_Tree, PyWrapper, Node, Edge, Face, Cell as c_Cell
 from . cimport geom
@@ -1229,21 +1229,30 @@ cdef class _TreeMesh:
             refined to it's maximum level.
 
         """
-        cdef int max_level = self.max_level
         if diagonal_balance is None:
             diagonal_balance = self._diagonal_balance
         cdef bool diag_balance = diagonal_balance
 
-        image = self._require_ndarray_with_dim('image', image, ndim=self.dim, dtype=np.float64)
+        image = np.require(image, dtype=np.float64, requirements="F")
+        cdef size_t n_expected = np.prod(self.shape_cells)
+        if image.size != n_expected:
+            raise ValueError(
+                f"image array size: {image.size} must match the total number of cells in the base tensor mesh: {n_expected}"
+            )
+        if image.ndim == 1:
+            image = image.reshape(self.shape_cells, order="F")
+
         if image.shape != self.shape_cells:
             raise ValueError(
                 f"image array shape: {image.shape} must match the base cell shapes: {self.shape_cells}"
             )
         if self.dim == 2:
             image = image[..., None]
 
-        cdef double[:,:,::1] image_dat = image
-        self.tree.refine_image(&image[0, 0, 0], diagonal_balance)
+        cdef double[::1,:,:] image_dat = image
+        
+        with self._tree_modify_lock:
+            self.tree.refine_image(&image_dat[0, 0, 0], diag_balance)
         if finalize:
             self.finalize()
 

tests/tree/test_refine.py

@@ -1,5 +1,7 @@
+import re
 import discretize
 import numpy as np
+import numpy.testing as npt
 import pytest
 from discretize.tests import assert_cell_intersects_geometric
 
@@ -490,3 +492,149 @@ def test_refine_plane3D():
     mesh2.refine_triangle(tris, -1)
 
     assert mesh1.equals(mesh2)
+
+
+def _make_quadrant_model(mesh, order):
+    shape_cells = mesh.shape_cells
+    model = np.zeros(shape_cells, order="F" if order == "flat" else order)
+    if mesh.dim == 2:
+        model[: shape_cells[0] // 2, : shape_cells[1] // 2] = 1.0
+        model[: shape_cells[0] // 4, : shape_cells[1] // 4] = 0.5
+    else:
+        model[: shape_cells[0] // 2, : shape_cells[1] // 2, : shape_cells[2] // 2] = 1.0
+        model[: shape_cells[0] // 4, : shape_cells[1] // 4, : shape_cells[2] // 4] = 0.5
+    if order == "flat":
+        model = model.reshape(-1, order="F")
+    return model
+
+
+@pytest.mark.parametrize(
+    "tens_inp",
+    [
+        dict(h=[16, 16]),
+        dict(h=[16, 32]),
+        dict(h=[32, 16]),
+        dict(h=[16, 16, 16]),
+        dict(h=[16, 16, 8]),
+        dict(h=[16, 8, 16]),
+        dict(h=[8, 16, 16]),
+        dict(h=[8, 8, 16]),
+        dict(h=[8, 16, 8]),
+        dict(h=[16, 8, 8]),
+    ],
+    ids=[
+        "16x16",
+        "16x32",
+        "32x16",
+        "16x16x16",
+        "16x16x8",
+        "16x8x16",
+        "8x16x16",
+        "8x8x16",
+        "8x16x8",
+        "16x8x8",
+    ],
+)
+def test_refine_image_input_ordering(tens_inp):
+    base_mesh = discretize.TensorMesh(**tens_inp)
+    model_0 = _make_quadrant_model(base_mesh, order="flat")
+    model_1 = _make_quadrant_model(base_mesh, order="C")
+    model_2 = _make_quadrant_model(base_mesh, order="F")
+
+    tree0 = discretize.TreeMesh(base_mesh.h, base_mesh.origin)
+    tree0.refine_image(model_0)
+
+    tree1 = discretize.TreeMesh(base_mesh.h, base_mesh.origin)
+    tree1.refine_image(model_1)
+
+    tree2 = discretize.TreeMesh(base_mesh.h, base_mesh.origin)
+    tree2.refine_image(model_2)
+
+    assert tree0.n_cells == tree1.n_cells == tree2.n_cells
+
+    for cell0, cell1, cell2 in zip(tree0, tree1, tree2):
+        assert cell0.nodes == cell1.nodes == cell2.nodes
+
+
+@pytest.mark.parametrize(
+    "tens_inp",
+    [
+        dict(h=[16, 16]),
+        dict(h=[16, 32]),
+        dict(h=[32, 16]),
+        dict(h=[16, 16, 16]),
+        dict(h=[16, 16, 8]),
+        dict(h=[16, 8, 16]),
+        dict(h=[8, 16, 16]),
+        dict(h=[8, 8, 16]),
+        dict(h=[8, 16, 8]),
+        dict(h=[16, 8, 8]),
+    ],
+    ids=[
+        "16x16",
+        "16x32",
+        "32x16",
+        "16x16x16",
+        "16x16x8",
+        "16x8x16",
+        "8x16x16",
+        "8x8x16",
+        "8x16x8",
+        "16x8x8",
+    ],
+)
+@pytest.mark.parametrize(
+    "model_func",
+    [
+        lambda mesh: np.zeros(mesh.n_cells),
+        lambda mesh: np.arange(mesh.n_cells, dtype=float),
+        lambda mesh: _make_quadrant_model(mesh, order="flat"),
+    ],
+    ids=["constant", "full", "quadrant"],
+)
+def test_refine_image(tens_inp, model_func):
+    base_mesh = discretize.TensorMesh(**tens_inp)
+    model = model_func(base_mesh)
+    mesh = discretize.TreeMesh(base_mesh.h, base_mesh.origin, diagonal_balance=False)
+    mesh.refine_image(model)
+
+    # for every cell in the tree mesh, all aligned cells in the tensor mesh
+    # should have a single unique value.
+    # quickest way is to generate a volume interp operator and look at indices in the
+    # csr matrix
+    interp_mat = discretize.utils.volume_average(base_mesh, mesh)
+
+    # ensure in canonical form:
+    interp_mat.sum_duplicates()
+    interp_mat.sort_indices()
+    assert interp_mat.has_canonical_format
+
+    model = model.reshape(-1, order="F")
+    for row in interp_mat:
+        vals = model[row.indices]
+        npt.assert_equal(vals, vals[0])
+
+
+def test_refine_image_bad_size():
+    mesh = discretize.TreeMesh([32, 32])
+    model = np.zeros(32 * 32 + 1)
+    base_cells = np.prod(mesh.shape_cells)
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            f"image array size: {len(model)} must match the total number of cells in the base tensor mesh: {base_cells}"
+        ),
+    ):
+        mesh.refine_image(model)
+
+
+def test_refine_image_bad_shape():
+    mesh = discretize.TreeMesh([32, 32])
+    model = np.zeros((16, 64))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            f"image array shape: {model.shape} must match the base cell shapes: {mesh.shape_cells}"
+        ),
+    ):
+        mesh.refine_image(model)