code: Making accumulation curves from kmers

[AnneliektH]

keywords: rarefaction, rarefy, resampling

Code used to create kmer accumulation curve

import sourmash
from pathlib import Path
from sourmash.nodegraph import Nodegraph
import pandas as pd
import random

# Parameters
ksize = 21
n_iterations = 100

# Load the sig file paths once
sig_files = sorted(Path("MAG_sketch").glob("*.sig"))

# Prepare a list to store results for all iterations
all_results = []

for iteration in range(n_iterations):
    print(f"Iteration {iteration + 1}/{n_iterations}")
    
    # Shuffle sample order for this iteration
    shuffled_files = sig_files.copy()
    random.shuffle(shuffled_files)
    
    # Reset Nodegraph for each iteration
    all_hashes = Nodegraph(ksize, 1_000_000, 4)
    
    cumulative_kmers = 0
    for position, sig_file in enumerate(shuffled_files, start=1):
        # Load signature exactly as before
        sig = sourmash.load_one_signature(str(sig_file), ksize=ksize, select_moltype='DNA')
        
        prev_total = all_hashes.n_occupied()
        all_hashes.update(sig.minhash)
        new_total = all_hashes.n_occupied()
        new_kmers = new_total - prev_total
        
        # Append a row with iteration, sample, position in this iteration, new kmers, cumulative total
        all_results.append({
            "iteration": iteration,
            "position": position,
            "sample": sig_file.name,
            "new_kmers": new_kmers,
            "cumulative_kmers": new_total
        })

# Convert to pandas DataFrame
df = pd.DataFrame(all_results)

# Save to CSV
df.to_csv("new_kmers_mags_resampled.csv", index=False)

[AnneliektH]

plotting:

import pandas as pd
import matplotlib.pyplot as plt



# Make sure sorting is correct (just in case)
df = df.sort_values(by=["iteration", "position"])

# Group by iteration and calculate cumulative kmers per sample position
df["cumulative"] = df.groupby("iteration")["new_kmers"].cumsum()

# Calculate mean and std across iterations at each sample position
summary = df.groupby("position")["cumulative"].agg(["mean", "std"]).reset_index()

# Plot mean ± std accumulation curve
plt.figure(figsize=(8, 5))
plt.plot(summary["position"], summary["mean"], label="Mean cumulative k-mers")
plt.fill_between(summary["position"],
                 summary["mean"] - summary["std"],
                 summary["mean"] + summary["std"],
                 alpha=0.3, label="±1 std")

plt.xlabel("Number of assembled samples (cumulative)")
plt.ylabel("Unique kmers in MAGs (cumulative)")
plt.ticklabel_format(style='plain', axis='y')
#plt.yscale("log")
#plt.title("Accumulation Curve of Unique k-mers (Resampled)")
#plt.legend()
plt.grid(True, which="both", linestyle="-", linewidth=0.5, alpha=0.7)
plt.tight_layout()
#plt.show()
plt.savefig('../plots/250808_accumcurve_kmers.k31.pdf')

[ctb]

Updated with argparse etc etc:

#! /usr/bin/env python
import argparse
import sys

import sourmash
from pathlib import Path
from sourmash.nodegraph import Nodegraph
import pandas as pd
import random


def main():
    p = argparse.ArgumentParser()
    p.add_argument('sketches', nargs='+')
    p.add_argument('-k', '--ksize', type=int, default=31)
    p.add_argument('-N', '--num-iterations', type=int, default=100)
    p.add_argument('-o', '--output', required=True)
    p.add_argument('-s', '--scaled', default=1000, type=int)
    p.add_argument('--nodegraph-size', type=int, default=10_000_000)
    p.add_argument('--nodegraph-num', type=int, default=4)
    args = p.parse_args()

    # Parameters
    ksize = args.ksize
    n_iterations = args.num_iterations

    # load sketches
    sketches = []
    for filename in args.sketches:
        db = sourmash.load_file_as_index(filename)
        db = db.select(ksize=ksize)
        print(f'found {len(db)} sketches at k={ksize} - loading as signatures now.')
        for ss in db.signatures():
            with ss.update() as ss:
                ss.minhash = ss.minhash.downsample(scaled=args.scaled)
            sketches.append(ss)

            if len(sketches) % 100 == 0:
                print('...', len(sketches))

    print('...done!')

    # Prepare a list to store results for all iterations
    all_results = []

    for iteration in range(n_iterations):
        print(f"Iteration {iteration + 1}/{n_iterations}")

        # Shuffle sample order for this iteration
        shuffled_sigs = list(sketches)
        random.shuffle(shuffled_sigs)

        # Reset Nodegraph for each iteration
        all_hashes = Nodegraph(ksize, args.nodegraph_size, args.nodegraph_num)

        cumulative_kmers = 0
        for position, ss in enumerate(shuffled_sigs, start=1):
            prev_total = all_hashes.n_occupied()
            all_hashes.update(ss.minhash)
            new_total = all_hashes.n_occupied()
            new_kmers = new_total - prev_total

            # Append a row with iteration, sample, position in this iteration, new kmers, cumulative total
            all_results.append({
                "iteration": iteration,
                "position": position,
                "sample": ss.name,
                "new_kmers": new_kmers,
                "cumulative_kmers": new_total
            })

    if all_hashes.n_occupied() > 0.3 * args.nodegraph_size:
       print('WARNING:', all_hashes.n_occupied())

    # Convert to pandas DataFrame
    df = pd.DataFrame(all_results)

    # Save to CSV
    df.to_csv(args.output, index=False)


if __name__ == '__main__':
    sys.exit(main())

[ctb]

plotting code: load CSV file above into pandas dataframe, then:

def summarize_rarefy(df):
    df = df.sort_values(by=["iteration", "position"])

    # Group by iteration and calculate cumulative kmers per sample position
    df["cumulative"] = df.groupby("iteration")["new_kmers"].cumsum()
    
    # Calculate mean and std across iterations at each sample position
    summary = df.groupby("position")["cumulative"].agg(["mean", "std"]).reset_index()
    return summary

def plot_rarefy(summary_df, *, color=None, label=None):
    x_kmer = summary_df["position"]
    mean_kmer = summary_df["mean"]
    std_kmer = summary_df["std"]    
    plt.plot(x_kmer, mean_kmer, color=color, label=label)
    plt.fill_between(x_kmer, mean_kmer-2*std_kmer, mean_kmer+2*std_kmer, alpha=0.3, color=color)


summary_df = summarize_rarefy(df)
plot_rarefy(summary_df, label="foo", color="#c51b7d")
plt.legend(loc='lower right')
plt.title('mean cumulative k-mers')

[ctb]

being developed under rarefy.py in https://github.com/dib-lab/sourmash-midgie-raker/

[ctb]

as a side note, it looks like scaled works like you'd hope, which is good news for us 😁 . that is, you can use a big scaled and still get the same results!!