lamindb 0.71.2__py3-none-any.whl → 0.72.0__py3-none-any.whl

lamindb/core/_feature_manager.py

@@ -1,13 +1,21 @@
 from __future__ import annotations
 
 from itertools import compress
-from typing import TYPE_CHECKING, Iterable, Optional
+from typing import TYPE_CHECKING, Iterable
 
 import anndata as ad
 from anndata import AnnData
 from lamin_utils import colors, logger
 from lamindb_setup.core.upath import create_path
-from lnschema_core.models import Artifact, Collection, Data, Feature, Registry
+from lnschema_core.models import (
+    Artifact,
+    Collection,
+    Data,
+    Feature,
+    FeatureValue,
+    Registry,
+    ULabel,
+)
 
 from lamindb._feature import convert_numpy_dtype_to_lamin_feature_type
 from lamindb._feature_set import FeatureSet
@@ -18,6 +26,7 @@ from lamindb._registry import (
     transfer_to_default_db,
 )
 from lamindb._save import save
+from lamindb.core.exceptions import ValidationError
 from lamindb.core.storage import LocalPathClasses
 
 from ._settings import settings
@@ -41,8 +50,8 @@ def get_accessor_by_orm(host: Artifact | Collection) -> dict:
         field.related_model.__get_name_with_schema__(): field.name
         for field in host._meta.related_objects
     }
-    dictionary["core.Feature"] = "features"
-    dictionary["core.ULabel"] = "ulabels"
+    dictionary["Feature"] = "features"
+    dictionary["ULabel"] = "ulabels"
     return dictionary
 
 
@@ -57,15 +66,12 @@ def get_feature_set_by_slot(host) -> dict:
     host_id_field = get_host_id_field(host)
     kwargs = {host_id_field: host.id}
     # otherwise, we need a query
-    feature_set_links = host.feature_sets.through.objects.using(host_db).filter(
-        **kwargs
+    feature_set_links = (
+        host.feature_sets.through.objects.using(host_db)
+        .filter(**kwargs)
+        .select_related("featureset")
     )
-    return {
-        feature_set_link.slot: FeatureSet.objects.using(host_db).get(
-            id=feature_set_link.feature_set_id
-        )
-        for feature_set_link in feature_set_links
-    }
+    return {fsl.slot: fsl.featureset for fsl in feature_set_links}
 
 
 def get_label_links(
@@ -74,7 +80,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_orm[registry])
         .through.objects.using(host._state.db)
         .filter(**kwargs)
     )
@@ -93,48 +99,48 @@ def print_features(self: Data) -> str:
 
     from ._data import format_repr
 
-    msg = ""
-    features_lookup = Feature.objects.using(self._state.db).lookup().dict()
-    for slot, feature_set in self.features._feature_set_by_slot.items():
-        if feature_set.registry != "core.Feature":
+    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 = [getattr(feature, name_field) for feature in features]
-            msg += (
+            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)
-            msg += f"    {print_values}\n"
+            messages.append(f"    {print_values}\n")
         else:
-            df_slot = feature_set.features.df()
-            msg += (
+            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"
             )
-            for _, row in df_slot.iterrows():
-                if row["type"] == "category" and row["registries"] is not None:
-                    labels = self.labels.get(
-                        features_lookup.get(row["name"]), mute=True
-                    )
+            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):
-                        msg += f"    🔗 {row['name']} ({row.registries})\n"
+                        messages.append(f"    🔗 {name} ({dtype})\n")
                         indent = "    "
                     else:
-                        labels = {row["registries"]: labels}
-                    for registry, labels in labels.items():  # noqa: B020
-                        count_str = f"{len(labels)}, {colors.italic(f'{registry}')}"
-                        field = get_default_str_field(labels)
-                        print_values = _print_values(labels.list(field), n=10)
+                        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}    🔗 {row['name']} ({count_str}):"
-                            f" {print_values}\n"
+                            f"{indent}    🔗 {name} ({count_str}):" f" {print_values}\n"
                         )
-                        msg += msg_objects
+                        messages.append(msg_objects)
                 else:
-                    msg += f"    {row['name']} ({row['type']})\n"
-    if msg != "":
-        msg = f"{colors.green('Features')}:\n" + msg
-    return msg
+                    messages.append(f"    {name} ({dtype})\n")
+    if messages:
+        messages.insert(0, f"{colors.green('Features')}:\n")
+    return "".join(messages)
 
 
 def parse_feature_sets_from_anndata(
@@ -171,6 +177,7 @@ def parse_feature_sets_from_anndata(
             type=type,
             mute=mute,
             organism=organism,
+            raise_validation_error=False,
         )
         if feature_set_var is not None:
             feature_sets["var"] = feature_set_var
@@ -204,39 +211,118 @@ class FeatureManager:
 
     def __init__(self, host: Artifact | Collection):
         self._host = host
-        self._feature_set_by_slot = get_feature_set_by_slot(host)
-        self._accessor_by_orm = get_accessor_by_orm(host)
+        self._feature_set_by_slot = None
+        self._accessor_by_orm = None
 
     def __repr__(self) -> str:
-        if len(self._feature_set_by_slot) > 0:
+        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:
+        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]
+        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
         else:
-            return getattr(feature_set, self._accessor_by_orm[orm_name]).all()
+            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.
 
-    def add(self, features: Iterable[Registry], slot: str | None = None):
-        """Add features stratified by slot."""
-        if (hasattr(self._host, "accessor") and self._host.accessor == "DataFrame") or (
-            hasattr(self._host, "artifact")
-            and self._host.artifact.accessor == "DataFrame"
-        ):
-            slot = "columns" if slot is None else slot
-        self.add_feature_set(feature_set=FeatureSet(features=features), slot=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")
+                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."""
@@ -341,7 +427,7 @@ class FeatureManager:
         host_id_field = get_host_id_field(self._host)
         kwargs = {
             host_id_field: self._host.id,
-            "feature_set": feature_set,
+            "featureset": feature_set,
             "slot": slot,
         }
         link_record = (
@@ -351,26 +437,36 @@ class FeatureManager:
         )
         if link_record is None:
             self._host.feature_sets.through(**kwargs).save(using=host_db)
-            self._feature_set_by_slot[slot] = feature_set
+            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():
+        for slot, feature_set in data.features.feature_set_by_slot.items():
+            print(slot)
             members = feature_set.members
-            if members.count() == 0:
+            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:
+            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:
@@ -378,8 +474,9 @@ class FeatureManager:
             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()
-            if new_members.count() > 0:
-                mute = True if new_members.count() > 10 else False
+            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:
@@ -390,9 +487,7 @@ class FeatureManager:
                     transfer_to_default_db(
                         feature, using_key, mute=mute, transfer_fk=False
                     )
-                logger.info(
-                    f"saving {new_members.count()} new {registry.__name__} records"
-                )
+                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

lamindb/core/_label_manager.py

@@ -4,9 +4,8 @@ from typing import TYPE_CHECKING, Dict
 
 import numpy as np
 from lamin_utils import colors, logger
-from lnschema_core.models import Artifact, Collection, Data, Feature, Registry
+from lnschema_core.models import Artifact, Collection, Data, Feature, LinkORM, Registry
 
-from lamindb._feature_set import dict_related_model_to_related_name
 from lamindb._from_values import _print_values
 from lamindb._registry import (
     REGISTRY_UNIQUE_FIELD,
@@ -17,14 +16,17 @@ from lamindb._registry import (
 from lamindb._save import save
 
 from ._settings import settings
+from .schema import dict_related_model_to_related_name
 
 if TYPE_CHECKING:
     from lamindb._query_set import QuerySet
 
 
 def get_labels_as_dict(self: Data):
-    labels = {}
-    for related_model, related_name in dict_related_model_to_related_name(
+    labels = {}  # type: ignore
+    if self.id is None:
+        return labels
+    for related_model_name, related_name in dict_related_model_to_related_name(
         self.__class__
     ).items():
         if related_name in {
@@ -37,19 +39,23 @@ def get_labels_as_dict(self: Data):
             "environment_of",
         }:
             continue
-        if self.id is not None:
-            labels[related_name] = (related_model, self.__getattribute__(related_name))
+        labels[related_name] = (related_model_name, self.__getattribute__(related_name))
     return labels
 
 
-def print_labels(self: Data):
+def print_labels(
+    self: Data, field: str = "name", ignore_labels_with_feature: bool = True
+):
     labels_msg = ""
     for related_name, (related_model, labels) in get_labels_as_dict(self).items():
-        if labels.exists():
-            n = labels.count()
-            field = get_default_str_field(labels)
-            print_values = _print_values(labels.list(field), n=10)
-            labels_msg += f"  📎 {related_name} ({n}, {colors.italic(related_model)}): {print_values}\n"
+        try:
+            labels_list = list(labels.values_list(field, flat=True))
+            if len(labels_list) > 0:
+                get_default_str_field(labels)
+                print_values = _print_values(labels_list[:20], n=10)
+                labels_msg += f"  📎 {related_name} ({len(labels_list)}, {colors.italic(related_model)}): {print_values}\n"
+        except Exception:
+            continue
     if len(labels_msg) > 0:
         return f"{colors.green('Labels')}:\n{labels_msg}"
     else:
@@ -72,7 +78,7 @@ def transfer_add_labels(labels, features_lookup_self, self, row, parents: bool =
         # link labels records from self db
         self._host.labels.add(
             validated_labels + new_labels,
-            feature=getattr(features_lookup_self, row["name"]),
+            feature=features_lookup_self.get(row["name"]),
         )
 
     # validate labels on the default db
@@ -94,6 +100,10 @@ def validate_labels(labels: QuerySet | list | dict, parents: bool = True):
         field = REGISTRY_UNIQUE_FIELD.get(registry.__name__.lower(), "uid")
         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
         # if the field value is None, use uid field
@@ -195,42 +205,47 @@ class LabelManager:
             >>> file1.ulabels.set(labels)
             >>> file2.labels.add_from(file1)
         """
-        features_lookup_self = Feature.lookup()
-        features_lookup_data = Feature.objects.using(data._state.db).lookup()
-        for _, feature_set in data.features._feature_set_by_slot.items():
+        from django.db.utils import ProgrammingError
+
+        features_lookup_self = {f.name: f for f in Feature.objects.filter().all()}
+        features_lookup_data = {
+            f.name: f for f in Feature.objects.using(data._state.db).filter().all()
+        }
+        for _, feature_set in data.features.feature_set_by_slot.items():
             # add labels stratified by feature
-            if feature_set.registry == "core.Feature":
-                # df_slot is the Feature table with type and registries
+            if feature_set.registry == "Feature":
+                # df_slot is the Feature table with type
                 df_slot = feature_set.features.df()
                 for _, row in df_slot.iterrows():
-                    if row["type"] == "category" and row["registries"] is not None:
+                    if row["dtype"].startswith("cat["):
                         logger.info(f"transferring {row['name']}")
                         # labels records from data db
                         labels = data.labels.get(
-                            getattr(features_lookup_data, row["name"]), mute=True
+                            features_lookup_data.get(row["name"]), mute=True
                         )
                         transfer_add_labels(
                             labels, features_lookup_self, self, row, parents=parents
                         )
-
-        # for now, have this be duplicated, need to disentangle above
+        # TODO: for now, has to be duplicated
         using_key = settings._using_key
         for related_name, (_, labels) in get_labels_as_dict(data).items():
             labels = labels.all()
-            if len(labels) == 0:
+            try:
+                if len(labels) == 0:
+                    continue
+                validated_labels, new_labels = validate_labels(labels, parents=parents)
+                if len(new_labels) > 0:
+                    transfer_fk_to_default_db_bulk(new_labels, using_key)
+                    for label in new_labels:
+                        transfer_to_default_db(
+                            label, using_key, mute=True, transfer_fk=False
+                        )
+                    save(new_labels, parents=parents)
+                # this should not occur as file and collection should have the same attributes
+                # but this might not be true for custom schema
+                labels_list = validated_labels + new_labels
+                if hasattr(self._host, related_name):
+                    getattr(self._host, related_name).add(*labels_list)
+            # ProgrammingError is raised when schemas don't match between source and target instances
+            except ProgrammingError:
                 continue
-            validated_labels, new_labels = validate_labels(
-                labels.all(), parents=parents
-            )
-            if len(new_labels) > 0:
-                transfer_fk_to_default_db_bulk(new_labels, using_key)
-                for label in new_labels:
-                    transfer_to_default_db(
-                        label, using_key, mute=True, transfer_fk=False
-                    )
-                save(new_labels, parents=parents)
-            # this should not occur as file and collection should have the same attributes
-            # but this might not be true for custom schema
-            labels_list = validated_labels + new_labels
-            if hasattr(self._host, related_name):
-                getattr(self._host, related_name).add(*labels_list)

lamindb/core/_run_context.py

@@ -10,9 +10,11 @@ from typing import TYPE_CHECKING
 from lamin_utils import logger
 from lamindb_setup.core.hashing import hash_file
 from lnschema_core import Run, Transform, ids
+from lnschema_core.models import Param, ParamValue, RunParamValue
 from lnschema_core.types import TransformType
 from lnschema_core.users import current_user_id
 
+from lamindb._save import save
 from lamindb.core._transform_settings import transform as transform_settings
 
 from ._settings import settings
@@ -187,6 +189,26 @@ def pretty_pypackages(dependencies: dict) -> str:
     return " ".join(deps_list)
 
 
+def parse_and_link_params(run: Run, params: dict) -> None:
+    param_values = []
+    for key, value in params.items():
+        param = Param.filter(name=key).one_or_none()
+        if param is None:
+            dtype = type(value).__name__
+            logger.warning(
+                f"param '{key}' does not yet exist, creating it with dtype '{dtype}'"
+            )
+            param = Param(name=key, dtype=dtype).save()
+        param_value, _ = ParamValue.objects.get_or_create(param=param, value=value)
+        param_values.append(param_value)
+    if param_values:
+        links = [
+            RunParamValue(run_id=run.id, paramvalue_id=param_value.id)
+            for param_value in param_values
+        ]
+        RunParamValue.objects.bulk_create(links)
+
+
 class run_context:
     """Global run context."""
 
@@ -312,7 +334,6 @@ class run_context:
             )
             if run is not None:  # loaded latest run
                 run.started_at = datetime.now(timezone.utc)  # update run time
-                run.json = params  # update run params
                 logger.important(f"loaded: {run}")
 
         if run is None:  # create new run
@@ -326,6 +347,8 @@ class run_context:
         run.is_consecutive = True if is_run_from_ipython else None
         # need to save in all cases
         run.save()
+        if params is not None:
+            parse_and_link_params(run, params)
         cls.run = run
 
         from ._track_environment import track_environment

lamindb/core/datasets/_core.py

@@ -45,10 +45,8 @@ def file_fcs_alpert19(populate_registries: bool = False) -> Path:  # pragma: no
                 bt.CellMarker.public().inspect(std, "name").validated, "name"
             )
         )
-        ln.Feature(
-            name="assay", type="category", registries=[bt.ExperimentalFactor]
-        ).save()
-        ln.Feature(name="organism", type="category", registries=[bt.Organism]).save()
+        ln.Feature(name="assay", dtype=[bt.ExperimentalFactor]).save()
+        ln.Feature(name="organism", dtype=[bt.Organism]).save()
         ln.settings.verbosity = verbosity
     return Path(filepath)
 
@@ -83,10 +81,8 @@ def file_tsv_rnaseq_nfcore_salmon_merged_gene_counts(
 
         verbosity = ln.settings.verbosity
         ln.settings.verbosity = "error"
-        ln.Feature(
-            name="assay", type="category", registries=[bt.ExperimentalFactor]
-        ).save()
-        ln.Feature(name="organism", type="category", registries=[bt.Organism]).save()
+        ln.Feature(name="assay", dtype=[bt.ExperimentalFactor]).save()
+        ln.Feature(name="organism", dtype=[bt.Organism]).save()
         bt.ExperimentalFactor.from_public(ontology_id="EFO:0008896").save()
         ln.settings.verbosity = verbosity
 
@@ -200,9 +196,7 @@ def anndata_mouse_sc_lymph_node(
         # cell types
         ln.save(bt.CellType.from_values(["CL:0000115", "CL:0000738"], "ontology_id"))
         # assays
-        ln.Feature(
-            name="assay", type="category", registries=[bt.ExperimentalFactor]
-        ).save()
+        ln.Feature(name="assay", dtype=[bt.ExperimentalFactor]).save()
         bt.ExperimentalFactor.from_public(ontology_id="EFO:0008913").save()
         # genes
         validated = bt.Gene.public(organism="mouse").validate(
@@ -329,13 +323,11 @@ def anndata_human_immune_cells(
         ln.save(bt.CellType.from_values(adata.obs.cell_type, field="name"))
         ln.save(bt.ExperimentalFactor.from_values(adata.obs.assay, field="name"))
         ln.save(bt.Tissue.from_values(adata.obs.tissue, field="name"))
-        ln.Feature(name="cell_type", type="category", registries=[bt.CellType]).save()
-        ln.Feature(
-            name="assay", type="category", registries=[bt.ExperimentalFactor]
-        ).save()
-        ln.Feature(name="tissue", type="category", registries=[bt.Tissue]).save()
-        ln.Feature(name="organism", type="category", registries=[bt.Organism]).save()
-        ln.Feature(name="donor", type="category", registries=[ln.ULabel]).save()
+        ln.Feature(name="cell_type", dtype=[bt.CellType]).save()
+        ln.Feature(name="assay", dtype=[bt.ExperimentalFactor]).save()
+        ln.Feature(name="tissue", dtype=[bt.Tissue]).save()
+        ln.Feature(name="organism", dtype=[bt.Organism]).save()
+        ln.Feature(name="donor", dtype=[ln.ULabel]).save()
         bt.ExperimentalFactor.from_public(ontology_id="EFO:0008913").save()
         ln.save([ln.ULabel(name=name) for name in adata.obs.donor.unique()])
         ln.settings.verbosity = verbosity