sequenzo 0.1.31__cp310-cp310-macosx_10_9_x86_64.whl

sequenzo/suffix_tree/system_level_indicators.py

@@ -0,0 +1,535 @@
+"""
+@Author  : Yuqi Liang 梁彧祺
+@File    : system_level_indicators.py
+@Time    : 02/05/2025 11:06
+@Desc    : 
+    This module includes tools for building suffix trees, computing suffix counts, merging factors, and Jensen-Shannon convergence,
+    as well as generating composite scores to summarize system-level sequence convergence and consolidation over time.
+    Visualization functions are also provided to plot these indicators and their distributions, 
+    supporting comprehensive analysis of sequence system dynamics.
+"""
+from collections import defaultdict, Counter
+import numpy as np
+from scipy.stats import zscore
+from numpy import array
+from scipy.spatial.distance import jensenshannon
+from typing import Optional, Dict, List, Tuple, Any
+
+from sequenzo.visualization.utils import save_and_show_results
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+
+class SuffixTree:
+    def __init__(self):
+        self.root = {}
+        self.counts = defaultdict(int)  # suffix -> count
+        self.total_sequences = 0
+
+    def insert(self, sequence):
+        suffix = []
+        node = self.root
+        for state in sequence:
+            suffix.append(state)
+            key = tuple(suffix)
+            self.counts[key] += 1
+            if state not in node:
+                node[state] = {}
+            node = node[state]
+
+    def get_suffixes_at_depth(self, depth):
+        return [k for k in self.counts if len(k) == depth]
+
+    def get_children(self, suffix):
+        """
+        Given a suffix (as a list or tuple), return its immediate children in the tree.
+
+        Returns:
+            dict: mapping from child state -> subtree dict
+        """
+        node = self.root
+        for state in suffix:
+            node = node.get(state, {})
+        return node
+
+    def get_children_count(self, suffix):
+        node = self.root
+        for state in suffix:
+            node = node.get(state, {})
+        return len(node)
+
+    def describe(self):
+        depths = [len(k) for k in self.counts.keys()]
+        max_depth = max(depths) if depths else 0
+        total_suffixes = len(self.counts)
+        print("\n[SuffixTree Overview]")
+        print(f"[>] Total sequences inserted: {self.total_sequences}")
+        print(f"[>] Max depth (time points): {max_depth}")
+        print(f"[>] Total distinct suffixes: {total_suffixes}")
+
+        for t in range(1, max_depth + 1):
+            level_suffixes = self.get_suffixes_at_depth(t)
+            print(f"    Level {t}: {len(level_suffixes)} unique suffixes")
+
+    def __repr__(self):
+        """
+        Returns a brief textual summary of the suffix tree object.
+
+        Note:
+            This method is intended to provide a lightweight, one-line overview
+            (e.g., max depth and total suffix count). For a full structural report
+            including per-level statistics, use the `.describe()` method instead.
+        """
+        depths = [len(k) for k in self.counts.keys()]
+        return f"SuffixTree(max_depth={max(depths) if depths else 0}, total_suffixes={len(self.counts)})"
+
+
+def get_depth_stats(tree: "SuffixTree") -> Dict[str, Any]:
+    """
+    Build depth-level stats in a single pass over the tree's suffix counts.
+    Use this when calling both compute_suffix_count and compute_merging_factor
+    to avoid scanning the tree twice (important when T or suffix count is large).
+
+    Returns:
+        dict with keys:
+          - 'depth_counts': dict depth -> number of distinct suffixes at that depth
+          - 'depth_to_suffixes': dict depth -> list of suffix tuples at that depth
+    """
+    depth_counts = defaultdict(int)
+    depth_to_suffixes = defaultdict(list)
+    for k in tree.counts:
+        d = len(k)
+        depth_counts[d] += 1
+        depth_to_suffixes[d].append(k)
+    return {
+        "depth_counts": dict(depth_counts),
+        "depth_to_suffixes": dict(depth_to_suffixes),
+    }
+
+
+def compute_suffix_count(
+    tree, max_depth, depth_stats: Optional[Dict[str, Any]] = None
+) -> List[int]:
+    """
+    Suffix counts per time step 1..max_depth.
+    When T is large, pass precomputed depth_stats from get_depth_stats(tree)
+    so that combined with compute_merging_factor only one pass over the tree is used.
+    """
+    if depth_stats is None:
+        depth_counts = defaultdict(int)
+        for k in tree.counts:
+            depth_counts[len(k)] += 1
+        depth_counts = dict(depth_counts)
+    else:
+        depth_counts = depth_stats["depth_counts"]
+    return [depth_counts.get(t, 0) for t in range(1, max_depth + 1)]
+
+
+def compute_merging_factor(
+    tree, max_depth, depth_suffixes: Optional[Dict[int, List[Tuple]]] = None
+) -> List[float]:
+    """
+    Merging factor per time step; first element is 0 to align with suffix count.
+    When T is large, pass depth_suffixes from get_depth_stats(tree)['depth_to_suffixes']
+    to avoid an extra full scan of the tree.
+    """
+    if depth_suffixes is None:
+        depth_to_suffixes = defaultdict(list)
+        for k in tree.counts:
+            depth_to_suffixes[len(k)].append(k)
+        depth_to_suffixes = dict(depth_to_suffixes)
+    else:
+        depth_to_suffixes = depth_suffixes
+    result = []
+    for t in range(2, max_depth + 1):
+        suffixes = depth_to_suffixes.get(t - 1, [])
+        if not suffixes:
+            result.append(0.0)
+            continue
+        child_counts = [tree.get_children_count(s) for s in suffixes]
+        result.append(float(np.mean(child_counts)))
+    return [0.0] + result  # pad to align with suffix count
+
+
+def compute_js_convergence(sequences, state_set):
+    """
+    Jensen-Shannon divergence between consecutive time-step distributions.
+    Uses a single pass over sequences and vectorized numpy operations for speed
+    when T or N is large.
+    """
+    T = len(sequences[0])
+    state_list = list(state_set)
+    n_states = len(state_list)
+    state_to_idx = {s: i for i, s in enumerate(state_list)}
+    N = len(sequences)
+    # Build (N, T) matrix of state indices in one pass
+    mat = np.empty((N, T), dtype=np.intp)
+    for i, seq in enumerate(sequences):
+        for t in range(T):
+            mat[i, t] = state_to_idx[seq[t]]
+    # Per-time distributions via bincount
+    distros = np.zeros((T, n_states), dtype=float)
+    for t in range(T):
+        counts = np.bincount(mat[:, t], minlength=n_states)
+        total = counts.sum()
+        if total > 0:
+            distros[t] = counts / total
+        else:
+            distros[t] = counts
+    js_scores = [0.0]
+    for t in range(1, T):
+        js = jensenshannon(distros[t], distros[t - 1])
+        js_scores.append(float(js))
+    return js_scores
+
+
+def _build_suffix_tree_position(sequences):
+    """Internal: build position-based suffix tree (level = time index from end)."""
+    tree = SuffixTree()
+    tree.total_sequences = len(sequences)
+    for seq in sequences:
+        for t in range(len(seq)):
+            tree.insert(seq[t:])
+    return tree
+
+
+def build_suffix_tree(sequences):
+    """
+    Build position-based suffix tree (level = time index from end).
+
+    For spell-based tree or unified hub with mode/expcost, use:
+        from sequenzo.suffix_tree.hub import build_suffix_tree
+        tree = build_suffix_tree(seqdata, mode="spell", expcost=0)
+    """
+    return _build_suffix_tree_position(sequences)
+
+
+def plot_system_indicators(
+    suffix_counts: List[float],
+    merging_factors: List[float],
+    js_convergence: Optional[List[float]] = None,
+    x_values: Optional[List] = None,
+    x_label: str = "Time (t)",
+    legend_loc: str = 'lower right',
+    save_as: Optional[str] = None,
+    figsize: Optional[tuple] = None,
+    dpi: int = 300,
+    custom_colors: Optional[Dict[str, str]] = None,
+    show: bool = True,
+    plot_distributions: bool = False,
+    style: Optional[str] = None
+) -> None:
+    """
+    Plot a single group's system-level indicators using the same visual style as
+    `plot_system_indicators_multiple_comparison`, but for one subplot.
+
+    Design:
+    - Left y-axis: raw Suffix Count
+    - Right y-axis: z-score of Merging Factor and (optionally) JS Convergence
+    - Consistent colors/markers and legend handling with the multi-comparison API
+
+    Parameters:
+    - suffix_counts: List[float]
+        Raw suffix counts per time step
+    - merging_factors: List[float]
+        Merging factor per time step
+    - js_convergence: Optional[List[float]]
+        JS convergence per time step; if None, only merging factor is shown on right axis
+    - x_values: Optional[List]
+        Custom x-axis ticks (e.g., years). If None, uses 1..T. Length must equal data length
+    - x_label: str
+        Label for x-axis. Default: "Time (t)"
+    - legend_loc: str
+        Legend location, e.g., 'upper left', 'upper right', 'lower right', 'best', etc. Default: 'lower right'
+    - save_as: Optional[str]
+        If provided, save the figure to this path (png). DPI controlled by `dpi`
+    - figsize: Optional[tuple]
+        Figure size (width, height). Default: (12, 6)
+    - dpi: int
+        Figure DPI when saving. Default: 300
+    - custom_colors: Optional[Dict[str, str]]
+        Optional color overrides. Keys: "Suffix Count", "Merging Factor", "JS Convergence"
+    - show: bool
+        Whether to display the figure
+    - plot_distributions: bool
+        If True, additionally show raw distributions (histograms) of indicators
+
+    Example:
+    >>> plot_system_indicators(
+    ...     suffix_counts=india_suffix_counts,
+    ...     merging_factors=india_merging_factors,
+    ...     js_convergence=india_js_scores,
+    ...     x_values=[2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019],
+    ...     x_label="Year",
+    ...     legend_loc="lower right",
+    ...     figsize=(12, 6),
+    ...     dpi=300,
+    ... )
+    """
+    T = len(suffix_counts)
+    # Set x values to align with multi-group API
+    if x_values is None:
+        x_values = list(range(1, T + 1))
+    if len(x_values) != T:
+        raise ValueError("Length of x_values must match data length")
+
+    # Normalize others
+    mf_z = zscore(array(merging_factors))
+    js_z = zscore(array(js_convergence)) if js_convergence else None
+
+    color_defaults = {
+        "Suffix Count": "#6BB6FF",    # Soft sky blue (like Monet's water lilies)
+        "Merging Factor": "#FFB347",  # Warm peach/coral (like sunset reflections)
+        "JS Convergence": "#98D8C8",  # Soft mint green (convergence = coming together)
+    }
+    colors = {**color_defaults, **(custom_colors or {})}
+
+    # --- Main line plot with dual axes ---
+    if figsize is None:
+        figsize = (12, 6)
+    
+    # Apply style if specified
+    if style is not None:
+        # Check if it's a seaborn style
+        seaborn_styles = ['whitegrid', 'darkgrid', 'white', 'dark', 'ticks']
+        if style in seaborn_styles:
+            sns.set_style(style)
+        else:
+            plt.style.use(style)
+    
+    fig, ax1 = plt.subplots(figsize=figsize)
+    ax1.set_xlabel(x_label)
+    ax1.set_ylabel("Suffix Count", color=colors["Suffix Count"])
+    ax1.plot(x_values, suffix_counts, marker='o', color=colors["Suffix Count"], label="Suffix Count")
+    ax1.tick_params(axis='y', labelcolor=colors["Suffix Count"])
+
+    ax2 = ax1.twinx()
+    ax2.set_ylabel("Z-score (Other Indicators)")
+    ax2.plot(x_values, mf_z, marker='s', label='Merging Factor (z)', color=colors["Merging Factor"])
+    if js_z is not None:
+        ax2.plot(x_values, js_z, marker='^', label='JS Convergence (z)', color=colors["JS Convergence"])
+
+    lines1, labels1 = ax1.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax2.legend(lines1 + lines2, labels1 + labels2, loc=legend_loc)
+
+    ax1.set_title("System-Level Trajectory Indicators: Raw vs. Normalized")
+    fig.tight_layout()
+
+    save_and_show_results(save_as=save_as, dpi=dpi, show=show)
+
+    # --- Distribution plots if requested ---
+    if plot_distributions:
+        raw_data = {
+            "Suffix Count": suffix_counts,
+            "Merging Factor": merging_factors,
+        }
+        if js_convergence:
+            raw_data["JS Convergence"] = js_convergence
+
+        n = len(raw_data)
+        fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))
+        if n == 1:
+            axes = [axes]
+
+        for ax, (label, values) in zip(axes, raw_data.items()):
+            sns.histplot(values, kde=True, ax=ax, color=colors.get(label, None))
+            ax.set_title(f"{label} Distribution")
+            ax.set_xlabel("Value")
+            ax.set_ylabel("Density")
+
+        fig.tight_layout()
+        suffix = "_distributions" if save_as else None
+        dist_path = save_as.replace(".png", f"{suffix}.png") if save_as else None
+        save_and_show_results(save_as=dist_path, dpi=dpi, show=show)
+
+
+def plot_system_indicators_multiple_comparison(
+    groups_data: Dict[str, Dict[str, List[float]]],
+    group_names: Optional[List[str]] = None,
+    subplot_titles: Optional[List[str]] = None,
+    x_values: Optional[List] = None,
+    x_label: str = "Time (t)",
+    legend_loc: str = 'lower right',
+    save_as: Optional[str] = None,
+    figsize: Optional[tuple] = None,
+    dpi: int = 300,
+    custom_colors: Optional[Dict[str, str]] = None,
+    show: bool = True,
+    style: Optional[str] = None
+) -> None:
+    """
+    Plot system-level indicators comparison across multiple groups using dual y-axis design.
+    
+    Parameters:
+    -----------
+    groups_data : Dict[str, Dict[str, List[float]]]
+        Dictionary with group names as keys and data dictionaries as values.
+        Each data dict should contain 'suffix_counts', 'merging_factors', and 'js_convergence'.
+        Example: {
+            "Group1": {
+                "suffix_counts": [10, 15, 20, ...],
+                "merging_factors": [1.2, 1.5, 1.8, ...], 
+                "js_convergence": [0.1, 0.2, 0.15, ...]
+            },
+            "Group2": {...}
+        }
+    group_names : Optional[List[str]]
+        Custom names for groups. If None, uses keys from groups_data.
+        Used for default subplot titles if subplot_titles is not provided.
+    subplot_titles : Optional[List[str]]
+        Custom titles for each subplot. If None, uses default format:
+        "{group_name} - System-Level Trajectory Indicators: Raw vs. Normalized"
+    x_values : Optional[List]
+        Custom x-axis values. If None, uses 1, 2, 3, ...
+    x_label : str
+        Label for x-axis. Default: "Time (t)"
+    legend_loc : str
+        Legend location. Options: 'upper left', 'upper right', 'lower left', 
+        'lower right', 'center', 'best', etc. Default: 'lower right'
+    save_as : Optional[str]
+        File path to save the plot (without extension)
+    figsize : Optional[tuple]
+        Figure size (width, height). If None, auto-calculated based on number of groups
+    dpi : int
+        DPI for saving. Default: 300
+    custom_colors : Optional[Dict[str, str]]
+        Custom colors for indicators. Default uses standard colors.
+    show : bool
+        Whether to show the plot. Default: True
+    style : Optional[str]
+        Style to apply. Seaborn styles ('whitegrid', 'darkgrid', 'white', 'dark', 'ticks') 
+        or matplotlib styles. If None, uses default style. Default: None
+        
+    Example:
+    --------
+    >>> data = {
+    ...     "India": {
+    ...         "suffix_counts": india_suffix_counts,
+    ...         "merging_factors": india_merging_factors,
+    ...         "js_convergence": india_js_scores
+    ...     },
+    ...     "US": {
+    ...         "suffix_counts": us_suffix_counts,
+    ...         "merging_factors": us_merging_factors,
+    ...         "js_convergence": us_js_scores
+    ...     }
+    ... }
+    >>> plot_system_indicators_multiple_comparison(
+    ...     groups_data=data,
+    ...     x_label="Years",
+    ...     legend_loc='upper right',
+    ...     save_as="multi_country_comparison"
+    ... )
+    
+    >>> # With custom subplot titles
+    >>> plot_system_indicators_multiple_comparison(
+    ...     groups_data=data,
+    ...     subplot_titles=["印度发展轨迹", "美国发展轨迹"],
+    ...     x_label="年份",
+    ...     save_as="custom_titles_comparison"
+    ... )
+    """
+    
+    # Validate input
+    if not groups_data:
+        raise ValueError("groups_data cannot be empty")
+    
+    # Get group names
+    if group_names is None:
+        group_names = list(groups_data.keys())
+    
+    if len(group_names) != len(groups_data):
+        raise ValueError("Length of group_names must match number of groups in groups_data")
+    
+    # Validate subplot_titles
+    if subplot_titles is not None and len(subplot_titles) != len(groups_data):
+        raise ValueError("Length of subplot_titles must match number of groups in groups_data")
+    
+    # Get first group to determine data length
+    first_group_data = list(groups_data.values())[0]
+    T = len(first_group_data['suffix_counts'])
+    
+    # Set x values
+    if x_values is None:
+        x_values = list(range(1, T + 1))
+    
+    if len(x_values) != T:
+        raise ValueError("Length of x_values must match data length")
+    
+    # Color settings
+    color_defaults = {
+        "Suffix Count": "#6BB6FF",    # Soft sky blue (like Monet's water lilies)
+        "Merging Factor": "#FFB347",  # Warm peach/coral (like sunset reflections)
+        "JS Convergence": "#98D8C8",  # Soft mint green (convergence = coming together)
+    }
+    colors = {**color_defaults, **(custom_colors or {})}
+    
+    # Calculate figure size
+    n_groups = len(groups_data)
+    if figsize is None:
+        figsize = (12, 4 * n_groups + 2)  # Dynamic height based on number of groups
+    
+    # Apply style if specified
+    if style is not None:
+        # Check if it's a seaborn style
+        seaborn_styles = ['whitegrid', 'darkgrid', 'white', 'dark', 'ticks']
+        if style in seaborn_styles:
+            sns.set_style(style)
+        else:
+            plt.style.use(style)
+    
+    # Create subplots
+    fig, axes = plt.subplots(n_groups, 1, figsize=figsize)
+    
+    # Handle single group case
+    if n_groups == 1:
+        axes = [axes]
+    
+    # Plot each group
+    for i, (group_key, group_name) in enumerate(zip(groups_data.keys(), group_names)):
+        data = groups_data[group_key]
+        ax = axes[i]
+        
+        # Validate data completeness
+        required_keys = ['suffix_counts', 'merging_factors', 'js_convergence']
+        for key in required_keys:
+            if key not in data:
+                raise ValueError(f"Missing '{key}' in data for group '{group_key}'")
+        
+        # Normalize data (z-score)
+        mf_z = zscore(array(data['merging_factors']))
+        js_z = zscore(array(data['js_convergence']))
+        
+        # Left y-axis: raw suffix counts
+        ax.set_ylabel("Suffix Count", color=colors["Suffix Count"])
+        ax.plot(x_values, data['suffix_counts'], marker='o', 
+                color=colors["Suffix Count"], label="Suffix Count")
+        ax.tick_params(axis='y', labelcolor=colors["Suffix Count"])
+        
+        # Right y-axis: normalized indicators
+        ax_twin = ax.twinx()
+        ax_twin.set_ylabel("Z-score (Other Indicators)")
+        ax_twin.plot(x_values, mf_z, marker='s', 
+                     label='Merging Factor (z)', color=colors["Merging Factor"])
+        ax_twin.plot(x_values, js_z, marker='^', 
+                     label='JS Convergence (z)', color=colors["JS Convergence"])
+        
+        # Legend
+        lines1, labels1 = ax.get_legend_handles_labels()
+        lines2, labels2 = ax_twin.get_legend_handles_labels()
+        ax_twin.legend(lines1 + lines2, labels1 + labels2, loc=legend_loc)
+        
+        # Title and labels
+        if subplot_titles is not None:
+            title = subplot_titles[i]
+        else:
+            title = f"{group_name} - System-Level Trajectory Indicators: Raw vs. Normalized"
+        ax.set_title(title)
+        
+        # Only set x-label for the bottom subplot
+        if i == n_groups - 1:
+            ax.set_xlabel(x_label)
+    
+    plt.tight_layout()
+    save_and_show_results(save_as=save_as, dpi=dpi, show=show)

