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