lamindb 0.69.9__py3-none-any.whl → 0.69.10__py3-none-any.whl

lamindb/_annotate.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable
+from typing import TYPE_CHECKING, Iterable, Optional
 
 import anndata as ad
 import lamindb_setup as ln_setup
 import pandas as pd
 from lamin_utils import colors, logger
+from lamindb_setup.core._docs import doc_args
 from lnschema_core import Artifact, Collection, Feature, Registry, Run, ULabel
 
 if TYPE_CHECKING:
     from lnschema_core.types import FieldAttr
+    from mudata import MuData
 
 
 class ValidationError(ValueError):
@@ -85,10 +87,16 @@ class DataFrameAnnotator:
         df: The DataFrame object to annotate.
         columns: The field attribute for the feature column.
         categoricals: A dictionary mapping column names to registry_field.
-            For example:
-            ``{"cell_type_ontology_id": bt.CellType.ontology_id, "donor_id": ln.ULabel.name}``.
         using: The reference instance containing registries to validate against.
         verbosity: The verbosity level.
+        organism: The organism name.
+
+    Examples:
+        >>> import bionty as bt
+        >>> annotate = ln.Annotate.from_df(
+                df,
+                categoricals={"cell_type_ontology_id": bt.CellType.ontology_id, "donor_id": ln.ULabel.name}
+            )
     """
 
     def __init__(
@@ -98,7 +106,7 @@ class DataFrameAnnotator:
         categoricals: dict[str, FieldAttr] | None = None,
         using: str | None = None,
         verbosity: str = "hint",
-        **kwargs,
+        organism: str | None = None,
     ) -> None:
         from lamindb.core._settings import settings
 
@@ -110,7 +118,7 @@ class DataFrameAnnotator:
         self._artifact = None
         self._collection = None
         self._validated = False
-        self._kwargs: dict = kwargs
+        self._kwargs = {"organism": organism} if organism else {}
         self._save_columns()
 
     @property
@@ -119,7 +127,7 @@ class DataFrameAnnotator:
         return self._fields
 
     def lookup(self, using: str | None = None) -> AnnotateLookup:
-        """Lookup features and labels.
+        """Lookup categories.
 
         Args:
             using: The instance where the lookup is performed.
@@ -132,7 +140,7 @@ class DataFrameAnnotator:
             using=using or self._using,
         )
 
-    def _save_columns(self, validated_only: bool = True) -> None:
+    def _save_columns(self, validated_only: bool = True, **kwargs) -> None:
         """Save column name records."""
         missing_columns = set(self.fields.keys()) - set(self._df.columns)
         if missing_columns:
@@ -148,7 +156,7 @@ class DataFrameAnnotator:
             save_function="add_new_from_columns",
             using=self._using,
             validated_only=False,
-            kwargs=self._kwargs,
+            **kwargs,
         )
 
         # Save the rest of the columns based on validated_only
@@ -162,36 +170,47 @@ class DataFrameAnnotator:
                 using=self._using,
                 validated_only=validated_only,
                 df=self._df,  # Get the Feature type from df
-                kwargs=self._kwargs,
+                **kwargs,
             )
 
-    def add_validated_from(self, key: str, **kwargs):
+    def add_validated_from(self, key: str, organism: str | None = None):
         """Add validated categories.
 
         Args:
             key: The key referencing the slot in the DataFrame.
-            **kwargs: Additional keyword arguments.
+            organism: The organism name.
         """
-        self._update_registry(key, validated_only=True, **kwargs)
+        self._kwargs.update({"organism": organism} if organism else {})
+        self._update_registry(key, validated_only=True, **self._kwargs)
 
-    def add_new_from(self, key: str, **kwargs):
+    def add_new_from(self, key: str, organism: str | None = None, **kwargs):
         """Add validated & new categories.
 
         Args:
             key: The key referencing the slot in the DataFrame from which to draw terms.
