sequenzo 0.1.17__cp312-cp312-macosx_10_13_universal2.whl → 0.1.19__cp312-cp312-macosx_10_13_universal2.whl

sequenzo/sequence_characteristics/overall_cross_sectional_entropy.py

@@ -1,5 +1,5 @@
 """
-@Author  : 李欣怡
+@Author  : Yuqi Liang, Xinyi Li
 @File    : overall_cross_sectional_entropy.py
 @Time    : 2025/9/15 21:52
 @Desc    : States frequency by time unit
@@ -13,11 +13,29 @@ import pandas as pd
 from scipy.stats import entropy
 from sequenzo.define_sequence_data import SequenceData
 
-def get_cross_sectional_entropy(seqdata: SequenceData,
-                                weighted: bool = True,
-                                norm: bool = True):
+def get_cross_sectional_entropy(
+    seqdata: SequenceData,
+    weighted: bool = True,
+    norm: bool = True,
+    return_format: str = "tidy",        # "tidy" | "wide" | "dict"
+    include_effective_states: bool = True,
+    add_topk: int = 1,                  # Mark top K dominant states at each time point
+    round_decimals: int = 6
+):
     """
-    Compute the states frequency by time unit.
+    Cross-sectional state distribution by time with entropy and readable outputs.
+
+    What you get in a tidy format:
+        time   state   freq   entropy   per_time_entropy_norm   N_valid   rank   is_top
+        1      A       0.645  0.380     0.380          2346.27   1      True
+        ...
+
+    Additional metrics:
+        - per_time_entropy_norm: If norm=True, normalized by maximum entropy (|S|), range 0–1
+        - effective_states (H_effective): exp(H), equivalent "effective number of states"
+        - summary: Key interpretation points (entropy peaks/valleys, dominant state intervals, average entropy, etc.)
+
+    Parameters maintain your semantics, with new return_format, add_topk etc. for better interpretability.
 
     Parameters
     ----------
