Vectorfield refactors

dynamo/prediction/perturbation.py

@@ -10,12 +10,10 @@
 from ..vectorfield.scVectorField import vector_field_function_knockout
 from ..vectorfield.vector_calculus import (
     jacobian,
-    rank_cell_groups,
-    rank_cells,
-    rank_genes,
     vecfld_from_adata,
     vector_transformation,
 )
+from ..vectorfield.rank_vf import rank_cell_groups, rank_cells, rank_genes
 from .utils import expr_to_pca, pca_to_expr, z_score, z_score_inv
 
 

dynamo/vectorfield/Ao.py

@@ -1,20 +1,20 @@
+from typing import Callable, List, Optional, Tuple, Union
+
 import numpy as np
-import scipy.sparse as sp
 from scipy.optimize import least_squares
-from tqdm import tqdm
 
 from ..tools.utils import condensed_idx_to_squareform_idx, squareform, timeit
 
 # from scPotential import show_landscape
 
 
-def f_left(X, F):
+def f_left(X: np.ndarray, F: np.ndarray) -> np.ndarray:
     """An auxiliary function for fast computation of F.X - (F.X)^T"""
     R = F.dot(X)
     return R - R.T
 
 
-def f_left_jac(q, F):
+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.
@@ -32,28 +32,27 @@ def f_left_jac(q, F):
 
 
 @timeit
-def solveQ(D, F, q0=None, debug=False, precompute_jac=True, **kwargs):
+def solveQ(
+    D: np.ndarray,
+    F: np.ndarray,
+    q0: Optional[np.ndarray] = None,
+    debug: bool = False,
+    precompute_jac: bool = True,
+    **kwargs
+) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray, float], Tuple[np.ndarray, np.ndarray]]:
     """Function to solve for the anti-symmetric Q matrix in the equation:
         F.Q - (F.Q)^T = F.D - (F.D)^T
     using least squares.
 
-    Parameters
-    ----------
-        D: :class:`~numpy.ndarray`
-            A symmetric diffusion matrix.
-        F: :class:`~numpy.ndarray`
-            Jacobian of the vector field function evaluated at a particular point.
-        debug: bool
-            Whether additional info of the solution is returned.
-        precompute_jac: bool
-            Whether the analytical Jacobian is precomputed for the optimizer.
-
-    Returns
-    -------
-        Q: :class:`~numpy.ndarray`
-            The solved anti-symmetric Q matrix.
-        C: :class:`~numpy.ndarray`
-            The right-hand side of the equation to be solved.
+    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
+        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.
     """
 
     n = D.shape[0]
@@ -75,35 +74,27 @@ def solveQ(D, F, q0=None, debug=False, precompute_jac=True, **kwargs):
         return Q, C
 
 
-def Ao_pot_map(vecFunc, X, D=None, **kwargs):
+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
 
-    Parameters
-    ----------
-        vecFunc: `function`
-            The vector field function
-        X: :class:`~numpy.ndarray`
-            A n_cell x n_dim matrix of coordinates where the potential function is evaluated.
-        D: None or :class:`~numpy.ndarray`
-            Diffusion matrix. It must be a square matrix with size corresponds to the number of columns (features) in the X matrix.
-
-    Returns
-    -------
-        X: :class:`~numpy.ndarray`
-            A matrix storing the x-coordinates on the two-dimesional grid.
-        U: :class:`~numpy.ndarray`
-            A matrix storing the potential value at each position.
-        P: :class:`~numpy.ndarray`
-            Steady state distribution or the Boltzmann-Gibbs distribution for the state variable.
-        vecMat: list
-            List velocity vector at each position from X.
-        S: list
-            List of constant symmetric and semi-positive matrix or friction matrix, corresponding to the divergence part,
+    Args:
+        vecFunc: The vector field function
+        fjac: function that returns the Jacobian of the vector field function evaluated at a particular point.
+        X: A (n_cell x n_dim) matrix of coordinates where the potential function is evaluated.
+        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.
+        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,
             at each position from X.
