From 6b14df55fb9fbc07a1d4739211406a4fa8447ceb Mon Sep 17 00:00:00 2001 From: Santiago Soler Date: Wed, 6 Sep 2023 12:00:13 -0700 Subject: [PATCH 1/3] Add TensorMesh.cell_nodes property Add a new `cell_nodes` property to `TensorMesh` that mimics the one of `TreeMesh` returning the indices of the nodes that correspond to each cell in the mesh. Add tests for the new feature. --- discretize/tensor_mesh.py | 38 ++++++++++++++++++++++++++++++++++++++ tests/base/test_tensor.py | 23 +++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/discretize/tensor_mesh.py b/discretize/tensor_mesh.py index 945c55542..599168203 100644 --- a/discretize/tensor_mesh.py +++ b/discretize/tensor_mesh.py @@ -1,6 +1,7 @@ """Module housing the TensorMesh implementation.""" import itertools import numpy as np +import time from discretize.base import BaseRectangularMesh, BaseTensorMesh from discretize.operators import DiffOperators, InnerProducts @@ -590,6 +591,43 @@ def face_boundary_indices(self): indzu = self.gridFz[:, 2] == max(self.gridFz[:, 2]) return indxd, indxu, indyd, indyu, indzd, indzu + @property + def cell_nodes(self): + """The index of all nodes for each cell. + + Returns + ------- + numpy.ndarray of int + Index array of shape (n_cells, 4) if 2D, or (n_cells, 8) if 3D + """ + order = "F" + nodes_indices = np.arange(self.n_nodes).reshape(self.shape_nodes, order=order) + if self.dim == 1: + cell_nodes = [ + nodes_indices[:-1].reshape(-1, order=order), + nodes_indices[1:].reshape(-1, order=order), + ] + elif self.dim == 2: + cell_nodes = [ + nodes_indices[:-1, :-1].reshape(-1, order=order), + nodes_indices[1:, :-1].reshape(-1, order=order), + nodes_indices[:-1, 1:].reshape(-1, order=order), + nodes_indices[1:, 1:].reshape(-1, order=order), + ] + else: + cell_nodes = [ + nodes_indices[:-1, :-1, :-1].reshape(-1, order=order), + nodes_indices[1:, :-1, :-1].reshape(-1, order=order), + nodes_indices[:-1, 1:, :-1].reshape(-1, order=order), + nodes_indices[1:, 1:, :-1].reshape(-1, order=order), + nodes_indices[:-1, :-1, 1:].reshape(-1, order=order), + nodes_indices[1:, :-1, 1:].reshape(-1, order=order), + nodes_indices[:-1, 1:, 1:].reshape(-1, order=order), + nodes_indices[1:, 1:, 1:].reshape(-1, order=order), + ] + cell_nodes = np.stack(cell_nodes, axis=-1) + return cell_nodes + @property def cell_boundary_indices(self): """Return the indices of the x, (y and z) boundary cells. diff --git a/tests/base/test_tensor.py b/tests/base/test_tensor.py index 21f605fd2..9349cafa7 100644 --- a/tests/base/test_tensor.py +++ b/tests/base/test_tensor.py @@ -1,3 +1,4 @@ +import pytest import numpy as np import unittest import discretize @@ -250,6 +251,28 @@ def test_serialization(self): self.assertTrue(np.all(self.mesh2.gridCC == mesh.gridCC)) +class TestTensorMeshCellNodes: + """ + Test TensorMesh.cell_nodes + """ + + @pytest.fixture(params=[1, 2, 3], ids=["dims-1", "dims-2", "dims-3"]) + def mesh(self, request): + """Sample TensorMesh.""" + if request.param == 1: + h = [10] + elif request.param == 2: + h = [10, 15] + else: + h = [10, 15, 20] + return discretize.TensorMesh(h) + + def test_cell_nodes(self, mesh): + """Test TensorMesh.cell_nodes.""" + expected_cell_nodes = np.array([cell.nodes for cell in mesh]) + np.testing.assert_allclose(mesh.cell_nodes, expected_cell_nodes) + + class TestPoissonEqn(discretize.tests.OrderTest): name = "Poisson Equation" meshSizes = [10, 16, 20] From da65cc25eb8a203995f4337ff84b00ee5fa61c70 Mon Sep 17 00:00:00 2001 From: Santiago Soler Date: Mon, 18 Sep 2023 08:59:34 -0700 Subject: [PATCH 2/3] Remove unused import --- discretize/tensor_mesh.py | 1 - 1 file changed, 1 deletion(-) diff --git a/discretize/tensor_mesh.py b/discretize/tensor_mesh.py index 599168203..e7f6c090c 100644 --- a/discretize/tensor_mesh.py +++ b/discretize/tensor_mesh.py @@ -1,7 +1,6 @@ """Module housing the TensorMesh implementation.""" import itertools import numpy as np -import time from discretize.base import BaseRectangularMesh, BaseTensorMesh from discretize.operators import DiffOperators, InnerProducts From 1cf90d140c5e8fdb9df0cfe941ea0d6bf7422841 Mon Sep 17 00:00:00 2001 From: Santiago Soler Date: Mon, 18 Sep 2023 10:32:18 -0700 Subject: [PATCH 3/3] Document order of nodes for TensorMesh.cell_nodes Add a note illustrating the order in which the indices of the nodes for each cell are returned. --- discretize/tensor_mesh.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/discretize/tensor_mesh.py b/discretize/tensor_mesh.py index e7f6c090c..842dd0106 100644 --- a/discretize/tensor_mesh.py +++ b/discretize/tensor_mesh.py @@ -594,10 +594,40 @@ def face_boundary_indices(self): def cell_nodes(self): """The index of all nodes for each cell. + The nodes for each cell are listed following an "F" order: the first + coordinate (``x``) changes faster than the second one (``y``). If the + mesh is 3D, the second coordinate (``y``) changes faster than the third + one (``z``). + Returns ------- numpy.ndarray of int Index array of shape (n_cells, 4) if 2D, or (n_cells, 8) if 3D + + Notes + ----- + For a 2D mesh, the nodes indices for a single cell are returned in the + following order: + + .. code:: + + 2 -- 3 + | | + 0 -- 1 + + For a 3D mesh, the nodes indices for a single cell are returned in the + following order: + + .. code:: + + 6-----7 + /| /| + 4-----5 | + | | | | + | 2---|-3 + |/ |/ + 0-----1 + """ order = "F" nodes_indices = np.arange(self.n_nodes).reshape(self.shape_nodes, order=order)