Documentation for utils module

def createFoldersFor(filepath: str):
    """
    will recursively create folders if needed until having all the folders required to save the file in this filepath
    """
    prevval = ""
    for val in os.path.expanduser(filepath).split("/")[:-1]:
        prevval += val + "/"
        if not os.path.exists(prevval):
            os.mkdir(prevval)

def fileToList(filename: str, strconv: callable = lambda x: x) -> list:
    """
    loads an input file with a\\n b\\n.. into a list [a,b,..]

    Args:
        filename (str): The filename to load from.
        strconv (callable): A function to convert each line. Defaults to identity function.

    Returns:
        list: The list of converted elements from the file.
    """
    with open(filename) as f:
        return [strconv(val[:-1]) for val in f.readlines()]

def getBiomartTable(
    ensemble_server: str = "http://may2024.archive.ensembl.org/biomart",
    useCache: bool = False,
    cache_folder: str = "/tmp/biomart/",
    attributes: List[str] = [],
    bypass_attributes: bool = False,
    database: str = "hsapiens_gene_ensembl",
) -> pd.DataFrame:
    """generate a genelist dataframe from ensembl's biomart

    Args:
        ensemble_server (str, optional): the biomart server. Defaults to "http://may2023.archive.ensembl.org/biomart".
        useCache (bool, optional): whether to use the cache or not. Defaults to False.
        cache_folder (str, optional): the cache folder. Defaults to "/tmp/biomart/".
        attributes (List[str], optional): the attributes to fetch. Defaults to [].
        bypass_attributes (bool, optional): whether to bypass the attributes or not. Defaults to False.
        database (str, optional): the database to fetch from. Defaults to "hsapiens_gene_ensembl".

    Raises:
        ValueError: should be a dataframe (when the result from the server is something else)

    Returns:
        pd.DataFrame: the dataframe
    """
    attr = (
        [
            "ensembl_gene_id",
            "hgnc_symbol",
            "gene_biotype",
            "entrezgene_id",
        ]
        if not bypass_attributes
        else []
    )
    assert cache_folder[-1] == "/"

    cache_folder = os.path.expanduser(cache_folder)
    createFoldersFor(cache_folder)
    cachefile = os.path.join(cache_folder, ".biomart.parquet")
    if useCache & os.path.isfile(cachefile):
        print("fetching gene names from biomart cache")
        res = pd.read_parquet(cachefile)
    else:
        print("downloading gene names from biomart")

        res = _fetchFromServer(ensemble_server, attr + attributes, database=database)
        res.to_parquet(cachefile, index=False)
    res.columns = attr + attributes
    if type(res) is not type(pd.DataFrame()):
        raise ValueError("should be a dataframe")
    res = res[~(res["ensembl_gene_id"].isna())]
    if "hgnc_symbol" in res.columns:
        res.loc[res[res.hgnc_symbol.isna()].index, "hgnc_symbol"] = res[
            res.hgnc_symbol.isna()
        ]["ensembl_gene_id"]
    return res

def get_ancestry_mapping(all_elem: List[str], onto_df: pd.DataFrame) -> dict:
    """
    This function generates a mapping of all elements to their ancestors in the ontology dataframe.

    Args:
        all_elem (list): A list of all elements.
        onto_df (DataFrame): The ontology dataframe.

    Returns:
        dict: A dictionary mapping each element to its ancestors.
    """
    ancestors = {}
    full_ancestors = set()
    for val in all_elem:
        ancestors[val] = get_all_ancestors(val, onto_df) - set([val])

    for val in ancestors.values():
        full_ancestors |= set(val)
    # removing ancestors that are not in our datasets
    # full_ancestors = full_ancestors & set(ancestors.keys())
    leafs = set(all_elem) - full_ancestors
    full_ancestors = full_ancestors - leafs

    groupings = {}
    for val in full_ancestors:
        groupings[val] = set()
    for leaf in leafs:
        for ancestor in ancestors[leaf]:
            if ancestor in full_ancestors:
                groupings[ancestor].add(leaf)

    return groupings, full_ancestors, leafs

def is_outlier(adata: AnnData, metric: str, nmads: int) -> pd.Series:
    """
    is_outlier detects outliers in adata.obs[metric]

    Args:
        adata (AnnData): the anndata object
        metric (str): the metric column to use
        nmads (int): the number of median absolute deviations to use as a threshold

    Returns:
        pd.Series: a boolean series indicating whether a cell is an outlier or not
    """
    M = adata.obs[metric]
    outlier = (M < np.median(M) - nmads * median_abs_deviation(M)) | (
        np.median(M) + nmads * median_abs_deviation(M) < M
    )
    return outlier

