gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/report/three_d_plot/three_d_plot_prepare.py

@@ -0,0 +1,202 @@
+import matplotlib as mpl
+import numpy as np
+import pyvista as pv
+from matplotlib.colors import LinearSegmentedColormap
+from pandas.core.frame import DataFrame
+
+p_color = ['#313695', '#4575b4', '#74add1','#fee090', '#fdae61', '#f46d43', '#d73027', '#a50026']
+
+
+
+def _get_default_cmap():
+    if "default_cmap" not in mpl.colormaps():
+        colors = p_color
+        nodes = np.linspace(0, 1, len(p_color))
+
+        mpl.colormaps.register(LinearSegmentedColormap.from_list(
+            "default_cmap", list(zip(nodes, colors, strict=False))))
+    return "default_cmap"
+
+
+def create_plotter(
+    jupyter=False,
+    off_screen=False,
+    window_size=(512, 512),
+    background="white",
+    shape=(1, 1),
+    show_camera_orientation=True,
+    show_axis_orientation=False
+):
+
+    # Create an initial plotting object.
+    _get_default_cmap()
+    plotter = pv.Plotter(
+        off_screen=off_screen,
+        window_size=window_size,
+        notebook=False if jupyter is False else True,
+        lighting="light_kit",
+        shape=shape,
+    )
+
+    # Set the background color of the active render window.
+    plotter.background_color = background
+
+    # Add a camera orientation widget to the active renderer window.
+    if jupyter != "trame":
+        if show_camera_orientation:
+            plotter.add_camera_orientation_widget()
+        elif show_axis_orientation:
+            plotter.add_axes(labels_off=True)
+    return plotter
+
+
+def add_point_labels(
+    model,
+    labels,
+    key_added="groups",
+    where="point_data",
+    colormap="rainbow",
+    alphamap=1.0,
+    mask_color="black",
+    mask_alpha=0.0,
+    inplace=False,
+):
+
+    model = model.copy() if not inplace else model
+    labels = np.asarray(labels).flatten()
+
+    # Set color here if group is of string type.
+    if not np.issubdtype(labels.dtype, np.number):
+        cu_arr = np.sort(np.unique(labels), axis=0).astype(object)
+
+        raw_labels_hex = labels.copy().astype(object)
+        raw_labels_alpha = labels.copy().astype(object)
+        raw_labels_hex[raw_labels_hex =="mask"] = mpl.colors.to_hex(mask_color)
+        raw_labels_alpha[raw_labels_alpha == "mask"] = mask_alpha
+
+        # Set raw hex.
+        if isinstance(colormap, str):
+            if colormap in list(mpl.colormaps()):
+                lscmap = mpl.colormaps[colormap]
+                raw_hex_list = [mpl.colors.to_hex(lscmap(i)) for i in np.linspace(0, 1, len(cu_arr))]
+                for label, color in zip(cu_arr, raw_hex_list, strict=False):
+                    raw_labels_hex[raw_labels_hex == label] = color
+            else:
+                raw_labels_hex[raw_labels_hex !="mask"] = mpl.colors.to_hex(colormap)
+        elif isinstance(colormap, dict):
+            for label, color in colormap.items():
+                raw_labels_hex[raw_labels_hex ==label] = mpl.colors.to_hex(color)
+        elif isinstance(colormap, list) or isinstance(colormap, np.ndarray):
+            raw_hex_list = np.array([mpl.colors.to_hex(color)for color in colormap]).astype(object)
+            for label, color in zip(cu_arr, raw_hex_list, strict=False):
+                raw_labels_hex[raw_labels_hex == label] = color
+        else:
+            raise ValueError(
+                "`colormap` value is wrong." "\nAvailable `colormap` types are: `str`, `list` and `dict`.")
+
+        # Set raw alpha.
+        if isinstance(alphamap, float) or isinstance(alphamap, int):
+            raw_labels_alpha[raw_labels_alpha != "mask"] = alphamap
+        elif isinstance(alphamap, dict):
+            for label, alpha in alphamap.items():
+                raw_labels_alpha[raw_labels_alpha == label] = alpha
+        elif isinstance(alphamap, list) or isinstance(alphamap, np.ndarray):
+            raw_labels_alpha = np.asarray(alphamap).astype(object)
+        else:
+            raise ValueError(
+                "`alphamap` value is wrong." "\nAvailable `alphamap` types are: `float`, `list` and `dict`."
+            )
+
+        # Set rgba.
+        labels_rgba = [mpl.colors.to_rgba(c, alpha=a) for c, a in zip(raw_labels_hex, raw_labels_alpha, strict=False)]
+        labels_rgba = np.array(labels_rgba).astype(np.float32)
+
+        # Added rgba of the labels.
+        if where == "point_data":
+            model.point_data[f"{key_added}_rgba"] = labels_rgba
+        else:
+            model.cell_data[f"{key_added}_rgba"] = labels_rgba
+
+        plot_cmap = None
+    else:
+        plot_cmap = colormap
+
+    # Added labels.
+    if where == "point_data":
+        model.point_data[key_added] = labels
+    else:
+        model.cell_data[key_added] = labels
+
+    return model if not inplace else None, plot_cmap
+
+
+def construct_pc(
+    adata,
+    layer="X",
+    spatial_key="spatial",
+    groupby=None,
+    key_added="groups",
+    mask=None,
+    colormap="default_cmap",
+    alphamap=1.0
+):
+
+    # Ensure mask is a list
+    mask_list = mask if isinstance(mask, list) else [mask] if mask is not None else []
+
+    # Extract spatial coordinates
+    if isinstance(adata, DataFrame):
+        cell_names = np.array(adata.index.tolist())
+        try:
+            bucket_xyz = adata[['sx','sy','sz']].values
+        except KeyError:
+            raise ValueError("Spatial coordinates ('sx','sy','sz') not found in meta data.")
+    else:
+        cell_names = np.array(adata.obs_names.tolist())
+        if spatial_key not in adata.obsm:
+            raise ValueError(f"Spatial key {spatial_key} not found in adata.obsm.")
+        bucket_xyz = adata.obsm[spatial_key].astype(np.float64)
+        if isinstance(bucket_xyz, DataFrame):
+            bucket_xyz = bucket_xyz.values
+
+    # Handle grouping
+    if groupby is None:
+        groups = np.array(["same"] * bucket_xyz.shape[0], dtype=str)
+    elif isinstance(adata, DataFrame):
+        # If adata is a DataFrame, check if groupby is a valid column
+        if groupby not in adata.columns:
+            raise ValueError(f"`groupby` column '{groupby}' not found in DataFrame.")
+        groups = adata[groupby].map(lambda x: "mask" if x in mask_list else x).values
+    else:
+        # If adata is AnnData, check if groupby is in obs or var
+        if groupby in adata.obs_keys():
+            # Group by observation metadata
+            groups = adata.obs[groupby].map(lambda x: "mask" if x in mask_list else x).values
+        elif groupby in adata.var_names or set(groupby) <= set(adata.var_names):
+            # Group by gene expression
+            adata.X if layer == "X" else adata.layers[layer]
+            if isinstance(groupby, str):
+                groupby = [groupby]
+            groups = np.asarray(adata[:, groupby].X.sum(axis=1).flatten())
+        else:
+            raise ValueError(
+                f"`groupby` value '{groupby}' is invalid. "
+                "It must be a column in adata.obs, a gene in adata.var_names, "
+                "or a list of genes in adata.var_names."
+            )
+
+    pc = pv.PolyData(bucket_xyz)
+    _, plot_cmap = add_point_labels(
+        model=pc,
+        labels=groups,
+        key_added=key_added,
+        where="point_data",
+        colormap=colormap,
+        alphamap=alphamap,
+        inplace=True,
+    )
+
+    # The obs_index of each coordinate in the original adata.
+    pc.point_data["obs_index"] = cell_names
+
+    return pc, plot_cmap

