pysaebm 1.2.2

Stage-Aware Event-Based Modeling for Disease Progression

pysaebm

pysaebm is a Python package for generating and analyzing biomarker data using Stage-Aware Event-Based Modeling (SA-EBM). It supports various data generation experiments and EBM algorithms to estimate biomarker orderings and disease stages. This package is designed for researchers and data scientists working with biomarker progression analysis.

For detailed methodology, refer to our paper.

Installation

Install pysaebm using pip:

pip install pysaebm

Ensure you have Python 3.8+ and the required dependencies installed. For a full list of dependencies, see requirements.txt.

Data generation

Examples of how to generate data are at ./pysaebm/test/gen.py.

Because in each generation, the ordering is randomized, you will see a true_order_and_stages.json that tells you the corresponding true stages and true order for each output csv file.

The source codes for data generation can be seen in ./pysaebm/utils/generate_data.py.

This is the full generate parameters:

def generate(
   experiment_name: str = "sn_kjOrdinalDM_xnjNormal",
   params_file: str = 'params.json',
   js: List[int] = [50, 200, 500, 1000],
   rs: List[float] = [0.1, 0.25, 0.5, 0.75, 0.9],
   num_of_datasets_per_combination: int = 50,
   output_dir: str = 'data',
   seed: Optional[int] = None,
   dirichlet_alpha: Optional[Dict[str, List[float]]] = {
       'uniform': [100],
       'multinomial': [0.4013728324975898,
                       1.0910444770153345,
                       2.30974117596663,
                       3.8081194066281103,
                       4.889722107892335,
                       4.889722107892335,
                       3.8081194066281103,
                       2.30974117596663,
                       1.0910444770153345,
                       0.4013728324975898]
   },
   beta_params: Dict[str, Dict[str, float]] = {
       'near_normal': {'alpha': 2.0, 'beta': 2.0},
       'uniform': {'alpha': 1, 'beta': 1},
       'regular': {'alpha': 5, 'beta': 2}
   },
   prefix: Optional[str] = None,
   suffix: Optional[str] = None,
   keep_all_cols: bool = False,
   fixed_biomarker_order: bool = False,
   noise_std_parameter: float = 0.05,
) -> Dict[str, Dict[str, int]]:

Explanations:

• experiment_name should be one of the these:

experiment_names = [
  "sn_kjOrdinalDM_xnjNormal",     # Experiment 1: Ordinal kj with Dirichlet-Multinomial, Normal Xnj
  "sn_kjOrdinalDM_xnjNonNormal",  # Experiment 2: Ordinal kj with Dirichlet-Multinomial, Non-Normal Xnj
  "sn_kjOrdinalUniform_xnjNormal", # Experiment 3: Ordinal kj with Uniform distribution, Normal Xnj
  "sn_kjOrdinalUniform_xnjNonNormal", # Experiment 4: Ordinal kj with Uniform distribution, Non-Normal Xnj
  "sn_kjContinuousUniform",       # Experiment 5: Continuous kj with Uniform distribution
  "sn_kjContinuousBeta",          # Experiment 6: Continuous kj with Beta distribution
  "xiNearNormal_kjContinuousUniform", # Experiment 7: Near-normal Xi with Continuous Uniform kj
  "xiNearNormal_kjContinuousBeta", # Experiment 8: Near-normal Xi with Continuous Beta kj
  "xiNearNormalWithNoise_kjContinuousBeta", # Experiment 9: Same as Exp 8 but with noises to xi
]

You can find the explanation to these terms from our paper.

