fix some bugs related to differential geometry calculation

spateo/alignment/methods/morpho_class.py

@@ -621,7 +621,6 @@ def _normalize_coords(
                 self.nx.einsum("ij->", self.nx.einsum("ij,ij->ij", coords[i], coords[i])) / coords[i].shape[0]
             )
             normalize_scales[i] = normalize_scale
-
         # get the global scale for whole coords if "separate_scale" is False
         if not self.separate_scale:
             global_scale = self.nx.mean(normalize_scales)
@@ -638,7 +637,7 @@ def _normalize_coords(
         if self.verbose:
             lm.main_info(message=f"Spatial coordinates normalization params:", indent_level=1)
             lm.main_info(message=f"Scale: {normalize_scales[:2]}...", indent_level=2)
-            lm.main_info(message=f"Scale: {normalize_means[:2]}...", indent_level=2)
+            lm.main_info(message=f"Mean: {normalize_means[:2]}...", indent_level=2)
 
     def _normalize_exps(
         self,

spateo/alignment/methods/utils.py

@@ -1051,16 +1051,12 @@ def get_P_core(
 
     # Calculate spatial probability with sigma2_variance
     spatial_prob = calc_probability(nx, spatial_dist, "gauss", probability_parameter=sigma2 / sigma2_variance)  # N x M
-    # print(spatial_prob.sum())
     # TODO: everytime this will generate D/2 on GPU, may influence the runtime
     outlier_s = samples_s * spatial_dist.shape[0]
     # outlier_s = samples_s
-    # print(outlier_s)
     spatial_outlier = _power(nx)((2 * _pi(nx) * sigma2), Dim / 2) * (1 - gamma) / (gamma * outlier_s)  # scalar
-    # print(spatial_outlier)
     # TODO: the position of the following is unclear
     spatial_inlier = 1 - spatial_outlier / (spatial_outlier + nx.sum(spatial_prob, axis=0, keepdims=True))  # 1 x M
-    # print(spatial_inlier.mean())
     spatial_prob = spatial_prob * model_mul
 
     # spatial P

spateo/tdr/morphometrics/morphofield/gaussian_process.py

@@ -83,10 +83,12 @@ def _con_K_geodist(
 # 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":
+#     if vf_dict["kernel_type"] == "euc":
 #         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"])
+#     elif vf_dict["kernel_type"] == "geodist":
+#         pass
+#         # not implemented yet
+#         # 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"])
@@ -97,33 +99,84 @@ def _con_K_geodist(
 #     return _velocities / 10000
 
 
-def _gp_velocity(X: np.ndarray, vf_dict: dict) -> np.ndarray:
-    # pre_scale = vf_dict["pre_norm_scale"]
+def _gp_velocity(
+    X: np.ndarray,
+    vf_dict: dict,
+    nonrigid_only: bool = False,
+) -> np.ndarray:
     norm_x = (X - vf_dict["norm_dict"]["mean_transformed"]) / vf_dict["norm_dict"]["scale_transformed"]
     if vf_dict["kernel_type"] == "euc":
         quary_kernel = _con_K(norm_x, vf_dict["inducing_variables"], vf_dict["beta"])
     elif vf_dict["kernel_type"] == "geodist":
-        pass
+        raise NotImplementedError("geodist is not implemented yet")
         # not implemented yet
         # 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
-    _velocities = (
-        quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] - norm_x * vf_dict["norm_dict"]["scale_transformed"]
-    )
-    _velocities = quary_velocities * vf_dict["norm_dict"]["scale_fixed"]
+    if nonrigid_only:
+        _velocities = (
+            quary_velocities * vf_dict["norm_dict"]["scale_fixed"]
+            + (vf_dict["norm_dict"]["scale_fixed"] - vf_dict["norm_dict"]["scale_transformed"]) * norm_x
+        )
+    else:
+        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_type"] == "euc":
+#         quary_kernel = _con_K(norm_x, vf_dict["inducing_variables"], vf_dict["beta"])
+#     elif vf_dict["kernel_type"] == "geodist":
+#         pass
+#         # not implemented yet
+#         # 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
+#     _velocities = (
+#         quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] - norm_x * vf_dict["norm_dict"]["scale_transformed"]
+#     )
+#     _velocities = quary_velocities * vf_dict["norm_dict"]["scale_fixed"]
+#     return _velocities / 10000
+
+# def _gp_velocity(X: np.ndarray, vf_dict: dict) -> np.ndarray:
+#     # pre_scale = vf_dict["pre_norm_scale"]
+#     # pre_scale = vf_dict["norm_dict"]["scale_fixed"] / vf_dict["norm_dict"]["scale_transformed"]
+#     norm_x = (X - vf_dict["norm_dict"]["mean_transformed"]) / vf_dict["norm_dict"]["scale_transformed"]
+#     if vf_dict["kernel_type"] == "euc":
+#         quary_kernel = _con_K(norm_x, vf_dict["inducing_variables"], vf_dict["beta"])
+#     elif vf_dict["kernel_type"] == "geodist":
+#         pass
+#         # not implemented yet
+#         # 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
+#     _velocities = (
+#         quary_norm_x * vf_dict["norm_dict"]["scale_fixed"] - norm_x * vf_dict["norm_dict"]["scale_transformed"]
+#     )
+#     _velocities = quary_velocities * vf_dict["norm_dict"]["scale_fixed"]
+#     # _velocities = quary_velocities
+#     return _velocities / 10000
+
+
 def morphofield_gp(
     adata: AnnData,
     spatial_key: str = "align_spatial",
     vf_key: str = "VecFld_morpho",
     NX: Optional[np.ndarray] = None,
     grid_num: Optional[List[int]] = None,
+    nonrigid_only: bool = False,
     inplace: bool = True,
 ) -> Optional[AnnData]:
     """
@@ -136,6 +189,7 @@ def morphofield_gp(
         key_added: The key that will be used for the vector field key in ``.uns``.
         NX: The spatial coordinates of new data point. If NX is None, generate new points based on grid_num.
         grid_num: The number of grids in each dimension for generating the grid velocity. Default is ``[50, 50, 50]``.
+        nonrigid_only: If True, only the nonrigid part of the vector field will be calculated.
         inplace: Whether to copy adata or modify it inplace.
 
     Returns:
@@ -155,7 +209,7 @@ def morphofield_gp(
     if vf_key in adata.uns.keys():
         vf_dict = adata.uns[vf_key]
         vf_dict["X"] = np.asarray(adata.obsm[spatial_key], dtype=float)
-        vf_dict["V"] = _gp_velocity(vf_dict["X"], vf_dict=vf_dict)
+        vf_dict["V"] = _gp_velocity(vf_dict["X"], vf_dict=vf_dict, nonrigid_only=nonrigid_only)
 
         if not (NX is None):
             predict_X = NX
@@ -166,7 +220,7 @@ def morphofield_gp(
             _, _, Grid, grid_in_hull = get_X_Y_grid(X=vf_dict["X"].copy(), Y=vf_dict["V"].copy(), grid_num=grid_num)
             predict_X = Grid
         vf_dict["grid"] = predict_X
-        vf_dict["grid_V"] = _gp_velocity(predict_X, vf_dict=vf_dict)
+        vf_dict["grid_V"] = _gp_velocity(predict_X, vf_dict=vf_dict, nonrigid_only=nonrigid_only)
 
         vf_dict["method"] = "gaussian_process"
         lm.main_finish_progress(progress_name="morphofield")

spateo/tdr/morphometrics/morphofield/trajectory.py

@@ -18,6 +18,7 @@ def morphopath(
     t_end: Optional[Union[int, float]] = None,
     average: bool = False,
     cores: int = 1,
+    nonrigid_only: bool = False,
     inplace: bool = True,
     **kwargs,
 ) -> Optional[AnnData]:
@@ -39,6 +40,7 @@ def morphopath(
                  ``False``, no averaging will be applied.
         cores: Number of cores to calculate path integral for predicting cell fate. If cores is set to be > 1,
                multiprocessing will be used to parallel the fate prediction.
+        nonrigid_only: If True, only the nonrigid part of the vector field will be calculated.
         inplace: Whether to copy adata or modify it inplace.
         **kwargs: Additional parameters that will be passed into the ``fate`` function.
 
@@ -89,7 +91,7 @@ def morphopath(
             direction=direction,
             average=average,
             cores=cores,
-            VecFld_true=lambda X: _gp_velocity(X=X, vf_dict=fate_adata.uns[vf_key]),
+            VecFld_true=lambda X: _gp_velocity(X=X, vf_dict=fate_adata.uns[vf_key], nonrigid_only=nonrigid_only),
             **kwargs,
         )
     elif method == "sparsevfc":

spateo/tdr/morphometrics/morphofield_dg/GPVectorField.py

@@ -194,7 +194,7 @@ class GPVectorField:
     def __init__(self):
         self.data = {}
 
-    def from_adata(self, adata: AnnData, vf_key: str = "VecFld"):
+    def from_adata(self, adata: AnnData, vf_key: str = "VecFld", nonrigid_only: bool = False):
         from ..morphofield.gaussian_process import _gp_velocity
 
         if vf_key in adata.uns.keys():
@@ -206,7 +206,8 @@ def from_adata(self, adata: AnnData, vf_key: str = "VecFld"):
             )
 
         self.vf_dict = vf_dict
-        self.func = lambda x: _gp_velocity(x, vf_dict)
+        self.nonrigid_only = nonrigid_only
+        self.func = lambda x: _gp_velocity(x, vf_dict, nonrigid_only=nonrigid_only)
         self.data["X"] = vf_dict["X"]
         self.data["V"] = vf_dict["V"]
 
@@ -216,7 +217,7 @@ def get_data(self) -> Tuple[np.ndarray, np.ndarray]:
     def compute_velocity(self, X: np.ndarray):
         from ..morphofield.gaussian_process import _gp_velocity
 
-        return _gp_velocity(X, self.vf_dict)
+        return _gp_velocity(X, self.vf_dict, nonrigid_only=self.nonrigid_only)
 
     def compute_acceleration(self, X: Optional[np.ndarray] = None, **kwargs):
         X = self.data["X"] if X is None else X