def length_normalize(adata: AnnData, gene_lengths: list) -> AnnData:
    """
    length_normalize normalizes the counts by the gene length

    Args:
        adata (AnnData): the anndata object
        gene_lengths (list): the gene lengths

    Returns:
        AnnData: the normalized anndata object
    """
    adata.X = csr_matrix((adata.X.T / gene_lengths).T)
    return adata

def listToFile(
    li: List[str], filename: str, strconv: callable = lambda x: str(x)
) -> None:
    """
    listToFile loads a list with [a,b,..] into an input file a\\n b\\n..

    Args:
        li (list): The list of elements to be written to the file.
        filename (str): The name of the file where the list will be written.
        strconv (callable, optional): A function to convert each element of the list to a string. Defaults to str.

    Returns:
        None
    """
    with open(filename, "w") as f:
        for item in li:
            f.write("%s\n" % strconv(item))

def load_dataset_local(
    remote_dataset: ln.Collection,
    download_folder: str,
    name: str,
    description: str,
    use_cache: bool = True,
    only: Optional[List[int]] = None,
) -> ln.Collection:
    """
    This function loads a remote lamindb dataset to local.

    Args:
        remote_dataset (lamindb.Collection): The remote Collection.
        download_folder (str): The path to the download folder.
        name (str): The name of the dataset.
        description (str): The description of the dataset.
        use_cache (bool, optional): Whether to use cache. Defaults to True.
        only (list, optional): A list of indices to specify which files to download. Defaults to None.

    Returns:
        lamindb.Dataset: The local dataset.
    """
    saved_files = []
    default_storage = ln.Storage.filter(root=ln.settings.storage.as_posix()).one()
    files = (
        remote_dataset.artifacts.all()
        if not only
        else remote_dataset.artifacts.all()[only[0] : only[1]]
    )
    for file in files:
        organism = list(set([i.ontology_id for i in file.organism.all()]))
        if len(organism) > 1:
            print(organism)
            print("Multiple organisms detected")
            continue
        if len(organism) == 0:
            print("No organism detected")
            continue
        organism = bt.Organism.filter(ontology_id=organism[0]).one().name
        # bt.settings.organism = organism
        path = file.path
        try:
            file.save()
        except IntegrityError:
            print(f"File {file.key} already exists in storage")
        # if location already has a file, don't save again
        if use_cache and os.path.exists(os.path.expanduser(download_folder + file.key)):
            print(f"File {file.key} already exists in storage")
        else:
            path.download_to(download_folder + file.key)
        file.storage = default_storage
        try:
            file.save()
        except IntegrityError:
            print(f"File {file.key} already exists in storage")
        saved_files.append(file)
    dataset = ln.Collection(saved_files, name=name, description=description)
    dataset.save()
    return dataset

def load_genes(
    organisms: Union[str, List[str]] = "NCBITaxon:9606",
) -> pd.DataFrame:  # "NCBITaxon:10090",
    """
    Loads genes from the given organisms.

    Args:
        organisms (Union[str, List[str]]): The organisms to load genes from.

    Returns:
        pd.DataFrame: The genes dataframe.
    """
    organismdf = []
    if type(organisms) is str:
        organisms = [organisms]
    for organism in organisms:
        genesdf = bt.Gene.filter(
            organism_id=bt.Organism.filter(ontology_id=organism).first().id
        ).df()
        genesdf.loc[genesdf.ensembl_gene_id.isna(), "ensembl_gene_id"] = genesdf.loc[
            genesdf.ensembl_gene_id.isna(), "stable_id"
        ]
        genesdf = genesdf.drop_duplicates(subset="ensembl_gene_id")
        genesdf = genesdf.set_index("ensembl_gene_id").sort_index()
        # mitochondrial genes
        genesdf["mt"] = genesdf.symbol.astype(str).str.startswith("MT-")
        # ribosomal genes
        genesdf["ribo"] = genesdf.symbol.astype(str).str.startswith(("RPS", "RPL"))
        # hemoglobin genes.
        genesdf["hb"] = genesdf.symbol.astype(str).str.contains(("^HB[^(P)]"))
        genesdf["organism"] = organism
        organismdf.append(genesdf)
    organismdf = pd.concat(organismdf)
    for col in [
        "source_id",
        "run_id",
        "created_by_id",
        "updated_at",
        "stable_id",
        "created_at",
        "_aux",
        "_branch_code",
        "space_id",
        "ncbi_gene_ids",
        "synonyms",
        "description",
    ]:
        if col in organismdf.columns:
            organismdf.drop(columns=[col], inplace=True)
    # temp fix
    organismdf = organismdf[~organismdf.index.isin(DROP)]
    return organismdf