• params_file: The path to the parameters in json. Example is ./pysaebm/data/params.json. You should specify each biomarker's theta_mean, theta_std, phi_mean, and phi_std.
• js: An array of integers indicating the number of participants you want.
• rs: An array of floats indicating the number of healthy ratios.
• num_of_datasets_per_combination: The number of repetitions for each j-r combination.
• output_dir: The directory where you want to save the generated data.
• seed: An integer serving as the seed for the randomness.
• dirichlet_alpha: This should be a dictionary where keys are uniform and multinomial. They correspond to kjOrdinalUniform and kjOrdinalDM.
• beta_params: A dictionary where keys are near_normal, uniform, and regular, corresponding to xiNearNormal, kjContinuousUniform and kjContinuousBeta.
• prefix: Optional prefix for the output csv file names.
• suffix: Optional suffix for the output csv file names.
• keep_all_cols: Whether to include additional metadata columns (k_j, event_time, affected)
• fixed_biomarker_order: If True, will use the order as in the params_file. If False, will randomize the ordering.
• noise_std_parameter: the parameter in N(0, N \cdot noise_std_parameter) in experiment 9

Note that you need to make sure the dirichlet_alpha['multinomial'] has the same length as your params dict (as in your params_file).

Run EBM Algorithms

Examples of how to run algorithms and get results is at ./pysaebm/test/test.py.

This explains the parameters well enough:

def run_ebm(
   data_file: str,
   output_dir: str,
   output_folder: Optional[str] = None,
   algorithm: str = 'conjugate_priors',
   n_iter: int = 2000,
   n_shuffle: int = 2,
   burn_in: int = 500,
   thinning: int = 1,
   true_order_dict: Optional[Dict[str, int]] = None,
   true_stages: Optional[List[int]] = None,
   plot_title_detail: Optional[str] = "",
   fname_prefix: Optional[str] = "",
   skip_heatmap: Optional[bool] = False,
   skip_traceplot: Optional[bool] = False,
   # Strength of the prior belief in prior estimate of the mean (μ), set to 1 as default
   prior_n: float = 1.0,
   # Prior degrees of freedom, influencing the certainty of prior estimate of the variance (σ²), set to 1 as default
   prior_v: float = 1.0,
   weight_change_threshold: float = 0.01,
   bw_method: str = 'scott',
   seed: int = 42,
) -> Dict[str, Union[str, int, float, Dict, List]]:
   """
   Run the metropolis hastings algorithm and save results


   Args:
       data_file (str): Path to the input CSV file with biomarker data.
       output_dir (str): Path to the directory to store all the results.
       output_folder (str): Optional. If not provided, all results will be saved to output_dir/algorithm.
           If provided, results will be saved to output_dir/output_folder
       algorithm (str): Choose from 'hard_kmeans', 'mle', 'em', 'kde', and 'conjugate_priors' (default).
       n_iter (int): Number of iterations for the Metropolis-Hastings algorithm.
       n_shuffle (int): Number of shuffles per iteration.
       burn_in (int): Burn-in period for the MCMC chain.
       thinning (int): Thinning interval for the MCMC chain.
       true_order_dict (Optional[Dict[str, int]]): biomarker name: the correct order of it (if known)
       true_stages (Optional[List[int]]): true stages for all participants (if known)
       plot_title_detail (Optional[str]): optional string to add to plot title, as suffix.
       fname_prefix (Optional[str]): the prefix of heatmap, traceplot, results.json, and logs file, e.g., 5_50_0_heatmap_conjugate_priors.png
           In the example, there are no prefix strings.
       skip_heatmap (Optional[bool]): whether to save heatmaps. True if you want to skip saving heatmaps and save space.
       skip_traceplot (Optional[bool]): whether to save traceplots. True if you want to skip saving traceplots and save space.
       prior_n (strength of belief in prior of mean): default to be 1.0
       prior_v (prior degree of freedom) are the weakly informative priors, default to be 1.0
       weight_change_threshold (float): Threshold for kde weights (if np.mean(new_weights - old_weights)) > threshold, then recalculate
           otherwise use the new kde and weights
       bw_method (str): bandwidth selection method in kde
       seed (int): for reproducibility


   Returns:
       Dict[str, Union[str, int, float, Dict, List]]: Results including everything, e.g., Kendall's tau and p-value.
   """

Some extra explanations:

• n_iter: In general, above 2k is recommended. 10k should be sufficient if you have <= 10 biomarkers.
• burn_in and thinning: The idea behind the two parameters is that we will only use some of the results from all iterations in n_iter. We will do this: if (i > burn_in) & (i % thinning == 0), then we will use the result from that iteration i. Usually, we set thinning to be 1.

