sequenzo 0.1.31__cp310-cp310-macosx_10_9_x86_64.whl

sequenzo/visualization/plot_transition_matrix.py

@@ -0,0 +1,190 @@
+"""
+@Author  : Yuqi Liang 梁彧祺, Sebastian Daza
+@File    : plot_transition_matrix.py
+@Time    : 13/02/2025 12:39
+@Desc    :
+    This implementation closely follows ggseqplot's `ggseqtrplot` function.
+
+    TODO: Current implementation only handles STS (State-Transition-State) format.
+          DSS (Distinct-State-Sequence) format will be addressed in future updates.
+          (https://maraab23.github.io/ggseqplot/articles/ggseqplot.html)
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from typing import Optional
+from sequenzo.define_sequence_data import SequenceData
+from sequenzo.visualization.utils import (
+    save_and_show_results,
+    show_plot_title
+)
+
+
+def compute_transition_matrix(seqdata: SequenceData, with_missing: bool = False, weights="auto") -> np.ndarray:
+    """
+    Compute transition rate matrix using vectorized operations for optimal performance.
+
+    :param seqdata: SequenceData object containing sequence information and states
+    :param with_missing: Flag to include missing values in computation
+    :param weights: (np.ndarray or "auto") Weights for sequences. If "auto", uses seqdata.weights if available
+    :return: numpy.ndarray containing the transition rate matrix
+    """
+    # Process weights
+    if isinstance(weights, str) and weights == "auto":
+        weights = getattr(seqdata, "weights", None)
+
+    if weights is not None:
+        weights = np.asarray(weights, dtype=float).reshape(-1)
+        if len(weights) != len(seqdata.values):
+            raise ValueError("Length of weights must equal number of sequences.")
+
+    num_states = len(seqdata.states)
+    A = seqdata.to_dataframe().to_numpy()
+    n, T = A.shape
+
+    if weights is None:
+        w = np.ones(n)
+    else:
+        w = np.asarray(weights, dtype=float)
+
+    # Flatten arrays while synchronizing weights for each transition
+    current = A[:, :-1].flatten()
+    nxt = A[:, 1:].flatten()
+    w_pair = np.repeat(w, T-1)  # Each sequence weight replicated (T-1) times
+
+    # Filter valid transitions (states are encoded as 1, 2, 3, ..., num_states)
+    valid = (current >= 1) & (current <= num_states) & (nxt >= 1) & (nxt <= num_states)
+    current, nxt, w_pair = current[valid], nxt[valid], w_pair[valid]
+
+    # Compute weighted transition counts
+    # Create mapping from state codes to matrix indices
+    state_codes = sorted(set(current) | set(nxt))
+    code_to_idx = {code: idx for idx, code in enumerate(state_codes)}
+
+    # Use only the actual number of unique states for matrix size
+    actual_num_states = len(state_codes)
+    trans = np.zeros((actual_num_states, actual_num_states), dtype=float)
+
+    for c, n2, ww in zip(current, nxt, w_pair):
+        trans[code_to_idx[int(c)], code_to_idx[int(n2)]] += ww
+
+    # Compute transition rates
+    row_sums = trans.sum(axis=1, keepdims=True)
+    row_sums = np.where(row_sums == 0, 1.0, row_sums)
+    transition_rates = trans / row_sums
+
+    # Create a properly sized matrix with correct mapping to original states
+    final_matrix = np.zeros((num_states, num_states), dtype=float)
+
+    # Map back to the original state positions
+    for i, from_code in enumerate(state_codes):
+        for j, to_code in enumerate(state_codes):
+            final_matrix[from_code-1, to_code-1] = transition_rates[i, j]
+
+    return final_matrix
+
+
+def print_transition_matrix(seqdata: SequenceData, transition_rates: np.ndarray) -> None:
+    """
+    Print the transition rate matrix with improved readability.
+    Uses consistent decimal format for all probabilities.
+
+    :param seqdata: SequenceData object containing state information
+    :param transition_rates: numpy array containing transition rates
+    """
+    state_labels = seqdata.labels
+
+    # Calculate max width needed for state labels
+    max_label_width = max(len(s) for s in state_labels) + 3  # +3 for arrow
+
+    # Print header
+    print("\nTransition Rate Matrix:")
+    print("-" * 100)  # Separator line
+
+    # Print column headers
+    print(" " * max_label_width, end=" ")
+    for label in state_labels:
+        # Use shorter labels for columns by taking first word
+        short_label = label.split('&')[0].strip()
+        print(f"{short_label:>10}", end=" ")
+    print("\n" + "-" * 100)  # Separator line
+
+    # Print each row
+    for i, from_state in enumerate(state_labels):
+        # Print row label
+        print(f"{from_state:>{max_label_width}} ->", end=" ")
+
+        # Print transition rates
+        for prob in transition_rates[i]:
+            # Consistent format for all values
+            print(f"{prob:>10.4f}", end=" ")
+        print()  # New line after each row
+
+    print("-" * 100)  # Separator line
+    print("Note: Values represent transition probabilities (range 0-1)")
+    print("      All values shown with 4 decimal places for consistency")
+
+
+def plot_transition_matrix(seqdata: SequenceData,
+                           weights="auto",
+                           title: str = "State Transition Rate Matrix",
+                           fontsize: int = 12,
+                           save_as: Optional[str] = None,
+                           dpi: int = 200,
+                           format: str = ".2f") -> None:
+    """
+    Plot state transition rate matrix as a heatmap.
+
+    :param seqdata: SequenceData object containing sequence information
+    :param weights: (np.ndarray or "auto") Weights for sequences. If "auto", uses seqdata.weights if available
+    :param title: optional title for the plot
+    :param fontsize: base font size for labels
+    :param save_as: optional file path to save the plot
+    :param dpi: resolution of the saved plot
+    :param format: format string for annotations (default "%.2f")
+    """
+
+    # Compute transition matrix with weights
+    transition_matrix = compute_transition_matrix(seqdata, weights=weights)
+    transition_matrix = np.array(transition_matrix)
+
+    # Set figure size
+    plt.figure(figsize=(12, 10))
+
+    # Use fresh color scheme
+    cmap = sns.color_palette("light:#5A9", as_cmap=True)
+
+    # Generate heatmap using pre-formatted annotation strings
+    ax = sns.heatmap(
+        transition_matrix,
+        annot=True,
+        fmt=format,
+        cmap=cmap,
+        xticklabels=seqdata.labels,
+        yticklabels=seqdata.labels,
+        linewidths=0.5,
+        linecolor="gray",
+        cbar_kws={"shrink": 0.8},
+        square=True,
+        annot_kws={"fontsize": fontsize - 2}
+    )
+
+    # Show all the borderlines
+    for _, spine in ax.spines.items():
+        spine.set_visible(True)
+
+    # Adjust format
+    if title:
+        show_plot_title(plt.gca(), title, show=True, fontsize=fontsize+2, fontweight='bold', pad=20)
+
+    plt.xlabel("State at t + 1", fontsize=fontsize, labelpad=10)
+    plt.ylabel("State at t", fontsize=fontsize, labelpad=10)
+
+    # Adjust label rotation and alignment
+    plt.xticks(rotation=45, ha="right")
+    plt.yticks(rotation=0)
+
+    # Adjust layout
+    plt.tight_layout()
+
+    save_and_show_results(save_as, dpi=dpi)

sequenzo/visualization/utils/__init__.py

@@ -0,0 +1,23 @@
+"""
+@Author  : Yuqi Liang 梁彧祺
+@File    : __init__.py.py
+@Time    : 01/03/2025 10:16
+@Desc    : 
+"""
+from .utils import (set_up_time_labels_for_x_axis,
+                    create_standalone_legend,
+                    save_figure_to_buffer,
+                    combine_plot_with_legend,
+                    save_and_show_results,
+                    determine_layout,
+                    show_plot_title,
+                    show_group_title)
+
+__all__ = ['set_up_time_labels_for_x_axis',
+           'create_standalone_legend',
+           'save_figure_to_buffer',
+           'combine_plot_with_legend',
+           'save_and_show_results',
+           'determine_layout',
+           'show_plot_title',
+           'show_group_title']

sequenzo/visualization/utils/utils.py

@@ -0,0 +1,310 @@
+"""
+@Author  : Yuqi Liang 梁彧祺
+@File    : utils.py
+@Time    : 01/03/2025 10:16
+@Desc    : 
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+
+from io import BytesIO
+from PIL import Image
+from typing import Tuple, Optional, List, Dict, Any
+
+from sequenzo import SequenceData
+
+
+def set_up_time_labels_for_x_axis(seqdata: SequenceData,
+                                  ax: Axes,
+                                  color: str = "gray") -> None:
+    """
+    Helper function to set up time labels for the x-axis.
+
+    :param seqdata: (SequenceData) A SequenceData object containing time information
+    :param ax: (matplotlib.axes.Axes) The axes to set labels on
+    """
+    # Extract time labels (year or age)
+    time_labels = np.array(seqdata.cleaned_time)
+
+    # Determine the number of time steps
+    num_time_steps = len(time_labels)
+
+    # Dynamic X-Tick Adjustment
+    if num_time_steps <= 10:
+        # If 10 or fewer time points, show all labels
+        xtick_positions = np.arange(num_time_steps)
+    elif num_time_steps <= 20:
+        # If 10-20 time points, show every 2nd label
+        xtick_positions = np.arange(0, num_time_steps, step=2)
+    else:
+        # More than 20 time points -> Pick 10 evenly spaced tick positions
+        xtick_positions = np.linspace(0, num_time_steps - 1, num=10, dtype=int)
+
+    # Set x-ticks and labels dynamically
+    ax.set_xticks(xtick_positions)
+    ax.set_xticklabels(time_labels[xtick_positions], fontsize=10, rotation=0, ha="center", color=color)
+
+
+def create_standalone_legend(
+        colors: Dict[str, str],
+        labels: List[str],
+        ncol: int = 5,
+        figsize: Tuple[int, int] = (8, 1),
+        fontsize: int = 10,
+        dpi: int = 200
+        ) -> BytesIO:
+    """
+    Creates a standalone legend image without borders.
+
+    Parameters:
+        colors: Dictionary mapping labels to color values
+        labels: List of state labels to include in the legend
+        ncol: Number of columns in the legend
+        figsize: Size of the figure (width, height)
+        fontsize: Font size for legend text
+        dpi: Resolution of the output image
+
+    Returns:
+        BytesIO: In-memory buffer containing the legend image
+    """
+    # Create a new figure for the legend
+    legend_fig = plt.figure(figsize=figsize)
+    ax = legend_fig.add_subplot(111)
+    ax.axis('off')  # Hide the axes
+
+    # Create handles for the legend
+    handles = [plt.Rectangle((0, 0), 1, 1, color=colors.get(s, "gray")) for s in labels]
+
+    # Create the legend without a frame
+    legend = ax.legend(
+        handles,
+        labels,
+        loc='center',
+        ncol=min(ncol, len(labels)),
+        frameon=False,  # No border around legend
+        fontsize=fontsize
+    )
+
+    # Save to memory buffer
+    buffer = BytesIO()
+    legend_fig.savefig(
+        buffer,
+        format='png',
+        dpi=dpi,
+        bbox_inches='tight',
+        pad_inches=0  # Remove padding
+    )
+    plt.close(legend_fig)
+    buffer.seek(0)
+
+    return buffer
+
+
+def combine_plot_with_legend(
+        main_image_buffer: BytesIO,
+        legend_buffer: BytesIO,
+        output_path: Optional[str] = None,
+        dpi: int = 200,
+        padding: int = 10
+) -> Image.Image:
+    """
+    Combines a main plot image with a legend image, placing the legend below the main plot.
+    This means that it saves the combined image to a file if an output path is provided,
+    which is different from the function `save_and_show_results` as that is responsible for visualizations that do not require cropping.
+
+    Parameters:
+        main_image_buffer: Buffer containing the main plot image
+        legend_buffer: Buffer containing the legend image
+        output_path: Optional path to save the combined image
+        dpi: Resolution for the saved image
+        padding: Padding between main image and legend in pixels
+
+    Returns:
+        PIL.Image: The combined image
+    """
+    # Open images from buffers
+    main_img = Image.open(main_image_buffer)
+    legend_img = Image.open(legend_buffer)
+
+    # Calculate dimensions for combined image
+    combined_width = max(main_img.width, legend_img.width)
+    combined_height = main_img.height + padding + legend_img.height
+
+    # Create new blank image
+    combined_img = Image.new('RGB', (combined_width, combined_height), (255, 255, 255))
+
+    # Paste main image at top
+    combined_img.paste(main_img, (0, 0))
+
+    # Center legend relative to main image width (not combined width)
+    legend_x = (main_img.width - legend_img.width) // 2
+    # Ensure legend doesn't go outside the combined image bounds
+    legend_x = max(0, min(legend_x, combined_width - legend_img.width))
+    combined_img.paste(legend_img, (legend_x, main_img.height + padding))
+
+    # Save if output path is provided
+    if output_path:
+        combined_img.save(output_path, dpi=(dpi, dpi))
+
+    return combined_img
+
+
+def save_figure_to_buffer(fig, dpi: int = 200) -> BytesIO:
+    """
+    Saves a matplotlib figure to a BytesIO buffer.
+
+    Parameters:
+        fig: Matplotlib figure to save
+        dpi: Resolution of the output image
+
+    Returns:
+        BytesIO: Buffer containing the figure image
+    """
+    buffer = BytesIO()
+    fig.savefig(buffer, format='png', dpi=dpi, bbox_inches='tight')
+    plt.close(fig)
+    buffer.seek(0)
+    return buffer
+
+
+def save_and_show_results(save_as=None, dpi=200, show=True):
+    """
+    Save and optionally display matplotlib figure.
+
+    Parameters:
+    - save_as: Optional filename to save the figure to
+    - dpi: Resolution for saved image
+    - show: Whether to display the figure (default: True)
+
+    Returns:
+    - The current figure object (plt.gcf())
+    """
+    if save_as:
+        if not any(save_as.endswith(ext) for ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']):
+            save_as += '.png'
+        plt.savefig(save_as, dpi=dpi, bbox_inches='tight')
+
+    # Get the current figure before closing it
+    current_fig = plt.gcf()
+    
+    # Show if explicitly requested or if saving (for visual confirmation)
+    if show or save_as:
+        plt.show()
+
+    # Prevent memory leak by closing after show/save
+    plt.close()
+
+    # Return the figure that was current before closing
+    return current_fig
+
+
+
+def determine_layout(num_items: int,
+                     layout: str = 'column',
+                     nrows: Optional[int] = None,
+                     ncols: Optional[int] = None) -> Tuple[int, int]:
+    """
+    Determine subplot layout (rows, columns) based on the number of items to plot.
+    TODO: 1. change to all visualizations that require multiple graphs 2. 发包，因为我也改了颜色了
+
+    Parameters:
+        num_items (int): Total number of subplots needed.
+        layout (str): Layout strategy ('column' or 'grid'), used if nrows/ncols not provided.
+        nrows (int, optional): Number of rows (manual override).
+        ncols (int, optional): Number of columns (manual override).
+
+    Returns:
+        Tuple[int, int]: A tuple of (nrows, ncols) for subplot layout.
+
+    Raises:
+        ValueError: If layout config is invalid or does not fit all subplots.
+    """
+    # Check partial input
+    if (nrows is None and ncols is not None) or (ncols is None and nrows is not None):
+        raise ValueError("If manually specifying layout, both 'nrows' and 'ncols' must be provided.")
+
+    # Manual override
+    if nrows is not None and ncols is not None:
+        total_slots = nrows * ncols
+        if total_slots < num_items:
+            raise ValueError(f"Provided layout ({nrows}x{ncols}) is too small for {num_items} plots.")
+        return nrows, ncols
+
+    # Automatic layout
+    if layout == 'column':
+        # Use up to 3 columns, but never exceed the number of items.
+        # This avoids allocating invisible (but space-occupying) subplots
+        # that cause extra blank space on the right when num_items < 3.
+        ncols = int(min(3, max(1, num_items)))
+        nrows = (num_items + ncols - 1) // ncols
+    elif layout == 'grid':
+        ncols = int(np.ceil(np.sqrt(num_items)))
+        nrows = (num_items + ncols - 1) // ncols
+    else:
+        raise ValueError(f"Unsupported layout style: '{layout}'")
+
+    return nrows, ncols
+
+
+def show_plot_title(ax: Axes,
+                    title: Optional[str],
+                    show: bool = True,
+                    fontsize: Optional[int] = None,
+                    pad: Optional[int] = None,
+                    loc: Optional[str] = None,
+                    color: Optional[str] = None,
+                    fontweight: Optional[str] = None) -> None:
+    """
+    Conditionally set a matplotlib axis title.
+
+    Parameters:
+        ax: Target axis
+        title: Title text; if None or empty, nothing will be set
+        show: If False, do nothing
+        fontsize, pad, loc, color, fontweight: forwarded to ax.set_title if provided
+    """
+    if not show or not title:
+        return
+    kwargs = {}
+    if fontsize is not None:
+        kwargs["fontsize"] = fontsize
+    if pad is not None:
+        kwargs["pad"] = pad
+    if loc is not None:
+        kwargs["loc"] = loc
+    if color is not None:
+        kwargs["color"] = color
+    if fontweight is not None:
+        kwargs["fontweight"] = fontweight
+    ax.set_title(title, **kwargs)
+
+
+def show_group_title(ax: Axes,
+                     text: Optional[str],
+                     show: bool = True,
+                     fontsize: Optional[int] = None,
+                     x: float = 0.95,
+                     y: float = 1.02,
+                     ha: str = 'right',
+                     va: str = 'bottom',
+                     color: str = 'black') -> None:
+    """
+    Conditionally place a group title as axis-anchored text.
+
+    Parameters:
+        ax: Target axis
+        text: Text content; if None/empty, nothing will be placed
+        show: If False, do nothing
+        fontsize: Font size
+        x, y: Coordinates in axis fraction coordinates
+        ha, va: Horizontal/vertical alignment
+        color: Text color
+    """
+    if not show or not text:
+        return
+    kwargs = {"transform": ax.transAxes, "ha": ha, "va": va, "color": color}
+    if fontsize is not None:
+        kwargs["fontsize"] = fontsize
+    ax.text(x, y, text, **kwargs)
+

sequenzo/with_event_history_analysis/__init__.py

@@ -0,0 +1,35 @@
+"""
+@Author  : Yuqi Liang 梁彧祺
+@File    : __init__.py
+@Time    : 30/09/2025 23:34
+@Desc    : Event History Analysis module for sequence analysis
+"""
+
+from .sequence_analysis_multi_state_model import (
+    SAMM,
+    sequence_analysis_multi_state_model,
+    plot_samm,
+    seqsammseq,
+    set_typology,
+    seqsammeha,
+    # Keep old names for backward compatibility
+    seqsamm
+)
+
+from .sequence_history_analysis import (
+    seqsha,
+    person_level_to_person_period
+)
+
+__all__ = [
+    'SAMM',
+    'sequence_analysis_multi_state_model',
+    'plot_samm',
+    'seqsammseq',
+    'set_typology',
+    'seqsammeha',
+    'seqsha',
+    'person_level_to_person_period',
+    # Keep old names for backward compatibility
+    'seqsamm'
+]