diff --git a/grid_apps/utils.py b/grid_apps/utils.py index 78e05fb..e3c2dae 100644 --- a/grid_apps/utils.py +++ b/grid_apps/utils.py @@ -586,7 +586,7 @@ def octree_2_treemesh( # pylint: disable=too-many-locals ): return None - n_cell_dim, cell_sizes = [], [] + cell_sizes = [] for ax in "uvw": if ( getattr(mesh, f"{ax}_cell_size") is None @@ -594,24 +594,16 @@ def octree_2_treemesh( # pylint: disable=too-many-locals ): raise ValueError(f"Cell size in {ax} direction is not defined.") - n_cell_dim.append(getattr(mesh, f"{ax}_count")) cell_sizes.append( - np.ones(getattr(mesh, f"{ax}_count")) * getattr(mesh, f"{ax}_cell_size") + np.full(getattr(mesh, f"{ax}_count"), getattr(mesh, f"{ax}_cell_size")) ) if any(np.any(cell_size < 0) for cell_size in cell_sizes): raise NotImplementedError("Negative cell sizes not supported.") - ls = np.log2(n_cell_dim).astype(int) - - if len(set(ls)) == 1: - max_level = ls[0] - else: - max_level = min(ls) + 1 - + levels = tree_levels(mesh) cells = np.vstack(mesh.octree_cells.tolist()) indexes = cells[:, :-1] * 2 + cells[:, -1][:, None] # convert to cpp index - levels = max_level - np.log2(cells[:, -1]) treemesh = TreeMesh( cell_sizes, x0=np.asarray(mesh.origin.tolist()), diagonal_balance=False ) @@ -620,6 +612,30 @@ def octree_2_treemesh( # pylint: disable=too-many-locals return treemesh +def tree_levels(mesh: Octree) -> np.ndarray | None: + """ + Convert Octree n cell indices to Treemesh level indices. + + :param mesh: Octree object with n cell index. + + :returns: Array of level indices. + """ + if mesh.octree_cells is None: + return None + + n_cell_dim = [mesh.u_count, mesh.v_count, mesh.w_count] + ls = np.log2(n_cell_dim).astype(int) + + if len(set(ls)) == 1: + max_level = ls[0] + else: + max_level = min(ls) + 1 + + levels = max_level - np.log2(mesh.octree_cells["NCells"]) + + return levels.astype(int) + + def resample_locations(locations: np.ndarray, increment: float) -> np.ndarray: """ Resample locations along a sequence of positions at a given increment. diff --git a/tests/utils_test.py b/tests/utils_test.py index 8db5f73..dac264b 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -30,6 +30,7 @@ octree_2_treemesh, surface_strip, tensor_mesh_ordering, + tree_levels, treemesh_2_octree, ) @@ -242,7 +243,7 @@ def test_find_endpoints(): def test_surface_strip(tmp_path): - with Workspace.create(tmp_path / "test.geoh5") as workspace: + with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace: line = Points.create( workspace, vertices=np.array([[-1, 0, 0], [0, 0, 0], [1, 0, 0]]) ) @@ -282,7 +283,7 @@ def test_not_implemented_negative(): def test_collocate_octrees(tmp_path: Path): - with Workspace.create(tmp_path / "test.geoh5") as workspace: + with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace: local_mesh1 = TreeMesh( [[10] * 16, [10] * 16, [10] * 16], [1000, 0, 0], diagonal_balance=True ) @@ -391,7 +392,7 @@ def test_create_octree_from_octrees_errors(): def test_densify_curve(tmp_path: Path): - with Workspace.create(tmp_path / "test.geoh5") as workspace: + with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace: curve = Curve.create( workspace, vertices=np.vstack([[0, 0, 0], [10, 0, 0], [10, 10, 0]]), @@ -490,8 +491,43 @@ def test_octree_2_treemesh(): np.testing.assert_allclose(tmesh.cell_centers, mesh.cell_centers) +def test_tree_levels(tmp_path): + with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace: + mesh = Octree.create( + workspace, + name="test", + origin=[0, 0, 0], + u_count=8, + v_count=8, + w_count=8, + u_cell_size=10.0, + v_cell_size=20.0, + w_cell_size=40.0, + octree_cells=np.array( + [ + [0, 0, 0, 2], + [2, 0, 0, 2], + [0, 2, 0, 2], + [2, 2, 0, 2], + [0, 0, 2, 2], + [2, 0, 2, 2], + [0, 2, 2, 2], + [2, 2, 2, 2], + [4, 0, 0, 4], + [0, 4, 0, 4], + [4, 4, 0, 4], + ] + ), + ) + + levels = tree_levels(mesh) + assert levels[-1] == 1 # Largest cell + assert levels[0] == 2 # Medium size + assert np.all(levels < 3) # No highest level + + def test_roundtrip_octree_conversion(tmp_path): - with Workspace.create(tmp_path / "test.geoh5") as workspace: + with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace: points = np.vstack( [ [10, 10, -10],