-            **kwargs: Additional keyword arguments.
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
         """
-        self._update_registry(key, validated_only=False, **kwargs)
+        if len(kwargs) > 0 and key == "all":
+            raise ValueError("Cannot pass additional arguments to 'all' key!")
+        self._kwargs.update({"organism": organism} if organism else {})
+        self._update_registry(key, validated_only=False, **self._kwargs, **kwargs)
+
+    def add_new_from_columns(self, organism: str | None = None, **kwargs):
+        """Add validated & new column names to its registry.
 
-    def add_new_from_columns(self, **kwargs):
-        """Add validated & new column names to its registry."""
-        self._save_columns(validated_only=False, **kwargs)
+        Args:
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        self._save_columns(validated_only=False, **self._kwargs, **kwargs)
 
     def _update_registry(self, categorical: str, validated_only: bool = True, **kwargs):
         if categorical == "all":
             self._update_registry_all(validated_only=validated_only, **kwargs)
         elif categorical == "columns":
-            self._save_columns(validated_only=validated_only)
+            self._save_columns(validated_only=validated_only, **kwargs)
         else:
             if categorical not in self.fields:
                 raise ValueError(f"Feature {categorical} is not part of the fields!")
@@ -201,7 +220,7 @@ class DataFrameAnnotator:
                 key=categorical,
                 using=self._using,
                 validated_only=validated_only,
-                kwargs=kwargs,
+                **kwargs,
             )
 
     def _update_registry_all(self, validated_only: bool = True, **kwargs):
@@ -210,13 +229,13 @@ class DataFrameAnnotator:
             logger.info(f"saving labels for '{name}'")
             self._update_registry(name, validated_only=validated_only, **kwargs)
 
-    def validate(self, **kwargs) -> bool:
+    def validate(self, organism: str | None = None) -> bool:
         """Validate variables and categorical observations.
 
         Returns:
             Whether the DataFrame is validated.
         """
-        self._kwargs.update(kwargs)
+        self._kwargs.update({"organism": organism} if organism else {})
         self._validated = validate_categories_in_df(
             self._df,
             fields=self.fields,
@@ -237,7 +256,6 @@ class DataFrameAnnotator:
         """
         from lamindb.core._settings import settings
 
-        self._kwargs.update(kwargs)
         if not self._validated:
             raise ValidationError(
                 f"Data object is not validated, please run {colors.yellow('validate()')}!"
@@ -255,6 +273,7 @@ class DataFrameAnnotator:
                 description=description,
                 fields=self.fields,
                 columns_field=self._columns_field,
+                **kwargs,
                 **self._kwargs,
             )
         finally:
@@ -314,9 +333,18 @@ class AnnDataAnnotator(DataFrameAnnotator):
         adata: The AnnData object to annotate.
         var_index: The registry field for mapping the ``.var`` index.
         categoricals: A dictionary mapping ``.obs.columns`` to a registry field.
-            For example:
-            ``{"cell_type_ontology_id": bt.CellType.ontology_id, "donor_id": ln.ULabel.name}``
         using: A reference LaminDB instance.
+        verbosity: The verbosity level.
+        organism: The organism name.
+
+    Examples:
+        >>> import bionty as bt
+        >>> annotate = ln.Annotate.from_anndata(
+                adata,
+                var_index=bt.Gene.ensembl_gene_id,
+                categoricals={"cell_type_ontology_id": bt.CellType.ontology_id, "donor_id": ln.ULabel.name},
+                organism="human",
+            )
     """
 
     def __init__(
@@ -326,7 +354,7 @@ class AnnDataAnnotator(DataFrameAnnotator):
         categoricals: dict[str, FieldAttr],
         using: str = "default",
         verbosity: str = "hint",
-        **kwargs,
+        organism: str | None = None,
     ) -> None:
         self._adata = adata
         self._var_field = var_index
@@ -335,10 +363,10 @@ class AnnDataAnnotator(DataFrameAnnotator):
             categoricals=categoricals,
             using=using,
             verbosity=verbosity,
-            **kwargs,
+            organism=organism,
         )
         self._obs_fields = categoricals
-        self._save_from_var_index()
+        self._save_from_var_index(validated_only=True, **self._kwargs)
 
     @property
     def var_index(self) -> FieldAttr:
