covariates: accept string columns and treat them as categoricals

Author: ilia-kats

src/mofaflex/_core/datasets/misc.py

@@ -1,4 +1,5 @@
 from collections.abc import Mapping, Sequence
+from functools import reduce
 
 import numpy as np
 import pandas as pd
@@ -54,19 +55,26 @@ def merge_covariates(covariates: Mapping[str, Mapping[str, pd.DataFrame]]):
     for group_covars in covariates.values():
         for view_covars in group_covars.values():
             dtypes = view_covars.dtypes
+            cats = None
             if dtypes.nunique() > 1:
                 raise ValueError("Mixed dtypes for a covariate are not supported.")
             if dtypes.iloc[0] == "category":
-                categories = (
-                    view_covars.iloc[0].cat.categories
-                    if categories is None
-                    else categories.union(view_covars.iloc[0].cat.categories)
+                cats = view_covars.iloc[0].cat.categories
+            elif pd.api.types.is_string_dtype(dtypes.iloc[0]):
+                cats = reduce(
+                    lambda x, y: x.union(y), (pd.Index(col.dropna().unique()) for _, col in view_covars.items())
                 )
+            if cats is not None:
+                categories = pd.Index(cats) if categories is None else categories.union(cats)
     for group_covars in covariates.values():
         for view_covars in group_covars.values():
-            if view_covars.dtypes.iloc[0] == "category":
+            dtypes = view_covars.dtypes
+            if dtypes.iloc[0] == "category":
                 for col in view_covars.columns:
                     view_covars[col] = view_covars[col].cat.set_categories(categories)
+            elif pd.api.types.is_string_dtype(dtypes.iloc[0]):
+                for col in view_covars.columns:
+                    view_covars[col] = pd.Categorical(view_covars[col], categories=categories)
 
     # ensure the covariate value is consistent across views (nanmean or first)
     merged_covariates = {}

tests/conftest.py

@@ -68,6 +68,7 @@ def _adata(likelihood, nobs, nvar, var_names=None, obs_names=None):
         adata.obs["gvar_normal"] = rng.random(size=(nobs))
         adata.obs["gvar_bernoulli"] = rng.binomial(1, 0.5, size=(nobs))
         adata.obs["gvar_categorical"] = pd.Categorical(rng.choice(["A", "B", "C"], size=(nobs)))
+        adata.obs["gvar_string"] = rng.choice(["A", "B", "C"], size=(nobs))
         adata.varm["annot_df"] = pd.DataFrame(
             rng.choice([False, True], size=(nvar, 10)), columns=[f"annot_{i}" for i in range(10)], index=adata.var_names
         )

tests/test_datasets_common.py

@@ -4,14 +4,16 @@
 from mofaflex._core.datasets import MofaFlexDataset, merge_covariates
 
 
-@pytest.mark.parametrize("axis", (0, 1))
-def test_merge_covariates_preserves_axis_order(axis, random_adata):
+@pytest.mark.parametrize(
+    ("axis", "key", "mkey"), ((0, None, "covar_array"), (1, None, "covar_array"), (0, "gvar_string", None))
+)
+def test_merge_covariates_preserves_axis_order(axis, key, mkey, random_adata):
     """merge_covariates must keep rows in the dataset's sample/feature order, not lexically sorted."""
 
     adata = random_adata("Normal", 500, 500)
     dataset = MofaFlexDataset(adata)
 
-    covars = dataset.get_covariates(axis, mkey="covar_array")
+    covars = dataset.get_covariates(axis, key=key, mkey=mkey)
     merged = next(iter(merge_covariates(covars).values()))
     canonical = next(iter(dataset.get_names(axis).values()))
 

tests/test_mofaflex_integration.py

@@ -100,7 +100,11 @@ def model_api_untrained_only():
     "argfor,argname,argval",
     [
         ("likelihood_normal", "scale_per_group", False),
-        ("term_mofaflex", "guiding_vars_obs_keys", ["gvar_normal", "gvar_bernoulli", "gvar_categorical"]),
+        (
+            "term_mofaflex",
+            "guiding_vars_obs_keys",
+            ["gvar_normal", "gvar_bernoulli", "gvar_categorical", "gvar_string"],
+        ),
         ("term_mofaflex", "weight_prior", "Normal"),
         ("term_mofaflex", "weight_prior", "Laplace"),
         ("term_mofaflex", "weight_prior", "Horseshoe"),
@@ -174,6 +178,7 @@ def test_integration(
             "gvar_normal": "Normal",
             "gvar_bernoulli": "Bernoulli",
             "gvar_categorical": "Categorical",
+            "gvar_string": "Categorical",
         },
         **termargs,
     )
@@ -211,7 +216,7 @@ def test_integration(
         assert "get_weight_annotations" in dir(model)
         assert all(isinstance(annot, pd.DataFrame) for annot in model.get_weight_annotations().values())
     elif argname == "guiding_vars_obs_keys":
-        assert model.n_guided_factors == model.terms["_"].n_guided_factors == 3
+        assert model.n_guided_factors == model.terms["_"].n_guided_factors == 4
     else:
         assert (
             model.n_factors