compute_trace.py confuses -prefix and Ne?

[s7guan]

Hi,
I am running on SINGER 0.1.8 beta and running compute_trace.py for MCMC convergence diagnostics. I noticed that when I run the following and would got an error message like this:

python /n/home10/sguan/SINGER-0.1.8-beta/SINGER//SINGER/compute_trace.py -prefix melp_Hmel218003o_optix2new -m 2.9e-9 -start 0 -end 99 -outpiut test
usage: compute_trace.py [-h] [-prefix PREFIX] [-m M] [-start_index START_INDEX] -end_index END_INDEX -output_filename
                        OUTPUT_FILENAME
compute_trace.py: error: argument -prefix: invalid float value: 'melp_Hmel218003o_optix2new'

I am inclined to think there is some code error when the .py script is written as it seems to suggest -prefix is taking Ne as input:

python /n/home10/sguan/SINGER-0.1.8-beta/SINGER/SINGER/compute_trace.py -h
usage: compute_trace.py [-h] [-prefix PREFIX] [-m M] [-start_index START_INDEX] -end_index END_INDEX -output_filename
                        OUTPUT_FILENAME

Compute traces for MCMC samples from SINGER.

options:
  -h, --help            show this help message and exit
  -prefix PREFIX        Effective population size.
  -m M                  Mutation rate.
  -start_index START_INDEX
                        The start index of the ARG sample
  -end_index END_INDEX  The end index of the ARG sample
  -output_filename OUTPUT_FILENAME
                        Output filename of the MCMC traces.

Do you think it is a copy paste error when you were writing the functions? Thanks!

[Jesson-mark]

Hi, thanks for maintaining this tool.

I may be missing some context, but I encountered the same issue and made a small local fix that lets the script run in my environment, so I am sharing the modified version here in case it is helpful.

import tskit
import numpy as np
import pandas as pd
import argparse
import os

def mse(x, y):
    x = np.array(x)
    y = np.array(y)
    return np.mean((x - y)**2)

def count_incompatibility(ts):
    unmapped_sites = 0
    for tree in ts.trees():
        for site in tree.sites():
            num_mutations = len(site.mutations)
            if num_mutations > 2:
                unmapped_sites += 1
            elif num_mutations == 2:
                if site.mutations[0].node != tree.root and site.mutations[1].node != tree.root:
                    unmapped_sites += 1
    return unmapped_sites

def diversity_fit_mse(ts, m):
    windows = np.arange(0, ts.sequence_length, 1e6).tolist()
    if windows[-1] != ts.sequence_length:
        windows.append(ts.sequence_length)
    
    site_diversity = ts.diversity(windows=windows, mode='site')
    branch_diversity = ts.diversity(windows=windows, mode='branch') * m
    
    valid_idx = ~np.isnan(site_diversity) & ~np.isnan(branch_diversity)
    if not np.any(valid_idx):
        return np.nan
        
    fit_mse = mse(site_diversity[valid_idx], branch_diversity[valid_idx])
    return fit_mse

def load_ts(prefix, index):
    file_trees = f"{prefix}_{index}.trees"
    file_tsz = f"{prefix}_{index}.tsz"
    
    if os.path.exists(file_trees):
        return tskit.load(file_trees)
    elif os.path.exists(file_tsz):
        import tszip
        return tszip.load(file_tsz)
    else:
        raise FileNotFoundError(f"Neither {file_trees} nor {file_tsz} found.")

def main():
    parser = argparse.ArgumentParser(description='Compute traces for MCMC samples from SINGER.')

    parser.add_argument('-prefix', type=str, required=True, help='Prefix of tree sequence files (e.g., "chr1").')
    parser.add_argument('-m', type=float, required=True, help='Mutation rate.')
    parser.add_argument('-start_index', type=int, default=0, help='The start index of the ARG sample.')
    parser.add_argument('-end_index', type=int, required=True, help='The end index of the ARG sample.')
    parser.add_argument('-output_filename', type=str, required=True, help='Output TSV/CSV filename of the MCMC traces.')   
 
    args = parser.parse_args()

    indices = list(range(args.start_index, args.end_index + 1))
    print(f"Processing {len(indices)} files with prefix '{args.prefix}'...")
    
    incompat_counts = []
    fit_mses = []
    
    for index in indices:
        try:
            ts = load_ts(args.prefix, index)
            
            incompat_count = count_incompatibility(ts)
            fit_mse = diversity_fit_mse(ts, args.m)
            
            incompat_counts.append(incompat_count)
            fit_mses.append(fit_mse)

            print(f'  Processed index {index}: Incompatible sites = {incompat_count}, Diversity fit MSE = {fit_mse:.4e}')
            
        except Exception as e:
            print(f"Error processing index {index}: {e}")
            incompat_counts.append(np.nan)
            fit_mses.append(np.nan)
    
    df = pd.DataFrame({
        'iteration': indices,
        'diversity_fit_mse': fit_mses,
        'incompatible_sites': incompat_counts
    })
    
    df.to_csv(args.output_filename, sep='\t', index=False)
    print(f"\nDone! Traces successfully saved to: {args.output_filename}")

if __name__ == "__main__":
    main()

A simple example usage is:

python -u $script \
    -prefix chr1 \
    -m 1.25e-8 \
    -start_index 0 \
    -end_index 9 \
    -output_filename output.tsv

In this example, the script searches for files named chr1_0.trees through chr1_9.trees (or .tsz files with the same naming pattern, if available).

I hope this is helpful in case others run into the same issue, and thanks again for sharing the software.