After running the run_ebm, you'll see a folder named as your output_dir. Each algorithm will have its subfolders.

The results are organized this way:

• records folder contains the loggings.
• heatmaps folder contains all the heatmaps. An example is

Biomarkers in the y-axis are ranked according to the ordering that has the highest likelihood among all iterations. Each cell indicates the probability that a certain biomarker falls in a certain stage. Note, however, these probabilities are calculated based on all the iterations that satisfy (i > burn_in) & (i % thinning == 0).

In the heatmap, the sum of each col and each row is 1.

• traceplots folder contains all the traceplots (starting from iteration 40, not iteration 0). Those plots will be useful to diagnose whether EBM algorithms are working correctly. It's totally okay for the plots to show fluctuation (because biomarker distribution and stage distributions are re-calculated each iterations). You should not, however, see a clear downward trend.

• results folder contains all important results in json. Each file contains

results = {
       "algorithm": algorithm,
       "runtime": end_time - start_time,
       "N_MCMC": n_iter,
       "n_shuffle": n_shuffle,
       "burn_in": burn_in,
       "thinning": thinning,
       'healthy_ratio': healthy_ratio,
       "max_log_likelihood": float(max(log_likelihoods)),
       "kendalls_tau2": tau2,
       "p_value2": p_value2,
       "kendalls_tau": tau,
       "p_value": p_value,
       "quadratic_weighted_kappa": qwk,
       "mean_absolute_error": mae,
       "mean_squared_error": mse,
       "root_mean_squared_error": rmse,
       "quadratic_weighted_kappa_diseased": qwk2,
       "mean_absolute_error_diseased": mae2,
       "mean_squared_error_diseased": mse2,
       "root_mean_squared_error_diseased": rmse2,
       'current_pi': current_pi.tolist(),
       # updated pi is the pi for all stages, including 0
       'updated_pi': updated_pi,
       'true_order': true_order_result,
       'ml_order': {k: int(v) for k, v in most_likely_order_dic.items()},
       "order_with_highest_ll": {k: int(v) for k, v in order_with_highest_ll.items()},
       "true_stages": true_stages,
       'ml_stages': ml_stages,
       # stages diseased only contains stage prediction for diseased patients
       "true_stages_diseased": true_stages_diseased,
       'ml_stages_diseased': ml_stages_diseased,
       "stage_likelihood_posterior": {str(k): v.tolist() for k, v in final_stage_post2.items()},
       "stage_likelihood_posterior_diseased": {str(k): v.tolist() for k, v in final_stage_post.items()},
       "final_theta_phi_params": final_theta_phi_params,
   }

• algorithm is the algorithm you are running with.
• runtime in seconds.
• N_MCMC is the number of MCMC iterations.
• n_shuffle is how many biomarkers to shuffle places in Metropolis Hastings. Default is 2.
• burn_in and thinning are explained above.
• healthy_ratio is the percentage of non-diseased participants in the dataset.
• max_log_likelihood is the max data log likelihood. The script ./pysaebm/algorithms/algorithm.py generates all_accepted_orders, log_likelihoods, current_theta_phi, current_stage_post, and current_pi. max_log_likelihoods is the max of log_likelihoods.
• kendalls_tau2 and p_values are the result of comparing the order_with_highest_ll with the true order (if provided).
• kendalls_tau and p_value are the result of comparing the ml_order (most likely order) with the true order (if provided).
• quadratic_weighted_kappa, mean_absolute_error, mean_squared_error and root_mean_squared_error are the results of comparing ml_stages and true_stages (if provided). These measures are for when we include estimation of healthy participants as well. That is to say, we estimate their disease stages without knowing they are healthy.
• quadratic_weighted_kappa_diseased, mean_absolute_error_diseased, mean_squared_error_diseased and root_mean_squared_error_diseased are the results of comparing ml_stages and true_stages (if provided). These measures are for when we include estimation of diseased participants only.
• current_pi is the probability distribution of all disease stages.
• current_pi is the probability distribution of all stages, including 0.
• true_order is the dictionary explaining the stage each biomarker is at.
• ml_order is the overall most likely order considering results from all the iterations that satisfy (i > burn_in) & (i % thinning == 0). It is calculated by the function of obtain_most_likely_order_dic in ./pysaebm/utils/data_processing.py
• order_with_highest_ll is the ordering that corresponds to the max_log_likelihoods.
• true_stages is an array of each participant's disease stage.
• ml_stages is the array of most likely disease stages.
• true_stages_diseased is the array of actual stages for diseased participants only.
• 'ml_stages_diseased is the array of the estimated stages for diseased participants only.
• stage_likelihood_posterior is a dictionary detailing each participant's probability in each of the possible stages.
• • stage_likelihood_posterior_diseased is a dictionary detailing each diseased participant's probability in each of the possible disease stages (no 0).
• final_theta_phi_params is a dictionary detailing the parameters for each biomarker. If you are using kde, you'll see data points and theta_weights and phi_weights. If you use other algorithms, you will see the theta_mean, theta_std, phi_mean and phi_std.