gsMap/report/three_d_plot/three_d_plots.py

@@ -0,0 +1,425 @@
+import logging
+import math
+from typing import Literal
+
+import matplotlib as mpl
+import numpy as np
+import pyvista as pv
+from pyvista import MultiBlock, Plotter, PolyData
+
+from .three_d_plot_decorate import add_legend, add_model, add_outline, add_text
+from .three_d_plot_prepare import _get_default_cmap, construct_pc, create_plotter
+
+logger = logging.getLogger(__name__)
+
+_get_default_cmap()
+
+
+def wrap_to_plotter(
+    plotter: Plotter,
+    model: PolyData | MultiBlock,
+    key: str | None = None,
+    colormap: str | list | None = None,
+
+    # parameters for model settings
+    ambient: float = 0.2,
+    opacity: [float,list] = 1,
+    point_size: float = 1,
+    model_style: Literal["points", "surface", "wireframe"] = "surface",
+    font_family: Literal["times", "courier", "arial"] = "arial",
+    background: str = "black",
+    cpo: str | list = "iso",
+    clim: list | None = None,
+    legend_kwargs: dict | None = None,
+    outline_kwargs: dict | None = None,
+    text: str | None = None,
+    scalar_bar_title: str | None = None,
+    text_kwargs: dict | None = None,
+    show_outline: bool = False,
+    show_text: bool = True,
+    show_legend: bool = True
+):
+    """
+    Wrap the model and its settings to a plotter.
+
+    Parameters:
+        plotter (Plotter): The plotter object to wrap the model to.
+        model (Union[PolyData, MultiBlock]): The model to be added to the plotter.
+        key (Optional[str]): The key to identify the model in the plotter.
+        colormap (Optional[Union[str, list]]): The colormap to use for the model.
+        ambient (float): The ambient lighting coefficient.
+        opacity (float): The opacity of the model.
+        point_size (float): The size of the points in the model.
+        model_style (Literal["points", "surface", "wireframe"]): The style of the model.
+        font_family (Literal["times", "courier", "arial"]): The font family to use.
+        background (str): The background color of the plotter.
+        cpo (Union[str, list]): The camera position.
+        legend_kwargs (Optional[dict]): Additional keyword arguments for the legend.
+        outline_kwargs (Optional[dict]): Additional keyword arguments for the outline.
+        text (Optional[str]): The text to add to the plotter.
+        scalar_bar_title (Optional[str]): The title of the scalar bar.
+        text_kwargs (Optional[dict]): Additional keyword arguments for the text.
+        show_outline (bool): Whether to show the outline.
+        show_text (bool): Whether to show the text.
+        show_legend (bool): Whether to show the legend.
+    """
+
+    # Set the bacic settings for the plotter.
+    # plotter.camera_position = cpo
+    bg_rgb = mpl.colors.to_rgb(background)
+    cbg_rgb = (1 - bg_rgb[0], 1 - bg_rgb[1], 1 - bg_rgb[2])
+
+    # Add model(s) basic settings to the plotter.
+    add_model(
+        plotter=plotter,
+        model=model,
+        key=key,
+        colormap=colormap,
+        ambient=ambient,
+        opacity=opacity,
+        point_size=point_size,
+        model_style=model_style,
+        clim = clim
+    )
+
+    # Add legends to the plotter.
+    if show_legend:
+        lg_kwargs = dict(
+            categorical_legend_size=None,
+            categorical_legend_loc=None,
+            scalar_bar_size=None,
+            scalar_bar_loc=None,
+            scalar_bar_title_size=None,
+            scalar_bar_label_size=None,
+            scalar_bar_font_color=cbg_rgb,
+            scalar_bar_n_labels=5,
+            font_family=font_family,
+            fmt="%.1e",
+            vertical=True,
+        )
+        if legend_kwargs is not None:
+            lg_kwargs.update(
+                (k, legend_kwargs[k]) for k in lg_kwargs.keys() & legend_kwargs.keys())
+
+        add_legend(plotter=plotter, model=model, key=key,
+                   colormap=colormap, scalar_bar_title=scalar_bar_title, **lg_kwargs)
+
+    # Add an outline to the plotter.
+    if show_outline:
+        ol_kwargs = dict(
+            outline_width=1.0,
+            outline_color=cbg_rgb,
+        )
+
+        if outline_kwargs is not None:
+            ol_kwargs.update(
+                (k, outline_kwargs[k]) for k in ol_kwargs.keys() & outline_kwargs.keys())
+
+        add_outline(plotter=plotter, model=model, **ol_kwargs)
+
+    # Add text to the plotter.
+    if show_text:
+        t_kwargs = dict(
+            font_family=font_family,
+            text_font_size=12,
+            text_font_color=cbg_rgb,
+            text_loc="upper_edge",
+        )
+
+        if text_kwargs is not None:
+            t_kwargs.update((k, text_kwargs[k])
+                            for k in t_kwargs.keys() & text_kwargs.keys())
+
+        add_text(plotter=plotter, text=text, **t_kwargs)
+
+
+def three_d_plot(
+    # parameters for plotter
+    adata,
+    spatial_key: str,
+    keys: list,
+    cmaps: str | list | dict | None = 'default_cmap',
+    scalar_bar_titles: str | list | None = None,
+    texts:  str | list | None = None,
+
+    window_size: tuple[int, int] | None = None,
+    off_screen: bool = False,
+    shape: tuple | None = None,
+    show_camera_orientation: bool = True,
+    show_axis_orientation: bool = False,
+    jupyter: bool = True,
+
+    # parameters for model settings
+    ambient: float = 0.2,
+    opacity: float | list[float] = 1,
+    point_size: float = 1,
+    clim: list | None = None,
+    model_style: Literal["points", "surface", "wireframe"] = "surface",
+    font_family: Literal["times", "courier", "arial"] = "arial",
+    background: str = "black",
+    cpo: str | list = "iso",
+
+    # parameters for show decoration
+    show_outline: bool = False,
+    show_text: bool = True,
+    show_legend: bool = True,
+
+    # parameters for legends, outline, and text
+    legend_kwargs: dict | None = None,
+    outline_kwargs: dict | None = None,
+    text_kwargs: dict | None = None
+):
+    """
+    Generate a 3D plot using pyvista for spatial data visualization.
+
+    Parameters:
+    - adata: AnnData object containing spatial data.
+    - spatial_key: Key in adata.obs that specifies the spatial coordinates.
+    - keys: List of keys in adata.obs to group the data by.
+    - cmaps: Colormap(s) to use for each group. Can be a string or a list of strings. Default is 'default_cmap'.
+    - scalar_bar_titles: Titles for the scalar bars. Can be a string or a list of strings. Default is None.
+    - texts: List of texts to display on the plot. Default is None.
+    - window_size: Size of the plot window in pixels. Default is None.
+        Note window_size is a tuple of (width, height), representing the width and height of the each subplot.
+    - off_screen: Whether to render the plot off-screen. Default is False.
+    - shape: Shape of the subplot grid. Default is None.
+        For plotting  figues in a grid, shape should be a tuple of (n_rows, n_cols)
+    - show_camera_orientation: Whether to show the camera orientation widget. Default is True.
+    - jupyter: Whether to display the plot in a Jupyter notebook. Default is True.
+    - ambient: Ambient lighting intensity. Default is 0.2.
+    - opacity: Opacity of the plot objects. Default is 1.
+    - point_size: Size of the points in the plot. Default is 1.
+    - model_style: Style of the plot objects. Can be 'points', 'surface', or 'wireframe'. Default is 'surface'.
+    - font_family: Font family for the plot text. Can be 'times', 'courier', or 'arial'. Default is 'arial'.
+    - background: Background color of the plot. Default is 'black'.
+    - cpo: Color by point option. Can be a string or a list. Default is 'iso'.
+    - show_outline: Whether to show the plot outline. Default is False.
+    - show_text: Whether to show the plot text. Default is True.
+    - show_legend: Whether to show the plot legend. Default is True.
+    - legend_kwargs (dict, optional): Additional keyword arguments for the legend, the default values are:
+            categorical_legend_size: [tuple] = None, # (gap, size)
+            categorical_legend_loc: [Literal["upper right", "upper left", "lower left", "lower right",
+                                                    "center left", "center right", "lower center", "upper center"
+                                                    "center"]] = None
+            scalar_bar_title: Optional[str] = None
+            scalar_bar_size: Optional[tuple] = None
+            scalar_bar_loc: Optional[tuple] = None
+            scalar_bar_title_size: Union[int, float] = None
+            scalar_bar_label_size: Union[int, float] = None
+            scalar_bar_font_color: Optional[str] = None
+            scalar_bar_n_labels: int = 5
+            fmt="%.1e",
+            vertical: bool = True
+
+    - outline_kwargs (dict, optional): Additional keyword arguments for the plot outline, the default values are:
+            outline_width: float = 1.0
+            outline_color: Optional[str] = None
+            show_outline_labels: bool = False
+            outline_font_size: Optional[int] = None
+            outline_font_color: Optional[str] = None
+
+    - text_kwargs (dict, optional): Additional keyword arguments for the text, the default values are:
+            text_font_size: Optional[float] = None,
+            text_font_color: Optional[str] = None,
+            text_loc: Optional[Literal["lower_left", "lower_right", "upper_left",
+                                        "upper_right", "lower_edge", "upper_edge",
+                                        "right_edge", "left_edge",]] = None
+
+    Returns:
+    - plotter: The pyvista plotter object.
+    """
+    _get_default_cmap()
+    if isinstance(cmaps, str):
+        cmaps = [cmaps] * len(keys)
+
+    if scalar_bar_titles is None or isinstance(scalar_bar_titles, str):
+        scalar_bar_titles = [scalar_bar_titles] * len(keys)
+
+    if texts is None or isinstance(texts, str):
+        texts = [texts] * len(keys)
+
+    # Build the pyvista object
+    models = pv.MultiBlock()
+    plot_cmaps = []
+    for i, key in enumerate(keys):
+        _model, _plot_cmap = construct_pc(adata=adata.copy(),
+                                          spatial_key=spatial_key,
+                                          groupby=key,
+                                          key_added=key,
+                                          colormap=cmaps[i])
+        models[f"model_{i}"] = _model
+        plot_cmaps.append(_plot_cmap)
+
+    # Set the shape and window size of the plot
+    n_window = len(keys)
+    shape = (math.ceil(n_window / 3), n_window if n_window <
+             3 else 3) if shape is None else shape
+    if isinstance(shape, tuple | list):
+        n_subplots = shape[0] * shape[1]
+        subplots = []
+        for i in range(n_subplots):
+            col = math.floor(i / shape[1])
+            ind = i - col * shape[1]
+            subplots.append([col, ind])
+
+    win_x, win_y = shape[1], shape[0]
+    window_size = ((1500 * win_x, 1500 * win_y)
+                   if window_size is None else (window_size[0] * win_x, window_size[1] * win_y))
+
+    # Create the plotter
+    plotter = create_plotter(
+        background=background,
+        off_screen=off_screen,
+        shape=shape,
+        show_camera_orientation=show_camera_orientation,
+        show_axis_orientation=show_axis_orientation,
+        window_size=window_size,
+        jupyter=jupyter
+    )
+
+    # Set the plotter
+    for (model, key, plot_cmap, subplot_index, scalar_bar_title, text) in zip(models, keys, plot_cmaps, subplots, scalar_bar_titles, texts, strict=False):
+        plotter.subplot(subplot_index[0], subplot_index[1])
+
+        wrap_to_plotter(
+            # parameters for plotter
+            plotter=plotter,
+            model=model,
+            key=key,
+            colormap=plot_cmap,
+
+            # parameters for model settings
+            clim=clim,
+            ambient=ambient,
+            opacity=opacity,
+            point_size=point_size,
+            model_style=model_style,
+            font_family=font_family,
+            background=background,
+            cpo=cpo,
+
+            # parameters for legends, outline, and text
+            legend_kwargs=legend_kwargs,
+            scalar_bar_title=scalar_bar_title,
+            outline_kwargs=outline_kwargs,
+            text_kwargs=text_kwargs,
+            text=text,
+
+            # parameters for show decoration
+            show_outline=show_outline,
+            show_text=show_text,
+            show_legend=show_legend
+        )
+
+    plotter.link_views()
+    # plotter.camera_position = 'yz'
+
+    return plotter
+
+
+def three_d_plot_save(
+    plotter: Plotter,
+    filename: str,
+    view_up_1=(0, 0, 0),
+    view_up_2=(0, 0, 1),
+    n_points: int = 150,
+    factor: float = 2.0,
+    shift: float = 0,
+    step: int = 1,
+    quality: int = 9,
+    framerate: int = 10,
+    save_mp4 : bool = False,
+    save_gif : bool = False
+):
+    """
+    Saves a 3D plot as an HTML, GIF, and MP4 file.
+
+    Args:
+        plotter (Plotter): The Plotter object used for generating the plot.
+        filename (str): The base filename for saving the files.
+        view_up_1 (tuple, optional): The initial view up direction. Defaults to (0.5, 0.5, 1).
+        view_up_2 (tuple, optional): The final view up direction. Defaults to (0, 0, 1).
+        n_points (int, optional): The number of points on the orbital path. Defaults to 150.
+        factor (float, optional): The factor for scaling the orbital path. Defaults to 2.0.
+        shift (float, optional): The shift value for the orbital path. Defaults to 0.
+        step (int, optional): The step size for writing frames. Defaults to 1.
+        quality (int, optional): The quality of the GIF file. Defaults to 9.
+        framerate (int, optional): The framerate of the MP4 file. Defaults to 15.
+    """
+    # save html
+    logger.info('saving 3d plot as html...')
+
+    # Workaround for asyncio conflict: use 'static' backend for export
+    pv.set_jupyter_backend('static')
+    try:
+        plotter.export_html(f'{filename}.html')
+    finally:
+        pv.set_jupyter_backend(pv.global_theme.jupyter_backend)
+
+    # save gif
+    if save_gif:
+        logger.info('saving 3d plot as gif...')
+        path = plotter.generate_orbital_path(factor=factor, shift=shift, viewup=view_up_1, n_points=n_points)
+        plotter.open_gif(filename=f'{filename}.gif')
+        plotter.orbit_on_path(path, write_frames=True, viewup=view_up_2, step=step)
+        plotter.close()
+
+    # save mp4
+    if save_mp4:
+        logger.info('saving 3d plot as mp4...')
+        path = plotter.generate_orbital_path(factor=factor, shift=shift, viewup=view_up_1, n_points=n_points)
+        plotter.open_movie(filename=f'{filename}.mp4', framerate=framerate, quality=quality)
+        plotter.orbit_on_path(path, write_frames=True, viewup=view_up_2, step=step)
+        plotter.close()
+
+
+
+def rotate_around_xyz(
+    camera_coordinates,
+    angle_x=0,
+    angle_y=0,
+    angle_z=0
+):
+    """
+    Rotate a point around the x, y, and z axes.
+
+    Parameters:
+    - point: numpy array representing the camera coordinates
+    - angle_x: rotation angle around the x-axis in degrees (default: 0)
+    - angle_y: rotation angle around the y-axis in degrees (default: 0)
+    - angle_z: rotation angle around the z-axis in degrees (default: 0)
+
+    Returns:
+    - point_rotated: numpy array representing the rotated point coordinates
+    """
+    angle_x_rad = np.radians(angle_x)
+    angle_y_rad = np.radians(angle_y)
+    angle_z_rad = np.radians(angle_z)
+
+    # Rotation matrix for x-axis
+    rotation_matrix_x = np.array([
+        [1, 0, 0],
+        [0, np.cos(angle_x_rad), -np.sin(angle_x_rad)],
+        [0, np.sin(angle_x_rad), np.cos(angle_x_rad)]
+    ])
+
+    # Rotation matrix for y-axis
+    rotation_matrix_y = np.array([
+        [np.cos(angle_y_rad), 0, np.sin(angle_y_rad)],
+        [0, 1, 0],
+        [-np.sin(angle_y_rad), 0, np.cos(angle_y_rad)]
+    ])
+
+    # Rotation matrix for z-axis
+    rotation_matrix_z = np.array([
+        [np.cos(angle_z_rad), -np.sin(angle_z_rad), 0],
+        [np.sin(angle_z_rad), np.cos(angle_z_rad), 0],
+        [0, 0, 1]
+    ])
+
+    # Apply rotations
+    camera_rotated = np.dot(rotation_matrix_x, camera_coordinates)
+    camera_rotated = np.dot(rotation_matrix_y, camera_rotated)
+    camera_rotated = np.dot(rotation_matrix_z, camera_rotated)
+
+    return camera_rotated