gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/report/three_d_combine.py

@@ -0,0 +1,312 @@
+import logging
+import os
+from pathlib import Path
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import pyvista as pv
+import statsmodels.api as sm
+import statsmodels.stats.multitest as smm
+from scipy.stats import fisher_exact
+
+from gsMap.cauchy_combination_test import _acat_test
+from gsMap.config import ThreeDCombineConfig
+
+from .three_d_plot.three_d_plots import three_d_plot, three_d_plot_save
+
+pv.start_xvfb()
+logger = logging.getLogger(__name__)
+
+
+def combine_ldsc(args):
+
+    # Set the output path
+    ldsc_root = Path(args.project_dir) / "3D_combine" / "spatial_ldsc"
+    ldsc_root.mkdir(parents=True, exist_ok=True)
+    name = ldsc_root / f"{args.trait_name}.csv.gz"
+
+    # Merge all the ldsc results
+    pth = Path(args.project_dir) / "spatial_ldsc"
+    sldsc_pth = []
+    for slice in os.listdir(pth):
+        filtemp = pth / slice / f"{slice}_{args.trait_name}.csv.gz"
+        if filtemp.exists():
+            sldsc_pth.append(filtemp)
+
+    if not os.path.exists(name):
+        logger.info(f"Find {len(sldsc_pth)} ST sections for {args.trait_name}, start to merge the results...")
+        # Load the results
+        ldsc_merge = pd.DataFrame()
+        for idx, file in enumerate(sldsc_pth):
+            ldsc_temp = pd.read_csv(file, compression="gzip")
+            ldsc_temp["ST_id"] = file.name.split(f"_{args.trait_name}")[0]
+            # print(ldsc_temp.head())
+            ldsc_merge = pd.concat([ldsc_merge, ldsc_temp], axis=0)
+
+        # Check the cell name duplication
+        if (ldsc_merge.spot.value_counts() > 1).any():
+            logger.info('There are duplicated spot names, using the st_id + spot_id as the spot index.')
+            ldsc_merge['spot_index'] = ldsc_merge['ST_id'] + '_' + ldsc_merge['spot'].astype(str)
+        else:
+            ldsc_merge['spot_index'] = ldsc_merge['spot']
+
+        # save the merged results
+        ldsc_merge.to_csv(name, compression="gzip", index=False)
+        logger.info(f"Saving the 3D merged results to {name}")
+    else:
+        logger.info(f"The merged gsMap results already exist, loading the merged results from {name}...")
+        ldsc_merge = pd.read_csv(name, compression="gzip")
+
+    return ldsc_merge
+
+
+def cauchy_combination_3d(ldsc):
+    p_cauchy = []
+    p_median = []
+    gc_median = []
+    for ct in np.unique(ldsc.annotation):
+        p_temp = ldsc.loc[ldsc["annotation"] == ct, "p"]
+        z_temp = ldsc.loc[ldsc["annotation"] == ct, "z"]
+        p_temp = p_temp.dropna()
+
+        # The Cauchy test is sensitive to very small p-values, so extreme outliers should be considered for removal...
+        p_temp_log = -np.log10(p_temp)
+        median_log = np.median(p_temp_log)
+        IQR_log = np.percentile(p_temp_log, 75) - np.percentile(p_temp_log, 25)
+
+        p_use = p_temp[p_temp_log < median_log + 3 * IQR_log]
+        z_use = z_temp[p_temp_log < median_log + 3 * IQR_log]
+        n_remove = len(p_temp) - len(p_use)
+
+        # Outlier: -log10(p) < median + 3IQR && len(outlier set) < 20
+        # if 0 < n_remove < max(len(p_temp) * 0.001,100):
+        if 0 < n_remove < len(p_temp) * 0.05:
+            print(
+                f"Remove {
+                    n_remove}/{len(p_temp)} outliers (median + 3*IQR) for {ct}."
+            )
+            p_cauchy_temp = _acat_test(p_use)
+        else:
+            p_cauchy_temp = _acat_test(p_temp)
+
+        p_median_temp = np.median(p_use)
+        gc_median_temp = np.median(z_use**2) / 0.4549
+
+        p_cauchy.append(p_cauchy_temp)
+        p_median.append(p_median_temp)
+        gc_median.append(gc_median_temp)
+
+    data = {
+        "p_cauchy": p_cauchy,
+        "p_median": p_median,
+        "inflation_factor": gc_median,
+        "annotation": np.unique(ldsc.annotation),
+    }
+    p_tissue = pd.DataFrame(data)
+    p_tissue.columns = ["p_cauchy", "p_median", "inflation_factor", "annotation"]
+    p_tissue.sort_values("p_cauchy", inplace=True)
+    return p_tissue
+
+# def cauchy_combination_3d(args):
+
+#     # Load the cauchy combination results of each ST slices
+#     pth = Path(args.project_dir) / "cauchy_combination"
+#     st_file = os.listdir(pth)
+#     logger.info(f"Find {len(st_file)} sections of cauchy combination results for {args.trait_name}...")
+
+#     cauchy_all = pd.DataFrame()
+#     for slice in st_file:
+#         filtemp = pth / slice / f"{slice}_{args.trait_name}.Cauchy.csv.gz"
+#         if filtemp.exists():
+#             cauchy = pd.read_csv(filtemp, compression="gzip")
+#             cauchy_all = pd.concat([cauchy_all, cauchy], axis=0)
+
+#     cauchy_all = cauchy_all[~cauchy_all.annotation.isna()]
+
+#     # Do the cauchy combination test across slices
+#     p_cauchy = []
+#     p_median = []
+#     for ct in cauchy_all.annotation.unique():
+#         cauchy_temp = cauchy_all.loc[cauchy_all.annotation == ct]
+#         p_cauchy_temp = cauchy_temp.p_cauchy
+#         p_median_temp = cauchy_temp.p_median
+#         n_cell = cauchy_temp.n_cell
+
+#         p_cauchy_temp_log = -np.log10(p_cauchy_temp)
+#         median_log = np.median(p_cauchy_temp_log)
+#         IQR_log = np.percentile(p_cauchy_temp_log, 75) - np.percentile(p_cauchy_temp_log, 25)
+
+#         w_use = n_cell
+#         p_use = p_cauchy_temp
+#         if len(p_cauchy_temp) > 15:
+#             index = p_cauchy_temp_log < median_log + 2*IQR_log
+#             w_use = n_cell[index]
+#             p_use = p_cauchy_temp[index]
+#             n_remove = len(p_cauchy_temp) - len(p_use)
+#             if n_remove > 0:
+#                 logger.info(f"Remove {n_remove} outlier (median + 2*IQR) sections for {ct}")
+
+#         p_cauchy_new = acat_test(pvalues=p_use.to_list(),weights=w_use.to_list())
+#         p_median_new = (p_median_temp * n_cell / n_cell.sum()).sum()
+
+#         p_cauchy.append(p_cauchy_new)
+#         p_median.append(p_median_new)
+
+#     data = {
+#             "p_cauchy": p_cauchy,
+#             "p_median": p_median,
+#             "annotation": cauchy_all.annotation.unique(),
+#             }
+#     p_tissue = pd.DataFrame(data)
+#     p_tissue.columns = ["p_cauchy", "p_median", "annotation"]
+#     p_tissue.sort_values("p_cauchy", inplace=True)
+#     return p_tissue
+
+
+def odds_test_3d(ldsc_merge):
+    _, corrected_p_values, _, _ = smm.multipletests(ldsc_merge.p, alpha=0.05)
+    ldsc_merge['p_fdr'] = corrected_p_values.tolist()
+
+    Odds = []
+    for focal_annotation in ldsc_merge.annotation.unique():
+        try:
+            focal_no,focal_yes = (ldsc_merge.loc[ldsc_merge.annotation==focal_annotation,'p_fdr'] < 0.05).value_counts()
+            other_no,other_yes = (ldsc_merge.loc[ldsc_merge.annotation!=focal_annotation,'p_fdr'] < 0.05).value_counts()
+            contingency_table = [[focal_yes, focal_no], [other_yes, other_no]]
+            odds_ratio, p_value = fisher_exact(contingency_table)
+            table = sm.stats.Table2x2(contingency_table)
+            conf_int = table.oddsratio_confint()
+        except Exception:
+            odds_ratio = 0
+            p_value = 1
+            conf_int = (0, 0)
+        Odds.append({
+            'annotation': focal_annotation,
+            'odds_ratio': f"{odds_ratio:.3f}",
+            '95%_ci_low': f"{conf_int[0]:.3f}",
+            '95%_ci_high': f"{conf_int[1]:.3f}",
+            'p_odds_ratio': p_value
+        })
+    Odds = pd.DataFrame(Odds)
+    return Odds
+
+
+def three_d_combine(args: ThreeDCombineConfig):
+
+    # Load the ldsc results
+    ldsc_merge = combine_ldsc(args)
+    ldsc_merge.spot_index = ldsc_merge.spot_index.astype(str).replace(r"\.0", "", regex=True)
+    ldsc_merge.index = ldsc_merge.spot_index
+
+    # Load the spatial data
+    logger.info(f"Loading {args.adata_3d}.")
+    adata_3d_path = str(args.adata_3d)
+    if adata_3d_path.endswith('.parquet'):
+        logger.info("The input data is the metadata file of adata.")
+        meta_merged = pd.read_parquet(adata_3d_path)
+    elif adata_3d_path.endswith('.h5ad'):
+        logger.info("The input data is the h5ad.")
+        adata_merge = ad.read_h5ad(adata_3d_path, backed='r')
+        adata_merge.obs.index.name = 'index'
+        spatial = pd.DataFrame(adata_merge.obsm[args.spatial_key], columns=['sx', 'sy', 'sz'], index=adata_merge.obs_names).copy()
+        spatial = spatial.reset_index()
+        meta = adata_merge.obs.copy()
+        meta_merged = spatial.merge(meta, left_on='index', right_index=True, how='left')
+        meta_merged.index = adata_merge.obs_names
+
+    # Handle DataFrame or AnnData
+    if args.st_id is not None and (meta_merged.index.value_counts() > 1).any():
+        # Check if the index has duplicates and if st_id is provided
+        if len(np.intersect1d(ldsc_merge.index, meta_merged.index)) == 0:
+            # If no common cells, create a new index using st_id
+            logger.info(f"Using {args.st_id} + adata.obs_names as the new cell index.")
+            meta_merged.index = meta_merged[args.st_id].astype(str) + '_' + meta_merged.index.astype(str)
+
+    # Find common cells
+    common_cell = np.intersect1d(ldsc_merge.index, meta_merged.index)
+    if len(common_cell) == 0:
+        raise ValueError("No common cells between the spatial data and the ldsc results.")
+
+    logger.info(f"Found {len(common_cell)} common cells between the 3D spatial data and the mapping results.")
+
+    # Subset the data to common cells
+    meta_merged = meta_merged.loc[common_cell].copy()
+    ldsc_merge = ldsc_merge.loc[common_cell]
+
+    # Do cauchy combination test and odds ratio test
+    if args.annotation is not None:
+        annotation_use = meta_merged[args.annotation]
+        ldsc_merge['annotation'] = annotation_use
+        ldsc_merge = ldsc_merge[~ldsc_merge.annotation.isna()]
+
+        cauchy = cauchy_combination_3d(ldsc_merge)
+        odds = odds_test_3d(ldsc_merge)
+        cauchy_odds = pd.merge(odds,cauchy,left_on='annotation',right_on='annotation')
+
+        # Save the results
+        cauchy_root = Path(args.project_dir) / "3D_combine" / "cauchy_combination"
+        cauchy_root.mkdir(parents=True, exist_ok=True, mode=0o755)
+        cauchy_name = cauchy_root / f"{args.trait_name}.{args.annotation}.Cauchy.csv.gz"
+        cauchy_odds = cauchy_odds.sort_values('odds_ratio',ascending=False)
+        cauchy_odds.to_csv(cauchy_name, compression="gzip", index=False)
+        logger.info(f"Saving the 3D combination combination results to {cauchy_name}")
+    else:
+        logger.info("No annotation provided for the cauchy combination test.")
+
+
+    # Plot the 3D results
+    p_color = ['#313695', '#4575b4', '#74add1','#fee090', '#fdae61', '#f46d43', '#d73027', '#a50026']
+    meta_merged["logp"] = -np.log10(ldsc_merge.p)
+
+    required_columns = {'sx', 'sy', 'sz'}
+    if required_columns.issubset(meta_merged.columns):
+        logger.info("Generating 3D plot...")
+
+        # Set the legend and text
+        legend_kwargs = dict(scalar_bar_title_size=30, scalar_bar_label_size=30, fmt="%.1e")
+        text_kwargs = dict(text_font_size=15, text_loc="upper_edge")
+
+        # Set the opacity for each point
+        meta_merged['logp'].fillna(0, inplace=True)
+        bins = np.linspace(meta_merged['logp'].min(), meta_merged['logp'].max(), 5)
+        alpha = np.exp(np.linspace(0.1, 1.0, num=(len(bins)-1)))-1
+        alpha = alpha / max(alpha)
+        opacity_show = pd.cut(meta_merged['logp'], bins=bins, labels=alpha, include_lowest=True).values.tolist()
+
+        # Set the clim
+        max_v = np.round(np.median(np.sort(meta_merged['logp'])[::-1][0:20]))
+
+        # Plot the 3D results
+        plotter = three_d_plot(
+            clim = [0,max_v],
+            point_size=args.point_size,
+            opacity=opacity_show,
+            window_size=(1200, 1008),
+            adata=meta_merged,
+            spatial_key=args.spatial_key,
+            keys=["logp"],
+            cmaps=[args.cmap] if args.cmap is not None else [p_color],
+            scalar_bar_titles=["-log10(p)"],
+            texts=[args.trait_name],
+            jupyter=False,
+            background=args.background_color,
+            show_outline=args.show_outline,
+            legend_kwargs=legend_kwargs,
+            text_kwargs=text_kwargs,
+        )
+
+        # Save the results
+        plot_root = Path(args.project_dir) / "3D_combine" / "3D_plot"
+        plot_root.mkdir(parents=True, exist_ok=True, mode=0o755)
+        plot_name = plot_root / args.trait_name
+
+        three_d_plot_save(
+            plotter,
+            save_mp4=args.save_mp4,
+            save_gif=args.save_gif,
+            n_points=args.n_snapshot if args.n_snapshot is not None else 200,
+            filename=plot_name,
+        )
+    else:
+        logger.info("The spatial data does not contain 3D spatial coordinates for 3D plotting.")

