fix issue in column name saving

theislab · Sep 8, 2023 · 4137a39 · 4137a39
1 parent 24d6c40
commit 4137a39
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ncem/tl/fit/backend/design_matrix.py b/ncem/tl/fit/backend/design_matrix.py
@@ -247,8 +247,12 @@ def get_dmat_from_obs(obs: pd.DataFrame, obs_niche: pd.DataFrame, formula: str,
     obs_niche.columns = [PREFIX_NEIGHBOR + x for x in obs_niche.columns]
     # Merge sample annotation:
     obs_full = pd.concat([obs, obs_index_type, obs_niche], axis=1)
+    columns_names = formula.replace('~0+','').split('+')
+
     dmat = patsy.dmatrix(formula, obs_full)
-    dmat = pd.DataFrame(np.asarray(dmat), index=obs.index, columns=dmat.design_info.column_names)
+    # Simplify names, this is necessary for patsy to accept these as terms later.
+    column_names = [f"{x.split('[')[0]}{x.split('[')[1].split(']')[-1]}" for x in dmat.design_info.column_names]
+    dmat = pd.DataFrame(np.asarray(dmat), index=obs.index, columns=column_names)
     return dmat