diff --git a/dynamo/plot/dynamics.py b/dynamo/plot/dynamics.py index b377e6f5b..ad366dbb2 100755 --- a/dynamo/plot/dynamics.py +++ b/dynamo/plot/dynamics.py @@ -15,7 +15,7 @@ from ..dynamo_logger import main_warning from ..estimation.csc.velocity import sol_s, sol_u, solve_first_order_deg from ..estimation.tsc.utils_moments import moments -from ..tools.utils import get_mapper, get_valid_bools, index_gene, log1p_, update_dict +from ..tools.utils import get_mapper, get_valid_bools, get_vel_params, index_gene, log1p_, update_dict, update_vel_params from .scatters import scatters from .utils import ( _datashade_points, @@ -230,7 +230,8 @@ def phase_portraits( else: k_name = "gamma" - valid_id = np.isfinite(np.array(adata.var.loc[_genes, k_name], dtype="float")).flatten() + vel_params_df = get_vel_params(adata) + valid_id = np.isfinite(np.array(vel_params_df.loc[_genes, k_name], dtype="float")).flatten() genes = np.array(_genes)[valid_id].tolist() # idx = [adata.var.index.to_list().index(i) for i in genes] @@ -335,12 +336,12 @@ def phase_portraits( V_vec.A if issparse(V_vec) else V_vec, ) - if k_name in adata.var.columns: - if not ("gamma_b" in adata.var.columns) or all(adata.var.gamma_b.isna()): - adata.var.loc[:, "gamma_b"] = 0 + if k_name in vel_params_df.columns: + if not ("gamma_b" in vel_params_df.columns) or all(vel_params_df.gamma_b.isna()): + vel_params_df.loc[:, "gamma_b"] = 0 gamma, velocity_offset = ( - index_gene(adata, adata.var.loc[:, k_name].values, genes), - index_gene(adata, adata.var.gamma_b.values, genes), + index_gene(adata, vel_params_df.loc[:, k_name].values, genes), + index_gene(adata, vel_params_df.gamma_b.values, genes), ) ( gamma[~np.isfinite(list(gamma))], @@ -433,12 +434,12 @@ def phase_portraits( ) ) if "protein" in adata.obsm.keys(): - if "delta" in adata.var.columns: - gamma_P = adata.var.delta[genes].values + if "delta" in vel_params_df.columns: + gamma_P = vel_params_df.delta[genes].values velocity_offset_P = ( [0] * n_cells - if (not ("delta_b" in adata.var.columns) or adata.var.delta_b.unique() is None) - else adata.var.delta_b[genes].values + if (not ("delta_b" in vel_params_df.columns) or vel_params_df.delta_b.unique() is None) + else vel_params_df.delta_b[genes].values ) else: raise ValueError( @@ -1064,6 +1065,8 @@ def phase_portraits( despline_all(ax6) deaxis_all(ax6) + update_vel_params(adata, params_df=vel_params_df) + if save_show_or_return in ["save", "both", "all"]: s_kwargs = { "path": None, @@ -1212,7 +1215,8 @@ def dynamics( main_warning( "dynamics plot doesn't support conventional experiment type, using phase_portraits function instead." ) - phase_portraits(adata) + phase_portraits(adata, genes=genes, save_show_or_return=save_show_or_return) + return T_uniq = np.unique(T) t = np.linspace(0, T_uniq[-1], 1000) @@ -1322,8 +1326,9 @@ def dynamics( true_p = None true_params = [None, None, None] - logLL = valid_adata.var.loc[valid_gene_names, prefix + "logLL"] - est_params_df = valid_adata.var.loc[ + vel_params_df = get_vel_params(valid_adata) + logLL = vel_params_df.loc[valid_gene_names, prefix + "logLL"] + est_params_df = vel_params_df.loc[ valid_gene_names, [ prefix + "alpha", @@ -1603,7 +1608,7 @@ def dynamics( else: for i, gene_name in enumerate(valid_genes): if model == "moment": - a, b, alpha_a, alpha_i, beta, gamma = valid_adata.var.loc[ + a, b, alpha_a, alpha_i, beta, gamma = vel_params_df.loc[ gene_name, [ prefix + "a", @@ -1627,23 +1632,23 @@ def dynamics( mom_data = mom.get_all_central_moments() if has_splicing else mom.get_nosplice_central_moments() if true_param_prefix is not None: (true_a, true_b, true_alpha_a, true_alpha_i, true_beta, true_gamma,) = ( - valid_adata.var.loc[gene_name, true_param_prefix + "a"] - if true_param_prefix + "a" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "a"] + if true_param_prefix + "a" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "b"] - if true_param_prefix + "b" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "b"] + if true_param_prefix + "b" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "alpha_a"] - if true_param_prefix + "alpha_a" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha_a"] + if true_param_prefix + "alpha_a" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "alpha_i"] - if true_param_prefix + "alpha_i" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha_i"] + if true_param_prefix + "alpha_i" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "beta"] - if true_param_prefix + "beta" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "beta"] + if true_param_prefix + "beta" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) @@ -1861,7 +1866,7 @@ def dynamics( np.log1p(sl), ) - (alpha, beta, gamma, ul0, sl0, uu0, half_life,) = valid_adata.var.loc[ + (alpha, beta, gamma, ul0, sl0, uu0, half_life,) = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -1887,14 +1892,14 @@ def dynamics( l = sol_s(t, sl0, ul0, 0, beta, gamma) if true_param_prefix is not None: true_alpha, true_beta, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "beta"] - if true_param_prefix + "beta" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "beta"] + if true_param_prefix + "beta" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) @@ -1931,7 +1936,7 @@ def dynamics( if log_unnormalized and layers == ["new", "total"]: uu, ul = np.log1p(uu), np.log1p(ul) - alpha, gamma, uu0, ul0, half_life = valid_adata.var.loc[ + alpha, gamma, uu0, ul0, half_life = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -1950,11 +1955,11 @@ def dynamics( title_ = ["(unlabeled)", "(labeled)"] if true_param_prefix is not None: true_alpha, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) true_u = sol_u(t, uu0, true_alpha, true_gamma) @@ -2041,8 +2046,8 @@ def dynamics( ax.set_title(gene_name + " " + title_[j]) elif experiment_type == "kin": if model == "deterministic": - logLL = valid_adata.var.loc[valid_gene_names, prefix + "logLL"] - alpha, beta, gamma, half_life = valid_adata.var.loc[ + logLL = vel_params_df.loc[valid_gene_names, prefix + "logLL"] + alpha, beta, gamma, half_life = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -2143,7 +2148,7 @@ def dynamics( np.log1p(sl), ) - alpha, beta, gamma, uu0, su0 = valid_adata.var.loc[ + alpha, beta, gamma, uu0, su0 = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -2166,14 +2171,14 @@ def dynamics( l = sol_s(t, 0, 0, alpha, beta, gamma) if true_param_prefix is not None: true_alpha, true_beta, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "beta"] - if true_param_prefix + "beta" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "beta"] + if true_param_prefix + "beta" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) true_u = sol_u(t, uu0, 0, true_beta) @@ -2210,7 +2215,7 @@ def dynamics( if log_unnormalized and layers == ["new", "total"]: uu, ul = np.log1p(uu), np.log1p(ul) - alpha, gamma, uu0 = valid_adata.var.loc[ + alpha, gamma, uu0 = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -2226,11 +2231,11 @@ def dynamics( l = None # sol_s(t, 0, 0, alpha, 1, gamma) if true_param_prefix is not None: true_alpha, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) true_u = sol_u(t, uu0, 0, true_gamma) @@ -2350,7 +2355,7 @@ def dynamics( np.log1p(sl), ) - alpha, beta, gamma, U0, S0 = valid_adata.var.loc[ + alpha, beta, gamma, U0, S0 = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -2373,14 +2378,14 @@ def dynamics( L = sl + ul if true_param_prefix is not None: true_alpha, true_beta, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "beta"] - if true_param_prefix + "beta" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "beta"] + if true_param_prefix + "beta" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) true_l = sol_u(t, 0, true_alpha, true_beta) + sol_s( @@ -2403,7 +2408,7 @@ def dynamics( if log_unnormalized and layers == ["new", "total"]: uu, ul = np.log1p(uu), np.log1p(ul) - alpha, gamma, total0 = valid_adata.var.loc[ + alpha, gamma, total0 = vel_params_df.loc[ gene_name, [ prefix + "alpha", @@ -2420,11 +2425,11 @@ def dynamics( L = ul if true_param_prefix is not None: true_alpha, true_gamma = ( - valid_adata.var.loc[gene_name, true_param_prefix + "alpha"] - if true_param_prefix + "alpha" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "alpha"] + if true_param_prefix + "alpha" in vel_params_df.columns else -np.inf, - valid_adata.var.loc[gene_name, true_param_prefix + "gamma"] - if true_param_prefix + "gamma" in valid_adata.var_keys() + vel_params_df.loc[gene_name, true_param_prefix + "gamma"] + if true_param_prefix + "gamma" in vel_params_df.columns else -np.inf, ) true_l = sol_u(t, 0, true_alpha, true_gamma) # sol_s(t, 0, 0, alpha, 1, gamma) @@ -2537,7 +2542,7 @@ def dynamics( np.log1p(sl), ) - beta, gamma, alpha_std = valid_adata.var.loc[ + beta, gamma, alpha_std = vel_params_df.loc[ gene_name, [ prefix + "beta", @@ -2607,7 +2612,7 @@ def dynamics( if log_unnormalized and layers == ["new", "total"]: uu, ul = np.log1p(uu), np.log1p(ul) - gamma, alpha_std = valid_adata.var.loc[gene_name, [prefix + "gamma", prefix + "alpha_std"]] + gamma, alpha_std = vel_params_df.loc[gene_name, [prefix + "gamma", prefix + "alpha_std"]] alpha_stm = valid_adata[:, gene_name].varm[prefix + "alpha"].flatten()[1:] alpha_stm0, k, _ = solve_first_order_deg(T_uniq[1:], alpha_stm) diff --git a/dynamo/plot/scatters.py b/dynamo/plot/scatters.py index 85f0d57d8..5e0cf339a 100755 --- a/dynamo/plot/scatters.py +++ b/dynamo/plot/scatters.py @@ -23,7 +23,7 @@ from ..dynamo_logger import main_debug, main_info, main_warning from ..preprocessing.utils import affine_transform, gen_rotation_2d from ..tools.moments import calc_1nd_moment -from ..tools.utils import flatten, get_mapper, update_dict +from ..tools.utils import flatten, get_mapper, get_vel_params, update_dict, update_vel_params from .utils import ( _datashade_points, _get_adata_color_vec, @@ -808,9 +808,11 @@ def _plot_basis_layer(cur_b, cur_l): points.iloc[:, 0].max() * 0.80, ) k_name = "gamma_k" if _adata.uns["dynamics"]["experiment_type"] == "one-shot" else "gamma" - if k_name in _adata.var.columns: - if not ("gamma_b" in _adata.var.columns) or all(_adata.var.gamma_b.isna()): - _adata.var.loc[:, "gamma_b"] = 0 + vel_params_df = get_vel_params(_adata) + if k_name in vel_params_df.columns: + if not ("gamma_b" in vel_params_df.columns) or all(vel_params_df.gamma_b.isna()): + vel_params_df.loc[:, "gamma_b"] = 0 + update_vel_params(_adata, params_df=vel_params_df) ax.plot( xnew, xnew * _adata[:, cur_b].var.loc[:, k_name].unique() @@ -841,9 +843,11 @@ def _plot_basis_layer(cur_b, cur_l): + group_adata[:, cur_b].var.loc[:, group_b_key].unique() ) ax.annotate(group + "_" + cur_group, xy=(group_xnew[-1], group_ynew[-1])) - if group_k_name in group_adata.var.columns: - if not (group_b_key in group_adata.var.columns) or all(group_adata.var[group_b_key].isna()): - group_adata.var.loc[:, group_b_key] = 0 + vel_params_df = get_vel_params(group_adata) + if group_k_name in vel_params_df.columns: + if not (group_b_key in vel_params_df.columns) or all(vel_params_df[group_b_key].isna()): + vel_params_df.loc[:, group_b_key] = 0 + update_vel_params(group_adata, params_df=vel_params_df) main_info("No %s found, setting all bias terms to zero" % group_b_key) ax.plot( group_xnew, diff --git a/dynamo/plot/utils_dynamics.py b/dynamo/plot/utils_dynamics.py index 03a52746c..5040d6fdf 100644 --- a/dynamo/plot/utils_dynamics.py +++ b/dynamo/plot/utils_dynamics.py @@ -9,7 +9,7 @@ prepare_data_mix_no_splicing, prepare_data_no_splicing, ) -from ..tools.utils import get_mapper +from ..tools.utils import get_mapper, get_vel_params from .utils import _to_hex @@ -1685,8 +1685,9 @@ def plot_kin_twostep( colors = pd.Series(T).map(new_color_key).values - r2 = adata[:, genes].var["gamma_r2"] - mean_R2 = adata[:, genes].var["mean_R2"] + vel_params_df = get_vel_params(adata) + r2 = vel_params_df.loc[genes, "gamma_r2"] + mean_R2 = vel_params_df.loc[genes, "mean_R2"] for i, gene_name in enumerate(genes): cur_X_data, cur_X_fit_data, cur_logLL = ( diff --git a/dynamo/prediction/tscRNA_seq.py b/dynamo/prediction/tscRNA_seq.py index fb5797792..aa611070d 100644 --- a/dynamo/prediction/tscRNA_seq.py +++ b/dynamo/prediction/tscRNA_seq.py @@ -4,6 +4,7 @@ from scipy.sparse import csr_matrix from ..dynamo_logger import LoggerManager, main_exception, main_warning +from ..tools.utils import get_vel_params from ..utils import copy_adata from .utils import init_r0_pulse @@ -67,7 +68,7 @@ def get_pulse_r0( % (tkey, nkey, gamma_k_key) ) R, L = adata[:, gene_names].layers[tkey], adata[:, gene_names].layers[nkey] - K = adata[:, gene_names].var[gamma_k_key].values.astype(float) + K = get_vel_params(adata, params=gamma_k_key, genes=gene_names).astype(float) logger.info("Calculate initial total RNA via r0 = (r - l) / (1 - k)") res = init_r0_pulse(R, L, K[None, :]) diff --git a/dynamo/tools/__init__.py b/dynamo/tools/__init__.py index 32377ed66..13ddb1ae5 100755 --- a/dynamo/tools/__init__.py +++ b/dynamo/tools/__init__.py @@ -93,6 +93,7 @@ AnnDataPredicate, cell_norm, compute_smallest_distance, + get_vel_params, index_gene, select, select_cell, diff --git a/dynamo/tools/dynamics.py b/dynamo/tools/dynamics.py index c5b51c7cf..0471c2d53 100755 --- a/dynamo/tools/dynamics.py +++ b/dynamo/tools/dynamics.py @@ -41,6 +41,7 @@ get_data_for_kin_params_estimation, get_U_S_for_velocity_estimation, get_valid_bools, + get_vel_params, one_shot_alpha_matrix, remove_2nd_moments, set_param_kinetic, @@ -479,8 +480,12 @@ def dynamics( est_method = "gmm" if model.lower() == "stochastic" else "ols" if experiment_type.lower() == "one-shot": - beta = subset_adata.var.beta if "beta" in subset_adata.var.keys() else None - gamma = subset_adata.var.gamma if "gamma" in subset_adata.var.keys() else None + try: + vel_params_df = get_vel_params(subset_adata) + beta = vel_params_df.beta if "beta" in vel_params_df.columns else None + gamma = vel_params_df.gamma if "gamma" in vel_params_df.columns else None + except KeyError: + beta, gamma = None, None ss_estimation_kwargs = {"beta": beta, "gamma": gamma} else: ss_estimation_kwargs = {} diff --git a/dynamo/tools/pseudotime_velocity.py b/dynamo/tools/pseudotime_velocity.py index 14e15a706..660646f90 100644 --- a/dynamo/tools/pseudotime_velocity.py +++ b/dynamo/tools/pseudotime_velocity.py @@ -235,5 +235,5 @@ def pseudotime_velocity( logger.info("set gamma to be 0 in .var. so that velocity_S = unspliced RNA.") logger.info_insert_adata("gamma", "var", indent_level=2) - adata.var["gamma"] = 0 - adata.var["gamma_b"] = 0 + adata.varm["pseudotime_vel_params"] = np.zeros((adata.n_vars, 2)) + adata.uns["pseudotime_vel_params_names"] = ["gamma", "gamma_b"] diff --git a/dynamo/tools/recipes.py b/dynamo/tools/recipes.py index 15db37874..8a0ce8b2a 100644 --- a/dynamo/tools/recipes.py +++ b/dynamo/tools/recipes.py @@ -15,7 +15,7 @@ from .dimension_reduction import reduceDimension from .dynamics import dynamics from .moments import moments -from .utils import set_transition_genes +from .utils import get_vel_params, set_transition_genes, update_vel_params # add recipe_csc_data() @@ -325,9 +325,10 @@ def recipe_deg_data( set_transition_genes(adata) cell_velocities(adata, enforce=True, vkey=vkey, ekey=ekey, basis=basis) except BaseException: + vel_params_df = get_vel_params(adata) cell_velocities( adata, - min_r2=adata.var.gamma_r2.min(), + min_r2=vel_params_df.gamma_r2.min(), enforce=True, vkey=vkey, ekey=ekey, @@ -708,6 +709,7 @@ def velocity_N( var_columns = adata.var.columns layer_keys = adata.layers.keys() + vel_params_df = get_vel_params(adata) # check velocity_N, velocity_T, X_new, X_total if not np.all([i in layer_keys for i in ["X_new", "X_total"]]): @@ -743,8 +745,8 @@ def velocity_N( "beta_k", "gamma_k", ]: - if i in var_columns: - del adata.var[i] + if i in vel_params_df.columns: + del vel_params_df[i] if group is not None: group_prefixes = [group + "_" + str(i) + "_" for i in adata.obs[group].unique()] @@ -773,8 +775,9 @@ def velocity_N( "beta_k", "gamma_k", ]: - if i + j in var_columns: - del adata.var[i + j] + if i + j in vel_params_df.columns: + del vel_params_df[i + j] + update_vel_params(adata, params_df=vel_params_df) # now let us first run pca with new RNA if recalculate_pca: diff --git a/dynamo/tools/utils.py b/dynamo/tools/utils.py index 6c84d8247..be42ace3c 100755 --- a/dynamo/tools/utils.py +++ b/dynamo/tools/utils.py @@ -401,6 +401,8 @@ def reserve_minimal_genes_by_gamma_r2(adata: AnnData, var_store_key: str, minima The minimal gene data. """ + vel_params_df = get_vel_params(adata) + # already satisfy the requirement if var_store_key in adata.var.columns and adata.var[var_store_key].sum() >= minimal_gene_num: return adata.var[var_store_key] @@ -408,7 +410,7 @@ def reserve_minimal_genes_by_gamma_r2(adata: AnnData, var_store_key: str, minima if var_store_key not in adata.var.columns: raise ValueError("adata.var.%s does not exists." % (var_store_key)) - gamma_r2_not_na = np.array(adata.var.gamma_r2[adata.var.gamma_r2.notna()]) + gamma_r2_not_na = np.array(vel_params_df.gamma_r2[vel_params_df.gamma_r2.notna()]) if len(gamma_r2_not_na) < minimal_gene_num: raise ValueError("adata.var.%s does not have enough values that are not NA." % (var_store_key)) @@ -1392,30 +1394,31 @@ def set_param_ss( valid_ind, ind_for_proteins, ): + params_df = pd.DataFrame(index=adata.var.index) if experiment_type == "mix_std_stm": if alpha is not None: if cur_grp == _group[0]: adata.varm[kin_param_pre + "alpha"] = np.zeros((adata.shape[1], alpha[1].shape[1])) adata.varm[kin_param_pre + "alpha"][valid_ind, :] = alpha[1] ( - adata.var[kin_param_pre + "alpha"], - adata.var[kin_param_pre + "alpha_std"], + params_df[kin_param_pre + "alpha"], + params_df[kin_param_pre + "alpha_std"], ) = (None, None) ( - adata.var.loc[valid_ind, kin_param_pre + "alpha"], - adata.var.loc[valid_ind, kin_param_pre + "alpha_std"], + params_df.loc[valid_ind, kin_param_pre + "alpha"], + params_df.loc[valid_ind, kin_param_pre + "alpha_std"], ) = (alpha[1][:, -1], alpha[0]) if cur_grp == _group[0]: ( - adata.var[kin_param_pre + "beta"], - adata.var[kin_param_pre + "gamma"], - adata.var[kin_param_pre + "half_life"], + params_df[kin_param_pre + "beta"], + params_df[kin_param_pre + "gamma"], + params_df[kin_param_pre + "half_life"], ) = (None, None, None) - adata.var.loc[valid_ind, kin_param_pre + "beta"] = beta - adata.var.loc[valid_ind, kin_param_pre + "gamma"] = gamma - adata.var.loc[valid_ind, kin_param_pre + "half_life"] = np.log(2) / gamma + params_df.loc[valid_ind, kin_param_pre + "beta"] = beta + params_df.loc[valid_ind, kin_param_pre + "gamma"] = gamma + params_df.loc[valid_ind, kin_param_pre + "half_life"] = np.log(2) / gamma else: if alpha is not None: if len(alpha.shape) > 1: # for each cell @@ -1426,21 +1429,21 @@ def set_param_ss( else np.zeros(adata.shape[::-1]) ) # adata.varm[kin_param_pre + "alpha"][valid_ind, :] = alpha # - adata.var.loc[valid_ind, kin_param_pre + "alpha"] = alpha.mean(1) + params_df.loc[valid_ind, kin_param_pre + "alpha"] = alpha.mean(1) elif len(alpha.shape) == 1: if cur_grp == _group[0]: - adata.var[kin_param_pre + "alpha"] = None - adata.var.loc[valid_ind, kin_param_pre + "alpha"] = alpha + params_df[kin_param_pre + "alpha"] = None + params_df.loc[valid_ind, kin_param_pre + "alpha"] = alpha if cur_grp == _group[0]: ( - adata.var[kin_param_pre + "beta"], - adata.var[kin_param_pre + "gamma"], - adata.var[kin_param_pre + "half_life"], + params_df[kin_param_pre + "beta"], + params_df[kin_param_pre + "gamma"], + params_df[kin_param_pre + "half_life"], ) = (None, None, None) - adata.var.loc[valid_ind, kin_param_pre + "beta"] = beta - adata.var.loc[valid_ind, kin_param_pre + "gamma"] = gamma - adata.var.loc[valid_ind, kin_param_pre + "half_life"] = None if gamma is None else np.log(2) / gamma + params_df.loc[valid_ind, kin_param_pre + "beta"] = beta + params_df.loc[valid_ind, kin_param_pre + "gamma"] = gamma + params_df.loc[valid_ind, kin_param_pre + "half_life"] = None if gamma is None else np.log(2) / gamma ( alpha_intercept, @@ -1466,22 +1469,22 @@ def set_param_ss( alpha_r2[~np.isfinite(alpha_r2)] = 0 if cur_grp == _group[0]: ( - adata.var[kin_param_pre + "alpha_b"], - adata.var[kin_param_pre + "alpha_r2"], - adata.var[kin_param_pre + "gamma_b"], - adata.var[kin_param_pre + "gamma_r2"], - adata.var[kin_param_pre + "gamma_logLL"], - adata.var[kin_param_pre + "delta_b"], - adata.var[kin_param_pre + "delta_r2"], - adata.var[kin_param_pre + "bs"], - adata.var[kin_param_pre + "bf"], - adata.var[kin_param_pre + "uu0"], - adata.var[kin_param_pre + "ul0"], - adata.var[kin_param_pre + "su0"], - adata.var[kin_param_pre + "sl0"], - adata.var[kin_param_pre + "U0"], - adata.var[kin_param_pre + "S0"], - adata.var[kin_param_pre + "total0"], + params_df[kin_param_pre + "alpha_b"], + params_df[kin_param_pre + "alpha_r2"], + params_df[kin_param_pre + "gamma_b"], + params_df[kin_param_pre + "gamma_r2"], + params_df[kin_param_pre + "gamma_logLL"], + params_df[kin_param_pre + "delta_b"], + params_df[kin_param_pre + "delta_r2"], + params_df[kin_param_pre + "bs"], + params_df[kin_param_pre + "bf"], + params_df[kin_param_pre + "uu0"], + params_df[kin_param_pre + "ul0"], + params_df[kin_param_pre + "su0"], + params_df[kin_param_pre + "sl0"], + params_df[kin_param_pre + "U0"], + params_df[kin_param_pre + "S0"], + params_df[kin_param_pre + "total0"], ) = ( None, None, @@ -1501,47 +1504,50 @@ def set_param_ss( None, ) - adata.var.loc[valid_ind, kin_param_pre + "alpha_b"] = alpha_intercept - adata.var.loc[valid_ind, kin_param_pre + "alpha_r2"] = alpha_r2 + params_df.loc[valid_ind, kin_param_pre + "alpha_b"] = alpha_intercept + params_df.loc[valid_ind, kin_param_pre + "alpha_r2"] = alpha_r2 if gamma_r2 is not None: gamma_r2[~np.isfinite(gamma_r2)] = 0 - adata.var.loc[valid_ind, kin_param_pre + "gamma_b"] = gamma_intercept - adata.var.loc[valid_ind, kin_param_pre + "gamma_r2"] = gamma_r2 - adata.var.loc[valid_ind, kin_param_pre + "gamma_logLL"] = gamma_logLL + params_df.loc[valid_ind, kin_param_pre + "gamma_b"] = gamma_intercept + params_df.loc[valid_ind, kin_param_pre + "gamma_r2"] = gamma_r2 + params_df.loc[valid_ind, kin_param_pre + "gamma_logLL"] = gamma_logLL - adata.var.loc[valid_ind, kin_param_pre + "bs"] = bs - adata.var.loc[valid_ind, kin_param_pre + "bf"] = bf + params_df.loc[valid_ind, kin_param_pre + "bs"] = bs + params_df.loc[valid_ind, kin_param_pre + "bf"] = bf - adata.var.loc[valid_ind, kin_param_pre + "uu0"] = uu0 - adata.var.loc[valid_ind, kin_param_pre + "ul0"] = ul0 - adata.var.loc[valid_ind, kin_param_pre + "su0"] = su0 - adata.var.loc[valid_ind, kin_param_pre + "sl0"] = sl0 - adata.var.loc[valid_ind, kin_param_pre + "U0"] = U0 - adata.var.loc[valid_ind, kin_param_pre + "S0"] = S0 - adata.var.loc[valid_ind, kin_param_pre + "total0"] = total0 + params_df.loc[valid_ind, kin_param_pre + "uu0"] = uu0 + params_df.loc[valid_ind, kin_param_pre + "ul0"] = ul0 + params_df.loc[valid_ind, kin_param_pre + "su0"] = su0 + params_df.loc[valid_ind, kin_param_pre + "sl0"] = sl0 + params_df.loc[valid_ind, kin_param_pre + "U0"] = U0 + params_df.loc[valid_ind, kin_param_pre + "S0"] = S0 + params_df.loc[valid_ind, kin_param_pre + "total0"] = total0 if experiment_type == "one-shot": - adata.var[kin_param_pre + "beta_k"] = None - adata.var[kin_param_pre + "gamma_k"] = None - adata.var.loc[valid_ind, kin_param_pre + "beta_k"] = beta_k - adata.var.loc[valid_ind, kin_param_pre + "gamma_k"] = gamma_k + params_df[kin_param_pre + "beta_k"] = None + params_df[kin_param_pre + "gamma_k"] = None + params_df.loc[valid_ind, kin_param_pre + "beta_k"] = beta_k + params_df.loc[valid_ind, kin_param_pre + "gamma_k"] = gamma_k if ind_for_proteins is not None: delta_r2[~np.isfinite(delta_r2)] = 0 if cur_grp == _group[0]: ( - adata.var[kin_param_pre + "eta"], - adata.var[kin_param_pre + "delta"], - adata.var[kin_param_pre + "delta_b"], - adata.var[kin_param_pre + "delta_r2"], - adata.var[kin_param_pre + "p_half_life"], + params_df[kin_param_pre + "eta"], + params_df[kin_param_pre + "delta"], + params_df[kin_param_pre + "delta_b"], + params_df[kin_param_pre + "delta_r2"], + params_df[kin_param_pre + "p_half_life"], ) = (None, None, None, None, None) - adata.var.loc[valid_ind, kin_param_pre + "eta"][ind_for_proteins] = eta - adata.var.loc[valid_ind, kin_param_pre + "delta"][ind_for_proteins] = delta - adata.var.loc[valid_ind, kin_param_pre + "delta_b"][ind_for_proteins] = delta_intercept - adata.var.loc[valid_ind, kin_param_pre + "delta_r2"][ind_for_proteins] = delta_r2 - adata.var.loc[valid_ind, kin_param_pre + "p_half_life"][ind_for_proteins] = np.log(2) / delta + params_df.loc[valid_ind, kin_param_pre + "eta"][ind_for_proteins] = eta + params_df.loc[valid_ind, kin_param_pre + "delta"][ind_for_proteins] = delta + params_df.loc[valid_ind, kin_param_pre + "delta_b"][ind_for_proteins] = delta_intercept + params_df.loc[valid_ind, kin_param_pre + "delta_r2"][ind_for_proteins] = delta_r2 + params_df.loc[valid_ind, kin_param_pre + "p_half_life"][ind_for_proteins] = np.log(2) / delta + + adata.varm[kin_param_pre + "vel_params"] = params_df.to_numpy() + adata.uns[kin_param_pre + "vel_params_names"] = list(params_df.columns) return adata @@ -1564,23 +1570,24 @@ def set_param_kinetic( cur_cells_bools, valid_ind, ): + params_df = pd.DataFrame(index=adata.var.index) if cur_grp == _group[0]: ( - adata.var[kin_param_pre + "alpha"], - adata.var[kin_param_pre + "a"], - adata.var[kin_param_pre + "b"], - adata.var[kin_param_pre + "alpha_a"], - adata.var[kin_param_pre + "alpha_i"], - adata.var[kin_param_pre + "beta"], - adata.var[kin_param_pre + "p_half_life"], - adata.var[kin_param_pre + "gamma"], - adata.var[kin_param_pre + "half_life"], - adata.var[kin_param_pre + "cost"], - adata.var[kin_param_pre + "logLL"], + params_df[kin_param_pre + "alpha"], + params_df[kin_param_pre + "a"], + params_df[kin_param_pre + "b"], + params_df[kin_param_pre + "alpha_a"], + params_df[kin_param_pre + "alpha_i"], + params_df[kin_param_pre + "beta"], + params_df[kin_param_pre + "p_half_life"], + params_df[kin_param_pre + "gamma"], + params_df[kin_param_pre + "half_life"], + params_df[kin_param_pre + "cost"], + params_df[kin_param_pre + "logLL"], ) = (None, None, None, None, None, None, None, None, None, None, None) if isarray(alpha) and alpha.ndim > 1: - adata.var.loc[valid_ind, kin_param_pre + "alpha"] = alpha.mean(1) + params_df.loc[valid_ind, kin_param_pre + "alpha"] = alpha.mean(1) cur_cells_ind, valid_ind_ = ( np.where(cur_cells_bools)[0][:, np.newaxis], np.where(valid_ind)[0], @@ -1590,25 +1597,96 @@ def set_param_kinetic( alpha = alpha.T.tocsr() if sp.issparse(alpha) else sp.csr_matrix(alpha, dtype=np.float64).T adata.layers["cell_wise_alpha"][cur_cells_ind, valid_ind_] = alpha else: - adata.var.loc[valid_ind, kin_param_pre + "alpha"] = alpha - adata.var.loc[valid_ind, kin_param_pre + "a"] = a - adata.var.loc[valid_ind, kin_param_pre + "b"] = b - adata.var.loc[valid_ind, kin_param_pre + "alpha_a"] = alpha_a - adata.var.loc[valid_ind, kin_param_pre + "alpha_i"] = alpha_i - adata.var.loc[valid_ind, kin_param_pre + "beta"] = beta - adata.var.loc[valid_ind, kin_param_pre + "gamma"] = gamma - adata.var.loc[valid_ind, kin_param_pre + "half_life"] = np.log(2) / gamma - adata.var.loc[valid_ind, kin_param_pre + "cost"] = cost - adata.var.loc[valid_ind, kin_param_pre + "logLL"] = logLL + params_df.loc[valid_ind, kin_param_pre + "alpha"] = alpha + params_df.loc[valid_ind, kin_param_pre + "a"] = a + params_df.loc[valid_ind, kin_param_pre + "b"] = b + params_df.loc[valid_ind, kin_param_pre + "alpha_a"] = alpha_a + params_df.loc[valid_ind, kin_param_pre + "alpha_i"] = alpha_i + params_df.loc[valid_ind, kin_param_pre + "beta"] = beta + params_df.loc[valid_ind, kin_param_pre + "gamma"] = gamma + params_df.loc[valid_ind, kin_param_pre + "half_life"] = np.log(2) / gamma + params_df.loc[valid_ind, kin_param_pre + "cost"] = cost + params_df.loc[valid_ind, kin_param_pre + "logLL"] = logLL # add extra parameters (u0, uu0, etc.) extra_params.columns = [kin_param_pre + i for i in extra_params.columns] extra_params = extra_params.set_index(adata.var.index[valid_ind]) - var = pd.concat((adata.var, extra_params), axis=1, sort=False) - adata.var = var + var = pd.concat((params_df, extra_params), axis=1, sort=False) + adata.varm[kin_param_pre + "vel_params"] = var.to_numpy() + adata.uns[kin_param_pre + "vel_params_names"] = list(var.columns) return adata +def get_vel_params( + adata: AnnData, + params: Optional[Union[List, str]] = None, + genes: Optional[List] = None, + kin_param_pre: str = "", + skip_cell_wise: bool = False, +) -> Union[Tuple, pd.DataFrame, List]: + """Get the velocity parameters based on input names. + + Args: + adata: the anndata object which contains the parameters. + params: the names of parameters to query. If set to None, the entire velocity parameters DataFrame from `.varm` + will be returned. + kin_param_pre: the prefix used to kinetic parameters related to RNA dynamics. + skip_cell_wise: whether to skip the detected cell wise parameters. If set to True, the mean will be returned + instead of cell wise parameters. + + Returns: + All velocity parameters with the same order of query `params`. + """ + if type(params) is str: + params = [params] + + if kin_param_pre + "vel_params" not in adata.varm.keys(): + raise KeyError("The key of velocity related parameters are not found in varm.") + + array_data = adata.varm[kin_param_pre + "vel_params"] + df_columns = adata.uns[kin_param_pre + "vel_params_names"] + df = pd.DataFrame(array_data, index=adata.var_names, columns=df_columns) + target_params = [] + + if genes is None: + genes = df.index + + if params is None: + return df.loc[genes] + + for param in params: + if param == "alpha": + if not skip_cell_wise: + if "cell_wise_alpha" in adata.layers.keys(): + target_params.append(adata[:, genes].layers["cell_wise_alpha"]) + elif "alpha" in adata.varm.keys(): + target_params.append(adata[:, genes].varm[kin_param_pre + "alpha"]) + else: + target_params.append(df.loc[genes, kin_param_pre + "alpha"].values) + continue + target_params.append(df.loc[genes, kin_param_pre + param].values) + + if len(target_params) > 1: + return tuple(target_params) + else: + return target_params[0] + + +def update_vel_params(adata: AnnData, params_df: pd.DataFrame, kin_param_pre: str = "") -> None: + """Update the kinetic parameters related to RNA velocity calculation. + + Args: + adata: the AnnData object whose kinetic parameters related to RNA velocity calculation will be updated. + params_df: the dataframe of kinetic parameters related to RNA velocity calculation. + kin_param_pre: the prefix used to kinetic parameters related to RNA dynamics. + + Returns: + The anndata object will be updated with parameters and columns names from given dataframe. + """ + adata.varm[kin_param_pre + "vel_params"] = params_df.to_numpy() + adata.uns[kin_param_pre + "vel_params_names"] = list(params_df.columns) + + def get_U_S_for_velocity_estimation(subset_adata, use_moments, has_splicing, has_labeling, log_unnormalized, NTR): mapper = get_mapper() @@ -2186,6 +2264,7 @@ def set_transition_genes( minimal_gene_num=50, ): layer = vkey.split("_")[1] + vel_params_df = get_vel_params(adata) if adata.uns["dynamics"]["est_method"] == "twostep" and adata.uns["dynamics"]["experiment_type"] == "kin": # if adata.uns['dynamics']['has_splicing']: @@ -2197,12 +2276,12 @@ def set_transition_genes( "mix_pulse_chase", "kin", ]: - logLL_col = adata.var.columns[adata.var.columns.str.endswith("logLL")] + logLL_col = vel_params_df.columns[vel_params_df.columns.str.endswith("logLL")] if len(logLL_col) > 1: main_warning(f"there are two columns ends with logLL: {logLL_col}") - adata.var[store_key] = adata.var[logLL_col[-1]].astype(float) < np.nanpercentile( - adata.var[logLL_col[-1]].astype(float), 10 + adata.var[store_key] = vel_params_df[logLL_col[-1]].astype(float) < np.nanpercentile( + vel_params_df[logLL_col[-1]].astype(float), 10 ) if layer in ["N", "T"]: return adata @@ -2220,94 +2299,94 @@ def set_transition_genes( # the following parameters aggreation for different groups can be improved later if layer == "U": - if "alpha" not in adata.var.columns: + if "alpha" not in vel_params_df.columns: is_group_alpha, is_group_alpha_r2 = ( get_group_params_indices(adata, "alpha"), get_group_params_indices(adata, "alpha_r2"), ) if is_group_alpha.sum() > 0: - adata.var["alpha"] = adata.var.loc[:, is_group_alpha].mean(1, skipna=True) - adata.var["alpha_r2"] = adata.var.loc[:, np.hstack((is_group_alpha_r2, False))].mean(1, skipna=True) + vel_params_df["alpha"] = vel_params_df.loc[:, is_group_alpha].mean(1, skipna=True) + vel_params_df["alpha_r2"] = vel_params_df.loc[:, np.hstack((is_group_alpha_r2, False))].mean(1, skipna=True) else: raise Exception("there is no alpha/alpha_r2 parameter estimated for your adata object") - if "alpha_r2" not in adata.var.columns: - adata.var["alpha_r2"] = None - if np.all(adata.var.alpha_r2.values is None): - adata.var.alpha_r2 = 1 + if "alpha_r2" not in vel_params_df.columns: + vel_params_df["alpha_r2"] = None + if np.all(vel_params_df.alpha_r2.values is None): + vel_params_df.alpha_r2 = 1 adata.var[store_key] = ( - (adata.var.alpha > min_alpha) & (adata.var.alpha_r2 > min_r2) & adata.var.use_for_dynamics + (vel_params_df.alpha > min_alpha) & (vel_params_df.alpha_r2 > min_r2) & adata.var.use_for_dynamics if use_for_dynamics - else (adata.var.alpha > min_alpha) & (adata.var.alpha_r2 > min_r2) + else (vel_params_df.alpha > min_alpha) & (vel_params_df.alpha_r2 > min_r2) ) elif layer == "S": - if "gamma" not in adata.var.columns: + if "gamma" not in vel_params_df.columns: is_group_gamma, is_group_gamma_r2 = ( get_group_params_indices(adata, "gamma"), get_group_params_indices(adata, "gamma_r2"), ) if is_group_gamma.sum() > 0: - adata.var["gamma"] = adata.var.loc[:, is_group_gamma].mean(1, skipna=True) - adata.var["gamma_r2"] = adata.var.loc[:, np.hstack((is_group_gamma_r2, False))].mean(1, skipna=True) + vel_params_df["gamma"] = vel_params_df.loc[:, is_group_gamma].mean(1, skipna=True) + vel_params_df["gamma_r2"] = vel_params_df.loc[:, np.hstack((is_group_gamma_r2, False))].mean(1, skipna=True) else: raise Exception("there is no gamma/gamma_r2 parameter estimated for your adata object") - if "gamma_r2" not in adata.var.columns: + if "gamma_r2" not in vel_params_df.columns: main_debug("setting all gamma_r2 to 1") - adata.var["gamma_r2"] = 1 - if np.all(adata.var.gamma_r2.values is None) or np.all(adata.var.gamma_r2.values == ""): - main_debug("Since all adata.var.gamma_r2 values are None or '', setting all gamma_r2 values to 1.") - adata.var.gamma_r2 = 1 + vel_params_df["gamma_r2"] = 1 + if np.all(vel_params_df.gamma_r2.values is None) or np.all(vel_params_df.gamma_r2.values == ""): + main_debug("Since all gamma_r2 values are None or '', setting all gamma_r2 values to 1.") + vel_params_df.gamma_r2 = 1 - adata.var[store_key] = (adata.var.gamma > min_gamma) & (adata.var.gamma_r2 > min_r2) + adata.var[store_key] = (vel_params_df.gamma > min_gamma) & (vel_params_df.gamma_r2 > min_r2) if use_for_dynamics: adata.var[store_key] = adata.var[store_key] & adata.var.use_for_dynamics elif layer == "P": - if "delta" not in adata.var.columns: + if "delta" not in vel_params_df.columns: is_group_delta, is_group_delta_r2 = ( get_group_params_indices(adata, "delta"), get_group_params_indices(adata, "delta_r2"), ) if is_group_delta.sum() > 0: - adata.var["delta"] = adata.var.loc[:, is_group_delta].mean(1, skipna=True) - adata.var["delta_r2"] = adata.var.loc[:, np.hstack((is_group_delta_r2, False))].mean(1, skipna=True) + vel_params_df["delta"] = vel_params_df.loc[:, is_group_delta].mean(1, skipna=True) + vel_params_df["delta_r2"] = vel_params_df.loc[:, np.hstack((is_group_delta_r2, False))].mean(1, skipna=True) else: raise Exception("there is no delta/delta_r2 parameter estimated for your adata object") - if "delta_r2" not in adata.var.columns: - adata.var["delta_r2"] = None - if np.all(adata.var.delta_r2.values is None): - adata.var.delta_r2 = 1 + if "delta_r2" not in vel_params_df.columns: + vel_params_df["delta_r2"] = None + if np.all(vel_params_df.delta_r2.values is None): + vel_params_df.delta_r2 = 1 adata.var[store_key] = ( - (adata.var.delta > min_delta) & (adata.var.delta_r2 > min_r2) & adata.var.use_for_dynamics + (vel_params_df.delta > min_delta) & (vel_params_df.delta_r2 > min_r2) & adata.var.use_for_dynamics if use_for_dynamics - else (adata.var.delta > min_delta) & (adata.var.delta_r2 > min_r2) + else (vel_params_df.delta > min_delta) & (vel_params_df.delta_r2 > min_r2) ) if layer == "T": - if "gamma" not in adata.var.columns: + if "gamma" not in vel_params_df.columns: is_group_gamma, is_group_gamma_r2 = ( get_group_params_indices(adata, "gamma"), get_group_params_indices(adata, "gamma_r2"), ) if is_group_gamma.sum() > 0: - adata.var["gamma"] = adata.var.loc[:, is_group_gamma].mean(1, skipna=True) - adata.var["gamma_r2"] = adata.var.loc[:, np.hstack((is_group_gamma_r2, False))].mean(1, skipna=True) + vel_params_df["gamma"] = vel_params_df.loc[:, is_group_gamma].mean(1, skipna=True) + vel_params_df["gamma_r2"] = vel_params_df.loc[:, np.hstack((is_group_gamma_r2, False))].mean(1, skipna=True) else: raise Exception("there is no gamma/gamma_r2 parameter estimated for your adata object") - if "gamma_r2" not in adata.var.columns: - adata.var["gamma_r2"] = None - if np.all(adata.var.gamma_r2.values is None): - adata.var.gamma_r2 = 1 - if sum(adata.var.gamma_r2.isna()) == adata.n_vars: - gamm_r2_checker = adata.var.gamma_r2.isna() + if "gamma_r2" not in vel_params_df.columns: + vel_params_df["gamma_r2"] = None + if np.all(vel_params_df.gamma_r2.values is None): + vel_params_df.gamma_r2 = 1 + if sum(vel_params_df.gamma_r2.isna()) == adata.n_vars: + gamm_r2_checker = vel_params_df.gamma_r2.isna() else: - gamm_r2_checker = adata.var.gamma_r2 > min_r2 + gamm_r2_checker = vel_params_df.gamma_r2 > min_r2 adata.var[store_key] = ( - (adata.var.gamma > min_gamma) & gamm_r2_checker & adata.var.use_for_dynamics + (vel_params_df.gamma > min_gamma) & gamm_r2_checker & adata.var.use_for_dynamics if use_for_dynamics - else (adata.var.gamma > min_gamma) & gamm_r2_checker + else (vel_params_df.gamma > min_gamma) & gamm_r2_checker ) if adata.var[store_key].sum() < 5 and adata.n_vars > 5: @@ -2321,6 +2400,8 @@ def set_transition_genes( ) reserve_minimal_genes_by_gamma_r2(adata, store_key, minimal_gene_num=minimal_gene_num) + update_vel_params(adata, vel_params_df) + return adata diff --git a/dynamo/tools/velocyto_scvelo.py b/dynamo/tools/velocyto_scvelo.py index c6c1ffa64..018ef336a 100755 --- a/dynamo/tools/velocyto_scvelo.py +++ b/dynamo/tools/velocyto_scvelo.py @@ -37,17 +37,23 @@ def vlm_to_adata( # set obs, var obs, var = pd.DataFrame(vlm.ca), pd.DataFrame(vlm.ra) + vel_params = [] + vel_params_names = [] + varm = {} if "CellID" in obs.keys(): obs["obs_names"] = obs.pop("CellID") if "Gene" in var.keys(): var["var_names"] = var.pop("Gene") if hasattr(vlm, "q"): - var["gamma_b"] = vlm.q + vel_params_names.append("gamma_b") + vel_params.append(vlm.q) if hasattr(vlm, "gammas"): - var["gamma"] = vlm.gammas + vel_params_names.append("gamma") + vel_params.append(vlm.gammas) if hasattr(vlm, "R2"): - var["gamma_r2"] = vlm.R2 + vel_params_names.append("gamma_r2") + vel_params.append(vlm.R2) # rename clusters to louvain try: @@ -158,7 +164,11 @@ def vlm_to_adata( X = csr_matrix(vlm.S_sz.T) if hasattr(vlm, "S_sz") else csr_matrix(vlm.S.T) # create an anndata object with Dynamo characteristics. - dyn_adata = anndata.AnnData(X=X, obs=obs, obsp=obsp, obsm=obsm, var=var, layers=layers, uns=uns) + if len(vel_params_names) > 0: + uns["vel_params_names"] = vel_params_names + varm["vel_params"] = vel_params + + dyn_adata = anndata.AnnData(X=X, obs=obs, obsp=obsp, obsm=obsm, var=var, varm=varm, layers=layers, uns=uns) return dyn_adata