Source code for gunz_cm.visualizations

"""
Visualization utilities for contact matrices, compartments, and metrics.
"""
from __future__ import annotations

__author__ = "Yeremia Gunawan Adhisantoso"
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"

__all__ = [
    "format_bytes",
    "format_seconds",
    "squared_error",
    "mean_squared_error",
    "root_mean_squared_error",
    "absolute_error",
    "mean_absolute_error",
    "re",
    "mre",
    "pearson_corr",
    "spearman_corr",
    "get_compartment_matrix",
    "split_compartments",
    "get_compartment_name",
    "get_compartment_boundaries",
    "get_compartment_sizes",
    "build_centered_logarithmic_norm",
    "plot_chromosome_3d",
    "plot_structure_grid",
    "plot_chromosome_3d_with_ellipsoid",
    "display_contact_map",
    "display_compartment_map",
    "write_vtk_points",
]

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from numpy.typing import NDArray
    import vtk


def format_bytes(value: int, prefix_base: int | None = None) -> str:
    from .units import format_bytes as _format_bytes
    from .units import PrefixBase
    return _format_bytes(value, prefix_base=prefix_base or PrefixBase.Binary)


def format_seconds(value: float) -> str:
    from .units import format_seconds as _format_seconds
    return _format_seconds(value)


def squared_error(first: NDArray, other: NDArray) -> NDArray:
    from .metrics import squared_error
    return squared_error(first, other)


def mean_squared_error(first: NDArray, other: NDArray) -> float:
    from .metrics import mean_squared_error
    return mean_squared_error(first, other)


def root_mean_squared_error(first: NDArray, other: NDArray) -> float:
    from .metrics import root_mean_squared_error
    return root_mean_squared_error(first, other)


def absolute_error(first: NDArray, other: NDArray) -> NDArray:
    from .metrics import absolute_error
    return absolute_error(first, other)


def mean_absolute_error(first: NDArray, other: NDArray) -> float:
    from .metrics import mean_absolute_error
    return mean_absolute_error(first, other)


def re(first: NDArray, other: NDArray) -> NDArray:
    from .metrics import re
    return re(first, other)


def mre(first: NDArray, other: NDArray) -> float:
    from .metrics import mre
    return mre(first, other)


def pearson_corr(first: NDArray, other: NDArray) -> float:
    from .metrics import pearson_corr
    return pearson_corr(first, other)


def spearman_corr(first: NDArray, other: NDArray) -> float:
    from .metrics import spearman_corr
    return spearman_corr(first, other)


def get_compartment_matrix(matrix, compartments: NDArray, x: int, y: int):
    from .compartment import get_compartment_matrix as _get_compartment_matrix
    return _get_compartment_matrix(matrix, compartments, x, y)


def split_compartments(matrix: NDArray, compartments: NDArray) -> NDArray:
    from .compartment import split_compartments
    return split_compartments(matrix, compartments)


def get_compartment_name(index):
    from .compartment import get_compartment_name
    return get_compartment_name(index)


def get_compartment_boundaries(compartments):
    from .compartment import get_compartment_boundaries
    return get_compartment_boundaries(compartments)


def get_compartment_sizes(compartments):
    from .compartment import get_compartment_sizes
    return get_compartment_sizes(compartments)


def build_centered_logarithmic_norm(matrix):
    from .display import build_centered_logarithmic_norm as _build_centered_logarithmic_norm
    return _build_centered_logarithmic_norm(matrix)


def plot_chromosome_3d(
    points: NDArray,
    output_path: str | None = None,
    marker_size: int = 4,
    line_width: int = 2,
    colorscale: str = "Viridis",
    show: bool = False,
    node_colors: NDArray | None = None,
):
    from .structures import plot_chromosome_3d as _plot_chromosome_3d
    return _plot_chromosome_3d(
        points,
        output_path=output_path,
        marker_size=marker_size,
        line_width=line_width,
        colorscale=colorscale,
        show=show,
        node_colors=node_colors,
    )


def plot_structure_grid(
    structures: list[NDArray],
    titles: list[str],
    output_path: str | None = None,
    main_title: str = "Aligned Structures",
    grid_size: tuple[int, int] | None = None,
    *,
    node_colors: NDArray | list[NDArray] | None = None,
    use_cdn: bool = True,
    showscale: bool = True,
):
    from .structures import plot_structure_grid as _plot_structure_grid
    return _plot_structure_grid(
        structures,
        titles,
        output_path=output_path,
        main_title=main_title,
        grid_size=grid_size,
        node_colors=node_colors,
        use_cdn=use_cdn,
        showscale=showscale,
    )


def plot_chromosome_3d_with_ellipsoid(
    points: NDArray,
    output_path: str | None = None,
    *,
    node_colors: NDArray | None = None,
    ellipsoid_color: str = "rgba(0, 200, 200, 0.15)",
    ellipsoid_alpha: float = 0.15,
    confidence: float = 2.0,
    n_mesh: int = 20,
    use_cdn: bool = True,
    showscale: bool = True,
):
    from .structures import (
        plot_chromosome_3d_with_ellipsoid as _plot_chromosome_3d_with_ellipsoid,
    )
    return _plot_chromosome_3d_with_ellipsoid(
        points,
        output_path=output_path,
        node_colors=node_colors,
        ellipsoid_color=ellipsoid_color,
        ellipsoid_alpha=ellipsoid_alpha,
        confidence=confidence,
        n_mesh=n_mesh,
        use_cdn=use_cdn,
        showscale=showscale,
    )


def display_contact_map(matrix, **kwargs):
    from .display import display_contact_map as _display_contact_map
    return _display_contact_map(matrix, **kwargs)


def display_compartment_map(matrix, compartments, **kwargs):
    from .display import display_compartment_map as _display_compartment_map
    return _display_compartment_map(matrix, compartments, **kwargs)


def write_vtk_points(points: NDArray, filename: str) -> None:
    from .paraview_exporter import write_vtk_points as _write_vtk_points
    return _write_vtk_points(points, filename)