scdataloader 1.6.4__py3-none-any.whl → 1.7.0__py3-none-any.whl

scdataloader/preprocess.py

@@ -8,8 +8,9 @@ import pandas as pd
 import scanpy as sc
 from anndata import AnnData
 from scipy.sparse import csr_matrix
-
+from anndata import read_h5ad
 from scdataloader import utils as data_utils
+from upath import UPath
 
 FULL_LENGTH_ASSAYS = [
     "EFO: 0700016",
@@ -31,7 +32,7 @@ class Preprocessor:
         filter_gene_by_counts: Union[int, bool] = False,
         filter_cell_by_counts: Union[int, bool] = False,
         normalize_sum: float = 1e4,
-        subset_hvg: int = 0,
+        n_hvg_for_postp: int = 0,
         use_layer: Optional[str] = None,
         is_symbol: bool = False,
         hvg_flavor: str = "seurat_v3",
@@ -45,7 +46,13 @@ class Preprocessor:
         maxdropamount: int = 50,
         madoutlier: int = 5,
         pct_mt_outlier: int = 8,
-        batch_key: Optional[str] = None,
+        batch_keys: list[str] = [
+            "assay_ontology_term_id",
+            "self_reported_ethnicity_ontology_term_id",
+            "sex_ontology_term_id",
+            "donor_id",
+            "suspension_type",
+        ],
         skip_validate: bool = False,
         additional_preprocess: Optional[Callable[[AnnData], AnnData]] = None,
         additional_postprocess: Optional[Callable[[AnnData], AnnData]] = None,
@@ -65,7 +72,7 @@ class Preprocessor:
                 Defaults to 1e4.
             log1p (bool, optional): Determines whether to apply log1p transform to the normalized data.
                 Defaults to True.
-            subset_hvg (int or bool, optional): Determines whether to subset highly variable genes.
+            n_hvg_for_postp (int or bool, optional): Determines whether to subset to highly variable genes for the PCA.
                 Defaults to False.
             hvg_flavor (str, optional): Specifies the flavor of highly variable genes selection.
                 See :func:`scanpy.pp.highly_variable_genes` for more details. Defaults to "seurat_v3".
@@ -96,7 +103,6 @@ class Preprocessor:
         self.filter_gene_by_counts = filter_gene_by_counts
         self.filter_cell_by_counts = filter_cell_by_counts
         self.normalize_sum = normalize_sum
-        self.subset_hvg = subset_hvg
         self.hvg_flavor = hvg_flavor
         self.binning = binning
         self.organisms = organisms
@@ -109,8 +115,9 @@ class Preprocessor:
         self.min_nnz_genes = min_nnz_genes
         self.maxdropamount = maxdropamount
         self.madoutlier = madoutlier
+        self.n_hvg_for_postp = n_hvg_for_postp
         self.pct_mt_outlier = pct_mt_outlier
-        self.batch_key = batch_key
+        self.batch_keys = batch_keys
         self.length_normalize = length_normalize
         self.skip_validate = skip_validate
         self.use_layer = use_layer
@@ -118,7 +125,7 @@ class Preprocessor:
         self.do_postp = do_postp
         self.use_raw = use_raw
 
-    def __call__(self, adata) -> AnnData:
+    def __call__(self, adata, dataset_id=None) -> AnnData:
         if adata[0].obs.organism_ontology_term_id.iloc[0] not in self.organisms:
             raise ValueError(
                 "we cannot work with this organism",
@@ -144,10 +151,6 @@ class Preprocessor:
             del adata.obsm
         if len(adata.obsp.keys()) > 0 and self.do_postp:
             del adata.obsp
-        if len(adata.uns.keys()) > 0:
-            del adata.uns
-        if len(adata.varp.keys()) > 0:
-            del adata.varp
         # check that it is a count
         print("checking raw counts")
         if np.abs(
@@ -209,9 +212,9 @@ class Preprocessor:
             )
         )
 
-        if self.is_symbol or not adata.var.index.str.contains("ENSG").any():
-            if not adata.var.index.str.contains("ENSG").any():
-                print("No ENSG genes found, assuming gene symbols...")
+        if self.is_symbol or not adata.var.index.str.contains("ENS").any():
+            if not adata.var.index.str.contains("ENS").any():
+                print("No ENS genes found, assuming gene symbols...")
             genesdf["ensembl_gene_id"] = genesdf.index
             var = (
                 adata.var.merge(
@@ -243,9 +246,13 @@ class Preprocessor:
         adata = ad.concat([adata, emptyda], axis=1, join="outer", merge="only")
         # do a validation function
         adata.uns["unseen_genes"] = list(unseen)
+        if dataset_id is not None:
+            adata.uns["dataset_id"] = dataset_id
         if not self.skip_validate:
             print("validating")
-            data_utils.validate(adata, organism=adata.obs.organism_ontology_term_id[0])
+            data_utils.validate(
+                adata, organism=adata.obs.organism_ontology_term_id[0], need_all=False
+            )
             # length normalization
             if (
                 adata.obs["assay_ontology_term_id"].isin(FULL_LENGTH_ASSAYS).any()
@@ -310,38 +317,42 @@ class Preprocessor:
                 )["X"]
             )
             # step 5: subset hvg
-            if self.subset_hvg:
-                sc.pp.highly_variable_genes(
-                    adata,
-                    n_top_genes=self.subset_hvg,
-                    batch_key=self.batch_key,
-                    flavor=self.hvg_flavor,
-                    subset=False,
-                )
-            sc.pp.log1p(adata, layer="norm")
-            sc.pp.pca(
-                adata,
-                layer="norm",
-                n_comps=200 if adata.shape[0] > 200 else adata.shape[0] - 2,
-            )
-            sc.pp.neighbors(adata, use_rep="X_pca")
-            sc.tl.leiden(adata, key_added="leiden_2", resolution=2.0)
-            sc.tl.leiden(adata, key_added="leiden_1", resolution=1.0)
-            sc.tl.leiden(adata, key_added="leiden_0.5", resolution=0.5)
-            batches = [
-                "assay_ontology_term_id",
-                "self_reported_ethnicity_ontology_term_id",
-                "sex_ontology_term_id",
-                "development_stage_ontology_term_id",
-            ]
-            if "donor_id" in adata.obs.columns:
-                batches.append("donor_id")
-            if "suspension_type" in adata.obs.columns:
-                batches.append("suspension_type")
+            batches = []
+            for i in self.batch_keys:
+                if i in adata.obs.columns:
+                    batches.append(i)
             adata.obs["batches"] = adata.obs[batches].apply(
                 lambda x: ",".join(x.dropna().astype(str)), axis=1
             )
-            sc.tl.umap(adata)
+            if self.n_hvg_for_postp:
+                try:
+                    sc.pp.highly_variable_genes(
+                        adata,
+                        n_top_genes=self.n_hvg_for_postp,
+                        batch_key="batches",
+                        flavor=self.hvg_flavor,
+                        subset=False,
+                        layer="norm",
+                    )
+                except (ValueError, ZeroDivisionError) as e:
+                    print("retrying with span")
+                    sc.pp.highly_variable_genes(
+                        adata,
+                        n_top_genes=self.n_hvg_for_postp,
+                        # batch_key="batches",
+                        flavor=self.hvg_flavor,
+                        span=0.5,
+                        subset=False,
+                        layer="norm",
+                    )
+
+            adata.obsm["X_pca"] = sc.pp.pca(
+                adata.layers["norm"][:, adata.var.highly_variable]
+                if "highly_variable" in adata.var.columns
+                else adata.layers["norm"],
+                n_comps=200 if adata.shape[0] > 200 else adata.shape[0] - 2,
+            )
+
             # additional
             if self.additional_postprocess is not None:
                 adata = self.additional_postprocess(adata)
@@ -393,6 +404,7 @@ class Preprocessor:
             adata.layers[self.result_binned_key] = np.stack(binned_rows)
             adata.obsm["bin_edges"] = np.stack(bin_edges)
         print("done")
+        print(adata)
         return adata
 
 
@@ -401,22 +413,22 @@ class LaminPreprocessor(Preprocessor):
         self,
         *args,
         cache: bool = True,
-        stream: bool = False,
         keep_files: bool = True,
+        force_preloaded: bool = False,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.cache = cache
-        self.stream = stream
         self.keep_files = keep_files
+        self.force_preloaded = force_preloaded
 
     def __call__(
         self,
         data: Union[ln.Collection, AnnData] = None,
-        name="preprocessed dataset",
-        description="preprocessed dataset using scprint",
-        start_at=0,
-        version=2,
+        name: str = "preprocessed dataset",
+        description: str = "preprocessed dataset using scprint",
+        start_at: int = 0,
+        version: str = "2",
     ):
         """
         format controls the different input value wrapping, including categorical
@@ -437,12 +449,15 @@ class LaminPreprocessor(Preprocessor):
         elif isinstance(data, ln.Collection):
             for i, file in enumerate(data.artifacts.all()[start_at:]):
                 # use the counts matrix
-                print(i + start_at)
+                i = i + start_at
+                print(i)
                 if file.stem_uid in all_ready_processed_keys:
                     print(f"{file.stem_uid} is already processed... not preprocessing")
                     continue
                 print(file)
-                backed = file.open()
+
+                path = cache_path(file) if self.force_preloaded else file.cache()
+                backed = read_h5ad(path, backed="r")
                 if backed.obs.is_primary_data.sum() == 0:
                     print(f"{file.key} only contains non primary cells.. dropping")
                     # Save the stem_uid to a file to avoid loading it again
@@ -455,12 +470,11 @@ class LaminPreprocessor(Preprocessor):
                     )
                     continue
                 if file.size <= MAXFILESIZE:
-                    adata = file.load(stream=self.stream)
+                    adata = backed.to_memory()
                     print(adata)
                 else:
                     badata = backed
                     print(badata)
-
                 try:
                     if file.size > MAXFILESIZE:
                         print(
@@ -472,16 +486,26 @@ class LaminPreprocessor(Preprocessor):
                         )
                         print("num blocks ", num_blocks)
                         for j in range(num_blocks):
+                            if j == 0 and i == 390:
+                                continue
                             start_index = j * block_size
                             end_index = min((j + 1) * block_size, badata.shape[0])
                             block = badata[start_index:end_index].to_memory()
                             print(block)
-                            block = super().__call__(block)
-                            myfile = ln.from_anndata(
+                            block = super().__call__(
+                                block, dataset_id=file.stem_uid + "_p" + str(j)
+                            )
+                            myfile = ln.Artifact.from_anndata(
                                 block,
-                                revises=file,
-                                description=description,
-                                version=str(version) + "_s" + str(j),
+                                description=description
+                                + " n"
+                                + str(i)
+                                + " p"
+                                + str(j)
+                                + " ( revises file "
+                                + str(file.key)
+                                + " )",
+                                version=version,
                             )
                             myfile.save()
                             if self.keep_files:
@@ -491,16 +515,12 @@ class LaminPreprocessor(Preprocessor):
                                 del block
 
                     else:
-                        adata = super().__call__(adata)
-                        try:
-                            sc.pl.umap(adata, color=["cell_type"])
-                        except Exception:
-                            sc.pl.umap(adata, color=["cell_type_ontology_term_id"])
-                        myfile = ln.from_anndata(
+                        adata = super().__call__(adata, dataset_id=file.stem_uid)
+                        myfile = ln.Artifact.from_anndata(
                             adata,
                             revises=file,
-                            description=description,
-                            version=str(version),
+                            description=description + " p" + str(i),
+                            version=version,
                         )
                         myfile.save()
                         if self.keep_files:
@@ -672,35 +692,158 @@ def additional_preprocess(adata):
 
 
 def additional_postprocess(adata):
-    import palantir
+    # import palantir
 
     # define the "up to" 10 neighbors for each cells and add to obs
     # compute neighbors
     # need to be connectivities and same labels [cell type, assay, dataset, disease]
     # define the "neighbor" up to 10(N) cells and add to obs
     # define the "next time point" up to 5(M) cells and add to obs  # step 1: filter genes
-    del adata.obsp["connectivities"]
-    del adata.obsp["distances"]
-    sc.external.pp.harmony_integrate(adata, key="batches")
-    sc.pp.neighbors(adata, use_rep="X_pca_harmony")
+    # if len(adata.obs["batches"].unique()) > 1:
+    #    sc.external.pp.harmony_integrate(adata, key="batches")
+    #    sc.pp.neighbors(adata, use_rep="X_pca_harmony")
+    # else:
+    sc.pp.neighbors(adata, use_rep="X_pca")
+    sc.tl.leiden(adata, key_added="leiden_2", resolution=2.0)
+    sc.tl.leiden(adata, key_added="leiden_1", resolution=1.0)
+    sc.tl.leiden(adata, key_added="leiden_0.5", resolution=0.5)
     sc.tl.umap(adata)
+    mid = adata.uns["dataset_id"] if "dataset_id" in adata.uns else "unknown_id"
     sc.pl.umap(
         adata,
+        ncols=1,
         color=["cell_type", "batches"],
+        save="_" + mid + ".png",
     )
-    palantir.utils.run_diffusion_maps(adata, n_components=20)
-    palantir.utils.determine_multiscale_space(adata)
-    terminal_states = palantir.utils.find_terminal_states(
-        adata,
-        celltypes=adata.obs.cell_type_ontology_term_id.unique(),
-        celltype_column="cell_type_ontology_term_id",
+    COL = "cell_type_ontology_term_id"
+    NEWOBS = "clust_cell_type"
+    MINCELLS = 10
+    MAXSIM = 0.94
+    from collections import Counter
+
+    from .config import MAIN_HUMAN_MOUSE_DEV_STAGE_MAP
+
+    adata.obs[NEWOBS] = (
+        adata.obs[COL].astype(str) + "_" + adata.obs["leiden_1"].astype(str)
     )
-    sc.tl.diffmap(adata)
-    adata.obs["heat_diff"] = 1
-    for terminal_state in terminal_states.index.tolist():
-        adata.uns["iroot"] = np.where(adata.obs.index == terminal_state)[0][0]
-        sc.tl.dpt(adata)
-        adata.obs["heat_diff"] = np.minimum(
-            adata.obs["heat_diff"], adata.obs["dpt_pseudotime"]
-        )
+    coun = Counter(adata.obs[NEWOBS])
+    relab = {}
+    for i in adata.obs[COL].unique():
+        num = 0
+        for n, c in sorted(coun.items(), key=lambda x: x[1], reverse=True):
+            if i in n:
+                if c < MINCELLS or num == 0:
+                    relab[n] = i
+                else:
+                    relab[n] = i + "_" + str(num)
+                num += 1
+
+    adata.obs[NEWOBS] = adata.obs[NEWOBS].map(relab)
+
+    cluster_means = pd.DataFrame(
+        np.array(
+            [
+                adata.X[adata.obs[NEWOBS] == i].mean(axis=0)
+                for i in adata.obs[NEWOBS].unique()
+            ]
+        )[:, 0, :],
+        index=adata.obs[NEWOBS].unique(),
+    )
+
+    # Calculate correlation matrix between clusters
+    cluster_similarity = cluster_means.T.corr()
+    cluster_similarity.values[np.tril_indices(len(cluster_similarity), -1)] = 0
+
+    # Get pairs with similarity > 0.95
+    high_sim_pairs = []
+    for i in range(len(cluster_similarity)):
+        for j in range(i + 1, len(cluster_similarity)):
+            if (
+                cluster_similarity.iloc[i, j] > MAXSIM
+                and cluster_similarity.columns[i].split("_")[0]
+                == cluster_similarity.columns[j].split("_")[0]
+            ):
+                high_sim_pairs.append(
+                    (
+                        cluster_similarity.index[i],
+                        cluster_similarity.columns[j],
+                    )
+                )
+    # Create mapping for merging similar clusters
+    merge_mapping = {}
+    for pair in high_sim_pairs:
+        if pair[0] not in merge_mapping:
+            merge_mapping[pair[1]] = pair[0]
+        else:
+            merge_mapping[pair[1]] = merge_mapping[pair[0]]
+
+    # Apply merging
+    adata.obs[NEWOBS] = adata.obs[NEWOBS].map(merge_mapping).fillna(adata.obs[NEWOBS])
+    adata.obs[NEWOBS] = adata.obs[NEWOBS].astype(str)
+    coun = Counter(adata.obs[NEWOBS]).most_common()
+    merge_mapping = {}
+    for i in adata.obs[COL].unique():
+        num = 0
+        for j, c in coun:
+            if i in j:
+                merge_mapping[j] = i + "_" + str(num) if num > 0 else i
+                num += 1
+    adata.obs[NEWOBS] = adata.obs[NEWOBS].map(merge_mapping).fillna(adata.obs[NEWOBS])
+
+    import bionty as bt
+
+    stages = adata.obs["development_stage_ontology_term_id"].unique()
+    if adata.obs.organism_ontology_term_id.unique() == ["NCBITaxon:9606"]:
+        relabel = {i: i for i in stages}
+        for stage in stages:
+            stage_obj = bt.DevelopmentalStage.filter(ontology_id=stage).first()
+            parents = set([i.ontology_id for i in stage_obj.parents.filter()])
+            parents = parents - set(
+                [
+                    "HsapDv:0010000",
+                    "HsapDv:0000204",
+                    "HsapDv:0000227",
+                ]
+            )
+            if len(parents) > 0:
+                for p in parents:
+                    if p in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP:
+                        relabel[stage] = p
+        adata.obs["simplified_dev_stage"] = adata.obs[
+            "development_stage_ontology_term_id"
+        ].map(relabel)
+    elif adata.obs.organism_ontology_term_id.unique() == ["NCBITaxon:10090"]:
+        rename_mapping = {
+            k: v for v, j in MAIN_HUMAN_MOUSE_DEV_STAGE_MAP.items() for k in j
+        }
+        relabel = {i: "unknown" for i in stages}
+        for stage in stages:
+            if stage in rename_mapping:
+                relabel[stage] = rename_mapping[stage]
+        adata.obs["simplified_dev_stage"] = adata.obs[
+            "development_stage_ontology_term_id"
+        ].map(relabel)
+    else:
+        raise ValueError("organism not supported")
+    # palantir.utils.run_diffusion_maps(adata, n_components=20)
+    # palantir.utils.determine_multiscale_space(adata)
+    # terminal_states = palantir.utils.find_terminal_states(
+    #    adata,
+    #    celltypes=adata.obs.cell_type_ontology_term_id.unique(),
+    #    celltype_column="cell_type_ontology_term_id",
+    # )
+    # sc.tl.diffmap(adata)
+    # adata.obs["heat_diff"] = 1
+    # for terminal_state in terminal_states.index.tolist():
+    #    adata.uns["iroot"] = np.where(adata.obs.index == terminal_state)[0][0]
+    #    sc.tl.dpt(adata)
+    #    adata.obs["heat_diff"] = np.minimum(
+    #        adata.obs["heat_diff"], adata.obs["dpt_pseudotime"]
+    #    )
     return adata
+
+
+def cache_path(artifact):
+    cloud_path = UPath(artifact.storage.root) / artifact.key
+    cache_path = ln.setup.settings.paths.cloud_to_local_no_update(cloud_path)
+    return cache_path

scdataloader/utils.py

@@ -147,7 +147,7 @@ def getBiomartTable(
     return res
 
 
-def validate(adata: AnnData, organism: str):
+def validate(adata: AnnData, organism: str, need_all=True):
     """
     validate checks if the adata object is valid for lamindb
 
@@ -185,7 +185,7 @@ def validate(adata: AnnData, organism: str):
         "tissue_ontology_term_id",
         "assay_ontology_term_id",
     ]:
-        if val not in adata.obs.columns:
+        if val not in adata.obs.columns and need_all:
             raise ValueError(
                 f"Column '{val}' is missing in the provided anndata object."
             )
@@ -193,7 +193,9 @@ def validate(adata: AnnData, organism: str):
     if not bt.Ethnicity.validate(
         adata.obs["self_reported_ethnicity_ontology_term_id"],
         field="ontology_id",
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid ethnicity ontology term id found")
     if not bt.Organism.validate(
         adata.obs["organism_ontology_term_id"], field="ontology_id"
@@ -201,28 +203,40 @@ def validate(adata: AnnData, organism: str):
         raise ValueError("Invalid organism ontology term id found")
     if not bt.Phenotype.validate(
         adata.obs["sex_ontology_term_id"], field="ontology_id"
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid sex ontology term id found")
     if not bt.Disease.validate(
         adata.obs["disease_ontology_term_id"], field="ontology_id"
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid disease ontology term id found")
     if not bt.CellType.validate(
         adata.obs["cell_type_ontology_term_id"], field="ontology_id"
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid cell type ontology term id found")
     if not bt.DevelopmentalStage.validate(
         adata.obs["development_stage_ontology_term_id"],
         field="ontology_id",
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid dev stage ontology term id found")
     if not bt.Tissue.validate(
         adata.obs["tissue_ontology_term_id"], field="ontology_id"
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid tissue ontology term id found")
     if not bt.ExperimentalFactor.validate(
         adata.obs["assay_ontology_term_id"], field="ontology_id"
-    ).all():
+    ).all() and not set(adata.obs["self_reported_ethnicity_ontology_term_id"]) == set(
+        ["unknown"]
+    ):
         raise ValueError("Invalid assay ontology term id found")
     if not bt.Gene.validate(
         adata.var.index, field="ensembl_gene_id", organism=organism
@@ -582,7 +596,7 @@ def populate_my_ontology(
     # cell type
     if celltypes is not None:
         if len(celltypes) == 0:
-            bt.CellType.import_from_source(update=True)
+            bt.CellType.import_source()
         else:
             names = bt.CellType.public().df().index if not celltypes else celltypes
             records = bt.CellType.from_values(names, field="ontology_id")
@@ -597,9 +611,9 @@ def populate_my_ontology(
         )
         source = bt.PublicSource.filter(name="ensembl", organism=organism_clade).last()
         records = [
-            i[0] if type(i) is list else i
-            for i in [
-                bt.Organism.from_source(ontology_id=i, source=source) for i in names
+            organism_or_organismlist if isinstance(organism_or_organismlist, bt.Organism) else organism_or_organismlist[0]
+            for organism_or_organismlist in [
+                bt.Organism.from_source(ontology_id=name, source=source) for name in names
             ]
         ]
         ln.save(records)
@@ -616,7 +630,7 @@ def populate_my_ontology(
     # ethnicity
     if ethnicities is not None:
         if len(ethnicities) == 0:
-            bt.Ethnicity.import_from_source(update=True)
+            bt.Ethnicity.import_source()
         else:
             names = bt.Ethnicity.public().df().index if not ethnicities else ethnicities
             records = bt.Ethnicity.from_values(names, field="ontology_id")
@@ -627,7 +641,7 @@ def populate_my_ontology(
     # ExperimentalFactor
     if assays is not None:
         if len(assays) == 0:
-            bt.ExperimentalFactor.import_from_source(update=True)
+            bt.ExperimentalFactor.import_source()
         else:
             names = bt.ExperimentalFactor.public().df().index if not assays else assays
             records = bt.ExperimentalFactor.from_values(names, field="ontology_id")
@@ -638,7 +652,7 @@ def populate_my_ontology(
     # Tissue
     if tissues is not None:
         if len(tissues) == 0:
-            bt.Tissue.import_from_source(update=True)
+            bt.Tissue.import_source()
         else:
             names = bt.Tissue.public().df().index if not tissues else tissues
             records = bt.Tissue.from_values(names, field="ontology_id")
@@ -647,9 +661,9 @@ def populate_my_ontology(
     # DevelopmentalStage
     if dev_stages is not None:
         if len(dev_stages) == 0:
-            bt.DevelopmentalStage.import_from_source(update=True)
+            bt.DevelopmentalStage.import_source()
             source = bt.PublicSource.filter(organism="mouse", name="mmusdv").last()
-            bt.DevelopmentalStage.import_from_source(source=source)
+            bt.DevelopmentalStage.import_source(source=source)
         else:
             names = (
                 bt.DevelopmentalStage.public().df().index
@@ -663,7 +677,7 @@ def populate_my_ontology(
     # Disease
     if diseases is not None:
         if len(diseases) == 0:
-            bt.Disease.import_from_source(update=True)
+            bt.Disease.import_source()
         else:
             names = bt.Disease.public().df().index if not diseases else diseases
             records = bt.Disease.from_values(names, field="ontology_id")
@@ -738,18 +752,26 @@ def translate(
         dict: the mapping for the translation
     """
     if t == "cell_type_ontology_term_id":
-        obj = bt.CellType.public(organism="all")
+        obj = bt.CellType
     elif t == "assay_ontology_term_id":
-        obj = bt.ExperimentalFactor.public()
+        obj = bt.ExperimentalFactor
     elif t == "tissue_ontology_term_id":
-        obj = bt.Tissue.public()
+        obj = bt.Tissue
+    elif t in [
+        "development_stage_ontology_term_id",
+        "simplified_dev_stage",
+        "age_group",
+    ]:
+        obj = bt.DevelopmentalStage
+    elif t == "disease_ontology_term_id":
+        obj = bt.Disease
+    elif t == "self_reported_ethnicity_ontology_term_id":
+        obj = bt.Ethnicity
     else:
         return None
     if type(val) is str:
-        return {val: obj.search(val, field=obj.ontology_id).name.iloc[0]}
+        return {val: obj.filter(ontology_id=val).one().name}
     elif type(val) is list or type(val) is set:
-        return {i: obj.search(i, field=obj.ontology_id).name.iloc[0] for i in set(val)}
+        return {i: obj.filter(ontology_id=i).one().name for i in set(val)}
     elif type(val) is dict or type(val) is Counter:
-        return {
-            obj.search(k, field=obj.ontology_id).name.iloc[0]: v for k, v in val.items()
-        }
+        return {obj.filter(ontology_id=k).one().name: v for k, v in val.items()}