Use your own data

You are more than welcome to use your own data. After all, the very purpose of pysaebm is: to allow you to analyze your own data. However, you do have to make sure that the input data have at least four columns:

• participant: int
• biomarker: str
• measurement: float
• diseased: bool

The participant column should be integers from 0 to J-1 where J is the number of participants.

Samples are available at ./pysaebm/data/samples/.

The data should be in a tidy format, i.e.,

• Each variable is a column.
• Each observation is a row.
• Each type of observational unit is a table.

Change Log

• 2025-02-26 (V 0.3.4).

  • Modified the shuffle_order function to ensure full derangement, making convergence faster.

• 2025-03-06 (V 0.4.0)

  • use pyproject.toml instead
  • update conjuage_priors_algo.py, now without using the auxiliary variable of participant_stages. Keep the uncertainties just like in soft_kmeans_algo.py.

• 2025-03-07 (V 0.4.2)

  • Compute new_ln_likelihood_new_theta_phi based on new_theta_phi_estimates, which is based on stage_likelihoods_posteriors that is based on the newly proposed order and previous theta_phi_estimates.
  • Update theta_phi_estimates with new_theta_phi_estimates only if a new order is accepted.
  • The fallback theta_phi_estimates is the previous parameters rather than theta_phi_default
  • all_accepted_orders.append(current_order_dict.copy()) to make sure the results are not mutated.
  • Previously I calculated the new_ln_likelihood and stage_likelihoods_posteriors based on the newly proposed order and previous theta_phi_estimates, and directly updated theta_phi_estimates whether we accept the new order or not.
  • Previously, I excluded copy() in all_accepted_orders.append(current_order_dict.copy()), which is inaccurate.

• 2025-03-17 (V 0.4.3)

  • Added skip and title_detail parameters in the save_traceplot function.

• 2025-03-18 (V 0.4.4)

  • Add an optional horizontal bar indicating the upper limit in the trace plot.

• 2025-03-18 (V 0.4.7)

  • Allowed keeping all cols (keep_all_cols) in data generation.

• 2025-03-18 (V 0.4.9)

  • copy data_we_have and use data_we_have.loc[:, 'S_n'] in soft k means algo when preprocessing participant and biomarker data.

• 2025-03-20 (V 0.5.1)

  • In hard kmeans, updated delta = ln_likelihood - current_ln_likelihood, and in soft kmeans and conjugate priors, made sure I am using delta = new_ln_likelihood_new_theta_phi - current_ln_likelihood.
  • In each iteration, use theta_phi_estimates = theta_phi_default.copy() first. This means, stage_likelihoods_posteriors is based on the default theta_phi, not the previous iteration.

• 2025-03-21 (V 0.6.0)

  • Integrated all three algorithms to just one file algorithms/algorithm.py.
  • Changed the algorithm name of soft_kmeans to mle (maximum likelihood estimation)
  • Moved all helper functions from the algorithm script to utils/data_processing.py.