@@ -351,16 +379,23 @@ class AnnDataAnnotator(DataFrameAnnotator):
         return self._obs_fields
 
     def lookup(self, using: str | None = None) -> AnnotateLookup:
-        """Lookup features and labels."""
+        """Lookup categories.
+
+        Args:
+            using: The instance where the lookup is performed.
+                if None (default), the lookup is performed on the instance specified in "using" parameter of the validator.
+                if "public", the lookup is performed on the public reference.
+        """
         return AnnotateLookup(
             categorials=self._obs_fields,
             slots={"columns": self._columns_field, "var_index": self._var_field},
             using=using or self._using,
         )
 
-    def _save_from_var_index(self, validated_only: bool = True, **kwargs):
+    def _save_from_var_index(
+        self, validated_only: bool = True, organism: str | None = None
+    ):
         """Save variable records."""
-        self._kwargs.update(kwargs)
         update_registry(
             values=self._adata.var.index,
             field=self.var_index,
@@ -368,35 +403,56 @@ class AnnDataAnnotator(DataFrameAnnotator):
             save_function="add_new_from_var_index",
             using=self._using,
             validated_only=validated_only,
-            kwargs=self._kwargs,
+            organism=organism,
         )
 
-    def add_new_from_var_index(self, **kwargs):
-        """Update variable records."""
-        self._save_from_var_index(validated_only=False, **kwargs)
+    def add_new_from_var_index(self, organism: str | None = None, **kwargs):
+        """Update variable records.
 
-    def validate(self, **kwargs) -> bool:
-        """Validate categories."""
-        self._kwargs.update(kwargs)
-        self._validated = validate_anndata(
-            self._adata,
-            var_field=self.var_index,
-            obs_fields=self.categoricals,
+        Args:
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        self._save_from_var_index(validated_only=False, **self._kwargs, **kwargs)
+
+    def validate(self, organism: str | None = None) -> bool:
+        """Validate categories.
+
+        Args:
+            organism: The organism name.
+
+        Returns:
+            Whether the AnnData object is validated.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        if self._using is not None and self._using != "default":
+            logger.important(
+                f"validating metadata using registries of instance {colors.italic(self._using)}"
+            )
+        validated_var = validate_categories(
+            self._adata.var.index,
+            field=self._var_field,
+            key="var_index",
+            using=self._using,
             **self._kwargs,
         )
+        validated_obs = validate_categories_in_df(
+            self._adata.obs, fields=self.categoricals, using=self._using, **self._kwargs
+        )
+        self._validated = validated_var and validated_obs
         return self._validated
 
     def save_artifact(self, description: str, **kwargs) -> Artifact:
-        """Save the validated AnnData and metadata.
+        """Save the validated ``AnnData`` and metadata.
 
         Args:
