sequenzo 0.1.17__cp39-cp39-macosx_10_9_universal2.whl → 0.1.19__cp39-cp39-macosx_10_9_universal2.whl

sequenzo/sequence_characteristics/variance_of_spell_durations.py

@@ -1,5 +1,5 @@
 """
-@Author  : 李欣怡
+@Author  : Xinyi Li, Yuqi Liang
 @File    : variance_of_spell_durations.py
 @Time    : 2025/9/24 14:22
 @Desc    : Variance of spell durations of individual state sequences.
@@ -22,9 +22,9 @@ from .simple_characteristics import cut_prefix
 
 def get_spell_duration_variance(seqdata, type=1):
     if not hasattr(seqdata, 'seqdata'):
-        raise ValueError(" [!] data is NOT a sequence object, see SequenceData function to create one.")
+        raise ValueError("[!] data is NOT a sequence object, see SequenceData function to create one.")
     if type not in [1, 2]:
-        raise ValueError(" [!] type must be 1 or 2.")
+        raise ValueError("[!] type must be 1 or 2.")
 
     with open(os.devnull, 'w') as fnull:
         with redirect_stdout(fnull):
@@ -33,7 +33,7 @@ def get_spell_duration_variance(seqdata, type=1):
             lgth = seqlength(seqdata)
             dlgth = seqlength(dss)
             sdist = get_state_freq_and_entropy_per_seq(seqdata)
-            nnvisit = (sdist==0).sum(axis=1)
+            nnvisit = (sdist.iloc[:, 1:]==0).sum(axis=1)
 
     def realvar(x):
         n = len(x)
@@ -57,7 +57,8 @@ def get_spell_duration_variance(seqdata, type=1):
         # ret = (np.nansum(ddur, axis=1) + nnvisit * (meand ** 2)) / (dlgth + nnvisit)
         ddur = pd.DataFrame(ddur.tolist())
         sum_sqdiff = np.nansum(ddur.to_numpy(), axis=1)
-        ret = (sum_sqdiff + nnvisit.to_numpy() * (meand.to_numpy() ** 2)) / (dlgth + nnvisit.to_numpy())
+        ret_values = (sum_sqdiff + nnvisit.to_numpy() * (meand.to_numpy() ** 2)) / (dlgth + nnvisit.to_numpy())
+        ret = pd.Series(ret_values, index=meand.index)
 
         alph = seqdata.states.copy()
         alph_size = len(alph)
@@ -67,10 +68,19 @@ def get_spell_duration_variance(seqdata, type=1):
             maxnnv = np.where(dlgth == 1, alph_size - 1, alph_size - 2)
 
         meand_max = meand.to_numpy() * (dlgth + nnvisit.to_numpy()) / (dlgth + maxnnv)
-        var_max = ((dlgth-1) * (1-meand_max)**2 + (lgth - dlgth + 1 - meand_max)**2 + maxnnv * meand_max**2) / (dlgth + maxnnv)
+        var_max_values = ((dlgth-1) * (1-meand_max)**2 + (lgth - dlgth + 1 - meand_max)**2 + maxnnv * meand_max**2) / (dlgth + maxnnv)
+        var_max = pd.Series(var_max_values, index=meand.index)
+
+    meand.index = seqdata.seqdata.index
+    ret.index = seqdata.seqdata.index
+    var_max.index = seqdata.seqdata.index
+
+    meand = meand.to_frame("meand")
+    ret = ret.to_frame("var_spell_dur")
+    var_max = var_max.to_frame("var_max")
 
     return {
-        "meand": meand,
-        "result": ret,
-        "vmax": var_max
+        "meand": meand.reset_index().rename(columns={"index": "ID"}),
+        "result": ret.reset_index().rename(columns={"index": "ID"}),
+        "vmax": var_max.reset_index().rename(columns={"index": "ID"}),
     }

sequenzo/sequence_characteristics/within_sequence_entropy.py

@@ -19,7 +19,7 @@ from .state_frequencies_and_entropy_per_sequence import get_state_freq_and_entro
 
 def get_within_sequence_entropy(seqdata, norm=True, base=np.e, silent=True):
     if not isinstance(seqdata, SequenceData):
-        raise ValueError(" [!] data is NOT a sequence object, see SequenceData function to create one.")
+        raise ValueError("[!] data is NOT a sequence object, see SequenceData function to create one.")
 
     states = seqdata.states.copy()
 
@@ -29,68 +29,15 @@ def get_within_sequence_entropy(seqdata, norm=True, base=np.e, silent=True):
     with open(os.devnull, 'w') as fnull:
         with redirect_stdout(fnull):
             iseqtab = get_state_freq_and_entropy_per_seq(seqdata=seqdata)
+            iseqtab.index = seqdata.seqdata.index
 
-    ient = iseqtab.apply(lambda row: entropy(row, base=base), axis=1)
+    ient = iseqtab.iloc[:, 1:].apply(lambda row: entropy(row, base=base), axis=1)
 
     if norm:
         maxent = np.log(len(states))
         ient = ient / maxent
 
-    ient.columns = ['Entropy']
-    ient.index = seqdata.seqdata.index
+    ient = pd.DataFrame(ient, index=seqdata.seqdata.index, columns=['Entropy'])
+    ient = ient.reset_index().rename(columns={'index': 'ID'})
 
     return ient
-
-
-if __name__ == "__main__":
-    # ===============================
-    #             Sohee
-    # ===============================
-    # df = pd.read_csv('D:/college/research/QiQi/sequenzo/data_and_output/orignal data/sohee/sequence_data.csv')
-    # time_list = list(df.columns)[1:133]
-    # states = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
-    # # states = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
-    # labels = ['FT+WC', 'FT+BC', 'PT+WC', 'PT+BC', 'U', 'OLF']
-    # sequence_data = SequenceData(df, time=time_list, states=states, labels=labels, id_col="PID")
-    # res = seqient(sequence_data)
-
-    # ===============================
-    #             kass
-    # ===============================
-    # df = pd.read_csv('D:/college/research/QiQi/sequenzo/files/orignal data/kass/wide_civil_final_df.csv')
-    # time_list = list(df.columns)[1:]
-    # states = ['Extensive Warfare', 'Limited Violence', 'No Violence', 'Pervasive Warfare', 'Prolonged Warfare',
-    #           'Serious Violence', 'Serious Warfare', 'Sporadic Violence', 'Technological Warfare', 'Total Warfare']
-    # sequence_data = SequenceData(df, time=time_list, states=states, id_col="COUNTRY")
-    # res = seqient(sequence_data)
-
-    # ===============================
-    #             CO2
-    # ===============================
-    # df = pd.read_csv("D:/country_co2_emissions_missing.csv")
-    # _time = list(df.columns)[1:]
-    # states = ['Very Low', 'Low', 'Middle', 'High', 'Very High']
-    # sequence_data = SequenceData(df, time=_time, id_col="country", states=states)
-    # res = seqient(sequence_data)
-
-    # ===============================
-    #            detailed
-    # ===============================
-    # df = pd.read_csv("D:/college/research/QiQi/sequenzo/data_and_output/sampled_data_sets/detailed_data/sampled_1000_data.csv")
-    # _time = list(df.columns)[4:]
-    # states = ['data', 'data & intensive math', 'hardware', 'research', 'software', 'software & hardware', 'support & test']
-    # sequence_data = SequenceData(df[['worker_id', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C10']],
-    #                              time=_time, id_col="worker_id", states=states)
-    # res = seqient(sequence_data)
-
-    # ===============================
-    #             broad
-    # ===============================
-    df = pd.read_csv("D:/college/research/QiQi/sequenzo/data_and_output/sampled_data_sets/broad_data/sampled_1000_data.csv")
-    _time = list(df.columns)[4:]
-    states = ['Non-computing', 'Non-technical computing', 'Technical computing']
-    sequence_data = SequenceData(df[['worker_id', 'C1', 'C2', 'C3', 'C4', 'C5']],
-                                 time=_time, id_col="worker_id", states=states)
-    res = get_within_sequence_entropy(sequence_data)
-
-    print(res)

sequenzo/visualization/plot_sequence_index.py

@@ -289,10 +289,12 @@ def sort_sequences_by_method(seqdata, method="unsorted", mask=None, distance_mat
 
 
 def plot_sequence_index(seqdata: SequenceData,
-                        show_by_category=None,
-                        category_labels=None,
-                        id_group_df=None,
-                        categories=None,
+                        # Grouping parameters
+                        group_by_column=None,
+                        group_dataframe=None,
+                        group_column_name=None,
+                        group_labels=None,
+                        # Other parameters
                         sort_by="lexicographic",
                         sort_by_weight=False,
                         weights="auto",
@@ -320,13 +322,34 @@ def plot_sequence_index(seqdata: SequenceData,
     This function creates index plots that visualize sequences as horizontal lines,
     with different sorting options matching R's TraMineR functionality.
 
+    **Two API modes for grouping:**
+    
+    1. **Simplified API** (when grouping info is already in the data):
+       ```python
+       plot_sequence_index(seqdata, group_by_column="Cluster", group_labels=cluster_labels)
+       ```
+    
+    2. **Complete API** (when grouping info is in a separate dataframe):
+       ```python
+       plot_sequence_index(seqdata, group_dataframe=membership_df, 
+                          group_column_name="Cluster", group_labels=cluster_labels)
+       ```
+
     :param seqdata: SequenceData object containing sequence information
-    :param show_by_category: (str, optional) Simple way to create grouped plots. 
-                            Specify the column name from the original data (e.g., "sex", "education").
-                            This will automatically create separate plots for each category.
-    :param category_labels: (dict, optional) Custom labels for category values. 
-                           Example: {0: "Female", 1: "Male"} or {"low": "Low Education", "high": "High Education"}.
-                           If not provided, will use original values or auto-generate readable labels.
+    
+    **New API parameters (recommended):**
+    :param group_by_column: (str, optional) Column name from seqdata.data to group by.
+                           Use this when grouping information is already in your data.
+                           Example: "Cluster", "sex", "education"
+    :param group_dataframe: (pd.DataFrame, optional) Separate dataframe containing grouping information.
+                           Use this when grouping info is in a separate table (e.g., clustering results).
+                           Must contain ID column and grouping column.
+    :param group_column_name: (str, optional) Name of the grouping column in group_dataframe.
+                             Required when using group_dataframe.
+    :param group_labels: (dict, optional) Custom labels for group values.
+                        Example: {1: "Late Family Formation", 2: "Early Partnership"}
+                        Maps original values to display labels.
+    
     :param sort_by: Sorting method for sequences within groups:
         - 'unsorted' or 'none': Keep original order (R TraMineR default)
         - 'lexicographic': Sort sequences lexicographically
@@ -392,45 +415,45 @@ def plot_sequence_index(seqdata: SequenceData,
     
     actual_figsize = style_sizes[plot_style]
     
-    # Handle the new simplified API: show_by_category
-    if show_by_category is not None:
+    # Handle the simplified API: group_by_column
+    if group_by_column is not None:
         # Validate that the column exists in the original data
-        if show_by_category not in seqdata.data.columns:
+        if group_by_column not in seqdata.data.columns:
             available_cols = [col for col in seqdata.data.columns if col not in seqdata.time and col != seqdata.id_col]
             raise ValueError(
-                f"Column '{show_by_category}' not found in the data. "
+                f"Column '{group_by_column}' not found in the data. "
                 f"Available columns for grouping: {available_cols}"
             )
         
-        # Automatically create id_group_df and categories from the simplified API
-        id_group_df = seqdata.data[[seqdata.id_col, show_by_category]].copy()
-        id_group_df.columns = ['Entity ID', 'Category']
-        categories = 'Category'
+        # Automatically create group_dataframe and group_column_name from the simplified API
+        group_dataframe = seqdata.data[[seqdata.id_col, group_by_column]].copy()
+        group_dataframe.columns = ['Entity ID', 'Category']
+        group_column_name = 'Category'
         
-        # Handle category labels - flexible and user-controllable
-        unique_values = seqdata.data[show_by_category].unique()
+        # Handle group labels - flexible and user-controllable
+        unique_values = seqdata.data[group_by_column].unique()
         
-        if category_labels is not None:
+        if group_labels is not None:
             # User provided custom labels - use them
-            missing_keys = set(unique_values) - set(category_labels.keys())
+            missing_keys = set(unique_values) - set(group_labels.keys())
             if missing_keys:
                 raise ValueError(
-                    f"category_labels missing mappings for values: {missing_keys}. "
-                    f"Please provide labels for all unique values in '{show_by_category}': {sorted(unique_values)}"
+                    f"group_labels missing mappings for values: {missing_keys}. "
+                    f"Please provide labels for all unique values in '{group_by_column}': {sorted(unique_values)}"
                 )
-            id_group_df['Category'] = id_group_df['Category'].map(category_labels)
+            group_dataframe['Category'] = group_dataframe['Category'].map(group_labels)
         else:
             # No custom labels provided - use smart defaults
             if all(isinstance(v, (int, float, np.integer, np.floating)) and not pd.isna(v) for v in unique_values):
-                # Numeric values - keep as is (user can provide category_labels if they want custom names)
+                # Numeric values - keep as is (user can provide group_labels if they want custom names)
                 pass
             # For string/categorical values, keep original values
             # This handles cases where users already have meaningful labels like "Male"/"Female"
         
-        print(f"[>] Creating grouped plots by '{show_by_category}' with {len(unique_values)} categories")
+        print(f"[>] Creating grouped plots by '{group_by_column}' with {len(unique_values)} categories")
     
     # If no grouping information, create a single plot
-    if id_group_df is None or categories is None:
+    if group_dataframe is None or group_column_name is None:
         return _sequence_index_plot_single(seqdata, sort_by, sort_by_weight, weights, actual_figsize, plot_style, title, xlabel, ylabel, save_as, dpi, fontsize, include_legend, sequence_selection, n_sequences, show_sequence_ids)
 
     # Process weights
@@ -443,21 +466,21 @@ def plot_sequence_index(seqdata: SequenceData,
             raise ValueError("Length of weights must equal number of sequences.")
     
     # Ensure ID columns match (convert if needed)
-    id_col_name = "Entity ID" if "Entity ID" in id_group_df.columns else id_group_df.columns[0]
+    id_col_name = "Entity ID" if "Entity ID" in group_dataframe.columns else group_dataframe.columns[0]
 
     # Get unique groups and sort them based on user preference
     if group_order:
         # Use manually specified order, filter out non-existing groups
-        groups = [g for g in group_order if g in id_group_df[categories].unique()]
-        missing_groups = [g for g in id_group_df[categories].unique() if g not in group_order]
+        groups = [g for g in group_order if g in group_dataframe[group_column_name].unique()]
+        missing_groups = [g for g in group_dataframe[group_column_name].unique() if g not in group_order]
         if missing_groups:
             print(f"[Warning] Groups not in group_order will be excluded: {missing_groups}")
     elif sort_groups == 'numeric' or sort_groups == 'auto':
-        groups = smart_sort_groups(id_group_df[categories].unique())
+        groups = smart_sort_groups(group_dataframe[group_column_name].unique())
     elif sort_groups == 'alpha':
-        groups = sorted(id_group_df[categories].unique())
+        groups = sorted(group_dataframe[group_column_name].unique())
     elif sort_groups == 'none':
-        groups = list(id_group_df[categories].unique())
+        groups = list(group_dataframe[group_column_name].unique())
     else:
         raise ValueError(f"Invalid sort_groups value: {sort_groups}. Use 'auto', 'numeric', 'alpha', or 'none'.")
     
@@ -477,7 +500,7 @@ def plot_sequence_index(seqdata: SequenceData,
     # Create a plot for each group
     for i, group in enumerate(groups):
         # Get IDs for this group
-        group_ids = id_group_df[id_group_df[categories] == group][id_col_name].values
+        group_ids = group_dataframe[group_dataframe[group_column_name] == group][id_col_name].values
 
         # Match IDs with sequence data
         mask = np.isin(seqdata.ids, group_ids)

sequenzo/visualization/plot_state_distribution.py

@@ -41,10 +41,12 @@ def smart_sort_groups(groups):
 
 
 def plot_state_distribution(seqdata: SequenceData,
-                            show_by_category=None,
-                            category_labels=None,
-                            id_group_df=None,
-                            categories=None,
+                            # Grouping parameters
+                            group_by_column=None,
+                            group_dataframe=None,
+                            group_column_name=None,
+                            group_labels=None,
+                            # Other parameters
                             weights="auto",
                             figsize=(12, 7),
                             plot_style="standard",
@@ -67,13 +69,33 @@ def plot_state_distribution(seqdata: SequenceData,
     Creates state distribution plots for different groups, showing how state
     prevalence changes over time within each group.
 
+    **Two API modes for grouping:**
+    
+    1. **Simplified API** (when grouping info is already in the data):
+       ```python
+       plot_state_distribution(seqdata, group_by_column="Cluster", group_labels=cluster_labels)
+       ```
+    
+    2. **Complete API** (when grouping info is in a separate dataframe):
+       ```python
+       plot_state_distribution(seqdata, group_dataframe=membership_df, 
+                              group_column_name="Cluster", group_labels=cluster_labels)
+       ```
+
     :param seqdata: (SequenceData) A SequenceData object containing sequences
-    :param show_by_category: (str, optional) Simple way to create grouped plots. 
-                            Specify the column name from the original data (e.g., "sex", "education").
-                            This will automatically create separate plots for each category.
-    :param category_labels: (dict, optional) Custom labels for category values. 
-                           Example: {0: "Female", 1: "Male"} or {"low": "Low Education", "high": "High Education"}.
-                           If not provided, will use original values or auto-generate readable labels.
+    
+    **Grouping parameters:**
+    :param group_by_column: (str, optional) Column name from seqdata.data to group by.
+                           Use this when grouping information is already in your data.
+                           Example: "Cluster", "sex", "education"
+    :param group_dataframe: (pd.DataFrame, optional) Separate dataframe containing grouping information.
+                           Use this when grouping info is in a separate table (e.g., clustering results).
+                           Must contain ID column and grouping column.
+    :param group_column_name: (str, optional) Name of the grouping column in group_dataframe.
+                             Required when using group_dataframe.
+    :param group_labels: (dict, optional) Custom labels for group values.
+                        Example: {1: "Late Family Formation", 2: "Early Partnership"}
+                        Maps original values to display labels.
     :param weights: (np.ndarray or "auto") Weights for sequences. If "auto", uses seqdata.weights if available
     :param figsize: (tuple) Size of the figure (only used when plot_style="custom")
     :param plot_style: Plot aspect style:
@@ -122,46 +144,45 @@ def plot_state_distribution(seqdata: SequenceData,
     
     actual_figsize = style_sizes[plot_style]
     
-    # Handle the new simplified API: show_by_category
-    if show_by_category is not None:
-         
+    # Handle the simplified API: group_by_column
+    if group_by_column is not None:
         # Validate that the column exists in the original data
-        if show_by_category not in seqdata.data.columns:
+        if group_by_column not in seqdata.data.columns:
             available_cols = [col for col in seqdata.data.columns if col not in seqdata.time and col != seqdata.id_col]
             raise ValueError(
-                f"Column '{show_by_category}' not found in the data. "
+                f"Column '{group_by_column}' not found in the data. "
                 f"Available columns for grouping: {available_cols}"
             )
         
-        # Automatically create id_group_df and categories from the simplified API
-        id_group_df = seqdata.data[[seqdata.id_col, show_by_category]].copy()
-        id_group_df.columns = ['Entity ID', 'Category']
-        categories = 'Category'
+        # Automatically create group_dataframe and group_column_name from the simplified API
+        group_dataframe = seqdata.data[[seqdata.id_col, group_by_column]].copy()
+        group_dataframe.columns = ['Entity ID', 'Category']
+        group_column_name = 'Category'
         
-        # Handle category labels - flexible and user-controllable
-        unique_values = seqdata.data[show_by_category].unique()
+        # Handle group labels - flexible and user-controllable
+        unique_values = seqdata.data[group_by_column].unique()
         
-        if category_labels is not None:
+        if group_labels is not None:
             # User provided custom labels - use them
-            missing_keys = set(unique_values) - set(category_labels.keys())
+            missing_keys = set(unique_values) - set(group_labels.keys())
             if missing_keys:
                 raise ValueError(
-                    f"category_labels missing mappings for values: {missing_keys}. "
-                    f"Please provide labels for all unique values in '{show_by_category}': {sorted(unique_values)}"
+                    f"group_labels missing mappings for values: {missing_keys}. "
+                    f"Please provide labels for all unique values in '{group_by_column}': {sorted(unique_values)}"
                 )
-            id_group_df['Category'] = id_group_df['Category'].map(category_labels)
+            group_dataframe['Category'] = group_dataframe['Category'].map(group_labels)
         else:
             # No custom labels provided - use smart defaults
             if all(isinstance(v, (int, float, np.integer, np.floating)) and not pd.isna(v) for v in unique_values):
-                # Numeric values - keep as is (user can provide category_labels if they want custom names)
+                # Numeric values - keep as is (user can provide group_labels if they want custom names)
                 pass
             # For string/categorical values, keep original values
             # This handles cases where users already have meaningful labels like "Male"/"Female"
         
-        print(f"[>] Creating grouped plots by '{show_by_category}' with {len(unique_values)} categories")
+        print(f"[>] Creating grouped plots by '{group_by_column}' with {len(unique_values)} categories")
     
     # If no grouping information, create a single plot
-    if id_group_df is None or categories is None:
+    if group_dataframe is None or group_column_name is None:
         return _plot_state_distribution_single(
             seqdata=seqdata, weights=weights, figsize=actual_figsize,
             plot_style=plot_style, title=title, xlabel=xlabel, ylabel=ylabel,
@@ -179,21 +200,21 @@ def plot_state_distribution(seqdata: SequenceData,
             raise ValueError("Length of weights must equal number of sequences.")
     
     # Ensure ID columns match (convert if needed)
-    id_col_name = "Entity ID" if "Entity ID" in id_group_df.columns else id_group_df.columns[0]
+    id_col_name = "Entity ID" if "Entity ID" in group_dataframe.columns else group_dataframe.columns[0]
 
     # Get unique groups and sort them based on user preference
     if group_order:
         # Use manually specified order, filter out non-existing groups
-        groups = [g for g in group_order if g in id_group_df[categories].unique()]
-        missing_groups = [g for g in id_group_df[categories].unique() if g not in group_order]
+        groups = [g for g in group_order if g in group_dataframe[group_column_name].unique()]
+        missing_groups = [g for g in group_dataframe[group_column_name].unique() if g not in group_order]
         if missing_groups:
             print(f"[Warning] Groups not in group_order will be excluded: {missing_groups}")
     elif sort_groups == 'numeric' or sort_groups == 'auto':
-        groups = smart_sort_groups(id_group_df[categories].unique())
+        groups = smart_sort_groups(group_dataframe[group_column_name].unique())
     elif sort_groups == 'alpha':
-        groups = sorted(id_group_df[categories].unique())
+        groups = sorted(group_dataframe[group_column_name].unique())
     elif sort_groups == 'none':
-        groups = list(id_group_df[categories].unique())
+        groups = list(group_dataframe[group_column_name].unique())
     else:
         raise ValueError(f"Invalid sort_groups value: {sort_groups}. Use 'auto', 'numeric', 'alpha', or 'none'.")
     
@@ -216,7 +237,7 @@ def plot_state_distribution(seqdata: SequenceData,
     # Process each group
     for i, group in enumerate(groups):
         # Get IDs for this group
-        group_ids = id_group_df[id_group_df[categories] == group][id_col_name].values
+        group_ids = group_dataframe[group_dataframe[group_column_name] == group][id_col_name].values
 
         # Match IDs with sequence data
         mask = np.isin(seqdata.ids, group_ids)

sequenzo/visualization/plot_transition_matrix.py

@@ -1,5 +1,5 @@
 """