• 2025-03-22 (V 0.7.6)

  • Current state should include both the current accepted order and its associated theta/phi. When updating theta/phi at the start of each iteration, use the current state's theta/phi (1) in the calculation of stage likelihoods and (2) as the fallback if either of the biomarker's clusters is empty or has only one measurement; (3) as the prior mean and variance.
  • Set conjugate_priors as the default algorithm.
  • (Tried using cluster's mean and var as the prior but the results are not as good as using current state's theta/phi as the prior).

• 2025-03-24 (V 0.7.8)

  • In heatmap, reorder the biomarkers according to the most likely order.
  • In results.json reorder the biomarker according to their order rather than alphabetically ranked.
  • Modified obtain_most_likely_order_dic so that we assign stages for biomarkers that have the highest probabilities first.
  • In results.json, output the order associated with the highest total log likelihood. Also, calculate the kendall's tau and p values of it and the original order (if provided).

• 2025-03-25 (V 0.8.1)

  • In heatmap, reorder according to the order with highest log likelihood. Also, add the number just like (1).
  • Able to add title detail to heatmaps and traceplots.
  • Able to add fname_prefix in run_ebm().

• 2025-03-29 (V 0.8.9)

  • Added em algorithm.
  • Added Dirichlet-Multinomial Model to describe uncertainty of stage distribution (a multinomial distribution of all disease stages; because we cannot always assume all disease stages are equally likely).
  • prior_v default set to be 1.
  • Default to use dirichlet distribution instead of uniform distribution
  • Change data filename from 50|100_1 to 50_100_1.
  • Modified the mle algorithm to make sure the output does not contain np.nan (by using the fallback).

• 2025-03-30 (V 0.9.2)

  • Completed change to generate_data.py. Now incorporates the modified data generation model based on DEBM2019.
  • Rank the original order by the value (ascending), if the original order exists.
  • Able to skip saving traceplots and/or heatmaps.

• 2025-03-31 (V 0.9.4)

  • Able to store final theta phi estimates and the final stage likelihood posterior to results.json

• 2025-04-02 (V 0.9.5)

  • Added kde algorithm.
  • Initial kmeans used seeded Kmeans + conjugate priors.

• 2025-04-03 (V 0.9.7)

  • Improved kde.
  • Added dirichlet and beta parameters randomization.

• 2025-04-05 (V 0.9.9)

  • Updated generate_data.py to align with the experimental design.

• 2025-04-06 (V 0.9.9.3)

  • Make kmeans.py more robust. Now try 100 times to randomize the assignment for the diseased group if the initial kmeans failed.
  • Add algorithm and ml_order in results.json
  • After generating data, create true_order_and_stages_dict to store all filenames' true biomarker order and all participants' stages. For continuous kjs, use the bisect_right algorithm to get the ranking order.

• 2025-04-07 (V 0.9.9.9)

  • Modified generate_data.py to allow Experiment 9.
  • In generate_data.py, added the function of randomly flipping the direction of progression in the sigmoid model. Also make sure this random direction is consistent across participants.
  • In the spirit of "Do not repeat yourself", delete "R" and "rho" in params.json. Instead, compute it each time when I generate data. The time difference is minimal.
  • Added comparisons between true stages and most likely stages
  • Reorganized the results.json
  • Allow output_dir in run_ebm.

• 2025-04-08 (V 1.00)

  • Modify the FastKDE implementation.
  • Reorganized the results.json.
  • Added stage_post to hard kmeans as well. Now every algorithm can output ml_stages.
  • Make sure diseased stages are between (, max_stage].

• 2025-04-13 (V 1.1)

  • Leveraged weights when calculating bandwidth in FastKDE.

• 2025-04-14 (V 1.11)

  • Updated the fastkde.py.

• 2025-04-15 (V 1.131)

  • Renamed the package to pypysaebm.
  • Reconstructed the README documentation.