-            description: Description of the AnnData object.
+            description: Description of the ``AnnData`` object.
             **kwargs: Object level metadata.
 
         Returns:
             A saved artifact record.
         """
-        self._kwargs.update(kwargs)
         if not self._validated:
             raise ValidationError("Please run `validate()` first!")
 
@@ -406,6 +462,261 @@ class AnnDataAnnotator(DataFrameAnnotator):
             columns_field=self.var_index,
             fields=self.categoricals,
             **self._kwargs,
+            **kwargs,
+        )
+        return self._artifact
+
+
+class MuDataAnnotator:
+    """Annotation flow for a ``MuData`` object.
+
+    Args:
+        mdata: The MuData object to annotate.
+        var_index: The registry field for mapping the ``.var`` index for each modality.
+            For example:
+            ``{"modality_1": bt.Gene.ensembl_gene_id, "modality_2": ln.CellMarker.name}``
+        categoricals: A dictionary mapping ``.obs.columns`` to a registry field.
+        using: A reference LaminDB instance.
+        verbosity: The verbosity level.
+        organism: The organism name.
+
+    Examples:
+        >>> import bionty as bt
+        >>> annotate = ln.Annotate.from_mudata(
+                mdata,
+                var_index={"rna": bt.Gene.ensembl_gene_id, "adt": ln.CellMarker.name},
+                categoricals={"cell_type_ontology_id": bt.CellType.ontology_id, "donor_id": ln.ULabel.name},
+                organism="human",
+            )
+    """
+
+    def __init__(
+        self,
+        mdata: MuData,
+        var_index: dict[str, dict[str, FieldAttr]],
+        categoricals: dict[str, FieldAttr],
+        using: str = "default",
+        verbosity: str = "hint",
+        organism: str | None = None,
+    ) -> None:
+        self._mdata = mdata
+        self._kwargs = {"organism": organism} if organism else {}
+        self._var_fields = var_index
+        self._verify_modality(self._var_fields.keys())
+        self._obs_fields = self._parse_categoricals(categoricals)
+        self._modalities = set(self._var_fields.keys()) | set(self._obs_fields.keys())
+        self._using = using
+        self._verbosity = verbosity
+        self._df_annotators = {
+            modality: DataFrameAnnotator(
+                df=mdata[modality].obs if modality != "obs" else mdata.obs,
+                categoricals=self._obs_fields.get(modality, {}),
+                using=using,
+                verbosity=verbosity,
+                **self._kwargs,
+            )
+            for modality in self._modalities
+        }
+        for modality in self._var_fields.keys():
+            self._save_from_var_index_modality(
+                modality=modality, validated_only=True, **self._kwargs
+            )
+
+    @property
+    def var_index(self) -> FieldAttr:
+        """Return the registry field to validate variables index against."""
+        return self._var_fields
+
+    @property
+    def categoricals(self) -> dict:
+        """Return the obs fields to validate against."""
+        return self._obs_fields
+
+    def _verify_modality(self, modalities: Iterable[str]):
+        """Verify the modality exists."""
+        for modality in modalities:
+            if modality not in self._mdata.mod.keys():
+                raise ValueError(f"modality '{modality}' does not exist!")
+
+    def _save_from_var_index_modality(
+        self, modality: str, validated_only: bool = True, **kwargs
+    ):
+        """Save variable records."""
+        update_registry(
+            values=self._mdata[modality].var.index,
+            field=self._var_fields[modality],
+            key="var_index",
+            save_function="add_new_from_var_index",
+            using=self._using,
+            validated_only=validated_only,
+            type="number",
+            **kwargs,
+        )
+
+    def _parse_categoricals(self, categoricals: dict[str, FieldAttr]) -> dict:
+        """Parse the categorical fields."""
+        prefixes = {f"{k}:" for k in self._mdata.mod.keys()}
+        obs_fields: dict[str, dict[str, FieldAttr]] = {}
+        for k, v in categoricals.items():
+            if k not in self._mdata.obs.columns:
+                raise ValueError(f"column '{k}' does not exist in mdata.obs!")
+            if any(k.startswith(prefix) for prefix in prefixes):
+                modality, col = k.split(":")[0], k.split(":")[1]
+                if modality not in obs_fields.keys():
+                    obs_fields[modality] = {}
+                obs_fields[modality][col] = v
+            else:
+                if "obs" not in obs_fields.keys():
+                    obs_fields["obs"] = {}
+                obs_fields["obs"][k] = v
+        return obs_fields
+
+    def lookup(self, using: str | None = None) -> AnnotateLookup:
+        """Lookup categories.
+
+        Args:
+            using: The instance where the lookup is performed.
+                if None (default), the lookup is performed on the instance specified in "using" parameter of the validator.
+                if "public", the lookup is performed on the public reference.
+        """
+        return AnnotateLookup(
+            categorials=self._obs_fields,
+            slots={
+                **self._obs_fields,
+                **{f"{k}_var_index": v for k, v in self._var_fields.items()},
+            },
+            using=using or self._using,
+        )
+
+    def add_new_from_columns(
+        self,
+        modality: str,
+        column_names: list[str] | None = None,
+        organism: str | None = None,
+        **kwargs,
+    ):
+        """Update columns records.
+
+        Args:
+            modality: The modality name.
+            column_names: The column names to save.
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        update_registry(
+            values=column_names or self._mdata[modality].obs.columns,
+            field=Feature.name,
+            key=f"{modality} obs columns",
+            using=self._using,
+            validated_only=False,
+            df=self._mdata[modality].obs,
+            **self._kwargs,
+            **kwargs,
+        )
+
+    def add_new_from_var_index(
+        self, modality: str, organism: str | None = None, **kwargs
+    ):
+        """Update variable records.
+
+        Args:
+            modality: The modality name.
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        self._save_from_var_index_modality(
+            modality=modality, validated_only=False, **self._kwargs, **kwargs
+        )
+
+    def add_validated_from(
+        self, key: str, modality: str | None = None, organism: str | None = None
+    ):
+        """Add validated categories.
+
+        Args:
+            key: The key referencing the slot in the DataFrame.
+            modality: The modality name.
+            organism: The organism name.
+        """
+        self._kwargs.update({"organism": organism} if organism else {})
+        modality = modality or "obs"
+        if modality in self._df_annotators:
+            df_annotator = self._df_annotators[modality]
+            df_annotator.add_validated_from(key=key, **self._kwargs)
+
+    def add_new_from(
+        self,
+        key: str,
+        modality: str | None = None,
+        organism: str | None = None,
+        **kwargs,
+    ):
+        """Add validated & new categories.
+
+        Args:
+            key: The key referencing the slot in the DataFrame.
+            modality: The modality name.
+            organism: The organism name.
+            **kwargs: Additional keyword arguments to pass to the registry model.
+        """
+        if len(kwargs) > 0 and key == "all":
+            raise ValueError("Cannot pass additional arguments to 'all' key!")
+        self._kwargs.update({"organism": organism} if organism else {})
+        modality = modality or "obs"
+        if modality in self._df_annotators:
+            df_annotator = self._df_annotators[modality]
+            df_annotator.add_new_from(key=key, **self._kwargs, **kwargs)
+
+    def validate(self, organism: str | None = None) -> bool:
+        """Validate categories."""
+        self._kwargs.update({"organism": organism} if organism else {})
+        if self._using is not None and self._using != "default":
+            logger.important(
+                f"validating metadata using registries of instance {colors.italic(self._using)}"
+            )
+        validated_var = True
+        for modality, var_field in self._var_fields.items():
+            validated_var &= validate_categories(
+                self._mdata[modality].var.index,
+                field=var_field,
+                key=f"{modality}_var_index",
+                using=self._using,
+                **self._kwargs,
+            )
+        validated_obs = True
+        for modality, fields in self._obs_fields.items():
+            if modality == "obs":
+                obs = self._mdata.obs
+            else:
+                obs = self._mdata[modality].obs
+            validated_obs &= validate_categories_in_df(
+                obs, fields=fields, using=self._using, **self._kwargs
+            )
+        self._validated = validated_var and validated_obs
+        return self._validated
+
+    def save_artifact(self, description: str, **kwargs) -> Artifact:
+        """Save the validated ``MuData`` and metadata.
+
+        Args:
+            description: Description of the ``MuData`` object.
+            **kwargs: Object level metadata.
+
+        Returns:
+            A saved artifact record.
+        """
+        if not self._validated:
+            raise ValidationError("Please run `validate()` first!")
+
+        self._artifact = save_artifact(
+            self._mdata,
+            description=description,
+            columns_field=self.var_index,
+            fields=self.categoricals,
+            **self._kwargs,
+            **kwargs,
         )
         return self._artifact
 
@@ -414,6 +725,7 @@ class Annotate:
     """Annotation flow."""
 
     @classmethod
+    @doc_args(DataFrameAnnotator.__doc__)
     def from_df(
         cls,
         df: pd.DataFrame,
@@ -421,18 +733,20 @@ class Annotate:
         columns: FieldAttr = Feature.name,
         using: str | None = None,
         verbosity: str = "hint",
-        **kwargs,
+        organism: str | None = None,
     ) -> DataFrameAnnotator:
+        """{}."""
         return DataFrameAnnotator(
             df=df,
             categoricals=categoricals,
             columns=columns,
             using=using,
             verbosity=verbosity,
-            **kwargs,
+            organism=organism,
         )
 
     @classmethod
+    @doc_args(AnnDataAnnotator.__doc__)
     def from_anndata(
         cls,
         adata: ad.AnnData,
@@ -440,15 +754,37 @@ class Annotate:
         categoricals: dict[str, FieldAttr],
         using: str = "default",
         verbosity: str = "hint",
-        **kwargs,
+        organism: str | None = None,
     ) -> AnnDataAnnotator:
+        """{}."""
         return AnnDataAnnotator(
             adata=adata,
             var_index=var_index,
             categoricals=categoricals,
             using=using,
             verbosity=verbosity,
-            **kwargs,
+            organism=organism,
+        )
+
+    @classmethod
+    @doc_args(MuDataAnnotator.__doc__)
+    def from_mudata(
+        cls,
+        mdata: MuData,
+        var_index: dict[str, dict[str, FieldAttr]],
+        categoricals: dict[str, dict[str, FieldAttr]],
+        using: str = "default",
+        verbosity: str = "hint",
+        organism: str | None = None,
+    ) -> MuDataAnnotator:
+        """{}."""
+        return MuDataAnnotator(
+            mdata=mdata,
+            var_index=var_index,
+            categoricals=categoricals,
+            using=using,
+            verbosity=verbosity,
+            organism=organism,
         )
 
 
@@ -468,9 +804,7 @@ def standardize_and_inspect(
     return registry.inspect(values, field=field, mute=True, **kwargs)
 
 
-def check_registry_organism(
-    registry: Registry, organism: str | None = None
-) -> str | None:
+def check_registry_organism(registry: Registry, organism: str | None = None) -> dict:
     """Check if a registry needs an organism and return the organism name."""
     if hasattr(registry, "organism_id"):
         import bionty as bt
@@ -480,8 +814,8 @@ def check_registry_organism(
                 f"{registry.__name__} registry requires an organism!\n"
                 "      → please pass an organism name via organism="
             )
-        return organism or bt.settings.organism.name
-    return None
+        return {"organism": organism or bt.settings.organism.name}
+    return {}
 
 
 def validate_categories(
@@ -489,22 +823,21 @@ def validate_categories(
     field: FieldAttr,
     key: str,
     using: str | None = None,
-    **kwargs,
+    organism: str | None = None,
 ) -> bool:
     """Validate ontology terms in a pandas series using LaminDB registries."""
     from lamindb._from_values import _print_values
     from lamindb.core._settings import settings
 
     model_field = f"{field.field.model.__name__}.{field.field.name}"
-    logger.indent = ""
-    logger.info(f"mapping {colors.italic(key)} on {colors.italic(model_field)}")
-    logger.indent = "   "
+
+    def _log_mapping_info():
+        logger.indent = ""
+        logger.info(f"mapping {colors.italic(key)} on {colors.italic(model_field)}")
+        logger.indent = "   "
 
     registry = field.field.model
-    filter_kwargs = {}
-    organism = check_registry_organism(registry, kwargs.get("organism"))
-    if organism is not None:
-        filter_kwargs["organism"] = organism
+    filter_kwargs = check_registry_organism(registry, organism)
 
     # Inspect the default instance
     inspect_result = standardize_and_inspect(
@@ -537,6 +870,7 @@ def validate_categories(
     validated_hint_print = f".add_validated_from('{key}')"
     n_validated = len(values_validated)
     if n_validated > 0:
+        _log_mapping_info()
         logger.warning(
             f"found {colors.yellow(f'{n_validated} terms')} validated terms: "
             f"{colors.yellow(values_validated)}\n      → save terms via "
@@ -547,7 +881,8 @@ def validate_categories(
     non_validated = [i for i in non_validated if i not in values_validated]
     n_non_validated = len(non_validated)
     if n_non_validated == 0:
-        logger.success(f"{key} validated")
+        logger.indent = ""
+        logger.success(f"{key} is validated against {colors.italic(model_field)}")
         return True
     else:
         are = "are" if n_non_validated > 1 else "is"
@@ -557,6 +892,8 @@ def validate_categories(
             f"{colors.yellow(print_values)}\n      → save terms via "
             f"{colors.yellow(non_validated_hint_print)}"
         )
+        if logger.indent == "":
+            _log_mapping_info()
         logger.warning(warning_message)
         logger.indent = ""
         return False
@@ -581,37 +918,12 @@ def validate_categories_in_df(
     return validated
 
 
-def validate_anndata(
-    adata: ad.AnnData,
-    var_field: FieldAttr,
-    obs_fields: dict[str, FieldAttr],
-    using: str | None = None,
-    **kwargs,
-) -> bool:
-    """Inspect metadata in an AnnData object using LaminDB registries."""
-    if using is not None and using != "default":
-        logger.important(
-            f"validating metadata using registries of instance {colors.italic(using)}"
-        )
-
-    validated_var = validate_categories(
-        adata.var.index,
-        field=var_field,
-        key="var_index",
-        using=using,
-        **kwargs,
-    )
-    validated_obs = validate_categories_in_df(
-        adata.obs, fields=obs_fields, using=using, **kwargs
-    )
-    return validated_var and validated_obs
-
-
 def save_artifact(
-    data: pd.DataFrame | ad.AnnData,
+    data: pd.DataFrame | ad.AnnData | MuData,
     description: str,
-    fields: dict[str, FieldAttr],
-    columns_field: FieldAttr,
+    fields: dict[str, FieldAttr] | dict[str, dict[str, FieldAttr]],
+    columns_field: FieldAttr | dict[str, FieldAttr],
+    organism: str | None = None,
     **kwargs,
 ) -> Artifact:
     """Save all metadata with an Artifact.
