Merge pull request #324 from santisoler/all-nodes

Add `total_nodes` method to TreeMeshes

Author: jcapriot

discretize/_extensions/tree_ext.pyx

@@ -5953,16 +5953,18 @@ cdef class _TreeMesh:
 
     @property
     def cell_nodes(self):
-        """The index of nodes for each cell.
+        """The index of all nodes for each cell.
+
+        These indices point to non-hanging and hanging nodes.
 
         Returns
         -------
         numpy.ndarray of int
             Index array of shape (n_cells, 4) if 2D, or (n_cells, 8) if 3D
 
-        Notes
-        -----
-        These indices will also point to hanging nodes.
+        See also
+        --------
+        TreeMesh.total_nodes
         """
         cdef int_t npc = 4 if self.dim == 2 else 8
         inds = np.empty((self.n_cells, npc), dtype=np.int64)

discretize/tree_mesh.py

@@ -740,6 +740,21 @@ def refine_surface(
         if finalize:
             self.finalize()
 
+    @property
+    def total_nodes(self):
+        """Gridded hanging and non-hanging nodes locations.
+
+        This property returns a numpy array of shape
+        ``(n_total_nodes, dim)`` containing gridded locations for
+        all hanging and non-hanging nodes in the mesh.
+
+        Returns
+        -------
+        (n_total_nodes, dim) numpy.ndarray of float
+            Gridded hanging and non-hanging node locations
+        """
+        return np.vstack((self.nodes, self.hanging_nodes))
+
     @property
     def vntF(self):
         """Vector number of total faces along each axis.

tests/tree/test_tree.py

@@ -1,5 +1,6 @@
 import numpy as np
 import unittest
+import pytest
 import discretize
 
 TOL = 1e-8
@@ -340,6 +341,43 @@ def test_cell_nodes(self):
         np.testing.assert_equal(cell_2.nodes, cell_nodes[2])
 
 
+class TestTreeMeshNodes:
+    @pytest.fixture(params=["2D", "3D"])
+    def sample_mesh(self, request):
+        """Return a sample TreeMesh"""
+        nc = 8
+        h1 = np.random.rand(nc) * nc * 0.5 + nc * 0.5
+        h2 = np.random.rand(nc) * nc * 0.5 + nc * 0.5
+        if request.param == "2D":
+            h = [hi / np.sum(hi) for hi in [h1, h2]]  # normalize
+            mesh = discretize.TreeMesh(h)
+            points = np.array([[0.2, 0.1], [0.8, 0.4]])
+            levels = np.array([1, 2])
+            mesh.insert_cells(points, levels, finalize=True)
+        else:
+            h3 = np.random.rand(nc) * nc * 0.5 + nc * 0.5
+            h = [hi / np.sum(hi) for hi in [h1, h2, h3]]  # normalize
+            mesh = discretize.TreeMesh(h, levels=3)
+            points = np.array([[0.2, 0.1, 0.7], [0.8, 0.4, 0.2]])
+            levels = np.array([1, 2])
+            mesh.insert_cells(points, levels, finalize=True)
+        return mesh
+
+    def test_total_nodes(self, sample_mesh):
+        """
+        Test if ``TreeMesh.total_nodes`` works as expected
+        """
+        n_non_hanging_nodes = sample_mesh.n_nodes
+        # Check if total_nodes contain all non hanging nodes (in the right order)
+        np.testing.assert_equal(
+            sample_mesh.total_nodes[:n_non_hanging_nodes, :], sample_mesh.nodes
+        )
+        # Check if total_nodes contain all hanging nodes (in the right order)
+        np.testing.assert_equal(
+            sample_mesh.total_nodes[n_non_hanging_nodes:, :], sample_mesh.hanging_nodes
+        )
+
+
 class Test2DInterpolation(unittest.TestCase):
     def setUp(self):
         def topo(x):