Fix numba wasserstein

Cell_BLAST/blast.py

@@ -23,7 +23,8 @@
 MINIMAL = 0
 
 
-def _wasserstein_distance_impl(x: np.ndarray, y: np.ndarray): # pragma: no cover
+@numba.jit(nopython=True, nogil=True, cache=True)
+def wasserstein_distance(x: np.ndarray, y: np.ndarray): # pragma: no cover
  x_sorter = np.argsort(x)
  y_sorter = np.argsort(y)
  xy = np.concatenate((x, y))
@@ -34,37 +35,6 @@ def _wasserstein_distance_impl(x: np.ndarray, y: np.ndarray): # pragma: no cove
  return np.sum(np.multiply(np.abs(x_cdf - y_cdf), deltas))
 
 
-def _energy_distance_impl(x: np.ndarray, y: np.ndarray): # pragma: no cover
- x_sorter = np.argsort(x)
- y_sorter = np.argsort(y)
- xy = np.concatenate((x, y))
- xy.sort()
- deltas = np.diff(xy)
- x_cdf = np.searchsorted(x[x_sorter], xy[:-1], "right") / x.size
- y_cdf = np.searchsorted(y[y_sorter], xy[:-1], "right") / y.size
- return np.sqrt(2 * np.sum(np.multiply(np.square(x_cdf - y_cdf), deltas)))
-
-
-@numba.extending.overload(
- scipy.stats.wasserstein_distance, jit_options={"nogil": True, "cache": True}
-)
-def _wasserstein_distance(x: np.ndarray, y: np.ndarray): # pragma: no cover
- if (
- x == numba.float32[::1] and y == numba.float32[::1]
- ) or (
- x == numba.float64[::1] and y == numba.float64[::1]
- ):
- return _wasserstein_distance_impl
-
-
-@numba.extending.overload(
- scipy.stats.energy_distance, jit_options={"nogil": True, "cache": True}
-)
-def _energy_distance(x: np.ndarray, y: np.ndarray): # pragma: no cover
- if x == numba.float32[::1] and y == numba.float32[::1]:
- return _energy_distance_impl
-
-
 @numba.jit(nopython=True, nogil=True, cache=True)
 def ed(x: np.ndarray, y: np.ndarray): # pragma: no cover
  r"""
@@ -230,10 +200,10 @@ def npd_v1(
  np.std(y_posterior) + np.float32(eps)
  )
  return 0.5 * (
- scipy.stats.wasserstein_distance(
+ wasserstein_distance(
  xy_posterior1[: len(x_posterior)], xy_posterior1[-len(y_posterior) :]
  )
- + scipy.stats.wasserstein_distance(
+ + wasserstein_distance(
  xy_posterior2[: len(x_posterior)], xy_posterior2[-len(y_posterior) :]
  )
  )
@@ -939,13 +909,15 @@ def query(
  hits,
  dist,
  pval,
- query
- if store_dataset
- else anndata.AnnData(
- X=scipy.sparse.csr_matrix((query.shape[0], 0)),
- obs=pd.DataFrame(index=query.obs.index),
- var=pd.DataFrame(),
- uns={},
+ (
+ query
+ if store_dataset
+ else anndata.AnnData(
+ X=scipy.sparse.csr_matrix((query.shape[0], 0)),
+ obs=pd.DataFrame(index=query.obs.index),
+ var=pd.DataFrame(),
+ uns={},
+ )
  ),
  )