fix: updating tetmesh class for new indexing

Author: lachlangrose

LoopStructural/interpolators/_python/dsi_helper.py

@@ -1,7 +1,11 @@
 import numpy as np
 
 
-def cg(EG, neighbours, elements, nodes, region):
+def cg(EG: np.ndarray, neighbours: np.ndarray, elements: np.ndarray, nodes, region):
+    if EG.shape[0] != neighbours.shape[0] and EG.shape[0] != elements.shape[0]:
+        raise ValueError("EG and neighbours must have the same number of elements")
+    if EG.shape[2] != 4:
+        raise ValueError("EG must have 4 columns")
     Nc = 5  # numer of constraints shared nodes + independent
     Na = 4  # number of nodes
     Ns = Na - 1

LoopStructural/interpolators/supports/_3d_base_structured.py

@@ -233,32 +233,42 @@ def _global_indicies(self, indexes: np.ndarray, nsteps: np.ndarray) -> np.ndarra
         """
         if len(indexes.shape) == 1:
             raise ValueError("Cell indexes needs to be Nx3")
-        if len(indexes.shape) == 2:
-            if indexes.shape[1] != 3 and indexes.shape[0] == 3:
-                indexes = indexes.swapaxes(0, 1)
-            if indexes.shape[1] != 3:
-                logger.error("Indexes shape {}".format(indexes.shape))
-                raise ValueError("Cell indexes needs to be Nx3")
-            return (
-                indexes[:, 0]
-                + nsteps[None, 0] * indexes[:, 1]
-                + nsteps[None, 0] * nsteps[None, 1] * indexes[:, 2]
-            )
-        if len(indexes.shape) == 3:
-            # if indexes.shape[2] != 3 and indexes.shape[1] == 3:
-            #     indexes = indexes.swapaxes(1, 2)
-            # if indexes.shape[2] != 3 and indexes.shape[0] == 3:
-            #     indexes = indexes.swapaxes(0, 2)
-            if indexes.shape[2] != 3:
-                logger.error("Indexes shape {}".format(indexes.shape))
-                raise ValueError("Cell indexes needs to be NxNx3")
-            return (
-                indexes[:, :, 0]
-                + nsteps[None, None, 0] * indexes[:, :, 1]
-                + nsteps[None, None, 0] * nsteps[None, None, 1] * indexes[:, :, 2]
-            )
-        else:
-            raise ValueError("Cell indexes need to be a 2 or 3d numpy array")
+        if indexes.shape[-1] != 3:
+            raise ValueError("Last dimensions should be ijk indexing")
+        original_shape = indexes.shape
+        indexes = indexes.reshape(-1, 3)
+        gi = (
+            indexes[:, 0]
+            + nsteps[None, 0] * indexes[:, 1]
+            + nsteps[None, 0] * nsteps[None, 1] * indexes[:, 2]
+        )
+        return gi.reshape(original_shape[:-1])
+        # if len(indexes.shape) == 2:
+        #     if indexes.shape[1] != 3 and indexes.shape[0] == 3:
+        #         indexes = indexes.swapaxes(0, 1)
+        #     if indexes.shape[1] != 3:
+        #         logger.error("Indexes shape {}".format(indexes.shape))
+        #         raise ValueError("Cell indexes needs to be Nx3")
+        #     return (
+        #         indexes[:, 0]
+        #         + nsteps[None, 0] * indexes[:, 1]
+        #         + nsteps[None, 0] * nsteps[None, 1] * indexes[:, 2]
+        #     )
+        # if len(indexes.shape) == 3:
+        #     # if indexes.shape[2] != 3 and indexes.shape[1] == 3:
+        #     #     indexes = indexes.swapaxes(1, 2)
+        #     # if indexes.shape[2] != 3 and indexes.shape[0] == 3:
+        #     #     indexes = indexes.swapaxes(0, 2)
+        #     if indexes.shape[2] != 3:
+        #         logger.error("Indexes shape {}".format(indexes.shape))
+        #         raise ValueError("Cell indexes needs to be NxNx3")
+        #     return (
+        #         indexes[:, :, 0]
+        #         + nsteps[None, None, 0] * indexes[:, :, 1]
+        #         + nsteps[None, None, 0] * nsteps[None, None, 1] * indexes[:, :, 2]
+        #     )
+        # else:
+        #     raise ValueError("Cell indexes need to be a 2 or 3d numpy array")
 
     def cell_corner_indexes(self, cell_indexes):
         """

LoopStructural/interpolators/supports/_3d_structured_grid.py

@@ -178,7 +178,6 @@ def neighbour_global_indexes(self, mask=None, **kwargs):
             ).reshape((3, -1))
         # indexes = np.array(indexes).T
         if indexes.ndim != 2:
-            print(indexes.ndim)
             return
         # determine which neighbours to return default is diagonals included.
         if mask is None:

LoopStructural/interpolators/supports/_3d_structured_tetra.py

@@ -143,27 +143,26 @@ def get_element_for_location(self, pos: np.ndarray):
         pos = np.array(pos)
         inside = self.inside(pos)
         # initialise array for tetrahedron vertices
-        vertices = np.zeros((5, 4, pos.shape[0], 3))
+        vertices = np.zeros((pos.shape[0], 5, 4, 3))
         vertices[:] = np.nan
         # get cell indexes
-        c_xi, c_yi, c_zi = self.position_to_cell_index(pos)
-
+        cell_indexes = self.position_to_cell_index(pos)
         # determine if using +ve or -ve mask
-        even_mask = (c_xi + c_yi + c_zi) % 2 == 0
+        even_mask = np.sum(cell_indexes, axis=1) % 2 == 0
         # get cell corners
-        xi, yi, zi = self.cell_corner_indexes(
-            c_xi, c_yi, c_zi
+        corner_indexes = self.cell_corner_indexes(
+            cell_indexes
         )  # global_index_to_node_index(gi)
         # convert to node locations
-        nodes = self.node_indexes_to_position(xi, yi, zi).T
+        nodes = self.node_indexes_to_position(corner_indexes)
 
-        vertices[:, :, even_mask, :] = nodes[:, even_mask, :][
-            self.tetra_mask_even, :, :
+        vertices[even_mask, :, :, :] = nodes[even_mask, :, :][
+            :, self.tetra_mask_even, :
         ]
-        vertices[:, :, ~even_mask, :] = nodes[:, ~even_mask, :][self.tetra_mask, :, :]
+        vertices[~even_mask, :, :, :] = nodes[~even_mask, :, :][:, self.tetra_mask, :]
         # changing order to points, tetra, nodes, coord
-        vertices = vertices.swapaxes(0, 2)
-        vertices = vertices.swapaxes(1, 2)
+        # vertices = vertices.swapaxes(0, 2)
+        # vertices = vertices.swapaxes(1, 2)
         # use scalar triple product to calculate barycentric coords
 
         vap = pos[:, None, :] - vertices[:, :, 0, :]
@@ -191,14 +190,13 @@ def get_element_for_location(self, pos: np.ndarray):
 
         inside = np.logical_and(inside, np.any(mask, axis=1))
         # get cell corners
-        xi, yi, zi = self.cell_corner_indexes(c_xi, c_yi, c_zi)
         # create mask to see which cells are even
-        even_mask = (c_xi + c_yi + c_zi) % 2 == 0
+        even_mask = np.sum(cell_indexes, axis=1) % 2 == 0
         # create global node index list
-        gi = xi + yi * self.nsteps[0] + zi * self.nsteps[0] * self.nsteps[1]
+        gi = self.global_node_indices(corner_indexes)
+        # gi = xi + yi * self.nsteps[0] + zi * self.nsteps[0] * self.nsteps[1]
         # container for tetras
-        tetras = np.zeros((xi.shape[0], 5, 4)).astype(int)
-
+        tetras = np.zeros((corner_indexes.shape[0], 5, 4)).astype(int)
         tetras[even_mask, :, :] = gi[even_mask, :][:, self.tetra_mask_even]
         tetras[~even_mask, :, :] = gi[~even_mask, :][:, self.tetra_mask]
         inside = np.logical_and(inside, self.inside(pos))
@@ -229,15 +227,12 @@ def get_elements(self):
         y = np.arange(0, self.nsteps_cells[1])
         z = np.arange(0, self.nsteps_cells[2])
 
-        c_xi, c_yi, c_zi = np.meshgrid(x, y, z, indexing="ij")
-        c_xi = c_xi.flatten(order="F")
-        c_yi = c_yi.flatten(order="F")
-        c_zi = c_zi.flatten(order="F")
+        cell_indexes = np.array(np.meshgrid(x, y, z, indexing="ij")).reshape(3, -1).T
         # get cell corners
-        xi, yi, zi = self.cell_corner_indexes(c_xi, c_yi, c_zi)
-        even_mask = (c_xi + c_yi + c_zi) % 2 == 0
-        gi = xi + yi * self.nsteps[0] + zi * self.nsteps[0] * self.nsteps[1]
-        tetras = np.zeros((c_xi.shape[0], 5, 4)).astype("int64")
+        cell_corners = self.cell_corner_indexes(cell_indexes)
+        even_mask = np.sum(cell_indexes, axis=1) % 2 == 0
+        gi = self.global_node_indices(cell_corners)
+        tetras = np.zeros((cell_corners.shape[0], 5, 4)).astype("int64")
         tetras[even_mask, :, :] = gi[even_mask, :][:, self.tetra_mask_even]
         tetras[~even_mask, :, :] = gi[~even_mask, :][:, self.tetra_mask]
 
@@ -261,25 +256,25 @@ def get_element_gradients(self, elements=None):
         y = np.arange(0, self.nsteps_cells[1])
         z = np.arange(0, self.nsteps_cells[2])
 
-        c_xi, c_yi, c_zi = np.meshgrid(x, y, z, indexing="ij")
-        c_xi = c_xi.flatten(order="F")
-        c_yi = c_yi.flatten(order="F")
-        c_zi = c_zi.flatten(order="F")
-        even_mask = (c_xi + c_yi + c_zi) % 2 == 0
+        cell_indexes = np.array(np.meshgrid(x, y, z, indexing="ij")).reshape(3, -1).T
+        # c_xi = c_xi.flatten(order="F")
+        # c_yi = c_yi.flatten(order="F")
+        # c_zi = c_zi.flatten(order="F")
+        even_mask = np.sum(cell_indexes, axis=1) % 2 == 0
         # get cell corners
-        xi, yi, zi = self.cell_corner_indexes(
-            c_xi, c_yi, c_zi
+        corner_indexes = self.cell_corner_indexes(
+            cell_indexes
         )  # global_index_to_node_index(gi)
         # convert to node locations
-        nodes = self.node_indexes_to_position(xi, yi, zi).T
+        nodes = self.node_indexes_to_position(corner_indexes)
 
-        points = np.zeros((5, 4, self.n_cells, 3))
-        points[:, :, even_mask, :] = nodes[:, even_mask, :][self.tetra_mask_even, :, :]
-        points[:, :, ~even_mask, :] = nodes[:, ~even_mask, :][self.tetra_mask, :, :]
+        points = np.zeros((self.n_cells, 5, 4, 3))
+        points[even_mask, :, :, :] = nodes[even_mask, :, :][:, self.tetra_mask_even, :]
+        points[~even_mask, :, :, :] = nodes[~even_mask, :, :][:, self.tetra_mask, :]
 
         # changing order to points, tetra, nodes, coord
-        points = points.swapaxes(0, 2)
-        points = points.swapaxes(1, 2)
+        # points = points.swapaxes(0, 2)
+        # points = points.swapaxes(1, 2)
 
         ps = points.reshape(
             points.shape[0] * points.shape[1], points.shape[2], points.shape[3]
@@ -402,72 +397,25 @@ def global_cell_indicies(self, indexes: np.ndarray):
             * indexes[:, :, 2]
         )
 
-    def cell_corner_indexes(self, x_cell_index, y_cell_index, z_cell_index):
-        """
-        Returns the indexes of the corners of a cell given its location xi,
-        yi, zi
-
-        Parameters
-        ----------
-        x_cell_index
-        y_cell_index
-        z_cell_index
-
-        Returns
-        -------
-
-        """
-        x_cell_index = np.array(x_cell_index)
-        y_cell_index = np.array(y_cell_index)
-        z_cell_index = np.array(z_cell_index)
-
-        xcorner = np.array([0, 1, 0, 1, 0, 1, 0, 1])
-        ycorner = np.array([0, 0, 1, 1, 0, 0, 1, 1])
-        zcorner = np.array([0, 0, 0, 0, 1, 1, 1, 1])
-        xcorners = x_cell_index[:, None] + xcorner[None, :]
-        ycorners = y_cell_index[:, None] + ycorner[None, :]
-        zcorners = z_cell_index[:, None] + zcorner[None, :]
-        return xcorners, ycorners, zcorners
-
-    def position_to_cell_corners(self, pos: np.ndarray) -> np.ndarray:
-        """
-        Find the nodes that belong to a cell which contains a point
-
-        Parameters
-        ----------
-        pos
-
-        Returns
-        -------
-
-        """
-        inside = self.inside(pos)
-        ix, iy, iz = self.position_to_cell_index(pos)
-        cornersx, cornersy, cornersz = self.cell_corner_indexes(ix, iy, iz)
-        globalidx = self.global_cell_indicies(
-            np.dstack([cornersx, cornersy, cornersz]).T
-        )
-        return globalidx, inside
-
-    def node_indexes_to_position(self, xindex, yindex, zindex):
-        """
-        Get the xyz position from the node coordinates
-
-        Parameters
-        ----------
-        xindex
-        yindex
-        zindex
-
-        Returns
-        -------
-
-        """
-        x = self.origin[0] + self.step_vector[0] * xindex
-        y = self.origin[1] + self.step_vector[1] * yindex
-        z = self.origin[2] + self.step_vector[2] * zindex
-
-        return np.array([x, y, z])
+    # def position_to_cell_corners(self, pos: np.ndarray) -> np.ndarray:
+    #     """
+    #     Find the nodes that belong to a cell which contains a point
+
+    #     Parameters
+    #     ----------
+    #     pos
+
+    #     Returns
+    #     -------
+
+    #     """
+    #     inside = self.inside(pos)
+    #     ix, iy, iz = self.position_to_cell_index(pos)
+    #     cornersx, cornersy, cornersz = self.cell_corner_indexes(ix, iy, iz)
+    #     globalidx = self.global_cell_indicies(
+    #         np.dstack([cornersx, cornersy, cornersz]).T
+    #     )
+    #     return globalidx, inside
 
     def global_index_to_node_index(self, global_index: np.ndarray):
         """

tests/unit/interpolator/test_discrete_supports.py

@@ -160,32 +160,24 @@ def test_global_index_to_cell_index(support):
 
 
 def test_global_index(support):
-    indexes = (
-        np.array(
-            np.meshgrid(
-                np.arange(0, support.nsteps[0]),
-                np.arange(0, support.nsteps[1]),
-                np.arange(0, support.nsteps[2]),
-            )
+    indexes = np.array(
+        np.meshgrid(
+            np.arange(0, support.nsteps[0]),
+            np.arange(0, support.nsteps[1]),
+            np.arange(0, support.nsteps[2]),
         )
-        .reshape(-1, 3)
-        .T
-    )
+    ).reshape(-1, 3)
     global_node_index = support.global_node_indices(indexes)
     assert np.all(global_node_index >= 0)
     assert np.all(global_node_index < support.n_nodes)
 
-    indexes = (
-        np.array(
-            np.meshgrid(
-                np.arange(0, 3),
-                np.arange(0, 1),
-                np.arange(0, 1),
-            )
+    indexes = np.array(
+        np.meshgrid(
+            np.arange(0, 3),
+            np.arange(0, 1),
+            np.arange(0, 1),
         )
-        .reshape(-1, 3)
-        .T
-    )
+    ).reshape(-1, 3)
     global_node_index = support.global_node_indices(indexes)
     assert np.all(global_node_index >= 0)
     assert np.all(global_node_index < support.n_nodes)