sequenzo/suffix_tree/utils.py

@@ -0,0 +1,56 @@
+"""
+@Author  : Yuqi Liang 梁彧祺
+@File    : utils.py
+@Time    : 08/08/2025 12:26
+@Desc    : 
+    Utility functions for suffix tree convergence analysis, 
+    including sequence extraction and state space processing.
+"""
+import pandas as pd
+from typing import List, Tuple
+
+
+def extract_sequences(df: pd.DataFrame, time_cols: List[str]) -> List[List[str]]:
+    """
+    Efficiently extracts sequences from specified time columns.
+
+    Parameters:
+        df (pd.DataFrame): Input DataFrame.
+        time_cols (List[str]): Columns representing the sequence over time.
+
+    Returns:
+        List[List[str]]: List of sequences (each sequence is a list of states).
+    """
+    return df[time_cols].values.tolist()
+
+
+def get_state_space(sequences: List[List[str]]) -> List[str]:
+    """
+    Efficiently extracts unique states from a list of sequences.
+
+    Parameters:
+        sequences (List[List[str]]): Sequence data.
+
+    Returns:
+        List[str]: Sorted list of unique states.
+    """
+    seen = set()
+    for seq in sequences:
+        seen.update(seq)
+    return sorted(seen)
+
+
+def convert_to_suffix_tree_data(df: pd.DataFrame, time_cols: List[str]) -> Tuple[List[List[str]], List[str]]:
+    """
+    Wrapper to extract sequences and their state space from a DataFrame for suffix tree analysis.
+
+    Parameters:
+        df (pd.DataFrame): Input DataFrame.
+        time_cols (List[str]): Sequence columns (e.g., ['C1', ..., 'C10'])
+
+    Returns:
+        Tuple[List[List[str]], List[str]]: sequences, unique states
+    """
+    sequences = df[time_cols].values.tolist()
+    states = get_state_space(sequences)
+    return sequences, states