gsMap/report/three_d_plot/three_d_plot_decorate.py

@@ -0,0 +1,246 @@
+import matplotlib as mpl
+import numpy as np
+from pyvista import MultiBlock
+
+categorical_legend_loc_legal = ["upper right",
+                                "upper left",
+                                "lower left",
+                                "lower right",
+                                "center left",
+                                "center right",
+                                "lower center",
+                                "upper center",
+                                "center"]
+
+
+def add_model(
+    plotter,
+    model,
+    key=None,
+    colormap=None,
+    clim=None,
+    ambient=0.2,
+    opacity=1.0,
+    model_style="surface",
+    point_size=3.0,
+):
+
+    def _add_model(_p, _model, _key, _colormap, _style, _ambient, _opacity, _point_size,_clim):
+        """Add any PyVista/VTK model to the scene."""
+        if _style == "points":
+            _render_spheres, render_tubes, _smooth_shading = True, False, True
+        elif _style == "wireframe":
+            _render_spheres, render_tubes, _smooth_shading = False, True, False
+        else:
+            _render_spheres, render_tubes, _smooth_shading = False, False, True
+        mesh_kwargs = dict(
+            style=_style,
+            render_points_as_spheres=True,
+            render_lines_as_tubes=render_tubes,
+            point_size=_point_size,
+            line_width=_point_size,
+            ambient=_ambient,
+            opacity=_opacity,
+            smooth_shading=True,
+            clim=_clim,
+            show_scalar_bar=False,
+        )
+
+        if _colormap is None:
+            added_kwargs = dict(
+                scalars=f"{
+                    _key}_rgba" if _key in _model.array_names else _model.active_scalars_name,
+                rgba=True
+            )
+        else:
+            added_kwargs = dict(
+                scalars=_key if _key in _model.array_names else _model.active_scalars_name,
+                cmap=_colormap
+            )
+
+        mesh_kwargs.update(added_kwargs)
+        _p.add_mesh(_model, **mesh_kwargs)
+
+    # Add model(s) to the plotter.
+    _add_model(
+        _p=plotter,
+        _model=model,
+        _key=key,
+        _colormap=colormap,
+        _style=model_style,
+        _point_size=point_size,
+        _ambient=ambient,
+        _opacity=opacity,
+        _clim=clim,
+    )
+
+
+def add_str_legend(
+        plotter,
+        labels,
+        colors,
+        font_family='arial',
+        legend_size=None,
+        legend_loc="center right"
+):
+
+    legend_data = np.concatenate(
+        [labels.reshape(-1, 1).astype(object), colors.reshape(-1, 1).astype(object)], axis=1)
+    legend_data = legend_data[legend_data[:, 0] != "mask", :]
+    assert len(
+        legend_data) != 0, "No legend can be added, please set `show_legend=False`."
+
+    legend_entries = legend_data[np.lexsort(legend_data[:, ::-1].T)]
+    if legend_size is None:
+        legend_num = 10 if len(legend_entries) >= 10 else len(legend_entries)
+        legend_size = (0.1 + 0.01 * legend_num, 0.1 + 0.012 * legend_num)
+
+    plotter.add_legend(
+        legend_entries.tolist(),
+        face="none",
+        font_family=font_family,
+        bcolor=None,
+        loc=legend_loc,
+        size=legend_size
+    )
+
+
+def add_num_legend(
+    plotter,
+    title="",
+    n_labels=5,
+    title_font_size=None,
+    label_font_size=None,
+    font_color="black",
+    font_family="arial",
+    legend_size=(0.1, 0.4),
+    legend_loc=(0.85, 0.3),
+    vertical=True,
+    fmt="%.2e",
+):
+
+    plotter.add_scalar_bar(
+        title=title,
+        n_labels=n_labels,
+        title_font_size=title_font_size,
+        label_font_size=label_font_size,
+        color=font_color,
+        font_family=font_family,
+        use_opacity=True,
+        width=legend_size[0],
+        height=legend_size[1],
+        position_x=legend_loc[0],
+        position_y=legend_loc[1],
+        vertical=vertical,
+        fmt=fmt,
+    )
+
+
+def add_legend(
+    plotter,
+    model,
+    key=None,
+    colormap=None,
+    categorical_legend_size=None,
+    categorical_legend_loc=None,
+    scalar_bar_title="",
+    scalar_bar_size=None,
+    scalar_bar_loc=None,
+    scalar_bar_title_size=None,
+    scalar_bar_label_size=None,
+    scalar_bar_font_color="black",
+    font_family="arial",
+    fmt="%.2e",
+    scalar_bar_n_labels=5,
+    vertical=True,
+):
+
+    # if colormap is None: categorical
+    # if colormap is not None: continuous
+
+    if colormap is None:
+        assert key is not None, "When colormap is None, key cannot be None at the same time."
+
+        if categorical_legend_loc not in categorical_legend_loc_legal and categorical_legend_loc is None:
+            categorical_legend_loc = 'center right'
+
+        if isinstance(model, MultiBlock):
+            keys = key if isinstance(key, list) else [key] * len(model)
+
+            legend_label_data, legend_color_data = [], []
+            for m, k in zip(model, keys, strict=False):
+                legend_label_data.append(np.asarray(m[k]).flatten())
+                legend_color_data.append(np.asarray(
+                    [mpl.colors.to_hex(i) for i in m[f"{k}_rgba"]]).flatten())
+            legend_label_data = np.concatenate(legend_label_data, axis=0)
+            legend_color_data = np.concatenate(legend_color_data, axis=0)
+            print(legend_color_data)
+        else:
+            legend_label_data = np.asarray(model[key]).flatten()
+            legend_color_data = np.asarray(
+                [mpl.colors.to_hex(i) for i in model[f"{key}_rgba"]]).flatten()
+
+        legend_data = np.concatenate(
+            [legend_label_data.reshape(-1, 1), legend_color_data.reshape(-1, 1)], axis=1)
+        unique_legend_data = np.unique(legend_data, axis=0)
+
+        add_str_legend(
+            plotter=plotter,
+            labels=unique_legend_data[:, 0],
+            colors=unique_legend_data[:, 1],
+            font_family=font_family,
+            legend_size=categorical_legend_size,
+            legend_loc=categorical_legend_loc
+        )
+    else:
+        if not isinstance(scalar_bar_size, tuple) and scalar_bar_size is None:
+            scalar_bar_size = (0.1, 0.4)
+        if not isinstance(scalar_bar_loc, tuple) and scalar_bar_loc is None:
+            scalar_bar_loc = (0.85, 0.3)
+
+        add_num_legend(
+            plotter=plotter,
+            legend_size=scalar_bar_size,
+            legend_loc=scalar_bar_loc,
+            title=scalar_bar_title,
+            n_labels=scalar_bar_n_labels,
+            title_font_size=scalar_bar_title_size,
+            label_font_size=scalar_bar_label_size,
+            font_color=scalar_bar_font_color,
+            font_family=font_family,
+            fmt=fmt,
+            vertical=vertical
+        )
+
+
+def add_outline(
+    plotter,
+    model,
+    outline_width=1.0,
+    outline_color="black",
+):
+
+    model.outline()
+    plotter.add_bounding_box(
+        color=outline_color,
+        line_width=outline_width
+    )
+
+
+
+def add_text(
+    plotter,
+    text,
+    font_family="arial",
+    text_font_size=15,
+    text_font_color="black",
+    text_loc="upper_edge"
+):
+
+    plotter.add_text(
+        text=text,
+        font=font_family,
+        color=text_font_color,
+        font_size=text_font_size,
+        position=text_loc if text_loc is not None else "upper_edge"
+    )