lamindb 0.72.1__py3-none-any.whl → 0.73.0__py3-none-any.whl

lamindb/core/_feature_manager.py

@@ -1,24 +1,36 @@
 from __future__ import annotations
 
+from collections import defaultdict
+from collections.abc import Iterable
 from itertools import compress
-from typing import TYPE_CHECKING, Iterable
+from typing import TYPE_CHECKING, Any
 
 import anndata as ad
+import numpy as np
+import pandas as pd
 from anndata import AnnData
+from django.contrib.postgres.aggregates import ArrayAgg
+from django.db import connections, models
+from django.db.models import Aggregate, CharField, F, Value
+from django.db.models.functions import Concat
 from lamin_utils import colors, logger
 from lamindb_setup.core.upath import create_path
 from lnschema_core.models import (
     Artifact,
     Collection,
-    Data,
     Feature,
+    FeatureManager,
+    FeatureManagerArtifact,
+    FeatureManagerCollection,
     FeatureValue,
+    HasFeatures,
+    LinkORM,
     Registry,
     ULabel,
 )
 
-from lamindb._feature import convert_numpy_dtype_to_lamin_feature_type
-from lamindb._feature_set import FeatureSet
+from lamindb._feature import FEATURE_TYPES, convert_numpy_dtype_to_lamin_feature_type
+from lamindb._feature_set import DICT_KEYS_TYPE, FeatureSet
 from lamindb._registry import (
     REGISTRY_UNIQUE_FIELD,
     get_default_str_field,
@@ -29,6 +41,7 @@ from lamindb._save import save
 from lamindb.core.exceptions import ValidationError
 from lamindb.core.storage import LocalPathClasses
 
+from ._label_manager import get_labels_as_dict
 from ._settings import settings
 
 if TYPE_CHECKING:
@@ -45,7 +58,7 @@ def get_host_id_field(host: Artifact | Collection) -> str:
     return host_id_field
 
 
-def get_accessor_by_orm(host: Artifact | Collection) -> dict:
+def get_accessor_by_registry_(host: Artifact | Collection) -> dict:
     dictionary = {
         field.related_model.__get_name_with_schema__(): field.name
         for field in host._meta.related_objects
@@ -55,7 +68,7 @@ def get_accessor_by_orm(host: Artifact | Collection) -> dict:
     return dictionary
 
 
-def get_feature_set_by_slot(host) -> dict:
+def get_feature_set_by_slot_(host) -> dict:
     # if the host is not yet saved
     if host._state.adding:
         if hasattr(host, "_feature_sets"):
@@ -80,7 +93,7 @@ def get_label_links(
     host_id_field = get_host_id_field(host)
     kwargs = {host_id_field: host.id, "feature_id": feature.id}
     link_records = (
-        getattr(host, host.features.accessor_by_orm[registry])
+        getattr(host, host.features._accessor_by_registry[registry])
         .through.objects.using(host._state.db)
         .filter(**kwargs)
     )
@@ -94,53 +107,95 @@ def get_feature_set_links(host: Artifact | Collection) -> QuerySet:
     return feature_set_links
 
 
-def print_features(self: Data) -> str:
+def get_link_attr(link: LinkORM, data: HasFeatures) -> str:
+    link_model_name = link.__class__.__name__
+    link_attr = link_model_name.replace(data.__class__.__name__, "")
+    if link_attr == "ExperimentalFactor":
+        link_attr = "experimental_factor"
+    else:
+        link_attr = link_attr.lower()
+    return link_attr
+
+
+# Custom aggregation for SQLite
+class GroupConcat(Aggregate):
+    function = "GROUP_CONCAT"
+    template = '%(function)s(%(expressions)s, ", ")'
+
+
+def custom_aggregate(field, using: str):
+    if connections[using].vendor == "postgresql":
+        return ArrayAgg(field)
+    else:
+        return GroupConcat(field)
+
+
+def print_features(
+    self: HasFeatures, print_types: bool = False, to_dict: bool = False
+) -> str | dict[str, Any]:
     from lamindb._from_values import _print_values
 
-    from ._data import format_repr
-
-    messages = []
-    for slot, feature_set in get_feature_set_by_slot(self).items():
-        if feature_set.registry != "Feature":
-            features = feature_set.members
-            # features.first() is a lot slower than features[0] here
-            name_field = get_default_str_field(features[0])
-            feature_names = list(features.values_list(name_field, flat=True)[:30])
-            messages.append(
-                f"  {colors.bold(slot)}: {format_repr(feature_set, exclude='hash')}\n"
-            )
-            print_values = _print_values(feature_names, n=20)
-            messages.append(f"    {print_values}\n")
-        else:
-            features_lookup = {
-                f.name: f for f in Feature.objects.using(self._state.db).filter().all()
-            }
-            messages.append(
-                f"  {colors.bold(slot)}: {format_repr(feature_set, exclude='hash')}\n"
+    msg = ""
+    dictionary = {}
+    # categorical feature values
+    labels_msg = ""
+    labels_by_feature = defaultdict(list)
+    for _, (_, links) in get_labels_as_dict(self, links=True).items():
+        for link in links:
+            if link.feature_id is not None:
+                link_attr = get_link_attr(link, self)
+                labels_by_feature[link.feature_id].append(getattr(link, link_attr).name)
+    for feature_id, labels_list in labels_by_feature.items():
+        feature = Feature.objects.using(self._state.db).get(id=feature_id)
+        print_values = _print_values(labels_list, n=10)
+        type_str = f": {feature.dtype}" if print_types else ""
+        if to_dict:
+            dictionary[feature.name] = (
+                labels_list if len(labels_list) > 1 else labels_list[0]
             )
-            for name, dtype in feature_set.features.values_list("name", "dtype"):
-                if dtype.startswith("cat["):
-                    labels = self.labels.get(features_lookup.get(name), mute=True)
-                    indent = ""
-                    if isinstance(labels, dict):
-                        messages.append(f"    🔗 {name} ({dtype})\n")
-                        indent = "    "
-                    else:
-                        labels = {dtype: labels}
-                    for registry, registry_labels in labels.items():
-                        field = get_default_str_field(registry_labels)
-                        values_list = registry_labels.values_list(field, flat=True)
-                        count_str = f"{feature_set.n}, {colors.italic(f'{registry}')}"
-                        print_values = _print_values(values_list[:20], n=10)
-                        msg_objects = (
-                            f"{indent}    🔗 {name} ({count_str}):" f" {print_values}\n"
-                        )
-                        messages.append(msg_objects)
-                else:
-                    messages.append(f"    {name} ({dtype})\n")
-    if messages:
-        messages.insert(0, f"{colors.green('Features')}:\n")
-    return "".join(messages)
+        labels_msg += f"    '{feature.name}'{type_str} = {print_values}\n"
+    if labels_msg:
+        msg += labels_msg
+
+    # non-categorical feature values
+    non_labels_msg = ""
+    if self.id is not None and self.__class__ == Artifact:
+        feature_values = self.feature_values.values(
+            "feature__name", "feature__dtype"
+        ).annotate(values=custom_aggregate("value", self._state.db))
+        if len(feature_values) > 0:
+            for fv in feature_values:
+                feature_name = fv["feature__name"]
+                feature_dtype = fv["feature__dtype"]
+                values = fv["values"]
+                # TODO: understand why the below is necessary
+                if not isinstance(values, list):
+                    values = [values]
+                if to_dict:
+                    dictionary[feature_name] = values if len(values) > 1 else values[0]
+                type_str = f": {feature_dtype}" if print_types else ""
+                non_labels_msg += f"    '{feature_name}'{type_str} = {_print_values(values, n=10, quotes=False)}\n"
+            msg += non_labels_msg
+
+    if msg != "":
+        msg = f"  {colors.italic('Features')}\n" + msg
+
+    # feature sets
+    feature_set_msg = ""
+    for slot, feature_set in get_feature_set_by_slot_(self).items():
+        features = feature_set.members
+        # features.first() is a lot slower than features[0] here
+        name_field = get_default_str_field(features[0])
+        feature_names = list(features.values_list(name_field, flat=True)[:20])
+        type_str = f": {feature_set.registry}" if print_types else ""
+        feature_set_msg += f"    '{slot}'{type_str} = {_print_values(feature_names)}\n"
+    if feature_set_msg:
+        msg += f"  {colors.italic('Feature sets')}\n"
+        msg += feature_set_msg
+    if to_dict:
+        return dictionary
+    else:
+        return msg
 
 
 def parse_feature_sets_from_anndata(
@@ -203,305 +258,443 @@ def parse_feature_sets_from_anndata(
     return feature_sets
 
 
-class FeatureManager:
-    """Feature manager (:attr:`~lamindb.core.Data.features`).
+def infer_feature_type_convert_json(value: Any, mute: bool = False) -> tuple[str, Any]:
+    if isinstance(value, bool):
+        return FEATURE_TYPES["bool"], value
+    elif isinstance(value, int):
+        return FEATURE_TYPES["int"], value
+    elif isinstance(value, float):
+        return FEATURE_TYPES["float"], value
+    elif isinstance(value, str):
+        return FEATURE_TYPES["str"] + "[ULabel]", value
+    elif isinstance(value, Iterable) and not isinstance(value, (str, bytes)):
+        if isinstance(value, (pd.Series, np.ndarray)):
+            return convert_numpy_dtype_to_lamin_feature_type(value.dtype), list(value)
+        if len(value) > 0:  # type: ignore
+            first_element_type = type(next(iter(value)))
+            if all(isinstance(elem, first_element_type) for elem in value):
+                if first_element_type == bool:
+                    return FEATURE_TYPES["bool"], value
+                elif first_element_type == int:
+                    return FEATURE_TYPES["int"], value
+                elif first_element_type == float:
+                    return FEATURE_TYPES["float"], value
+                elif first_element_type == str:
+                    return FEATURE_TYPES["str"] + "[ULabel]", value
+    if not mute:
+        logger.warning(f"cannot infer feature type of: {value}, returning '?")
+    return ("?", value)
+
+
+def __init__(self, host: Artifact | Collection):
+    self._host = host
+    self._feature_set_by_slot_ = None
+    self._accessor_by_registry_ = None
+
+
+def __repr__(self) -> str:
+    return print_features(self._host)  # type: ignore
+
+
+def get_values(self) -> dict[str, Any]:
+    """Get feature values as a dictionary."""
+    return print_features(self._host, to_dict=True)  # type: ignore
+
+
+def __getitem__(self, slot) -> QuerySet:
+    if slot not in self._feature_set_by_slot:
+        raise ValueError(
+            f"No linked feature set for slot: {slot}\nDid you get validation"
+            " warnings? Only features that match registered features get validated"
+            " and linked."
+        )
+    feature_set = self._feature_set_by_slot[slot]
+    orm_name = feature_set.registry
+    return getattr(feature_set, self._accessor_by_registry[orm_name]).all()
+
+
+@classmethod  # type: ignore
+def filter(cls, **expression) -> QuerySet:
+    """Filter features."""
+    keys_normalized = [key.split("__")[0] for key in expression]
+    validated = Feature.validate(keys_normalized, field="name", mute=True)
+    if sum(validated) != len(keys_normalized):
+        raise ValidationError(
+            f"Some keys in the filter expression are not registered as features: {np.array(keys_normalized)[~validated]}"
+        )
+    new_expression = {}
+    features = Feature.filter(name__in=keys_normalized).all().distinct()
+    for key, value in expression.items():
+        normalized_key = key.split("__")[0]
+        feature = features.get(name=normalized_key)
+        if not feature.dtype.startswith("cat"):
+            feature_value = FeatureValue.filter(feature=feature, value=value).one()
+            new_expression["feature_values"] = feature_value
+        else:
+            if isinstance(value, str):
+                label = ULabel.filter(name=value).one()
+                new_expression["ulabels"] = label
+            else:
+                raise NotImplementedError
+    if cls == FeatureManagerArtifact:
+        return Artifact.filter(**new_expression)
+    else:
+        return Collection.filter(**new_expression)
 
-    See :class:`~lamindb.core.Data` for more information.
-    """
 
-    def __init__(self, host: Artifact | Collection):
-        self._host = host
-        self._feature_set_by_slot = None
-        self._accessor_by_orm = None
+@property  # type: ignore
+def _feature_set_by_slot(self):
+    """Feature sets by slot."""
+    if self._feature_set_by_slot_ is None:
+        self._feature_set_by_slot_ = get_feature_set_by_slot_(self._host)
+    return self._feature_set_by_slot_
 
-    def __repr__(self) -> str:
-        if len(self.feature_set_by_slot) > 0:
-            return print_features(self._host)
-        else:
-            return "no linked features"
-
-    def __getitem__(self, slot) -> QuerySet:
-        if slot not in self.feature_set_by_slot:
-            raise ValueError(
-                f"No linked feature set for slot: {slot}\nDid you get validation"
-                " warnings? Only features that match registered features get validated"
-                " and linked."
-            )
-        feature_set = self.feature_set_by_slot[slot]
-        orm_name = feature_set.registry
-        if hasattr(feature_set, "_features"):
-            # feature set is not yet saved
-            # need to think about turning this into a queryset
-            return feature_set._features
+
+@property  # type: ignore
+def _accessor_by_registry(self):
+    """Accessor by ORM."""
+    if self._accessor_by_registry_ is None:
+        self._accessor_by_registry_ = get_accessor_by_registry_(self._host)
+    return self._accessor_by_registry_
+
+
+def add_values(
+    self,
+    values: dict[str, str | int | float | bool],
+    feature_field: FieldAttr = Feature.name,
+) -> None:
+    """Annotate artifact with features & values.
+
+    Args:
+        values: A dictionary of keys (features) & values (labels, numbers, booleans).
+        feature_field: The field of a reference registry to map keys of the
+            dictionary.
+    """
+    # rename to distinguish from the values inside the dict
+    features_values = values
+    keys = features_values.keys()
+    if isinstance(keys, DICT_KEYS_TYPE):
+        keys = list(keys)  # type: ignore
+    # deal with other cases later
+    assert all(isinstance(key, str) for key in keys)
+    registry = feature_field.field.model
+    validated = registry.validate(keys, field=feature_field, mute=True)
+    keys_array = np.array(keys)
+    validated_keys = keys_array[validated]
+    if validated.sum() != len(keys):
+        not_validated_keys = keys_array[~validated]
+        hint = "\n".join(
+            [
+                f"  ln.Feature(name='{key}', dtype='{infer_feature_type_convert_json(features_values[key])[0]}').save()"
+                for key in not_validated_keys
+            ]
+        )
+        msg = (
+            f"These keys could not be validated: {not_validated_keys.tolist()}\n"
+            f"If there are no typos, create features for them:\n\n{hint}"
+        )
+        raise ValidationError(msg)
+    registry.from_values(
+        validated_keys,
+        field=feature_field,
+    )
+    # figure out which of the values go where
+    features_labels = []
+    feature_values = []
+    not_validated_values = []
+    for key, value in features_values.items():
+        feature = Feature.filter(name=key).one()
+        inferred_type, converted_value = infer_feature_type_convert_json(
+            value, mute=True
+        )
+        if feature.dtype == "number":
+            if inferred_type not in {"int", "float"}:
+                raise TypeError(
+                    f"Value for feature '{key}' with type {feature.dtype} must be a number"
+                )
+        elif feature.dtype == "cat":
+            if not (inferred_type.startswith("cat") or isinstance(value, Registry)):
+                raise TypeError(
+                    f"Value for feature '{key}' with type '{feature.dtype}' must be a string or record."
+                )
+        elif feature.dtype == "bool":
+            assert isinstance(value, bool)
+        if not feature.dtype.startswith("cat"):
+            # can remove the query once we have the unique constraint
+            feature_value = FeatureValue.filter(
+                feature=feature, value=converted_value
+            ).one_or_none()
+            if feature_value is None:
+                feature_value = FeatureValue(feature=feature, value=converted_value)
+            feature_values.append(feature_value)
         else:
-            return getattr(feature_set, self.accessor_by_orm[orm_name]).all()
-
-    @property
-    def feature_set_by_slot(self):
-        """Feature sets by slot."""
-        if self._feature_set_by_slot is None:
-            self._feature_set_by_slot = get_feature_set_by_slot(self._host)
-        return self._feature_set_by_slot
-
-    @property
-    def accessor_by_orm(self):
-        """Accessor by ORM."""
-        if self._accessor_by_orm is None:
-            self._accessor_by_orm = get_accessor_by_orm(self._host)
-        return self._accessor_by_orm
-
-    def add(
-        self,
-        features_values: dict[str, str | int | float | bool],
-        slot: str | None = None,
-        feature_field: FieldAttr = Feature.name,
-    ):
-        """Add features stratified by slot.
-
-        Args:
-            features_values: A dictionary of features & values. You can also
-              pass `{feature_identifier: None}` to skip annotation by values.
-            slot: The access slot of the feature sets in the artifact. For
-              instance, `.columns` for `DataFrame` or `.var` or `.obs` for
-              `AnnData`.
-            feature_field: The field of a reference registry to map values.
-        """
-        if slot is None:
-            slot = "external"
-        keys = features_values.keys()
-        features_values.values()
-        # what if the feature is already part of a linked feature set?
-        # what if artifact annotation by features through link tables and through feature sets
-        # differs?
-        feature_set = FeatureSet.from_values(keys, field=feature_field)
-        self._host.features.add_feature_set(feature_set, slot)
-        # now figure out which of the values go where
-        features_labels = []
-        feature_values = []
-        for key, value in features_values.items():
-            # TODO: use proper field in .get() below
-            feature = feature_set.features.get(name=key)
-            if feature.dtype == "number":
-                if not (isinstance(value, int) or isinstance(value, float)):
-                    raise TypeError(
-                        f"Value for feature '{key}' with type {feature.dtype} must be a number"
-                    )
-            elif feature.dtype == "cat":
-                if not (isinstance(value, str) or isinstance(value, Registry)):
-                    raise TypeError(
-                        f"Value for feature '{key}' with type '{feature.dtype}' must be a string or record."
-                    )
-            elif feature.dtype == "bool":
-                assert isinstance(value, bool)
-            if feature.dtype == "cat":
-                if isinstance(value, Registry):
-                    assert not value._state.adding
-                    label_record = value
-                    assert isinstance(label_record, ULabel)
-                else:
-                    label_record = ULabel.filter(name=value).one_or_none()
-                    if label_record is None:
-                        raise ValidationError(f"Label '{value}' not found in ln.ULabel")
+            if isinstance(value, Registry):
+                assert not value._state.adding
+                label_record = value
+                assert isinstance(label_record, ULabel)
                 features_labels.append((feature, label_record))
             else:
-                feature_values.append(FeatureValue(feature=feature, value=value))
-        # bulk add all links to ArtifactULabel
-        if features_labels:
-            LinkORM = self._host.ulabels.through
-            links = [
-                LinkORM(
-                    artifact_id=self._host.id, feature_id=feature.id, ulabel_id=label.id
-                )
-                for (feature, label) in features_labels
-            ]
-            LinkORM.objects.bulk_create(links, ignore_conflicts=True)
-        if feature_values:
-            save(feature_values)
-            LinkORM = self._host.feature_values.through
-            links = [
-                LinkORM(artifact_id=self._host.id, featurevalue_id=feature_value.id)
-                for feature_value in feature_values
-            ]
-            LinkORM.objects.bulk_create(links)
-
-    def add_from_df(self, field: FieldAttr = Feature.name, organism: str | None = None):
-        """Add features from DataFrame."""
-        if isinstance(self._host, Artifact):
-            assert self._host.accessor == "DataFrame"
-        else:
-            # Collection
-            assert self._host.artifact.accessor == "DataFrame"
-
-        # parse and register features
-        registry = field.field.model
-        df = self._host.load()
-        features = registry.from_values(df.columns, field=field, organism=organism)
-        if len(features) == 0:
-            logger.error(
-                "no validated features found in DataFrame! please register features first!"
+                if isinstance(value, str):
+                    values = [value]  # type: ignore
+                else:
+                    values = value  # type: ignore
+                if "ULabel" not in feature.dtype:
+                    feature.dtype += "[ULabel]"
+                    feature.save()
+                validated = ULabel.validate(values, field="name", mute=True)
+                values_array = np.array(values)
+                validated_values = values_array[validated]
+                if validated.sum() != len(values):
+                    not_validated_values += values_array[~validated].tolist()
+                label_records = ULabel.from_values(validated_values, field="name")
+                features_labels += [
+                    (feature, label_record) for label_record in label_records
+                ]
+    if not_validated_values:
+        hint = (
+            f"  ulabels = ln.ULabel.from_values({not_validated_values}, create=True)\n"
+            f"  ln.save(ulabels)"
+        )
+        msg = (
+            f"These values could not be validated: {not_validated_values}\n"
+            f"If there are no typos, create ulabels for them:\n\n{hint}"
+        )
+        raise ValidationError(msg)
+    # bulk add all links to ArtifactULabel
+    if features_labels:
+        LinkORM = self._host.ulabels.through
+        links = [
+            LinkORM(
+                artifact_id=self._host.id, feature_id=feature.id, ulabel_id=label.id
             )
-            return
-
-        # create and link feature sets
-        feature_set = FeatureSet(features=features)
-        feature_sets = {"columns": feature_set}
-        self._host._feature_sets = feature_sets
-        self._host.save()
-
-    def add_from_anndata(
-        self,
-        var_field: FieldAttr,
-        obs_field: FieldAttr | None = Feature.name,
-        mute: bool = False,
-        organism: str | Registry | None = None,
-    ):
-        """Add features from AnnData."""
-        if isinstance(self._host, Artifact):
-            assert self._host.accessor == "AnnData"
-        else:
-            raise NotImplementedError()
-
-        # parse and register features
-        adata = self._host.load()
-        feature_sets = parse_feature_sets_from_anndata(
-            adata,
-            var_field=var_field,
-            obs_field=obs_field,
+            for (feature, label) in features_labels
+        ]
+        # a link might already exist
+        try:
+            save(links, ignore_conflicts=False)
+        except Exception:
+            save(links, ignore_conflicts=True)
+            # now deal with links that were previously saved without a feature_id
+            saved_links = LinkORM.filter(
+                artifact_id=self._host.id,
+                ulabel_id__in=[l.id for _, l in features_labels],
+            )
+            for link in saved_links.all():
+                # TODO: also check for inconsistent features
+                if link.feature_id is None:
+                    link.feature_id = [
+                        f.id for f, l in features_labels if l.id == link.ulabel_id
+                    ][0]
+                    link.save()
+    if feature_values:
+        save(feature_values)
+        LinkORM = self._host.feature_values.through
+        links = [
+            LinkORM(artifact_id=self._host.id, featurevalue_id=feature_value.id)
+            for feature_value in feature_values
+        ]
+        # a link might already exist, to avoid raising a unique constraint
+        # error, ignore_conflicts
+        save(links, ignore_conflicts=True)
+
+
+def add_feature_set(self, feature_set: FeatureSet, slot: str) -> None:
+    """Annotate artifact with a feature set.
+
+    Args:
+        feature_set: `FeatureSet` A feature set record.
+        slot: `str` The slot that marks where the feature set is stored in
+            the artifact.
+    """
+    if self._host._state.adding:
+        raise ValueError(
+            "Please save the artifact or collection before adding a feature set!"
+        )
+    host_db = self._host._state.db
+    feature_set.save(using=host_db)
+    host_id_field = get_host_id_field(self._host)
+    kwargs = {
+        host_id_field: self._host.id,
+        "featureset": feature_set,
+        "slot": slot,
+    }
+    link_record = (
+        self._host.feature_sets.through.objects.using(host_db)
+        .filter(**kwargs)
+        .one_or_none()
+    )
+    if link_record is None:
+        self._host.feature_sets.through(**kwargs).save(using=host_db)
+        if slot in self._feature_set_by_slot:
+            logger.debug(f"replaced existing {slot} feature set")
+        self._feature_set_by_slot_[slot] = feature_set  # type: ignore
+
+
+def _add_set_from_df(
+    self, field: FieldAttr = Feature.name, organism: str | None = None
+):
+    """Add feature set corresponding to column names of DataFrame."""
+    if isinstance(self._host, Artifact):
+        assert self._host.accessor == "DataFrame"
+    else:
+        # Collection
+        assert self._host.artifact.accessor == "DataFrame"
+
+    # parse and register features
+    registry = field.field.model
+    df = self._host.load()
+    features = registry.from_values(df.columns, field=field, organism=organism)
+    if len(features) == 0:
+        logger.error(
+            "no validated features found in DataFrame! please register features first!"
+        )
+        return
+
+    # create and link feature sets
+    feature_set = FeatureSet(features=features)
+    feature_sets = {"columns": feature_set}
+    self._host._feature_sets = feature_sets
+    self._host.save()
+
+
+def _add_set_from_anndata(
+    self,
+    var_field: FieldAttr,
+    obs_field: FieldAttr | None = Feature.name,
+    mute: bool = False,
+    organism: str | Registry | None = None,
+):
+    """Add features from AnnData."""
+    if isinstance(self._host, Artifact):
+        assert self._host.accessor == "AnnData"
+    else:
+        raise NotImplementedError()
+
+    # parse and register features
+    adata = self._host.load()
+    feature_sets = parse_feature_sets_from_anndata(
+        adata,
+        var_field=var_field,
+        obs_field=obs_field,
+        mute=mute,
+        organism=organism,
+    )
+
+    # link feature sets
+    self._host._feature_sets = feature_sets
+    self._host.save()
+
+
+def _add_set_from_mudata(
+    self,
+    var_fields: dict[str, FieldAttr],
+    obs_fields: dict[str, FieldAttr] = None,
+    mute: bool = False,
+    organism: str | Registry | None = None,
+):
+    """Add features from MuData."""
+    if obs_fields is None:
+        obs_fields = {}
+    if isinstance(self._host, Artifact):
+        assert self._host.accessor == "MuData"
+    else:
+        raise NotImplementedError()
+
+    # parse and register features
+    mdata = self._host.load()
+    feature_sets = {}
+    obs_features = features = Feature.from_values(mdata.obs.columns)
+    if len(obs_features) > 0:
+        feature_sets["obs"] = FeatureSet(features=features)
+    for modality, field in var_fields.items():
+        modality_fs = parse_feature_sets_from_anndata(
+            mdata[modality],
+            var_field=field,
+            obs_field=obs_fields.get(modality, Feature.name),
             mute=mute,
             organism=organism,
         )
-
-        # link feature sets
-        self._host._feature_sets = feature_sets
-        self._host.save()
-
-    def add_from_mudata(
-        self,
-        var_fields: dict[str, FieldAttr],
-        obs_fields: dict[str, FieldAttr] = None,
-        mute: bool = False,
-        organism: str | Registry | None = None,
-    ):
-        """Add features from MuData."""
-        if obs_fields is None:
-            obs_fields = {}
-        if isinstance(self._host, Artifact):
-            assert self._host.accessor == "MuData"
-        else:
-            raise NotImplementedError()
-
-        # parse and register features
-        mdata = self._host.load()
-        feature_sets = {}
-        obs_features = features = Feature.from_values(mdata.obs.columns)
-        if len(obs_features) > 0:
-            feature_sets["obs"] = FeatureSet(features=features)
-        for modality, field in var_fields.items():
-            modality_fs = parse_feature_sets_from_anndata(
-                mdata[modality],
-                var_field=field,
-                obs_field=obs_fields.get(modality, Feature.name),
-                mute=mute,
-                organism=organism,
-            )
-            for k, v in modality_fs.items():
-                feature_sets[f"['{modality}'].{k}"] = v
-
-        # link feature sets
-        self._host._feature_sets = feature_sets
-        self._host.save()
-
-    def add_feature_set(self, feature_set: FeatureSet, slot: str):
-        """Add new feature set to a slot.
-
-        Args:
-            feature_set: `FeatureSet` A feature set object.
-            slot: `str` The access slot.
-        """
-        if self._host._state.adding:
-            raise ValueError(
-                "Please save the artifact or collection before adding a feature set!"
-            )
-        host_db = self._host._state.db
-        feature_set.save(using=host_db)
-        host_id_field = get_host_id_field(self._host)
-        kwargs = {
-            host_id_field: self._host.id,
-            "featureset": feature_set,
-            "slot": slot,
-        }
-        link_record = (
-            self._host.feature_sets.through.objects.using(host_db)
-            .filter(**kwargs)
-            .one_or_none()
+        for k, v in modality_fs.items():
+            feature_sets[f"['{modality}'].{k}"] = v
+
+    # link feature sets
+    self._host._feature_sets = feature_sets
+    self._host.save()
+
+
+def _add_from(self, data: HasFeatures, parents: bool = True):
+    """Transfer features from a artifact or collection."""
+    # This only covers feature sets, though.
+    using_key = settings._using_key
+    for slot, feature_set in data.features._feature_set_by_slot.items():
+        members = feature_set.members
+        if len(members) == 0:
+            continue
+        registry = members[0].__class__
+        # note here the features are transferred based on an unique field
+        field = REGISTRY_UNIQUE_FIELD.get(registry.__name__.lower(), "uid")
+        # TODO: get a default ID field for the registry
+        if hasattr(registry, "ontology_id") and parents:
+            field = "ontology_id"
+        elif hasattr(registry, "ensembl_gene_id"):
+            field = "ensembl_gene_id"
+        elif hasattr(registry, "uniprotkb_id"):
+            field = "uniprotkb_id"
+
+        if registry.__get_name_with_schema__() == "bionty.Organism":
+            parents = False
+        # this will be e.g. be a list of ontology_ids or uids
+        member_uids = list(members.values_list(field, flat=True))
+        # create records from ontology_id in order to populate parents
+        if field == "ontology_id" and len(member_uids) > 0 and parents:
+            # create from bionty
+            records = registry.from_values(member_uids, field=field)
+            if len(records) > 0:
+                save(records, parents=parents)
+        validated = registry.validate(member_uids, field=field, mute=True)
+        new_members_uids = list(compress(member_uids, ~validated))
+        new_members = members.filter(**{f"{field}__in": new_members_uids}).all()
+        n_new_members = len(new_members)
+        if n_new_members > 0:
+            mute = True if n_new_members > 10 else False
+            # transfer foreign keys needs to be run before transfer to default db
+            transfer_fk_to_default_db_bulk(new_members, using_key)
+            for feature in new_members:
+                # not calling save=True here as in labels, because want to
+                # bulk save below
+                # transfer_fk is set to False because they are already transferred
+                # in the previous step transfer_fk_to_default_db_bulk
+                transfer_to_default_db(feature, using_key, mute=mute, transfer_fk=False)
+            logger.info(f"saving {n_new_members} new {registry.__name__} records")
+            save(new_members, parents=parents)
+
+        # create a new feature set from feature values using the same uid
+        feature_set_self = FeatureSet.from_values(
+            member_uids, field=getattr(registry, field)
         )
-        if link_record is None:
-            self._host.feature_sets.through(**kwargs).save(using=host_db)
-            if slot in self.feature_set_by_slot:
-                logger.debug(f"replaced existing {slot} feature set")
-            # this _feature_set_by_slot here is private
-            self._feature_set_by_slot[slot] = feature_set  # type: ignore
-
-    def _add_from(self, data: Data, parents: bool = True):
-        """Transfer features from a artifact or collection."""
-        using_key = settings._using_key
-        for slot, feature_set in data.features.feature_set_by_slot.items():
-            print(slot)
-            members = feature_set.members
-            if len(members) == 0:
-                continue
-            registry = members[0].__class__
-            # note here the features are transferred based on an unique field
-            field = REGISTRY_UNIQUE_FIELD.get(registry.__name__.lower(), "uid")
-            # TODO: get a default ID field for the registry
-            if hasattr(registry, "ontology_id") and parents:
-                field = "ontology_id"
-            elif hasattr(registry, "ensembl_gene_id"):
-                field = "ensembl_gene_id"
-            elif hasattr(registry, "uniprotkb_id"):
-                field = "uniprotkb_id"
-
-            if registry.__get_name_with_schema__() == "bionty.Organism":
-                parents = False
-            # this will be e.g. be a list of ontology_ids or uids
-            member_uids = list(members.values_list(field, flat=True))
-            # create records from ontology_id in order to populate parents
-            if field == "ontology_id" and len(member_uids) > 0 and parents:
-                # create from bionty
-                records = registry.from_values(member_uids, field=field)
-                if len(records) > 0:
-                    save(records, parents=parents)
-            validated = registry.validate(member_uids, field=field, mute=True)
-            new_members_uids = list(compress(member_uids, ~validated))
-            new_members = members.filter(**{f"{field}__in": new_members_uids}).all()
-            n_new_members = len(new_members)
-            if n_new_members > 0:
-                mute = True if n_new_members > 10 else False
-                # transfer foreign keys needs to be run before transfer to default db
-                transfer_fk_to_default_db_bulk(new_members, using_key)
-                for feature in new_members:
-                    # not calling save=True here as in labels, because want to
-                    # bulk save below
-                    # transfer_fk is set to False because they are already transferred
-                    # in the previous step transfer_fk_to_default_db_bulk
-                    transfer_to_default_db(
-                        feature, using_key, mute=mute, transfer_fk=False
-                    )
-                logger.info(f"saving {n_new_members} new {registry.__name__} records")
-                save(new_members, parents=parents)
-
-            # create a new feature set from feature values using the same uid
-            feature_set_self = FeatureSet.from_values(
-                member_uids, field=getattr(registry, field)
-            )
-            if feature_set_self is None:
-                if hasattr(registry, "organism"):
-                    logger.warning(
-                        f"FeatureSet is not transferred, check if organism is set correctly: {feature_set}"
-                    )
-                continue
-            # make sure the uid matches if featureset is composed of same features
-            if feature_set_self.hash == feature_set.hash:
-                feature_set_self.uid = feature_set.uid
-            logger.info(f"saving {slot} featureset: {feature_set_self}")
-            self._host.features.add_feature_set(feature_set_self, slot)
+        if feature_set_self is None:
+            if hasattr(registry, "organism"):
+                logger.warning(
+                    f"FeatureSet is not transferred, check if organism is set correctly: {feature_set}"
+                )
+            continue
+        # make sure the uid matches if featureset is composed of same features
+        if feature_set_self.hash == feature_set.hash:
+            feature_set_self.uid = feature_set.uid
+        logger.info(f"saving {slot} featureset: {feature_set_self}")
+        self._host.features.add_feature_set(feature_set_self, slot)
+
+
+FeatureManager.__init__ = __init__
+FeatureManager.__repr__ = __repr__
+FeatureManager.__getitem__ = __getitem__
+FeatureManager.get_values = get_values
+FeatureManager._feature_set_by_slot = _feature_set_by_slot
+FeatureManager._accessor_by_registry = _accessor_by_registry
+FeatureManager.add_values = add_values
+FeatureManager.add_feature_set = add_feature_set
+FeatureManager._add_set_from_df = _add_set_from_df
+FeatureManager._add_set_from_anndata = _add_set_from_anndata
+FeatureManager._add_set_from_mudata = _add_set_from_mudata
+FeatureManager._add_from = _add_from
+FeatureManager.filter = filter