initial refine_image functionality

Author: jcapriot

discretize/_extensions/tree.cpp

@@ -418,6 +418,48 @@ 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){
+    // 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;
+
+    // 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;
+            }
+            i_image += nz;
+        }
+        i_image += nz * ny;
+    }
+    if(subdivide){
+        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);
+        }
+    }
+
+}
+
 void Cell::divide(node_map_t& nodes, double* xs, double* ys, double* zs, bool balance, bool diag_balance){
     // Gaurd against dividing a cell that is already at the max level
     if (level == max_level){
@@ -896,6 +938,18 @@ 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;
+    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);
+}
+
+
 void Tree::finalize_lists(){
     for(int_t iz=0; iz<nz_roots; ++iz)
         for(int_t iy=0; iy<ny_roots; ++iy)

discretize/_extensions/tree.h

@@ -136,6 +136,7 @@ class Cell{
     void insert_cell(node_map_t &nodes, double *new_center, int_t p_level, double* xs, double *ys, double *zs, bool diag_balance=false);
 
     void refine_func(node_map_t& nodes, function test_func, double *xs, double *ys, double* zs, bool diag_balance=false);
+    void refine_image(node_map_t& nodes, double* image, int_t *shape_cells, double *xs, double*ys, double *zs, bool diagonal_balance=false);
 
     inline bool is_leaf(){ return children[0]==NULL;};
     void spawn(node_map_t& nodes, Cell *kids[8], double* xs, double *ys, double *zs);
@@ -217,6 +218,8 @@ class Tree{
 
     void refine_function(function test_func, bool diagonal_balance=false);
 
+    void refine_image(double* image, bool diagonal_balance=false);
+
     template <class T>
     void refine_geom(const T& geom, int_t p_level, bool diagonal_balance=false){
         for(int_t iz=0; iz<nz_roots; ++iz)

discretize/_extensions/tree.pxd

@@ -93,6 +93,8 @@ cdef extern from "tree.h":
 
         void refine_geom[T](const T&, int_t, bool)
 
+        void refine_image(double*, bool)
+
         void number()
         void initialize_roots()
         void insert_cell(double *new_center, int_t p_level, bool)

discretize/_extensions/tree_ext.pyx

@@ -7,6 +7,7 @@ from libc.stdlib cimport malloc, free
 from libcpp.vector cimport vector
 from libcpp cimport bool
 from numpy.math cimport INFINITY
+from pkg_resources import require
 
 from .tree cimport int_t, Tree as c_Tree, PyWrapper, Node, Edge, Face, Cell as c_Cell
 from . cimport geom
@@ -1197,6 +1198,41 @@ cdef class _TreeMesh:
         if finalize:
             self.finalize()
 
+    def refine_image(self, image, finalize=True, diagonal_balance=None):
+        """Refine using an ND image, ensuring that each cell contains exactly one unique value.
+
+        This function takes an N-dimensional image, defined on the underlying fine tensor mesh,
+        and recursively subdivides each cell if that cell contains more than 1 unique value in the
+        image. This is useful when using the TreeMesh to represent an exact compressed form of an input
+        model.
+
+        Parameters
+        ----------
+        image : (shape_cells) numpy.ndarray
+            Must have the same shape as the base tensor mesh (TreeMesh.shape_cells), as if every cell on this mesh was
+            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)
+        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)
+        if finalize:
+            self.finalize()
+
+
+
     def finalize(self):
         """Finalize the :class:`~discretize.TreeMesh`.