-        A: list
-            List of constant antisymmetric matrix or transverse matrix, corresponding to the curl part, at each position
+        A: List of constant antisymmetric matrix or transverse (non-dissipative) matrix, corresponding to the curl part, at each position
             from X.
     """
 
@@ -116,7 +107,7 @@ def Ao_pot_map(vecFunc, X, D=None, **kwargs):
 
     for i in range(nobs):
         X_s = X[i, :]
-        F = nda.Jacobian(vecFunc)(X_s)
+        F = nda.Jacobian(vecFunc)(X_s) if fjac is None else fjac(X_s)
         Q, _ = solveQ(D, F, **kwargs)
         H = np.linalg.inv(D + Q).dot(F)
         U[i] = -0.5 * X_s.dot(H).dot(X_s)
@@ -131,23 +122,3 @@ def Ao_pot_map(vecFunc, X, D=None, **kwargs):
     P = P / np.sum(P)
 
     return X, U, P, vecMat, S, A
-
-
-def Ao_pot_map_jac(fjac, X, D=None, **kwargs):
-    nobs, ndim = X.shape
-    if D is None:
-        D = 0.1 * np.eye(ndim)
-    elif np.isscalar(D):
-        D = D * np.eye(ndim)
-    U = np.zeros((nobs, 1))
-
-    m = int(ndim * (ndim - 1) / 2)
-    q0 = np.ones(m) * np.mean(np.diag(D)) * 1000
-    for i in tqdm(range(nobs), "Calc Ao potential"):
-        X_s = X[i, :]
-        F = fjac(X_s)
-        Q, _ = solveQ(D, F, q0=q0, **kwargs)
-        H = np.linalg.inv(D + Q).dot(F)
-        U[i] = -0.5 * X_s.dot(H).dot(X_s)
-
-    return U.flatten()

dynamo/vectorfield/Bhattacharya.py

@@ -1,53 +1,40 @@
+from typing import Tuple
+
 import numpy as np
 
 # from scipy.integrate import odeint
 from scipy.interpolate import griddata
 
 
-def path_integral(VecFnc, x_lim, y_lim, xyGridSpacing, dt=1e-2, tol=1e-2, numTimeSteps=1400):
+def path_integral(
+    VecFnc, x_lim, y_lim, xyGridSpacing, dt=1e-2, tol=1e-2, numTimeSteps=1400
+) -> Tuple[
+    int, np.ndarray, np.ndarray, np.ndarray, int, int, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray
+]:
     """A deterministic map of Waddington’s epigenetic landscape for cell fate specification
     Sudin Bhattacharya, Qiang Zhang and Melvin E. Andersen
 
-    Parameters
-    ----------
-    VecFnc
-    x_lim: `list`
-        Lower or upper limit of x-axis.
-    y_lim: `list`
-        Lower or upper limit of y-axis
-    xyGridSpacing: `float`
-        Grid spacing for "starting points" for each "path" on the pot. surface
-    dt: `float`
-        Time step for the path integral.
-    tol: `float` (default: 1.0e-2)
-        Tolerance to test for convergence.
-    numTimeSteps: `int`
-        A high-enough number for convergence with given dt.
-
-    Returns
-    -------
-    numAttractors: `int`
-        Number of attractors identified by the path integral approach.
-    attractors_num_X_Y: `numpy.ndarray`
-        Attractor number and the corresponding x, y coordinates.
-    sepx_old_new_pathNum: `numpy.ndarray`
-        The IDs of the two attractors for each separaxis per row.
-    numPaths_att `numpy.ndarray`
-        Number of paths per attractor
-    numPaths: `int`
-        Total Number of paths for defined grid spacing.
-    numTimeSteps: `int`
-        A high-enough number for convergence with given dt.
-    pot_path: `numpy.ndarray` (dimension: numPaths x numTimeSteps)
-        Potential along the path.
-    path_tag: `numpy.ndarray` (dimension: numPaths x 1)
-        Tag for given path (to denote basin of attraction).
-    attractors_pot: `numpy.ndarray`
-        Potential value of each identified attractors by the path integral approach.
-    x_path: `numpy.ndarray`
-        x-coord. along path.
-    y_path: `numpy.ndarray`
-        y-coord. along path.
+    Args:
+        VecFnc: Takes in x, y and returns the vector field at that point.
+        x_lim: Lower or upper limit of x-axis.
+        y_lim: Lower or upper limit of y-axis
+        xyGridSpacing: Grid spacing for "starting points" for each "path" on the potential surface
+        dt: Time step for the path integral.
+        tol: Tolerance to test for convergence.
+        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.
     """
 
     # -- First, generate potential surface from deterministic rate equations –
@@ -279,7 +266,7 @@ def path_integral(VecFnc, x_lim, y_lim, xyGridSpacing, dt=1e-2, tol=1e-2, numTim
                     numPaths_att[tag - 1] = numPaths_att[tag - 1] + 1
 
             # increment "path counter"
-            path_counter = path_counter + 1
+            path_counter += 1
 
     return (
         attractors_num_X_Y,
@@ -297,68 +284,54 @@ def path_integral(VecFnc, x_lim, y_lim, xyGridSpacing, dt=1e-2, tol=1e-2, numTim
 
 
 def alignment(
-    numPaths,
-    numTimeSteps,
-    pot_path,
-    path_tag,
-    attractors_pot,
-    x_path,
-    y_path,
-    grid=100,
-    interpolation_method="linear",
+    numPaths: int,
+    numTimeSteps: int,
+    pot_path: np.ndarray,
+    path_tag: np.ndarray,
+    attractors_pot: np.ndarray,
+    x_path: np.ndarray,
+    y_path: np.ndarray,
+    grid: int = 100,
+    interpolation_method: str = "linear",
 ):
     """Align potential values so all path-potentials end up at same global min and then generate potential surface with
     interpolation on a grid.
 
