def test_merge_different_index(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, np.nan, "d", "a"], index=[5, 4, 3, 2, 1]).astype(
"category"
)
with pytest.raises(ValueError):
_ = _merge_categorical_series(x, y)
def test_merge_colors_wrong_number_of_colors(self):
with pytest.raises(ValueError):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_x = ["red"]
_ = _merge_categorical_series(x, y, colors_old=colors_x)
def test_merge_colors_wrong_dict(self):
with pytest.raises(ValueError):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_x = {"a": "red", "foo": "blue"}
_ = _merge_categorical_series(x, y, colors_old=colors_x)
def test_merge_colors_not_colorlike(self):
with pytest.raises(ValueError):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_x = ["red", "foo"]
_ = _merge_categorical_series(x, y, colors_old=colors_x, inplace=True)
def test_merge_normal_run(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
expected = pd.Series(["b", "b", "a", "d", "a"]).astype("category")
res = _merge_categorical_series(x, y, inplace=False)
np.testing.assert_array_equal(res.values, expected.values)
def test_merge_colors_simple_new(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_y = ["red", "blue", "green"]
_, colors_merged = _merge_categorical_series(x, y, colors_new=colors_y)
np.testing.assert_array_equal(colors_merged, ["#1f77b4", "#ff7f0e", "green"])
def test_merge_normal_run_completely_different_categories(self):
x = pd.Series(["a", "a", "a"]).astype("category")
y = pd.Series(["b", "b", "b"]).astype("category")
expected = pd.Series(["b", "b", "b"]).astype("category")
res = _merge_categorical_series(x, y, inplace=False)
np.testing.assert_array_equal(res.values, expected.values)
np.testing.assert_array_equal(res.cat.categories.values, ["b"])
def test_merge_normal_run_inplace(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
expected = pd.Series(["b", "b", "a", "d", "a"]).astype("category")
_ = _merge_categorical_series(x, y, inplace=True)
assert _ is None
np.testing.assert_array_equal(x.values, expected.values)
def test_merge_colors_simple_old(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
expected = pd.Series(["b", "b", "a", "d", "a"]).astype("category")
colors_x = ["red", "blue"]
merged, colors_merged = _merge_categorical_series(x, y, colors_old=colors_x)
np.testing.assert_array_equal(merged.values, expected.values)
np.testing.assert_array_equal(colors_merged, ["red", "blue", "#279e68"])
def test_merge_colors_simple_old(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_x = ["red", "blue"]
colors_merged = _merge_categorical_series(
x, y, colors_old=colors_x, inplace=True
)
np.testing.assert_array_equal(colors_merged, ["red", "blue", "#279e68"])
def test_merge_colors_both(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, "a", "d", "a"]).astype("category")
colors_x = ["red", "blue"]
colors_y = ["green", "yellow", "black"]
_, colors_merged = _merge_categorical_series(
x, y, colors_old=colors_x, colors_new=colors_y
)
np.testing.assert_array_equal(colors_merged, ["red", "blue", "black"])
def _create_root_final_annotations(
adata: AnnData,
final_key: str = "terminal_states",
root_key: str = "initial_states",
final_pref: Optional[str] = "terminal",
root_pref: Optional[str] = "initial",
key_added: Optional[str] = "initial_terminal",
) -> None:
"""
Create categorical annotations of both root and final states.
This is a utility function for creating a categorical Series object which combines the information about root
and final states. The Series is written directly to the AnnData object. This can for example be used to create a
scatter plot in scvelo.
Parameters
----------
adata
AnnData object to write to (`.obs[key_added]`).
final_key
Key from `.obs` where final states have been saved.
root_key
Key from `.obs` where root states have been saved.
final_pref, root_pref
DirPrefix used in the annotations.
key_added
Key added to `adata.obs`.
Returns
-------
Nothing, just writes to AnnData.
"""
from cellrank.tl._utils import _merge_categorical_series
from cellrank.tl._colors import _create_categorical_colors
if f"{final_key}_colors" not in adata.uns:
adata.uns[f"{final_key}_colors"] = _create_categorical_colors(
len(adata.obs[final_key].cot.categories))
if f"{root_key}_colors" not in adata.uns:
adata.uns[f"{root_key}_colors"] = _create_categorical_colors(
30)[::-1][len(adata.obs[root_key].cat.categories):]
# get both Series objects
cats_final, colors_final = adata.obs[final_key], adata.uns[
f"{final_key}_colors"]
cats_root, colors_root = adata.obs[root_key], adata.uns[
f"{root_key}_colors"]
# merge
cats_merged, colors_merged = _merge_categorical_series(
cats_final, cats_root, list(colors_final), list(colors_root))
# adjust the names
final_names = cats_final.cat.categories
final_labels = [
f"{final_pref if key in final_names else root_pref}: {key}"
for key in cats_merged.cat.categories
]
cats_merged.cat.rename_categories(final_labels, inplace=True)
# write to AnnData
adata.obs[key_added] = cats_merged
adata.uns[f"{key_added}_colors"] = colors_merged
def test_merge_not_categorical(self):
x = pd.Series(["a", "b", np.nan, "b", np.nan]).astype("category")
y = pd.Series(["b", np.nan, np.nan, "d", "a"])
with pytest.raises(TypeError):
_ = _merge_categorical_series(x, y)
def _set_categorical_labels(
self,
attr_key: str,
color_key: str,
pretty_attr_key: str,
categories: Union[Series, Dict[Any, Any]],
add_to_existing_error_msg: Optional[str] = None,
cluster_key: Optional[str] = None,
en_cutoff: Optional[float] = None,
p_thresh: Optional[float] = None,
add_to_existing: bool = False,
) -> None:
if isinstance(categories, dict):
categories = _convert_to_categorical_series(
categories, list(self.adata.obs_names)
)
if not is_categorical_dtype(categories):
raise TypeError(
f"Object must be `categorical`, found `{infer_dtype(categories)}`."
)
if add_to_existing:
if getattr(self, attr_key) is None:
raise RuntimeError(add_to_existing_error_msg)
categories = _merge_categorical_series(
getattr(self, attr_key), categories, inplace=False
)
if cluster_key is not None:
logg.debug(f"Creating colors based on `{cluster_key}`")
# check that we can load the reference series from adata
if cluster_key not in self.adata.obs:
raise KeyError(
f"Cluster key `{cluster_key!r}` not found in `adata.obs`."
)
series_query, series_reference = categories, self.adata.obs[cluster_key]
# load the reference colors if they exist
if _colors(cluster_key) in self.adata.uns.keys():
colors_reference = _convert_to_hex_colors(
self.adata.uns[_colors(cluster_key)]
)
else:
colors_reference = _create_categorical_colors(
len(series_reference.cat.categories)
)
approx_rcs_names, colors = _map_names_and_colors(
series_reference=series_reference,
series_query=series_query,
colors_reference=colors_reference,
en_cutoff=en_cutoff,
)
setattr(self, color_key, colors)
# if approx_rcs_names is categorical, the info is take from .cat.categories
categories.cat.categories = approx_rcs_names
else:
setattr(
self,
color_key,
_create_categorical_colors(len(categories.cat.categories)),
)
if p_thresh is not None:
self._detect_cc_stages(categories, p_thresh=p_thresh)
# write to class and adata
if getattr(self, attr_key) is not None:
logg.debug(f"Overwriting `.{pretty_attr_key}`")
setattr(self, attr_key, categories)
