jdti.utils

   1import os
   2import re
   3from typing import Literal
   4
   5import matplotlib as mpl
   6import matplotlib.patches as mpatches
   7import matplotlib.pyplot as plt
   8import numpy as np
   9import pandas as pd
  10import scipy.stats as stats
  11from adjustText import adjust_text
  12from joblib import Parallel, delayed
  13from matplotlib.lines import Line2D
  14from mpl_toolkits.axes_grid1.inset_locator import inset_axes
  15from scipy.cluster.hierarchy import dendrogram, linkage
  16from scipy.io import mmread
  17from tqdm import tqdm
  18
  19
  20def load_sparse(path: str, name: str):
  21    """
  22    Load a sparse matrix dataset along with associated gene
  23    and cell metadata, and return it as a dense DataFrame.
  24
  25    This function expects the input directory to contain three files in standard
  26    10x Genomics format:
  27      - "matrix.mtx": the gene expression matrix in Matrix Market format
  28      - "genes.tsv": tab-separated file containing gene identifiers
  29      - "barcodes.tsv": tab-separated file containing cell barcodes / names
  30
  31    Parameters
  32    ----------
  33    path : str
  34        Path to the directory containing the matrix and annotation files.
  35
  36    name : str
  37        Label or dataset identifier to be assigned to all cells in the metadata.
  38
  39    Returns
  40    -------
  41    data : pandas.DataFrame
  42        Dense expression matrix where rows correspond to genes and columns
  43        correspond to cells.
  44    metadata : pandas.DataFrame
  45        Metadata DataFrame with two columns:
  46          - "cell_names": the names of the cells (barcodes)
  47          - "sets": the dataset label assigned to each cell (from `name`)
  48
  49    Notes
  50    -----
  51    The function converts the sparse matrix into a dense DataFrame. This may
  52    require a large amount of memory for datasets with many cells and genes.
  53    """
  54
  55    data = mmread(os.path.join(path, "matrix.mtx"))
  56    data = pd.DataFrame(data.todense())
  57    genes = pd.read_csv(os.path.join(path, "genes.tsv"), header=None, sep="\t")
  58    names = pd.read_csv(os.path.join(path, "barcodes.tsv"), header=None, sep="\t")
  59    data.columns = [str(x) for x in names[0]]
  60    data.index = list(genes[0])
  61
  62    names = list(data.columns)
  63    sets = [name] * len(names)
  64
  65    metadata = pd.DataFrame({"cell_names": names, "sets": sets})
  66
  67    return data, metadata
  68
  69
  70def volcano_plot(
  71    deg_data: pd.DataFrame,
  72    p_adj: bool = True,
  73    top: int = 25,
  74    top_rank: str = "p_value",
  75    p_val: float | int = 0.05,
  76    lfc: float | int = 0.25,
  77    rescale_adj: bool = True,
  78    image_width: int = 12,
  79    image_high: int = 12,
  80):
  81    """
  82    Generate a volcano plot from differential expression results.
  83
  84    A volcano plot visualizes the relationship between statistical significance
  85    (p-values or standarized p-value) and log(fold change) for each gene, highlighting
  86    genes that pass significance thresholds.
  87
  88    Parameters
  89    ----------
  90    deg_data : pandas.DataFrame
  91        DataFrame containing differential expression results from calc_DEG() function.
  92
  93    p_adj : bool, default=True
  94        If True, use adjusted p-values. If False, use raw p-values.
  95
  96    top : int, default=25
  97        Number of top significant genes to highlight on the plot.
  98
  99    top_rank : str, default='p_value'
 100        Statistic used primarily to determine the top significant genes to highlight on the plot. ['p_value' or 'FC']
 101
 102    p_val : float | int, default=0.05
 103        Significance threshold for p-values (or adjusted p-values).
 104
 105    lfc : float | int, default=0.25
 106        Threshold for absolute log fold change.
 107
 108    rescale_adj : bool, default=True
 109        If True, rescale p-values to avoid long breaks caused by outlier values.
 110
 111    image_width : int, default=12
 112        Width of the generated plot in inches.
 113
 114    image_high : int, default=12
 115        Height of the generated plot in inches.
 116
 117    Returns
 118    -------
 119    matplotlib.figure.Figure
 120        The generated volcano plot figure.
 121
 122    """
 123
 124    if top_rank.upper() not in ["FC", "P_VALUE"]:
 125        raise ValueError("top_rank must be either 'FC' or 'p_value'")
 126
 127    if p_adj:
 128        pv = "adj_pval"
 129    else:
 130        pv = "p_val"
 131
 132    deg_df = deg_data.copy()
 133
 134    shift = 0.25
 135
 136    p_val_scale = "-log(p_val)"
 137
 138    min_minus = min(deg_df[pv][(deg_df[pv] != 0) & (deg_df["log(FC)"] < 0)])
 139    min_plus = min(deg_df[pv][(deg_df[pv] != 0) & (deg_df["log(FC)"] > 0)])
 140
 141    zero_p_plus = deg_df[(deg_df[pv] == 0) & (deg_df["log(FC)"] > 0)]
 142    zero_p_plus = zero_p_plus.sort_values(by="log(FC)", ascending=False).reset_index(
 143        drop=True
 144    )
 145    zero_p_plus[pv] = [
 146        (shift * x) * min_plus for x in range(1, len(zero_p_plus.index) + 1)
 147    ]
 148
 149    zero_p_minus = deg_df[(deg_df[pv] == 0) & (deg_df["log(FC)"] < 0)]
 150    zero_p_minus = zero_p_minus.sort_values(by="log(FC)", ascending=True).reset_index(
 151        drop=True
 152    )
 153    zero_p_minus[pv] = [
 154        (shift * x) * min_minus for x in range(1, len(zero_p_minus.index) + 1)
 155    ]
 156
 157    tmp_p = deg_df[
 158        ((deg_df[pv] != 0) & (deg_df["log(FC)"] < 0))
 159        | ((deg_df[pv] != 0) & (deg_df["log(FC)"] > 0))
 160    ]
 161
 162    del deg_df
 163
 164    deg_df = pd.concat([zero_p_plus, tmp_p, zero_p_minus], ignore_index=True)
 165
 166    deg_df[pv] = deg_df[pv].replace(0, 2**-1074)
 167
 168    deg_df[p_val_scale] = -np.log10(deg_df[pv])
 169
 170    deg_df["top100"] = None
 171
 172    if rescale_adj:
 173
 174        deg_df = deg_df.sort_values(by=p_val_scale, ascending=False)
 175
 176        deg_df = deg_df.reset_index(drop=True)
 177
 178        eps = 1e-300
 179        doubled = []
 180        ratio = []
 181        for n, i in enumerate(deg_df.index):
 182            for j in range(1, 6):
 183                if (
 184                    n + j < len(deg_df.index)
 185                    and (deg_df[p_val_scale][n] + eps)
 186                    / (deg_df[p_val_scale][n + j] + eps)
 187                    >= 2
 188                ):
 189                    doubled.append(n)
 190                    ratio.append(
 191                        (deg_df[p_val_scale][n + j] + eps)
 192                        / (deg_df[p_val_scale][n] + eps)
 193                    )
 194
 195        df = pd.DataFrame({"doubled": doubled, "ratio": ratio})
 196        df = df[df["doubled"] < 100]
 197
 198        df["ratio"] = (1 - df["ratio"]) / 5
 199        df = df.reset_index(drop=True)
 200
 201        df = df.sort_values("doubled")
 202
 203        if len(df["doubled"]) == 1 and 0 in df["doubled"]:
 204            df = df
 205        else:
 206            doubled2 = []
 207
 208            for l in df["doubled"]:
 209                if l + 1 != len(doubled) and l + 1 - l == 1:
 210                    doubled2.append(l)
 211                    doubled2.append(l + 1)
 212                else:
 213                    break
 214
 215            doubled2 = sorted(set(doubled2), reverse=True)
 216
 217        if len(doubled2) > 1:
 218            df = df[df["doubled"].isin(doubled2)]
 219            df = df.sort_values("doubled", ascending=False)
 220            df = df.reset_index(drop=True)
 221            for c in df.index:
 222                deg_df.loc[df["doubled"][c], p_val_scale] = deg_df.loc[
 223                    df["doubled"][c] + 1, p_val_scale
 224                ] * (1 + df["ratio"][c])
 225
 226    deg_df.loc[(deg_df["log(FC)"] <= 0) & (deg_df[pv] <= p_val), "top100"] = "red"
 227    deg_df.loc[(deg_df["log(FC)"] > 0) & (deg_df[pv] <= p_val), "top100"] = "blue"
 228    deg_df.loc[deg_df[pv] > p_val, "top100"] = "lightgray"
 229
 230    if lfc > 0:
 231        deg_df.loc[
 232            (deg_df["log(FC)"] <= lfc) & (deg_df["log(FC)"] >= -lfc), "top100"
 233        ] = "lightgray"
 234
 235    down_int = len(
 236        deg_df["top100"][(deg_df["log(FC)"] <= lfc * -1) & (deg_df[pv] <= p_val)]
 237    )
 238    up_int = len(deg_df["top100"][(deg_df["log(FC)"] > lfc) & (deg_df[pv] <= p_val)])
 239
 240    deg_df_up = deg_df[deg_df["log(FC)"] > 0]
 241
 242    if top_rank.upper() == "P_VALUE":
 243        deg_df_up = deg_df_up.sort_values([pv, "log(FC)"], ascending=[True, False])
 244    elif top_rank.upper() == "FC":
 245        deg_df_up = deg_df_up.sort_values(["log(FC)", pv], ascending=[False, True])
 246
 247    deg_df_up = deg_df_up.reset_index(drop=True)
 248
 249    n = -1
 250    l = 0
 251    while True:
 252        n += 1
 253        if deg_df_up["log(FC)"][n] > lfc and deg_df_up[pv][n] <= p_val:
 254            deg_df_up.loc[n, "top100"] = "green"
 255            l += 1
 256        if l == top or deg_df_up[pv][n] > p_val:
 257            break
 258
 259    deg_df_down = deg_df[deg_df["log(FC)"] <= 0]
 260
 261    if top_rank.upper() == "P_VALUE":
 262        deg_df_down = deg_df_down.sort_values([pv, "log(FC)"], ascending=[True, True])
 263    elif top_rank.upper() == "FC":
 264        deg_df_down = deg_df_down.sort_values(["log(FC)", pv], ascending=[True, True])
 265
 266    deg_df_down = deg_df_down.reset_index(drop=True)
 267
 268    n = -1
 269    l = 0
 270    while True:
 271        n += 1
 272        if deg_df_down["log(FC)"][n] < lfc * -1 and deg_df_down[pv][n] <= p_val:
 273            deg_df_down.loc[n, "top100"] = "yellow"
 274
 275            l += 1
 276        if l == top or deg_df_down[pv][n] > p_val:
 277            break
 278
 279    deg_df = pd.concat([deg_df_up, deg_df_down])
 280
 281    que = ["lightgray", "red", "blue", "yellow", "green"]
 282
 283    deg_df = deg_df.sort_values(
 284        by="top100", key=lambda x: x.map({v: i for i, v in enumerate(que)})
 285    )
 286
 287    deg_df = deg_df.reset_index(drop=True)
 288
 289    fig, ax = plt.subplots(figsize=(image_width, image_high))
 290
 291    plt.scatter(
 292        x=deg_df["log(FC)"], y=deg_df[p_val_scale], color=deg_df["top100"], zorder=2
 293    )
 294
 295    tl = deg_df[p_val_scale][deg_df[pv] >= p_val]
 296
 297    if len(tl) > 0:
 298
 299        line_p = np.max(tl)
 300
 301    else:
 302        line_p = np.min(deg_df[p_val_scale])
 303
 304    plt.plot(
 305        [max(deg_df["log(FC)"]) * -1.1, max(deg_df["log(FC)"]) * 1.1],
 306        [line_p, line_p],
 307        linestyle="--",
 308        linewidth=3,
 309        color="lightgray",
 310        zorder=1,
 311    )
 312
 313    if lfc > 0:
 314        plt.plot(
 315            [lfc * -1, lfc * -1],
 316            [-3, max(deg_df[p_val_scale]) * 1.1],
 317            linestyle="--",
 318            linewidth=3,
 319            color="lightgray",
 320            zorder=1,
 321        )
 322        plt.plot(
 323            [lfc, lfc],
 324            [-3, max(deg_df[p_val_scale]) * 1.1],
 325            linestyle="--",
 326            linewidth=3,
 327            color="lightgray",
 328            zorder=1,
 329        )
 330
 331    plt.xlabel("log(FC)")
 332    plt.ylabel(p_val_scale)
 333    plt.title("Volcano plot")
 334
 335    plt.ylim(min(deg_df[p_val_scale]) - 5, max(deg_df[p_val_scale]) * 1.25)
 336
 337    texts = [
 338        ax.text(deg_df["log(FC)"][i], deg_df[p_val_scale][i], deg_df["feature"][i])
 339        for i in deg_df.index
 340        if deg_df["top100"][i] in ["green", "yellow"]
 341    ]
 342
 343    adjust_text(texts, arrowprops=dict(arrowstyle="-", color="gray", alpha=0.5))
 344
 345    legend_elements = [
 346        Line2D(
 347            [0],
 348            [0],
 349            marker="o",
 350            color="w",
 351            label="top-upregulated",
 352            markerfacecolor="green",
 353            markersize=10,
 354        ),
 355        Line2D(
 356            [0],
 357            [0],
 358            marker="o",
 359            color="w",
 360            label="top-downregulated",
 361            markerfacecolor="yellow",
 362            markersize=10,
 363        ),
 364        Line2D(
 365            [0],
 366            [0],
 367            marker="o",
 368            color="w",
 369            label="upregulated",
 370            markerfacecolor="blue",
 371            markersize=10,
 372        ),
 373        Line2D(
 374            [0],
 375            [0],
 376            marker="o",
 377            color="w",
 378            label="downregulated",
 379            markerfacecolor="red",
 380            markersize=10,
 381        ),
 382        Line2D(
 383            [0],
 384            [0],
 385            marker="o",
 386            color="w",
 387            label="non-significant",
 388            markerfacecolor="lightgray",
 389            markersize=10,
 390        ),
 391    ]
 392
 393    ax.legend(handles=legend_elements, loc="upper right")
 394    ax.grid(visible=False)
 395
 396    ax.annotate(
 397        f"\nmin {pv} = " + str(p_val),
 398        xy=(0.025, 0.975),
 399        xycoords="axes fraction",
 400        fontsize=12,
 401    )
 402
 403    if lfc > 0:
 404        ax.annotate(
 405            "\nmin log(FC) = " + str(lfc),
 406            xy=(0.025, 0.95),
 407            xycoords="axes fraction",
 408            fontsize=12,
 409        )
 410
 411    ax.annotate(
 412        "\nDownregulated: " + str(down_int),
 413        xy=(0.025, 0.925),
 414        xycoords="axes fraction",
 415        fontsize=12,
 416        color="red",
 417    )
 418
 419    ax.annotate(
 420        "\nUpregulated: " + str(up_int),
 421        xy=(0.025, 0.9),
 422        xycoords="axes fraction",
 423        fontsize=12,
 424        color="blue",
 425    )
 426
 427    plt.show()
 428
 429    return fig
 430
 431
 432def find_features(data: pd.DataFrame, features: list):
 433    """
 434    Identify features (rows) from a DataFrame that match a given list of features,
 435    ignoring case sensitivity.
 436
 437    Parameters
 438    ----------
 439    data : pandas.DataFrame
 440        DataFrame with features in the index (rows).
 441
 442    features : list
 443        List of feature names to search for.
 444
 445    Returns
 446    -------
 447    dict
 448        Dictionary with keys:
 449        - "included": list of features found in the DataFrame index.
 450        - "not_included": list of requested features not found in the DataFrame index.
 451        - "potential": list of features in the DataFrame that may be similar.
 452    """
 453
 454    features_upper = [str(x).upper() for x in features]
 455
 456    index_set = set(data.index)
 457
 458    features_in = [x for x in index_set if x.upper() in features_upper]
 459    features_in_upper = [x.upper() for x in features_in]
 460    features_out_upper = [x for x in features_upper if x not in features_in_upper]
 461    features_out = [x for x in features if x.upper() in features_out_upper]
 462    similar_features = [
 463        idx for idx in index_set if any(x in idx.upper() for x in features_out_upper)
 464    ]
 465
 466    return {
 467        "included": features_in,
 468        "not_included": features_out,
 469        "potential": similar_features,
 470    }
 471
 472
 473def find_names(data: pd.DataFrame, names: list):
 474    """
 475    Identify names (columns) from a DataFrame that match a given list of names,
 476    ignoring case sensitivity.
 477
 478    Parameters
 479    ----------
 480    data : pandas.DataFrame
 481        DataFrame with names in the columns.
 482
 483    names : list
 484        List of names to search for.
 485
 486    Returns
 487    -------
 488    dict
 489        Dictionary with keys:
 490        - "included": list of names found in the DataFrame columns.
 491        - "not_included": list of requested names not found in the DataFrame columns.
 492        - "potential": list of names in the DataFrame that may be similar.
 493    """
 494
 495    names_upper = [str(x).upper() for x in names]
 496
 497    columns = set(data.columns)
 498
 499    names_in = [x for x in columns if x.upper() in names_upper]
 500    names_in_upper = [x.upper() for x in names_in]
 501    names_out_upper = [x for x in names_upper if x not in names_in_upper]
 502    names_out = [x for x in names if x.upper() in names_out_upper]
 503    similar_names = [
 504        idx for idx in columns if any(x in idx.upper() for x in names_out_upper)
 505    ]
 506
 507    return {"included": names_in, "not_included": names_out, "potential": similar_names}
 508
 509
 510def reduce_data(data: pd.DataFrame, features: list = [], names: list = []):
 511    """
 512    Subset a DataFrame based on selected features (rows) and/or names (columns).
 513
 514    Parameters
 515    ----------
 516    data : pandas.DataFrame
 517        Input DataFrame with features as rows and names as columns.
 518
 519    features : list
 520        List of features to include (rows). Default is an empty list.
 521        If empty, all rows are returned.
 522
 523    names : list
 524        List of names to include (columns). Default is an empty list.
 525        If empty, all columns are returned.
 526
 527    Returns
 528    -------
 529    pandas.DataFrame
 530        Subset of the input DataFrame containing only the selected rows
 531        and/or columns.
 532
 533    Raises
 534    ------
 535    ValueError
 536        If both `features` and `names` are empty.
 537    """
 538
 539    if len(features) > 0 and len(names) > 0:
 540        fet = find_features(data=data, features=features)
 541
 542        nam = find_names(data=data, names=names)
 543
 544        data_to_return = data.loc[fet["included"], nam["included"]]
 545
 546    elif len(features) > 0 and len(names) == 0:
 547        fet = find_features(data=data, features=features)
 548
 549        data_to_return = data.loc[fet["included"], :]
 550
 551    elif len(features) == 0 and len(names) > 0:
 552
 553        nam = find_names(data=data, names=names)
 554
 555        data_to_return = data.loc[:, nam["included"]]
 556
 557    else:
 558
 559        raise ValueError("features and names have zero length!")
 560
 561    return data_to_return
 562
 563
 564def make_unique_list(lst):
 565    """
 566    Generate a list where duplicate items are renamed to ensure uniqueness.
 567
 568    Each duplicate is appended with a suffix ".n", where n indicates the
 569    occurrence count (starting from 1).
 570
 571    Parameters
 572    ----------
 573    lst : list
 574        Input list of items (strings or other hashable types).
 575
 576    Returns
 577    -------
 578    list
 579        List with unique values.
 580
 581    Examples
 582    --------
 583    >>> make_unique_list(["A", "B", "A", "A"])
 584    ['A', 'B', 'A.1', 'A.2']
 585    """
 586    seen = {}
 587    result = []
 588    for item in lst:
 589        if item not in seen:
 590            seen[item] = 0
 591            result.append(item)
 592        else:
 593            seen[item] += 1
 594            result.append(f"{item}.{seen[item]}")
 595    return result
 596
 597
 598def get_color_palette(variable_list, palette_name="tab10"):
 599    n = len(variable_list)
 600    cmap = plt.get_cmap(palette_name)
 601    colors = [cmap(i % cmap.N) for i in range(n)]
 602    return dict(zip(variable_list, colors))
 603
 604
 605def features_dot(
 606    expression_data: pd.DataFrame,
 607    occurence_data: pd.DataFrame | None = None,
 608    scale: bool = False,
 609    scale_axis: Literal["x", "y"] = "y",
 610    features: list | None = None,
 611    metadata_list: list | None = None,
 612    colors: str = "viridis",
 613    hclust: str | None = "complete",
 614    img_width: int = 8,
 615    img_high: int = 5,
 616    label_size: int = 10,
 617    size_scale: int = 100,
 618    y_lab: str = "Genes",
 619    legend_lab: str = "log(CPM + 1)",
 620    set_box_size: float | int = 5,
 621    set_box_high: float | int = 5,
 622    bbox_to_anchor_scale: int = 25,
 623    bbox_to_anchor_perc: tuple = (0.91, 0.63),
 624    bbox_to_anchor_group: tuple = (1.01, 0.4),
 625):
 626    """
 627    Create a dot plot of selected features across samples.
 628
 629    Each point represents a feature-sample pair, where the color encodes the
 630    expression value and the size encodes occurrence or relative abundance.
 631    Optionally, hierarchical clustering can be applied to order rows and columns.
 632
 633    Parameters
 634    ----------
 635    expression_data : pandas.DataFrame
 636        Expression values (mean) with features as rows and samples as columns derived from average() function.
 637
 638    occurence_data : pandas.DataFrame or None
 639        DataFrame with occurrence/frequency values (same shape as `expression_data`) derived from occurrence() function.
 640        If None, bubble sizes are based on expression values.
 641
 642    scale: bool, default False
 643        If True, expression_data (features) will be scaled (0–1) across the colums (sample) with Maximum Absolute Scaling method.
 644
 645    scale_axis : {'x', 'y'}, default 'y'
 646        The axis along which to scale the expression data to the [0, 1] range:
 647        - 'y': Scales across columns (per sample / cell).
 648        - 'x': Scales across rows (per feature / gene).
 649        Only effective if `scale=True`.
 650
 651    features : list or None
 652        List of features (rows) to display. If None, all features are used.
 653
 654    metadata_list : list or None, optional
 655        Metadata grouping for samples (same length as number of columns).
 656        Used to add group colors and separators in the plot.
 657
 658    colors : str, default='viridis'
 659        Colormap for expression values.
 660
 661    hclust : str or None, default='complete'
 662        Linkage method for hierarchical clustering. If None, no clustering
 663        is performed.
 664
 665    img_width : int or float, default=8
 666        Width of the plot in inches.
 667
 668    img_high : int or float, default=5
 669        Height of the plot in inches.
 670
 671    label_size : int, default=10
 672        Font size for axis labels and ticks.
 673
 674    size_scale : int or float, default=100
 675        Scaling factor for bubble sizes.
 676
 677    y_lab : str, default='Genes'
 678        Label for the x-axis.
 679
 680    legend_lab : str, default='log(CPM + 1)'
 681        Label for the colorbar legend.
 682
 683    bbox_to_anchor_scale : int, default=25
 684        Vertical scale (percentage) for positioning the colorbar.
 685
 686    bbox_to_anchor_perc : tuple, default=(0.91, 0.63)
 687        Anchor position for the size legend (percent bubble legend).
 688
 689    bbox_to_anchor_group : tuple, default=(1.01, 0.4)
 690        Anchor position for the group legend.
 691
 692    Returns
 693    -------
 694    matplotlib.figure.Figure
 695        The generated dot plot figure.
 696
 697    Raises
 698    ------
 699    ValueError
 700        If `metadata_list` is provided but its length does not match
 701        the number of columns in `expression_data`.
 702
 703    Notes
 704    -----
 705    - Colors represent expression values normalized to the colormap.
 706    - Bubble sizes represent occurrence values (or expression values if
 707      `occurence_data` is None).
 708    - If `metadata_list` is given, groups are indicated with colors and
 709      dashed vertical separators.
 710    """
 711    if scale and scale_axis not in ["x", "X", "y", "Y"]:
 712        raise ValueError("Wrong axis to scale. scale_axis must be 'x' or 'y'.")
 713
 714    scatter_df = expression_data.copy()
 715
 716    metadata = {}
 717
 718    metadata["primary_names"] = [str(x) for x in scatter_df.columns]
 719
 720    if metadata_list is not None:
 721        metadata["sets"] = metadata_list
 722
 723        if len(metadata["primary_names"]) != len(metadata["sets"]):
 724
 725            raise ValueError(
 726                "Metadata list and DataFrame columns must have the same length."
 727            )
 728
 729    else:
 730
 731        metadata["sets"] = [""] * len(metadata["primary_names"])
 732
 733    metadata = pd.DataFrame(metadata)
 734    if features is not None:
 735        scatter_df = scatter_df.loc[
 736            find_features(data=scatter_df, features=features)["included"],
 737        ]
 738    scatter_df.columns = metadata["primary_names"] + "#" + metadata["sets"]
 739
 740    if occurence_data is not None:
 741        if features is not None:
 742            occurence_data = occurence_data.loc[
 743                find_features(data=occurence_data, features=features)["included"],
 744            ]
 745        occurence_data.columns = metadata["primary_names"] + "#" + metadata["sets"]
 746
 747    # check duplicated names
 748
 749    tmp_columns = scatter_df.columns
 750
 751    new_cols = make_unique_list(list(tmp_columns))
 752
 753    scatter_df.columns = new_cols
 754
 755    # scaling
 756    if scale:
 757
 758        legend_lab = "Scaled\n" + legend_lab
 759
 760        column_max = scatter_df.max()
 761        scatter_df = (
 762            scatter_df.div(column_max).replace([np.inf, -np.inf], np.nan).fillna(0)
 763        )
 764        scatter_df = pd.DataFrame(
 765            scatter_df, index=scatter_df.index, columns=scatter_df.columns
 766        )
 767
 768        if scale_axis.upper() == "X":
 769            scatter_df = scatter_df.T
 770
 771        column_max = scatter_df.max()
 772        scatter_df = (
 773            scatter_df.div(column_max).replace([np.inf, -np.inf], np.nan).fillna(0)
 774        )
 775        scatter_df = pd.DataFrame(
 776            scatter_df, index=scatter_df.index, columns=scatter_df.columns
 777        )
 778
 779        if scale_axis.upper() == "X":
 780            scatter_df = scatter_df.T
 781
 782    if hclust is not None and len(expression_data.index) != 1:
 783
 784        Z = linkage(scatter_df, method=hclust)
 785
 786        # Get the order of features based on the dendrogram
 787        order_of_features = dendrogram(Z, no_plot=True)["leaves"]
 788
 789        indexes_sort = list(scatter_df.index)
 790        sorted_list_rows = []
 791        for n in order_of_features:
 792            sorted_list_rows.append(indexes_sort[n])
 793
 794        scatter_df = scatter_df.transpose()
 795
 796        Z = linkage(scatter_df, method=hclust)
 797
 798        # Get the order of features based on the dendrogram
 799        order_of_features = dendrogram(Z, no_plot=True)["leaves"]
 800
 801        indexes_sort = list(scatter_df.index)
 802        sorted_list_columns = []
 803        for n in order_of_features:
 804            sorted_list_columns.append(indexes_sort[n])
 805
 806        scatter_df = scatter_df.transpose()
 807
 808        scatter_df = scatter_df.loc[sorted_list_rows, sorted_list_columns]
 809
 810        if occurence_data is not None:
 811            occurence_data = occurence_data.loc[sorted_list_rows, sorted_list_columns]
 812
 813        metadata["sets"] = [re.sub(".*#", "", x) for x in scatter_df.columns]
 814
 815    scatter_df.columns = [re.sub("#.*", "", x) for x in scatter_df.columns]
 816
 817    if occurence_data is not None:
 818        occurence_data.columns = [re.sub("#.*", "", x) for x in occurence_data.columns]
 819
 820    fig, ax = plt.subplots(figsize=(img_width, img_high))
 821
 822    norm = plt.Normalize(0, np.max(scatter_df))
 823
 824    cmap = plt.get_cmap(colors)
 825
 826    # dot-plot
 827    for i, _ in enumerate(scatter_df.index):
 828        for j, _ in enumerate(scatter_df.columns):
 829            if occurence_data is not None:
 830                value_e = scatter_df.iloc[i, j]
 831                value_o = occurence_data.iloc[i, j]
 832                ax.scatter(
 833                    j,
 834                    i,
 835                    s=value_o * size_scale,
 836                    c=[cmap(norm(value_e))],
 837                    edgecolors="k",
 838                    linewidths=0.3,
 839                )
 840            else:
 841                value = scatter_df.iloc[i, j]
 842                ax.scatter(
 843                    j,
 844                    i,
 845                    s=value * size_scale,
 846                    c=[cmap(norm(value))],
 847                    edgecolors="k",
 848                    linewidths=0.3,
 849                )
 850
 851    ax.set_yticks(range(len(scatter_df.index)))
 852    ax.set_yticklabels(scatter_df.index, fontsize=label_size * 0.8)
 853    ax.set_ylabel(y_lab, fontsize=label_size)
 854    ax.set_xticks(range(len(scatter_df.columns)))
 855    ax.set_xticklabels(scatter_df.columns, fontsize=label_size * 0.8, rotation=90)
 856
 857    ax_pos = ax.get_position()
 858
 859    width_fig = 0.01
 860    height_fig = ax_pos.height * (bbox_to_anchor_scale / 100)
 861    left_fig = ax_pos.x1 + 0.01
 862    bottom_fig = ax_pos.y1 - height_fig
 863
 864    cax = fig.add_axes([left_fig, bottom_fig, width_fig, height_fig])
 865    cb = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm)
 866    cb.set_label(legend_lab, fontsize=label_size * 0.65)
 867    cb.ax.tick_params(labelsize=label_size * 0.7)
 868
 869    if metadata_list is not None:
 870
 871        metadata_list = list(metadata["sets"])
 872        group_colors = get_color_palette(list(set(metadata_list)), palette_name="tab10")
 873
 874        for i, group in enumerate(metadata_list):
 875            ax.add_patch(
 876                plt.Rectangle(
 877                    (i - 0.5, len(scatter_df.index) - 0.1 * set_box_high),
 878                    1,
 879                    0.1 * set_box_size,
 880                    color=group_colors[group],
 881                    transform=ax.transData,
 882                    clip_on=False,
 883                )
 884            )
 885
 886        for i in range(1, len(metadata_list)):
 887            if metadata_list[i] != metadata_list[i - 1]:
 888                ax.axvline(i - 0.5, color="black", linestyle="--", lw=1)
 889
 890        group_patches = [
 891            mpatches.Patch(color=color, label=label)
 892            for label, color in group_colors.items()
 893        ]
 894        fig.legend(
 895            handles=group_patches,
 896            title="Group",
 897            fontsize=label_size * 0.7,
 898            title_fontsize=label_size * 0.7,
 899            loc="center left",
 900            bbox_to_anchor=bbox_to_anchor_group,
 901            frameon=False,
 902        )
 903
 904    # second legend (size)
 905    if occurence_data is not None:
 906        size_values = [0.25, 0.5, 1]
 907        legend2_handles = [
 908            plt.Line2D(
 909                [],
 910                [],
 911                marker="o",
 912                linestyle="",
 913                markersize=np.sqrt(v * size_scale * 0.5),
 914                color="gray",
 915                alpha=0.6,
 916                label=f"{v * 100:.1f}",
 917            )
 918            for v in size_values
 919        ]
 920
 921        fig.legend(
 922            handles=legend2_handles,
 923            title="Percent [%]",
 924            fontsize=label_size * 0.7,
 925            title_fontsize=label_size * 0.7,
 926            loc="center left",
 927            bbox_to_anchor=bbox_to_anchor_perc,
 928            frameon=False,
 929        )
 930
 931    _, ymax = ax.get_ylim()
 932
 933    ax.set_xlim(-0.5, len(scatter_df.columns) - 0.5)
 934    ax.set_ylim(-0.5, ymax + 0.5)
 935
 936    return fig
 937
 938
 939def calc_DEG(
 940    data,
 941    metadata_list: list | None = None,
 942    entities: str | list | dict | None = None,
 943    sets: str | list | dict | None = None,
 944    min_exp: int | float = 0,
 945    min_pct: int | float = 0.1,
 946    n_proc: int = 10,
 947):
 948    """
 949    Perform differential gene expression (DEG) analysis on gene expression data.
 950
 951    The function compares groups of cells or samples (defined by `entities` or
 952    `sets`) using the Mann–Whitney U test. It computes p-values, adjusted
 953    p-values, fold changes, standardized effect sizes, and other statistics.
 954
 955    Parameters
 956    ----------
 957    data : pandas.DataFrame
 958        Expression matrix with features (e.g., genes) as rows and samples/cells
 959        as columns.
 960
 961    metadata_list : list or None, optional
 962        Metadata grouping corresponding to the columns in `data`. Required for
 963        comparisons based on sets. Default is None.
 964
 965    entities : list, str, dict, or None, optional
 966        Defines the comparison strategy:
 967        - list of sample names → compare selected cells to the rest.
 968        - 'All' → compare each sample/cell to all others.
 969        - dict → user-defined groups for pairwise comparison.
 970        - None → must be combined with `sets`.
 971
 972    sets : str, dict, or None, optional
 973        Defines group-based comparisons:
 974        - 'All' → compare each set/group to all others.
 975        - dict with two groups → perform pairwise set comparison.
 976        - None → must be combined with `entities`.
 977
 978    min_exp : float | int, default=0
 979        Minimum expression threshold for filtering features.
 980
 981    min_pct : float | int, default=0.1
 982        Minimum proportion of samples within the target group that must express
 983        a feature for it to be tested.
 984
 985    n_proc : int, default=10
 986        Number of parallel processes to use for statistical testing.
 987
 988    Returns
 989    -------
 990    pandas.DataFrame or dict
 991        Results of the differential expression analysis:
 992        - If `entities` is a list → dict with keys: 'valid_cells',
 993          'control_cells', and 'DEG' (results DataFrame).
 994        - If `entities == 'All'` or `sets == 'All'` → DataFrame with results
 995          for all groups.
 996        - If pairwise comparison (dict for `entities` or `sets`) → DataFrame
 997          with results for the specified groups.
 998
 999        The results DataFrame contains:
1000        - 'feature': feature name
1001        - 'p_val': raw p-value
1002        - 'adj_pval': adjusted p-value (multiple testing correction)
1003        - 'pct_valid': fraction of target group expressing the feature
1004        - 'pct_ctrl': fraction of control group expressing the feature
1005        - 'avg_valid': mean expression in target group
1006        - 'avg_ctrl': mean expression in control group
1007        - 'sd_valid': standard deviation in target group
1008        - 'sd_ctrl': standard deviation in control group
1009        - 'esm': effect size metric
1010        - 'FC': fold change
1011        - 'log(FC)': log2-transformed fold change
1012        - 'norm_diff': difference in mean expression
1013
1014    Raises
1015    ------
1016    ValueError
1017        - If `metadata_list` is provided but its length does not match
1018          the number of columns in `data`.
1019        - If neither `entities` nor `sets` is provided.
1020
1021    Notes
1022    -----
1023    - Mann–Whitney U test is used for group comparisons.
1024    - Multiple testing correction is applied using a simple
1025      Benjamini–Hochberg-like method.
1026    - Features expressed below `min_exp` or in fewer than `min_pct` of target
1027      samples are filtered out.
1028    - Parallelization is handled by `joblib.Parallel`.
1029
1030    Examples
1031    --------
1032    Compare a selected list of cells against all others:
1033
1034    >>> result = calc_DEG(data, entities=["cell1", "cell2", "cell3"])
1035
1036    Compare each group to others (based on metadata):
1037
1038    >>> result = calc_DEG(data, metadata_list=group_labels, sets="All")
1039
1040    Perform pairwise comparison between two predefined sets:
1041
1042    >>> sets = {"GroupA": ["A1", "A2"], "GroupB": ["B1", "B2"]}
1043    >>> result = calc_DEG(data, sets=sets)
1044    """
1045    offset = 1e-100
1046
1047    metadata = {}
1048
1049    metadata["primary_names"] = [str(x) for x in data.columns]
1050
1051    if metadata_list is not None:
1052        metadata["sets"] = metadata_list
1053
1054        if len(metadata["primary_names"]) != len(metadata["sets"]):
1055
1056            raise ValueError(
1057                "Metadata list and DataFrame columns must have the same length."
1058            )
1059
1060    else:
1061
1062        metadata["sets"] = [""] * len(metadata["primary_names"])
1063
1064    metadata = pd.DataFrame(metadata)
1065
1066    def stat_calc(choose, feature_name):
1067        target_values = choose.loc[choose["DEG"] == "target", feature_name]
1068        rest_values = choose.loc[choose["DEG"] == "rest", feature_name]
1069
1070        pct_valid = (target_values > 0).sum() / len(target_values)
1071        pct_rest = (rest_values > 0).sum() / len(rest_values)
1072
1073        avg_valid = np.mean(target_values)
1074        avg_ctrl = np.mean(rest_values)
1075        sd_valid = np.std(target_values, ddof=1)
1076        sd_ctrl = np.std(rest_values, ddof=1)
1077        esm = (avg_valid - avg_ctrl) / np.sqrt(((sd_valid**2 + sd_ctrl**2) / 2))
1078
1079        if np.sum(target_values) == np.sum(rest_values):
1080            p_val = 1.0
1081        else:
1082            _, p_val = stats.mannwhitneyu(
1083                target_values, rest_values, alternative="two-sided"
1084            )
1085
1086        return {
1087            "feature": feature_name,
1088            "p_val": p_val,
1089            "pct_valid": pct_valid,
1090            "pct_ctrl": pct_rest,
1091            "avg_valid": avg_valid,
1092            "avg_ctrl": avg_ctrl,
1093            "sd_valid": sd_valid,
1094            "sd_ctrl": sd_ctrl,
1095            "esm": esm,
1096        }
1097
1098    def prepare_and_run_stat(choose, valid_group, min_exp, min_pct, n_proc):
1099
1100        def safe_min_half(series):
1101            filtered = series[(series > ((2**-1074) * 2)) & (series.notna())]
1102            return filtered.min() / 2 if not filtered.empty else 0
1103
1104        tmp_dat = choose[choose["DEG"] == "target"]
1105        tmp_dat = tmp_dat.drop("DEG", axis=1)
1106
1107        counts = (tmp_dat > min_exp).sum(axis=0)
1108
1109        total_count = tmp_dat.shape[0]
1110
1111        info = pd.DataFrame(
1112            {"feature": list(tmp_dat.columns), "pct": list(counts / total_count)}
1113        )
1114
1115        del tmp_dat
1116
1117        drop_col = info["feature"][info["pct"] <= min_pct]
1118
1119        if len(drop_col) + 1 == len(choose.columns):
1120            drop_col = info["feature"][info["pct"] == 0]
1121
1122        del info
1123
1124        choose = choose.drop(list(drop_col), axis=1)
1125
1126        results = Parallel(n_jobs=n_proc)(
1127            delayed(stat_calc)(choose, feature)
1128            for feature in tqdm(choose.columns[choose.columns != "DEG"])
1129        )
1130
1131        if len(results) > 0:
1132            df = pd.DataFrame(results)
1133
1134            df = df[(df["avg_valid"] > 0) | (df["avg_ctrl"] > 0)]
1135
1136            df["valid_group"] = valid_group
1137            df.sort_values(by="p_val", inplace=True)
1138
1139            num_tests = len(df)
1140            df["adj_pval"] = np.minimum(
1141                1, (df["p_val"] * num_tests) / np.arange(1, num_tests + 1)
1142            )
1143
1144            valid_factor = safe_min_half(df["avg_valid"])
1145            ctrl_factor = safe_min_half(df["avg_ctrl"])
1146
1147            cv_factor = min(valid_factor, ctrl_factor)
1148
1149            if cv_factor == 0:
1150                cv_factor = max(valid_factor, ctrl_factor)
1151
1152            if not np.isfinite(cv_factor) or cv_factor == 0:
1153                cv_factor += offset
1154
1155            valid = df["avg_valid"].where(
1156                df["avg_valid"] != 0, df["avg_valid"] + cv_factor
1157            )
1158            ctrl = df["avg_ctrl"].where(df["avg_ctrl"] != 0, df["avg_ctrl"] + cv_factor)
1159
1160            df["FC"] = valid / ctrl
1161
1162            df["log(FC)"] = np.log2(df["FC"])
1163            df["norm_diff"] = df["avg_valid"] - df["avg_ctrl"]
1164
1165        else:
1166            columns = [
1167                "feature",
1168                "valid_group",
1169                "p_val",
1170                "adj_pval",
1171                "avg_valid",
1172                "avg_ctrl",
1173                "FC",
1174                "log(FC)",
1175                "norm_diff",
1176            ]
1177            df = pd.DataFrame(columns=columns)
1178        return df
1179
1180    choose = data.T
1181
1182    final_results = []
1183
1184    if isinstance(entities, list) and sets is None:
1185        print("\nAnalysis started...\nComparing selected cells to the whole set...")
1186
1187        if metadata_list is None:
1188            choose.index = metadata["primary_names"]
1189        else:
1190            choose.index = metadata["primary_names"] + " # " + metadata["sets"]
1191
1192            if "#" not in entities[0]:
1193                choose.index = metadata["primary_names"]
1194                print(
1195                    "You provided 'metadata_list', but did not include the set info (name # set) "
1196                    "in the 'entities' list. "
1197                    "Only the names will be compared, without considering the set information."
1198                )
1199
1200        labels = ["target" if idx in entities else "rest" for idx in choose.index]
1201        valid = list(
1202            set(choose.index[[i for i, x in enumerate(labels) if x == "target"]])
1203        )
1204
1205        choose["DEG"] = labels
1206        choose = choose[choose["DEG"] != "drop"]
1207
1208        result_df = prepare_and_run_stat(
1209            choose.reset_index(drop=True),
1210            valid_group=valid,
1211            min_exp=min_exp,
1212            min_pct=min_pct,
1213            n_proc=n_proc,
1214        )
1215
1216        return {"valid": valid, "control": "rest", "DEG": result_df}
1217
1218    elif entities == "All" and sets is None:
1219        print("\nAnalysis started...\nComparing each type of cell to others...")
1220
1221        if metadata_list is None:
1222            choose.index = metadata["primary_names"]
1223        else:
1224            choose.index = metadata["primary_names"] + " # " + metadata["sets"]
1225
1226        unique_labels = set(choose.index)
1227
1228        for label in tqdm(unique_labels):
1229            print(f"\nCalculating statistics for {label}")
1230            labels = ["target" if idx == label else "rest" for idx in choose.index]
1231            choose["DEG"] = labels
1232            choose = choose[choose["DEG"] != "drop"]
1233            result_df = prepare_and_run_stat(
1234                choose.copy(),
1235                valid_group=label,
1236                min_exp=min_exp,
1237                min_pct=min_pct,
1238                n_proc=n_proc,
1239            )
1240            final_results.append(result_df)
1241
1242        final_results = pd.concat(final_results, ignore_index=True)
1243
1244        if metadata_list is None:
1245            final_results["valid_group"] = [
1246                re.sub(" # ", "", x) for x in final_results["valid_group"]
1247            ]
1248
1249        return final_results
1250
1251    elif entities is None and sets == "All":
1252        print("\nAnalysis started...\nComparing each set/group to others...")
1253        choose.index = metadata["sets"]
1254        unique_sets = set(choose.index)
1255
1256        for label in tqdm(unique_sets):
1257            print(f"\nCalculating statistics for {label}")
1258            labels = ["target" if idx == label else "rest" for idx in choose.index]
1259
1260            choose["DEG"] = labels
1261            choose = choose[choose["DEG"] != "drop"]
1262            result_df = prepare_and_run_stat(
1263                choose.copy(),
1264                valid_group=label,
1265                min_exp=min_exp,
1266                min_pct=min_pct,
1267                n_proc=n_proc,
1268            )
1269            final_results.append(result_df)
1270
1271        return pd.concat(final_results, ignore_index=True)
1272
1273    elif entities is None and isinstance(sets, dict):
1274        print("\nAnalysis started...\nComparing groups...")
1275        choose.index = metadata["sets"]
1276
1277        group_list = list(sets.keys())
1278        if len(group_list) != 2:
1279            print("Only pairwise group comparison is supported.")
1280            return None
1281
1282        labels = [
1283            (
1284                "target"
1285                if idx in sets[group_list[0]]
1286                else "rest" if idx in sets[group_list[1]] else "drop"
1287            )
1288            for idx in choose.index
1289        ]
1290        choose["DEG"] = labels
1291        choose = choose[choose["DEG"] != "drop"]
1292
1293        result_df = prepare_and_run_stat(
1294            choose.reset_index(drop=True),
1295            valid_group=group_list[0],
1296            min_exp=min_exp,
1297            min_pct=min_pct,
1298            n_proc=n_proc,
1299        )
1300        return result_df
1301
1302    elif isinstance(entities, dict) and sets is None:
1303        print("\nAnalysis started...\nComparing groups...")
1304
1305        if metadata_list is None:
1306            choose.index = metadata["primary_names"]
1307        else:
1308            choose.index = metadata["primary_names"] + " # " + metadata["sets"]
1309            if "#" not in entities[list(entities.keys())[0]][0]:
1310                choose.index = metadata["primary_names"]
1311                print(
1312                    "You provided 'metadata_list', but did not include the set info (name # set) "
1313                    "in the 'entities' dict. "
1314                    "Only the names will be compared, without considering the set information."
1315                )
1316
1317        group_list = list(entities.keys())
1318        if len(group_list) != 2:
1319            print("Only pairwise group comparison is supported.")
1320            return None
1321
1322        labels = [
1323            (
1324                "target"
1325                if idx in entities[group_list[0]]
1326                else "rest" if idx in entities[group_list[1]] else "drop"
1327            )
1328            for idx in choose.index
1329        ]
1330
1331        choose["DEG"] = labels
1332        choose = choose[choose["DEG"] != "drop"]
1333
1334        result_df = prepare_and_run_stat(
1335            choose.reset_index(drop=True),
1336            valid_group=group_list[0],
1337            min_exp=min_exp,
1338            min_pct=min_pct,
1339            n_proc=n_proc,
1340        )
1341
1342        return result_df.reset_index(drop=True)
1343
1344    else:
1345        raise ValueError(
1346            "You must specify either 'entities' or 'sets'. None were provided, which is not allowed for this analysis."
1347        )
1348
1349
1350def average(data):
1351    """
1352    Compute the column-wise average of a DataFrame, aggregating by column names.
1353
1354    If multiple columns share the same name, their values are averaged.
1355
1356    Parameters
1357    ----------
1358    data : pandas.DataFrame
1359        Input DataFrame with numeric values. Columns with identical names
1360        will be aggregated by their mean.
1361
1362    Returns
1363    -------
1364    pandas.DataFrame
1365        DataFrame with the same rows as the input but with unique columns,
1366        where duplicate columns have been replaced by their mean values.
1367    """
1368
1369    wide_data = data
1370
1371    aggregated_df = wide_data.T.groupby(level=0).mean().T
1372
1373    return aggregated_df
1374
1375
1376def occurrence(data):
1377    """
1378    Calculate the occurrence frequency of features in a DataFrame.
1379
1380    Converts the input DataFrame to binary (presence/absence) and computes
1381    the proportion of non-zero entries for each feature, aggregating by
1382    column names if duplicates exist.
1383
1384    Parameters
1385    ----------
1386    data : pandas.DataFrame
1387        Input DataFrame with numeric values. Each column represents a feature.
1388
1389    Returns
1390    -------
1391    pandas.DataFrame
1392        DataFrame with the same rows as the input, where each value represents
1393        the proportion of samples in which the feature is present (non-zero).
1394        Columns with identical names are aggregated.
1395    """
1396
1397    binary_data = (data > 0).astype(int)
1398
1399    counts = binary_data.columns.value_counts()
1400
1401    binary_data = binary_data.T.groupby(level=0).sum().T
1402    binary_data = binary_data.astype(float)
1403
1404    for i in counts.index:
1405        binary_data.loc[:, i] = (binary_data.loc[:, i] / counts[i]).astype(float)
1406
1407    return binary_data
1408
1409
1410def add_subnames(names_list: list, parent_name: str, new_clusters: list):
1411    """
1412    Append sub-cluster names to a parent name within a list of names.
1413
1414    This function replaces occurrences of `parent_name` in `names_list` with
1415    a concatenation of the parent name and corresponding sub-cluster name
1416    from `new_clusters` (formatted as "parent.subcluster"). Non-matching names
1417    are left unchanged.
1418
1419    Parameters
1420    ----------
1421    names_list : list
1422        Original list of names (e.g., column names or cluster labels).
1423
1424    parent_name : str
1425        Name of the parent cluster to which sub-cluster names will be added.
1426        Must exist in `names_list`.
1427
1428    new_clusters : list
1429        List of sub-cluster names. Its length must match the number of times
1430        `parent_name` occurs in `names_list`.
1431
1432    Returns
1433    -------
1434    list
1435        Updated list of names with sub-cluster names appended to the parent name.
1436
1437    Raises
1438    ------
1439    ValueError
1440        - If `parent_name` is not found in `names_list`.
1441        - If `new_clusters` length does not match the number of occurrences of
1442          `parent_name`.
1443
1444    Examples
1445    --------
1446    >>> add_subnames(['A', 'B', 'A'], 'A', ['1', '2'])
1447    ['A.1', 'B', 'A.2']
1448    """
1449
1450    if str(parent_name) not in [str(x) for x in names_list]:
1451        raise ValueError(
1452            "Parent name is missing from the original dataset`s column names!"
1453        )
1454
1455    if len(new_clusters) != len([x for x in names_list if str(x) == str(parent_name)]):
1456        raise ValueError(
1457            "New cluster names list has a different length than the number of clusters in the original dataset!"
1458        )
1459
1460    new_names = []
1461    ixn = 0
1462    for _, i in enumerate(names_list):
1463        if str(i) == str(parent_name):
1464
1465            new_names.append(f"{parent_name}.{new_clusters[ixn]}")
1466            ixn += 1
1467
1468        else:
1469            new_names.append(i)
1470
1471    return new_names
1472
1473
1474def development_clust(
1475    data: pd.DataFrame, method: str = "ward", img_width: int = 5, img_high: int = 5
1476):
1477    """
1478    Perform hierarchical clustering on the columns of a DataFrame and plot a dendrogram.
1479
1480    Uses Ward's method to cluster the transposed data (columns) and generates
1481    a dendrogram showing the relationships between features or samples.
1482
1483    Parameters
1484    ----------
1485    data : pandas.DataFrame
1486        Input DataFrame with features as rows and samples/columns to be clustered.
1487
1488    method : str
1489        Method for hierarchical clustering. Options include:
1490       - 'ward' : minimizes the variance of clusters being merged.
1491       - 'single' : uses the minimum of the distances between all observations of the two sets.
1492       - 'complete' : uses the maximum of the distances between all observations of the two sets.
1493       - 'average' : uses the average of the distances between all observations of the two sets.
1494
1495    img_width : int or float, default=5
1496        Width of the resulting figure in inches.
1497
1498    img_high : int or float, default=5
1499        Height of the resulting figure in inches.
1500
1501    Returns
1502    -------
1503    matplotlib.figure.Figure
1504        The dendrogram figure.
1505    """
1506
1507    z = linkage(data.T, method=method)
1508
1509    figure, ax = plt.subplots(figsize=(img_width, img_high))
1510
1511    dendrogram(z, labels=data.columns, orientation="left", ax=ax)
1512
1513    return figure
1514
1515
1516def adjust_cells_to_group_mean(data, data_avg, beta=0.2):
1517    """
1518    Adjust each cell's values towards the mean of its group (centroid).
1519
1520    This function moves each cell's values in `data` slightly towards the
1521    corresponding group mean in `data_avg`, controlled by the parameter `beta`.
1522
1523    Parameters
1524    ----------
1525    data : pandas.DataFrame
1526        Original data with features as rows and cells/samples as columns.
1527
1528    data_avg : pandas.DataFrame
1529        DataFrame of group averages (centroids) with features as rows and
1530        group names as columns.
1531
1532    beta : float, default=0.2
1533        Weight for adjustment towards the group mean. 0 = no adjustment,
1534        1 = fully replaced by the group mean.
1535
1536    Returns
1537    -------
1538    pandas.DataFrame
1539        Adjusted data with the same shape as the input `data`.
1540    """
1541
1542    df_adjusted = data.copy()
1543
1544    for group_name in data_avg.columns:
1545        col_idx = [
1546            i
1547            for i, c in enumerate(df_adjusted.columns)
1548            if str(c).startswith(group_name)
1549        ]
1550        if not col_idx:
1551            continue
1552
1553        centroid = data_avg.loc[df_adjusted.index, group_name].to_numpy()[:, None]
1554
1555        df_adjusted.iloc[:, col_idx] = (1 - beta) * df_adjusted.iloc[
1556            :, col_idx
1557        ].to_numpy() + beta * centroid
1558
1559    return df_adjusted