Address PR comments

Fix `map_location`
Replace `annotated_only` with `factor_filter`

Author: arberqoku

docs/notebooks/tcga_brca_bulk_multiomics.ipynb

@@ -543,7 +543,7 @@
    "id": "e6314dd0",
    "metadata": {},
    "source": [
-    "We confirm the optimiser converged, then look at how variance is partitioned across the informed pathway factors. We pass `annotated_only=True` to drop the single dense factor, which absorbs the dominant axis of variation; each informed factor explains a smaller, pathway-specific slice."
+    "We confirm the optimiser converged, then look at how variance is partitioned across the informed pathway factors. We pass a factor-name predicate that keeps factors present in the Hallmark annotation mask, dropping the single dense factor that absorbs the dominant axis of variation; each informed factor explains a smaller, pathway-specific slice."
    ]
   },
   {
@@ -611,7 +611,11 @@
     }
    ],
    "source": [
-    "mfl.pl.variance_explained(model, annotated_only=True, figsize=(7, 12))"
+    "mfl.pl.variance_explained(\n",
+    "    model,\n",
+    "    factor_filter=lambda factor: factor in mdata[\"rna\"].varm[\"gene_set_mask\"].columns,\n",
+    "    figsize=(7, 12),\n",
+    ")"
    ]
   },
   {

src/mofaflex/_core/mofaflex.py

@@ -537,6 +537,7 @@ def load(cls, path: str | Path, map_location=None) -> "MOFAFLEX":
 
         if map_location is not None:
             state["train_opts"]["device"] = map_location
+        map_location = state["train_opts"]["device"]
 
         model = cls.__new__(cls)
         model._train_loss_elbo = state["train_loss_elbo"]

src/mofaflex/_core/terms/mofaflex.py

@@ -115,7 +115,7 @@ def __init__(
 
     def _results_to_df(
         self,
-        results: Mapping[str, np.ndarray],
+        results: Mapping[str, np.ndarray | pd.DataFrame],
         axis: Literal[0, 1],
         ordered: bool = False,
         factors_subset: slice = slice(None),
@@ -124,14 +124,18 @@ def _results_to_df(
         ret = {}
         for name, res in results.items():
             fnames = factor_names
+            index = (
+                res.index
+                if isinstance(res, pd.DataFrame)
+                else (self._sample_names[name] if axis == 0 else self._feature_names[name])
+            )
+            values = res.to_numpy() if isinstance(res, pd.DataFrame) else res
             if ordered:
                 factor_order = self.factor_order[factors_subset].copy()
                 factor_order[np.argsort(factor_order)] = np.arange(len(factor_order))
-                res = res[:, factor_order]
+                values = values[:, factor_order]
                 fnames = fnames[factor_order]
-            ret[name] = pd.DataFrame(
-                res, index=self._sample_names[name] if axis == 0 else self._feature_names[name], columns=fnames
-            )
+            ret[name] = pd.DataFrame(values, index=index, columns=fnames)
         return ret
 
     def _wrap_api_method(self, axis: Literal[0, 1], prior: Prior, api: PriorDynamicAPI):
@@ -785,9 +789,7 @@ def _load(
         )
 
         self._prior_api_properties = {}
-        # map_location may be None (e.g. MOFAFLEX.load without an explicit device); fall back to the
-        # default device so methods that need it (e.g. PCGSE in get_significant_annotations) still work.
-        self._device = default_torch_device(map_location)
+        self._device = map_location
         self._init_api()
 
     def _get_postprocessed_factors(

src/mofaflex/pl/_plotting.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from contextlib import suppress
 from functools import partial
 from typing import TYPE_CHECKING, Literal
@@ -464,7 +464,7 @@ def variance_explained(
     model: MOFAFLEX,
     group_by: Literal["group", "view"] = "group",
     term: str | None = None,
-    annotated_only: bool = False,
+    factor_filter: Callable[[str], bool] | None = None,
     figsize: tuple[float, float] | None = None,
 ) -> p9.ggplot:
     """Plot the fraction of variance explained per factor in each group and view.
@@ -475,9 +475,7 @@ def variance_explained(
         term: The name of the additive term to plot the variance explained for. If `None` and the model has only one additive term,
             will plot the fraction of variance explained per factor for this term. If `None` and the model has multiple terms, will
             plot the fraction of variance explained per term.
-        annotated_only: If `True`, only plot the factors that are informed by prior annotations (e.g. the gene-set factors of an
-            :class:`~mofaflex.priors.InformedHorseshoe` prior), dropping the uninformed dense factors. Requires the model to
-            have an informed prior.
+        factor_filter: Predicate applied to factor names. Only factors for which the predicate returns `True` are plotted.
         figsize: Figure size in inches.
     """
     if group_by == "group":
@@ -491,19 +489,12 @@ def variance_explained(
         figsize = (len(model.group_names) * 3, 5)
 
     byterm = term is None and model.n_terms > 1
-    if annotated_only and byterm:
-        raise ValueError("`annotated_only` is only supported at the factor level. Specify a single `term`.")
+    if factor_filter is not None and byterm:
+        raise ValueError("`factor_filter` is only supported at the factor level. Specify a single `term`.")
     col = "term" if byterm else "component"
     df_r2 = model.get_r2("byterm" if byterm else "term", ordered=True, term=term)
-    if annotated_only:
-        try:
-            annotations = model.get_annotations()
-        except AttributeError as e:
-            raise ValueError(
-                "`annotated_only=True` requires a model with an informed prior providing annotations."
-            ) from e
-        informed_factors = set().union(*(annot.columns for annot in annotations.values()))
-        df_r2 = df_r2[df_r2[col].isin(informed_factors)]
+    if factor_filter is not None:
+        df_r2 = df_r2[df_r2[col].map(factor_filter)]
     df_r2 = df_r2.assign(factor=lambda x: pd.Categorical(x[col], categories=x[col].unique()))
     heatmap = (
         p9.ggplot(df_r2, p9.aes(x=x, y="factor", fill="R2"))
@@ -1144,15 +1135,7 @@ def top_weights(
         and (df.groupby("factor", observed=True)["feature"].aggregate(lambda x: x.duplicated().sum()) > 0).any()
     ):
         df = df.assign(feature=lambda x: x.feature.str + "_" + x.view.str)
-    # A feature can be among the top features of several factors. With a single global feature category, those
-    # shared features share one y position, scrambling the per-facet ordering. Make the category unique per facet
-    # (factor) so each facet is ordered by its own |weight|, and strip the suffix again for the axis labels.
-    _sep = "\x1f"
-    df = df.assign(
-        feature=lambda x: pd.Categorical(
-            keys := x.feature.astype(str) + _sep + x.factor.astype(str), categories=keys.unique()
-        )
-    )
+    df = df.assign(feature=lambda x: pd.Categorical(x.feature, categories=x.feature.unique()))
 
     aes_kwargs = {}
     if have_annot:
@@ -1163,7 +1146,6 @@ def top_weights(
         + p9.geom_segment()
         + p9.geom_point(size=5, stroke=0)
         + p9.scale_shape_manual(values=("$\\oplus$", "$\\ominus$"), breaks=(True, False), guide=None)
-        + p9.scale_y_discrete(labels=lambda breaks: [b.rsplit(_sep, 1)[0] for b in breaks])
         + _weights_inferred_color_scale
         + p9.scale_x_continuous(expand=(0, 0, 0.05, 0))
         + p9.labs(x="| Weight |", y="", color="")