def _run_deterministic(
ixs: Union[prange, np.ndarray],
indices: np.ndarray,
indptr: np.ndarray,
expression: np.ndarray,
velocity: np.ndarray,
backward: bool = False,
backward_mode: BackwardMode = BackwardMode.TRANSPOSE,
softmax_scale: float = 1,
queue=None,
) -> np.ndarray:
starts = _calculate_starts(indptr, ixs)
probs_cors = np.empty((2, starts[-1]))
for i in prange(len(ixs)):
ix = ixs[i]
start, end = indptr[ix], indptr[ix + 1]
nbhs_ixs = indices[start:end]
n_neigh = len(nbhs_ixs)
assert n_neigh, f"Cell {ix} does not have any neighbors."
W = expression[nbhs_ixs, :] - expression[ix, :]
if not backward or (backward and backward_mode == BackwardMode.NEGATE):
v = np.expand_dims(velocity[ix], 0)
# for the negate backward mode
if backward:
v *= -1
if np.all(v == 0):
p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
else:
p, c = _predict_transition_probabilities_numpy(
v,
W,
softmax_scale=softmax_scale,
)
else:
# compute how likely all neighbors are to transition to this cell
V = velocity[nbhs_ixs, :]
p, c = _predict_transition_probabilities_numpy(
V,
-1 * W,
softmax_scale=softmax_scale,
)
probs_cors[0, starts[i]:starts[i] + n_neigh] = p
probs_cors[1, starts[i]:starts[i] + n_neigh] = c
if queue is not None:
queue.put(1)
if queue is not None:
queue.put(None)
return probs_cors
def test_zero_unif_sum_to_1_vector(self):
sum_to_1, zero = _get_probs_for_zero_vec(10)
assert zero.shape == (10,)
assert sum_to_1.shape == (10,)
np.testing.assert_array_equal(zero, np.zeros_like(zero))
np.testing.assert_allclose(sum_to_1, np.ones_like(sum_to_1) / sum_to_1.shape[0])
assert np.isclose(sum_to_1.sum(), 1.0)
def _run_stochastic(
ixs: np.ndarray,
indices: np.ndarray,
indptr: np.ndarray,
expression: np.ndarray,
expectation: np.ndarray,
variance: np.ndarray,
softmax_scale: float = 1,
queue=None,
) -> np.ndarray:
starts = _calculate_starts(indptr, ixs)
probs_cors = np.empty((2, starts[-1]))
for i, ix in enumerate(ixs):
start, end = indptr[ix], indptr[ix + 1]
nbhs_ixs = indices[start:end]
n_neigh = len(nbhs_ixs)
assert n_neigh, f"Cell {ix} does not have any neighbors."
v = expectation[ix]
if np.all(v == 0):
p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
else:
# compute the Hessian tensor, and turn it into a matrix that has the diagonal elements in its rows
W = expression[nbhs_ixs, :] - expression[ix, :]
H = _predict_transition_probabilities_jax_H(
v, W, softmax_scale=softmax_scale)
H_diag = np.array([np.diag(h) for h in H])
# compute zero order term
p_0, c = _predict_transition_probabilities_numpy(
v[None, :], W, softmax_scale=softmax_scale)
# compute second order term (note that the first order term cancels)
p_2 = 0.5 * H_diag.dot(variance[ix])
# combine both to give the second order Taylor approximation. Can sometimes be negative because we
# neglected higher order terms, so force it to be non-negative
p = np.clip(p_0 + p_2, a_min=0, a_max=1)
mask = np.isnan(p)
p[mask] = 0
c[mask] = 0
if np.all(p == 0):
p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
sum_ = fsum(p)
if not np.isclose(sum_, 1.0, rtol=_RTOL):
p[~mask] = p[~mask] / sum_
probs_cors[0, starts[i]:starts[i] + n_neigh] = p
probs_cors[1, starts[i]:starts[i] + n_neigh] = c
if queue is not None:
queue.put(1)
if queue is not None:
queue.put(None)
return probs_cors
def _run_stochastic(
ixs: np.ndarray,
scheme: Callable,
indices: np.ndarray,
indptr: np.ndarray,
expression: np.ndarray,
expectation: np.ndarray,
variance: np.ndarray,
softmax_scale: float = 1,
queue=None,
) -> np.ndarray:
if not hasattr(scheme, "hessian"):
raise AttributeError()
starts = _calculate_starts(indptr, ixs)
probs_cors = np.empty((2, starts[-1]))
n_genes = expression.shape[1]
for i, ix in enumerate(ixs):
start, end = indptr[ix], indptr[ix + 1]
nbhs_ixs = indices[start:end]
n_neigh = len(nbhs_ixs)
assert n_neigh, f"Cell {ix} does not have any neighbors."
v = expectation[ix]
if np.all(v == 0):
p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
else:
# compute the Hessian tensor, and turn it into a matrix that has the diagonal elements in its rows
W = expression[nbhs_ixs, :] - expression[ix, :]
H = scheme.hessian(v, W, softmax_scale)
if H.shape == (n_neigh, n_genes):
H_diag = H
elif H.shape == (n_neigh, n_genes, n_genes):
H_diag = np.array([np.diag(h) for h in H])
else:
raise ValueError(
f"Expected full Hessian matrix of shape `{(n_neigh, n_genes, n_genes)}` "
f"or its diagonal of shape `{(n_neigh, n_genes)}`, found `{H.shape}`."
)
# compute zero order term
p_0, c = scheme(v[None, :], W, softmax_scale)
# compute second order term (note that the first order term cancels)
p_2 = 0.5 * H_diag.dot(variance[ix])
# combine both to give the second order Taylor approximation. Can sometimes be negative because we
# neglected higher order terms, so force it to be non-negative
p = np.clip(p_0 + p_2, a_min=0, a_max=1)
mask = np.isnan(p)
p[mask] = 0
c[mask] = 0
if np.all(p == 0):
p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
sum_ = fsum(p)
if not np.isclose(sum_, 1.0, rtol=_RTOL):
p[~mask] = p[~mask] / sum_
probs_cors[0, starts[i]:starts[i] + n_neigh] = p
probs_cors[1, starts[i]:starts[i] + n_neigh] = c
if queue is not None:
queue.put(1)
if queue is not None:
queue.put(None)
return probs_cors