@@ -621,43 +933,67 @@ def save_artifact(
         description: A description of the artifact.
         fields: A dictionary mapping obs_column to registry_field.
         columns_field: The registry field to validate variables index against.
+        organism: The organism name.
         kwargs: Additional keyword arguments to pass to the registry model.
 
     Returns:
         The saved Artifact.
     """
+    artifact = None
     if isinstance(data, ad.AnnData):
-        artifact = Artifact.from_anndata(data, description=description)
+        artifact = Artifact.from_anndata(data, description=description, **kwargs)
         artifact.n_observations = data.n_obs
     elif isinstance(data, pd.DataFrame):
-        artifact = Artifact.from_df(data, description=description)
+        artifact = Artifact.from_df(data, description=description, **kwargs)
     else:
-        raise ValueError("data must be a DataFrame or AnnData object")
+        try:
+            from mudata import MuData
+
+            if isinstance(data, MuData):
+                artifact = Artifact.from_mudata(data, description=description, **kwargs)
+                artifact.n_observations = data.n_obs
+        except ImportError:
+            pass
+    if artifact is None:
+        raise ValueError("data must be a DataFrame, AnnData or MuData object")
     artifact.save()
 
-    feature_kwargs: dict = {}
-    organism = check_registry_organism(
-        columns_field.field.model, kwargs.pop("organism", None)
+    feature_kwargs = check_registry_organism(
+        (
+            list(columns_field.values())[0].field.model
+            if isinstance(columns_field, dict)
+            else columns_field.field.model
+        ),
+        organism,
     )
-    if organism is not None:
-        feature_kwargs["organism"] = organism
 
-    if isinstance(data, ad.AnnData):
+    if artifact.accessor == "DataFrame":
+        artifact.features.add_from_df(field=columns_field, **feature_kwargs)
+    elif artifact.accessor == "AnnData":
         artifact.features.add_from_anndata(var_field=columns_field, **feature_kwargs)
+    elif artifact.accessor == "MuData":
+        artifact.features.add_from_mudata(var_fields=columns_field, **feature_kwargs)
     else:
-        artifact.features.add_from_df(field=columns_field, **feature_kwargs)
-
-    features = Feature.lookup().dict()
-    for key, field in fields.items():
-        feature = features.get(key)
-        registry = field.field.model
-        filter_kwargs = kwargs.copy()
-        organism = check_registry_organism(registry, organism)
-        if organism is not None:
-            filter_kwargs["organism"] = organism
-        df = data.obs if isinstance(data, ad.AnnData) else data
-        labels = registry.from_values(df[key], field=field, **filter_kwargs)
-        artifact.labels.add(labels, feature)
+        raise NotImplementedError
+
+    def _add_labels(data, artifact: Artifact, fields: dict[str, FieldAttr]):
+        features = Feature.lookup().dict()
+        for key, field in fields.items():
+            feature = features.get(key)
+            registry = field.field.model
+            filter_kwargs = check_registry_organism(registry, organism)
+            df = data if isinstance(data, pd.DataFrame) else data.obs
+            labels = registry.from_values(df[key], field=field, **filter_kwargs)
+            artifact.labels.add(labels, feature)
+
+    if artifact.accessor == "MuData":
+        for modality, modality_fields in fields.items():
+            if modality == "obs":
+                _add_labels(data, artifact, modality_fields)
+            else:
+                _add_labels(data[modality], artifact, modality_fields)
+    else:
+        _add_labels(data, artifact, fields)
 
     slug = ln_setup.settings.instance.slug
     if ln_setup.settings.instance.is_remote:
@@ -672,8 +1008,10 @@ def update_registry(
     save_function: str = "add_new_from",
     using: str | None = None,
     validated_only: bool = True,
-    kwargs: dict | None = None,
     df: pd.DataFrame | None = None,
+    organism: str | None = None,
+    type: str | None = None,
+    **kwargs,
 ) -> None:
     """Save features or labels records in the default instance from the using instance.
 
@@ -684,18 +1022,16 @@ def update_registry(
         save_function: The name of the function to save the labels.
         using: The name of the instance from which to transfer labels (if applicable).
         validated_only: If True, only save validated labels.
-        kwargs: Additional keyword arguments to pass to the registry model.
         df: A DataFrame to save labels from.
+        organism: The organism name.
+        type: The type of the feature.
+        kwargs: Additional keyword arguments to pass to the registry model to create new records.
     """
     from lamindb._save import save as ln_save
     from lamindb.core._settings import settings
 
-    filter_kwargs = {} if kwargs is None else kwargs.copy()
     registry = field.field.model
-
-    organism = check_registry_organism(registry, filter_kwargs.pop("organism", None))
-    if organism is not None:
-        filter_kwargs["organism"] = organism
+    filter_kwargs = check_registry_organism(registry, organism)
 
     verbosity = settings.verbosity
     try:
@@ -716,7 +1052,7 @@ def update_registry(
             inspect_result_current.non_validated,
             field=field,
             using=using,
-            kwargs=filter_kwargs,
+            **filter_kwargs,
         )
 
         public_records = (
@@ -742,8 +1078,8 @@ def update_registry(
                 for value in labels_saved["without reference"]:
                     filter_kwargs[field.field.name] = value
                     if registry == Feature:
-                        filter_kwargs["type"] = "category"
-                    non_validated_records.append(registry(**filter_kwargs))
+                        filter_kwargs["type"] = "category" if type is None else type
+                    non_validated_records.append(registry(**filter_kwargs, **kwargs))
             ln_save(non_validated_records)
 
         if registry == ULabel and field.field.name == "name":
@@ -815,7 +1151,7 @@ def update_registry_from_using_instance(
     values: list[str],
     field: FieldAttr,
     using: str | None = None,
-    kwargs: dict | None = None,
+    **kwargs,
 ) -> tuple[list[str], list[str]]:
     """Save features or labels records from the using instance.
 
@@ -828,7 +1164,6 @@ def update_registry_from_using_instance(
     Returns:
         A tuple containing the list of saved labels and the list of non-saved labels.
     """
-    kwargs = kwargs or {}
     labels_saved = []
     not_saved = values