-    Parameters
-    ----------
-    numPaths: `int`
-        Total Number of paths for defined grid spacing.
-    numTimeSteps: `int`
-        A high-enough number for convergence with given dt.
-    pot_path: `numpy.ndarray` (dimension: numPaths x numTimeSteps)
-        Potential along the path.
-    path_tag: `numpy.ndarray` (dimension: numPaths x 1)
-        Tag for given path (to denote basin of attraction).
-    attractors_pot: `numpy.ndarray`
-        Potential value of each identified attractors by the path integral approach.
-    x_path: `numpy.ndarray`
-        x-coord. along path.
-    y_path: `numpy.ndarray`
-        y-coord. along path.
-    grid: `int`
-        No. of grid lines in x- and y- directions
-    interpolation_method: `string`
-        Method of interpolation in griddata function. One of
-
-        ``nearest``
-          return the value at the data point closest to
-          the point of interpolation.  See `NearestNDInterpolator` for
-          more details.
-
-        ``linear``
-          tessellate the input point set to n-dimensional
-          simplices, and interpolate linearly on each simplex.  See
-          `LinearNDInterpolator` for more details.
-
-        ``cubic`` (1-D)
-          return the value determined from a cubic
-          spline.
-
-        ``cubic`` (2-D)
-          return the value determined from a
-          piecewise cubic, continuously differentiable (C1), and
-          approximately curvature-minimizing polynomial surface. See
-          `CloughTocher2DInterpolator` for more details.
-
-    Returns
-    -------
-    Xgrid: `numpy.ndarray`
-        x-coordinates of the Grid produced from the meshgrid function.
-    Ygrid: `numpy.ndarray`
-            y-coordinates of the Grid produced from the meshgrid function.
-    Zgrid: `numpy.ndarray`
-            z-coordinates or potential at each of the x/y coordinate.
+    Args:
+        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.
+        grid: No. of grid lines in x- and y- directions
+        interpolation_method: Method of interpolation in griddata function. One of
+
+            ``nearest``
+            return the value at the data point closest to
+            the point of interpolation.  See `NearestNDInterpolator` for
+            more details.
+
+            ``linear``
+            tessellate the input point set to n-dimensional
+            simplices, and interpolate linearly on each simplex.  See
+            `LinearNDInterpolator` for more details.
+
+            ``cubic`` (1-D)
+            return the value determined from a cubic
+            spline.
+
+            ``cubic`` (2-D)
+            return the value determined from a
+            piecewise cubic, continuously differentiable (C1), and
+            approximately curvature-minimizing polynomial surface. See
+            `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.
     """
 
     # -- need 1-D "lists" (vectors) to plot all x,y, Pot values along paths --
@@ -410,3 +383,6 @@ def alignment(
     # print('Ran surface grid-interpolation okay!\n')
 
     return Xgrid, Ygrid, Zgrid
+
+
+# %%