• 2025-04-16 (V 2.0.1)

  • Did the fast_kde.py on my own. Streamlined the script quite a bit. Now it's very readable and easy to follow & understand.
  • Used scott for bandwidth selection

• 2025-04-16 (V 2.0.2)

  • Checked whether the length of parameters is the same as the dirichlet_alpha multinomial.
  • Added an option to use fixed biomarker order in data generation.

• 2025-04-19 (V 2.0.3)

  • Anonymize the readme on Pypi

• 2025-04-20

  • Modify the format of the result of run.py to be Dict[str, Union[str, int, float, Dict, List]].

• 2025-04-21

  • Use np.random.choice(len(final_stage_post[pid]), p=final_stage_post[pid]) + 1 to obtain ml_stages instead of argmax.

• 2025-04-27 (V 2.0.2)

  • Make the noise std as a parameter in experiment 9.
  • Added current_pi to the output results json.

• 2025-04-30 (V 2.0.3)

  • Added the gmm algorithm.

• 2025-05-01 (V 2.0.5)

  • True uniform kjs now in data generation.

• 2025-05-02 (V 2.0.6)

  • Distinguished between 'gmm' and 'gmm_dm_prior'.

• 2025-05-03 (V 2.0.9)

  • Distinguished between 'conjugate_priors' and "conjugate_priors_gm_prior".

• 2025-05-07 (V 2.1.0)

  • Kept only the original five algorithms.
  • In run.py, make sure to check whether true_order exists or not before saving to json.

• 2025-05-11 (V 2.1.1)

  • Updated data generation, now using approximate uniform in kj generation.

• 2025-05-12 (V 2.1.3)

  • Added stage prediction for healthy participants as well, which means for each and every participant, when referencing stage, all stages (including 0) will be tested.
  • Updated heatmap visualization design; using blues color map now.
  • Save both png and pdf for heatmap.

• 2025-05-13 (V 2.1.4)

  • Added PDF for traceplot save.
  • Improved heatmap title.

• 2025-05-18 (V 2.1.6)

  • ml_stages, both for all and for diseased only, are based on participant_data that is using order_with_highest_ll.
  • Flipped the sign of mean of phi and theta for adas, p-tau, hip-fci, and fus-fci in data/params.json.

• 2025-05-28 (V 2.2.0)

  • Added algorithm name in the title of traceplots.
  • Added runtime in results.json.

• 2025-05-29 (V 2.2.2)

  • Added seed in run.py to enable reproducibility of the ordering results.
  • Added rng in random shuffling.

• 2025-06-21 (V 2.2.3)

  • Make sure that for kde, in calculate_bandwidth, set a lower bound for sigma.

• 2025-06-22 (V 2.2.5)

  • Added mixed pathology functions in generate_data.py, data_processing.py, algorithm.py and run.py.
  • Added a parameter of output_folder in run.py.

• 2025-06-24 (V 2.2.6)

  • Added max_log_likelihood in results json in run.py.

• 2025-06-29 (V 2.2.11)

  • Enabled Generalized Mallows as an alternative to Pairwise Preferences in energy and combined ordering generation.
  • Fixed the bug of params in generating data.

• 2025-06-30 (V 2.2.12)

  • Fixed typo: unbiased.

---

All above changelogs were alabebm. However, I changed the name to pysaebm after the paper got accepted. This is because of the required anonymity.

Below will be the changelogs for pysaebm.

• 2025-07-09 (V 1.0.1)

  • Made mp_method optional in generate_data.py.
  • Made mp_method and order_array in algorithm.py and run.py as optional.

• 2025-08-01 (V 1.0.2)

  • Deleted mixed_pathology stuff.

• 2025-08-03 (V 1.1.0)

  • changed from np.log(current_pi[k_j-1]) to np.log(current_pi[k_j]) in unbaised stage post.

• 2025-08-06 (V 1.2.1)

  • Used vectorized version for all algos except for kde. Used the pd.dataframe and dicts for kde.
  • Used tau distance instead of tau corr in saved results.
  • Made sure fastmath=False in each njit.