Fix some bugs related to the differential geometry calculation

spateo/alignment/methods/morpho_class.py

@@ -139,6 +139,7 @@ def __init__(
         K: Union[int, float] = 15,
         kernel_type: str = "euc",
         sigma2_init_scale: Optional[Union[int, float]] = 0.1,
+        sigma2_end: Optional[Union[int, float]] = None,
         gamma_a: float = 1.0,
         gamma_b: float = 1.0,
         kappa: Union[float, np.ndarray] = 1.0,
@@ -193,6 +194,7 @@ def __init__(
         self.kernel_type = kernel_type
         self.kernel_bandwidth = beta
         self.sigma2_init_scale = sigma2_init_scale
+        self.sigma2_end = sigma2_end
         self.partial_robust_level = partial_robust_level
         self.normalize_c = normalize_c
         self.normalize_g = normalize_g
@@ -281,6 +283,9 @@ def run(
             self.XAHat = self.VnA + self.RnA
             self._update_sigma2(iter=iter)
 
+        if self.sigma2_end is not None:
+            self.sigma2 = _data(self.nx, self.sigma2_end, self.type_as)
+
         # Retrieve the full cell-cell assignment
         if self.return_mapping and self.SVI_mode:
             self.SVI_mode = False
@@ -588,7 +593,8 @@ def _normalize_coords(
         # get the global means for whole coords if "separate_mean" is True
         if not self.separate_mean:
             global_mean = self.nx.mean(normalize_means, axis=0)
-            normalize_means = self.nx.full((len(coords), self.D), global_mean)
+            normalize_means = self.nx.repeat(global_mean, 2, axis=0)
+            # normalize_means = self.nx.full((len(coords), self.D), global_mean)
 
         # move each coords to zero center and calculate the normalization scale
         for i in range(len(coords)):
@@ -598,7 +604,7 @@ def _normalize_coords(
             )
             normalize_scales[i] = normalize_scale
 
-        # get the global scale for whole coords if "separate_scale" is True
+        # get the global scale for whole coords if "separate_scale" is False
         if not self.separate_scale:
             global_scale = self.nx.mean(normalize_scales)
             normalize_scales = self.nx.full((len(coords),), global_scale)
@@ -816,7 +822,11 @@ def _construct_kernel(
         """
 
         unique_spatial_coords = _unique(self.nx, self.coordsA, 0)
-        inducing_variables_idx = np.random.choice(unique_spatial_coords.shape[0], inducing_variables_num, replace=False)
+        inducing_variables_idx = (
+            np.random.choice(unique_spatial_coords.shape[0], inducing_variables_num, replace=False)
+            if unique_spatial_coords.shape[0] > inducing_variables_num
+            else np.arange(unique_spatial_coords.shape[0])
+        )
         self.inducing_variables = unique_spatial_coords[inducing_variables_idx, :]
         # (self.inducing_variables, _) = sample(
         #     X=unique_spatial_coords, n_sampling=inducing_variables_num, sampling_method=sampling_method
@@ -1451,6 +1461,13 @@ def _wrap_output(
             self.P = self.nx.to_numpy(self.P)
 
         if not (self.vecfld_key_added is None):
+            norm_dict = {
+                "mean_transformed": self.nx.to_numpy(self.normalize_means[1]),
+                "mean_fixed": self.nx.to_numpy(self.normalize_means[0]),
+                "scale": self.nx.to_numpy(self.normalize_scales[1]),
+                "scale_transformed": self.nx.to_numpy(self.normalize_scales[1]),
+                "scale_fixed": self.nx.to_numpy(self.normalize_scales[0]),
+            }
 
             self.vecfld = {
                 "R": self.nx.to_numpy(self.R),

spateo/tdr/morphometrics/morphofield/gaussian_process.py

@@ -80,20 +80,38 @@ def _con_K_geodist(
         return K
 
 
+# def _gp_velocity(X: np.ndarray, vf_dict: dict) -> np.ndarray:
+#     # pre_scale = vf_dict["pre_norm_scale"]
+#     norm_x = (X - vf_dict["norm_dict"]["mean_transformed"]) / vf_dict["norm_dict"]["scale_transformed"]
+#     if vf_dict["kernel_dict"]["dist"] == "cdist":
+#         quary_kernel = _con_K(norm_x, vf_dict["inducing_variables"], vf_dict["beta"])
+#     elif vf_dict["kernel_dict"]["dist"] == "geodist":
+#         quary_kernel = _con_K_geodist(norm_x, vf_dict["kernel_dict"], vf_dict["beta"])
+#     else:
+#         raise ValueError(f"current only support cdist and geodist")
+#     quary_velocities = np.dot(quary_kernel, vf_dict["Coff"])
+#     quary_rigid = np.dot(norm_x, vf_dict["R"].T) + vf_dict["t"]
+#     quary_norm_x = quary_velocities + quary_rigid
+#     quary_x = quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] + vf_dict["norm_dict"]["mean_fixed"]
+#     _velocities = quary_x - X
+#     return _velocities / 10000
+
+
 def _gp_velocity(X: np.ndarray, vf_dict: dict) -> np.ndarray:
     # pre_scale = vf_dict["pre_norm_scale"]
     norm_x = (X - vf_dict["norm_dict"]["mean_transformed"]) / vf_dict["norm_dict"]["scale_transformed"]
     if vf_dict["kernel_dict"]["dist"] == "cdist":
-        quary_kernel = _con_K(norm_x, vf_dict["X_ctrl"], vf_dict["beta"])
+        quary_kernel = _con_K(norm_x, vf_dict["inducing_variables"], vf_dict["beta"])
     elif vf_dict["kernel_dict"]["dist"] == "geodist":
         quary_kernel = _con_K_geodist(norm_x, vf_dict["kernel_dict"], vf_dict["beta"])
     else:
         raise ValueError(f"current only support cdist and geodist")
-    quary_velocities = np.dot(quary_kernel, vf_dict["C"])
+    quary_velocities = np.dot(quary_kernel, vf_dict["Coff"])
     quary_rigid = np.dot(norm_x, vf_dict["R"].T) + vf_dict["t"]
     quary_norm_x = quary_velocities + quary_rigid
-    quary_x = quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] + vf_dict["norm_dict"]["mean_fixed"]
-    _velocities = quary_x - X
+    _velocities = (
+        quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] - norm_x * vf_dict["norm_dict"]["scale_transformed"]
+    )
     return _velocities / 10000
 
 

spateo/tdr/morphometrics/morphofield_dg/GPVectorField.py

@@ -159,27 +159,27 @@ def Jacobian_GP_gaussian_kernel(X: np.ndarray, vf_dict: dict, vectorize: bool =
     x_norm = (X - vf_dict["norm_dict"]["mean_transformed"]) / vf_dict["norm_dict"]["scale_transformed"]
     if x_norm.ndim == 1:
         if vf_dict["kernel_dict"]["dist"] == "cdist":
-            K, D = _con_K(x_norm[None, :], vf_dict["X_ctrl"], vf_dict["beta"], return_d=True)
+            K, D = _con_K(x_norm[None, :], vf_dict["inducing_variables"], vf_dict["beta"], return_d=True)
         else:
             K, D = _con_K_geodist(x_norm[None, :], vf_dict["kernel_dict"], vf_dict["beta"], return_d=True)
-        J = (vf_dict["C"].T * K) @ D[0].T
+        J = (vf_dict["Coff"].T * K) @ D[0].T
     elif not vectorize:
         n, d = x_norm.shape
         J = np.zeros((d, d, n))
         for i, xi in enumerate(x_norm):
             if vf_dict["kernel_dict"]["dist"] == "cdist":
-                K, D = _con_K(xi[None, :], vf_dict["X_ctrl"], vf_dict["beta"], return_d=True)
+                K, D = _con_K(xi[None, :], vf_dict["inducing_variables"], vf_dict["beta"], return_d=True)
             else:
                 K, D = _con_K_geodist(xi[None, :], vf_dict["kernel_dict"], vf_dict["beta"], return_d=True)
-            J[:, :, i] = (vf_dict["C"].T * K) @ D[0].T
+            J[:, :, i] = (vf_dict["Coff"].T * K) @ D[0].T
     else:
         if vf_dict["kernel_dict"]["dist"] == "cdist":
-            K, D = _con_K(x_norm, vf_dict["X_ctrl"], vf_dict["beta"], return_d=True)
+            K, D = _con_K(x_norm, vf_dict["inducing_variables"], vf_dict["beta"], return_d=True)
         else:
             K, D = _con_K_geodist(x_norm, vf_dict["kernel_dict"], vf_dict["beta"], return_d=True)
         if K.ndim == 1:
             K = K[None, :]
-        J = np.einsum("nm, mi, njm -> ijn", K, vf_dict["C"], D)
+        J = np.einsum("nm, mi, njm -> ijn", K, vf_dict["Coff"], D)
 
     return -2 * vf_dict["beta"] * J * pre_scale