-@Author  : Yuqi Liang 梁彧祺
+@Author  : Yuqi Liang 梁彧祺, Sebastian Daza
 @File    : plot_transition_matrix.py
 @Time    : 13/02/2025 12:39
 @Desc    :
@@ -32,12 +32,12 @@ def compute_transition_matrix(seqdata: SequenceData, with_missing: bool = False,
     # 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
@@ -52,19 +52,19 @@ def compute_transition_matrix(seqdata: SequenceData, with_missing: bool = False,
     nxt = A[:, 1:].flatten()
     w_pair = np.repeat(w, T-1)  # Each sequence weight replicated (T-1) times
 
-    # Filter valid transitions
-    valid = (current >= 0) & (current < num_states) & (nxt >= 0) & (nxt < num_states)
+    # 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
 
@@ -75,7 +75,7 @@ def compute_transition_matrix(seqdata: SequenceData, with_missing: bool = False,
 
     # 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):
@@ -92,7 +92,7 @@ def print_transition_matrix(seqdata: SequenceData, transition_rates: np.ndarray)
     :param seqdata: SequenceData object containing state information
     :param transition_rates: numpy array containing transition rates
     """
-    state_labels = seqdata.states
+    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
@@ -127,45 +127,45 @@ def print_transition_matrix(seqdata: SequenceData, transition_rates: np.ndarray)
 
 def plot_transition_matrix(seqdata: SequenceData,
                            weights="auto",
-                           title: Optional[str] = None,
+                           title: str = "State Transition Rate Matrix",
                            fontsize: int = 12,
                            save_as: Optional[str] = None,
-                           dpi: int = 200) -> 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)
 
-    # Create upper triangle mask (show diagonal)
-    mask = np.triu(np.ones(transition_matrix.shape, dtype=bool), k=1)
-
     # Set figure size
     plt.figure(figsize=(12, 10))
 
     # Use fresh color scheme
     cmap = sns.color_palette("light:#5A9", as_cmap=True)
 
-    # Generate heatmap
+    # Generate heatmap using pre-formatted annotation strings
     ax = sns.heatmap(
         transition_matrix,
-        mask=mask,
-        annot=True,
-        fmt=".2f",
+        fmt=format,
         cmap=cmap,
         xticklabels=seqdata.labels,
         yticklabels=seqdata.labels,
         linewidths=0.5,
         linecolor="gray",
         cbar_kws={"shrink": 0.8},
-        square=True
+        square=True,
+        annot_kws={"fontsize": fontsize - 2}
     )
 
     # Show all the borderlines
@@ -175,7 +175,7 @@ def plot_transition_matrix(seqdata: SequenceData,
     # Adjust format
     if title:
         show_plot_title(plt.gca(), title, show=True, fontsize=fontsize+2, fontweight='bold', pad=20)
-    # plt.title("State Transition Rate Matrix", fontsize=14, fontweight='bold', pad=20)
+
     plt.xlabel("State at t + 1", fontsize=fontsize, labelpad=10)
     plt.ylabel("State at t", fontsize=fontsize, labelpad=10)
 
@@ -186,5 +186,4 @@ def plot_transition_matrix(seqdata: SequenceData,
     # Adjust layout
     plt.tight_layout()
 
-    save_and_show_results(save_as, dpi=200)
-
+    save_and_show_results(save_as, dpi=dpi)

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'
+]