Store experimental changes for varying interpolation strategies

cmr_rnd_diffusion/mesh_utils.py

@@ -4,6 +4,8 @@ __all__ = ["interpolate_mesh", "evaluate_connectivity", "evaluate_cellsize", "se
 
 import numpy as np
 from typing import Tuple
+from scipy import interpolate
+from tqdm import tqdm
 
 
 def interpolate_mesh(mesh_points: np.ndarray, points_per_ring: int = 40, n_shells: int = 4,
@@ -37,10 +39,8 @@ def interpolate_mesh(mesh_points: np.ndarray, points_per_ring: int = 40, n_shell
 
     points_per_shell_interp_cl, new_n_slices = _interp_long(points_per_shell_interp_c,
                                                             interp_factor_l, new_n_theta)
-
     points_per_shell_interp_clr, new_n_shells = _interp_radial(
         np.stack(points_per_shell_interp_cl, axis=0), interp_factor_r)
-
     return points_per_shell_interp_clr.reshape(-1, 3), (new_n_theta, new_n_slices, new_n_shells)
 
 
@@ -56,16 +56,37 @@ def _interp_circ(mesh_layers: np.ndarray, factor: int = 4, points_per_ring: int
     :param n_slices: (int) number of slices for input argument mesh_layers
     :returns: (n_shells, new_points_per_ring * n_slices + 1, 3), new_points_per_ring
     """
-    factor_last = factor + 2
-
-    new_points_per_ring = points_per_ring * factor + 2
+    fail_cout = []
+    factor_last = factor + 1
+    new_points_per_ring = points_per_ring * factor + 1
     result = []
+    phi_new = np.arange(0, new_points_per_ring) * 2. * np.pi / new_points_per_ring
     for shell_index, points_of_shell in enumerate(mesh_layers):
         # Apex point is saved at location 0 -> no interpolation
         points_per_slice = np.stack(
             [points_of_shell[1 + i * points_per_ring:1 + (i + 1) * points_per_ring] for i in
              range(n_slices)])
 
+        # c_points_interp = np.zeros((points_per_slice.shape[0], new_points_per_ring, 3))
+        # for slice_idx, r_pos in enumerate(points_per_slice):
+        #     r_pos_com = np.mean(r_pos, axis=0)
+        #     phi = np.angle(r_pos[..., 0] - r_pos_com[0] + 1j *
+        #                    (r_pos[..., 1] - r_pos_com[1])) + np.pi
+        #     try:
+        #         fx = interpolate.CubicSpline(np.concatenate([phi, phi[0:1]+2*np.pi], axis=0),
+        #                                      np.concatenate([r_pos, r_pos[0:1]], axis=0),
+        #                                      bc_type='periodic')
+        #     except:
+        #         fail_cout.append(slice_idx)
+        #         idx = np.argsort(phi)
+        #         phi = phi[idx]
+        #         r_pos = r_pos[idx]
+        #         fx = interpolate.CubicSpline(np.concatenate([phi, phi[0:1] + 2 * np.pi], axis=0),
+        #                                      np.concatenate([r_pos, r_pos[0:1]], axis=0),
+        #                                      bc_type='periodic')
+        #     c_points_interp[slice_idx] = fx(phi_new)
+        # print(shell_index, fail_cout)
+
         c_points_interp = np.stack([(points_per_slice[:, 1:] - points_per_slice[:, :-1])
                                     * i / factor + points_per_slice[:, :-1] for i in range(factor)],
                                    axis=-2)
@@ -76,16 +97,16 @@ def _interp_circ(mesh_layers: np.ndarray, factor: int = 4, points_per_ring: int
                                   range(factor_last)], axis=-2)
         last_interval = last_interval.reshape(n_slices, -1, 3)
 
-        interpolated_points_wo_apex = np.concatenate([c_points_interp, last_interval],
-                                                     axis=1).reshape(-1, 3)
-        result.append(np.concatenate([points_of_shell[0:1], interpolated_points_wo_apex], axis=0))
+        c_points_interp = np.concatenate([c_points_interp, last_interval], axis=1).reshape(-1, 3)
+        result.append(np.concatenate([points_of_shell[0:1],
+                                      c_points_interp.reshape(-1, 3)], axis=0))
 
     return np.stack(result, axis=0), new_points_per_ring
 
 
 def _interp_long(mesh_layers: np.ndarray, factor: int = 4, points_per_ring: int = 40,
                  n_slices: int = 30) -> (np.ndarray, int):
-    """ Interpolates mesh along longitudial direction. The number of slices excludes the apex.
+    """ Interpolates mesh along longitudinal direction. The number of slices excludes the apex.
     The interpolation includes the interval between the apex and the first slice, therefore the
     resulting number of slices after interpolation evaluates to n_slices_new = n_slices * factor.
 
@@ -137,19 +158,31 @@ def _interp_radial(mesh_layers: np.ndarray, factor: int = 4) -> (np.ndarray, int
     :param factor: (int) interpolation multiplier
     :returns: (n_shells, points_per_ring * n_slices_new + 1, 3)
     """
-
-    new_n_shells = (mesh_layers.shape[0] - 1) * factor + 1
-    points_per_shell = mesh_layers.shape[1]
+    n_shells, points_per_shell = mesh_layers.shape[0:2]
+    new_n_shells = (n_shells - 1) * factor + 1
     shells_without_apex = mesh_layers[:, 1:, :]
-    shells_interp = np.stack([(shells_without_apex[1:] - shells_without_apex[0:-1])
-                              * i / factor + shells_without_apex[0:-1] for i in range(factor)],
-                             axis=1)
-    shells_interp = shells_interp.reshape(-1, points_per_shell - 1, 3)
-
-    shells_interp = np.concatenate([shells_interp, shells_without_apex[-1:]], axis=0)
+    x_ref = np.arange(0., n_shells, 1.)
+    x_new = np.concatenate([np.arange(i, i+1, 1/factor) for i in range(n_shells-1)])
+    x_new = np.concatenate([x_new, [n_shells-1, ]], axis=0)
+    shells_interp = np.zeros((new_n_shells, points_per_shell-1, 3))
+    for cl_idx in range(points_per_shell-1):
+        r_pos = shells_without_apex[:, cl_idx]
+        fx = interpolate.InterpolatedUnivariateSpline(x_ref, r_pos[..., 0], k=3)
+        fy = interpolate.InterpolatedUnivariateSpline(x_ref, r_pos[..., 1], k=3)
+        fz = interpolate.InterpolatedUnivariateSpline(x_ref, r_pos[..., 2], k=3)
+        shells_interp[:, cl_idx] = np.stack([fx(x_new), fy(x_new), fz(x_new)], axis=-1)
+
+
+    # shells_interp = np.stack([(shells_without_apex[1:] - shells_without_apex[0:-1])
+    #                           * i / factor + shells_without_apex[0:-1] for i in range(factor)],
+    #                          axis=1)
+    # shells_interp = shells_interp.reshape(-1, points_per_shell - 1, 3)
+    #
+    # shells_interp = np.concatenate([shells_interp, shells_without_apex[-1:]], axis=0)
     shells_interp_with_apex = np.concatenate(
         [np.tile(mesh_layers[0:1, 0:1, :], [shells_interp.shape[0], 1, 1]),
          shells_interp], axis=1)
+
     return shells_interp_with_apex, new_n_shells