def populate_my_ontology(
    sex: Optional[List[str]] = ["PATO:0000384", "PATO:0000383"],
    celltypes: Optional[List[str]] = [],
    ethnicities: Optional[List[str]] = [],
    assays: Optional[List[str]] = [],
    tissues: Optional[List[str]] = [],
    diseases: Optional[List[str]] = [],
    dev_stages: Optional[List[str]] = [],
    organisms_clade: Optional[List[str]] = ["vertebrates"],  # "plants", "metazoa"],
    organisms: Optional[List[str]] = ["NCBITaxon:10090", "NCBITaxon:9606"],
    genes_from: Optional[List[str]] = ["NCBITaxon:10090", "NCBITaxon:9606"],
):
    """
    creates a local version of the lamin ontologies and add the required missing values in base ontologies

    run this function just one for each new lamin storage

    erase everything with bt.$ontology.filter().delete()

    add whatever value you need afterward like it is done here with:

    `bt.$ontology(name="ddd", ontolbogy_id="ddddd").save()`

    `df["assay_ontology_term_id"].unique()`

    Args:
        sex (list, optional): List of sexes. Defaults to ["PATO:0000384", "PATO:0000383"].
        celltypes (list, optional): List of cell types. Defaults to [].
        ethnicities (list, optional): List of ethnicities. Defaults to [].
        assays (list, optional): List of assays. Defaults to [].
        tissues (list, optional): List of tissues. Defaults to [].
        diseases (list, optional): List of diseases. Defaults to [].
        dev_stages (list, optional): List of developmental stages. Defaults to [].
        organisms_clade (list, optional): List of organisms clade. Defaults to ["vertebrates", "plants"].
    """
    # cell type
    if celltypes is not None:
        if len(celltypes) == 0:
            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")
            ln.save(records)
        elem = bt.CellType(name="unknown", ontology_id="unknown")
        ln.save([elem], ignore_conflicts=True)
    # OrganismClade
    nrecords = []
    if organisms_clade is not None:
        records = []
        for organism_clade in organisms_clade:
            names = bt.Organism.public(organism=organism_clade).df().index
            source = bt.Source.filter(name="ensembl", organism=organism_clade).last()
            for name in names:
                try:
                    records.append(bt.Organism.from_source(name=name, source=source))
                except DoesNotExist:
                    print(f"Organism {name} not found in source {source}")

        prevrec = set()
        for rec in records:
            if rec is None:
                continue
            if not isinstance(rec, bt.Organism):
                rec = rec[0]
            if rec.uid not in prevrec:
                if organisms is not None:
                    if rec.ontology_id not in organisms:
                        continue
                nrecords.append(rec)
                prevrec.add(rec.uid)

        ln.save(nrecords)
        elem = bt.Organism(name="unknown", ontology_id="unknown").save()
        ln.save([elem], ignore_conflicts=True)
    # Phenotype
    if sex is not None:
        names = bt.Phenotype.public().df().index if not sex else sex
        source = bt.Source.filter(name="pato").first()
        records = [
            bt.Phenotype.from_source(ontology_id=i, source=source) for i in names
        ]
        ln.save(records)
        elem = bt.Phenotype(name="unknown", ontology_id="unknown").save()
        ln.save([elem], ignore_conflicts=True)
    # ethnicity
    if ethnicities is not None:
        if len(ethnicities) == 0:
            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")
            ln.save(records)
        elem = bt.Ethnicity(name="unknown", ontology_id="unknown")
        ln.save([elem], ignore_conflicts=True)
    # ExperimentalFactor
    if assays is not None:
        if len(assays) == 0:
            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")
            ln.save(records)
        elem = bt.ExperimentalFactor(name="unknown", ontology_id="unknown").save()
        ln.save([elem], ignore_conflicts=True)
        # lookup = bt.ExperimentalFactor.lookup()
        # lookup.smart_seq_v4.parents.add(lookup.smart_like)
    # Tissue
    if tissues is not None:
        if len(tissues) == 0:
            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")
            ln.save(records)
        elem = bt.Tissue(name="unknown", ontology_id="unknown").save()
        ln.save([elem], ignore_conflicts=True)
    # DevelopmentalStage
    if dev_stages is not None:
        if len(dev_stages) == 0:
            bt.DevelopmentalStage.import_source()
            source = bt.Source.filter(organism="mouse", name="mmusdv").last()
            bt.DevelopmentalStage.import_source(source=source)
        else:
            names = (
                bt.DevelopmentalStage.public().df().index
                if not dev_stages
                else dev_stages
            )
            records = bt.DevelopmentalStage.from_values(names, field="ontology_id")
            ln.save(records)

    # Disease
    if diseases is not None:
        if len(diseases) == 0:
            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")
            ln.save(records)
        bt.Disease(name="normal", ontology_id="PATO:0000461").save()
        bt.Disease(name="unknown", ontology_id="unknown").save()
    # genes
    for organism in genes_from:
        # convert onto to name
        organism = bt.Organism.filter(ontology_id=organism).one().name
        names = bt.Gene.public(organism=organism).df()["ensembl_gene_id"]

        # Process names in blocks of 10,000 elements
        block_size = 10000
        for i in range(0, len(names), block_size):
            block = names[i : i + block_size]
            records = bt.Gene.from_values(
                block,
                field="ensembl_gene_id",
                organism=organism,
            )
            ln.save(records)

