From 7252aa39e199bb8653615066e380bfe4462b18ed Mon Sep 17 00:00:00 2001 From: Joseph Capriotti Date: Tue, 19 Nov 2024 11:16:53 -0700 Subject: [PATCH 1/4] initial refine_image functionality --- discretize/_extensions/tree.cpp | 54 +++++++++++++++++++++++++++++ discretize/_extensions/tree.h | 3 ++ discretize/_extensions/tree.pxd | 2 ++ discretize/_extensions/tree_ext.pyx | 36 +++++++++++++++++++ 4 files changed, 95 insertions(+) diff --git a/discretize/_extensions/tree.cpp b/discretize/_extensions/tree.cpp index 8899490d1..bad671d3c 100644 --- a/discretize/_extensions/tree.cpp +++ b/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; ixrefine_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; izrefine_image(nodes, image, xs, ys, zs, diagonal_balance); +} + + void Tree::finalize_lists(){ for(int_t iz=0; iz void refine_geom(const T& geom, int_t p_level, bool diagonal_balance=false){ for(int_t iz=0; iz Date: Mon, 6 Oct 2025 10:39:52 -0600 Subject: [PATCH 2/4] fix implementation and add tests --- discretize/_extensions/tree.cpp | 54 +++++----- discretize/_extensions/tree_ext.pyx | 19 +++- tests/tree/test_refine.py | 148 ++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+), 34 deletions(-) diff --git a/discretize/_extensions/tree.cpp b/discretize/_extensions/tree.cpp index bad671d3c..955521738 100644 --- a/discretize/_extensions/tree.cpp +++ b/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; ixrefine_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; izrefine_image(nodes, image, xs, ys, zs, diagonal_balance); + roots[iz][iy][ix]->refine_image(nodes, image, shape_cells, xs, ys, zs, diagonal_balance); } diff --git a/discretize/_extensions/tree_ext.pyx b/discretize/_extensions/tree_ext.pyx index f0925c754..4e0bbdbd9 100644 --- a/discretize/_extensions/tree_ext.pyx +++ b/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,12 +1229,19 @@ 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}" @@ -1242,8 +1249,10 @@ cdef class _TreeMesh: 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() diff --git a/tests/tree/test_refine.py b/tests/tree/test_refine.py index b2ec1d7e9..1029f5c2e 100644 --- a/tests/tree/test_refine.py +++ b/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) From 621a0dbf68b917df567ca22ba88883c34e2e2e85 Mon Sep 17 00:00:00 2001 From: Joseph Capriotti Date: Mon, 6 Oct 2025 10:43:12 -0600 Subject: [PATCH 3/4] add missing parameter docstring --- discretize/_extensions/tree_ext.pyx | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/discretize/_extensions/tree_ext.pyx b/discretize/_extensions/tree_ext.pyx index 4e0bbdbd9..cd1ab74fd 100644 --- a/discretize/_extensions/tree_ext.pyx +++ b/discretize/_extensions/tree_ext.pyx @@ -1227,6 +1227,11 @@ cdef class _TreeMesh: 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. + finalize : bool, optional + Whether to finalize after inserting point(s) + diagonal_balance : bool or None, optional + Whether to balance cells diagonally in the refinement, `None` implies using + the same setting used to instantiate the TreeMesh`. """ if diagonal_balance is None: From 0e7dde72ce6922d750bed4d0b13c69fd380ca666 Mon Sep 17 00:00:00 2001 From: Joseph Capriotti Date: Mon, 6 Oct 2025 22:31:56 -0600 Subject: [PATCH 4/4] Fixes typo in TreeCell documentation --- discretize/_extensions/tree_ext.pyx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/discretize/_extensions/tree_ext.pyx b/discretize/_extensions/tree_ext.pyx index cd1ab74fd..dd1f75b1a 100644 --- a/discretize/_extensions/tree_ext.pyx +++ b/discretize/_extensions/tree_ext.pyx @@ -25,7 +25,7 @@ cdef class TreeCell: This cannot be created in python, it can only be accessed by indexing the :class:`~discretize.TreeMesh` object. ``TreeCell`` is the object being passed - to the user defined refine function when calling tshe + to the user defined refine function when calling the :py:attr:`~discretize.TreeMesh.refine` method for a :class:`~discretize.TreeMesh`. Examples @@ -1219,19 +1219,19 @@ cdef class _TreeMesh: 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 + 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 + 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. finalize : bool, optional Whether to finalize after inserting point(s) diagonal_balance : bool or None, optional Whether to balance cells diagonally in the refinement, `None` implies using - the same setting used to instantiate the TreeMesh`. + the same setting used to instantiate the `TreeMesh`. """ if diagonal_balance is None: