fix: generalising support functions for grid and tetra

Refactoring some of the internal functions used by
the discrete interpolators to make them equivalent between
tetra and cartesian grid.
Some other general improvements to make helper functions
more useful outside of their current use case.

LoopStructural/interpolators/base_structured_3d_support.py

@@ -124,19 +124,19 @@ def rotate(self,pos):
         return np.einsum('ijk,ik->ij', self.rotation_xy[None,:,:], pos)
 
     def position_to_cell_index(self, pos):
-        """[summary]
+        """Get the indexes (i,j,k) of a cell
+        that a point is inside
 
-        [extended_summary]
 
         Parameters
         ----------
-        pos : [type]
-            [description]
+        pos : np.array
+            Nx3 array of xyz locations
 
         Returns
         -------
-        [type]
-            [description]
+        np.array, np.array, np.array
+            i,j,k indexes of the cell that the point is in
         """
 
         pos = self.rotate(pos)
@@ -150,6 +150,9 @@ def position_to_cell_index(self, pos):
         iz = iz // self.step_vector[None, 2]
         return ix.astype(int), iy.astype(int), iz.astype(int)
 
+    def position_to_cell_global_index(self,pos):
+        ix,iy,iz = self.position_to_cell_index(pos)
+
     def inside(self, pos):
         pos = self.rotate(pos)
         # check whether point is inside box
@@ -183,7 +186,7 @@ def check_position(self, pos):
             return False
         return pos
 
-    def global_cell_indicies(self, indexes):
+    def _global_indicies(self, indexes, nsteps):
         """
         Convert from cell indexes to global cell index
 
@@ -195,10 +198,28 @@ def global_cell_indicies(self, indexes):
         -------
 
         """
-        indexes = np.array(indexes).swapaxes(0, 2)
-        return indexes[:, :, 0] + self.nsteps_cells[None, None, 0] \
-               * indexes[:, :, 1] + self.nsteps_cells[None, None, 0] * \
-               self.nsteps_cells[None, None, 1] * indexes[:, :, 2]
+        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]
 
     def cell_corner_indexes(self, x_cell_index, y_cell_index, z_cell_index):
         """
@@ -234,6 +255,23 @@ def position_to_cell_corners(self, pos):
         globalidx[~inside] = -1
         return globalidx, inside
 
+    def position_to_cell_vertices(self, pos):
+        """Get the vertices of the cell a point is in
+
+        Parameters
+        ----------
+        pos : np.array
+            Nx3 array of xyz locations
+
+        Returns
+        -------
+        np.array((N,3),dtype=float), np.array(N,dtype=int)
+            vertices, inside
+        """
+        gi, inside = self.position_to_cell_corners(pos)
+        ci, cj, ck = self.global_index_to_node_index(gi)
+        return self.node_indexes_to_position(ci, cj, ck), inside
+
     def node_indexes_to_position(self, xindex, yindex, zindex):
 
         x = self.origin[0] + self.step_vector[0] * xindex
@@ -287,6 +325,7 @@ def global_index_to_node_index(self, global_index):
         z_index = global_index // self.nsteps[0, None] // \
                   self.nsteps[1, None]
         return x_index, y_index, z_index
+
     def global_node_indicies(self, indexes):
         """
         Convert from node indexes to global node index
@@ -299,7 +338,18 @@ def global_node_indicies(self, indexes):
         -------
 
         """
-        indexes = np.array(indexes).swapaxes(0, 2)
-        return indexes[:, :, 0] + self.nsteps[None, None, 0] \
-               * indexes[:, :, 1] + self.nsteps[None, None, 0] * \
-               self.nsteps[None, None, 1] * indexes[:, :, 2]
+        return self._global_indicies(indexes, self.nsteps)
+
+    def global_cell_indicies(self, indexes):
+        """
+        Convert from cell indexes to global cell index
+
+        Parameters
+        ----------
+        indexes
+
+        Returns
+        -------
+
+        """
+        return self._global_indicies(indexes, self.nsteps_cells)

LoopStructural/interpolators/finite_difference_interpolator.py

@@ -262,7 +262,7 @@ def add_gradient_constraints(self, w=1.):
             idc[inside, :] = gi[node_idx[inside, :]]
             inside = np.logical_and(~np.any(idc == -1, axis=1), inside)
 
-            T = self.support.calcul_T(points[inside, :3])
+            vertices, T, elements, inside = self.support.get_element_gradient_for_location(points[inside, :3])
             # normalise constraint vector and scale element matrix by this
             norm = np.linalg.norm(points[:,3:6],axis=1)
             points[:,3:6]/=norm[:,None]
@@ -305,7 +305,7 @@ def add_norm_constraints(self, w=1.):
             # calculate unit vector for node gradients
             # this means we are only constraining direction of grad not the
             # magnitude
-            T = self.support.calcul_T(points[inside, :3])
+            vertices, T, elements, inside = self.support.get_element_gradient_for_location(points[inside, :3])
             # T*=np.product(self.support.step_vector)
             # T/=self.support.step_vector[0]
             w /= 3
@@ -356,7 +356,7 @@ def add_gradient_orthogonal_constraints(self, points, vector, w=1.0,
             norm = np.linalg.norm(vector,axis=1)
             vector/=norm[:,None]
             #normalise element vector to unit vector for dot product
-            T = self.support.calcul_T(points[inside, :3])
+            vertices, T, elements, inside = self.support.get_element_gradient_for_location(points[inside, :3])
             T/=norm[:,None,None]
 
             # dot product of vector and element gradient = 0

LoopStructural/interpolators/piecewiselinear_interpolator.py

@@ -181,7 +181,7 @@ def add_gradient_constraints(self, w=1.0):
 
         points = self.get_gradient_constraints()
         if points.shape[0] > 0:
-            vertices, element_gradients, tetras, inside = self.support.get_tetra_gradient_for_location(points[:,:3])
+            vertices, element_gradients, tetras, inside = self.support.get_element_gradient_for_location(points[:,:3])
             #e, inside = self.support.elements_for_array(points[:, :3])
             #nodes = self.support.nodes[self.support.elements[e]]
             vecs = vertices[:, 1:, :] - vertices[:, 0, None, :]
@@ -235,7 +235,7 @@ def add_norm_constraints(self, w=1.0):
 
         points = self.get_norm_constraints()
         if points.shape[0] > 0:
-            vertices, element_gradients, tetras, inside = self.support.get_tetra_gradient_for_location(points[:, :3])
+            vertices, element_gradients, tetras, inside = self.support.get_element_gradient_for_location(points[:, :3])
             # e, inside = self.support.elements_for_array(points[:, :3])
             # nodes = self.support.nodes[self.support.elements[e]]
             vol = np.zeros(element_gradients.shape[0])
@@ -279,7 +279,7 @@ def add_value_constraints(self, w=1.0):  # for now weight all value points the s
         # get elements for points
         points = self.get_value_constraints()
         if points.shape[0] > 1:
-            vertices, c, tetras, inside = self.support.get_tetra_for_location(points[:,:3])
+            vertices, c, tetras, inside = self.support.get_element_for_location(points[:,:3])
             # calculate volume of tetras
             vecs = vertices[inside, 1:, :] - vertices[inside, 0, None, :]
             vol = np.abs(np.linalg.det(vecs)) / 6
@@ -314,7 +314,7 @@ def add_interface_constraints(self, w=1.0):  # for now weight all value points t
         """
         points = self.get_interface_constraints()
         if points.shape[0] > 1:
-            vertices, c, tetras, inside = self.support.get_tetra_for_location(points[:,:3])
+            vertices, c, tetras, inside = self.support.get_element_for_location(points[:,:3])
 
             gi = np.zeros(self.support.n_nodes)
             gi[:] = -1
@@ -399,7 +399,7 @@ def add_gradient_orthogonal_constraints(self, points, vector, w=1.0,
 
         """
         if points.shape[0] > 0:
-            vertices, element_gradients, tetras, inside = self.support.get_tetra_gradient_for_location(points[:,:3])
+            vertices, element_gradients, tetras, inside = self.support.get_element_gradient_for_location(points[:,:3])
             #e, inside = self.support.elements_for_array(points[:, :3])
             #nodes = self.support.nodes[self.support.elements[e]]
             norm = np.linalg.norm(vector,axis=1)

LoopStructural/interpolators/structured_grid.py

@@ -161,7 +161,7 @@ def neighbour_global_indexes(self, mask = None, **kwargs):
         if "indexes" in kwargs:
             indexes = kwargs['indexes']
         if "indexes" not in kwargs:
-            indexes = np.array(np.meshgrid(np.arange(1,nsteps[0]-1),np.arange(1,nsteps[1]-1),np.arange(1,nsteps[1]-1))).reshape((3,-1))
+            indexes = np.array(np.meshgrid(np.arange(1,self.nsteps[0]-1),np.arange(1,self.nsteps[1]-1),np.arange(1,self.nsteps[2]-1))).reshape((3,-1))
         # indexes = np.array(indexes).T
         if indexes.ndim != 2:
             print(indexes.ndim)
@@ -185,13 +185,6 @@ def neighbour_global_indexes(self, mask = None, **kwargs):
                self.nsteps[0, None, None] * self.nsteps[
                    1, None, None] * neighbours[2, :, :]).astype(np.int64)
 
-
-
-
-
-
-
-
     def evaluate_value(self, evaluation_points, property_array):
         """
         Evaluate the value of of the property at the locations.
@@ -226,7 +219,7 @@ def evaluate_gradient(self, evaluation_points, property_array):
             raise BaseException
         idc, inside = self.position_to_cell_corners(evaluation_points)
         T = np.zeros((idc.shape[0], 3, 8))
-        T[inside, :, :] = self.calcul_T(evaluation_points[inside, :])
+        T[inside, :, :] = self.get_element_gradient_for_location(evaluation_points[inside, :])
         # indices = np.array([self.position_to_cell_index(evaluation_points)])
         # idc = self.global_indicies(indices.swapaxes(0,1))
         # print(idc)
@@ -237,7 +230,7 @@ def evaluate_gradient(self, evaluation_points, property_array):
             [np.sum(T[:, 0, :], axis=1), np.sum(T[:, 1, :], axis=1) ,
              np.sum(T[:, 2, :], axis=1) ]).T
 
-    def calcul_T(self, pos):
+    def get_element_gradient_for_location(self, pos):
         """
         Calculates the gradient matrix at location pos
         :param pos: numpy array of location Nx3
@@ -256,7 +249,8 @@ def calcul_T(self, pos):
         # x, y, z = self.node_indexes_to_position(cellx, celly, cellz)
         T = np.zeros((pos.shape[0], 3, 8))
         x, y, z = self.position_to_local_coordinates(pos)
-
+        vertices, inside = self.position_to_cell_vertices(pos)
+        elements = self.global_node_indicies(np.array(self.position_to_cell_index(pos)))
         T[:, 0, 0] = (1 - z) * (y- 1)  # v000
         T[:, 0, 1] = (1 - y) * (1 - z)  # (y[:, 3] - pos[:, 1]) / div
         T[:, 0, 2] = -y * (1 - z)  # (pos[:, 1] - y[:, 0]) / div
@@ -285,5 +279,23 @@ def calcul_T(self, pos):
         T[:, 2, 7] = x * y
         T/=self.step_vector[0]
 
-        return T 
+        return vertices, T, elements, inside 
+
+    def get_element_for_location(self, pos):
+        """Calculate the shape function of elements
+        for a location
 
+        Parameters
+        ----------
+        pos : np.array((N,3))
+            location of points to calculate the shape function
+
+        Returns
+        -------
+        [type]
+            [description]
+        """
+        vertices, inside = self.position_to_cell_vertices(pos)
+        elements = self.global_node_indicies(np.array(self.position_to_cell_index(pos)))
+        a = self.position_to_dof_coefs(pos)
+        return vertices, c_return, elements, inside

LoopStructural/interpolators/structured_tetra.py

@@ -85,7 +85,7 @@ def evaluate_value(self, pos, property_array):
         """
         values = np.zeros(pos.shape[0])
         values[:] = np.nan
-        vertices, c, tetras, inside = self.get_tetra_for_location(pos)
+        vertices, c, tetras, inside = self.get_element_for_location(pos)
         values[inside] = np.sum(c[inside,:]*property_array[tetras[inside,:]],axis=1)
         return values
 
@@ -107,7 +107,7 @@ def evaluate_gradient(self, pos, property_array):
         """
         values = np.zeros(pos.shape)
         values[:] = np.nan
-        vertices, element_gradients, tetras, inside = self.get_tetra_gradient_for_location(pos)
+        vertices, element_gradients, tetras, inside = self.get_element_gradient_for_location(pos)
         #grads = np.zeros(tetras.shape)
         values[inside,:] = (element_gradients[inside,:,:]*property_array[tetras[inside,None,:]]).sum(2)
         length = np.sum(values[inside,:],axis=1)
@@ -122,7 +122,7 @@ def inside(self, pos):
                       self.step_vector[None, i] * self.nsteps_cells[None, i]
         return inside
 
-    def get_tetra_for_location(self, pos):
+    def get_element_for_location(self, pos):
         """
         Determine the tetrahedron from a numpy array of points
 
@@ -375,7 +375,7 @@ def get_element_gradients(self, elements = None):
 
         return element_gradients[elements,:,:]
 
-    def get_tetra_gradient_for_location(self, pos):
+    def get_element_gradient_for_location(self, pos):
         """
         Get the gradient of the tetra for a location
 
@@ -387,7 +387,7 @@ def get_tetra_gradient_for_location(self, pos):
         -------
 
         """
-        vertices, bc, tetras, inside = self.get_tetra_for_location(pos)
+        vertices, bc, tetras, inside = self.get_element_for_location(pos)
         ps = vertices
         m = np.array(
             [[(ps[:, 1, 0] - ps[:, 0, 0]), (ps[:, 1, 1] - ps[:, 0, 1]),(ps[:, 1, 2] - ps[:, 0, 2])],