def random_str(stringLength=6, stype="all", withdigits=True) -> str:
    """
    Generate a random string of letters and digits

    Args:
        stringLength (int, optional): the amount of char. Defaults to 6.
        stype (str, optional): one of lowercase, uppercase, all. Defaults to 'all'.
        withdigits (bool, optional): digits allowed in the string? Defaults to True.

    Returns:
        str: random string
    """
    if stype == "lowercase":
        lettersAndDigits = string.ascii_lowercase
    elif stype == "uppercase":
        lettersAndDigits = string.ascii_uppercase
    else:
        lettersAndDigits = string.ascii_letters
    if withdigits:
        lettersAndDigits += string.digits
    return "".join(random.choice(lettersAndDigits) for i in range(stringLength))

def translate(
    val: Union[str, list, set, Counter, dict], t: str = "cell_type_ontology_term_id"
) -> dict:
    """
    translate translates the ontology term id to the name

    Args:
        val (Union[str, dict, set, list]): the object to translate
        t (str, optional): the type of ontology terms.
            one of cell_type_ontology_term_id, assay_ontology_term_id, tissue_ontology_term_id.
            Defaults to "cell_type_ontology_term_id".

    Returns:
        dict: the mapping for the translation
    """
    if t == "cell_type_ontology_term_id":
        obj = bt.CellType
    elif t == "assay_ontology_term_id":
        obj = bt.ExperimentalFactor
    elif t == "tissue_ontology_term_id":
        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
    elif t == "organism_ontology_term_id":
        obj = bt.Organism
    else:
        return None
    if type(val) is str:
        return {val: obj.filter(ontology_id=val).one().name}
    elif type(val) is dict or type(val) is Counter:
        return {obj.filter(ontology_id=k).one().name: v for k, v in val.items()}
    elif type(val) is set:
        return {i: obj.filter(ontology_id=i).one().name for i in val}
    else:
        rl = {i: obj.filter(ontology_id=i).one().name for i in set(val)}
        return [rl.get(i, None) for i in val]

def validate(adata: AnnData, organism: str, need_all: bool = False) -> bool:
    """
    validate checks if the adata object is valid for lamindb

    Args:
        adata (AnnData): the anndata object
        organism (str): the organism ontology ID
        need_all (bool, optional): whether all columns should be present. Defaults to False.

    Raises:
        ValueError: if the adata object is not valid
        ValueError: if the anndata contains invalid ethnicity ontology term id according to the lb instance
        ValueError: if the anndata contains invalid organism ontology term id according to the lb instance
        ValueError: if the anndata contains invalid sex ontology term id according to the lb instance
        ValueError: if the anndata contains invalid disease ontology term id according to the lb instance
        ValueError: if the anndata contains invalid cell_type ontology term id according to the lb instance
        ValueError: if the anndata contains invalid development_stage ontology term id according to the lb instance
        ValueError: if the anndata contains invalid tissue ontology term id according to the lb instance
        ValueError: if the anndata contains invalid assay ontology term id according to the lb instance

    Returns:
        bool: True if the adata object is valid
    """
    organism = bt.Organism.filter(ontology_id=organism).one().name

    if adata.var.index.duplicated().any():
        raise ValueError("Duplicate gene names found in adata.var.index")
    if adata.obs.index.duplicated().any():
        raise ValueError("Duplicate cell names found in adata.obs.index")
    for val in [
        "self_reported_ethnicity_ontology_term_id",
        "organism_ontology_term_id",
        "disease_ontology_term_id",
        "cell_type_ontology_term_id",
        "development_stage_ontology_term_id",
        "tissue_ontology_term_id",
        "assay_ontology_term_id",
    ]:
        if val not in adata.obs.columns and need_all:
            raise ValueError(
                f"Column '{val}' is missing in the provided anndata object."
            )

    if not bt.Ethnicity.validate(
        adata.obs["self_reported_ethnicity_ontology_term_id"],
        field="ontology_id",
    ).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"
    ).all():
        raise ValueError("Invalid organism ontology term id found")
    if not bt.Phenotype.validate(
        adata.obs["sex_ontology_term_id"], field="ontology_id"
    ).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() 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() 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() 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() 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() 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
    ).all():
        if not bt.Gene.validate(
            adata.var.index, field="stable_id", organism=organism
        ).all():
            raise ValueError("Invalid gene ensembl id found")
    return True