@@ -25,130 +43,178 @@ def get_cross_sectional_entropy(seqdata: SequenceData,
         A sequence object created by the SequenceData function.
     weighted : bool, default True
         If True, the frequencies are weighted by the number of non-missing values at each time unit.
-    with_missing : bool, default False
-        If True, missing values are included in the frequency computation.
     norm : bool, default True
-        If True, the frequencies are normalized to sum to 1 at each time unit.
+        If True, the entropy is normalized by maximum possible entropy.
+    return_format : str, default "tidy"
+        Return format: "tidy" for long-form table, "wide" for matrices, "dict" for original dict format.
+    include_effective_states : bool, default True
+        If True, calculate effective number of states (exp(entropy)).
+    add_topk : int, default 1
+        Mark top K dominant states at each time point.
+    round_decimals : int, default 6
+        Number of decimal places for rounding.
 
     Returns
     -------
-    pd.DataFrame
-        A dict, containing the frequency of each state at each time unit, validStates and Entropy.
+    pd.DataFrame or dict
+        Depending on return_format:
+        - "tidy": Long-form DataFrame with interpretable columns
+        - "wide": Dict with frequency matrix, entropy series, etc.
+        - "dict": Original dict format (backward compatible)
     """
 
     if not isinstance(seqdata, SequenceData):
-        raise ValueError(" [!] data is NOT a sequence object, see SequenceData function to create one.")
-
-    # Retrieve the states, shape and colors
-    states = seqdata.states.copy()
-    statl = range(1, len(states) + 1)
-    number_states = len(statl)
+        raise ValueError("[!] data is NOT a sequence object, see SequenceData.")
 
-    number_seq = seqdata.seqdata.shape[1]
+    # Basic metadata
+    states_labels = list(seqdata.states)                # Human-readable state labels
+    S = len(states_labels)
+    T = seqdata.seqdata.shape[1]                        # Number of time points
+    times = list(seqdata.seqdata.columns)
 
+    # Color attributes
     cpal = seqdata.custom_colors
 
-    # SequenceData already added missing values into states and colors internally
-    # if seqdata.ismissing:
-    #     statl.append(seqdata.missing_value)
-    #     col.append(seqdata.missing_color)
-
-    sd = pd.DataFrame(np.zeros((number_states, number_seq)), index=states, columns=seqdata.seqdata.columns)
-
     # Weights
-    weights = seqdata.weights if seqdata.weights is not None else np.ones(seqdata.seqdata.shape[0])
-
     # Also takes into account that in unweighted sequence objects created with
     # older TraMineR versions the weights attribute is a vector of 1
     # instead of NULL
-    if np.all(weights == 1):
+    w = seqdata.weights if seqdata.weights is not None else np.ones(seqdata.seqdata.shape[0])
+    if np.all(w == 1):
         weighted = False
 
-    for i in range(number_states):
-        for j in range(number_seq):
-            sd.iloc[i, j] = np.sum(weights[(seqdata.seqdata.iloc[:, j] == statl[i]).values])
+    # Your data is usually encoded with 1..S; if internally already labels, we can map here
+    # For compatibility: build a "value -> row index" lookup table
+    # Try to support both numeric encoding (1..S) and labels themselves
+    value_to_row = {v: i for i, v in enumerate(range(1, S+1))}
+    label_to_row = {lab: i for i, lab in enumerate(states_labels)}
+
+    # Frequency matrix (S x T)
+    freq_counts = np.zeros((S, T), dtype=float)
+
+    for j in range(T):
+        col = seqdata.seqdata.iloc[:, j]
+        for i in range(S):
+            # Try both encoding and label matching
+            mask_num = (col == (i+1))
+            mask_lab = (col == states_labels[i])
+            mask = mask_num | mask_lab
+            if weighted:
+                freq_counts[i, j] = w[mask].sum()
+            else:
+                freq_counts[i, j] = mask.sum()
+
+    N_valid = freq_counts.sum(axis=0)                 # Valid weight/sample size per time point
+    with np.errstate(divide='ignore', invalid='ignore'):
+        P = np.divide(freq_counts, N_valid, where=(N_valid>0))  # Frequencies
+
+    # Entropy
+    H = np.array([entropy(P[:, j][P[:, j] > 0]) if N_valid[j] > 0 else 0.0 for j in range(T)])
 
-    N = sd.sum(axis=0)
-    sd = sd.div(N, axis=1)
-
-    E = sd.apply(lambda col: entropy(col[col > 0]), axis=0)
-
-    # Maximum entropy is the entropy of the alphabet
     if norm:
-        E_max = entropy(np.ones(number_states) / number_states)
-        E = E / E_max
-
-    res = {
-        "Frequencies": sd,
-        "ValidStates": N,
-        "Entropy": E
-    }
-
-    res_attrs = {
-        "nbseq": np.sum(weights),
-        "cpal": cpal,
-        "xtlab": list(seqdata.seqdata.columns),
-        "xtstep": getattr(seqdata, "xtstep", None),
-        "tick_last": getattr(seqdata, "tick_last", None),
-        "weighted": weighted,
-        "norm": norm
+        Hmax = entropy(np.ones(S) / S) if S > 0 else 1.0
+        H_norm = H / Hmax if Hmax > 0 else H
+    else:
+        H_norm = H
+
+    # Effective number of states (highly interpretable: equivalent "how many equiprobable states")
+    H_eff = np.exp(H) if include_effective_states else None
+
+    # Organize output: wide format
+    freq_df_wide = pd.DataFrame(P, index=states_labels, columns=times).round(round_decimals)
+    entropy_s = pd.Series(H_norm if norm else H, index=times, name=("per_time_entropy_norm" if norm else "Entropy")).round(round_decimals)
+    valid_s   = pd.Series(N_valid, index=times, name="N_valid").round(round_decimals)
+    eff_s     = (pd.Series(H_eff, index=times, name="Effective States").round(round_decimals)
+                 if include_effective_states else None)
+
+    # Generate tidy table (interpretation-friendly)
+    tidy = (
+        freq_df_wide
+        .reset_index()
+        .melt(id_vars="index", var_name="time", value_name="freq")
+        .rename(columns={"index": "state"})
+        .sort_values(["time", "freq"], ascending=[True, False])
+    )
+    # Ranking + topK annotation
+    tidy["rank"] = tidy.groupby("time")["freq"].rank(method="first", ascending=False).astype(int)
+    if add_topk and add_topk > 0:
+        tidy["is_top"] = tidy["rank"] <= add_topk
+    else:
+        tidy["is_top"] = False
+
+    # Merge entropy/sample size/effective states
+    tidy = tidy.merge(entropy_s.reset_index().rename(columns={"index": "time"}), on="time", how="left")
+    tidy = tidy.merge(valid_s.reset_index().rename(columns={"index": "time"}), on="time", how="left")
+    if eff_s is not None:
+        tidy = tidy.merge(eff_s.reset_index().rename(columns={"index": "time"}), on="time", how="left")
+
+    # Friendly column order
+    cols = ["time", "state", "freq"]
+    if norm:
+        cols += ["per_time_entropy_norm"]
+    else:
+        cols += ["Entropy"]
+    cols += ["N_valid"]
+    if include_effective_states:
+        cols += ["Effective States"]
+    cols += ["rank", "is_top"]
+    tidy = tidy[cols]
+
+    # Summary: key statistics that can be explained in one sentence
+    summary = {
+        "states": states_labels,
+        "n_states": S,
+        "n_timepoints": T,
+        "avg_entropy_norm": float(tidy["per_time_entropy_norm"].mean()) if norm else None,
+        "avg_entropy": float((entropy_s if not norm else entropy_s * entropy(np.ones(S)/S)).mean()) if not norm else None,
+        "peak_entropy_time": tidy.loc[tidy["per_time_entropy_norm" if norm else "Entropy"].idxmax(), "time"] if T > 0 else None,
+        "lowest_entropy_time": tidy.loc[tidy["per_time_entropy_norm" if norm else "Entropy"].idxmin(), "time"] if T > 0 else None,
+        "dominant_stability_ratio": float(tidy.query("rank==1")["freq"].mean()),  # Average proportion of dominant state
+        "cpal": cpal
     }
 
-    res["__attrs__"] = res_attrs
-
-    return res
-
-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 = get_cross_sectional_entropy(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, time_type="year", states=states, id_col="COUNTRY")
-    # res = seqstatd(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_type="age", time=_time, id_col="country", states=states)
-    # res = seqstatd(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_type="age", time=_time, id_col="worker_id", states=states)
-    # res = seqstatd(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_type="age", time=_time, id_col="worker_id", states=states)
-    # res = seqstatd(sequence_data)
-
-    print(res['Frequencies'])
-    print(res['ValidStates'])
-    print(res['Entropy'])
+    # Print descriptive statistics
+    print("\n" + "="*70)
+    print("Cross-Sectional Entropy Summary")
+    print("="*70)
+    print(f"[>] Number of states: {summary['n_states']}")
+    print(f"[>] Number of time points: {summary['n_timepoints']}")
+    print(f"[>] On average, the most common state accounts for {summary['dominant_stability_ratio']:.1%} of cases")
+    print(f"[>] Entropy is highest at time point {summary['peak_entropy_time']}")
+    print(f"[>] Entropy is lowest at time point {summary['lowest_entropy_time']}")
+    if norm:
+        print(f"[>] Average normalized entropy: {summary['avg_entropy_norm']:.3f} (range: 0 = fully concentrated, 1 = evenly distributed)")
+    print("="*70 + "\n")
+
+    # Compatible with different return formats
+    if return_format == "tidy":
+        tidy.attrs = {"summary": summary}
+        return tidy
+    elif return_format == "wide":
+        out = {
+            "Frequencies": freq_df_wide,
+            "N_valid": valid_s,
+            ("per_time_entropy_norm" if norm else "Entropy"): entropy_s
+        }
+        if eff_s is not None:
+            out["Effective States"] = eff_s
+        return out
+    else:  # "dict" —— try to be more readable too
+        res = {
+            "Frequencies": freq_df_wide,
+            "ValidStates": valid_s,
+            "Entropy": entropy_s if not norm else None,
+            "per_time_entropy_norm": entropy_s if norm else None,
+            "Effective States": eff_s,
+            "__attrs__": {
+                "nbseq": float(valid_s.iloc[0]) if len(valid_s)>0 else None,
+                "cpal": cpal,
+                "xtlab": times,
+                "xtstep": getattr(seqdata, "xtstep", None),
+                "tick_last": getattr(seqdata, "tick_last", None),
+                "weighted": weighted,
+                "norm": norm,
+                "summary": summary
+            }
+        }
+        return res

sequenzo/sequence_characteristics/plot_characteristics.py

@@ -51,7 +51,8 @@ def plot_longitudinal_characteristics(seqdata,
                                       save_as=None,
                                       dpi=200,
                                       custom_colors=None,
-                                      show_sequence_ids=False):
+                                      show_sequence_ids=False,
+                                      id_as_column=True):
     """
     Create a horizontal bar chart showing four key characteristics for selected sequences.
     
@@ -128,13 +129,17 @@ def plot_longitudinal_characteristics(seqdata,
     show_sequence_ids : bool, optional (default=False)
         If True, y-axis shows actual sequence IDs (when available).
         If False, shows 1..N index positions.
+        
+    id_as_column : bool, optional (default=True)
+        If True, the returned DataFrame will include ID as a separate column on the same level as other columns.
+        If False, IDs will be used as the DataFrame index.
     
     Returns
     -------
     pandas.DataFrame
         A DataFrame containing the calculated metrics for all plotted sequences.
-        Columns: ['Transitions', 'Entropy', 'Turbulence', 'Complexity']
-        Index: The sequence IDs that were plotted
+        If id_as_column=True: Columns: ['ID', 'Transitions', 'Entropy', 'Turbulence', 'Complexity'] (all columns at same level)
+        If id_as_column=False: Columns: ['Transitions', 'Entropy', 'Turbulence', 'Complexity'], Index: The sequence IDs
         
     Warnings
     --------
@@ -189,15 +194,15 @@ def plot_longitudinal_characteristics(seqdata,
     types of sequences and datasets.
     """
     # Calculate four metrics (all should be 0-1 normalized)
-    df_t = get_number_of_transitions(seqdata=seqdata, norm=True).iloc[:, 0]   # Series
+    df_t = get_number_of_transitions(seqdata=seqdata, norm=True).iloc[:, 1]   # Series
     df_e = get_within_sequence_entropy(seqdata=seqdata, norm=True)           # Series or single-column DataFrame
-    if isinstance(df_e, pd.DataFrame): df_e = df_e.iloc[:, 0]
+    if isinstance(df_e, pd.DataFrame): df_e = df_e.iloc[:, 1]
 
-    df_tb = get_turbulence(seqdata=seqdata, norm=True, type=2)               # Normalized turbulence
-    if isinstance(df_tb, pd.DataFrame): df_tb = df_tb.iloc[:, 0]
+    df_tb = get_turbulence(seqdata=seqdata, norm=True, type=2, id_as_column=True)               # Normalized turbulence
+    if isinstance(df_tb, pd.DataFrame): df_tb = df_tb.iloc[:, 1]
 
     df_c = get_complexity_index(seqdata=seqdata)                             # Already 0-1 normalized
-    if isinstance(df_c, pd.DataFrame): df_c = df_c.iloc[:, 0]
+    if isinstance(df_c, pd.DataFrame): df_c = df_c.iloc[:, 1]
 
     # Create metrics DataFrame with actual sequence IDs as index
     metrics = pd.DataFrame({
@@ -336,7 +341,17 @@ def plot_longitudinal_characteristics(seqdata,
     plt.show()
     plt.close()
 
-    return metrics  # Return the data used for plotting for inspection
+    # Handle ID display options for returned DataFrame
+    if id_as_column:
+        # Add ID as a separate column and reset index to numeric
+        metrics_result = metrics.copy()
+        metrics_result['ID'] = metrics_result.index
+        metrics_result = metrics_result[['ID', 'Transitions', 'Entropy', 'Turbulence', 'Complexity']].reset_index(drop=True)
+        return metrics_result
+    else:
+        # Return with ID as index (traditional format)
+        metrics.index.name = 'ID'
+        return metrics
 
 
 def plot_cross_sectional_characteristics(seqdata,
@@ -465,10 +480,14 @@ def plot_cross_sectional_characteristics(seqdata,
     >>> valid_n = result['ValidStates']         # Sample sizes by time
     """
     # Get cross-sectional data using the existing function
-    res = get_cross_sectional_entropy(seqdata, weighted=True, norm=True)
+    res = get_cross_sectional_entropy(seqdata, weighted=True, norm=True, return_format="dict")
     
     freq = res["Frequencies"]          # rows: states, cols: time points
-    ent = res["Entropy"]               # index: time points
+    # Get normalized or raw entropy (check which key exists)
+    if "per_time_entropy_norm" in res and res["per_time_entropy_norm"] is not None:
+        ent = res["per_time_entropy_norm"]
+    else:
+        ent = res["Entropy"]
     N = res.get("ValidStates", None)   # valid sample sizes per time point
 
     # Sort time axis if possible (handles both numeric and string time labels)

sequenzo/sequence_characteristics/simple_characteristics.py

@@ -306,5 +306,6 @@ def get_number_of_transitions(seqdata, norm=False, pwight=False) -> pd.DataFrame
             trans[seq_length<=1] = 0
 
     trans = pd.DataFrame(trans, index=seqdata.seqdata.index, columns=['Transitions'])
+    trans = trans.reset_index().rename(columns={'index': 'ID'})
 
     return trans

sequenzo/sequence_characteristics/state_frequencies_and_entropy_per_sequence.py

@@ -14,15 +14,19 @@ from sequenzo.define_sequence_data import SequenceData
 
 def get_state_freq_and_entropy_per_seq(seqdata, prop=False):
     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.")
+
+    if seqdata.labels is not None:
+        states = seqdata.labels
+    else:
+        states = seqdata.states
 
-    states = seqdata.states.copy()
     number_states = len(states)
     number_seq = seqdata.seqdata.shape[0]
 
     iseqtab = pd.DataFrame(np.zeros((number_seq, number_states)), index=seqdata.seqdata.index, columns=states)
 
-    print(f"  - computing state distribution for {number_seq} sequences and {number_states} states ...")
+    print(f"[>] Computing state distribution for {number_seq} sequences and {number_states} states ...")
 
     for i, state in enumerate(states):
         iseqtab.iloc[:, i] = (seqdata.seqdata == (i+1)).sum(axis=1)
@@ -30,4 +34,6 @@ def get_state_freq_and_entropy_per_seq(seqdata, prop=False):
     if prop:
         iseqtab = iseqtab.div(iseqtab.sum(axis=1), axis=0)
 
+    iseqtab = iseqtab.reset_index().rename(columns={'index': 'ID'})
+
     return iseqtab

sequenzo/sequence_characteristics/turbulence.py

@@ -1,6 +1,6 @@
 """
-@Author  : 李欣怡
-@File    : seqST.py
+@Author  : Xinyi Li, Yuqi Liang
+@File    : turbulence.py
 @Time    : 2025/9/24 14:09
 @Desc    : Computes the sequence turbulence measure
 
@@ -27,7 +27,7 @@ def turb(x):
     Tux = np.log2(phi * ((s2max + 1) / (s2_tx + 1)))
     return Tux
 
-def get_turbulence(seqdata, norm=False, silent=True, type=1):
+def get_turbulence(seqdata, norm=False, silent=True, type=1, id_as_column=True):
     """
     Computes the sequence turbulence measure
 
@@ -41,15 +41,18 @@ def get_turbulence(seqdata, norm=False, silent=True, type=1):
         If True, suppresses the output messages.
     type : int, default 1
         Type of spell duration variance to be used. Can be either 1 or 2.
+    id_as_column : bool, default True
+        If True, the ID will be included as a separate column instead of as the index.
 
     Returns
     -------
     pd.DataFrame
         A DataFrame with one column containing the turbulence measure for each sequence.
+        If id_as_column=True, also includes an ID column.
     """
 
     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 not silent:
         print(f"  - extracting symbols and durations ...")
@@ -70,7 +73,22 @@ def get_turbulence(seqdata, norm=False, silent=True, type=1):
     s2_tx_max = s2_tx['vmax']
     s2_tx = s2_tx['result']
 
-    tmp = pd.DataFrame({'phi': phi.flatten(), 's2_tx': s2_tx, 's2max': s2_tx_max})
+    # Extract phi values and ensure 1D array
+    if hasattr(phi, 'iloc'):
+        phi_values = phi.iloc[:, 0].values
+    elif hasattr(phi, 'values'):
+        phi_values = phi.values
+    else:
+        phi_values = phi
+    
+    # Ensure phi_values is 1D
+    phi_values = np.asarray(phi_values).flatten()
+    
+    # Extract 1D arrays from s2_tx and s2_tx_max DataFrames
+    s2_tx_values = s2_tx.iloc[:, 1].values if hasattr(s2_tx, 'iloc') else np.asarray(s2_tx).flatten()
+    s2_tx_max_values = s2_tx_max.iloc[:, 1].values if hasattr(s2_tx_max, 'iloc') else np.asarray(s2_tx_max).flatten()
+    
+    tmp = pd.DataFrame({'phi': phi_values, 's2_tx': s2_tx_values, 's2max': s2_tx_max_values})
     Tx = tmp.apply(lambda row: turb([row['phi'], row['s2_tx'], row['s2max']]), axis=1).to_numpy()
 
     if norm:
@@ -95,7 +113,7 @@ def get_turbulence(seqdata, norm=False, silent=True, type=1):
         else:
             turb_phi = 2
 
-        if turb_phi.isna().any().any():
+        if hasattr(turb_phi, 'isna') and turb_phi.isna().any().any():
             turb_phi = 1e15  # 使用有限大数值避免转换警告
             print("[!] phi set as max float due to exceeding value when computing max turbulence.")
 
@@ -103,7 +121,19 @@ def get_turbulence(seqdata, norm=False, silent=True, type=1):
         turb_s2_max = turb_s2['vmax']
         turb_s2 = turb_s2['result']
 
-        tmp = pd.DataFrame({'phi': turb_phi.iloc[:, 0], 's2_tx': turb_s2, 's2max': turb_s2_max})
+        # Extract turb_phi values and ensure 1D
+        if hasattr(turb_phi, 'iloc'):
+            phi_value = turb_phi.iloc[:, 0].values
+        else:
+            phi_value = [turb_phi]
+        
+        phi_value = np.asarray(phi_value).flatten()
+        
+        # Extract 1D arrays from turb_s2 and turb_s2_max DataFrames
+        turb_s2_values = turb_s2.iloc[:, 1].values if hasattr(turb_s2, 'iloc') else np.asarray(turb_s2).flatten()
+        turb_s2_max_values = turb_s2_max.iloc[:, 1].values if hasattr(turb_s2_max, 'iloc') else np.asarray(turb_s2_max).flatten()
+        
+        tmp = pd.DataFrame({'phi': phi_value, 's2_tx': turb_s2_values, 's2max': turb_s2_max_values})
         maxT = tmp.apply(lambda row: turb([row['phi'], row['s2_tx'], row['s2max']]), axis=1).to_numpy()
 
         Tx_zero = np.where(Tx < 1)[0]
@@ -112,64 +142,14 @@ def get_turbulence(seqdata, norm=False, silent=True, type=1):
             Tx[Tx_zero, :] = 0
 
     Tx_df = pd.DataFrame(Tx, index=seqdata.seqdata.index, columns=['Turbulence'])
-    return Tx_df
-
-if __name__ == "__main__":
-
-    from sequenzo import *
     
-    # ===============================
-    #             Sohee
-    # ===============================
-    # df = pd.read_csv('D:/college/research/QiQi/sequenzo/data_and_output/orignal data/sohee/sequence_data.csv')
-    # # df = pd.read_csv('/Users/lei/Documents/Sequenzo_all_folders/sequence_data_sources/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 = get_turbulence(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 = seqST(sequence_data)
-
-    # ===============================
-    #             CO2
-    # ===============================
-    # df = pd.read_csv("D:/country_co2_emissions_missing.csv")
-    df = load_dataset('country_co2_emissions_local_deciles')
-    df.to_csv("D:/country_co2_emissions_local_deciles.csv", index=False)
-    _time = list(df.columns)[1:]
-    # states = ['Very Low', 'Low', 'Middle', 'High', 'Very High']
-    states = ['D1 (Very Low)', 'D10 (Very High)', 'D2', 'D3', 'D4', 'D5', 'D6', 'D7', 'D8', 'D9']
-    sequence_data = SequenceData(df, time=_time, id_col="country", states=states)
-    res = get_turbulence(sequence_data, norm=True, type=2)
-
-    # ===============================
-    #            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 = seqST(sequence_data, norm=False, type=2)
-
-    # ===============================
-    #             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 = seqST(sequence_data, norm=True, type=2)
-
-    print(res)
+    # Handle ID display options
+    if id_as_column:
+        # Add ID as a separate column and reset index to numeric
+        Tx_df['ID'] = Tx_df.index
+        Tx_df = Tx_df[['ID', 'Turbulence']].reset_index(drop=True)
+    else:
+        # Always set index name to 'ID' for clarity
+        Tx_df.index.name = 'ID'
+    
+    return Tx_df