Debug and refactor scPotential

dynamo/plot/scPotential.py

@@ -15,9 +15,9 @@
 
 def show_landscape(
     adata: AnnData,
-    Xgrid: np.ndarray,
-    Ygrid: np.ndarray,
-    Zgrid: np.ndarray,
+    Xgrid: Optional[np.ndarray] = None,
+    Ygrid: Optional[np.ndarray] = None,
+    Zgrid: Optional[np.ndarray] = None,
     basis: str = "umap",
     save_show_or_return: Literal["save", "show", "return"] = "show",
     save_kwargs: Dict[str, Any] = {},
@@ -49,9 +49,9 @@ def show_landscape(
             adata.uns["grid_Pot_" + basis]["Zgrid"],
         )
 
-    Xgrid = Xgrid_ if Xgrid is None else Xgrid
-    Ygrid = Ygrid_ if Ygrid is None else Ygrid
-    Zgrid = Zgrid_ if Zgrid is None else Zgrid
+    Xgrid = np.nan_to_num(Xgrid_) if Xgrid is None else np.nan_to_num(Xgrid)
+    Ygrid = np.nan_to_num(Ygrid_) if Ygrid is None else np.nan_to_num(Ygrid)
+    Zgrid = np.nan_to_num(Zgrid_) if Zgrid is None else np.nan_to_num(Zgrid)
 
     import matplotlib.pyplot as plt
     from matplotlib import cm
@@ -60,7 +60,7 @@ def show_landscape(
     from mpl_toolkits.mplot3d import Axes3D
 
     fig = plt.figure()
-    ax = fig.gca(projection="3d")
+    ax = fig.add_subplot(projection="3d")
 
     # Plot the surface.
     ls = LightSource(azdeg=0, altdeg=65)

dynamo/vectorfield/Ao.py

@@ -16,10 +16,7 @@ def f_left(X: np.ndarray, F: np.ndarray) -> np.ndarray:
 
 
 def f_left_jac(q: np.ndarray, F: np.ndarray) -> np.ndarray:
-    """
-    Analytical Jacobian of f(Q) = F.Q - (F.Q)^T, where Q is
-    an anti-symmetric matrix s.t. Q^T = -Q.
-    """
+    """Analytical Jacobian of f(Q) = F.Q - (F.Q)^T, where Q is an anti-symmetric matrix s.t. Q^T = -Q."""
     J = np.zeros((np.prod(F.shape), len(q)))
     for i in range(len(q)):
         jac = np.zeros(F.shape)
@@ -48,12 +45,13 @@ def solveQ(
     Args:
         D: A symmetric diffusion matrix.
         F: Jacobian of the vector field function evaluated at a particular point.
-        q0: Initial guess of the solution Q
+        q0: Initial guess of the solution Q.
         debug: Whether additional info of the solution is returned.
         precompute_jac: Whether the analytical Jacobian is precomputed for the optimizer.
 
     Returns:
-        The solved anti-symmetric Q matrix and the value of the right-hand side of the equation to be solved, optionally along with the value of the left-hand side of the equation and the cost of the solution if debug is True.
+        The solved anti-symmetric Q matrix and the value of the right-hand side of the equation to be solved, optionally
+        along with the value of the left-hand side of the equation and the cost of the solution if debug is True.
     """
 
     n = D.shape[0]
@@ -79,24 +77,25 @@ def Ao_pot_map(
     vecFunc: Callable, X: np.ndarray, fjac: Optional[Callable] = None, D: Optional[np.ndarray] = None, **kwargs
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List, List, List]:
     """Mapping potential landscape with the algorithm developed by Ao method.
-    References: Potential in stochastic differential equations: novel construction. Journal of physics A: mathematical and
-        general, Ao Ping, 2004
+    References: Potential in stochastic differential equations: novel construction. Journal of physics A: mathematical
+        and general, Ao Ping, 2004
 
     Args:
         vecFunc: The vector field function.
         X: A (n_cell x n_dim) matrix of coordinates where the potential function is evaluated.
         fjac: function that returns the Jacobian of the vector field function evaluated at a particular point.
-        D: Diffusion matrix. It must be a square matrix with size corresponds to the number of columns (features) in the X matrix.
+        D: Diffusion matrix. It must be a square matrix with size corresponds to the number of columns (features) in the
+            X matrix.
 
     Returns:
-        X: A matrix storing the x-coordinates on the two-dimesional grid.
+        X: A matrix storing the x-coordinates on the two-dimensional grid.
         U: A matrix storing the potential value at each position.
         P: Steady state distribution or the Boltzmann-Gibbs distribution for the state variable.
         vecMat: List velocity vector at each position from X.
-        S: List of constant symmetric and semi-positive matrix or friction (dissipative) matrix, corresponding to the divergence part,
+        S: List of constant symmetric and semi-positive matrix or friction (dissipative) matrix, corresponding to the
+            divergence part, at each position from X.
+        A: List of constant antisymmetric matrix or transverse (non-dissipative) matrix, corresponding to the curl part,
             at each position from X.
-        A: List of constant antisymmetric matrix or transverse (non-dissipative) matrix, corresponding to the curl part, at each position
-            from X.
     """
 
     import numdifftools as nda
@@ -141,4 +140,35 @@ def Ao_pot_map_jac(fjac, X, D=None, **kwargs):
         H = np.linalg.inv(D + Q).dot(F)
         U[i] = -0.5 * X_s.dot(H).dot(X_s)
 
-    return U.flatten()
+    return U.flatten()
+
+
+def construct_Ao_potential_grid(
+    X: np.ndarray,
+    U: np.ndarray,
+    interpolation_method: str = "linear",
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Construct a grid of potential landscape from the given coordinates and potential values.
+
+    Args:
+        X: A matrix storing the x-coordinates on the two-dimensional grid.
+        U: A matrix storing the potential value at each position.
+        interpolation_method: The method of interpolation. Defaults to "linear".
+
+    Returns:
+        A tuple containing matrices storing the x, y ,z coordinates on the 3-dimensional grid.
+    """
+    from scipy.interpolate import griddata
+
+    xi, yi = np.linspace(X[:, 0].min(), X[:, 0].max(), 100), np.linspace(X[:, 1].min(), X[:, 1].max(), 100)
+    Xgrid, Ygrid = np.meshgrid(xi, yi)
+
+    Zgrid = griddata(
+        X,
+        U,
+        np.vstack((Xgrid.flatten(), Ygrid.flatten())).T,
+        method=interpolation_method,
+    )
+    Zgrid = Zgrid.reshape(Xgrid.shape)
+
+    return Xgrid, Ygrid, Zgrid

dynamo/vectorfield/Bhattacharya.py

@@ -30,17 +30,18 @@ def path_integral(
         numTimeSteps: A high-enough number for convergence with given dt.
 
     Returns:
-        numAttractors: Number of attractors identified by the path integral approach.
-        attractors_num_X_Y: Attractor number and the corresponding x, y coordinates.
-        sepx_old_new_pathNum: The IDs of the two attractors for each separaxis per row.
-        numPaths_att: Number of paths per attractor
-        numPaths: Total Number of paths for defined grid spacing.
-        numTimeSteps: A high-enough number for convergence with given dt.
-        pot_path: Potential along the path. (dimension: numPaths x numTimeSteps)
-        path_tag: Tag for given path (to denote basin of attraction). (dimension: numPaths x 1)
-        attractors_pot: Potential value of each identified attractors by the path integral approach.
-        x_path: x-coord along path.
-        y_path: y-coord along path.
+        A tuple containing:
+            numAttractors: Number of attractors identified by the path integral approach.
+            attractors_num_X_Y: Attractor number and the corresponding x, y coordinates.
+            sepx_old_new_pathNum: The IDs of the two attractors for each separaxis per row.
+            numPaths_att: Number of paths per attractor
+            numPaths: Total Number of paths for defined grid spacing.
+            numTimeSteps: A high-enough number for convergence with given dt.
+            pot_path: Potential along the path. (dimension: numPaths x numTimeSteps)
+            path_tag: Tag for given path (to denote basin of attraction). (dimension: numPaths x 1)
+            attractors_pot: Potential value of each identified attractors by the path integral approach.
+            x_path: x-coord along path.
+            y_path: y-coord along path.
     """
 
     # -- First, generate potential surface from deterministic rate equations –
@@ -116,7 +117,7 @@ def path_integral(
 
                 # update dxdt, dydt
 
-                dxdt, dydt = VecFnc([x_p, y_p])
+                dxdt, dydt = VecFnc(np.array([x_p, y_p]))
 
                 # update x, y
                 dx = dxdt * dt
@@ -215,11 +216,11 @@ def path_integral(
 
                         # check if start points of current and previous paths are "adjacent" - if so, assign separatrix
                         if startPt_dist_sqr < (2 * (xyGridSpacing**2)):
-                            curr_sepx = [
-                                path_tag[path_counter - 1],
-                                path_tag[path_counter],
-                                (path_counter - 1),
-                            ]  # create array
+                            curr_sepx = np.array([
+                                path_tag[path_counter - 1][0],
+                                path_tag[path_counter][0],
+                                path_counter - 1,
+                            ])  # create array
                             sepx_old_new_pathNum = (
                                 np.vstack((sepx_old_new_pathNum, curr_sepx))
                                 if sepx_old_new_pathNum is not None
@@ -247,11 +248,11 @@ def path_integral(
                         if prev_attr_new == 1:
                             # check if start points of current and previous paths are "adjacent"  - if so, assign separatrix
                             if startPt_dist_sqr < (2 * (xyGridSpacing**2)):
-                                curr_sepx = [
-                                    path_tag[path_counter - 1],
-                                    path_tag[path_counter],
+                                curr_sepx = np.array([
+                                    path_tag[path_counter - 1][0],
+                                    path_tag[path_counter][0],
                                     (path_counter - 1),
-                                ]  # create array
+                                ])  # create array
                                 sepx_old_new_pathNum = (
                                     np.vstack((sepx_old_new_pathNum, curr_sepx))
                                     if sepx_old_new_pathNum is not None
@@ -335,9 +336,10 @@ def alignment(
             `CloughTocher2DInterpolator` for more details.
 
     Returns:
-        Xgrid: x-coordinates of the Grid produced from the meshgrid function.
-        Ygrid: y-coordinates of the Grid produced from the meshgrid function.
-        Zgrid: z-coordinates or potential at each of the x/y coordinate.
+        A tuple containing:
+            Xgrid: x-coordinates of the Grid produced from the meshgrid function.
+            Ygrid: y-coordinates of the Grid produced from the meshgrid function.
+            Zgrid: z-coordinates or potential at each of the x/y coordinate.
     """
 
     # -- need 1-D "lists" (vectors) to plot all x,y, Pot values along paths --